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.
This commit is contained in:
Sunil Srivatsa
2026-08-04 16:18:45 -04:00
parent 644ad30d6e
commit 49415712f2
14 changed files with 192 additions and 48 deletions
+4
View File
@@ -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
+35 -3
View File
@@ -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<MnemopiSessionState> {
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),
@@ -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.
@@ -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)}`;
}
+15 -6
View File
@@ -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<string | null> {
async complete(modelKey: string, prompt: string, options: TinyModelCompletionOptions = {}): Promise<string | null> {
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);
@@ -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 =
+7 -15
View File
@@ -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<string | null> {
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;
@@ -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.");
});
});
@@ -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<Api> {
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[] = [];
+4
View File
@@ -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
+4 -1
View File
@@ -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);
+1
View File
@@ -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;
}
@@ -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 = (
+7
View File
@@ -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);
});