From 49415712f231a2637d4d0b2e321f3dc653f1c968 Mon Sep 17 00:00:00 2001 From: Sunil Srivatsa Date: Tue, 4 Aug 2026 16:18:45 -0400 Subject: [PATCH] fix(memory): separate extraction instructions from user input The memory-extraction prompt concatenated its instructions, few-shot examples, and the user message into a single user turn, so a small local model could not distinguish instructions from input and frequently echoed the Globex/weather examples instead of extracting facts. Send the instructions as a real system turn and the raw text as the user turn. The tiny worker protocol gains a systemPrompt field, and Mnemopi completion input carries task metadata so the backend selects the right prompt per call. Drop the code-built MEMORY_EXTRACTION_TEMPLATE rather than porting it: prompt text belongs in .md files, and resolveMemoryCompletionInput already overrides that template for every extraction call, so Mnemopi rendered it only for the result to be discarded. Measured on ONNX q4 CPU, LFM2.5-1.2B memory extraction improved from 1/8 to 5/8 once the roles were separated. --- packages/coding-agent/CHANGELOG.md | 4 ++ packages/coding-agent/src/mnemopi/backend.ts | 38 +++++++++++++-- .../system/memory-extraction-system.md | 27 ++--------- .../src/tiny/completion-prompt.ts | 16 +++++++ .../coding-agent/src/tiny/title-client.ts | 21 ++++++--- .../coding-agent/src/tiny/title-protocol.ts | 9 +++- packages/coding-agent/src/tiny/worker.ts | 22 +++------ .../test/mnemopi-completion-input.test.ts | 34 ++++++++++++++ .../test/tiny-title-generator.test.ts | 46 +++++++++++++++++++ packages/mnemopi/CHANGELOG.md | 4 ++ packages/mnemopi/src/core/extraction.ts | 5 +- packages/mnemopi/src/core/local-llm.ts | 1 + packages/mnemopi/src/core/runtime-options.ts | 6 +++ packages/mnemopi/test/extraction.test.ts | 7 +++ 14 files changed, 192 insertions(+), 48 deletions(-) create mode 100644 packages/coding-agent/src/tiny/completion-prompt.ts create mode 100644 packages/coding-agent/test/mnemopi-completion-input.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 2f7b500ff..703208833 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed memory extraction sending its instructions, few-shot examples, and the user's message as a single user turn, which caused small local models to echo the examples instead of extracting facts; instructions now travel as a system turn and the raw text as the user turn + ## [17.3.7] - 2026-08-17 ### Changed diff --git a/packages/coding-agent/src/mnemopi/backend.ts b/packages/coding-agent/src/mnemopi/backend.ts index a84a8b359..31b7f4c79 100644 --- a/packages/coding-agent/src/mnemopi/backend.ts +++ b/packages/coding-agent/src/mnemopi/backend.ts @@ -3,6 +3,7 @@ import * as path from "node:path"; import { type ApiKeyResolver, completeSimple, retryTransientCompletion } from "@oh-my-pi/pi-ai"; import { hostMatchesUrl } from "@oh-my-pi/pi-catalog/hosts"; import type { Mnemopi } from "@oh-my-pi/pi-mnemopi"; +import type { MnemopiLlmCompleteOptions } from "@oh-my-pi/pi-mnemopi/core/runtime-options"; import type * as MnemopiDiagnoseNs from "@oh-my-pi/pi-mnemopi/diagnose"; import type { DiagnosticSummary } from "@oh-my-pi/pi-mnemopi/diagnose"; import { logger } from "@oh-my-pi/pi-utils"; @@ -62,6 +63,27 @@ const STATIC_INSTRUCTIONS = [ "", ].join("\n"); +/** Prompt turns for one Mnemopi completion. */ +export interface MemoryCompletionInput { + prompt: string; + systemPrompt?: string; +} + +/** Maps a Mnemopi completion into instruction and input turns. + * + * Extraction is the only task with its own instructions, and it always supplies + * the raw text, so the instructions become the system turn and the text becomes + * the user turn. Every other task keeps the prompt Mnemopi rendered. */ +export function resolveMemoryCompletionInput( + prompt: string, + options?: MnemopiLlmCompleteOptions, +): MemoryCompletionInput { + if (options?.task?.kind === "memory-extraction") { + return { prompt: options.task.input, systemPrompt: memoryExtractionPrompt }; + } + return { prompt }; +} + async function installMnemopiState(session: AgentSession, config: MnemopiBackendConfig): Promise { const state = new MnemopiSessionState({ sessionId: session.sessionId, config, session }); const previous = setMnemopiSessionState(session, state); @@ -506,8 +528,16 @@ async function resolveMnemopiProviderOptions( return { ...base, llm: { - complete: (prompt, opts) => tinyModelClient.complete(memoryModel, prompt, { maxTokens: opts?.maxTokens }), - extractionPrompt: memoryExtractionPrompt, + complete: (prompt, opts) => { + const request = resolveMemoryCompletionInput(prompt, opts); + return tinyModelClient.complete(memoryModel, request.prompt, { + maxTokens: opts?.maxTokens, + systemPrompt: request.systemPrompt, + }); + }, + // No `extractionPrompt`: resolveMemoryCompletionInput supplies the + // instructions as a system turn for every extraction call, so anything + // rendered here would be built in code and then discarded. consolidationPrompt: memoryConsolidationPrompt, }, }; @@ -537,6 +567,7 @@ async function resolveMnemopiProviderOptions( return { ...base, llm: async (prompt, opts) => { + const request = resolveMemoryCompletionInput(prompt, opts); const hasApiKey = await modelRegistry.getApiKey(model, sessionId); if (!hasApiKey) { logger.warn("Mnemopi: smol completion requested but no current API key is available.", { @@ -549,7 +580,8 @@ async function resolveMnemopiProviderOptions( completeSimple( model, { - messages: [{ role: "user", content: prompt, timestamp: Date.now() }], + ...(request.systemPrompt ? { systemPrompt: [request.systemPrompt] } : {}), + messages: [{ role: "user", content: request.prompt, timestamp: Date.now() }], }, { apiKey: modelRegistry.resolver(model, sessionId), diff --git a/packages/coding-agent/src/prompts/system/memory-extraction-system.md b/packages/coding-agent/src/prompts/system/memory-extraction-system.md index 3e86ab858..cbefacdc7 100644 --- a/packages/coding-agent/src/prompts/system/memory-extraction-system.md +++ b/packages/coding-agent/src/prompts/system/memory-extraction-system.md @@ -1,26 +1,9 @@ -Extract durable, long-term memory items from the user message below. +You are a precise long-term memory extractor. -Output ONE item per line as a short plain-text statement: no JSON, no bullets, no numbering, no field labels. -Capture only persistent, reusable information: -- facts (name, role, employer, config, ports, versions, numbers) -- explicit instructions to the assistant -- stable preferences -- dated events or deadlines +Extract only persistent information explicitly stated in the user message: stable facts, explicit instructions to the assistant, stable preferences, dates, deadlines, paths, ports, and versions. -Keep names, numbers, versions, and dates exact, in the message's original language. When a value is updated, output only the latest value. Ignore greetings, acknowledgements, small talk, weather, and one-off remarks. -If nothing qualifies, output exactly: NO_FACTS +Never infer, explain, invent, or copy information from another message. Ignore greetings, acknowledgements, weather, and one-off plans. When a value is corrected, output only the latest value. -Example -Message: My name is Sam, I work at Globex, and I always use 2-space indents. -Items: -name is Sam -works at Globex -prefers 2-space indents +Preserve names, numbers, paths, versions, dates, and the original language exactly. Output one short plain-text fact per line with no bullets, numbering, labels, JSON, or commentary. -Example -Message: lol nice weather today, might grab a coffee later -Items: -NO_FACTS - -Message: {text} -Items: +If nothing qualifies, output exactly NO_FACTS. diff --git a/packages/coding-agent/src/tiny/completion-prompt.ts b/packages/coding-agent/src/tiny/completion-prompt.ts new file mode 100644 index 000000000..78b180320 --- /dev/null +++ b/packages/coding-agent/src/tiny/completion-prompt.ts @@ -0,0 +1,16 @@ +import type { TextGenerationPipeline } from "@huggingface/transformers"; + +export function buildCompletionPrompt( + tokenizer: TextGenerationPipeline["tokenizer"], + promptText: string, + systemPrompt?: string, +): string { + const userMessage = { role: "user", content: promptText }; + const chat = systemPrompt?.trim() ? [{ role: "system", content: systemPrompt.trim() }, userMessage] : [userMessage]; + const chatTemplateOptions = { + add_generation_prompt: true, + tokenize: false, + enable_thinking: false, + }; + return `${tokenizer.apply_chat_template(chat, chatTemplateOptions)}`; +} diff --git a/packages/coding-agent/src/tiny/title-client.ts b/packages/coding-agent/src/tiny/title-client.ts index 414c31dd4..ef98a1acb 100644 --- a/packages/coding-agent/src/tiny/title-client.ts +++ b/packages/coding-agent/src/tiny/title-client.ts @@ -52,6 +52,12 @@ export interface TinyTitleGenerateOptions { systemPrompt?: string; } +export interface TinyModelCompletionOptions { + maxTokens?: number; + signal?: AbortSignal; + systemPrompt?: string; +} + function normalizeTinyTitleGenerateOptions( options: AbortSignal | TinyTitleGenerateOptions | undefined, ): TinyTitleGenerateOptions { @@ -260,11 +266,7 @@ export class TinyTitleClient { } } - async complete( - modelKey: string, - prompt: string, - options: { maxTokens?: number; signal?: AbortSignal } = {}, - ): Promise { + async complete(modelKey: string, prompt: string, options: TinyModelCompletionOptions = {}): Promise { if (!isTinyMemoryLocalModelKey(modelKey)) return null; if (options.signal?.aborted || this.#failedModels.has(modelKey)) return null; @@ -281,7 +283,14 @@ export class TinyTitleClient { }; options.signal?.addEventListener("abort", abort, { once: true }); try { - worker.send({ type: "complete", id, modelKey, prompt, maxTokens: options.maxTokens }); + worker.send({ + type: "complete", + id, + modelKey, + prompt, + maxTokens: options.maxTokens, + systemPrompt: options.systemPrompt, + }); return await promise; } finally { options.signal?.removeEventListener("abort", abort); diff --git a/packages/coding-agent/src/tiny/title-protocol.ts b/packages/coding-agent/src/tiny/title-protocol.ts index 5bc86f525..080aaef30 100644 --- a/packages/coding-agent/src/tiny/title-protocol.ts +++ b/packages/coding-agent/src/tiny/title-protocol.ts @@ -30,7 +30,14 @@ export interface TinyTitleProgressEvent { export type TinyTitleWorkerInbound = | { type: "ping"; id: string } | { type: "generate"; id: string; modelKey: TinyTitleLocalModelKey; message: string; systemPrompt?: string } - | { type: "complete"; id: string; modelKey: TinyLocalModelKey; prompt: string; maxTokens?: number } + | { + type: "complete"; + id: string; + modelKey: TinyLocalModelKey; + prompt: string; + maxTokens?: number; + systemPrompt?: string; + } | { type: "download"; id: string; modelKey: TinyLocalModelKey }; export type TinyTitleWorkerOutbound = diff --git a/packages/coding-agent/src/tiny/worker.ts b/packages/coding-agent/src/tiny/worker.ts index ca056baa6..56a0c0c4b 100644 --- a/packages/coding-agent/src/tiny/worker.ts +++ b/packages/coding-agent/src/tiny/worker.ts @@ -19,6 +19,7 @@ import { sendProgress, type TransformersRuntimeMetadata, } from "../subprocess/worker-runtime"; +import { buildCompletionPrompt } from "./completion-prompt"; import { resolveTinyModelDevicePreference, type TinyModelDevice, tinyModelDeviceLoadOrder } from "./device"; import { resolveTinyModelDtypeOverride, type TinyModelDtype } from "./dtype"; import { formatTitleUserMessage } from "./message-preproc"; @@ -265,21 +266,10 @@ async function generateTitle( return extractTinyTitle(output[0]?.generated_text ?? "", message); } -function buildCompletionPrompt(generator: TextGenerationPipeline, promptText: string): string { - const chat = [{ role: "user", content: promptText }]; - const chatTemplateOptions = { - add_generation_prompt: true, - tokenize: false, - enable_thinking: false, - }; - return `${generator.tokenizer.apply_chat_template(chat, chatTemplateOptions)}`; -} - /** - * Generic single-turn completion used by Mnemopi memory tasks (fact extraction - * and consolidation). The caller (Mnemopi) supplies the full task prompt; we - * wrap it as the user turn, decode greedily, and return the raw text for the - * caller's own parser. Output is capped to keep local inference latency bounded. + * Completion path for Mnemopi memory tasks. Extraction can carry a dedicated + * system prompt and user payload; consolidation retains the generic user-only + * prompt. Output is capped to keep local inference latency bounded. */ async function generateCompletion( transport: TinyTitleTransport, @@ -287,9 +277,10 @@ async function generateCompletion( modelKey: TinyLocalModelKey, promptText: string, maxTokens: number | undefined, + systemPrompt: string | undefined, ): Promise { const generator = await loadPipeline(modelKey, transport, requestId); - const text = buildCompletionPrompt(generator, promptText); + const text = buildCompletionPrompt(generator.tokenizer, promptText, systemPrompt); const requested = maxTokens ?? MEMORY_COMPLETION_DEFAULT_MAX_NEW_TOKENS; const maxNewTokens = Math.min(Math.max(1, requested), COMPLETION_MAX_NEW_TOKENS); const output = (await generator(text, { @@ -332,6 +323,7 @@ async function handleQueuedRequest( request.modelKey, request.prompt, request.maxTokens, + request.systemPrompt, ); transport.send({ type: "completion", id: request.id, text }); return; diff --git a/packages/coding-agent/test/mnemopi-completion-input.test.ts b/packages/coding-agent/test/mnemopi-completion-input.test.ts new file mode 100644 index 000000000..72c123bd8 --- /dev/null +++ b/packages/coding-agent/test/mnemopi-completion-input.test.ts @@ -0,0 +1,34 @@ +import { describe, expect, it } from "bun:test"; +import { resolveMemoryCompletionInput } from "../src/mnemopi/backend"; +import memoryExtractionPrompt from "../src/prompts/system/memory-extraction-system.md" with { type: "text" }; + +describe("resolveMemoryCompletionInput", () => { + it("splits an extraction call into instruction and input turns", () => { + const rendered = "whatever Mnemopi rendered for the prompt slot"; + const request = resolveMemoryCompletionInput(rendered, { + task: { kind: "memory-extraction", input: "Sam works at Globex." }, + }); + expect(request.systemPrompt).toBe(memoryExtractionPrompt); + expect(request.prompt).toBe("Sam works at Globex."); + // The rendered prompt is deliberately discarded: instructions belong in the + // system turn and the user turn must carry only the text to extract from. + expect(request.prompt).not.toContain("rendered"); + }); + + it("keeps the rendered prompt and adds no system turn without an extraction task", () => { + // Consolidation reaches the same completion fn with no task, so it must keep + // the prompt Mnemopi rendered from consolidationPrompt. + const rendered = "Summarize these memories faithfully."; + expect(resolveMemoryCompletionInput(rendered)).toEqual({ prompt: rendered }); + expect(resolveMemoryCompletionInput(rendered, {})).toEqual({ prompt: rendered }); + expect(resolveMemoryCompletionInput(rendered, { maxTokens: 256 })).toEqual({ prompt: rendered }); + }); + + it("does not leak the extraction instructions into the user turn", () => { + const request = resolveMemoryCompletionInput("ignored", { + task: { kind: "memory-extraction", input: "Sam prefers dark mode." }, + }); + expect(request.prompt).toBe("Sam prefers dark mode."); + expect(memoryExtractionPrompt).not.toContain("Sam prefers dark mode."); + }); +}); diff --git a/packages/coding-agent/test/tiny-title-generator.test.ts b/packages/coding-agent/test/tiny-title-generator.test.ts index 2583e8cbb..3e114e582 100644 --- a/packages/coding-agent/test/tiny-title-generator.test.ts +++ b/packages/coding-agent/test/tiny-title-generator.test.ts @@ -30,6 +30,7 @@ import { import type { TinyTitleWorkerInbound, TinyTitleWorkerOutbound } from "@oh-my-pi/pi-coding-agent/tiny/title-protocol"; import { generateSessionTitle } from "@oh-my-pi/pi-coding-agent/utils/title-generator"; import type { Subprocess } from "bun"; +import { buildCompletionPrompt } from "../src/tiny/completion-prompt"; function getModelOrThrow(id: string): Model { const model = getBundledModel("anthropic", id); @@ -246,6 +247,51 @@ function createFakeTinyWorker(): FakeTinyWorker { return worker; } +describe("tiny memory completion prompts", () => { + it("renders extraction instructions as a system turn separate from user input", () => { + const applyChatTemplate = vi.fn(() => "rendered prompt"); + const tokenizer = { apply_chat_template: applyChatTemplate }; + + expect(buildCompletionPrompt(tokenizer as never, "actual user input", " extraction instructions ")).toBe( + "rendered prompt", + ); + expect(applyChatTemplate).toHaveBeenCalledWith( + [ + { role: "system", content: "extraction instructions" }, + { role: "user", content: "actual user input" }, + ], + { + add_generation_prompt: true, + tokenize: false, + enable_thinking: false, + }, + ); + }); + + it("carries the extraction system prompt over the worker protocol", async () => { + const worker = createFakeTinyWorker(); + const client = new TinyTitleClient(() => worker.handle); + + const completion = client.complete("lfm2-1.2b", "actual user input", { + maxTokens: 64, + systemPrompt: "extraction instructions", + }); + const request = worker.sent.find(message => message.type === "complete"); + expect(request).toEqual({ + type: "complete", + id: expect.any(String), + modelKey: "lfm2-1.2b", + prompt: "actual user input", + maxTokens: 64, + systemPrompt: "extraction instructions", + }); + worker.emit({ type: "completion", id: request?.id ?? "", text: "extracted fact" }); + + expect(await completion).toBe("extracted fact"); + await client.terminate(); + }); +}); + describe("tiny title prewarm", () => { it("spawns one idle worker that the first generate reuses (issue #6462)", async () => { const workers: FakeTinyWorker[] = []; diff --git a/packages/mnemopi/CHANGELOG.md b/packages/mnemopi/CHANGELOG.md index 3e80f0647..d0292e5b6 100644 --- a/packages/mnemopi/CHANGELOG.md +++ b/packages/mnemopi/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added optional task metadata to the runtime LLM completion interface so hosts can tell an extraction call from a consolidation call and choose the matching prompt + ## [17.3.5] - 2026-08-16 ### Fixed diff --git a/packages/mnemopi/src/core/extraction.ts b/packages/mnemopi/src/core/extraction.ts index 99cd3f52e..77c0ad845 100644 --- a/packages/mnemopi/src/core/extraction.ts +++ b/packages/mnemopi/src/core/extraction.ts @@ -387,7 +387,10 @@ export async function extractFactCategories( if (configuredLlmWillHandleCall()) { diag.recordAttempt("host"); try { - const raw = await callConfiguredCompletion(prompt, 0, { maxTokens: llmMaxTokens() }); + const raw = await callConfiguredCompletion(prompt, 0, { + maxTokens: llmMaxTokens(), + task: { kind: "memory-extraction", input: text }, + }); if (typeof raw === "string" && raw.trim() !== "") { const extracted = parseExtractedFactCategories(raw); const count = countExtractedFactCategories(extracted); diff --git a/packages/mnemopi/src/core/local-llm.ts b/packages/mnemopi/src/core/local-llm.ts index 36615359f..5922d8834 100644 --- a/packages/mnemopi/src/core/local-llm.ts +++ b/packages/mnemopi/src/core/local-llm.ts @@ -180,6 +180,7 @@ export async function callConfiguredCompletion( timeout: opts.timeout, provider: opts.provider, model: opts.model, + task: opts.task, }); return typeof raw === "string" ? raw : null; } diff --git a/packages/mnemopi/src/core/runtime-options.ts b/packages/mnemopi/src/core/runtime-options.ts index ca6ee1594..198a0f8e3 100644 --- a/packages/mnemopi/src/core/runtime-options.ts +++ b/packages/mnemopi/src/core/runtime-options.ts @@ -1,12 +1,18 @@ import { AsyncLocalStorage } from "node:async_hooks"; import type { Api, ApiKey, Model } from "@oh-my-pi/pi-ai"; +export type MnemopiLlmCompletionTask = { + kind: "memory-extraction"; + input: string; +}; + export interface MnemopiLlmCompleteOptions { maxTokens?: number; temperature?: number; timeout?: number; provider?: string | null; model?: string | null; + task?: MnemopiLlmCompletionTask; } export type MnemopiLlmCompletion = ( diff --git a/packages/mnemopi/test/extraction.test.ts b/packages/mnemopi/test/extraction.test.ts index acf501acc..8c89167c9 100644 --- a/packages/mnemopi/test/extraction.test.ts +++ b/packages/mnemopi/test/extraction.test.ts @@ -14,6 +14,7 @@ import { setHostLlmBackend, } from "@oh-my-pi/pi-mnemopi/core/llm-backends"; import { + type MnemopiLlmCompletionTask, type ResolvedMnemopiRuntimeOptions, withMnemopiRuntimeOptions, } from "@oh-my-pi/pi-mnemopi/core/runtime-options"; @@ -136,6 +137,7 @@ describe("structured extraction", () => { process.env.MNEMOPI_LLM_ENABLED = "true"; let capturedPrompt = ""; let capturedTemperature = -1; + let capturedTask: MnemopiLlmCompletionTask | undefined; const resolved: ResolvedMnemopiRuntimeOptions = { llm: { enabled: true, @@ -143,6 +145,7 @@ describe("structured extraction", () => { complete: (prompt, opts) => { capturedPrompt = prompt; capturedTemperature = opts?.temperature ?? -1; + capturedTask = opts?.task; return "Sam works at Globex\nSam prefers dark mode"; }, }, @@ -155,6 +158,10 @@ describe("structured extraction", () => { expect(facts).toEqual(["Sam works at Globex", "Sam prefers dark mode"]); expect(capturedPrompt).toContain("ONLY-LINES for: Sam works at Globex and prefers dark mode."); expect(capturedTemperature).toBe(0); + expect(capturedTask).toEqual({ + kind: "memory-extraction", + input: "Sam works at Globex and prefers dark mode.", + }); expect(getExtractionStats().by_tier.host.successes).toBe(1); });