Update main.py

main.py

@@ -80,89 +80,161 @@ if "deepinfra" in EMB_BACKEND_ORDER and not DI_TOKEN:
     LOG.warning("DEEPINFRA_API_KEY manquant — tentatives DeepInfra échoueront.")
 
 # ---------- Vector store abstraction ----------
-class VectorStoreBase:
+# ---------- Vector Stores ----------
+from typing import Dict, Any, List, Optional, Tuple
+import numpy as np
+import logging
+
+LOG = logging.getLogger("remote_indexer")
+
+try:
+    from qdrant_client import QdrantClient
+    from qdrant_client.http.models import VectorParams, Distance, PointStruct
+except Exception:
+    QdrantClient = None
+    PointStruct = None
+
+
+class BaseStore:
     def ensure_collection(self, name: str, dim: int): ...
-    def upsert(self, name: str, vectors: np.ndarray, payloads: List[dict]) -> int: ...
-    def search(self, name: str, query_vec: np.ndarray, limit: int):
-        """return list of objects with .score and .payload"""
-        ...
+    def upsert(self, name: str, vectors: np.ndarray, payloads: List[Dict[str, Any]]) -> int: ...
+    def search(self, name: str, query_vec: np.ndarray, top_k: int) -> List[Dict[str, Any]]: ...
     def wipe(self, name: str): ...
 
-class MemoryHit:
-    def __init__(self, score: float, payload: dict):
-        self.score = score
-        self.payload = payload
 
-class MemoryStore(VectorStoreBase):
-    """Simple store en mémoire (cosine sur vecteurs normalisés). Persistance: vie du process."""
+class MemoryStore(BaseStore):
+    """Store en mémoire (volatile)."""
     def __init__(self):
-        self.data: Dict[str, Dict[str, Any]] = {}  # {col: {"dim": d, "vecs": np.ndarray [N,d], "payloads": List[dict]}}
-        LOG.warning("Vector store: MEMORY (fallback). Les données sont volatiles (perdues au restart).")
+        # { collection: {"vecs": [np.ndarray], "payloads": [dict]} }
+        self.db: Dict[str, Dict[str, List[Any]]] = {}
 
     def ensure_collection(self, name: str, dim: int):
-        col = self.data.get(name)
-        if not col:
-            self.data[name] = {"dim": dim, "vecs": np.zeros((0, dim), dtype=np.float32), "payloads": []}
-
-    def upsert(self, name: str, vectors: np.ndarray, payloads: List[dict]) -> int:
-        self.ensure_collection(name, vectors.shape[1])
-        col = self.data[name]
-        if vectors.ndim != 2 or vectors.shape[1] != col["dim"]:
-            raise RuntimeError(f"MemoryStore: bad shape {vectors.shape}, expected (*,{col['dim']})")
-        col["vecs"] = np.vstack([col["vecs"], vectors.astype(np.float32)])
-        col["payloads"].extend(payloads)
-        return vectors.shape[0]
-
-    def search(self, name: str, query_vec: np.ndarray, limit: int):
-        col = self.data.get(name)
-        if not col or col["vecs"].shape[0] == 0:
+        self.db.setdefault(name, {"vecs": [], "payloads": [], "dim": dim})
+
+    def upsert(self, name: str, vectors: np.ndarray, payloads: List[Dict[str, Any]]) -> int:
+        if name not in self.db:
+            raise RuntimeError(f"MemoryStore: collection {name} inconnue")
+        if len(vectors) != len(payloads):
+            raise ValueError("MemoryStore.upsert: tailles vectors/payloads incohérentes")
+        self.db[name]["vecs"].extend([v.astype(np.float32) for v in vectors])
+        self.db[name]["payloads"].extend(payloads)
+        return len(vectors)
+
+    def search(self, name: str, query_vec: np.ndarray, top_k: int) -> List[Dict[str, Any]]:
+        if name not in self.db or not self.db[name]["vecs"]:
             return []
-        V = col["vecs"]  # [N,d], déjà normalisés
-        q = query_vec.reshape(1, -1)  # [1,d]
-        scores = (V @ q.T).ravel()  # cos sim
-        idx = np.argsort(-scores)[:limit]
-        return [MemoryHit(float(scores[i]), col["payloads"][i]) for i in idx]
+        mat = np.vstack(self.db[name]["vecs"])  # [N, dim]
+        q = query_vec.reshape(1, -1).astype(np.float32)      # [1, dim]
+        # cosine similarity sur vecteurs normalisés
+        # (on suppose que les embeddings sont déjà normalisés en amont)
+        sims = (mat @ q.T).ravel()  # [N]
+        top_idx = np.argsort(-sims)[:top_k]
+        out = []
+        for i in top_idx:
+            pl = dict(self.db[name]["payloads"][i])
+            pl["_score"] = float(sims[i])
+            out.append(pl)
+        return out
 
     def wipe(self, name: str):
-        if name in self.data:
-            del self.data[name]
+        self.db.pop(name, None)
 
