Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions docs/cli/configuration.md
Original file line number Diff line number Diff line change
Expand Up @@ -476,6 +476,10 @@ the `excludedProjectEnvVars` setting in your `settings.json` file.
- If set to `bearer`, the API key will be sent in the
`Authorization: Bearer <key>` header.
- Example: `export GEMINI_API_KEY_AUTH_MECHANISM="bearer"`
- **`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"`
- **`GOOGLE_API_KEY`**:
- Your Google Cloud API key.
- Required for using Vertex AI in express mode.
Expand Down
146 changes: 146 additions & 0 deletions packages/core/src/core/contentGenerator.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand Down
2 changes: 2 additions & 0 deletions packages/core/src/core/contentGenerator.ts
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,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,
Expand Down Expand Up @@ -179,6 +180,7 @@ export async function createContentGenerator(
apiKey: config.apiKey === '' ? undefined : config.apiKey,
vertexai: config.vertexai,
httpOptions,
...(apiVersionEnv && { apiVersion: apiVersionEnv }),
});
return new LoggingContentGenerator(googleGenAI.models, gcConfig);
}
Expand Down