mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-03-11 22:51:00 -07:00
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:
@@ -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> {
|
||||
|
||||
Reference in New Issue
Block a user