Update main.py

main.py

@@ -9,6 +9,7 @@ import tarfile
 import logging
 import hashlib
 from typing import Dict, Any, List, Tuple, Optional
+from concurrent.futures import ThreadPoolExecutor
 
 import numpy as np
 import faiss
@@ -22,18 +23,25 @@ import gradio as gr
 # =============================================================================
 # LOGGING
 # =============================================================================
-LOG = logging.getLogger("remote-indexer-space")
+LOG = logging.getLogger("remote-indexer-async")
 if not LOG.handlers:
     h = logging.StreamHandler()
     h.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(message)s"))
     LOG.addHandler(h)
 LOG.setLevel(logging.INFO)
 
+DBG = logging.getLogger("remote-indexer-async.debug")
+if not DBG.handlers:
+    hd = logging.StreamHandler()
+    hd.setFormatter(logging.Formatter("[DEBUG] %(asctime)s - %(message)s"))
+    DBG.addHandler(hd)
+DBG.setLevel(logging.DEBUG)
+
 # =============================================================================
 # CONFIG (via ENV)
 # =============================================================================
 PORT = int(os.getenv("PORT", "7860"))
-DATA_ROOT = os.getenv("DATA_ROOT", "/tmp/data")  # stockage interne du Space
+DATA_ROOT = os.getenv("DATA_ROOT", "/tmp/data")  # stockage interne du Space (volatile en Free)
 os.makedirs(DATA_ROOT, exist_ok=True)
 
 # Provider d'embeddings:
@@ -45,11 +53,13 @@ EMB_MODEL = os.getenv("EMB_MODEL", "sentence-transformers/paraphrase-multilingua
 EMB_BATCH = int(os.getenv("EMB_BATCH", "32"))
 EMB_DIM = int(os.getenv("EMB_DIM", "128"))  # utilisé pour dummy
 
+# Taille du pool de workers (asynchrone)
+MAX_WORKERS = int(os.getenv("MAX_WORKERS", "1"))
+
 # =============================================================================
-# CACHE DIRECTORIES (crucial pour éviter PermissionError: '/.cache')
+# CACHE DIRECTORIES (évite PermissionError: '/.cache')
 # =============================================================================
 def _setup_cache_dirs() -> Dict[str, str]:
-    # HOME peut être vide -> expanduser('~') => '/' -> '/.cache' -> Permission denied
     os.environ.setdefault("HOME", "/home/user")
 
     CACHE_ROOT = os.getenv("CACHE_ROOT", "/tmp/.cache").rstrip("/")
@@ -68,15 +78,12 @@ def _setup_cache_dirs() -> Dict[str, str]:
         except Exception as e:
             LOG.warning("Impossible de créer %s : %s", p, e)
 
-    # Variables standard HF/Transformers/Torch/ST
     os.environ["HF_HOME"] = paths["hf_home"]
     os.environ["HF_HUB_CACHE"] = paths["hf_hub"]
     os.environ["TRANSFORMERS_CACHE"] = paths["hf_tf"]
     os.environ["TORCH_HOME"] = paths["torch"]
     os.environ["SENTENCE_TRANSFORMERS_HOME"] = paths["st"]
-    os.environ["MPLCONFIGDIR"] = paths["mpl"]  # évite les warnings matplotlib
-
-    # Qualité de vie
+    os.environ["MPLCONFIGDIR"] = paths["mpl"]
     os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")
     os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
 
@@ -122,6 +129,7 @@ def _proj_dirs(project_id: str) -> Tuple[str, str, str]:
 def _add_msg(st: JobState, msg: str):
     st.messages.append(f"[{_now()}] {msg}")
     LOG.info("[%s] %s", st.job_id, msg)
+    DBG.debug("[%s] %s", st.job_id, msg)
 
 def _set_stage(st: JobState, stage: str):
     st.stage = stage
@@ -183,14 +191,6 @@ def _emb_st(texts: List[str]) -> np.ndarray:
     ).astype("float32")
     return vecs
 
-def _st_dim() -> int:
-    model = _get_st_model()
-    try:
-        return int(model.get_sentence_embedding_dimension())
-    except Exception:
-        v = model.encode(["dimension probe"], convert_to_numpy=True)
-        return int(v.shape[1])
-
 # ----------------------- PROVIDER: Transformers (HF) --------------------------
 def _get_hf_model():
     global _HF_TOKENIZER, _HF_MODEL
