# https://python.langchain.com/docs/integrations/toolkits/csv

import os
from pathlib import Path

import streamlit as st

from langchain.agents import create_csv_agent
from langchain.agents.agent_types import AgentType
from langchain.chat_models import ChatOpenAI
from langchain.llms.openai import OpenAI


OPENAI_API_KEY = st.secrets.get("openai_api_key", None)
MODEL_NAME = "gpt-3.5-turbo-0613"

DATA_PATH = Path("data")
CSV_PATH = DATA_PATH / "dataset.csv"
METADATA_PATH = DATA_PATH / "metadata.md"
METADATA = METADATA_PATH.read_text()


llm = ChatOpenAI(model_name=MODEL_NAME, openai_api_key=OPENAI_API_KEY, temperature=0)

agent = create_csv_agent(
    llm=llm,
    path=str(CSV_PATH),
    verbose=True,
    agent_type=AgentType.OPENAI_FUNCTIONS,
    prefix=METADATA,
)

# agent_executor = create_sql_agent(
#     llm=llm,
#     toolkit=toolkit,
#     verbose=True,
#     agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
# )

st.title("💬 Dataoot")

if "messages" not in st.session_state:
    st.session_state["messages"] = []

for msg in st.session_state.messages:
    st.chat_message(msg["role"]).write(msg["content"])

if prompt := st.chat_input():
    st.session_state.messages.append({"role": "user", "content": prompt})
    st.chat_message("user").write(prompt)
    response = agent.run(prompt)
    st.session_state.messages.append({"role": "assistant", "content": response})
    st.chat_message("assistant").write(response)