Funk-lab/stt/server/stt_server/longterm.py

"""Mémoire long-terme sémantique (Qdrant).

Stocke les tours utilisateur comme vecteurs dans la collection `stt-memory` et retrouve
les souvenirs pertinents pour les injecter dans le prompt. Embeddings via une instance
llama-server dédiée (`nomic-embed-text`, gpu-01 `:1238`, dim 768). **Tout dégrade proprement** :
si Qdrant ou l'endpoint d'embedding est injoignable, `recall` renvoie [] et `remember` ne
fait rien (la mémoire court-terme de session continue de fonctionner).

> Le modèle d'embedding est configurable (`STT_EMBED_URL`/`STT_EMBED_MODEL`). Si la dimension
> change, `_ensure_collection` recrée automatiquement la collection. Voir admin/ia/rag.md.
"""

from __future__ import annotations

import time
import uuid

import httpx

from stt_server.config import settings


class LongTermMemory:
    def __init__(self) -> None:
        self.qdrant = settings.qdrant_url.rstrip("/")
        self.collection = settings.qdrant_collection
        self.embed_url = settings.embed_url
        self.embed_model = settings.embed_model
        self.top_k = settings.memory_top_k
        self.recall_timeout = settings.memory_recall_timeout
        self.store_timeout = settings.memory_store_timeout
        self._ready = False
        # Client persistant : pooling + keep-alive (évite un handshake TCP par appel).
        self._http: httpx.AsyncClient | None = None

    def _client(self) -> httpx.AsyncClient:
        if self._http is None or self._http.is_closed:
            self._http = httpx.AsyncClient()
        return self._http

    async def aclose(self) -> None:
        if self._http is not None and not self._http.is_closed:
            await self._http.aclose()

    async def _embed(self, text: str, timeout: float) -> list[float]:
        r = await self._client().post(
            self.embed_url,
            json={"model": self.embed_model, "input": text},
            timeout=timeout,
        )
        r.raise_for_status()
        return r.json()["data"][0]["embedding"]

    async def _ensure_collection(self, dim: int) -> None:
        if self._ready:
            return
        client = self._client()
        r = await client.get(
            f"{self.qdrant}/collections/{self.collection}", timeout=self.store_timeout
        )
        if r.status_code == 200:
            # Collection existante : si la dimension a changé (modèle d'embedding
            # différent, ex. Qwen3 4096 → nomic 768), on recrée — les anciens vecteurs
            # sont incomparables dans le nouvel espace.
            vectors = (
                r.json().get("result", {}).get("config", {}).get("params", {}).get("vectors", {})
            )
            size = vectors.get("size") if isinstance(vectors, dict) else None
            if size == dim:
                self._ready = True
                return
            await client.delete(f"{self.qdrant}/collections/{self.collection}")
        await client.put(
            f"{self.qdrant}/collections/{self.collection}",
            json={"vectors": {"size": dim, "distance": "Cosine"}},
            timeout=self.store_timeout,
        )
        self._ready = True

    async def recall(self, text: str) -> tuple[list[str], list[float] | None]:
        """Souvenirs pertinents + le vecteur de la requête (réutilisable pour `store`).

        Sur le chemin de réponse → timeout serré (`recall_timeout`) : si l'embed ou Qdrant
        traîne, on dégrade vite en `([], None)` plutôt que de faire patienter le client.
        Le vecteur renvoyé évite de ré-embedder le même texte au moment du `store`.
        """
        try:
            client = self._client()
            vec = await self._embed(text, timeout=self.recall_timeout)
            r = await client.post(
                f"{self.qdrant}/collections/{self.collection}/points/search",
                json={"vector": vec, "limit": self.top_k, "with_payload": True},
                timeout=self.recall_timeout,
            )
            if r.status_code == 404:  # collection pas encore créée → rien à rappeler
                return [], vec
            r.raise_for_status()
            pts = r.json().get("result", [])
            texts = [p["payload"]["text"] for p in pts if p.get("payload", {}).get("text")]
            return texts, vec
        except Exception:  # noqa: BLE001 — dégrade silencieusement
            return [], None

    async def health(self) -> dict:
        """Diagnostic actif : sonde embeddings + Qdrant + collection, sans rien avaler.

        Contrairement à recall/store (qui dégradent en silence), expose les erreurs
        pour pouvoir déboguer la mémoire long-terme sans `kubectl exec`.
        """
        out: dict = {
            "enabled": True,
            "qdrant_url": self.qdrant,
            "embed_url": self.embed_url,
            "collection": self.collection,
            "embed": {"ok": False},
            "qdrant": {"ok": False},
        }
        client = self._client()
        # 1) embeddings
        try:
            vec = await self._embed("ping mémoire", timeout=self.store_timeout)
            out["embed"] = {"ok": True, "dim": len(vec)}
        except Exception as e:  # noqa: BLE001 — on veut l'erreur
            out["embed"] = {"ok": False, "error": f"{type(e).__name__}: {e}"}
        # 2) Qdrant + collection
        try:
            r = await client.get(f"{self.qdrant}/collections", timeout=self.store_timeout)
            r.raise_for_status()
            names = [c["name"] for c in r.json().get("result", {}).get("collections", [])]
            exists = self.collection in names
            qdrant: dict = {"ok": True, "collections": names, "collection_exists": exists}
            if exists:
                cr = await client.get(
                    f"{self.qdrant}/collections/{self.collection}", timeout=self.store_timeout
                )
                if cr.status_code == 200:
                    qdrant["points_count"] = cr.json().get("result", {}).get("points_count")
            out["qdrant"] = qdrant
        except Exception as e:  # noqa: BLE001
            out["qdrant"] = {"ok": False, "error": f"{type(e).__name__}: {e}"}
        return out

    async def store(
        self, session_id: str, text: str, vec: list[float] | None = None
    ) -> None:
        """Mémorise un tour. Pensé pour tourner **en tâche de fond** (hors chemin de réponse).

        `vec` : si fourni (réutilisé depuis `recall`), on évite un 2ᵉ embed du même texte.
        Pas de `?wait=true` : on n'attend pas le flush disque de Qdrant — le client a déjà
        sa réponse, ce point peut être indexé en arrière-plan.
        """
        try:
            if vec is None:
                vec = await self._embed(text, timeout=self.store_timeout)
            await self._ensure_collection(len(vec))
            await self._client().put(
                f"{self.qdrant}/collections/{self.collection}/points",
                json={
                    "points": [
                        {
                            "id": str(uuid.uuid4()),
                            "vector": vec,
                            "payload": {
                                "text": text,
                                "session_id": session_id,
                                "ts": time.time(),
                            },
                        }
                    ]
                },
                timeout=self.store_timeout,
            )
        except Exception:  # noqa: BLE001 — dégrade silencieusement
            pass