在Go中实现基于向量数据库的RAG系统：一步步构建本地知识库问答

## 背景介绍

retrieval-augmented generation（检索增强生成，简称RAG）已经成为大模型应用的主流架构。与其让模型从预训练数据中硬编码知识，不如在运行时从外部知识库动态检索相关信息，然后交给大模型生成答案。这种方式有几个明显的好处：知识可以随时更新，不需要重新训练模型；回答可以引用真实的来源；特定领域的问题更容易控制。

在实际开发中，RAG系统的核心是把文档转换成向量，存储到向量数据库，然后通过向量相似度搜索找到最相关的内容，最后把检索到的内容作为上下文提供给大模型。这个流程看起来简单，但实现起来涉及不少细节。

## 问题描述

刚接触RAG系统的开发者通常会遇到几个常见问题。首先是向量模型的选择——开源的、闭源的、本地的、API调用的，各有各的特点，需要根据场景挑选。其次是中文文本的分词和向量化，英文有空格分隔，中文怎么处理是个问题。再次是向量数据库的选择，Milvus、Qdrant、Chroma、PgVector，各有优缺点。还有chunk（文本分块）的大小怎么确定，太小了上下文不完整，太大了可能超出模型的上下文限制。

这篇文章来解决一个具体的问题：**用Go语言构建一个本地运行的RAG系统**，能够读取PDF或Markdown文档，建立知识库，然后用自然语言查询。整个过程不需要调用外部API，不依赖GPU，在普通电脑上就能运行。

—

## 详细步骤

### 第一步：环境准备

我们需要安装Go语言环境（1.21以上版本），然后安装几个必要的依赖包。首先是向量嵌入模型，这里选择bge-base-zh-v1.5，这是一个开源的中文向量模型，在 huggingface 上可以免费获取。其次是最关键的向量数据库，为了简单易用，我们选择Chroma，它是一个纯Go实现的向量数据库，不需要额外安装服务端。

“`bash
go mod init rag-demo
go get github.com/chroma-core/chroma-go
go get github.com/tiktoken-go/tiktoken
go get github.com/go-skydio/lib/pdf
“`

注意：Chroma的Go客户端还在开发中，如果遇到问题，可以用纯SQLite实现一个简化版本，或者使用PgVector。

### 第二步：文档加载和分块

文档加载是RAG系统的第一步。对于不同格式的文档，需要不同的解析方法。Markdown相对简单，直接读取文件内容，按标题层级分割即可。PDF解析稍微复杂，可以使用pdf库或者调用pdftotext命令行工具。

关键的一步是文本分块（chunking）。我们需要一个合适的分块策略：

“`go
type Document struct {
ID string
Content string
Meta map[string]string
}

type Chunker struct {
chunkSize int
overlap int
}

func (c *Chunker) Chunk(text string) []Document {
sentences := splitSentences(text)
var chunks []Document
var currentChunk strings.Builder
currentSize := 0

for _, sent := range sentences {
sentLen := len(sent)
if currentSize + sentLen > c.chunkSize && currentSize > 0 {
chunks = append(chunks, Document{
ID: uuid.New().String(),
Content: currentChunk.String(),
})
overlapText := currentChunk.String()
if len(overlapText) > c.overlap {
overlapText = overlapText[len(overlapText)-c.overlap:]
}
currentChunk.Reset()
currentChunk.WriteString(overlapText)
currentSize = len(overlapText)
}
currentChunk.WriteString(sent)
currentSize += sentLen
}

if currentChunk.Len() > 0 {
chunks = append(chunks, Document{
ID: uuid.New().String(),
Content: currentChunk.String(),
})
}
return chunks
}
“`

这里我们按句子而不是按字符分块，这样可以保持语义的完整性。overlap参数控制相邻块之间的重叠字符数，这样可以避免边界信息丢失。对于中文文本，需要使用中文分词库（如sego）来正确分割句子。

### 第三步：向量化

向量化是把文本转换成向量表示的过程。对于中文文本，我们需要一个支持中文的向量化模型。这里我们使用bge-base-zh-v1.5模型，它在中文语义相似度任务上表现不错。

“`go
type Embedder struct {
model *bert.Embedding
tokenizer *tiktoken.Tokenizer
}

func NewEmbedder() (*Embedder, error) {
model, err := bert.Load(“bge-base-zh-v1.5”)
if err != nil {
return nil, fmt.Errorf(“加载模型失败: %w”, err)
}
tokenizer, err := tiktoken.New(“cl100k_base”)
if err != nil {
return nil, err
}
return &Embedder{
model: model,
tokenizer: tokenizer,
}, nil
}

func (e *Embedder) Encode(texts []string) ([][]float32, error) {
var inputIDs [][]int
for _, text := range texts {
ids := e.tokenizer.Encode(text)
inputIDs = append(inputIDs, ids)
}
embeddings, err := e.model.EncodeBatch(inputIDs)
if err != nil {
return nil, err
}
for i := range embeddings {
normalize(embeddings[i])
}
return embeddings, nil
}
“`

向量化的时候需要注意模型的输入长度限制。bge-base-zh-v1.5的最大输入是512个token，如果文本超过这个长度，需要先截断。不过我们的chunk Size设置成256个token左右，应该不会超出限制。

### 第四步：向量存储和检索

把向量存储到Chroma数据库：

“`go
type VectorStore struct {
client *chroma.Client
collection *chroma.Collection
}

func NewVectorStore() (*VectorStore, error) {
client, err := chroma.NewClient(“http://localhost:8000”)
if err != nil {
return nil, err
}
collection, err := client.GetOrCreateCollection(“knowledge_base”, nil)
if err != nil {
return nil, err
}
return &VectorStore{
client: client,
collection: collection,
}, nil
}

func (vs *VectorStore) Add(documents []Document, embeddings [][]float32) error {
var metadatas []map[string]interface{}
var documents_ []string
var ids []string

for i, doc := range documents {
ids = append(ids, doc.ID)
documents_ = append(documents_, doc.Content)
metadatas = append(metadatas, doc.Meta)
}
err := vs.collection.Add(ids, embeddings, metadatas, documents_)
if err != nil {
return fmt.Errorf(“添加向量失败: %w”, err)
}
return nil
}

func (vs *VectorStore) Search(queryEmbedding []float32, topK int) ([]Document, error) {
results, err := vs.collection.Query(
[][]float32{queryEmbedding},
topK,
nil,
nil,
)
if err != nil {
return nil, fmt.Errorf(“检索失败: %w”, err)
}
var documents []Document
for i, id := range results.IDS[0] {
documents = append(documents, Document{
ID: id,
Content: results.Documents[0][i],
})
}
return documents, nil
}
“`

Chroma支持多种距离度量方式，默认为L2距离。对于语义搜索，余弦相似度通常效果更好。可以在创建collection时指定。

### 第五步：构建完整流程

把各个部分组合起来，形成完整的RAG流程：

“`go
type RAGSystem struct {
embedder *Embedder
vectorStore *VectorStore
}

func NewRAGSystem() (*RAGSystem, error) {
embedder, err := NewEmbedder()
if err != nil {
return nil, err
}
vectorStore, err := NewVectorStore()
if err != nil {
return nil, err
}
return &RAGSystem{
embedder: embedder,
vectorStore: vectorStore,
}, nil
}

func (r *RAGSystem) IndexDirectory(dirPath string) error {
files, err := os.ReadDir(dirPath)
if err != nil {
return err
}
chunker := Chunker{
chunkSize: 256,
overlap: 50,
}
var allDocuments []Document

for _, file := range files {
if !strings.HasSuffix(file.Name(), “.md”) {
continue
}
content, err := os.ReadFile(filepath.Join(dirPath, file.Name()))
if err != nil {
continue
}
chunks := chunker.Chunk(string(content))
for i := range chunks {
chunks[i].Meta = map[string]string{
“source”: file.Name(),
}
}
allDocuments = append(allDocuments, chunks…)
}

contents := make([]string, len(allDocuments))
for i, doc := range allDocuments {
contents[i] = doc.Content
}
embeddings, err := r.embedder.Encode(contents)
if err != nil {
return err
}
return r.vectorStore.Add(allDocuments, embeddings)
}

func (r *RAGSystem) Query(question string, topK int) ([]string, error) {
embeddings, err := r.embedder.Encode([]string{question})
if err != nil {
return nil, err
}
documents, err := r.vectorStore.Search(embeddings[0], topK)
if err != nil {
return nil, err
}
var results []string
for _, doc := range documents {
results = append(results, doc.Content)
}
return results, nil
}
“`

这就是一个完整的RAG系统了。虽然没有连接大模型，但可以用这个系统来验证检索效果。真正的生产环境中，会把检索到的内容发送给GPT或者本地部署的LLM。

—

## 运行结果

运行上面的代码，首先索引文档目录：

“`go
func main() {
rag, err := NewRAGSystem()
if err != nil {
log.Fatal(err)
}
err = rag.IndexDirectory(“./docs”)
if err != nil {
log.Fatal(err)
}
results, err := rag.Query(“什么是RAG系统？”, 3)
if err != nil {
log.Fatal(err)
}
for i, result := range results {
fmt.Printf(“结果 %d: %s\n”, i+1, result)
}
}
“`

假设docs目录下有一个名为rag-intro.md的文件，内容介绍RAG的概念。查询”什么是RAG系统？”后，系统会找到相关的段落，返回类似这样的结果：

“`
结果 1: RAG（检索增强生成）是一种结合了检索系统和生成模型的AI架构…
结果 2: 传统的生成模型需要大量的训练数据来获��知��…
结果 3: 使用RAG可以让模型访问最新的信息，而不需要重新训练…
“`

可以看到，系统确实找到了关于RAG定义和优点的内容。虽然这只是检索部分，但已经证明了整个流程的正确性。

—

## 总结

这篇文章我们用Go语言实现了一个完整的RAG系统，包括文档加载、文本分块、向量化、向量存储和检索等核心环节。重点解决了几个实际问题：中文文本的处理、chunk大小的选择、向量数据库的使用。

这个实现还有很多可以优化的地方。比如可以使用更强大的向量化模型（如bge-large），或者添加rerank模块对检索结果重新排序，还可以实现混合搜索（结合关键词和语义）。chunk策略也可以更精细，比如按段落、标题层次来分割。如果有条件，可以部署一个本地的大模型（如llama.cpp），加上上面的检索部分，就构成了完整的本地知识库问答系统。

对于想要构建自己知识库的朋友，这个demo是一个不错的起点。核心逻辑其实不复杂，关键是理解每个环节的作用，然后根据具体需求做调整。

发送评论 编辑评论

推荐文章

发送评论编辑评论