Update app.py

app.py

@@ -1,12 +1,14 @@
-import os
+import os 
 import streamlit as st
 import pandas as pd
 import sqlite3
 import logging
-from langchain.chat_models import ChatOpenAI
-from langchain.agents import create_sql_agent
-from langchain.agents.agent_toolkits import SQLDatabaseToolkit
-from langchain.sql_database import SQLDatabase
+import ast  # For parsing string representations of lists
+
+from langchain_community.chat_models import ChatOpenAI
+from langchain_community.agent_toolkits.sql.base import create_sql_agent
+from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit
+from langchain_community.utilities import SQLDatabase
 from langchain.prompts import PromptTemplate
 from langchain.chains import LLMChain
 
@@ -57,7 +59,8 @@ toolkit = SQLDatabaseToolkit(db=db, llm=llm)
 agent_executor = create_sql_agent(
     llm=llm,
     toolkit=toolkit,
-    verbose=True
+    verbose=True,
+    return_intermediate_steps=True
 )
 
 # Step 4: Define the callback function
@@ -71,10 +74,35 @@ def process_input():
 
             # Use the agent to get the response
             with st.spinner("Processing..."):
-                response = agent_executor.run(user_prompt)
+                response = agent_executor(user_prompt)
+
+            # Extract the final answer and the data from intermediate steps
+            final_answer = response['output']
+            intermediate_steps = response['intermediate_steps']
+
+            # Initialize an empty list for SQL result
+            sql_result = []
+
+            # Find the SQL query result
+            for step in intermediate_steps:
+                if step[0].tool == 'sql_db_query':
+                    # The result is a string representation of a list
+                    sql_result = ast.literal_eval(step[1])
+                    break
+
+            # Convert the result to a DataFrame for better formatting
+            if sql_result:
+                # Adjust the column names based on your query
+                df_result = pd.DataFrame(sql_result, columns=['Region', 'Total_Sales'])
+                sql_result_formatted = df_result.to_markdown(index=False)
+            else:
+                sql_result_formatted = "No results found."
+
+            # Include the data in the final answer
+            assistant_response = f"{final_answer}\n\n**Query Result:**\n{sql_result_formatted}"
 
             # Append the assistant's response to the history
-            st.session_state.history.append({"role": "assistant", "content": response})
+            st.session_state.history.append({"role": "assistant", "content": assistant_response})
 
             # Generate insights based on the response
             insights_template = """
@@ -90,7 +118,7 @@ def process_input():
             insights_prompt = PromptTemplate(template=insights_template, input_variables=['question', 'response'])
             insights_chain = LLMChain(llm=llm, prompt=insights_prompt)
 
-            insights = insights_chain.run({'question': user_prompt, 'response': response})
+            insights = insights_chain.run({'question': user_prompt, 'response': assistant_response})
 
             # Append the assistant's insights to the history
             st.session_state.history.append({"role": "assistant", "content": insights})