fix(token-usage): fix token usage recording for failed API calls (#7667)

Co-authored-by: Bryan Morgan <bryanmorgan@google.com>
This commit is contained in:
Nick Salerni
2025-09-05 11:22:51 -07:00
committed by GitHub
parent 5634be70f1
commit a7bfab4d96

View File

@@ -151,19 +151,26 @@ export class LoggingContentGenerator implements ContentGenerator {
startTime: number,
userPromptId: string,
): AsyncGenerator<GenerateContentResponse> {
let lastResponse: GenerateContentResponse | undefined;
const responses: GenerateContentResponse[] = [];
let lastUsageMetadata: GenerateContentResponseUsageMetadata | undefined;
try {
for await (const response of stream) {
responses.push(response);
lastResponse = response;
if (response.usageMetadata) {
lastUsageMetadata = response.usageMetadata;
}
yield response;
}
// Only log successful API response if no error occurred
const durationMs = Date.now() - startTime;
this._logApiResponse(
durationMs,
responses[0]?.modelVersion || '',
userPromptId,
lastUsageMetadata,
JSON.stringify(responses),
);
} catch (error) {
const durationMs = Date.now() - startTime;
this._logApiError(
@@ -174,16 +181,6 @@ export class LoggingContentGenerator implements ContentGenerator {
);
throw error;
}
const durationMs = Date.now() - startTime;
if (lastResponse) {
this._logApiResponse(
durationMs,
responses[0]?.modelVersion || '',
userPromptId,
lastUsageMetadata,
JSON.stringify(responses),
);
}
}
async countTokens(req: CountTokensParameters): Promise<CountTokensResponse> {