cloud-sean/csv-chat

1import streamlit as st2import pandas as pd3import json4 5from agent import query_agent, create_agent6 7 8def decode_response(response: str) -> dict:9    """This function converts the string response from the model to a dictionary object.10 11    Args:12        response (str): response from the model13 14    Returns:15        dict: dictionary with response data16    """17    return json.loads(response)18 19def write_response(response_dict: dict):20    """21    Write a response from an agent to a Streamlit app.22 23    Args:24        response_dict: The response from the agent.25 26    Returns:27        None.28    """29 30    # Check if the response is an answer.31    if "answer" in response_dict:32        st.write(response_dict["answer"])33 34    # Check if the response is a bar chart.35    if "bar" in response_dict:36        data = response_dict["bar"]37        df = pd.DataFrame(data)38        df.set_index("columns", inplace=True)39        st.bar_chart(df)40 41    # Check if the response is a line chart.42    if "line" in response_dict:43        data = response_dict["line"]44        df = pd.DataFrame(data)45        df.set_index("columns", inplace=True)46        st.line_chart(df)47 48    # Check if the response is a table.49    if "table" in response_dict:50        data = response_dict["table"]51        df = pd.DataFrame(data["data"], columns=data["columns"])52        st.table(df)53 54st.title("👨‍💻 Chat with your CSV")55 56st.write("Please upload your CSV file below.")57 58data = st.file_uploader("Upload a CSV")59 60query = st.text_area("Insert your query")61 62if st.button("Submit Query", type="primary"):63    # Create an agent from the CSV file.64    agent = create_agent(data)65 66    # Query the agent.67    response = query_agent(agent=agent, query=query)68 69    # Decode the response.70    decoded_response = decode_response(response)71 72    # Write the response to the Streamlit app.73    write_response(decoded_response)