feat(google): add filtering options to vertexRagStore tool

koistya · koistya · commit 231f70206d73 · 2025-12-08T22:27:49.000+01:00
diff --git a/.changeset/feat-google-rag-file-ids.md b/.changeset/feat-google-rag-file-ids.md
@@ -0,0 +1,6 @@
+---
+'@ai-sdk/google': patch
+'@ai-sdk/google-vertex': patch
+---
+
+feat(google): add ragFileIds, metadataFilter, and threshold options to vertexRagStore
diff --git a/content/providers/01-ai-sdk-providers/15-google-generative-ai.mdx b/content/providers/01-ai-sdk-providers/15-google-generative-ai.mdx
@@ -767,11 +767,31 @@ The `vertexRagStore` tool accepts the following configuration options:
   - The RagCorpus resource name in the format: `projects/{project}/locations/{location}/ragCorpora/{rag_corpus}`
   - This identifies your specific RAG corpus to search against
 
+- **`ragFileIds`** (`string[]`, optional)
+
+  - File IDs within the corpus to search
+  - When specified, only these files are searched instead of the entire corpus
+  - For large-scale filtering, prefer using `metadataFilter` instead
+
 - **`topK`** (`number`, optional)
 
   - The number of top contexts to retrieve from your RAG corpus
   - Defaults to the corpus configuration if not specified
 
+- **`metadataFilter`** (`string`, optional)
+
+  - Filter expression for metadata assigned to files when uploaded
+  - Example: `'user_id = "user-123" AND project = "acme"'`
+
+- **`vectorSimilarityThreshold`** (`number`, optional)
+
+  - Only return results with vector similarity larger than this threshold
+  - Value should be between 0 and 1
+
+- **`vectorDistanceThreshold`** (`number`, optional)
+
+  - Only return results with vector distance smaller than this threshold
+
 ### Image Outputs
 
 Gemini models with image generation capabilities (`gemini-2.5-flash-image-preview`) support image generation. Images are exposed as files in the response.
diff --git a/packages/google/src/google-generative-ai-language-model.test.ts b/packages/google/src/google-generative-ai-language-model.test.ts
@@ -1560,17 +1560,62 @@ describe('doGenerate', () => {
           {
             retrieval: {
               vertex_rag_store: {
-                rag_resources: {
-                  rag_corpus:
-                    'projects/my-project/locations/us-central1/ragCorpora/my-rag-corpus',
-                },
+                rag_resources: [
+                  {
+                    rag_corpus:
+                      'projects/my-project/locations/us-central1/ragCorpora/my-rag-corpus',
+                  },
+                ],
                 similarity_top_k: 5,
               },
             },
           },
         ],
       });
     });
