From 50e1e49b6dd9413ab2eeacaea766ae2130b1d34b Mon Sep 17 00:00:00 2001 From: Lucas Winther Date: Fri, 30 Oct 2026 03:18:21 +0100 Subject: [PATCH] feat(ai): BM25 course retrieval over lessons, transcripts, code, deck text and notes Co-Authored-By: Claude Opus 5.5 --- packages/feature-ai/src/index.ts | 2 + packages/feature-ai/src/retrieval/bm25.ts | 79 ++++++ .../feature-ai/src/retrieval/course-index.ts | 236 ++++++++++++++++++ packages/feature-ai/test/retrieval.test.ts | 78 ++++++ 4 files changed, 395 insertions(+) create mode 100644 packages/feature-ai/src/retrieval/bm25.ts create mode 100644 packages/feature-ai/src/retrieval/course-index.ts create mode 100644 packages/feature-ai/test/retrieval.test.ts diff --git a/packages/feature-ai/src/index.ts b/packages/feature-ai/src/index.ts index 1bc5947..9489187 100644 --- a/packages/feature-ai/src/index.ts +++ b/packages/feature-ai/src/index.ts @@ -4,3 +4,5 @@ export { createSecretBox, type SecretBox } from "./crypto"; export const migrationsDirs: string[] = []; export * from "./settings"; export * from "./providers"; +export * from "./retrieval/course-index"; +export { Bm25, tokenize } from "./retrieval/bm25"; diff --git a/packages/feature-ai/src/retrieval/bm25.ts b/packages/feature-ai/src/retrieval/bm25.ts new file mode 100644 index 0000000..601ac06 --- /dev/null +++ b/packages/feature-ai/src/retrieval/bm25.ts @@ -0,0 +1,79 @@ +const STOP = new Set( + "a an and are as at be but by for from has have i if in into is it its of on or so that the their then there these this to was were what when which who why will with you your".split(" "), +); + +/** Light plural folding so "apples" matches "apple". */ +const stem = (w: string): string => + w.length > 4 && w.endsWith("ies") ? `${w.slice(0, -3)}y` : w.length > 3 && w.endsWith("s") && !w.endsWith("ss") ? w.slice(0, -1) : w; + +/** Lowercased word tokens without stopwords; camelCase and snake_case identifiers also add their parts. */ +export function tokenize(text: string): string[] { + const out: string[] = []; + for (const raw of text.match(/[\p{L}\p{N}_]+/gu) ?? []) { + const lower = raw.toLowerCase(); + if (STOP.has(lower) || lower.length > 64) continue; + out.push(stem(lower)); + const parts = raw.split(/_|(?<=[\p{Ll}\p{N}])(?=\p{Lu})/u).filter(Boolean); + if (parts.length > 1) for (const p of parts) if (p.length > 1) out.push(stem(p.toLowerCase())); + } + return out; +} + +type Doc = { item: T; tf: Map; len: number }; + +const termFreq = (tokens: string[]) => { + const tf = new Map(); + for (const t of tokens) tf.set(t, (tf.get(t) ?? 0) + 1); + return tf; +}; + +/** Okapi BM25 (k1 = 1.2, b = 0.75) over an in-memory document list. */ +export class Bm25 { + private docs: Doc[] = []; + private df = new Map(); + private totalLen = 0; + + add(text: string, item: T): void { + const tokens = tokenize(text); + const tf = termFreq(tokens); + for (const t of tf.keys()) this.df.set(t, (this.df.get(t) ?? 0) + 1); + this.docs.push({ item, tf, len: tokens.length }); + this.totalLen += tokens.length; + } + + get size(): number { return this.docs.length; } + + private idf(term: string): number { + const n = this.docs.length; + const df = this.df.get(term) ?? 0; + return Math.log(1 + (n - df + 0.5) / (df + 0.5)); + } + + private score(q: string[], tf: Map, len: number): number { + const avg = this.docs.length ? this.totalLen / this.docs.length : Math.max(1, len); + let s = 0; + for (const term of q) { + const f = tf.get(term); + if (!f) continue; + s += this.idf(term) * ((f * 2.2) / (f + 1.2 * (0.25 + (0.75 * len) / avg))); + } + return s; + } + + /** Top-k matches. `extra` documents are scored with this index's statistics without being added to it. */ + search(query: string, k: number, extra: { text: string; item: T }[] = []): { item: T; score: number }[] { + const q = [...new Set(tokenize(query))]; + if (!q.length) return []; + const results: { item: T; score: number }[] = []; + for (const d of this.docs) { + const s = this.score(q, d.tf, d.len); + if (s > 0) results.push({ item: d.item, score: s }); + } + for (const e of extra) { + const tokens = tokenize(e.text); + const s = this.score(q, termFreq(tokens), tokens.length); + if (s > 0) results.push({ item: e.item, score: s }); + } + return results.sort((a, b) => b.score - a.score).slice(0, k); + } +} diff --git a/packages/feature-ai/src/retrieval/course-index.ts b/packages/feature-ai/src/retrieval/course-index.ts new file mode 100644 index 0000000..616c498 --- /dev/null +++ b/packages/feature-ai/src/retrieval/course-index.ts @@ -0,0 +1,236 @@ +import { readdir, stat } from "node:fs/promises"; +import { basename, extname, join, relative, sep } from "node:path"; +import { slugify, type LoadedCourse } from "@stc/course-format"; +import { Bm25 } from "./bm25"; + +export type ChunkKind = "lesson" | "transcript" | "markdown" | "code" | "deck" | "note"; +export interface Chunk { + path: string; // course-relative path (for cache text: the deck or lesson it belongs to) + kind: ChunkKind; + text: string; + heading?: string; + lessonId?: string; + lessonTitle?: string; +} +export interface CourseIndex { size: number; search(query: string, k: number, extra?: Chunk[]): Chunk[] } + +const SKIP_DIRS = new Set([".git", "node_modules", "dist", "build", "target", ".astro", ".next", "coverage", "bin", "obj", ".venv", "__pycache__"]); +const MD = new Set([".md", ".markdown"]); +const CODE = new Set([".js", ".mjs", ".cjs", ".ts", ".tsx", ".jsx", ".py", ".rs", ".scala", ".sc", ".fs", ".fsx", ".go", ".java", ".kt", ".rb", ".php", ".c", ".h", ".cpp", ".cs", ".sh", ".bash", ".ps1", ".sql", ".yaml", ".yml", ".toml", ".json", ".tf", ".html", ".css"]); +const CODE_NAMES = new Set(["Dockerfile", "Makefile", "docker-compose.yml", "Containerfile"]); +const CACHE_TEXT = new Set([".txt", ".md", ".vtt", ".srt"]); +const SKIP_FILES = new Set(["package-lock.json", "bun.lock", "yarn.lock", "pnpm-lock.yaml", "Cargo.lock"]); +const MAX_FILE = 200_000; +const MAX_FILES = 2_000; +const MAX_CHUNK = 1_500; +const SOLUTION_SEGMENTS = new Set(["solution", "solutions", "reference"]); + +const posix = (p: string) => p.split(sep).join("/"); + +/** + * Files only authors may see: manifest `browser.solution` files and anything under a `solution`, `solutions` + * or `reference` folder (Pure Snake `js/solutions/`, Systemic Design `reference/` and `expansion/solution/`). + * The assistant never retrieves or reads them. + */ +export function isSolutionPath(course: LoadedCourse, rel: string): boolean { + const p = posix(rel).replace(/^\.\//, ""); + if (p.split("/").slice(0, -1).some((s) => SOLUTION_SEGMENTS.has(s.toLowerCase()))) return true; + if (/(^|\/)solution\.[a-z]+$/i.test(p)) return true; + for (const m of course.manifest?.modules ?? []) for (const l of m.lessons) if (l.browser?.solution && posix(l.browser.solution) === p) return true; + return false; +} + +async function walk(root: string, onFile: (abs: string, rel: string) => Promise, depth = 0, base = root): Promise { + if (depth > 8) return true; + let entries; + try { entries = await readdir(root, { withFileTypes: true }); } catch { return true; } + entries.sort((a, b) => a.name.localeCompare(b.name)); + for (const e of entries) { + if (e.name.startsWith(".") && e.name !== ".") continue; + const abs = join(root, e.name); + if (e.isDirectory()) { + if (SKIP_DIRS.has(e.name)) continue; + if (!(await walk(abs, onFile, depth + 1, base))) return false; + } else if (e.isFile()) { + if (!(await onFile(abs, posix(relative(base, abs))))) return false; + } + } + return true; +} + +/** Splits Markdown at headings (outside code fences), then long sections at paragraph breaks. */ +export function splitMarkdown(text: string): { heading?: string; text: string }[] { + const sections: { heading?: string; lines: string[] }[] = [{ lines: [] }]; + let fence = false; + for (const line of text.split("\n")) { + if (/^\s*(```|~~~)/.test(line)) fence = !fence; + const h = !fence && /^(#{1,6})\s+(.*)$/.exec(line); + if (h) sections.push({ heading: h[2]!.replace(/#+\s*$/, "").trim(), lines: [line] }); + else sections[sections.length - 1]!.lines.push(line); + } + const out: { heading?: string; text: string }[] = []; + for (const s of sections) { + const body = s.lines.join("\n").trim(); + if (!body) continue; + out.push(...splitLong(body).map((t) => ({ heading: s.heading, text: t }))); + } + return out; +} + +function splitLong(text: string): string[] { + if (text.length <= MAX_CHUNK) return [text]; + const out: string[] = []; + let cur = ""; + for (const para of text.split(/\n{2,}/)) { + if (cur && cur.length + para.length > MAX_CHUNK) { out.push(cur); cur = ""; } + if (para.length > MAX_CHUNK) { + for (let i = 0; i < para.length; i += MAX_CHUNK) out.push(para.slice(i, i + MAX_CHUNK)); + } else cur = cur ? `${cur}\n\n${para}` : para; + } + if (cur) out.push(cur); + return out; +} + +type LessonRef = { lessonId: string; lessonTitle: string; anchor?: string; transcript: boolean }; + +function lessonRefs(course: LoadedCourse): Map { + const refs = new Map(); + const add = (p: string | undefined, r: LessonRef) => { + if (!p) return; + const key = posix(p).replace(/^\.\//, ""); + refs.set(key, [...(refs.get(key) ?? []), r]); + }; + for (const m of course.manifest?.modules ?? []) { + for (const l of m.lessons) { + const base = { lessonId: l.id, lessonTitle: l.title }; + if (l.type !== "video") add(l.file, { ...base, anchor: l.anchor, transcript: false }); + add(l.transcript, { ...base, transcript: true }); + if (l.hints?.file && l.hints.file !== l.file) add(l.hints.file, { ...base, transcript: false }); + for (const t of Object.values(l.tracks ?? {})) add(t.file, { ...base, transcript: false }); + } + } + return refs; +} + +function pickLesson(refs: LessonRef[] | undefined, heading?: string): LessonRef | undefined { + if (!refs?.length) return undefined; + if (heading) { + const slug = slugify(heading); + const hit = refs.find((r) => r.anchor && r.anchor === slug); + if (hit) return hit; + } + return refs.find((r) => !r.anchor) ?? refs[0]; +} + +function lessonTitleOf(course: LoadedCourse, lessonId: string | undefined): string | undefined { + return lessonId ? course.lessonIndex.get(lessonId)?.lesson.title : undefined; +} + +export async function buildCourseIndex(course: LoadedCourse, dataDir: string): Promise { + const bm = new Bm25(); + const seen = new Set(); + const refs = lessonRefs(course); + const codeDirs = (course.manifest?.modules ?? []).flatMap((m) => (m.code ? [posix(m.code.dir).replace(/\/$/, "") + "/"] : [])); + const add = (c: Chunk) => { + const key = `${c.kind}\u0000${c.text}`; + if (seen.has(key)) return; + seen.add(key); + bm.add([c.heading ?? "", c.lessonTitle ?? "", c.path, c.text].join("\n"), c); + }; + + let files = 0; + if (course.status === "ok") { + await walk(course.root, async (abs, rel) => { + if (files >= MAX_FILES) return false; + const ext = extname(rel).toLowerCase(); + const name = basename(rel); + const isMd = MD.has(ext); + const isCode = !isMd && (CODE.has(ext) || CODE_NAMES.has(name) || codeDirs.some((d) => rel.startsWith(d))); + if ((!isMd && !isCode) || SKIP_FILES.has(name) || isSolutionPath(course, rel)) return true; + let size = 0; + try { size = (await stat(abs)).size; } catch { return true; } + if (size > MAX_FILE || size === 0) return true; + const text = await Bun.file(abs).text(); + if (isCode && text.includes("\u0000")) return true; + files++; + const r = refs.get(rel); + if (isMd) { + for (const part of splitMarkdown(text)) { + const ref = pickLesson(r, part.heading); + add({ + path: rel, + kind: ref ? (ref.transcript ? "transcript" : "lesson") : "markdown", + text: part.text, + heading: part.heading, + lessonId: ref?.lessonId, + lessonTitle: ref?.lessonTitle, + }); + } + } else { + const lines = text.split("\n"); + for (let i = 0; i < lines.length; i += 60) { + add({ path: rel, kind: "code", text: lines.slice(i, i + 60).join("\n"), heading: lines.length > 60 ? `lines ${i + 1}-${Math.min(lines.length, i + 60)}` : undefined }); + } + } + return true; + }); + } + + // Derived text from the importer and transcription jobs: /cache//**. + const decks = new Map((course.manifest?.decks ?? []).map((d) => [d.id, d])); + const cacheRoot = join(dataDir, "cache", course.id); + await walk(cacheRoot, async (abs, rel) => { + if (!CACHE_TEXT.has(extname(rel).toLowerCase())) return true; + const size = (await stat(abs).catch(() => null))?.size ?? 0; + if (!size || size > MAX_FILE * 5) return true; + const text = await Bun.file(abs).text(); + const segs = rel.split("/"); + const deck = segs.map((s) => decks.get(s)).find(Boolean); + const lessonSeg = segs.map((s) => s.replace(/\.[^.]+$/, "")).find((s) => course.lessonIndex.has(s)); + const page = /(?:page|slide)[-_ ]?(\d+)/i.exec(basename(rel))?.[1]; + if (deck) { + const lessonId = course.lessonIndex.has(deck.attachTo) ? deck.attachTo : undefined; + for (const part of splitLong(text.trim())) { + add({ path: posix(deck.file), kind: "deck", text: part, heading: page ? `Slide ${Number(page)}` : deck.title, lessonId, lessonTitle: lessonTitleOf(course, lessonId) }); + } + } else { + const kind: ChunkKind = /transcript|\.vtt$|\.srt$/i.test(rel) || lessonSeg ? "transcript" : "deck"; + for (const part of splitLong(text.trim())) { + add({ path: `cache/${rel}`, kind, text: part, lessonId: lessonSeg, lessonTitle: lessonTitleOf(course, lessonSeg) }); + } + } + return true; + }); + + return { + get size() { return bm.size; }, + search(query, k, extra = []) { + return bm.search(query, k, extra.map((c) => ({ text: [c.heading ?? "", c.lessonTitle ?? "", c.text].join("\n"), item: c }))).map((r) => r.item); + }, + }; +} + +/** The user's notes as chunks, scored at query time and never cached. */ +export function noteChunks(course: LoadedCourse, notes: { id: string; lessonId: string; body: string; anchorKind: string; anchorValue: string | null }[]): Chunk[] { + return notes.map((n) => ({ + path: "notes", + kind: "note" as const, + text: n.body, + heading: n.anchorKind === "video" && n.anchorValue ? `at ${n.anchorValue}s` : n.anchorKind === "slide" && n.anchorValue ? `slide ${n.anchorValue}` : undefined, + lessonId: n.lessonId, + lessonTitle: lessonTitleOf(course, n.lessonId), + })); +} + +const TTL_MS = 5 * 60_000; +const cache = new WeakMap }>(); + +/** Cached per LoadedCourse object: a repository reload yields new objects, and entries expire after 5 minutes. */ +export function getCourseIndex(course: LoadedCourse, dataDir: string): Promise { + const hit = cache.get(course); + if (hit && Date.now() - hit.at < TTL_MS) return hit.index; + const index = buildCourseIndex(course, dataDir); + cache.set(course, { at: Date.now(), index }); + index.catch(() => cache.delete(course)); + return index; +} diff --git a/packages/feature-ai/test/retrieval.test.ts b/packages/feature-ai/test/retrieval.test.ts new file mode 100644 index 0000000..de15aef --- /dev/null +++ b/packages/feature-ai/test/retrieval.test.ts @@ -0,0 +1,78 @@ +import { afterAll, beforeAll, expect, test } from "bun:test"; +import { mkdtempSync, mkdirSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { createCourseRepository, type LoadedCourse } from "@stc/course-format"; +import { Bm25, tokenize } from "../src/retrieval/bm25"; +import { buildCourseIndex, getCourseIndex, isSolutionPath, noteChunks } from "../src/retrieval/course-index"; + +const FIXTURES = join(import.meta.dir, "../../../fixtures/courses"); +let course: LoadedCourse; +let dataDir: string; + +beforeAll(async () => { + const repo = await createCourseRepository([FIXTURES]); + course = repo.get("sample")!; + dataDir = mkdtempSync(join(tmpdir(), "stc-ai-")); + mkdirSync(join(dataDir, "cache/sample/decks/d1"), { recursive: true }); + writeFileSync(join(dataDir, "cache/sample/decks/d1/page-3.txt"), "Container images are layered filesystems built from a Dockerfile."); + mkdirSync(join(dataDir, "cache/sample/transcripts"), { recursive: true }); + writeFileSync(join(dataDir, "cache/sample/transcripts/m1-video.txt"), "In this clip the presenter explains the test pattern calibration."); +}); +afterAll(() => rmSync(dataDir, { recursive: true, force: true })); + +test("tokenize lowercases, drops stopwords and splits identifiers", () => { + expect(tokenize("The inBounds function, and dropLast!")).toEqual(["inbound", "in", "bound", "function", "droplast", "drop", "last"]); + expect(tokenize("Æble på dansk")).toEqual(["æble", "på", "dansk"]); +}); + +test("BM25 ranks the more relevant document first", () => { + const b = new Bm25(); + b.add("snake moves one cell per tick", "a"); + b.add("apples spawn at random free cells", "b"); + b.add("the snake grows when it eats an apple apple", "c"); + expect(b.search("apple", 5).map((r) => r.item)).toEqual(["c", "b"]); + expect(b.search("zebra", 5)).toEqual([]); +}); + +test("finds lesson text and maps it to a lesson", async () => { + const idx = await buildCourseIndex(course, dataDir); + const hits = idx.search("sum of its arguments", 3); + expect(hits[0]!.path).toBe("lessons/04-exercise.md"); + expect(["m2-text", "m2-ex-add"]).toContain(hits[0]!.lessonId!); + expect(hits[0]!.lessonTitle).toBeTruthy(); +}); + +test("indexes code, transcripts and deck text from the cache", async () => { + const idx = await buildCourseIndex(course, dataDir); + const code = idx.search("hello.sh", 5).find((c) => c.kind === "code"); + expect(code?.path).toBe("code/S01 - First steps/hello.sh"); + const deck = idx.search("Dockerfile layered images", 3)[0]!; + expect(deck.kind).toBe("deck"); + expect(deck.path).toBe("media/deck.pdf"); + expect(deck.heading).toBe("Slide 3"); + expect(deck.lessonId).toBe("m1-video"); + const tr = idx.search("calibration presenter", 3)[0]!; + expect(tr.kind).toBe("transcript"); + expect(tr.lessonId).toBe("m1-video"); +}); + +test("solution files are never indexed", async () => { + const idx = await buildCourseIndex(course, dataDir); + for (const q of ["add", "return a + b", "export const add"]) { + expect(idx.search(q, 50).some((c) => c.path.includes("solution"))).toBe(false); + } + expect(isSolutionPath(course, "exercises/add/solution.js")).toBe(true); + expect(isSolutionPath(course, "js/solutions/rules.js")).toBe(true); + expect(isSolutionPath(course, "course/01-handles/reference/ts/src/engine.ts")).toBe(true); + expect(isSolutionPath(course, "course/01-handles/expansion/solution/x.rs")).toBe(true); + expect(isSolutionPath(course, "exercises/add/starter.js")).toBe(false); +}); + +test("notes are searched alongside the course", async () => { + const idx = await getCourseIndex(course, dataDir); + expect(await getCourseIndex(course, dataDir)).toBe(idx); + const extra = noteChunks(course, [{ id: "n1", lessonId: "m1-intro", body: "Remember: zebras are striped horses", anchorKind: "none", anchorValue: null }]); + const hits = idx.search("zebras striped", 3, extra); + expect(hits[0]).toMatchObject({ kind: "note", lessonId: "m1-intro", lessonTitle: "Introduction" }); +});