Skip to content

Commit 231f702

Browse files
koistyakoistya
committed
feat(google): add filtering options to vertexRagStore tool
1 parent d8e2081 commit 231f702

File tree

6 files changed

+329
-9
lines changed

6 files changed

+329
-9
lines changed

‎.changeset/feat-google-rag-file-ids.md‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
---
2+
'@ai-sdk/google': patch
3+
'@ai-sdk/google-vertex': patch
4+
---
5+
6+
feat(google): add ragFileIds, metadataFilter, and threshold options to vertexRagStore

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -767,11 +767,31 @@ The `vertexRagStore` tool accepts the following configuration options:
767767
- The RagCorpus resource name in the format: `projects/{project}/locations/{location}/ragCorpora/{rag_corpus}`
768768
- This identifies your specific RAG corpus to search against
769769

770+
- **`ragFileIds`** (`string[]`, optional)
771+
772+
- File IDs within the corpus to search
773+
- When specified, only these files are searched instead of the entire corpus
774+
- For large-scale filtering, prefer using `metadataFilter` instead
775+
770776
- **`topK`** (`number`, optional)
771777

772778
- The number of top contexts to retrieve from your RAG corpus
773779
- Defaults to the corpus configuration if not specified
774780

781+
- **`metadataFilter`** (`string`, optional)
782+
783+
- Filter expression for metadata assigned to files when uploaded
784+
- Example: `'user_id = "user-123" AND project = "acme"'`
785+
786+
- **`vectorSimilarityThreshold`** (`number`, optional)
787+
788+
- Only return results with vector similarity larger than this threshold
789+
- Value should be between 0 and 1
790+
791+
- **`vectorDistanceThreshold`** (`number`, optional)
792+
793+
- Only return results with vector distance smaller than this threshold
794+
775795
### Image Outputs
776796

777797
Gemini models with image generation capabilities (`gemini-2.5-flash-image-preview`) support image generation. Images are exposed as files in the response.

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

Lines changed: 49 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1560,17 +1560,62 @@ describe('doGenerate', () => {
15601560
{
15611561
retrieval: {
15621562
vertex_rag_store: {
1563-
rag_resources: {
1564-
rag_corpus:
1565-
'projects/my-project/locations/us-central1/ragCorpora/my-rag-corpus',
1566-
},
1563+
rag_resources: [
1564+
{
1565+
rag_corpus:
1566+
'projects/my-project/locations/us-central1/ragCorpora/my-rag-corpus',
1567+
},
1568+
],
15671569
similarity_top_k: 5,
15681570
},
15691571
},
15701572
},
15711573
],
15721574
});
15731575
});
1576+
1577+
it('should use vertexRagStore with ragFileIds for gemini-2.0-pro', async () => {
1578+
prepareJsonResponse({
1579+
url: TEST_URL_GEMINI_2_0_PRO,
1580+
});
1581+
1582+
const gemini2Pro = provider.languageModel('gemini-2.0-pro');
1583+
await gemini2Pro.doGenerate({
1584+
prompt: TEST_PROMPT,
1585+
tools: [
1586+
{
1587+
type: 'provider',
1588+
id: 'google.vertex_rag_store',
1589+
name: 'vertex_rag_store',
1590+
args: {
1591+
ragCorpus:
1592+
'projects/my-project/locations/us-central1/ragCorpora/my-rag-corpus',
1593+
ragFileIds: ['file-1', 'file-2'],
1594+
topK: 10,
1595+
},
1596+
},
1597+
],
1598+
});
1599+
1600+
expect(await server.calls[0].requestBodyJson).toMatchObject({
1601+
tools: [
1602+
{
1603+
retrieval: {
1604+
vertex_rag_store: {
1605+
rag_resources: [
1606+
{
1607+
rag_corpus:
1608+
'projects/my-project/locations/us-central1/ragCorpora/my-rag-corpus',
1609+
rag_file_ids: ['file-1', 'file-2'],
1610+
},
1611+
],
1612+
similarity_top_k: 10,
1613+
},
1614+
},
1615+
},
1616+
],
1617+
});
1618+
});
15741619
});
15751620

