Update app.py

app.py

@@ -1,63 +1,68 @@
 import streamlit as st
 import pandas as pd
 import sqlite3
-import plotly.express as px
-import json
-
-# Set paths to the default files
-DEFAULT_PROMPT_PATH = "prompt_engineering.json"
-DEFAULT_METADATA_PATH = "default_metadata.csv"
-DEFAULT_DATA_PATH = "default_data.csv"
-
-# Load the prompt engineering JSON file (use default if no user-uploaded prompt file)
-with open(DEFAULT_PROMPT_PATH) as f:
-    prompt_data = json.load(f)
-
-# Function to find a query based on the user prompt
-def get_query_from_prompt(user_prompt):
-    for item in prompt_data['prompts']:
-        if item['question'].lower() in user_prompt.lower():
-            return item['query']
-    return None  # Return None if no matching query is found
-
-# Step 1: Upload metadata.csv file (or use default)
-metadata_file = st.file_uploader("Upload your metadata.csv file", type=["csv"])
-if metadata_file is None:
-    metadata = pd.read_csv(DEFAULT_METADATA_PATH)
-    st.write("Using default metadata.csv file.")
-else:
-    metadata = pd.read_csv(metadata_file)
-    st.write("Metadata loaded successfully!")
-    st.dataframe(metadata)
+import openai
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from langchain import OpenAI
+from langchain.agents import create_sql_agent
+from langchain.sql_database import SQLDatabase
+from langchain.chains import RetrievalQA
+from langchain.document_loaders import CSVLoader
+from langchain.vectorstores import FAISS
+from langchain.embeddings.openai import OpenAIEmbeddings
+
+# OpenAI API key (ensure it's stored securely)
+openai.api_key = os.getenv("OPENAI_API_KEY")
 
-# Step 2: Upload CSV data file (or use default)
+# 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_PATH)
+    data = pd.read_csv("default_data.csv")  # Using default CSV
     st.write("Using default data.csv file.")
 else:
     data = pd.read_csv(csv_file)
     st.write("Data Preview:")
     st.dataframe(data.head())
 
-# Step 3: Load CSV data into a SQLite database (SQL agent)
-conn = sqlite3.connect(':memory:')  # Use an in-memory SQLite database
+# Step 2: Load CSV data into SQLite database (SQL agent)
+conn = sqlite3.connect(':memory:')  # In-memory SQLite database
 data.to_sql('sales_data', conn, index=False, if_exists='replace')
 
-# Step 4: Get user prompt and map to SQL query
-user_prompt = st.text_input("Enter your natural language prompt:")
+# Create a SQL database connection for LangChain
+db = SQLDatabase.from_uri('sqlite:///:memory:')
+db.raw_connection = conn
+
+# Step 3: Use LLaMA for context retrieval (RAG)
+tokenizer = AutoTokenizer.from_pretrained("huggyllama/llama-7b")
+llama_model = AutoModelForCausalLM.from_pretrained("huggyllama/llama-7b")
+
+# Load and vectorize documents for retrieval
+embeddings = OpenAIEmbeddings()  # Using OpenAI embeddings, but you can swap this out for another one
+loader = CSVLoader(file_path=csv_file.name if csv_file else "default_data.csv")
+documents = loader.load()
 
-# Step 5: Process the prompt and generate SQL query dynamically
+# Use FAISS to create a retriever from the documents
+vector_store = FAISS.from_documents(documents, embeddings)
+retriever = vector_store.as_retriever()
+
+# Step 4: Create a RAG (Retrieval-Augmented Generation) chain
+rag_chain = RetrievalQA.from_chain_type(llama_model, retriever=retriever)
+
+# Step 5: Use OpenAI for SQL query generation
+openai_llm = OpenAI(temperature=0)  # OpenAI LLM for SQL query generation
+sql_agent = create_sql_agent(openai_llm, db, verbose=True)
+
+# Step 6: Get user prompt and augment with RAG retrieval before SQL generation
+user_prompt = st.text_input("Enter your natural language prompt:")
 if user_prompt:
-    query = get_query_from_prompt(user_prompt)
-    if query:
-        result = pd.read_sql(query, conn)
-        st.write("Query Results:")
-        st.dataframe(result)
-        
-        # If the query involves plotting (like "plot sales"), show the chart
-        if "plot" in user_prompt.lower():
-            fig = px.bar(result, x='Date', y='Sales', title="Sales Over Time")
-            st.plotly_chart(fig)
-    else:
-        st.write("Sorry, I couldn't find a matching query for your prompt.")
+    try:
+        # Step 7: Retrieve context using LLaMA-based RAG
+        rag_result = rag_chain.run(user_prompt)
+        st.write(f"Retrieved Context from LLaMA RAG: {rag_result}")
+
+        # Step 8: Generate and execute SQL query using OpenAI based on prompt and retrieved context
+        query_input = f"{user_prompt} {rag_result}"
+        response = sql_agent.run(query_input)
+        st.write(f"Generated SQL Query Results: {response}")
+    except Exception as e:
+        st.write(f"An error occurred: {e}")