@@ -219,18 +219,11 @@ def _emb_hf(texts: List[str]) -> np.ndarray:
             enc = tok(batch, padding=True, truncation=True, return_tensors="pt")
             out = mod(**enc)
             last = out.last_hidden_state  # (b, t, h)
-            pooled = _mean_pool(last.numpy(), enc["attention_mask"].numpy())  # numpy
+            pooled = _mean_pool(last.numpy(), enc["attention_mask"].numpy())
             all_vecs.append(pooled.astype("float32"))
     vecs = np.concatenate(all_vecs, axis=0)
     return _l2_normalize(vecs)
 
-def _hf_dim() -> int:
-    try:
-        _, mod = _get_hf_model()
-        return int(getattr(mod.config, "hidden_size", 768))
-    except Exception:
-        return 768
-
 # ---------------------------- DATASET / FAISS ---------------------------------
 def _save_dataset(ds_dir: str, rows: List[Dict[str, Any]]):
     os.makedirs(ds_dir, exist_ok=True)
@@ -278,74 +271,44 @@ def _tar_dir_to_bytes(dir_path: str) -> bytes:
     return bio.read()
 
 # =============================================================================
-# FASTAPI
+# WORKER POOL (asynchrone)
 # =============================================================================
-fastapi_app = FastAPI(title="remote-indexer", version="2.1.0")
-fastapi_app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"],
-)
-
-class FileItem(BaseModel):
-    path: str
-    text: str
-
-class IndexRequest(BaseModel):
-    project_id: str
-    files: List[FileItem]
-    chunk_size: int = 200
-    overlap: int = 20
-    batch_size: int = 32
-    store_text: bool = True
-
-@fastapi_app.get("/health")
-def health():
-    info = {
-        "ok": True,
-        "service": "remote-indexer",
-        "provider": EMB_PROVIDER,
-        "model": EMB_MODEL if EMB_PROVIDER != "dummy" else None,
-        "cache_root": os.getenv("CACHE_ROOT", "/tmp/.cache"),
-    }
-    return info
-
-@fastapi_app.get("/")
-def root_redirect():
-    return {"ok": True, "service": "remote-indexer", "ui": "/ui"}
-
-@fastapi_app.post("/index")
-def index(req: IndexRequest):
-    job_id = hashlib.sha1(f"{req.project_id}{time.time()}".encode()).hexdigest()[:12]
-    st = JobState(job_id=job_id, project_id=req.project_id, stage="pending", messages=[])
-    JOBS[job_id] = st
-    _add_msg(st, f"Job {job_id} créé pour project {req.project_id}")
-    _add_msg(st, f"Index start project={req.project_id} files={len(req.files)} chunk_size={req.chunk_size} overlap={req.overlap} batch_size={req.batch_size} store_text={req.store_text} provider={EMB_PROVIDER} model={EMB_MODEL if EMB_PROVIDER!='dummy' else '-'}")
+EXECUTOR = ThreadPoolExecutor(max_workers=max(1, MAX_WORKERS))
+LOG.info("ThreadPoolExecutor initialisé : max_workers=%s", MAX_WORKERS)
+
+def _do_index_job(st: JobState, files: List[Dict[str, str]], chunk_size: int, overlap: int, batch_size: int, store_text: bool) -> None:
+    """
+    Tâche lourde lancée dans un worker thread.
+    Met à jour l'état 'st' tout au long du pipeline.
+    """
     try:
-        base, ds_dir, fx_dir = _proj_dirs(req.project_id)
+        base, ds_dir, fx_dir = _proj_dirs(st.project_id)
 
         # 1) Chunking
         _set_stage(st, "chunking")
         rows: List[Dict[str, Any]] = []
