Update app.py

app.py

@@ -1,114 +1,198 @@
 import os
 import gradio as gr
-from pypdf import PdfReader
-from groq import Groq
-
-def extract_text_from_pdf(pdf_path):
-    reader = PdfReader(pdf_path)
-    pages = []
-    for i, page in enumerate(reader.pages):
-        text = page.extract_text()
-        if text and text.strip():
-            pages.append(f"[Page {i+1}]\n{text.strip()}")
-    return "\n\n".join(pages) if pages else None
-
-def truncate_text(text, max_chars=6000):
-    if len(text) <= max_chars:
-        return text
-    return text[:4000] + "\n\n...[middle truncated]...\n\n" + text[-2000:]
-
-def ask_groq(pdf_text, question, model, history):
-    client = Groq(api_key=os.environ["GROQ_API_KEY"])
-    safe_pdf_text = truncate_text(pdf_text)
-
-    messages = [{
-        "role": "system",
-        "content": (
-            "You are a helpful assistant. Answer questions ONLY from the provided PDF. "
-            "If the answer is not in the document, say so. Mention page numbers when possible."
-        )
-    }]
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain_groq import ChatGroq
+from langchain_core.prompts import PromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnableLambda
+
+# ── Global state ──────────────────────────────────────────────
+vectorstore    = None
+qa_chain       = None
+retrieved_docs = {}
+
+# ── Embedding model (loaded once) ─────────────────────────────
+embeddings = HuggingFaceEmbeddings(
+    model_name="sentence-transformers/all-MiniLM-L6-v2",
+    model_kwargs={"device": "cpu"},
+    encode_kwargs={"normalize_embeddings": True}
+)
+
+# ── Helpers ───────────────────────────────────────────────────
+def format_docs(docs):
+    return "\n\n".join(
+        f"[Page {doc.metadata.get('page', '?') + 1}]\n{doc.page_content}"
+        for doc in docs
+    )
+
+def build_chain(groq_api_key: str):
+    """Build the LCEL RAG chain with the provided Groq key."""
+    llm = ChatGroq(
+        api_key=groq_api_key,
+        model="llama-3.3-70b-versatile",
+        temperature=0.2,
+        max_tokens=1024,
+    )
 
-    for user_msg, bot_msg in history:
-        messages.append({"role": "user",      "content": user_msg})
-        messages.append({"role": "assistant", "content": bot_msg})
+    PROMPT = PromptTemplate(
+        template="""You are a helpful assistant. Use the context below to answer the question.
+If the answer is not in the context, say "I don't have enough information to answer that."
 
-    if not history:
-        content = f"PDF Content:\n\n{safe_pdf_text}\n\n---\nQuestion: {question}"
-    else:
-        content = question
+Context:
+{context}
 
-    messages.append({"role": "user", "content": content})
+Question: {question}
 
-    response = client.chat.completions.create(
-        model=model,
-        messages=messages,
-        temperature=0.2,
-        max_tokens=1024,
+Answer:""",
+        input_variables=["context", "question"]
     )
-    return response.choices[0].message.content
 
-def load_pdf(pdf_file, state):
+    retriever = vectorstore.as_retriever(
+        search_type="similarity",
+        search_kwargs={"k": 4}
+    )
+
+    def retrieve_and_format(input_dict):
+        question = input_dict["question"]
+        docs = retriever.invoke(question)
+        retrieved_docs["docs"] = docs
+        return {
+            "context": format_docs(docs),
+            "question": question
+        }
+
+    chain = (
+        RunnableLambda(retrieve_and_format)
+        | PROMPT
+        | llm
+        | StrOutputParser()
+    )
+    return chain
+
+# ── Core functions ────────────────────────────────────────────
+def process_pdf(pdf_file, groq_api_key, progress=gr.Progress()):
+    global vectorstore, qa_chain
+
     if pdf_file is None:
-        return state, "⚠️ Please upload a PDF.", gr.update(interactive=False)
-    pdf_text = extract_text_from_pdf(pdf_file)
-    if not pdf_text:
-        return state, "❌ Could not extract text. Try a non-scanned PDF.", gr.update(interactive=False)
-    state["pdf_text"] = pdf_text
-    state["history"]  = []
-    pages = pdf_text.count("[Page ")
-    words = len(pdf_text.split())
-    return state, f"✅ Loaded! {pages} page(s), ~{words:,} words.\nNow ask your questions!", gr.update(interactive=True)
-
-def ask_question(question, model, state, chat_history):
+        return "⚠️ Please upload a PDF file."
+    if not groq_api_key.strip():
+        return "⚠️ Please enter your Groq API key."
+
+    try:
+        progress(0.1, desc="Loading PDF...")
+        loader = PyPDFLoader(pdf_file.name)
+        pages = loader.load()
+
+        progress(0.3, desc="Splitting into chunks...")
+        splitter = RecursiveCharacterTextSplitter(
+            chunk_size=800,
+            chunk_overlap=100,
+            separators=["\n\n", "\n", ".", " "]
+        )
+        chunks = splitter.split_documents(pages)
+
+        progress(0.6, desc="Building FAISS index...")
+        vectorstore = FAISS.from_documents(chunks, embeddings)
+
+        progress(0.9, desc="Setting up RAG chain...")
+        qa_chain = build_chain(groq_api_key.strip())
+
+        progress(1.0, desc="Done!")
+        return f"✅ Ready! Loaded **{len(pages)} pages** → **{len(chunks)} chunks**."
+
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+
+
+def answer_question(question, history):
+    if vectorstore is None or qa_chain is None:
+        history.append((question, "⚠️ Please upload a PDF and enter your Groq API key first."))
+        return "", history
     if not question.strip():
