mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-30 15:04:16 -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,
|
startTime: number,
|
||||||
userPromptId: string,
|
userPromptId: string,
|
||||||
): AsyncGenerator<GenerateContentResponse> {
|
): AsyncGenerator<GenerateContentResponse> {
|
||||||
let lastResponse: GenerateContentResponse | undefined;
|
|
||||||
const responses: GenerateContentResponse[] = [];
|
const responses: GenerateContentResponse[] = [];
|
||||||
|
|
||||||
let lastUsageMetadata: GenerateContentResponseUsageMetadata | undefined;
|
let lastUsageMetadata: GenerateContentResponseUsageMetadata | undefined;
|
||||||
try {
|
try {
|
||||||
for await (const response of stream) {
|
for await (const response of stream) {
|
||||||
responses.push(response);
|
responses.push(response);
|
||||||
lastResponse = response;
|
|
||||||
if (response.usageMetadata) {
|
if (response.usageMetadata) {
|
||||||
lastUsageMetadata = response.usageMetadata;
|
lastUsageMetadata = response.usageMetadata;
|
||||||
}
|
}
|
||||||
yield response;
|
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) {
|
} catch (error) {
|
||||||
const durationMs = Date.now() - startTime;
|
const durationMs = Date.now() - startTime;
|
||||||
this._logApiError(
|
this._logApiError(
|
||||||
@@ -174,16 +181,6 @@ export class LoggingContentGenerator implements ContentGenerator {
|
|||||||
);
|
);
|
||||||
throw error;
|
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> {
|
async countTokens(req: CountTokensParameters): Promise<CountTokensResponse> {
|
||||||
|
|||||||
Reference in New Issue
Block a user