File size: 3,865 Bytes
4e23b01 | 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 | import { Disposable } from '#/_base/di/lifecycle';
import type { AgentContext } from '#/agent/agentContext/agentContext';
import { agentSpaceOf } from '#/agent/agentContext/agentSpace';
import { TurnEnded } from '#/agent/loop/turnOps';
import { IConfigService } from '#/app/config/config';
import { ISessionEventBus } from '#/app/event/eventBus';
import {
TOKEN_COUNTING_SECTION,
type TokenCountingConfig,
} from '#/agent/tokenCounting/configSection';
import type {
ContextSize,
TokenCountingRequest,
TokenCountingStrategy,
} from '#/agent/tokenCounting/tokenCounting';
import type { Message } from '#/llm-adapter/contract/message';
import type { ToolDescription as Tool } from '#human/llm/message';
import {
estimateTokens,
estimateTokensForMessage,
estimateTokensForMessages,
estimateTokensForTools,
} from '#/llm-adapter/contract/tokens';
import type { TokenUsage } from '#human/llm/usage';
import { IAgentLifecycleService } from '#/session/agentLifecycle/agentLifecycle';
import {
ISessionTokenCountingService,
type TokenCountingRebaseInput,
} from './sessionTokenCounting';
import { TokenCountingAgentModelDefinition } from './tokenCountingAgentModel';
export class SessionTokenCountingService extends Disposable implements ISessionTokenCountingService {
declare readonly _serviceBrand: undefined;
constructor(
@IConfigService private readonly config: IConfigService,
@ISessionEventBus eventBus: ISessionEventBus,
@IAgentLifecycleService agentLifecycle: IAgentLifecycleService,
) {
super();
this._register(
eventBus.subscribe(TurnEnded, (event) => {
const agent = agentLifecycle.get(event.agentId);
if (agent === undefined) return;
void agentSpaceOf(agent).use(
TokenCountingAgentModelDefinition,
(model) => model.recordTurn(event.turnId, this.strategy),
);
}),
);
}
get strategy(): TokenCountingStrategy {
return (
this.config.get<TokenCountingConfig>(TOKEN_COUNTING_SECTION)?.strategy ??
'measured+estimated'
);
}
get(agent: AgentContext, start?: number, end?: number): ContextSize {
return agentSpaceOf(agent).use(TokenCountingAgentModelDefinition, (model) =>
model.get(start, end),
);
}
measured(
agent: AgentContext,
input: readonly Message[],
output: readonly Message[],
usage: TokenUsage,
): void {
void agentSpaceOf(agent).use(TokenCountingAgentModelDefinition, (model) =>
model.measured(input, output, usage),
);
}
latestMeasured(agent: AgentContext): number {
return agentSpaceOf(agent).use(TokenCountingAgentModelDefinition, (model) =>
model.latestMeasured(),
);
}
statusSize(agent: AgentContext): number {
return agentSpaceOf(agent).use(TokenCountingAgentModelDefinition, (model) =>
model.statusSize(this.strategy),
);
}
recordTruncation(agent: AgentContext, cutIndex: number): void {
void agentSpaceOf(agent).use(TokenCountingAgentModelDefinition, (model) =>
model.recordTruncation(cutIndex),
);
}
rebase(agent: AgentContext, input: TokenCountingRebaseInput): void {
void agentSpaceOf(agent).use(TokenCountingAgentModelDefinition, (model) =>
model.rebase(input),
);
}
requestSize(request: TokenCountingRequest): number {
return (
this.estimateText(request.systemPrompt) +
this.estimateTools(request.tools) +
this.estimateMessages(request.messages)
);
}
estimateText(text: string): number {
return estimateTokens(text);
}
estimateMessage(message: Message): number {
return estimateTokensForMessage(message);
}
estimateMessages(messages: readonly Message[]): number {
return estimateTokensForMessages(messages);
}
estimateTools(tools: readonly Tool[]): number {
return estimateTokensForTools(tools);
}
}
|