-        return chat_history, state, ""
-    if not state.get("pdf_text"):
-        chat_history.append((question, "⚠️ Please load a PDF first."))
-        return chat_history, state, ""
+        return "", history
+
     try:
-        answer = ask_groq(state["pdf_text"], question, model, state["history"])
-        state["history"].append((question, answer))
-        chat_history.append((question, answer))
+        answer = qa_chain.invoke({"question": question})
+
+        docs = retrieved_docs.get("docs", [])
+        if docs:
+            pages = sorted(set(
+                doc.metadata.get("page", 0) + 1
+                for doc in docs
+            ))
+            answer += f"\n\n📄 *Sources: pages {pages}*"
+
     except Exception as e:
-        chat_history.append((question, f"❌ Error: {e}"))
-    return chat_history, state, ""
+        answer = f"❌ Error generating answer: {str(e)}"
 
-def clear_chat(state):
-    state["history"] = []
-    return [], state
+    history.append((question, answer))
+    return "", history
 
-MODELS = [
-    "llama-3.3-70b-versatile",
-    "llama-3.1-8b-instant",
-    "llama3-70b-8192",
-    "gemma2-9b-it",
-]
 
-with gr.Blocks(title="PDF Q&A — Groq", theme=gr.themes.Soft()) as demo:
+def clear_all():
+    global vectorstore, qa_chain, retrieved_docs
+    vectorstore    = None
+    qa_chain       = None
+    retrieved_docs = {}
+    return [], "", "🗑️ Cleared. Upload a new PDF to start again."
 
-    state = gr.State({"pdf_text": "", "history": []})
+# ── Gradio UI ─────────────────────────────────────────────────
+with gr.Blocks(title="PDF RAG Chatbot", theme=gr.themes.Soft()) as demo:
 
-    gr.Markdown("# 📄 PDF Q&A App\n### Upload a PDF → Ask questions → Get AI answers")
+    gr.Markdown("""
+    # 📚 PDF RAG Chatbot
+    Upload a PDF, enter your Groq API key, then ask questions about the document.
+    """)
 
     with gr.Row():
         with gr.Column(scale=1):
-            pdf_input  = gr.File(label="Upload PDF", file_types=[".pdf"], type="filepath")
-            upload_btn = gr.Button("📤 Load PDF", variant="primary", size="lg")
-            status     = gr.Textbox(label="Status", value="Upload a PDF to begin...", lines=3, interactive=False)
-            model_drop = gr.Dropdown(choices=MODELS, value=MODELS[0], label="🤖 Groq Model")
+            gr.Markdown("### ⚙️ Setup")
+            groq_key_box = gr.Textbox(
+                label="Groq API Key",
+                placeholder="gsk_...",
+                type="password"
+            )
+            pdf_upload = gr.File(
+                label="Upload PDF",
+                file_types=[".pdf"]
+            )
+            process_btn = gr.Button("📥 Process PDF", variant="primary")
+            status_box  = gr.Markdown("*Upload a PDF to begin.*")
 
         with gr.Column(scale=2):
-            chatbot = gr.Chatbot(label="Conversation", height=420)
+            gr.Markdown("### 💬 Chat")
+            chatbot = gr.Chatbot(height=480, bubble_full_width=False)
             with gr.Row():
-                q_input = gr.Textbox(label="Your Question", placeholder="e.g. Summarize this document", lines=2, interactive=False, scale=4)
-                with gr.Column(scale=1):
-                    ask_btn   = gr.Button("🔍 Ask",   variant="primary")
-                    clear_btn = gr.Button("🗑️ Clear", variant="secondary")
-
-    upload_btn.click(load_pdf,     [pdf_input, state],                    [state, status, q_input])
-    ask_btn.click(ask_question,    [q_input, model_drop, state, chatbot], [chatbot, state, q_input])
-    q_input.submit(ask_question,   [q_input, model_drop, state, chatbot], [chatbot, state, q_input])
-    clear_btn.click(clear_chat,    [state],                               [chatbot, state])
+                question_box = gr.Textbox(
+                    placeholder="Ask a question about your PDF...",
+                    show_label=False,
+                    scale=4
+                )
+                submit_btn = gr.Button("Send", variant="primary", scale=1)
+            clear_btn = gr.Button("🗑️ Clear Chat & Reset")
+
+    # ── Event handlers ─────────────────────────────────────────
+    process_btn.click(
+        process_pdf,
+        inputs=[pdf_upload, groq_key_box],
+        outputs=[status_box]
+    )
+    submit_btn.click(
+        answer_question,
+        inputs=[question_box, chatbot],
+        outputs=[question_box, chatbot]
+    )
+    question_box.submit(
+        answer_question,
+        inputs=[question_box, chatbot],
+        outputs=[question_box, chatbot]
+    )
+    clear_btn.click(
+        clear_all,
+        outputs=[chatbot, question_box, status_box]
+    )
 
-demo.launch()  # ← No share=True needed on Hugging Face
+demo.launch()