Spaces:
Running
Running
| """FastAPI backend over DuckDB - ICMR + HITEK combined dataset (2.5B rows, 11 cols). | |
| Index-aware routing: | |
| - idx_phone.parquet (sorted by phoneNumber) -> fast exact phone/other lookups | |
| - idx_aadhar.parquet (sorted by aadharNumber) -> fast exact aadhar lookups | |
| - idx_name.parquet (sorted by name) -> fast name prefix/exact | |
| - Remote Parquet loading via Hugging Face Datasets with Token Authentication | |
| Dedup: max 2 rows per person. source column included (icmr / inddata "connected docs"). | |
| """ | |
| import asyncio | |
| import json | |
| import os | |
| import threading | |
| from concurrent.futures import ThreadPoolExecutor | |
| from typing import Any | |
| import duckdb | |
| from fastapi import FastAPI, HTTPException, Query, Response | |
| from pydantic import BaseModel | |
| BASE = os.path.dirname(os.path.abspath(__file__)) | |
| # โโ Updated Database, Token & Remote Dataset URL โโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| HF_TOKEN = os.environ.get("HF_TOKEN", "") | |
| HF_INDEX_BASE = os.environ.get( | |
| "ICMR_HF_INDEX_BASE", | |
| "https://huggingface.co/datasets/HiTeckGroup/HiTeckgroop/resolve/main", | |
| ).rstrip("/") | |
| INDEX_SOURCE = os.environ.get("ICMR_INDEX_SOURCE", "remote").lower() | |
| PARALLELISM = int(os.environ.get("ICMR_PARALLEL", "2")) | |
| THREADS_PER_CONN = int(os.environ.get("ICMR_THREADS_PER_CONN", "2")) | |
| DUPLICATE_CAP = 2 # max 2 copies of the same person in results | |
| SEARCH_FIELDS = [ | |
| "name", "fathersName", "phoneNumber", "aadharNumber", "otherNumber", | |
| "address", "district", "pincode", "state", "town", "source", | |
| ] | |
| NUMBER_FIELDS = ["phoneNumber", "aadharNumber", "otherNumber"] | |
| # Remote Parquet Files List | |
| REMOTE_INDEXES = { | |
| "phone": [f"{HF_INDEX_BASE}/idx_phone.{i}.parquet" for i in range(7)], | |
| "aadhar": [f"{HF_INDEX_BASE}/idx_aadhar.{i}.parquet" for i in range(7)], | |
| } | |
| # ๐ API KEY (secret เคธเฅ load) | |
| API_KEY = os.environ.get("API_KEY_1", "") | |
| VALID_KEYS = {} | |
| if API_KEY: | |
| VALID_KEYS[API_KEY] = {} | |
| app = FastAPI(title="ICMR + HITEK Search API", | |
| description="DuckDB-backed search over 2.5B records (11 cols, dedup max 2)") | |
| # ---- DuckDB connection pool (one connection per worker thread) ---- | |
| _conns: list[duckdb.DuckDBPyConnection] = [] | |
| _conns_lock = threading.Lock() | |
| _thread_local = threading.local() | |
| pool = ThreadPoolExecutor(max_workers=PARALLELISM, thread_name_prefix="duck") | |
| def _idx_ready(kind: str) -> bool: | |
| return kind in REMOTE_INDEXES | |
| def _new_conn() -> duckdb.DuckDBPyConnection: | |
| con = duckdb.connect() # in-memory | |
| con.execute("INSTALL parquet; LOAD parquet;") | |
| con.execute("INSTALL httpfs; LOAD httpfs;") | |
| # โโ PRIVATE DATASET AUTHENTICATION โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| if HF_TOKEN: | |
| con.execute(f"SET http_headers = {{'Authorization': 'Bearer {HF_TOKEN}'}};") | |
| # ๐ฅ๐ฅ๐ฅ SPEED OPTIMIZATIONS (เคจเคฏเคพ) ๐ฅ๐ฅ๐ฅ | |
| con.execute("SET enable_http_metadata_cache = true;") # metadata cache | |
| con.execute("SET http_keep_alive = true;") # connection reuse | |
| con.execute("SET http_timeout = 600000;") # 10 เคฎเคฟเคจเค timeout | |
| # Create sorted index views from remote Hugging Face dataset parts | |
| 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}") | |
| return con | |
| def _thread_id() -> int: | |
| tid = getattr(_thread_local, "id", None) | |
| if tid is None: | |
| with _conns_lock: | |
| tid = len(_conns) | |
| _thread_local.id = tid | |
| return tid | |
| def _get_conn() -> duckdb.DuckDBPyConnection: | |
| ident = _thread_id() | |
| with _conns_lock: | |
| while len(_conns) <= ident: | |
| _conns.append(_new_conn()) | |
| return _conns[ident] | |
| # ---- Dedup: max 2 copies per person ---- | |
| def _person_key(row: dict[str, Any]) -> tuple[str, ...]: | |
| ph = (row.get("phoneNumber") or "").strip() | |
| ad = (row.get("aadharNumber") or "").strip() | |
| if ph or ad: | |
| return (ph, ad) | |
| return (row.get("name") or "").strip(), (row.get("fathersName") or "").strip() | |
| def _cap_duplicates(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: | |
| """Keep at most DUPLICATE_CAP rows per person (same phone+aadhar identity).""" | |
| seen: dict[tuple[str, ...], int] = {} | |
| out: list[dict[str, Any]] = [] | |
| for r in rows: | |
| k = _person_key(r) | |
| n = seen.get(k, 0) | |
| if n < DUPLICATE_CAP: | |
| seen[k] = n + 1 | |
| out.append(r) | |
| return out | |
| def _source_clause(source: str | None) -> str: | |
| if source and source in ("icmr", "hitek", "inddata"): | |
| src = "hitek" if source == "hitek" else source | |
| return f" AND source = '{src}'" | |
| return "" | |
| def _run_field_search(field: str, value: str, mode: str, limit: int, | |
| source: str | None = None) -> dict[str, Any]: | |
| if field not in SEARCH_FIELDS: | |
| raise ValueError(f"Unknown field: {field}") | |
| v = value.replace("'", "''") | |
| # Use remote sorted index view when available | |
| if mode == "exact": | |
| if field == "phoneNumber" and _idx_ready("phone"): | |
| view = "people_phone" | |
| elif field == "aadharNumber" and _idx_ready("aadhar"): | |
| view = "people_aadhar" | |
| else: | |
| return {"field": field, "value": value, "mode": mode, "count": 0, "results": []} | |
| sql = (f"SELECT * FROM {view} WHERE {field} = '{v}'" | |
| f"{_source_clause(source)} LIMIT {limit * DUPLICATE_CAP + 20}") | |
| elif mode == "contains": | |
| v2 = v.replace("%", r"\%").replace("_", r"\_") | |
| sql = (f"SELECT * FROM people_phone WHERE {field} ILIKE '%{v2}%' ESCAPE '\\'" | |
| f"{_source_clause(source)} 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] | |
| return {"field": field, "value": value, "mode": mode, "count": len(results), | |
| "results": results} | |
| # ๐ KEY CHECK FUNCTION | |
| def _check_key(key: str | None) -> tuple[bool, dict | None]: | |
| if not key or key not in VALID_KEYS: | |
| return False, {"message": "Provide correct key", "contact": "@dmforbans"} | |
| return True, VALID_KEYS[key] | |
| async def _parallel(queries: list[tuple[str, str, str, int]]) -> list[dict[str, Any]]: | |
| loop = asyncio.get_running_loop() | |
| async def one(t: tuple[str, str, str, int]) -> dict[str, Any]: | |
| field, value, mode, limit = t | |
| return await loop.run_in_executor(pool, _run_field_search, field, value, mode, limit) | |
| return await asyncio.gather(*[one(t) for t in queries]) | |
| async def _unified_search(q: str, limit: int, source: str | None = None) -> dict[str, Any]: | |
| """Smart unified search: uses remote sorted indexes per field, merges, dedups.""" | |
| q = q.strip() | |
| is_num = q.isdigit() and len(q) >= 8 | |
| loop = asyncio.get_running_loop() | |
| if is_num: | |
| all_rows: list[dict[str, Any]] = [] | |
| searched: list[str] = [] | |
| if _idx_ready("phone"): | |
| r = await loop.run_in_executor(pool, _run_field_search, "phoneNumber", q, | |
| "exact", limit, source) | |
| all_rows.extend(r["results"]) | |
| searched.append("phoneNumber") | |
| if not all_rows and _idx_ready("aadhar"): | |
| r = await loop.run_in_executor(pool, _run_field_search, "aadharNumber", q, | |
| "exact", limit, source) | |
| all_rows.extend(r["results"]) | |
| searched.append("aadharNumber") | |
| all_rows = _cap_duplicates(all_rows)[:limit] | |
| return {"query": q, "searched_fields": searched or list(NUMBER_FIELDS), | |
| "dedup": DUPLICATE_CAP, "count": len(all_rows), "results": all_rows} | |
| else: | |
| return {"query": q, "searched_fields": [], "dedup": DUPLICATE_CAP, | |
| "count": 0, "results": []} | |
| def _pretty(data: dict[str, Any], pretty: bool) -> Response: | |
| if pretty: | |
| return Response(content=json.dumps(data, indent=2, ensure_ascii=False), | |
| media_type="application/json") | |
| return Response(content=json.dumps(data, ensure_ascii=False), | |
| media_type="application/json") | |
| # ---- API Routes ---- | |
| class BatchRequest(BaseModel): | |
| queries: list[dict[str, Any]] | |
| limit: int = 10 | |
| key: str | None = None # ๐ KEY | |
| def root(): | |
| return { | |
| "app": "ICMR + HITEK Search API", | |
| "status": "protected", | |
| "contact": "@dmforbans", | |
| } | |
| def health(): | |
| return { | |
| "status": "ok", | |
| "protected": True, | |
| } | |
| async def search( | |
| q: str = Query(..., description="Search term - phone number ya identification number"), | |
| key: str = Query(..., description="API Key"), # ๐ KEY REQUIRED | |
| field: str | None = Query(None, description=f"Restrict to one field: {SEARCH_FIELDS}"), | |
| mode: str = Query("exact", pattern="^(exact|contains)$"), | |
| limit: int = Query(10, ge=1, le=1000), | |
| source: str | None = Query(None, pattern="^(icmr|hitek|inddata)$", | |
| description="Filter by source: icmr, hitek (=inddata), inddata"), | |
| pretty: bool = Query(True, description="Pretty-print JSON"), | |
| ): | |
| """Unified search across remote HF Parquet files.""" | |
| # ๐ KEY CHECK | |
| valid, err = _check_key(key) | |
| if not valid: | |
| return Response( | |
| content=json.dumps(err, ensure_ascii=False), | |
| status_code=403, | |
| media_type="application/json" | |
| ) | |
| if field: | |
| try: | |
| data = _run_field_search(field, q, mode, limit, source) | |
| except ValueError as e: | |
| raise HTTPException(400, str(e)) | |
| data = {"query": q, "field": field, "mode": mode, "dedup": DUPLICATE_CAP, | |
| "count": data["count"], "results": data["results"]} | |
| else: | |
| data = await _unified_search(q, limit, source) | |
| return _pretty(data, pretty) | |
| async def search_parallel(req: BatchRequest): | |
| """Run up to 50 searches in parallel across DuckDB connection pool.""" | |
| # ๐ KEY CHECK | |
| valid, err = _check_key(req.key) | |
| if not valid: | |
| return Response( | |
| content=json.dumps(err, ensure_ascii=False), | |
| status_code=403, | |
| media_type="application/json" | |
| ) | |
| 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") | |
| results = await _parallel([(item.get("field", "phoneNumber"), item.get("value", ""), | |
| item.get("mode", "exact"), int(item.get("limit", req.limit))) | |
| for item in req.queries]) | |
| out = [{"error": str(r)} if isinstance(r, Exception) else r for r in results] | |
| return _pretty({"searches": len(req.queries), "parallelism": PARALLELISM, "results": out}, True) |