@vellumai/assistant 0.8.6 → 0.8.7-dev.202606052118.34cd356

package/src/plugins/defaults/memory-v3-shadow/router.ts

@@ -0,0 +1,190 @@
+/**
+ * Memory v3 — L1 leaf router.
+ *
+ * One forced-tool LLM call per turn that picks which leaves of the topic tree
+ * to open for the next reply. The design mirrors `../v2/router.ts`:
+ *   - resolve the configured provider via `getConfiguredProvider`,
+ *   - call `provider.sendMessage` with a forced `tool_choice`,
+ *   - validate the tool input via Zod,
+ *   - map numbered IDs back to leaf paths.
+ *
+ * Cache strategy. The numbered leaf block is the single largest input and is
+ * STABLE across turns — it changes only when leaves are added/removed/edited.
+ * We render it first in the user message and tag it with an ephemeral
+ * `cache_control` breakpoint so the provider can serve it from the prompt
+ * cache turn after turn. The trailing recent-context / current-message block
+ * changes every turn, so it carries no breakpoint.
+ *
+ * Failure handling. A *model-call* failure is not the same as the model
+ * choosing to open everything, so the two no longer share an outcome:
+ *   - explicit `ids` → open exactly those leaves,
+ *   - explicit empty array (`ids: []`) → open nothing (deliberate abstention),
+ *   - omitted `ids` → open nothing: the router must name the leaves it wants,
+ *     never the whole tree (~137 leaves would fan out a full L2 pass per turn),
+ *   - infrastructure failure (provider unavailable, a throw that survived the
+ *     provider's own retries, no usable `tool_use`, or a schema mismatch) →
+ *     open nothing after a short re-prompt retry, degrading to the deterministic
+ *     recall lanes (always-on core, the BM25 needle, the carry-forward working
+ *     set) that the orchestrator unions in regardless.
+ */
+
+import { z } from "zod";
+
+import {
+  extractToolUse,
+  getConfiguredProvider,
+} from "../../../providers/provider-send-message.js";
+import type { Message, ToolDefinition } from "../../../providers/types.js";
+import { getLogger } from "../../../util/logger.js";
+import { retryForResult } from "./llm-retry.js";
+import { cachedTextBlock } from "./provider-blocks.js";
+import type { LeafPath, LeafTree, TurnContext } from "./types.js";
+
+const log = getLogger("memory-v3-router");
+
+/** Tool name forced via `tool_choice`. Shared constant so tests can match it. */
+const OPEN_LEAVES_TOOL_NAME = "open_leaves";
+
+const OpenLeavesSchema = z.object({
+  // Optional so the field can be absent on the wire, but an omitted `ids` opens
+  // nothing — the router must name the leaves it wants, never the whole tree.
+  ids: z.array(z.number().int()).optional(),
+});
+
+const OPEN_LEAVES_TOOL: ToolDefinition = {
+  name: OPEN_LEAVES_TOOL_NAME,
+  description:
+    "Open the leaves whose contents could plausibly bear on the next reply. " +
+    "Lean toward inclusion — a missed relevant leaf is a worse error than an " +
+    "unused one. Pass the chosen IDs explicitly; return `[]` only when nothing " +
+    "in the tree could possibly help.",
+  input_schema: {
+    type: "object",
+    properties: {
+      ids: {
+        type: "array",
+        items: { type: "integer" },
+      },
+    },
+  },
+};
+
+const SYSTEM_PROMPT = `You route a conversation turn to the leaves of a topic tree that should be opened for the next reply.
+
+Each leaf has a numbered ID, a path, and a description of what it holds. Decide which leaves to open by weighing four signals:
+
+- Topic — entities, projects, and events named or implied by the turn.
+- Register — the affect and mode of the message (e.g. playful, distressed, formal). A register signal is enough to open a leaf even when no entity is named.
+- Recent context — the immediately preceding exchange, which resolves references like "this", "that", or "the same thing" to concrete topics.
+- Situation — the current date and a live scratchpad of what is salient right now. A date or state cue can make a leaf relevant even when the message never names it (e.g. a person whose anniversary is today, an active thread).
+
+Include on doubt: open every leaf that could plausibly hold something useful. Missing a relevant leaf is a worse error than opening an unused one. Call \`open_leaves\` with the chosen IDs explicitly; return \`[]\` only when nothing in the tree could possibly help.`;
+
+/** Leaves sorted deterministically by path so the numbered block is stable. */
+function sortedLeaves(tree: LeafTree): LeafPath[] {
+  return [...tree.leaves.keys()].sort();
+}
+
+/**
+ * Render the STATIC numbered leaf block from a pre-sorted path list. Identical
+ * across turns for any given tree, which is what makes the prompt-cache
+ * breakpoint pay off.
+ */
+function renderLeafBlockFromPaths(tree: LeafTree, paths: LeafPath[]): string {
+  const lines = paths.map((path, i) => {
+    const description = tree.leaves.get(path)?.description ?? "";
+    // Collapse the (possibly multi-line) description to a single line so each
+    // leaf is exactly one numbered entry.
+    const oneLine = description.replace(/\s+/g, " ").trim();
+    return `[${i + 1}] ${path} — ${oneLine}`;
+  });
+  return `<leaves>\n${lines.join("\n")}\n</leaves>`;
+}
+
+/**
+ * Render the static numbered leaf block for a tree. Exported for the test that
+ * locks the byte-identical cache invariant; `routeL1` renders from its already
+ * sorted path list to avoid sorting twice.
+ */
+export function renderLeafBlock(tree: LeafTree): string {
+  return renderLeafBlockFromPaths(tree, sortedLeaves(tree));
+}
+
+/**
+ * Run the L1 router for one turn. Returns the leaf paths to open — only ever the
+ * leaves the model names explicitly. An omitted `ids`, an explicit `[]`, or an
+ * infrastructure failure (after a short re-prompt retry) all open nothing,
+ * degrading to the deterministic recall lanes the orchestrator unions in.
+ */
+export async function routeL1(
+  turn: TurnContext,
+  tree: LeafTree,
+): Promise<LeafPath[]> {
+  const paths = sortedLeaves(tree);
+  if (paths.length === 0) return [];
+
+  const provider = await getConfiguredProvider("memoryV3RouteL1");
+  if (!provider) {
+    log.warn(
+      "L1 router provider unavailable; degrading to deterministic lanes",
+    );
+    return [];
+  }
+
+  const userMsg: Message = {
+    role: "user",
+    content: [
+      cachedTextBlock(renderLeafBlockFromPaths(tree, paths)),
+      {
+        type: "text",
+        text:
+          (turn.situationalContext
+            ? `<situation>${turn.situationalContext}</situation>\n`
+            : "") +
+          `<recent_context>${turn.recentContext}</recent_context>\n` +
+          `<current_message>${turn.currentMessage}</current_message>`,
+      },
+    ],
+  };
+
+  // One forced-tool call, retried a few times so a transient malformed response
+  // (no usable tool_use, or tool input that fails the schema) re-prompts before
+  // we give up. `null` from an attempt means "unusable, retry"; the provider
+  // layer already backs off transient throws, so this loop adds no delay.
+  const parsed = await retryForResult(async () => {
+    const response = await provider.sendMessage([userMsg], {
+      tools: [OPEN_LEAVES_TOOL],
+      systemPrompt: SYSTEM_PROMPT,
+      config: {
+        callSite: "memoryV3RouteL1" as const,
+        tool_choice: { type: "tool" as const, name: OPEN_LEAVES_TOOL_NAME },
+      },
+    });
+    const toolBlock = extractToolUse(response);
+    if (!toolBlock || toolBlock.name !== OPEN_LEAVES_TOOL_NAME) return null;
+    const result = OpenLeavesSchema.safeParse(toolBlock.input);
+    return result.success ? result.data : null;
+  });
+
+  if (parsed === null) {
+    log.warn(
+      "L1 router could not obtain a selection after retries; degrading to deterministic lanes",
+    );
+    return [];
+  }
+
+  // An omitted `ids` field means the model named no leaves — open nothing rather
+  // than the whole tree. Only explicitly listed IDs open leaves.
+  if (parsed.ids === undefined) return [];
+
+  // Map 1-based IDs back to leaf paths, dropping out-of-range IDs without
+  // throwing. De-duplicate while preserving model-returned order.
+  const seen = new Set<number>();
+  const selected: LeafPath[] = [];
+  for (const id of parsed.ids) {
+    if (id < 1 || id > paths.length || seen.has(id)) continue;
+    seen.add(id);
+    selected.push(paths[id - 1]);
+  }
+  return selected;
+}

