Mastra RAG:createVectorQueryTool 数据库专属配置、多租户动态解析与 Bedrock 知识库工具
2026/9/13 16:42:45 网站建设 项目流程

Mastra RAG:createVectorQueryTool 数据库专属配置、多租户动态解析与 Bedrock 知识库工具

【免费下载链接】mastraMastra is the modern TypeScript framework for AI-powered applications and agents.项目地址: https://gitcode.com/GitHub_Trending/ma/mastra

本文围绕@mastra/rag包中的向量查询工具 createVectorQueryTool 展开,系统讲解databaseConfig数据库专属配置(Pinecone、pgVector、Chroma、MongoDB、Turbopuffer)的用法与底层参数映射机制、基于RequestContext的运行时配置覆盖、面向多租户应用的动态向量库解析器(VectorStoreResolver),以及配套的createBedrockKBTool托管知识库工具。读完本文,你能够针对不同向量数据库精确调优检索行为,在请求粒度动态切换 namespace/租户数据,并理解这些能力在源码中的真实调用链与默认值。

工具的基本形态与输入输出契约

createVectorQueryTool@mastra/rag/tools导出的工厂函数,用于把一个向量库索引包装成可被 Agent 调用的 Mastra Tool。它支持两种互斥的向量库定位方式(判别联合类型):

  • 按名称:传vectorStoreName,从 Mastra 实例中查找已注册的向量库(mastra.getVector(name));
  • 直接实例或解析函数:传vectorStore,可以是MastraVector实例,也可以是动态解析器函数(多租户场景,详见下文)。

工具的输入输出 schema 定义在 tool-schemas.ts 中:

  • 输入基础 schema(baseSchema):queryText(检索文本)与topK(返回条数,经z.coerce.number()强制转为数字);
  • 当构造时传入enableFilter: true时,输入 schema 切换为filterSchema,额外暴露一个filter字段(字符串形式的 JSON 过滤条件);
  • 输出 schema 固定为{ relevantContext, sources }relevantContext是结果元数据数组,sources是完整的检索结果对象数组,每项包含idmetadatavectorscoredocument字段。

构造工具时的完整选项类型是 VectorQueryToolOptions,关键字段如下:

字段说明默认值
indexName向量库内索引名(必填)
model用于把查询文本转为向量的嵌入模型(必填)
vectorStoreName/vectorStore二选一:注册名 或 实例/解析器
id工具 IDVectorQuery {storeName} {indexName} Tool
description给 LLM 看的工具描述内置默认描述
enableFilter是否在输入 schema 中暴露filter字段false
includeVectors结果是否包含嵌入向量false
includeSources响应是否包含sourcestrue
reranker重排器配置(RerankConfig
databaseConfig数据库专属配置,详见下节
providerOptions嵌入模型的提供方选项(注意:仅 AI SDK v2 模型可用,v1 模型应在创建模型时配置)

databaseConfig:为不同向量数据库传递专属参数

不同向量数据库的检索 API 各有独特的可调参数——Pinecone 关心 namespace 与稀疏向量,pgVector 关心 HNSW/IVFFlat 搜索参数,Chroma 关心元数据过滤。databaseConfig字段就是为此设计的:一个以数据库名为键的对象,把各库专属参数透传给vectorStore.query()

各数据库配置类型全部在 types.ts 中定义并导出:

Pinecone 配置

PineconeConfig支持两个字段:namespace(Pinecone 命名空间)与sparseVectorindices/values数组对,用于混合检索)。

