Skip to content

Commit 9c0540a

Browse files
matthew29tangcopybara-github
matthew29tang
authored and
copybara-github
committed
feat: Allow api key + proj/location for enterprise mode
PiperOrigin-RevId: 955366889
1 parent d221cf3 commit 9c0540a

4 files changed

Lines changed: 48 additions & 29 deletions

File tree

‎src/_api_client.ts‎

Lines changed: 3 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -178,15 +178,6 @@ export class ApiClient {
178178

179179
this.customBaseUrl = opts.httpOptions?.baseUrl;
180180

181-
if (this.clientOptions.vertexai) {
182-
if (this.clientOptions.project && this.clientOptions.location) {
183-
this.clientOptions.apiKey = undefined;
184-
} else if (this.clientOptions.apiKey) {
185-
this.clientOptions.project = undefined;
186-
this.clientOptions.location = undefined;
187-
}
188-
}
189-
190181
const initHttpOptions: types.HttpOptions = {};
191182

192183
if (this.clientOptions.vertexai) {
@@ -216,7 +207,7 @@ export class ApiClient {
216207
this.clientOptions.project = undefined;
217208
this.clientOptions.location = undefined;
218209
} else if (
219-
this.clientOptions.apiKey ||
210+
(this.clientOptions.apiKey && !this.clientOptions.project) ||
220211
this.clientOptions.location === 'global'
221212
) {
222213
// Vertex Express or global endpoint case.
@@ -382,10 +373,10 @@ export class ApiClient {
382373
) {
383374
return false;
384375
}
385-
if (this.clientOptions.apiKey) {
376+
if (!this.clientOptions.vertexai) {
386377
return false;
387378
}
388-
if (!this.clientOptions.vertexai) {
379+
if (!this.clientOptions.project || !this.clientOptions.location) {
389380
return false;
390381
}
391382
if (request.path.startsWith('projects/')) {

‎src/node/node_client.ts‎

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -207,15 +207,17 @@ export class GoogleGenAI {
207207
this._triggers = new GeminiNextGenTriggers(this.apiClient);
208208
return this._triggers;
209209
}
210+
210211
constructor(options: GoogleGenAIOptions) {
212+
this.vertexai = resolveCloudFlag(options);
213+
211214
// Validate explicitly set initializer values.
212-
if ((options.project || options.location) && options.apiKey) {
215+
if ((options.project || options.location) && !this.vertexai) {
213216
throw new Error(
214-
'Project/location and API key are mutually exclusive in the client initializer.',
217+
'Project and location are not supported for Gemini API backend.',
215218
);
216219
}
217220

218-
this.vertexai = resolveCloudFlag(options);
219221
const envApiKey = getApiKeyFromEnv();
220222
const envProject = getEnv('GOOGLE_CLOUD_PROJECT');
221223
const envLocation = getEnv('GOOGLE_CLOUD_LOCATION');
@@ -238,23 +240,37 @@ export class GoogleGenAI {
238240
);
239241
this.apiKey = undefined;
240242
}
241-
// Explicit api_key and explicit project/location already handled above.
242-
if ((envProject || envLocation) && options.apiKey) {
243+
if (
244+
!options.project &&
245+
!options.location &&
246+
(envProject || envLocation) &&
247+
options.apiKey
248+
) {
243249
// Explicit api_key takes precedence over implicit project/location.
244250
console.debug(
245251
'The user provided Vertex AI API key will take precedence over' +
246252
' the project/location from the environment variables.',
247253
);
248254
this.project = undefined;
249255
this.location = undefined;
250-
} else if ((options.project || options.location) && envApiKey) {
256+
} else if (
257+
(options.project || options.location) &&
258+
!options.apiKey &&
259+
envApiKey
260+
) {
251261
// Explicit project/location takes precedence over implicit api_key.
252262
console.debug(
253263
'The user provided project/location will take precedence over' +
254264
' the API key from the environment variables.',
255265
);
256266
this.apiKey = undefined;
257-
} else if ((envProject || envLocation) && envApiKey) {
267+
} else if (
268+
!options.project &&
269+
!options.location &&
270+
!options.apiKey &&
271+
(envProject || envLocation) &&
272+
envApiKey
273+
) {
258274
// Implicit project/location takes precedence over implicit api_key.
259275
console.debug(
260276
'The project/location from the environment variables will take' +

‎test/unit/api_client_test.ts‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -456,7 +456,7 @@ describe('ApiClient', () => {
456456
expect(client.isVertexAI()).toBe(true);
457457
expect(client.getProject()).toBe('vertex-project');
458458
expect(client.getLocation()).toBe('vertex-location');
459-
expect(client.getApiKey()).toBeUndefined(); // API key is ignored when setting opts.vertexai
459+
expect(client.getApiKey()).toBe('apikey-from-opts');
460460
expect(client.getRequestUrl()).toBe(
461461
'https://vertex-location-aiplatform.googleapis.com/v1beta1',
462462
);
@@ -478,7 +478,7 @@ describe('ApiClient', () => {
478478
expect(client.isVertexAI()).toBe(true);
479479
expect(client.getProject()).toBe('vertex-project');
480480
expect(client.getLocation()).toBe('us');
481-
expect(client.getApiKey()).toBeUndefined();
481+
expect(client.getApiKey()).toBe('apikey-from-opts');
482482
expect(client.getRequestUrl()).toBe(
483483
'https://aiplatform.us.rep.googleapis.com/v1beta1',
484484
);
@@ -500,7 +500,7 @@ describe('ApiClient', () => {
500500
expect(client.isVertexAI()).toBe(true);
501501
expect(client.getProject()).toBe('vertex-project');
502502
expect(client.getLocation()).toBe('eu');
503-
expect(client.getApiKey()).toBeUndefined();
503+
expect(client.getApiKey()).toBe('apikey-from-opts');
504504
expect(client.getRequestUrl()).toBe(
505505
'https://aiplatform.eu.rep.googleapis.com/v1beta1',
506506
);
@@ -523,12 +523,12 @@ describe('ApiClient', () => {
523523
expect(client.isVertexAI()).toBe(true);
524524
expect(client.getProject()).toBe('vertex-project');
525525
expect(client.getLocation()).toBe('us');
526-
expect(client.getApiKey()).toBeUndefined();
526+
expect(client.getApiKey()).toBe('apikey-from-opts');
527527
expect(client.getRequestUrl()).toBe('https://my-custom-url.com/v1beta1');
528528
expect(client.getApiVersion()).toBe('v1beta1');
529529
});
530530

531-
it('should not have api key if project/location is provided for vertexai', () => {
531+
it('should preserve api key if project/location is provided for vertexai', () => {
532532
const client = new ApiClient({
533533
auth: new FakeAuth(),
534534
project: 'vertex-project',
@@ -542,7 +542,7 @@ describe('ApiClient', () => {
542542
expect(client.isVertexAI()).toBe(true);
543543
expect(client.getProject()).toBe('vertex-project');
544544
expect(client.getLocation()).toBe('vertex-location');
545-
expect(client.getApiKey()).toBeUndefined();
545+
expect(client.getApiKey()).toBe('apikey-from-opts');
546546
expect(client.getRequestUrl()).toBe(
547547
'https://vertex-location-aiplatform.googleapis.com/v1beta1',
548548
);

‎test/unit/node/client_test.ts‎

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -123,18 +123,30 @@ describe('Client', () => {
123123
expect(client['project']).toBe('constructor_project');
124124
expect(client['location']).toBe('constructor_location');
125125
});
126-
it('should not allow both project and apikey in constructor', () => {
126+
it('should not allow project or location for Gemini API backend', () => {
127127
expect(() => {
128128
new GoogleGenAI({
129129
apiKey: 'constructor_api_key',
130-
vertexai: true,
130+
vertexai: false,
131131
project: 'constructor_project',
132132
location: 'constructor_location',
133133
});
134134
}).toThrowError(
135-
'Project/location and API key are mutually exclusive in the client initializer.',
135+
'Project and location are not supported for Gemini API backend.',
136136
);
137137
});
138+
it('should allow explicit apiKey with project and location when vertexai is true', () => {
139+
const client = new GoogleGenAI({
140+
apiKey: 'constructor_api_key',
141+
vertexai: true,
142+
project: 'constructor_project',
143+
location: 'constructor_location',
144+
});
145+
expect(client.vertexai).toBe(true);
146+
expect(client['apiKey']).toBe('constructor_api_key');
147+
expect(client['project']).toBe('constructor_project');
148+
expect(client['location']).toBe('constructor_location');
149+
});
138150
it('should prioritize credentials over implicit api key', () => {
139151
process.env['GOOGLE_API_KEY'] = '';
140152

@@ -604,7 +616,7 @@ describe('Client', () => {
604616
it('env enterprise should take precedence over env vertexai', () => {
605617
process.env['GOOGLE_GENAI_USE_ENTERPRISE'] = 'false';
606618
process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'true';
607-
const client = new GoogleGenAI({project: 'p', location: 'l'});
619+
const client = new GoogleGenAI({apiKey: 'key'});
608620
expect(client.vertexai).toBeFalse();
609621
});
610622

0 commit comments

Comments
 (0)