Update app.py

app.py

@@ -1,20 +1,20 @@
 import os
 import gradio as gr
+import tempfile
+from typing import List, Optional
+
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import Chroma
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.document_loaders import PyPDFLoader
 from langchain.chains import RetrievalQA
 from langchain.llms.base import LLM
-from typing import List, Optional
 from groq import Groq
-import tempfile
-import shutil
 
-# Custom LLM using Groq
+# ---- Custom GroqLLM class using LangChain LLM base ----
 class GroqLLM(LLM):
     model: str = "llama3-8b-8192"
-    api_key: str = os.environ.get("GROQ_API_KEY")  # Use env var for security
+    api_key: str = os.environ.get("GROQ_API_KEY")  # Load from HF secrets
     temperature: float = 0.7
 
     def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
@@ -34,64 +34,68 @@ class GroqLLM(LLM):
     def _llm_type(self) -> str:
         return "groq-llm"
 
-# Global cache to reuse vectorstore during the session
-vectorstore_cache = {}
+# Global cache for vectorstore
+rag_context = {"retriever": None}
+
+# ---- Step 1: Upload & Embed PDF ----
+def process_pdf(file):
+    if file is None:
+        return "❌ Please upload a PDF."
 
-def process_pdf(file_obj):
-    # Save uploaded PDF to temp directory
     with tempfile.TemporaryDirectory() as temp_dir:
-        file_path = os.path.join(temp_dir, file_obj.name)
-        with open(file_path, "wb") as f:
-            f.write(file_obj.read())
+        temp_pdf_path = os.path.join(temp_dir, file.name)
+        with open(temp_pdf_path, "wb") as f:
+            f.write(file.read())
 
-        # Load and split
-        loader = PyPDFLoader(file_path)
+        # Load and split PDF
+        loader = PyPDFLoader(temp_pdf_path)
         documents = loader.load()
 
         text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-        docs = text_splitter.split_documents(documents)
+        chunks = text_splitter.split_documents(documents)
 
+        # Embeddings and vectorstore
         embedding = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
-
-        # Create persistent Chroma DB
-        persist_dir = os.path.join(temp_dir, "chroma_db")
-        vectorstore = Chroma.from_documents(docs, embedding, persist_directory=persist_dir)
+        vectorstore = Chroma.from_documents(chunks, embedding, persist_directory=temp_dir)
         vectorstore.persist()
 
-        # Store for session use
-        vectorstore_cache["retriever"] = vectorstore.as_retriever()
-    
-    return "PDF processed and ready. You can now ask questions."
+        # Store retriever in session
+        rag_context["retriever"] = vectorstore.as_retriever()
 
+    return "✅ PDF processed and ready. Ask your questions!"
+
+# ---- Step 2: Ask questions to the RAG chain ----
 def ask_question(query):
-    if "retriever" not in vectorstore_cache:
-        return "Please upload a PDF first."
+    retriever = rag_context.get("retriever")
+    if retriever is None:
+        return "❌ Please upload and process a PDF first."
 
     llm = GroqLLM()
     qa_chain = RetrievalQA.from_chain_type(
         llm=llm,
-        retriever=vectorstore_cache["retriever"],
+        retriever=retriever,
         return_source_documents=True
     )
+
     result = qa_chain({"query": query})
     answer = result["result"]
-    sources = "\n".join([doc.metadata.get("source", "No metadata") for doc in result["source_documents"]])
-    return f"### Answer:\n{answer}\n\n### Sources:\n{sources}"
+    return f"### Answer:\n{answer}"
 
+# ---- Gradio UI ----
 with gr.Blocks() as demo:
-    gr.Markdown("## 📄 PDF Question Answering Bot (Groq + HuggingFace + LangChain)")
+    gr.Markdown("# 📚 RAG Chatbot with Groq & LangChain\nUpload a PDF, then ask questions about it!")
 
     with gr.Row():
-        pdf_file = gr.File(label="Upload your PDF")
+        pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"])
         upload_btn = gr.Button("Process PDF")
+    upload_status = gr.Textbox(label="Status", interactive=False)
 
-    upload_output = gr.Textbox(label="Status", interactive=False)
-    upload_btn.click(process_pdf, inputs=pdf_file, outputs=upload_output)
+    upload_btn.click(process_pdf, inputs=pdf_input, outputs=upload_status)
 
-    query = gr.Textbox(label="Ask a question")
+    query_input = gr.Textbox(label="Ask a question")
+    ask_btn = gr.Button("Get Answer")
     answer_output = gr.Markdown()
-    query_btn = gr.Button("Get Answer")
 
-    query_btn.click(ask_question, inputs=query, outputs=answer_output)
+    ask_btn.click(ask_question, inputs=query_input, outputs=answer_output)
 
 demo.launch()