From d96af8bacd84d065310dddb49abb7561c4a7b059 Mon Sep 17 00:00:00 2001 From: jerop Date: Wed, 11 Jun 2025 16:50:24 +0000 Subject: [PATCH] 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. --- docs/core/telemetry.md | 2 + packages/cli/src/nonInteractiveCli.test.ts | 1 + packages/cli/src/nonInteractiveCli.ts | 1 + .../cli/src/ui/hooks/useReactToolScheduler.ts | 2 +- packages/core/src/core/client.ts | 8 +-- packages/core/src/core/geminiChat.test.ts | 12 +++-- packages/core/src/core/geminiChat.ts | 12 ++--- .../core/nonInteractiveToolExecutor.test.ts | 8 +++ .../src/core/nonInteractiveToolExecutor.ts | 2 + packages/core/src/telemetry/loggers.ts | 49 +++++++++++++++++-- packages/core/src/telemetry/metrics.test.ts | 31 ++++++++---- packages/core/src/telemetry/metrics.ts | 35 +++++++++++-- packages/core/src/telemetry/sdk.ts | 2 +- .../core/src/utils/nextSpeakerChecker.test.ts | 2 +- 14 files changed, 129 insertions(+), 38 deletions(-) diff --git a/docs/core/telemetry.md b/docs/core/telemetry.md index bd448776..aec63e80 100644 --- a/docs/core/telemetry.md +++ b/docs/core/telemetry.md @@ -252,6 +252,8 @@ Use this method if you prefer not to use Docker. ## Data Reference: Logs & Metrics +A `sessionId` is included as a common attribute on all logs and metrics. + ### Logs These are timestamped records of specific events. diff --git a/packages/cli/src/nonInteractiveCli.test.ts b/packages/cli/src/nonInteractiveCli.test.ts index 4e7fd727..959cf03f 100644 --- a/packages/cli/src/nonInteractiveCli.test.ts +++ b/packages/cli/src/nonInteractiveCli.test.ts @@ -134,6 +134,7 @@ describe('runNonInteractive', () => { expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2); expect(mockCoreExecuteToolCall).toHaveBeenCalledWith( + mockConfig, expect.objectContaining({ callId: 'fc1', name: 'testTool' }), mockToolRegistry, expect.any(AbortSignal), diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts index 6b828dc8..069a7be8 100644 --- a/packages/cli/src/nonInteractiveCli.ts +++ b/packages/cli/src/nonInteractiveCli.ts @@ -85,6 +85,7 @@ export async function runNonInteractive( }; const toolResponse = await executeToolCall( + config, requestInfo, toolRegistry, abortController.signal, diff --git a/packages/cli/src/ui/hooks/useReactToolScheduler.ts b/packages/cli/src/ui/hooks/useReactToolScheduler.ts index 8dbcfb87..4e55cba4 100644 --- a/packages/cli/src/ui/hooks/useReactToolScheduler.ts +++ b/packages/cli/src/ui/hooks/useReactToolScheduler.ts @@ -122,7 +122,7 @@ export function useReactToolScheduler( } duration = call.durationMs || 0; - logToolCall({ + logToolCall(config, { function_name: call.request.name, function_args: call.request.args, duration_ms: duration, diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 349caf1f..a2deca4e 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -160,9 +160,9 @@ export class GeminiClient { const systemInstruction = getCoreSystemPrompt(userMemory); return new GeminiChat( + this.config, await this.contentGenerator, this.model, - this.config.getSessionId(), { systemInstruction, ...this.generateContentConfig, @@ -214,7 +214,7 @@ export class GeminiClient { } private _logApiRequest(model: string, inputTokenCount: number): void { - logApiRequest({ + logApiRequest(this.config, { model, input_token_count: inputTokenCount, duration_ms: 0, // Duration is not known at request time @@ -239,7 +239,7 @@ export class GeminiClient { responseError = `Finished with reason: ${finishReason}`; } - logApiResponse({ + logApiResponse(this.config, { model, duration_ms: durationMs, attempt, @@ -277,7 +277,7 @@ export class GeminiClient { } } - logApiError({ + logApiError(this.config, { model, error: errorMessage, status_code: statusCode, diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 49341939..03e933cc 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -7,6 +7,7 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; import { Content, Models, GenerateContentConfig, Part } from '@google/genai'; import { GeminiChat } from './geminiChat.js'; +import { Config } from '../config/config.js'; // Mocks const mockModelsModule = { @@ -17,16 +18,19 @@ const mockModelsModule = { batchEmbedContents: vi.fn(), } as unknown as Models; +const mockConfig = { + getSessionId: () => 'test-session-id', +} as unknown as Config; + describe('GeminiChat', () => { let chat: GeminiChat; const model = 'gemini-pro'; const config: GenerateContentConfig = {}; - const sessionId = 'test-session-id'; beforeEach(() => { vi.clearAllMocks(); // 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(() => { @@ -121,7 +125,7 @@ describe('GeminiChat', () => { 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 - chat = new GeminiChat(mockModelsModule, model, sessionId, config, []); + chat = new GeminiChat(mockConfig, mockModelsModule, model, config, []); const firstUserInput: Content = { role: 'user', parts: [{ text: 'First user input' }], @@ -164,7 +168,7 @@ describe('GeminiChat', () => { role: 'model', parts: [{ text: 'Initial model answer.' }], }; - chat = new GeminiChat(mockModelsModule, model, sessionId, config, [ + chat = new GeminiChat(mockConfig, mockModelsModule, model, config, [ initialUser, initialModel, ]); diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index a78a7884..2a81aca8 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -18,7 +18,7 @@ import { import { retryWithBackoff } from '../utils/retry.js'; import { isFunctionResponse } from '../utils/messageInspectors.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. @@ -118,17 +118,15 @@ export class GeminiChat { // A promise to represent the current state of the message being sent to the // model. private sendPromise: Promise = Promise.resolve(); - private logger: Logger; constructor( + private readonly config: Config, private readonly contentGenerator: ContentGenerator, private readonly model: string, - sessionId: string, - private readonly config: GenerateContentConfig = {}, + private readonly generationConfig: GenerateContentConfig = {}, private history: Content[] = [], ) { validateHistory(history); - this.logger = new Logger(sessionId); } /** @@ -161,7 +159,7 @@ export class GeminiChat { this.contentGenerator.generateContent({ model: this.model, contents: this.getHistory(true).concat(userContent), - config: { ...this.config, ...params.config }, + config: { ...this.generationConfig, ...params.config }, }); const responsePromise = retryWithBackoff(apiCall); @@ -230,7 +228,7 @@ export class GeminiChat { this.contentGenerator.generateContentStream({ model: this.model, 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 diff --git a/packages/core/src/core/nonInteractiveToolExecutor.test.ts b/packages/core/src/core/nonInteractiveToolExecutor.test.ts index be9294a8..07946af5 100644 --- a/packages/core/src/core/nonInteractiveToolExecutor.test.ts +++ b/packages/core/src/core/nonInteractiveToolExecutor.test.ts @@ -12,9 +12,12 @@ import { ToolResult, Tool, ToolCallConfirmationDetails, + Config, } from '../index.js'; import { Part, Type } from '@google/genai'; +const mockConfig = {} as unknown as Config; + describe('executeToolCall', () => { let mockToolRegistry: ToolRegistry; let mockTool: Tool; @@ -68,6 +71,7 @@ describe('executeToolCall', () => { vi.mocked(mockTool.execute).mockResolvedValue(toolResult); const response = await executeToolCall( + mockConfig, request, mockToolRegistry, abortController.signal, @@ -99,6 +103,7 @@ describe('executeToolCall', () => { vi.mocked(mockToolRegistry.getTool).mockReturnValue(undefined); const response = await executeToolCall( + mockConfig, request, mockToolRegistry, abortController.signal, @@ -134,6 +139,7 @@ describe('executeToolCall', () => { vi.mocked(mockTool.execute).mockRejectedValue(executionError); const response = await executeToolCall( + mockConfig, request, mockToolRegistry, abortController.signal, @@ -184,6 +190,7 @@ describe('executeToolCall', () => { abortController.abort(); // Abort before calling const response = await executeToolCall( + mockConfig, request, mockToolRegistry, abortController.signal, @@ -211,6 +218,7 @@ describe('executeToolCall', () => { vi.mocked(mockTool.execute).mockResolvedValue(toolResult); const response = await executeToolCall( + mockConfig, request, mockToolRegistry, abortController.signal, diff --git a/packages/core/src/core/nonInteractiveToolExecutor.ts b/packages/core/src/core/nonInteractiveToolExecutor.ts index 87d17c2b..5f7ed4ff 100644 --- a/packages/core/src/core/nonInteractiveToolExecutor.ts +++ b/packages/core/src/core/nonInteractiveToolExecutor.ts @@ -10,6 +10,7 @@ import { ToolRegistry, ToolResult, } from '../index.js'; +import { Config } from '../config/config.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. */ export async function executeToolCall( + config: Config, toolCallRequest: ToolCallRequestInfo, toolRegistry: ToolRegistry, abortSignal?: AbortSignal, diff --git a/packages/core/src/telemetry/loggers.ts b/packages/core/src/telemetry/loggers.ts index 66be0fca..d2b01f65 100644 --- a/packages/core/src/telemetry/loggers.ts +++ b/packages/core/src/telemetry/loggers.ts @@ -34,10 +34,17 @@ import { isTelemetrySdkInitialized } from './sdk.js'; const shouldLogUserPrompts = (config: Config): boolean => config.getTelemetryLogUserPromptsEnabled() ?? false; +function getCommonAttributes(config: Config): LogAttributes { + return { + 'session.id': config.getSessionId(), + }; +} + export function logCliConfiguration(config: Config): void { if (!isTelemetrySdkInitialized()) return; const attributes: LogAttributes = { + ...getCommonAttributes(config), 'event.name': EVENT_CLI_CONFIG, 'event.timestamp': new Date().toISOString(), model: config.getModel(), @@ -69,6 +76,7 @@ export function logUserPrompt( if (!isTelemetrySdkInitialized()) return; const { prompt, ...restOfEventArgs } = event; const attributes: LogAttributes = { + ...getCommonAttributes(config), ...restOfEventArgs, 'event.name': EVENT_USER_PROMPT, 'event.timestamp': new Date().toISOString(), @@ -85,10 +93,12 @@ export function logUserPrompt( } export function logToolCall( + config: Config, event: Omit, ): void { if (!isTelemetrySdkInitialized()) return; const attributes: LogAttributes = { + ...getCommonAttributes(config), ...event, 'event.name': EVENT_TOOL_CALL, 'event.timestamp': new Date().toISOString(), @@ -106,14 +116,21 @@ export function logToolCall( attributes, }; 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( + config: Config, event: Omit, ): void { if (!isTelemetrySdkInitialized()) return; const attributes: LogAttributes = { + ...getCommonAttributes(config), ...event, 'event.name': EVENT_API_REQUEST, 'event.timestamp': new Date().toISOString(), @@ -124,14 +141,21 @@ export function logApiRequest( attributes, }; logger.emit(logRecord); - recordTokenUsageMetrics(event.model, event.input_token_count, 'input'); + recordTokenUsageMetrics( + config, + event.model, + event.input_token_count, + 'input', + ); } export function logApiError( + config: Config, event: Omit, ): void { if (!isTelemetrySdkInitialized()) return; const attributes: LogAttributes = { + ...getCommonAttributes(config), ...event, 'event.name': EVENT_API_ERROR, 'event.timestamp': new Date().toISOString(), @@ -152,6 +176,7 @@ export function logApiError( }; logger.emit(logRecord); recordApiErrorMetrics( + config, event.model, event.duration_ms, event.status_code, @@ -160,10 +185,12 @@ export function logApiError( } export function logApiResponse( + config: Config, event: Omit, ): void { if (!isTelemetrySdkInitialized()) return; const attributes: LogAttributes = { + ...getCommonAttributes(config), ...event, 'event.name': EVENT_API_RESPONSE, 'event.timestamp': new Date().toISOString(), @@ -183,17 +210,29 @@ export function logApiResponse( }; logger.emit(logRecord); recordApiResponseMetrics( + config, event.model, event.duration_ms, event.status_code, event.error, ); - recordTokenUsageMetrics(event.model, event.output_token_count, 'output'); recordTokenUsageMetrics( + config, + event.model, + event.output_token_count, + 'output', + ); + recordTokenUsageMetrics( + config, event.model, event.cached_content_token_count, 'cache', ); - recordTokenUsageMetrics(event.model, event.thoughts_token_count, 'thought'); - recordTokenUsageMetrics(event.model, event.tool_token_count, 'tool'); + recordTokenUsageMetrics( + config, + event.model, + event.thoughts_token_count, + 'thought', + ); + recordTokenUsageMetrics(config, event.model, event.tool_token_count, 'tool'); } diff --git a/packages/core/src/telemetry/metrics.test.ts b/packages/core/src/telemetry/metrics.test.ts index b7864131..7e24b9ad 100644 --- a/packages/core/src/telemetry/metrics.test.ts +++ b/packages/core/src/telemetry/metrics.test.ts @@ -7,6 +7,7 @@ import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest'; import { Counter, Meter, metrics } from '@opentelemetry/api'; import { initializeMetrics, recordTokenUsageMetrics } from './metrics.js'; +import { Config } from '../config/config.js'; const mockCounter = { add: vi.fn(), @@ -33,51 +34,61 @@ describe('Telemetry Metrics', () => { }); describe('recordTokenUsageMetrics', () => { + const mockConfig = { + getSessionId: () => 'test-session-id', + } as unknown as Config; + it('should not record metrics if not initialized', () => { - recordTokenUsageMetrics('gemini-pro', 100, 'input'); + recordTokenUsageMetrics(mockConfig, 'gemini-pro', 100, 'input'); expect(mockCounter.add).not.toHaveBeenCalled(); }); it('should record token usage with the correct attributes', () => { - initializeMetrics(); - recordTokenUsageMetrics('gemini-pro', 100, 'input'); + initializeMetrics(mockConfig); + recordTokenUsageMetrics(mockConfig, 'gemini-pro', 100, 'input'); expect(mockCounter.add).toHaveBeenCalledWith(100, { + 'session.id': 'test-session-id', model: 'gemini-pro', type: 'input', }); }); it('should record token usage for different types', () => { - initializeMetrics(); - recordTokenUsageMetrics('gemini-pro', 50, 'output'); + initializeMetrics(mockConfig); + recordTokenUsageMetrics(mockConfig, 'gemini-pro', 50, 'output'); expect(mockCounter.add).toHaveBeenCalledWith(50, { + 'session.id': 'test-session-id', model: 'gemini-pro', type: 'output', }); - recordTokenUsageMetrics('gemini-pro', 25, 'thought'); + recordTokenUsageMetrics(mockConfig, 'gemini-pro', 25, 'thought'); expect(mockCounter.add).toHaveBeenCalledWith(25, { + 'session.id': 'test-session-id', model: 'gemini-pro', type: 'thought', }); - recordTokenUsageMetrics('gemini-pro', 75, 'cache'); + recordTokenUsageMetrics(mockConfig, 'gemini-pro', 75, 'cache'); expect(mockCounter.add).toHaveBeenCalledWith(75, { + 'session.id': 'test-session-id', model: 'gemini-pro', type: 'cache', }); - recordTokenUsageMetrics('gemini-pro', 125, 'tool'); + recordTokenUsageMetrics(mockConfig, 'gemini-pro', 125, 'tool'); expect(mockCounter.add).toHaveBeenCalledWith(125, { + 'session.id': 'test-session-id', model: 'gemini-pro', type: 'tool', }); }); it('should handle different models', () => { - initializeMetrics(); - recordTokenUsageMetrics('gemini-ultra', 200, 'input'); + initializeMetrics(mockConfig); + recordTokenUsageMetrics(mockConfig, 'gemini-ultra', 200, 'input'); expect(mockCounter.add).toHaveBeenCalledWith(200, { + 'session.id': 'test-session-id', model: 'gemini-ultra', type: 'input', }); diff --git a/packages/core/src/telemetry/metrics.ts b/packages/core/src/telemetry/metrics.ts index 69ac91e9..93aa2189 100644 --- a/packages/core/src/telemetry/metrics.ts +++ b/packages/core/src/telemetry/metrics.ts @@ -21,6 +21,7 @@ import { METRIC_TOKEN_USAGE, METRIC_SESSION_COUNT, } from './constants.js'; +import { Config } from '../config/config.js'; let cliMeter: Meter | undefined; let toolCallCounter: Counter | undefined; @@ -30,6 +31,12 @@ let apiRequestLatencyHistogram: Histogram | undefined; let tokenUsageCounter: Counter | undefined; let isMetricsInitialized = false; +function getCommonAttributes(config: Config): Attributes { + return { + 'session.id': config.getSessionId(), + }; +} + export function getMeter(): Meter | undefined { if (!cliMeter) { cliMeter = metrics.getMeter(SERVICE_NAME); @@ -37,7 +44,7 @@ export function getMeter(): Meter | undefined { return cliMeter; } -export function initializeMetrics(): void { +export function initializeMetrics(config: Config): void { if (isMetricsInitialized) return; const meter = getMeter(); @@ -73,11 +80,12 @@ export function initializeMetrics(): void { description: 'Count of CLI sessions started.', valueType: ValueType.INT, }); - sessionCounter.add(1); + sessionCounter.add(1, getCommonAttributes(config)); isMetricsInitialized = true; } export function recordToolCallMetrics( + config: Config, functionName: string, durationMs: number, success: boolean, @@ -86,25 +94,33 @@ export function recordToolCallMetrics( return; const metricAttributes: Attributes = { + ...getCommonAttributes(config), function_name: functionName, success, }; toolCallCounter.add(1, metricAttributes); toolCallLatencyHistogram.record(durationMs, { + ...getCommonAttributes(config), function_name: functionName, }); } export function recordTokenUsageMetrics( + config: Config, model: string, tokenCount: number, type: 'input' | 'output' | 'thought' | 'cache' | 'tool', ): void { if (!tokenUsageCounter || !isMetricsInitialized) return; - tokenUsageCounter.add(tokenCount, { model, type }); + tokenUsageCounter.add(tokenCount, { + ...getCommonAttributes(config), + model, + type, + }); } export function recordApiResponseMetrics( + config: Config, model: string, durationMs: number, statusCode?: number | string, @@ -117,14 +133,19 @@ export function recordApiResponseMetrics( ) return; const metricAttributes: Attributes = { + ...getCommonAttributes(config), model, status_code: statusCode ?? (error ? 'error' : 'ok'), }; apiRequestCounter.add(1, metricAttributes); - apiRequestLatencyHistogram.record(durationMs, { model }); + apiRequestLatencyHistogram.record(durationMs, { + ...getCommonAttributes(config), + model, + }); } export function recordApiErrorMetrics( + config: Config, model: string, durationMs: number, statusCode?: number | string, @@ -137,10 +158,14 @@ export function recordApiErrorMetrics( ) return; const metricAttributes: Attributes = { + ...getCommonAttributes(config), model, status_code: statusCode ?? 'error', error_type: errorType ?? 'unknown', }; apiRequestCounter.add(1, metricAttributes); - apiRequestLatencyHistogram.record(durationMs, { model }); + apiRequestLatencyHistogram.record(durationMs, { + ...getCommonAttributes(config), + model, + }); } diff --git a/packages/core/src/telemetry/sdk.ts b/packages/core/src/telemetry/sdk.ts index 704661c5..61f501a6 100644 --- a/packages/core/src/telemetry/sdk.ts +++ b/packages/core/src/telemetry/sdk.ts @@ -112,7 +112,7 @@ export function initializeTelemetry(config: Config): void { sdk.start(); console.log('OpenTelemetry SDK started successfully.'); telemetryInitialized = true; - initializeMetrics(); + initializeMetrics(config); logCliConfiguration(config); } catch (error) { console.error('Error starting OpenTelemetry SDK:', error); diff --git a/packages/core/src/utils/nextSpeakerChecker.test.ts b/packages/core/src/utils/nextSpeakerChecker.test.ts index a19045d1..83ce97fd 100644 --- a/packages/core/src/utils/nextSpeakerChecker.test.ts +++ b/packages/core/src/utils/nextSpeakerChecker.test.ts @@ -69,9 +69,9 @@ describe('checkNextSpeaker', () => { // GeminiChat will receive the mocked instances via the mocked GoogleGenAI constructor chatInstance = new GeminiChat( + mockConfigInstance, mockModelsInstance, // This is the instance returned by mockGoogleGenAIInstance.getGenerativeModel 'gemini-pro', // model name - 'test-session-id', {}, [], // initial history );