diff --git a/src/core/config.ts b/src/core/config.ts index e81954886..644df2c2a 100644 --- a/src/core/config.ts +++ b/src/core/config.ts @@ -928,6 +928,8 @@ export const KNOWN_CONFIG_KEYS: readonly string[] = [ 'models.tier.subagent', 'models.aliases', 'models.dream.synthesize', + 'models.dream.extract_atoms', + 'cycle.extract_atoms.budget_usd', 'models.dream.patterns', 'models.dream.synthesize_verdict', 'models.drift', diff --git a/src/core/cycle/extract-atoms.ts b/src/core/cycle/extract-atoms.ts index 4c73ea400..08565b343 100644 --- a/src/core/cycle/extract-atoms.ts +++ b/src/core/cycle/extract-atoms.ts @@ -51,13 +51,15 @@ import type { BrainEngine } from '../engine.ts'; import type { PhaseResult } from '../cycle.ts'; import type { GBrainConfig } from '../config.ts'; import type { ProgressReporter } from '../progress.ts'; -import { chat as gatewayChat } from '../ai/gateway.ts'; +import { chat as gatewayChat, withBudgetTracker } from '../ai/gateway.ts'; +import { BudgetExhausted, BudgetTracker } from '../budget/budget-tracker.ts'; import { writeReceipt } from '../extract/receipt-writer.ts'; import { upsertExtractRollup } from '../extract/rollup-writer.ts'; import { createHash } from 'crypto'; import { slugifySegment } from '../sync.ts'; const DEFAULT_BUDGET_USD = 0.3; +const DEFAULT_EXTRACT_ATOMS_MODEL = 'anthropic:claude-haiku-4-5'; // v0.42+ TODO: read atom_type enum from active pack manifest at runtime. const ATOM_TYPES = [ @@ -500,7 +502,24 @@ export async function runPhaseExtractAtoms( let pagesSkipped = 0; const failures: Array<{ source: string; error: string }> = []; let estimatedSpendUsd = 0; - const budgetCap = DEFAULT_BUDGET_USD; + let budgetExhausted = false; + let extractModel = DEFAULT_EXTRACT_ATOMS_MODEL; + let budgetCap = DEFAULT_BUDGET_USD; + try { + const configuredModel = await engine.getConfig('models.dream.extract_atoms'); + if (configuredModel) extractModel = configuredModel; + const configuredBudget = await engine.getConfig('cycle.extract_atoms.budget_usd'); + if (configuredBudget) { + const n = Number(configuredBudget); + if (Number.isFinite(n) && n > 0) budgetCap = n; + } + } catch { + // Keep safe defaults: Haiku + $0.30. + } + const budgetTracker = new BudgetTracker({ + maxCostUsd: budgetCap, + label: 'cycle.extract_atoms', + }); // v0.41.19.0 (T3): throttled yield helper. Fires `opts.yieldDuringPhase` // every 30s. Cycle.ts threads `buildYieldDuringPhase(lock, outer)` so @@ -525,9 +544,10 @@ export async function runPhaseExtractAtoms( } } + await withBudgetTracker(budgetTracker, async () => { for (const item of work) { await maybeYield(); - if (estimatedSpendUsd >= budgetCap) { + if (budgetExhausted || budgetTracker.totalSpent >= budgetCap) { if (item.kind === 'transcript') transcriptsSkipped++; else pagesSkipped++; continue; @@ -536,6 +556,7 @@ export async function runPhaseExtractAtoms( const originLabel = item.kind === 'transcript' ? item.filePath : item.slug; try { const result = await chat({ + model: extractModel, system: EXTRACT_PROMPT, messages: [ { @@ -550,9 +571,7 @@ export async function runPhaseExtractAtoms( // actual refresh rate so this is cheap when calls are fast. await maybeYield(); - // Rough cost estimate — Haiku at ~$0.80/M input + $4/M output - estimatedSpendUsd += - (result.usage.input_tokens * 0.8 + result.usage.output_tokens * 4.0) / 1_000_000; + estimatedSpendUsd = budgetTracker.totalSpent; const atoms = parseAtomsResponse(result.text); if (atoms.length === 0) { @@ -605,12 +624,20 @@ export async function runPhaseExtractAtoms( // Reporter rate-limits to ~1 line/sec; safe to tick every iter. opts.progress?.tick(1, `${totalAtomsExtracted} atoms / ${duplicatesSkipped} skipped`); } catch (err) { + if (err instanceof BudgetExhausted) { + budgetExhausted = true; + if (item.kind === 'transcript') transcriptsSkipped++; + else pagesSkipped++; + continue; + } failures.push({ source: originLabel, error: err instanceof Error ? err.message : String(err), }); } } + }); + estimatedSpendUsd = budgetTracker.totalSpent; // v0.42 Wave B2: write extract receipt + rollup row when the phase // actually extracted atoms. Both are best-effort per F-OUT-19 — @@ -668,6 +695,8 @@ export async function runPhaseExtractAtoms( failures, estimated_spend_usd: estimatedSpendUsd, budget_usd: budgetCap, + model: extractModel, + budget_exhausted: budgetExhausted, source_id: sourceId, dry_run: opts.dryRun ?? false, }, diff --git a/test/cycle/extract-atoms-progress.test.ts b/test/cycle/extract-atoms-progress.test.ts index 0852555ce..62a9a5b10 100644 --- a/test/cycle/extract-atoms-progress.test.ts +++ b/test/cycle/extract-atoms-progress.test.ts @@ -117,6 +117,43 @@ describe('extract_atoms progress wiring (T4)', () => { expect(ticks[0].note).toMatch(/atoms.*skipped/); }); + test('passes an explicit Haiku model to chat calls', async () => { + const seenModels: Array = []; + const validAtomJson = JSON.stringify([ + { title: 'A', atom_type: 'insight', body: 'body a' }, + ]); + await runPhaseExtractAtoms(engine, { + sourceId: 'default', + _transcripts: [ + { filePath: '/tmp/t1.txt', content: 'transcript 1 body', contentHash: 'h1'.repeat(8) }, + ], + _pages: [], + _chat: async (o: ChatOpts) => { + seenModels.push(o.model); + return stubChat(validAtomJson)(o); + }, + }); + expect(seenModels).toEqual(['anthropic:claude-haiku-4-5']); + }); + + test('DB config can override the extract_atoms budget and model', async () => { + await engine.setConfig('models.dream.extract_atoms', 'anthropic:claude-haiku-4-5-20251001'); + await engine.setConfig('cycle.extract_atoms.budget_usd', '0.12'); + const validAtomJson = JSON.stringify([ + { title: 'A', atom_type: 'insight', body: 'body a' }, + ]); + const result = await runPhaseExtractAtoms(engine, { + sourceId: 'default', + _transcripts: [ + { filePath: '/tmp/t1.txt', content: 'transcript 1 body', contentHash: 'h1'.repeat(8) }, + ], + _pages: [], + _chat: stubChat(validAtomJson), + }); + expect(result.details.model).toBe('anthropic:claude-haiku-4-5-20251001'); + expect(result.details.budget_usd).toBe(0.12); + }); + test('no progress wiring required — opts.progress is optional', async () => { // Sanity: phase works without a reporter. const result = await runPhaseExtractAtoms(engine, {