cloud-sean/csv-chat

1# agent.py2from langchain import OpenAI3from langchain.chat_models import ChatOpenAI4from langchain.agents import create_pandas_dataframe_agent5import pandas as pd6 7import os8 9 10API_KEY = os.environ["OPENAI_API_KEY"] 11 12 13 14def create_agent(filename: str):15    """16    Create an agent that can access and use a large language model (LLM).17 18    Args:19        filename: The path to the CSV file that contains the data.20 21    Returns:22        An agent that can access and use the LLM.23    """24 25    # Create an OpenAI object.26    llm = ChatOpenAI(openai_api_key=API_KEY, model="gpt-4")27 28    # Read the CSV file into a Pandas DataFrame.29    df = pd.read_csv(filename)30 31    # Create a Pandas DataFrame agent.32    return create_pandas_dataframe_agent(llm, df, verbose=False)33 34 35def query_agent(agent, query):36    """37    Query an agent and return the response as a string.38 39    Args:40        agent: The agent to query.41        query: The query to ask the agent.42 43    Returns:44        The response from the agent as a string.45    """46 47    prompt = (48        """49            For the following query, if it requires drawing a table, reply as follows:50            {"table": {"columns": ["column1", "column2", ...], "data": [[value1, value2, ...], [value1, value2, ...], ...]}}51 52            If the query requires creating a bar chart, reply as follows:53            {"bar": {"columns": ["A", "B", "C", ...], "data": [25, 24, 10, ...]}}54 55            If the query requires creating a line chart, reply as follows:56            {"line": {"columns": ["A", "B", "C", ...], "data": [25, 24, 10, ...]}}57 58            There can only be two types of chart, "bar" and "line".59 60            If it is just asking a question that requires neither, reply as follows:61            {"answer": "answer"}62            Example:63            {"answer": "The title with the highest rating is 'Gilead'"}64 65            If you do not know the answer, reply as follows:66            {"answer": "I do not know."}67 68            Return all output as a string.69 70            All strings in "columns" list and data list, should be in double quotes,71 72            For example: {"columns": ["title", "ratings_count"], "data": [["Gilead", 361], ["Spider's Web", 5164]]}73 74            Lets think step by step.75 76            Below is the query.77            Query: 78            """79        + query80    )81 82    # Run the prompt through the agent.83    response = agent.run(prompt)84 85    # Convert the response to a string.86    return response.__str__()87 88