Skip to content

Commit d32abbd

Browse files
shapershaper
authored
feat (provider/google-vertex): Add gemini 2 models. (#4122)
1 parent 03cf4e2 commit d32abbd

File tree

7 files changed

+154
-119
lines changed

7 files changed

+154
-119
lines changed

‎.changeset/gorgeous-boats-add.md

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
'@ai-sdk/google-vertex': patch
3+
---
4+
5+
feat (provider/google-vertex): Add gemini 2 models.

‎content/docs/02-foundations/02-providers-and-models.mdx

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@ Here are the capabilities of popular models:
8181
| [Google Generative AI](/providers/ai-sdk-providers/google-generative-ai) | `gemini-2.0-flash-exp` | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
8282
| [Google Generative AI](/providers/ai-sdk-providers/google-generative-ai) | `gemini-1.5-flash` | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
8383
| [Google Generative AI](/providers/ai-sdk-providers/google-generative-ai) | `gemini-1.5-pro` | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
84+
| [Google Vertex](/providers/ai-sdk-providers/google-vertex) | `gemini-2.0-flash-exp` | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
8485
| [Google Vertex](/providers/ai-sdk-providers/google-vertex) | `gemini-1.5-flash` | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
8586
| [Google Vertex](/providers/ai-sdk-providers/google-vertex) | `gemini-1.5-pro` | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
8687
| [xAI Grok](/providers/ai-sdk-providers/xai) | `grok-2-1212` | <Cross size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |

‎content/providers/01-ai-sdk-providers/11-google-vertex.mdx

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -513,10 +513,11 @@ const result = await generateObject({
513513

514514
### Model Capabilities
515515

516-
| Model | Image Input | Object Generation | Tool Usage | Tool Streaming |
517-
| ------------------ | ------------------- | ------------------- | ------------------- | ------------------- |
518-
| `gemini-1.5-flash` | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
519-
| `gemini-1.5-pro` | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
516+
| Model | Image Input | Object Generation | Tool Usage | Tool Streaming |
517+
| ---------------------- | ------------------- | ------------------- | ------------------- | ------------------- |
518+
| `gemini-2.0-flash-exp` | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
519+
| `gemini-1.5-flash` | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
520+
| `gemini-1.5-pro` | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
520521

521522
<Note>
522523
The table above lists popular models. Please see the [Google Vertex AI

‎content/providers/01-ai-sdk-providers/index.mdx

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ Not all providers support all AI SDK features. Here's a quick comparison of the
3535
| [Google Generative AI](/providers/ai-sdk-providers/google-generative-ai) | `gemini-2.0-flash-exp` | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
3636
| [Google Generative AI](/providers/ai-sdk-providers/google-generative-ai) | `gemini-1.5-flash` | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
3737
| [Google Generative AI](/providers/ai-sdk-providers/google-generative-ai) | `gemini-1.5-pro` | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
38+
| [Google Vertex](/providers/ai-sdk-providers/google-vertex) | `gemini-2.0-flash-exp` | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
3839
| [Google Vertex](/providers/ai-sdk-providers/google-vertex) | `gemini-1.5-flash` | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
3940
| [Google Vertex](/providers/ai-sdk-providers/google-vertex) | `gemini-1.5-pro` | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
4041
| [xAI Grok](/providers/ai-sdk-providers/xai) | `grok-2-1212` | <Cross size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |

‎examples/ai-core/src/e2e/google-vertex.test.ts

Lines changed: 68 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,8 @@ const LONG_TEST_MILLIS = 10000;
2020
const MODEL_VARIANTS = {
2121
chat: [
2222
'gemini-1.5-flash',
23-
// Pro models have low quota limits and can only be used if you have a
24-
// Google Cloud account with appropriate billing enabled.
23+
// Gemini 2.0 and Pro models have low quota limits and may require billing enabled.
24+
// 'gemini-2.0-flash-exp',
2525
// 'gemini-1.5-pro-001',
2626
// 'gemini-1.0-pro-001',
2727
],
@@ -40,6 +40,36 @@ const RUNTIME_VARIANTS = {
4040
},
4141
} as const;
4242

43+
// Add these helper functions near the top of the file
44+
const verifyGroundingMetadata = (groundingMetadata: any) => {
45+
expect(Array.isArray(groundingMetadata?.webSearchQueries)).toBe(true);
46+
expect(groundingMetadata?.webSearchQueries?.length).toBeGreaterThan(0);
47+
48+
// Verify search entry point exists
49+
expect(groundingMetadata?.searchEntryPoint?.renderedContent).toBeDefined();
50+
51+
// Verify grounding supports
52+
expect(Array.isArray(groundingMetadata?.groundingSupports)).toBe(true);
53+
const support = groundingMetadata?.groundingSupports?.[0];
54+
expect(support?.segment).toBeDefined();
55+
expect(Array.isArray(support?.groundingChunkIndices)).toBe(true);
56+
expect(Array.isArray(support?.confidenceScores)).toBe(true);
57+
};
58+
59+
const verifySafetyRatings = (safetyRatings: any[]) => {
60+
expect(Array.isArray(safetyRatings)).toBe(true);
61+
expect(safetyRatings?.length).toBeGreaterThan(0);
62+
63+
// Verify each safety rating has required properties
64+
safetyRatings?.forEach(rating => {
65+
expect(rating.category).toBeDefined();
66+
expect(rating.probability).toBeDefined();
67+
expect(typeof rating.probabilityScore).toBe('number');
68+
expect(rating.severity).toBeDefined();
69+
expect(typeof rating.severityScore).toBe('number');
70+
});
71+
};
72+
4373
describe.each(Object.values(RUNTIME_VARIANTS))(
4474
'Google Vertex E2E Tests - $name',
4575
({ vertex }) => {
@@ -82,7 +112,7 @@ describe.each(Object.values(RUNTIME_VARIANTS))(
82112
expect(result.usage?.totalTokens).toBeGreaterThan(0);
83113
});
84114

85-
it('should generate text with search grounding metadata in response when search grounding is enabled', async () => {
115+
it('should include search grounding metadata in response when search grounding is enabled', async () => {
86116
const model = vertex(modelId, {
87117
useSearchGrounding: true,
88118
});
@@ -96,43 +126,26 @@ describe.each(Object.values(RUNTIME_VARIANTS))(
96126
expect(result.text.toLowerCase()).toContain('tokyo');
97127
expect(result.usage?.totalTokens).toBeGreaterThan(0);
98128

99-
// Verify specific grounding metadata fields
100129
const metadata = result.experimental_providerMetadata?.google as
101130
| GoogleGenerativeAIProviderMetadata
102131
| undefined;
103-
const groundingMetadata = metadata?.groundingMetadata;
104-
expect(Array.isArray(groundingMetadata?.webSearchQueries)).toBe(true);
105-
expect(groundingMetadata?.webSearchQueries?.length).toBeGreaterThan(0);
106-
107-
// Verify search entry point exists
108-
expect(
109-
groundingMetadata?.searchEntryPoint?.renderedContent,
110-
).toBeDefined();
111-
112-
// Verify grounding supports
113-
expect(Array.isArray(groundingMetadata?.groundingSupports)).toBe(true);
114-
const support = groundingMetadata?.groundingSupports?.[0];
115-
expect(support?.segment).toBeDefined();
116-
expect(Array.isArray(support?.groundingChunkIndices)).toBe(true);
117-
expect(Array.isArray(support?.confidenceScores)).toBe(true);
118-
119-
// Verify safety ratings
120-
const safetyRatings = metadata?.safetyRatings;
121-
expect(Array.isArray(safetyRatings)).toBe(true);
122-
expect(safetyRatings?.length).toBeGreaterThan(0);
123-
124-
// Verify each safety rating has required properties
125-
safetyRatings?.forEach(rating => {
126-
expect(rating.category).toBeDefined();
127-
expect(rating.probability).toBeDefined();
128-
expect(typeof rating.probabilityScore).toBe('number');
129-
expect(rating.severity).toBeDefined();
130-
expect(typeof rating.severityScore).toBe('number');
132+
verifyGroundingMetadata(metadata?.groundingMetadata);
133+
});
134+
135+
it('should include safety ratings in response when search grounding is enabled', async () => {
136+
const model = vertex(modelId, {
137+
useSearchGrounding: true,
131138
});
132139

133-
// Basic response checks
134-
expect(result.text).toBeTruthy();
135-
expect(result.usage?.totalTokens).toBeGreaterThan(0);
140+
const result = await generateText({
141+
model,
142+
prompt: 'What is the current population of Tokyo?',
143+
});
144+
145+
const metadata = result.experimental_providerMetadata?.google as
146+
| GoogleGenerativeAIProviderMetadata
147+
| undefined;
148+
verifySafetyRatings(metadata?.safetyRatings ?? []);
136149
});
137150

138151
it('should generate text with PDF input', async () => {
@@ -234,7 +247,7 @@ describe.each(Object.values(RUNTIME_VARIANTS))(
234247
expect((await result.usage)?.totalTokens).toBeGreaterThan(0);
235248
});
236249

237-
it('should stream text with search grounding metadata when search grounding is enabled', async () => {
250+
it('should include search grounding metadata when streaming with search grounding enabled', async () => {
238251
const model = vertex(modelId, {
239252
useSearchGrounding: true,
240253
});
@@ -249,49 +262,35 @@ describe.each(Object.values(RUNTIME_VARIANTS))(
249262
chunks.push(chunk);
250263
}
251264

252-
// Get the complete response metadata
253265
const metadata = (await result.experimental_providerMetadata)
254266
?.google as GoogleGenerativeAIProviderMetadata | undefined;
255-
const groundingMetadata = metadata?.groundingMetadata;
256267

257268
const completeText = chunks.join('');
258269
expect(completeText).toBeTruthy();
259270
expect(completeText.toLowerCase()).toContain('tokyo');
260271
expect((await result.usage)?.totalTokens).toBeGreaterThan(0);
261272

262-
// Verify specific grounding metadata fields
263-
expect(Array.isArray(groundingMetadata?.webSearchQueries)).toBe(true);
264-
expect(groundingMetadata?.webSearchQueries?.length).toBeGreaterThan(0);
265-
266-
// Verify search entry point exists
267-
expect(
268-
groundingMetadata?.searchEntryPoint?.renderedContent,
269-
).toBeDefined();
270-
271-
// Verify grounding supports
272-
expect(Array.isArray(groundingMetadata?.groundingSupports)).toBe(true);
273-
const support = groundingMetadata?.groundingSupports?.[0];
274-
expect(support?.segment).toBeDefined();
275-
expect(Array.isArray(support?.groundingChunkIndices)).toBe(true);
276-
expect(Array.isArray(support?.confidenceScores)).toBe(true);
277-
278-
// Verify safety ratings
279-
const safetyRatings = metadata?.safetyRatings;
280-
expect(Array.isArray(safetyRatings)).toBe(true);
281-
expect(safetyRatings?.length).toBeGreaterThan(0);
282-
283-
// Verify each safety rating has required properties
284-
safetyRatings?.forEach(rating => {
285-
expect(rating.category).toBeDefined();
286-
expect(rating.probability).toBeDefined();
287-
expect(typeof rating.probabilityScore).toBe('number');
288-
expect(rating.severity).toBeDefined();
289-
expect(typeof rating.severityScore).toBe('number');
273+
verifyGroundingMetadata(metadata?.groundingMetadata);
274+
});
275+
276+
it('should include safety ratings when streaming with search grounding enabled', async () => {
277+
const model = vertex(modelId, {
278+
useSearchGrounding: true,
290279
});
291280

292-
// Basic response checks
293-
expect(chunks.join('')).toBeTruthy();
294-
expect((await result.usage)?.totalTokens).toBeGreaterThan(0);
281+
const result = streamText({
282+
model,
283+
prompt: 'What is the current population of Tokyo?',
284+
});
285+
286+
for await (const _ of result.textStream) {
287+
// consume the stream
288+
}
289+
290+
const metadata = (await result.experimental_providerMetadata)
291+
?.google as GoogleGenerativeAIProviderMetadata | undefined;
292+
293+
verifySafetyRatings(metadata?.safetyRatings ?? []);
295294
});
296295

297296
it('should stream object', async () => {

0 commit comments

Comments
 (0)