Update app.py

app.py

@@ -3,6 +3,7 @@ from fastapi.middleware.cors import CORSMiddleware
 from datasets import load_dataset
 from huggingface_hub import HfApi, login
 import os
+import time
 
 os.environ["HF_HOME"] = "/tmp/hf"
 
@@ -11,7 +12,7 @@ HF_TOKEN = os.environ.get("HF_TOKEN")
 if HF_TOKEN:
     login(token=HF_TOKEN)
 
-app = FastAPI()
+app = FastAPI(title="Telegram UID Search API")
 
 app.add_middleware(
     CORSMiddleware,
@@ -21,62 +22,87 @@ app.add_middleware(
 )
 
 DATASET = "CodeXDevloper/MergedTgDataset"
+PART_SIZE = 100000  # har part mein 1 lakh rows
+MAX_PARTS = 50      # ek request mein max 50 parts
+TIME_LIMIT = 25     # 25 second ke andar ruk jao
 
-def get_files():
+def get_all_files():
     api = HfApi()
     files = api.list_repo_files(DATASET, repo_type="dataset")
-    return [f for f in files if f.endswith(".parquet")]
+    return sorted([f for f in files if f.endswith(".parquet")])
 
 @app.get("/")
 def root():
-    return {"status": "running", "usage": "/search?uid=XXX&offset=0"}
+    return {
+        "status": "running",
+        "usage": "/search?uid=385167100",
+        "note": "Poore dataset mein search (parts mein divided)"
+    }
 
 @app.get("/search")
-def search_uid(
-    uid: str = Query(...),
-    offset: int = Query(0, ge=0)
-):
+def search_uid(uid: str = Query(..., description="Telegram UID")):
+    """
+    Poore dataset mein UID dhoondho - parts mein divided, fast.
+    """
     try:
-        files = get_files()
+        files = get_all_files()
         
+        if not files:
+            return {"error": "Koi parquet file nahi mili"}
+        
+        # Saari files ko parts mein load karo
         ds = load_dataset(
             DATASET,
-            data_files={"train": files[:50]},  # 50 files per request
+            data_files={"train": files},
             split="train",
             streaming=True
         )
         
-        # Offset skip
-        iterator = iter(ds)
-        for _ in range(offset):
-            try:
-                next(iterator)
-            except StopIteration:
-                return {"found": False, "uid": uid, "message": "End", "next_offset": None}
-        
-        # 2 lakh rows scan
+        start_time = time.time()
         scanned = 0
-        max_scan = 200000
         
-        for row in iterator:
+        for row in ds:
             scanned += 1
+            
+            # UID match karo
             if uid == str(row.get("user_id", "")).strip():
                 return {
                     "found": True,
                     "uid": uid,
-                    "offset": offset,
                     "scanned": scanned,
-                    "data": row
+                    "time_taken": round(time.time() - start_time, 2),
+                    "data": {
+                        "user_id": row.get("user_id"),
+                        "username": row.get("username"),
+                        "first_name": row.get("first_name"),
+                        "last_name": row.get("last_name"),
+                        "phone": row.get("phone"),
+                        "email": row.get("email"),
+                        "status": row.get("status"),
+                        "linked_id": row.get("linked_id"),
+                        "linked_name": row.get("linked_name"),
+                        "linked_handle": row.get("linked_handle"),
+                    }
                 }
-            if scanned >= max_scan:
-                break
+            
+            # Time limit check (har 5000 rows)
+            if scanned % 5000 == 0:
+                if time.time() - start_time > TIME_LIMIT:
+                    return {
+                        "found": False,
+                        "uid": uid,
+                        "scanned": scanned,
+                        "time_taken": round(time.time() - start_time, 2),
+                        "message": f"{TIME_LIMIT} second mein nahi mila. Phir try karo."
+                    }
         
         return {
             "found": False,
             "uid": uid,
-            "offset": offset,
             "scanned": scanned,
-            "next_offset": offset + scanned
+            "time_taken": round(time.time() - start_time, 2),
+            "message": "Poore dataset mein nahi mila"
         }
+    
     except Exception as e:
-        return {"error": str(e)}
+        return {"error": str(e), "uid": uid}