diff --git a/core/src/models/google_llm.ts b/core/src/models/google_llm.ts index 0bc30049e..5ae52c079 100644 --- a/core/src/models/google_llm.ts +++ b/core/src/models/google_llm.ts @@ -13,7 +13,7 @@ import { LiveServerMessage, } from '@google/genai'; -import {getBooleanEnvVar, isBrowser} from '../utils/env_aware_utils.js'; +import {isBrowser, isEnterpriseModeEnabled} from '../utils/env_aware_utils.js'; import {logger} from '../utils/logger.js'; import {GoogleLLMVariant} from '../utils/variant_utils.js'; @@ -373,7 +373,7 @@ export function geminiInitParams({ params.vertexai = !!vertexai; if (!params.vertexai && !isBrowser()) { - params.vertexai = getBooleanEnvVar('GOOGLE_GENAI_USE_VERTEXAI'); + params.vertexai = isEnterpriseModeEnabled(); } if (params.vertexai) { diff --git a/core/src/utils/env_aware_utils.ts b/core/src/utils/env_aware_utils.ts index b20557332..23671293b 100644 --- a/core/src/utils/env_aware_utils.ts +++ b/core/src/utils/env_aware_utils.ts @@ -4,6 +4,11 @@ * SPDX-License-Identifier: Apache-2.0 */ +import {logger} from './logger.js'; + +const ENTERPRISE_MODE_ENV_VAR = 'GOOGLE_GENAI_USE_ENTERPRISE'; +const DEPRECATED_ENTERPRISE_MODE_ENV_VAR = 'GOOGLE_GENAI_USE_VERTEXAI'; + /** * Returns true if the environment is a browser. */ @@ -105,3 +110,34 @@ export function getBooleanEnvVar(envVar: string): boolean { return ['true', '1'].includes(envVarValue); } + +let warnedDeprecatedEnterpriseModeEnvVar = false; + +/** + * Returns whether Google GenAI enterprise mode is enabled. + * + * `GOOGLE_GENAI_USE_ENTERPRISE` takes precedence whenever it is set, even when + * it is set to a falsy value. `GOOGLE_GENAI_USE_VERTEXAI` is only consulted + * when `GOOGLE_GENAI_USE_ENTERPRISE` is absent, and using it logs a deprecation + * warning. This is read per request, so the warning is logged only once. + * + * @return True if enterprise mode is enabled, false otherwise. + */ +export function isEnterpriseModeEnabled(): boolean { + if (process.env?.[ENTERPRISE_MODE_ENV_VAR] !== undefined) { + return getBooleanEnvVar(ENTERPRISE_MODE_ENV_VAR); + } + + if (process.env?.[DEPRECATED_ENTERPRISE_MODE_ENV_VAR] !== undefined) { + if (!warnedDeprecatedEnterpriseModeEnvVar) { + warnedDeprecatedEnterpriseModeEnvVar = true; + logger.warn( + `${DEPRECATED_ENTERPRISE_MODE_ENV_VAR} is deprecated, please use ` + + `${ENTERPRISE_MODE_ENV_VAR} instead`, + ); + } + return getBooleanEnvVar(DEPRECATED_ENTERPRISE_MODE_ENV_VAR); + } + + return false; +} diff --git a/core/src/utils/variant_utils.ts b/core/src/utils/variant_utils.ts index ba3c08543..bc4c53fee 100644 --- a/core/src/utils/variant_utils.ts +++ b/core/src/utils/variant_utils.ts @@ -4,7 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {getBooleanEnvVar} from './env_aware_utils.js'; +import {isEnterpriseModeEnabled} from './env_aware_utils.js'; /** * The Google LLM variant to use. @@ -27,7 +27,7 @@ export enum GoogleLLMVariant { * Gets the Google LLM variant to use. */ export function getGoogleLlmVariant() { - return getBooleanEnvVar('GOOGLE_GENAI_USE_VERTEXAI') + return isEnterpriseModeEnabled() ? GoogleLLMVariant.VERTEX_AI : GoogleLLMVariant.GEMINI_API; } diff --git a/core/src/utils/vertex_ai_utils.ts b/core/src/utils/vertex_ai_utils.ts index b4ff3cc07..8c0f922a1 100644 --- a/core/src/utils/vertex_ai_utils.ts +++ b/core/src/utils/vertex_ai_utils.ts @@ -4,11 +4,15 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {getBooleanEnvVar} from './env_aware_utils.js'; +import {isEnterpriseModeEnabled} from './env_aware_utils.js'; /** * Validates and returns the API key for Express Mode. * + * The key is only returned when enterprise mode is enabled via + * `GOOGLE_GENAI_USE_ENTERPRISE` (or the deprecated + * `GOOGLE_GENAI_USE_VERTEXAI`). + * * @param project The project id. * @param location The location. * @param expressModeApiKey The API key for Express Mode. @@ -26,7 +30,7 @@ export function getExpressModeApiKey( ); } - if (getBooleanEnvVar('GOOGLE_GENAI_USE_VERTEXAI')) { + if (isEnterpriseModeEnabled()) { return expressModeApiKey || process.env.GOOGLE_API_KEY; } diff --git a/core/test/models/apigee_llm_test.ts b/core/test/models/apigee_llm_test.ts index 4a88d0324..94b67cfc2 100644 --- a/core/test/models/apigee_llm_test.ts +++ b/core/test/models/apigee_llm_test.ts @@ -25,6 +25,7 @@ describe('ApigeeLlm', () => { delete process.env['APIGEE_PROXY_URL']; delete process.env['GOOGLE_GENAI_API_KEY']; delete process.env['GOOGLE_GENAI_USE_VERTEXAI']; + delete process.env['GOOGLE_GENAI_USE_ENTERPRISE']; delete process.env['GOOGLE_CLOUD_PROJECT']; delete process.env['GOOGLE_CLOUD_LOCATION']; }); diff --git a/core/test/models/google_llm_test.ts b/core/test/models/google_llm_test.ts index 1d745c403..cc9c2e405 100644 --- a/core/test/models/google_llm_test.ts +++ b/core/test/models/google_llm_test.ts @@ -58,6 +58,7 @@ describe('GoogleLlm', () => { delete process.env['GOOGLE_GENAI_API_KEY']; delete process.env['GEMINI_API_KEY']; delete process.env['GOOGLE_GENAI_USE_VERTEXAI']; + delete process.env['GOOGLE_GENAI_USE_ENTERPRISE']; delete process.env['GOOGLE_CLOUD_AGENT_ENGINE_ID']; }; @@ -330,6 +331,28 @@ describe('GoogleLlm', () => { expect(params.vertexai).toBe(true); }); + it('should detect Vertex AI from GOOGLE_GENAI_USE_ENTERPRISE', () => { + process.env['GOOGLE_GENAI_USE_ENTERPRISE'] = 'true'; + process.env['GOOGLE_CLOUD_PROJECT'] = 'env-project'; + process.env['GOOGLE_CLOUD_LOCATION'] = 'env-location'; + const input = { + model: 'gemini-1.5-flash', + }; + const params = geminiInitParams(input); + expect(params.vertexai).toBe(true); + }); + + it('should not use Vertex AI when GOOGLE_GENAI_USE_ENTERPRISE disables it', () => { + process.env['GOOGLE_GENAI_USE_ENTERPRISE'] = 'false'; + process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'true'; + const input = { + model: 'gemini-1.5-flash', + apiKey: 'test-key', + }; + const params = geminiInitParams(input); + expect(params.vertexai).toBe(false); + }); + it('should throw error if project is missing for Vertex AI', () => { const input = { model: 'gemini-1.5-flash', diff --git a/core/test/utils/env_aware_utils_test.ts b/core/test/utils/env_aware_utils_test.ts index bc1a21dc3..2cc8ef56d 100644 --- a/core/test/utils/env_aware_utils_test.ts +++ b/core/test/utils/env_aware_utils_test.ts @@ -4,8 +4,17 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {afterEach, describe, expect, it} from 'vitest'; +import { + afterEach, + beforeEach, + describe, + expect, + it, + vi, + type MockInstance, +} from 'vitest'; import {getBooleanEnvVar} from '../../src/utils/env_aware_utils.js'; +import type {Logger} from '../../src/utils/logger.js'; describe('env_aware_utils', () => { describe('getBooleanEnvVar', () => { @@ -50,4 +59,78 @@ describe('env_aware_utils', () => { expect(getBooleanEnvVar('NON_EXISTENT_VAR')).toBe(false); }); }); + + describe('isEnterpriseModeEnabled', () => { + const originalEnv = process.env; + let isEnterpriseModeEnabled: () => boolean; + let warnSpy: MockInstance; + + beforeEach(async () => { + process.env = {...originalEnv}; + delete process.env['GOOGLE_GENAI_USE_ENTERPRISE']; + delete process.env['GOOGLE_GENAI_USE_VERTEXAI']; + // The deprecation warning is emitted once per module instance, so each + // test needs its own copy of the module and of the logger it warns to. + vi.resetModules(); + const {logger} = await import('../../src/utils/logger.js'); + warnSpy = vi.spyOn(logger, 'warn').mockImplementation(() => {}); + ({isEnterpriseModeEnabled} = + await import('../../src/utils/env_aware_utils.js')); + }); + + afterEach(() => { + process.env = originalEnv; + vi.restoreAllMocks(); + }); + + it('should return false and not warn when neither variable is set', () => { + expect(isEnterpriseModeEnabled()).toBe(false); + expect(warnSpy).not.toHaveBeenCalled(); + }); + + it('should return true when GOOGLE_GENAI_USE_ENTERPRISE is enabled', () => { + process.env['GOOGLE_GENAI_USE_ENTERPRISE'] = 'true'; + expect(isEnterpriseModeEnabled()).toBe(true); + expect(warnSpy).not.toHaveBeenCalled(); + }); + + it('should prefer an enabled GOOGLE_GENAI_USE_ENTERPRISE over GOOGLE_GENAI_USE_VERTEXAI', () => { + process.env['GOOGLE_GENAI_USE_ENTERPRISE'] = 'true'; + process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'false'; + expect(isEnterpriseModeEnabled()).toBe(true); + expect(warnSpy).not.toHaveBeenCalled(); + }); + + it('should not fall back to GOOGLE_GENAI_USE_VERTEXAI when GOOGLE_GENAI_USE_ENTERPRISE is set but disabled', () => { + process.env['GOOGLE_GENAI_USE_ENTERPRISE'] = ''; + process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'true'; + expect(isEnterpriseModeEnabled()).toBe(false); + expect(warnSpy).not.toHaveBeenCalled(); + }); + + it('should fall back to an enabled GOOGLE_GENAI_USE_VERTEXAI and warn', () => { + process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'true'; + expect(isEnterpriseModeEnabled()).toBe(true); + expect(warnSpy).toHaveBeenCalledOnce(); + expect(warnSpy).toHaveBeenCalledWith( + expect.stringContaining( + 'GOOGLE_GENAI_USE_VERTEXAI is deprecated, please use ' + + 'GOOGLE_GENAI_USE_ENTERPRISE', + ), + ); + }); + + it('should warn when GOOGLE_GENAI_USE_VERTEXAI is present but disabled', () => { + process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'false'; + expect(isEnterpriseModeEnabled()).toBe(false); + expect(warnSpy).toHaveBeenCalledOnce(); + }); + + it('should warn only once even when read repeatedly', () => { + process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'true'; + expect(isEnterpriseModeEnabled()).toBe(true); + expect(isEnterpriseModeEnabled()).toBe(true); + expect(warnSpy).toHaveBeenCalledOnce(); + }); + }); }); diff --git a/core/test/utils/variant_utils_test.ts b/core/test/utils/variant_utils_test.ts index 4f949e433..a7915ff62 100644 --- a/core/test/utils/variant_utils_test.ts +++ b/core/test/utils/variant_utils_test.ts @@ -5,34 +5,53 @@ */ import {GoogleLLMVariant} from '@google/adk'; -import {afterEach, describe, expect, it} from 'vitest'; +import {afterEach, beforeEach, describe, expect, it, vi} from 'vitest'; +import {logger} from '../../src/utils/logger.js'; import {getGoogleLlmVariant} from '../../src/utils/variant_utils.js'; describe('variant_utils', () => { describe('getGoogleLlmVariant', () => { const originalEnv = process.env; + beforeEach(() => { + process.env = {...originalEnv}; + delete process.env['GOOGLE_GENAI_USE_ENTERPRISE']; + delete process.env['GOOGLE_GENAI_USE_VERTEXAI']; + vi.spyOn(logger, 'warn').mockImplementation(() => {}); + }); + afterEach(() => { process.env = originalEnv; + vi.restoreAllMocks(); }); it('should return GEMINI_API by default (when env var is not set)', () => { - delete process.env['GOOGLE_GENAI_USE_VERTEXAI']; expect(getGoogleLlmVariant()).toBe(GoogleLLMVariant.GEMINI_API); }); it('should return VERTEX_AI when GOOGLE_GENAI_USE_VERTEXAI is "true"', () => { - process.env = {...originalEnv, 'GOOGLE_GENAI_USE_VERTEXAI': 'true'}; + process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'true'; expect(getGoogleLlmVariant()).toBe(GoogleLLMVariant.VERTEX_AI); }); it('should return VERTEX_AI when GOOGLE_GENAI_USE_VERTEXAI is "1"', () => { - process.env = {...originalEnv, 'GOOGLE_GENAI_USE_VERTEXAI': '1'}; + process.env['GOOGLE_GENAI_USE_VERTEXAI'] = '1'; expect(getGoogleLlmVariant()).toBe(GoogleLLMVariant.VERTEX_AI); }); it('should return GEMINI_API when GOOGLE_GENAI_USE_VERTEXAI is "false"', () => { - process.env = {...originalEnv, 'GOOGLE_GENAI_USE_VERTEXAI': 'false'}; + process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'false'; + expect(getGoogleLlmVariant()).toBe(GoogleLLMVariant.GEMINI_API); + }); + + it('should return VERTEX_AI when GOOGLE_GENAI_USE_ENTERPRISE is "true"', () => { + process.env['GOOGLE_GENAI_USE_ENTERPRISE'] = 'true'; + expect(getGoogleLlmVariant()).toBe(GoogleLLMVariant.VERTEX_AI); + }); + + it('should return GEMINI_API when GOOGLE_GENAI_USE_ENTERPRISE is "false" and GOOGLE_GENAI_USE_VERTEXAI is "true"', () => { + process.env['GOOGLE_GENAI_USE_ENTERPRISE'] = 'false'; + process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'true'; expect(getGoogleLlmVariant()).toBe(GoogleLLMVariant.GEMINI_API); }); }); diff --git a/core/test/utils/vertex_ai_utils_test.ts b/core/test/utils/vertex_ai_utils_test.ts index d397f143c..287d2fbeb 100644 --- a/core/test/utils/vertex_ai_utils_test.ts +++ b/core/test/utils/vertex_ai_utils_test.ts @@ -4,7 +4,8 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {afterEach, beforeEach, describe, expect, it} from 'vitest'; +import {afterEach, beforeEach, describe, expect, it, vi} from 'vitest'; +import {logger} from '../../src/utils/logger.js'; import {getExpressModeApiKey} from '../../src/utils/vertex_ai_utils.js'; describe('vertex_ai_utils', () => { @@ -13,10 +14,15 @@ describe('vertex_ai_utils', () => { beforeEach(() => { process.env = {...originalEnv}; + delete process.env['GOOGLE_GENAI_USE_ENTERPRISE']; + delete process.env['GOOGLE_GENAI_USE_VERTEXAI']; + delete process.env['GOOGLE_API_KEY']; + vi.spyOn(logger, 'warn').mockImplementation(() => {}); }); afterEach(() => { process.env = originalEnv; + vi.restoreAllMocks(); }); it('should throw when both project and expressModeApiKey are provided', () => { @@ -37,14 +43,15 @@ describe('vertex_ai_utils', () => { ).toThrow(); }); - it('should return undefined when GOOGLE_GENAI_USE_VERTEXAI is not set', () => { - delete process.env['GOOGLE_GENAI_USE_VERTEXAI']; + it('should return undefined when neither enterprise mode variable is set', () => { + process.env['GOOGLE_API_KEY'] = 'env-api-key'; const result = getExpressModeApiKey(); expect(result).toBeUndefined(); }); it('should return undefined when GOOGLE_GENAI_USE_VERTEXAI is false', () => { process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'false'; + process.env['GOOGLE_API_KEY'] = 'env-api-key'; const result = getExpressModeApiKey(); expect(result).toBeUndefined(); }); @@ -64,7 +71,28 @@ describe('vertex_ai_utils', () => { it('should return undefined when GOOGLE_GENAI_USE_VERTEXAI is true but no key available', () => { process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'true'; - delete process.env['GOOGLE_API_KEY']; + const result = getExpressModeApiKey(); + expect(result).toBeUndefined(); + }); + + it('should return expressModeApiKey when GOOGLE_GENAI_USE_ENTERPRISE is true', () => { + process.env['GOOGLE_GENAI_USE_ENTERPRISE'] = 'true'; + const result = getExpressModeApiKey(undefined, undefined, 'my-api-key'); + expect(result).toBe('my-api-key'); + }); + + it('should prefer an enabled GOOGLE_GENAI_USE_ENTERPRISE over GOOGLE_GENAI_USE_VERTEXAI', () => { + process.env['GOOGLE_GENAI_USE_ENTERPRISE'] = 'true'; + process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'false'; + process.env['GOOGLE_API_KEY'] = 'env-api-key'; + const result = getExpressModeApiKey(); + expect(result).toBe('env-api-key'); + }); + + it('should not fall back to GOOGLE_GENAI_USE_VERTEXAI when GOOGLE_GENAI_USE_ENTERPRISE is set but disabled', () => { + process.env['GOOGLE_GENAI_USE_ENTERPRISE'] = ''; + process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'true'; + process.env['GOOGLE_API_KEY'] = 'env-api-key'; const result = getExpressModeApiKey(); expect(result).toBeUndefined(); });