golang-cookbook

Appearance

在 Go 中使用 Chroma 向量数据库的基本过程通常包括连接数据库、存储向量数据、执行向量检索以及查询相似的文档。Chroma 是一个开源的向量数据库，专门用于高效地存储和检索向量，特别适合用于机器学习和自然语言处理中的相似度搜索。

1. 安装 Chroma 和 Go 客户端 ​

要在 Go 中使用 Chroma，首先需要确保你的 Chroma 服务已经在运行，通常它是通过 HTTP API 提供服务的。你还需要安装 Go 客户端库。

安装 Chroma Go 客户端 ​

目前，Chroma 并没有直接提供一个 Go 客户端库，而是可以通过 HTTP API 与 Chroma 服务进行交互。所以，你可以通过 Go 的 HTTP 客户端来访问 Chroma API，或者如果有相关的第三方库，也可以直接使用。

以下是通过 Go 的 net/http 库与 Chroma 交互的一个基本示例。

bash
go get github.com/chroma-core/chroma-go

2. 与 Chroma 进行交互 ​

启动 Chroma 服务 ​

首先，你需要在本地或服务器上启动 Chroma 服务。假设你已经安装并配置了 Chroma，可以通过 Docker 启动它：

bash
docker run -d -p 8000:8000 chromadb/chroma:latest

这样，Chroma 服务就会在 http://localhost:8000 上运行，提供 HTTP API 进行向量存储和检索。

使用 Go 客户端与 Chroma 交互 ​

以下是如何通过 Go 客户端与 Chroma 进行交互的示例：

go
package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
)

type ChromaDocument struct {
	ID        string    `json:"id"`
	Text      string    `json:"text"`
	Embedding []float32 `json:"embedding"`
}

type ChromaQuery struct {
	Embedding []float32 `json:"embedding"`
	Limit     int       `json:"limit"`
}

func main() {
	// Chroma API 服务地址
	chromaURL := "http://localhost:8000"

	// 插入文档到 Chroma
	document := ChromaDocument{
		ID:        "1",
		Text:      "This is an example document.",
		Embedding: []float32{0.1, 0.2, 0.3, 0.4, 0.5}, // 示例向量
	}

	// 向 Chroma 插入数据
	insertDocument(chromaURL, document)

	// 执行查询
	queryVector := []float32{0.1, 0.2, 0.3, 0.4, 0.5} // 查询向量
	query := ChromaQuery{
		Embedding: queryVector,
		Limit:     5, // 获取最相似的 5 个文档
	}
	queryResults := queryChroma(chromaURL, query)

	// 输出查询结果
	for _, result := range queryResults {
		fmt.Printf("Document ID: %s, Text: %s\n", result.ID, result.Text)
	}
}

// 插入文档到 Chroma
func insertDocument(chromaURL string, document ChromaDocument) {
	// 构造插入请求
	url := fmt.Sprintf("%s/collections/default/documents", chromaURL)
	body, err := json.Marshal(document)
	if err != nil {
		log.Fatal("Failed to marshal document:", err)
	}

	resp, err := http.Post(url, "application/json", bytes.NewBuffer(body))
	if err != nil {
		log.Fatal("Failed to insert document:", err)
	}
	defer resp.Body.Close()

	// 读取响应
	respBody, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("Failed to read response:", err)
	}

	fmt.Printf("Insert response: %s\n", string(respBody))
}

// 查询 Chroma 数据库
func queryChroma(chromaURL string, query ChromaQuery) []ChromaDocument {
	// 构造查询请求
	url := fmt.Sprintf("%s/collections/default/query", chromaURL)
	body, err := json.Marshal(query)
	if err != nil {
		log.Fatal("Failed to marshal query:", err)
	}

	resp, err := http.Post(url, "application/json", bytes.NewBuffer(body))
	if err != nil {
		log.Fatal("Failed to query Chroma:", err)
	}
	defer resp.Body.Close()

	// 读取查询响应
	respBody, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatal("Failed to read response:", err)
	}

	// 解析查询结果
	var queryResults []ChromaDocument
	err = json.Unmarshal(respBody, &queryResults)
	if err != nil {
		log.Fatal("Failed to unmarshal response:", err)
	}

	return queryResults
}

3. 如何工作 ​

  • 插入文档：insertDocument 函数将包含向量数据和文本的文档插入到 Chroma 数据库中。你可以在插入文档时指定文本以及与之关联的向量（例如，文本的嵌入向量）。

  • 查询相似文档：queryChroma 函数通过一个查询向量从 Chroma 数据库中查询最相似的文档。你可以指定查询向量以及希望返回的结果数量（例如，最相似的 5 个文档）。

4. Chroma API 端点 ​

  • 插入文档：POST /collections/{collection_name}/documents，该端点用于将文档和它的向量插入到指定的集合中。
  • 查询文档：POST /collections/{collection_name}/query，该端点用于查询与指定向量相似的文档。

5. 更多功能 ​

  • 集合管理：你可以创建、删除和管理不同的集合（/collections），每个集合可以包含多文档。可以为每个集合定义不同的配置。
  • 元数据：除了文本和向量数据，Chroma 还支持将元数据（如标签、类别等）与向量一起存储，以便更精细的查询。

6. 性能优化 ​

  • 批量插入：对于大量文档的插入，Chroma 支持批量插入操作，以提高性能。
  • 索引优化：根据文档的数量和查询的频繁程度，可以调整 Chroma 的索引策略以提升查询效率。

总结 ​

通过 Go 客户端与 Chroma 向量数据库交互，你可以实现高效的向量存储、查询和相似度计算等功能。与 Chroma 服务的交互主要是通过 HTTP API 进行的，借助 Go 的 net/http 包，你可以很容易地向数据库中插入数据并进行相似度查询。