diff --git a/docs/get-started/configuration.md b/docs/get-started/configuration.md index c8c5d31b09..fd44bb36f4 100644 --- a/docs/get-started/configuration.md +++ b/docs/get-started/configuration.md @@ -1185,6 +1185,10 @@ the `advanced.excludedEnvVars` setting in your `settings.json` file. - **Description:** The path to your Google Application Credentials JSON file. - **Example:** `export GOOGLE_APPLICATION_CREDENTIALS="/path/to/your/credentials.json"` +- **`GOOGLE_GENAI_API_VERSION`**: + - Specifies the API version to use for Gemini API requests. + - When set, overrides the default API version used by the SDK. + - Example: `export GOOGLE_GENAI_API_VERSION="v1"` - **`OTLP_GOOGLE_CLOUD_PROJECT`**: - Your Google Cloud Project ID for Telemetry in Google Cloud - Example: `export OTLP_GOOGLE_CLOUD_PROJECT="YOUR_PROJECT_ID"`. diff --git a/packages/core/src/core/contentGenerator.test.ts b/packages/core/src/core/contentGenerator.test.ts index 9e10558a18..f7c5a6d8d8 100644 --- a/packages/core/src/core/contentGenerator.test.ts +++ b/packages/core/src/core/contentGenerator.test.ts @@ -338,6 +338,152 @@ describe('createContentGenerator', () => { new LoggingContentGenerator(mockGenerator.models, mockConfig), ); }); + + it('should pass apiVersion to GoogleGenAI when GOOGLE_GENAI_API_VERSION is set', async () => { + const mockConfig = { + getModel: vi.fn().mockReturnValue('gemini-pro'), + getProxy: vi.fn().mockReturnValue(undefined), + getUsageStatisticsEnabled: () => false, + getPreviewFeatures: vi.fn().mockReturnValue(false), + } as unknown as Config; + + const mockGenerator = { + models: {}, + } as unknown as GoogleGenAI; + vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never); + vi.stubEnv('GOOGLE_GENAI_API_VERSION', 'v1'); + + await createContentGenerator( + { + apiKey: 'test-api-key', + authType: AuthType.USE_GEMINI, + }, + mockConfig, + ); + + expect(GoogleGenAI).toHaveBeenCalledWith({ + apiKey: 'test-api-key', + vertexai: undefined, + httpOptions: { + headers: expect.objectContaining({ + 'User-Agent': expect.any(String), + }), + }, + apiVersion: 'v1', + }); + }); + + it('should not include apiVersion when GOOGLE_GENAI_API_VERSION is not set', async () => { + const mockConfig = { + getModel: vi.fn().mockReturnValue('gemini-pro'), + getProxy: vi.fn().mockReturnValue(undefined), + getUsageStatisticsEnabled: () => false, + getPreviewFeatures: vi.fn().mockReturnValue(false), + } as unknown as Config; + + const mockGenerator = { + models: {}, + } as unknown as GoogleGenAI; + vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never); + + await createContentGenerator( + { + apiKey: 'test-api-key', + authType: AuthType.USE_GEMINI, + }, + mockConfig, + ); + + expect(GoogleGenAI).toHaveBeenCalledWith({ + apiKey: 'test-api-key', + vertexai: undefined, + httpOptions: { + headers: expect.objectContaining({ + 'User-Agent': expect.any(String), + }), + }, + }); + + expect(GoogleGenAI).toHaveBeenCalledWith( + expect.not.objectContaining({ + apiVersion: expect.any(String), + }), + ); + }); + + it('should not include apiVersion when GOOGLE_GENAI_API_VERSION is an empty string', async () => { + const mockConfig = { + getModel: vi.fn().mockReturnValue('gemini-pro'), + getProxy: vi.fn().mockReturnValue(undefined), + getUsageStatisticsEnabled: () => false, + getPreviewFeatures: vi.fn().mockReturnValue(false), + } as unknown as Config; + + const mockGenerator = { + models: {}, + } as unknown as GoogleGenAI; + vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never); + vi.stubEnv('GOOGLE_GENAI_API_VERSION', ''); + + await createContentGenerator( + { + apiKey: 'test-api-key', + authType: AuthType.USE_GEMINI, + }, + mockConfig, + ); + + expect(GoogleGenAI).toHaveBeenCalledWith({ + apiKey: 'test-api-key', + vertexai: undefined, + httpOptions: { + headers: expect.objectContaining({ + 'User-Agent': expect.any(String), + }), + }, + }); + + expect(GoogleGenAI).toHaveBeenCalledWith( + expect.not.objectContaining({ + apiVersion: expect.any(String), + }), + ); + }); + + it('should pass apiVersion for Vertex AI when GOOGLE_GENAI_API_VERSION is set', async () => { + const mockConfig = { + getModel: vi.fn().mockReturnValue('gemini-pro'), + getProxy: vi.fn().mockReturnValue(undefined), + getUsageStatisticsEnabled: () => false, + getPreviewFeatures: vi.fn().mockReturnValue(false), + } as unknown as Config; + + const mockGenerator = { + models: {}, + } as unknown as GoogleGenAI; + vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never); + vi.stubEnv('GOOGLE_GENAI_API_VERSION', 'v1alpha'); + + await createContentGenerator( + { + apiKey: 'test-api-key', + vertexai: true, + authType: AuthType.USE_VERTEX_AI, + }, + mockConfig, + ); + + expect(GoogleGenAI).toHaveBeenCalledWith({ + apiKey: 'test-api-key', + vertexai: true, + httpOptions: { + headers: expect.objectContaining({ + 'User-Agent': expect.any(String), + }), + }, + apiVersion: 'v1alpha', + }); + }); }); describe('createContentGeneratorConfig', () => { diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index eb45c9f218..77d0413349 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -132,6 +132,7 @@ export async function createContentGenerator( const customHeadersMap = parseCustomHeaders(customHeadersEnv); const apiKeyAuthMechanism = process.env['GEMINI_API_KEY_AUTH_MECHANISM'] || 'x-goog-api-key'; + const apiVersionEnv = process.env['GOOGLE_GENAI_API_VERSION']; const baseHeaders: Record = { ...customHeadersMap, @@ -181,6 +182,7 @@ export async function createContentGenerator( apiKey: config.apiKey === '' ? undefined : config.apiKey, vertexai: config.vertexai, httpOptions, + ...(apiVersionEnv && { apiVersion: apiVersionEnv }), }); return new LoggingContentGenerator(googleGenAI.models, gcConfig); }