Update app.py

app.py

@@ -1,95 +1,98 @@
+import os
 import streamlit as st
 import pandas as pd
-import os
-import tempfile
-from PyPDF2 import PdfReader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from sentence_transformers import SentenceTransformer
-import faiss
 import openai
+import torch
+import matplotlib.pyplot as plt
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from dotenv import load_dotenv
+import anthropic
 
-# OpenAI API key configuration
-st.set_page_config(page_title="RAG Chatbot with Files", layout="centered")
-openai.api_key = st.sidebar.text_input("Enter OpenAI API Key:", type="password")
-
-# Initialize FAISS and embedding model
-embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
-faiss_index = None
-data_chunks = []
-chunk_mapping = {}
-
-# File Upload and Processing
-def load_files(uploaded_files):
-    global data_chunks, chunk_mapping, faiss_index
-    data_chunks = []
-    chunk_mapping = {}
-    for uploaded_file in uploaded_files:
-        file_type = uploaded_file.name.split('.')[-1]
-        with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
-            tmp_file.write(uploaded_file.read())
-            tmp_file_path = tmp_file.name
-
-        if file_type == "csv":
-            df = pd.read_csv(tmp_file_path)
-            content = "\n".join(df.astype(str).values.flatten())
-        elif file_type == "xlsx":
-            df = pd.read_excel(tmp_file_path)
-            content = "\n".join(df.astype(str).values.flatten())
-        elif file_type == "pdf":
-            reader = PdfReader(tmp_file_path)
-            content = "".join([page.extract_text() for page in reader.pages])
-        else:
-            st.error(f"Unsupported file type: {file_type}")
-            continue
-
-        # Split into chunks
-        splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-        chunks = splitter.split_text(content)
-        data_chunks.extend(chunks)
-        chunk_mapping.update({i: (uploaded_file.name, chunk) for i, chunk in enumerate(chunks)})
-
-    # Create FAISS index
-    embeddings = embedding_model.encode(data_chunks)
-    faiss_index = faiss.IndexFlatL2(embeddings.shape[1])
-    faiss_index.add(embeddings)
+# Load environment variables
+load_dotenv()
+os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
+os.environ["ANTHROPIC_API_KEY"] = os.getenv("ANTHROPIC_API_KEY")
 
-# Query Processing
-def handle_query(query):
-    if not faiss_index:
-        return "No data available. Please upload files first."
+st.title("Excel Q&A Chatbot 📊")
 
-    # Generate embedding for the query
-    query_embedding = embedding_model.encode([query])
-    distances, indices = faiss_index.search(query_embedding, k=5)
-    relevant_chunks = [chunk_mapping[idx][1] for idx in indices[0]]
+# Model Selection
+model_choice = st.selectbox("Select LLM Model", ["OpenAI GPT-3.5", "Claude 3 Haiku", "Mistral-7B"])
 
-    # Use OpenAI for summarization
-    prompt = "Summarize the following information:\n" + "\n".join(relevant_chunks)
-    response = openai.Completion.create(
-        engine="text-davinci-003",
-        prompt=prompt,
-        max_tokens=150
-    )
-    return response['choices'][0]['text']
+# Load appropriate model based on selection
+if model_choice == "Mistral-7B":
+    model_name = "mistralai/Mistral-7B-Instruct"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
+    def ask_mistral(query):
+        inputs = tokenizer(query, return_tensors="pt").to("cuda")
+        output = model.generate(**inputs)
+        return tokenizer.decode(output[0])
 
-# Streamlit UI
-def main():
-    st.title("RAG Chatbot with Files")
-    st.sidebar.title("Options")
-    uploaded_files = st.sidebar.file_uploader("Upload files (CSV, Excel, PDF):", type=["csv", "xlsx", "pdf"], accept_multiple_files=True)
+elif model_choice == "Claude 3 Haiku":
+    client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
+    def ask_claude(query):
+        response = client.messages.create(
+            model="claude-3-haiku",
+            messages=[{"role": "user", "content": query}]
+        )
+        return response.content
 
-    if uploaded_files:
-        load_files(uploaded_files)
-        st.sidebar.success("Files loaded successfully!")
+else:
+    client = openai.OpenAI()
+    def ask_gpt(query):
+        response = client.chat.completions.create(
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": query}]
+        )
+        return response.choices[0].message.content
 
-    query = st.text_input("Ask a question about the data:")
-    if st.button("Get Answer"):
-        if openai.api_key and query:
-            answer = handle_query(query)
-            st.subheader("Answer:")
-            st.write(answer)
-        else:
-            st.error("Please provide a valid API key and query.")
+# File Upload
+uploaded_file = st.file_uploader("Upload an Excel file", type=["csv", "xlsx"])
 
-if __name__ == "__main__":
-    main()
+if uploaded_file is not None:
+    file_extension = uploaded_file.name.split(".")[-1].lower()
+    df = pd.read_csv(uploaded_file) if file_extension == "csv" else pd.read_excel(uploaded_file)
+    st.write("### Preview of Data:")
+    st.write(df.head())
+    
+    # Extract metadata
+    column_names = df.columns.tolist()
+    data_types = df.dtypes.apply(lambda x: x.name).to_dict()
+    missing_values = df.isnull().sum().to_dict()
+    
+    # Display metadata
+    st.write("### Column Details:")
+    st.write(pd.DataFrame({"Column": column_names, "Type": data_types.values(), "Missing Values": missing_values.values()}))
+    
+    # User Query
+    query = st.text_input("Ask a question about this data:")
+    
+    if st.button("Submit Query"):
+        if query:
+            # Interpret the query using selected LLM
+            if model_choice == "Mistral-7B":
+                parsed_query = ask_mistral(f"Convert this question into a Pandas operation: {query}")
+            elif model_choice == "Claude 3 Haiku":
+                parsed_query = ask_claude(f"Convert this question into a Pandas operation: {query}")
+            else:
+                parsed_query = ask_gpt(f"Convert this question into a Pandas operation: {query}")
+            
+            # Execute the query
+            try:
+                result = eval(f"df.{parsed_query}")
+                st.write("### Result:")
+                st.write(result if isinstance(result, pd.DataFrame) else str(result))
+                
+                # If numerical data, show a visualization
+                if isinstance(result, pd.Series) and result.dtype in ["int64", "float64"]:
+                    fig, ax = plt.subplots()
+                    result.plot(kind="bar", ax=ax)
+                    st.pyplot(fig)
+                
+            except Exception as e:
+                st.error(f"Error executing query: {str(e)}")
+    
+    # Memory for context retention
+    if "query_history" not in st.session_state:
+        st.session_state.query_history = []
+    st.session_state.query_history.append(query)