Update app.py

app.py

@@ -1,6 +1,21 @@
+# ============================================================
+# HfFolder monkey-patch (MUST be before gradio import)
+# ============================================================
+import os
+try:
+    import huggingface_hub
+    if not hasattr(huggingface_hub, "HfFolder"):
+        class _HfFolder:
+            @staticmethod
+            def get_token():
+                return os.getenv("HF_TOKEN", "")
+        huggingface_hub.HfFolder = _HfFolder
+except ImportError:
+    pass
+# ============================================================
+
 import asyncio
 import json
-import os
 import threading
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any
@@ -16,21 +31,19 @@ HF_INDEX_BASE = os.environ.get(
     "PK_HF_INDEX_BASE",
     "https://huggingface.co/datasets/parthislive/pk-db/resolve/main/parts",
 ).rstrip("/")
-INDEX_SOURCE = os.environ.get("PK_INDEX_SOURCE", "remote").lower()
-PARALLELISM = int(os.environ.get("PK_PARALLEL", "2"))
+PARALLELISM = int(os.environ.get("PK_PARALLEL", "1"))
 THREADS_PER_CONN = int(os.environ.get("PK_THREADS_PER_CONN", "2"))
 DUPLICATE_CAP = 2
 
 SEARCH_FIELDS = ["NUMBER", "CNIC", "NAME", "ADDRESS"]
 NUMBER_FIELDS = ["NUMBER"]
 
-# Remote parquet index files (17 parts)
 REMOTE_INDEXES = {
     "main": [f"{HF_INDEX_BASE}/data_{i}.parquet" for i in range(17)],
 }
 
 # ── DuckDB Connection Pool ──────────────────────────────────────────────────
-_conns: list[duckdb.DuckDBPyConnection] = []
+_conns: list = []
 _conns_lock = threading.Lock()
 _thread_local = threading.local()
 pool = ThreadPoolExecutor(max_workers=PARALLELISM, thread_name_prefix="duck")
@@ -40,20 +53,26 @@ def _idx_ready(kind: str) -> bool:
     return kind in REMOTE_INDEXES
 
 
-def _new_conn() -> duckdb.DuckDBPyConnection:
+def _new_conn():
     con = duckdb.connect()
     con.execute("SET home_directory='/tmp'")
     con.execute("SET extension_directory='/tmp/duckdb_extensions'")
     con.execute("INSTALL parquet; LOAD parquet;")
     con.execute("INSTALL httpfs; LOAD httpfs;")
-    
-    # Create view from remote HF parts
+
+    con.execute("SET memory_limit='400MB'")
+    con.execute("SET enable_object_cache=true")
+    con.execute("SET http_keep_alive=true")
+    con.execute("SET http_timeout=60000")
+    con.execute("SET threads=2")
+
     for kind, urls in REMOTE_INDEXES.items():
         view = f"people_{kind}"
         lst = ", ".join(f"'{u}'" for u in urls)
-        con.execute(f"CREATE OR REPLACE VIEW {view} AS SELECT * FROM read_parquet([{lst}])")
-    
-    con.execute(f"SET threads = {THREADS_PER_CONN}")
+        con.execute(
+            f"CREATE OR REPLACE VIEW {view} AS "
+            f"SELECT NUMBER, CNIC, NAME, ADDRESS FROM read_parquet([{lst}])"
+        )
     return con
 
 
@@ -66,7 +85,7 @@ def _thread_id() -> int:
     return tid
 
 
-def _get_conn() -> duckdb.DuckDBPyConnection:
+def _get_conn():
     ident = _thread_id()
     with _conns_lock:
         while len(_conns) <= ident:
@@ -74,16 +93,14 @@ def _get_conn() -> duckdb.DuckDBPyConnection:
     return _conns[ident]
 
 
-# ── Dedup & Connected Records ───────────────────────────────────────────────
+# ── Dedup ───────────────────────────────────────────────────────────────────
 def _person_key(row: dict) -> tuple:
     cnic = (row.get("CNIC") or "").strip()
-    num = (row.get("NUMBER") or "").strip()
-    if cnic or num:
-        return (str(cnic), str(num))
-    return (row.get("NAME") or "").strip(), ""
+    num = str(row.get("NUMBER") or "").strip()
+    return (cnic, num)
 
 
-def _connected_numbers(row: dict) -> list[dict]:
+def _connected_numbers(row: dict) -> list:
     connected, seen = [], set()
     for field in NUMBER_FIELDS:
         raw = row.get(field)
@@ -97,9 +114,8 @@ def _connected_numbers(row: dict) -> list[dict]:
     return connected
 
 
-def _cap_duplicates(rows: list[dict]) -> list[dict]:
-    seen: dict[tuple, int] = {}
-    out = []
+def _cap_duplicates(rows: list) -> list:
+    seen, out = {}, []
     for r in rows:
         k = _person_key(r)
         n = seen.get(k, 0)
@@ -111,33 +127,40 @@ def _cap_duplicates(rows: list[dict]) -> list[dict]:
     return out
 
 
+# ── Fetch Without Pandas ────────────────────────────────────────────────────
+def _fetch_rows(con, sql):
+    cur = con.execute(sql)
+    cols = [d[0] for d in cur.description]
+    rows = cur.fetchall()
+    return [dict(zip(cols, r)) for r in rows]
+
+
 # ── Search Logic ────────────────────────────────────────────────────────────
 def _run_field_search(field: str, value: str, mode: str, limit: int) -> dict:
     if field not in SEARCH_FIELDS:
         raise ValueError(f"Unknown field: {field}")
-    
+
     v = str(value).replace("'", "''")
     view = "people_main"
+    cols = "NUMBER, CNIC, NAME, ADDRESS"
 
     if mode == "exact":
         if field == "NUMBER":
             try:
-                v = int(value)
-                sql = f"SELECT * FROM {view} WHERE NUMBER = {v} LIMIT {limit * DUPLICATE_CAP + 20}"
+                sql = f"SELECT {cols} FROM {view} WHERE NUMBER = {int(value)} LIMIT {limit * DUPLICATE_CAP + 20}"
             except:
                 return {"field": field, "value": value, "mode": mode, "count": 0, "results": []}
         else:
-            sql = f"SELECT * FROM {view} WHERE {field} = '{v}' LIMIT {limit * DUPLICATE_CAP + 20}"
+            sql = f"SELECT {cols} FROM {view} WHERE {field} = '{v}' LIMIT {limit * DUPLICATE_CAP + 20}"
     elif mode == "contains":
         v2 = v.replace("%", r"\%").replace("_", r"\_")
-        sql = f"SELECT * FROM {view} WHERE {field} ILIKE '%{v2}%' ESCAPE '\\' LIMIT {limit * DUPLICATE_CAP + 20}"
+        sql = f"SELECT {cols} FROM {view} WHERE {field} ILIKE '%{v2}%' ESCAPE '\\' LIMIT {limit * DUPLICATE_CAP + 20}"
     else:
         raise ValueError(f"Unknown mode: {mode}")
 
     con = _get_conn()
-    rows = con.execute(sql).fetchall()
-    cols = [d[0] for d in con.description]
-    results = _cap_duplicates([dict(zip(cols, r)) for r in rows])[:limit]
+    raw_rows = _fetch_rows(con, sql)
+    results = _cap_duplicates(raw_rows)[:limit]
     return {"field": field, "value": value, "mode": mode, "count": len(results), "results": results}
 
 
@@ -145,36 +168,27 @@ def _unified_search(q: str, limit: int = 10) -> dict:
     q = q.strip()
     if not q:
         return {"query": q, "searched_fields": [], "count": 0, "results": []}
-    
+
     is_num = q.isdigit() and len(q) >= 8
-    all_rows = []
-    searched = []
-    
+    all_rows, searched = [], []
+
     if is_num:
-        # NUMBER search
         r = _run_field_search("NUMBER", q, "exact", limit)
         all_rows.extend(r["results"])
         searched.append("NUMBER")