import { createVectorQueryTool } from '@mastra/rag/tools'; const pineconeVectorTool = createVectorQueryTool({ id: 'pinecone-search', indexName: 'my-index', vectorStoreName: 'pinecone', model: embedModel, databaseConfig: { pinecone: { namespace: 'my-namespace', // Pinecone namespace sparseVector: { // For hybrid search indices: [0, 1, 2], values: [0.1, 0.2, 0.3], }, }, }, });

pgVector 配置

PgVectorConfig支持三个字段:minScore(最低相似度分数)、ef(HNSW 搜索参数)、probes(IVFFlat 探测参数)。

const pgVectorTool = createVectorQueryTool({ id: 'pgvector-search', indexName: 'my-index', vectorStoreName: 'postgres', model: embedModel, databaseConfig: { pgvector: { minScore: 0.7, // Minimum similarity score ef: 200, // HNSW search parameter probes: 10, // IVFFlat probe parameter }, }, });

Chroma 配置

ChromaConfig支持where(元数据过滤,类型完整建模了 Chroma 的$and/$or/$in/$gt等运算符体系)与whereDocument(文档内容过滤,支持$contains/$not_contains)。

const chromaTool = createVectorQueryTool({ id: 'chroma-search', indexName: 'my-index', vectorStoreName: 'chroma', model: embedModel, databaseConfig: { chroma: { where: { // Metadata filtering category: 'documents', }, whereDocument: { // Document content filtering $contains: 'important', }, }, }, });

注意这与工具级enableFilter+ 输入filter字段是两套机制:where/whereDocument是 Chroma 原生的过滤语法,直接透传;而filter字段走的是 Mastra 通用的 VectorFilter 解析路径。

MongoDB 与 Turbopuffer 配置

除文档示例的三种数据库外,DatabaseConfig还内置了另外两种(见 types.ts,并有对应测试用例覆盖):

const mongoTool = createVectorQueryTool({ vectorStoreName: 'mongodb', indexName: 'my-index', model: embedModel, databaseConfig: { mongodb: { numCandidates: 500, // HNSW 候选数,须 >= topK;默认 20 * topK,上限 10000 }, }, }); const turboTool = createVectorQueryTool({ vectorStoreName: 'turbopuffer', indexName: 'my-index', model: embedModel, databaseConfig: { turbopuffer: { consistency: 'eventual', // 'strong'(默认)或 'eventual'(更低延迟) }, }, });

DatabaseConfig类型本身带[key: string]: any索引签名,允许为未来新数据库任意扩展键:

export type DatabaseConfig = { pinecone?: PineconeConfig; pgvector?: PgVectorConfig; chroma?: ChromaConfig; mongodb?: MongoDBConfig; turbopuffer?: TurbopufferConfig; // Add other database configs as needed [key: string]: any; // Allow for future database extensions };

参数如何抵达 query 调用:源码级映射

配置并不是原样透传的。在 vector-search.ts 的databaseSpecificParams()中,框架按数据库名把嵌套配置“摊平”为vectorStore.query()能直接识别的顶层参数:

  • pinecone.namespacenamespacepinecone.sparseVectorsparseVector
  • pgvector.minScore/ef/probes→ 同名顶层参数;
  • chroma.where/whereDocument→ 同名顶层参数;
  • mongodb.numCandidatesnumCandidates
  • turbopuffer.consistencyconsistency

最终在 vector-search.ts 处合并进查询参数:

results = await vectorStore.query({ ...queryParams, ...databaseSpecificParams(databaseConfig) });

其中queryParamsindexNamequeryVector(嵌入结果)、topKfilterincludeVector组成。

运行时覆盖:用 RequestContext 按请求改写配置

工具级databaseConfig是静态的;若需要按请求动态调整(例如切换 Pinecone namespace 到不同环境的数据分区),可以在RequestContext中设置同名键databaseConfig。在 vector-query.ts 中,几乎所有运行时变量都遵循“requestContext 优先、options 兜底”的取值顺序:

const indexName: string = requestContext?.get('indexName') ?? options.indexName; const databaseConfig = requestContext?.get('databaseConfig') ?? options.databaseConfig; const model: MastraEmbeddingModel<string> = requestContext?.get('model') ?? options.model; const topK: number = requestContext?.get('topK') ?? (inputData.topK as number) ?? 10; // includeVectors、includeSources、reranker、filter、providerOptions 同理

可覆盖的键包括indexNamevectorStoreNameincludeVectorsincludeSourcesrerankerdatabaseConfigmodelproviderOptionstopKfilter。基于测试用例 vector-query-database-config.test.ts 验证过的运行时覆盖写法如下:

import { RequestContext } from '@mastra/core/request-context'; const tool = createVectorQueryTool({ vectorStoreName: 'pinecone', indexName: 'testIndex', model: embedModel, databaseConfig: { pinecone: { namespace: 'initial-namespace' } }, }); // 运行时覆盖 Pinecone namespace const requestContext = new RequestContext(); requestContext.set('databaseConfig', { pinecone: { namespace: 'runtime-namespace' }, }); const result = await tool.execute( { queryText: 'test query', topK: 5 }, { mastra, requestContext }, );

一个需要注意的语义细节:由于取值是??而非深合并,requestContext中一旦设置了databaseConfig,将整体替换工具级配置,而不是逐字段合并。测试用例明确断言了运行时配置会完整取代初始配置(databaseConfig: runtimeConfig)。

多租户应用:动态向量库解析器 VectorStoreResolver

对于每个租户数据隔离的场景(例如各租户使用独立的 PostgreSQL schema),vectorStore除了接收静态实例外,还可以接收一个解析器函数。类型定义在 types.ts:

export interface VectorStoreResolverContext { requestContext?: RequestContext; mastra?: MastraUnion; } export type VectorStoreResolver = ( context: VectorStoreResolverContext ) => MastraVector | Promise<MastraVector>;

解析器在每次execute时收到requestContextmastra,据此返回当次请求应使用的向量库:

import { createVectorQueryTool, VectorStoreResolver } from '@mastra/rag/tools'; import { PgVector } from '@mastra/pg'; // Resolver function receives requestContext and mastra instance const vectorStoreResolver: VectorStoreResolver = async ({ requestContext }) => { const tenantId = requestContext?.get('tenantId'); return new PgVector({ id: `pg-vector-${tenantId}`, connectionString: process.env.POSTGRES_CONNECTION_STRING!, schemaName: `tenant_${tenantId}`, // Each tenant has their own schema }); }; const vectorQueryTool = createVectorQueryTool({ indexName: 'embeddings', model: embedModel, vectorStore: vectorStoreResolver, // Dynamic resolution! }); // Usage with tenant context const requestContext = new RequestContext(); requestContext.set('tenantId', 'acme-corp'); const result = await vectorQueryTool.execute( { queryText: 'search query', topK: 5 }, { requestContext }, );

从源码看,解析逻辑集中在 tool-helpers.ts 的resolveVectorStore()中:

  1. vectorStore是函数,则以{ requestContext, mastra }调用并await其结果;
  2. 返回值经过isValidMastraVector运行时校验(非 null/undefined 的对象),resolver 返回无效值时会抛出带上下文的错误(错误信息会附带vectorStoreNameschemaIdtenantId等诊断信息,便于排查);
  3. 若未提供vectorStore,则回退到mastra.getVector(vectorStoreName)

同样的解析机制对 GraphRAG 工具同样生效——createGraphRAGTool 接受相同的判别联合选项(见 GraphRagToolOptions,含dimension默认 1536、randomWalkSteps默认 100、restartProb默认 0.15、threshold默认 0.7):

import { createGraphRAGTool } from '@mastra/rag/tools'; const graphTool = createGraphRAGTool({ indexName: 'embeddings', model: embedModel, vectorStore: vectorStoreResolver, });

执行流程与容错行为(源码级)

理解工具在execute中做了什么,有助于判断异常时的行为边界。vector-query.ts 的执行链是:

  1. 解析运行时变量(如上节,requestContext 优先);
  2. coerceTopKtopK规整为有限正数,无效值回退默认10(见 tool-helpers.ts);
  3. resolveVectorStore解析向量库。若解析结果为undefined(按名称找不到已注册向量库),工具记录 error 日志并优雅降级——返回{ relevantContext: [], sources: [] }而非抛错;
  4. 嵌入查询文本。vectorQuerySearch 按model.specificationVersion分派到embedV3/embedV2/embedV1,并创建RAG_EMBEDDING观测 span(记录模型、提供方、维度与 token 用量);
  5. 向量查询。携带RAG_VECTOR_OPERATIONspan 调用vectorStore.query(...)
  6. 可选重排。若配置了reranker,则走 rerank/rerankWithScorer(重排后relevantContext取重排结果的元数据,sources由重排结果转换);
  7. 异常兜底。整个流程包在 try/catch 中,任何未预期异常都会记录 error 日志并返回空结果,避免单次检索失败中断整个 Agent 对话。

此外,filter输入如果是字符串,会经 parseFilterValue 做JSON.parse并校验必须是普通对象,解析失败会抛错并记录日志。

扩展新数据库

系统为扩展预留了两条路径(与 README 的说明一致):

  1. 添加类型:为新库定义配置接口并加入DatabaseConfig[key: string]: any索引签名已允许直接加键):
export interface NewDatabaseConfig { customParam1?: string; customParam2?: number; } export type DatabaseConfig = { pinecone?: PineconeConfig; pgvector?: PgVectorConfig; chroma?: ChromaConfig; newdatabase?: NewDatabaseConfig; // Add your config here [key: string]: any; };
  1. 参数透传databaseSpecificParams()的兜底分支(vector-search.ts)会把不在内置五库枚举中的键对应的配置对象整体平铺合并进查询参数:
Object.keys(databaseConfig).forEach(dbName => { if (!DATABASE_TYPE_MAP.includes(dbName)) { // For unknown database types, merge the config directly const config = databaseConfig[dbName]; if (config && typeof config === 'object') { Object.assign(databaseSpecificParams, config); } } });

也就是说,新数据库只要其query()接受的参数名与配置对象的键一致,无需改动核心代码即可透传——类型安全由你自行补充的接口保证。

createBedrockKBTool:Amazon Bedrock 托管知识库

@mastra/rag还导出了 createBedrockKBTool,直接对接 Amazon Bedrock 托管知识库——向量存储、索引与检索基础设施全部由 AWS 托管,无需自行维护任何向量库:

import { createBedrockKBTool } from '@mastra/rag'; const kbTool = createBedrockKBTool({ knowledgeBaseId: 'ABCDEFGHIJ', region: 'us-west-2', }); const results = await kbTool.execute({ queryText: 'What are our policies?' });

也可以像普通工具一样挂到 Agent 上(参考 BEDROCK_MANAGED_KB.md):

import { Agent } from '@mastra/core'; const agent = new Agent({ name: 'research-agent', tools: { kb: kbTool }, instructions: 'Use the knowledge base to answer questions.', });

选项与默认值

选项/环境变量说明默认值
knowledgeBaseIdBedrock 知识库 ID(必填)
regionAWS 区域(对应环境变量AWS_REGIONAWS_REGIONus-east-1
numberOfResults最大返回条数5
useAgenticRetrieval是否启用 Agentic Retrieval(对应环境变量USE_AGENTIC_RETRIEVAL,设为false即关闭)true
userId访问控制用的默认 AWS 用户 ID;请求上下文requestContext.get('userId')优先

凭证走标准 AWS 环境变量(AWS_ACCESS_KEY_IDAWS_SECRET_ACCESS_KEY)。

Agentic Retrieval 与自动回退

从源码(bedrock-knowledge-base.ts)看,useAgenticRetrieval为 true 时调用AgenticRetrieveStreamCommand,配置foundationModelType: 'MANAGED'rerankingModelType: 'MANAGED'——即由 Bedrock 托管模型自动完成查询分解与重排,流式收集结果;一旦失败则catch中自动回退到标准RetrieveCommandmanagedRetrieve),并打印警告。这也解释了 README 中“Agentic Retrieval 默认开启、失败自动降级为标准检索”的描述。关闭方式:

export USE_AGENTIC_RETRIEVAL=false # disable agentic, use standard retrieve

其他实现细节:

  • 工具 ID 固定为bedrock_knowledge_base_${knowledgeBaseId},输入 schema 仅queryText一个字段;
  • 输出的source字段会从检索结果的location中提取来源 URI,支持 S3、Web、Confluence、Salesforce、SharePoint、自定义文档六类来源(见 getSourceUri);
  • 所需 IAM 权限(来自 BEDROCK_MANAGED_KB.md):
{ "Effect": "Allow", "Action": ["bedrock:Retrieve", "bedrock:AgenticRetrieveStream"], "Resource": "arn:aws:bedrock:<region>:<account-id>:knowledge-base/<kb-id>" }
  • SDK 要求:@aws-sdk/client-bedrock-agent-runtime需 3.1000 以上(AgenticRetrieveStreamCommand依赖该版本能力)。

向后兼容:旧代码无需改动

databaseConfig是完全增量式的可选项:不传时行为与旧版本一致。测试用例专门验证了“无databaseConfigvectorQuerySearch收到databaseConfig: undefined”这一向后兼容路径。给既有工具补充数据库配置只需加一个字段:

const vectorTool = createVectorQueryTool({ indexName: 'my-index', vectorStoreName: 'pinecone', model: embedModel, + databaseConfig: { + pinecone: { + namespace: 'my-namespace' + } + } });

测试覆盖与延伸阅读

本文所有关键行为均有仓库内测试佐证:

  • vector-query-database-config.test.ts:覆盖 Pinecone/pgVector/MongoDB/Turbopuffer 配置透传、requestContext 覆盖、无配置时的向后兼容、多库配置共存;
  • vector-query.test.ts:向量查询工具的主流程测试;
  • bedrock-knowledge-base.test.ts:Bedrock 知识库工具的检索与回退测试;
  • 类型与默认值:types.ts、default-settings.ts;
  • 模块导出入口:index.ts。

需要注意的适用前提:databaseConfig的专属参数只有当对应向量库驱动的query()实现支持该参数时才会实际生效,框架只负责透传;providerOptions仅对 AI SDK v2 嵌入模型有效。

【免费下载链接】mastraMastra is the modern TypeScript framework for AI-powered applications and agents.项目地址: https://gitcode.com/GitHub_Trending/ma/mastra

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询