-        st.total_files = len(req.files)
-        for it in req.files:
-            txt = it.text or ""
-            chunks = _chunk_text(txt, size=req.chunk_size, overlap=req.overlap)
-            _add_msg(st, f"{it.path}: len(text)={len(txt)} chunks={len(chunks)}")
-            for ci, ck in enumerate(chunks):
-                rows.append({"path": it.path, "text": ck, "chunk_id": ci})
+        st.total_files = len(files)
+        for it in files:
+            path = (it.get("path") or "unknown").strip()
+            txt = it.get("text") or ""
+            chks = _chunk_text(txt, size=int(chunk_size), overlap=int(overlap))
+            _add_msg(st, f"{path}: len(text)={len(txt)} chunks={len(chks)}")
+            for ci, ck in enumerate(chks):
+                rows.append({"path": path, "text": ck, "chunk_id": ci})
         st.total_chunks = len(rows)
         _add_msg(st, f"Total chunks = {st.total_chunks}")
 
         # 2) Embedding
         _set_stage(st, "embedding")
+        texts = [r["text"] for r in rows]
         if EMB_PROVIDER == "dummy":
-            xb = _emb_dummy([r["text"] for r in rows], dim=EMB_DIM)
+            xb = _emb_dummy(texts, dim=EMB_DIM)
             dim = xb.shape[1]
         elif EMB_PROVIDER == "st":
-            xb = _emb_st([r["text"] for r in rows])
+            xb = _emb_st(texts)
             dim = xb.shape[1]
-        else:  # "hf"
-            xb = _emb_hf([r["text"] for r in rows])
+        else:
+            xb = _emb_hf(texts)
             dim = xb.shape[1]
 
         st.embedded = xb.shape[0]
@@ -367,17 +330,86 @@ def index(req: IndexRequest):
         _save_faiss(fx_dir, xb, meta=faiss_meta)
         st.indexed = int(xb.shape[0])
         _add_msg(st, f"FAISS écrit sur {os.path.join(fx_dir, 'emb.faiss')}")
-        _add_msg(st, f"OK — dataset+index prêts (projet={req.project_id})")
+        _add_msg(st, f"OK — dataset+index prêts (projet={st.project_id})")
 
         _set_stage(st, "done")
         st.finished_at = time.time()
-        return {"job_id": job_id}
     except Exception as e:
-        LOG.exception("index failed")
+        LOG.exception("Job %s failed", st.job_id)
         st.errors.append(str(e))
         _add_msg(st, f"❌ Exception: {e}")
         st.stage = "failed"
         st.finished_at = time.time()
+
+def _submit_job(project_id: str, files: List[Dict[str, str]], chunk_size: int, overlap: int, batch_size: int, store_text: bool) -> str:
+    job_id = hashlib.sha1(f"{project_id}{time.time()}".encode()).hexdigest()[:12]
+    st = JobState(job_id=job_id, project_id=project_id, stage="pending", messages=[])
+    JOBS[job_id] = st
+    _add_msg(st, f"Job {job_id} créé pour project {project_id}")
+    _add_msg(st, f"Index start project={project_id} files={len(files)} chunk_size={chunk_size} overlap={overlap} batch_size={batch_size} store_text={store_text} provider={EMB_PROVIDER} model={EMB_MODEL if EMB_PROVIDER!='dummy' else '-'}")
+
+    # Soumission au pool (retour immédiat)
+    EXECUTOR.submit(_do_index_job, st, files, chunk_size, overlap, batch_size, store_text)
+    _set_stage(st, "queued")
+    return job_id
+
+# =============================================================================
+# FASTAPI
+# =============================================================================
+fastapi_app = FastAPI(title="remote-indexer-async", version="3.0.0")
+fastapi_app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"],
+)
+
+class FileItem(BaseModel):
+    path: str
+    text: str
+
+class IndexRequest(BaseModel):
+    project_id: str
+    files: List[FileItem]
+    chunk_size: int = 200
+    overlap: int = 20
+    batch_size: int = 32
+    store_text: bool = True
+
+@fastapi_app.get("/health")
+def health():
+    info = {
+        "ok": True,
+        "service": "remote-indexer-async",
+        "provider": EMB_PROVIDER,
+        "model": EMB_MODEL if EMB_PROVIDER != "dummy" else None,
+        "cache_root": os.getenv("CACHE_ROOT", "/tmp/.cache"),
+        "workers": MAX_WORKERS,
+        "data_root": DATA_ROOT,
+    }
+    return info
+
+@fastapi_app.get("/")
+def root_redirect():
+    return {"ok": True, "service": "remote-indexer-async", "ui": "/ui"}
+
+@fastapi_app.post("/index")
+def index(req: IndexRequest):
+    """
+    ASYNCHRONE : retourne immédiatement un job_id.
+    Le traitement est effectué en arrière-plan par le pool de threads.
+    """
+    try:
+        files = [fi.model_dump() for fi in req.files]
+        job_id = _submit_job(
+            project_id=req.project_id,
+            files=files,
+            chunk_size=int(req.chunk_size),
+            overlap=int(req.overlap),
+            batch_size=int(req.batch_size),
+            store_text=bool(req.store_text),
+        )
+        return {"job_id": job_id}
+    except Exception as e:
+        LOG.exception("index failed (submit)")
         raise HTTPException(status_code=500, detail=str(e))
 
 @fastapi_app.get("/status/{job_id}")
