STCourses/packages/feature-ai/test/retrieval.test.ts

import { afterAll, beforeAll, expect, test } from "bun:test";
import { cpSync, mkdtempSync, mkdirSync, readFileSync, 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<string>();
  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("the mentor file and the manifest are not retrievable", async () => {
  const idx = await buildCourseIndex(course, dataDir);
  expect(idx.search("mentor hints never full solutions", 50).some((c) => c.path === "MENTOR.md" || c.path === "stcourse.yaml")).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" });
});

// The importer's layout (feature-import): <dataDir>/cache/decks/<courseId>/<deckId>/text.json ([{ page, text }])
// and <dataDir>/cache/transcripts/<courseId>/<lessonId>.json ({ lessonId, text, segments: [{ start, end, text }] }).
test("reads the importer's deck text.json and transcript cache", async () => {
  const dir = mkdtempSync(join(tmpdir(), "stc-ai-imp-"));
  try {
    const courses = join(dir, "courses");
    cpSync(join(FIXTURES, "sample"), join(courses, "sample"), { recursive: true });
    const yaml = join(courses, "sample/stcourse.yaml");
    writeFileSync(yaml, readFileSync(yaml, "utf8").replace(", transcript: lessons/02-video.md", ""));
    const c = (await createCourseRepository([courses])).get("sample")!;
    expect(c.status).toBe("ok");
    const data = join(dir, "data");
    mkdirSync(join(data, "cache/decks/sample/d1"), { recursive: true });
    writeFileSync(join(data, "cache/decks/sample/d1/text.json"), JSON.stringify([
      { page: 1, text: "Title slide" },
      { page: 2, text: "Orchestration with Kubernetes pods and replica sets." },
      { page: 3, text: "" },
    ]));
    mkdirSync(join(data, "cache/transcripts/sample"), { recursive: true });
    writeFileSync(join(data, "cache/transcripts/sample/m1-video.json"), JSON.stringify({
      lessonId: "m1-video",
      text: "Welcome. Today the presenter shows a waveform monitor.",
      segments: [{ start: 0, end: 4, text: "Welcome." }, { start: 65, end: 70, text: "Today the presenter shows a waveform monitor." }],
    }));
    const idx = await buildCourseIndex(c, data);
    const deck = idx.search("Kubernetes replica sets", 3)[0]!;
    expect(deck).toMatchObject({ kind: "deck", path: "media/deck.pdf", heading: "Slide 2", lessonId: "m1-video" });
    const tr = idx.search("waveform monitor", 3)[0]!;
    expect(tr).toMatchObject({ kind: "transcript", lessonId: "m1-video", lessonTitle: "Watch the clip" });
    expect(tr.text).toContain("waveform monitor");
  } finally {
    rmSync(dir, { recursive: true, force: true });
  }
});

test("a cached transcript is skipped when the lesson already has a transcript file", async () => {
  const data = mkdtempSync(join(tmpdir(), "stc-ai-tr-"));
  try {
    mkdirSync(join(data, "cache/transcripts/sample"), { recursive: true });
    writeFileSync(join(data, "cache/transcripts/sample/m1-video.json"), JSON.stringify({ lessonId: "m1-video", text: "zebracorn unique words", segments: [] }));
    const idx = await buildCourseIndex(course, data);
    expect(idx.search("zebracorn", 3)).toEqual([]);
  } finally {
    rmSync(data, { recursive: true, force: true });
  }
});