Skip to content
92 linesCodeBlameRaw
1import assert from "node:assert/strict";
2import { test } from "node:test";
3
4import { MAX_IN_FILTER, MEANING_FLOOR, QueryCache, TOP_K, TOP_K_UNFILTERED, fuseRanks, pickPassages, queryKey, recallLimit, requiredSpaces, vectorQueryPlan, type Candidate } from "./recall.ts";
5import { ftsAnyQuery } from "./search.ts";
6
7const c = (id: string, space: string, score: number, by: Candidate["by"] = "meaning"): Candidate => ({ id, doc_id: id.split(":")[0]!, space_id: space, score, by });
8
9test("the index is filtered by the allowed spaces, or by workspace and after when there are too many", () => {
10 assert.equal(vectorQueryPlan("w1", []), null);
11 assert.deepEqual(vectorQueryPlan("w1", ["s1", "s2", "s1"]), { topK: TOP_K, filter: { workspace_id: "w1", space_ids: ["s1", "s2"] }, filterAfter: false });
12 const many = Array.from({ length: MAX_IN_FILTER + 1 }, (_, i) => `s${i}`);
13 assert.deepEqual(vectorQueryPlan("w1", many), { topK: TOP_K_UNFILTERED, filter: { workspace_id: "w1" }, filterAfter: true });
14});
15
16test("required reading never widens what may be read", () => {
17 assert.deepEqual(requiredSpaces(["s1", "s2"], ["s2", "s3", "s2"]), ["s2"]);
18 assert.deepEqual(requiredSpaces(["s1"], null), []);
19 assert.deepEqual(requiredSpaces(["s1"], "s1"), []);
20});
21
22test("limits: five by default, ten at most", () => {
23 assert.equal(recallLimit(undefined), 5);
24 assert.equal(recallLimit(0), 5);
25 assert.equal(recallLimit(3), 3);
26 assert.equal(recallLimit(50), 10);
27 assert.equal(recallLimit("x"), 5);
28});
29
30test("only close enough, only allowed, at most two per page, best first", () => {
31 const picked = pickPassages(
32 [
33 c("p1:0", "s1", 0.81),
34 c("p1:1", "s1", 0.8),
35 c("p1:2", "s1", 0.79),
36 c("p2:0", "s1", 0.7),
37 c("p3:0", "secret", 0.95),
38 c("p4:0", "s1", MEANING_FLOOR - 0.01),
39 ],
40 { allowed: new Set(["s1"]), limit: 10 },
41 );
42 assert.deepEqual(
43 picked.map((p) => p.id),
44 ["p1:0", "p1:1", "p2:0"],
45 );
46});
47
48test("required spaces come first, then the rest", () => {
49 const picked = pickPassages([c("a:0", "s1", 0.9), c("b:0", "s2", 0.75), c("c:0", "s1", 0.85)], { allowed: new Set(["s1", "s2"]), required: ["s2"], limit: 2 });
50 assert.deepEqual(
51 picked.map((p) => p.id),
52 ["b:0", "a:0"],
53 );
54});
55
56test("words fill in only when meaning finds too few, and never twice", () => {
57 const candidates = [c("a:0", "s1", 0.9), c("a:0", "s1", 0.5, "words"), c("b:0", "s1", 0.5, "words"), c("c:0", "s1", 0.5, "words"), c("d:0", "nope", 0.5, "words")];
58 assert.deepEqual(
59 pickPassages(candidates, { allowed: new Set(["s1"]), limit: 3 }).map((p) => `${p.id}/${p.by}`),
60 ["a:0/meaning", "b:0/words", "c:0/words"],
61 );
62 assert.deepEqual(
63 pickPassages(candidates, { allowed: new Set(["s1"]), limit: 1 }).map((p) => p.id),
64 ["a:0"],
65 );
66 // Nothing close and no words: nothing.
67 assert.deepEqual(pickPassages([c("x:0", "s1", 0.3)], { allowed: new Set(["s1"]), limit: 5 }), []);
68});
69
70test("hybrid search fuses ranks: found both ways first", () => {
71 assert.deepEqual(fuseRanks(["a", "b", "c"], ["c", "d"]), ["c", "a", "b", "d"]);
72 assert.deepEqual(fuseRanks([], ["x", "y"]), ["x", "y"]);
73});
74
75test("query embeddings are kept a minute", () => {
76 const cache = new QueryCache(60_000, 2);
77 cache.set("a", [1], 0);
78 assert.deepEqual(cache.get("a", 59_000), [1]);
79 assert.equal(cache.get("a", 61_000), null);
80 cache.set("a", [1], 0);
81 cache.set("b", [2], 1);
82 cache.set("c", [3], 2);
83 assert.equal(cache.get("a", 3), null);
84 assert.deepEqual(cache.get("c", 3), [3]);
85 assert.equal(queryKey(" How do we Deploy? "), "how do we deploy?");
86});
87
88test("recall's word fallback matches any meaningful word", () => {
89 assert.equal(ftsAnyQuery("How do we roll back a deploy?"), '"roll" OR "back" OR "deploy"');
90 assert.equal(ftsAnyQuery("is it a"), null);
91 assert.equal(ftsAnyQuery('NEAR("x") deploy-key'), '"near" OR "deploy" OR "key"');
92});