-    
-    # CNIC search (always try)
+
     if not all_rows:
         r = _run_field_search("CNIC", q, "exact", limit)
         all_rows.extend(r["results"])
         searched.append("CNIC")
-    
-    # NAME search (if not a number)
+
     if not all_rows and not is_num:
         r = _run_field_search("NAME", q, "contains", limit)
         all_rows.extend(r["results"])
         searched.append("NAME")
-    
+
     all_rows = _cap_duplicates(all_rows)[:limit]
-    return {
-        "query": q,
-        "searched_fields": searched,
-        "count": len(all_rows),
-        "results": all_rows,
-    }
+    return {"query": q, "searched_fields": searched, "count": len(all_rows), "results": all_rows}
 
 
 # ── FastAPI ─────────────────────────────────────────────────────────────────
@@ -182,30 +196,18 @@ fastapi_app = FastAPI(title="PK DB Search API")
 
 
 class BatchRequest(BaseModel):
-    queries: list[dict[str, Any]]
+    queries: list
     limit: int = 10
 
 
 @fastapi_app.get("/")
 def root():
-    return {
-        "app": "PK DB Search API",
-        "records": 218_697_206,
-        "indexes": {"main": _idx_ready("main")},
-        "index_source": INDEX_SOURCE,
-        "columns": SEARCH_FIELDS,
-        "docs": "/docs",
-        "developer": "@UnknownGuy9876 | channel @SGCodexs",
-    }
+    return {"app": "PK DB Search API", "records": 218_697_206, "columns": SEARCH_FIELDS}
 
 
 @fastapi_app.get("/health")
 def health():
-    return {
-        "status": "ok",
-        "indexes": {"main": _idx_ready("main")},
-        "index_source": INDEX_SOURCE
-    }
+    return {"status": "ok"}
 
 
 @fastapi_app.get("/search")
