diff --git a/apps/memos-local-plugin/core/capture/ALGORITHMS.md b/apps/memos-local-plugin/core/capture/ALGORITHMS.md index 49c98a105..15a15ac64 100644 --- a/apps/memos-local-plugin/core/capture/ALGORITHMS.md +++ b/apps/memos-local-plugin/core/capture/ALGORITHMS.md @@ -78,8 +78,7 @@ priority once reward arrives. ## V7 §3.2 batched variant — `batch-scorer.ts` The per-step path (`reflection-synth.ts` + `alpha-scorer.ts`) issues 2N -LLM calls per N-step episode. `batch-scorer.ts` collapses up to -`batchThreshold` steps into one call: +LLM calls per N-step episode. `batch-scorer.ts` collapses them into ONE: ``` inputs = [{idx, state, action, outcome, reflection, synth_allowed}, …] @@ -92,8 +91,8 @@ Dispatch (in `capture.ts`): | `cfg.batchMode` | `cfg.batchThreshold` | behavior | |-------------------|----------------------|----------| | `per_step` | (ignored) | legacy: 2N calls | -| `per_episode` | chunk size | batch when `N ≤ threshold`; else chunk-batch | -| `auto` (default) | `12` | batch when `N ≤ 12`; else chunk-batch | +| `per_episode` | (ignored) | always batch | +| `auto` (default) | `12` | batch when `N ≤ 12`; else per-step | The dispatcher also refuses to batch when no LLM is wired — same fallback path as missing-LLM in per-step mode. @@ -108,15 +107,15 @@ Failure handling: - LLM throws / facade gives up after `malformedRetries=1` → capture catches in `runBatchScoring`, surfaces a `{stage: "batch"}` warning, - and the per-step path runs as a fallback for that chunk. + and the per-step path runs as a fallback. - Validator rejects on length mismatch, missing/non-numeric `alpha`, non-boolean `usable`, non-string `reflection_text`. Same fallback. Bookkeeping (`CaptureResult.llmCalls`): -- `batchedReflection`: number of successful batch/chunk calls. +- `batchedReflection`: 0 or 1 per episode (1 on a successful batch). - `reflectionSynth` / `alphaScoring`: only nonzero when the per-step path - ran (either selected directly, or as fallback after a chunk failure). + ran (either selected directly, or as fallback after a batch failure). Stable prompt fingerprint: diff --git a/apps/memos-local-plugin/core/capture/batch-scorer.ts b/apps/memos-local-plugin/core/capture/batch-scorer.ts index da434c3b2..e7b8ab50f 100644 --- a/apps/memos-local-plugin/core/capture/batch-scorer.ts +++ b/apps/memos-local-plugin/core/capture/batch-scorer.ts @@ -16,12 +16,11 @@ * `transferability` axes benefit directly. * * Trade-offs (encoded in capture.ts dispatch): - * - Prompt grows linearly with N steps. Each call is capped at - * `batchThreshold`; long episodes run as several bounded chunks. - * - One bad chunk forces a single batched retry for that chunk instead - * of N isolated retries — but the facade already does - * `malformedRetries` for us, and on hard failure capture.ts falls - * back to per-step for that chunk only. + * - Prompt grows linearly with N steps. Capped via `batchThreshold`; + * long episodes degrade to the per-step path automatically. + * - One bad output value forces a single batched retry instead of N + * isolated retries — but the facade already does `malformedRetries` + * for us, and on hard failure capture.ts falls back to per-step. * * Wire format ↔ prompt: * Send `{ host_context?, task_context?, steps: [{idx, state, action, outcome, reflection, synth_allowed}] }`. @@ -171,7 +170,6 @@ export async function batchScoreReflections( validate: (v) => validateBatchPayload(v, inputs.length), malformedRetries: 1, temperature: 0, - maxTokens: batchMaxTokens(inputs.length), }, ); @@ -323,15 +321,6 @@ function validateBatchPayload(v: unknown, expected: number): void { } } -function batchMaxTokens(stepCount: number): number { - // Batch output scales with step count; keep a per-step budget but cap below - // the 16k range that triggered avoidable reasoning spend on mimo replay. - const perStepOutputBudget = 512; - const baseBudget = 768; - const ceiling = 8_192; - return Math.min(ceiling, baseBudget + Math.max(1, stepCount) * perStepOutputBudget); -} - function lastToolOutcome(step: NormalizedStep, max: number): string { const last = step.toolCalls[step.toolCalls.length - 1]; if (!last) return "(assistant-only step)"; diff --git a/apps/memos-local-plugin/core/capture/capture.ts b/apps/memos-local-plugin/core/capture/capture.ts index bc9c7a812..d5b62aa38 100644 --- a/apps/memos-local-plugin/core/capture/capture.ts +++ b/apps/memos-local-plugin/core/capture/capture.ts @@ -463,14 +463,14 @@ export function createCaptureRunner(deps: CaptureDeps): CaptureRunner { } // Batch reflection + α across every step of the now-closed - // episode. Long episodes are chunk-batched at `batchThreshold`; - // failed chunks fall back to per-step scoring. The reflect pass uses + // episode. Falls back to per-step scoring when over the threshold + // or when batching fails / no LLM is wired. The reflect pass uses // `reflectLlm` (skill-evolver model when configured) for higher // quality reflections; per-turn lite capture still uses `llm`. const reflectStart = now(); const rLlm = deps.reflectLlm ?? deps.llm; - const scoringPlan = planScoring(deps.cfg, normalized.length, rLlm !== null); - const contextEnabled = contextModeFor(deps.cfg, scoringPlan, normalized.length); + const useBatch = shouldBatch(deps.cfg, normalized.length, rLlm !== null); + const contextEnabled = contextModeFor(deps.cfg, useBatch, normalized.length); const taskSummary = contextEnabled.includeTask ? buildTaskReflectionSummary(input.episode, normalized, deps.cfg.taskContextMaxChars) : null; @@ -481,10 +481,7 @@ export function createCaptureRunner(deps: CaptureDeps): CaptureRunner { episodeId: input.episode.id, sessionId: input.episode.sessionId, steps: normalized.length, - mode: scoringPlan === "per_step" && contextEnabled.includeDownstream ? "per_step_downstream" : scoringPlan, - chunks: scoringPlan === "chunk_batch" - ? Math.ceil(normalized.length / Math.max(1, deps.cfg.batchThreshold)) - : undefined, + mode: useBatch ? "batch" : contextEnabled.includeDownstream ? "per_step_downstream" : "per_step", reflectionContextMode: deps.cfg.reflectionContextMode, downstreamPreview: contextEnabled.includeDownstream, provider: rLlm?.provider ?? "none", @@ -492,13 +489,10 @@ export function createCaptureRunner(deps: CaptureDeps): CaptureRunner { taskSummary: taskSummary ? taskSummary.slice(0, 240) : null, }); let scored: ScoredStep[] = []; - if (scoringPlan === "batch") { + if (useBatch) { scored = await runBatchScoring(normalized, rLlm!, deps, warnings, llmCalls, input.episode.id, taskSummary); } - if (scoringPlan === "chunk_batch") { - scored = await runChunkedBatchScoring(normalized, rLlm!, deps, warnings, llmCalls, input.episode.id, taskSummary); - } - if (scoringPlan === "per_step" || scored.length === 0) { + if (!useBatch || scored.length === 0) { scored = await runPerStepScoring( normalized, rLlm, @@ -1068,30 +1062,30 @@ export function createCaptureRunner(deps: CaptureDeps): CaptureRunner { // ─── helpers ──────────────────────────────────────────────────────────────── /** - * Decide which reflection+α path to use. + * Decide whether to use the batched reflection+α path. * * `per_step` → never (legacy path). - * `per_episode` → batch up to threshold, then chunk-batch. - * `auto` → batch up to threshold, then chunk-batch. + * `per_episode` → always, when an LLM is available. + * `auto` → batch when step count fits inside `batchThreshold`. */ -type ScoringPlan = "per_step" | "batch" | "chunk_batch"; - -function planScoring(cfg: CaptureConfig, stepCount: number, hasLlm: boolean): ScoringPlan { - if (!hasLlm) return "per_step"; - if (stepCount === 0) return "per_step"; - if (cfg.batchMode === "per_step") return "per_step"; - return stepCount <= Math.max(1, cfg.batchThreshold) ? "batch" : "chunk_batch"; +function shouldBatch(cfg: CaptureConfig, stepCount: number, hasLlm: boolean): boolean { + if (!hasLlm) return false; + if (stepCount === 0) return false; + if (cfg.batchMode === "per_step") return false; + if (cfg.batchMode === "per_episode") return true; + // "auto" + return stepCount <= cfg.batchThreshold; } function contextModeFor( cfg: CaptureConfig, - scoringPlan: ScoringPlan, + useBatch: boolean, stepCount: number, ): { includeTask: boolean; includeDownstream: boolean } { const mode = cfg.reflectionContextMode; const includeTask = mode === "task" || mode === "task_downstream"; const wantsDownstream = mode === "downstream" || mode === "task_downstream"; - const longPerStep = scoringPlan === "per_step" && stepCount > cfg.batchThreshold; + const longPerStep = !useBatch && stepCount > cfg.batchThreshold; const includeDownstream = wantsDownstream && cfg.longEpisodeReflectMode === "per_step_downstream" && @@ -1151,37 +1145,6 @@ async function runBatchScoring( } } -async function runChunkedBatchScoring( - normalized: NormalizedStep[], - llm: LlmClient, - deps: CaptureDeps, - warnings: CaptureResult["warnings"], - llmCalls: { reflectionSynth: number; alphaScoring: number; batchedReflection: number }, - episodeId: string, - taskSummary: string | null, -): Promise { - const chunkSize = Math.max(1, deps.cfg.batchThreshold); - const chunks: NormalizedStep[][] = []; - for (let start = 0; start < normalized.length; start += chunkSize) { - chunks.push(normalized.slice(start, start + chunkSize)); - } - const concurrency = Math.max(1, deps.cfg.llmConcurrency); - const scoredChunks = await runConcurrently(chunks, concurrency, async (chunk): Promise => { - const scored = await runBatchScoring(chunk, llm, deps, warnings, llmCalls, episodeId, taskSummary); - if (scored.length > 0) return scored; - return runPerStepScoring( - chunk, - llm, - deps, - warnings, - llmCalls, - episodeId, - buildReflectionContexts(chunk, taskSummary, chunk.map(() => [])), - ); - }); - return scoredChunks.flat(); -} - async function runPerStepScoring( normalized: NormalizedStep[], llm: LlmClient | null, diff --git a/apps/memos-local-plugin/tests/helpers/fake-llm.ts b/apps/memos-local-plugin/tests/helpers/fake-llm.ts index ec2c00261..22d9fc1a3 100644 --- a/apps/memos-local-plugin/tests/helpers/fake-llm.ts +++ b/apps/memos-local-plugin/tests/helpers/fake-llm.ts @@ -20,7 +20,7 @@ export interface FakeLlmScript { complete?: Record string | Promise)>; completeJson?: Record< string, - unknown | ((input: unknown, opts?: unknown) => unknown | Promise) + unknown | ((input: unknown) => unknown | Promise) >; /** Override the served-by identifier. */ servedBy?: LlmProviderName | "host_fallback"; @@ -64,7 +64,7 @@ export function fakeLlm(script: FakeLlmScript = {}): LlmClient { throw new Error(`fakeLlm: no completeJson mock for op="${op}"`); } const value = (typeof entry === "function" - ? await (entry as (x: unknown, o?: unknown) => unknown)(input, opts) + ? await (entry as (x: unknown) => unknown)(input) : entry) as T; if (o?.validate) o.validate(value); return { diff --git a/apps/memos-local-plugin/tests/unit/capture/capture-batch.test.ts b/apps/memos-local-plugin/tests/unit/capture/capture-batch.test.ts index bc4b76f28..d86290517 100644 --- a/apps/memos-local-plugin/tests/unit/capture/capture-batch.test.ts +++ b/apps/memos-local-plugin/tests/unit/capture/capture-batch.test.ts @@ -7,15 +7,14 @@ * 2. existing reflections are preserved verbatim; * 3. synth-disabled steps stay at α=0 even when the LLM tries to write * one for them; - * 4. `auto` mode chunk-batches when stepCount > batchThreshold; - * 5. a malformed chunk degrades only that chunk into the per-step path - * instead of dropping the whole episode to per-step. + * 4. `auto` mode falls back to per-step when stepCount > batchThreshold; + * 5. a malformed batched response degrades into the per-step path + * instead of crashing capture. */ import { afterEach, beforeAll, beforeEach, describe, expect, it } from "vitest"; import { createCaptureRunner, type CaptureRunner } from "../../../core/capture/capture.js"; -import { batchScoreReflections } from "../../../core/capture/batch-scorer.js"; import { createCaptureEventBus } from "../../../core/capture/events.js"; import { BATCH_REFLECTION_PROMPT, @@ -313,31 +312,20 @@ describe("capture/pipeline (batched ρ+α path)", () => { expect(t.alpha).toBe(0); // V7 disabledScore semantics }); - it("auto mode chunk-batches when stepCount > batchThreshold", async () => { - const batchStates: string[][] = []; + it("auto mode falls back to per-step when stepCount > batchThreshold", async () => { const llm = fakeLlm({ completeJson: { - [batchOp]: (input) => { - const messages = input as Array<{ role: string; content: string }>; - const payload = JSON.parse(messages[messages.length - 1]!.content) as { - steps: Array<{ idx: number; state: string }>; - }; - batchStates.push(payload.steps.map((s) => s.state)); - return { - scores: payload.steps.map((step) => ({ - idx: step.idx, - reflection_text: `reflection ${step.state}`, - alpha: step.idx === 0 ? 0.2 : 0.4, - usable: true, - reason: "ok", - })), - }; - }, + // ONLY per-step alpha mock; if batched gets called, the test fails + // with "no completeJson mock for op=...batch...". + [alphaOp]: { alpha: 0.5, usable: true, reason: "ok" }, + }, + complete: { + "capture.reflection.synth": "I made this decision deliberately.", }, }); const runner = buildRunner({ batchMode: "auto", batchThreshold: 2 }, llm); - // 3 steps → above threshold → two bounded batch chunks. + // 3 steps → above threshold → per-step path. const ep = episodeSnapshot({ id: "ep_1", sessionId: "se_1", @@ -353,17 +341,10 @@ describe("capture/pipeline (batched ρ+α path)", () => { const result = await runCapture(runner, ep); expect(result.traceIds).toHaveLength(3); - expect(batchStates).toEqual([["a", "b"], ["c"]]); - expect(result.llmCalls.batchedReflection).toBe(2); - expect(result.llmCalls.reflectionSynth).toBe(0); - expect(result.llmCalls.alphaScoring).toBe(0); - - const rows = result.traceIds.map((id) => tmp.repos.traces.getById(id)!); - expect(rows.map((row) => row.reflection)).toEqual([ - "reflection a", - "reflection b", - "reflection c", - ]); + expect(result.llmCalls.batchedReflection).toBe(0); + // 3 synth + 3 alpha calls in per-step mode. + expect(result.llmCalls.reflectionSynth).toBe(3); + expect(result.llmCalls.alphaScoring).toBe(3); }); it("long per-step downstream mode injects up to three following steps", async () => { @@ -387,7 +368,7 @@ describe("capture/pipeline (batched ρ+α path)", () => { }); const runner = buildRunner( { - batchMode: "per_step", + batchMode: "auto", batchThreshold: 2, reflectionContextMode: "task_downstream", longEpisodeReflectMode: "per_step_downstream", @@ -434,27 +415,16 @@ describe("capture/pipeline (batched ρ+α path)", () => { expect(step3Prompt).not.toContain("[step+3]"); }); - it("per_episode mode chunk-batches when step count is large", async () => { - const chunkSizes: number[] = []; + it("per_episode mode batches even when step count is large", async () => { + const scores = Array.from({ length: 5 }, (_, i) => ({ + idx: i, + reflection_text: `reflection #${i}`, + alpha: 0.4, + usable: true, + reason: "ok", + })); const llm = fakeLlm({ - completeJson: { - [batchOp]: (input) => { - const messages = input as Array<{ role: string; content: string }>; - const payload = JSON.parse(messages[messages.length - 1]!.content) as { - steps: Array<{ idx: number; state: string }>; - }; - chunkSizes.push(payload.steps.length); - return { - scores: payload.steps.map((step) => ({ - idx: step.idx, - reflection_text: `reflection ${step.state}`, - alpha: 0.4, - usable: true, - reason: "ok", - })), - }; - }, - }, + completeJson: { [batchOp]: { scores } }, }); const runner = buildRunner({ batchMode: "per_episode", batchThreshold: 2 }, llm); @@ -466,111 +436,10 @@ describe("capture/pipeline (batched ρ+α path)", () => { const ep = episodeSnapshot({ id: "ep_1", sessionId: "se_1", turns }); const result = await runCapture(runner, ep); expect(result.traceIds).toHaveLength(5); - expect(chunkSizes).toEqual([2, 2, 1]); - expect(result.llmCalls.batchedReflection).toBe(3); + expect(result.llmCalls.batchedReflection).toBe(1); expect(result.llmCalls.alphaScoring).toBe(0); }); - it("chunk-batch falls back to per-step only for the failed chunk", async () => { - const llm = fakeLlm({ - completeJson: { - [batchOp]: (input) => { - const messages = input as Array<{ role: string; content: string }>; - const payload = JSON.parse(messages[messages.length - 1]!.content) as { - steps: Array<{ idx: number; state: string }>; - }; - if (payload.steps[0]?.state === "q2") { - throw new Error("chunk failed"); - } - return { - scores: payload.steps.map((step) => ({ - idx: step.idx, - reflection_text: `batch ${step.state}`, - alpha: step.state === "q4" ? 0.5 : 0.2, - usable: true, - reason: "ok", - })), - }; - }, - [alphaOp]: { alpha: 0.9, usable: true, reason: "fallback" }, - }, - complete: { - "capture.reflection.synth": "per-step fallback reflection", - }, - }); - const runner = buildRunner({ batchMode: "auto", batchThreshold: 2 }, llm); - - const turns: EpisodeTurn[] = []; - for (let i = 0; i < 5; i++) { - turns.push(turn("user", `q${i}`, 1_000 + i * 100)); - turns.push(turn("assistant", `a${i}`, 1_050 + i * 100)); - } - const ep = episodeSnapshot({ id: "ep_1", sessionId: "se_1", turns }); - - const result = await runCapture(runner, ep); - expect(result.traceIds).toHaveLength(5); - expect(result.llmCalls.batchedReflection).toBe(2); - expect(result.llmCalls.reflectionSynth).toBe(2); - expect(result.llmCalls.alphaScoring).toBe(2); - expect(result.warnings.filter((w) => w.stage === "batch")).toHaveLength(1); - - const rows = result.traceIds.map((id) => tmp.repos.traces.getById(id)!); - expect(rows.map((row) => row.reflection)).toEqual([ - "batch q0", - "batch q1", - "per-step fallback reflection", - "per-step fallback reflection", - "batch q4", - ]); - expect(rows.map((row) => row.alpha)).toEqual([0.2, 0.2, 0.9, 0.9, 0.5]); - }); - - it("batch scorer passes an explicit maxTokens budget", async () => { - let seenMaxTokens: number | undefined; - const llm = fakeLlm({ - completeJson: { - [batchOp]: (_input, opts) => { - seenMaxTokens = (opts as { maxTokens?: number }).maxTokens; - return { - scores: [ - { - idx: 0, - reflection_text: "I made a useful choice.", - alpha: 0.5, - usable: true, - reason: "ok", - }, - ], - }; - }, - }, - }); - - await batchScoreReflections( - llm, - [ - { - step: { - key: "s1", - ts: 1_000 as EpochMs, - type: "text", - userText: "q", - agentText: "a", - agentThinking: null, - toolCalls: [], - rawReflection: null, - meta: {}, - }, - existingReflection: null, - }, - ], - { synthReflections: true }, - ); - - expect(seenMaxTokens).toBeGreaterThan(0); - expect(seenMaxTokens).toBeLessThan(16_384); - }); - it("malformed batched response → falls back to per-step + emits warning", async () => { const llm = fakeLlm({ completeJson: {