Update main.py

main.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 """
 Remote Indexer (HF Space) — Qdrant + embeddings (HF ou dummy)
-Version: worker-thread (pas d'asyncio.create_task dans l'UI), robust logging.
+Version: worker-thread + ensure_collection_hard (retries + verify) + count tolérant.
 
 Endpoints:
 - GET  /            → redirige vers UI_PATH (défaut: /ui)
@@ -14,6 +14,7 @@ Endpoints:
 - GET  /status/{job_id}
 - GET  /collections/{project_id}/count
 - POST /query
+- POST /collections/{project_id}/ensure?dim=XXX  (debug)
 
 ENV:
 - QDRANT_URL, QDRANT_API_KEY      (requis pour upsert)
@@ -153,38 +154,65 @@ def chunk_text(text: str, chunk_size: int, overlap: int) -> List[Tuple[int, int,
 # ------------------------------------------------------------------------------
 # Qdrant
 # ------------------------------------------------------------------------------
-async def ensure_collection(client: httpx.AsyncClient, coll: str, vector_size: int) -> None:
+async def ensure_collection_hard(client: httpx.AsyncClient, coll: str, vector_size: int, job: Optional[JobState] = None) -> None:
+    """
+    Crée la collection si absente, ou la recrée si dimension ≠, puis
+    vérifie qu'elle existe bien (poll GET avec retries).
+    """
+    def _j(msg: str): 
+        if job: job.log(msg)
+        else: LOG.debug(msg)
+
     url = f"{QDRANT_URL}/collections/{coll}"
     r = await client.get(url, headers={"api-key": QDRANT_API_KEY}, timeout=20)
     recreate = False
     if r.status_code == 200:
         data = r.json()
         existing_size = data.get("result", {}).get("vectors", {}).get("size")
+        _j(f"GET collection '{coll}' → 200 dim={existing_size}")
         if existing_size and int(existing_size) != int(vector_size):
-            LOG.warning(f"Collection {coll} dim={existing_size} ≠ attendu {vector_size} → recréation")
-            await client.delete(url, headers={"api-key": QDRANT_API_KEY}, timeout=20)
+            _j(f"Dimension ≠ ({existing_size} ≠ {vector_size}) → suppression + recréation")
+            rd = await client.delete(url, headers={"api-key": QDRANT_API_KEY}, timeout=20)
+            _j(f"DELETE collection '{coll}' → {rd.status_code}")
             recreate = True
-        else:
-            LOG.debug(f"Collection {coll} existante (dim={existing_size})")
+    elif r.status_code == 404:
+        _j(f"GET collection '{coll}' → 404 (à créer)")
+    else:
+        _j(f"GET collection '{coll}' → {r.status_code} {r.text}")
+
     if r.status_code != 200 or recreate:
         body = {"vectors": {"size": vector_size, "distance": "Cosine"}}
         r2 = await client.put(url, headers={"api-key": QDRANT_API_KEY}, json=body, timeout=30)
+        _j(f"PUT create collection '{coll}' dim={vector_size} → {r2.status_code}")
         if r2.status_code not in (200, 201):
             raise HTTPException(status_code=500, detail=f"Qdrant PUT collection a échoué: {r2.text}")
 
-async def qdrant_upsert(client: httpx.AsyncClient, coll: str, points: List[Dict[str, Any]]) -> int:
+    # Poll jusqu'à visibilité
+    for i in range(10):
+        rg = await client.get(url, headers={"api-key": QDRANT_API_KEY}, timeout=20)
+        if rg.status_code == 200:
+            _j(f"Collection '{coll}' disponible (try={i+1})")
+            return
+        await asyncio.sleep(0.3)
+    raise HTTPException(status_code=500, detail=f"Collection '{coll}' non visible après création.")
+
+async def qdrant_upsert(client: httpx.AsyncClient, coll: str, points: List[Dict[str, Any]], job: Optional[JobState] = None) -> int:
     if not points:
         return 0
     url = f"{QDRANT_URL}/collections/{coll}/points?wait=true"
     body = {"points": points}
     r = await client.put(url, headers={"api-key": QDRANT_API_KEY}, json=body, timeout=60)
     if r.status_code not in (200, 202):
+        if job: job.log(f"Upsert → {r.status_code}: {r.text[:200]}")
         raise HTTPException(status_code=500, detail=f"Qdrant upsert échoué: {r.text}")
     return len(points)
 
 async def qdrant_count(client: httpx.AsyncClient, coll: str) -> int:
     url = f"{QDRANT_URL}/collections/{coll}/points/count"
     r = await client.post(url, headers={"api-key": QDRANT_API_KEY}, json={"exact": True}, timeout=20)
+    if r.status_code == 404 and "doesn't exist" in (r.text or ""):
+        # Tolérant: collection absente → 0 (utile pour l'UI)
+        return 0
     if r.status_code != 200:
         raise HTTPException(status_code=500, detail=f"Qdrant count échoué: {r.text}")
     return int(r.json().get("result", {}).get("count", 0))
@@ -293,9 +321,9 @@ async def run_index_job(job: JobState, req: IndexRequest) -> None:
             vec_dim = len(warmup_vec)
             job.log(f"Warmup embeddings dim={vec_dim}")
 
-            # Collection
+            # Collection (hard ensure)
             coll = f"{COLLECTION_PREFIX}{req.project_id}"
-            await ensure_collection(client, coll, vector_size=vec_dim)
+            await ensure_collection_hard(client, coll, vector_size=vec_dim, job=job)
             job.log(f"Collection prête: {coll} (dim={vec_dim})")
 
             # Upsert
@@ -306,7 +334,7 @@ async def run_index_job(job: JobState, req: IndexRequest) -> None:
                 nonlocal batch_points
                 if not batch_points:
                     return 0
-                added = await qdrant_upsert(client, coll, batch_points)
+                added = await qdrant_upsert(client, coll, batch_points, job=job)
                 job.upserted += added
                 job.log(f"+{added} points upsert (total={job.upserted})")
                 batch_points = []
@@ -433,8 +461,12 @@ async def coll_count(project_id: str):
         raise HTTPException(status_code=400, detail="QDRANT_URL / QDRANT_API_KEY requis")
     coll = f"{COLLECTION_PREFIX}{project_id}"
     async with httpx.AsyncClient() as client:
-        cnt = await qdrant_count(client, coll)
-    return {"project_id": project_id, "collection": coll, "count": cnt}
+        try:
+            cnt = await qdrant_count(client, coll)
+            return {"project_id": project_id, "collection": coll, "count": cnt}
+        except HTTPException as he:
+            # On remonte le 500 (autre qu'un 404 not found géré dans qdrant_count)
+            raise he
 
 @fastapi_app.post("/query")
 async def query(req: QueryRequest):
@@ -446,8 +478,18 @@ async def query(req: QueryRequest):
         data = await qdrant_search(client, coll, vec, limit=req.top_k)
     return data
 
+@fastapi_app.post("/collections/{project_id}/ensure")
+async def http_ensure(project_id: str, dim: int = Query(..., ge=16, le=8192)):
+    """Endpoint debug pour forcer la création d'une collection à une dimension donnée."""
+    if not QDRANT_URL or not QDRANT_API_KEY:
+        raise HTTPException(status_code=400, detail="QDRANT_URL / QDRANT_API_KEY requis")
+    coll = f"{COLLECTION_PREFIX}{project_id}"
+    async with httpx.AsyncClient() as client:
+        await ensure_collection_hard(client, coll, vector_size=dim, job=None)
+    return {"ok": True, "collection": coll, "dim": dim}
+
 # ------------------------------------------------------------------------------
