Update app.py

app.py

@@ -1,108 +1,172 @@
-# =======================================
-# 📘 RAG App – Gemini + Local Embeddings
-# =======================================
-
-#!pip install google-generativeai sentence-transformers chromadb beautifulsoup4 PyPDF2 gradio
-
-import os, textwrap, re
-import google.generativeai as genai
-from sentence_transformers import SentenceTransformer
-from bs4 import BeautifulSoup
-import requests
-import chromadb
-from PyPDF2 import PdfReader
-import gradio as gr
-
-# ======================
-# 🔹 API Key Setup
-# ======================
-genai.configure(api_key="AIzaSyDr2X5N-hHt9EqUNy7JCm58aG1FpeGVpgs")  # apni key yahan daalo
-
-MODEL = 'gemini-2.5-flash'
-embedder = SentenceTransformer('all-MiniLM-L6-v2')  # local free embedding model
-chroma_client = chromadb.Client()
-collection = chroma_client.create_collection(name="rag_store")
-
-# ======================
-# 🔹 Helper Functions
-# ======================
-
-def chunk_text(text, size=1000, overlap=100):
-    chunks = []
-    for i in range(0, len(text), size - overlap):
-        chunks.append(text[i:i+size])
-    return chunks
-
-def clean_text(text):
-    text = re.sub(r'\s+', ' ', text)
-    return text.strip()
-
-def ingest_source(source, from_url=True):
-    """
-    ✅ Web URL ya PDF se text nikaalo aur Chroma me store karo
-    """
-    if from_url:
-        html = requests.get(source).text
-        soup = BeautifulSoup(html, "html.parser")
-        text = clean_text(soup.get_text())
-    else:
-        reader = PdfReader(source)
-        text = " ".join([page.extract_text() for page in reader.pages])
-
-    chunks = chunk_text(text)
-    embeddings = embedder.encode(chunks).tolist()
-    
-    for i, emb in enumerate(embeddings):
-        collection.add(ids=[f"doc_{i}"], embeddings=[emb], documents=[chunks[i]])
-    print(f"✅ Ingested {len(chunks)} chunks into Chroma DB")
-
-def rag_query(query):
-    """
-    ✅ Query kare aur best-matched chunks Gemini ko dekar answer banaye
-    """
-    q_emb = embedder.encode([query]).tolist()
-    results = collection.query(query_embeddings=q_emb, n_results=4)
-    context = " ".join(results['documents'][0])
-
-    prompt = f"""
-    You are an AI assistant. Use the context below to answer clearly:
-    Context: {context}
-    Question: {query}
-    Answer:
-    """
-    response = genai.GenerativeModel(MODEL).generate_content(prompt)
-    return textwrap.fill(response.text, width=100)
-
-# ======================
-# 🔹 Gradio UI
-# ======================
-def web_ingest_ui(url):
-    ingest_source(url, from_url=True)
-    return f"✅ Website data added: {url}"
-
-def pdf_ingest_ui(file):
-    ingest_source(file.name, from_url=False)
-    return f"✅ PDF data added: {file.name}"
-
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="teal")) as demo:
-    gr.Markdown("## 🤖 RAG App (Gemini + Local Embeddings)")
-    
-    with gr.Tab("🌐 Ingest Website"):
-        url_in = gr.Textbox(label="Enter Website URL")
-        url_btn = gr.Button("Ingest Website")
-        url_out = gr.Textbox(label="Status")
-        url_btn.click(fn=web_ingest_ui, inputs=url_in, outputs=url_out)
-
-    with gr.Tab("📄 Ingest PDF"):
-        pdf_in = gr.File(label="Upload PDF")
-        pdf_btn = gr.Button("Ingest PDF")
-        pdf_out = gr.Textbox(label="Status")
-        pdf_btn.click(fn=pdf_ingest_ui, inputs=pdf_in, outputs=pdf_out)
-
-    with gr.Tab("💬 Ask Questions"):
-        q_in = gr.Textbox(label="Ask anything from ingested sources")
-        q_btn = gr.Button("Ask")
-        q_out = gr.Markdown(label="Answer")  # ✅ Markdown shows multi-line output
-        q_btn.click(fn=rag_query, inputs=q_in, outputs=q_out)
-
-demo.launch()
+# =======================================
+# 📘 RAG App Pro – Gemini + Smart Embeddings (Multi-User Safe)
+# =======================================
+
+import os, re, shutil, textwrap, requests, uuid
+from bs4 import BeautifulSoup
+import google.generativeai as genai
+from sentence_transformers import SentenceTransformer
+import chromadb
+import gradio as gr
+from langchain_community.document_loaders import UnstructuredPDFLoader
+import camelot
+
+# ======================
+# 🔹 Gemini + Local Setup
+# ======================
+genai.configure(api_key="AIzaSyDr2X5N-hHt9EqUNy7JCm58aG1FpeGVpgs")  # 🔑 apni Gemini key
+MODEL = "gemini-2.5-flash"
+
+embedder = SentenceTransformer("all-MiniLM-L6-v2")
+chroma_client = chromadb.Client()
+
+# ======================
+# 🔹 Utils
+# ======================
+def clean_text(text):
+    return re.sub(r"\s+", " ", text).strip()
+
+def adaptive_chunk_text(text):
+    length = len(text)
+    if length < 3000:
+        size = 500
+    elif length < 10000:
+        size = 1000
+    else:
+        size = 1500
+    chunks = []
+    for i in range(0, len(text), size - 150):
+        chunks.append(text[i:i + size])
+    return chunks
+
+def extract_pdf_text(pdf_path):
+    """Smart PDF extractor (tables + text)"""
+    full_text = ""
+    try:
+        tables = camelot.read_pdf(pdf_path, pages="all")
+        for i, table in enumerate(tables):
+            full_text += f"\n\n[Table {i+1}]\n" + table.df.to_string(index=False)
+    except Exception:
+        pass
+
+    try:
+        loader = UnstructuredPDFLoader(pdf_path)
+        docs = loader.load()
+        full_text += "\n\n".join([doc.page_content for doc in docs])
+    except Exception as e:
+        full_text += f"\n\n[Error extracting text: {e}]"
+
+    return clean_text(full_text)
+
+# ======================
+# 🔹 Session Handling
+# ======================
+def create_user_collection():
+    """Each user/session gets unique collection"""
+    session_id = f"user_{str(uuid.uuid4())[:8]}"
+    collection = chroma_client.create_collection(name=session_id)
+    return session_id, collection
+
+def reset_collection(collection_name):
+    """Delete previous data for same user"""
+    try:
+        chroma_client.delete_collection(name=collection_name)
+    except Exception:
+        pass
+    return chroma_client.create_collection(name=collection_name)
+
+# ======================
+# 🔹 Ingestion Logic
+# ======================
+def ingest_source(source, from_url, collection_name):
+    # Delete previous user data
+    collection = reset_collection(collection_name)
+
+    if from_url:
+        html = requests.get(source, timeout=15).text
+        soup = BeautifulSoup(html, "html.parser")
+        text = clean_text(soup.get_text())
+    else:
+        text = extract_pdf_text(source)
+
+    if not text.strip():
+        return "⚠️ No readable text found (maybe image-only PDF)."
+
+    chunks = adaptive_chunk_text(text)
+    embeddings = embedder.encode(chunks).tolist()
+
+    for i, emb in enumerate(embeddings):
+        collection.add(ids=[f"{collection_name}_{i}"], embeddings=[emb], documents=[chunks[i]])
+
+    return f"✅ [{collection_name}] Ingested {len(chunks)} chunks successfully!"
+
+# ======================
+# 🔹 Query Logic
+# ======================
+def rag_query(query, collection_name):
+    try:
+        collection = chroma_client.get_collection(name=collection_name)
+        q_emb = embedder.encode([query]).tolist()
+        results = collection.query(query_embeddings=q_emb, n_results=4)
+        if not results["documents"]:
+            return "⚠️ No context found. Try ingesting data first."
+
+        context = "\n\n".join(results["documents"][0])
+        prompt = f"""
+You are a knowledgeable AI assistant.
+Use the context below to answer clearly and in multiple lines.
+
+Context:
+{context}
+
+Question: {query}
+Answer:
+"""
+        response = genai.GenerativeModel(MODEL).generate_content(prompt)
+        ans = response.text.replace(". ", ".\n")
+        return ans
+    except Exception as e:
+        return f"⚠️ Error: {e}"
+
+# ======================
+# 🔹 Gradio UI
+# ======================
+def start_new_session():
+    session_id, _ = create_user_collection()
+    return session_id
+
+session_id = start_new_session()
+
+def ingest_website(url):
+    return ingest_source(url, True, session_id)
+
+def ingest_pdf(file):
+    return ingest_source(file.name, False, session_id)
+
+def query_ask(q):
+    return rag_query(q, session_id)
+
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="emerald")) as demo:
+    gr.Markdown("# 🤖 Smart RAG App Pro (Gemini + Adaptive PDF + Multi-User Mode)")
+
+    gr.Markdown(f"🆕 **Private Session ID:** `{session_id}` – Your data is isolated and auto-clears on refresh.")
+
+    with gr.Tab("🌐 Ingest Website"):
+        url_in = gr.Textbox(label="Enter Website URL")
+        url_btn = gr.Button("Ingest Website")
+        url_out = gr.Textbox(label="Status")
+        url_btn.click(fn=ingest_website, inputs=url_in, outputs=url_out)
+
+    with gr.Tab("📄 Ingest PDF"):
+        pdf_in = gr.File(label="Upload PDF")
+        pdf_btn = gr.Button("Ingest PDF")
+        pdf_out = gr.Textbox(label="Status")
+        pdf_btn.click(fn=ingest_pdf, inputs=pdf_in, outputs=pdf_out)
+
+    with gr.Tab("💬 Ask Questions"):
+        q_in = gr.Textbox(label="Ask anything from ingested data")
+        q_btn = gr.Button("Ask Gemini")
+        q_out = gr.Markdown(label="Answer")
+        q_btn.click(fn=query_ask, inputs=q_in, outputs=q_out)
+
+demo.launch()