Create app.py

app.py

@@ -0,0 +1,466 @@
+# ══════════════════════════════════════════════════════════════════
+#  CRITICAL SHIM — MUST BE THE FIRST EXECUTABLE CODE IN THIS FILE
+# ══════════════════════════════════════════════════════════════════
+import huggingface_hub
+
+if not hasattr(huggingface_hub, "HfFolder"):
+    class _HfFolder:
+        @staticmethod
+        def get_token():
+            try: return huggingface_hub.get_token()
+            except Exception: return None
+        @staticmethod
+        def save_token(token):
+            try: huggingface_hub.login(token)
+            except Exception: pass
+        @staticmethod
+        def delete_token():
+            try: huggingface_hub.logout()
+            except Exception: pass
+    huggingface_hub.HfFolder = _HfFolder
+
+# ── Now safe to import everything else ────────────────────────────
+import os
+import asyncio
+import html
+import time
+import uuid
+import threading
+from typing import Optional
+
+import duckdb
+import httpx
+import gradio as gr
+from fastapi import Request, HTTPException
+
+# ─────────────────────────────────────────────
+# CONFIG
+# ─────────────────────────────────────────────
+BOT_TOKEN      = os.environ.get("BOT_TOKEN", "").strip()
+HF_TOKEN       = os.environ.get("HF_TOKEN", "").strip()
+WEBHOOK_SECRET = os.environ.get("WEBHOOK_SECRET", "1221").strip()
+SPACE_URL      = os.environ.get("SPACE_URL", "").strip().rstrip("/")
+TELEGRAM_PROXY = os.environ.get("TELEGRAM_PROXY", "").strip().rstrip("/")
+
+if TELEGRAM_PROXY:
+    TELEGRAM_API = f"{TELEGRAM_PROXY}/bot{BOT_TOKEN}"
+    print(f"🔀 Using Telegram proxy: {TELEGRAM_PROXY}", flush=True)
+else:
+    TELEGRAM_API = f"https://api.telegram.org/bot{BOT_TOKEN}"
+
+# Startup sanity check
+print(f"🔑 BOT_TOKEN set: {bool(BOT_TOKEN)}", flush=True)
+print(f"🔑 SPACE_URL: {SPACE_URL or '(missing!)'}", flush=True)
+print(f"🔑 WEBHOOK_SECRET set: {bool(WEBHOOK_SECRET)}", flush=True)
+
+DATASET = "Kzr0xx/Icmr-and-hitek"
+PARQUET_URLS = [
+    f"https://huggingface.co/datasets/{DATASET}/resolve/main/part1.parquet",
+    f"https://huggingface.co/datasets/{DATASET}/resolve/main/part2a.parquet",
+    f"https://huggingface.co/datasets/{DATASET}/resolve/main/part2b_new.parquet",
+]
+
+PAGE_SIZE     = 4
+FETCH_LIMIT   = 40
+QUERY_TIMEOUT = 300
+USER_COOLDOWN = 8
+CACHE_TTL     = 3600
+
+ALLOWED_FIELDS = {
+    "name", "fathersName", "phoneNumber", "aadharNumber", "otherNumber",
+    "address", "district", "pincode", "state", "town", "source",
+}
+NUMERIC_FIELDS = {"phoneNumber", "aadharNumber", "otherNumber"}
+
+TRUNC = {
+    "name": 80, "fathersName": 80, "address": 200,
+    "phoneNumber": 40, "aadharNumber": 40, "otherNumber": 40,
+    "district": 40, "pincode": 15, "state": 40, "town": 40, "source": 30,
+}
+
+# ─────────────────────────────────────────────
+# DUCKDB
+# ─────────────────────────────────────────────
+_con = None
+_con_lock = threading.Lock()
+_con_ready = threading.Event()
+
+def get_con() -> duckdb.DuckDBPyConnection:
+    global _con
+    with _con_lock:
+        if _con is not None:
+            return _con
+        print("🔧 Initializing DuckDB...", flush=True)
+        c = duckdb.connect(":memory:")
+        c.execute("INSTALL httpfs; LOAD httpfs;")
+        if HF_TOKEN:
+            c.execute(f"CREATE SECRET hf (TYPE HUGGINGFACE, TOKEN '{HF_TOKEN}');")
+        c.execute("SET enable_object_cache=true;")
+        c.execute("SET preserve_insertion_order=false;")
+        c.execute("SET threads=4;")
+        c.execute("SET memory_limit='10GB';")
+        urls = ", ".join(f"'{u}'" for u in PARQUET_URLS)
+        try:
+            c.execute(f"SELECT COUNT(*) FROM read_parquet([{urls}], union_by_name=true) WHERE 1=0")
+            print("✅ DuckDB ready", flush=True)
+        except Exception as e:
+            print(f"⚠️ Warm-up warning: {e}", flush=True)
+        _con = c
+        _con_ready.set()
+        return _con
+
+def run_search(field: str, value: str, limit: int) -> list[dict]:
+    if field not in ALLOWED_FIELDS:
+        raise ValueError(f"Field not allowed: {field}")
+    con = get_con()
+    urls = ", ".join(f"'{u}'" for u in PARQUET_URLS)
+    if field in NUMERIC_FIELDS:
+        where, param = f"CAST({field} AS VARCHAR) = ?", value
+    else:
+        where, param = f"{field} ILIKE ?", f"%{value}%"
+    sql = f"""
+        SELECT name, fathersName, phoneNumber, aadharNumber, otherNumber,
+               address, district, pincode, state, town, source
+        FROM read_parquet([{urls}], union_by_name=true)
+        WHERE {where}
+        LIMIT {int(limit)}
+    """
+    cur = con.execute(sql, [param])
+    cols = [d[0] for d in cur.description]
+    return [dict(zip(cols, r)) for r in cur.fetchall()]
+
+# ─────────────────────────────────────────────
+# INPUT DETECTION & FORMATTING
+# ─────────────────────────────────────────────
+def detect_field(text: str) -> tuple[str, str]:
+    t = text.strip()
+    if t.isdigit() and len(t) == 10:
+        return "phoneNumber", t
+    if t.isdigit() and len(t) == 12:
+        return "aadharNumber", t
+    if t.isdigit() and len(t) >= 6:
+        return "phoneNumber", t
+    has_alpha = any(c.isalpha() for c in t)
+    has_digit = any(c.isdigit() for c in t)
+    if has_alpha and has_digit and len(t) >= 15:
+        return "address", t
+    if len(t) >= 40:
+        return "address", t
+    return "name", t
+
+def _tv(v, n: int) -> str:
+    if v is None:
+        return ""
+    s = str(v).strip()
+    if s.lower() in ("", "null", "none", "nan"):
+        return ""
+    return s if len(s) <= n else s[: n - 1] + "…"
+
+def format_result(idx: int, row: dict) -> str:
+    lines = [f"<b>── #{idx} ──</b>"]
+    order = ["name", "fathersName", "phoneNumber", "aadharNumber", "otherNumber",
+             "address", "district", "pincode", "state", "town", "source"]
+    for k in order:
+        v = _tv(row.get(k), TRUNC.get(k, 60))
+        if v:
+            lines.append(f"<b>{k}:</b> <code>{html.escape(v)}</code>")
+    return "\n".join(lines)
+
+def render_page(results, page, total_pages, query, elapsed):
+    start = page * PAGE_SIZE
+    chunk = results[start:start + PAGE_SIZE]
+    head = (
+        f"🔍 <b>Query:</b> <code>{html.escape(query)}</code>\n"
+        f"📄 <b>Page {page+1}/{total_pages}</b> · "
+        f"{len(results)} result(s) · ⏱️ {elapsed:.1f}s\n"
+    )
+    body = "\n\n".join(format_result(start + i + 1, r) for i, r in enumerate(chunk))
+    return f"{head}\n{body}"
+
+def page_keyboard(search_id, page, total_pages):
+    nav = []
+    if page > 0:
+        nav.append({"text": "◀ Prev", "callback_data": f"p:{search_id}:{page-1}"})
+    nav.append({"text": f"{page+1} / {total_pages}", "callback_data": "noop"})
+    if page < total_pages - 1:
+        nav.append({"text": "Next ▶", "callback_data": f"p:{search_id}:{page+1}"})
+    return {"inline_keyboard": [nav]}
+
+# ─────────────────────────────────────────────
+# TELEGRAM API
+# ─────────────────────────────────────────────
+async def tg(method: str, payload: dict, timeout: float = 30) -> dict:
+    try:
+        async with httpx.AsyncClient(timeout=timeout) as client:
+            r = await client.post(f"{TELEGRAM_API}/{method}", json=payload)
+            return r.json()
+    except Exception as e:
+        print(f"❌ tg({method}) error: {e!r}", flush=True)
+        return {"ok": False, "error": str(e)}
+
+async def send(chat_id, text, **kw):
+    res = await tg("sendMessage", {"chat_id": chat_id, "text": text, **kw})
+    if res.get("ok"):
+        return res["result"]["message_id"]
+    print(f"⚠️ sendMessage failed: {res}", flush=True)
+    return None
+
+async def edit(chat_id, message_id, text, **kw):
+    res = await tg("editMessageText",
+                   {"chat_id": chat_id, "message_id": message_id, "text": text, **kw})
+    if not res.get("ok") and "message is not modified" not in str(res):
+        print(f"⚠️ editMessageText failed: {res}", flush=True)
+
+# ─────────────────────────────────────────────
+# STATE
+# ─────────────────────────────────────────────
+SEARCH_CACHE: dict[str, dict] = {}
+USER_LAST: dict[int, float] = {}
+
+def cache_gc():
+    while True:
+        time.sleep(300)
+        now = time.time()
+        dead = [k for k, v in SEARCH_CACHE.items() if now - v["ts"] > CACHE_TTL]
+        for k in dead:
+            SEARCH_CACHE.pop(k, None)
+
+# ─────────────────────────────────────────────
+# HANDLERS
+# ─────────────────────────────────────────────
+WELCOME = (
+    "👋 <b>Phone Lookup Bot</b>\n\n"
+    "Send me any of:\n"
+    "• 10-digit phone number\n"
+    "• 12-digit Aadhaar number\n"
+    "• A person's name\n"
+    "• An address\n\n"
+    "I'll search the database and return matches with pagination."
+)
+
+async def handle_message(msg: dict):
+    chat_id = msg["chat"]["id"]
+    text = (msg.get("text") or "").strip()
+    print(f"💬 Message from {chat_id}: {text[:50]!r}", flush=True)
+    if not text:
+        return
+    if text in ("/start", "/help"):
+        await send(chat_id, WELCOME)
+        return
+    now = time.time()
+    if now - USER_LAST.get(chat_id, 0) < USER_COOLDOWN:
+        wait = USER_COOLDOWN - (now - USER_LAST.get(chat_id, 0))
+        await send(chat_id, f"⏳ Slow down — wait {wait:.0f}s before next search.")
+        return
+    USER_LAST[chat_id] = now
+    field, value = detect_field(text)
+    print(f"🔍 Detected field={field} value={value[:50]!r}", flush=True)
+    if len(value) < 3:
+        await send(chat_id, "❌ Query too short (min 3 chars).")
+        return
+    placeholder = await send(
+        chat_id,
+        f"🔍 <b>Searching…</b>\nField: <code>{field}</code>\n"
+        f"Value: <code>{html.escape(value[:100])}</code>\n\n"
+        f"<i>Cold-cache queries can take 30–120s.</i>",
+    )
+    if placeholder is None:
+        return
+    loop = asyncio.get_running_loop()
+    t0 = time.time()
+    try:
+        results = await asyncio.wait_for(
+            loop.run_in_executor(None, run_search, field, value, FETCH_LIMIT),
+            timeout=QUERY_TIMEOUT,
+        )
+    except asyncio.TimeoutError:
+        await edit(chat_id, placeholder, "⏱️ Search timed out.")
+        return
+    except Exception as e:
+        print(f"❌ search error: {e!r}", flush=True)
+        await edit(chat_id, placeholder,
+                   f"❌ Error: <code>{html.escape(str(e)[:200])}</code>")
+        return
+    elapsed = time.time() - t0
+    print(f"✅ Search returned {len(results)} rows in {elapsed:.1f}s", flush=True)
+    if not results:
+        await edit(chat_id, placeholder,
+                   f"❌ <b>No results</b>\nField: <code>{field}</code>\n"
+                   f"Value: <code>{html.escape(value[:100])}</code>\n⏱️ {elapsed:.1f}s")
+        return
+    search_id = uuid.uuid4().hex[:10]
+    total_pages = (len(results) + PAGE_SIZE - 1) // PAGE_SIZE
+    SEARCH_CACHE[search_id] = {
+        "ts": time.time(), "results": results,
+        "query": value, "elapsed": elapsed,
+    }
+    await edit(
+        chat_id, placeholder,
+        render_page(results, 0, total_pages, value, elapsed),
+        parse_mode="HTML",
+        reply_markup=page_keyboard(search_id, 0, total_pages),
+    )
+
+async def handle_callback(cb: dict):
+    data = cb.get("data", "")
+    cb_id = cb["id"]
+    msg = cb.get("message") or {}
+    chat_id = msg.get("chat", {}).get("id")
+    message_id = msg.get("message_id")
+    print(f"🔘 Callback: {data!r} from {chat_id}", flush=True)
+    if not chat_id or not message_id:
+        await tg("answerCallbackQuery", {"callback_query_id": cb_id})
+        return
+    if data == "noop" or not data.startswith("p:"):
+        await tg("answerCallbackQuery", {"callback_query_id": cb_id})
+        return
+    try:
+        _, search_id, page_str = data.split(":")
+        page = int(page_str)
+    except Exception:
+        await tg("answerCallbackQuery", {"callback_query_id": cb_id})
+        return
+    entry = SEARCH_CACHE.get(search_id)
+    if not entry:
+        await tg("answerCallbackQuery", {
+            "callback_query_id": cb_id,
+            "text": "Session expired — search again.",
+            "show_alert": True,
+        })
+        return
+    results = entry["results"]
+    total_pages = (len(results) + PAGE_SIZE - 1) // PAGE_SIZE
+    page = max(0, min(page, total_pages - 1))
+    await edit(
+        chat_id, message_id,
+        render_page(results, page, total_pages, entry["query"], entry["elapsed"]),
+        parse_mode="HTML",
+        reply_markup=page_keyboard(search_id, page, total_pages),
+    )
+    await tg("answerCallbackQuery", {"callback_query_id": cb_id})
+
+# ─────────────────────────────────────────────
+# WEBHOOK REGISTRATION
+# ─────────────────────────────────────────────
+async def set_webhook():
+    if not SPACE_URL:
+        print("⚠️ SPACE_URL not set — skipping setWebhook", flush=True)
+        return
+    url = f"{SPACE_URL}/webhook"
+    payload = {
+        "url": url,
+        "secret_token": WEBHOOK_SECRET,
+        "drop_pending_updates": True,
+        "allowed_updates": ["message", "callback_query"],
+    }
+    for attempt in range(1, 4):
+        try:
+            async with httpx.AsyncClient(timeout=30) as client:
+                r = await client.post(f"{TELEGRAM_API}/setWebhook", json=payload)
+                data = r.json()
+                print(f"🔗 setWebhook → {data}", flush=True)
+                if data.get("ok"):
+                    return
+        except Exception as e:
+            print(f"⚠️ setWebhook attempt {attempt}: {e}", flush=True)
+        await asyncio.sleep(3)
+
+# ─────────────────────────────────────────────
+# GRADIO UI + ZEROGPU REGISTRATION
+# ─────────────────────────────────────────────
+with gr.Blocks(title="Phone Lookup Bot") as demo:
+    gr.Markdown("# 📞 Phone Lookup Bot")
+    gr.Markdown(
+        "This Space hosts a Telegram bot that searches the "
+        "`Kzr0xx/Icmr-and-hitek` dataset (104 GB, streamed live).\n\n"
+        "**Open Telegram and message the bot to search.**"
+    )
+    gr.Markdown("### Status")
+    status_box = gr.JSON(label="Health")
+    demo.load(
+        fn=lambda: {
+            "status": "ok",
+            "duckdb_ready": _con_ready.is_set(),
+            "cached_searches": len(SEARCH_CACHE),
+        },
+        outputs=status_box,
+    )
+
+    # ZeroGPU registration
+    try:
+        import spaces
+
+        @spaces.GPU(duration=10)
+        def _zerogpu_register():
+            return "gpu-ready"
+
+        _zerogpu_btn = gr.Button("init-gpu", visible=False)
+        _zerogpu_btn.click(fn=_zerogpu_register, inputs=None, outputs=None)
+        print("✅ @spaces.GPU function registered", flush=True)
+    except Exception as e:
+        print(f"⚠️ Could not register @spaces.GPU: {e}", flush=True)
+
+# ─────────────────────────────────────────────
+# CUSTOM ROUTES ON GRADIO'S FASTAPI APP
+# ─────────────────────────────────────────────
+app = demo.app
+
+
+@app.get("/health")
+async def health():
+    return {
+        "status": "ok",
+        "duckdb_ready": _con_ready.is_set(),
+        "cached_searches": len(SEARCH_CACHE),
+        "telegram_proxy": bool(TELEGRAM_PROXY),
+        "bot_token_set": bool(BOT_TOKEN),
+    }
+
+
+@app.post("/webhook")
+async def webhook(request: Request):
+    print(f"📩 /webhook hit from {request.client}", flush=True)
+
+    if WEBHOOK_SECRET:
+        got = request.headers.get("x-telegram-bot-api-secret-token", "")
+        if got != WEBHOOK_SECRET:
+            print(f"❌ Bad secret: got={got!r} expected={WEBHOOK_SECRET!r}", flush=True)
+            raise HTTPException(status_code=403, detail="forbidden")
+
+    try:
+        update = await request.json()
+    except Exception as e:
+        print(f"⚠️ JSON parse error: {e}", flush=True)
+        return {"ok": True}
+
+    kind = next((k for k in ("message", "callback_query", "edited_message") if k in update), None)
+    print(f"📨 Update type: {kind}", flush=True)
+
+    if "message" in update:
+        asyncio.create_task(handle_message(update["message"]))
+    elif "callback_query" in update:
+        asyncio.create_task(handle_callback(update["callback_query"]))
+    return {"ok": True}
+
+# ─────────────────────────────────────────────
+# BACKGROUND THREADS + LAUNCH
+# ─────────────────────────────────────────────
+threading.Thread(target=cache_gc, daemon=True).start()
+threading.Thread(target=get_con, daemon=True).start()
+
+def _register_webhook_later():
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    loop.run_until_complete(set_webhook())
+
+threading.Thread(target=_register_webhook_later, daemon=True).start()
+
+# 🔥 CRITICAL: ssr_mode=False disables the Node SSR proxy.
+# Without this, external POSTs to /webhook hit the Node proxy on
+# port 7860 and never reach Python's FastAPI routes.
+demo.launch(
+    server_name="0.0.0.0",
+    server_port=7860,
+    share=False,
+    ssr_mode=False,
+)