Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 import gradio as gr
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import Chroma
@@ -5,23 +6,15 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.document_loaders import PyPDFLoader
 from langchain.chains import RetrievalQA
 from langchain.llms.base import LLM
-from groq import Groq
 from typing import List, Optional
+from groq import Groq
+import tempfile
+import shutil
 
-# Step 1: Load PDF and prepare vector store
-loader = PyPDFLoader("ivas103.pdf")
-documents = loader.load()
-
-text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-docs = text_splitter.split_documents(documents)
-
-embedding = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
-vectorstore = Chroma.from_documents(docs, embedding, persist_directory="rag_chroma_groq")
-
-# Step 2: Define custom LLM class using Groq
+# Custom LLM using Groq
 class GroqLLM(LLM):
     model: str = "llama3-8b-8192"
-    api_key: str = "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"  # Replace with your key
+    api_key: str = os.environ.get("GROQ_API_KEY")  # Use env var for security
     temperature: float = 0.7
 
     def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
@@ -41,30 +34,64 @@ class GroqLLM(LLM):
     def _llm_type(self) -> str:
         return "groq-llm"
 
-# Step 3: Create RetrievalQA chain
-retriever = vectorstore.as_retriever()
-groq_llm = GroqLLM(api_key="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")  # Replace with your key
+# Global cache to reuse vectorstore during the session
+vectorstore_cache = {}
+
+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())
+
+        # Load and split
+        loader = PyPDFLoader(file_path)
+        documents = loader.load()
+
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+        docs = text_splitter.split_documents(documents)
 
-qa_chain = RetrievalQA.from_chain_type(
-    llm=groq_llm,
-    retriever=retriever,
-    return_source_documents=True
-)
+        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.persist()
+
+        # Store for session use
+        vectorstore_cache["retriever"] = vectorstore.as_retriever()
+    
+    return "PDF processed and ready. You can now ask questions."
 
-# Step 4: Gradio interface function
 def ask_question(query):
+    if "retriever" not in vectorstore_cache:
+        return "Please upload a PDF first."
+
+    llm = GroqLLM()
+    qa_chain = RetrievalQA.from_chain_type(
+        llm=llm,
+        retriever=vectorstore_cache["retriever"],
+        return_source_documents=True
+    )
     result = qa_chain({"query": query})
     answer = result["result"]
-    sources = "\n\n".join([doc.metadata.get("source", "Unknown") for doc in result["source_documents"]])
+    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}"
 
-# Step 5: Launch Gradio UI
-iface = gr.Interface(
-    fn=ask_question,
-    inputs=gr.Textbox(label="Ask a question", placeholder="e.g., What is a chassis?"),
-    outputs=gr.Markdown(),
-    title="📄 PDF RAG Chatbot (Groq + LangChain)",
-    description="Ask questions based on the content of the PDF file."
-)
+with gr.Blocks() as demo:
+    gr.Markdown("## 📄 PDF Question Answering Bot (Groq + HuggingFace + LangChain)")
+
+    with gr.Row():
+        pdf_file = gr.File(label="Upload your PDF")
+        upload_btn = gr.Button("Process PDF")
+
+    upload_output = gr.Textbox(label="Status", interactive=False)
+    upload_btn.click(process_pdf, inputs=pdf_file, outputs=upload_output)
+
+    query = gr.Textbox(label="Ask a question")
+    answer_output = gr.Markdown()
+    query_btn = gr.Button("Get Answer")
+
+    query_btn.click(ask_question, inputs=query, outputs=answer_output)
 
-iface.launch()
+demo.launch()