@@ -395,9 +427,16 @@ class SearchRequest(BaseModel):
 @fastapi_app.post("/search")
 def search(req: SearchRequest):
     base, ds_dir, fx_dir = _proj_dirs(req.project_id)
+
+    # Si l'index n'existe pas encore, on répond 409 (conflit / pas prêt)
+    idx_path = os.path.join(fx_dir, "emb.faiss")
+    ds_path = os.path.join(ds_dir, "data.jsonl")
+    if not (os.path.isfile(idx_path) and os.path.isfile(ds_path)):
+        raise HTTPException(status_code=409, detail="Index non prêt (reviens plus tard)")
+
     rows = _load_dataset(ds_dir)
     if not rows:
-        raise HTTPException(status_code=404, detail="dataset introuvable (index pas encore construit ?)")
+        raise HTTPException(status_code=404, detail="dataset introuvable")
 
     # Embedding de la requête avec le MÊME provider
     if EMB_PROVIDER == "dummy":
@@ -443,7 +482,7 @@ def download_faiss(project_id: str):
     return StreamingResponse(io.BytesIO(buf), media_type="application/gzip", headers=headers)
 
 # =============================================================================
-# GRADIO UI (facultatif)
+# GRADIO UI (facultatif de test)
 # =============================================================================
 def _ui_index(project_id: str, sample_text: str):
     files = [{"path": "sample.txt", "text": sample_text}]
@@ -465,9 +504,9 @@ def _ui_search(project_id: str, query: str, k: int):
     except Exception as e:
         return f"Erreur search: {e}"
 
-with gr.Blocks(title="Remote Indexer (FAISS)", analytics_enabled=False) as ui:
-    gr.Markdown("## Remote Indexer — demo UI (API: `/index`, `/status/{job}`, `/search`, `/artifacts/...`).")
-    gr.Markdown(f"**Provider**: `{EMB_PROVIDER}` — **Model**: `{EMB_MODEL if EMB_PROVIDER!='dummy' else '-'}` — **Cache**: `{os.getenv('CACHE_ROOT', '/tmp/.cache')}`")
+with gr.Blocks(title="Remote Indexer (Async FAISS)", analytics_enabled=False) as ui:
+    gr.Markdown("## Remote Indexer — **Async** (API: `/index`, `/status/{job}`, `/search`, `/artifacts/...`).")
+    gr.Markdown(f"**Provider**: `{EMB_PROVIDER}` — **Model**: `{EMB_MODEL if EMB_PROVIDER!='dummy' else '-'}` — **Cache**: `{os.getenv('CACHE_ROOT', '/tmp/.cache')}` — **Workers**: `{MAX_WORKERS}`")
     with gr.Tab("Index"):
         pid = gr.Textbox(label="Project ID", value="DEEPWEB")
         sample = gr.Textbox(label="Texte d’exemple", value="Alpha bravo charlie delta echo foxtrot.", lines=4)
@@ -490,5 +529,5 @@ fastapi_app = gr.mount_gradio_app(fastapi_app, ui, path="/ui")
 # =============================================================================
 if __name__ == "__main__":
     import uvicorn
-    LOG.info("Démarrage Uvicorn sur 0.0.0.0:%s (UI_PATH=/ui)", PORT)
-    uvicorn.run(fastapi_app, host="0.0.0.0", port=PORT)
+    LOG.info("Démarrage Uvicorn sur 0.0.0.0:%s (UI_PATH=/ui) — async index", PORT)
+    uvicorn.run(fastapi_app, host="0.0.0.0", port=PORT)