Update app.py

app.py

@@ -3,13 +3,11 @@ import streamlit as st
 import pandas as pd
 import sqlite3
 import logging
-from langchain.agents import create_sql_agent, AgentType
+from langchain.agents import create_sql_agent
 from langchain.agents.agent_toolkits import SQLDatabaseToolkit
 from langchain.llms import OpenAI
 from langchain.sql_database import SQLDatabase
-from langchain.prompts import (
-    PromptTemplate,
-)
+from langchain.prompts import PromptTemplate
 from langchain.evaluation import load_evaluator
 
 # Initialize logging
@@ -19,7 +17,7 @@ logging.basicConfig(level=logging.INFO)
 if 'history' not in st.session_state:
     st.session_state.history = []
 
-# OpenAI API key (ensure it is securely stored)
+# OpenAI API key
 openai_api_key = os.getenv("OPENAI_API_KEY")
 
 # Check if the API key is set
@@ -33,7 +31,7 @@ st.write("Upload a CSV file to get started, or use the default dataset.")
 
 csv_file = st.file_uploader("Upload your CSV file", type=["csv"])
 if csv_file is None:
-    data = pd.read_csv("default_data.csv")  # Ensure this file exists in your working directory
+    data = pd.read_csv("default_data.csv")  # Ensure this file exists
     st.write("Using default_data.csv file.")
     table_name = "default_table"
 else:
@@ -42,19 +40,19 @@ else:
     st.write(f"Data Preview ({csv_file.name}):")
     st.dataframe(data.head())
 
-# Step 2: Load CSV data into a persistent SQLite database
+# Step 2: Load CSV data into SQLite database
 db_file = 'my_database.db'
 conn = sqlite3.connect(db_file)
 data.to_sql(table_name, conn, index=False, if_exists='replace')
 
-# SQL table metadata (for validation and schema)
+# SQL table metadata
 valid_columns = list(data.columns)
 st.write(f"Valid columns: {valid_columns}")
 
-# Create SQLDatabase instance with custom table info
+# Create SQLDatabase instance
 engine = SQLDatabase.from_uri(f"sqlite:///{db_file}", include_tables=[table_name])
 
-# Step 3: Define the few-shot examples for the prompt
+# Step 3: Define few-shot examples
 few_shot_examples = [
     {
         "input": "What is the total revenue for each category?",
@@ -78,12 +76,15 @@ for ex in few_shot_examples:
 # Prepare table information
 table_info = f"Table: {table_name}\nColumns: {', '.join(valid_columns)}"
 
+# Initialize the LLM
+llm = OpenAI(temperature=0, openai_api_key=openai_api_key)
+
 # Step 4: Define the prompt template
 system_message = """
 You are an expert data analyst who can convert natural language questions into SQL queries.
 
-Available tools:
-{tool_descriptions}
+You have access to the following tools:
+{tools}
 
 Follow these guidelines:
 1. Only use the columns and tables provided.
@@ -104,19 +105,16 @@ Question: {input}
 {agent_scratchpad}
 """
 
-# Initialize the LLM
-llm = OpenAI(temperature=0, openai_api_key=openai_api_key)
-
 # Step 5: Create the agent
 toolkit = SQLDatabaseToolkit(db=engine, llm=llm)
 tools = toolkit.get_tools()
-tool_names = [tool.name for tool in tools]
+tool_names = ", ".join([tool.name for tool in tools])
 tool_descriptions = "\n".join([f"{tool.name}: {tool.description}" for tool in tools])
 
 # Create the prompt
 agent_prompt = PromptTemplate(
     template=system_message,
-    input_variables=["input", "agent_scratchpad", "table_info", "few_shot_examples", "tool_descriptions"]
+    input_variables=["input", "agent_scratchpad", "table_info", "few_shot_examples", "tools", "tool_names"]
 )
 
 # Create the agent
@@ -146,14 +144,15 @@ def process_input():
                     table_info=table_info,
                     few_shot_examples=few_shot_str,
                     agent_scratchpad="",
-                    tool_descriptions=tool_descriptions
+                    tools=tool_descriptions,
+                    tool_names=tool_names
                 )
 
             # Extract the SQL query from the agent's response
             sql_query = response.strip()
             logging.info(f"Generated SQL Query: {sql_query}")
 
-            # Attempt to execute SQL query and handle exceptions
+            # Execute SQL query
             try:
                 result = pd.read_sql_query(sql_query, conn)
 
@@ -161,12 +160,12 @@ def process_input():
                     assistant_response = "The query returned no results. Please try a different question."
                     st.session_state.history.append({"role": "assistant", "content": assistant_response})
                 else:
-                    # Limit the result to first 10 rows for display
+                    # Display results
                     result_display = result.head(10)
                     st.session_state.history.append({"role": "assistant", "content": "Here are the results:"})
                     st.session_state.history.append({"role": "assistant", "content": result_display})
 
-                    # Generate insights based on the query result
+                    # Generate insights
                     insights_template = """
                     You are an expert data analyst. Based on the user's question and the SQL query result provided below, generate a concise analysis that includes key data insights and actionable recommendations. Limit the response to a maximum of 150 words.
 
@@ -183,7 +182,7 @@ def process_input():
                     result_str = result_display.to_string(index=False)
                     insights = insights_chain.run({'question': user_prompt, 'result': result_str})
 
-                    # Append the assistant's insights to the history
+                    # Append insights to history
                     st.session_state.history.append({"role": "assistant", "content": insights})
             except Exception as e:
                 logging.error(f"An error occurred during SQL execution: {e}")
@@ -194,10 +193,10 @@ def process_input():
             assistant_response = f"Error: {e}"
             st.session_state.history.append({"role": "assistant", "content": assistant_response})
 
-        # Reset the user_input in session state
+        # Reset user input
         st.session_state['user_input'] = ''
 
-# Step 7: Display the conversation history
+# Step 7: Display conversation history
 for message in st.session_state.history:
     if message['role'] == 'user':
         st.markdown(f"**User:** {message['content']}")
@@ -208,5 +207,5 @@ for message in st.session_state.history:
         else:
             st.markdown(f"**Assistant:** {message['content']}")
 
-# Place the input field at the bottom with the callback
+# Input field
 st.text_input("Enter your message:", key='user_input', on_change=process_input)