使用Ent、Atlas和pgvector构建Go语言RAG系统-优快云博客

使用Ent、Atlas和pgvector构建Go语言RAG系统

ent ent: 是一个基于 Go 语言的轻量级 ORM 库，用于处理关系型数据库。适合开发者使用 ent 进行数据库访问和操作。项目地址: https://gitcode.com/gh_mirrors/en/ent

什么是RAG系统

RAG（Retrieval-Augmented Generation，检索增强生成）是一种结合信息检索和文本生成的技术。它通过以下步骤工作：

从知识库中检索与问题相关的文档片段
将这些片段与问题一起输入生成模型
生成基于检索内容的准确回答

这种技术特别适用于需要结合特定领域知识或最新信息的问答系统。

项目架构

我们将使用以下技术栈构建RAG系统：

Ent：Go语言的实体框架，用于数据建模和数据库操作
Atlas：数据库schema管理工具，与Ent无缝集成
pgvector：PostgreSQL扩展，支持向量存储和相似性搜索

环境准备

数据库准备

首先启动一个包含pgvector扩展的PostgreSQL容器：

docker run --rm --name postgres -e POSTGRES_PASSWORD=pass -p 5432:5432 -d pgvector/pgvector:pg17

项目初始化

创建Go模块并初始化Ent项目：

go mod init myragproject
go run -mod=mod entgo.io/ent/cmd/ent new Embedding Chunk

数据模型设计

Chunk模型

Chunk模型表示文档分块：

type Chunk struct {
    ent.Schema
}

func (Chunk) Fields() []ent.Field {
    return []ent.Field{
        field.String("path"),      // 文档路径
        field.Int("nchunk"),       // 分块序号
        field.Text("data"),        // 分块内容
    }
}

func (Chunk) Edges() []ent.Edge {
    return []ent.Edge{
        edge.To("embedding", Embedding.Type).Unique(),
    }
}

Embedding模型

Embedding模型存储文本块的向量表示：

type Embedding struct {
    ent.Schema
}

func (Embedding) Fields() []ent.Field {
    return []ent.Field{
        field.Other("embedding", pgvector.Vector{}).
            SchemaType(map[string]string{
                dialect.Postgres: "vector(1536)",  // OpenAI嵌入维度
            }),
    }
}

func (Embedding) Indexes() []ent.Index {
    return []ent.Index{
        index.Fields("embedding").
            Annotations(
                entsql.IndexType("hnsw"),       // 高效近似最近邻搜索
                entsql.OpClass("vector_l2_ops"), // L2距离度量
            ),
    }
}

数据库迁移

使用Atlas管理数据库schema迁移：

安装Atlas
创建atlas.hcl配置文件
应用迁移：

atlas schema apply --env local

核心功能实现

文档加载

实现文档分块和存储：

func (cmd *LoadCmd) Run(ctx *CLI) error {
    return filepath.WalkDir(cmd.Path, func(path string, d fs.DirEntry, err error) error {
        if isMarkdown(path) {
            chunks := breakToChunks(path)
            for i, chunk := range chunks {
                client.Chunk.Create().
                    SetData(chunk).
                    SetPath(path).
                    SetNchunk(i).
                    SaveX(context.Background())
            }
        }
        return nil
    })
}

嵌入生成

使用OpenAI API生成文本嵌入：

func getEmbedding(data string) []float32 {
    client := openai.NewClient(os.Getenv("OPENAI_KEY"))
    resp, err := client.CreateEmbeddings(context.Background(), 
        openai.EmbeddingRequest{
            Input: []string{data},
            Model: openai.AdaEmbeddingV2,
        })
    if err != nil {
        log.Fatal(err)
    }
    return resp.Data[0].Embedding
}

问答功能

实现基于向量相似度的检索：

func (cmd *AskCmd) Run(ctx *CLI) error {
    embedding := getEmbedding(cmd.Text)
    client, _ := ctx.entClient()
    
    // 查找最相似的3个文本块
    chunks := client.Embedding.Query().
        Order(func(s *sql.Selector) {
            s.OrderBy(fmt.Sprintf("embedding <-> '%s'", 
                pgvector.NewVector(embedding).String()))
        }).
        Limit(3).
        QueryChunk().
        AllX(context.Background())
    
    // 构建提示词并调用生成模型
    prompt := buildPrompt(cmd.Text, chunks)
    answer := generateAnswer(prompt)
    
    fmt.Println(answer)
    return nil
}