summaryrefslogtreecommitdiff
path: root/packages/core/src
diff options
context:
space:
mode:
Diffstat (limited to 'packages/core/src')
-rw-r--r--packages/core/src/core/geminiChat.ts63
-rw-r--r--packages/core/src/core/nonInteractiveToolExecutor.ts26
-rw-r--r--packages/core/src/telemetry/index.ts1
-rw-r--r--packages/core/src/telemetry/loggers.test.ts36
-rw-r--r--packages/core/src/telemetry/loggers.ts60
-rw-r--r--packages/core/src/telemetry/types.ts1
-rw-r--r--packages/core/src/utils/requestUtils.ts15
7 files changed, 114 insertions, 88 deletions
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<void> {
- 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<void> {
- 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<ToolCallResponseInfo> {
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(
@@ -261,6 +259,12 @@ export function logApiResponse(
recordTokenUsageMetrics(
config,
event.model,
+ event.input_token_count,
+ 'input',
+ );
+ recordTokenUsageMetrics(
+ config,
+ event.model,
event.output_token_count,
'output',
);
@@ -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;
+}
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('');
-}
generated by cgit v1.2.3 (git 2.51.0) at 2025-11-08 23:29:49 +0000