src/memory/rag.py

"""Phase 11 — mémoire / « apprendre des cours », écrit à la main.

Deux modes, implémentés **séparément** :

**Mode A — RAG (réaliste).** Les cours (texte ; et pour les vidéos/audio, leurs
transcriptions + légendes de frames clés) sont découpés en passages, encodés en
vecteurs et indexés dans une base vectorielle maison (numPy, similarité cosinus).
À la question, on récupère les passages pertinents et on répond à partir d'eux.

**Mode B — apprentissage permanent (EXPÉRIMENTAL ET INSTABLE).** Mise à jour des
poids pendant l'« ingestion » d'un cours (quelques pas de gradient), avec un tampon
de rejeu pour limiter l'oubli catastrophique. Voir `src/memory/permanent.py` et
docs/architecture.md §10. On ne promet rien d'irréaliste : cela mémorise un texte
court, ne remplace pas un vrai fine-tuning, et peut dégrader les autres capacités.

Aucun modèle pré-entraîné, aucun service externe : embeddings par *feature hashing*
déterministe (pur Python), ou embeddings du décodeur maison si l'on veut.
"""

from __future__ import annotations

import hashlib
import re
from dataclasses import dataclass, field

import numpy as np

_WORD_RE = re.compile(r"[\wàâäéèêëîïôöùûüç'-]+", re.UNICODE)
_SENT_RE = re.compile(r"(?<=[.!?])\s+|\n+")


# ---------------------------------------------------------------------------
# Embeddings maison (feature hashing) — déterministe, sans entraînement
# ---------------------------------------------------------------------------


def _hash(word: str, dim: int) -> tuple[int, float]:
    """Hash stable d'un mot vers (index, signe) pour une projection aléatoire."""
    h = int.from_bytes(hashlib.blake2b(word.encode("utf-8"), digest_size=8).digest(), "big")
    return h % dim, 1.0 if (h >> 63) & 1 else -1.0


def hashed_embed(text: str, dim: int = 256, idf: dict[str, float] | None = None) -> np.ndarray:
    """Vecteur L2-normalisé d'un texte par feature hashing (sac de mots signé).

    Aucun paramètre appris : deux textes proches lexicalement le sont aussi
    vectoriellement. Suffisant et honnête pour un RAG jouet à petite échelle.
    """
    vec = np.zeros(dim, dtype=np.float32)
    for m in _WORD_RE.finditer(text.lower()):
        w = m.group(0)
        idx, sign = _hash(w, dim)
        weight = idf.get(w, 1.0) if idf else 1.0
        vec[idx] += sign * weight
    norm = np.linalg.norm(vec)
    return vec / norm if norm > 0 else vec


def compute_idf(texts: list[str]) -> dict[str, float]:
    n = max(1, len(texts))
    df: dict[str, int] = {}
    for t in texts:
        for w in {m.group(0) for m in _WORD_RE.finditer(t.lower())}:
            df[w] = df.get(w, 0) + 1
    return {w: float(np.log((1 + n) / (1 + c)) + 1.0) for w, c in df.items()}


# ---------------------------------------------------------------------------
# Base vectorielle maison (numpy)
# ---------------------------------------------------------------------------


class VectorIndex:
    """Index vectoriel par similarité cosinus (numPy), écrit à la main."""

    def __init__(self, dim: int = 256):
        self.dim = dim
        self._vectors: np.ndarray | None = None

    def add(self, vectors: np.ndarray) -> None:
        v = np.asarray(vectors, dtype=np.float32)
        if v.ndim == 1:
            v = v[None, :]
        self._vectors = v if self._vectors is None else np.vstack([self._vectors, v])

    def search(self, query: np.ndarray, k: int = 3) -> list[tuple[int, float]]:
        """Retourne ``(index, score)`` des ``k`` passages les plus proches."""
        if self._vectors is None or len(self._vectors) == 0:
            return []
        q = np.asarray(query, dtype=np.float32)
        q = q / (np.linalg.norm(q) + 1e-9)
        mat = self._vectors / (np.linalg.norm(self._vectors, axis=1, keepdims=True) + 1e-9)
        scores = mat @ q
        k = min(k, len(scores))
        top = np.argsort(-scores)[:k]
        return [(int(i), float(scores[i])) for i in top]

    def __len__(self) -> int:
        return 0 if self._vectors is None else len(self._vectors)


# ---------------------------------------------------------------------------
# Découpage en passages
# ---------------------------------------------------------------------------


@dataclass
class Passage:
    text: str
    source: str = "manual"
    start: float = 0.0  # position (s) dans la vidéo/audio source, si applicable


