From b20c8389f3f483f3972c254ec97fff4004b7c75f Mon Sep 17 00:00:00 2001 From: Jerop Kipruto Date: Fri, 13 Jun 2025 03:44:17 -0400 Subject: [PATCH] Handle telemetry in non-interactive mode (#1002) Changes: - Ensure proper shutdown in non-interactive mode - Ensures the initial user prompt is logged in non-interactive mode - Improve telemetry for streaming - handle chunks and input token count is now alongside other token counts in response To test: - Follow instructions in https://github.com/google-gemini/gemini-cli/blob/main/docs/core/telemetry.md#google-cloud - Run CLI in non-interactive mode and observe logs/metrics in GCP Logs Explorer and Metrics Explorer #750 --- docs/core/telemetry.md | 1 - packages/cli/src/gemini.tsx | 6 ++ packages/cli/src/nonInteractiveCli.ts | 6 ++ packages/core/src/core/geminiChat.ts | 63 +++++++++---------- .../src/core/nonInteractiveToolExecutor.ts | 26 ++++++++ packages/core/src/telemetry/index.ts | 1 + packages/core/src/telemetry/loggers.test.ts | 36 ++--------- packages/core/src/telemetry/loggers.ts | 60 +++++++++++++++--- packages/core/src/telemetry/types.ts | 1 - packages/core/src/utils/requestUtils.ts | 15 ----- 10 files changed, 126 insertions(+), 89 deletions(-) delete mode 100644 packages/core/src/utils/requestUtils.ts diff --git a/docs/core/telemetry.md b/docs/core/telemetry.md index 41bd9fb7..b7989483 100644 --- a/docs/core/telemetry.md +++ b/docs/core/telemetry.md @@ -296,7 +296,6 @@ These are timestamped records of specific events. - **Attributes**: - `model` - - `input_token_count` - `request_text` (optional) - `gemini_cli.api_error`: Fired if the API request fails. diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx index 0478a40e..34ec3bd0 100644 --- a/packages/cli/src/gemini.tsx +++ b/packages/cli/src/gemini.tsx @@ -33,6 +33,7 @@ import { WebSearchTool, WriteFileTool, sessionId, + logUserPrompt, } from '@gemini-cli/core'; export async function main() { @@ -120,6 +121,11 @@ export async function main() { process.exit(1); } + logUserPrompt(config, { + prompt: input, + prompt_length: input.length, + }); + // Non-interactive mode handled by runNonInteractive const nonInteractiveConfig = await loadNonInteractiveConfig( config, diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts index 069a7be8..bf60011f 100644 --- a/packages/cli/src/nonInteractiveCli.ts +++ b/packages/cli/src/nonInteractiveCli.ts @@ -9,6 +9,8 @@ import { ToolCallRequestInfo, executeToolCall, ToolRegistry, + shutdownTelemetry, + isTelemetrySdkInitialized, } from '@gemini-cli/core'; import { Content, @@ -119,5 +121,9 @@ export async function runNonInteractive( } catch (error) { console.error('Error processing input:', error); process.exit(1); + } finally { + if (isTelemetrySdkInitialized()) { + await shutdownTelemetry(); + } } } diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index e1f2738d..eb740c4a 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -14,6 +14,7 @@ import { SendMessageParameters, createUserContent, Part, + GenerateContentResponseUsageMetadata, } from '@google/genai'; import { retryWithBackoff } from '../utils/retry.js'; import { isFunctionResponse } from '../utils/messageInspectors.js'; @@ -23,13 +24,12 @@ import { logApiRequest, logApiResponse, logApiError, + combinedUsageMetadata, } from '../telemetry/loggers.js'; import { getResponseText, getResponseTextFromParts, } from '../utils/generateContentResponseUtilities.js'; -import { getErrorMessage } from '../utils/errors.js'; -import { getRequestTextFromContents } from '../utils/requestUtils.js'; /** * Returns true if the response is valid, false otherwise. @@ -140,51 +140,42 @@ export class GeminiChat { validateHistory(history); } + private _getRequestTextFromContents(contents: Content[]): string { + return contents + .flatMap((content) => content.parts ?? []) + .map((part) => part.text) + .filter(Boolean) + .join(''); + } + private async _logApiRequest( contents: Content[], model: string, ): Promise { - let inputTokenCount = 0; - try { - const { totalTokens } = await this.contentGenerator.countTokens({ - model, - contents, - }); - inputTokenCount = totalTokens || 0; - } catch (_e) { - console.warn( - `Failed to count tokens for model ${model}. Proceeding with inputTokenCount = 0. Error: ${getErrorMessage(_e)}`, - ); - inputTokenCount = 0; - } - const shouldLogUserPrompts = (config: Config): boolean => config.getTelemetryLogUserPromptsEnabled() ?? false; - const requestText = getRequestTextFromContents(contents); + const requestText = this._getRequestTextFromContents(contents); logApiRequest(this.config, { model, - input_token_count: inputTokenCount, request_text: shouldLogUserPrompts(this.config) ? requestText : undefined, }); } private async _logApiResponse( durationMs: number, - response: GenerateContentResponse, - fullStreamedText?: string, + usageMetadata?: GenerateContentResponseUsageMetadata, + responseText?: string, ): Promise { - const responseText = fullStreamedText ?? getResponseText(response); logApiResponse(this.config, { model: this.model, duration_ms: durationMs, status_code: 200, // Assuming 200 for success - input_token_count: response.usageMetadata?.promptTokenCount ?? 0, - output_token_count: response.usageMetadata?.candidatesTokenCount ?? 0, - cached_content_token_count: - response.usageMetadata?.cachedContentTokenCount ?? 0, - thoughts_token_count: response.usageMetadata?.thoughtsTokenCount ?? 0, - tool_token_count: response.usageMetadata?.toolUsePromptTokenCount ?? 0, + input_token_count: usageMetadata?.promptTokenCount ?? 0, + output_token_count: usageMetadata?.candidatesTokenCount ?? 0, + cached_content_token_count: usageMetadata?.cachedContentTokenCount ?? 0, + thoughts_token_count: usageMetadata?.thoughtsTokenCount ?? 0, + tool_token_count: usageMetadata?.toolUsePromptTokenCount ?? 0, response_text: responseText, }); } @@ -245,7 +236,11 @@ export class GeminiChat { response = await retryWithBackoff(apiCall); const durationMs = Date.now() - startTime; - await this._logApiResponse(durationMs, response); + await this._logApiResponse( + durationMs, + response.usageMetadata, + getResponseText(response), + ); this.sendPromise = (async () => { const outputContent = response.candidates?.[0]?.content; @@ -413,18 +408,18 @@ export class GeminiChat { startTime: number, ) { const outputContent: Content[] = []; - let lastChunk: GenerateContentResponse | undefined; + const chunks: GenerateContentResponse[] = []; let errorOccurred = false; try { for await (const chunk of streamResponse) { if (isValidResponse(chunk)) { + chunks.push(chunk); const content = chunk.candidates?.[0]?.content; if (content !== undefined) { outputContent.push(content); } } - lastChunk = chunk; yield chunk; } } catch (error) { @@ -434,7 +429,7 @@ export class GeminiChat { throw error; } - if (!errorOccurred && lastChunk) { + if (!errorOccurred) { const durationMs = Date.now() - startTime; const allParts: Part[] = []; for (const content of outputContent) { @@ -443,7 +438,11 @@ export class GeminiChat { } } const fullText = getResponseTextFromParts(allParts); - await this._logApiResponse(durationMs, lastChunk, fullText); + await this._logApiResponse( + durationMs, + combinedUsageMetadata(chunks), + fullText, + ); } this.recordHistory(inputContent, outputContent); } diff --git a/packages/core/src/core/nonInteractiveToolExecutor.ts b/packages/core/src/core/nonInteractiveToolExecutor.ts index 5f7ed4ff..f2174e06 100644 --- a/packages/core/src/core/nonInteractiveToolExecutor.ts +++ b/packages/core/src/core/nonInteractiveToolExecutor.ts @@ -5,6 +5,7 @@ */ import { + logToolCall, ToolCallRequestInfo, ToolCallResponseInfo, ToolRegistry, @@ -25,10 +26,19 @@ export async function executeToolCall( ): Promise { const tool = toolRegistry.getTool(toolCallRequest.name); + const startTime = Date.now(); if (!tool) { const error = new Error( `Tool "${toolCallRequest.name}" not found in registry.`, ); + const durationMs = Date.now() - startTime; + logToolCall(config, { + function_name: toolCallRequest.name, + function_args: toolCallRequest.args, + duration_ms: durationMs, + success: false, + error: error.message, + }); // Ensure the response structure matches what the API expects for an error return { callId: toolCallRequest.callId, @@ -55,6 +65,14 @@ export async function executeToolCall( // No live output callback for non-interactive mode ); + const durationMs = Date.now() - startTime; + logToolCall(config, { + function_name: toolCallRequest.name, + function_args: toolCallRequest.args, + duration_ms: durationMs, + success: true, + }); + const response = convertToFunctionResponse( toolCallRequest.name, toolCallRequest.callId, @@ -69,6 +87,14 @@ export async function executeToolCall( }; } catch (e) { const error = e instanceof Error ? e : new Error(String(e)); + const durationMs = Date.now() - startTime; + logToolCall(config, { + function_name: toolCallRequest.name, + function_args: toolCallRequest.args, + duration_ms: durationMs, + success: false, + error: error.message, + }); return { callId: toolCallRequest.callId, responseParts: [ diff --git a/packages/core/src/telemetry/index.ts b/packages/core/src/telemetry/index.ts index cbb7b4d2..e8248bf9 100644 --- a/packages/core/src/telemetry/index.ts +++ b/packages/core/src/telemetry/index.ts @@ -16,6 +16,7 @@ export { logApiRequest, logApiError, logApiResponse, + combinedUsageMetadata, } from './loggers.js'; export { UserPromptEvent, diff --git a/packages/core/src/telemetry/loggers.test.ts b/packages/core/src/telemetry/loggers.test.ts index 092d70f0..58ea3a10 100644 --- a/packages/core/src/telemetry/loggers.test.ts +++ b/packages/core/src/telemetry/loggers.test.ts @@ -113,7 +113,7 @@ describe('loggers', () => { logUserPrompt(mockConfig, event); expect(mockLogger.emit).toHaveBeenCalledWith({ - body: 'User prompt. Length: 11', + body: 'User prompt. Length: 11.', attributes: { 'session.id': 'test-session-id', 'event.name': EVENT_USER_PROMPT, @@ -137,7 +137,7 @@ describe('loggers', () => { logUserPrompt(mockConfig, event); expect(mockLogger.emit).toHaveBeenCalledWith({ - body: 'User prompt. Length: 11', + body: 'User prompt. Length: 11.', attributes: { 'session.id': 'test-session-id', 'event.name': EVENT_USER_PROMPT, @@ -250,70 +250,42 @@ describe('loggers', () => { getSessionId: () => 'test-session-id', } as Config; - const mockMetrics = { - recordTokenUsageMetrics: vi.fn(), - }; - - beforeEach(() => { - vi.spyOn(metrics, 'recordTokenUsageMetrics').mockImplementation( - mockMetrics.recordTokenUsageMetrics, - ); - }); - it('should log an API request with request_text', () => { const event = { model: 'test-model', - input_token_count: 123, request_text: 'This is a test request', }; logApiRequest(mockConfig, event); expect(mockLogger.emit).toHaveBeenCalledWith({ - body: 'API request to test-model. Tokens: 123.', + body: 'API request to test-model.', attributes: { 'session.id': 'test-session-id', 'event.name': EVENT_API_REQUEST, 'event.timestamp': '2025-01-01T00:00:00.000Z', model: 'test-model', - input_token_count: 123, request_text: 'This is a test request', }, }); - - expect(mockMetrics.recordTokenUsageMetrics).toHaveBeenCalledWith( - mockConfig, - 'test-model', - 123, - 'input', - ); }); it('should log an API request without request_text', () => { const event = { model: 'test-model', - input_token_count: 456, }; logApiRequest(mockConfig, event); expect(mockLogger.emit).toHaveBeenCalledWith({ - body: 'API request to test-model. Tokens: 456.', + body: 'API request to test-model.', attributes: { 'session.id': 'test-session-id', 'event.name': EVENT_API_REQUEST, 'event.timestamp': '2025-01-01T00:00:00.000Z', model: 'test-model', - input_token_count: 456, }, }); - - expect(mockMetrics.recordTokenUsageMetrics).toHaveBeenCalledWith( - mockConfig, - 'test-model', - 456, - 'input', - ); }); }); diff --git a/packages/core/src/telemetry/loggers.ts b/packages/core/src/telemetry/loggers.ts index 66f584e7..f6896def 100644 --- a/packages/core/src/telemetry/loggers.ts +++ b/packages/core/src/telemetry/loggers.ts @@ -31,6 +31,10 @@ import { } from './metrics.js'; import { isTelemetrySdkInitialized } from './sdk.js'; import { ToolConfirmationOutcome } from '../index.js'; +import { + GenerateContentResponse, + GenerateContentResponseUsageMetadata, +} from '@google/genai'; const shouldLogUserPrompts = (config: Config): boolean => config.getTelemetryLogUserPromptsEnabled() ?? false; @@ -119,7 +123,7 @@ export function logUserPrompt( const logger = logs.getLogger(SERVICE_NAME); const logRecord: LogRecord = { - body: `User prompt. Length: ${event.prompt_length}`, + body: `User prompt. Length: ${event.prompt_length}.`, attributes, }; logger.emit(logRecord); @@ -176,16 +180,10 @@ export function logApiRequest( }; const logger = logs.getLogger(SERVICE_NAME); const logRecord: LogRecord = { - body: `API request to ${event.model}. Tokens: ${event.input_token_count}.`, + body: `API request to ${event.model}.`, attributes, }; logger.emit(logRecord); - recordTokenUsageMetrics( - config, - event.model, - event.input_token_count, - 'input', - ); } export function logApiError( @@ -258,6 +256,12 @@ export function logApiResponse( event.status_code, event.error, ); + recordTokenUsageMetrics( + config, + event.model, + event.input_token_count, + 'input', + ); recordTokenUsageMetrics( config, event.model, @@ -278,3 +282,43 @@ export function logApiResponse( ); recordTokenUsageMetrics(config, event.model, event.tool_token_count, 'tool'); } + +export function combinedUsageMetadata( + chunks: GenerateContentResponse[], +): GenerateContentResponseUsageMetadata { + const metadataKeys: Array = [ + 'promptTokenCount', + 'candidatesTokenCount', + 'cachedContentTokenCount', + 'thoughtsTokenCount', + 'toolUsePromptTokenCount', + 'totalTokenCount', + ]; + + const totals: Record = { + promptTokenCount: 0, + candidatesTokenCount: 0, + cachedContentTokenCount: 0, + thoughtsTokenCount: 0, + toolUsePromptTokenCount: 0, + totalTokenCount: 0, + cacheTokensDetails: 0, + candidatesTokensDetails: 0, + promptTokensDetails: 0, + toolUsePromptTokensDetails: 0, + trafficType: 0, + }; + + for (const chunk of chunks) { + if (chunk.usageMetadata) { + for (const key of metadataKeys) { + const chunkValue = chunk.usageMetadata[key]; + if (typeof chunkValue === 'number') { + totals[key] += chunkValue; + } + } + } + } + + return totals as unknown as GenerateContentResponseUsageMetadata; +} diff --git a/packages/core/src/telemetry/types.ts b/packages/core/src/telemetry/types.ts index e9d397a8..bed9b16d 100644 --- a/packages/core/src/telemetry/types.ts +++ b/packages/core/src/telemetry/types.ts @@ -29,7 +29,6 @@ export interface ApiRequestEvent { 'event.name': 'api_request'; 'event.timestamp': string; // ISO 8601 model: string; - input_token_count: number; request_text?: string; } diff --git a/packages/core/src/utils/requestUtils.ts b/packages/core/src/utils/requestUtils.ts deleted file mode 100644 index 4978e968..00000000 --- a/packages/core/src/utils/requestUtils.ts +++ /dev/null @@ -1,15 +0,0 @@ -/** - * @license - * Copyright 2025 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { Content } from '@google/genai'; - -export function getRequestTextFromContents(contents: Content[]): string { - return contents - .flatMap((content) => content.parts ?? []) - .map((part) => part.text) - .filter(Boolean) - .join(''); -}