-# Gradio UI (avec auto-refresh)
+# Gradio UI (avec auto-refresh + bouton Ensure debug)
 # ------------------------------------------------------------------------------
 def _default_two_docs() -> List[Dict[str, str]]:
     a = "Alpha bravo charlie delta echo foxtrot golf hotel india. " * 3
@@ -497,8 +539,11 @@ async def ui_status(job_id: str):
 
 async def ui_count(project: str):
     try:
-        resp = await coll_count(project)
-        return f"📊 Count — collection={resp['collection']} → {resp['count']} points"
+        # Count tolérant (0 si collection absente)
+        async with httpx.AsyncClient() as client:
+            coll = f"{COLLECTION_PREFIX}{project}"
+            cnt = await qdrant_count(client, coll)
+            return f"📊 Count — collection={coll} → {cnt} points"
     except Exception as e:
         LOG.exception("count UI error")
         return f"❌ Count erreur: {e}"
@@ -522,6 +567,14 @@ async def ui_query(project: str, text: str, topk: int):
         LOG.exception("query UI error")
         return f"❌ Query erreur: {e}"
 
+async def ui_ensure(project: str, dim: int):
+    try:
+        resp = await http_ensure(project, dim)
+        return f"🛠️ Ensure — collection={resp['collection']} dim={resp['dim']} OK"
+    except Exception as e:
+        LOG.exception("ensure UI error")
+        return f"❌ Ensure erreur: {e}"
+
 with gr.Blocks(title="Remote Indexer — Tests sans console", analytics_enabled=False) as ui:
     gr.Markdown("## 🔬 Remote Indexer — Tests sans console\n"
                 "Wipe → Index 2 docs → Status → Count → Query\n"
@@ -547,6 +600,10 @@ with gr.Blocks(title="Remote Indexer — Tests sans console", analytics_enabled=
         status_btn = gr.Button("📡 Status (refresh)")
         auto_chk = gr.Checkbox(False, label="⏱️ Auto-refresh status (2 s)")
 
+    with gr.Row():
+        ensure_dim = gr.Slider(16, 2048, value=128, step=16, label="ensure dim (debug)")
+        ensure_btn = gr.Button("🛠️ Ensure collection (debug)")
+
     with gr.Row():
         query_tb = gr.Textbox(label="Query text", value="alpha bravo")
         topk = gr.Slider(1, 20, value=5, step=1, label="top_k")
@@ -562,6 +619,8 @@ with gr.Blocks(title="Remote Indexer — Tests sans console", analytics_enabled=
     timer.tick(ui_status, inputs=[jobid_tb], outputs=[out_log])
     auto_chk.change(lambda x: gr.update(active=x), inputs=auto_chk, outputs=timer)
 
+    ensure_btn.click(ui_ensure, inputs=[project_tb, ensure_dim], outputs=[out_log])
+
 # Monte l'UI Gradio
 app = gr.mount_gradio_app(fastapi_app, ui, path=UI_PATH)