refactor(telemetry): pass config object to telemetry functions
This commit refactors the telemetry system to pass a object to various logging and metrics functions. This change centralizes configuration management within the telemetry system, making it more modular and easier to maintain. The constructor and various tool execution functions have been updated to accept the object, which is then passed down to the telemetry functions. This eliminates the need to pass individual configuration values, such as , through multiple layers of the application.
This commit is contained in:
parent
9c5b5ff823
commit
d96af8bacd
|
@ -252,6 +252,8 @@ Use this method if you prefer not to use Docker.
|
||||||
|
|
||||||
## Data Reference: Logs & Metrics
|
## Data Reference: Logs & Metrics
|
||||||
|
|
||||||
|
A `sessionId` is included as a common attribute on all logs and metrics.
|
||||||
|
|
||||||
### Logs
|
### Logs
|
||||||
|
|
||||||
These are timestamped records of specific events.
|
These are timestamped records of specific events.
|
||||||
|
|
|
@ -134,6 +134,7 @@ describe('runNonInteractive', () => {
|
||||||
|
|
||||||
expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2);
|
expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2);
|
||||||
expect(mockCoreExecuteToolCall).toHaveBeenCalledWith(
|
expect(mockCoreExecuteToolCall).toHaveBeenCalledWith(
|
||||||
|
mockConfig,
|
||||||
expect.objectContaining({ callId: 'fc1', name: 'testTool' }),
|
expect.objectContaining({ callId: 'fc1', name: 'testTool' }),
|
||||||
mockToolRegistry,
|
mockToolRegistry,
|
||||||
expect.any(AbortSignal),
|
expect.any(AbortSignal),
|
||||||
|
|
|
@ -85,6 +85,7 @@ export async function runNonInteractive(
|
||||||
};
|
};
|
||||||
|
|
||||||
const toolResponse = await executeToolCall(
|
const toolResponse = await executeToolCall(
|
||||||
|
config,
|
||||||
requestInfo,
|
requestInfo,
|
||||||
toolRegistry,
|
toolRegistry,
|
||||||
abortController.signal,
|
abortController.signal,
|
||||||
|
|
|
@ -122,7 +122,7 @@ export function useReactToolScheduler(
|
||||||
}
|
}
|
||||||
duration = call.durationMs || 0;
|
duration = call.durationMs || 0;
|
||||||
|
|
||||||
logToolCall({
|
logToolCall(config, {
|
||||||
function_name: call.request.name,
|
function_name: call.request.name,
|
||||||
function_args: call.request.args,
|
function_args: call.request.args,
|
||||||
duration_ms: duration,
|
duration_ms: duration,
|
||||||
|
|
|
@ -160,9 +160,9 @@ export class GeminiClient {
|
||||||
const systemInstruction = getCoreSystemPrompt(userMemory);
|
const systemInstruction = getCoreSystemPrompt(userMemory);
|
||||||
|
|
||||||
return new GeminiChat(
|
return new GeminiChat(
|
||||||
|
this.config,
|
||||||
await this.contentGenerator,
|
await this.contentGenerator,
|
||||||
this.model,
|
this.model,
|
||||||
this.config.getSessionId(),
|
|
||||||
{
|
{
|
||||||
systemInstruction,
|
systemInstruction,
|
||||||
...this.generateContentConfig,
|
...this.generateContentConfig,
|
||||||
|
@ -214,7 +214,7 @@ export class GeminiClient {
|
||||||
}
|
}
|
||||||
|
|
||||||
private _logApiRequest(model: string, inputTokenCount: number): void {
|
private _logApiRequest(model: string, inputTokenCount: number): void {
|
||||||
logApiRequest({
|
logApiRequest(this.config, {
|
||||||
model,
|
model,
|
||||||
input_token_count: inputTokenCount,
|
input_token_count: inputTokenCount,
|
||||||
duration_ms: 0, // Duration is not known at request time
|
duration_ms: 0, // Duration is not known at request time
|
||||||
|
@ -239,7 +239,7 @@ export class GeminiClient {
|
||||||
responseError = `Finished with reason: ${finishReason}`;
|
responseError = `Finished with reason: ${finishReason}`;
|
||||||
}
|
}
|
||||||
|
|
||||||
logApiResponse({
|
logApiResponse(this.config, {
|
||||||
model,
|
model,
|
||||||
duration_ms: durationMs,
|
duration_ms: durationMs,
|
||||||
attempt,
|
attempt,
|
||||||
|
@ -277,7 +277,7 @@ export class GeminiClient {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
logApiError({
|
logApiError(this.config, {
|
||||||
model,
|
model,
|
||||||
error: errorMessage,
|
error: errorMessage,
|
||||||
status_code: statusCode,
|
status_code: statusCode,
|
||||||
|
|
|
@ -7,6 +7,7 @@
|
||||||
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
||||||
import { Content, Models, GenerateContentConfig, Part } from '@google/genai';
|
import { Content, Models, GenerateContentConfig, Part } from '@google/genai';
|
||||||
import { GeminiChat } from './geminiChat.js';
|
import { GeminiChat } from './geminiChat.js';
|
||||||
|
import { Config } from '../config/config.js';
|
||||||
|
|
||||||
// Mocks
|
// Mocks
|
||||||
const mockModelsModule = {
|
const mockModelsModule = {
|
||||||
|
@ -17,16 +18,19 @@ const mockModelsModule = {
|
||||||
batchEmbedContents: vi.fn(),
|
batchEmbedContents: vi.fn(),
|
||||||
} as unknown as Models;
|
} as unknown as Models;
|
||||||
|
|
||||||
|
const mockConfig = {
|
||||||
|
getSessionId: () => 'test-session-id',
|
||||||
|
} as unknown as Config;
|
||||||
|
|
||||||
describe('GeminiChat', () => {
|
describe('GeminiChat', () => {
|
||||||
let chat: GeminiChat;
|
let chat: GeminiChat;
|
||||||
const model = 'gemini-pro';
|
const model = 'gemini-pro';
|
||||||
const config: GenerateContentConfig = {};
|
const config: GenerateContentConfig = {};
|
||||||
const sessionId = 'test-session-id';
|
|
||||||
|
|
||||||
beforeEach(() => {
|
beforeEach(() => {
|
||||||
vi.clearAllMocks();
|
vi.clearAllMocks();
|
||||||
// Reset history for each test by creating a new instance
|
// Reset history for each test by creating a new instance
|
||||||
chat = new GeminiChat(mockModelsModule, model, sessionId, config, []);
|
chat = new GeminiChat(mockConfig, mockModelsModule, model, config, []);
|
||||||
});
|
});
|
||||||
|
|
||||||
afterEach(() => {
|
afterEach(() => {
|
||||||
|
@ -121,7 +125,7 @@ describe('GeminiChat', () => {
|
||||||
chat.recordHistory(userInput, newModelOutput); // userInput here is for the *next* turn, but history is already primed
|
chat.recordHistory(userInput, newModelOutput); // userInput here is for the *next* turn, but history is already primed
|
||||||
|
|
||||||
// Reset and set up a more realistic scenario for merging with existing history
|
// Reset and set up a more realistic scenario for merging with existing history
|
||||||
chat = new GeminiChat(mockModelsModule, model, sessionId, config, []);
|
chat = new GeminiChat(mockConfig, mockModelsModule, model, config, []);
|
||||||
const firstUserInput: Content = {
|
const firstUserInput: Content = {
|
||||||
role: 'user',
|
role: 'user',
|
||||||
parts: [{ text: 'First user input' }],
|
parts: [{ text: 'First user input' }],
|
||||||
|
@ -164,7 +168,7 @@ describe('GeminiChat', () => {
|
||||||
role: 'model',
|
role: 'model',
|
||||||
parts: [{ text: 'Initial model answer.' }],
|
parts: [{ text: 'Initial model answer.' }],
|
||||||
};
|
};
|
||||||
chat = new GeminiChat(mockModelsModule, model, sessionId, config, [
|
chat = new GeminiChat(mockConfig, mockModelsModule, model, config, [
|
||||||
initialUser,
|
initialUser,
|
||||||
initialModel,
|
initialModel,
|
||||||
]);
|
]);
|
||||||
|
|
|
@ -18,7 +18,7 @@ import {
|
||||||
import { retryWithBackoff } from '../utils/retry.js';
|
import { retryWithBackoff } from '../utils/retry.js';
|
||||||
import { isFunctionResponse } from '../utils/messageInspectors.js';
|
import { isFunctionResponse } from '../utils/messageInspectors.js';
|
||||||
import { ContentGenerator } from './contentGenerator.js';
|
import { ContentGenerator } from './contentGenerator.js';
|
||||||
import { Logger } from './logger.js';
|
import { Config } from '../config/config.js';
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Returns true if the response is valid, false otherwise.
|
* Returns true if the response is valid, false otherwise.
|
||||||
|
@ -118,17 +118,15 @@ export class GeminiChat {
|
||||||
// A promise to represent the current state of the message being sent to the
|
// A promise to represent the current state of the message being sent to the
|
||||||
// model.
|
// model.
|
||||||
private sendPromise: Promise<void> = Promise.resolve();
|
private sendPromise: Promise<void> = Promise.resolve();
|
||||||
private logger: Logger;
|
|
||||||
|
|
||||||
constructor(
|
constructor(
|
||||||
|
private readonly config: Config,
|
||||||
private readonly contentGenerator: ContentGenerator,
|
private readonly contentGenerator: ContentGenerator,
|
||||||
private readonly model: string,
|
private readonly model: string,
|
||||||
sessionId: string,
|
private readonly generationConfig: GenerateContentConfig = {},
|
||||||
private readonly config: GenerateContentConfig = {},
|
|
||||||
private history: Content[] = [],
|
private history: Content[] = [],
|
||||||
) {
|
) {
|
||||||
validateHistory(history);
|
validateHistory(history);
|
||||||
this.logger = new Logger(sessionId);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
|
@ -161,7 +159,7 @@ export class GeminiChat {
|
||||||
this.contentGenerator.generateContent({
|
this.contentGenerator.generateContent({
|
||||||
model: this.model,
|
model: this.model,
|
||||||
contents: this.getHistory(true).concat(userContent),
|
contents: this.getHistory(true).concat(userContent),
|
||||||
config: { ...this.config, ...params.config },
|
config: { ...this.generationConfig, ...params.config },
|
||||||
});
|
});
|
||||||
|
|
||||||
const responsePromise = retryWithBackoff(apiCall);
|
const responsePromise = retryWithBackoff(apiCall);
|
||||||
|
@ -230,7 +228,7 @@ export class GeminiChat {
|
||||||
this.contentGenerator.generateContentStream({
|
this.contentGenerator.generateContentStream({
|
||||||
model: this.model,
|
model: this.model,
|
||||||
contents: this.getHistory(true).concat(userContent),
|
contents: this.getHistory(true).concat(userContent),
|
||||||
config: { ...this.config, ...params.config },
|
config: { ...this.generationConfig, ...params.config },
|
||||||
});
|
});
|
||||||
|
|
||||||
// Note: Retrying streams can be complex. If generateContentStream itself doesn't handle retries
|
// Note: Retrying streams can be complex. If generateContentStream itself doesn't handle retries
|
||||||
|
|
|
@ -12,9 +12,12 @@ import {
|
||||||
ToolResult,
|
ToolResult,
|
||||||
Tool,
|
Tool,
|
||||||
ToolCallConfirmationDetails,
|
ToolCallConfirmationDetails,
|
||||||
|
Config,
|
||||||
} from '../index.js';
|
} from '../index.js';
|
||||||
import { Part, Type } from '@google/genai';
|
import { Part, Type } from '@google/genai';
|
||||||
|
|
||||||
|
const mockConfig = {} as unknown as Config;
|
||||||
|
|
||||||
describe('executeToolCall', () => {
|
describe('executeToolCall', () => {
|
||||||
let mockToolRegistry: ToolRegistry;
|
let mockToolRegistry: ToolRegistry;
|
||||||
let mockTool: Tool;
|
let mockTool: Tool;
|
||||||
|
@ -68,6 +71,7 @@ describe('executeToolCall', () => {
|
||||||
vi.mocked(mockTool.execute).mockResolvedValue(toolResult);
|
vi.mocked(mockTool.execute).mockResolvedValue(toolResult);
|
||||||
|
|
||||||
const response = await executeToolCall(
|
const response = await executeToolCall(
|
||||||
|
mockConfig,
|
||||||
request,
|
request,
|
||||||
mockToolRegistry,
|
mockToolRegistry,
|
||||||
abortController.signal,
|
abortController.signal,
|
||||||
|
@ -99,6 +103,7 @@ describe('executeToolCall', () => {
|
||||||
vi.mocked(mockToolRegistry.getTool).mockReturnValue(undefined);
|
vi.mocked(mockToolRegistry.getTool).mockReturnValue(undefined);
|
||||||
|
|
||||||
const response = await executeToolCall(
|
const response = await executeToolCall(
|
||||||
|
mockConfig,
|
||||||
request,
|
request,
|
||||||
mockToolRegistry,
|
mockToolRegistry,
|
||||||
abortController.signal,
|
abortController.signal,
|
||||||
|
@ -134,6 +139,7 @@ describe('executeToolCall', () => {
|
||||||
vi.mocked(mockTool.execute).mockRejectedValue(executionError);
|
vi.mocked(mockTool.execute).mockRejectedValue(executionError);
|
||||||
|
|
||||||
const response = await executeToolCall(
|
const response = await executeToolCall(
|
||||||
|
mockConfig,
|
||||||
request,
|
request,
|
||||||
mockToolRegistry,
|
mockToolRegistry,
|
||||||
abortController.signal,
|
abortController.signal,
|
||||||
|
@ -184,6 +190,7 @@ describe('executeToolCall', () => {
|
||||||
|
|
||||||
abortController.abort(); // Abort before calling
|
abortController.abort(); // Abort before calling
|
||||||
const response = await executeToolCall(
|
const response = await executeToolCall(
|
||||||
|
mockConfig,
|
||||||
request,
|
request,
|
||||||
mockToolRegistry,
|
mockToolRegistry,
|
||||||
abortController.signal,
|
abortController.signal,
|
||||||
|
@ -211,6 +218,7 @@ describe('executeToolCall', () => {
|
||||||
vi.mocked(mockTool.execute).mockResolvedValue(toolResult);
|
vi.mocked(mockTool.execute).mockResolvedValue(toolResult);
|
||||||
|
|
||||||
const response = await executeToolCall(
|
const response = await executeToolCall(
|
||||||
|
mockConfig,
|
||||||
request,
|
request,
|
||||||
mockToolRegistry,
|
mockToolRegistry,
|
||||||
abortController.signal,
|
abortController.signal,
|
||||||
|
|
|
@ -10,6 +10,7 @@ import {
|
||||||
ToolRegistry,
|
ToolRegistry,
|
||||||
ToolResult,
|
ToolResult,
|
||||||
} from '../index.js';
|
} from '../index.js';
|
||||||
|
import { Config } from '../config/config.js';
|
||||||
import { convertToFunctionResponse } from './coreToolScheduler.js';
|
import { convertToFunctionResponse } from './coreToolScheduler.js';
|
||||||
|
|
||||||
/**
|
/**
|
||||||
|
@ -17,6 +18,7 @@ import { convertToFunctionResponse } from './coreToolScheduler.js';
|
||||||
* It does not handle confirmations, multiple calls, or live updates.
|
* It does not handle confirmations, multiple calls, or live updates.
|
||||||
*/
|
*/
|
||||||
export async function executeToolCall(
|
export async function executeToolCall(
|
||||||
|
config: Config,
|
||||||
toolCallRequest: ToolCallRequestInfo,
|
toolCallRequest: ToolCallRequestInfo,
|
||||||
toolRegistry: ToolRegistry,
|
toolRegistry: ToolRegistry,
|
||||||
abortSignal?: AbortSignal,
|
abortSignal?: AbortSignal,
|
||||||
|
|
|
@ -34,10 +34,17 @@ import { isTelemetrySdkInitialized } from './sdk.js';
|
||||||
const shouldLogUserPrompts = (config: Config): boolean =>
|
const shouldLogUserPrompts = (config: Config): boolean =>
|
||||||
config.getTelemetryLogUserPromptsEnabled() ?? false;
|
config.getTelemetryLogUserPromptsEnabled() ?? false;
|
||||||
|
|
||||||
|
function getCommonAttributes(config: Config): LogAttributes {
|
||||||
|
return {
|
||||||
|
'session.id': config.getSessionId(),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
export function logCliConfiguration(config: Config): void {
|
export function logCliConfiguration(config: Config): void {
|
||||||
if (!isTelemetrySdkInitialized()) return;
|
if (!isTelemetrySdkInitialized()) return;
|
||||||
|
|
||||||
const attributes: LogAttributes = {
|
const attributes: LogAttributes = {
|
||||||
|
...getCommonAttributes(config),
|
||||||
'event.name': EVENT_CLI_CONFIG,
|
'event.name': EVENT_CLI_CONFIG,
|
||||||
'event.timestamp': new Date().toISOString(),
|
'event.timestamp': new Date().toISOString(),
|
||||||
model: config.getModel(),
|
model: config.getModel(),
|
||||||
|
@ -69,6 +76,7 @@ export function logUserPrompt(
|
||||||
if (!isTelemetrySdkInitialized()) return;
|
if (!isTelemetrySdkInitialized()) return;
|
||||||
const { prompt, ...restOfEventArgs } = event;
|
const { prompt, ...restOfEventArgs } = event;
|
||||||
const attributes: LogAttributes = {
|
const attributes: LogAttributes = {
|
||||||
|
...getCommonAttributes(config),
|
||||||
...restOfEventArgs,
|
...restOfEventArgs,
|
||||||
'event.name': EVENT_USER_PROMPT,
|
'event.name': EVENT_USER_PROMPT,
|
||||||
'event.timestamp': new Date().toISOString(),
|
'event.timestamp': new Date().toISOString(),
|
||||||
|
@ -85,10 +93,12 @@ export function logUserPrompt(
|
||||||
}
|
}
|
||||||
|
|
||||||
export function logToolCall(
|
export function logToolCall(
|
||||||
|
config: Config,
|
||||||
event: Omit<ToolCallEvent, 'event.name' | 'event.timestamp'>,
|
event: Omit<ToolCallEvent, 'event.name' | 'event.timestamp'>,
|
||||||
): void {
|
): void {
|
||||||
if (!isTelemetrySdkInitialized()) return;
|
if (!isTelemetrySdkInitialized()) return;
|
||||||
const attributes: LogAttributes = {
|
const attributes: LogAttributes = {
|
||||||
|
...getCommonAttributes(config),
|
||||||
...event,
|
...event,
|
||||||
'event.name': EVENT_TOOL_CALL,
|
'event.name': EVENT_TOOL_CALL,
|
||||||
'event.timestamp': new Date().toISOString(),
|
'event.timestamp': new Date().toISOString(),
|
||||||
|
@ -106,14 +116,21 @@ export function logToolCall(
|
||||||
attributes,
|
attributes,
|
||||||
};
|
};
|
||||||
logger.emit(logRecord);
|
logger.emit(logRecord);
|
||||||
recordToolCallMetrics(event.function_name, event.duration_ms, event.success);
|
recordToolCallMetrics(
|
||||||
|
config,
|
||||||
|
event.function_name,
|
||||||
|
event.duration_ms,
|
||||||
|
event.success,
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
export function logApiRequest(
|
export function logApiRequest(
|
||||||
|
config: Config,
|
||||||
event: Omit<ApiRequestEvent, 'event.name' | 'event.timestamp'>,
|
event: Omit<ApiRequestEvent, 'event.name' | 'event.timestamp'>,
|
||||||
): void {
|
): void {
|
||||||
if (!isTelemetrySdkInitialized()) return;
|
if (!isTelemetrySdkInitialized()) return;
|
||||||
const attributes: LogAttributes = {
|
const attributes: LogAttributes = {
|
||||||
|
...getCommonAttributes(config),
|
||||||
...event,
|
...event,
|
||||||
'event.name': EVENT_API_REQUEST,
|
'event.name': EVENT_API_REQUEST,
|
||||||
'event.timestamp': new Date().toISOString(),
|
'event.timestamp': new Date().toISOString(),
|
||||||
|
@ -124,14 +141,21 @@ export function logApiRequest(
|
||||||
attributes,
|
attributes,
|
||||||
};
|
};
|
||||||
logger.emit(logRecord);
|
logger.emit(logRecord);
|
||||||
recordTokenUsageMetrics(event.model, event.input_token_count, 'input');
|
recordTokenUsageMetrics(
|
||||||
|
config,
|
||||||
|
event.model,
|
||||||
|
event.input_token_count,
|
||||||
|
'input',
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
export function logApiError(
|
export function logApiError(
|
||||||
|
config: Config,
|
||||||
event: Omit<ApiErrorEvent, 'event.name' | 'event.timestamp'>,
|
event: Omit<ApiErrorEvent, 'event.name' | 'event.timestamp'>,
|
||||||
): void {
|
): void {
|
||||||
if (!isTelemetrySdkInitialized()) return;
|
if (!isTelemetrySdkInitialized()) return;
|
||||||
const attributes: LogAttributes = {
|
const attributes: LogAttributes = {
|
||||||
|
...getCommonAttributes(config),
|
||||||
...event,
|
...event,
|
||||||
'event.name': EVENT_API_ERROR,
|
'event.name': EVENT_API_ERROR,
|
||||||
'event.timestamp': new Date().toISOString(),
|
'event.timestamp': new Date().toISOString(),
|
||||||
|
@ -152,6 +176,7 @@ export function logApiError(
|
||||||
};
|
};
|
||||||
logger.emit(logRecord);
|
logger.emit(logRecord);
|
||||||
recordApiErrorMetrics(
|
recordApiErrorMetrics(
|
||||||
|
config,
|
||||||
event.model,
|
event.model,
|
||||||
event.duration_ms,
|
event.duration_ms,
|
||||||
event.status_code,
|
event.status_code,
|
||||||
|
@ -160,10 +185,12 @@ export function logApiError(
|
||||||
}
|
}
|
||||||
|
|
||||||
export function logApiResponse(
|
export function logApiResponse(
|
||||||
|
config: Config,
|
||||||
event: Omit<ApiResponseEvent, 'event.name' | 'event.timestamp'>,
|
event: Omit<ApiResponseEvent, 'event.name' | 'event.timestamp'>,
|
||||||
): void {
|
): void {
|
||||||
if (!isTelemetrySdkInitialized()) return;
|
if (!isTelemetrySdkInitialized()) return;
|
||||||
const attributes: LogAttributes = {
|
const attributes: LogAttributes = {
|
||||||
|
...getCommonAttributes(config),
|
||||||
...event,
|
...event,
|
||||||
'event.name': EVENT_API_RESPONSE,
|
'event.name': EVENT_API_RESPONSE,
|
||||||
'event.timestamp': new Date().toISOString(),
|
'event.timestamp': new Date().toISOString(),
|
||||||
|
@ -183,17 +210,29 @@ export function logApiResponse(
|
||||||
};
|
};
|
||||||
logger.emit(logRecord);
|
logger.emit(logRecord);
|
||||||
recordApiResponseMetrics(
|
recordApiResponseMetrics(
|
||||||
|
config,
|
||||||
event.model,
|
event.model,
|
||||||
event.duration_ms,
|
event.duration_ms,
|
||||||
event.status_code,
|
event.status_code,
|
||||||
event.error,
|
event.error,
|
||||||
);
|
);
|
||||||
recordTokenUsageMetrics(event.model, event.output_token_count, 'output');
|
|
||||||
recordTokenUsageMetrics(
|
recordTokenUsageMetrics(
|
||||||
|
config,
|
||||||
|
event.model,
|
||||||
|
event.output_token_count,
|
||||||
|
'output',
|
||||||
|
);
|
||||||
|
recordTokenUsageMetrics(
|
||||||
|
config,
|
||||||
event.model,
|
event.model,
|
||||||
event.cached_content_token_count,
|
event.cached_content_token_count,
|
||||||
'cache',
|
'cache',
|
||||||
);
|
);
|
||||||
recordTokenUsageMetrics(event.model, event.thoughts_token_count, 'thought');
|
recordTokenUsageMetrics(
|
||||||
recordTokenUsageMetrics(event.model, event.tool_token_count, 'tool');
|
config,
|
||||||
|
event.model,
|
||||||
|
event.thoughts_token_count,
|
||||||
|
'thought',
|
||||||
|
);
|
||||||
|
recordTokenUsageMetrics(config, event.model, event.tool_token_count, 'tool');
|
||||||
}
|
}
|
||||||
|
|
|
@ -7,6 +7,7 @@
|
||||||
import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
|
import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
|
||||||
import { Counter, Meter, metrics } from '@opentelemetry/api';
|
import { Counter, Meter, metrics } from '@opentelemetry/api';
|
||||||
import { initializeMetrics, recordTokenUsageMetrics } from './metrics.js';
|
import { initializeMetrics, recordTokenUsageMetrics } from './metrics.js';
|
||||||
|
import { Config } from '../config/config.js';
|
||||||
|
|
||||||
const mockCounter = {
|
const mockCounter = {
|
||||||
add: vi.fn(),
|
add: vi.fn(),
|
||||||
|
@ -33,51 +34,61 @@ describe('Telemetry Metrics', () => {
|
||||||
});
|
});
|
||||||
|
|
||||||
describe('recordTokenUsageMetrics', () => {
|
describe('recordTokenUsageMetrics', () => {
|
||||||
|
const mockConfig = {
|
||||||
|
getSessionId: () => 'test-session-id',
|
||||||
|
} as unknown as Config;
|
||||||
|
|
||||||
it('should not record metrics if not initialized', () => {
|
it('should not record metrics if not initialized', () => {
|
||||||
recordTokenUsageMetrics('gemini-pro', 100, 'input');
|
recordTokenUsageMetrics(mockConfig, 'gemini-pro', 100, 'input');
|
||||||
expect(mockCounter.add).not.toHaveBeenCalled();
|
expect(mockCounter.add).not.toHaveBeenCalled();
|
||||||
});
|
});
|
||||||
|
|
||||||
it('should record token usage with the correct attributes', () => {
|
it('should record token usage with the correct attributes', () => {
|
||||||
initializeMetrics();
|
initializeMetrics(mockConfig);
|
||||||
recordTokenUsageMetrics('gemini-pro', 100, 'input');
|
recordTokenUsageMetrics(mockConfig, 'gemini-pro', 100, 'input');
|
||||||
expect(mockCounter.add).toHaveBeenCalledWith(100, {
|
expect(mockCounter.add).toHaveBeenCalledWith(100, {
|
||||||
|
'session.id': 'test-session-id',
|
||||||
model: 'gemini-pro',
|
model: 'gemini-pro',
|
||||||
type: 'input',
|
type: 'input',
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
it('should record token usage for different types', () => {
|
it('should record token usage for different types', () => {
|
||||||
initializeMetrics();
|
initializeMetrics(mockConfig);
|
||||||
recordTokenUsageMetrics('gemini-pro', 50, 'output');
|
recordTokenUsageMetrics(mockConfig, 'gemini-pro', 50, 'output');
|
||||||
expect(mockCounter.add).toHaveBeenCalledWith(50, {
|
expect(mockCounter.add).toHaveBeenCalledWith(50, {
|
||||||
|
'session.id': 'test-session-id',
|
||||||
model: 'gemini-pro',
|
model: 'gemini-pro',
|
||||||
type: 'output',
|
type: 'output',
|
||||||
});
|
});
|
||||||
|
|
||||||
recordTokenUsageMetrics('gemini-pro', 25, 'thought');
|
recordTokenUsageMetrics(mockConfig, 'gemini-pro', 25, 'thought');
|
||||||
expect(mockCounter.add).toHaveBeenCalledWith(25, {
|
expect(mockCounter.add).toHaveBeenCalledWith(25, {
|
||||||
|
'session.id': 'test-session-id',
|
||||||
model: 'gemini-pro',
|
model: 'gemini-pro',
|
||||||
type: 'thought',
|
type: 'thought',
|
||||||
});
|
});
|
||||||
|
|
||||||
recordTokenUsageMetrics('gemini-pro', 75, 'cache');
|
recordTokenUsageMetrics(mockConfig, 'gemini-pro', 75, 'cache');
|
||||||
expect(mockCounter.add).toHaveBeenCalledWith(75, {
|
expect(mockCounter.add).toHaveBeenCalledWith(75, {
|
||||||
|
'session.id': 'test-session-id',
|
||||||
model: 'gemini-pro',
|
model: 'gemini-pro',
|
||||||
type: 'cache',
|
type: 'cache',
|
||||||
});
|
});
|
||||||
|
|
||||||
recordTokenUsageMetrics('gemini-pro', 125, 'tool');
|
recordTokenUsageMetrics(mockConfig, 'gemini-pro', 125, 'tool');
|
||||||
expect(mockCounter.add).toHaveBeenCalledWith(125, {
|
expect(mockCounter.add).toHaveBeenCalledWith(125, {
|
||||||
|
'session.id': 'test-session-id',
|
||||||
model: 'gemini-pro',
|
model: 'gemini-pro',
|
||||||
type: 'tool',
|
type: 'tool',
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
it('should handle different models', () => {
|
it('should handle different models', () => {
|
||||||
initializeMetrics();
|
initializeMetrics(mockConfig);
|
||||||
recordTokenUsageMetrics('gemini-ultra', 200, 'input');
|
recordTokenUsageMetrics(mockConfig, 'gemini-ultra', 200, 'input');
|
||||||
expect(mockCounter.add).toHaveBeenCalledWith(200, {
|
expect(mockCounter.add).toHaveBeenCalledWith(200, {
|
||||||
|
'session.id': 'test-session-id',
|
||||||
model: 'gemini-ultra',
|
model: 'gemini-ultra',
|
||||||
type: 'input',
|
type: 'input',
|
||||||
});
|
});
|
||||||
|
|
|
@ -21,6 +21,7 @@ import {
|
||||||
METRIC_TOKEN_USAGE,
|
METRIC_TOKEN_USAGE,
|
||||||
METRIC_SESSION_COUNT,
|
METRIC_SESSION_COUNT,
|
||||||
} from './constants.js';
|
} from './constants.js';
|
||||||
|
import { Config } from '../config/config.js';
|
||||||
|
|
||||||
let cliMeter: Meter | undefined;
|
let cliMeter: Meter | undefined;
|
||||||
let toolCallCounter: Counter | undefined;
|
let toolCallCounter: Counter | undefined;
|
||||||
|
@ -30,6 +31,12 @@ let apiRequestLatencyHistogram: Histogram | undefined;
|
||||||
let tokenUsageCounter: Counter | undefined;
|
let tokenUsageCounter: Counter | undefined;
|
||||||
let isMetricsInitialized = false;
|
let isMetricsInitialized = false;
|
||||||
|
|
||||||
|
function getCommonAttributes(config: Config): Attributes {
|
||||||
|
return {
|
||||||
|
'session.id': config.getSessionId(),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
export function getMeter(): Meter | undefined {
|
export function getMeter(): Meter | undefined {
|
||||||
if (!cliMeter) {
|
if (!cliMeter) {
|
||||||
cliMeter = metrics.getMeter(SERVICE_NAME);
|
cliMeter = metrics.getMeter(SERVICE_NAME);
|
||||||
|
@ -37,7 +44,7 @@ export function getMeter(): Meter | undefined {
|
||||||
return cliMeter;
|
return cliMeter;
|
||||||
}
|
}
|
||||||
|
|
||||||
export function initializeMetrics(): void {
|
export function initializeMetrics(config: Config): void {
|
||||||
if (isMetricsInitialized) return;
|
if (isMetricsInitialized) return;
|
||||||
|
|
||||||
const meter = getMeter();
|
const meter = getMeter();
|
||||||
|
@ -73,11 +80,12 @@ export function initializeMetrics(): void {
|
||||||
description: 'Count of CLI sessions started.',
|
description: 'Count of CLI sessions started.',
|
||||||
valueType: ValueType.INT,
|
valueType: ValueType.INT,
|
||||||
});
|
});
|
||||||
sessionCounter.add(1);
|
sessionCounter.add(1, getCommonAttributes(config));
|
||||||
isMetricsInitialized = true;
|
isMetricsInitialized = true;
|
||||||
}
|
}
|
||||||
|
|
||||||
export function recordToolCallMetrics(
|
export function recordToolCallMetrics(
|
||||||
|
config: Config,
|
||||||
functionName: string,
|
functionName: string,
|
||||||
durationMs: number,
|
durationMs: number,
|
||||||
success: boolean,
|
success: boolean,
|
||||||
|
@ -86,25 +94,33 @@ export function recordToolCallMetrics(
|
||||||
return;
|
return;
|
||||||
|
|
||||||
const metricAttributes: Attributes = {
|
const metricAttributes: Attributes = {
|
||||||
|
...getCommonAttributes(config),
|
||||||
function_name: functionName,
|
function_name: functionName,
|
||||||
success,
|
success,
|
||||||
};
|
};
|
||||||
toolCallCounter.add(1, metricAttributes);
|
toolCallCounter.add(1, metricAttributes);
|
||||||
toolCallLatencyHistogram.record(durationMs, {
|
toolCallLatencyHistogram.record(durationMs, {
|
||||||
|
...getCommonAttributes(config),
|
||||||
function_name: functionName,
|
function_name: functionName,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
export function recordTokenUsageMetrics(
|
export function recordTokenUsageMetrics(
|
||||||
|
config: Config,
|
||||||
model: string,
|
model: string,
|
||||||
tokenCount: number,
|
tokenCount: number,
|
||||||
type: 'input' | 'output' | 'thought' | 'cache' | 'tool',
|
type: 'input' | 'output' | 'thought' | 'cache' | 'tool',
|
||||||
): void {
|
): void {
|
||||||
if (!tokenUsageCounter || !isMetricsInitialized) return;
|
if (!tokenUsageCounter || !isMetricsInitialized) return;
|
||||||
tokenUsageCounter.add(tokenCount, { model, type });
|
tokenUsageCounter.add(tokenCount, {
|
||||||
|
...getCommonAttributes(config),
|
||||||
|
model,
|
||||||
|
type,
|
||||||
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
export function recordApiResponseMetrics(
|
export function recordApiResponseMetrics(
|
||||||
|
config: Config,
|
||||||
model: string,
|
model: string,
|
||||||
durationMs: number,
|
durationMs: number,
|
||||||
statusCode?: number | string,
|
statusCode?: number | string,
|
||||||
|
@ -117,14 +133,19 @@ export function recordApiResponseMetrics(
|
||||||
)
|
)
|
||||||
return;
|
return;
|
||||||
const metricAttributes: Attributes = {
|
const metricAttributes: Attributes = {
|
||||||
|
...getCommonAttributes(config),
|
||||||
model,
|
model,
|
||||||
status_code: statusCode ?? (error ? 'error' : 'ok'),
|
status_code: statusCode ?? (error ? 'error' : 'ok'),
|
||||||
};
|
};
|
||||||
apiRequestCounter.add(1, metricAttributes);
|
apiRequestCounter.add(1, metricAttributes);
|
||||||
apiRequestLatencyHistogram.record(durationMs, { model });
|
apiRequestLatencyHistogram.record(durationMs, {
|
||||||
|
...getCommonAttributes(config),
|
||||||
|
model,
|
||||||
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
export function recordApiErrorMetrics(
|
export function recordApiErrorMetrics(
|
||||||
|
config: Config,
|
||||||
model: string,
|
model: string,
|
||||||
durationMs: number,
|
durationMs: number,
|
||||||
statusCode?: number | string,
|
statusCode?: number | string,
|
||||||
|
@ -137,10 +158,14 @@ export function recordApiErrorMetrics(
|
||||||
)
|
)
|
||||||
return;
|
return;
|
||||||
const metricAttributes: Attributes = {
|
const metricAttributes: Attributes = {
|
||||||
|
...getCommonAttributes(config),
|
||||||
model,
|
model,
|
||||||
status_code: statusCode ?? 'error',
|
status_code: statusCode ?? 'error',
|
||||||
error_type: errorType ?? 'unknown',
|
error_type: errorType ?? 'unknown',
|
||||||
};
|
};
|
||||||
apiRequestCounter.add(1, metricAttributes);
|
apiRequestCounter.add(1, metricAttributes);
|
||||||
apiRequestLatencyHistogram.record(durationMs, { model });
|
apiRequestLatencyHistogram.record(durationMs, {
|
||||||
|
...getCommonAttributes(config),
|
||||||
|
model,
|
||||||
|
});
|
||||||
}
|
}
|
||||||
|
|
|
@ -112,7 +112,7 @@ export function initializeTelemetry(config: Config): void {
|
||||||
sdk.start();
|
sdk.start();
|
||||||
console.log('OpenTelemetry SDK started successfully.');
|
console.log('OpenTelemetry SDK started successfully.');
|
||||||
telemetryInitialized = true;
|
telemetryInitialized = true;
|
||||||
initializeMetrics();
|
initializeMetrics(config);
|
||||||
logCliConfiguration(config);
|
logCliConfiguration(config);
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
console.error('Error starting OpenTelemetry SDK:', error);
|
console.error('Error starting OpenTelemetry SDK:', error);
|
||||||
|
|
|
@ -69,9 +69,9 @@ describe('checkNextSpeaker', () => {
|
||||||
|
|
||||||
// GeminiChat will receive the mocked instances via the mocked GoogleGenAI constructor
|
// GeminiChat will receive the mocked instances via the mocked GoogleGenAI constructor
|
||||||
chatInstance = new GeminiChat(
|
chatInstance = new GeminiChat(
|
||||||
|
mockConfigInstance,
|
||||||
mockModelsInstance, // This is the instance returned by mockGoogleGenAIInstance.getGenerativeModel
|
mockModelsInstance, // This is the instance returned by mockGoogleGenAIInstance.getGenerativeModel
|
||||||
'gemini-pro', // model name
|
'gemini-pro', // model name
|
||||||
'test-session-id',
|
|
||||||
{},
|
{},
|
||||||
[], // initial history
|
[], // initial history
|
||||||
);
|
);
|
||||||
|
|
Loading…
Reference in New Issue