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); });