Update app.py

app.py

@@ -1,90 +1,89 @@
-import os
 import gradio as gr
-from typing import List
+import tempfile
+import os
+
 from langchain_community.document_loaders import PyPDFLoader
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_google_genai import GoogleGenerativeAIEmbeddings, ChatGoogleGenerativeAI
 from langchain_community.vectorstores import FAISS
-import google.generativeai as genai
-
-# Path to save vector index
-INDEX_DIR = "rag_multi_pdf_index"
-
-# Step 1: Create knowledge base from PDFs
-def create_knowledge_base(pdf_files: List[gr.File]) -> str:
-    if not pdf_files:
-        return "❌ No PDFs uploaded."
-
-    all_chunks = []
-    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
-
-    for file in pdf_files:
-        loader = PyPDFLoader(file.name)
-        try:
-            docs = loader.load()
-            chunks = splitter.split_documents(docs)
-            all_chunks.extend(chunks)
-        except Exception as e:
-            return f"❌ Error reading {file.name}: {str(e)}"
-
-    if not all_chunks:
-        return "❌ No content extracted from PDFs."
-
-    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
-    vectorstore = FAISS.from_documents(all_chunks, embeddings)
-    vectorstore.save_local(INDEX_DIR)
-
-    return f"✅ Knowledge base created with {len(all_chunks)} chunks from {len(pdf_files)} PDFs."
-
-# Step 2: Load vectorstore
-def load_vectorstore() -> FAISS:
-    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
-    return FAISS.load_local(INDEX_DIR, embeddings)
-
-# Step 3: Ask question via Gemini using retrieved context
-def chat_with_rag(api_key: str, question: str) -> str:
-    if not api_key or not api_key.startswith("AI"):
-        return "❌ Invalid Gemini API Key. It should start with 'AI'."
-
+from langchain.chains import RetrievalQA
+from langchain.callbacks.base import BaseCallbackHandler
+
+# Global state
+kb = None
+retriever = None
+qa = None
+
+class StreamHandler(BaseCallbackHandler):
+    def __init__(self, update_fn):
+        self.text = ""
+        self.update_fn = update_fn
+
+    def on_llm_new_token(self, token: str, **kwargs):
+        self.text += token
+        self.update_fn(self.text)
+
+def save_pdfs(pdf_list):
+    paths = []
+    for pdf in pdf_list:
+        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
+        tmp.write(pdf.read())
+        tmp.close()
+        paths.append(tmp.name)
+    return paths
+
+def create_kb(api_key, pdf_list):
+    global retriever, qa
     try:
-        genai.configure(api_key=api_key)
-        model = genai.GenerativeModel("gemini-pro")
+        pdf_paths = save_pdfs(pdf_list)
+        docs = []
+        for path in pdf_paths:
+            loader = PyPDFLoader(path)
+            docs.extend(loader.load())
+        splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+        chunks = splitter.split_documents(docs)
+        embeddings = GoogleGenerativeAIEmbeddings(model="models/text-embedding-004", google_api_key=api_key)
+        db = FAISS.from_documents(chunks, embeddings)
+        retriever = db.as_retriever(search_kwargs={"k": 3})
+        qa = RetrievalQA.from_chain_type(llm=None, retriever=retriever)  # llm passed later
+        return "✅ Knowledge base created."
     except Exception as e:
-        return f"❌ Gemini configuration error: {str(e)}"
-
-    try:
-        vs = load_vectorstore()
-        top_docs = vs.similarity_search(question, k=3)
-        context = "\n\n".join([doc.page_content for doc in top_docs])
-    except Exception as e:
-        return f"❌ Error loading vectorstore or retrieving context: {str(e)}"
-
-    prompt = f"""Use the following context to answer the question:\n\n{context}\n\nQuestion: {question}"""
-
-    try:
-        response = model.generate_content(prompt)
-        return response.text
-    except Exception as e:
-        return f"❌ Gemini error: {str(e)}"
-
-# Step 4: Gradio UI
-with gr.Blocks(title="📚 RAG Q&A with Gemini") as demo:
-    gr.Markdown("## 📄 Upload multiple PDFs → 🧠 Build Knowledge Base → 🤖 Ask Questions with Gemini")
-
-    api_key = gr.Textbox(label="🔐 Gemini API Key", placeholder="Enter your Gemini API Key", type="password")
-
-    pdfs = gr.File(label="📂 Upload PDFs", file_types=[".pdf"], file_count="multiple")
-    create_btn = gr.Button("📄 Create Knowledge Base")
-    kb_status = gr.Textbox(label="📦 Knowledge Base Status", interactive=False)
-
-    create_btn.click(fn=create_knowledge_base, inputs=[pdfs], outputs=[kb_status])
-
-    question = gr.Textbox(label="❓ Ask a Question")
-    answer = gr.Textbox(label="💬 Gemini Answer", lines=10, interactive=False)
-    ask_btn = gr.Button("🚀 Ask")
-
-    ask_btn.click(fn=chat_with_rag, inputs=[api_key, question], outputs=[answer])
-
-# Step 5: Launch app
-if __name__ == "__main__":
-    demo.launch()
+        return f"❌ Error creating KB: {e}"
+
+def ask_question(api_key, question, chat_history, set_stream):
+    global retriever, qa
+    if retriever is None:
+        return chat_history, "❌ Create KB first."
+    handler = StreamHandler(lambda txt: set_stream(txt))
+    llm = ChatGoogleGenerativeAI(model="models/gemini-1.5-pro-latest",
+                                 google_api_key=api_key,
+                                 streaming=True,
+                                 callbacks=[handler])
+    qa.llm = llm
+    chat_history = chat_history or []
+    chat_history.append({"role": "user", "content": question})
+    result = qa.invoke({"query": question})
+    chat_history.append({"role": "assistant", "content": handler.text})
+    return chat_history, ""
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 📚 Multi‑PDF RAG Chat with Gemini")
+
+    with gr.Column():
+        api_key = gr.Textbox(show_label=False, placeholder="Enter your Gemini API Key", type="password")
+        pdfs = gr.File(file_types=[".pdf"], label="Upload PDFs", file_count="multiple")
+        kb_status = gr.Textbox(label="Status")
+        create_btn = gr.Button("▶️ Create Knowledge Base")
+
+    create_btn.click(create_kb, inputs=[api_key, pdfs], outputs=kb_status)
+
+    chatbot = gr.Chatbot(label="🧠 Assistant", type="messages")
+    question = gr.Textbox(show_label=False, placeholder="Ask a question")
+    stream_output = gr.State("")  # to capture stream text
+    send = gr.Button("🔍 Ask")
+
+    send.click(fn=ask_question,
+               inputs=[api_key, question, chatbot, stream_output],
+               outputs=[chatbot, ""])
+
+demo.launch()