Update app.py

app.py

@@ -1,15 +1,14 @@
-import os
+import os 
 import streamlit as st
 import pandas as pd
 import sqlite3
 import logging
-from langchain.agents import create_sql_agent
-from langchain.agents.agent_toolkits import SQLDatabaseToolkit
-from langchain.agents.agent_types import AgentType
 from langchain.llms import OpenAI
-from langchain.sql_database import SQLDatabase
 from langchain.chat_models import ChatOpenAI
-from langchain.evaluation import load_evaluator
+from langchain.chains import SQLDatabaseChain
+from langchain.prompts import PromptTemplate
+from langchain.chains import LLMChain
+from langchain.sql_database import SQLDatabase
 
 # Initialize logging
 logging.basicConfig(level=logging.INFO)
@@ -52,96 +51,20 @@ engine = SQLDatabase.from_uri(f"sqlite:///{db_file}", include_tables=[table_name
 # Initialize the LLM
 llm = ChatOpenAI(temperature=0, openai_api_key=openai_api_key)
 
-# Step 3: Create the agent
-toolkit = SQLDatabaseToolkit(db=engine, llm=llm)
-
-sql_agent = create_sql_agent(
-    llm=llm,
-    toolkit=toolkit,
-    verbose=True,
-    agent_type=AgentType.OPENAI_FUNCTIONS,
-    max_iterations=5
-)
+# Initialize the SQLDatabaseChain
+sql_chain = SQLDatabaseChain(llm=llm, database=engine, verbose=True)
 
 # Step 4: Define the callback function
 def process_input():
-    user_prompt = st.session_state['user_input']
-
-    if user_prompt:
-        try:
-            # Append user message to history
-            st.session_state.history.append({"role": "user", "content": user_prompt})
-
-            # Use the agent to generate the SQL query and get the response
-            with st.spinner("Processing..."):
-                response = sql_agent.run(user_prompt)
-
-            # Check if the response contains a SQL query
-            if "```sql" in response:
-                # Extract the SQL query
-                start_index = response.find("```sql") + len("```sql")
-                end_index = response.find("```", start_index)
-                sql_query = response[start_index:end_index].strip()
-            else:
-                # If no SQL code is found, assume the entire response is the SQL query
-                sql_query = response.strip()
-
-            logging.info(f"Generated SQL Query: {sql_query}")
-
-            # Attempt to execute SQL query and handle exceptions
-            try:
-                result = pd.read_sql_query(sql_query, conn)
-
-                if result.empty:
-                    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
-                    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
-                    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.
-
-                    User's Question: {question}
-
-                    SQL Query Result:
-                    {result}
-
-                    Concise Analysis:
-                    """
-                    insights_prompt = PromptTemplate(template=insights_template, input_variables=['question', 'result'])
-                    insights_chain = LLMChain(llm=llm, prompt=insights_prompt)
-
-                    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
-                    st.session_state.history.append({"role": "assistant", "content": insights})
-            except Exception as e:
-                logging.error(f"An error occurred during SQL execution: {e}")
-                assistant_response = f"Error executing SQL query: {e}"
-                st.session_state.history.append({"role": "assistant", "content": assistant_response})
-        except Exception as e:
-            logging.error(f"An error occurred: {e}")
-            assistant_response = f"Error: {e}"
-            st.session_state.history.append({"role": "assistant", "content": assistant_response})
-
-        # Reset user input
-        st.session_state['user_input'] = ''
+    # (Use the updated process_input function provided earlier)
+    # ...
 
 # Step 5: Display conversation history
 for message in st.session_state.history:
     if message['role'] == 'user':
         st.markdown(f"**User:** {message['content']}")
     elif message['role'] == 'assistant':
-        if isinstance(message['content'], pd.DataFrame):
-            st.markdown("**Assistant:** Query Results:")
-            st.dataframe(message['content'])
-        else:
-            st.markdown(f"**Assistant:** {message['content']}")
+        st.markdown(f"**Assistant:** {message['content']}")
 
 # Input field
 st.text_input("Enter your message:", key='user_input', on_change=process_input)