diff --git a/.changeset/feat-google-rag-file-ids.md b/.changeset/feat-google-rag-file-ids.md new file mode 100644 index 000000000000..f3492c0f546f --- /dev/null +++ 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 index 7dfa5bc7b0ae..8e26f39d4ffb 100644 --- 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 index 98e6e722e11a..1f5f06830d41 100644 --- a/packages/google/src/google-generative-ai-language-model.test.ts +++ b/packages/google/src/google-generative-ai-language-model.test.ts @@ -1571,10 +1571,12 @@ 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, }, }, @@ -1582,6 +1584,49 @@ describe('doGenerate', () => { ], }); }); + + 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 index 8102f9e3e2f2..53d421606542 100644 --- 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 index 13b895f5aabc..68753d7f4ca1 100644 --- a/packages/google/src/google-prepare-tools.ts +++ b/packages/google/src/google-prepare-tools.ts @@ -136,13 +136,43 @@ 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 index 4c592651e8b8..dd449d34f129 100644 --- 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(), }), });