Update main.py

main.py

@@ -1,6 +1,15 @@
 # -*- coding: utf-8 -*-
-from __future__ import annotations
+"""
+Version optimisée du module FAISS :
+- Réduction de la dimension des vecteurs (EMB_DIM, configurable)
+- Index quantisé **IVF‑PQ** (faible empreinte disque)
+- Chargement *on‑disk* (mmap) pour limiter la RAM
+- Option `store_text` : ne pas persister le texte brut dans le dataset
+- Compression gzip des artefacts exportés
+- Paramètres contrôlables via variables d’environnement
+"""
 
+from __future__ import annotations
 import os
 import io
 import json
@@ -8,8 +17,7 @@ import time
 import tarfile
 import logging
 import hashlib
-from typing import Dict, Any, List, Tuple, Optional
-from concurrent.futures import ThreadPoolExecutor
+from typing import List, Dict, Any, Tuple, Optional
 
 import numpy as np
 import faiss
@@ -18,151 +26,132 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel
 
-import gradio as gr
-
-# =============================================================================
-# LOGGING
-# =============================================================================
-LOG = logging.getLogger("remote-indexer-async")
+# --------------------------------------------------------------------------- #
+#   CONFIGURATION (variables d’environnement – modifiable à la volée)
+# --------------------------------------------------------------------------- #
+EMB_PROVIDER = os.getenv("EMB_PROVIDER", "dummy").strip().lower()
+EMB_MODEL   = os.getenv("EMB_MODEL", "sentence-transformers/all-mpnet-base-v2").strip()
+EMB_BATCH   = int(os.getenv("EMB_BATCH", "32"))
+EMB_DIM     = int(os.getenv("EMB_DIM", "64"))          # ← dimension réduite (ex. 64)
+
+# FAISS quantisation
+FAISS_TYPE  = os.getenv("FAISS_TYPE", "IVF_PQ").upper()   # FLAT ou IVF_PQ
+FAISS_NLIST = int(os.getenv("FAISS_NLIST", "100"))       # nb de centroides (IVF)
+FAISS_M     = int(os.getenv("FAISS_M", "8"))             # sous‑vecteurs (PQ)
+FAISS_NBITS = int(os.getenv("FAISS_NBITS", "8"))         # bits / sous‑vecteur
+
+# Stockage du texte brut dans le dataset ? (False → économise disque)
+STORE_TEXT  = os.getenv("STORE_TEXT", "false").lower() in ("1", "true", "yes")
+
+# --------------------------------------------------------------------------- #
+#   LOGGING
+# --------------------------------------------------------------------------- #
+LOG = logging.getLogger("appli_v1")
 if not LOG.handlers:
     h = logging.StreamHandler()
