Update app.py

app.py

@@ -1,104 +1,159 @@
-import streamlit as st
-import faiss
-import numpy as np
-from sentence_transformers import SentenceTransformer
-from groq import Groq
-import os
-import pypdf
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-
-# Set Groq API key
-GROQ_API_KEY = os.getenv("GROQ_API_KEY", "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")
-if not GROQ_API_KEY:
-    st.error("⚠️ GROQ_API_KEY is missing! Please set it in your environment variables.")
-    st.stop()
-
-# Load embedding model
-embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
-
-# Set up Groq client
-client = Groq(api_key=GROQ_API_KEY)
-
-# Function to extract text from PDF
-def extract_text_from_pdf(uploaded_file):
-    reader = pypdf.PdfReader(uploaded_file)
-    extracted_text = [page.extract_text() for page in reader.pages if page.extract_text()]
-    return "\n".join(extracted_text) if extracted_text else "No text could be extracted from this PDF."
-
-# Function to create text chunks
-def create_chunks(text, chunk_size=500, chunk_overlap=100):
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
-    return text_splitter.split_text(text)
-
-# Function to create and save FAISS index
-def create_faiss_index(chunks):
-    embeddings = embedding_model.encode(chunks, convert_to_numpy=True)
-
-    # Create FAISS index
-    dimension = embeddings.shape[1]
-    index = faiss.IndexFlatL2(dimension)
-    index.add(embeddings)
-
-    # Save FAISS index and chunks
-    faiss.write_index(index, "faiss_index.bin")
-    np.save("chunks.npy", np.array(chunks, dtype=object))
-
-    return index, chunks
-
-# Function to search FAISS
-def search_faiss(query, index, chunks, top_k=2):
-    if index is None:
-        return ["No database found. Please upload a PDF first."]
-
-    query_embedding = embedding_model.encode([query], convert_to_numpy=True)
-    distances, indices = index.search(query_embedding, top_k)
-    return [chunks[i] for i in indices[0] if i < len(chunks)]
-
-# Function to query Groq
-def query_groq(query):
-    chat_completion = client.chat.completions.create(
-        messages=[{"role": "user", "content": query}],
-        model="llama-3.3-70b-versatile",
-    )
-    return chat_completion.choices[0].message.content
-
-# Streamlit UI
-st.set_page_config(page_title="RAG Chatbot", page_icon="🤖")
-st.title("📄 RAG-Based Chatbot with FAISS & Groq")
-
-# Upload PDF
-uploaded_file = st.file_uploader("📤 Upload a PDF file", type="pdf")
-
-if uploaded_file:
-    with st.spinner("🔄 Processing PDF..."):
-        text = extract_text_from_pdf(uploaded_file)
-        if text.strip():
-            chunks = create_chunks(text)
-
-            # Create FAISS index
-            index, chunks = create_faiss_index(chunks)
-
-            # Store in session state
-            st.session_state["faiss_index"] = index
-            st.session_state["chunks"] = chunks
-
-            st.success("✅ PDF processed successfully!")
-        else:
-            st.error("❌ No text found in the uploaded PDF.")
-
-# Load FAISS index if available
-index = st.session_state.get("faiss_index", None)
-chunks = st.session_state.get("chunks", [])
-
-# User query input
-user_query = st.text_input("💬 Ask me something about the document:")
-
-if st.button("Search") and user_query:
-    with st.spinner("🔎 Retrieving response..."):
-        retrieved_text = search_faiss(user_query, index, chunks)
-        if retrieved_text:
-            response = query_groq("\n".join(retrieved_text))
-            st.write("### 🤖 AI Response:")
-            st.write(response)
-        else:
-            st.error("⚠️ No relevant information found.")
+import streamlit as st
+import faiss
+import numpy as np
+from sentence_transformers import SentenceTransformer
+from groq import Groq
+import os
+import pypdf
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+# Initialize session state variables
+if "faiss_index" not in st.session_state:
+    st.session_state["faiss_index"] = None
+if "chunks" not in st.session_state:
+    st.session_state["chunks"] = []
+
+# Set Groq API key - Consider using st.secrets for better security
+GROQ_API_KEY = os.getenv("GROQ_API_KEY") or st.secrets.get("GROQ_API_KEY", "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")
+if not GROQ_API_KEY:
+    st.error("⚠️ GROQ_API_KEY is missing! Please set it in your environment variables or secrets.toml file.")
+    st.stop()
+
+# Load embedding model with error handling
+try:
+    embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
+except Exception as e:
+    st.error(f"❌ Failed to load embedding model: {str(e)}")
+    st.stop()
+
+# Set up Groq client with error handling
+try:
+    client = Groq(api_key=GROQ_API_KEY)
+except Exception as e:
+    st.error(f"❌ Failed to initialize Groq client: {str(e)}")
+    st.stop()
+
+# Function to extract text from PDF with error handling
+def extract_text_from_pdf(uploaded_file):
+    try:
+        reader = pypdf.PdfReader(uploaded_file)
+        extracted_text = [page.extract_text() for page in reader.pages if page.extract_text()]
+        return "\n".join(extracted_text) if extracted_text else ""
+    except Exception as e:
+        st.error(f"❌ Error extracting text from PDF: {str(e)}")
+        return ""
+
+# Function to create text chunks
+def create_chunks(text, chunk_size=500, chunk_overlap=100):
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
+        separators=["\n\n", "\n", " ", ""]  # Added separators for better splitting
+    )
+    return text_splitter.split_text(text)
+
+# Function to create and save FAISS index
+def create_faiss_index(chunks):
+    try:
+        embeddings = embedding_model.encode(chunks, convert_to_numpy=True)
+        
+        # Create FAISS index
+        dimension = embeddings.shape[1]
+        index = faiss.IndexFlatL2(dimension)
+        index.add(embeddings)
+        
+        return index, chunks
+    except Exception as e:
+        st.error(f"❌ Error creating FAISS index: {str(e)}")
+        return None, []
+
+# Function to search FAISS
+def search_faiss(query, index, chunks, top_k=2):
+    if index is None or not chunks:
+        return []
+    
+    try:
+        query_embedding = embedding_model.encode([query], convert_to_numpy=True)
+        distances, indices = index.search(query_embedding, top_k)
+        return [chunks[i] for i in indices[0] if i < len(chunks)]
+    except Exception as e:
+        st.error(f"❌ Search error: {str(e)}")
+        return []
+
+# Function to query Groq with enhanced prompt
+def query_groq(query, context=None):
+    try:
+        prompt = f"""Use the following context to answer the question. 
+        If you don't know the answer, say you don't know. Don't make up answers.
+        
+        Context: {context if context else 'No specific context provided'}
+        
+        Question: {query}
+        
+        Answer:"""
+        
+        chat_completion = client.chat.completions.create(
+            messages=[{"role": "user", "content": prompt}],
+            model="llama-3-70b-8192",  # Updated model name
+            temperature=0.3,
+            max_tokens=1024
+        )
+        return chat_completion.choices[0].message.content
+    except Exception as e:
+        return f"Error querying Groq: {str(e)}"
+
+# Streamlit UI
+st.set_page_config(page_title="RAG Chatbot", page_icon="🤖", layout="wide")
+st.title("📄 RAG-Based Chatbot with FAISS & Groq")
+
+# Sidebar for settings
+with st.sidebar:
+    st.header("Settings")
+    top_k = st.slider("Number of chunks to retrieve", 1, 5, 2)
+    chunk_size = st.slider("Chunk size (characters)", 200, 1000, 500)
+    chunk_overlap = st.slider("Chunk overlap (characters)", 0, 200, 100)
+
+# Upload PDF
+uploaded_file = st.file_uploader("📤 Upload a PDF file", type="pdf")
+
+if uploaded_file:
+    with st.spinner("🔄 Processing PDF..."):
+        text = extract_text_from_pdf(uploaded_file)
+        if text.strip():
+            chunks = create_chunks(text, chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+            
+            # Create FAISS index
+            index, chunks = create_faiss_index(chunks)
+            
+            # Store in session state
+            st.session_state["faiss_index"] = index
+            st.session_state["chunks"] = chunks
+            
+            st.success(f"✅ PDF processed successfully! Created {len(chunks)} chunks.")
+        else:
+            st.error("❌ No text found in the uploaded PDF.")
+
+# Chat interface
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Display chat messages
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# User query input
+if prompt := st.chat_input("💬 Ask me something about the document:"):
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user"):
+        st.markdown(prompt)
+    
+    with st.spinner("🔎 Retrieving response..."):
+        retrieved_text = search_faiss(prompt, st.session_state["faiss_index"], st.session_state["chunks"], top_k=top_k)
+        context = "\n".join(retrieved_text) if retrieved_text else "No relevant context found."
+        
+        response = query_groq(prompt, context)
+        
+        st.session_state.messages.append({"role": "assistant", "content": response})
+        with st.chat_message("assistant"):
+            st.markdown(response)