Skip to content

Commit 4e22c2c

Browse files
Backport: feat(provider/google): add support for service tier parameter (#13916)
This is an automated backport of #13915 to the release-v6.0 branch. FYI @felixarntz Co-authored-by: Felix Arntz <felix.arntz@vercel.com>
1 parent d1eccb5 commit 4e22c2c

File tree

9 files changed

+212
-0
lines changed

9 files changed

+212
-0
lines changed

.changeset/cyan-pears-allow.md

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"@ai-sdk/google": patch
3+
---
4+
5+
feat(provider/google): add support for service tier parameter

content/providers/01-ai-sdk-providers/15-google-generative-ai.mdx

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -245,6 +245,12 @@ The following optional provider options are available for Google Generative AI m
245245
Optional. Defines labels used in billing reports. Available on Vertex AI only.
246246
See [Google Cloud labels documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/add-labels-to-api-calls).
247247

248+
- **serviceTier** _'SERVICE_TIER_STANDARD' | 'SERVICE_TIER_FLEX' | 'SERVICE_TIER_PRIORITY'_
249+
250+
Optional. The service tier to use for the request.
251+
Set to 'SERVICE_TIER_FLEX' for 50% cheaper processing at the cost of increased latency.
252+
Set to 'SERVICE_TIER_PRIORITY' for ultra-low latency at a 75-100% price premium over 'SERVICE_TIER_STANDARD'.
253+
248254
- **threshold** _string_
249255

250256
Optional. Standalone threshold setting that can be used independently of `safetySettings`.
Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
import { google, type GoogleLanguageModelOptions } from '@ai-sdk/google';
2+
import { generateText } from 'ai';
3+
import { run } from '../../lib/run';
4+
5+
run(async () => {
6+
const result = await generateText({
7+
model: google('gemini-3.1-pro-preview'),
8+
prompt: 'What color is the sky in one word?',
9+
providerOptions: {
10+
google: {
11+
serviceTier: 'SERVICE_TIER_FLEX',
12+
} satisfies GoogleLanguageModelOptions,
13+
},
14+
});
15+
16+
console.log(result.text);
17+
console.log('serviceTier:', result.providerMetadata?.google?.serviceTier);
18+
});
Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
import { google, type GoogleLanguageModelOptions } from '@ai-sdk/google';
2+
import { streamText } from 'ai';
3+
import { run } from '../../lib/run';
4+
5+
run(async () => {
6+
const result = streamText({
7+
model: google('gemini-3.1-pro-preview'),
8+
prompt: 'What color is the sky in one word?',
9+
providerOptions: {
10+
google: {
11+
serviceTier: 'SERVICE_TIER_FLEX',
12+
} satisfies GoogleLanguageModelOptions,
13+
},
14+
});
15+
16+
await result.consumeStream();
17+
18+
console.log(await result.text);
19+
console.log(
20+
'serviceTier:',
21+
(await result.providerMetadata)?.google?.serviceTier,
22+
);
23+
});

packages/google/src/__snapshots__/google-generative-ai-language-model.test.ts.snap

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ Here is the breakdown: st**r**awbe**rr**y.",
2525
"groundingMetadata": null,
2626
"promptFeedback": null,
2727
"safetyRatings": null,
28+
"serviceTier": null,
2829
"urlContextMetadata": null,
2930
"usageMetadata": {
3031
"candidatesTokenCount": 29,
@@ -63,6 +64,7 @@ Here is the breakdown: st**r**awbe**rr**y.",
6364
},
6465
"labels": undefined,
6566
"safetySettings": undefined,
67+
"serviceTier": undefined,
6668
"systemInstruction": undefined,
6769
"toolConfig": undefined,
6870
"tools": undefined,
@@ -155,6 +157,7 @@ Here is the breakdown: st**r**awbe**rr**y.",
155157
"groundingMetadata": null,
156158
"promptFeedback": null,
157159
"safetyRatings": null,
160+
"serviceTier": null,
158161
"urlContextMetadata": null,
159162
"usageMetadata": {
160163
"candidatesTokenCount": 28,
@@ -193,6 +196,7 @@ Here is the breakdown: st**r**awbe**rr**y.",
193196
},
194197
"labels": undefined,
195198
"safetySettings": undefined,
199+
"serviceTier": undefined,
196200
"systemInstruction": undefined,
197201
"toolConfig": undefined,
198202
"tools": undefined,
@@ -307,6 +311,7 @@ exports[`doGenerate > tool-call > should extract tool calls 1`] = `
307311
"groundingMetadata": null,
308312
"promptFeedback": null,
309313
"safetyRatings": null,
314+
"serviceTier": null,
310315
"urlContextMetadata": null,
311316
"usageMetadata": {
312317
"candidatesTokenCount": 15,
@@ -345,6 +350,7 @@ exports[`doGenerate > tool-call > should extract tool calls 1`] = `
345350
},
346351
"labels": undefined,
347352
"safetySettings": undefined,
353+
"serviceTier": undefined,
348354
"systemInstruction": undefined,
349355
"toolConfig": undefined,
350356
"tools": [
@@ -461,6 +467,7 @@ exports[`doGenerate > tool-call-gemini3 > should extract tool call with thoughtS
461467
"groundingMetadata": null,
462468
"promptFeedback": null,
463469
"safetyRatings": null,
470+
"serviceTier": null,
464471
"urlContextMetadata": null,
465472
"usageMetadata": {
466473
"candidatesTokenCount": 15,
@@ -499,6 +506,7 @@ exports[`doGenerate > tool-call-gemini3 > should extract tool call with thoughtS
499506
},
500507
"labels": undefined,
501508
"safetySettings": undefined,
509+
"serviceTier": undefined,
502510
"systemInstruction": undefined,
503511
"toolConfig": undefined,
504512
"tools": undefined,
@@ -620,6 +628,7 @@ Here is the breakdown: st**r**awbe**rr**y.",
620628
"groundingMetadata": null,
621629
"promptFeedback": null,
622630
"safetyRatings": null,
631+
"serviceTier": null,
623632
"urlContextMetadata": null,
624633
"usageMetadata": {
625634
"candidatesTokenCount": 29,
@@ -703,6 +712,7 @@ exports[`doStream > reasoning-gemini3 > should stream reasoning with thoughtSign
703712
"groundingMetadata": null,
704713
"promptFeedback": null,
705714
"safetyRatings": null,
715+
"serviceTier": null,
706716
"urlContextMetadata": null,
707717
"usageMetadata": {
708718
"candidatesTokenCount": 23,
@@ -891,6 +901,7 @@ st**r**awbe**rr**y",
891901
"groundingMetadata": null,
892902
"promptFeedback": null,
893903
"safetyRatings": null,
904+
"serviceTier": null,
894905
"urlContextMetadata": null,
895906
"usageMetadata": {
896907
"candidatesTokenCount": 23,
@@ -981,6 +992,7 @@ exports[`doStream > tool-call > should stream tool call 1`] = `
981992
"groundingMetadata": null,
982993
"promptFeedback": null,
983994
"safetyRatings": null,
995+
"serviceTier": null,
984996
"urlContextMetadata": null,
985997
"usageMetadata": {
986998
"candidatesTokenCount": 15,
@@ -1071,6 +1083,7 @@ exports[`doStream > tool-call-gemini3 > should stream tool call with thoughtSign
10711083
"groundingMetadata": null,
10721084
"promptFeedback": null,
10731085
"safetyRatings": null,
1086+
"serviceTier": null,
10741087
"urlContextMetadata": null,
10751088
"usageMetadata": {
10761089
"candidatesTokenCount": 15,

packages/google/src/google-generative-ai-language-model.test.ts

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -555,6 +555,74 @@ describe('doGenerate', () => {
555555
expect(providerMetadata?.google.finishMessage).toBeNull();
556556
});
557557

558+
it('should send serviceTier in request body when specified', async () => {
559+
prepareJsonResponse({ content: 'test response' });
560+
561+
await model.doGenerate({
562+
prompt: TEST_PROMPT,
563+
providerOptions: {
564+
google: {
565+
serviceTier: 'SERVICE_TIER_FLEX',
566+
},
567+
},
568+
});
569+
570+
expect(await server.calls[0].requestBodyJson).toMatchObject({
571+
serviceTier: 'SERVICE_TIER_FLEX',
572+
});
573+
});
574+
575+
it('should not send serviceTier in request body when not specified', async () => {
576+
prepareJsonResponse({ content: 'test response' });
577+
578+
await model.doGenerate({
579+
prompt: TEST_PROMPT,
580+
});
581+
582+
const body = await server.calls[0].requestBodyJson;
583+
expect(body).not.toHaveProperty('serviceTier');
584+
});
585+
586+
it('should expose serviceTier in provider metadata', async () => {
587+
server.urls[TEST_URL_GEMINI_PRO].response = {
588+
type: 'json-value',
589+
body: {
590+
candidates: [
591+
{
592+
content: {
593+
parts: [{ text: 'test response' }],
594+
role: 'model',
595+
},
596+
finishReason: 'STOP',
597+
safetyRatings: SAFETY_RATINGS,
598+
},
599+
],
600+
usageMetadata: {
601+
promptTokenCount: 1,
602+
candidatesTokenCount: 2,
603+
totalTokenCount: 3,
604+
},
605+
serviceTier: 'SERVICE_TIER_FLEX',
606+
},
607+
};
608+
609+
const { providerMetadata } = await model.doGenerate({
610+
prompt: TEST_PROMPT,
611+
});
612+
613+
expect(providerMetadata?.google.serviceTier).toBe('SERVICE_TIER_FLEX');
614+
});
615+
616+
it('should expose null serviceTier in provider metadata when not present', async () => {
617+
prepareJsonResponse({ content: 'test response' });
618+
619+
const { providerMetadata } = await model.doGenerate({
620+
prompt: TEST_PROMPT,
621+
});
622+
623+
expect(providerMetadata?.google.serviceTier).toBeNull();
624+
});
625+
558626
describe('tool-call', () => {
559627
beforeEach(() => {
560628
prepareJsonFixtureResponse('google-tool-call');
@@ -3745,6 +3813,60 @@ describe('doStream', () => {
37453813
).toBeNull();
37463814
});
37473815

3816+
it('should expose serviceTier in provider metadata on finish', async () => {
3817+
server.urls[TEST_URL_GEMINI_PRO].response = {
3818+
type: 'stream-chunks',
3819+
chunks: [
3820+
`data: ${JSON.stringify({
3821+
candidates: [
3822+
{
3823+
content: {
3824+
parts: [{ text: 'test response' }],
3825+
role: 'model',
3826+
},
3827+
finishReason: 'STOP',
3828+
safetyRatings: SAFETY_RATINGS,
3829+
},
3830+
],
3831+
usageMetadata: {
3832+
promptTokenCount: 1,
3833+
candidatesTokenCount: 2,
3834+
totalTokenCount: 3,
3835+
},
3836+
serviceTier: 'SERVICE_TIER_FLEX',
3837+
})}\n\n`,
3838+
],
3839+
};
3840+
3841+
const { stream } = await model.doStream({
3842+
prompt: TEST_PROMPT,
3843+
});
3844+
3845+
const events = await convertReadableStreamToArray(stream);
3846+
const finishEvent = events.find(event => event.type === 'finish');
3847+
3848+
expect(
3849+
finishEvent?.type === 'finish' &&
3850+
finishEvent.providerMetadata?.google.serviceTier,
3851+
).toBe('SERVICE_TIER_FLEX');
3852+
});
3853+
3854+
it('should expose null serviceTier in provider metadata on finish when not present', async () => {
3855+
prepareStreamResponse({ content: ['test'] });
3856+
3857+
const { stream } = await model.doStream({
3858+
prompt: TEST_PROMPT,
3859+
});
3860+
3861+
const events = await convertReadableStreamToArray(stream);
3862+
const finishEvent = events.find(event => event.type === 'finish');
3863+
3864+
expect(
3865+
finishEvent?.type === 'finish' &&
3866+
finishEvent.providerMetadata?.google.serviceTier,
3867+
).toBeNull();
3868+
});
3869+
37483870
it('should stream code execution tool calls and results', async () => {
37493871
server.urls[TEST_URL_GEMINI_2_0_PRO].response = {
37503872
type: 'stream-chunks',
@@ -4354,6 +4476,7 @@ describe('doStream', () => {
43544476
"probability": "NEGLIGIBLE",
43554477
},
43564478
],
4479+
"serviceTier": null,
43574480
"urlContextMetadata": null,
43584481
"usageMetadata": null,
43594482
},
@@ -4815,6 +4938,7 @@ describe('doStream', () => {
48154938
"groundingMetadata": null,
48164939
"promptFeedback": null,
48174940
"safetyRatings": null,
4941+
"serviceTier": null,
48184942
"urlContextMetadata": null,
48194943
"usageMetadata": {
48204944
"candidatesTokenCount": 18,
@@ -4977,6 +5101,7 @@ describe('doStream', () => {
49775101
"probability": "NEGLIGIBLE",
49785102
},
49795103
],
5104+
"serviceTier": null,
49805105
"urlContextMetadata": null,
49815106
"usageMetadata": null,
49825107
},

packages/google/src/google-generative-ai-language-model.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -207,6 +207,7 @@ export class GoogleGenerativeAILanguageModel implements LanguageModelV3 {
207207
: googleToolConfig,
208208
cachedContent: googleOptions?.cachedContent,
209209
labels: googleOptions?.labels,
210+
serviceTier: googleOptions?.serviceTier,
210211
},
211212
warnings: [...warnings, ...toolWarnings],
212213
providerOptionsName,
@@ -364,6 +365,7 @@ export class GoogleGenerativeAILanguageModel implements LanguageModelV3 {
364365
safetyRatings: candidate.safetyRatings ?? null,
365366
usageMetadata: usageMetadata ?? null,
366367
finishMessage: candidate.finishMessage ?? null,
368+
serviceTier: response.serviceTier ?? null,
367369
} satisfies GoogleGenerativeAIProviderMetadata,
368370
},
369371
request: { body: args },
@@ -405,6 +407,7 @@ export class GoogleGenerativeAILanguageModel implements LanguageModelV3 {
405407
let providerMetadata: SharedV3ProviderMetadata | undefined = undefined;
406408
let lastGroundingMetadata: GroundingMetadataSchema | null = null;
407409
let lastUrlContextMetadata: UrlContextMetadataSchema | null = null;
410+
let serviceTier: string | null = null;
408411

409412
const generateId = this.config.generateId;
410413
let hasToolCalls = false;
@@ -447,6 +450,10 @@ export class GoogleGenerativeAILanguageModel implements LanguageModelV3 {
447450
usage = usageMetadata;
448451
}
449452

453+
if (value.serviceTier != null) {
454+
serviceTier = value.serviceTier;
455+
}
456+
450457
const candidate = value.candidates?.[0];
451458

452459
// sometimes the API returns an empty candidates array
@@ -685,6 +692,7 @@ export class GoogleGenerativeAILanguageModel implements LanguageModelV3 {
685692
safetyRatings: candidate.safetyRatings ?? null,
686693
usageMetadata: usageMetadata ?? null,
687694
finishMessage: candidate.finishMessage ?? null,
695+
serviceTier,
688696
} satisfies GoogleGenerativeAIProviderMetadata,
689697
};
690698
}
@@ -1028,6 +1036,7 @@ const responseSchema = lazySchema(() =>
10281036
safetyRatings: z.array(getSafetyRatingSchema()).nullish(),
10291037
})
10301038
.nullish(),
1039+
serviceTier: z.string().nullish(),
10311040
}),
10321041
),
10331042
);
@@ -1083,6 +1092,7 @@ const chunkSchema = lazySchema(() =>
10831092
safetyRatings: z.array(getSafetyRatingSchema()).nullish(),
10841093
})
10851094
.nullish(),
1095+
serviceTier: z.string().nullish(),
10861096
}),
10871097
),
10881098
);

0 commit comments

Comments
 (0)