Spaces:
Running
Running
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." | |