Update app.py

app.py

@@ -1,40 +1,40 @@
-import streamlit as st
-import requests
 import os
+import tempfile
+import gc
+import logging
+import streamlit as st
+from groq import Groq, APIError
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import Chroma
+import torch
+
+# ---------------- CONFIGURATION ----------------
+logging.basicConfig(level=logging.INFO)
+
+# Load API key from Hugging Face secrets
+GROQ_API_KEY = st.secrets.get("GROQ_API_KEY", os.environ.get("GROQ_API_KEY"))
+GROQ_MODEL = "llama-3.1-8b-instant"
+
+# Initialize Groq client
+client = None
+if GROQ_API_KEY:
+    try:
+        client = Groq(api_key=GROQ_API_KEY)
+        st.success("✅ Groq client initialized successfully.")
+    except Exception as e:
+        st.error(f"❌ Failed to initialize Groq client: {e}")
+        client = None
+else:
+    st.warning("⚠️ GROQ_API_KEY not found. Please add it to Hugging Face secrets.")
 
-# ============== BACKEND URL CONFIGURATION ==============
-# For Hugging Face deployment
-import socket
-
-def get_backend_url():
-    """Determine backend URL based on environment."""
-    # Check if we're on Hugging Face
-    if os.environ.get("SPACE_ID"):
-        # On Hugging Face, backend is on localhost:8000
-        return "http://localhost:8000"
-    else:
-        # Local development
-        return "http://localhost:8000"
-
-BACKEND_URL = get_backend_url()
-# ======================================================
-
+# ---------------- STREAMLIT UI SETUP ----------------
 st.set_page_config(page_title="PDF Assistant", page_icon="📘", layout="wide")
 
-# Test backend connection on startup
-try:
-    health_check = requests.get(f"{BACKEND_URL}/health", timeout=2)
-    if health_check.status_code == 200:
-        st.sidebar.success("✅ Backend connected")
-    else:
-        st.sidebar.warning(f"⚠️ Backend status: {health_check.status_code}")
-except:
-    st.sidebar.error("❌ Cannot connect to backend")
-
-# ---------------- CSS (Dark Theme) ----------------
+# ---------------- CSS (Your exact UI) ----------------
 st.markdown("""
 <style>
-/* Streamlit standard setup for dark theme adherence */
 :root {
     --primary-color: #1e3a8a;
     --background-color: #0e1117;
@@ -42,7 +42,6 @@ st.markdown("""
     --text-color: #f0f2f6;
 }
 