15761621
it('should extract image file outputs', async () => {

‎packages/google/src/google-prepare-tools.test.ts‎

Lines changed: 191 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -402,3 +402,194 @@ it('should handle url context tool alone', () => {
402402
expect(result.toolConfig).toBeUndefined();
403403
expect(result.toolWarnings).toEqual([]);
404404
});
405+
406+
it('should handle vertex rag store tool', () => {
407+
const result = prepareTools({
408+
tools: [
409+
{
410+
type: 'provider',
411+
id: 'google.vertex_rag_store',
412+
name: 'vertex_rag_store',
413+
args: {
414+
ragCorpus:
415+
'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
416+
topK: 5,
417+
},
418+
},
419+
],
420+
modelId: 'gemini-2.5-flash',
421+
});
422+
expect(result.tools).toStrictEqual([
423+
{
424+
retrieval: {
425+
vertex_rag_store: {
426+
rag_resources: [
427+
{
428+
rag_corpus:
429+
'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
430+
},
431+
],
432+
similarity_top_k: 5,
433+
},
434+
},
435+
},
436+
]);
437+
expect(result.toolConfig).toBeUndefined();
438+
expect(result.toolWarnings).toEqual([]);
439+
});
440+
441+
it('should handle vertex rag store tool with ragFileIds', () => {
442+
const result = prepareTools({
443+
tools: [
444+
{
445+
type: 'provider',
446+
id: 'google.vertex_rag_store',
447+
name: 'vertex_rag_store',
448+
args: {
449+
ragCorpus:
450+
'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
451+
ragFileIds: ['file-1', 'file-2'],
452+
topK: 10,
453+
},
454+
},
455+
],
456+
modelId: 'gemini-2.5-flash',
457+
});
458+
expect(result.tools).toStrictEqual([
459+
{
460+
retrieval: {
461+
vertex_rag_store: {
462+
rag_resources: [
463+
{
464+
rag_corpus:
465+
'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
466+
rag_file_ids: ['file-1', 'file-2'],
467+
},
468+
],
469+
similarity_top_k: 10,
470+
},
471+
},
472+
},
473+
]);
474+
expect(result.toolConfig).toBeUndefined();
475+
expect(result.toolWarnings).toEqual([]);
476+
});
477+
478+
it('should handle vertex rag store tool with metadataFilter', () => {
479+
const result = prepareTools({
480+
tools: [
481+
{
482+
type: 'provider',
483+
id: 'google.vertex_rag_store',
484+
name: 'vertex_rag_store',
485+
args: {
486+
ragCorpus:
487+
'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
488+
metadataFilter: 'user_id = "user-123" AND project = "acme"',
489+
topK: 10,
490+
},
491+
},
492+
],
493+
modelId: 'gemini-2.5-flash',
494+
});
495+
expect(result.tools).toStrictEqual([
496+
{
497+
retrieval: {
498+
vertex_rag_store: {
499+
rag_resources: [
500+
{
501+
rag_corpus:
502+
'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
503+
},
504+
],
505+
similarity_top_k: 10,
506+
rag_retrieval_config: {
507+
filter: {
508+
metadata_filter: 'user_id = "user-123" AND project = "acme"',
509+
},
510+
},
511+
},
512+
},
513+
},
514+
]);
515+
expect(result.toolConfig).toBeUndefined();
516+
expect(result.toolWarnings).toEqual([]);
517+
});
518+
519+
it('should handle vertex rag store tool with vectorSimilarityThreshold', () => {
520+
const result = prepareTools({
521+
tools: [
522+
{
523+
type: 'provider',
524+
id: 'google.vertex_rag_store',
525+
name: 'vertex_rag_store',
526+
args: {
527+
ragCorpus:
528+
'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
529+
vectorSimilarityThreshold: 0.7,
530+
},
531+
},
532+
],
533+
modelId: 'gemini-2.5-flash',
534+
});
535+
expect(result.tools).toStrictEqual([
536+
{
537+
retrieval: {
538+
vertex_rag_store: {
539+
rag_resources: [
540+
{
541+
rag_corpus:
542+
'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
543+
},
544+
],
545+
rag_retrieval_config: {
546+
filter: {
547+
vector_similarity_threshold: 0.7,
548+
},
549+
},
550+
},
551+
},
552+
},
553+
]);
554+
expect(result.toolConfig).toBeUndefined();
555+
expect(result.toolWarnings).toEqual([]);
556+
});
557+
558+
it('should handle vertex rag store tool with vectorDistanceThreshold', () => {
559+
const result = prepareTools({
560+
tools: [
561+
{
562+
type: 'provider',
563+
id: 'google.vertex_rag_store',
564+
name: 'vertex_rag_store',
565+
args: {
566+
ragCorpus:
567+
'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
568+
vectorDistanceThreshold: 0.5,
569+
},
570+
},
571+
],
572+
modelId: 'gemini-2.5-flash',
573+
});
574+
expect(result.tools).toStrictEqual([
575+
{
576+
retrieval: {
577+
vertex_rag_store: {
578+
rag_resources: [
579+
{
580+
rag_corpus:
581+
'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
582+
},
583+
],
584+
rag_retrieval_config: {
585+
filter: {
586+
vector_distance_threshold: 0.5,
587+
},
588+
},
589+
},
590+
},
591+
},
592+
]);
593+
expect(result.toolConfig).toBeUndefined();
594+
expect(result.toolWarnings).toEqual([]);
595+
});

‎packages/google/src/google-prepare-tools.ts‎

Lines changed: 33 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -136,13 +136,42 @@ export function prepareTools({
136136
break;
137137
case 'google.vertex_rag_store':
138138
if (isGemini2orNewer) {
139+
const hasFilter =
140+
tool.args.metadataFilter ||
141+
tool.args.vectorSimilarityThreshold !== undefined ||
142+
tool.args.vectorDistanceThreshold !== undefined;
143+
139144
googleTools.push({
140145
retrieval: {
141146
vertex_rag_store: {
142-
rag_resources: {
143-
rag_corpus: tool.args.ragCorpus,
144-
},
145-
similarity_top_k: tool.args.topK as number | undefined,
147+
rag_resources: [
148+
{
149+
rag_corpus: tool.args.ragCorpus,
150+
...(tool.args.ragFileIds && {
151+
rag_file_ids: tool.args.ragFileIds,
152+
}),
153+
},
154+
],
155+
...(tool.args.topK !== undefined && {
156+
similarity_top_k: tool.args.topK,
157+
}),
158+
...(hasFilter && {
159+
rag_retrieval_config: {
160+
filter: {
161+
...(tool.args.metadataFilter && {
162+
metadata_filter: tool.args.metadataFilter,
163+
}),
164+
...(tool.args.vectorSimilarityThreshold !== undefined && {
165+
vector_similarity_threshold:
166+
tool.args.vectorSimilarityThreshold,
167+
}),
168+
...(tool.args.vectorDistanceThreshold !== undefined && {
169+
vector_distance_threshold:
170+
tool.args.vectorDistanceThreshold,
171+
}),
172+
},
173+
},
174+
}),
146175
},
147176
},
148177
});

‎packages/google/src/tool/vertex-rag-store.ts‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,19 +13,48 @@ export const vertexRagStore = createProviderToolFactory<
1313
{},
1414
{
1515
/**
16-
* RagCorpus resource names, eg: projects/{project}/locations/{location}/ragCorpora/{rag_corpus}
16+
* RagCorpus resource name, eg: projects/{project}/locations/{location}/ragCorpora/{rag_corpus}
1717
*/
1818
ragCorpus: string;
1919

20+
/**
21+
* File IDs within the corpus to search. When specified, only these files
22+
* are searched instead of the entire corpus. For large-scale filtering,
23+
* prefer using `metadataFilter` instead.
24+
*/
25+
ragFileIds?: string[];
26+
2027
/**
2128
* The number of top contexts to retrieve.
2229
*/
2330
topK?: number;
31+
32+
/**
33+
* Filter expression for metadata. Use this to filter results by tags/metadata
34+
* assigned to files when they were uploaded.
35+
* @example 'user_id = "user-123" AND project = "acme"'
36+
*/
37+
metadataFilter?: string;
38+
39+
/**
40+
* Only return results with vector similarity larger than this threshold.
41+
* Value should be between 0 and 1.
42+
*/
43+
vectorSimilarityThreshold?: number;
44+
45+
/**
46+
* Only return results with vector distance smaller than this threshold.
47+
*/
48+
vectorDistanceThreshold?: number;
2449
}
2550
>({
2651
id: 'google.vertex_rag_store',
2752
inputSchema: z.object({
2853
ragCorpus: z.string(),
54+
ragFileIds: z.array(z.string()).optional(),
2955
topK: z.number().optional(),
56+
metadataFilter: z.string().optional(),
57+
vectorSimilarityThreshold: z.number().optional(),
58+
vectorDistanceThreshold: z.number().optional(),
3059
}),
3160
});

0 commit comments

Comments
 (0)