chromem-go：Go 语言 RAG 应用的高效轻量级向量数据库

前言

在开发 RAG（Retrieval-Augmented Generation）应用时，起初你可能更倾向于选择一款轻量级的向量数据库，而非复杂的大型数据库。例如，在关系型数据库的选择上，许多人会更愿意使用 SQLite 而不是 PostgreSQL 或 MySQL，以减少额外的配置和维护成本。

在 Go 语言中，chromem-go 提供了一种简洁高效的解决方案------它是一款可嵌入到 Go 程序中的轻量级向量数据库，无需单独部署数据库服务，极大地降低了使用门槛。

本文将深入探讨 chromem-go 的使用方法，帮助你快速上手并应用到实际项目中。

准备好了吗？准备一杯你最喜欢的咖啡或茶，随着本文一探究竟吧。

chromem-go

介绍

chromem-go 是一个 Go 语言中可嵌入的向量数据库，提供类似 Chroma 的接口且无第三方依赖。支持内存存储，并可选持久化功能。

由于 chromem-go 是可嵌入的，因此你能够在 Go 应用中添加检索增强生成（Retrieval Augmented Generation, RAG）以及类似的基于嵌入（embeddings）的功能，而无需运行一个单独的数据库。这就像使用 SQLite 而不是 PostgreSQL/MySQL 等数据库一样。

这是一款独立的向量数据库，专为 Go 语言设计，并非用于连接 Chroma 的库，也不是 Chroma 的 Go 语言重新实现。

安装

在 Go 项目根目录下，执行以下命令：

bash 复制代码

go get github.com/philippgille/chromem-go

快速开始

go 复制代码

package main

import (
	"context"
	"fmt"
	"runtime"

	"github.com/philippgille/chromem-go"
)

func main() {
	ctx := context.Background()

	db := chromem.NewDB()

	c, err := db.CreateCollection("knowledge-base", nil, nil)
	if err != nil {
		panic(err)
	}

	err = c.AddDocuments(ctx, []chromem.Document{
		{
			ID:      "1",
			Content: "The sky is blue because of Rayleigh scattering.",
		},
		{
			ID:      "2",
			Content: "Leaves are green because chlorophyll absorbs red and blue light.",
		},
	}, runtime.NumCPU())
	if err != nil {
		panic(err)
	}

	res, err := c.Query(ctx, "Why is the sky blue?", 1, nil, nil)
	if err != nil {
		panic(err)
	}

	fmt.Printf("ID: %v\nSimilarity: %v\nContent: %v\n", res[0].ID, res[0].Similarity, res[0].Content)
}

输出：

bash 复制代码

ID: 1
Similarity: 0.6833369
Content: The sky is blue because of Rayleigh scattering.

注意事项：

在使用 db.CreateCollection 创建集合时，如果未显式指定 EmbeddingFunc 参数，chromem-go 将默认使用 OpenAI 的嵌入实现。这意味着，你需要在环境变量中设置 OPENAI_API_KEY，否则数据插入操作将会失败。

数据库的创建

chromem-go 支持两种数据库的创建方式：内存数据库 和 持久化的数据库。

内存数据库

go 复制代码

db := chromem.NewDB()

持久化的数据库

go 复制代码

db, err := chromem.NewPersistentDB("./chromem-go", true)

第一个参数为保存数据库文件的路径参数，如果传递 ""，那么默认值为 "./chromem-go"。

第二个参数表示是否使用 gzip 格式压缩数据。

集合操作

创建集合

创建新集合

go 复制代码

c, err := db.CreateCollection("knowledge-base", nil, nil)

"knowledge-base"：集合的名称。

其余的参数分别表示集合的 metadata 参数以及 Embedding 构建器的函数实现参数。

获取集合

获取单个集合

go 复制代码

c := db.GetCollection("knowledge-base", nil)

第二个参数为 Embedding 构建器的函数实现。

获取或创建新集合

go 复制代码

c, err := db.GetOrCreateCollection("knowledge-base", nil, nil)

获取指定的集合，如果集合不存在则会创建新集合。该方法的参数和 db.CreateCollection 一样。

获取集合列表

go 复制代码

collections := db.ListCollections()

该方法的返回值类型为 map[string]*Collection。

删除集合

go 复制代码

err := db.DeleteCollection("knowledge-base")

集合数据操作

添加集合数据

go 复制代码

err := c.AddDocument(ctx, chromem.Document{})
err = c.AddDocuments(ctx, []chromem.Document{
	{
		ID:      "1",
		Content: "The sky is blue because of Rayleigh scattering.",
	},
	{
		ID:      "2",
		Content: "Leaves are green because chlorophyll absorbs red and blue light.",
	},
}, runtime.NumCPU())

AddDocument 和 AddDocuments 的作用分别为添加单条和多条集合数据。chromem.Document 对象除了 Content 参数以外，还有 Embedding 参数，如果指定了 Embedding 参数，则不会调用 embedding function 对 Content 参数进行转换。

查询集合数据

go 复制代码

c.GetByID(ctx, "1")
c.Query(ctx, "Why is the sky blue?", 1, nil, nil)
c.QueryEmbedding(ctx, []float32{0.1, 0.2, 0.3}, 1, nil, nil)

chromem-go 提供了三种查询集合数据的方法，我们可以在不同的场景下选择合适的方法进行使用。

删除集合数据

go 复制代码

c.Delete(ctx, metadataFilterMap, documentFilterMap, ids)

metadataFilterMap：基于 metadata 进行筛选的查询条件。
documentFilterMap：基于 document 内容进行筛选的查询条件。
ids：指定要删除的文档 ID 集合。

小结

本文详细介绍了 chromem-go 的使用方法，包括但不仅限于 chromem-go 的介绍 、安装方式 以及 数据库的创建 等内容。如果你正尝试开发一个 RAG 应用练练手时，相信 chromem-go 将是你向量数据库的最佳选择。

你好，我是陈明勇，一名热爱技术、乐于分享的开发者，同时也是开源爱好者。

我专注于分享 Go 语言相关的技术知识，同时也会深入探讨 AI 领域的前沿技术。

成功的路上并不拥挤，有没有兴趣结个伴？