Update app.py

app.py

@@ -3,7 +3,7 @@ import streamlit as st
 import pandas as pd
 import sqlite3
 import logging
-from langchain.agents import create_sql_agent
+from langchain.agents import create_sql_agent, AgentType
 from langchain.agents.agent_toolkits import SQLDatabaseToolkit
 from langchain.llms import OpenAI
 from langchain.sql_database import SQLDatabase
@@ -65,20 +65,20 @@ engine = SQLDatabase.from_uri(f"sqlite:///{db_file}", include_tables=[table_name
 few_shot_examples = [
     {
         "input": "What is the total revenue for each category?",
-        "query": f"SELECT category, SUM(revenue) FROM {table_name} GROUP BY category;"
+        "output": f"SELECT category, SUM(revenue) FROM {table_name} GROUP BY category;"
     },
     {
         "input": "Show the top 5 products by sales.",
-        "query": f"SELECT product_name, sales FROM {table_name} ORDER BY sales DESC LIMIT 5;"
+        "output": f"SELECT product_name, sales FROM {table_name} ORDER BY sales DESC LIMIT 5;"
     },
     {
         "input": "How many orders were placed in the last month?",
-        "query": f"SELECT COUNT(*) FROM {table_name} WHERE order_date >= DATE('now', '-1 month');"
+        "output": f"SELECT COUNT(*) FROM {table_name} WHERE order_date >= DATE('now', '-1 month');"
     }
 ]
 
 # Step 4: Define the prompt templates
-system_prefix = """
+system_message = """
 You are an expert data analyst who can convert natural language questions into SQL queries.
 Follow these guidelines:
 1. Only use the columns and tables provided.
@@ -86,37 +86,53 @@ Follow these guidelines:
 3. Ensure string comparisons are case-insensitive.
 4. Do not execute queries that could be harmful or unethical.
 5. Provide clear and concise SQL queries.
+
+Available tables and columns:
+{table_info}
+
+Use the following examples as a guide:
+{few_shot_examples}
 """
 
-few_shot_prompt = FewShotPromptTemplate(
-    example_prompt=PromptTemplate.from_template("Question: {input}\nSQL Query: {query}"),
-    examples=few_shot_examples,
-    prefix=system_prefix,
-    suffix="Question: {input}\nSQL Query:",
-    input_variables=["input", "agent_scratchpad"]
-)
+# Prepare few-shot examples as a string
+few_shot_str = ""
+for ex in few_shot_examples:
+    few_shot_str += f"Q: {ex['input']}\nA: {ex['output']}\n\n"
+
+# Prepare table information
+table_info = f"Table: {table_name}\nColumns: {', '.join(valid_columns)}"
 
-# Step 5: Initialize the LLM and toolkit
+# Initialize the LLM
 llm = ChatOpenAI(temperature=0, openai_api_key=openai_api_key)
+
+# Step 5: Create the agent
+
+# Get the list of tools from the toolkit
 toolkit = SQLDatabaseToolkit(db=engine, llm=llm)
+tools = toolkit.get_tools()
+tool_names = [tool.name for tool in tools]
 
-# Step 6: Create the agent using 'zero-shot-react-description' agent type
+# Create the agent prompt
 agent_prompt = ChatPromptTemplate.from_messages([
-    SystemMessagePromptTemplate(prompt=few_shot_prompt),
+    SystemMessagePromptTemplate.from_template(system_message),
     HumanMessagePromptTemplate.from_template("{input}"),
     MessagesPlaceholder(variable_name="agent_scratchpad")
 ])
 
+# Set input variables for the prompt
+agent_prompt.input_variables = ["input", "agent_scratchpad", "table_info", "few_shot_examples"]
+
+# Create the agent
 sql_agent = create_sql_agent(
     llm=llm,
     toolkit=toolkit,
     prompt=agent_prompt,
     verbose=True,
-    agent_type="zero-shot-react-description",
+    agent_type=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
     max_iterations=5
 )
 
-# Step 7: Define the callback function
+# Step 6: Define the callback function
 def process_input():
     user_prompt = st.session_state['user_input']
 
@@ -127,8 +143,14 @@ def process_input():
 
             # Use the agent to generate the SQL query
             with st.spinner("Generating SQL query..."):
-                response = sql_agent.run(user_prompt)
-            
+                # Run the agent with the necessary inputs
+                response = sql_agent.run(
+                    input=user_prompt,
+                    table_info=table_info,
+                    few_shot_examples=few_shot_str,
+                    agent_scratchpad=""
+                )
+
             # Extract the SQL query from the agent's response
             sql_query = response.strip()
             logging.info(f"Generated SQL Query: {sql_query}")
@@ -177,7 +199,7 @@ def process_input():
         # Reset the user_input in session state
         st.session_state['user_input'] = ''
 
-# Step 8: Display the conversation history
+# Step 7: Display the conversation history
 for message in st.session_state.history:
     if message['role'] == 'user':
         st.markdown(f"**User:** {message['content']}")