File size: 2,984 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 | import { z } from 'zod';
import type { AgentLLMRequestSource } from '#/agent/llmRequester/llmRequester';
import type { UsageStatus } from '#/agent/usage/usage';
import { AgentStatusUpdated } from '#/agent/usage/usageEvents';
import {
copyUsage,
UsageRecord,
type UsageModelState,
type UsageRecordScope,
} from '#/agent/usage/usageOps';
import { addUsage, type TokenUsage } from '#human/llm/usage';
import { AgentModel, defineAgentModel, type AgentModelContext } from '#/state/agentModel';
export interface UsageRecordInput {
readonly model: string;
readonly usage: TokenUsage;
readonly source?: AgentLLMRequestSource;
}
export class UsageAgentModel extends AgentModel<UsageModelState> {
private currentTurnId: number | undefined;
private currentTurn: TokenUsage | undefined;
constructor(context: AgentModelContext) {
super(context);
this.on(UsageRecord, (event) => {
const current = this.state.byModel[event.model];
this.state.byModel[event.model] =
current === undefined ? copyUsage(event.usage) : addUsage(current, event.usage);
});
}
record(input: UsageRecordInput): Promise<boolean> {
const firstRecord = Object.keys(this.state.byModel).length === 0;
const usageScope: UsageRecordScope = input.source?.type === 'turn' ? 'turn' : 'session';
const recorded = this.emit(
new UsageRecord({
agentId: this.agent.agentId,
model: input.model,
usage: input.usage,
usageScope,
}),
);
const turnId = input.source?.type === 'turn' ? input.source.turnId : undefined;
if (turnId !== undefined) {
if (this.currentTurnId !== turnId) {
this.currentTurnId = turnId;
this.currentTurn = copyUsage(input.usage);
} else {
this.currentTurn =
this.currentTurn === undefined
? copyUsage(input.usage)
: addUsage(this.currentTurn, input.usage);
}
}
const notified = this.emit(
new AgentStatusUpdated({ agentId: this.agent.agentId, usage: this.status() }),
);
return recorded.then(() => notified).then(() => firstRecord);
}
status(): UsageStatus {
const byModel = Object.fromEntries(
Object.entries(this.state.byModel).map(([model, usage]) => [model, copyUsage(usage)]),
);
const hasByModel = Object.keys(byModel).length > 0;
let total: TokenUsage | undefined;
if (hasByModel) {
for (const usage of Object.values(byModel)) {
total = total === undefined ? copyUsage(usage) : addUsage(total, usage);
}
}
return {
byModel: hasByModel ? byModel : undefined,
total,
currentTurn: this.currentTurn === undefined ? undefined : copyUsage(this.currentTurn),
};
}
}
export const UsageAgentModelDefinition = defineAgentModel({
id: 'usage',
model: UsageAgentModel,
state: {
initial: (): UsageModelState => ({ byModel: {} }),
schema: z.custom<UsageModelState>(),
},
events: [UsageRecord],
});
|