SPB Git forge

spb/immbot-ai

Public
1commits 1branches 0releases
1.5 MBsize
maindefault branch
20 days agolast push
TypeScript 98.3% CSS 0.9% Shell 0.7%
8.1 KB · 204 lines typescript
Raw Blame History
1// Recherche hybride : FTS5 (BM25) ∥ cosinus vectoriel → fusion RRF → boosts → expansion voisins.2import { all, get } from "../db/index.ts";3import { blobToVec, cosine, embedQuery } from "./embeddings.ts";45export type RetrievedChunk = {6  id: number;7  document_id: number;8  course_code: string | null;9  space: string;10  ref_type: string;11  ref_number: number | null;12  ref_label: string;13  section_title: string;14  title: string;15  content: string;16  box_types: string;17  week: number | null;18  doc_title: string;19  doc_path: string;20  filename: string;21  score: number;22};2324export type SearchOptions = {25  query: string;26  spaces: string[]; // espaces autorisés — TOUJOURS filtrés côté serveur27  conversationId?: number; // pour inclure les téléversements de cette conversation28  ownerUserId?: number;29  k?: number;30};3132const RRF_K = 60;3334// Cache mémoire des vecteurs par clé d'espace (invalidé par le nombre de fragments).35type VecEntry = { id: number; vec: Float32Array };36const vecCache = new Map<string, { count: number; entries: VecEntry[] }>();3738function spacePlaceholders(spaces: string[]): string {39  return spaces.map(() => "?").join(",");40}4142function loadVectors(spaces: string[], conversationId?: number, ownerUserId?: number): VecEntry[] {43  const key = spaces.sort().join("|") + `#${conversationId ?? 0}#${ownerUserId ?? 0}`;44  const where = buildScopeWhere(spaces, conversationId, ownerUserId);45  const countRow = get<{ n: number }>(`SELECT COUNT(*) as n FROM chunks WHERE ${where.sql}`, ...where.params);46  const count = countRow?.n ?? 0;47  const cached = vecCache.get(key);48  if (cached && cached.count === count) return cached.entries;49  const rows = all<{ id: number; embedding: Uint8Array | null }>(50    `SELECT id, embedding FROM chunks WHERE ${where.sql}`,51    ...where.params52  );53  const entries: VecEntry[] = [];54  for (const r of rows) if (r.embedding) entries.push({ id: r.id, vec: blobToVec(r.embedding) });55  vecCache.set(key, { count, entries });56  if (vecCache.size > 24) vecCache.delete(vecCache.keys().next().value as string);57  return entries;58}5960function buildScopeWhere(spaces: string[], conversationId?: number, ownerUserId?: number) {61  // Espaces officiels/privés : filtre simple. Espaces étudiants : restreints au propriétaire/à la conversation.62  const parts: string[] = [];63  const params: unknown[] = [];64  const official = spaces.filter((s) => !s.startsWith("student-"));65  if (official.length) {66    parts.push(`(space IN (${spacePlaceholders(official)}))`);67    params.push(...official);68  }69  if (spaces.includes("student-temporary-upload") && conversationId && ownerUserId) {70    parts.push(`(space = 'student-temporary-upload' AND conversation_id = ? AND owner_user_id = ?)`);71    params.push(conversationId, ownerUserId);72  }73  if (spaces.includes("student-persistent-files") && ownerUserId) {74    parts.push(`(space = 'student-persistent-files' AND owner_user_id = ?)`);75    params.push(ownerUserId);76  }77  if (!parts.length) return { sql: "0", params: [] as unknown[] };78  return { sql: `(${parts.join(" OR ")})`, params };79}8081function ftsQuery(query: string): string {82  const tokens = query83    .toLowerCase()84    .replace(/[«»"'’()\[\]{}:;,!?<>=+*/\\^~`|-]/g, " ")85    .split(/\s+/)86    .filter((t) => t.length > 1)87    .slice(0, 12);88  if (!tokens.length) return '""';89  return tokens.map((t) => `"${t}"`).join(" OR ");90}9192export async function hybridSearch(opts: SearchOptions): Promise<RetrievedChunk[]> {93  const k = opts.k ?? 10;94  const where = buildScopeWhere(opts.spaces, opts.conversationId, opts.ownerUserId);95  if (where.sql === "0") return [];9697  // 1) FTS5 BM2598  const ftsRows = all<{ rowid: number; rank: number }>(99    `SELECT f.rowid as rowid, bm25(chunks_fts, 3.0, 1.0) as rank100     FROM chunks_fts f JOIN chunks c ON c.id = f.rowid101     WHERE chunks_fts MATCH ? AND ${where.sql}102     ORDER BY rank LIMIT 30`,103    ftsQuery(opts.query),104    ...where.params105  );106107  // 2) Vectoriel108  const qvec = await embedQuery(opts.query);109  const entries = loadVectors(opts.spaces, opts.conversationId, opts.ownerUserId);110  const vecScored = entries111    .map((e) => ({ id: e.id, s: cosine(qvec, e.vec) }))112    .sort((a, b) => b.s - a.s)113    .slice(0, 30);114115  // 3) Fusion RRF116  const rrf = new Map<number, number>();117  ftsRows.forEach((r, i) => rrf.set(r.rowid, (rrf.get(r.rowid) ?? 0) + 1 / (RRF_K + i + 1)));118  vecScored.forEach((r, i) => rrf.set(r.id, (rrf.get(r.id) ?? 0) + 1 / (RRF_K + i + 1)));119  if (!rrf.size) return [];120121  const ids = [...rrf.keys()];122  const rows = all<RetrievedChunk>(123    `SELECT c.id, c.document_id, c.course_code, c.space, c.ref_type, c.ref_number, c.ref_label,124            c.section_title, c.title, c.content, c.box_types, c.week,125            d.title as doc_title, d.path as doc_path, d.filename, 0 as score126     FROM chunks c JOIN documents d ON d.id = c.document_id127     WHERE c.id IN (${ids.map(() => "?").join(",")})`,128    ...ids129  );130  const byId = new Map(rows.map((r) => [r.id, r]));131132  // 4) Boosts légers selon le type de question133  const q = opts.query.toLowerCase();134  const wantsDefinition = /\b(défini|definition|qu'est|c'est quoi|signifie)\b/.test(q);135  const wantsFormula = /\b(formule|calcul|comment calculer|équation)\b/.test(q);136  const scored = ids137    .map((id) => {138      const row = byId.get(id);139      if (!row) return null;140      let s = rrf.get(id)!;141      if (wantsDefinition && row.box_types.includes("defbox")) s *= 1.25;142      if (wantsFormula && (row.box_types.includes("importbox") || /formule|=/.test(row.content))) s *= 1.15;143      if (row.ref_type === "glossary") s *= wantsDefinition ? 1.2 : 1.0;144      return { ...row, score: s };145    })146    .filter((r): r is RetrievedChunk => !!r)147    .sort((a, b) => b.score - a.score);148149  // 5) Dédoublonnage (max 3 fragments par document) + équilibre150  const perDoc = new Map<number, number>();151  const selected: RetrievedChunk[] = [];152  for (const r of scored) {153    const n = perDoc.get(r.document_id) ?? 0;154    if (n >= 3) continue;155    perDoc.set(r.document_id, n + 1);156    selected.push(r);157    if (selected.length >= k) break;158  }159160  // 6) Expansion : diapositive suite (1/2 → 2/2) et voisines immédiates du meilleur résultat161  if (selected.length && selected[0].ref_type === "slide" && /\(\d\/\d\)|1\/2|\bsuite\b/i.test(selected[0].title)) {162    const neighbor = get<RetrievedChunk>(163      `SELECT c.id, c.document_id, c.course_code, c.space, c.ref_type, c.ref_number, c.ref_label,164              c.section_title, c.title, c.content, c.box_types, c.week,165              d.title as doc_title, d.path as doc_path, d.filename, 0 as score166       FROM chunks c JOIN documents d ON d.id = c.document_id167       WHERE c.document_id = ? AND c.ref_number = ? AND c.id != ?`,168      selected[0].document_id,169      (selected[0].ref_number ?? 0) + 1,170      selected[0].id171    );172    if (neighbor && !selected.some((s) => s.id === neighbor.id)) selected.push({ ...neighbor, score: selected[0].score * 0.9 });173  }174  return selected.slice(0, k + 2);175}176177/** Seuil de pertinence : si le meilleur score RRF est trop faible, le corpus ne couvre pas la question. */178export function isConfidentEnough(results: RetrievedChunk[]): boolean {179  if (!results.length) return false;180  return results[0].score >= 1 / (RRF_K + 8); // présent dans le top-8 d'au moins une des deux recherches181}182183export type ContextBlock = {184  text: string;185  sources: { tag: string; chunk: RetrievedChunk }[];186};187188/** Construit le contexte numéroté [S1..Sn] transmis au modèle. */189export function buildContext(results: RetrievedChunk[], maxChars = 14000): ContextBlock {190  const sources: { tag: string; chunk: RetrievedChunk }[] = [];191  const parts: string[] = [];192  let total = 0;193  results.forEach((r) => {194    if (total > maxChars) return;195    const tag = `S${sources.length + 1}`;196    const header = `[${tag}] (${r.course_code ?? "téléversement"} — ${r.doc_title}${r.ref_label ? " — " + r.ref_label : ""}${r.title ? " — « " + r.title + " »" : ""})`;197    const body = r.content.slice(0, 2400);198    parts.push(`${header}\n${body}`);199    total += header.length + body.length;200    sources.push({ tag, chunk: r });201  });202  return { text: parts.join("\n\n---\n\n"), sources };203}204