import os
import re

import streamlit as st
import pandas as pd

import langchain
from langchain.agents import AgentExecutor
from langchain.callbacks import StreamlitCallbackHandler
from langchain.chat_models import ChatOpenAI
from langchain.tools import PythonAstREPLTool
from langchain.schema import SystemMessage

from tools import PythonPlotTool
from agents import OpenAIFunctionsAgentFix  # https://github.com/langchain-ai/langchain/issues/6364


@st.cache_data
def load_csv(csv) -> pd.DataFrame:
    return pd.read_csv(csv)


def get_agent(df, openai_api_key, number_of_head_rows=5, outdir="./datavizqa/static"):
    SYSTEM_PROMPT = """You are working with a pandas dataframe in Python. The name of the dataframe is `df`.
    This is the result of `print(df.head())`:
    {df_head}""".format(df_head=str(df.head(number_of_head_rows).to_markdown()))

    tools = [
        PythonPlotTool(locals={"df": df}, outdir=outdir),
        PythonAstREPLTool(name="python", locals={"df": df}),
    ]
    llm = ChatOpenAI(
        model="gpt-3.5-turbo-0613",
        openai_api_key=openai_api_key,
        temperature=0,
        streaming=True,
    )
    agent = OpenAIFunctionsAgentFix.from_llm_and_tools(
        llm=llm,
        tools=tools,
        system_message=SystemMessage(content=SYSTEM_PROMPT),
    )
    agent_exe = AgentExecutor.from_agent_and_tools(agent, tools)
    return agent_exe

langchain.debug = os.getenv("LANGCHAIN_DEBUG")

RE_MARKDOWN_IMAGE = r"!\[(.*?)\]\((.*?)\)"


st.set_page_config(page_title="DataVizQA", page_icon="🤖")
st.title("QA on your data with visualizations")
custom_css = """
<style>
    p > img {
        width: 100%
    }
</style>
"""
st.markdown(custom_css, unsafe_allow_html=True)

ss = st.session_state

with st.sidebar:
    ss.openai_api_key = st.text_input("Your OpenAI API key", placeholder="sk-xxxx")
    ss.cot = st.radio(
        "Expand new thoughts", [False, True], format_func=lambda x: "Yes" if x else "No")


csv = st.file_uploader("Upload your CSV file", type=["csv"])
if csv is not None:
    df = load_csv(csv)
    st.dataframe(df.head())
    if key := ss.openai_api_key or os.getenv("OPENAI_API_KEY"):
        ss.agent = get_agent(df, openai_api_key=key)

if "agent" in ss:
    if "messages" not in ss:
        ss.messages = [{"role": "assistant", "content": "Data loaded! Ask me anything! I can also plot charts!"}]
    for message in ss.messages:
        st.chat_message(message["role"]).write(message["content"])

    if question := st.chat_input(placeholder="Your question"):
        ss.messages.append({"role": "user", "content": question})
        st.chat_message("user").write(question)
        with st.chat_message("assistant"):
            handler = StreamlitCallbackHandler(st.container(), expand_new_thoughts=ss.cot)
            output_image = ""
            for step in ss.agent.iter(question, callbacks=[handler]):
                if output := step.get("intermediate_step"):
                    action, value = output[0]
                    if action.tool == "python_plot":
                        output_image = value

            answer = step.get("output")
            if output_image:
                if re.search(RE_MARKDOWN_IMAGE, answer):
                    answer = re.sub(RE_MARKDOWN_IMAGE, f"![\g<1>]({output_image})", answer)
                else:
                    answer = answer + "\n" + f"![{output_image.split('/')[0]}]({output_image})"

            ss.messages.append({"role": "assistant", "content": answer})
            st.write(answer)
