import pandas as pd
from langchain_openai import ChatOpenAI
from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent
from langchain.agents.agent_types import AgentType
import streamlit as st
import os
from dotenv import load_dotenv

load_dotenv()

llm = ChatOpenAI(temperature=0.5, model="mistralai/mistral-7b-instruct:free",base_url="https://openrouter.ai/api/v1",api_key=os.getenv("OPENAI_API_KEY"))
st.set_page_config(
    page_title="Data Frame Agent"
)
st.title("Dataframe Agent 🏢")
with st.sidebar:
    st.header("Upload your CSV data and get insights from it.")
    data = st.file_uploader("Upload your CSV data", type='csv')
query = st.chat_input("Ask question from data")


def get_response(df):
    agent = create_pandas_dataframe_agent(
        llm,
        df,
        verbose=True,
        allow_dangerous_code=True
    )
    return agent.invoke(query)

if query:
    df = pd.read_csv(data)
    response = get_response(df)
    st.success(response['output'])