+
+    it('should use vertexRagStore with ragFileIds for gemini-2.0-pro', async () => {
+      prepareJsonResponse({
+        url: TEST_URL_GEMINI_2_0_PRO,
+      });
+
+      const gemini2Pro = provider.languageModel('gemini-2.0-pro');
+      await gemini2Pro.doGenerate({
+        prompt: TEST_PROMPT,
+        tools: [
+          {
+            type: 'provider',
+            id: 'google.vertex_rag_store',
+            name: 'vertex_rag_store',
+            args: {
+              ragCorpus:
+                'projects/my-project/locations/us-central1/ragCorpora/my-rag-corpus',
+              ragFileIds: ['file-1', 'file-2'],
+              topK: 10,
+            },
+          },
+        ],
+      });
+
+      expect(await server.calls[0].requestBodyJson).toMatchObject({
+        tools: [
+          {
+            retrieval: {
+              vertex_rag_store: {
+                rag_resources: [
+                  {
+                    rag_corpus:
+                      'projects/my-project/locations/us-central1/ragCorpora/my-rag-corpus',
+                    rag_file_ids: ['file-1', 'file-2'],
+                  },
+                ],
+                similarity_top_k: 10,
+              },
+            },
+          },
+        ],
+      });
+    });
   });
 
   it('should extract image file outputs', async () => {
diff --git a/packages/google/src/google-prepare-tools.test.ts b/packages/google/src/google-prepare-tools.test.ts
@@ -402,3 +402,194 @@ it('should handle url context tool alone', () => {
   expect(result.toolConfig).toBeUndefined();
   expect(result.toolWarnings).toEqual([]);
 });
+
+it('should handle vertex rag store tool', () => {
+  const result = prepareTools({
+    tools: [
+      {
+        type: 'provider',
+        id: 'google.vertex_rag_store',
+        name: 'vertex_rag_store',
+        args: {
+          ragCorpus:
+            'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
+          topK: 5,
+        },
+      },
+    ],
+    modelId: 'gemini-2.5-flash',
+  });
+  expect(result.tools).toStrictEqual([
+    {
+      retrieval: {
+        vertex_rag_store: {
+          rag_resources: [
+            {
+              rag_corpus:
+                'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
+            },
+          ],
+          similarity_top_k: 5,
+        },
+      },
+    },
+  ]);
+  expect(result.toolConfig).toBeUndefined();
+  expect(result.toolWarnings).toEqual([]);
+});
+
+it('should handle vertex rag store tool with ragFileIds', () => {
+  const result = prepareTools({
+    tools: [
+      {
+        type: 'provider',
+        id: 'google.vertex_rag_store',
+        name: 'vertex_rag_store',
+        args: {
+          ragCorpus:
+            'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
+          ragFileIds: ['file-1', 'file-2'],
+          topK: 10,
+        },
+      },
+    ],
+    modelId: 'gemini-2.5-flash',
+  });
+  expect(result.tools).toStrictEqual([
+    {
+      retrieval: {
+        vertex_rag_store: {
+          rag_resources: [
+            {
+              rag_corpus:
+                'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
+              rag_file_ids: ['file-1', 'file-2'],
+            },
+          ],
+          similarity_top_k: 10,
+        },
+      },
+    },
+  ]);
+  expect(result.toolConfig).toBeUndefined();
+  expect(result.toolWarnings).toEqual([]);
+});
+
+it('should handle vertex rag store tool with metadataFilter', () => {
+  const result = prepareTools({
+    tools: [
+      {
+        type: 'provider',
+        id: 'google.vertex_rag_store',
+        name: 'vertex_rag_store',
+        args: {
+          ragCorpus:
+            'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
+          metadataFilter: 'user_id = "user-123" AND project = "acme"',
+          topK: 10,
+        },
+      },
+    ],
+    modelId: 'gemini-2.5-flash',
+  });
+  expect(result.tools).toStrictEqual([
+    {
+      retrieval: {
+        vertex_rag_store: {
+          rag_resources: [
+            {
+              rag_corpus:
+                'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
+            },
+          ],
+          similarity_top_k: 10,
+          rag_retrieval_config: {
+            filter: {
+              metadata_filter: 'user_id = "user-123" AND project = "acme"',
+            },
+          },
+        },
+      },
+    },
+  ]);
+  expect(result.toolConfig).toBeUndefined();
+  expect(result.toolWarnings).toEqual([]);
+});
+
+it('should handle vertex rag store tool with vectorSimilarityThreshold', () => {
+  const result = prepareTools({
+    tools: [
+      {
+        type: 'provider',
+        id: 'google.vertex_rag_store',
+        name: 'vertex_rag_store',
+        args: {
+          ragCorpus:
+            'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
+          vectorSimilarityThreshold: 0.7,
+        },
+      },
+    ],
+    modelId: 'gemini-2.5-flash',
+  });
+  expect(result.tools).toStrictEqual([
+    {
+      retrieval: {
+        vertex_rag_store: {
+          rag_resources: [
+            {
+              rag_corpus:
+                'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
+            },
+          ],
+          rag_retrieval_config: {
+            filter: {
+              vector_similarity_threshold: 0.7,
+            },
+          },
+        },
+      },
+    },
+  ]);
+  expect(result.toolConfig).toBeUndefined();
+  expect(result.toolWarnings).toEqual([]);
+});
+
+it('should handle vertex rag store tool with vectorDistanceThreshold', () => {
+  const result = prepareTools({
+    tools: [
+      {
+        type: 'provider',
+        id: 'google.vertex_rag_store',
+        name: 'vertex_rag_store',
+        args: {
+          ragCorpus:
+            'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
+          vectorDistanceThreshold: 0.5,
+        },
+      },
+    ],
+    modelId: 'gemini-2.5-flash',
+  });
+  expect(result.tools).toStrictEqual([
+    {
+      retrieval: {
+        vertex_rag_store: {
+          rag_resources: [
+            {
+              rag_corpus:
+                'projects/my-project/locations/us-central1/ragCorpora/my-corpus',
+            },
+          ],
+          rag_retrieval_config: {
+            filter: {
+              vector_distance_threshold: 0.5,
+            },
+          },
+        },
+      },
+    },
+  ]);
+  expect(result.toolConfig).toBeUndefined();
+  expect(result.toolWarnings).toEqual([]);
+});
diff --git a/packages/google/src/google-prepare-tools.ts b/packages/google/src/google-prepare-tools.ts
@@ -136,13 +136,42 @@ export function prepareTools({
           break;
         case 'google.vertex_rag_store':
           if (isGemini2orNewer) {
+            const hasFilter =
+              tool.args.metadataFilter ||
+              tool.args.vectorSimilarityThreshold !== undefined ||
+              tool.args.vectorDistanceThreshold !== undefined;
+
             googleTools.push({
               retrieval: {
                 vertex_rag_store: {
-                  rag_resources: {
-                    rag_corpus: tool.args.ragCorpus,
-                  },
-                  similarity_top_k: tool.args.topK as number | undefined,
+                  rag_resources: [
+                    {
+                      rag_corpus: tool.args.ragCorpus,
+                      ...(tool.args.ragFileIds && {
+                        rag_file_ids: tool.args.ragFileIds,
+                      }),
+                    },
+                  ],
+                  ...(tool.args.topK !== undefined && {
+                    similarity_top_k: tool.args.topK,
+                  }),
+                  ...(hasFilter && {
+                    rag_retrieval_config: {
+                      filter: {
+                        ...(tool.args.metadataFilter && {
+                          metadata_filter: tool.args.metadataFilter,
+                        }),
+                        ...(tool.args.vectorSimilarityThreshold !== undefined && {
+                          vector_similarity_threshold:
+                            tool.args.vectorSimilarityThreshold,
+                        }),
+                        ...(tool.args.vectorDistanceThreshold !== undefined && {
+                          vector_distance_threshold:
+                            tool.args.vectorDistanceThreshold,
+                        }),
+                      },
+                    },
+                  }),
                 },
               },
             });
diff --git a/packages/google/src/tool/vertex-rag-store.ts b/packages/google/src/tool/vertex-rag-store.ts
@@ -13,19 +13,48 @@ export const vertexRagStore = createProviderToolFactory<
   {},
   {
     /**
-     * RagCorpus resource names, eg: projects/{project}/locations/{location}/ragCorpora/{rag_corpus}
+     * RagCorpus resource name, eg: projects/{project}/locations/{location}/ragCorpora/{rag_corpus}
      */
     ragCorpus: string;
 
+    /**
+     * File IDs within the corpus to search. When specified, only these files
+     * are searched instead of the entire corpus. For large-scale filtering,
+     * prefer using `metadataFilter` instead.
+     */
+    ragFileIds?: string[];
+
     /**
      * The number of top contexts to retrieve.
      */
     topK?: number;
+
+    /**
+     * Filter expression for metadata. Use this to filter results by tags/metadata
+     * assigned to files when they were uploaded.
+     * @example 'user_id = "user-123" AND project = "acme"'
+     */
+    metadataFilter?: string;
+
+    /**
+     * Only return results with vector similarity larger than this threshold.
+     * Value should be between 0 and 1.
+     */
+    vectorSimilarityThreshold?: number;
+
+    /**
+     * Only return results with vector distance smaller than this threshold.
+     */
+    vectorDistanceThreshold?: number;
   }
 >({
   id: 'google.vertex_rag_store',
   inputSchema: z.object({
     ragCorpus: z.string(),
+    ragFileIds: z.array(z.string()).optional(),
     topK: z.number().optional(),
+    metadataFilter: z.string().optional(),
+    vectorSimilarityThreshold: z.number().optional(),
+    vectorDistanceThreshold: z.number().optional(),
   }),
 });
