summaryrefslogtreecommitdiff
path: root/packages/core/src
diff options
context:
space:
mode:
Diffstat (limited to 'packages/core/src')
-rw-r--r--packages/core/src/core/client.ts8
-rw-r--r--packages/core/src/core/geminiChat.test.ts12
-rw-r--r--packages/core/src/core/geminiChat.ts12
-rw-r--r--packages/core/src/core/nonInteractiveToolExecutor.test.ts8
-rw-r--r--packages/core/src/core/nonInteractiveToolExecutor.ts2
-rw-r--r--packages/core/src/telemetry/loggers.ts49
-rw-r--r--packages/core/src/telemetry/metrics.test.ts31
-rw-r--r--packages/core/src/telemetry/metrics.ts35
-rw-r--r--packages/core/src/telemetry/sdk.ts2
-rw-r--r--packages/core/src/utils/nextSpeakerChecker.test.ts2
10 files changed, 124 insertions, 37 deletions
diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts
index 349caf1f..a2deca4e 100644
--- a/packages/core/src/core/client.ts
+++ b/packages/core/src/core/client.ts
@@ -160,9 +160,9 @@ export class GeminiClient {
const systemInstruction = getCoreSystemPrompt(userMemory);
return new GeminiChat(
+ this.config,
await this.contentGenerator,
this.model,
- this.config.getSessionId(),
{
systemInstruction,
...this.generateContentConfig,
@@ -214,7 +214,7 @@ export class GeminiClient {
}
private _logApiRequest(model: string, inputTokenCount: number): void {
- logApiRequest({
+ logApiRequest(this.config, {
model,
input_token_count: inputTokenCount,
duration_ms: 0, // Duration is not known at request time
@@ -239,7 +239,7 @@ export class GeminiClient {
responseError = `Finished with reason: ${finishReason}`;
}
- logApiResponse({
+ logApiResponse(this.config, {
model,
duration_ms: durationMs,
attempt,
@@ -277,7 +277,7 @@ export class GeminiClient {
}
}
- logApiError({
+ logApiError(this.config, {
model,
error: errorMessage,
status_code: statusCode,
diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts
index 49341939..03e933cc 100644
--- a/packages/core/src/core/geminiChat.test.ts
+++ b/packages/core/src/core/geminiChat.test.ts
@@ -7,6 +7,7 @@
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { Content, Models, GenerateContentConfig, Part } from '@google/genai';
import { GeminiChat } from './geminiChat.js';
+import { Config } from '../config/config.js';
// Mocks
const mockModelsModule = {
@@ -17,16 +18,19 @@ const mockModelsModule = {
batchEmbedContents: vi.fn(),
} as unknown as Models;
+const mockConfig = {
+ getSessionId: () => 'test-session-id',
+} as unknown as Config;
+
describe('GeminiChat', () => {
let chat: GeminiChat;
const model = 'gemini-pro';
const config: GenerateContentConfig = {};
- const sessionId = 'test-session-id';
beforeEach(() => {
vi.clearAllMocks();
// Reset history for each test by creating a new instance
- chat = new GeminiChat(mockModelsModule, model, sessionId, config, []);
+ chat = new GeminiChat(mockConfig, mockModelsModule, model, config, []);
});
afterEach(() => {
@@ -121,7 +125,7 @@ describe('GeminiChat', () => {
chat.recordHistory(userInput, newModelOutput); // userInput here is for the *next* turn, but history is already primed
// Reset and set up a more realistic scenario for merging with existing history
- chat = new GeminiChat(mockModelsModule, model, sessionId, config, []);
+ chat = new GeminiChat(mockConfig, mockModelsModule, model, config, []);
const firstUserInput: Content = {
role: 'user',
parts: [{ text: 'First user input' }],
@@ -164,7 +168,7 @@ describe('GeminiChat', () => {
role: 'model',
parts: [{ text: 'Initial model answer.' }],
};
- chat = new GeminiChat(mockModelsModule, model, sessionId, config, [
+ chat = new GeminiChat(mockConfig, mockModelsModule, model, config, [
initialUser,
initialModel,
]);
diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts
index a78a7884..2a81aca8 100644
--- a/packages/core/src/core/geminiChat.ts
+++ b/packages/core/src/core/geminiChat.ts
@@ -18,7 +18,7 @@ import {
import { retryWithBackoff } from '../utils/retry.js';
import { isFunctionResponse } from '../utils/messageInspectors.js';
import { ContentGenerator } from './contentGenerator.js';
-import { Logger } from './logger.js';
+import { Config } from '../config/config.js';
/**
* Returns true if the response is valid, false otherwise.
@@ -118,17 +118,15 @@ export class GeminiChat {
// A promise to represent the current state of the message being sent to the
// model.
private sendPromise: Promise<void> = Promise.resolve();
- private logger: Logger;
constructor(
+ private readonly config: Config,
private readonly contentGenerator: ContentGenerator,
private readonly model: string,
- sessionId: string,
- private readonly config: GenerateContentConfig = {},
+ private readonly generationConfig: GenerateContentConfig = {},
private history: Content[] = [],
) {
validateHistory(history);
- this.logger = new Logger(sessionId);
}
/**
@@ -161,7 +159,7 @@ export class GeminiChat {
this.contentGenerator.generateContent({
model: this.model,
contents: this.getHistory(true).concat(userContent),
- config: { ...this.config, ...params.config },
+ config: { ...this.generationConfig, ...params.config },
});
const responsePromise = retryWithBackoff(apiCall);
@@ -230,7 +228,7 @@ export class GeminiChat {
this.contentGenerator.generateContentStream({
model: this.model,
contents: this.getHistory(true).concat(userContent),
- config: { ...this.config, ...params.config },
+ config: { ...this.generationConfig, ...params.config },
});
// Note: Retrying streams can be complex. If generateContentStream itself doesn't handle retries
diff --git a/packages/core/src/core/nonInteractiveToolExecutor.test.ts b/packages/core/src/core/nonInteractiveToolExecutor.test.ts
index be9294a8..07946af5 100644
--- a/packages/core/src/core/nonInteractiveToolExecutor.test.ts
+++ b/packages/core/src/core/nonInteractiveToolExecutor.test.ts
@@ -12,9 +12,12 @@ import {
ToolResult,
Tool,
ToolCallConfirmationDetails,
+ Config,
} from '../index.js';
import { Part, Type } from '@google/genai';
+const mockConfig = {} as unknown as Config;
+
describe('executeToolCall', () => {
let mockToolRegistry: ToolRegistry;
let mockTool: Tool;
@@ -68,6 +71,7 @@ describe('executeToolCall', () => {
vi.mocked(mockTool.execute).mockResolvedValue(toolResult);
const response = await executeToolCall(
+ mockConfig,
request,
mockToolRegistry,
abortController.signal,
@@ -99,6 +103,7 @@ describe('executeToolCall', () => {
vi.mocked(mockToolRegistry.getTool).mockReturnValue(undefined);
const response = await executeToolCall(
+ mockConfig,
request,
mockToolRegistry,
abortController.signal,
@@ -134,6 +139,7 @@ describe('executeToolCall', () => {
vi.mocked(mockTool.execute).mockRejectedValue(executionError);
const response = await executeToolCall(
+ mockConfig,
request,
mockToolRegistry,
abortController.signal,
@@ -184,6 +190,7 @@ describe('executeToolCall', () => {
abortController.abort(); // Abort before calling
const response = await executeToolCall(
+ mockConfig,
request,
mockToolRegistry,
abortController.signal,
@@ -211,6 +218,7 @@ describe('executeToolCall', () => {
vi.mocked(mockTool.execute).mockResolvedValue(toolResult);
const response = await executeToolCall(
+ mockConfig,
request,
mockToolRegistry,
abortController.signal,
diff --git a/packages/core/src/core/nonInteractiveToolExecutor.ts b/packages/core/src/core/nonInteractiveToolExecutor.ts
index 87d17c2b..5f7ed4ff 100644
--- a/packages/core/src/core/nonInteractiveToolExecutor.ts
+++ b/packages/core/src/core/nonInteractiveToolExecutor.ts
@@ -10,6 +10,7 @@ import {
ToolRegistry,
ToolResult,
} from '../index.js';
+import { Config } from '../config/config.js';
import { convertToFunctionResponse } from './coreToolScheduler.js';
/**
@@ -17,6 +18,7 @@ import { convertToFunctionResponse } from './coreToolScheduler.js';
* It does not handle confirmations, multiple calls, or live updates.
*/
export async function executeToolCall(
+ config: Config,
toolCallRequest: ToolCallRequestInfo,
toolRegistry: ToolRegistry,
abortSignal?: AbortSignal,
diff --git a/packages/core/src/telemetry/loggers.ts b/packages/core/src/telemetry/loggers.ts
index 66be0fca..d2b01f65 100644
--- a/packages/core/src/telemetry/loggers.ts
+++ b/packages/core/src/telemetry/loggers.ts
@@ -34,10 +34,17 @@ import { isTelemetrySdkInitialized } from './sdk.js';
const shouldLogUserPrompts = (config: Config): boolean =>
config.getTelemetryLogUserPromptsEnabled() ?? false;
+function getCommonAttributes(config: Config): LogAttributes {
+ return {
+ 'session.id': config.getSessionId(),
+ };
+}
+
export function logCliConfiguration(config: Config): void {
if (!isTelemetrySdkInitialized()) return;
const attributes: LogAttributes = {
+ ...getCommonAttributes(config),
'event.name': EVENT_CLI_CONFIG,
'event.timestamp': new Date().toISOString(),
model: config.getModel(),
@@ -69,6 +76,7 @@ export function logUserPrompt(
if (!isTelemetrySdkInitialized()) return;
const { prompt, ...restOfEventArgs } = event;
const attributes: LogAttributes = {
+ ...getCommonAttributes(config),
...restOfEventArgs,
'event.name': EVENT_USER_PROMPT,
'event.timestamp': new Date().toISOString(),
@@ -85,10 +93,12 @@ export function logUserPrompt(
}
export function logToolCall(
+ config: Config,
event: Omit<ToolCallEvent, 'event.name' | 'event.timestamp'>,
): void {
if (!isTelemetrySdkInitialized()) return;
const attributes: LogAttributes = {
+ ...getCommonAttributes(config),
...event,
'event.name': EVENT_TOOL_CALL,
'event.timestamp': new Date().toISOString(),
@@ -106,14 +116,21 @@ export function logToolCall(
attributes,
};
logger.emit(logRecord);
- recordToolCallMetrics(event.function_name, event.duration_ms, event.success);
+ recordToolCallMetrics(
+ config,
+ event.function_name,
+ event.duration_ms,
+ event.success,
+ );
}
export function logApiRequest(
+ config: Config,
event: Omit<ApiRequestEvent, 'event.name' | 'event.timestamp'>,
): void {
if (!isTelemetrySdkInitialized()) return;
const attributes: LogAttributes = {
+ ...getCommonAttributes(config),
...event,
'event.name': EVENT_API_REQUEST,
'event.timestamp': new Date().toISOString(),
@@ -124,14 +141,21 @@ export function logApiRequest(
attributes,
};
logger.emit(logRecord);
- recordTokenUsageMetrics(event.model, event.input_token_count, 'input');
+ recordTokenUsageMetrics(
+ config,
+ event.model,
+ event.input_token_count,
+ 'input',
+ );
}
export function logApiError(
+ config: Config,
event: Omit<ApiErrorEvent, 'event.name' | 'event.timestamp'>,
): void {
if (!isTelemetrySdkInitialized()) return;
const attributes: LogAttributes = {
+ ...getCommonAttributes(config),
...event,
'event.name': EVENT_API_ERROR,
'event.timestamp': new Date().toISOString(),
@@ -152,6 +176,7 @@ export function logApiError(
};
logger.emit(logRecord);
recordApiErrorMetrics(
+ config,
event.model,
event.duration_ms,
event.status_code,
@@ -160,10 +185,12 @@ export function logApiError(
}
export function logApiResponse(
+ config: Config,
event: Omit<ApiResponseEvent, 'event.name' | 'event.timestamp'>,
): void {
if (!isTelemetrySdkInitialized()) return;
const attributes: LogAttributes = {
+ ...getCommonAttributes(config),
...event,
'event.name': EVENT_API_RESPONSE,
'event.timestamp': new Date().toISOString(),
@@ -183,17 +210,29 @@ export function logApiResponse(
};
logger.emit(logRecord);
recordApiResponseMetrics(
+ config,
event.model,
event.duration_ms,
event.status_code,
event.error,
);
- recordTokenUsageMetrics(event.model, event.output_token_count, 'output');
recordTokenUsageMetrics(
+ config,
+ event.model,
+ event.output_token_count,
+ 'output',
+ );
+ recordTokenUsageMetrics(
+ config,
event.model,
event.cached_content_token_count,
'cache',
);
- recordTokenUsageMetrics(event.model, event.thoughts_token_count, 'thought');
- recordTokenUsageMetrics(event.model, event.tool_token_count, 'tool');
+ recordTokenUsageMetrics(
+ config,
+ event.model,
+ event.thoughts_token_count,
+ 'thought',
+ );
+ recordTokenUsageMetrics(config, event.model, event.tool_token_count, 'tool');
}
diff --git a/packages/core/src/telemetry/metrics.test.ts b/packages/core/src/telemetry/metrics.test.ts
index b7864131..7e24b9ad 100644
--- a/packages/core/src/telemetry/metrics.test.ts
+++ b/packages/core/src/telemetry/metrics.test.ts
@@ -7,6 +7,7 @@
import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest';
import { Counter, Meter, metrics } from '@opentelemetry/api';
import { initializeMetrics, recordTokenUsageMetrics } from './metrics.js';
+import { Config } from '../config/config.js';
const mockCounter = {
add: vi.fn(),
@@ -33,51 +34,61 @@ describe('Telemetry Metrics', () => {
});
describe('recordTokenUsageMetrics', () => {
+ const mockConfig = {
+ getSessionId: () => 'test-session-id',
+ } as unknown as Config;
+
it('should not record metrics if not initialized', () => {
- recordTokenUsageMetrics('gemini-pro', 100, 'input');
+ recordTokenUsageMetrics(mockConfig, 'gemini-pro', 100, 'input');
expect(mockCounter.add).not.toHaveBeenCalled();
});
it('should record token usage with the correct attributes', () => {
- initializeMetrics();
- recordTokenUsageMetrics('gemini-pro', 100, 'input');
+ initializeMetrics(mockConfig);
+ recordTokenUsageMetrics(mockConfig, 'gemini-pro', 100, 'input');
expect(mockCounter.add).toHaveBeenCalledWith(100, {
+ 'session.id': 'test-session-id',
model: 'gemini-pro',
type: 'input',
});
});
it('should record token usage for different types', () => {
- initializeMetrics();
- recordTokenUsageMetrics('gemini-pro', 50, 'output');
+ initializeMetrics(mockConfig);
+ recordTokenUsageMetrics(mockConfig, 'gemini-pro', 50, 'output');
expect(mockCounter.add).toHaveBeenCalledWith(50, {
+ 'session.id': 'test-session-id',
model: 'gemini-pro',
type: 'output',
});
- recordTokenUsageMetrics('gemini-pro', 25, 'thought');
+ recordTokenUsageMetrics(mockConfig, 'gemini-pro', 25, 'thought');
expect(mockCounter.add).toHaveBeenCalledWith(25, {
+ 'session.id': 'test-session-id',
model: 'gemini-pro',
type: 'thought',
});
- recordTokenUsageMetrics('gemini-pro', 75, 'cache');
+ recordTokenUsageMetrics(mockConfig, 'gemini-pro', 75, 'cache');
expect(mockCounter.add).toHaveBeenCalledWith(75, {
+ 'session.id': 'test-session-id',
model: 'gemini-pro',
type: 'cache',
});
- recordTokenUsageMetrics('gemini-pro', 125, 'tool');
+ recordTokenUsageMetrics(mockConfig, 'gemini-pro', 125, 'tool');
expect(mockCounter.add).toHaveBeenCalledWith(125, {
+ 'session.id': 'test-session-id',
model: 'gemini-pro',
type: 'tool',
});
});
it('should handle different models', () => {
- initializeMetrics();
- recordTokenUsageMetrics('gemini-ultra', 200, 'input');
+ initializeMetrics(mockConfig);
+ recordTokenUsageMetrics(mockConfig, 'gemini-ultra', 200, 'input');
expect(mockCounter.add).toHaveBeenCalledWith(200, {
+ 'session.id': 'test-session-id',
model: 'gemini-ultra',
type: 'input',
});
diff --git a/packages/core/src/telemetry/metrics.ts b/packages/core/src/telemetry/metrics.ts
index 69ac91e9..93aa2189 100644
--- a/packages/core/src/telemetry/metrics.ts
+++ b/packages/core/src/telemetry/metrics.ts
@@ -21,6 +21,7 @@ import {
METRIC_TOKEN_USAGE,
METRIC_SESSION_COUNT,
} from './constants.js';
+import { Config } from '../config/config.js';
let cliMeter: Meter | undefined;
let toolCallCounter: Counter | undefined;
@@ -30,6 +31,12 @@ let apiRequestLatencyHistogram: Histogram | undefined;
let tokenUsageCounter: Counter | undefined;
let isMetricsInitialized = false;
+function getCommonAttributes(config: Config): Attributes {
+ return {
+ 'session.id': config.getSessionId(),
+ };
+}
+
export function getMeter(): Meter | undefined {
if (!cliMeter) {
cliMeter = metrics.getMeter(SERVICE_NAME);
@@ -37,7 +44,7 @@ export function getMeter(): Meter | undefined {
return cliMeter;
}
-export function initializeMetrics(): void {
+export function initializeMetrics(config: Config): void {
if (isMetricsInitialized) return;
const meter = getMeter();
@@ -73,11 +80,12 @@ export function initializeMetrics(): void {
description: 'Count of CLI sessions started.',
valueType: ValueType.INT,
});
- sessionCounter.add(1);
+ sessionCounter.add(1, getCommonAttributes(config));
isMetricsInitialized = true;
}
export function recordToolCallMetrics(
+ config: Config,
functionName: string,
durationMs: number,
success: boolean,
@@ -86,25 +94,33 @@ export function recordToolCallMetrics(
return;
const metricAttributes: Attributes = {
+ ...getCommonAttributes(config),
function_name: functionName,
success,
};
toolCallCounter.add(1, metricAttributes);
toolCallLatencyHistogram.record(durationMs, {
+ ...getCommonAttributes(config),
function_name: functionName,
});
}
export function recordTokenUsageMetrics(
+ config: Config,
model: string,
tokenCount: number,
type: 'input' | 'output' | 'thought' | 'cache' | 'tool',
): void {
if (!tokenUsageCounter || !isMetricsInitialized) return;
- tokenUsageCounter.add(tokenCount, { model, type });
+ tokenUsageCounter.add(tokenCount, {
+ ...getCommonAttributes(config),
+ model,
+ type,
+ });
}
export function recordApiResponseMetrics(
+ config: Config,
model: string,
durationMs: number,
statusCode?: number | string,
@@ -117,14 +133,19 @@ export function recordApiResponseMetrics(
)
return;
const metricAttributes: Attributes = {
+ ...getCommonAttributes(config),
model,
status_code: statusCode ?? (error ? 'error' : 'ok'),
};
apiRequestCounter.add(1, metricAttributes);
- apiRequestLatencyHistogram.record(durationMs, { model });
+ apiRequestLatencyHistogram.record(durationMs, {
+ ...getCommonAttributes(config),
+ model,
+ });
}
export function recordApiErrorMetrics(
+ config: Config,
model: string,
durationMs: number,
statusCode?: number | string,
@@ -137,10 +158,14 @@ export function recordApiErrorMetrics(
)
return;
const metricAttributes: Attributes = {
+ ...getCommonAttributes(config),
model,
status_code: statusCode ?? 'error',
error_type: errorType ?? 'unknown',
};
apiRequestCounter.add(1, metricAttributes);
- apiRequestLatencyHistogram.record(durationMs, { model });
+ apiRequestLatencyHistogram.record(durationMs, {
+ ...getCommonAttributes(config),
+ model,
+ });
}
diff --git a/packages/core/src/telemetry/sdk.ts b/packages/core/src/telemetry/sdk.ts
index 704661c5..61f501a6 100644
--- a/packages/core/src/telemetry/sdk.ts
+++ b/packages/core/src/telemetry/sdk.ts
@@ -112,7 +112,7 @@ export function initializeTelemetry(config: Config): void {
sdk.start();
console.log('OpenTelemetry SDK started successfully.');
telemetryInitialized = true;
- initializeMetrics();
+ initializeMetrics(config);
logCliConfiguration(config);
} catch (error) {
console.error('Error starting OpenTelemetry SDK:', error);
diff --git a/packages/core/src/utils/nextSpeakerChecker.test.ts b/packages/core/src/utils/nextSpeakerChecker.test.ts
index a19045d1..83ce97fd 100644
--- a/packages/core/src/utils/nextSpeakerChecker.test.ts
+++ b/packages/core/src/utils/nextSpeakerChecker.test.ts
@@ -69,9 +69,9 @@ describe('checkNextSpeaker', () => {
// GeminiChat will receive the mocked instances via the mocked GoogleGenAI constructor
chatInstance = new GeminiChat(
+ mockConfigInstance,
mockModelsInstance, // This is the instance returned by mockGoogleGenAIInstance.getGenerativeModel
'gemini-pro', // model name
- 'test-session-id',
{},
[], // initial history
);