| 1 | import assert from "node:assert/strict"; |
| 2 | import { test } from "node:test"; |
| 3 | |
| 4 | import { MAX_IN_FILTER, MEANING_FLOOR, QueryCache, TOP_K, TOP_K_UNFILTERED, fuseRanks, pickPassages, queryKey, recallLimit, requiredSpaces, vectorQueryPlan, type Candidate } from "./recall.ts"; |
| 5 | import { ftsAnyQuery } from "./search.ts"; |
| 6 | |
| 7 | const c = (id: string, space: string, score: number, by: Candidate["by"] = "meaning"): Candidate => ({ id, doc_id: id.split(":")[0]!, space_id: space, score, by }); |
| 8 | |
| 9 | test("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 | |
| 16 | test("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 | |
| 22 | test("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 | |
| 30 | test("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 | |
| 48 | test("required spaces come first, then the rest", () => { |
| 49 | const picked = pickPassages([c("a:0", "s1", 0.9), c("b:0", "s2", 0.65), 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 | |
| 56 | test("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 | |
| 70 | test("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 | |
| 75 | test("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 | |
| 88 | test("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 | }); |