@vellumai/assistant 0.3.4 → 0.3.6

package/src/memory/retriever.ts

@@ -1,8 +1,9 @@
-import { eq } from 'drizzle-orm';
+import { inArray } from 'drizzle-orm';
 import type { AssistantConfig } from '../config/types.js';
 import { estimateTextTokens } from '../context/token-estimator.js';
 import { getLogger } from '../util/logger.js';
 import { embedWithBackend, getMemoryBackendStatus, logMemoryEmbeddingWarning } from './embedding-backend.js';
+import { computeRetryDelay, isRetryableNetworkError, abortableSleep } from '../util/retry.js';
 import { getDb } from './db.js';
 import { memoryItemSources } from './schema.js';
 import type {
@@ -36,6 +37,37 @@ export {
 
 const log = getLogger('memory-retriever');
 
+const EMBED_MAX_RETRIES = 3;
+const EMBED_BASE_DELAY_MS = 500;
+
+/**
+ * Wrap embedWithBackend with retry + exponential backoff for transient failures
+ * (network errors, 429s, 5xx). Aborts immediately if the caller's signal fires.
+ */
+async function embedWithRetry(
+  config: AssistantConfig,
+  texts: string[],
+  opts?: { signal?: AbortSignal },
+): ReturnType<typeof embedWithBackend> {
+  let lastError: unknown;
+  for (let attempt = 0; attempt <= EMBED_MAX_RETRIES; attempt++) {
+    try {
+      return await embedWithBackend(config, texts, opts);
+    } catch (err) {
+      lastError = err;
+      if (opts?.signal?.aborted || isAbortError(err)) throw err;
+      const isTransient = isRetryableNetworkError(err)
+        || isHttpStatusError(err);
+      if (!isTransient || attempt === EMBED_MAX_RETRIES) throw err;
+      const delay = computeRetryDelay(attempt, EMBED_BASE_DELAY_MS);
+      log.warn({ err, attempt: attempt + 1, delayMs: Math.round(delay) }, 'Transient embedding failure, retrying');
+      await abortableSleep(delay, opts?.signal);
+      if (opts?.signal?.aborted) throw err;
+    }
+  }
+  throw lastError;
+}
+
 /**
  * Build the list of scope IDs to include in queries.
  * - If a `scopePolicyOverride` is provided, it takes precedence over both
@@ -108,26 +140,70 @@ async function collectAndMergeCandidates(
     : [];
 
   // Direct item search supplements FTS with LIKE-based matching.
-  // Overfetch when exclusions are present so that filtering doesn't
-  // silently drop valid candidates just below the SQL LIMIT cutoff.
+  // When exclusions are present, adaptively increase the fetch size until
+  // we collect directLimit valid (non-excluded) items or exhaust the DB.
   const directLimit = Math.max(10, config.memory.retrieval.lexicalTopK);
-  const directFetchLimit = excludeMessageIds.length > 0
-    ? Math.max(directLimit + 24, directLimit * 2)
-    : directLimit;
-  let directItems = directItemSearch(query, directFetchLimit, scopeIds);
-
-  // Filter direct items by excluded message IDs -- items whose only evidence
-  // comes from excluded messages should not leak back into recall.
-  if (excludeMessageIds.length > 0 && directItems.length > 0) {
+
+  // Helper: filter fetched direct items to those with at least one non-excluded source.
+  const filterDirectItems = (items: Candidate[]): Candidate[] => {
+    if (items.length === 0) return items;
     const db = getDb();
-    directItems = directItems.filter((candidate) => {
-      const sources = db.select().from(memoryItemSources)
-        .where(eq(memoryItemSources.memoryItemId, candidate.id)).all();
-      if (sources.length === 0) return true;
-      const nonExcluded = sources.filter((s) => !excludeMessageIds.includes(s.messageId));
-      return nonExcluded.length > 0;
-    });
-    directItems = directItems.slice(0, directLimit);
+    const excludedSet = new Set(excludeMessageIds);
+    const allSources = db.select({
+      memoryItemId: memoryItemSources.memoryItemId,
+      messageId: memoryItemSources.messageId,
+    }).from(memoryItemSources)
+      .where(inArray(memoryItemSources.memoryItemId, items.map((c) => c.id)))
+      .all();
+    const hasNonExcluded = new Set<string>();
+    const hasSources = new Set<string>();
+    for (const s of allSources) {
+      hasSources.add(s.memoryItemId);
+      if (!excludedSet.has(s.messageId)) {
+        hasNonExcluded.add(s.memoryItemId);
+      }
+    }
+    return items.filter((c) => !hasSources.has(c.id) || hasNonExcluded.has(c.id));
+  };
+
+  let directItems: Candidate[];
+  if (excludeMessageIds.length > 0) {
+    const MAX_FETCH = directLimit * 8;
+
+    // Probe: fetch directLimit items and measure how many survive filtering.
+    const probe = directItemSearch(query, directLimit, scopeIds);
+    const probeFiltered = filterDirectItems(probe);
+    const probeExhausted = probe.length < directLimit;
+
+    if (probeFiltered.length >= directLimit || probeExhausted) {
+      directItems = probeFiltered.slice(0, directLimit);
+    } else {
+      // Compute exclusion ratio from probe and extrapolate the fetch size
+      // needed to yield directLimit surviving items in a single query.
+      const exclusionRatio = probe.length > 0
+        ? 1 - probeFiltered.length / probe.length
+        : 0;
+      // Fetch enough to compensate for the observed exclusion rate, with
+      // a 1.5x safety margin to avoid a second round in most cases.
+      const estimatedFetch = exclusionRatio < 1
+        ? Math.ceil((directLimit / (1 - exclusionRatio)) * 1.5)
+        : MAX_FETCH;
+      let fetchSize = Math.min(Math.max(estimatedFetch, directLimit + 24), MAX_FETCH);
+
+      let fetched = directItemSearch(query, fetchSize, scopeIds);
+      directItems = filterDirectItems(fetched).slice(0, directLimit);
+
+      // Retry loop: when the estimate under-fetched (uneven exclusion
+      // distribution), keep increasing fetchSize until quota is met or
+      // the DB is exhausted.
+      while (directItems.length < directLimit && fetched.length === fetchSize && fetchSize < MAX_FETCH) {
+        fetchSize = Math.min(fetchSize * 2, MAX_FETCH);
+        fetched = directItemSearch(query, fetchSize, scopeIds);
+        directItems = filterDirectItems(fetched).slice(0, directLimit);
+      }
+    }
+  } else {
+    directItems = directItemSearch(query, directLimit, scopeIds);
   }
 
   // -- Early termination check --
@@ -161,38 +237,50 @@ async function collectAndMergeCandidates(
     && cheapCandidates.filter((c) => c.lexical >= etConfig.confidenceThreshold).length >= etConfig.minHighConfidence;
 
   // -- Phase 2: expensive searches (skipped on early termination) --
+  // Semantic search (async Qdrant network call) and entity search (sync
+  // SQLite graph traversal) are independent. Start the network call first,
+  // run the sync work while it's in flight, then await the result.
   let semantic: Candidate[] = [];
   let semanticSearchFailed = false;
-  if (queryVector && !canTerminateEarly) {
-    try {
-      semantic = await semanticSearch(queryVector, opts?.provider ?? 'unknown', opts?.model ?? 'unknown', config.memory.retrieval.semanticTopK, excludeMessageIds, scopeIds);
-    } catch (err) {
-      semanticSearchFailed = true;
-      if (isQdrantConnectionError(err)) {
-        log.warn({ err }, 'Qdrant is unavailable — semantic search disabled, memory recall will be degraded');
-      } else {
-        log.warn({ err }, 'Semantic search failed, continuing with other retrieval methods');
-      }
-    }
-  }
-
   let entity: Candidate[] = [];
+  let candidateDepths: Map<string, number> | undefined;
   let relationSeedEntityCount = 0;
   let relationTraversedEdgeCount = 0;
   let relationNeighborEntityCount = 0;
   let relationExpandedItemCount = 0;
-  if (config.memory.entity.enabled && !canTerminateEarly) {
-    const entitySearchResult = entitySearch(
-      query,
-      config.memory.entity,
-      scopeIds,
-      excludeMessageIds,
-    );
-    entity = entitySearchResult.candidates;
-    relationSeedEntityCount = entitySearchResult.relationSeedEntityCount;
-    relationTraversedEdgeCount = entitySearchResult.relationTraversedEdgeCount;
-    relationNeighborEntityCount = entitySearchResult.relationNeighborEntityCount;
-    relationExpandedItemCount = entitySearchResult.relationExpandedItemCount;
+
+  if (!canTerminateEarly) {
+    const semanticPromise = queryVector
+      ? semanticSearch(queryVector, opts?.provider ?? 'unknown', opts?.model ?? 'unknown', config.memory.retrieval.semanticTopK, excludeMessageIds, scopeIds)
+          .catch((err): Candidate[] => {
+            semanticSearchFailed = true;
+            if (isQdrantConnectionError(err)) {
+              log.warn({ err }, 'Qdrant is unavailable — semantic search disabled, memory recall will be degraded');
+            } else {
+              log.warn({ err }, 'Semantic search failed, continuing with other retrieval methods');
+            }
+            return [];
+          })
+      : null;
+
+    if (config.memory.entity.enabled) {
+      const entitySearchResult = entitySearch(
+        query,
+        config.memory.entity,
+        scopeIds,
+        excludeMessageIds,
+      );
+      entity = entitySearchResult.candidates;
+      candidateDepths = entitySearchResult.candidateDepths;
+      relationSeedEntityCount = entitySearchResult.relationSeedEntityCount;
+      relationTraversedEdgeCount = entitySearchResult.relationTraversedEdgeCount;
+      relationNeighborEntityCount = entitySearchResult.relationNeighborEntityCount;
+      relationExpandedItemCount = entitySearchResult.relationExpandedItemCount;
+    }
+
+    if (semanticPromise) {
+      semantic = await semanticPromise;
+    }
   }
 
   if (canTerminateEarly) {
@@ -205,7 +293,10 @@ async function collectAndMergeCandidates(
   const relationScoreMultiplier = config.memory.entity.enabled && config.memory.entity.relationRetrieval.enabled
     ? config.memory.entity.relationRetrieval.neighborScoreMultiplier
     : undefined;
-  const merged = mergeCandidates(lexical, semantic, recency, [...entity, ...directItems], config.memory.retrieval.freshness, relationScoreMultiplier);
+  const depthMap = config.memory.entity.enabled && config.memory.entity.relationRetrieval.depthDecay
+    ? candidateDepths
+    : undefined;
+  const merged = mergeCandidates(lexical, semantic, recency, [...entity, ...directItems], config.memory.retrieval.freshness, relationScoreMultiplier, depthMap);
 
   return {
     lexical,
@@ -256,7 +347,7 @@ export async function buildMemoryRecall(
 
   if (backendStatus.provider) {
     try {
-      const embedded = await embedWithBackend(config, [query], { signal });
+      const embedded = await embedWithRetry(config, [query], { signal });
       queryVector = embedded.vectors[0] ?? null;
       provider = embedded.provider;
       model = embedded.model;
@@ -608,7 +699,7 @@ export async function searchMemoryItems(
   const backendStatus = getMemoryBackendStatus(config);
   if (backendStatus.provider) {
     try {
-      const embedded = await embedWithBackend(config, [trimmed]);
+      const embedded = await embedWithRetry(config, [trimmed]);
       queryVector = embedded.vectors[0] ?? null;
       provider = embedded.provider;
       model = embedded.model;
@@ -680,3 +771,30 @@ function isAbortError(err: unknown): boolean {
   if (!(err instanceof Error)) return false;
   return err.name === 'AbortError' || err.name === 'APIUserAbortError';
 }
+
+/**
+ * Check if an error represents a retryable HTTP status (429 or 5xx).
+ * Checks the error's `status` or `statusCode` property first (set by most
+ * HTTP/API clients), then falls back to looking for "status <code>" patterns
+ * in the message. This avoids false positives from dimension numbers like 512.
+ */
+function getErrorStatusCode(err: Error): unknown {
+  if ('status' in err) return (err as { status: unknown }).status;
+  if ('statusCode' in err) return (err as { statusCode: unknown }).statusCode;
+  return undefined;
+}
+
+function isHttpStatusError(err: unknown): boolean {
+  if (!(err instanceof Error)) return false;
+  const status = getErrorStatusCode(err);
+  if (typeof status === 'number') {
+    return status === 429 || (status >= 500 && status < 600);
+  }
+  // Fall back to message matching, but only for patterns that clearly
+  // indicate an HTTP status code rather than arbitrary numbers.
+  // Matches: "status 503", "HTTP 500", "status code: 502", parenthesized
+  // codes like "failed (503)" from Gemini/Ollama (requires "failed" or
+  // "error" context to avoid false positives from dimension numbers like
+  // 512), and bare "429" (rate-limit).
+  return /\b429\b|(?:failed|error)\s*\((?:429|5\d{2})\)|(?:status|http)\s*(?:code\s*)?:?\s*5\d{2}\b/i.test(err.message);
+}