@vellumai/assistant 0.6.3 → 0.6.5

package/src/memory/graph/retriever.ts

@@ -15,7 +15,11 @@ import {
 import type { ContentBlock, ImageContent } from "../../providers/types.js";
 import { getLogger } from "../../util/logger.js";
 import { embedWithRetry } from "../embed.js";
-import { selectedBackendSupportsMultimodal } from "../embedding-backend.js";
+import {
+  generateSparseEmbedding,
+  selectedBackendSupportsMultimodal,
+} from "../embedding-backend.js";
+import type { QdrantSparseVector } from "../qdrant-client.js";
 import { searchGraphNodes } from "./graph-search.js";
 import type { InContextTracker } from "./injection.js";
 import {
@@ -25,9 +29,15 @@ import {
   computeTemporalBoost,
   PER_TURN_WEIGHTS,
   scoreCandidate,
+  weightsForContextLoad,
 } from "./scoring.js";
 import { sampleSerendipity } from "./serendipity.js";
-import { getEdgesForNode, getNodesByIds, queryNodes } from "./store.js";
+import {
+  getEdgesForNode,
+  getNodesByIds,
+  queryCapabilityNodes,
+  queryNodes,
+} from "./store.js";
 import { getActiveTriggersByType } from "./store.js";
 import {
   evaluateEventTriggers,
@@ -45,6 +55,13 @@ import { isCapabilityNode } from "./types.js";
 
 const log = getLogger("graph-retriever");
 
+function extractCapabilityId(node: MemoryNode): string | null {
+  const match = node.content.match(
+    /^skill:(\S+)\n|^cli:(\S+)\n|^\s*The ".*?" skill \(([^)]+)\)|^\s*The "assistant (\S+)" CLI command/,
+  );
+  return match?.[1] ?? match?.[2] ?? match?.[3] ?? match?.[4] ?? null;
+}
+
 // ---------------------------------------------------------------------------
 // LLM re-ranking + deduplication
 // ---------------------------------------------------------------------------
@@ -80,7 +97,7 @@ async function rerankAndDedup(
   if (candidates.length <= maxNodes) return candidates;
 
   try {
-    const provider = await getConfiguredProvider();
+    const provider = await getConfiguredProvider("memoryRetrieval");
     if (!provider) return candidates.slice(0, maxNodes);
 
     // Numbered listing for the LLM: index + age + full content
@@ -110,7 +127,7 @@ Your job:
 3. Return the IDs in order of importance (most important first).`,
       {
         config: {
-          modelIntent: "quality-optimized" as const,
+          callSite: "memoryRetrieval" as const,
           tool_choice: { type: "tool" as const, name: "select_memories" },
           thinking: { type: "disabled" },
           temperature: 0,
@@ -179,7 +196,7 @@ async function dedupForTurn(
   query: string,
 ): Promise<{ nodes: ScoredNode[]; llmApplied: boolean }> {
   try {
-    const provider = await getConfiguredProvider();
+    const provider = await getConfiguredProvider("memoryRetrieval");
     if (!provider)
       return { nodes: candidates.slice(0, maxNodes), llmApplied: false };
 
@@ -201,7 +218,7 @@ async function dedupForTurn(
       `Dedupe + rerank the following numbered items. Pick the most relevant items to the query. Call the select_items tool.\n\nBe aggressive on dedup — when multiple items describe the same event, fact, or status, keep ONLY the richest version. But be generous on relevance — only cut items that are completely irrelevant to the query. If it's even tangentially related, keep it.`,
       {
         config: {
-          modelIntent: "latency-optimized" as const,
+          callSite: "memoryRetrieval" as const,
           tool_choice: { type: "tool" as const, name: "select_items" },
           thinking: { type: "disabled" },
           temperature: 0,
@@ -272,7 +289,7 @@ async function dedupCrossCategory(
   maxNodes: number,
 ): Promise<ScoredNode[]> {
   try {
-    const provider = await getConfiguredProvider();
+    const provider = await getConfiguredProvider("memoryRetrieval");
     if (!provider) return candidates.slice(0, maxNodes);
 
     const now = Date.now();
@@ -293,7 +310,7 @@ async function dedupCrossCategory(
       `Deduplicate the following numbered items. When multiple items describe the same event, fact, or status, keep ONLY the richest version. Keep ALL items that are not duplicates — do not filter by relevance or topic. Call the select_items tool with every item that survives dedup.`,
       {
         config: {
-          modelIntent: "latency-optimized" as const,
+          callSite: "memoryRetrieval" as const,
           tool_choice: { type: "tool" as const, name: "select_items" },
           thinking: { type: "disabled" },
           temperature: 0,
@@ -346,6 +363,13 @@ export interface ContextLoadOpts {
   serendipitySlots?: number;
   /** Maximum nodes to return (default 40). */
   maxNodes?: number;
+  /**
+   * Optional dedicated user-message query text. When present and non-empty,
+   * `loadContextMemory` embeds this text independently of
+   * `recentSummaries` and uses the resulting vector to rank capability
+   * reserve slots. Leave `undefined` to use summary-only retrieval.
+   */
+  userQuery?: string;
 }
 
 export interface ContextLoadResult {
@@ -354,6 +378,37 @@ export interface ContextLoadResult {
   triggeredNodes: TriggeredResult[];
   latencyMs: number;
   metrics: RetrievalMetrics;
+  /**
+   * Dense query vector computed from `recentSummaries`. Surfaced so downstream
+   * callers (e.g. the PKB hint retriever) can reuse the same embedding for a
+   * second Qdrant query without paying for another embedding call. `undefined`
+   * when no summaries were provided or embedding failed (circuit breaker).
+   */
+  queryVector?: number[];
+  /**
+   * Optional sparse vector passed into `searchGraphNodes` alongside the dense
+   * query vector. Currently always `undefined` — reserved for future hybrid
+   * retrieval that produces a sparse vector at the call site.
+   */
+  sparseVector?: QdrantSparseVector;
+  /**
+   * Dense query vector computed from `opts.userQuery`. Surfaced so
+   * downstream callers (PKB hint search) can prefer it over the
+   * summary-based `queryVector` for user-intent-aligned retrieval.
+   * `undefined` when `userQuery` was not provided, was effectively empty,
+   * or the dedicated embed call was skipped/failed.
+   */
+  userQueryVector?: number[];
+  /**
+   * Sparse (TF-IDF) vector of `opts.userQuery`. Surfaced so PKB hint search
+   * can pair it with `userQueryVector` to run a hybrid dense+sparse query —
+   * RRF fusion captures lexical matches (exact filenames, proper nouns,
+   * uncommon tokens) that pure dense embeddings wash out. Computed locally
+   * (no embedding-service call), so it's cheap to produce whenever the user
+   * query is non-empty. `undefined` on the same conditions as
+   * `userQueryVector`.
+   */
+  userQuerySparseVector?: QdrantSparseVector;
 }
 
 /**
@@ -379,6 +434,7 @@ export async function loadContextMemory(
 
   // 1. Embed recent conversation summaries as retrieval queries
   let queryVector: number[] | null = null;
+  const sparseVector: QdrantSparseVector | undefined = undefined;
   let embeddingProvider: string | null = null;
   let embeddingModel: string | null = null;
   let contextQueryText: string | null = null;
@@ -399,15 +455,51 @@ export async function loadContextMemory(
     }
   }
 
+  // 1b. Dedicated user-query embedding. Always embed the user query
+  //     independently when present. Summaries and the user query are
+  //     disjoint signals, so both vectors carry unique retrieval value —
+  //     especially in workloads with short summaries and a substantive
+  //     user question.
+  let userQueryVector: number[] | null = null;
+  let userQuerySparseVector: QdrantSparseVector | undefined = undefined;
+  const userQueryCandidateIds = new Map<string, number>(); // nodeId → score
+  const trimmedUserQuery = opts.userQuery?.trim() ?? "";
+  const shouldEmbedUserQuery = trimmedUserQuery.length > 0;
+  if (shouldEmbedUserQuery) {
+    try {
+      const result = await embedWithRetry(opts.config, [trimmedUserQuery], {
+        signal: opts.signal,
+      });
+      userQueryVector = result.vectors[0] ?? null;
+      if (!embeddingProvider) {
+        embeddingProvider = result.provider;
+        embeddingModel = result.model;
+      }
+    } catch (err) {
+      log.warn({ err }, "Failed to embed userQuery for context load");
+    }
+    // Sparse embedding is a local TF-IDF computation — no network call, so
+    // compute it independently of the dense embed. Even if the dense call
+    // failed, a sparse vector is still useful for downstream consumers that
+    // can operate on it alone.
+    const sparse = generateSparseEmbedding(trimmedUserQuery);
+    if (sparse.indices.length > 0) {
+      userQuerySparseVector = sparse;
+    }
+  }
+
   // 2. Hybrid retrieval from Qdrant (dense search on graph_node points)
   const semanticCandidateIds = new Map<string, number>(); // nodeId → score
   let hybridSearchLatencyMs = 0;
   if (queryVector) {
     const searchStart = Date.now();
     try {
-      const results = await searchGraphNodes(queryVector, maxNodes * 3, [
-        opts.scopeId,
-      ]);
+      const results = await searchGraphNodes(
+        queryVector,
+        maxNodes * 3,
+        [opts.scopeId],
+        sparseVector,
+      );
       for (const r of results) {
         semanticCandidateIds.set(r.nodeId, r.score);
       }
@@ -419,6 +511,30 @@ export async function loadContextMemory(
   }
   const pureSemanticHits = semanticCandidateIds.size;
 
+  // 2b. Run a parallel Qdrant search against the user-query vector and merge
+  //     the results into the organic scoring pool via max-score union: a node
+  //     hit by only the user-query vector still participates in downstream
+  //     scoring, and a node hit by both vectors retains the higher score.
+  if (userQueryVector) {
+    try {
+      const results = await searchGraphNodes(
+        userQueryVector,
+        maxNodes * 3,
+        [opts.scopeId],
+        undefined,
+      );
+      for (const r of results) {
+        userQueryCandidateIds.set(r.nodeId, r.score);
+        const existing = semanticCandidateIds.get(r.nodeId);
+        if (existing === undefined || r.score > existing) {
+          semanticCandidateIds.set(r.nodeId, r.score);
+        }
+      }
+    } catch (err) {
+      log.warn({ err }, "Qdrant search failed for userQuery vector");
+    }
+  }
+
   // Also include top-significance nodes as a fallback
   const topSignificance = queryNodes({
     scopeId: opts.scopeId,
@@ -427,7 +543,7 @@ export async function loadContextMemory(
   });
   for (const node of topSignificance) {
     if (!semanticCandidateIds.has(node.id)) {
-      semanticCandidateIds.set(node.id, 0); // no semantic score, ranked by significance
+      semanticCandidateIds.set(node.id, 0); // no score from either Qdrant query, ranked by significance only
     }
   }
 
@@ -530,77 +646,107 @@ export async function loadContextMemory(
     const recency = computeRecencyBoost(node, nowMs);
 
     scored.push(
-      scoreCandidate(node, {
-        semanticSimilarity: semanticSim,
-        effectiveSignificance: effectiveSig,
-        emotionalIntensity: node.emotionalCharge.intensity,
-        temporalBoost: normalizedTemporal,
-        recencyBoost: recency,
-        triggerBoost,
-        activationBoost: activation,
-      }),
+      scoreCandidate(
+        node,
+        {
+          semanticSimilarity: semanticSim,
+          effectiveSignificance: effectiveSig,
+          emotionalIntensity: node.emotionalCharge.intensity,
+          temporalBoost: normalizedTemporal,
+          recencyBoost: recency,
+          triggerBoost,
+          activationBoost: activation,
+        },
+        weightsForContextLoad(node),
+      ),
     );
   }
 
   // Sort by score descending
   scored.sort((a, b) => b.score - a.score);
 
-  // 5b. Reserve slots for skill/CLI capabilities. Queried directly from
-  // SQLite — no Qdrant vectors needed — so capabilities surface even on
-  // fresh assistants whose embedding jobs haven't completed yet.
+  // 5b. Reserve slots for skill/CLI capabilities.
+  //
+  // Source candidates from the hydrated semantic-search set (the same
+  // strategy `retrieveForTurn` uses) so ranking reflects query relevance.
+  // For cold-start cases (capability nodes exist in SQLite but their
+  // embeddings haven't landed in Qdrant yet), fall back to a narrow
+  // SQL pull that matches only capability-shaped content so organic
+  // procedurals can't crowd the pool.
   const capabilityReserve = ctxLoadCfg.capabilityReserve;
-  const rawCapabilityNodes =
-    capabilityReserve > 0
-      ? queryNodes({
-          scopeId: opts.scopeId,
-          types: ["procedural"],
-          fidelityNot: ["gone"],
-          limit: capabilityReserve * 4,
-        })
-      : [];
+  const capabilityEntries: { node: MemoryNode; sim: number }[] = [];
+  if (capabilityReserve > 0) {
+    const uniqueCapabilityIds = new Set<string>();
+    let untaggedCount = 0;
+    for (const [nodeId, node] of nodeMap) {
+      if (node.fidelity === "gone") continue;
+      if (!isCapabilityNode(node)) continue;
+      const sim =
+        userQueryCandidateIds.get(nodeId) ??
+        semanticCandidateIds.get(nodeId) ??
+        0;
+      capabilityEntries.push({ node, sim });
+      const capId = extractCapabilityId(node);
+      if (capId) uniqueCapabilityIds.add(capId);
+      else untaggedCount++;
+    }
+
+    // Gate the fallback on distinct capability IDs (plus any entries
+    // whose content didn't match a known ID pattern), not raw entry
+    // count — duplicate capability-ID seeding formats can otherwise push
+    // `capabilityEntries.length` past the threshold and then collapse
+    // below it during the dedup pass below.
+    const distinctCount = uniqueCapabilityIds.size + untaggedCount;
+    if (distinctCount < capabilityReserve) {
+      const alreadySeen = new Set(capabilityEntries.map((e) => e.node.id));
+      const fallback = queryCapabilityNodes(
+        opts.scopeId,
+        capabilityReserve * 4,
+      );
+      for (const node of fallback) {
+        if (alreadySeen.has(node.id)) continue;
+        const sim =
+          userQueryCandidateIds.get(node.id) ??
+          semanticCandidateIds.get(node.id) ??
+          0;
+        capabilityEntries.push({ node, sim });
+      }
+    }
+  }
+
+  capabilityEntries.sort((a, b) => b.sim - a.sim);
 
   // Dedup: both seeding systems may create nodes for the same capability.
   // Extract capability ID from content and keep only the first node per ID.
   const seenCapabilityIds = new Set<string>();
-  const capabilityNodes = rawCapabilityNodes
-    .filter(isCapabilityNode)
-    .filter((node) => {
-      const match = node.content.match(
-        /^skill:(\S+)\n|^cli:(\S+)\n|^\s*The ".*?" skill \(([^)]+)\)|^\s*The "assistant (\S+)" CLI command/,
-      );
-      const capId = match?.[1] ?? match?.[2] ?? match?.[3] ?? match?.[4];
-      if (capId) {
-        if (seenCapabilityIds.has(capId)) return false;
-        seenCapabilityIds.add(capId);
-      }
-      return true;
-    });
-
-  // Rank by semantic similarity when a query vector exists
-  let selectedCapabilities: MemoryNode[];
-  if (queryVector && capabilityNodes.length > capabilityReserve) {
-    selectedCapabilities = capabilityNodes
-      .map((node) => ({ node, sim: semanticCandidateIds.get(node.id) ?? 0 }))
-      .sort((a, b) => b.sim - a.sim)
-      .slice(0, capabilityReserve)
-      .map((e) => e.node);
-  } else {
-    selectedCapabilities = capabilityNodes.slice(0, capabilityReserve);
+  const selectedCapabilities: MemoryNode[] = [];
+  for (const { node } of capabilityEntries) {
+    if (selectedCapabilities.length >= capabilityReserve) break;
+    const capId = extractCapabilityId(node);
+    if (capId) {
+      if (seenCapabilityIds.has(capId)) continue;
+      seenCapabilityIds.add(capId);
+    }
+    selectedCapabilities.push(node);
   }
 
   const reservedCapabilities: ScoredNode[] = selectedCapabilities.map(
     (node) => {
       const existing = scored.find((s) => s.node.id === node.id);
       if (existing) return existing;
-      return scoreCandidate(node, {
-        semanticSimilarity: semanticCandidateIds.get(node.id) ?? 0,
-        effectiveSignificance: computeEffectiveSignificance(node, nowMs),
-        emotionalIntensity: node.emotionalCharge.intensity,
-        temporalBoost: (computeTemporalBoost(node, now) + 1) / 2,
-        recencyBoost: 0,
-        triggerBoost: 0,
-        activationBoost: 0,
-      });
+      return scoreCandidate(
+        node,
+        {
+          semanticSimilarity: semanticCandidateIds.get(node.id) ?? 0,
+          effectiveSignificance: computeEffectiveSignificance(node, nowMs),
+          emotionalIntensity: node.emotionalCharge.intensity,
+          temporalBoost: (computeTemporalBoost(node, now) + 1) / 2,
+          recencyBoost: 0,
+          triggerBoost: 0,
+          activationBoost: 0,
+        },
+        weightsForContextLoad(node),
+      );
     },
   );
 
@@ -688,6 +834,10 @@ export async function loadContextMemory(
       queryContext: contextQueryText,
       topCandidates,
     },
+    queryVector: queryVector ?? undefined,
+    sparseVector,
+    userQueryVector: userQueryVector ?? undefined,
+    userQuerySparseVector,
   };
 }
 
@@ -717,6 +867,21 @@ export interface TurnRetrievalResult {
   triggeredNodes: TriggeredResult[];
   latencyMs: number;
   metrics: RetrievalMetrics;
+  /**
+   * Dense query vector computed from the last-exchange text (assistant +
+   * user message). Surfaced so downstream callers (e.g. the PKB hint
+   * retriever in `applyRuntimeInjections`) can reuse the same embedding
+   * for a second Qdrant query without paying for another embedding call.
+   * `undefined` when no text was embedded (image-only turn) or embedding
+   * failed (circuit breaker).
+   */
+  queryVector?: number[];
+  /**
+   * Optional sparse vector passed alongside `queryVector`. Currently always
+   * `undefined` — reserved for future hybrid retrieval that produces a
+   * sparse vector at the call site.
+   */
+  sparseVector?: QdrantSparseVector;
 }
 
 /**
@@ -757,6 +922,24 @@ export async function retrieveForTurn(
     .filter((m) => m.length > 0)
     .join("\n\n");
 
+  // Sparse (TF-IDF) vector of the user's last message only. Surfaced so PKB
+  // 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.
+  const trimmedUserLast = opts.userLastMessage.trim();
+  let perTurnSparseVector: QdrantSparseVector | undefined = undefined;
+  if (trimmedUserLast.length > 0) {
+    const sparse = generateSparseEmbedding(trimmedUserLast);
+    if (sparse.indices.length > 0) {
+      perTurnSparseVector = sparse;
+    }
+  }
+
   // Image-to-image search: embed incoming user images as queries
   // Runs before the text-empty early return so image-only turns are handled
   const imageBlocks = (opts.userLastMessageBlocks ?? []).filter(
@@ -819,6 +1002,8 @@ export async function retrieveForTurn(
         embeddingModel,
         queryContext: queryText || null,
       },
+      queryVector: undefined,
+      sparseVector: perTurnSparseVector,
     };
   }
 
@@ -891,6 +1076,8 @@ export async function retrieveForTurn(
             embeddingModel,
             queryContext: queryText || null,
           },
+          queryVector: undefined,
+          sparseVector: perTurnSparseVector,
         };
       }
     }
@@ -943,6 +1130,8 @@ export async function retrieveForTurn(
         embeddingModel,
         queryContext: queryText || null,
       },
+      queryVector: queryEmbeddings[0],
+      sparseVector: perTurnSparseVector,
     };
   }
 
@@ -1125,5 +1314,7 @@ export async function retrieveForTurn(
       queryContext: queryText || null,
       topCandidates,
     },
+    queryVector: queryEmbeddings[0],
+    sparseVector: perTurnSparseVector,
   };
 }