From bae922a6327a8ae97ca53ac1829ff385ca025460 Mon Sep 17 00:00:00 2001 From: shishu314 Date: Thu, 7 Aug 2025 19:58:18 -0400 Subject: [PATCH] fix(cli) - Move logging into CodeAssistServer (#5781) Co-authored-by: Shi Shu --- .../core/src/code_assist/converter.test.ts | 55 +++++++++++++++ packages/core/src/code_assist/converter.ts | 2 +- packages/core/src/code_assist/server.test.ts | 6 +- .../core/src/core/contentGenerator.test.ts | 17 +++-- packages/core/src/core/contentGenerator.ts | 16 +++-- packages/core/src/core/geminiChat.ts | 31 +-------- .../core/src/core/loggingContentGenerator.ts | 68 +++++++++++++++++++ .../src/core/nonInteractiveToolExecutor.ts | 2 +- packages/core/src/telemetry/loggers.test.ts | 2 +- .../core/src/telemetry/tool-call-decision.ts | 32 +++++++++ packages/core/src/telemetry/types.ts | 30 ++------ .../core/src/telemetry/uiTelemetry.test.ts | 8 +-- packages/core/src/telemetry/uiTelemetry.ts | 8 +-- 13 files changed, 195 insertions(+), 82 deletions(-) create mode 100644 packages/core/src/core/loggingContentGenerator.ts create mode 100644 packages/core/src/telemetry/tool-call-decision.ts diff --git a/packages/core/src/code_assist/converter.test.ts b/packages/core/src/code_assist/converter.test.ts index 3d3a8ef3..ab040fe6 100644 --- a/packages/core/src/code_assist/converter.test.ts +++ b/packages/core/src/code_assist/converter.test.ts @@ -9,8 +9,10 @@ import { toGenerateContentRequest, fromGenerateContentResponse, CaGenerateContentResponse, + toContents, } from './converter.js'; import { + ContentListUnion, GenerateContentParameters, GenerateContentResponse, FinishReason, @@ -295,4 +297,57 @@ describe('converter', () => { ); }); }); + + describe('toContents', () => { + it('should handle Content', () => { + const content: ContentListUnion = { + role: 'user', + parts: [{ text: 'hello' }], + }; + expect(toContents(content)).toEqual([ + { role: 'user', parts: [{ text: 'hello' }] }, + ]); + }); + + it('should handle array of Contents', () => { + const contents: ContentListUnion = [ + { role: 'user', parts: [{ text: 'hello' }] }, + { role: 'model', parts: [{ text: 'hi' }] }, + ]; + expect(toContents(contents)).toEqual([ + { role: 'user', parts: [{ text: 'hello' }] }, + { role: 'model', parts: [{ text: 'hi' }] }, + ]); + }); + + it('should handle Part', () => { + const part: ContentListUnion = { text: 'a part' }; + expect(toContents(part)).toEqual([ + { role: 'user', parts: [{ text: 'a part' }] }, + ]); + }); + + it('should handle array of Parts', () => { + const parts = [{ text: 'part 1' }, 'part 2']; + expect(toContents(parts)).toEqual([ + { role: 'user', parts: [{ text: 'part 1' }] }, + { role: 'user', parts: [{ text: 'part 2' }] }, + ]); + }); + + it('should handle string', () => { + const str: ContentListUnion = 'a string'; + expect(toContents(str)).toEqual([ + { role: 'user', parts: [{ text: 'a string' }] }, + ]); + }); + + it('should handle array of strings', () => { + const strings: ContentListUnion = ['string 1', 'string 2']; + expect(toContents(strings)).toEqual([ + { role: 'user', parts: [{ text: 'string 1' }] }, + { role: 'user', parts: [{ text: 'string 2' }] }, + ]); + }); + }); }); diff --git a/packages/core/src/code_assist/converter.ts b/packages/core/src/code_assist/converter.ts index ffd471da..a71c70da 100644 --- a/packages/core/src/code_assist/converter.ts +++ b/packages/core/src/code_assist/converter.ts @@ -157,7 +157,7 @@ function toVertexGenerateContentRequest( }; } -function toContents(contents: ContentListUnion): Content[] { +export function toContents(contents: ContentListUnion): Content[] { if (Array.isArray(contents)) { // it's a Content[] or a PartsUnion[] return contents.map(toContent); diff --git a/packages/core/src/code_assist/server.test.ts b/packages/core/src/code_assist/server.test.ts index 3fc1891f..bc3978fb 100644 --- a/packages/core/src/code_assist/server.test.ts +++ b/packages/core/src/code_assist/server.test.ts @@ -4,7 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { describe, it, expect, vi } from 'vitest'; +import { beforeEach, describe, it, expect, vi } from 'vitest'; import { CodeAssistServer } from './server.js'; import { OAuth2Client } from 'google-auth-library'; import { UserTierId } from './types.js'; @@ -12,6 +12,10 @@ import { UserTierId } from './types.js'; vi.mock('google-auth-library'); describe('CodeAssistServer', () => { + beforeEach(() => { + vi.resetAllMocks(); + }); + it('should be able to be constructed', () => { const auth = new OAuth2Client(); const server = new CodeAssistServer( diff --git a/packages/core/src/core/contentGenerator.test.ts b/packages/core/src/core/contentGenerator.test.ts index 78eee386..e89b5bc7 100644 --- a/packages/core/src/core/contentGenerator.test.ts +++ b/packages/core/src/core/contentGenerator.test.ts @@ -9,10 +9,12 @@ import { createContentGenerator, AuthType, createContentGeneratorConfig, + ContentGenerator, } from './contentGenerator.js'; import { createCodeAssistContentGenerator } from '../code_assist/codeAssist.js'; import { GoogleGenAI } from '@google/genai'; import { Config } from '../config/config.js'; +import { LoggingContentGenerator } from './loggingContentGenerator.js'; vi.mock('../code_assist/codeAssist.js'); vi.mock('@google/genai'); @@ -21,7 +23,7 @@ const mockConfig = {} as unknown as Config; describe('createContentGenerator', () => { it('should create a CodeAssistContentGenerator', async () => { - const mockGenerator = {} as unknown; + const mockGenerator = {} as unknown as ContentGenerator; vi.mocked(createCodeAssistContentGenerator).mockResolvedValue( mockGenerator as never, ); @@ -33,13 +35,15 @@ describe('createContentGenerator', () => { mockConfig, ); expect(createCodeAssistContentGenerator).toHaveBeenCalled(); - expect(generator).toBe(mockGenerator); + expect(generator).toEqual( + new LoggingContentGenerator(mockGenerator, mockConfig), + ); }); it('should create a GoogleGenAI content generator', async () => { const mockGenerator = { models: {}, - } as unknown; + } as unknown as GoogleGenAI; vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never); const generator = await createContentGenerator( { @@ -58,7 +62,12 @@ describe('createContentGenerator', () => { }, }, }); - expect(generator).toBe((mockGenerator as GoogleGenAI).models); + expect(generator).toEqual( + new LoggingContentGenerator( + (mockGenerator as GoogleGenAI).models, + mockConfig, + ), + ); }); }); diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index 797bad73..ac716ac3 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -18,6 +18,7 @@ import { DEFAULT_GEMINI_MODEL } from '../config/models.js'; import { Config } from '../config/config.js'; import { getEffectiveModel } from './modelCheck.js'; import { UserTierId } from '../code_assist/types.js'; +import { LoggingContentGenerator } from './loggingContentGenerator.js'; /** * Interface abstracting the core functionalities for generating content and counting tokens. @@ -121,11 +122,14 @@ export async function createContentGenerator( config.authType === AuthType.LOGIN_WITH_GOOGLE || config.authType === AuthType.CLOUD_SHELL ) { - return createCodeAssistContentGenerator( - httpOptions, - config.authType, + return new LoggingContentGenerator( + await createCodeAssistContentGenerator( + httpOptions, + config.authType, + gcConfig, + sessionId, + ), gcConfig, - sessionId, ); } @@ -138,10 +142,8 @@ export async function createContentGenerator( vertexai: config.vertexai, httpOptions, }); - - return googleGenAI.models; + return new LoggingContentGenerator(googleGenAI.models, gcConfig); } - throw new Error( `Error creating contentGenerator: Unsupported authType: ${config.authType}`, ); diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 085bc993..27b34d87 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -21,16 +21,8 @@ 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 { - logApiRequest, - logApiResponse, - logApiError, -} from '../telemetry/loggers.js'; -import { - ApiErrorEvent, - ApiRequestEvent, - ApiResponseEvent, -} from '../telemetry/types.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'; @@ -139,22 +131,6 @@ export class GeminiChat { validateHistory(history); } - private _getRequestTextFromContents(contents: Content[]): string { - return JSON.stringify(contents); - } - - private async _logApiRequest( - contents: Content[], - model: string, - prompt_id: string, - ): Promise { - const requestText = this._getRequestTextFromContents(contents); - logApiRequest( - this.config, - new ApiRequestEvent(model, prompt_id, requestText), - ); - } - private async _logApiResponse( durationMs: number, prompt_id: string, @@ -273,8 +249,6 @@ export class GeminiChat { const userContent = createUserContent(params.message); const requestContents = this.getHistory(true).concat(userContent); - this._logApiRequest(requestContents, this.config.getModel(), prompt_id); - const startTime = Date.now(); let response: GenerateContentResponse; @@ -386,7 +360,6 @@ export class GeminiChat { await this.sendPromise; const userContent = createUserContent(params.message); const requestContents = this.getHistory(true).concat(userContent); - this._logApiRequest(requestContents, this.config.getModel(), prompt_id); const startTime = Date.now(); diff --git a/packages/core/src/core/loggingContentGenerator.ts b/packages/core/src/core/loggingContentGenerator.ts new file mode 100644 index 00000000..c9069bab --- /dev/null +++ b/packages/core/src/core/loggingContentGenerator.ts @@ -0,0 +1,68 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + Content, + CountTokensParameters, + CountTokensResponse, + EmbedContentParameters, + EmbedContentResponse, + GenerateContentParameters, + GenerateContentResponse, +} from '@google/genai'; +import { ApiRequestEvent } from '../telemetry/types.js'; +import { Config } from '../config/config.js'; +import { logApiRequest } from '../telemetry/loggers.js'; +import { ContentGenerator } from './contentGenerator.js'; +import { toContents } from '../code_assist/converter.js'; + +/** + * A decorator that wraps a ContentGenerator to add logging to API calls. + */ +export class LoggingContentGenerator implements ContentGenerator { + constructor( + private readonly wrapped: ContentGenerator, + private readonly config: Config, + ) {} + + private logApiRequest( + contents: Content[], + model: string, + promptId: string, + ): void { + const requestText = JSON.stringify(contents); + logApiRequest( + this.config, + new ApiRequestEvent(model, promptId, requestText), + ); + } + + async generateContent( + req: GenerateContentParameters, + userPromptId: string, + ): Promise { + this.logApiRequest(toContents(req.contents), req.model, userPromptId); + return this.wrapped.generateContent(req, userPromptId); + } + + async generateContentStream( + req: GenerateContentParameters, + userPromptId: string, + ): Promise> { + this.logApiRequest(toContents(req.contents), req.model, userPromptId); + return this.wrapped.generateContentStream(req, userPromptId); + } + + async countTokens(req: CountTokensParameters): Promise { + return this.wrapped.countTokens(req); + } + + async embedContent( + req: EmbedContentParameters, + ): Promise { + return this.wrapped.embedContent(req); + } +} diff --git a/packages/core/src/core/nonInteractiveToolExecutor.ts b/packages/core/src/core/nonInteractiveToolExecutor.ts index 43061f83..29d79152 100644 --- a/packages/core/src/core/nonInteractiveToolExecutor.ts +++ b/packages/core/src/core/nonInteractiveToolExecutor.ts @@ -14,7 +14,7 @@ import { } from '../index.js'; import { Config } from '../config/config.js'; import { convertToFunctionResponse } from './coreToolScheduler.js'; -import { ToolCallDecision } from '../telemetry/types.js'; +import { ToolCallDecision } from '../telemetry/tool-call-decision.js'; /** * Executes a single tool call non-interactively. diff --git a/packages/core/src/telemetry/loggers.test.ts b/packages/core/src/telemetry/loggers.test.ts index 14de83a9..e1c4e65b 100644 --- a/packages/core/src/telemetry/loggers.test.ts +++ b/packages/core/src/telemetry/loggers.test.ts @@ -35,11 +35,11 @@ import { logToolCall, logFlashFallback, } from './loggers.js'; +import { ToolCallDecision } from './tool-call-decision.js'; import { ApiRequestEvent, ApiResponseEvent, StartSessionEvent, - ToolCallDecision, ToolCallEvent, UserPromptEvent, FlashFallbackEvent, diff --git a/packages/core/src/telemetry/tool-call-decision.ts b/packages/core/src/telemetry/tool-call-decision.ts new file mode 100644 index 00000000..167df10a --- /dev/null +++ b/packages/core/src/telemetry/tool-call-decision.ts @@ -0,0 +1,32 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { ToolConfirmationOutcome } from '../tools/tools.js'; + +export enum ToolCallDecision { + ACCEPT = 'accept', + REJECT = 'reject', + MODIFY = 'modify', + AUTO_ACCEPT = 'auto_accept', +} + +export function getDecisionFromOutcome( + outcome: ToolConfirmationOutcome, +): ToolCallDecision { + switch (outcome) { + case ToolConfirmationOutcome.ProceedOnce: + return ToolCallDecision.ACCEPT; + case ToolConfirmationOutcome.ProceedAlways: + case ToolConfirmationOutcome.ProceedAlwaysServer: + case ToolConfirmationOutcome.ProceedAlwaysTool: + return ToolCallDecision.AUTO_ACCEPT; + case ToolConfirmationOutcome.ModifyWithEditor: + return ToolCallDecision.MODIFY; + case ToolConfirmationOutcome.Cancel: + default: + return ToolCallDecision.REJECT; + } +} diff --git a/packages/core/src/telemetry/types.ts b/packages/core/src/telemetry/types.ts index db84e2da..766d5f47 100644 --- a/packages/core/src/telemetry/types.ts +++ b/packages/core/src/telemetry/types.ts @@ -7,33 +7,11 @@ import { GenerateContentResponseUsageMetadata } from '@google/genai'; import { Config } from '../config/config.js'; import { CompletedToolCall } from '../core/coreToolScheduler.js'; -import { ToolConfirmationOutcome } from '../tools/tools.js'; import { AuthType } from '../core/contentGenerator.js'; - -export enum ToolCallDecision { - ACCEPT = 'accept', - REJECT = 'reject', - MODIFY = 'modify', - AUTO_ACCEPT = 'auto_accept', -} - -export function getDecisionFromOutcome( - outcome: ToolConfirmationOutcome, -): ToolCallDecision { - switch (outcome) { - case ToolConfirmationOutcome.ProceedOnce: - return ToolCallDecision.ACCEPT; - case ToolConfirmationOutcome.ProceedAlways: - case ToolConfirmationOutcome.ProceedAlwaysServer: - case ToolConfirmationOutcome.ProceedAlwaysTool: - return ToolCallDecision.AUTO_ACCEPT; - case ToolConfirmationOutcome.ModifyWithEditor: - return ToolCallDecision.MODIFY; - case ToolConfirmationOutcome.Cancel: - default: - return ToolCallDecision.REJECT; - } -} +import { + getDecisionFromOutcome, + ToolCallDecision, +} from './tool-call-decision.js'; export class StartSessionEvent { 'event.name': 'cli_config'; diff --git a/packages/core/src/telemetry/uiTelemetry.test.ts b/packages/core/src/telemetry/uiTelemetry.test.ts index ac9727f1..cd509a8e 100644 --- a/packages/core/src/telemetry/uiTelemetry.test.ts +++ b/packages/core/src/telemetry/uiTelemetry.test.ts @@ -6,12 +6,8 @@ import { describe, it, expect, vi, beforeEach } from 'vitest'; import { UiTelemetryService } from './uiTelemetry.js'; -import { - ApiErrorEvent, - ApiResponseEvent, - ToolCallEvent, - ToolCallDecision, -} from './types.js'; +import { ToolCallDecision } from './tool-call-decision.js'; +import { ApiErrorEvent, ApiResponseEvent, ToolCallEvent } from './types.js'; import { EVENT_API_ERROR, EVENT_API_RESPONSE, diff --git a/packages/core/src/telemetry/uiTelemetry.ts b/packages/core/src/telemetry/uiTelemetry.ts index fbf5b8dc..8d1b044f 100644 --- a/packages/core/src/telemetry/uiTelemetry.ts +++ b/packages/core/src/telemetry/uiTelemetry.ts @@ -11,12 +11,8 @@ import { EVENT_TOOL_CALL, } from './constants.js'; -import { - ApiErrorEvent, - ApiResponseEvent, - ToolCallEvent, - ToolCallDecision, -} from './types.js'; +import { ToolCallDecision } from './tool-call-decision.js'; +import { ApiErrorEvent, ApiResponseEvent, ToolCallEvent } from './types.js'; export type UiEvent = | (ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE })