Update app_webhook.py

app_webhook.py

@@ -0,0 +1,222 @@
+import os
+import glob
+import hashlib
+import numpy as np
+import pandas as pd
+from typing import List, Tuple
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse, PlainTextResponse
+from dotenv import load_dotenv
+from openai import OpenAI
+from docx import Document as DocxDocument
+from pypdf import PdfReader
+import faiss
+from telegram import Update
+from telegram.ext import Application, CommandHandler, MessageHandler, ContextTypes, AIORateLimiter, filters
+
+# ---------- Load environment ----------
+load_dotenv()
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
+PUBLIC_URL = os.getenv("PUBLIC_URL", "")  # e.g., https://username-space.hf.space
+OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4o-mini")
+EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "text-embedding-3-small")
+STRICT_DOC_MODE = os.getenv("STRICT_DOC_MODE", "true").lower() == "true"
+DOCS_DIR = os.getenv("DOCS_DIR", "wedding_docs")
+INDEX_PATH = os.getenv("INDEX_PATH", "wedding.index")
+META_CSV = os.getenv("META_CSV", "wedding_chunks.csv")
+
+client = OpenAI(api_key=OPENAI_API_KEY)
+
+# ---------- Doc loaders ----------
+def read_txt_md(path: str) -> str:
+    with open(path, "r", encoding="utf-8", errors="ignore") as f:
+        return f.read()
+
+def read_docx(path: str) -> str:
+    doc = DocxDocument(path)
+    return "\n".join(p.text for p in doc.paragraphs)
+
+def read_pdf(path: str) -> str:
+    reader = PdfReader(path)
+    return "\n".join((p.extract_text() or "") for p in reader.pages)
+
+def load_all_docs(folder: str) -> List[Tuple[str, str]]:
+    paths = []
+    for ext in ("*.md", "*.txt", "*.docx", "*.pdf"):
+        paths.extend(glob.glob(os.path.join(folder, ext)))
+    docs = []
+    for p in paths:
+        if p.endswith((".md", ".txt")):
+            text = read_txt_md(p)
+        elif p.endswith(".docx"):
+            text = read_docx(p)
+        elif p.endswith(".pdf"):
+            text = read_pdf(p)
+        else:
+            continue
+        if text.strip():
+            docs.append((p, text))
+    return docs
+
+# ---------- Index ----------
+def chunk_text(text: str, source: str, chunk_size: int = 900, overlap: int = 150):
+    words = text.split()
+    i = 0
+    while i < len(words):
+        chunk = " ".join(words[i:i + chunk_size])
+        yield {"source": source, "chunk": chunk, "hash": hashlib.md5((source + str(i)).encode()).hexdigest()}
+        i += (chunk_size - overlap)
+
+def embed_texts(texts: list[str]) -> np.ndarray:
+    resp = client.embeddings.create(model=EMBEDDING_MODEL, input=texts)
+    vecs = [d.embedding for d in resp.data]
+    return np.array(vecs).astype("float32")
+
+class RAGIndex:
+    def __init__(self):
+        self.index = None
+        self.df = None
+        self.dim = None
+
+    def load_or_build(self, force=False):
+        docs = load_all_docs(DOCS_DIR)
+        if not docs:
+            raise RuntimeError(f"No docs in {DOCS_DIR}/")
+        index_exists = os.path.exists(INDEX_PATH) and os.path.exists(META_CSV)
+        need = force or not index_exists
+
+        if index_exists and not need:
+            df = pd.read_csv(META_CSV)
+            vecs = np.load(INDEX_PATH)
+            faiss.normalize_L2(vecs)
+            idx = faiss.IndexFlatIP(vecs.shape[1])
+            idx.add(vecs)
+            self.index, self.df, self.dim = idx, df, vecs.shape[1]
+            return
+
+        # build new index
+        chunks = []
+        for p, t in docs:
+            for c in chunk_text(t, p):
+                chunks.append(c)
+        df = pd.DataFrame(chunks)
+        vecs = embed_texts(df["chunk"].tolist())
+        faiss.normalize_L2(vecs)
+        np.save(INDEX_PATH, vecs)
+        df.to_csv(META_CSV, index=False)
+        idx = faiss.IndexFlatIP(vecs.shape[1])
+        idx.add(vecs)
+        self.index, self.df, self.dim = idx, df, vecs.shape[1]
+
+    def retrieve(self, q: str, k=6):
+        qv = embed_texts([q])
+        faiss.normalize_L2(qv)
+        D, I = self.index.search(qv, k)
+        out = []
+        for s, i in zip(D[0], I[0]):
+            if int(i) >= 0:
+                row = self.df.iloc[int(i)]
+                out.append({"score": float(s), "source": row["source"], "chunk": row["chunk"]})
+        return out
+
+RAG = RAGIndex()
+
+SYSTEM_PROMPT = (
+    "You are a concise wedding assistant for Samson’s brother’s wedding. "
+    "Use ONLY the provided context. If missing, say so and suggest contacting Overall IC. "
+    "Keep answers under 150 words."
+)
+
+async def answer_with_rag(q: str) -> str:
+    ctx = RAG.retrieve(q, k=6)
+    blocks = []
+    for r in ctx:
+        t = r["chunk"]
+        if len(t) > 800:
+            t = t[:800] + "…"
+        blocks.append(f"[Source: {os.path.basename(r['source'])}]\n{t}")
+
+    context_text = "\n\n".join(blocks)
+    completion = client.chat.completions.create(
+        model=OPENAI_MODEL,
+        messages=[
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": f"Context from docs:\n\n{context_text}\n\nQuestion: {q}"}
+        ],
+        temperature=0.2,
+    )
+    a = completion.choices[0].message.content.strip()
+    if STRICT_DOC_MODE and not blocks:
+        return (
+            "I couldn’t find this in the docs. Please check the playbook or ask the Overall IC. "
+            "Admins can trigger /refresh after updating docs."
+        )
+    return a
+
+# ---------- Telegram App (Webhook) ----------
+app = FastAPI()
+telegram_app: Application | None = None
+
+async def start_telegram():
+    global telegram_app
+    if telegram_app is not None:
+        return telegram_app
+
+    RAG.load_or_build(force=False)
+
+    application = (
+        Application.builder()
+        .token(TELEGRAM_BOT_TOKEN)
+        .rate_limiter(AIORateLimiter())
+        .build()
+    )
+
+    async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
+        await update.message.reply_text(
+            "Hello! Ask me anything about roles, timings, addresses, and logistics.\n"
+            "Admins can use /refresh after updating docs."
+        )
+
+    async def help_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE):
+        await update.message.reply_text("Use /refresh or just ask your question in plain text.")
+
+    async def refresh_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE):
+        RAG.load_or_build(force=True)
+        await update.message.reply_text("✅ Refreshed from latest documents.")
+
+    async def on_message(update: Update, context: ContextTypes.DEFAULT_TYPE):
+        q = update.message.text or ""
+        a = await answer_with_rag(q)
+        await update.message.reply_text(a)
+
+    application.add_handler(CommandHandler("start", start))
+    application.add_handler(CommandHandler("help", help_cmd))
+    application.add_handler(CommandHandler("refresh", refresh_cmd))
+    application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, on_message))
+
+    # Configure webhook
+    if not PUBLIC_URL:
+        raise RuntimeError("PUBLIC_URL is not set. Add it in Space secrets.")
+    webhook_path = f"/webhook/{TELEGRAM_BOT_TOKEN}"
+    await application.bot.set_webhook(url=PUBLIC_URL.rstrip("/") + webhook_path)
+
+    telegram_app = application
+    return telegram_app
+
+@app.on_event("startup")
+async def _startup():
+    await start_telegram()
+
+@app.get("/")
+async def root():
+    return {"status": "ok", "msg": "Wedding bot running", "docs": os.listdir(DOCS_DIR)}
+
+@app.post("/webhook/{token}")
+async def telegram_webhook(token: str, request: Request):
+    if token != TELEGRAM_BOT_TOKEN:
+        return PlainTextResponse("forbidden", status_code=403)
+    data = await request.json()
+    update = Update.de_json(data, (await start_telegram()).bot)
+    await (await start_telegram()).process_update(update)
+    return JSONResponse({"ok": True})