-class QdrantStore(VectorStoreBase):
-    def __init__(self, url: str, api_key: Optional[str]):
-        if QdrantClient is None:
-            raise RuntimeError("qdrant-client non installé.")
+
+class QdrantStore(BaseStore):
+    """Store Qdrant avec gestion d'IDs séquentiels par collection."""
+    def __init__(self, url: str, api_key: Optional[str] = None):
+        if QdrantClient is None or PointStruct is None:
+            raise RuntimeError("qdrant_client non disponible")
         self.client = QdrantClient(url=url, api_key=api_key if api_key else None)
-        # ping rapide
+        # compteur d'IDs par collection
+        self._next_ids: Dict[str, int] = {}
+
+    def _init_next_id(self, name: str):
+        # on cherche le count exact des points existants pour démarrer l'ID à count
         try:
-            _ = self.client.get_collections()
-            LOG.info("Connecté à Qdrant.")
-        except Exception as e:
-            raise RuntimeError(f"Connexion Qdrant impossible: {e}")
+            cnt = self.client.count(collection_name=name, exact=True).count
+        except Exception:
+            # si count échoue (collection vide juste créée), on démarre à 0
+            cnt = 0
+        self._next_ids[name] = int(cnt)
 
     def ensure_collection(self, name: str, dim: int):
+        # si existe déjà, rien à faire ; sinon, création
         try:
-            self.client.get_collection(name); return
+            self.client.get_collection(name)
         except Exception:
-            pass
-        self.client.create_collection(
-            collection_name=name,
-            vectors_config=VectorParams(size=dim, distance=Distance.COSINE),
-        )
-
-    def upsert(self, name: str, vectors: np.ndarray, payloads: List[dict]) -> int:
-        points = [
-            PointStruct(id=None, vector=v.tolist(), payload=payloads[i])
+            self.client.create_collection(
+                collection_name=name,
+                vectors_config=VectorParams(size=dim, distance=Distance.COSINE),
+            )
+        # initialiser le prochain id si absent
+        if name not in self._next_ids:
+            self._init_next_id(name)
+
+    def upsert(self, name: str, vectors: np.ndarray, payloads: List[Dict[str, Any]]) -> int:
+        if vectors is None or len(vectors) == 0:
+            return 0
+        if len(vectors) != len(payloads):
+            raise ValueError("QdrantStore.upsert: tailles vectors/payloads incohérentes")
+
+        if name not in self._next_ids:
+            self._init_next_id(name)
+
+        start = self._next_ids[name]
+        # construction des points avec IDs séquentiels (int)
+        pts = [
+            PointStruct(id=start + i, vector=v.astype(np.float32).tolist(), payload=payloads[i])
             for i, v in enumerate(vectors)
         ]
-        self.client.upsert(collection_name=name, points=points)
-        return len(points)
-
-    def search(self, name: str, query_vec: np.ndarray, limit: int):
-        res = self.client.search(collection_name=name, query_vector=query_vec.tolist(), limit=limit)
-        return res
+        self.client.upsert(collection_name=name, points=pts)
+        added = len(pts)
+        self._next_ids[name] += added
+        LOG.debug(f"QdrantStore.upsert: +{added} points (next_id={self._next_ids[name]})")
+        return added
+
+    def search(self, name: str, query_vec: np.ndarray, top_k: int) -> List[Dict[str, Any]]:
+        if query_vec.ndim == 2:
+            qv = query_vec[0].astype(np.float32).tolist()
+        else:
+            qv = query_vec.astype(np.float32).tolist()
+        res = self.client.search(collection_name=name, query_vector=qv, limit=int(top_k))
+        out = []
+        for p in res:
+            pl = p.payload or {}
+            pl["_score"] = float(p.score) if hasattr(p, "score") else None
+            out.append(pl)
+        return out
 
     def wipe(self, name: str):
-        self.client.delete_collection(name)
+        try:
+            self.client.delete_collection(name)
+        except Exception:
+            pass
+        self._next_ids.pop(name, None)
+
+
+# ---------- Initialisation du store actif ----------
+import os
+
+VECTOR_STORE = os.getenv("VECTOR_STORE", "qdrant").strip().lower()
+QDRANT_URL = os.getenv("QDRANT_URL", "").strip()
+QDRANT_API = os.getenv("QDRANT_API_KEY", "").strip()
+
+try:
+    if VECTOR_STORE == "qdrant" and QDRANT_URL:
+        STORE: BaseStore = QdrantStore(QDRANT_URL, api_key=QDRANT_API)
+        # test léger: liste des collections
+        _ = STORE.client.get_collections()
+        LOG.info("Connecté à Qdrant.")
+        VECTOR_STORE_ACTIVE = "QdrantStore"
+    else:
+        raise RuntimeError("Qdrant non configuré, fallback mémoire.")
+except Exception as e:
+    LOG.error(f"Qdrant indisponible ({e}) — fallback en mémoire.")
+    STORE = MemoryStore()
+    VECTOR_STORE_ACTIVE = "MemoryStore"
+    LOG.warning("Vector store: MEMORY (fallback). Les données sont volatiles (perdues au restart).")
+
 
 # Sélection / auto-fallback du store
 STORE: VectorStoreBase