def chunk_text(text: str, max_words: int = 40, source: str = "manual") -> list[Passage]:
    """Découpe un texte en passages (phrases regroupées jusqu'à ``max_words``)."""
    sentences = [s.strip() for s in _SENT_RE.split(text) if s.strip()]
    passages: list[Passage] = []
    buf: list[str] = []
    count = 0
    for s in sentences:
        n = len(_WORD_RE.findall(s))
        if buf and count + n > max_words:
            passages.append(Passage(" ".join(buf), source=source))
            buf, count = [], 0
        buf.append(s)
        count += n
    if buf:
        passages.append(Passage(" ".join(buf), source=source))
    return passages


# ---------------------------------------------------------------------------
# Mémoire de cours (mode A — RAG)
# ---------------------------------------------------------------------------


@dataclass
class LessonMemory:
    """Base RAG maison : ingestion de cours -> passages -> réponses extractives."""

    dim: int = 256
    max_words: int = 40
    passages: list[Passage] = field(default_factory=list)
    index: VectorIndex = field(default_factory=VectorIndex)
    _idf: dict[str, float] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if self.index.dim != self.dim:
            self.index = VectorIndex(self.dim)

    # -- ingestion ---------------------------------------------------------
    def add_text(self, text: str, source: str = "manual") -> None:
        for p in chunk_text(text, max_words=self.max_words, source=source):
            self.passages.append(p)
        self._reindex()

    def add_transcript(self, segments: list[tuple[float, str]], source: str = "video") -> None:
        """Ajoute une transcription horodatée (vidéo/audio) : ``[(t, texte), ...]``."""
        for t, seg in segments:
            for p in chunk_text(seg, max_words=self.max_words, source=source):
                p.start = float(t)
                self.passages.append(p)
        self._reindex()

    def add_video(self, transcript: list[tuple[float, str]],
                  frame_captions: list[tuple[float, str]] | None = None,
                  source: str = "video") -> None:
        """Un « cours vidéo » = transcription + légendes de frames clés."""
        self.add_transcript(transcript, source=source)
        if frame_captions:
            self.add_transcript(frame_captions, source=f"{source}:frames")

    def _reindex(self) -> None:
        texts = [p.text for p in self.passages]
        self._idf = compute_idf(texts)
        self.index = VectorIndex(self.dim)
        if texts:
            self.index.add(np.stack([hashed_embed(t, self.dim, self._idf) for t in texts]))

    # -- requête -----------------------------------------------------------
    def retrieve(self, query: str, k: int = 3) -> list[tuple[Passage, float]]:
        qv = hashed_embed(query, self.dim, self._idf)
        return [(self.passages[i], s) for i, s in self.index.search(qv, k)]

    def answer(self, query: str, k: int = 3) -> dict:
        """Réponse **extractive** : renvoie le meilleur passage + les sources.

        Mode A réaliste : pas d'hallucination, la réponse est un passage réel du cours.
        La confiance est le score cosinus du meilleur passage.
        """
        hits = self.retrieve(query, k=k)
        if not hits:
            return {"answer": "", "sources": [], "confidence": 0.0, "retrieved": []}
        best, score = hits[0]
        return {
            "answer": best.text,
            "sources": sorted({p.source for p, _ in hits}),
            "confidence": round(score, 4),
            "retrieved": [{"text": p.text, "source": p.source, "start": p.start,
                           "score": round(s, 4)} for p, s in hits],
        }

    def answer_with_model(self, query: str, model, tokenizer, k: int = 3,
                          max_new_tokens: int = 60) -> dict:
        """Réponse **générative** : le contexte récupéré est mis dans le prompt.

        Le modèle maison génère la réponse à partir des passages retrouvés. On
        renvoie aussi les passages (traçabilité) : si la génération est faible, la
        réponse extractive reste disponible via :meth:`answer`.
        """
        from src.generate import generate_text

        retrieved = self.retrieve(query, k=k)
        context = " ".join(p.text for p, _ in retrieved)
        prompt = f"Contexte : {context}\nQuestion : {query}\nRéponse :"
        text = generate_text(model, tokenizer, prompt, max_new_tokens=max_new_tokens,
                             greedy=True)
        return {
            "answer": text.split("Réponse :")[-1].strip(),
            "sources": sorted({p.source for p, _ in retrieved}),
            "confidence": round(retrieved[0][1], 4) if retrieved else 0.0,
            "retrieved": [{"text": p.text, "source": p.source, "score": round(s, 4)}
                          for p, s in retrieved],
        }