Create app.py

app.py

@@ -0,0 +1,211 @@
+import os
+import asyncio
+import json
+import hashlib
+import shutil
+from io import BytesIO
+from typing import List, Tuple
+
+import gradio as gr
+import numpy as np
+import faiss
+import requests
+from sentence_transformers import SentenceTransformer
+import fitz  # PyMuPDF
+
+# ---------------- Config ----------------
+OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
+OPENROUTER_MODEL = "nvidia/nemotron-nano-12b-v2-vl:free"
+EMBEDDING_MODEL_NAME = "all-MiniLM-L6-v2"
+CACHE_DIR = "./cache"
+SYSTEM_PROMPT = "You are a helpful assistant."
+
+os.makedirs(CACHE_DIR, exist_ok=True)
+
+embedder = SentenceTransformer(EMBEDDING_MODEL_NAME)
+
+DOCS: List[str] = []
+FILENAMES: List[str] = []
+EMBEDDINGS: np.ndarray = None
+FAISS_INDEX = None
+CURRENT_CACHE_KEY: str = ""
+
+
+# ---------------- Periodic cache cleanup ----------------
+async def clear_cache_every_5min():
+    while True:
+        await asyncio.sleep(300)
+        try:
+            if os.path.exists(CACHE_DIR):
+                shutil.rmtree(CACHE_DIR)
+            os.makedirs(CACHE_DIR, exist_ok=True)
+            print("🧹 Cache cleared.")
+        except Exception as e:
+            print(f"[Cache cleanup error] {e}")
+
+asyncio.get_event_loop().create_task(clear_cache_every_5min())
+
+
+# ---------------- PDF extraction ----------------
+def extract_text_from_pdf(file_bytes: bytes) -> str:
+    try:
+        doc = fitz.open(stream=file_bytes, filetype="pdf")
+        return "\n".join(page.get_text() for page in doc)
+    except Exception as e:
+        return f"[PDF extraction error] {e}"
+
+
+# ---------------- Cache + FAISS helpers ----------------
+def make_cache_key(files: List[Tuple[str, bytes]]) -> str:
+    h = hashlib.sha256()
+    for name, b in sorted(files, key=lambda x: x[0]):
+        h.update(name.encode())
+        h.update(str(len(b)).encode())
+        h.update(hashlib.sha256(b).digest())
+    return h.hexdigest()
+
+def cache_save(cache_key: str, embeddings: np.ndarray, filenames: List[str]):
+    np.savez_compressed(os.path.join(CACHE_DIR, f"{cache_key}.npz"),
+                        embeddings=embeddings, filenames=np.array(filenames))
+
+def cache_load(cache_key: str):
+    path = os.path.join(CACHE_DIR, f"{cache_key}.npz")
+    if not os.path.exists(path): return None
+    try:
+        data = np.load(path, allow_pickle=True)
+        return data["embeddings"], data["filenames"].tolist()
+    except:
+        return None
+
+def build_faiss(emb: np.ndarray):
+    global FAISS_INDEX
+    if emb is None or len(emb) == 0:
+        FAISS_INDEX = None
+        return None
+    emb = emb.astype("float32")
+    index = faiss.IndexFlatL2(emb.shape[1])
+    index.add(emb)
+    FAISS_INDEX = index
+    return index
+
+def search(query: str, k: int = 3):
+    if FAISS_INDEX is None:
+        return []
+    q_emb = embedder.encode([query], convert_to_numpy=True).astype("float32")
+    D, I = FAISS_INDEX.search(q_emb, k)
+    return [
+        {"index": int(i), "distance": float(d), "text": DOCS[i], "source": FILENAMES[i]}
+        for d, i in zip(D[0], I[0]) if i >= 0
+    ]
+
+
+# ---------------- OpenRouter API ----------------
+def call_openrouter(prompt: str):
+    if not OPENROUTER_API_KEY:
+        return "[OpenRouter error] Missing OPENROUTER_API_KEY."
+
+    url = "https://openrouter.ai/api/v1/chat/completions"
+    headers = {
+        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
+        "Content-Type": "application/json",
+    }
+
+    payload = {
+        "model": OPENROUTER_MODEL,
+        "messages": [
+            {"role": "system",
+             "content": SYSTEM_PROMPT + " Always respond in plain text. Avoid markdown."},
+            {"role": "user", "content": prompt},
+        ],
+    }
+
+    try:
+        r = requests.post(url, headers=headers, json=payload, timeout=60)
+        r.raise_for_status()
+        obj = r.json()
+
+        if "choices" in obj and obj["choices"]:
+            text = obj["choices"][0]["message"]["content"]
+            return text.strip().replace("```", "")
+        return "[Unexpected OpenRouter response]"
+    except Exception as e:
+        return f"[OpenRouter request error] {e}"
+
+
+# ---------------- PDF Upload & Index ----------------
+def upload_and_index(files):
+    global DOCS, FILENAMES, EMBEDDINGS, CURRENT_CACHE_KEY
+
+    if not files:
+        return "No PDF uploaded.", ""
+
+    processed = []
+    for f in files:
+        name = os.path.basename(f.name)
+        b = f.read()
+        processed.append((name, b))
+
+    preview = [{"name": n, "size": len(b)} for n, b in processed]
+
+    cache_key = make_cache_key(processed)
+    CURRENT_CACHE_KEY = cache_key
+
+    cached = cache_load(cache_key)
+    if cached:
+        EMBEDDINGS, FILENAMES = cached
+        EMBEDDINGS = np.array(EMBEDDINGS)
+        DOCS = [extract_text_from_pdf(b) for _, b in processed]
+        build_faiss(EMBEDDINGS)
+        return f"Loaded cached embeddings ({len(FILENAMES)} PDFs).", json.dumps(preview)
+
+    DOCS = [extract_text_from_pdf(b) for _, b in processed]
+    FILENAMES = [n for n, _ in processed]
+
+    EMBEDDINGS = embedder.encode(DOCS, convert_to_numpy=True).astype("float32")
+    cache_save(cache_key, EMBEDDINGS, FILENAMES)
+    build_faiss(EMBEDDINGS)
+
+    return f"Uploaded + indexed {len(DOCS)} PDFs.", json.dumps(preview)
+
+
+# ---------------- Question Answering ----------------
+def ask(question: str):
+    if not question:
+        return "Please enter a question."
+    if not DOCS:
+        return "No PDFs indexed."
+
+    results = search(question)
+
+    if not results:
+        return "No relevant text found."
+
+    context = "\n".join(
+        f"Source: {r['source']}\n\n{r['text'][:15000]}\n---\n"
+        for r in results
+    )
+
+    prompt = f"Use this context to answer briefly:\n\n{context}\nQuestion: {question}\nAnswer:"
+    return call_openrouter(prompt)
+
+
+# ---------------- Gradio UI ----------------
+with gr.Blocks(title="PDF RAG Bot") as demo:
+    gr.Markdown("# 📄 PDF-Only RAG Bot\nUpload PDFs → Ask Questions → AI Answers from PDF content.")
+
+    file_input = gr.File(label="Upload PDF files", file_count="multiple", file_types=[".pdf"])
+    upload_btn = gr.Button("Upload & Index")
+    status = gr.Textbox(label="Status", interactive=False)
+    preview = gr.Textbox(label="Upload preview (JSON)", interactive=False)
+
+    upload_btn.click(upload_and_index, inputs=[file_input], outputs=[status, preview])
+
+    gr.Markdown("### Ask a Question")
+    q = gr.Textbox(label="Your question", lines=3)
+    ask_btn = gr.Button("Ask PDF Bot")
+    answer = gr.Textbox(label="Answer", lines=15)
+
+    ask_btn.click(ask, inputs=[q], outputs=[answer])
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)