Create app.py

app.py

@@ -0,0 +1,334 @@
+# ─────────────────────────────────────────────────────────────
+# 0. HfFolder shim (must be first executable code)
+# ─────────────────────────────────────────────────────────────
+import huggingface_hub
+if not hasattr(huggingface_hub, "HfFolder"):
+    class _HfFolder:
+        @staticmethod
+        def get_token(): return huggingface_hub.get_token()
+        @staticmethod
+        def save_token(token): huggingface_hub.login(token)
+        @staticmethod
+        def delete_token(): huggingface_hub.logout()
+    huggingface_hub.HfFolder = _HfFolder
+
+# ─────────────────────────────────────────────────────────────
+# 1. Imports
+# ─────────────────────────────────────────────────────────────
+import os, time, threading, asyncio, html, logging
+from typing import Optional
+import httpx
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+from starlette.routing import Route as StarletteRoute
+import gradio as gr
+import spaces
+
+# ─────────────────────────────────────────────────────────────
+# 2. Config
+# ─────────────────────────────────────────────────────────────
+BOT_TOKEN       = os.environ["BOT_TOKEN"]
+WEBHOOK_SECRET  = os.environ.get("WEBHOOK_SECRET", "1221")
+SPACE_URL       = os.environ["SPACE_URL"].rstrip("/")
+EXTERNAL_API    = os.environ["EXTERNAL_API_URL"].rstrip("/")
+EXTERNAL_KEY    = os.environ.get("EXTERNAL_API_KEY", "@kzr0x")
+TELEGRAM_PROXY  = os.environ.get("TELEGRAM_PROXY", "").rstrip("/")
+
+TG_API = (TELEGRAM_PROXY + "/bot" + BOT_TOKEN) if TELEGRAM_PROXY else f"https://api.telegram.org/bot{BOT_TOKEN}"
+
+RATE_LIMIT_SEC = 5
+PAGE_SIZE      = 4
+QUERY_TIMEOUT  = 30
+HTTP_TIMEOUT   = httpx.Timeout(30.0, connect=10.0)
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+log = logging.getLogger("bot")
+
+# ─────────────────────────────────────────────────────────────
+# 3. External API client
+# ─────────────────────────────────────────────────────────────
+_api_client: Optional[httpx.AsyncClient] = None
+
+async def _get_api_client() -> httpx.AsyncClient:
+    global _api_client
+    if _api_client is None or _api_client.is_closed:
+        _api_client = httpx.AsyncClient(timeout=HTTP_TIMEOUT, follow_redirects=True)
+    return _api_client
+
+async def search_phone(number: str) -> list[dict]:
+    """Call the external API with the required key parameter."""
+    client = await _get_api_client()
+    params = {
+        "key": EXTERNAL_KEY,       # <-- REQUIRED: fixes the 422 error
+        "number": number,          # <-- matches original app.py's parameter name
+        "limit": 10,
+    }
+    try:
+        r = await client.get(f"{EXTERNAL_API}/search", params=params)
+        log.info("External API %s → %s", r.url, r.status_code)
+        r.raise_for_status()
+        data = r.json()
+    except httpx.TimeoutException:
+        raise RuntimeError("External API timed out")
+    except httpx.HTTPStatusError as e:
+        raise RuntimeError(f"External API error {e.response.status_code}")
+    except Exception as e:
+        raise RuntimeError(f"External API call failed: {e}")
+
+    # Normalise response — the API may return various shapes
+    if isinstance(data, dict):
+        results = data.get("results") or data.get("data") or []
+        # If the API returns icmr_data / hitek_data separately, merge them
+        if not results:
+            results = (data.get("icmr_data") or []) + (data.get("hitek_data") or [])
+    elif isinstance(data, list):
+        results = data
+    else:
+        results = []
+
+    normalised = []
+    for row in results:
+        if not isinstance(row, dict):
+            continue
+        normalised.append({
+            "name":         row.get("name") or row.get("Name") or "",
+            "fathersName":  row.get("fathersName") or row.get("father_name") or "",
+            "phoneNumber":  row.get("phoneNumber") or row.get("phone") or number,
+            "aadharNumber": row.get("aadharNumber") or row.get("aadhaar") or "",
+            "otherNumber":  row.get("otherNumber") or row.get("other_number") or "",
+            "address":      row.get("address") or "",
+            "district":     row.get("district") or "",
+            "pincode":      row.get("pincode") or "",
+            "state":        row.get("state") or "",
+            "town":         row.get("town") or "",
+            "source":       row.get("source") or "",
+        })
+    return normalised
+
+# ─────────────────────────────────────────────────────────────
+# 4. Telegram helpers
+# ─────────────────────────────────────────────────────────────
+async def tg(method: str, **payload):
+    client = await _get_api_client()
+    url = f"{TG_API}/{method}"
+    r = await client.post(url, json=payload)
+    try:
+        r.raise_for_status()
+    except httpx.HTTPStatusError:
+        log.warning("Telegram %s failed: %s %s", method, r.status_code, r.text[:200])
+    return r.json() if r.headers.get("content-type", "").startswith("application/json") else {}
+
+async def send(chat_id: int, text: str, **kw):
+    kw.setdefault("parse_mode", "HTML")
+    kw.setdefault("disable_web_page_preview", True)
+    return await tg("sendMessage", chat_id=chat_id, text=text, **kw)
+
+async def edit(chat_id: int, message_id: int, text: str, **kw):
+    kw.setdefault("parse_mode", "HTML")
+    kw.setdefault("disable_web_page_preview", True)
+    return await tg("editMessageText", chat_id=chat_id, message_id=message_id, text=text, **kw)
+
+async def answer_callback(callback_id: str, **kw):
+    return await tg("answerCallbackQuery", callback_query_id=callback_id, **kw)
+
+# ─────────────────────────────────────────────────────────────
+# 5. Formatting
+# ─────────────────────────────────────────────────────────────
+def _esc(v) -> str:
+    return html.escape(str(v)) if v is not None else ""
+
+def render_record(rec: dict) -> str:
+    lines = [
+        f"👤 <b>Name:</b> {_esc(rec['name']) or '—'}",
+        f"👨 <b>Father:</b> {_esc(rec['fathersName']) or '—'}",
+        f"📱 <b>Phone:</b> {_esc(rec['phoneNumber'])}",
+        f"🆔 <b>Aadhaar:</b> {_esc(rec['aadharNumber']) or '—'}",
+        f"📞 <b>Other:</b> {_esc(rec['otherNumber']) or '—'}",
+        f"🏠 <b>Address:</b> {_esc(rec['address']) or '—'}",
+        f"🏙 <b>Town:</b> {_esc(rec['town']) or '—'}",
+        f"📍 <b>District:</b> {_esc(rec['district']) or '—'}",
+        f"📮 <b>Pincode:</b> {_esc(rec['pincode']) or '—'}",
+        f"🗺 <b>State:</b> {_esc(rec['state']) or '—'}",
+        f"🗂 <b>Source:</b> {_esc(rec['source']) or '—'}",
+    ]
+    return "\n".join(lines)
+
+def render_page(number: str, records: list[dict], page: int) -> str:
+    total = len(records)
+    start = page * PAGE_SIZE
+    chunk = records[start:start + PAGE_SIZE]
+    if not chunk:
+        return "No more results."
+    header = f"🔎 <b>Results for</b> <code>{_esc(number)}</code>  —  page {page + 1}/{(total - 1) // PAGE_SIZE + 1}\n\n"
+    body = "\n\n━━━━━━━━━━━━━━━━━━\n\n".join(render_record(r) for r in chunk)
+    return header + body
+
+def pagination_kb(number: str, page: int, total: int):
+    nav = []
+    if page > 0:
+        nav.append({"text": "◀ Prev", "callback_data": f"pg|{number}|{page - 1}"})
+    if (page + 1) * PAGE_SIZE < total:
+        nav.append({"text": "Next ▶", "callback_data": f"pg|{number}|{page + 1}"})
+    return {"inline_keyboard": [nav]} if nav else None
+
+# ─────────────────────────────────────────────────────────────
+# 6. State
+# ─────────────────────────────────────────────────────────────
+_last_search: dict[int, float] = {}
+_search_lock = threading.Lock()
+
+# ─────────────────────────────────────────────────────────────
+# 7. Handlers
+# ─────────────────────────────────────────────────────────────
+async def handle_message(update: dict):
+    msg = update.get("message") or update.get("edited_message")
+    if not msg:
+        return
+    chat_id = msg["chat"]["id"]
+    text = (msg.get("text") or "").strip()
+
+    if text.startswith("/"):
+        cmd = text.split()[0].lower()
+        if cmd == "/start":
+            await send(chat_id,
+                "👋 <b>Welcome!</b>\n\n"
+                "Send me a <b>10‑digit mobile number</b> and I'll look it up."
+            )
+            return
+        if cmd == "/status":
+            await send(chat_id,
+                "✅ <b>Bot alive</b>\n"
+                f"⏱ Rate limit: {RATE_LIMIT_SEC}s\n"
+                f"📄 Page size: {PAGE_SIZE}"
+            )
+            return
+        return
+
+    if not (text.isdigit() and len(text) == 10):
+        await send(chat_id, "❗ Please send a <b>10‑digit mobile number</b> only.")
+        return
+
+    now = time.time()
+    with _search_lock:
+        last = _last_search.get(chat_id, 0)
+        if now - last < RATE_LIMIT_SEC:
+            wait = int(RATE_LIMIT_SEC - (now - last)) + 1
+            await send(chat_id, f"⏳ Please wait <b>{wait}s</b> before searching again.")
+            return
+        _last_search[chat_id] = now
+
+    await send(chat_id, f"🔍 <b>Searching</b> <code>{_esc(text)}</code> …")
+
+    try:
+        records = await asyncio.wait_for(search_phone(text), timeout=QUERY_TIMEOUT)
+    except asyncio.TimeoutError:
+        await send(chat_id, "⏰ Search timed out. Try again in a moment.")
+        return
+    except RuntimeError as e:
+        await send(chat_id, f"⚠️ Search failed: <code>{_esc(str(e))}</code>")
+        return
+
+    if not records:
+        await send(chat_id, f"❌ No records found for <code>{_esc(text)}</code>.")
+        return
+
+    kb = pagination_kb(text, 0, len(records))
+    await send(chat_id, render_page(text, records, 0), reply_markup=kb)
+
+async def handle_callback(update: dict):
+    cb = update.get("callback_query")
+    if not cb:
+        return
+    data = cb.get("data") or ""
+    chat_id = cb["message"]["chat"]["id"]
+    message_id = cb["message"]["message_id"]
+    await answer_callback(cb["id"])
+
+    if not data.startswith("pg|"):
+        return
+    _, number, page_s = data.split("|", 2)
+    page = int(page_s)
+
+    try:
+        records = await asyncio.wait_for(search_phone(number), timeout=QUERY_TIMEOUT)
+    except Exception:
+        await edit(chat_id, message_id, "⚠️ Could not load that page. Try again.")
+        return
+
+    if not records:
+        await edit(chat_id, message_id, "❌ No records found.")
+        return
+
+    kb = pagination_kb(number, page, len(records))
+    await edit(chat_id, message_id, render_page(number, records, page), reply_markup=kb)
+
+# ─────────────────────────────────────────────────────────────
+# 8. FastAPI routes
+# ─────────────────────────────────────────────────────────────
+async def webhook_handler(request: Request):
+    if request.headers.get("x-telegram-bot-api-secret-token") != WEBHOOK_SECRET:
+        return JSONResponse({"ok": False}, status_code=403)
+    try:
+        update = await request.json()
+    except Exception:
+        return JSONResponse({"ok": True})
+
+    try:
+        if "message" in update or "edited_message" in update:
+            await handle_message(update)
+        elif "callback_query" in update:
+            await handle_callback(update)
+    except Exception:
+        log.exception("Handler error")
+    return JSONResponse({"ok": True})
+
+async def health_handler(request: Request):
+    return JSONResponse({"status": "ok", "time": time.time()})
+
+# ─────────────────────────────────────────────────────────────
+# 9. Gradio UI (ZeroGPU registration)
+# ─────────────────────────────────────────────────────────────
+@spaces.GPU(duration=10)
+def _zerogpu_register():
+    return "gpu-ready"
+
+with gr.Blocks(title="ICMR Lookup Bot") as demo:
+    gr.Markdown("## ICMR Lookup Bot\nTelegram webhook endpoint is live.")
+    _btn = gr.Button("init-gpu", visible=False)
+    _btn.click(fn=_zerogpu_register, inputs=None, outputs=None)
+
+# ─────────────────────────────────────────────────────────────
+# 10. Launch + inject routes
+# ─────────────────────────────────────────────────────────────
+demo.launch(
+    server_name="0.0.0.0",
+    server_port=7860,
+    share=False,
+    ssr_mode=False,
+    prevent_thread_lock=True,
+)
+
+live_app = demo.app
+live_app.router.routes.insert(0, StarletteRoute("/webhook", webhook_handler, methods=["POST"]))
+live_app.router.routes.insert(0, StarletteRoute("/health", health_handler, methods=["GET"]))
+live_app.build_middleware_stack()
+
+# ─────────────────────────────────────────────────────────────
+# 11. Set webhook
+# ─────────────────────────────────────────────────────────────
+async def _set_webhook():
+    await asyncio.sleep(3)
+    client = await _get_api_client()
+    try:
+        r = await client.post(f"{TG_API}/setWebhook", json={
+            "url": f"{SPACE_URL}/webhook",
+            "secret_token": WEBHOOK_SECRET,
+            "allowed_updates": ["message", "edited_message", "callback_query"],
+        })
+        log.info("setWebhook → %s", r.text)
+    except Exception:
+        log.exception("setWebhook failed")
+
+threading.Thread(target=lambda: asyncio.run(_set_webhook()), daemon=True).start()
+
+demo.block_thread()