diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 27b34d87..cef82d23 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -14,15 +14,12 @@ import { SendMessageParameters, createUserContent, Part, - GenerateContentResponseUsageMetadata, Tool, } from '@google/genai'; import { retryWithBackoff } from '../utils/retry.js'; import { isFunctionResponse } from '../utils/messageInspectors.js'; import { ContentGenerator, AuthType } from './contentGenerator.js'; import { Config } from '../config/config.js'; -import { logApiResponse, logApiError } from '../telemetry/loggers.js'; -import { ApiErrorEvent, ApiResponseEvent } from '../telemetry/types.js'; import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js'; import { hasCycleInSchema } from '../tools/tools.js'; import { StructuredError } from './turn.js'; @@ -131,46 +128,6 @@ export class GeminiChat { validateHistory(history); } - private async _logApiResponse( - durationMs: number, - prompt_id: string, - usageMetadata?: GenerateContentResponseUsageMetadata, - responseText?: string, - ): Promise { - logApiResponse( - this.config, - new ApiResponseEvent( - this.config.getModel(), - durationMs, - prompt_id, - this.config.getContentGeneratorConfig()?.authType, - usageMetadata, - responseText, - ), - ); - } - - private _logApiError( - durationMs: number, - error: unknown, - prompt_id: string, - ): void { - const errorMessage = error instanceof Error ? error.message : String(error); - const errorType = error instanceof Error ? error.name : 'unknown'; - - logApiError( - this.config, - new ApiErrorEvent( - this.config.getModel(), - errorMessage, - durationMs, - prompt_id, - this.config.getContentGeneratorConfig()?.authType, - errorType, - ), - ); - } - /** * Handles falling back to Flash model when persistent 429 errors occur for OAuth users. * Uses a fallback handler if provided by the config; otherwise, returns null. @@ -249,7 +206,6 @@ export class GeminiChat { const userContent = createUserContent(params.message); const requestContents = this.getHistory(true).concat(userContent); - const startTime = Date.now(); let response: GenerateContentResponse; try { @@ -290,13 +246,6 @@ export class GeminiChat { await this.handleFlashFallback(authType, error), authType: this.config.getContentGeneratorConfig()?.authType, }); - const durationMs = Date.now() - startTime; - await this._logApiResponse( - durationMs, - prompt_id, - response.usageMetadata, - JSON.stringify(response), - ); this.sendPromise = (async () => { const outputContent = response.candidates?.[0]?.content; @@ -324,8 +273,6 @@ export class GeminiChat { }); return response; } catch (error) { - const durationMs = Date.now() - startTime; - this._logApiError(durationMs, error, prompt_id); this.sendPromise = Promise.resolve(); throw error; } @@ -361,8 +308,6 @@ export class GeminiChat { const userContent = createUserContent(params.message); const requestContents = this.getHistory(true).concat(userContent); - const startTime = Date.now(); - try { const apiCall = () => { const modelToUse = this.config.getModel(); @@ -413,16 +358,9 @@ export class GeminiChat { .then(() => undefined) .catch(() => undefined); - const result = this.processStreamResponse( - streamResponse, - userContent, - startTime, - prompt_id, - ); + const result = this.processStreamResponse(streamResponse, userContent); return result; } catch (error) { - const durationMs = Date.now() - startTime; - this._logApiError(durationMs, error, prompt_id); this.sendPromise = Promise.resolve(); throw error; } @@ -483,17 +421,6 @@ export class GeminiChat { this.generationConfig.tools = tools; } - getFinalUsageMetadata( - chunks: GenerateContentResponse[], - ): GenerateContentResponseUsageMetadata | undefined { - const lastChunkWithMetadata = chunks - .slice() - .reverse() - .find((chunk) => chunk.usageMetadata); - - return lastChunkWithMetadata?.usageMetadata; - } - async maybeIncludeSchemaDepthContext(error: StructuredError): Promise { // Check for potentially problematic cyclic tools with cyclic schemas // and include a recommendation to remove potentially problematic tools. @@ -525,8 +452,6 @@ export class GeminiChat { private async *processStreamResponse( streamResponse: AsyncGenerator, inputContent: Content, - startTime: number, - prompt_id: string, ) { const outputContent: Content[] = []; const chunks: GenerateContentResponse[] = []; @@ -549,25 +474,16 @@ export class GeminiChat { } } catch (error) { errorOccurred = true; - const durationMs = Date.now() - startTime; - this._logApiError(durationMs, error, prompt_id); throw error; } if (!errorOccurred) { - const durationMs = Date.now() - startTime; const allParts: Part[] = []; for (const content of outputContent) { if (content.parts) { allParts.push(...content.parts); } } - await this._logApiResponse( - durationMs, - prompt_id, - this.getFinalUsageMetadata(chunks), - JSON.stringify(chunks), - ); } this.recordHistory(inputContent, outputContent); } diff --git a/packages/core/src/core/loggingContentGenerator.ts b/packages/core/src/core/loggingContentGenerator.ts index c9069bab..305b19a4 100644 --- a/packages/core/src/core/loggingContentGenerator.ts +++ b/packages/core/src/core/loggingContentGenerator.ts @@ -11,11 +11,20 @@ import { EmbedContentParameters, EmbedContentResponse, GenerateContentParameters, + GenerateContentResponseUsageMetadata, GenerateContentResponse, } from '@google/genai'; -import { ApiRequestEvent } from '../telemetry/types.js'; +import { + ApiRequestEvent, + ApiResponseEvent, + ApiErrorEvent, +} from '../telemetry/types.js'; import { Config } from '../config/config.js'; -import { logApiRequest } from '../telemetry/loggers.js'; +import { + logApiError, + logApiRequest, + logApiResponse, +} from '../telemetry/loggers.js'; import { ContentGenerator } from './contentGenerator.js'; import { toContents } from '../code_assist/converter.js'; @@ -40,20 +49,117 @@ export class LoggingContentGenerator implements ContentGenerator { ); } + private _logApiResponse( + durationMs: number, + prompt_id: string, + usageMetadata?: GenerateContentResponseUsageMetadata, + responseText?: string, + ): void { + logApiResponse( + this.config, + new ApiResponseEvent( + this.config.getModel(), + durationMs, + prompt_id, + this.config.getContentGeneratorConfig()?.authType, + usageMetadata, + responseText, + ), + ); + } + + private _logApiError( + durationMs: number, + error: unknown, + prompt_id: string, + ): void { + const errorMessage = error instanceof Error ? error.message : String(error); + const errorType = error instanceof Error ? error.name : 'unknown'; + + logApiError( + this.config, + new ApiErrorEvent( + this.config.getModel(), + errorMessage, + durationMs, + prompt_id, + this.config.getContentGeneratorConfig()?.authType, + errorType, + ), + ); + } + async generateContent( req: GenerateContentParameters, userPromptId: string, ): Promise { + const startTime = Date.now(); this.logApiRequest(toContents(req.contents), req.model, userPromptId); - return this.wrapped.generateContent(req, userPromptId); + try { + const response = await this.wrapped.generateContent(req, userPromptId); + const durationMs = Date.now() - startTime; + this._logApiResponse( + durationMs, + userPromptId, + response.usageMetadata, + JSON.stringify(response), + ); + return response; + } catch (error) { + const durationMs = Date.now() - startTime; + this._logApiError(durationMs, error, userPromptId); + throw error; + } } async generateContentStream( req: GenerateContentParameters, userPromptId: string, ): Promise> { + const startTime = Date.now(); this.logApiRequest(toContents(req.contents), req.model, userPromptId); - return this.wrapped.generateContentStream(req, userPromptId); + + let stream: AsyncGenerator; + try { + stream = await this.wrapped.generateContentStream(req, userPromptId); + } catch (error) { + const durationMs = Date.now() - startTime; + this._logApiError(durationMs, error, userPromptId); + throw error; + } + + return this.loggingStreamWrapper(stream, startTime, userPromptId); + } + + private async *loggingStreamWrapper( + stream: AsyncGenerator, + startTime: number, + userPromptId: string, + ): AsyncGenerator { + let lastResponse: GenerateContentResponse | undefined; + let lastUsageMetadata: GenerateContentResponseUsageMetadata | undefined; + try { + for await (const response of stream) { + lastResponse = response; + if (response.usageMetadata) { + lastUsageMetadata = response.usageMetadata; + } + yield response; + } + } catch (error) { + const durationMs = Date.now() - startTime; + this._logApiError(durationMs, error, userPromptId); + throw error; + } + const durationMs = Date.now() - startTime; + if (lastResponse) { + this._logApiResponse( + durationMs, + userPromptId, + lastUsageMetadata, + JSON.stringify(lastResponse), + ); + } } async countTokens(req: CountTokensParameters): Promise {