chromem-go

适用于 Go 语言的可嵌入向量数据库，提供与 Chroma 类似的接口，并且零第三方依赖。支持内存存储及可选持久化。

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

它并不是用于连接到 Chroma 的库，也不是在 Go 中对其的重新实现。它本身就是一个独立的数据库。

该项目的重点并非大规模（如数百万文档）或功能数量，而是针对最常见的用例提供简单性和高性能。在一台中端 2020 年款 Intel 笔记本电脑上，您可以以 0.3 毫秒查询 1,000 条文档，以 40 毫秒查询 100,000 条文档，且内存分配极少、体积很小。详情请参阅 基准测试。

⚠️ 该项目目前处于 beta 阶段，开发仍在进行中，在 v1.0.0 之前发布的版本中可能会引入破坏性更改。所有更改均记录在 CHANGELOG 中。

目录

1. 用例
2. 接口
3. 特性 + 路线图
4. 安装
5. 使用
6. 基准测试
7. 开发
8. 动机
9. 相关项目

用例

借助向量数据库，您可以实现多种功能：

• 检索增强生成（RAG）、问答系统（Q&A）
• 文本和代码搜索
• 推荐系统
• 分类
• 聚类

让我们更详细地探讨 RAG 的用例：

RAG

大型语言模型（LLM）的知识——即使是那些拥有 300 亿、700 亿参数甚至更多参数的模型——也是有限的。它们不了解训练结束之后发生的事情，也不了解未在其训练数据中出现的内容（例如贵公司的内网、Jira/缺陷跟踪器、维基或其他类型的知识库）。即便它们确实知道某些内容，也往往无法精确复现，反而会产生“幻觉”。

对 LLM 进行微调可以在一定程度上有所帮助，但它主要是为了提升模型在特定主题上的推理能力，或者模仿某种文本或代码的风格。微调并不能将知识以“一对一”的方式直接注入模型中。细节往往会丢失或混淆，而且也无法解决知识截止问题（即无法处理微调之后发生的事情）。

因此，向量数据库可以作为 LLM 的最新、精确知识来源：

1. 您将希望 LLM 了解的相关文档存储在数据库中。
2. 数据库会同时存储文档的嵌入表示，这些嵌入可以由您提供，也可以由特定的“嵌入模型”（如 OpenAI 的 text-embedding-3-small）自动生成。
   • chromem-go 可以为您完成这一过程，并开箱即用支持多种嵌入提供商和模型。
3. 当您需要与 LLM 对话时，首先将问题发送到向量数据库，以查找“相似”或“相关”的内容。这被称为“最近邻搜索”。
4. 在向 LLM 提出问题时，您将这些相关内容与问题一同提供。
5. LLM 在回答时可以将这些最新、精确的内容纳入考虑。

请查看 示例代码，了解其实际应用！

接口

我们的最初灵感来源于 Chroma 的接口，其核心 API 如下（摘自他们的 README)：

import chromadb
# 设置内存中的 Chroma，便于快速原型设计。也可轻松添加持久化！
client = chromadb.Client()

# 创建集合。还可以获取集合、创建或获取集合、删除集合！
collection = client.create_collection("all-my-documents")

# 向集合中添加文档。未来还将支持更新和删除操作。
# 支持按行操作的 API 即将推出！
collection.add(
    documents=["这是文档1", "这是文档2"], # 我们会自动处理分词、嵌入和索引。您也可以跳过这些步骤，自行提供嵌入。
    metadatas=[{"source": "notion"}, {"source": "google-docs"}], # 可以根据这些元数据进行过滤！
    ids=["doc1", "doc2"], # 每个文档都有唯一 ID
)

# 查询/搜索最相似的 2 个结果。也可以按 ID 获取。
results = collection.query(
    query_texts=["这是一份查询文档"],
    n_results=2,
    # where={"metadata_field": "is_equal_to_this"}, # 可选过滤条件
    # where_document={"$contains":"search_string"}  # 可选过滤条件
)

我们的 Go 库暴露了相同的接口：

package main

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

func main() {
    // 设置内存中的 chromem-go，便于快速原型设计。也可轻松添加持久化！
    // 我们称之为 DB 而不是 client，因为这里没有客户端-服务器分离。DB 是嵌入式的。
    db := chromem.NewDB()

    // 创建集合。还可以获取集合、创建或获取集合、删除集合！
    collection, _ := db.CreateCollection("all-my-documents", nil, nil)

    // 向集合中添加文档。未来将支持更新和删除操作。
    // 可以使用 AddConcurrently() 方法进行多线程操作！
    // 这里展示的是类似 Chroma 的方法，但也有更符合 Go 语言习惯的方法！
    _ = collection.Add(ctx,
        []string{"doc1", "doc2"}, // 每个文档都有唯一 ID
        nil, // 我们会自动处理嵌入。您也可以跳过这一步，自行提供嵌入。
        []map[string]string{{"source": "notion"}, {"source": "google-docs"}}, // 可以根据这些元数据进行过滤！
        []string{"这是文档1", "这是文档2"},
    )

    // 查询/搜索最相似的 2 个结果。也可以按 ID 获取。
    results, _ := collection.Query(ctx,
        "这是一份查询文档",
        2,
        map[string]string{"metadata_field": "is_equal_to_this"}, // 可选过滤条件
        map[string]string{"$contains": "search_string"},         // 可选过滤条件
    )
}

最初，chromem-go 只实现了四个核心方法，但我们随着时间的推移又增加了更多功能。不过，我们并没有打算完全覆盖 Chroma 的全部 API 表面。相反，我们提供了更具 Go 语言风格的替代方法。

完整的接口请参阅 Godoc：https://pkg.go.dev/github.com/philippgille/chromem-go

特性

• 无第三方库依赖
• 可嵌入式（类似 SQLite，即无客户端-服务器模式，无需维护独立的数据库）
• 多线程处理（在添加和查询文档时），充分利用 Go 的原生并发特性
• 实验性 WebAssembly 绑定
• 嵌入模型提供者：
  • 托管服务：
    • OpenAI（默认）
    • Azure OpenAI
    • GCP Vertex AI
    • Cohere
    • Mistral
    • Jina
    • mixedbread.ai
  • 本地运行：
    • Ollama
    • LocalAI
  • 自定义实现（实现 chromem.EmbeddingFunc）
  • 也可以在向集合中添加文档时直接传入已有的嵌入向量，而不是让 chromem-go 自动生成
• 相似度搜索：
  • 使用余弦相似度进行穷举最近邻搜索（有时也称为精确搜索、暴力搜索或 FLAT 索引）
• 过滤器：
  • 文档过滤：$contains、$not_contains
  • 元数据过滤：精确匹配
• 存储：
  • 内存中
  • 可选即时持久化（为每个添加的集合和文档写入一个文件，使用 gob 编码，可选 gzip 压缩）
  • 备份：将整个数据库导出和导入到单个文件中（使用 gob 编码，可选 gzip 压缩和 AES-GCM 加密）
    • 包括针对通用 io.Writer/io.Reader 的方法，因此可以对接 S3 存储桶和其他对象存储，示例代码参见 examples/s3-export-import
• 数据类型：
  • 文档（文本）

路线图

• 性能：
  • 在支持的 CPU 上使用 SIMD 进行点积计算（PR 草案：#48）
  • 添加 roaring bitmaps 以加速全文过滤
• 嵌入模型提供者：
  • 添加一个 EmbeddingFunc，用于下载并调用 llamafile
• 相似度搜索：
  • 使用索引的近似最近邻搜索 (ANN)
    • 分层可导航小世界 (HNSW)
    • 倒排文件平铺 (IVFFlat)
• 过滤器：
  • 操作符（$and、$or 等）
• 存储：
  • 将 JSON 作为第二种编码格式
  • 将预写日志 (WAL) 作为第二种文件格式
  • 可选远程存储（S3、PostgreSQL 等）
