diff --git a/Memory/readme.md b/Memory/readme.md index 0770505c..14e79817 100644 --- a/Memory/readme.md +++ b/Memory/readme.md @@ -72,6 +72,12 @@ The `MEMMY_MEMORY_HOST`, `MEMMY_MEMORY_PORT`, and `MEMMY_MEMORY_DB` environment variables override the corresponding server settings. The `MEMORY_SERVICE_*` aliases are also accepted. +Known OpenAI embedding model names are tokenized and split automatically with +a safe 7,500-token budget. When an OpenAI-compatible endpoint uses an opaque +deployment alias, set `memmyMemory.embedding.maxInputTokens` or +`MEMMY_EMBEDDING_MAX_INPUT_TOKENS` to a safe per-input budget so the same +token-aware splitting is enabled. + When `storage.token`, `MEMMY_MEMORY_TOKEN`, or `MEMORY_SERVICE_TOKEN` is set, all HTTP routes except `GET /api/v1/health` require that token as a bearer token or `x-api-key`. diff --git a/Memory/src/config/index.ts b/Memory/src/config/index.ts index cfc9a21c..1f8bc3e7 100644 --- a/Memory/src/config/index.ts +++ b/Memory/src/config/index.ts @@ -88,6 +88,7 @@ export interface EmbeddingConfig { extraBody?: Record; actualModelContext?: ActualModelContext; selectionError?: "model_selection_unavailable"; + maxInputTokens?: number; batchSize: number; timeoutMs: number; maxRetries: number; @@ -571,6 +572,7 @@ function configFromEnv(): Record { endpoint: process.env.MEMMY_EMBEDDING_ENDPOINT, model: process.env.MEMMY_EMBEDDING_MODEL, apiKey: process.env.MEMMY_EMBEDDING_API_KEY, + maxInputTokens: numberEnv("MEMMY_EMBEDDING_MAX_INPUT_TOKENS"), batchSize: numberEnv("MEMMY_EMBEDDING_BATCH_SIZE"), timeoutMs: numberEnv("MEMMY_EMBEDDING_TIMEOUT_MS"), maxRetries: numberEnv("MEMMY_EMBEDDING_MAX_RETRIES") @@ -703,6 +705,7 @@ function normalizeEmbedding(input: Record): EmbeddingConfig { selectionError: input.selectionError === "model_selection_unavailable" ? input.selectionError : undefined, + maxInputTokens: positiveInteger(input.maxInputTokens), batchSize: numberValue(input.batchSize, DEFAULT_MEMMY_CONFIG.embedding.batchSize), timeoutMs: numberValue(input.timeoutMs, DEFAULT_MEMMY_CONFIG.embedding.timeoutMs), maxRetries: numberValue(input.maxRetries, DEFAULT_MEMMY_CONFIG.embedding.maxRetries), @@ -1419,6 +1422,12 @@ function numberValue(value: unknown, fallback: number): number { return typeof value === "number" && Number.isFinite(value) ? value : fallback; } +function positiveInteger(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) && value > 0 + ? Math.floor(value) + : undefined; +} + function booleanValue(value: unknown, fallback: boolean): boolean { return typeof value === "boolean" ? value : fallback; } diff --git a/Memory/src/model/embedder.ts b/Memory/src/model/embedder.ts index c4f35d47..a7aac71b 100644 --- a/Memory/src/model/embedder.ts +++ b/Memory/src/model/embedder.ts @@ -221,7 +221,7 @@ class HttpEmbedder implements Embedder { throw new Error(`${provider} embedding provider requires apiKey or endpoint`); } const plan = provider === "openai_compatible" - ? planOpenAiEmbeddingInputs(texts, this.config.model) + ? planOpenAiEmbeddingInputs(texts, this.config.model, this.config.maxInputTokens) : null; if (!plan) return this.requestOpenAiShape(texts, provider, url, role); diff --git a/Memory/src/model/openai-embedding-inputs.ts b/Memory/src/model/openai-embedding-inputs.ts index 7bd9cfc6..30e79632 100644 --- a/Memory/src/model/openai-embedding-inputs.ts +++ b/Memory/src/model/openai-embedding-inputs.ts @@ -16,21 +16,26 @@ export interface OpenAiEmbeddingPlan { let encoder: ReturnType | undefined; -export function planOpenAiEmbeddingInputs(texts: string[], model?: string): OpenAiEmbeddingPlan | null { - if (!isKnownOpenAiEmbeddingModel(model)) return null; +export function planOpenAiEmbeddingInputs( + texts: string[], + model?: string, + configuredMaxInputTokens?: number +): OpenAiEmbeddingPlan | null { + const inputTokenBudget = resolveInputTokenBudget(model, configuredMaxInputTokens); + if (!inputTokenBudget) return null; encoder ??= get_encoding("cl100k_base"); const encoded = texts.map((text) => Array.from(encoder!.encode(text, [], []))); const totalTokens = encoded.reduce((sum, tokens) => sum + tokens.length, 0); if (totalTokens <= OPENAI_EMBEDDING_BATCH_TOKEN_BUDGET && - encoded.every((tokens) => tokens.length <= OPENAI_EMBEDDING_INPUT_TOKEN_BUDGET)) return null; + encoded.every((tokens) => tokens.length <= inputTokenBudget)) return null; const chunks = encoded.flatMap((tokens, originalIndex) => { if (tokens.length === 0) return [{ originalIndex, tokens }]; const items: OpenAiEmbeddingChunk[] = []; - for (let offset = 0; offset < tokens.length; offset += OPENAI_EMBEDDING_INPUT_TOKEN_BUDGET) { + for (let offset = 0; offset < tokens.length; offset += inputTokenBudget) { items.push({ originalIndex, - tokens: tokens.slice(offset, offset + OPENAI_EMBEDDING_INPUT_TOKEN_BUDGET) + tokens: tokens.slice(offset, offset + inputTokenBudget) }); } return items; @@ -69,6 +74,14 @@ function isKnownOpenAiEmbeddingModel(model?: string): boolean { return /(?:^|[/.:])text-embedding-(?:3-(?:small|large)|ada-002)(?:$|[/.:])/i.test(model?.trim() ?? ""); } +function resolveInputTokenBudget(model?: string, configured?: number): number | null { + const explicit = typeof configured === "number" && Number.isFinite(configured) && configured > 0 + ? Math.floor(configured) + : undefined; + if (!isKnownOpenAiEmbeddingModel(model) && !explicit) return null; + return Math.min(explicit ?? OPENAI_EMBEDDING_INPUT_TOKEN_BUDGET, OPENAI_EMBEDDING_INPUT_TOKEN_BUDGET); +} + function batchChunks(chunks: OpenAiEmbeddingChunk[]): OpenAiEmbeddingChunk[][] { const batches: OpenAiEmbeddingChunk[][] = []; let current: OpenAiEmbeddingChunk[] = []; diff --git a/Memory/src/service/worker/worker-runner.ts b/Memory/src/service/worker/worker-runner.ts index b45cf8d8..76752e53 100644 --- a/Memory/src/service/worker/worker-runner.ts +++ b/Memory/src/service/worker/worker-runner.ts @@ -440,7 +440,7 @@ export class WorkerRunner { } } catch (error) { const classification = classifyProcessingError(error); - if (classification.code === "model_input_too_long" || classification.code === "invalid_model_request") { + if (classification.code === "model_input_too_long" && batch.length > 1) { for (const item of batch) results.push(await this.runLeasedEmbeddingItem(item)); continue; } @@ -626,7 +626,7 @@ export class WorkerRunner { } } catch (error) { const classification = classifyProcessingError(error); - if (classification.code === "model_input_too_long" || classification.code === "invalid_model_request") { + if (classification.code === "model_input_too_long" && batch.length > 1) { for (const item of batch) { results.push(await this.runClaimedEmbeddingRetryItem(item.retry, item.claim, item.attemptNo)); } diff --git a/Memory/tests/config.test.ts b/Memory/tests/config.test.ts index 304d27c5..054c0757 100644 --- a/Memory/tests/config.test.ts +++ b/Memory/tests/config.test.ts @@ -99,6 +99,26 @@ describe("memmy memory config", () => { expect(loadMemmyConfig(configPath).config.summary.timeoutMs).toBe(240_000); }); + it("preserves an explicit embedding token budget and allows an environment override", () => { + const root = tempRoot(); + const configPath = join(root, "config.yaml"); + writeFileSync(configPath, YAML.stringify({ + memmyMemory: { + embedding: { + mode: "custom", + provider: "openai_compatible", + endpoint: "https://embedding.example/v1", + model: "deployment-alias", + maxInputTokens: 1_200 + } + } + })); + + expect(loadMemmyConfig(configPath).config.embedding.maxInputTokens).toBe(1_200); + setEnv("MEMMY_EMBEDDING_MAX_INPUT_TOKENS", "640"); + expect(loadMemmyConfig(configPath).config.embedding.maxInputTokens).toBe(640); + }); + it("expands home-relative sqlite paths from config files", () => { const root = tempRoot(); const configPath = join(root, "config.yaml"); diff --git a/Memory/tests/embedder.test.ts b/Memory/tests/embedder.test.ts index 1daed46b..3f94e71e 100644 --- a/Memory/tests/embedder.test.ts +++ b/Memory/tests/embedder.test.ts @@ -112,6 +112,32 @@ describe("embedder", () => { expect(vectors[1]).toEqual([0, 2]); }); + it("uses an explicit token budget for an OpenAI-compatible deployment alias", async () => { + const sentInputs: number[][] = []; + vi.stubGlobal("fetch", vi.fn(async (_url, init) => { + const body = JSON.parse(String(init?.body)) as { input: number[][] }; + sentInputs.push(...body.input); + return new Response(JSON.stringify({ + data: body.input.map(() => ({ embedding: [1, 0] })) + }), { status: 200, headers: { "content-type": "application/json" } }); + })); + const embedder = createEmbedder({ + ...DEFAULT_MEMMY_CONFIG.embedding, + provider: "openai_compatible", + endpoint: "https://api.example.test/v1", + model: "production-embedding-deployment", + apiKey: "sk-test", + maxInputTokens: 512, + cache: false, + maxRetries: 0 + }); + + await expect(embedder.embedOne(" memory".repeat(600))).resolves.toEqual([1, 0]); + + expect(sentInputs.length).toBeGreaterThan(1); + expect(sentInputs.every((input) => input.length <= 512)).toBe(true); + }); + it("keeps chunked OpenAI embedding request batches below the aggregate token budget", async () => { const requestTokenCounts: number[] = []; vi.stubGlobal("fetch", vi.fn(async (_url, init) => { diff --git a/Memory/tests/service/embedding/embedding-processing.test.ts b/Memory/tests/service/embedding/embedding-processing.test.ts index bd9ce7ee..519cc334 100644 --- a/Memory/tests/service/embedding/embedding-processing.test.ts +++ b/Memory/tests/service/embedding/embedding-processing.test.ts @@ -243,7 +243,8 @@ describe("MemoryService / embedding / processing", () => { }); it("does not enqueue a legacy retry for a processing-less deterministic worker failure", async () => { - const { db, service } = createTestService({ embedder: createSelectiveFailureEmbedder() }); + const calls = { batch: 0, single: 0 }; + const { db, service } = createTestService({ embedder: createSelectiveFailureEmbedder(calls) }); const repositories = new Repositories(db.db); const memory = skillMemory(undefined, { id: "skill_deterministic_worker_failure", @@ -275,6 +276,35 @@ describe("MemoryService / embedding / processing", () => { expect(db.db.prepare( `SELECT id FROM embedding_retry_queue WHERE target_id = ?` ).all(memory.id)).toEqual([]); + expect(calls).toEqual({ batch: 1, single: 0 }); + db.close(); + }); + + it("does not re-request a single legacy retry after a token-limit failure", async () => { + const calls = { batch: 0, single: 0 }; + const { db, service } = createTestService({ embedder: createSelectiveFailureEmbedder(calls) }); + const repositories = new Repositories(db.db); + const memory = skillMemory(undefined, { + id: "skill_legacy_token_limit", + content: "BAD_EMBEDDING_ITEM" + }); + repositories.memories.insert(memory); + const retry = repositories.runtime.enqueueEmbeddingRetry({ + targetKind: "skill", + targetId: memory.id, + vectorField: "vec", + sourceText: embeddingTextForMemory(memory), + embedRole: "query", + now: Date.now() - 1 + }); + + await service.runWorkerOnce(10); + + expect(repositories.runtime.getEmbeddingRetry(retry.id)).toMatchObject({ + status: "failed", + attempts: 1 + }); + expect(calls).toEqual({ batch: 1, single: 0 }); db.close(); }); @@ -594,7 +624,7 @@ function createFlakyEmbedder(): Embedder { }; } -function createSelectiveFailureEmbedder(): Embedder { +function createSelectiveFailureEmbedder(calls?: { batch: number; single: number }): Embedder { const inputTooLong = () => new ModelHttpError( "openai_compatible HTTP 400: maximum context length exceeded", "openai_compatible", @@ -612,11 +642,13 @@ function createSelectiveFailureEmbedder(): Embedder { return true; }, async embed(texts: string[]) { + if (calls) calls.batch += 1; if (texts.length > 1) throw inputTooLong(); if (texts[0]?.includes("BAD_EMBEDDING_ITEM")) throw inputTooLong(); return texts.map((text) => stableTestVector(text)); }, async embedOne(text: string) { + if (calls) calls.single += 1; if (text.includes("BAD_EMBEDDING_ITEM")) throw inputTooLong(); return stableTestVector(text); }, diff --git a/Memory/viewer/src/stores/i18n.ts b/Memory/viewer/src/stores/i18n.ts index f6998f6f..3ca389b7 100644 --- a/Memory/viewer/src/stores/i18n.ts +++ b/Memory/viewer/src/stores/i18n.ts @@ -898,7 +898,7 @@ const en = { "Configure an embedding provider before repairing or rebuilding vectors.", "settings.embedding.maxInputTokens.label": "Maximum input tokens", "settings.embedding.maxInputTokens.hint": - "Defaults to 1024; use 0 for no client-side limit. Longer inputs are sampled into chunks and pooled; rebuild vectors after changing it.", + "Known OpenAI embedding models use a safe 7,500-token budget automatically. Set a lower budget for custom deployment aliases; rebuild vectors after changing it.", "settings.embedding.providerBatchSize.label": "Embedding API batch size", "settings.embedding.providerBatchSize.hint": "Maximum texts per provider request. Rejected oversized batches are split automatically.", @@ -1854,7 +1854,7 @@ const zh: Record = { "可用 {ready}/{total};缺失 {missing};维度不匹配 {mismatch};当前维度 {dim}。", "settings.embedding.maintenance.unavailable": "请先配置嵌入模型,再修复或重建向量。", "settings.embedding.maxInputTokens.label": "单条输入最大 Token 数", - "settings.embedding.maxInputTokens.hint": "默认 1024;设为 0 表示不启用客户端限制。超长输入会分块采样并聚合向量,修改后请重建向量。", + "settings.embedding.maxInputTokens.hint": "已知 OpenAI Embedding 模型自动使用安全的 7500 Token 预算;自定义部署别名可设置更低预算,修改后请重建向量。", "settings.embedding.providerBatchSize.label": "Embedding API 批量大小", "settings.embedding.providerBatchSize.hint": "单次模型请求最多发送的文本数;超限失败时会自动拆批。", "settings.embedding.repair": "修复缺失/错维",