chromem-go

适用于 Go 的可嵌入向量数据库，具有类 Chroma 接口，零第三方依赖。内存存储，可选持久化。

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

它不是连接 Chroma 的库，也不是用 Go 重新实现 Chroma。它是一个独立的数据库。

重点不是规模（数百万文档）或功能数量，而是最常见用例的简单性和性能。在 2020 年中端英特尔笔记本 CPU 上，您可以在 0.3 毫秒内查询 1,000 个文档，在 40 毫秒内查询 100,000 个文档，内存分配很少且很小。详情请参见基准测试。

⚠️ 该项目处于测试阶段，正在大规模构建中，在 v1.0.0 发布之前的版本中可能会引入重大更改。所有更改都记录在 CHANGELOG 中。

使用场景

使用向量数据库，您可以做各种事情：

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

让我们详细看看 RAG 使用场景：

RAG

大型语言模型（LLM）的知识是有限的 - 即使是那些拥有 300 亿、700 亿甚至更多参数的模型。它们不知道训练结束后发生的事情，不知道未经训练的数据（如您公司的内网、Jira / 错误跟踪器、维基或其他类型的知识库），即使是它们确实知道的数据，它们通常也无法准确重现，反而会开始产生"幻觉"。

微调 LLM 可以稍微有所帮助，但它更多是为了改善 LLM 对特定主题的推理，或重现文本或代码的风格。微调并不能将知识一对一地添加到模型中。细节会丢失或混淆。而且知识截止（关于微调后发生的任何事情）的问题也没有解决。

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

您将想要 LLM 了解的相关文档存储在数据库中。
数据库将嵌入向量与文档一起存储，您可以提供这些嵌入向量，也可以由特定的"嵌入模型"（如 OpenAI 的 text-embedding-3-small）创建。
- chromem-go 可以为您完成这项工作，并支持多个嵌入提供商和模型。
之后，当您想与 LLM 对话时，首先将问题发送到向量数据库以找到相似/相关的内容。这称为"最近邻搜索"。
在向 LLM 提问时，您将这些内容与您的问题一起提供。
LLM 在回答时可以考虑这些最新的精确内容。

查看示例代码以了解其运行情况！

接口

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

<details><summary>Chroma 核心接口</summary>

import chromadb
# 在内存中设置 Chroma，便于原型设计。可以轻松添加持久性！
client = chromadb.Client()

# 创建集合。还可使用 get_collection、get_or_create_collection、delete_collection！
collection = client.create_collection("all-my-documents")
# 向集合添加文档。也可以更新和删除。基于行的 API 即将推出！
collection.add(
    documents=["这是文档1", "这是文档2"], # 我们自动处理分词、嵌入和索引。你也可以跳过这些步骤，添加自己的嵌入
    metadatas=[{"来源": "notion"}, {"来源": "google-docs"}], # 可以基于这些进行过滤！
    ids=["doc1", "doc2"], # 每个文档唯一
)

# 查询/搜索2个最相似的结果。你也可以通过 id 使用 .get 方法获取
results = collection.query(
    query_texts=["这是一个查询文档"],
    n_results=2,
    # where={"元数据字段": "等于这个值"}, # 可选过滤器
    # where_document={"$contains":"搜索字符串"}  # 可选过滤器
)

</details>

我们的 Go 库提供了相同的接口：

<details><summary>chromem-go 等效代码</summary>

package main

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

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

    // 创建集合。也可以使用 GetCollection、GetOrCreateCollection、DeleteCollection！
    collection, _ := db.CreateCollection("我的所有文档", nil, nil)

    // 向集合添加文档。未来将添加更新和删除功能。
    // 可以使用 AddConcurrently() 进行多线程处理！
    // 这里我们展示了类似 Chroma 的方法，但也提供了更符合 Go 习惯的方法！
    _ = collection.Add(ctx,
        []string{"doc1", "doc2"}, // 每个文档的唯一 ID
        nil, // 我们自动处理嵌入。你也可以跳过这步，添加自己的嵌入。
        []map[string]string{{"来源": "notion"}, {"来源": "google-docs"}}, // 可以基于这些进行过滤！
        []string{"这是文档1", "这是文档2"},
    )

    // 查询/搜索2个最相似的结果。未来将添加通过 ID 获取的功能。
    results, _ := collection.Query(ctx,
        "这是一个查询文档",
        2,
        map[string]string{"元数据字段": "等于这个值"}, // 可选过滤器
        map[string]string{"$contains": "搜索字符串"},         // 可选过滤器
    )
}

</details>

最初，chromem-go 只有四个核心方法，但随着时间推移我们添加了更多功能。不过我们有意不覆盖 Chroma 100% 的 API 表面。我们提供了一些更符合 Go 习惯的替代方法。

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

功能

路线图

性能：
- 在支持的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中），请参见示例代码。

快速入门

以下内容摘自"最小"示例：

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: "天空呈蓝色是因为瑞利散射。",
    },
    {
      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-03-17进行的基准测试，配置如下：

计算机：Framework Laptop 13（第一代，2021）
CPU：第11代英特尔酷睿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: 11th Gen 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