summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
Diffstat
-rw-r--r--packages/core/src/code_assist/converter.test.ts55
-rw-r--r--packages/core/src/code_assist/converter.ts2
-rw-r--r--packages/core/src/code_assist/server.test.ts6
-rw-r--r--packages/core/src/core/contentGenerator.test.ts17
-rw-r--r--packages/core/src/core/contentGenerator.ts16
-rw-r--r--packages/core/src/core/geminiChat.ts31
-rw-r--r--packages/core/src/core/loggingContentGenerator.ts68
-rw-r--r--packages/core/src/core/nonInteractiveToolExecutor.ts2
-rw-r--r--packages/core/src/telemetry/loggers.test.ts2
-rw-r--r--packages/core/src/telemetry/tool-call-decision.ts32
-rw-r--r--packages/core/src/telemetry/types.ts30
-rw-r--r--packages/core/src/telemetry/uiTelemetry.test.ts8
-rw-r--r--packages/core/src/telemetry/uiTelemetry.ts8
13 files changed, 195 insertions, 82 deletions
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<void> {
- 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<GenerateContentResponse> {
+ this.logApiRequest(toContents(req.contents), req.model, userPromptId);
+ return this.wrapped.generateContent(req, userPromptId);
+ }
+
+ async generateContentStream(
+ req: GenerateContentParameters,
+ userPromptId: string,
+ ): Promise<AsyncGenerator<GenerateContentResponse>> {
+ this.logApiRequest(toContents(req.contents), req.model, userPromptId);
+ return this.wrapped.generateContentStream(req, userPromptId);
+ }
+
+ async countTokens(req: CountTokensParameters): Promise<CountTokensResponse> {
+ return this.wrapped.countTokens(req);
+ }
+
+ async embedContent(
+ req: EmbedContentParameters,
+ ): Promise<EmbedContentResponse> {
+ 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 })
generated by cgit v1.2.3 (git 2.51.0) at 2025-11-05 16:43:10 +0000