1_📊_Chat With Your Data.py

from typing import Any, Dict, Optional

import matplotlib.pyplot as plt
import pandas as pd
import streamlit as st
from dotenv import load_dotenv
from langchain.agents import AgentType
from langchain_openai import ChatOpenAI

from dagpt.agents.pandas_agent import DataAnalysisAgent
from dagpt.models.llms import load_llm
from dagpt.utils.utils import BaseLogger

# Load environment variables
load_dotenv()

logger = BaseLogger()

def execute_plot_code(code: str, df: pd.DataFrame, fig_size: tuple = (10, 6)) -> Optional[plt.Figure]:
    """Execute the plot code safely.

    Args:
        code (str): The code to execute.
        df (pd.DataFrame): The dataframe to use.
        fig_size (tuple): The size of the figure (width, height).

    Returns:
        fig: The matplotlib figure generated by the code.
    """
    try:
        plt.figure(figsize=fig_size)  # Set the figure size here
        local_vars = {"plt": plt, "df": df}
        compiled_code = compile(code, "<string>", "exec")
        exec(compiled_code, globals(), local_vars)
        return plt.gcf()
    except Exception as e:
        st.error(f"Error executing plot code: {e}")
        return None

def display_data(df: pd.DataFrame):
    """Display the dataframe in the Streamlit app."""
    st.write("### DataFrame Head:")
    st.write(df.head())

def get_agent_type() -> AgentType:
    """Get the agent type selected by the user."""
    return st.selectbox(
        "Select Agent Type",
        [AgentType.OPENAI_FUNCTIONS, AgentType.ZERO_SHOT_REACT_DESCRIPTION],
    )

def initialize_agent(
    df: pd.DataFrame, llm: ChatOpenAI, agent_type: AgentType, verbose: bool
) -> DataAnalysisAgent:
    """Initialize the Data Analysis Agent."""
    daagent = DataAnalysisAgent(
        df=df,
        llm=llm,
        agent_type=agent_type,
        verbose=verbose,
        return_intermediate_steps=True,
    )
    return daagent.create_agent()

def process_query(agent: DataAnalysisAgent, query: str):
    """Process the user query using the agent and display the response."""
    response = agent(query)
    
    if response["intermediate_steps"]:
        logger.info("### Intermediate steps ###")
        action = response["intermediate_steps"][-1][0].tool_input["query"]

        if "plt" in action:
            st.write(response["output"])
            fig = execute_plot_code(action, st.session_state.df, fig_size=(12, 8))  # Adjust the figure size here
            if fig:
                st.pyplot(fig)
            st.markdown("**Executed the code:**")
            st.code(action)
            
            to_display_str = response["output"] + "\n" + f"```python\n{action}\n```"
            st.session_state.history.append((query, to_display_str))
        else:
            st.write(response["output"])
            st.session_state.history.append((query, response["output"]))
    else:
        logger.info("### No intermediate steps ###")
        st.write(response["output"])
        st.session_state.history.append((query, response["output"]))

def display_chat_history():
    """Display the chat history."""
    st.markdown("### Chat History:")
    for i, (q, r) in enumerate(st.session_state.history):
        st.markdown(f"**Query {i+1}:** {q}")
        st.markdown(f"**Response {i+1}:** {r}")
        st.markdown("---")

def main():
    """Main function to run the Streamlit app."""
    
    # Set the page configuration
    st.set_page_config(
        page_title="Smart Data Analysis Agent",
        page_icon="📊",
        layout="centered",
    )

    # Add some CSS styling for a better interface
    st.markdown("""
        <style>
        .main {
            # background-color: #f0f2f6;
            padding: 20px;
        }
        .title {
            font-size: 2.5em;
            color: #4a4a4a;
            text-align: center;
            margin-bottom: 10px;
        }
        .subtitle {
            font-size: 1.5em;
            color: #4a4a4a;
            text-align: center;
            margin-bottom: 20px;
        }
        </style>
        """, unsafe_allow_html=True)

    # Initialize the language model
    llm = load_llm("gpt-3.5")

    # Sidebar for uploading the CSV file
    with st.sidebar:
        uploaded_file = st.file_uploader("Upload your CSV here:", type="csv")

    st.markdown('<div class="main">', unsafe_allow_html=True)
    st.markdown('<div class="title">📊 Data Analysis Agent</div>', unsafe_allow_html=True)
    st.markdown('<div class="subtitle">Welcome! Upload your CSV file, enter your query, \
        and let the agent handle the rest.</div>', unsafe_allow_html=True)

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

    if uploaded_file is not None:
        st.session_state.df = pd.read_csv(uploaded_file)
        display_data(st.session_state.df)
        agent_type = get_agent_type()
        verbose = False
        agent = initialize_agent(st.session_state.df, llm, agent_type, verbose)

        query = st.text_input("Enter your query:")

        if st.button("Run Query"):
            with st.spinner("Processing..."):
                process_query(agent, query)

        st.divider()
        display_chat_history()

    st.markdown('</div>', unsafe_allow_html=True)

if __name__ == "__main__":
    main()