Introduce ContentGeneratorConfig (#826)

9 files changed, 56 insertions, 52 deletions

diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts
index 0ce8f5ee..1c1996bf 100644
--- a/packages/core/src/config/config.test.ts
+++ b/packages/core/src/config/config.test.ts
@@ -50,15 +50,17 @@ describe('Server Config (config.ts)', () => {
   const TELEMETRY = false;
   const EMBEDDING_MODEL = 'gemini-embedding';
   const baseParams: ConfigParameters = {
-    apiKey: API_KEY,
-    model: MODEL,
+    contentGeneratorConfig: {
+      apiKey: API_KEY,
+      model: MODEL,
+      userAgent: USER_AGENT,
+    },
     embeddingModel: EMBEDDING_MODEL,
     sandbox: SANDBOX,
     targetDir: TARGET_DIR,
     debugMode: DEBUG_MODE,
     question: QUESTION,
     fullContext: FULL_CONTEXT,
-    userAgent: USER_AGENT,
     userMemory: USER_MEMORY,
     telemetry: TELEMETRY,
   };
@@ -73,10 +75,7 @@ describe('Server Config (config.ts)', () => {
 
     expect(config.getUserMemory()).toBe(USER_MEMORY);
     // Verify other getters if needed
-    expect(config.getApiKey()).toBe(API_KEY);
-    expect(config.getModel()).toBe(MODEL);
     expect(config.getTargetDir()).toBe(path.resolve(TARGET_DIR)); // Check resolved path
-    expect(config.getUserAgent()).toBe(USER_AGENT);
   });
 
   it('Config constructor should default userMemory to empty string if not provided', () => {
diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts
index 7d25c1bb..b8f4d41f 100644
--- a/packages/core/src/config/config.ts
+++ b/packages/core/src/config/config.ts
@@ -9,6 +9,7 @@ import * as fs from 'node:fs';
 import * as path from 'node:path';
 import process from 'node:process';
 import * as os from 'node:os';
+import { ContentGeneratorConfig } from '../core/contentGenerator.js';
 import { ToolRegistry } from '../tools/tool-registry.js';
 import { CodeParserTool } from '../tools/code_parser.js'; // Added CodeParserTool
 import { LSTool } from '../tools/ls.js';
@@ -17,7 +18,6 @@ import { GrepTool } from '../tools/grep.js';
 import { GlobTool } from '../tools/glob.js';
 import { EditTool } from '../tools/edit.js';
 import { ShellTool } from '../tools/shell.js';
-
 import { WebFetchTool } from '../tools/web-fetch.js';
 import { ReadManyFilesTool } from '../tools/read-many-files.js';
 import { MemoryTool, setGeminiMdFilename } from '../tools/memoryTool.js';
@@ -55,8 +55,7 @@ export class MCPServerConfig {
 }
 
 export interface ConfigParameters {
-  apiKey: string;
-  model: string;
+  contentGeneratorConfig: ContentGeneratorConfig;
   embeddingModel: string;
   sandbox: boolean | string;
   targetDir: string;
@@ -68,11 +67,9 @@ export interface ConfigParameters {
   toolCallCommand?: string;
   mcpServerCommand?: string;
   mcpServers?: Record<string, MCPServerConfig>;
-  userAgent: string;
   userMemory?: string;
   geminiMdFileCount?: number;
   approvalMode?: ApprovalMode;
-  vertexai?: boolean;
   showMemoryUsage?: boolean;
   contextFileName?: string;
   geminiIgnorePatterns?: string[];
@@ -85,8 +82,7 @@ export interface ConfigParameters {
 
 export class Config {
   private toolRegistry: Promise<ToolRegistry>;
-  private readonly apiKey: string;
-  private readonly model: string;
+  private readonly contentGeneratorConfig: ContentGeneratorConfig;
   private readonly embeddingModel: string;
   private readonly sandbox: boolean | string;
   private readonly targetDir: string;
@@ -98,11 +94,9 @@ export class Config {
   private readonly toolCallCommand: string | undefined;
   private readonly mcpServerCommand: string | undefined;
   private readonly mcpServers: Record<string, MCPServerConfig> | undefined;
-  private readonly userAgent: string;
   private userMemory: string;
   private geminiMdFileCount: number;
   private approvalMode: ApprovalMode;
-  private readonly vertexai: boolean | undefined;
   private readonly showMemoryUsage: boolean;
   private readonly accessibility: AccessibilitySettings;
   private readonly telemetry: boolean;
@@ -115,8 +109,7 @@ export class Config {
   private fileDiscoveryService: FileDiscoveryService | null = null;
 
   constructor(params: ConfigParameters) {
-    this.apiKey = params.apiKey;
-    this.model = params.model;
+    this.contentGeneratorConfig = params.contentGeneratorConfig;
     this.embeddingModel = params.embeddingModel;
     this.sandbox = params.sandbox;
     this.targetDir = path.resolve(params.targetDir);
@@ -128,11 +121,9 @@ export class Config {
     this.toolCallCommand = params.toolCallCommand;
     this.mcpServerCommand = params.mcpServerCommand;
     this.mcpServers = params.mcpServers;
-    this.userAgent = params.userAgent ?? 'GeminiCLI/unknown';
     this.userMemory = params.userMemory ?? '';
     this.geminiMdFileCount = params.geminiMdFileCount ?? 0;
     this.approvalMode = params.approvalMode ?? ApprovalMode.DEFAULT;
-    this.vertexai = params.vertexai;
     this.showMemoryUsage = params.showMemoryUsage ?? false;
     this.accessibility = params.accessibility ?? {};
     this.telemetry = params.telemetry ?? false;
@@ -160,12 +151,12 @@ export class Config {
     }
   }
 
-  getApiKey(): string {
-    return this.apiKey;
+  getContentGeneratorConfig(): ContentGeneratorConfig {
+    return this.contentGeneratorConfig;
   }
 
   getModel(): string {
-    return this.model;
+    return this.contentGeneratorConfig.model;
   }
 
   getEmbeddingModel(): string {
@@ -215,10 +206,6 @@ export class Config {
     return this.mcpServers;
   }
 
-  getUserAgent(): string {
-    return this.userAgent;
-  }
-
   getUserMemory(): string {
     return this.userMemory;
   }
@@ -243,10 +230,6 @@ export class Config {
     this.approvalMode = mode;
   }
 
-  getVertexAI(): boolean | undefined {
-    return this.vertexai;
-  }
-
   getShowMemoryUsage(): boolean {
     return this.showMemoryUsage;
   }
diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts
index 9c12423c..180d74bb 100644
--- a/packages/core/src/core/client.test.ts
+++ b/packages/core/src/core/client.test.ts
@@ -77,6 +77,11 @@ describe('Gemini Client (client.ts)', () => {
     const MockedConfig = vi.mocked(Config, true);
     MockedConfig.mockImplementation(() => {
       const mock = {
+        getContentGeneratorConfig: vi.fn().mockReturnValue({
+          model: 'test-model',
+          apiKey: 'test-key',
+          vertexai: false,
+        }),
         getToolRegistry: vi.fn().mockResolvedValue(mockToolRegistry),
         getModel: vi.fn().mockReturnValue('test-model'),
         getEmbeddingModel: vi.fn().mockReturnValue('test-embedding-model'),
diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts
index c4515f93..0f2a1b8a 100644
--- a/packages/core/src/core/client.ts
+++ b/packages/core/src/core/client.ts
@@ -8,7 +8,6 @@ import {
   EmbedContentResponse,
   EmbedContentParameters,
   GenerateContentConfig,
-  GoogleGenAI,
   Part,
   SchemaUnion,
   PartListUnion,
@@ -34,7 +33,10 @@ import {
   logApiResponse,
   logApiError,
 } from '../telemetry/index.js';
-import { ContentGenerator } from './contentGenerator.js';
+import {
+  ContentGenerator,
+  createContentGenerator,
+} from './contentGenerator.js';
 
 export class GeminiClient {
   private chat: Promise<GeminiChat>;
@@ -48,20 +50,9 @@ export class GeminiClient {
   private readonly MAX_TURNS = 100;
 
   constructor(private config: Config) {
-    const userAgent = config.getUserAgent();
-    const apiKeyFromConfig = config.getApiKey();
-    const vertexaiFlag = config.getVertexAI();
-
-    const googleGenAI = new GoogleGenAI({
-      apiKey: apiKeyFromConfig === '' ? undefined : apiKeyFromConfig,
-      vertexai: vertexaiFlag,
-      httpOptions: {
-        headers: {
-          'User-Agent': userAgent,
-        },
-      },
-    });
-    this.contentGenerator = googleGenAI.models;
+    this.contentGenerator = createContentGenerator(
+      this.config.getContentGeneratorConfig(),
+    );
     this.model = config.getModel();
     this.embeddingModel = config.getEmbeddingModel();
     this.chat = this.startChat();
diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts
index 955cd152..f9db121b 100644
--- a/packages/core/src/core/contentGenerator.ts
+++ b/packages/core/src/core/contentGenerator.ts
@@ -11,6 +11,7 @@ import {
   CountTokensParameters,
   EmbedContentResponse,
   EmbedContentParameters,
+  GoogleGenAI,
 } from '@google/genai';
 
 /**
@@ -29,3 +30,25 @@ export interface ContentGenerator {
 
   embedContent(request: EmbedContentParameters): Promise<EmbedContentResponse>;
 }
+
+export type ContentGeneratorConfig = {
+  model: string;
+  apiKey?: string;
+  vertexai?: boolean;
+  userAgent: string;
+};
+
+export function createContentGenerator(
+  config: ContentGeneratorConfig,
+): ContentGenerator {
+  const googleGenAI = new GoogleGenAI({
+    apiKey: config.apiKey === '' ? undefined : config.apiKey,
+    vertexai: config.vertexai,
+    httpOptions: {
+      headers: {
+        'User-Agent': config.userAgent,
+      },
+    },
+  });
+  return googleGenAI.models;
+}
diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts
index 1f58d97c..f69ac6fd 100644
--- a/packages/core/src/index.ts
+++ b/packages/core/src/index.ts
@@ -9,6 +9,7 @@ export * from './config/config.js';
 
 // Export Core Logic
 export * from './core/client.js';
+export * from './core/contentGenerator.js';
 export * from './core/geminiChat.js';
 export * from './core/logger.js';
 export * from './core/prompts.js';
diff --git a/packages/core/src/telemetry/loggers.ts b/packages/core/src/telemetry/loggers.ts
index ccab12ff..e8a6156c 100644
--- a/packages/core/src/telemetry/loggers.ts
+++ b/packages/core/src/telemetry/loggers.ts
@@ -45,7 +45,7 @@ export function logCliConfiguration(config: Config): void {
       typeof config.getSandbox() === 'string' ? true : config.getSandbox(),
     core_tools_enabled: (config.getCoreTools() ?? []).join(','),
     approval_mode: config.getApprovalMode(),
-    vertex_ai_enabled: config.getVertexAI() ?? false,
+    vertex_ai_enabled: !!config.getContentGeneratorConfig().vertexai,
     log_user_prompts_enabled: config.getTelemetryLogUserPromptsEnabled(),
     file_filtering_respect_git_ignore:
       config.getFileFilteringRespectGitIgnore(),
diff --git a/packages/core/src/telemetry/sdk.ts b/packages/core/src/telemetry/sdk.ts
index b55cb149..ad5a9c44 100644
--- a/packages/core/src/telemetry/sdk.ts
+++ b/packages/core/src/telemetry/sdk.ts
@@ -64,7 +64,7 @@ export function initializeTelemetry(config: Config): void {
     return;
   }
 
-  const geminiCliVersion = config.getUserAgent() || 'unknown';
+  const geminiCliVersion = config.getContentGeneratorConfig().userAgent;
   const resource = new Resource({
     [SemanticResourceAttributes.SERVICE_NAME]: SERVICE_NAME,
     [SemanticResourceAttributes.SERVICE_VERSION]: geminiCliVersion,
diff --git a/packages/core/src/tools/tool-registry.test.ts b/packages/core/src/tools/tool-registry.test.ts
index 0c23a74e..44e04a9d 100644
--- a/packages/core/src/tools/tool-registry.test.ts
+++ b/packages/core/src/tools/tool-registry.test.ts
@@ -124,17 +124,19 @@ class MockTool extends BaseTool<{ param: string }, ToolResult> {
 }
 
 const baseConfigParams: ConfigParameters = {
-  apiKey: 'test-api-key',
-  model: 'test-model',
+  contentGeneratorConfig: {
+    model: 'test-model',
+    apiKey: 'test-api-key',
+    vertexai: false,
+    userAgent: 'TestAgent/1.0',
+  },
   embeddingModel: 'test-embedding-model',
   sandbox: false,
   targetDir: '/test/dir',
   debugMode: false,
-  userAgent: 'TestAgent/1.0',
   userMemory: '',
   geminiMdFileCount: 0,
   approvalMode: ApprovalMode.DEFAULT,
-  vertexai: false,
 };
 
 describe('ToolRegistry', () => {