• 数据类型：
  • 图像
  • 视频

安装

go get github.com/philippgille/chromem-go@latest

使用

参考 Godoc：https://pkg.go.dev/github.com/philippgille/chromem-go

有关完整可用的示例，包括如何使用向量数据库进行检索增强生成 (RAG) 和语义搜索，以及如何使用 OpenAI 或本地运行的嵌入模型和 LLM（通过 Ollama），请参阅 示例代码。

快速入门

以下内容摘自 "minimal" 示例：

package main

import (
 "context"
 "fmt"
 "runtime"

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

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

  db := chromem.NewDB()

  // 传入 nil 作为嵌入函数会导致使用 OpenAI，并且需要设置 "OPENAI_API_KEY" 环境变量。其他提供商也同样支持。
  // 例如，传入 `chromem.NewEmbeddingFuncOllama(...)` 可以使用 Ollama。
  c, err := db.CreateCollection("knowledge-base", nil, nil)
  if err != nil {
    panic(err)
  }

  err = c.AddDocuments(ctx, []chromem.Document{
    {
      ID:      "1",
      Content: "天空之所以是蓝色，是因为瑞利散射。",
    },
    {
      ID:      "2",
      Content: "叶子之所以是绿色，是因为叶绿素吸收了红光和蓝光。",
    },
  }, runtime.NumCPU())
  if err != nil {
    panic(err)
  }

  res, err := c.Query(ctx, "为什么天空是蓝色的？", 1, nil, nil)
  if err != nil {
    panic(err)
  }

  fmt.Printf("ID: %v\n相似度: %v\n内容: %v\n", res[0].ID, res[0].Similarity, res[0].Content)
}

输出：

ID: 1
相似度: 0.6833369
内容: 天空之所以是蓝色，是因为瑞利散射。

基准测试

基准测试于 2024 年 3 月 17 日进行，配置如下：

• 计算机：Framework Laptop 13（第一代，2021 年）
• CPU：第 11 代 Intel Core i5-1135G7（2020 年）
• 内存：32 GB
• 操作系统：Fedora Linux 39
  • 内核：6.7

$ go test -benchmem -run=^$ -bench .
goos: linux
goarch: amd64
pkg: github.com/philippgille/chromem-go
cpu: 第 11 代 Intel(R) Core(TM) i5-1135G7 @ 2.40GHz
BenchmarkCollection_Query_NoContent_100-8          13164      90276 ns/op     5176 B/op       95 allocs/op
BenchmarkCollection_Query_NoContent_1000-8          2142     520261 ns/op    13558 B/op      141 allocs/op
BenchmarkCollection_Query_NoContent_5000-8           561    2150354 ns/op    47096 B/op      173 allocs/op
BenchmarkCollection_Query_NoContent_25000-8          120    9890177 ns/op   211783 B/op      208 allocs/op
BenchmarkCollection_Query_NoContent_100000-8          30   39574238 ns/op   810370 B/op      232 allocs/op
BenchmarkCollection_Query_100-8                    13225      91058 ns/op     5177 B/op       95 allocs/op
BenchmarkCollection_Query_1000-8                    2226     519693 ns/op    13552 B/op      140 allocs/op
BenchmarkCollection_Query_5000-8                     550    2128121 ns/op    47108 B/op      173 allocs/op
BenchmarkCollection_Query_25000-8                    100   10063260 ns/op   211705 B/op      205 allocs/op
BenchmarkCollection_Query_100000-8                    30   39404005 ns/op   810295 B/op      229 allocs/op
PASS
ok   github.com/philippgille/chromem-go 28.402s

开发

• 构建：go build ./...
• 测试：go test -v -race -count 1 ./...
• 基准测试：
  • go test -benchmem -run=^$ -bench .（可添加 > bench.out 等重定向到文件）
  • 含性能分析：go test -benchmem -run ^$ -cpuprofile cpu.out -bench .
    • （分析选项：-cpuprofile、-memprofile、-blockprofile、-mutexprofile）
