rag_core.py

"""
rag_core.py — Le coeur du moteur RAG.

Contient les 4 briques de base, dans l'ordre du pipeline :
  1. load_documents  : lire les fichiers (.txt, .md, .pdf)
  2. chunk_text      : decouper en morceaux qui se chevauchent
  3. Embedder        : transformer du texte en vecteurs
  4. VectorStore     : stocker les vecteurs et retrouver les plus proches

Aucune base de donnees externe : l'index est un simple fichier .npz.
C'est volontaire — on veut pouvoir expliquer chaque ligne.
"""

from __future__ import annotations

import json
import os
import re
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Iterable

import numpy as np

# --------------------------------------------------------------------------
# 1. LECTURE DES DOCUMENTS
# --------------------------------------------------------------------------

SUPPORTED = {".txt", ".md", ".pdf"}


@dataclass
class Chunk:
    """Un morceau de document, avec sa provenance (indispensable pour citer)."""
    text: str
    source: str        # nom du fichier
    chunk_id: int      # numero du morceau dans ce fichier
    page: int | None = None


def _read_pdf(path: Path) -> list[tuple[str, int]]:
    """Retourne [(texte_de_la_page, numero_de_page), ...]."""
    from pypdf import PdfReader

    reader = PdfReader(str(path))
    out = []
    for i, page in enumerate(reader.pages, start=1):
        txt = page.extract_text() or ""
        if txt.strip():
            out.append((txt, i))
    return out


def load_documents(folder: str | Path) -> list[tuple[str, str, int | None]]:
    """
    Lit tous les fichiers supportes d'un dossier.
    Retourne une liste de (texte, nom_du_fichier, page).
    """
    folder = Path(folder)
    if not folder.exists():
        raise FileNotFoundError(f"Dossier introuvable : {folder}")

    docs: list[tuple[str, str, int | None]] = []
    for path in sorted(folder.iterdir()):
        if path.suffix.lower() not in SUPPORTED:
            continue
        if path.suffix.lower() == ".pdf":
            for txt, page in _read_pdf(path):
                docs.append((txt, path.name, page))
        else:
            docs.append((path.read_text(encoding="utf-8", errors="ignore"), path.name, None))

    if not docs:
        raise ValueError(
            f"Aucun document lisible dans {folder}. "
            f"Formats acceptes : {', '.join(sorted(SUPPORTED))}"
        )
    return docs


# --------------------------------------------------------------------------
# 2. DECOUPAGE (CHUNKING)
# --------------------------------------------------------------------------

def clean_text(text: str) -> str:
    """Normalise les espaces et supprime les sauts de ligne parasites."""
    text = text.replace("\r", "\n")
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


def chunk_text(text: str, size: int = 800, overlap: int = 150) -> list[str]:
    """
    Decoupe le texte en morceaux d'environ `size` caracteres,
    avec un chevauchement de `overlap` caracteres.

    Pourquoi le chevauchement ? Si une information est a cheval sur deux
    morceaux, elle serait coupee en deux et deviendrait introuvable.
    Le chevauchement garantit qu'elle apparait entiere quelque part.

    On coupe en priorite sur une fin de phrase pour ne pas casser le sens.
    """
    text = clean_text(text)
    if len(text) <= size:
        return [text] if text else []

    chunks: list[str] = []
    start = 0
    while start < len(text):
        end = start + size

        if end < len(text):
            # on cherche une fin de phrase dans le dernier tiers du morceau
            window = text[start + int(size * 0.6): end]
            match = list(re.finditer(r"[.!?]\s", window))
            if match:
                end = start + int(size * 0.6) + match[-1].end()

        piece = text[start:end].strip()
        if piece:
            chunks.append(piece)

        start = end - overlap
        if start <= 0 or start >= len(text):
            break
    return chunks


def build_chunks(folder: str | Path, size: int = 800, overlap: int = 150) -> list[Chunk]:
    """Lit les documents et les transforme en liste de Chunk."""
    chunks: list[Chunk] = []
    for text, source, page in load_documents(folder):
        for i, piece in enumerate(chunk_text(text, size, overlap)):
            chunks.append(Chunk(text=piece, source=source, chunk_id=i, page=page))
    return chunks


# --------------------------------------------------------------------------
# 3. EMBEDDINGS
# --------------------------------------------------------------------------

class Embedder:
    """
    Transforme du texte en vecteurs, en local, via sentence-transformers.

    Le modele par defaut (all-MiniLM-L6-v2) fait ~90 Mo, tourne sur CPU,
    et produit des vecteurs de dimension 384. Il est telecharge une seule
    fois puis fonctionne hors ligne.
    """

    def __init__(self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2"):
        from sentence_transformers import SentenceTransformer

        self.model_name = model_name
        self.model = SentenceTransformer(model_name)

    def encode(self, texts: list[str]) -> np.ndarray:
        vecs = self.model.encode(
            texts,
            batch_size=32,
            show_progress_bar=len(texts) > 50,
            convert_to_numpy=True,
            normalize_embeddings=True,   # normalise -> produit scalaire = cosinus
        )
        return vecs.astype(np.float32)


# --------------------------------------------------------------------------
# 4. INDEX VECTORIEL
# --------------------------------------------------------------------------

class VectorStore:
    """
    Index vectoriel minimal : une matrice numpy + les metadonnees.

    La recherche est une simple similarite cosinus. Comme les vecteurs sont
    normalises, le cosinus se reduit a un produit scalaire : vectors @ query.
    Sur quelques milliers de morceaux, c'est instantane et parfaitement
    suffisant. Une vraie base vectorielle (Azure AI Search, FAISS, Chroma)
    devient utile a partir de centaines de milliers de vecteurs.
    """

    def __init__(self, vectors: np.ndarray, chunks: list[Chunk], model_name: str):
        self.vectors = vectors
        self.chunks = chunks
        self.model_name = model_name

    # --- persistance -------------------------------------------------------

    def save(self, path: str | Path = "index") -> None:
        path = Path(path)
        path.mkdir(exist_ok=True)
        np.save(path / "vectors.npy", self.vectors)
        (path / "chunks.json").write_text(
            json.dumps(
                {"model": self.model_name, "chunks": [asdict(c) for c in self.chunks]},
                ensure_ascii=False,
                indent=2,
            ),
            encoding="utf-8",
        )

    @classmethod
    def load(cls, path: str | Path = "index") -> "VectorStore":
        path = Path(path)
        if not (path / "vectors.npy").exists():
            raise FileNotFoundError(
                "Index introuvable. Lance d'abord :  python ingest.py"
            )
        vectors = np.load(path / "vectors.npy")
        data = json.loads((path / "chunks.json").read_text(encoding="utf-8"))
        chunks = [Chunk(**c) for c in data["chunks"]]
        return cls(vectors, chunks, data.get("model", "inconnu"))

    # --- recherche ---------------------------------------------------------

    def search(self, query_vector: np.ndarray, k: int = 4) -> list[tuple[Chunk, float]]:
        """Retourne les k morceaux les plus proches, avec leur score."""
        scores = self.vectors @ query_vector      # produit scalaire = cosinus
        k = min(k, len(self.chunks))
        top = np.argpartition(-scores, k - 1)[:k]
        top = top[np.argsort(-scores[top])]       # tri decroissant
        return [(self.chunks[i], float(scores[i])) for i in top]


def build_store(chunks: list[Chunk], embedder: Embedder) -> VectorStore:
    vectors = embedder.encode([c.text for c in chunks])
    return VectorStore(vectors, chunks, embedder.model_name)