@vellumai/assistant 0.6.6 → 0.7.0

package/src/memory/graph/retriever.test.ts

@@ -31,10 +31,12 @@ mock.module("../embed.js", () => ({
   ) => {
     embedCallCount++;
     if (embedShouldThrow) throw new Error("embedding backend down");
-    const firstText = typeof texts[0] === "string" ? texts[0] : "";
-    const vector = embedRouter ? embedRouter(firstText) : embedVector;
+    const vectors = texts.map((t) => {
+      const text = typeof t === "string" ? t : "";
+      return embedRouter ? embedRouter(text) : embedVector;
+    });
     return {
-      vectors: [vector],
+      vectors,
       provider: "test-provider",
       model: "test-model",
     };
@@ -71,7 +73,9 @@ mock.module("../../providers/provider-send-message.js", () => ({
 
 import { DEFAULT_CONFIG } from "../../config/defaults.js";
 import type { AssistantConfig } from "../../config/types.js";
-import { initializeDb, resetDb, resetTestTables } from "../db.js";
+import { resetDb } from "../db-connection.js";
+import { initializeDb } from "../db-init.js";
+import { resetTestTables } from "../raw-query.js";
 import { InContextTracker } from "./injection.js";
 import { loadContextMemory, retrieveForTurn } from "./retriever.js";
 import { createNode } from "./store.js";
@@ -238,6 +242,156 @@ describe("retrieveForTurn — query/sparse vector surfacing", () => {
   });
 });
 
+describe("retrieveForTurn — topic-pivot recovery", () => {
+  // Build a one-hot keyword router so the combined assistant+user query and
+  // the user-only query produce distinct vectors that can be routed to
+  // disjoint Qdrant result sets. The key for the bug fix is that a long
+  // shirt-themed assistant message + a short cake-themed user message
+  // currently produces a vector dominated by the shirt content (the
+  // user-only signal is washed out), so unless the user-only embedding
+  // also runs and feeds Qdrant, cake nodes never enter the candidate pool.
+  function keywordEmbedRouter(text: string): number[] {
+    const lowered = text.toLowerCase();
+    const cakeHits = (lowered.match(/cake/g) ?? []).length;
+    const shirtHits = (lowered.match(/shirt/g) ?? []).length;
+    if (cakeHits > shirtHits) return [0, 1, 0];
+    if (shirtHits > cakeHits) return [1, 0, 0];
+    return [0.1, 0.1, 0.1];
+  }
+
+  let cakeNodeId = "";
+  let shirtNodeId = "";
+
+  function vectorSearchRouter(vector: number[]): SearchHit[] {
+    const [a = 0, b = 0] = vector;
+    if (a === 1 && b === 0) return [{ nodeId: shirtNodeId, score: 0.9 }];
+    if (a === 0 && b === 1) return [{ nodeId: cakeNodeId, score: 0.9 }];
+    return [];
+  }
+
+  function makeEpisodicNode(content: string): NewNode {
+    const now = Date.now();
+    return {
+      content,
+      type: "episodic",
+      created: now,
+      lastAccessed: now,
+      lastConsolidated: now,
+      eventDate: null,
+      emotionalCharge: {
+        valence: 0,
+        intensity: 0,
+        decayCurve: "linear",
+        decayRate: 0,
+        originalIntensity: 0,
+      },
+      fidelity: "vivid",
+      confidence: 1,
+      significance: 0.5,
+      stability: 14,
+      reinforcementCount: 0,
+      lastReinforced: now,
+      sourceConversations: [],
+      sourceType: "direct",
+      narrativeRole: null,
+      partOfStory: null,
+      imageRefs: null,
+      scopeId: "test-scope",
+    };
+  }
+
+  beforeAll(() => {
+    initializeDb();
+  });
+
+  beforeEach(() => {
+    embedShouldThrow = false;
+    embedVector = [0.1, 0.2, 0.3];
+    embedCallCount = 0;
+    embedRouter = keywordEmbedRouter;
+    searchRouter = vectorSearchRouter;
+    resetTestTables(
+      "memory_graph_triggers",
+      "memory_graph_edges",
+      "memory_graph_nodes",
+    );
+    cakeNodeId = createNode(
+      makeEpisodicNode("Notes on the birthday cake from the bakery."),
+    ).id;
+    shirtNodeId = createNode(
+      makeEpisodicNode("Notes on the souvenir shirt from the conference."),
+    ).id;
+  });
+
+  test("user-only embed surfaces topic-pivot candidates the combined embed misses", async () => {
+    const tracker = new InContextTracker();
+    // Long shirt-themed assistant message that dominates the combined
+    // embedding, paired with a short cake-themed user pivot.
+    const assistantLast =
+      "shirt shirt shirt shirt shirt shirt shirt shirt " +
+      "shirt shirt shirt shirt shirt shirt shirt shirt";
+    const userLast = "the birthday cake";
+
+    const result = await retrieveForTurn({
+      assistantLastMessage: assistantLast,
+      userLastMessage: userLast,
+      scopeId: "test-scope",
+      config: TEST_CONFIG,
+      tracker,
+    });
+
+    // Both embeds ran (combined + user-only).
+    expect(embedCallCount).toBe(1);
+    // Combined query embedding (still surfaced for PKB hybrid search).
+    expect(result.queryVector).toEqual([1, 0, 0]);
+    // The cake node should appear in the candidate pool — it would be
+    // missing entirely if only the combined embedding drove the search.
+    const candidateIds = new Set(
+      result.metrics.topCandidates.map((c) => c.nodeId),
+    );
+    expect(candidateIds.has(cakeNodeId)).toBe(true);
+    expect(candidateIds.has(shirtNodeId)).toBe(true);
+  });
+
+  test("skips the user-only chunk when the user message is empty", async () => {
+    const tracker = new InContextTracker();
+    const result = await retrieveForTurn({
+      assistantLastMessage: "shirt shirt shirt shirt shirt",
+      userLastMessage: "",
+      scopeId: "test-scope",
+      config: TEST_CONFIG,
+      tracker,
+    });
+
+    // Only the combined (assistant-only) embedding ran; user-only is skipped.
+    expect(embedCallCount).toBe(1);
+    const candidateIds = new Set(
+      result.metrics.topCandidates.map((c) => c.nodeId),
+    );
+    expect(candidateIds.has(shirtNodeId)).toBe(true);
+    expect(candidateIds.has(cakeNodeId)).toBe(false);
+  });
+
+  test("does not duplicate the user-only chunk when it is already the only chunk", async () => {
+    const tracker = new InContextTracker();
+    // No assistant message → queryText === userLastMessage, so user-only is
+    // already the sole chunk. The dedup check should prevent re-adding it.
+    const result = await retrieveForTurn({
+      assistantLastMessage: "",
+      userLastMessage: "the birthday cake",
+      scopeId: "test-scope",
+      config: TEST_CONFIG,
+      tracker,
+    });
+
+    expect(embedCallCount).toBe(1);
+    const candidateIds = new Set(
+      result.metrics.topCandidates.map((c) => c.nodeId),
+    );
+    expect(candidateIds.has(cakeNodeId)).toBe(true);
+  });
+});
+
 describe("loadContextMemory — dual-query capability ranking", () => {
   // Capture seeded capability node IDs so the searchGraphNodes mock can
   // reference them by ID (the mock runs at call time, not seed time).

package/src/memory/graph/retriever.ts

@@ -933,11 +933,10 @@ export async function retrieveForTurn(
   // hint search can pair it with the per-turn dense vector, pulling in
   // lexical matches (exact filenames, proper nouns, uncommon tokens) that
   // pure dense embeddings wash out. Computed locally with no network call.
-  // The dense `queryVector` for per-turn is the combined assistant+user
-  // embedding (it drives graph search), so this pairing is slightly
-  // asymmetric — the sparse signal is user-only while the dense is mixed.
-  // That's acceptable for PKB's purpose: lexical matches are overwhelmingly
-  // driven by the user's own wording.
+  // The surfaced `queryVector` is the combined assistant+user embedding
+  // (drives PKB hybrid search alongside this sparse signal). The graph
+  // search itself runs against both the combined embedding and a separate
+  // user-only embedding — see the user-only chunk added below.
   const trimmedUserLast = opts.userLastMessage.trim();
   let perTurnSparseVector: QdrantSparseVector | undefined = undefined;
   if (trimmedUserLast.length > 0) {
@@ -1044,6 +1043,19 @@ export async function retrieveForTurn(
     }
   }
 
+  // Topic-pivot recovery: also embed the user message alone. When the
+  // assistant's prior message is long (e.g. includes <thinking>), the
+  // combined embedding is dominated by it and a short user pivot to a new
+  // topic gets drowned out. Searching with both vectors and unioning the
+  // candidates lets the pivot compete in scoring.
+  if (
+    trimmedUserLast.length > 0 &&
+    opts.userLastMessage.length <= maxQueryChars &&
+    !chunks.includes(opts.userLastMessage)
+  ) {
+    chunks.push(opts.userLastMessage);
+  }
+
   // 2. Embed chunks and search (parallel)
   let queryEmbeddings: number[][] = [];
 

package/src/memory/graph/store.test.ts

@@ -1,6 +1,7 @@
 import { beforeAll, beforeEach, describe, expect, test } from "bun:test";
 
-import { initializeDb, resetTestTables } from "../db.js";
+import { initializeDb } from "../db-init.js";
+import { resetTestTables } from "../raw-query.js";
 import {
   applyDiff,
   countNodes,

package/src/memory/graph/store.ts

@@ -5,7 +5,7 @@
 import { and, desc, eq, inArray, or, sql } from "drizzle-orm";
 import { v4 as uuid } from "uuid";
 
-import { getDb } from "../db.js";
+import { getDb } from "../db-connection.js";
 import { enqueueMemoryJob } from "../jobs-store.js";
 import {
   memoryGraphEdges,

package/src/memory/graph/tool-handlers.ts

@@ -1,252 +1,25 @@
 // ---------------------------------------------------------------------------
-// Memory Graph — Tool handlers for recall and remember
+// Memory Tool handlers
 //
-// recall: search the living graph or raw archive
-// remember: save facts to the PKB (buffer.md + daily archive)
+// remember: save facts to the PKB (buffer.md + daily archive) under the v1
+// path, or to memory/buffer.md + memory/archive/<today>.md when the
+// `memory-v2-enabled` feature flag is on.
 // ---------------------------------------------------------------------------
 
 import { appendFileSync, existsSync, mkdirSync } from "node:fs";
 import { join } from "node:path";
 
+import { isAssistantFeatureFlagEnabled } from "../../config/assistant-feature-flags.js";
 import type { AssistantConfig } from "../../config/types.js";
 import { getLogger } from "../../util/logger.js";
 import { getWorkspaceDir } from "../../util/platform.js";
-import { buildExcerpt, buildFtsMatchQuery } from "../conversation-queries.js";
-import { embedWithRetry } from "../embed.js";
-import { generateSparseEmbedding } from "../embedding-backend.js";
 import { enqueuePkbIndexJob } from "../jobs/embed-pkb-file.js";
 import { PKB_WORKSPACE_SCOPE } from "../pkb/types.js";
-import { searchGraphNodes } from "./graph-search.js";
-import { getNodesByIds } from "./store.js";
 
 const log = getLogger("graph-tool-handlers");
 
 // ---------------------------------------------------------------------------
-// recall handler
-// ---------------------------------------------------------------------------
-
-export interface RecallInput {
-  query: string;
-  mode?: "memory" | "archive";
-  num_results?: number;
-  filters?: {
-    types?: string[];
-    after?: string;
-    before?: string;
-  };
-}
-
-export interface RecallResult {
-  results: Array<{
-    id: string;
-    content: string;
-    type: string;
-    confidence: number;
-    significance: number;
-    score: number;
-    created: number;
-  }>;
-  mode: "memory" | "archive";
-  query: string;
-}
-
-export async function handleRecall(
-  input: RecallInput,
-  config: AssistantConfig,
-  scopeId: string,
-): Promise<RecallResult> {
-  const mode = input.mode ?? "memory";
-
-  if (mode === "archive") {
-    return handleArchiveRecall(input, scopeId);
-  }
-
-  return handleMemoryRecall(input, config, scopeId);
-}
-
-async function handleMemoryRecall(
-  input: RecallInput,
-  config: AssistantConfig,
-  scopeId: string,
-): Promise<RecallResult> {
-  // Embed the query
-  let queryVector: number[] | null = null;
-  try {
-    const result = await embedWithRetry(config, [input.query]);
-    queryVector = result.vectors[0] ?? null;
-  } catch (err) {
-    log.warn({ err }, "Failed to embed recall query");
-    return { results: [], mode: "memory", query: input.query };
-  }
-
-  if (!queryVector) {
-    return { results: [], mode: "memory", query: input.query };
-  }
-
-  // Generate sparse embedding for hybrid search (dense + sparse with RRF fusion)
-  const sparseVector = generateSparseEmbedding(input.query);
-
-  // Build date range filter for Qdrant-level filtering
-  const dateRange: { afterMs?: number; beforeMs?: number } = {};
-  if (input.filters?.after) {
-    const afterMs = new Date(input.filters.after).getTime();
-    if (!isNaN(afterMs)) dateRange.afterMs = afterMs;
-  }
-  if (input.filters?.before) {
-    const beforeMs = new Date(input.filters.before).getTime();
-    if (!isNaN(beforeMs)) dateRange.beforeMs = beforeMs;
-  }
-
-  // Search graph nodes
-  const limit = Math.max(1, Math.min(input.num_results ?? 20, 50));
-  const searchResults = await searchGraphNodes(
-    queryVector,
-    limit,
-    [scopeId],
-    sparseVector,
-    dateRange.afterMs != null || dateRange.beforeMs != null
-      ? dateRange
-      : undefined,
-  );
-  if (searchResults.length === 0) {
-    return { results: [], mode: "memory", query: input.query };
-  }
-
-  // Hydrate
-  const nodes = getNodesByIds(searchResults.map((r) => r.nodeId));
-  const nodeMap = new Map(nodes.map((n) => [n.id, n]));
-
-  // Apply filters
-  const results = searchResults.flatMap((r) => {
-    const node = nodeMap.get(r.nodeId);
-    if (!node || node.fidelity === "gone") return [];
-
-    // Type filter
-    if (input.filters?.types && input.filters.types.length > 0) {
-      if (!input.filters.types.includes(node.type)) return [];
-    }
-
-    return [
-      {
-        id: node.id,
-        content: node.content,
-        type: node.type,
-        confidence: node.confidence,
-        significance: node.significance,
-        score: r.score,
-        created: node.created,
-      },
-    ];
-  });
-
-  return { results, mode: "memory", query: input.query };
-}
-
-async function handleArchiveRecall(
-  input: RecallInput,
-  scopeId: string,
-): Promise<RecallResult> {
-  // Archive mode: search raw conversation transcripts via messages FTS
-  // This is a simple text search — no embedding needed
-  const { rawAll } = await import("../db.js");
-
-  try {
-    const limit = Math.max(1, Math.min(input.num_results ?? 20, 50));
-    const ftsMatch = buildFtsMatchQuery(input.query.trim(), {
-      allowFts5Syntax: true,
-    });
-
-    const afterMs = input.filters?.after
-      ? new Date(input.filters.after).getTime()
-      : NaN;
-    const beforeMs = input.filters?.before
-      ? new Date(input.filters.before).getTime()
-      : NaN;
-    const dateConditions: string[] = [];
-    const dateParams: number[] = [];
-    if (!isNaN(afterMs)) {
-      dateConditions.push("m.created_at >= ?");
-      dateParams.push(afterMs);
-    }
-    if (!isNaN(beforeMs)) {
-      dateConditions.push("m.created_at <= ?");
-      dateParams.push(beforeMs);
-    }
-    const dateClause =
-      dateConditions.length > 0 ? " AND " + dateConditions.join(" AND ") : "";
-
-    type ArchiveRow = {
-      id: string;
-      content: string;
-      role: string;
-      created_at: number;
-      conversation_id: string;
-    };
-
-    let rows: ArchiveRow[];
-
-    if (ftsMatch) {
-      // Use SQLite FTS on messages table, scoped to the active memory scope
-      rows = rawAll<ArchiveRow>(
-        `SELECT m.id, m.content, m.role, m.created_at, c.id as conversation_id
-         FROM messages_fts fts
-         JOIN messages m ON m.id = fts.message_id
-         JOIN conversations c ON c.id = m.conversation_id
-         WHERE messages_fts MATCH ?
-           AND c.memory_scope_id = ?${dateClause}
-         ORDER BY rank
-         LIMIT ?`,
-        ftsMatch,
-        scopeId,
-        ...dateParams,
-        limit,
-      );
-    } else if (!input.query.trim()) {
-      // Empty or whitespace-only query — return nothing rather than matching
-      // every message via a `%%` LIKE pattern.
-      rows = [];
-    } else {
-      // All FTS tokens dropped (non-ASCII, single-char, etc.) — fall back to
-      // LIKE-based search so queries like CJK characters still match.
-      const likePattern = `%${input.query
-        .replace(/\\/g, "\\\\")
-        .replace(/%/g, "\\%")
-        .replace(/_/g, "\\_")}%`;
-      rows = rawAll<ArchiveRow>(
-        `SELECT m.id, m.content, m.role, m.created_at, c.id as conversation_id
-         FROM messages m
-         JOIN conversations c ON c.id = m.conversation_id
-         WHERE m.content LIKE ? ESCAPE '\\' AND c.memory_scope_id = ?${dateClause}
-         ORDER BY m.created_at DESC
-         LIMIT ?`,
-        likePattern,
-        scopeId,
-        ...dateParams,
-        limit,
-      );
-    }
-
-    return {
-      results: rows.map((r) => ({
-        id: r.id,
-        content: buildExcerpt(r.content, input.query),
-        type: "archive",
-        confidence: 0,
-        significance: 0,
-        score: 0,
-        created: r.created_at,
-      })),
-      mode: "archive",
-      query: input.query,
-    };
-  } catch (err) {
-    log.warn({ err }, "Archive recall FTS failed");
-    return { results: [], mode: "archive", query: input.query };
-  }
-}
-
-// ---------------------------------------------------------------------------
-// remember handler — writes to PKB buffer + daily archive
+// remember handler — writes to PKB (v1) or memory/ (v2) buffer + daily archive
 // ---------------------------------------------------------------------------
 
 export interface RememberInput {
@@ -263,46 +36,99 @@ export function handleRemember(
   input: RememberInput,
   _conversationId: string,
   _scopeId: string,
+  config: AssistantConfig,
 ): RememberResult {
   if (!input.content || input.content.trim().length === 0) {
     return { success: false, message: "content is required" };
   }
 
   const workspaceDir = getWorkspaceDir();
+  const now = new Date();
+  const entry = formatRememberEntry(input.content.trim(), now);
+
+  if (isAssistantFeatureFlagEnabled("memory-v2-enabled", config)) {
+    appendBufferAndArchive({
+      rootDir: join(workspaceDir, "memory"),
+      entry,
+      now,
+    });
+    // v2 path skips the PKB re-index queue — embedding for memory v2 happens
+    // via the dedicated `embed_concept_page` job after consolidation, not on
+    // every remember() write.
+    return { success: true, message: "Saved to knowledge base." };
+  }
+
   const pkbDir = join(workspaceDir, "pkb");
-  const archiveDir = join(pkbDir, "archive");
+  const { bufferPath, archivePath } = appendBufferAndArchive({
+    rootDir: pkbDir,
+    entry,
+    now,
+  });
+  enqueuePkbReindex(pkbDir, bufferPath);
+  enqueuePkbReindex(pkbDir, archivePath);
 
-  // Ensure directories exist
-  mkdirSync(pkbDir, { recursive: true });
-  mkdirSync(archiveDir, { recursive: true });
+  return { success: true, message: "Saved to knowledge base." };
+}
 
-  // Build timestamped entry
-  const now = new Date();
+/**
+ * Build a timestamped bullet entry for `buffer.md` / `archive/<date>.md`.
+ *
+ * Format mirrors the long-standing v1 PKB layout so v2 buffers stay
+ * human-readable and downstream consumers (sweep, consolidation) can parse
+ * the same shape regardless of which path produced the entry.
+ *
+ * Exported so memory v2 sweep / extractor jobs format their auto-remembered
+ * entries identically to user-facing `remember()` calls.
+ */
+export function formatRememberEntry(content: string, now: Date): string {
   const month = now.toLocaleString("en-US", { month: "short" });
   const day = now.getDate();
   const hours = now.getHours();
   const minutes = String(now.getMinutes()).padStart(2, "0");
   const ampm = hours >= 12 ? "PM" : "AM";
   const displayHour = hours % 12 || 12;
-  const entry = `- [${month} ${day}, ${displayHour}:${minutes} ${ampm}] ${input.content.trim()}\n`;
+  return `- [${month} ${day}, ${displayHour}:${minutes} ${ampm}] ${content}\n`;
+}
 
-  // Append to buffer.md
-  const bufferPath = join(pkbDir, "buffer.md");
+/**
+ * Append `entry` to `<rootDir>/buffer.md` and `<rootDir>/archive/<today>.md`,
+ * creating the archive directory and seeding the archive header if missing.
+ *
+ * Returns the absolute paths of both files so callers can fan out follow-up
+ * work (e.g. PKB re-indexing in the v1 path).
+ *
+ * Exported so memory v2 background jobs (`sweep`, future LLM-driven
+ * extractors) can append to `memory/buffer.md` + `memory/archive/<today>.md`
+ * with exactly the same format `remember()` produces, keeping the two write
+ * paths byte-compatible for downstream consumers (consolidation, search).
+ */
+export function appendBufferAndArchive(args: {
+  rootDir: string;
+  entry: string;
+  now: Date;
+}): { bufferPath: string; archivePath: string } {
+  const { rootDir, entry, now } = args;
+  const archiveDir = join(rootDir, "archive");
+  mkdirSync(archiveDir, { recursive: true });
+
+  const bufferPath = join(rootDir, "buffer.md");
   appendFileSync(bufferPath, entry, "utf-8");
-  enqueuePkbReindex(pkbDir, bufferPath);
 
-  // Append to daily archive
   const yyyy = now.getFullYear();
   const mm = String(now.getMonth() + 1).padStart(2, "0");
   const dd = String(now.getDate()).padStart(2, "0");
   const archivePath = join(archiveDir, `${yyyy}-${mm}-${dd}.md`);
   if (!existsSync(archivePath)) {
-    appendFileSync(archivePath, `# ${month} ${day}, ${yyyy}\n\n`, "utf-8");
+    const month = now.toLocaleString("en-US", { month: "short" });
+    appendFileSync(
+      archivePath,
+      `# ${month} ${now.getDate()}, ${yyyy}\n\n`,
+      "utf-8",
+    );
   }
   appendFileSync(archivePath, entry, "utf-8");
-  enqueuePkbReindex(pkbDir, archivePath);
 
-  return { success: true, message: "Saved to knowledge base." };
+  return { bufferPath, archivePath };
 }
 
 /**