-/* Custom Chat Bubbles */
 .chat-user {
     background: #2d3748;
     padding: 12px;
@@ -62,7 +61,6 @@ st.markdown("""
     color: #ffffff;
 }
 
-/* Sources section styling */
 .sources {
     font-size: 0.8em;
     opacity: 0.7;
@@ -71,7 +69,6 @@ st.markdown("""
     padding-top: 5px;
 }
 
-/* Footer styling */
 .footer {
     position: fixed;
     left: 0;
@@ -92,13 +89,6 @@ st.markdown("""
 .footer a:hover {
     text-decoration: underline;
 }
-
-/* Debug info */
-.debug-info {
-    font-size: 0.7em;
-    color: #888;
-    margin-top: 5px;
-}
 </style>
 """, unsafe_allow_html=True)
 
@@ -106,47 +96,138 @@ st.markdown("""
 if "chat" not in st.session_state:
     st.session_state.chat = []
 
+if "vectorstore" not in st.session_state:
+    st.session_state.vectorstore = None
+
+if "retriever" not in st.session_state:
+    st.session_state.retriever = None
+
 if "uploaded_file_name" not in st.session_state:
     st.session_state.uploaded_file_name = None
 
 if "uploader_key" not in st.session_state:
     st.session_state.uploader_key = 0
 
-# Show backend URL for debugging
-st.sidebar.markdown(f"**Backend URL:** `{BACKEND_URL}`", unsafe_allow_html=True)
-
 # ---------------- FUNCTIONS ----------------
 def clear_chat_history():
     st.session_state.chat = []
 
 def clear_memory():
+    st.session_state.vectorstore = None
+    st.session_state.retriever = None
+    st.session_state.uploaded_file_name = None
+    st.session_state.uploader_key += 1
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    st.success("Memory cleared. Please upload a new PDF.")
+
+def process_pdf(uploaded_file):
+    """Process uploaded PDF and create vectorstore."""
     try:
-        res = requests.post(f"{BACKEND_URL}/clear")
-        if res.status_code == 200:
-            st.session_state.uploaded_file_name = None
-            st.session_state.uploader_key += 1
-            st.success("Memory cleared. Please upload a new PDF.")
-        else:
-            st.error(f"Failed to clear memory: {res.json().get('detail', 'Unknown error')}")
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
+            tmp.write(uploaded_file.getvalue())
+            path = tmp.name
+        
+        # Load PDF
+        loader = PyPDFLoader(path)
+        docs = loader.load()
+        
+        # Split into chunks
+        splitter = RecursiveCharacterTextSplitter(
+            chunk_size=800,
+            chunk_overlap=50
+        )
+        chunks = splitter.split_documents(docs)
+        
+        # Create embeddings
+        embeddings = HuggingFaceEmbeddings(
+            model_name="sentence-transformers/all-MiniLM-L6-v2",
+            model_kwargs={"device": "cpu"},
+            encode_kwargs={"normalize_embeddings": True}
+        )
+        
+        # Create vectorstore
+        vectorstore = Chroma.from_documents(chunks, embeddings)
+        retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
+        
+        # Store in session state
+        st.session_state.vectorstore = vectorstore
+        st.session_state.retriever = retriever
+        
+        # Cleanup
+        if os.path.exists(path):
+            os.unlink(path)
+        
+        return len(chunks)
+        
     except Exception as e:
-        st.error(f"Error clearing memory: {str(e)}")
+        st.error(f"Error processing PDF: {str(e)}")
+        return None
+
+def ask_question(question):
+    """Retrieve and generate answer for the question."""
+    if not client:
+        return None, 0, "Groq client is not initialized. Check API key setup."
+    
+    if not st.session_state.retriever:
+        return None, 0, "Upload PDF first to initialize the knowledge base."
+    
+    try:
+        # Retrieve relevant chunks
+        docs = st.session_state.retriever.invoke(question)
+        context = "\n\n".join(d.page_content for d in docs)
+        
+        # Build prompt
+        prompt = f"""
+You are a strict RAG Q&A assistant.
+Use ONLY the context provided. If the answer is not found, reply:
+"I cannot find this in the PDF."
+
+---------------- CONTEXT ----------------
+{context}
+-----------------------------------------
+
+QUESTION: {question}
+
+FINAL ANSWER:
+"""
+        
+        # Call Groq API
+        response = client.chat.completions.create(
+            model=GROQ_MODEL,
+            messages=[
+                {"role": "system",
+                 "content": "Use only the PDF content. If answer not found, say: 'I cannot find this in the PDF.'"},
+                {"role": "user", "content": prompt}
+            ],
+            temperature=0.0
+        )
+        
+        answer = response.choices[0].message.content.strip()
+        return answer, len(docs), None
+        
+    except APIError as e:
+        return None, 0, f"Groq API Error: {str(e)}"
+    except Exception as e:
+        return None, 0, f"General error: {str(e)}"
 
-# ---------------- SIDEBAR CONTROLS ----------------
+# ---------------- UI COMPONENTS ----------------
+st.title("📘 PDF Assistant")
+
+# Sidebar Controls
 with st.sidebar:
     st.header("Controls")
     st.button("🗑️ Clear Chat History", on_click=clear_chat_history, use_container_width=True)
     st.button("🔥 Clear PDF Memory", on_click=clear_memory, use_container_width=True)
-
+    
     st.markdown("---")
     if st.session_state.uploaded_file_name:
         st.success(f"✅ **Active PDF:**\n `{st.session_state.uploaded_file_name}`")
     else:
         st.warning("⬆️ Upload a PDF to start chatting!")
 
-# ---------------- MAIN APP ----------------
-st.title("📘 PDF Assistant")
-
-# ---------------- UPLOAD PDF ----------------
+# File Upload
 uploaded = st.file_uploader(
     "Upload your PDF",
     type=["pdf"],
@@ -156,33 +237,21 @@ uploaded = st.file_uploader(
 if uploaded and uploaded.name != st.session_state.uploaded_file_name:
     st.session_state.uploaded_file_name = None
     st.session_state.chat = []
-
+    
     with st.spinner(f"Processing '{uploaded.name}'..."):
-        try:
-            files = {"file": (uploaded.name, uploaded.getvalue(), "application/pdf")}
-            res = requests.post(f"{BACKEND_URL}/upload", files=files)
-
-            if res.status_code == 200:
-                chunks = res.json().get("chunks", 0)
-                st.success(f"✅ PDF processed successfully! {chunks} chunks created.")
-                st.session_state.uploaded_file_name = uploaded.name
-                st.markdown(f'<div class="debug-info">Backend: {BACKEND_URL}, Status: {res.status_code}</div>', unsafe_allow_html=True)
-            else:
-                error_msg = res.json().get("detail", f"HTTP {res.status_code}")
-                st.error(f"❌ Upload failed: {error_msg}")
-                st.session_state.uploaded_file_name = None
-
-        except requests.exceptions.ConnectionError:
-            st.error(f"🔌 Could not connect to backend at {BACKEND_URL}")
-            st.session_state.uploaded_file_name = None
-        except Exception as e:
-            st.error(f"⚠️ Unexpected error: {str(e)}")
+        chunks_count = process_pdf(uploaded)
+        
+        if chunks_count is not None:
+            st.success(f"✅ PDF processed successfully! {chunks_count} chunks created.")
+            st.session_state.uploaded_file_name = uploaded.name
+        else:
+            st.error("❌ Failed to process PDF")
             st.session_state.uploaded_file_name = None
-
+    
     st.rerun()
 
-# ---------------- CHAT INPUT ----------------
-disabled_input = st.session_state.uploaded_file_name is None
+# Chat Input
+disabled_input = st.session_state.uploaded_file_name is None or client is None
 question = st.text_input(
     "Ask a question about the loaded PDF:",
     key="question_input",
@@ -190,33 +259,22 @@ question = st.text_input(
 )
 
 if st.button("Send", disabled=disabled_input) and question:
+    # Add user query to chat history
     st.session_state.chat.append(("user", question))
-
+    
+    # Get answer
     with st.spinner("Thinking..."):
-        try:
-            res = requests.post(f"{BACKEND_URL}/ask", json={"question": question})
-
-            if res.status_code == 200:
-                data = res.json()
-                answer = data.get("answer", "No answer provided.")
-                sources = data.get("sources", 0)
-
-                bot_message = f"{answer}<div class='sources'>Context Chunks Used: {sources}</div>"
-                st.session_state.chat.append(("bot", bot_message))
-                st.markdown(f'<div class="debug-info">Backend: {BACKEND_URL}, Status: {res.status_code}</div>', unsafe_allow_html=True)
-
-            else:
-                error_detail = res.json().get("detail", f"HTTP {res.status_code}")
-                st.session_state.chat.append(("bot", f"🔴 **Error:** {error_detail}"))
-
-        except requests.exceptions.ConnectionError:
-            st.session_state.chat.append(("bot", f"🔴 **Error:** Could not connect to backend at {BACKEND_URL}"))
-        except Exception as e:
-            st.session_state.chat.append(("bot", f"🔴 **Unexpected error:** {str(e)}"))
-
+        answer, sources, error = ask_question(question)
+        
+        if answer:
+            bot_message = f"{answer}<div class='sources'>Context Chunks Used: {sources}</div>"
+            st.session_state.chat.append(("bot", bot_message))
+        else:
+            st.session_state.chat.append(("bot", f"🔴 **Error:** {error}"))
+    
     st.rerun()
 
-# ---------------- SHOW CHAT HISTORY ----------------
+# Display Chat History
 st.markdown("## Chat History")
 for role, msg in st.session_state.chat:
     if role == "user":
@@ -224,7 +282,7 @@ for role, msg in st.session_state.chat:
     else:
         st.markdown(f"<div class='chat-bot'>{msg}</div>", unsafe_allow_html=True)
 
-# ---------------- FOOTER ----------------
+# Footer
 footer_html = """
 <div class="footer">
     Created by <a href="https://www.linkedin.com/in/abhishek-iitr/" target="_blank">Abhishek Saxena</a>