langchain.agents.AgentExecutor.from_agent_and_tools langchain.tools.PythonAstREPLTool langchain.schema.SystemMessage langchain.chat_models.ChatOpenAI
[((1411, 1439), 'os.getenv', 'os.getenv', (['"""LANGCHAIN_DEBUG"""'], {}), "('LANGCHAIN_DEBUG')\n", (1420, 1439), False, 'import os\n'), ((1486, 1543), 'streamlit.set_page_config', 'st.set_page_config', ([], {'page_title': '"""DataVizQA"""', 'page_icon': '"""🤖"""'}), "(page_title='DataVizQA', page_icon='🤖')\n", (1504, 1543), True, 'import streamlit as st\n'), ((1544, 1591), 'streamlit.title', 'st.title', (['"""QA on your data with visualizations"""'], {}), "('QA on your data with visualizations')\n", (1552, 1591), True, 'import streamlit as st\n'), ((1670, 1717), 'streamlit.markdown', 'st.markdown', (['custom_css'], {'unsafe_allow_html': '(True)'}), '(custom_css, unsafe_allow_html=True)\n', (1681, 1717), True, 'import streamlit as st\n'), ((1964, 2018), 'streamlit.file_uploader', 'st.file_uploader', (['"""Upload your CSV file"""'], {'type': "['csv']"}), "('Upload your CSV file', type=['csv'])\n", (1980, 2018), True, 'import streamlit as st\n'), ((513, 529), 'pandas.read_csv', 'pd.read_csv', (['csv'], {}), '(csv)\n', (524, 529), True, 'import pandas as pd\n'), ((1005, 1109), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model': '"""gpt-3.5-turbo-0613"""', 'openai_api_key': 'openai_api_key', 'temperature': '(0)', 'streaming': '(True)'}), "(model='gpt-3.5-turbo-0613', openai_api_key=openai_api_key,\n    temperature=0, streaming=True)\n", (1015, 1109), False, 'from langchain.chat_models import ChatOpenAI\n'), ((1322, 1370), 'langchain.agents.AgentExecutor.from_agent_and_tools', 'AgentExecutor.from_agent_and_tools', (['agent', 'tools'], {}), '(agent, tools)\n', (1356, 1370), False, 'from langchain.agents import AgentExecutor\n'), ((1783, 1842), 'streamlit.text_input', 'st.text_input', (['"""Your OpenAI API key"""'], {'placeholder': '"""sk-xxxx"""'}), "('Your OpenAI API key', placeholder='sk-xxxx')\n", (1796, 1842), True, 'import streamlit as st\n'), ((1856, 1950), 'streamlit.radio', 'st.radio', (['"""Expand new thoughts"""', '[False, True]'], {'format_func': "(lambda x: 'Yes' if x else 'No')"}), "('Expand new thoughts', [False, True], format_func=lambda x: 'Yes' if\n    x else 'No')\n", (1864, 1950), True, 'import streamlit as st\n'), ((878, 926), 'tools.PythonPlotTool', 'PythonPlotTool', ([], {'locals': "{'df': df}", 'outdir': 'outdir'}), "(locals={'df': df}, outdir=outdir)\n", (892, 926), False, 'from tools import PythonPlotTool\n'), ((936, 987), 'langchain.tools.PythonAstREPLTool', 'PythonAstREPLTool', ([], {'name': '"""python"""', 'locals': "{'df': df}"}), "(name='python', locals={'df': df})\n", (953, 987), False, 'from langchain.tools import PythonAstREPLTool\n'), ((2488, 2530), 'streamlit.chat_input', 'st.chat_input', ([], {'placeholder': '"""Your question"""'}), "(placeholder='Your question')\n", (2501, 2530), True, 'import streamlit as st\n'), ((1262, 1298), 'langchain.schema.SystemMessage', 'SystemMessage', ([], {'content': 'SYSTEM_PROMPT'}), '(content=SYSTEM_PROMPT)\n', (1275, 1298), False, 'from langchain.schema import SystemMessage\n'), ((2125, 2152), 'os.getenv', 'os.getenv', (['"""OPENAI_API_KEY"""'], {}), "('OPENAI_API_KEY')\n", (2134, 2152), False, 'import os\n'), ((2659, 2687), 'streamlit.chat_message', 'st.chat_message', (['"""assistant"""'], {}), "('assistant')\n", (2674, 2687), True, 'import streamlit as st\n'), ((3507, 3523), 'streamlit.write', 'st.write', (['answer'], {}), '(answer)\n', (3515, 3523), True, 'import streamlit as st\n'), ((2409, 2441), 'streamlit.chat_message', 'st.chat_message', (["message['role']"], {}), "(message['role'])\n", (2424, 2441), True, 'import streamlit as st\n'), ((2606, 2629), 'streamlit.chat_message', 'st.chat_message', (['"""user"""'], {}), "('user')\n", (2621, 2629), True, 'import streamlit as st\n'), ((2736, 2750), 'streamlit.container', 'st.container', ([], {}), '()\n', (2748, 2750), True, 'import streamlit as st\n'), ((3173, 3209), 're.search', 're.search', (['RE_MARKDOWN_IMAGE', 'answer'], {}), '(RE_MARKDOWN_IMAGE, answer)\n', (3182, 3209), False, 'import re\n'), ((3240, 3303), 're.sub', 're.sub', (['RE_MARKDOWN_IMAGE', 'f"""![\\\\g<1>]({output_image})"""', 'answer'], {}), "(RE_MARKDOWN_IMAGE, f'![\\\\g<1>]({output_image})', answer)\n", (3246, 3303), False, 'import re\n')]