• 对比基准测试结果：
  1. 安装 benchstat：go install golang.org/x/perf/cmd/benchstat@latest
  2. 对比两个基准测试结果：benchstat before.out after.out

动机

2023年12月，当我想要在Go程序中尝试检索增强生成（RAG）时，我寻找一种可以嵌入到Go程序中的向量数据库，就像你可以嵌入SQLite一样，这样就不需要单独的数据库设置和维护了。然而，令我惊讶的是，在Go生态系统中虽然有许多嵌入式键值存储，却几乎找不到这样的向量数据库。

当时，许多流行的向量数据库，如Pinecone、Qdrant、Milvus、Chroma、Weaviate等，要么完全无法嵌入，要么仅支持Python或JavaScript/TypeScript语言。

后来，我发现了@eliben的博客文章和示例代码，它们展示了只需少量Go代码就能实现一个非常基础的向量数据库概念验证。

正是在那时，我决定基于ChromaDB的接口，构建一个可嵌入Go的向量数据库。ChromaDB之所以脱颖而出，是因为它支持嵌入式使用（Python），并且在其README文件和官网首页上仅用4个命令就展示了其核心API。

相关项目

• 特别感谢@eliben，他的博客文章和示例代码，启发了我启动这个项目！
• Chroma：对比Pinecone、Qdrant、Milvus、Weaviate等，Chroma以其在README和官网首页上仅用4个命令展示核心API而独树一帜。此外，Chroma还特别强调其嵌入式能力（Python）。
• 一些大型、功能完善的客户端-服务器架构向量数据库，旨在实现最大规模和最佳性能：
  • Pinecone：闭源
  • Qdrant：使用Rust编写，无法嵌入Go
  • Milvus：使用Go和C++编写，但截至2023年12月仍不支持嵌入式使用
  • Weaviate：使用Go编写，但截至2024年3月仍不支持Go嵌入（仅支持Python和JavaScript/TypeScript，且处于实验阶段）
• 一些非专门的SQL、NoSQL和键值存储数据库开始支持向量存储以及基于相似度的查询：
  • pgvector扩展，适用于PostgreSQL：客户端-服务器模式
  • Redis（1，2）：客户端-服务器模式
  • sqlite-vss扩展，适用于SQLite：嵌入式，但其Go绑定依赖CGO。虽然存在一个无CGO的SQLite Go库1，但该库并不包含向量搜索扩展。
  • DuckDB提供计算余弦相似度的函数（1）：嵌入式，但其Go绑定使用CGO
  • MongoDB的云平台提供向量搜索服务（1）：客户端-服务器模式
• 一些用于向量相似度搜索的库：
  • Faiss：使用C++编写；第三方Go绑定依赖CGO
  • Annoy：使用C++编写；Go绑定依赖CGO（1）
  • USearch：使用C++编写；Go绑定依赖CGO
• 一些编排库，受Python库LangChain启发，但内置的向量数据库功能要么缺失，要么非常简陋：
  • LangChain Go
  • LinGoose
  • GoLC

chromem-go 快速上手指南

环境准备

• 需要安装 Go 1.18+ 环境
• 无第三方依赖（零依赖）
• 使用 OpenAI 时需设置 OPENAI_API_KEY 环境变量

安装步骤

export GOPROXY=https://goproxy.cn,direct
go get -u github.com/philippgille/chromem-go@latest

基本使用

以下是最小化示例（使用 OpenAI 生成嵌入）：

package main

import (
	"context"
	"fmt"
	"runtime"

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

func main() {
	ctx := context.Background()
	db := chromem.NewDB()

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

	_ = 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())

	res, _ := c.Query(ctx, "Why is the sky blue?", 1, nil, nil)
	fmt.Printf("ID: %v\nSimilarity: %v\nContent: %v\n", res[0].ID, res[0].Similarity, res[0].Content)
}

运行后输出：

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