fix(cli) - Move logging into CodeAssistServer (#5781)
Co-authored-by: Shi Shu <shii@google.com>
This commit is contained in:
parent
60362e0329
commit
bae922a632
|
@ -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' }] },
|
||||
]);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
@ -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);
|
||||
|
|
|
@ -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(
|
||||
|
|
|
@ -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,
|
||||
),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
|
|
|
@ -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}`,
|
||||
);
|
||||
|
|
|
@ -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();
|
||||
|
||||
|
|
|
@ -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);
|
||||
}
|
||||
}
|
|
@ -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.
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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;
|
||||
}
|
||||
}
|
|
@ -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';
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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 })
|
||||
|
|
Loading…
Reference in New Issue