@vellumai/assistant 0.7.2 → 0.7.3

package/src/memory/v2/__tests__/reranker.test.ts

@@ -0,0 +1,338 @@
+/**
+ * Tests for `memory/v2/reranker.ts` — public `rerankCandidates` function.
+ *
+ * Mocks the underlying `LocalRerankBackend` and the `readPage` page reader so
+ * the test is hermetic (no subprocess, no filesystem). Verifies the public
+ * contract: scores keyed by slug, fail-open on backend failure, page-read
+ * failures drop slugs silently, LRU cache hits skip the backend.
+ */
+import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
+
+import { makeMockLogger } from "../../../__tests__/helpers/mock-logger.js";
+import type { AssistantConfig } from "../../../config/types.js";
+
+mock.module("../../../util/logger.js", () => ({
+  getLogger: () => makeMockLogger(),
+}));
+
+mock.module("../../../util/platform.js", () => ({
+  getWorkspaceDir: () => "/tmp/test-workspace",
+}));
+
+const backendState = {
+  scores: [] as number[],
+  shouldThrow: false,
+  calls: [] as Array<{ queries: string[]; passages: string[] }>,
+};
+mock.module("../../rerank-local.js", () => ({
+  getOrCreateRerankBackend: (_model: string, _dtype: string) => ({
+    score: async (queries: string[], passages: string[]): Promise<number[]> => {
+      backendState.calls.push({
+        queries: [...queries],
+        passages: [...passages],
+      });
+      if (backendState.shouldThrow) throw new Error("backend down");
+      return backendState.scores.slice(0, passages.length);
+    },
+  }),
+}));
+
+const pageState = {
+  pages: new Map<string, { body: string } | null>(),
+  failingSlugs: new Set<string>(),
+};
+// Partial mock — Bun's `mock.module` is process-wide, so we re-export every
+// real symbol and override only `readPage`. Without this, sibling test files
+// that import `listPages` etc. would crash with "Export not found".
+const realPageStore = await import("../page-store.js");
+mock.module("../page-store.js", () => ({
+  ...realPageStore,
+  readPage: async (_dir: string, slug: string) => {
+    if (pageState.failingSlugs.has(slug)) {
+      throw new Error("read failure");
+    }
+    return pageState.pages.get(slug) ?? null;
+  },
+}));
+
+const { rerankCandidates, _resetRerankCacheForTests } =
+  await import("../reranker.js");
+
+function configWithModel(model = "test-model"): AssistantConfig {
+  return {
+    memory: {
+      v2: {
+        rerank: {
+          model,
+          enabled: true,
+          top_k: 50,
+          alpha: 0.3,
+          dtype: "q8",
+        },
+      },
+    },
+  } as unknown as AssistantConfig;
+}
+
+function resetState() {
+  backendState.scores = [];
+  backendState.shouldThrow = false;
+  backendState.calls.length = 0;
+  pageState.pages.clear();
+  pageState.failingSlugs.clear();
+  _resetRerankCacheForTests();
+}
+
+beforeEach(resetState);
+afterEach(resetState);
+
+/**
+ * Convenience: run `rerankCandidates` with a single query and unwrap the
+ * returned array. Most tests below assert the contract for the
+ * legacy-equivalent single-query path.
+ */
+async function rerankSingle(
+  query: string,
+  candidates: readonly string[],
+  config: AssistantConfig,
+): Promise<Map<string, number>> {
+  const out = await rerankCandidates([query], candidates, config);
+  return out[0] ?? new Map();
+}
+
+describe("rerankCandidates", () => {
+  test("returns empty maps for empty candidates", async () => {
+    const out = await rerankCandidates(["query"], [], configWithModel());
+    expect(out).toHaveLength(1);
+    expect(out[0].size).toBe(0);
+    expect(backendState.calls).toHaveLength(0);
+  });
+
+  test("returns empty array when no queries are passed", async () => {
+    const out = await rerankCandidates([], ["a"], configWithModel());
+    expect(out).toHaveLength(0);
+    expect(backendState.calls).toHaveLength(0);
+  });
+
+  test("returns empty map for whitespace-only query", async () => {
+    pageState.pages.set("a", { body: "content" });
+    const out = await rerankSingle("   ", ["a"], configWithModel());
+    expect(out.size).toBe(0);
+    expect(backendState.calls).toHaveLength(0);
+  });
+
+  test("scores returned keyed by slug, in [0, 1]", async () => {
+    pageState.pages.set("a", { body: "first paragraph of a" });
+    pageState.pages.set("b", { body: "first paragraph of b" });
+    backendState.scores = [0.9, 0.1];
+
+    const out = await rerankSingle("query", ["a", "b"], configWithModel());
+
+    expect(out.get("a")).toBe(0.9);
+    expect(out.get("b")).toBe(0.1);
+  });
+
+  test("clamps scores to [0, 1]", async () => {
+    pageState.pages.set("a", { body: "x" });
+    pageState.pages.set("b", { body: "x" });
+    backendState.scores = [1.5, -0.2];
+
+    const out = await rerankSingle("query", ["a", "b"], configWithModel());
+
+    expect(out.get("a")).toBe(1);
+    expect(out.get("b")).toBe(0);
+  });
+
+  test("drops slugs whose page failed to read; others present", async () => {
+    pageState.pages.set("a", { body: "x" });
+    pageState.failingSlugs.add("b");
+    pageState.pages.set("c", { body: "y" });
+    backendState.scores = [0.5, 0.7];
+
+    const out = await rerankSingle("query", ["a", "b", "c"], configWithModel());
+
+    expect(out.has("b")).toBe(false);
+    expect(out.get("a")).toBe(0.5);
+    expect(out.get("c")).toBe(0.7);
+  });
+
+  test("drops slugs whose page is null (missing on disk)", async () => {
+    pageState.pages.set("a", { body: "x" });
+    pageState.pages.set("missing", null);
+    backendState.scores = [0.5];
+
+    const out = await rerankSingle(
+      "query",
+      ["a", "missing"],
+      configWithModel(),
+    );
+
+    expect(out.size).toBe(1);
+    expect(out.get("a")).toBe(0.5);
+    expect(out.has("missing")).toBe(false);
+  });
+
+  test("returns empty map when backend throws (fail-open)", async () => {
+    pageState.pages.set("a", { body: "x" });
+    backendState.shouldThrow = true;
+
+    const out = await rerankSingle("query", ["a"], configWithModel());
+
+    expect(out.size).toBe(0);
+  });
+
+  test("returns empty map when no pages load (no backend call)", async () => {
+    pageState.failingSlugs.add("a");
+
+    const out = await rerankSingle("query", ["a"], configWithModel());
+
+    expect(out.size).toBe(0);
+    expect(backendState.calls).toHaveLength(0);
+  });
+
+  test("LRU cache hit skips the backend on identical inputs", async () => {
+    pageState.pages.set("a", { body: "x" });
+    backendState.scores = [0.7];
+
+    const first = await rerankSingle("query", ["a"], configWithModel());
+    const second = await rerankSingle("query", ["a"], configWithModel());
+
+    expect(first.get("a")).toBe(0.7);
+    expect(second.get("a")).toBe(0.7);
+    // Backend called only once — second call hit the cache.
+    expect(backendState.calls).toHaveLength(1);
+  });
+
+  test("cache key insensitive to candidate order", async () => {
+    pageState.pages.set("a", { body: "x" });
+    pageState.pages.set("b", { body: "y" });
+    backendState.scores = [0.5, 0.6];
+
+    await rerankSingle("query", ["a", "b"], configWithModel());
+    await rerankSingle("query", ["b", "a"], configWithModel());
+
+    // Same query, same set of candidates — second call hits cache.
+    expect(backendState.calls).toHaveLength(1);
+  });
+
+  test("passage construction caps at 240 chars after slug newline", async () => {
+    const longBody = "a".repeat(500);
+    pageState.pages.set("slug", { body: longBody });
+    backendState.scores = [0.5];
+
+    await rerankSingle("q", ["slug"], configWithModel());
+
+    expect(backendState.calls).toHaveLength(1);
+    const passage = backendState.calls[0].passages[0];
+    // "slug\n" prefix + 240 chars of body
+    expect(passage.startsWith("slug\n")).toBe(true);
+    expect(passage.length).toBeLessThanOrEqual(5 + 240);
+  });
+
+  test("first paragraph is taken (body truncated at blank line)", async () => {
+    pageState.pages.set("slug", {
+      body: "first para line\n\nsecond para should not appear",
+    });
+    backendState.scores = [0.5];
+
+    await rerankSingle("q", ["slug"], configWithModel());
+
+    const passage = backendState.calls[0].passages[0];
+    expect(passage).toContain("first para line");
+    expect(passage).not.toContain("second para");
+  });
+
+  test("multiple queries are batched into one backend call", async () => {
+    pageState.pages.set("a", { body: "x" });
+    pageState.pages.set("b", { body: "y" });
+    // Two queries × two passages = four scores, query-major:
+    //   q1×a, q1×b, q2×a, q2×b
+    backendState.scores = [0.9, 0.1, 0.2, 0.8];
+
+    const out = await rerankCandidates(
+      ["user-text", "assistant-text"],
+      ["a", "b"],
+      configWithModel(),
+    );
+
+    expect(out).toHaveLength(2);
+    // Backend invoked exactly once with both queries' pairs.
+    expect(backendState.calls).toHaveLength(1);
+    expect(backendState.calls[0].queries).toEqual([
+      "user-text",
+      "user-text",
+      "assistant-text",
+      "assistant-text",
+    ]);
+    expect(backendState.calls[0].passages).toHaveLength(4);
+
+    expect(out[0].get("a")).toBe(0.9);
+    expect(out[0].get("b")).toBeCloseTo(0.1, 6);
+    expect(out[1].get("a")).toBe(0.2);
+    expect(out[1].get("b")).toBe(0.8);
+  });
+
+  test("partial cache hit skips the backend for the cached query", async () => {
+    pageState.pages.set("a", { body: "x" });
+    // Prime the cache with q1.
+    backendState.scores = [0.7];
+    await rerankCandidates(["q1"], ["a"], configWithModel());
+    expect(backendState.calls).toHaveLength(1);
+
+    // Now request both q1 (cached) and q2 (fresh). The backend should see
+    // only q2's pair, not q1's.
+    backendState.scores = [0.4];
+    const out = await rerankCandidates(["q1", "q2"], ["a"], configWithModel());
+
+    expect(backendState.calls).toHaveLength(2);
+    expect(backendState.calls[1].queries).toEqual(["q2"]);
+    expect(out[0].get("a")).toBe(0.7);
+    expect(out[1].get("a")).toBe(0.4);
+  });
+
+  test("forwards configured dtype to the backend factory", async () => {
+    pageState.pages.set("a", { body: "x" });
+    backendState.scores = [0.5];
+    const dtypes: string[] = [];
+
+    // Re-mock the factory just for this test to capture the dtype arg.
+    mock.module("../../rerank-local.js", () => ({
+      getOrCreateRerankBackend: (_model: string, dtype: string) => {
+        dtypes.push(dtype);
+        return {
+          score: async (
+            queries: string[],
+            passages: string[],
+          ): Promise<number[]> => {
+            backendState.calls.push({
+              queries: [...queries],
+              passages: [...passages],
+            });
+            return backendState.scores.slice(0, passages.length);
+          },
+        };
+      },
+    }));
+    const { rerankCandidates: freshRerank, _resetRerankCacheForTests: reset } =
+      await import("../reranker.js");
+    reset();
+
+    const config = {
+      memory: {
+        v2: {
+          rerank: {
+            model: "test-model",
+            enabled: true,
+            top_k: 50,
+            alpha: 0.3,
+            dtype: "fp32",
+          },
+        },
+      },
+    } as unknown as AssistantConfig;
+
+    await freshRerank(["q"], ["a"], config);
+
+    expect(dtypes).toContain("fp32");
+  });
+});

