RAG:几十行代码的企业知识库

一分钟能跑通的 RAG

import { MDocument } from '@mastra/rag';
import { PgVector } from '@mastra/pg';
import { openai } from '@ai-sdk/openai';
import { embedMany } from 'ai';

// 1. 读文档 + 切块
const doc = MDocument.fromMarkdown(rawMarkdown);
const chunks = await doc.chunk({
  strategy: 'recursive',
  size: 512,
  overlap: 50,
});

// 2. 批量 embedding
const { embeddings } = await embedMany({
  model: openai.embedding('text-embedding-3-small'),
  values: chunks.map(c => c.text),
});

// 3. 写入向量库
const store = new PgVector({ connectionString: DB });
await store.createIndex({ indexName: 'kb', dimension: 1536 });
await store.upsert({
  indexName: 'kb',
  vectors: embeddings,
  metadata: chunks.map(c => ({ text: c.text, source: c.metadata.source })),
});

MDocument:多源文档抽象

Chunk 策略

strategy	说明	适合
recursive	多级分隔符递归切,默认首选	通用文本
character	按字符数硬切	超长无结构文本
token	按 token 切(需 tiktoken)	严格控制 token 预算
markdown	按 heading/code/list 语义	技术文档、README
html	按 DOM 结构语义	网页爬虫
json	按对象路径	配置、schema

Vector Store 一致 API

Mastra 对所有向量库抽象成同一接口,createIndex / upsert / query / delete,换 provider 只改一行:

// 本地开发:libsql
import { LibSQLVector } from '@mastra/libsql';
const store = new LibSQLVector({ connectionUrl: 'file:./kb.db' });

// 生产:PostgreSQL + pgvector
import { PgVector } from '@mastra/pg';
const store = new PgVector({ connectionString: DB });

// 大规模:Pinecone
import { PineconeVector } from '@mastra/pinecone';
const store = new PineconeVector({ apiKey: KEY });

// 自托管:Qdrant
import { QdrantVector } from '@mastra/qdrant';
const store = new QdrantVector({ url: 'http://localhost:6333' });

查询

const { embedding } = await embed({
  model: openai.embedding('text-embedding-3-small'),
  value: '如何配置部署到 Cloudflare?',
});

const hits = await store.query({
  indexName: 'kb',
  queryVector: embedding,
  topK: 5,
  filter: { source: { $eq: 'docs' } },  // 元数据过滤,mongo 语法
});

for (const h of hits) {
  console.log(h.score, h.metadata.text);
}

把 RAG 做成 Tool

import { createVectorQueryTool } from '@mastra/rag';

export const kbSearchTool = createVectorQueryTool({
  id: 'kb-search',
  description: '搜索公司内部知识库',
  vectorStoreName: 'kb-store',
  indexName: 'kb',
  model: openai.embedding('text-embedding-3-small'),
});

export const mastra = new Mastra({
  agents: { qaAgent },
  vectors: { 'kb-store': new PgVector({ connectionString: DB }) },
});

Agent 挂上 kbSearchTool 后,就会在回答前先检索知识库,再用检索到的 chunks 作为上下文生成回答——这是最经典的 RAG pipeline。

Query Rewriter

用户问"上次的事怎么处理"——embedding 这种模糊 query 召回率低。Mastra 自带 rewriter,让 LLM 把原问题重写成检索友好的形式:

const tool = createVectorQueryTool({
  ...,
  enableFilter: true,      // 允许 LLM 推断 metadata 过滤
  reranker: {
    model: openai('gpt-4o-mini'),
    options: { topK: 3 },
  },
});

Rerank(两段式检索)

import { rerank } from '@mastra/rag';

// 第一段:向量召回 20 条
const candidates = await store.query({ ..., topK: 20 });

// 第二段:LLM 精排 topK=5
const reranked = await rerank({
  results: candidates,
  query: userQuery,
  model: openai('gpt-4o-mini'),
  options: { topK: 5, weights: { semantic: 0.6, vector: 0.4 } },
});

向量相似度 + LLM 判断组合打分。纯向量 "文本相近但实际答非所问" 的情况被压下去,整体精度普遍 +10-15%。

Graph RAG(关联图谱)

Mastra 的 createGraphRAGTool 会在 chunks 之间建邻接边(同文档、语义相似、显式引用),检索时沿边扩张,解决「需要多段凑才能答」的问题。

import { createGraphRAGTool } from '@mastra/rag';

export const kbGraphTool = createGraphRAGTool({
  id: 'kb-graph',
  vectorStoreName: 'kb-store',
  indexName: 'kb',
  model: openai.embedding('text-embedding-3-small'),
  graphOptions: {
    threshold: 0.75,
    dimension: 1536,
  },
});

索引更新策略

成本估算

text-embedding-3-small:$0.02 / 1M tokens;1MB 英文约 250k tokens。10GB 知识库一次全量 embedding 约 $50。查询端更便宜:每次 query 约 0.001 分。

本章小结