mirror of
https://github.com/QwenLM/qwen-code.git
synced 2025-12-19 09:33:53 +00:00
feat: Change /stats to include more detailed breakdowns (#2615)
This commit is contained in:
@@ -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 {
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -38,3 +38,4 @@ export {
|
||||
} from './types.js';
|
||||
export { SpanStatusCode, ValueType } from '@opentelemetry/api';
|
||||
export { SemanticAttributes } from '@opentelemetry/semantic-conventions';
|
||||
export * from './uiTelemetry.js';
|
||||
|
||||
@@ -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',
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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 = {
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
510
packages/core/src/telemetry/uiTelemetry.test.ts
Normal file
510
packages/core/src/telemetry/uiTelemetry.test.ts
Normal 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);
|
||||
});
|
||||
});
|
||||
});
|
||||
207
packages/core/src/telemetry/uiTelemetry.ts
Normal file
207
packages/core/src/telemetry/uiTelemetry.ts
Normal 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();
|
||||
Reference in New Issue
Block a user