Update app.py

app.py

@@ -3,8 +3,8 @@ import openai
 import pandas as pd
 from sqlalchemy import create_engine
 from langchain.chat_models import ChatOpenAI
-from langchain.utilities.sql_database import SQLDatabase
-from langchain.chains import SQLDatabaseChain
+from langchain.prompts import PromptTemplate
+from langchain.chains import LLMChain
 
 # Set OpenAI API Key
 openai.api_key = "sk-O7esHSo2XAWm-GXUGXp7_P9l4qXrQMn0CIGzs34ojLT3BlbkFJeXGSSvywppRTAvyT0zZkmZLZsj5cg7XkAkBTh8ZxoA"
@@ -15,8 +15,16 @@ engine = create_engine(DATABASE_URL)
 
 # Set up LangChain components
 llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.5)  # OpenAI's Chat model for LLM
-db = SQLDatabase(engine)  # Connect LangChain to the database
-sql_chain = SQLDatabaseChain.from_llm(llm, database=db)  # Create the SQL chain
+
+# Custom prompt to generate SQL query
+prompt_template = """
+You are an SQL expert. Write an SQL query to answer the following question:
+Question: {question}
+SQL Query:
+"""
+
+prompt = PromptTemplate(input_variables=["question"], template=prompt_template)
+llm_chain = LLMChain(llm=llm, prompt=prompt)
 
 # Streamlit UI setup
 st.title("SQL Data Chatbot with LangChain")
@@ -27,18 +35,18 @@ user_question = st.text_input("Your question:")
 
 # Process the question if provided
 if user_question:
-    # Generate the SQL query and answer using the SQL chain
+    # Generate the SQL query using the custom LLM chain
     try:
-        # Execute the question through the SQL chain
-        response = sql_chain.run(user_question)
+        # Generate SQL query with the language model
+        sql_query = llm_chain.run({"question": user_question}).strip()
 
-        # Display the generated SQL query and answer
-        st.subheader("Generated SQL Query and Answer")
-        st.write(response)
+        # Display the generated SQL query
+        st.subheader("Generated SQL Query")
+        st.code(sql_query, language="sql")
 
-        # Execute the SQL query to get results
+        # Execute SQL query to get the results
         with engine.connect() as conn:
-            result_df = pd.read_sql_query(response.query, conn)
+            result_df = pd.read_sql_query(sql_query, conn)
 
         # Show query results if any
         if not result_df.empty: