想让大模型回答你公司内部文档的问题?直接喂文档太占 Token,而且模型会「忘记」长内容。RAG(检索增强生成) 的做法是:把文档切块、转成向量存起来,用户提问时先检索相关片段,再把片段和问题一起发给模型,让它基于这些上下文回答。实现一套 RAG 服务,核心就三步:文档切块向量化检索


RAG 为什么有用?

大模型有两个局限:上下文长度有限(不能塞太多文档)、知识有时效性(训练数据可能过时)。RAG 的做法是:

  1. 预处理阶段:把文档切成小块,每块转成向量(embedding),存到向量数据库
  2. 查询阶段:用户提问时,把问题也转成向量,在向量库中找最相似的文档块
  3. 生成阶段:把检索到的文档块作为上下文,和问题一起发给模型,让它基于这些内容回答

这样既突破了上下文长度限制,又能随时更新文档库。下面用 Go 实现这三个步骤。


文档切块:把长文档分成小片段

文档切块的原则是:每块大小适中(通常 200-500 字)、尽量保持语义完整(不要从句子中间切断)、可以重叠(相邻块之间重叠一部分,避免边界信息丢失)。

最简单的方式是按字符数切,核心逻辑如下:

type Chunk struct {
    ID       string
    Content  string
    Metadata map[string]string
}

func SplitText(text string, chunkSize, overlap int) []Chunk {
    var chunks []Chunk
    start := 0
    for start < len(text) {
        end := start + chunkSize
        if end > len(text) {
            end = len(text)
        }
        chunks = append(chunks, Chunk{Content: text[start:end]})
        if end >= len(text) {
            break
        }
        start = end - overlap // 重叠处理
    }
    return chunks
}

实际项目中可以按段落切(遇到空行)、按句子切(用标点判断),或使用专门的文本分割库。


向量化:把文本转成向量

向量化就是调用 Embedding API(如 OpenAI 的 text-embedding-3-small),把文本转成一个固定长度的浮点数数组(如 1536 维)。语义相近的文本,向量也相近。

封装一个简单的客户端:

type EmbeddingClient struct {
    APIKey  string
    BaseURL string
}

func (c *EmbeddingClient) Embed(ctx context.Context, text string) ([]float32, error) {
    // POST {BaseURL}/embeddings
    // Body: {"model": "text-embedding-3-small", "input": text}
    // 返回: {"data": [{"embedding": [...]}]}
    // 实际用 http.Client 或第三方 SDK 实现
    return callEmbeddingAPI(text), nil
}

把每个文档块的 Content 传给 Embed,得到向量后存起来。批量调用时可以用 EmbedBatch 提高效率。


检索:用向量相似度找相关文档

检索就是:把用户问题也转成向量,然后和所有文档块的向量算相似度(常用余弦相似度),取最相似的几个块作为上下文。

实现一个简单的内存向量库(实际项目可以用 Pgvector、Milvus、Qdrant 等):

type VectorStore struct {
    chunks  []Chunk
    vectors [][]float32
}

func (vs *VectorStore) Add(chunk Chunk, vector []float32) {
    vs.chunks = append(vs.chunks, chunk)
    vs.vectors = append(vs.vectors, vector)
}

func (vs *VectorStore) Search(queryVector []float32, k int) []Chunk {
    // 遍历算 cosineSimilarity(queryVector, vec),按 score 降序排序后取前 k 个 chunk
    // cosineSimilarity(a,b) = 点积(a,b) / (|a|*|b|)
    return nil
}

实际项目中用专业的向量数据库会更高效(支持索引、批量插入等),但内存版本足够理解原理。


串起来:完整的 RAG 流程

下面把切块、向量化、检索串成一个完整的 RAG 服务:

type RAGService struct {
    embeddingClient *EmbeddingClient
    vectorStore     *VectorStore
}

func (r *RAGService) AddDocument(ctx context.Context, text string) error {
    chunks := SplitText(text, 500, 50)
    vectors, err := r.embeddingClient.EmbedBatch(ctx, extractTexts(chunks))
    if err != nil {
        return err
    }
    for i := range chunks {
        r.vectorStore.Add(chunks[i], vectors[i])
    }
    return nil
}

func (r *RAGService) Query(ctx context.Context, question string, topK int) ([]Chunk, error) {
    queryVector, err := r.embeddingClient.Embed(ctx, question)
    if err != nil {
        return nil, err
    }
    return r.vectorStore.Search(queryVector, topK), nil
}

使用示例:

rag := NewRAGService("api-key", "https://api.openai.com/v1")
rag.AddDocument(ctx, "这是一段很长的文档内容...")
chunks, _ := rag.Query(ctx, "文档里说了什么?", 3)
prompt := fmt.Sprintf("基于以下上下文回答问题:\n%s\n\n问题:%s", 
    chunks[0].Content, question)

注意事项与优化建议

切块策略

  • 按字符数切最简单,但可能切断句子
  • 按段落切(遇到空行)更自然,但块大小不均匀
  • 按句子切最精细,但需要分词库
  • 重叠很重要:相邻块重叠 10-20%,避免边界信息丢失

向量化优化

  • 批量调用:一次传多个文本,比逐个调用快
  • 缓存:相同文本的向量可以缓存,避免重复计算
  • 模型选择text-embedding-3-small 性价比高,text-embedding-3-large 效果更好但更贵

检索优化

  • Top-K 选择:通常取 3-5 个最相似的块,太少可能遗漏信息,太多可能引入噪音
  • 相似度阈值:可以设置最低相似度,过滤掉不相关的结果
  • 向量数据库:数据量大时用 Pgvector(PostgreSQL 插件)、Milvus、Qdrant 等,支持索引和高效检索

实际项目建议

  • 文档更新时,可以增量更新:只重新向量化变更的块
  • 添加元数据过滤:检索时可以按文件名、日期等过滤
  • 混合检索:结合关键词检索(BM25)和向量检索,效果更好

写在最后

RAG 的核心流程就是三步:

  1. 文档切块:把长文档切成 200-500 字的小块,可以重叠
  2. 向量化:调用 Embedding API,把每个块转成向量,存到向量库
  3. 检索:用户提问时,把问题也向量化,用余弦相似度找最相似的几个块,作为上下文发给模型