feat(core): Add GOOGLE_GENAI_API_VERSION environment variable support (#16177)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Bryan Morgan <bryanmorgan@google.com>
This commit is contained in:
Danielle Yim
2026-01-29 14:45:37 -08:00
committed by GitHub
parent 2de39311cf
commit fdda3a2399
3 changed files with 152 additions and 0 deletions

View File

@@ -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', () => {

View File

@@ -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<string, string> = {
...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);
}