Update src/agent.py

src/agent.py

@@ -1,66 +1,83 @@
-import streamlit as st
+from langchain.llms import OpenAI
+from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent
 import pandas as pd
-import json
+import environ
 
-from agent import query_agent, create_agent
+env = environ.Env()
+environ.Env.read_env()
+
+API_KEY = env("OPENAI_API_KEY")
+
+
+def create_agent(filename: str):
+    """
+    Create an agent that can access and use a large language model (LLM).
 
-def decode_response(response: str) -> dict:
-    """ This function converts the string response from the model to a dictionary.
     Args:
-        response (str): The string response from the model.
+        filename: The path to the CSV file that contains the data.
+
     Returns:
-        dict: The dictionary representation of the response.
+        An agent that can access and use the LLM.
     """
-    response_dict = json.loads(response)
-    return response_dict
 
-def write_response(response_dict: dict):
-    """ This function writes the response from the model to the Streamlit app.
+    # Create an OpenAI object.
+    llm = OpenAI(openai_api_key=API_KEY)
+
+    # Read the CSV file into a Pandas DataFrame.
+    df = pd.read_csv(filename)
+
+    # Create a Pandas DataFrame agent.
+    return create_pandas_dataframe_agent(llm, df, verbose=False)
+
+
+def query_agent(agent, query):
+    """
+    Query an agent and return the response as a string.
+
     Args:
-        response_dict (dict): The dictionary representation of the response.
+        agent: The agent to query.
+        query: The query to ask the agent.
+
+    Returns:
+        The response from the agent as a string.
     """
-    # Check if the response is an answer.
-    if "answer" in response_dict:
-        st.write(response_dict["answer"])
-
-    # Check if the response is a bar chart.
-    if "bar" in response_dict:
-        data = response_dict["bar"]
-        df = pd.DataFrame(data)
-        df.set_index("columns", inplace=True)
-        st.bar_chart(df)
-
-    # Check if the response is a line chart.
-    if "line" in response_dict:
-        data = response_dict["line"]
-        df = pd.DataFrame(data)
-        df.set_index("columns", inplace=True)
-        st.line_chart(df)
-        
-    # Check if the response is a table. 
-    if "table" in response_dict:
-        data = response_dict["table"]
-        df = pd.DataFrame(data["data"], columns=data["columns"])
-        st.table(df)
-
-st.title("👨‍💻 Chat with your CSV")
-
-st.write("Please upload your CSV file below.")
-
-data = st.file_uploader("Upload a CSV")
-
-query = st.text_area("Type your query here")
-
-if st.button("Submit Query", type="primary"):
-    # Create an agent from the CSV file.
-    agent = create_agent(data)
-
-    # Query the agent.
-    response = query_agent(agent=agent, query=query)
-
-    # Decode the response.
-    decoded_response = decode_response(response)
-    print(decoded_response)
-
-    # Write the response to the Streamlit app.
-    write_response(decoded_response)
+    prompt = (
+        """
+            For the following query, if it requires drawing a table, reply as follows:
+            {"table": {"columns": ["column1", "column2", ...], "data": [[value1, value2, ...], [value1, value2, ...], ...]}}
+
+            If the query requires creating a bar chart, reply as follows:
+            {"bar": {"columns": ["A", "B", "C", ...], "data": [25, 24, 10, ...]}}
+
+            If the query requires creating a line chart, reply as follows:
+            {"line": {"columns": ["A", "B", "C", ...], "data": [25, 24, 10, ...]}}
+
+            There can only be two types of chart, "bar" and "line".
+
+            If it is just asking a question that requires neither, reply as follows:
+            {"answer": "answer"}
+            Example:
+            {"answer": "I do not have sufficient information to provide an answer to your query."}
+
+            If you do not know the answer, reply as follows:
+            {"answer": "I do not know."}
+
+            Return all output as a string.
+
+            All strings in "columns" list and data list, should be in double quotes,
+
+            For example: {"columns": ["title", "ratings_count"], "data": [["Gilead", 361], ["Spider's Web", 5164]]}
+
+            Lets think step by step.
+
+            Below is the query.
+            Query: 
+            """
+        + query
+    )
+
+    # Run the prompt through the agent.
+    response = agent.run(prompt)
+
+    # Convert the response to a string.
+    return response.__str__()