spb/uqo-chat

1"""Hybrid retriever: BM25 (in-memory index over course_chunks) + optional embeddings."""23from __future__ import annotations45import asyncio6import math7from collections import Counter, defaultdict8from dataclasses import dataclass910import numpy as np11from sqlalchemy import select1213from app.core.logging import get_logger14from app.db import SessionLocal15from app.models import CourseChunk16from app.rag.chunking import tokenize1718log = get_logger("rag.retriever")192021@dataclass22class Hit:23    chunk_id: str24    course: str25    module: str26    section: str27    page: str28    url: str29    content: str30    score: float31    visibility: str3233    def source_label(self) -> str:34        sec = f", {self.section}" if self.section and self.section != self.module else ""35        return f"{self.course} – {self.module}{sec}"363738class BM25Index:39    def __init__(self, k1: float = 1.4, b: float = 0.75) -> None:40        self.k1, self.b = k1, b41        self.docs: list[dict] = []42        self.tf: list[Counter[str]] = []43        self.df: Counter[str] = Counter()44        self.doc_len: list[int] = []45        self.avg_len = 1.046        self.embeddings: np.ndarray | None = None47        self.postings: dict[str, list[int]] = defaultdict(list)4849    def build(self, rows: list[CourseChunk]) -> None:50        self.docs, self.tf, self.doc_len = [], [], []51        self.df = Counter()52        self.postings = defaultdict(list)53        embs: list[list[float]] = []54        have_emb = True55        for i, r in enumerate(rows):56            head = f"{r.module} {r.section} "57            toks = tokenize(head * 2 + r.content)58            tf = Counter(toks)59            self.docs.append({"id": r.id, "course": r.course_code, "module": r.module,60                              "section": r.section, "page": r.page, "url": r.url,61                              "content": r.content, "visibility": r.visibility})62            self.tf.append(tf)63            self.doc_len.append(len(toks))64            for t in tf:65                self.df[t] += 166                self.postings[t].append(i)67            if r.embedding:68                embs.append(r.embedding)69            else:70                have_emb = False71        self.avg_len = (sum(self.doc_len) / len(self.doc_len)) if self.doc_len else 1.072        self.embeddings = np.array(embs, dtype=np.float32) if (have_emb and embs) else None73        log.info("bm25_built", docs=len(self.docs), embeddings=self.embeddings is not None)7475    def _idf(self, t: str) -> float:76        n = len(self.docs)77        df = self.df.get(t, 0)78        return math.log(1 + (n - df + 0.5) / (df + 0.5))7980    def search(self, query: str, top_k: int = 8, courses: set[str] | None = None,81               boost_course: str | None = None, include_professor: bool = False,82               query_embedding: list[float] | None = None) -> list[Hit]:83        q = tokenize(query)84        if not q or not self.docs:85            return []86        scores: dict[int, float] = defaultdict(float)87        for t in set(q):88            idf = self._idf(t)89            for i in self.postings.get(t, ()):90                f = self.tf[i][t]91                denom = f + self.k1 * (1 - self.b + self.b * self.doc_len[i] / self.avg_len)92                scores[i] += idf * f * (self.k1 + 1) / denom93        if scores:94            mx = max(scores.values()) or 1.095            for i in scores:96                scores[i] /= mx97        if query_embedding is not None and self.embeddings is not None:98            qv = np.array(query_embedding, dtype=np.float32)99            sims = self.embeddings @ qv / (np.linalg.norm(self.embeddings, axis=1) *100                                           (np.linalg.norm(qv) or 1.0) + 1e-9)101            top = np.argsort(-sims)[: top_k * 4]102            for i in top:103                scores[int(i)] = 0.35 * scores.get(int(i), 0.0) + 0.65 * float(sims[i])104        hits: list[Hit] = []105        for i, s in scores.items():106            d = self.docs[i]107            if courses and d["course"] not in courses:108                continue109            if d["visibility"] == "professor_only" and not include_professor:110                continue111            if boost_course and d["course"] == boost_course:112                s *= 1.15113            hits.append(Hit(d["id"], d["course"], d["module"], d["section"], d["page"],114                            d["url"], d["content"], s, d["visibility"]))115        hits.sort(key=lambda h: h.score, reverse=True)116        # diversify: at most 3 hits per module117        out: list[Hit] = []118        per_module: Counter[str] = Counter()119        for h in hits:120            if per_module[h.module] >= 3:121                continue122            per_module[h.module] += 1123            out.append(h)124            if len(out) >= top_k:125                break126        return out127128129index = BM25Index()130_lock = asyncio.Lock()131132133async def rebuild_index() -> int:134    async with _lock:135        async with SessionLocal() as session:136            rows = list((await session.execute(select(CourseChunk))).scalars())137        index.build(rows)138        return len(rows)139140141def format_for_model(hits: list[Hit]) -> str:142    if not hits:143        return "Aucun passage pertinent trouvé dans le matériel du cours."144    parts = ["Passages du matériel de cours (cite-les avec l'identifiant entre crochets) :"]145    for i, h in enumerate(hits, 1):146        parts.append(f"[S{i}] Source: {h.source_label()}"147                     f"{' — ' + h.url if h.url else ''}\n<document>\n{h.content[:3000]}\n</document>")148    return "\n\n".join(parts)149