File size: 4,681 Bytes
7a1ad33 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | /**
* @license
* Copyright 2026 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'node:crypto';
import { ContextManager } from '../contextManager.js';
import { AgentChatHistory } from '../../core/agentChatHistory.js';
import type { Content } from '@google/genai';
import type { ContextProfile } from '../config/profiles.js';
import { ContextEnvironmentImpl } from '../pipeline/environmentImpl.js';
import { ContextTracer } from '../tracer.js';
import { ContextEventBus } from '../eventBus.js';
import { PipelineOrchestrator } from '../pipeline/orchestrator.js';
import type { BaseLlmClient } from '../../core/baseLlmClient.js';
import { StaticTokenCalculator } from '../utils/contextTokenCalculator.js';
import { NodeBehaviorRegistry } from '../graph/behaviorRegistry.js';
import { registerBuiltInBehaviors } from '../graph/builtinBehaviors.js';
export interface TurnSummary {
turnIndex: number;
tokensBeforeBackground: number;
tokensAfterBackground: number;
}
export class SimulationHarness {
readonly chatHistory: AgentChatHistory;
contextManager!: ContextManager;
env!: ContextEnvironmentImpl;
orchestrator!: PipelineOrchestrator;
readonly eventBus: ContextEventBus;
config!: ContextProfile;
private tracer!: ContextTracer;
private currentTurnIndex = 0;
private tokenTrajectory: TurnSummary[] = [];
static async create(
config: ContextProfile,
mockLlmClient: BaseLlmClient,
mockTempDir = '/tmp/sim',
): Promise<SimulationHarness> {
const harness = new SimulationHarness();
await harness.init(config, mockLlmClient, mockTempDir);
return harness;
}
private constructor() {
this.chatHistory = new AgentChatHistory();
this.eventBus = new ContextEventBus();
}
private async init(
config: ContextProfile,
mockLlmClient: BaseLlmClient,
mockTempDir: string,
) {
this.config = config;
this.tracer = new ContextTracer({
targetDir: mockTempDir,
sessionId: 'sim-session',
});
const behaviorRegistry = new NodeBehaviorRegistry();
registerBuiltInBehaviors(behaviorRegistry);
const calculator = new StaticTokenCalculator(1, behaviorRegistry);
this.env = new ContextEnvironmentImpl(
() => mockLlmClient,
'sim-prompt',
'sim-session',
mockTempDir,
mockTempDir,
this.tracer,
1, // 1 char per token average for estimation (but estimator uses 0.33)
this.eventBus,
calculator,
behaviorRegistry,
);
this.orchestrator = new PipelineOrchestrator(
config.buildPipelines(this.env),
config.buildAsyncPipelines(this.env),
this.env,
this.tracer,
);
this.contextManager = new ContextManager(
config,
this.env,
this.tracer,
this.orchestrator,
this.chatHistory,
calculator,
);
}
async simulateTurn(messages: Content[]) {
// In the new turn-based flow, we simulate the 'next' prompt or turn
// by calling renderHistory on the pending content.
// For the purpose of the simulation, we'll treat the first message as the 'pending' one
// if it hasn't been added to history yet.
const pendingContent = messages[messages.length - 1];
// 1. Render to trigger sync and management
const { processedNodes } = await this.contextManager.renderHistory({
id: randomUUID(),
content: pendingContent,
});
const tokensBefore =
this.env.tokenCalculator.calculateConcreteListTokens(processedNodes);
// 2. Append the new messages to durable history
const currentHistory = this.chatHistory.get();
const turns = messages.map((m) => ({ id: randomUUID(), content: m }));
this.chatHistory.set([...currentHistory, ...turns]);
// 3. Wait for any async pipelines triggered by the sync
await this.contextManager.waitForPipelines();
// 4. Measure tokens after background processors (requires another render or sync check)
// In the new model, we'd need to re-render to see the effect of async processors
// that might have finished.
const { processedNodes: nodesAfter } =
await this.contextManager.renderHistory();
const tokensAfter =
this.env.tokenCalculator.calculateConcreteListTokens(nodesAfter);
this.tokenTrajectory.push({
turnIndex: this.currentTurnIndex++,
tokensBeforeBackground: tokensBefore,
tokensAfterBackground: tokensAfter,
});
}
async getGoldenState() {
const { history: finalProjection, baseUnits } =
await this.contextManager.renderHistory();
return {
tokenTrajectory: this.tokenTrajectory,
finalProjection,
baseUnits,
};
}
}
|