Update app.py

app.py

@@ -3,31 +3,47 @@ import streamlit as st
 import pandas as pd
 import sqlite3
 import openai
-from langchain_openai import AzureChatOpenAI, OpenAIEmbeddings
-from langchain.agents import create_sql_agent
-from langchain.sql_database import SQLDatabase
-from langchain.vectorstores import FAISS
+from langchain_openai import AzureChatOpenAI
+from langchain_community.agent_toolkits.sql.base import create_sql_agent
+from langchain_community.utilities import SQLDatabase
+from langchain_community.document_loaders import CSVLoader
+from langchain_community.vectorstores import FAISS
+from langchain_openai.embeddings import AzureOpenAIEmbeddings
 from langchain.chains import RetrievalQA
-from langchain.document_loaders import CSVLoader
-from langchain.prompts import ChatPromptTemplate, FewShotPromptTemplate
-from langchain.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate
 import sqlparse
 import logging
+from dotenv import load_dotenv
 
-# Load environment variables for Azure OpenAI
-openai.api_key = os.getenv("OPENAI_API_KEY")
+# Load environment variables from .env file
+load_dotenv()
+
+# Set up API credentials and environment variables
 api_key = os.getenv("OPENAI_API_KEY")
-endpoint = os.getenv("azure_endpoint")
-api_type = os.getenv("OPENAI_API_TYPE")
-api_version = os.getenv("OPENAI_API_VERSION", "2023-05-15")
+azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
+api_version = os.getenv("OPENAI_API_VERSION", "2023-05-15")  # Set a default if not provided
+chat_model = os.getenv("CHAT_MODEL")
+chat_deployment = os.getenv("CHAT_DEPLOYMENT")
+embed_model = os.getenv("EMBED_MODEL")
+embed_deployment = os.getenv("EMBED_DEPLOYMENT")
+
+# Default to a specific endpoint if the environment variable is missing
+if not azure_endpoint:
+    azure_endpoint = "https://<your-azure-endpoint>.openai.azure.com"  # Replace with your actual endpoint
+
+# OpenAI API key (ensure it is securely stored)
+openai.api_key = api_key
 
-# Models
-chat_model = os.getenv("chat_model")
-embed_model = os.getenv("embed_model")
-chat_deployment = os.getenv("chat_deployment")
-embed_deployment = os.getenv("embed_deployment")
+# Initialize Azure OpenAI LLM (Language Model)
+llm = AzureChatOpenAI(
+    temperature=0,
+    model=chat_model,
+    deployment_name=chat_deployment,
+    api_key=api_key,
+    azure_endpoint=azure_endpoint,
+    api_version=api_version
+)
 
-# Load CSV file for data
+# Step 1: Upload CSV data file (or use default)
 csv_file = st.file_uploader("Upload your CSV file", type=["csv"])
 if csv_file is None:
     data = pd.read_csv("default_data.csv")  # Use default CSV if no file is uploaded
@@ -37,7 +53,8 @@ else:
     st.write(f"Data Preview ({csv_file.name}):")
     st.dataframe(data.head())
 
-# Use a persistent SQLite database instead of in-memory
+# Step 2: Load CSV data into a persistent SQLite database
+# Use a persistent database file instead of in-memory to retain schema context
 db_file = 'my_database.db'
 conn = sqlite3.connect(db_file)
 table_name = csv_file.name.split('.')[0] if csv_file else "default_table"
@@ -45,59 +62,79 @@ data.to_sql(table_name, conn, index=False, if_exists='replace')
 
 # SQL table metadata (for validation and schema)
 valid_columns = list(data.columns)
+
+# Debug: Display valid columns for user to verify
 st.write(f"Valid columns: {valid_columns}")
 
-# Set up the SQL Database for LangChain with AzureOpenAI configuration
-llm = AzureChatOpenAI(
-    temperature=0,
-    model=chat_model,
-    deployment_name=chat_deployment,
-    api_key=api_key,
-    azure_endpoint=endpoint,
-    api_version=api_version
-)
+# Step 3: Set up the SQL Database for LangChain
 db = SQLDatabase.from_uri(f'sqlite:///{db_file}')
-db.raw_connection = conn
+db.raw_connection = conn  # Use the persistent database connection for LangChain
 
-# Create the SQL agent with prompt and toolkit for SQL querying
+# Step 4: Create the SQL agent with increased iteration and time limits
 sql_agent = create_sql_agent(
-    llm=llm,
+    llm,
     db=db,
     verbose=True,
     max_iterations=20,  # Increased iteration limit
     max_execution_time=90  # Set timeout limit to 90 seconds
 )
 
-# Set up FAISS for retrieval using OpenAIEmbeddings
-embeddings = OpenAIEmbeddings()
+# Step 5: Use FAISS with RAG for context retrieval
+embeddings = AzureOpenAIEmbeddings(
+    model=embed_model,
+    deployment_name=embed_deployment,
+    azure_endpoint=azure_endpoint,
+    api_key=api_key,
+    api_version=api_version
+)
 loader = CSVLoader(file_path=csv_file.name if csv_file else "default_data.csv")
 documents = loader.load()
+
 vector_store = FAISS.from_documents(documents, embeddings)
 retriever = vector_store.as_retriever()
 rag_chain = RetrievalQA.from_chain_type(llm=llm, retriever=retriever)
 
-# Insight extraction and SQL query generation
+# Step 6: Define SQL validation helpers
+def validate_sql(query, valid_columns):
+    """Validates the SQL query by ensuring it references only valid columns."""
+    parsed = sqlparse.parse(query)
+    for token in parsed[0].tokens:
+        if token.ttype is None:  # If it's a column name
+            column_name = str(token).strip()
+            if column_name not in valid_columns:
+                return False
+    return True
+
+def validate_sql_with_sqlparse(query):
+    """Validates SQL syntax using sqlparse."""
+    parsed_query = sqlparse.parse(query)
+    return len(parsed_query) > 0
+
+# Step 7: Generate SQL query based on user input and run it with LangChain SQL Agent
 user_prompt = st.text_input("Enter your natural language prompt:")
 if user_prompt:
     try:
-        # Add column hints to the user prompt
+        # Step 8: Add valid column names to the prompt
         column_hints = f" Use only these columns: {', '.join(valid_columns)}"
         prompt_with_columns = user_prompt + column_hints
 
-        # Retrieve context using FAISS and RAG
+        # Step 9: Retrieve context using RAG
         context = rag_chain.run(prompt_with_columns)
         st.write(f"Retrieved Context: {context}")
 
-        # Generate SQL query using SQL agent
+        # Step 10: Generate SQL query using SQL agent
         generated_sql = sql_agent.run(f"{prompt_with_columns} {context}")
+        
+        # Debug: Display generated SQL query for inspection
         st.write(f"Generated SQL Query: {generated_sql}")
 
-        # Validate SQL query and execute
+        # Step 11: Validate SQL query
         if not validate_sql_with_sqlparse(generated_sql):
             st.write("Generated SQL is not valid.")
         elif not validate_sql(generated_sql, valid_columns):
             st.write("Generated SQL references invalid columns.")
         else:
+            # Step 12: Execute SQL query
             result = pd.read_sql(generated_sql, conn)
             st.write("Query Results:")
             st.dataframe(result)