package/src/memory/v2/__tests__/sim.test.ts

@@ -70,16 +70,6 @@ const state = {
       points: Array<{ score?: number; payload: Record<string, unknown> }>;
     }>,
   },
-  // Separate response queue for the dedicated skills collection so a test
-  // querying both concept pages and skills doesn't have to interleave.
-  skillQueryResponses: {
-    dense: [] as Array<{
-      points: Array<{ score?: number; payload: Record<string, unknown> }>;
-    }>,
-    sparse: [] as Array<{
-      points: Array<{ score?: number; payload: Record<string, unknown> }>;
-    }>,
-  },
   queryCalls: [] as Array<{
     collection: string;
     using: string;
@@ -147,11 +137,7 @@ class MockQdrantClient {
       filter: params.filter,
     });
     const channel = params.using as "dense" | "sparse";
-    const queue =
-      name === "memory_v2_skills"
-        ? state.skillQueryResponses[channel]
-        : state.queryResponses[channel];
-    return queue.shift() ?? { points: [] };
+    return state.queryResponses[channel].shift() ?? { points: [] };
   }
 }
 
@@ -159,10 +145,7 @@ mock.module("@qdrant/js-client-rest", () => ({
   QdrantClient: MockQdrantClient,
 }));
 
-const { simBatch, simSkillBatch, clamp01, effectiveWeights } =
-  await import("../sim.js");
-const { _resetMemoryV2SkillQdrantForTests } =
-  await import("../skill-qdrant.js");
+const { simBatch, clamp01, effectiveWeights } = await import("../sim.js");
 const { _resetMemoryV2QdrantForTests } = await import("../qdrant.js");
 
 // ---------------------------------------------------------------------------
