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."