diff --git a/extensions/memory-core/src/memory/manager-embedding-cache.ts b/extensions/memory-core/src/memory/manager-embedding-cache.ts index f24c3ff5f16..ba797492c1c 100644 --- a/extensions/memory-core/src/memory/manager-embedding-cache.ts +++ b/extensions/memory-core/src/memory/manager-embedding-cache.ts @@ -1,14 +1,10 @@ +import type { DatabaseSync, SQLInputValue } from "node:sqlite"; import { parseEmbedding, type MemoryChunk, } from "openclaw/plugin-sdk/memory-core-host-engine-storage"; -type EmbeddingCacheDb = { - prepare: (sql: string) => { - all: (...args: unknown[]) => Array<{ hash: string; embedding: string }>; - run: (...args: unknown[]) => unknown; - }; -}; +type EmbeddingCacheDb = Pick; type EmbeddingProviderRef = { id: string; @@ -19,11 +15,12 @@ export function loadMemoryEmbeddingCache(params: { db: EmbeddingCacheDb; enabled: boolean; provider: EmbeddingProviderRef | null; - providerKey: string; + providerKey: string | null; hashes: string[]; tableName?: string; }): Map { - if (!params.enabled || !params.provider || params.hashes.length === 0) { + const provider = params.provider; + if (!params.enabled || !provider || !params.providerKey || params.hashes.length === 0) { return new Map(); } const unique: string[] = []; @@ -41,7 +38,7 @@ export function loadMemoryEmbeddingCache(params: { const tableName = params.tableName ?? "embedding_cache"; const out = new Map(); - const baseParams = [params.provider.id, params.provider.model, params.providerKey]; + const baseParams: SQLInputValue[] = [provider.id, provider.model, params.providerKey]; const batchSize = 400; for (let start = 0; start < unique.length; start += batchSize) { const batch = unique.slice(start, start + batchSize); @@ -51,7 +48,7 @@ export function loadMemoryEmbeddingCache(params: { `SELECT hash, embedding FROM ${tableName}\n` + ` WHERE provider = ? AND model = ? AND provider_key = ? AND hash IN (${placeholders})`, ) - .all(...baseParams, ...batch); + .all(...baseParams, ...batch) as Array<{ hash: string; embedding: string }>; for (const row of rows) { out.set(row.hash, parseEmbedding(row.embedding)); } @@ -63,12 +60,13 @@ export function upsertMemoryEmbeddingCache(params: { db: EmbeddingCacheDb; enabled: boolean; provider: EmbeddingProviderRef | null; - providerKey: string; + providerKey: string | null; entries: Array<{ hash: string; embedding: number[] }>; now?: number; tableName?: string; }): void { - if (!params.enabled || !params.provider || params.entries.length === 0) { + const provider = params.provider; + if (!params.enabled || !provider || !params.providerKey || params.entries.length === 0) { return; } const tableName = params.tableName ?? "embedding_cache"; @@ -84,8 +82,8 @@ export function upsertMemoryEmbeddingCache(params: { for (const entry of params.entries) { const embedding = entry.embedding ?? []; stmt.run( - params.provider.id, - params.provider.model, + provider.id, + provider.model, params.providerKey, entry.hash, JSON.stringify(embedding), diff --git a/extensions/memory-core/src/memory/manager-embedding-ops.ts b/extensions/memory-core/src/memory/manager-embedding-ops.ts index 0e86e906cdc..8b1913ecee2 100644 --- a/extensions/memory-core/src/memory/manager-embedding-ops.ts +++ b/extensions/memory-core/src/memory/manager-embedding-ops.ts @@ -15,7 +15,11 @@ import { type MemoryFileEntry, type MemorySource, } from "openclaw/plugin-sdk/memory-core-host-engine-storage"; -import { recordMemoryBatchFailure, resetMemoryBatchFailureState } from "./manager-batch-state.js"; +import { + MEMORY_BATCH_FAILURE_LIMIT, + recordMemoryBatchFailure, + resetMemoryBatchFailureState, +} from "./manager-batch-state.js"; import { collectMemoryCachedEmbeddings, loadMemoryEmbeddingCache, @@ -156,8 +160,9 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps { _entry: MemoryFileEntry | SessionFileEntry, source: MemorySource, ): Promise { + const provider = this.provider; const batchEmbed = this.providerRuntime?.batchEmbed; - if (!this.provider || !batchEmbed) { + if (!provider || !batchEmbed) { return this.embedChunksInBatches(chunks); } if (chunks.length === 0) { @@ -170,7 +175,7 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps { const missingChunks = missing.map((item) => item.chunk); const batchResult = await this.runBatchWithFallback({ - provider: this.provider.id, + provider: provider.id, run: async () => await batchEmbed({ agentId: this.agentId, @@ -199,7 +204,7 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps { upsertMemoryEmbeddingCache({ db: this.db, enabled: this.cache.enabled, - provider: this.provider, + provider, providerKey: this.providerKey, entries: toCache, tableName: EMBEDDING_CACHE_TABLE, @@ -228,19 +233,20 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps { if (texts.length === 0) { return []; } - if (!this.provider) { + const provider = this.provider; + if (!provider) { throw new Error("Cannot embed batch in FTS-only mode (no embedding provider)"); } return await runMemoryEmbeddingRetryLoop({ run: async () => { const timeoutMs = this.resolveEmbeddingTimeout("batch"); log.debug("memory embeddings: batch start", { - provider: this.provider.id, + provider: provider.id, items: texts.length, timeoutMs, }); return await this.withTimeout( - this.provider.embedBatch(texts), + provider.embedBatch(texts), timeoutMs, `memory embeddings batch timed out after ${Math.round(timeoutMs / 1000)}s`, ); @@ -258,19 +264,21 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps { if (inputs.length === 0) { return []; } - if (!this.provider?.embedBatchInputs) { + const provider = this.provider; + const embedBatchInputs = provider?.embedBatchInputs; + if (!embedBatchInputs) { return await this.embedBatchWithRetry(inputs.map((input) => input.text)); } return await runMemoryEmbeddingRetryLoop({ run: async () => { const timeoutMs = this.resolveEmbeddingTimeout("batch"); log.debug("memory embeddings: structured batch start", { - provider: this.provider.id, + provider: provider.id, items: inputs.length, timeoutMs, }); return await this.withTimeout( - this.provider.embedBatchInputs(inputs), + embedBatchInputs(inputs), timeoutMs, `memory embeddings batch timed out after ${Math.round(timeoutMs / 1000)}s`, ); @@ -447,7 +455,7 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps { }); const suffix = failure.disabled ? "disabling batch" : "keeping batch enabled"; log.warn( - `memory embeddings: ${params.provider} batch failed (${failure.count}/${BATCH_FAILURE_LIMIT}); ${suffix}; falling back to non-batch embeddings: ${message}`, + `memory embeddings: ${params.provider} batch failed (${failure.count}/${MEMORY_BATCH_FAILURE_LIMIT}); ${suffix}; falling back to non-batch embeddings: ${message}`, ); return await params.fallback(); }