-    h.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(message)s"))
+    h.setFormatter(logging.Formatter("[%(levelname)s] %(asctime)s - %(message)s", "%H:%M:%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 (volatile en Free)
-os.makedirs(DATA_ROOT, exist_ok=True)
-
-# Provider d'embeddings:
-#   - "dummy" : vecteurs aléatoires déterministes (très rapide)
-#   - "st"    : Sentence-Transformers (CPU-friendly)
-#   - "hf"    : Transformers pur (AutoModel/AutoTokenizer)
-EMB_PROVIDER = os.getenv("EMB_PROVIDER", "dummy").strip().lower()
-EMB_MODEL = os.getenv("EMB_MODEL", "sentence-transformers/all-mpnet-base-v2").strip()
-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 (évite PermissionError: '/.cache')
-# =============================================================================
-def _setup_cache_dirs() -> Dict[str, str]:
-    os.environ.setdefault("HOME", "/home/user")
-
-    CACHE_ROOT = os.getenv("CACHE_ROOT", "/tmp/.cache").rstrip("/")
-    paths = {
-        "root": CACHE_ROOT,
-        "hf_home": f"{CACHE_ROOT}/huggingface",
-        "hf_hub": f"{CACHE_ROOT}/huggingface/hub",
-        "hf_tf": f"{CACHE_ROOT}/huggingface/transformers",
-        "torch": f"{CACHE_ROOT}/torch",
-        "st": f"{CACHE_ROOT}/sentence-transformers",
-        "mpl": f"{CACHE_ROOT}/matplotlib",
-    }
-    for p in paths.values():
-        try:
-            os.makedirs(p, exist_ok=True)
-        except Exception as e:
-            LOG.warning("Impossible de créer %s : %s", p, e)
-
-    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"]
-    os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")
-    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
-
-    LOG.info("Caches configurés: %s", json.dumps(paths, indent=2))
-    return paths
-
-CACHE_PATHS = _setup_cache_dirs()
-
-# Cache global lazy (pour les modèles)
-_ST_MODEL = None
-_HF_TOKENIZER = None
-_HF_MODEL = None
-
-# =============================================================================
-# JOB STATE
-# =============================================================================
-class JobState(BaseModel):
-    job_id: str
-    project_id: str
-    stage: str = "pending"          # pending -> chunking -> embedding -> indexing -> done/failed
-    total_files: int = 0
-    total_chunks: int = 0
-    embedded: int = 0
-    indexed: int = 0
-    errors: List[str] = []
-    messages: List[str] = []
-    started_at: float = time.time()
-    finished_at: Optional[float] = None
+# --------------------------------------------------------------------------- #
+#   UTILITAIRES
+# --------------------------------------------------------------------------- #
+def list_repo_files(repo_dir: str, top_k: int = 500) -> List[str]:
+    """
+    Retourne la liste des fichiers texte du dépôt, en respectant .gitignore
+    (via Git si disponible, sinon fallback os.walk).
+    """
+    if not os.path.isdir(repo_dir):
+        return []
 
-JOBS: Dict[str, JobState] = {}
+    files: List[str] = []
+    try:
+        from git import Repo
+        repo = Repo(repo_dir)
+
+        # fichiers trackés
+        tracked = repo.git.ls_files().splitlines()
+        files.extend(tracked)
+
+        # fichiers non‑trackés mais non ignorés
+        untracked = repo.git.ls_files(others=True, exclude_standard=True).splitlines()
+        files.extend(untracked)
+
+        # filtrage simple
+        files = [
+            f for f in files
+            if not f.startswith('.git/') and not any(p.startswith('.') for p in f.split(os.sep))
+        ]
+        files = sorted(set(files))[:top_k]
+    except Exception as e:
+        LOG.debug("Git indisponible / pas un dépôt → fallback os.walk : %s", e)
+        for root, _, names in os.walk(repo_dir):
+            for name in sorted(names):
+                if name.startswith('.'):
+                    continue
+                rel = os.path.relpath(os.path.join(root, name), repo_dir)
+                if rel.startswith('.git') or any(p.startswith('.') for p in rel.split(os.sep)):
+                    continue
+                files.append(rel)
+                if len(files) >= top_k:
+                    break
+            if len(files) >= top_k:
+                break
+        files = sorted(set(files))
+
+    return files
+
+
+def read_file_safe(file_path: str) -> str:
+    """Lit un fichier en UTF‑8, ignore les erreurs."""
+    try:
+        with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
+            return f.read()
+    except Exception as e:
+        LOG.error("Erreur lecture %s : %s", file_path, e)
+        return f"# Erreur lecture : {e}"
 
-def _now() -> str:
-    return time.strftime("%H:%M:%S")
 
-def _proj_dirs(project_id: str) -> Tuple[str, str, str]:
-    base = os.path.join(DATA_ROOT, project_id)
-    ds_dir = os.path.join(base, "dataset")
-    fx_dir = os.path.join(base, "faiss")
-    os.makedirs(ds_dir, exist_ok=True)
-    os.makedirs(fx_dir, exist_ok=True)
-    return base, ds_dir, fx_dir
+def write_file_safe(file_path: str, content: str) -> str:
+    """Écrit un fichier, crée les dossiers parents si besoin."""
+    try:
+        os.makedirs(os.path.dirname(file_path), exist_ok=True)
+        with open(file_path, "w", encoding="utf-8") as f:
+            f.write(content)
+        return f"✅ Fichier sauvegardé : {os.path.basename(file_path)}"
+    except Exception as e:
+        LOG.error("Erreur écriture %s : %s", file_path, e)
+        return f"❌ Erreur sauvegarde : {e}"
+
+
+# --------------------------------------------------------------------------- #
+#   FAKE / DUMMY FAISS (pour compatibilité)
+# --------------------------------------------------------------------------- #
+class DummyFAISS:
+    """Classe factice – aucune fonctionnalité réelle."""
+    pass
+
+
+def create_faiss_index(*_, **__) -> DummyFAISS:
+    LOG.warning("FAISS désactivé – utilisation du client distant")
+    return DummyFAISS()
+
+
+def search_faiss_index(*_, **__) -> List[Any]:
+    LOG.warning("FAISS désactivé – utilisation du client distant")
+    return []
+
+
+# --------------------------------------------------------------------------- #
+#   EMBEDDING PROVIDERS
+# --------------------------------------------------------------------------- #
+_ST_MODEL: Optional[Any] = None
+_HF_TOKENIZER: Optional[Any] = None
+_HF_MODEL: Optional[Any] = None
+
 
-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
-    _add_msg(st, f"stage={stage}")
-
-# =============================================================================
-# UTILS
-# =============================================================================
-def _chunk_text(text: str, size: int = 200, overlap: int = 20) -> List[str]:
-    text = (text or "").replace("\r\n", "\n")
-    tokens = list(text)
-    if size <= 0:
-        return [text] if text else []
-    if overlap < 0:
-        overlap = 0
-    chunks = []
-    i = 0
-    while i < len(tokens):
-        j = min(i + size, len(tokens))
-        chunk = "".join(tokens[i:j]).strip()
-        if chunk:
-            chunks.append(chunk)
-        if j == len(tokens):
-            break
-        i = j - overlap if (j - overlap) > i else j
-    return chunks
-
-def _l2_normalize(x: np.ndarray) -> np.ndarray:
-    n = np.linalg.norm(x, axis=1, keepdims=True) + 1e-12
-    return x / n
-
-# ----------------------- PROVIDER: DUMMY --------------------------------------
 def _emb_dummy(texts: List[str], dim: int = EMB_DIM) -> np.ndarray:
+    """Vecteurs aléatoires déterministes (SHA‑1 → seed)."""
     vecs = np.zeros((len(texts), dim), dtype="float32")
     for i, t in enumerate(texts):
         h = hashlib.sha1((t or "").encode("utf-8")).digest()
@@ -171,15 +160,16 @@ def _emb_dummy(texts: List[str], dim: int = EMB_DIM) -> np.ndarray:
         vecs[i] = v / (np.linalg.norm(v) + 1e-9)
     return vecs
 
-# ----------------- PROVIDER: Sentence-Transformers ----------------------------
+
 def _get_st_model():
     global _ST_MODEL
     if _ST_MODEL is None:
         from sentence_transformers import SentenceTransformer
-        _ST_MODEL = SentenceTransformer(EMB_MODEL, cache_folder=CACHE_PATHS["st"])
-        LOG.info("[st] modèle chargé: %s (cache=%s)", EMB_MODEL, CACHE_PATHS["st"])
+        _ST_MODEL = SentenceTransformer(EMB_MODEL, cache_folder=os.getenv("HF_HOME", "/tmp/.cache/huggingface"))
+        LOG.info("[st] modèle chargé : %s", EMB_MODEL)
     return _ST_MODEL
 
+
 def _emb_st(texts: List[str]) -> np.ndarray:
     model = _get_st_model()
     vecs = model.encode(
@@ -191,55 +181,72 @@ def _emb_st(texts: List[str]) -> np.ndarray:
     ).astype("float32")
     return vecs
 
-# ----------------------- PROVIDER: Transformers (HF) --------------------------
+
 def _get_hf_model():
     global _HF_TOKENIZER, _HF_MODEL
     if _HF_MODEL is None or _HF_TOKENIZER is None:
         from transformers import AutoTokenizer, AutoModel
-        _HF_TOKENIZER = AutoTokenizer.from_pretrained(EMB_MODEL, cache_dir=CACHE_PATHS["hf_tf"])
-        _HF_MODEL = AutoModel.from_pretrained(EMB_MODEL, cache_dir=CACHE_PATHS["hf_tf"])
+        _HF_TOKENIZER = AutoTokenizer.from_pretrained(EMB_MODEL, cache_dir=os.getenv("HF_HOME", "/tmp/.cache/huggingface"))
+        _HF_MODEL = AutoModel.from_pretrained(EMB_MODEL, cache_dir=os.getenv("HF_HOME", "/tmp/.cache/huggingface"))
         _HF_MODEL.eval()
-        LOG.info("[hf] modèle chargé: %s (cache=%s)", EMB_MODEL, CACHE_PATHS["hf_tf"])
+        LOG.info("[hf] modèle chargé : %s", EMB_MODEL)
     return _HF_TOKENIZER, _HF_MODEL
 
-def _mean_pool(last_hidden_state: "np.ndarray", attention_mask: "np.ndarray") -> "np.ndarray":
+
+def _mean_pool(last_hidden_state: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
     mask = attention_mask[..., None].astype(last_hidden_state.dtype)
     summed = (last_hidden_state * mask).sum(axis=1)
     counts = mask.sum(axis=1).clip(min=1e-9)
     return summed / counts
 
+
 def _emb_hf(texts: List[str]) -> np.ndarray:
     import torch
     tok, mod = _get_hf_model()
-    all_vecs = []
+    all_vecs: List[np.ndarray] = []
     bs = max(1, EMB_BATCH)
     with torch.no_grad():
         for i in range(0, len(texts), bs):
-            batch = texts[i:i+bs]
+            batch = texts[i:i + bs]
             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())
             all_vecs.append(pooled.astype("float32"))
-    vecs = np.concatenate(all_vecs, axis=0)
-    return _l2_normalize(vecs)
+    return np.concatenate(all_vecs, axis=0)
+
+
+def _reduce_dim(vectors: np.ndarray, target_dim: int = EMB_DIM) -> np.ndarray:
+    """PCA simple pour réduire la dimension (si target_dim < current)."""
+    if target_dim >= vectors.shape[1]:
+        return vectors
+    from sklearn.decomposition import PCA
+    pca = PCA(n_components=target_dim, random_state=0)
+    return pca.fit_transform(vectors).astype("float32")
 
-# ---------------------------- DATASET / FAISS ---------------------------------
-def _save_dataset(ds_dir: str, rows: List[Dict[str, Any]]):
+
+# --------------------------------------------------------------------------- #
+#   DATASET / FAISS I/O
+# --------------------------------------------------------------------------- #
+def _save_dataset(ds_dir: str, rows: List[Dict[str, Any]], store_text: bool = STORE_TEXT) -> None:
+    """Sauvegarde le dataset au format JSONL (optionnellement sans le texte)."""
     os.makedirs(ds_dir, exist_ok=True)
     data_path = os.path.join(ds_dir, "data.jsonl")
     with open(data_path, "w", encoding="utf-8") as f:
         for r in rows:
+            if not store_text:
+                r = {k: v for k, v in r.items() if k != "text"}
             f.write(json.dumps(r, ensure_ascii=False) + "\n")
     meta = {"format": "jsonl", "columns": ["path", "text", "chunk_id"], "count": len(rows)}
     with open(os.path.join(ds_dir, "meta.json"), "w", encoding="utf-8") as f:
         json.dump(meta, f, ensure_ascii=False, indent=2)
 
+
 def _load_dataset(ds_dir: str) -> List[Dict[str, Any]]:
     data_path = os.path.join(ds_dir, "data.jsonl")
     if not os.path.isfile(data_path):
         return []
-    out = []
+    out: List[Dict[str, Any]] = []
     with open(data_path, "r", encoding="utf-8") as f:
         for line in f:
             try:
@@ -248,155 +255,234 @@ def _load_dataset(ds_dir: str) -> List[Dict[str, Any]]:
                 continue
     return out
 
-def _save_faiss(fx_dir: str, xb: np.ndarray, meta: Dict[str, Any]):
+
+def _save_faiss(fx_dir: str, xb: np.ndarray, meta: Dict[str, Any]) -> None:
+    """Sauvegarde un index FAISS quantisé (IVF‑PQ) ou plat selon FAISS_TYPE."""
     os.makedirs(fx_dir, exist_ok=True)
     idx_path = os.path.join(fx_dir, "emb.faiss")
-    index = faiss.IndexFlatIP(xb.shape[1])  # cosine ~ inner product si embeddings normalisés
-    index.add(xb)
+
+    if FAISS_TYPE == "IVF_PQ":
+        # ---- IVF‑PQ ---------------------------------------------------------
+        quantizer = faiss.IndexFlatIP(xb.shape[1])          # base (inner‑product ≈ cosine)
+        index = faiss.IndexIVFPQ(quantizer, xb.shape[1], FAISS_NLIST, FAISS_M, FAISS_NBITS)
+
+        # entraînement sur un sous‑échantillon (max 10 k vecteurs)
+        rng = np.random.default_rng(0)
+        train = xb[rng.choice(xb.shape[0], min(10_000, xb.shape[0]), replace=False]
+        index.train(train)
+
+        index.add(xb)
+        meta.update({
+            "index_type": "IVF_PQ",
+            "nlist": FAISS_NLIST,
+            "m": FAISS_M,
+            "nbits": FAISS_NBITS,
+        })
+    else:  # FLAT (fallback)
+        index = faiss.IndexFlatIP(xb.shape[1])
+        index.add(xb)
+        meta.update({"index_type": "FLAT"})
+
     faiss.write_index(index, idx_path)
+
+    # meta.json (inclut le type d’index)
     with open(os.path.join(fx_dir, "meta.json"), "w", encoding="utf-8") as f:
         json.dump(meta, f, ensure_ascii=False, indent=2)
 
+
 def _load_faiss(fx_dir: str) -> faiss.Index:
+    """Charge l’index en mode mmap (lecture à la volée)."""
     idx_path = os.path.join(fx_dir, "emb.faiss")
     if not os.path.isfile(idx_path):
-        raise FileNotFoundError(f"FAISS index introuvable: {idx_path}")
-    return faiss.read_index(idx_path)
+        raise FileNotFoundError(f"FAISS index introuvable : {idx_path}")
+    # mmap minimise la RAM utilisée
+    return faiss.read_index(idx_path, faiss.IO_FLAG_MMAP)
+
 
 def _tar_dir_to_bytes(dir_path: str) -> bytes:
+    """Archive gzip du répertoire (compression maximale)."""
     bio = io.BytesIO()
-    with tarfile.open(fileobj=bio, mode="w:gz") as tar:
+    with tarfile.open(fileobj=bio, mode="w:gz", compresslevel=9) as tar:
         tar.add(dir_path, arcname=os.path.basename(dir_path))
     bio.seek(0)
     return bio.read()
 
-# =============================================================================
-# WORKER POOL (asynchrone)
-# =============================================================================
-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:
+# --------------------------------------------------------------------------- #
+#   WORKER POOL (asynchrone)
+# --------------------------------------------------------------------------- #
+from concurrent.futures import ThreadPoolExecutor
+
+MAX_WORKERS = max(1, int(os.getenv("MAX_WORKERS", "1")))
+EXECUTOR = ThreadPoolExecutor(max_workers=MAX_WORKERS)
+LOG.info("ThreadPoolExecutor initialisé : max_workers=%s", MAX_WORKERS)
+
+
+def _proj_dirs(project_id: str) -> Tuple[str, str, str]:
+    base = os.path.join(os.getenv("DATA_ROOT", "/tmp/data"), project_id)
+    ds_dir = os.path.join(base, "dataset")
+    fx_dir = os.path.join(base, "faiss")
+    os.makedirs(ds_dir, exist_ok=True)
+    os.makedirs(fx_dir, exist_ok=True)
+    return base, ds_dir, fx_dir
+
+
+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.
+    Pipeline complet :
+    1️⃣ Chunking
+    2️⃣ Embedding (dummy / st / hf)
+    3️⃣ Réduction de dimension (PCA) si EMB_DIM < dim du modèle
+    4️⃣ Sauvegarde dataset (optionnel texte)
+    5️⃣ Index FAISS quantisé + mmap
     """
     try:
         base, ds_dir, fx_dir = _proj_dirs(st.project_id)
 
-        # 1) Chunking
-        _set_stage(st, "chunking")
+        # ------------------------------------------------------------------- #
+        # 1️⃣ Chunking
+        # ------------------------------------------------------------------- #
         rows: List[Dict[str, Any]] = []
         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})
+
+        for f in files:
+            path = (f.get("path") or "unknown").strip()
+            txt = f.get("text") or ""
+            chunks = _chunk_text(txt, size=chunk_size, overlap=overlap)
+            for i, ck in enumerate(chunks):
+                rows.append({"path": path, "text": ck, "chunk_id": i})
+
         st.total_chunks = len(rows)
-        _add_msg(st, f"Total chunks = {st.total_chunks}")
+        LOG.info("Chunking terminé : %d chunks", st.total_chunks)
 
-        # 2) Embedding
-        _set_stage(st, "embedding")
+        # ------------------------------------------------------------------- #
+        # 2️⃣ Embedding
+        # ------------------------------------------------------------------- #
         texts = [r["text"] for r in rows]
         if EMB_PROVIDER == "dummy":
             xb = _emb_dummy(texts, dim=EMB_DIM)
-            dim = xb.shape[1]
         elif EMB_PROVIDER == "st":
             xb = _emb_st(texts)
-            dim = xb.shape[1]
         else:
             xb = _emb_hf(texts)
-            dim = xb.shape[1]
 
-        st.embedded = xb.shape[0]
-        _add_msg(st, f"Embeddings {st.embedded}/{st.total_chunks}")
-        _add_msg(st, f"Embeddings dim={dim}")
-
-        # 3) Sauvegarde dataset (texte)
-        _save_dataset(ds_dir, rows)
-        _add_msg(st, f"Dataset (sans index) sauvegardé dans {ds_dir}")
+        # ------------------------------------------------------------------- #
+        # 3️⃣ Réduction de dimension (si nécessaire)
+        # ------------------------------------------------------------------- #
+        if xb.shape[1] != EMB_DIM:
+            xb = _reduce_dim(xb, target_dim=EMB_DIM)
 
-        # 4) FAISS
-        _set_stage(st, "indexing")
-        faiss_meta = {
-            "dim": int(dim),
+        st.embedded = xb.shape[0]
+        LOG.info("Embedding terminé : %d vecteurs (dim=%d)", st.embedded, xb.shape[1])
+
+        # ------------------------------------------------------------------- #
+        # 4️⃣ Sauvegarde du dataset
+        # ------------------------------------------------------------------- #
+        _save_dataset(ds_dir, rows, store_text=store_text)
+
+        # ------------------------------------------------------------------- #
+        # 5️⃣ Index FAISS
+        # ------------------------------------------------------------------- #
+        meta = {
+            "dim": int(xb.shape[1]),
             "count": int(xb.shape[0]),
             "provider": EMB_PROVIDER,
-            "model": EMB_MODEL if EMB_PROVIDER != "dummy" else None
+            "model": EMB_MODEL if EMB_PROVIDER != "dummy" else None,
         }
-        _save_faiss(fx_dir, xb, meta=faiss_meta)
+        _save_faiss(fx_dir, xb, 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={st.project_id})")
+        LOG.info("FAISS (%s) écrit : %s", FAISS_TYPE, os.path.join(fx_dir, "emb.faiss"))
 
-        _set_stage(st, "done")
+        # ------------------------------------------------------------------- #
+        #   Finalisation
+        # ------------------------------------------------------------------- #
+        st.stage = "done"
         st.finished_at = time.time()
     except Exception as e:
-        LOG.exception("Job %s failed", st.job_id)
+        LOG.exception("Job %s échoué", 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:
+
+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")
+    LOG.info("Job %s créé – %d fichiers", job_id, len(files))
+
+    EXECUTOR.submit(
+        _do_index_job,
+        st,
+        files,
+        chunk_size,
+        overlap,
+        batch_size,
+        store_text,
+    )
+    st.stage = "queued"
     return job_id
 
-# =============================================================================
-# FASTAPI
-# =============================================================================
+
+# --------------------------------------------------------------------------- #
+#   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=["*"],
+    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
+    store_text: bool = STORE_TEXT  # ← configurable
+
 
 @fastapi_app.get("/health")
 def health():
-    info = {
+    return {
         "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,
+        "data_root": os.getenv("DATA_ROOT", "/tmp/data"),
+        "faiss_type": FAISS_TYPE,
+        "emb_dim": EMB_DIM,
     }
-    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(
@@ -409,9 +495,10 @@ def index(req: IndexRequest):
         )
         return {"job_id": job_id}
     except Exception as e:
-        LOG.exception("index failed (submit)")
+        LOG.exception("Erreur soumission index")
         raise HTTPException(status_code=500, detail=str(e))
 
+
 @fastapi_app.get("/status/{job_id}")
 def status(job_id: str):
     st = JOBS.get(job_id)
@@ -419,26 +506,26 @@ def status(job_id: str):
         raise HTTPException(status_code=404, detail="job inconnu")
     return JSONResponse(st.model_dump())
 
+
 class SearchRequest(BaseModel):
     project_id: str
     query: str
     k: int = 5
 
+
 @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)):
+    # Vérifier la présence de l'index
+    if not (os.path.isfile(os.path.join(fx_dir, "emb.faiss")) and os.path.isfile(os.path.join(ds_dir, "data.jsonl"))):
         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")
 
-    # Embedding de la requête avec le MÊME provider
+    # Embedding de la requête (même provider)
     if EMB_PROVIDER == "dummy":
         q = _emb_dummy([req.query], dim=EMB_DIM)[0:1, :]
     elif EMB_PROVIDER == "st":
@@ -446,10 +533,13 @@ def search(req: SearchRequest):
     else:
         q = _emb_hf([req.query])[0:1, :]
 
-    # FAISS
+    # Recherche FAISS (mmap)
     index = _load_faiss(fx_dir)
     if index.d != q.shape[1]:
-        raise HTTPException(status_code=500, detail=f"dim incompatibles: index.d={index.d} vs query={q.shape[1]}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"dim incompatibles : index.d={index.d} vs query={q.shape[1]}",
+        )
     scores, ids = index.search(q, int(max(1, req.k)))
     ids = ids[0].tolist()
     scores = scores[0].tolist()
@@ -462,72 +552,81 @@ def search(req: SearchRequest):
         out.append({"path": r.get("path"), "text": r.get("text"), "score": float(sc)})
     return {"results": out}
 
-# ----------- ARTIFACTS EXPORT -----------
+
+# --------------------------------------------------------------------------- #
+#   ARTIFACTS EXPORT (gzip)
+# --------------------------------------------------------------------------- #
 @fastapi_app.get("/artifacts/{project_id}/dataset")
 def download_dataset(project_id: str):
-    base, ds_dir, _ = _proj_dirs(project_id)
+    _, ds_dir, _ = _proj_dirs(project_id)
     if not os.path.isdir(ds_dir):
         raise HTTPException(status_code=404, detail="Dataset introuvable")
     buf = _tar_dir_to_bytes(ds_dir)
-    headers = {"Content-Disposition": f'attachment; filename="{project_id}_dataset.tgz"'}
-    return StreamingResponse(io.BytesIO(buf), media_type="application/gzip", headers=headers)
+    hdr = {"Content-Disposition": f'attachment; filename="{project_id}_dataset.tgz"'}
+    return StreamingResponse(io.BytesIO(buf), media_type="application/gzip", headers=hdr)
+
 
 @fastapi_app.get("/artifacts/{project_id}/faiss")
 def download_faiss(project_id: str):
-    base, _, fx_dir = _proj_dirs(project_id)
+    _, _, fx_dir = _proj_dirs(project_id)
     if not os.path.isdir(fx_dir):
         raise HTTPException(status_code=404, detail="FAISS introuvable")
     buf = _tar_dir_to_bytes(fx_dir)
-    headers = {"Content-Disposition": f'attachment; filename="{project_id}_faiss.tgz"'}
-    return StreamingResponse(io.BytesIO(buf), media_type="application/gzip", headers=headers)
+    hdr = {"Content-Disposition": f'attachment; filename="{project_id}_faiss.tgz"'}
+    return StreamingResponse(io.BytesIO(buf), media_type="application/gzip", headers=hdr)
+
 
-# =============================================================================
-# GRADIO UI (facultatif de test)
-# =============================================================================
+# --------------------------------------------------------------------------- #
+#   GRADIO UI (facultatif – simple test)
+# --------------------------------------------------------------------------- #
 def _ui_index(project_id: str, sample_text: str):
     files = [{"path": "sample.txt", "text": sample_text}]
-    from pydantic import ValidationError
     try:
         req = IndexRequest(project_id=project_id, files=[FileItem(**f) for f in files])
-    except ValidationError as e:
-        return f"Erreur: {e}"
+    except Exception as e:
+        return f"❌ Erreur validation : {e}"
     try:
         res = index(req)
-        return f"Job lancé: {res['job_id']}"
+        return f"✅ Job lancé : {res['job_id']}"
     except Exception as e:
-        return f"Erreur index: {e}"
+        return f"❌ Erreur index : {e}"
+
 
 def _ui_search(project_id: str, query: str, k: int):
     try:
         res = search(SearchRequest(project_id=project_id, query=query, k=int(k)))
         return json.dumps(res, ensure_ascii=False, indent=2)
     except Exception as e:
-        return f"Erreur search: {e}"
-
-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)
-        btn = gr.Button("Lancer index (sample)")
-        out = gr.Textbox(label="Résultat")
-        btn.click(_ui_index, inputs=[pid, sample], outputs=[out])
-
-    with gr.Tab("Search"):
-        pid2 = gr.Textbox(label="Project ID", value="DEEPWEB")
+        return f"❌ Erreur recherche : {e}"
+
+
+import gradio as gr
+
+with gr.Blocks(title="Remote Indexer (Async – Optimisé)", analytics_enabled=False) as ui:
+    gr.Markdown("## Remote Indexer – Optimisé (FAISS quantisé, mmap, texte optionnel)")
+    with gr.Row():
+        pid = gr.Textbox(label="Project ID", value="DEMO")
+        txt = gr.Textbox(label="Texte d’exemple", lines=4, value="Alpha bravo charlie delta echo foxtrot.")
+        btn_idx = gr.Button("Lancer index (sample)")
+        out_idx = gr.Textbox(label="Résultat")
+        btn_idx.click(_ui_index, inputs=[pid, txt], outputs=[out_idx])
+
+    with gr.Row():
         q = gr.Textbox(label="Query", value="alpha")
-        k = gr.Slider(1, 20, value=5, step=1, label="k")
-        btn2 = gr.Button("Rechercher")
-        out2 = gr.Code(label="Résultats")
-        btn2.click(_ui_search, inputs=[pid2, q, k], outputs=[out2])
+        k = gr.Slider(1, 20, value=5, step=1, label="Top‑K")
+        btn_q = gr.Button("Rechercher")
+        out_q = gr.Code(label="Résultats")
+        btn_q.click(_ui_search, inputs=[pid, q, k], outputs=[out_q])
 
 fastapi_app = gr.mount_gradio_app(fastapi_app, ui, path="/ui")
 
-# =============================================================================
-# MAIN
-# =============================================================================
+
+# --------------------------------------------------------------------------- #
+#   MAIN
+# --------------------------------------------------------------------------- #
 if __name__ == "__main__":
     import uvicorn
-    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)
+
+    PORT = int(os.getenv("PORT", "7860"))
+    LOG.info("Démarrage Uvicorn – port %s – UI à /ui", PORT)
+    uvicorn.run(fastapi_app, host="0.0.0.0", port=PORT)