Update app.py

app.py

@@ -6,7 +6,7 @@ import os
 
 os.environ["HF_HOME"] = "/tmp/hf"
 
-app = FastAPI(title="Telegram UID Search API")
+app = FastAPI(title="Telegram UID Multi-Search API")
 
 app.add_middleware(
     CORSMiddleware,
@@ -17,99 +17,125 @@ app.add_middleware(
 
 DATASET = "CodeXDevloper/MergedTgDataset"
 
-def get_all_parquet_files():
-    """Saari parquet files nikalo - saare folders se"""
+# Folders ki priority - pehle V1, phir V2, phir baaki
+FOLDER_PRIORITY = [
+    "TG_DATA_PARTS",       # V1
+    "Telegram_V2",         # V2
+    None,                  # None = baaki sab (root files)
+]
+
+def get_files_for_folder(folder):
+    """Kisi folder ki saari parquet files nikalo"""
     api = HfApi()
     files = api.list_repo_files(DATASET, repo_type="dataset")
-    parquet_files = [f for f in files if f.endswith(".parquet")]
-    return parquet_files
+    
+    if folder:
+        return [f for f in files if f.startswith(folder + "/") and f.endswith(".parquet")]
+    else:
+        # Root ki files (jo kisi folder mein nahi hain)
+        return [f for f in files if "/" not in f and f.endswith(".parquet")]
 
 @app.get("/")
 def root():
     return {
         "status": "running",
         "dataset": DATASET,
-        "usage": "/search?uid=385167100"
+        "usage": "/search?uid=385167100",
+        "search_order": ["TG_DATA_PARTS (V1)", "Telegram_V2", "Root files"]
     }
 
 @app.get("/search")
 def search_uid(
-    uid: str = Query(..., description="Telegram UID"),
-    offset: int = Query(0, ge=0, description="Kahan se start karna hai")
+    uid: str = Query(..., description="Telegram UID")
 ):
     """
-    POORE dataset mein UID dhoondho - saare folders se
+    AUTOMATIC MULTI-FOLDER SEARCH:
+    1. Pehle V1 (TG_DATA_PARTS) mein dhoondho
+    2. Nahi mila toh V2 (Telegram_V2) mein
+    3. Nahi mila toh root files mein
     """
     try:
-        # Saari parquet files nikalo (saare folders)
-        parquet_files = get_all_parquet_files()
-        
-        if not parquet_files:
-            return {"error": "Koi parquet file nahi mili"}
+        results = []  # Har folder ka result track karo
         
-        # Saari files ko load karo (streaming mode)
-        ds = load_dataset(
-            DATASET,
-            data_files=parquet_files,
-            split="train",
-            streaming=True
-        )
-        
-        # Offset tak skip karo
-        iterator = iter(ds)
-        skipped = 0
-        for _ in range(offset):
+        for folder in FOLDER_PRIORITY:
+            folder_name = folder if folder else "Root"
+            print(f"🔍 Searching in: {folder_name}")
+            
+            # Folder ki files nikalo
+            files = get_files_for_folder(folder)
+            
+            if not files:
+                results.append({
+                    "folder": folder_name,
+                    "status": "no_files",
+                    "scanned": 0
+                })
+                continue
+            
+            # Dataset load karo (streaming)
             try:
-                next(iterator)
-                skipped += 1
-            except StopIteration:
-                return {
-                    "found": False,
-                    "uid": uid,
-                    "offset": offset,
-                    "scanned": 0,
-                    "message": "End of dataset",
-                    "next_offset": None
-                }
-        
-        # Ab 5 lakh rows scan karo
-        scanned = 0
-        max_scan = 500000
-        
-        for row in iterator:
-            scanned += 1
-            if uid == str(row.get("user_id", "")).strip():
-                return {
-                    "found": True,
-                    "uid": uid,
-                    "offset": offset,
-                    "scanned": scanned,
-                    "total_scanned": offset + scanned,
-                    "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"),
+                ds = load_dataset(
+                    DATASET,
+                    data_files={"train": files},
+                    split="train",
+                    streaming=True
+                )
+            except Exception as e:
+                results.append({
+                    "folder": folder_name,
+                    "status": "load_error",
+                    "error": str(e)
+                })
+                continue
+            
+            # Is folder mein dhoondho
+            scanned = 0
+            max_scan = 200000  # har folder mein 2 lakh rows
+            
+            for row in ds:
+                scanned += 1
+                if uid == str(row.get("user_id", "")).strip():
+                    # ✅ MIL GAYA
+                    return {
+                        "found": True,
+                        "uid": uid,
+                        "found_in": folder_name,
+                        "scanned_in_folder": scanned,
+                        "search_log": results + [{
+                            "folder": folder_name,
+                            "status": "FOUND",
+                            "scanned": scanned
+                        }],
+                        "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
             
-            if scanned >= max_scan:
-                break
+            # Is folder mein nahi mila
+            results.append({
+                "folder": folder_name,
+                "status": "not_found",
+                "scanned": scanned
+            })
         
+        # Kahin nahi mila
         return {
             "found": False,
             "uid": uid,
-            "offset": offset,
-            "scanned": scanned,
-            "next_offset": offset + scanned,
-            "total_files": len(parquet_files),
-            "message": f"Is chunk mein nahi mila. Ab offset={offset + scanned} try karo."
+            "search_log": results,
+            "message": "Poore dataset mein nahi mila"
         }
     
     except Exception as e: