Update app.py

app.py

@@ -1,14 +1,14 @@
-import os 
+import os
 import streamlit as st
 import pandas as pd
 import sqlite3
 import logging
-from langchain.llms import OpenAI
 from langchain.chat_models import ChatOpenAI
-from langchain.chains import SQLChain as SQLDatabaseChain  # Updated import
+from langchain.agents import create_sql_agent
+from langchain.agents.agent_toolkits import SQLDatabaseToolkit
+from langchain.sql_database import SQLDatabase
 from langchain.prompts import PromptTemplate
 from langchain.chains import LLMChain
-from langchain.sql_database import SQLDatabase
 
 # Initialize logging
 logging.basicConfig(level=logging.INFO)
@@ -44,15 +44,21 @@ else:
 db_file = 'my_database.db'
 conn = sqlite3.connect(db_file)
 data.to_sql(table_name, conn, index=False, if_exists='replace')
+conn.close()
 
 # Create SQLDatabase instance
-engine = SQLDatabase.from_uri(f"sqlite:///{db_file}", include_tables=[table_name])
+db = SQLDatabase.from_uri(f"sqlite:///{db_file}", include_tables=[table_name])
 
 # Initialize the LLM
 llm = ChatOpenAI(temperature=0, openai_api_key=openai_api_key)
 
-# Initialize the SQLDatabaseChain
-sql_chain = SQLDatabaseChain(llm=llm, database=engine, verbose=True)
+# Initialize the SQL Agent
+toolkit = SQLDatabaseToolkit(db=db, llm=llm)
+agent_executor = create_sql_agent(
+    llm=llm,
+    toolkit=toolkit,
+    verbose=True
+)
 
 # Step 4: Define the callback function
 def process_input():
@@ -63,9 +69,9 @@ def process_input():
             # Append user message to history
             st.session_state.history.append({"role": "user", "content": user_prompt})
 
-            # Use the SQLDatabaseChain to get the response
+            # Use the agent to get the response
             with st.spinner("Processing..."):
-                response = sql_chain.run(user_prompt)
+                response = agent_executor.run(user_prompt)
 
             # Append the assistant's response to the history
             st.session_state.history.append({"role": "assistant", "content": response})