src/main.py

import ast
import json
import streamlit as st
import pandas as pd
from langchain.agents.agent_types import AgentType
from langchain_experimental.agents import create_csv_agent
from langchain_groq import ChatGroq
from langchain.memory import ChatMessageHistory
from groq import Groq
import os 

# Initialize Groq client and model
client = Groq(api_key=os.getenv('Groq_API'))
MODEL = 'llama3-70b-8192'

# Initialize chat history
history = ChatMessageHistory()
history.add_user_message("hi!")
history.add_ai_message("whats up?")

# Initialize language model
llm = ChatGroq(
    temperature=0,
    groq_api_key=os.getenv('Groq_API'),
    model_name='llama3-70b-8192'
)

# Create CSV agent
agent = create_csv_agent(
    llm,
    "data/Financial_data.csv",
    verbose=True,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    max_iterations=5,
    handle_parsing_errors=True
)

# Functions to handle conversations
def convo_agent(question, chat_history):
    response = 'I was built to answer questions related to financials MSFT, TSLA and AAPL. Let me know if you have any questions on these.'
    return {'answer': response}

def csv_agent(question, chat_history):
    prompt = (
        """
        Let's decode the way to respond to the queries. The responses depend on the type of information requested in the query. 

        Return just the data, don't take effort of creating plots, prints and all.
        No explanation needed. Return just the dict
        Always include units in response .

        1. If the query requires a table, format your answer like this:
           {"table": {"columns": ["column1", "column2", ...], "data": [[value1, value2, ...], [value1, value2, ...], ...]}}

        2. For a bar chart, respond like this:
           {"bar": {"columns": ["A", "B", "C", ...], "data": [25, 24, 10, ...]}}

        3. If a line chart is more appropriate, your reply should look like this:
           {"line": {"columns": ["A", "B", "C", ...], "data": [25, 24, 10, ...]}}

        Note: We only accommodate two types of charts: "bar" and "line".

        4. For a plain question that doesn't need a chart or table, your response should be:
           {"answer": "Your answer goes here"}

        For example:
           {"answer": "The Product with the highest Orders is '15143Exfo'"}

        5. If the answer is not known or available, respond with:
           {"answer": "I do not know."}

        Return all output as a string. Remember to encase all strings in the "columns" list and data list in double quotes. 
        For example: {"columns": ["Products", "Orders"], "data": [["51993Masc", 191], ["49631Foun", 152]]}

        Return all the numerical values in int format only.
        Now, let's tackle the query step by step. Here's the query for you to work on:"""
        + 
        question
    )

        

    response = agent.run(prompt)
    return ast.literal_eval(response)

# Define tools and function mapping
tool_convo_agent = {
    "type": "function",
    "function": {
        "name": "convo_agent",
        "description": "Answers questions like chit chat or simple friendly messages",
        "parameters": {
            "type": "object",
            "properties": {
                "question": {"type": "string", "description": "The user question"}
            },
            "required": ["question"],
        },
    },
}

tool_fin_agent = {
    "type": "function",
    "function": {
        "name": "csv_agent",
        "description": "Answers questions related to financial metrics of us Apple, Microsoft and Tesla.",
        "parameters": {
            "type": "object",
            "properties": {
                "question": {"type": "string", "description": "The user question"}
            },
            "required": ["question"],
        },
    },
}

tools = [tool_convo_agent, tool_fin_agent]

function_map = {
    "csv_agent": csv_agent,
    "convo_agent": convo_agent
}

# Conversation handling
def run_conversation(chat_history, user_prompt, tools):
    final_prompt = {'chat_history':{chat_history}, 'question':{user_prompt}}
    messages = [
        {"role": "system", "content": "You are an efficient agent that determines which function to use in order to answer user question."},
        {"role": "user", "content": str(final_prompt)},
    ]
    
    response = client.chat.completions.create(
        model=MODEL,
        messages=messages,
        tools=tools,
        tool_choice="auto",
        max_tokens=4096
    )

    response_message = response.choices[0].message
    tool_calls = response_message.tool_calls
    return tool_calls

def get_response(question):
    try:
        history.add_user_message(question)
        chat_history = str(history.messages)
        agents = run_conversation(chat_history, question, tools)
        func_to_call = agents[0].function.name
        
        
        if func_to_call in function_map:
            question_to_run = ast.literal_eval(agents[0].function.arguments)['question']
            result = function_map[func_to_call](question_to_run, chat_history)
        else:
            result = {"error": "Something went Wrong"}
        
        if 'error' in result:
            return "Something went wrong"
        print(result)
        history.add_ai_message(str(result))
        return result

    except Exception as e:
        return f"Something went wrong: {e}"

# Response writing for Streamlit
def write_answer(response_dict):
    if not isinstance(response_dict, dict):
        return "Invalid response format received."

    if "answer" in response_dict:
        return response_dict

    if "bar" in response_dict:
        data = response_dict["bar"]
        try:
            df_data = {col: [x[i] if isinstance(x, list) else x for x in data['data']] for i, col in enumerate(data['columns'])}
            df = pd.DataFrame(df_data)
            df.set_index("Year", inplace=True)
            st.bar_chart(df)
            return {'bar': ''}
        except ValueError:
            st.error(f"Couldn't create DataFrame from data: {data}")

    if "line" in response_dict:
        data = response_dict["line"]
        try:
            df_data = {col: [x[i] for x in data['data']] for i, col in enumerate(data['columns'])}
            df = pd.DataFrame(df_data)
            df.set_index("Year", inplace=True)
            st.line_chart(df)
            return {'line': ''}
        except ValueError:
            st.error(f"Couldn't create DataFrame from data: {data}")

    if "table" in response_dict:
        data = response_dict["table"]
        try:
            clean_data = [
                [int(x.replace(',', '')) if isinstance(x, str) and x.replace(',', '').isdigit() else x for x in row]
                for row in data["data"]
            ]
            df = pd.DataFrame(clean_data, columns=data["columns"])
            st.table(df)
            return {'table': ''}
        except ValueError as e:
            st.error(f"Couldn't create DataFrame from data: {data}. Error: {e}")

    return "No valid response type found."