diff --git a/apps/memos-local-plugin/core/llm/client.ts b/apps/memos-local-plugin/core/llm/client.ts index ee456ac12..f44e0cb25 100644 --- a/apps/memos-local-plugin/core/llm/client.ts +++ b/apps/memos-local-plugin/core/llm/client.ts @@ -45,6 +45,8 @@ import type { } from "./types.js"; const DEFAULT_MAX_TOKENS = 1024; +// Upper bound for the one-shot truncation retry budget (see completeJson). +const LENGTH_RETRY_MAX_TOKENS_CEILING = 32_768; // ─── Factory ───────────────────────────────────────────────────────────────── @@ -501,16 +503,51 @@ export function createLlmClientWithProvider( const messages = normalizeMessages(input); const systemHint = buildJsonSystemHint(opts.schemaHint); const msgs = ensureJsonWordInUserMessage(inject(messages, systemHint)); - const call = buildCallInput(opts, true); + let call = buildCallInput(opts, true); const op = opts.op ?? "complete.json"; - const maxMalformedRetries = Math.max(0, opts.malformedRetries ?? 1); + let maxMalformedRetries = Math.max(0, opts.malformedRetries ?? 1); let attempt = 0; let lastRaw = ""; let lastErr: unknown = null; + // Thinking models share one max_tokens budget between reasoning and the + // final answer, so a truncated response arrives as 200 OK with + // finish_reason="length". Without an explicit check it looks like plain + // malformed JSON and the retry re-sends the same doomed budget. On the + // first truncation, log diagnostics and double max_tokens for the next + // attempt (capped). + let truncatedOnce = false; while (attempt <= maxMalformedRetries) { attempt++; const { completion } = await callWithFallback(msgs, call, opts, op); + if (completion.finishReason === "length") { + jsonLog.warn("max_tokens_truncated", { + op, + attempt, + maxTokens: call.maxTokens, + completionChars: completion.text.length, + usage: completion.usage + ? { + completionTokens: completion.usage.completionTokens, + totalTokens: completion.usage.totalTokens, + } + : undefined, + }); + if (!truncatedOnce && (call.maxTokens ?? 0) < LENGTH_RETRY_MAX_TOKENS_CEILING) { + truncatedOnce = true; + // The upgraded budget is useless without at least one more attempt: + // a caller may pass malformedRetries: 0 for fail-fast parsing, and + // the truncation retry must not be silently skipped then. + if (attempt > maxMalformedRetries) maxMalformedRetries = attempt; + call = { + ...call, + maxTokens: Math.min( + LENGTH_RETRY_MAX_TOKENS_CEILING, + Math.max(2 * (call.maxTokens ?? DEFAULT_MAX_TOKENS), DEFAULT_MAX_TOKENS), + ), + }; + } + } lastRaw = completion.text; try { const parsed = opts.parse diff --git a/apps/memos-local-plugin/tests/unit/llm/length-retry-budget.test.ts b/apps/memos-local-plugin/tests/unit/llm/length-retry-budget.test.ts new file mode 100644 index 000000000..70a7060e6 --- /dev/null +++ b/apps/memos-local-plugin/tests/unit/llm/length-retry-budget.test.ts @@ -0,0 +1,128 @@ +/** + * Truncation retry budget: finish_reason="length" is a budget problem, not + * plain malformed output. The retry must double max_tokens (once, capped) + * instead of re-sending the same doomed budget. + */ +import { beforeAll, describe, expect, it } from "vitest"; + +import { MemosError } from "../../../agent-contract/errors.js"; +import { createLlmClientWithProvider } from "../../../core/llm/index.js"; +import { initTestLogger } from "../../../core/logger/index.js"; +import type { + LlmConfig, + LlmMessage, + LlmProvider, + LlmProviderCtx, + LlmProviderName, + ProviderCallInput, + ProviderCompletion, +} from "../../../core/llm/types.js"; + +beforeAll(async () => { + await initTestLogger(); +}); + +function cfg(partial: Partial = {}): LlmConfig { + return { + provider: "openai_compatible", + model: "gpt-test", + endpoint: "", + apiKey: "X", + temperature: 0.3, + fallbackToHost: false, + timeoutMs: 5_000, + maxRetries: 0, + ...partial, + }; +} + +class StubProvider implements LlmProvider { + public inputs: ProviderCallInput[] = []; + public readonly name: LlmProviderName; + constructor( + name: LlmProviderName, + private readonly responder: (n: number) => ProviderCompletion, + ) { + this.name = name; + } + async complete( + _messages: LlmMessage[], + opts: ProviderCallInput, + _ctx: LlmProviderCtx, + ): Promise { + this.inputs.push(opts); + return this.responder(this.inputs.length); + } +} + +const DEFAULT_MAX_TOKENS = 1024; + +describe("completeJson truncation retry budget", () => { + it("doubles max_tokens after finish_reason=length and parses the retry", async () => { + const stub = new StubProvider("openai_compatible", (n) => + n === 1 + ? { text: "{\"x\":", durationMs: 1, finishReason: "length" as const } + : { text: "{\"x\":1}", durationMs: 1, finishReason: "stop" as const }, + ); + const client = createLlmClientWithProvider(cfg(), stub); + const r = await client.completeJson<{ x: number }>("ask", { malformedRetries: 1 }); + expect(r.value.x).toBe(1); + expect(stub.inputs).toHaveLength(2); + expect(stub.inputs[0]!.maxTokens).toBe(DEFAULT_MAX_TOKENS); + expect(stub.inputs[1]!.maxTokens).toBe(2 * DEFAULT_MAX_TOKENS); + }); + + it("still retries once on truncation when malformedRetries is 0", async () => { + const stub = new StubProvider("openai_compatible", (n) => + n === 1 + ? { text: "{\"q\":", durationMs: 1, finishReason: "length" as const } + : { text: "{\"q\":2}", durationMs: 1, finishReason: "stop" as const }, + ); + const client = createLlmClientWithProvider(cfg(), stub); + const r = await client.completeJson<{ q: number }>("ask", { malformedRetries: 0 }); + expect(r.value.q).toBe(2); + expect(stub.inputs).toHaveLength(2); + expect(stub.inputs[1]!.maxTokens).toBe(2 * DEFAULT_MAX_TOKENS); + }); + + it("caps the doubled budget at 32768", async () => { + const stub = new StubProvider("openai_compatible", (n) => + n === 1 + ? { text: "truncated", durationMs: 1, finishReason: "length" as const } + : { text: "{\"y\":2}", durationMs: 1, finishReason: "stop" as const }, + ); + const client = createLlmClientWithProvider(cfg({ maxTokens: 16_384 }), stub); + const r = await client.completeJson<{ y: number }>("ask", { malformedRetries: 1 }); + expect(r.value.y).toBe(2); + expect(stub.inputs[0]!.maxTokens).toBe(16_384); + expect(stub.inputs[1]!.maxTokens).toBe(32_768); + }); + + it("doubles at most once across consecutive truncations", async () => { + const stub = new StubProvider("openai_compatible", () => ({ + text: "truncated", + durationMs: 1, + finishReason: "length" as const, + })); + const client = createLlmClientWithProvider(cfg({ maxTokens: 16_384 }), stub); + await expect( + client.completeJson("ask", { malformedRetries: 2 }), + ).rejects.toBeInstanceOf(MemosError); + expect(stub.inputs).toHaveLength(3); + expect(stub.inputs[0]!.maxTokens).toBe(16_384); + expect(stub.inputs[1]!.maxTokens).toBe(32_768); + expect(stub.inputs[2]!.maxTokens).toBe(32_768); // no further doubling + }); + + it("keeps the budget unchanged when finish_reason is stop", async () => { + const stub = new StubProvider("openai_compatible", (n) => + n === 1 + ? { text: "not json", durationMs: 1, finishReason: "stop" as const } + : { text: "{\"z\":3}", durationMs: 1, finishReason: "stop" as const }, + ); + const client = createLlmClientWithProvider(cfg(), stub); + const r = await client.completeJson<{ z: number }>("ask", { malformedRetries: 1 }); + expect(r.value.z).toBe(3); + expect(stub.inputs[1]!.maxTokens).toBe(stub.inputs[0]!.maxTokens); + }); +});