Update src/main.py

src/main.py

@@ -1,209 +1,209 @@
-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
-
-# Initialize Groq client and model
-client = Groq(api_key='gsk')
-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='gsk...',
-    model_name='llama3-70b-8192'
-)
-
-# Create CSV agent
-agent = create_csv_agent(
-    llm,
-    r"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."
+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
+
+# Initialize Groq client and model
+client = Groq(api_key='gsk')
+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='gsk...',
+    model_name='llama3-70b-8192'
+)
+
+# Create CSV agent
+agent = create_csv_agent(
+    llm,
+    r"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."