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
This commit is contained in:
Jerop Kipruto
2025-06-13 03:44:17 -04:00
committed by GitHub
parent 8bb6eca915
commit b20c8389f3
10 changed files with 126 additions and 89 deletions

View File

@@ -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<keyof GenerateContentResponseUsageMetadata> = [
'promptTokenCount',
'candidatesTokenCount',
'cachedContentTokenCount',
'thoughtsTokenCount',
'toolUsePromptTokenCount',
'totalTokenCount',
];
const totals: Record<keyof GenerateContentResponseUsageMetadata, number> = {
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;
}