@@ -175,15 +158,12 @@ function resetState(): void {
   state.embedReturn = [[0.1, 0.2, 0.3]];
   state.queryResponses.dense.length = 0;
   state.queryResponses.sparse.length = 0;
-  state.skillQueryResponses.dense.length = 0;
-  state.skillQueryResponses.sparse.length = 0;
   state.queryCalls.length = 0;
   // Bun's `mock.module` persists across files in the same process, so the
-  // qdrant modules' singletons may already hold a MockQdrantClient instance
-  // from a sibling test file. Reset both readiness caches so each test in
-  // this file gets a fresh `new QdrantClient()` resolved against our mock.
+  // qdrant module's singleton may already hold a MockQdrantClient instance
+  // from a sibling test file. Reset readiness so each test in this file
+  // gets a fresh `new QdrantClient()` resolved against our mock.
   _resetMemoryV2QdrantForTests();
-  _resetMemoryV2SkillQdrantForTests();
 }
 
 function configWithWeights(
@@ -533,177 +513,3 @@ describe("simBatch", () => {
     }
   });
 });
-
-// ---------------------------------------------------------------------------
-// simSkillBatch
-// ---------------------------------------------------------------------------
-
-/**
- * Stage a single hybrid response on the dedicated skills queues. Mirrors
- * `stageHybridResponse` but uses `payload.id` (skills' Qdrant payload key)
- * instead of `payload.slug`.
- */
-function stageSkillHybridResponse(
-  hits: Array<{ id: string; denseScore?: number; sparseScore?: number }>,
-): void {
-  state.skillQueryResponses.dense.push({
-    points: hits
-      .filter((h) => h.denseScore !== undefined)
-      .map((h) => ({ score: h.denseScore, payload: { id: h.id } })),
-  });
-  state.skillQueryResponses.sparse.push({
-    points: hits
-      .filter((h) => h.sparseScore !== undefined)
-      .map((h) => ({ score: h.sparseScore, payload: { id: h.id } })),
-  });
-}
-
-describe("simSkillBatch", () => {
-  test("empty id list returns empty map without touching backends", async () => {
-    const config = configWithWeights(0.7, 0.3);
-
-    const out = await simSkillBatch("anything", [], config);
-
-    expect(out.size).toBe(0);
-    expect(state.embedCalls).toHaveLength(0);
-    expect(state.sparseCalls).toHaveLength(0);
-    expect(state.queryCalls).toHaveLength(0);
-  });
-
-  test("empty text returns empty map without touching backends", async () => {
-    const config = configWithWeights(0.7, 0.3);
-
-    for (const text of ["", "   ", "\n\n"]) {
-      state.embedCalls.length = 0;
-      state.sparseCalls.length = 0;
-      state.queryCalls.length = 0;
-      const out = await simSkillBatch(text, ["example-skill-a"], config);
-      expect(out.size).toBe(0);
-      expect(state.embedCalls).toHaveLength(0);
-      expect(state.sparseCalls).toHaveLength(0);
-      expect(state.queryCalls).toHaveLength(0);
-    }
-  });
-
-  test("queries the dedicated skills collection and forwards an id-IN filter", async () => {
-    const config = configWithWeights(0.7, 0.3);
-    stageSkillHybridResponse([]);
-
-    await simSkillBatch(
-      "query",
-      ["example-skill-a", "example-skill-b"],
-      config,
-    );
-
-    expect(state.queryCalls).toHaveLength(2);
-    for (const call of state.queryCalls) {
-      expect(call.collection).toBe("memory_v2_skills");
-      // The candidate ids are forwarded as a Qdrant filter so Qdrant scores
-      // exactly the candidate set, not its global top-K. Without this,
-      // candidate ids absent from the global top-K silently score 0.
-      expect(call.filter).toEqual({
-        must: [
-          { key: "id", match: { any: ["example-skill-a", "example-skill-b"] } },
-        ],
-      });
-      // Limit equals the candidate count.
-      expect(call.limit).toBe(2);
-    }
-  });
-
-  test("fuses dense + sparse with the configured weight blend", async () => {
-    const config = configWithWeights(0.4, 0.6);
-    stageSkillHybridResponse([
-      { id: "example-skill-a", denseScore: 0.5, sparseScore: 4 }, // sparse-norm 1.0
-      { id: "example-skill-b", denseScore: 0.25, sparseScore: 2 }, // sparse-norm 0.5
-    ]);
-
-    const out = await simSkillBatch(
-      "query",
-      ["example-skill-a", "example-skill-b"],
-      config,
-    );
-
-    // example-skill-a: 0.4 * 0.5 + 0.6 * 1.0 = 0.8
-    // example-skill-b: 0.4 * 0.25 + 0.6 * 0.5 = 0.4
-    expect(out.get("example-skill-a")).toBeCloseTo(0.8, 6);
-    expect(out.get("example-skill-b")).toBeCloseTo(0.4, 6);
-  });
-
-  test("dense-only and sparse-only hits are handled symmetrically", async () => {
-    const config = configWithWeights(0.7, 0.3);
-    stageSkillHybridResponse([
-      { id: "example-skill-a", denseScore: 0.5 /* sparse omitted */ },
-      { id: "example-skill-b", sparseScore: 8 /* dense omitted */ },
-    ]);
-
-    const out = await simSkillBatch(
-      "query",
-      ["example-skill-a", "example-skill-b"],
-      config,
-    );
-
-    // example-skill-a: 0.7 * 0.5 + 0.3 * 0   = 0.35
-    // example-skill-b: 0.7 * 0   + 0.3 * 1.0 = 0.30 (sparse-norm = 8/8)
-    expect(out.get("example-skill-a")).toBeCloseTo(0.35, 6);
-    expect(out.get("example-skill-b")).toBeCloseTo(0.3, 6);
-  });
-
-  test("forwards candidate ids as the Qdrant restriction; only candidates in result", async () => {
-    // The bug we're guarding against: when the skills collection has more
-    // skills than `ids.length`, calling `hybridQuerySkills` without a filter
-    // returns Qdrant's global top-K. Candidate ids absent from that top-K
-    // would silently score 0. The fix is to forward the candidate ids as a
-    // server-side restriction so Qdrant scores exactly the candidate set.
-    const config = configWithWeights(0.7, 0.3);
-    stageSkillHybridResponse([
-      { id: "example-skill-a", denseScore: 0.5, sparseScore: 1 },
-      // `example-skill-c` would never be returned in production once the
-      // filter is applied; the post-filter in simSkillBatch defensively
-      // drops it even if a stale payload slips through.
-      { id: "example-skill-c", denseScore: 0.9, sparseScore: 1 },
-    ]);
-
-    const out = await simSkillBatch(
-      "query",
-      ["example-skill-a", "example-skill-b"],
-      config,
-    );
-
-    // The Qdrant filter was forwarded — both channels carry the id-IN
-    // restriction matching the caller's candidate set.
-    expect(state.queryCalls).toHaveLength(2);
-    for (const call of state.queryCalls) {
-      expect(call.filter).toEqual({
-        must: [
-          { key: "id", match: { any: ["example-skill-a", "example-skill-b"] } },
-        ],
-      });
-    }
-    // Only candidate ids appear in the result map.
-    expect(out.has("example-skill-a")).toBe(true);
-    expect(out.has("example-skill-c")).toBe(false);
-  });
-
-  test("returned scores are clamped into [0, 1]", async () => {
-    const config = configWithWeights(0.8, 0.5); // intentionally sums to > 1
-    stageSkillHybridResponse([
-      { id: "example-skill-a", denseScore: 1.0, sparseScore: 1 },
-    ]);
-
-    const out = await simSkillBatch("query", ["example-skill-a"], config);
-
-    expect(out.get("example-skill-a")).toBe(1);
-  });
-
-  test("embeds the query text exactly once via dense + sparse backends", async () => {
-    const config = configWithWeights(0.7, 0.3);
-    stageSkillHybridResponse([]);
-
-    await simSkillBatch("hello skill", ["example-skill-a"], config);
-
-    expect(state.embedCalls).toHaveLength(1);
-    expect(state.embedCalls[0].inputs).toEqual(["hello skill"]);
-    expect(state.sparseCalls).toEqual(["hello skill"]);
-  });
-});