package/src/plugins/defaults/memory-v3-shadow/selection-log-store.ts

@@ -0,0 +1,84 @@
+/**
+ * Read-side store for the inspector's Memory V3 section. Reads the persisted
+ * `memory_v3_selections` rows for a turn and re-renders the `<memory>` block
+ * the v3 working set selected, so the inspector can show what v3 chose (and,
+ * in live mode, what it actually injected) without re-running orchestration —
+ * which would be wrong anyway, since the working set is stateful (carry-forward
+ * across turns) and can't be reproduced after the fact.
+ */
+
+import type { MemoryV3SelectionLog } from "../../../api/responses/memory-v3-selection-log.js";
+import { isAssistantFeatureFlagEnabled } from "../../../config/assistant-feature-flags.js";
+import { getConfig } from "../../../config/loader.js";
+import { getDb, getSqliteFrom } from "../../../memory/db-connection.js";
+import { renderV3PageContent } from "./page-content.js";
+import { renderMemoryBlock } from "./render-injection.js";
+import type { Slug } from "./types.js";
+
+const MEMORY_V3_SHADOW = "memory-v3-shadow" as const;
+const MEMORY_V3_LIVE = "memory-v3-live" as const;
+
+interface SelectionRow {
+  slug: string;
+  source: string;
+  pinned: number;
+}
+
+function rowsForTurn(conversationId: string, turn: number): SelectionRow[] {
+  return getSqliteFrom(getDb())
+    .query(
+      /*sql*/ `
+      SELECT slug, source, pinned FROM memory_v3_selections
+      WHERE conversation_id = ? AND turn = ?
+      ORDER BY rowid
+    `,
+    )
+    .all(conversationId, turn) as SelectionRow[];
+}
+
+/**
+ * Build the inspector's v3 selection log for one turn of a conversation.
+ *
+ * `turn` is the inspected message's turn (the v2-activation log's turn). The
+ * selection is returned ONLY for that exact turn; when there are no v3 rows for
+ * it — including when `turn` is null (the message has no v2-activation turn) —
+ * this returns `null` rather than falling back to a different turn. A fallback
+ * would attribute another turn's selection (e.g. a later turn's pages and
+ * rendered block) to the inspected message, which corrupts shadow validation.
+ *
+ * Caveat: v3 rows are keyed by the orchestrator turn counter (`ctx.turnCount`)
+ * while `turn` here is the v2 memory-tracker turn. They coincide for normal
+ * turns, but if they diverge this simply yields `null` (no section) rather than
+ * wrong data. Tying v3 rows to a message id for exact per-message attribution
+ * regardless of counter drift is a documented follow-up.
+ *
+ * Selection rows are stored in `finalInjection` order (this turn's L2
+ * selections, then carry-forward), so rendering them in row order reproduces
+ * the block v3 would inject.
+ */
+export async function getMemoryV3SelectionForInspector(
+  conversationId: string,
+  turn: number | null | undefined,
+): Promise<MemoryV3SelectionLog | null> {
+  if (turn == null) return null;
+
+  const rows = rowsForTurn(conversationId, turn);
+  if (rows.length === 0) return null;
+
+  const config = getConfig();
+  const selections = rows.map((r) => ({
+    slug: r.slug,
+    source: r.source,
+    pinned: r.pinned === 1,
+  }));
+  const slugs: Slug[] = selections.map((s) => s.slug);
+  const injectedText = await renderMemoryBlock(slugs, renderV3PageContent);
+
+  return {
+    turn,
+    live: isAssistantFeatureFlagEnabled(MEMORY_V3_LIVE, config),
+    shadow: isAssistantFeatureFlagEnabled(MEMORY_V3_SHADOW, config),
+    selections,
+    injectedText,
+  };
+}

