feat: Change /stats to include more detailed breakdowns (#2615)

This commit is contained in:
Abhi
2025-06-29 20:44:33 -04:00
committed by GitHub
parent 0fd602eb43
commit 770f862832
36 changed files with 3218 additions and 758 deletions

View File

@@ -10,14 +10,8 @@ import {
GeminiEventType,
ServerGeminiToolCallRequestEvent,
ServerGeminiErrorEvent,
ServerGeminiUsageMetadataEvent,
} from './turn.js';
import {
GenerateContentResponse,
Part,
Content,
GenerateContentResponseUsageMetadata,
} from '@google/genai';
import { GenerateContentResponse, Part, Content } from '@google/genai';
import { reportError } from '../utils/errorReporting.js';
import { GeminiChat } from './geminiChat.js';
@@ -55,24 +49,6 @@ describe('Turn', () => {
};
let mockChatInstance: MockedChatInstance;
const mockMetadata1: GenerateContentResponseUsageMetadata = {
promptTokenCount: 10,
candidatesTokenCount: 20,
totalTokenCount: 30,
cachedContentTokenCount: 5,
toolUsePromptTokenCount: 2,
thoughtsTokenCount: 3,
};
const mockMetadata2: GenerateContentResponseUsageMetadata = {
promptTokenCount: 100,
candidatesTokenCount: 200,
totalTokenCount: 300,
cachedContentTokenCount: 50,
toolUsePromptTokenCount: 20,
thoughtsTokenCount: 30,
};
beforeEach(() => {
vi.resetAllMocks();
mockChatInstance = {
@@ -245,46 +221,6 @@ describe('Turn', () => {
);
});
it('should yield the last UsageMetadata event from the stream', async () => {
const mockResponseStream = (async function* () {
yield {
candidates: [{ content: { parts: [{ text: 'First response' }] } }],
usageMetadata: mockMetadata1,
} as unknown as GenerateContentResponse;
// Add a small delay to ensure apiTimeMs is > 0
await new Promise((resolve) => setTimeout(resolve, 10));
yield {
functionCalls: [{ name: 'aTool' }],
usageMetadata: mockMetadata2,
} as unknown as GenerateContentResponse;
})();
mockSendMessageStream.mockResolvedValue(mockResponseStream);
const events = [];
const reqParts: Part[] = [{ text: 'Test metadata' }];
for await (const event of turn.run(
reqParts,
new AbortController().signal,
)) {
events.push(event);
}
// There should be a content event, a tool call, and our metadata event
expect(events.length).toBe(3);
const metadataEvent = events[2] as ServerGeminiUsageMetadataEvent;
expect(metadataEvent.type).toBe(GeminiEventType.UsageMetadata);
// The value should be the *last* metadata object received.
expect(metadataEvent.value).toEqual(
expect.objectContaining(mockMetadata2),
);
expect(metadataEvent.value.apiTimeMs).toBeGreaterThan(0);
// Also check the public getter
expect(turn.getUsageMetadata()).toEqual(mockMetadata2);
});
it('should handle function calls with undefined name or args', async () => {
const mockResponseStream = (async function* () {
yield {

View File

@@ -9,7 +9,6 @@ import {
GenerateContentResponse,
FunctionCall,
FunctionDeclaration,
GenerateContentResponseUsageMetadata,
} from '@google/genai';
import {
ToolCallConfirmationDetails,
@@ -48,7 +47,6 @@ export enum GeminiEventType {
UserCancelled = 'user_cancelled',
Error = 'error',
ChatCompressed = 'chat_compressed',
UsageMetadata = 'usage_metadata',
Thought = 'thought',
}
@@ -129,11 +127,6 @@ export type ServerGeminiChatCompressedEvent = {
value: ChatCompressionInfo | null;
};
export type ServerGeminiUsageMetadataEvent = {
type: GeminiEventType.UsageMetadata;
value: GenerateContentResponseUsageMetadata & { apiTimeMs?: number };
};
// The original union type, now composed of the individual types
export type ServerGeminiStreamEvent =
| ServerGeminiContentEvent
@@ -143,14 +136,12 @@ export type ServerGeminiStreamEvent =
| ServerGeminiUserCancelledEvent
| ServerGeminiErrorEvent
| ServerGeminiChatCompressedEvent
| ServerGeminiUsageMetadataEvent
| ServerGeminiThoughtEvent;
// A turn manages the agentic loop turn within the server context.
export class Turn {
readonly pendingToolCalls: ToolCallRequestInfo[];
private debugResponses: GenerateContentResponse[];
private lastUsageMetadata: GenerateContentResponseUsageMetadata | null = null;
constructor(private readonly chat: GeminiChat) {
this.pendingToolCalls = [];
@@ -161,7 +152,6 @@ export class Turn {
req: PartListUnion,
signal: AbortSignal,
): AsyncGenerator<ServerGeminiStreamEvent> {
const startTime = Date.now();
try {
const responseStream = await this.chat.sendMessageStream({
message: req,
@@ -213,19 +203,6 @@ export class Turn {
yield event;
}
}
if (resp.usageMetadata) {
this.lastUsageMetadata =
resp.usageMetadata as GenerateContentResponseUsageMetadata;
}
}
if (this.lastUsageMetadata) {
const durationMs = Date.now() - startTime;
yield {
type: GeminiEventType.UsageMetadata,
value: { ...this.lastUsageMetadata, apiTimeMs: durationMs },
};
}
} catch (e) {
const error = toFriendlyError(e);
@@ -286,8 +263,4 @@ export class Turn {
getDebugResponses(): GenerateContentResponse[] {
return this.debugResponses;
}
getUsageMetadata(): GenerateContentResponseUsageMetadata | null {
return this.lastUsageMetadata;
}
}

View File

@@ -38,3 +38,4 @@ export {
} from './types.js';
export { SpanStatusCode, ValueType } from '@opentelemetry/api';
export { SemanticAttributes } from '@opentelemetry/semantic-conventions';
export * from './uiTelemetry.js';

View File

@@ -43,15 +43,22 @@ import * as metrics from './metrics.js';
import * as sdk from './sdk.js';
import { vi, describe, beforeEach, it, expect } from 'vitest';
import { GenerateContentResponseUsageMetadata } from '@google/genai';
import * as uiTelemetry from './uiTelemetry.js';
describe('loggers', () => {
const mockLogger = {
emit: vi.fn(),
};
const mockUiEvent = {
addEvent: vi.fn(),
};
beforeEach(() => {
vi.spyOn(sdk, 'isTelemetrySdkInitialized').mockReturnValue(true);
vi.spyOn(logs, 'getLogger').mockReturnValue(mockLogger);
vi.spyOn(uiTelemetry.uiTelemetryService, 'addEvent').mockImplementation(
mockUiEvent.addEvent,
);
vi.useFakeTimers();
vi.setSystemTime(new Date('2025-01-01T00:00:00.000Z'));
});
@@ -215,6 +222,7 @@ describe('loggers', () => {
cached_content_token_count: 10,
thoughts_token_count: 5,
tool_token_count: 2,
total_token_count: 0,
response_text: 'test-response',
},
});
@@ -233,6 +241,12 @@ describe('loggers', () => {
50,
'output',
);
expect(mockUiEvent.addEvent).toHaveBeenCalledWith({
...event,
'event.name': EVENT_API_RESPONSE,
'event.timestamp': '2025-01-01T00:00:00.000Z',
});
});
it('should log an API response with an error', () => {
@@ -263,6 +277,12 @@ describe('loggers', () => {
'error.message': 'test-error',
},
});
expect(mockUiEvent.addEvent).toHaveBeenCalledWith({
...event,
'event.name': EVENT_API_RESPONSE,
'event.timestamp': '2025-01-01T00:00:00.000Z',
});
});
});
@@ -417,6 +437,12 @@ describe('loggers', () => {
true,
ToolCallDecision.ACCEPT,
);
expect(mockUiEvent.addEvent).toHaveBeenCalledWith({
...event,
'event.name': EVENT_TOOL_CALL,
'event.timestamp': '2025-01-01T00:00:00.000Z',
});
});
it('should log a tool call with a reject decision', () => {
const call: ErroredToolCall = {
@@ -471,6 +497,12 @@ describe('loggers', () => {
false,
ToolCallDecision.REJECT,
);
expect(mockUiEvent.addEvent).toHaveBeenCalledWith({
...event,
'event.name': EVENT_TOOL_CALL,
'event.timestamp': '2025-01-01T00:00:00.000Z',
});
});
it('should log a tool call with a modify decision', () => {
@@ -527,6 +559,12 @@ describe('loggers', () => {
true,
ToolCallDecision.MODIFY,
);
expect(mockUiEvent.addEvent).toHaveBeenCalledWith({
...event,
'event.name': EVENT_TOOL_CALL,
'event.timestamp': '2025-01-01T00:00:00.000Z',
});
});
it('should log a tool call without a decision', () => {
@@ -581,6 +619,12 @@ describe('loggers', () => {
true,
undefined,
);
expect(mockUiEvent.addEvent).toHaveBeenCalledWith({
...event,
'event.name': EVENT_TOOL_CALL,
'event.timestamp': '2025-01-01T00:00:00.000Z',
});
});
it('should log a failed tool call with an error', () => {
@@ -641,6 +685,12 @@ describe('loggers', () => {
false,
undefined,
);
expect(mockUiEvent.addEvent).toHaveBeenCalledWith({
...event,
'event.name': EVENT_TOOL_CALL,
'event.timestamp': '2025-01-01T00:00:00.000Z',
});
});
});
});

View File

@@ -31,6 +31,7 @@ import {
recordToolCallMetrics,
} from './metrics.js';
import { isTelemetrySdkInitialized } from './sdk.js';
import { uiTelemetryService, UiEvent } from './uiTelemetry.js';
import { ClearcutLogger } from './clearcut-logger/clearcut-logger.js';
const shouldLogUserPrompts = (config: Config): boolean =>
@@ -98,6 +99,12 @@ export function logUserPrompt(config: Config, event: UserPromptEvent): void {
}
export function logToolCall(config: Config, event: ToolCallEvent): void {
const uiEvent = {
...event,
'event.name': EVENT_TOOL_CALL,
'event.timestamp': new Date().toISOString(),
} as UiEvent;
uiTelemetryService.addEvent(uiEvent);
ClearcutLogger.getInstance(config)?.logToolCallEvent(event);
if (!isTelemetrySdkInitialized()) return;
@@ -150,6 +157,12 @@ export function logApiRequest(config: Config, event: ApiRequestEvent): void {
}
export function logApiError(config: Config, event: ApiErrorEvent): void {
const uiEvent = {
...event,
'event.name': EVENT_API_ERROR,
'event.timestamp': new Date().toISOString(),
} as UiEvent;
uiTelemetryService.addEvent(uiEvent);
ClearcutLogger.getInstance(config)?.logApiErrorEvent(event);
if (!isTelemetrySdkInitialized()) return;
@@ -186,6 +199,12 @@ export function logApiError(config: Config, event: ApiErrorEvent): void {
}
export function logApiResponse(config: Config, event: ApiResponseEvent): void {
const uiEvent = {
...event,
'event.name': EVENT_API_RESPONSE,
'event.timestamp': new Date().toISOString(),
} as UiEvent;
uiTelemetryService.addEvent(uiEvent);
ClearcutLogger.getInstance(config)?.logApiResponseEvent(event);
if (!isTelemetrySdkInitialized()) return;
const attributes: LogAttributes = {

View File

@@ -183,6 +183,7 @@ export class ApiResponseEvent {
cached_content_token_count: number;
thoughts_token_count: number;
tool_token_count: number;
total_token_count: number;
response_text?: string;
constructor(
@@ -202,6 +203,7 @@ export class ApiResponseEvent {
this.cached_content_token_count = usage_data?.cachedContentTokenCount ?? 0;
this.thoughts_token_count = usage_data?.thoughtsTokenCount ?? 0;
this.tool_token_count = usage_data?.toolUsePromptTokenCount ?? 0;
this.total_token_count = usage_data?.totalTokenCount ?? 0;
this.response_text = response_text;
this.error = error;
}

View File

@@ -0,0 +1,510 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { UiTelemetryService } from './uiTelemetry.js';
import {
ApiErrorEvent,
ApiResponseEvent,
ToolCallEvent,
ToolCallDecision,
} from './types.js';
import {
EVENT_API_ERROR,
EVENT_API_RESPONSE,
EVENT_TOOL_CALL,
} from './constants.js';
import {
CompletedToolCall,
ErroredToolCall,
SuccessfulToolCall,
} from '../core/coreToolScheduler.js';
import { Tool, ToolConfirmationOutcome } from '../tools/tools.js';
const createFakeCompletedToolCall = (
name: string,
success: boolean,
duration = 100,
outcome?: ToolConfirmationOutcome,
error?: Error,
): CompletedToolCall => {
const request = {
callId: `call_${name}_${Date.now()}`,
name,
args: { foo: 'bar' },
isClientInitiated: false,
};
if (success) {
return {
status: 'success',
request,
tool: { name } as Tool, // Mock tool
response: {
callId: request.callId,
responseParts: {
functionResponse: {
id: request.callId,
name,
response: { output: 'Success!' },
},
},
error: undefined,
resultDisplay: 'Success!',
},
durationMs: duration,
outcome,
} as SuccessfulToolCall;
} else {
return {
status: 'error',
request,
response: {
callId: request.callId,
responseParts: {
functionResponse: {
id: request.callId,
name,
response: { error: 'Tool failed' },
},
},
error: error || new Error('Tool failed'),
resultDisplay: 'Failure!',
},
durationMs: duration,
outcome,
} as ErroredToolCall;
}
};
describe('UiTelemetryService', () => {
let service: UiTelemetryService;
beforeEach(() => {
service = new UiTelemetryService();
});
it('should have correct initial metrics', () => {
const metrics = service.getMetrics();
expect(metrics).toEqual({
models: {},
tools: {
totalCalls: 0,
totalSuccess: 0,
totalFail: 0,
totalDurationMs: 0,
totalDecisions: {
[ToolCallDecision.ACCEPT]: 0,
[ToolCallDecision.REJECT]: 0,
[ToolCallDecision.MODIFY]: 0,
},
byName: {},
},
});
expect(service.getLastPromptTokenCount()).toBe(0);
});
it('should emit an update event when an event is added', () => {
const spy = vi.fn();
service.on('update', spy);
const event = {
'event.name': EVENT_API_RESPONSE,
model: 'gemini-2.5-pro',
duration_ms: 500,
input_token_count: 10,
output_token_count: 20,
total_token_count: 30,
cached_content_token_count: 5,
thoughts_token_count: 2,
tool_token_count: 3,
} as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };
service.addEvent(event);
expect(spy).toHaveBeenCalledOnce();
const { metrics, lastPromptTokenCount } = spy.mock.calls[0][0];
expect(metrics).toBeDefined();
expect(lastPromptTokenCount).toBe(10);
});
describe('API Response Event Processing', () => {
it('should process a single ApiResponseEvent', () => {
const event = {
'event.name': EVENT_API_RESPONSE,
model: 'gemini-2.5-pro',
duration_ms: 500,
input_token_count: 10,
output_token_count: 20,
total_token_count: 30,
cached_content_token_count: 5,
thoughts_token_count: 2,
tool_token_count: 3,
} as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };
service.addEvent(event);
const metrics = service.getMetrics();
expect(metrics.models['gemini-2.5-pro']).toEqual({
api: {
totalRequests: 1,
totalErrors: 0,
totalLatencyMs: 500,
},
tokens: {
prompt: 10,
candidates: 20,
total: 30,
cached: 5,
thoughts: 2,
tool: 3,
},
});
expect(service.getLastPromptTokenCount()).toBe(10);
});
it('should aggregate multiple ApiResponseEvents for the same model', () => {
const event1 = {
'event.name': EVENT_API_RESPONSE,
model: 'gemini-2.5-pro',
duration_ms: 500,
input_token_count: 10,
output_token_count: 20,
total_token_count: 30,
cached_content_token_count: 5,
thoughts_token_count: 2,
tool_token_count: 3,
} as ApiResponseEvent & {
'event.name': typeof EVENT_API_RESPONSE;
};
const event2 = {
'event.name': EVENT_API_RESPONSE,
model: 'gemini-2.5-pro',
duration_ms: 600,
input_token_count: 15,
output_token_count: 25,
total_token_count: 40,
cached_content_token_count: 10,
thoughts_token_count: 4,
tool_token_count: 6,
} as ApiResponseEvent & {
'event.name': typeof EVENT_API_RESPONSE;
};
service.addEvent(event1);
service.addEvent(event2);
const metrics = service.getMetrics();
expect(metrics.models['gemini-2.5-pro']).toEqual({
api: {
totalRequests: 2,
totalErrors: 0,
totalLatencyMs: 1100,
},
tokens: {
prompt: 25,
candidates: 45,
total: 70,
cached: 15,
thoughts: 6,
tool: 9,
},
});
expect(service.getLastPromptTokenCount()).toBe(15);
});
it('should handle ApiResponseEvents for different models', () => {
const event1 = {
'event.name': EVENT_API_RESPONSE,
model: 'gemini-2.5-pro',
duration_ms: 500,
input_token_count: 10,
output_token_count: 20,
total_token_count: 30,
cached_content_token_count: 5,
thoughts_token_count: 2,
tool_token_count: 3,
} as ApiResponseEvent & {
'event.name': typeof EVENT_API_RESPONSE;
};
const event2 = {
'event.name': EVENT_API_RESPONSE,
model: 'gemini-2.5-flash',
duration_ms: 1000,
input_token_count: 100,
output_token_count: 200,
total_token_count: 300,
cached_content_token_count: 50,
thoughts_token_count: 20,
tool_token_count: 30,
} as ApiResponseEvent & {
'event.name': typeof EVENT_API_RESPONSE;
};
service.addEvent(event1);
service.addEvent(event2);
const metrics = service.getMetrics();
expect(metrics.models['gemini-2.5-pro']).toBeDefined();
expect(metrics.models['gemini-2.5-flash']).toBeDefined();
expect(metrics.models['gemini-2.5-pro'].api.totalRequests).toBe(1);
expect(metrics.models['gemini-2.5-flash'].api.totalRequests).toBe(1);
expect(service.getLastPromptTokenCount()).toBe(100);
});
});
describe('API Error Event Processing', () => {
it('should process a single ApiErrorEvent', () => {
const event = {
'event.name': EVENT_API_ERROR,
model: 'gemini-2.5-pro',
duration_ms: 300,
error: 'Something went wrong',
} as ApiErrorEvent & { 'event.name': typeof EVENT_API_ERROR };
service.addEvent(event);
const metrics = service.getMetrics();
expect(metrics.models['gemini-2.5-pro']).toEqual({
api: {
totalRequests: 1,
totalErrors: 1,
totalLatencyMs: 300,
},
tokens: {
prompt: 0,
candidates: 0,
total: 0,
cached: 0,
thoughts: 0,
tool: 0,
},
});
});
it('should aggregate ApiErrorEvents and ApiResponseEvents', () => {
const responseEvent = {
'event.name': EVENT_API_RESPONSE,
model: 'gemini-2.5-pro',
duration_ms: 500,
input_token_count: 10,
output_token_count: 20,
total_token_count: 30,
cached_content_token_count: 5,
thoughts_token_count: 2,
tool_token_count: 3,
} as ApiResponseEvent & {
'event.name': typeof EVENT_API_RESPONSE;
};
const errorEvent = {
'event.name': EVENT_API_ERROR,
model: 'gemini-2.5-pro',
duration_ms: 300,
error: 'Something went wrong',
} as ApiErrorEvent & { 'event.name': typeof EVENT_API_ERROR };
service.addEvent(responseEvent);
service.addEvent(errorEvent);
const metrics = service.getMetrics();
expect(metrics.models['gemini-2.5-pro']).toEqual({
api: {
totalRequests: 2,
totalErrors: 1,
totalLatencyMs: 800,
},
tokens: {
prompt: 10,
candidates: 20,
total: 30,
cached: 5,
thoughts: 2,
tool: 3,
},
});
});
});
describe('Tool Call Event Processing', () => {
it('should process a single successful ToolCallEvent', () => {
const toolCall = createFakeCompletedToolCall(
'test_tool',
true,
150,
ToolConfirmationOutcome.ProceedOnce,
);
service.addEvent({
...JSON.parse(JSON.stringify(new ToolCallEvent(toolCall))),
'event.name': EVENT_TOOL_CALL,
});
const metrics = service.getMetrics();
const { tools } = metrics;
expect(tools.totalCalls).toBe(1);
expect(tools.totalSuccess).toBe(1);
expect(tools.totalFail).toBe(0);
expect(tools.totalDurationMs).toBe(150);
expect(tools.totalDecisions[ToolCallDecision.ACCEPT]).toBe(1);
expect(tools.byName['test_tool']).toEqual({
count: 1,
success: 1,
fail: 0,
durationMs: 150,
decisions: {
[ToolCallDecision.ACCEPT]: 1,
[ToolCallDecision.REJECT]: 0,
[ToolCallDecision.MODIFY]: 0,
},
});
});
it('should process a single failed ToolCallEvent', () => {
const toolCall = createFakeCompletedToolCall(
'test_tool',
false,
200,
ToolConfirmationOutcome.Cancel,
);
service.addEvent({
...JSON.parse(JSON.stringify(new ToolCallEvent(toolCall))),
'event.name': EVENT_TOOL_CALL,
});
const metrics = service.getMetrics();
const { tools } = metrics;
expect(tools.totalCalls).toBe(1);
expect(tools.totalSuccess).toBe(0);
expect(tools.totalFail).toBe(1);
expect(tools.totalDurationMs).toBe(200);
expect(tools.totalDecisions[ToolCallDecision.REJECT]).toBe(1);
expect(tools.byName['test_tool']).toEqual({
count: 1,
success: 0,
fail: 1,
durationMs: 200,
decisions: {
[ToolCallDecision.ACCEPT]: 0,
[ToolCallDecision.REJECT]: 1,
[ToolCallDecision.MODIFY]: 0,
},
});
});
it('should process a ToolCallEvent with modify decision', () => {
const toolCall = createFakeCompletedToolCall(
'test_tool',
true,
250,
ToolConfirmationOutcome.ModifyWithEditor,
);
service.addEvent({
...JSON.parse(JSON.stringify(new ToolCallEvent(toolCall))),
'event.name': EVENT_TOOL_CALL,
});
const metrics = service.getMetrics();
const { tools } = metrics;
expect(tools.totalDecisions[ToolCallDecision.MODIFY]).toBe(1);
expect(tools.byName['test_tool'].decisions[ToolCallDecision.MODIFY]).toBe(
1,
);
});
it('should process a ToolCallEvent without a decision', () => {
const toolCall = createFakeCompletedToolCall('test_tool', true, 100);
service.addEvent({
...JSON.parse(JSON.stringify(new ToolCallEvent(toolCall))),
'event.name': EVENT_TOOL_CALL,
});
const metrics = service.getMetrics();
const { tools } = metrics;
expect(tools.totalDecisions).toEqual({
[ToolCallDecision.ACCEPT]: 0,
[ToolCallDecision.REJECT]: 0,
[ToolCallDecision.MODIFY]: 0,
});
expect(tools.byName['test_tool'].decisions).toEqual({
[ToolCallDecision.ACCEPT]: 0,
[ToolCallDecision.REJECT]: 0,
[ToolCallDecision.MODIFY]: 0,
});
});
it('should aggregate multiple ToolCallEvents for the same tool', () => {
const toolCall1 = createFakeCompletedToolCall(
'test_tool',
true,
100,
ToolConfirmationOutcome.ProceedOnce,
);
const toolCall2 = createFakeCompletedToolCall(
'test_tool',
false,
150,
ToolConfirmationOutcome.Cancel,
);
service.addEvent({
...JSON.parse(JSON.stringify(new ToolCallEvent(toolCall1))),
'event.name': EVENT_TOOL_CALL,
});
service.addEvent({
...JSON.parse(JSON.stringify(new ToolCallEvent(toolCall2))),
'event.name': EVENT_TOOL_CALL,
});
const metrics = service.getMetrics();
const { tools } = metrics;
expect(tools.totalCalls).toBe(2);
expect(tools.totalSuccess).toBe(1);
expect(tools.totalFail).toBe(1);
expect(tools.totalDurationMs).toBe(250);
expect(tools.totalDecisions[ToolCallDecision.ACCEPT]).toBe(1);
expect(tools.totalDecisions[ToolCallDecision.REJECT]).toBe(1);
expect(tools.byName['test_tool']).toEqual({
count: 2,
success: 1,
fail: 1,
durationMs: 250,
decisions: {
[ToolCallDecision.ACCEPT]: 1,
[ToolCallDecision.REJECT]: 1,
[ToolCallDecision.MODIFY]: 0,
},
});
});
it('should handle ToolCallEvents for different tools', () => {
const toolCall1 = createFakeCompletedToolCall('tool_A', true, 100);
const toolCall2 = createFakeCompletedToolCall('tool_B', false, 200);
service.addEvent({
...JSON.parse(JSON.stringify(new ToolCallEvent(toolCall1))),
'event.name': EVENT_TOOL_CALL,
});
service.addEvent({
...JSON.parse(JSON.stringify(new ToolCallEvent(toolCall2))),
'event.name': EVENT_TOOL_CALL,
});
const metrics = service.getMetrics();
const { tools } = metrics;
expect(tools.totalCalls).toBe(2);
expect(tools.totalSuccess).toBe(1);
expect(tools.totalFail).toBe(1);
expect(tools.byName['tool_A']).toBeDefined();
expect(tools.byName['tool_B']).toBeDefined();
expect(tools.byName['tool_A'].count).toBe(1);
expect(tools.byName['tool_B'].count).toBe(1);
});
});
});

View File

@@ -0,0 +1,207 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { EventEmitter } from 'events';
import {
EVENT_API_ERROR,
EVENT_API_RESPONSE,
EVENT_TOOL_CALL,
} from './constants.js';
import {
ApiErrorEvent,
ApiResponseEvent,
ToolCallEvent,
ToolCallDecision,
} from './types.js';
export type UiEvent =
| (ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE })
| (ApiErrorEvent & { 'event.name': typeof EVENT_API_ERROR })
| (ToolCallEvent & { 'event.name': typeof EVENT_TOOL_CALL });
export interface ToolCallStats {
count: number;
success: number;
fail: number;
durationMs: number;
decisions: {
[ToolCallDecision.ACCEPT]: number;
[ToolCallDecision.REJECT]: number;
[ToolCallDecision.MODIFY]: number;
};
}
export interface ModelMetrics {
api: {
totalRequests: number;
totalErrors: number;
totalLatencyMs: number;
};
tokens: {
prompt: number;
candidates: number;
total: number;
cached: number;
thoughts: number;
tool: number;
};
}
export interface SessionMetrics {
models: Record<string, ModelMetrics>;
tools: {
totalCalls: number;
totalSuccess: number;
totalFail: number;
totalDurationMs: number;
totalDecisions: {
[ToolCallDecision.ACCEPT]: number;
[ToolCallDecision.REJECT]: number;
[ToolCallDecision.MODIFY]: number;
};
byName: Record<string, ToolCallStats>;
};
}
const createInitialModelMetrics = (): ModelMetrics => ({
api: {
totalRequests: 0,
totalErrors: 0,
totalLatencyMs: 0,
},
tokens: {
prompt: 0,
candidates: 0,
total: 0,
cached: 0,
thoughts: 0,
tool: 0,
},
});
const createInitialMetrics = (): SessionMetrics => ({
models: {},
tools: {
totalCalls: 0,
totalSuccess: 0,
totalFail: 0,
totalDurationMs: 0,
totalDecisions: {
[ToolCallDecision.ACCEPT]: 0,
[ToolCallDecision.REJECT]: 0,
[ToolCallDecision.MODIFY]: 0,
},
byName: {},
},
});
export class UiTelemetryService extends EventEmitter {
#metrics: SessionMetrics = createInitialMetrics();
#lastPromptTokenCount = 0;
addEvent(event: UiEvent) {
switch (event['event.name']) {
case EVENT_API_RESPONSE:
this.processApiResponse(event);
break;
case EVENT_API_ERROR:
this.processApiError(event);
break;
case EVENT_TOOL_CALL:
this.processToolCall(event);
break;
default:
// We should not emit update for any other event metric.
return;
}
this.emit('update', {
metrics: this.#metrics,
lastPromptTokenCount: this.#lastPromptTokenCount,
});
}
getMetrics(): SessionMetrics {
return this.#metrics;
}
getLastPromptTokenCount(): number {
return this.#lastPromptTokenCount;
}
private getOrCreateModelMetrics(modelName: string): ModelMetrics {
if (!this.#metrics.models[modelName]) {
this.#metrics.models[modelName] = createInitialModelMetrics();
}
return this.#metrics.models[modelName];
}
private processApiResponse(event: ApiResponseEvent) {
const modelMetrics = this.getOrCreateModelMetrics(event.model);
modelMetrics.api.totalRequests++;
modelMetrics.api.totalLatencyMs += event.duration_ms;
modelMetrics.tokens.prompt += event.input_token_count;
modelMetrics.tokens.candidates += event.output_token_count;
modelMetrics.tokens.total += event.total_token_count;
modelMetrics.tokens.cached += event.cached_content_token_count;
modelMetrics.tokens.thoughts += event.thoughts_token_count;
modelMetrics.tokens.tool += event.tool_token_count;
this.#lastPromptTokenCount = event.input_token_count;
}
private processApiError(event: ApiErrorEvent) {
const modelMetrics = this.getOrCreateModelMetrics(event.model);
modelMetrics.api.totalRequests++;
modelMetrics.api.totalErrors++;
modelMetrics.api.totalLatencyMs += event.duration_ms;
}
private processToolCall(event: ToolCallEvent) {
const { tools } = this.#metrics;
tools.totalCalls++;
tools.totalDurationMs += event.duration_ms;
if (event.success) {
tools.totalSuccess++;
} else {
tools.totalFail++;
}
if (!tools.byName[event.function_name]) {
tools.byName[event.function_name] = {
count: 0,
success: 0,
fail: 0,
durationMs: 0,
decisions: {
[ToolCallDecision.ACCEPT]: 0,
[ToolCallDecision.REJECT]: 0,
[ToolCallDecision.MODIFY]: 0,
},
};
}
const toolStats = tools.byName[event.function_name];
toolStats.count++;
toolStats.durationMs += event.duration_ms;
if (event.success) {
toolStats.success++;
} else {
toolStats.fail++;
}
if (event.decision) {
tools.totalDecisions[event.decision]++;
toolStats.decisions[event.decision]++;
}
}
}
export const uiTelemetryService = new UiTelemetryService();