@@ -220,13 +222,11 @@ async def search(
     q_val = (q or number or "").strip()
     if not q_val:
         raise HTTPException(422, "Provide q or number")
-    
     loop = asyncio.get_running_loop()
     if field:
         data = await loop.run_in_executor(pool, _run_field_search, field, q_val, mode, limit)
     else:
         data = await loop.run_in_executor(pool, _unified_search, q_val, limit)
-    
     result = {"success": bool(data["count"]), **data, "number": q_val, "total": data["count"]}
     content = json.dumps(result, indent=2 if pretty else None, ensure_ascii=False)
     return Response(content=content, media_type="application/json")
@@ -260,31 +260,6 @@ async def api_address(address: str = Query(...)):
     return {"status": "success" if data["count"] else "not_found", **data}
 
 
-@fastapi_app.post("/search/parallel")
-async def search_parallel(req: BatchRequest):
-    if not req.queries:
-        raise HTTPException(400, "queries must not be empty")
-    if len(req.queries) > 50:
-        raise HTTPException(400, "max 50 queries per batch")
-    
-    loop = asyncio.get_running_loop()
-    tasks = [
-        loop.run_in_executor(
-            pool, _run_field_search,
-            item.get("field", "NUMBER"),
-            item.get("value", ""),
-            item.get("mode", "exact"),
-            int(item.get("limit", req.limit))
-        )
-        for item in req.queries
-    ]
-    results = await asyncio.gather(*tasks)
-    return Response(
-        content=json.dumps({"searches": len(req.queries), "results": list(results)}, indent=2, ensure_ascii=False),
-        media_type="application/json"
-    )
-
-
 # ── Pinger ──────────────────────────────────────────────────────────────────
 async def pinger():
     port = os.getenv("PORT", "7860")
@@ -293,10 +268,9 @@ async def pinger():
         while True:
             await asyncio.sleep(120)
             try:
-                resp = await client.get(url)
-                print(f"[Pinger] {resp.status_code}")
-            except Exception as e:
-                print(f"[Pinger] Error: {e}")
+                await client.get(url)
+            except Exception:
+                pass
 
 
 @fastapi_app.on_event("startup")
@@ -311,30 +285,22 @@ def format_result(row: dict) -> str:
         val = row.get(field, "")
         if val:
             lines.append(f"**{field}:** {val}")
-    cn = row.get("connected_numbers", [])
-    if cn:
-        nums = ", ".join(f"{c['field']}={c['value']}" for c in cn)
-        lines.append(f"**connected:** {nums}")
     return "\n\n".join(lines)
 
 
 def search_ui(query: str, limit: int) -> str:
     if not query or not query.strip():
         return "⚠️ Kuch toh search karo — NUMBER, CNIC, NAME ya ADDRESS daalo."
-    
     q = query.strip()
     try:
         data = _unified_search(q, int(limit))
     except Exception as e:
         return f"❌ Error: {str(e)}"
-    
     count = data["count"]
     results = data["results"]
     searched = ", ".join(data.get("searched_fields", []))
-    
     if not results:
         return f"🔍 **Query:** `{q}`\n**Searched:** {searched}\n\n❌ **No data found**."
-    
     header = f"🔍 **Query:** `{q}`  |  **Found:** {count} results  |  **Searched:** {searched}\n\n---\n\n"
     parts = []
     for i, row in enumerate(results, 1):
@@ -345,52 +311,24 @@ def search_ui(query: str, limit: int) -> str:
 def build_ui():
     with gr.Blocks(title="PK DB Search") as demo:
         gr.Markdown("# 🔍 PK DB Search")
-        gr.Markdown("Search **218M Pakistan records** — NUMBER, CNIC, NAME, ADDRESS")
-        
+        gr.Markdown("Search **218M Pakistan records**")
         with gr.Row():
             with gr.Column(scale=3):
-                query_input = gr.Textbox(
-                    label="Search Query",
-                    placeholder="NUMBER, CNIC, NAME ya ADDRESS daalo...",
-                    lines=1,
-                )
+                query_input = gr.Textbox(label="Search Query", placeholder="NUMBER, CNIC, NAME ya ADDRESS", lines=1)
             with gr.Column(scale=1):
-                limit_slider = gr.Slider(
-                    minimum=1, maximum=50, value=10, step=1,
-                    label="Max Results",
-                )
-        
+                limit_slider = gr.Slider(minimum=1, maximum=50, value=10, step=1, label="Max Results")
         search_btn = gr.Button("🔍 Search", variant="primary", size="lg")
         output = gr.Markdown(label="Results")
-        
         search_btn.click(fn=search_ui, inputs=[query_input, limit_slider], outputs=output)
         query_input.submit(fn=search_ui, inputs=[query_input, limit_slider], outputs=output)
-        
-        gr.Markdown("---")
-        with gr.Accordion("📡 API Info", open=False):
-            gr.Markdown("""
-**Endpoints:**
-- `GET /api/number?number=<num>` — Search by NUMBER
-- `GET /api/cnic?cnic=<cnic>` — Search by CNIC
-- `GET /api/name?name=<name>` — Search by NAME
-- `GET /api/address?address=<addr>` — Search by ADDRESS
-- `GET /search?q=<value>` — Unified search
-- `GET /health` — Health check
-- `GET /docs` — Swagger UI
-
-**Source:** [HF Dataset](https://huggingface.co/datasets/parthislive/pk-db)
-            """)
-        
-        gr.Markdown(
-            "---\n"
-            "<div style='text-align:center;color:#888;'>"
-            "👨‍💻 **Developer:** @UnknownGuy9876 | 📢 **Channel:** @SGCodexs"
-            "</div>"
-        )
-    
     return demo
 
 
-# ── Mount Gradio on FastAPI ─────────────────────────────────────────────────
 demo = build_ui()
-app = gr.mount_gradio_app(fastapi_app, demo, path="/ui")
+app = gr.mount_gradio_app(fastapi_app, demo, path="/ui")
+
+
+# ============================================================
+# HF Space Gradio SDK Launch (module level, NOT in __main__)
+# ============================================================
+demo.launch(server_name="0.0.0.0", server_port=7860)