package/src/plugins/defaults/memory-v3-shadow/selector.ts

@@ -0,0 +1,226 @@
+/**
+ * Memory v3 — L2 per-leaf page selector.
+ *
+ * After the L1 router (`./router.ts`) decides which leaves to open, the L2
+ * selector runs ONE forced-tool LLM call PER opened leaf to pick which pages
+ * inside that leaf are relevant for the next reply. `selectAcrossLeaves` fans
+ * the per-leaf calls out with bounded concurrency.
+ *
+ * Cache strategy. Each leaf's `<pages>` block — the numbered list of its member
+ * pages with their full summaries — is STABLE for that leaf: it changes only
+ * when pages are added/removed or a summary is rewritten, never per turn. We
+ * render it FIRST in the per-leaf user message and tag it with an ephemeral
+ * `cache_control` breakpoint so the provider serves it from the prompt cache
+ * turn after turn. The trailing recent-context / current-message block changes
+ * every turn and carries no breakpoint. This mirrors `./router.ts`.
+ *
+ * Failure handling. A deliberate "select everything" and a model-call failure
+ * are different events with different outcomes:
+ *   - explicit `ids` → select exactly those pages,
+ *   - explicit empty `ids: []` → select nothing (deliberate abstention),
+ *   - omitted `ids` → select ALL members of the leaf (the recall-safe "this
+ *     whole leaf is relevant" signal, e.g. "give me all of X"); bounded to one
+ *     leaf, so unlike the router this stays a select-all,
+ *   - infrastructure failure (provider unavailable, a throw that survived the
+ *     provider's own retries, no usable `tool_use`, or a schema mismatch) →
+ *     select nothing after a short re-prompt retry, degrading to the
+ *     deterministic recall lanes (core, needle, carry-forward working set) the
+ *     orchestrator unions in regardless.
+ */
+
+import { z } from "zod";
+
+import {
+  extractToolUse,
+  getConfiguredProvider,
+} from "../../../providers/provider-send-message.js";
+import type { Message, ToolDefinition } from "../../../providers/types.js";
+import { getLogger } from "../../../util/logger.js";
+import { mapLimit } from "../../../util/map-limit.js";
+import { retryForResult } from "./llm-retry.js";
+import { cachedTextBlock } from "./provider-blocks.js";
+import { membersOf } from "./tree.js";
+import type { LeafPath, LeafTree, Slug, TurnContext } from "./types.js";
+
+const log = getLogger("memory-v3-selector");
+
+/** A page selected from an opened leaf, with whether the turn centers on it. */
+export interface SelectedPage {
+  slug: Slug;
+  pinned: boolean;
+}
+
+/** Tool name forced via `tool_choice`. Shared constant so tests can match it. */
+const SELECT_PAGES_TOOL_NAME = "select_pages";
+
+const SelectPagesSchema = z.object({
+  // Optional: an omitted `ids` field is the recall-safe "select everything"
+  // signal, distinct from an explicit empty array (deliberate abstention).
+  ids: z.array(z.number().int()).optional(),
+  pinned_ids: z.array(z.number().int()).optional(),
+});
+
+const SELECT_PAGES_TOOL: ToolDefinition = {
+  name: SELECT_PAGES_TOOL_NAME,
+  description:
+    "Select the pages in this leaf whose content the reply would directly " +
+    "draw on. Be selective — prefer a few precisely-relevant pages over many " +
+    "loosely-related ones; a leaf opened on a weak signal may yield none. " +
+    "Pass `pinned_ids` for pages the conversation is centrally about. Omit " +
+    "`ids` only as a recall-safe fallback when you cannot judge the leaf " +
+    "(selects every page); return `[]` when pages are present but none are " +
+    "directly relevant.",
+  input_schema: {
+    type: "object",
+    properties: {
+      ids: {
+        type: "array",
+        items: { type: "integer" },
+      },
+      pinned_ids: {
+        type: "array",
+        items: { type: "integer" },
+      },
+    },
+  },
+};
+
+const SYSTEM_PROMPT = `This leaf of the topic tree is potentially relevant to the conversation. Select ONLY the pages whose content the reply to THIS message would directly draw on.
+
+Be selective: exclude pages that are merely topically adjacent, part of the ever-present background, or only loosely related. Most opened leaves should contribute a few precisely-relevant pages, not most of their contents — a leaf opened on a weak signal may yield none.
+
+A page can also be directly relevant because of the current situation — the date or the live scratchpad — not only the message: keep a page the situation makes pertinent (e.g. a person whose anniversary is today).
+
+If the conversation is centrally ABOUT a page (rather than only peripherally relevant to it), mark that page as pinned. Call \`select_pages\` with the chosen IDs. Omit \`ids\` only as a recall-safe fallback when you cannot judge the leaf (selects every page); return \`[]\` when the pages are present but none are directly relevant.`;
+
+/**
+ * Render the STATIC numbered `<pages>` block for a leaf from its member slugs.
+ * Identical across turns for a given leaf, which is what makes the per-leaf
+ * prompt-cache breakpoint pay off. Summaries are rendered in full (no
+ * truncation) so the selector sees each page's complete description.
+ */
+async function renderPagesBlock(
+  members: Slug[],
+  pageSummary: (slug: Slug) => Promise<string>,
+): Promise<string> {
+  const lines = await Promise.all(
+    members.map(async (slug, i) => {
+      const summary = await pageSummary(slug);
+      return `[${i + 1}] ${slug} — ${summary}`;
+    }),
+  );
+  return `<pages>\n${lines.join("\n")}\n</pages>`;
+}
+
+/**
+ * Run the L2 selector for a single opened leaf. Returns the pages to inject.
+ *
+ * An omitted `ids` selects ALL members (the recall-safe "whole leaf is
+ * relevant" signal); an explicit `[]` selects none; an infrastructure failure
+ * (after a short re-prompt retry) selects none, degrading to the deterministic
+ * recall lanes the orchestrator unions in.
+ */
+export async function selectFromLeaf(
+  leaf: LeafPath,
+  turn: TurnContext,
+  tree: LeafTree,
+  pageSummary: (slug: Slug) => Promise<string>,
+): Promise<SelectedPage[]> {
+  const members = membersOf(tree, leaf);
+  if (members.length === 0) return [];
+
+  const allMembers = (): SelectedPage[] =>
+    members.map((slug) => ({ slug, pinned: false }));
+
+  const provider = await getConfiguredProvider("memoryV3SelectL2");
+  if (!provider) {
+    log.warn(
+      { leaf },
+      "L2 selector provider unavailable; degrading to deterministic lanes",
+    );
+    return [];
+  }
+
+  const userMsg: Message = {
+    role: "user",
+    content: [
+      cachedTextBlock(
+        `<leaf>${leaf}</leaf>\n` +
+          (await renderPagesBlock(members, pageSummary)),
+      ),
+      {
+        type: "text",
+        text:
+          (turn.situationalContext
+            ? `<situation>${turn.situationalContext}</situation>\n`
+            : "") +
+          `<recent_context>${turn.recentContext}</recent_context>\n` +
+          `<current_message>${turn.currentMessage}</current_message>`,
+      },
+    ],
+  };
+
+  // One forced-tool call, retried a few times so a transient malformed response
+  // (no usable tool_use, or tool input that fails the schema) re-prompts before
+  // we give up. `null` from an attempt means "unusable, retry"; the provider
+  // layer already backs off transient throws, so this loop adds no delay.
+  const parsed = await retryForResult(async () => {
+    const response = await provider.sendMessage([userMsg], {
+      tools: [SELECT_PAGES_TOOL],
+      systemPrompt: SYSTEM_PROMPT,
+      config: {
+        callSite: "memoryV3SelectL2" as const,
+        tool_choice: { type: "tool" as const, name: SELECT_PAGES_TOOL_NAME },
+      },
+    });
+    const toolBlock = extractToolUse(response);
+    if (!toolBlock || toolBlock.name !== SELECT_PAGES_TOOL_NAME) return null;
+    const result = SelectPagesSchema.safeParse(toolBlock.input);
+    return result.success ? result.data : null;
+  });
+
+  if (parsed === null) {
+    log.warn(
+      { leaf },
+      "L2 selector could not obtain a selection after retries; degrading to deterministic lanes",
+    );
+    return [];
+  }
+
+  // Omitted `ids` is the recall-safe "this whole leaf is relevant" signal.
+  // Bounded to one leaf, so it stays a select-all (unlike the L1 router).
+  if (parsed.ids === undefined) return allMembers();
+
+  const pinned = new Set(parsed.pinned_ids ?? []);
+
+  // Map 1-based IDs back to member slugs, dropping out-of-range IDs without
+  // throwing. De-duplicate while preserving model-returned order.
+  const seen = new Set<number>();
+  const selected: SelectedPage[] = [];
+  for (const id of parsed.ids) {
+    if (id < 1 || id > members.length || seen.has(id)) continue;
+    seen.add(id);
+    selected.push({ slug: members[id - 1]!, pinned: pinned.has(id) });
+  }
+  return selected;
+}
+
+/**
+ * Run the L2 selector across every opened leaf with bounded concurrency and
+ * flatten the per-leaf selections.
+ *
+ * A page assigned to more than one opened leaf may appear more than once in the
+ * result; de-duplication across leaves is the orchestrator's job in a later PR.
+ */
+export async function selectAcrossLeaves(
+  leaves: LeafPath[],
+  turn: TurnContext,
+  tree: LeafTree,
+  pageSummary: (slug: Slug) => Promise<string>,
+  concurrency = 4,
+): Promise<SelectedPage[]> {
+  const perLeaf = await mapLimit(leaves, concurrency, (leaf) =>
+    selectFromLeaf(leaf, turn, tree, pageSummary),
+  );
+  return perLeaf.flat();
+}