diff --git a/src/core/record/l1-extractor.test.ts b/src/core/record/l1-extractor.test.ts new file mode 100644 index 0000000..cbeb83e --- /dev/null +++ b/src/core/record/l1-extractor.test.ts @@ -0,0 +1,99 @@ +import fs from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { afterEach, describe, expect, it } from "vitest"; +import { extractL1Memories } from "./l1-extractor.js"; +import type { ConversationMessage } from "../conversation/l0-recorder.js"; +import type { LLMRunner } from "../types.js"; + +const tempDirs: string[] = []; + +describe("extractL1Memories retry", () => { + afterEach(async () => { + await Promise.all(tempDirs.splice(0).map((dir) => fs.rm(dir, { recursive: true, force: true }))); + }); + + it("retries once when the first extraction response is unparsable", async () => { + const dataDir = await fs.mkdtemp(path.join(os.tmpdir(), "tdai-l1-retry-")); + tempDirs.push(dataDir); + const calls: Array<{ taskId?: string; systemPrompt?: string }> = []; + const llmRunner: LLMRunner = { + async run(params) { + calls.push({ taskId: params.taskId, systemPrompt: params.systemPrompt }); + if (calls.length === 1) { + return '[{"scene_name":"work","message_ids":["msg_1"],"memories":[{"content":"broken quote}]}]'; + } + return JSON.stringify([ + { + scene_name: "work", + message_ids: ["msg_1", "msg_2"], + memories: [ + { + content: "User prefers concise engineering status updates.", + type: "instruction", + priority: 80, + source_message_ids: ["msg_1"], + metadata: {}, + }, + ], + }, + ]); + }, + }; + + const result = await extractL1Memories({ + messages: sampleMessages(), + sessionKey: "retry-session", + baseDir: dataDir, + config: {}, + options: { enableDedup: false, llmRunner }, + }); + + expect(calls.map((c) => c.taskId)).toEqual(["l1-extraction", "l1-extraction-retry"]); + expect(calls[1]?.systemPrompt).toContain("Previous output was not valid JSON"); + expect(result.extractedCount).toBe(1); + expect(result.storedCount).toBe(1); + expect(result.records[0]?.content).toBe("User prefers concise engineering status updates."); + }); + + it("does not retry a valid empty extraction array", async () => { + const dataDir = await fs.mkdtemp(path.join(os.tmpdir(), "tdai-l1-retry-")); + tempDirs.push(dataDir); + let calls = 0; + const llmRunner: LLMRunner = { + async run() { + calls++; + return "[]"; + }, + }; + + const result = await extractL1Memories({ + messages: sampleMessages(), + sessionKey: "empty-session", + baseDir: dataDir, + config: {}, + options: { enableDedup: false, llmRunner }, + }); + + expect(calls).toBe(1); + expect(result.extractedCount).toBe(0); + expect(result.storedCount).toBe(0); + }); +}); + +function sampleMessages(): ConversationMessage[] { + return [ + { + id: "msg_1", + role: "user", + content: "Please remember that I prefer concise engineering status updates during implementation work.", + timestamp: Date.now(), + }, + { + id: "msg_2", + role: "assistant", + content: "Understood. I will keep engineering status updates concise during implementation work.", + timestamp: Date.now() + 1, + }, + ]; +} diff --git a/src/core/record/l1-extractor.ts b/src/core/record/l1-extractor.ts index d38fd4b..a1b1da7 100644 --- a/src/core/record/l1-extractor.ts +++ b/src/core/record/l1-extractor.ts @@ -43,6 +43,11 @@ interface SceneSegment { }>; } +interface ParsedExtractionResult { + scenes: SceneSegment[]; + parseFailed: boolean; +} + export interface L1ExtractionResult { /** Whether extraction succeeded */ success: boolean; @@ -316,41 +321,55 @@ async function callLlmExtraction(params: { `${TAG} [l1-debug] ENTRY taskId=l1-extraction, newMsgs=${newMessages.length}, bgMsgs=${backgroundMessages.length}, userPromptLen=${userPrompt.length}, sysPromptLen=${EXTRACT_MEMORIES_SYSTEM_PROMPT.length}, model=${model ?? "(default)"}, previousSceneName=${previousSceneName ? JSON.stringify(previousSceneName) : "(none)"}, runnerKind=${llmRunner ? "llmRunner" : "CleanContextRunner"}`, ); - let result: string; + let runner: CleanContextRunner | undefined; + const runAttempt = async (attempt: "initial" | "retry"): Promise => { + const taskId = attempt === "initial" ? "l1-extraction" : "l1-extraction-retry"; + const systemPrompt = attempt === "initial" + ? EXTRACT_MEMORIES_SYSTEM_PROMPT + : `${EXTRACT_MEMORIES_SYSTEM_PROMPT}\n\nPrevious output was not valid JSON. Return only the required JSON array, with no markdown or explanation.`; + + if (llmRunner) { + // Use the host-neutral LLMRunner interface + return llmRunner.run({ + prompt: userPrompt, + systemPrompt, + taskId, + timeoutMs: 180_000, + }); + } - if (llmRunner) { - // Use the host-neutral LLMRunner interface - result = await llmRunner.run({ - prompt: userPrompt, - systemPrompt: EXTRACT_MEMORIES_SYSTEM_PROMPT, - taskId: "l1-extraction", - timeoutMs: 180_000, - }); - } else { // Fallback: create CleanContextRunner (OpenClaw path) - const runner = new CleanContextRunner({ + runner ??= new CleanContextRunner({ config, modelRef: model, enableTools: false, logger, }); - result = await runner.run({ + return runner.run({ prompt: userPrompt, - systemPrompt: EXTRACT_MEMORIES_SYSTEM_PROMPT, - taskId: "l1-extraction", + systemPrompt, + taskId, timeoutMs: 180_000, }); + }; + + const result = await runAttempt("initial"); + const parsed = parseExtractionResult(result, logger); + if (!parsed.parseFailed) { + return parsed.scenes; } - return parseExtractionResult(result, logger); + logger?.warn?.(`${TAG} Retrying L1 extraction once after unparsable model output`); + const retryResult = await runAttempt("retry"); + return parseExtractionResult(retryResult, logger).scenes; } /** * Parse the LLM's JSON response into SceneSegment array. * Expected format: [{scene_name, message_ids, memories: [...]}] */ -function parseExtractionResult(raw: string, logger?: Logger): SceneSegment[] { +function parseExtractionResult(raw: string, logger?: Logger): ParsedExtractionResult { try { // Strip markdown code block wrappers if present let cleaned = raw.trim(); @@ -367,7 +386,7 @@ function parseExtractionResult(raw: string, logger?: Logger): SceneSegment[] { logger?.warn?.( `${TAG} [l1-debug] NO_JSON taskId=l1-extraction, rawLen=${raw.length}, cleanedLen=${cleaned.length}, rawFull=${JSON.stringify(rawPreview)}${raw.length > 2048 ? `…(+${raw.length - 2048})` : ""}`, ); - return []; + return { scenes: [], parseFailed: true }; } // Sanitize control characters inside JSON string literals that LLM may produce @@ -376,7 +395,7 @@ function parseExtractionResult(raw: string, logger?: Logger): SceneSegment[] { if (!Array.isArray(parsed)) { logger?.warn?.(`${TAG} Extraction response is not an array`); - return []; + return { scenes: [], parseFailed: true }; } const scenes: SceneSegment[] = []; @@ -401,10 +420,10 @@ function parseExtractionResult(raw: string, logger?: Logger): SceneSegment[] { }); } - return scenes; + return { scenes, parseFailed: false }; } catch (err) { logger?.warn?.(`${TAG} Failed to parse extraction result: ${err instanceof Error ? err.message : String(err)}`); - return []; + return { scenes: [], parseFailed: true }; } }