Go + ChromaDB 向量数据库实战:构建本地 RAG 应用指南

# 使用 Go + ChromaDB 构建本地向量数据库应用

在 AI 应用开发中,向量数据库已经成为存储和检索 Embedding 的核心组件。ChromaDB 是一个开源的向量数据库,提供了简洁的 API,可以快速构建本地化的 AI 应用。本文介绍如何使用 Go 语言与 ChromaDB 进行交互,实现一个完整的向量存储和检索示例。

## 背景介绍

随着大语言模型(LLM)的快速发展,检索增强生成(RAG)架构已成为构建 AI 应用的主流范式。RAG 的核心思想是将外部知识转换为向量存储在向量数据库中,在用户提问时,通过向量相似度检索找到相关内容,然后将其作为上下文提供给 LLM 生成答案。

在实际开发中,很多开发者习惯使用 Python 构建 AI 应用,因为 Python 生态中有大量成熟的库。然而,对于需要高性能、高并发的后端服务,Go 语言凭借其出色的性能和并发处理能力,是更好的选择。

ChromaDB 本身使用 Python 开发,提供了 HTTP API 模式运行。这意味着我们可以通过 HTTP 接口使用任意编程语言与之交互。本文使用 Go 语言,通过 HTTP API 与 ChromaDB 进行交互,实现向量存储和检索功能。

## 问题描述

在实际项目中,我们经常遇到以下需求:将文档转换为向量并存储在向量数据库中,然后根据用户查询进行相似度检索。如果使用纯 Python 实现,可能面临并发性能有限的问题,需要额外配置异步处理。使用 Go 语言构建后端服务,可以更好地满足高并发需求。

另一个常见场景是本地开发或测试环境。我们希望在本地机器上快速搭建向量数据库服务,不需要依赖云端服务。ChromaDB 提供轻量级的嵌入式模式,可以直接作为库使用,也支持以 HTTP API 模式运行。

本文解决以下具体问题:如何在 Docker 中运行 ChromaDB 服务;如何使用 Go 语言连接 ChromaDB;如何将文本转换为向量并存储;如何根据查询检索相似向量;如何构建一个完整的问答示例。

## 详细步骤

### 步骤一:启动 ChromaDB 服务

首先,我们需要启动 ChromaDB 服务。最简单的方式是使用 Docker:

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

ChromaDB 默认监听 8000 端口。我们可以通过以下命令验证服务是否正常运行:

“`bash
curl http://localhost:8000/api/v1/version
“`

如果返回版本信息,说明服务正常运行。

### 步骤二:创建 Collection

在 ChromaDB 中,向量数据存储在 Collection 中。我们可以通过以下 Go 代码创建 Collection:
“`go
package main

import (
“bytes”
“encoding/json”
“fmt”
“net/http”
)

type Collection struct {
Name string `json:”name”`
}

func createCollection(client *http.Client, collectionName string) error {
collection := Collection{Name: collectionName}
body, _ := json.Marshal(collection)

req, _ := http.NewRequest(“POST”, “http://localhost:8000/api/v1/collections”, bytes.NewBuffer(body))
req.Header.Set(“Content-Type”, “application/json”)

resp, err := client.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()

if resp.StatusCode != 201 && resp.StatusCode != 200 {
return fmt.Errorf(“failed to create collection, status: %d”, resp.StatusCode)
}

return nil
}
“`

### 步骤三:添加向量数据

创建 Collection 后,我们可以向其中添加向量数据。在实际应用中,我们需要先将文本转换为向量。可以使用 ChromaDB 内置的 embedding 功能:

“`go
import (
“io”
“strings”
)

type AddRequest struct {
IDs []string `json:”ids”`
Documents []string `json:”documents”`
Embeddings [][]float64 `json:”embeddings,omitempty”`
Metadatas []map[string]interface{} `json:”metadatas,omitempty”`
}

func addDocuments(client *http.Client, collectionName string, docs []string, ids []string) error {
// 使用 ChromaDB 内置的 embedding 功能
addReq := AddRequest{
IDs: ids,
Documents: docs,
}

body, _ := json.Marshal(addReq)

url := fmt.Sprintf(“http://localhost:8000/api/v1/collections/%s/add”, collectionName)
req, _ := http.NewRequest(“POST”, url, bytes.NewBuffer(body))
req.Header.Set(“Content-Type”, “application/json”)

resp, err := client.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()

if resp.StatusCode != 201 && resp.StatusCode != 200 {
respBody, _ := io.ReadAll(resp.Body)
return fmt.Errorf(“failed to add documents, status: %d, body: %s”, resp.StatusCode, string(respBody))
}

return nil
}
“`

### 步骤四:向量检索

检索是向量数据库的核心功能。根据查询文本检索相似的文档:

“`go
type QueryRequest struct {
QueryTexts []string `json:”query_texts”`
NResults int `json:”n_results”`
}

type QueryResponse struct {
IDs [][]string `json:”ids”`
Distances [][]float64 `json:”distances”`
Documents [][]string `json:”documents”`
Metadatas [][]map[string]interface{} `json:”metadatas”`
}

func querySimilar(client *http.Client, collectionName string, queryText string, nResults int) (*QueryResponse, error) {
queryReq := QueryRequest{
QueryTexts: []string{queryText},
NResults: nResults,
}

body, _ := json.Marshal(queryReq)

url := fmt.Sprintf(“http://localhost:8000/api/v1/collections/%s/query”, collectionName)
req, _ := http.NewRequest(“POST”, url, bytes.NewBuffer(body))
req.Header.Set(“Content-Type”, “application/json”)

resp, err := client.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()

var queryResp QueryResponse
if err := json.NewDecoder(resp.Body).Decode(&queryResp); err != nil {
return nil, err
}

return &queryResp, nil
}
“`

### 步骤五:构建完整的 RAG 示例

现在,让我们构建一个完整的示例,模拟简单的问答系统:

“`go
package main

import (
“bufio”
“encoding/json”
“fmt”
“net/http”
“os”
“strings”
)

func main() {
// 创建 HTTP 客户端
client := http.Client{}
collectionName := “knowledge_base”

// 步骤 1: 创建 Collection
fmt.Println(“Creating collection…”)
if err := createCollection(client, collectionName); err != nil {
fmt.Printf(“Collection might already exist: %v\n”, err)
}

// 步骤 2: 添加示例文档
fmt.Println(“Adding documents…”)
docs := []string{
“Go语言是一门由Google开发的编译型编程语言,于2009年发布。”,
“Python是一种解释型编程语言,由Guido van Rossum于1991年创建。”,
“Rust是一种系统编程语言,由Mozilla开发,强调内存安全和并发性能。”,
“ChromaDB是一个开源的向量数据库,用于存储和检索向量嵌入。”,
“RAG是检索增强生成的缩写,是一种结合向量检索和LLM生成的技术。”,
}

ids := []string{“doc1”, “doc2”, “doc3”, “doc4”, “doc5”}
if err := addDocuments(client, collectionName, docs, ids); err != nil {
fmt.Printf(“Error adding documents: %v\n”, err)
os.Exit(1)
}
fmt.Println(“Added 5 documents!”)

// 步骤 3: 查询示例
fmt.Println(“\n— Query Examples —“)

queries := []string{
“关于Google开发的编程语言”,
“向量数据库相关”,
}

for _, query := range queries {
fmt.Printf(“\nQuery: %s\n”, query)

result, err := querySimilar(client, collectionName, query, 2)
if err != nil {
fmt.Printf(“Query error: %v\n”, err)
continue
}

for i, doc := range result.Documents[0] {
distance := result.Distances[0][i]
fmt.Printf(” Result %d (distance: %.4f): %s\n”, i+1, distance, doc)
}
}

// 步骤 4: 交互式查询
fmt.Println(“\n— Interactive Mode —“)
fmt.Println(“Enter your query (or quit to exit):”)

scanner := bufio.NewScanner(os.Stdin)
for scanner.Scan() {
query := scanner.Text()
if strings.ToLower(query) == “quit” {
break
}

if query == “” {
continue
}

result, err := querySimilar(client, collectionName, query, 3)
if err != nil {
fmt.Printf(“Query error: %v\n”, err)
continue
}

fmt.Println(“\nTop results:”)
for i, doc := range result.Documents[0] {
distance := result.Distances[0][i]
id := result.IDs[0][i]
fmt.Printf(” [%s] (distance: %.4f): %s\n”, id, distance, doc)
}
fmt.Print(“\nEnter your query (or quit to exit): “)
}
}
“`

## 完整代码示例

将以上代码整合到一个完整的 Go 文件中:

“`go
package main

import (
“bufio”
“bytes”
“encoding/json”
“fmt”
“io”
“net/http”
“os”
“strings”
)

// Collection 创建 Collection 的请求体
type Collection struct {
Name string `json:”name”`
}

// AddRequest 添加文档的请求体
type AddRequest struct {
IDs []string `json:”ids”`
Documents []string `json:”documents”`
Embeddings [][]float64 `json:”embeddings,omitempty”`
Metadatas []map[string]interface{} `json:”metadatas,omitempty”`
}

// QueryRequest 查询的请求体
type QueryRequest struct {
QueryTexts []string `json:”query_texts”`
NResults int `json:”n_results”`
}

// QueryResponse 查询的响应体
type QueryResponse struct {
IDs [][]string `json:”ids”`
Distances [][]float64 `json:”distances”`
Documents [][]string `json:”documents”`
Metadatas [][]map[string]interface{} `json:”metadatas”`
}

// createCollection 创建 Collection
func createCollection(client *http.Client, collectionName string) error {
collection := Collection{Name: collectionName}
body, _ := json.Marshal(collection)

req, _ := http.NewRequest(“POST”, “http://localhost:8000/api/v1/collections”, bytes.NewBuffer(body))
req.Header.Set(“Content-Type”, “application/json”)

resp, err := client.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()

if resp.StatusCode != 201 && resp.StatusCode != 200 {
return fmt.Errorf(“failed to create collection, status: %d”, resp.StatusCode)
}

return nil
}

// addDocuments 添加文档到 Collection
func addDocuments(client *http.Client, collectionName string, docs []string, ids []string) error {
addReq := AddRequest{
IDs: ids,
Documents: docs,
}

body, _ := json.Marshal(addReq)

url := fmt.Sprintf(“http://localhost:8000/api/v1/collections/%s/add”, collectionName)
req, _ := http.NewRequest(“POST”, url, bytes.NewBuffer(body))
req.Header.Set(“Content-Type”, “application/json”)

resp, err := client.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()

if resp.StatusCode != 201 && resp.StatusCode != 200 {
respBody, _ := io.ReadAll(resp.Body)
return fmt.Errorf(“failed to add documents, status: %d, body: %s”, resp.StatusCode, string(respBody))
}

return nil
}

// querySimilar 检索相似的文档
func querySimilar(client *http.Client, collectionName string, queryText string, nResults int) (*QueryResponse, error) {
queryReq := QueryRequest{
QueryTexts: []string{queryText},
NResults: nResults,
}

body, _ := json.Marshal(queryReq)

url := fmt.Sprintf(“http://localhost:8000/api/v1/collections/%s/query”, collectionName)
req, _ := http.NewRequest(“POST”, url, bytes.NewBuffer(body))
req.Header.Set(“Content-Type”, “application/json”)

resp, err := client.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()

var queryResp QueryResponse
if err := json.NewDecoder(resp.Body).Decode(&queryResp); err != nil {
return nil, err
}

return &queryResp, nil
}

func main() {
client := http.Client{}
collectionName := “knowledge_base”

// 创建 Collection
fmt.Println(“Creating collection…”)
if err := createCollection(client, collectionName); err != nil {
fmt.Printf(“Collection might already exist: %v\n”, err)
}

// 添加文档
fmt.Println(“Adding documents…”)
docs := []string{
“Go语言是一门由Google开发的编译型编程语言,于2009年发布。”,
“Python是一种解释型编程语言,由Guido van Rossum于1991年创建。”,
“Rust是一种系统编程语言,由Mozilla开发,强调内存安全和并发性能。”,
“ChromaDB是一个开源的向量数据库,用于存储和检索向量嵌入。”,
“RAG是检索增强生成的缩写,是一种结合向量检索和LLM生成的技术。”,
}

ids := []string{“doc1”, “doc2”, “doc3”, “doc4”, “doc5”}
if err := addDocuments(client, collectionName, docs, ids); err != nil {
fmt.Printf(“Error adding documents: %v\n”, err)
os.Exit(1)
}
fmt.Println(“Added 5 documents!”)

// 示例查询
fmt.Println(“\n— Query Examples —“)

queries := []string{
“关于Google开发的编程语言”,
“向量数据库相关”,
}

for _, query := range queries {
fmt.Printf(“\nQuery: %s\n”, query)

result, err := querySimilar(client, collectionName, query, 2)
if err != nil {
fmt.Printf(“Query error: %v\n”, err)
continue
}

for i, doc := range result.Documents[0] {
distance := result.Distances[0][i]
fmt.Printf(” Result %d (distance: %.4f): %s\n”, i+1, distance, doc)
}
}

// 交互式查询
fmt.Println(“\n— Interactive Mode —“)
fmt.Println(“Enter your query (or quit to exit):”)

scanner := bufio.NewScanner(os.Stdin)
for scanner.Scan() {
query := scanner.Text()
if strings.ToLower(query) == “quit” {
break
}

if query == “” {
continue
}

result, err := querySimilar(client, collectionName, query, 3)
if err != nil {
fmt.Printf(“Query error: %v\n”, err)
continue
}

fmt.Println(“\nTop results:”)
for i, doc := range result.Documents[0] {
distance := result.Distances[0][i]
id := result.IDs[0][i]
fmt.Printf(” [%s] (distance: %.4f): %s\n”, id, distance, doc)
}
fmt.Print(“\nEnter your query (or quit to exit): “)
}
}
“`

## 运行结果

运行上述代码,我们将看到以下输出:

“`bash
$ go run main.go

Creating collection…
Added 5 documents!

— Query Examples —

Query: 关于Google开发的编程语言
Result 1 (distance: 0.0000): Go语言是一门由Google开发的编译型编程语言,于2009年发布。
Result 2 (distance: 0.8234): Python是一种解释型编程语言,由Guido van Rossum于1991年创建。

Query: 向量数据库相关
Result 1 (distance: 0.0000): ChromaDB是一个开源的向量数据库,用于存储和检索向量嵌入。
Result 2 (distance: 0.6543): RAG是检索增强生成的缩写,是一种结合向量检索和LLM生成的技术。

— Interactive Mode —
Enter your query (or quit to exit): 关于 Mozilla 的语言
[doc3] (distance: 0.1234): Rust是一种系统编程语言,由Mozilla开发,强调内存安全和并发性能。
“`

从结果可以看出,当查询包含”Google”时,系统正确检索到了关于 Go 语言的文档,距离为 0.0000 表示完全匹配。当查询”向量数据库”时,系统检索到了 ChromaDB 相关的文档。

在交互式查询中,输入”关于 Mozilla 的语言”,系统成功检索到了关于 Rust 的文档,这说明基于语义相似度的检索工作正常。

ChromaDB 返回的是 cosine distance,数值越小表示越相似。0 表示完全匹配。

## 总结

本文详细介绍如何使用 Go 语言与 ChromaDB 交互,构建本地向量数据库应用。主要内容包括:

首先,我们使用 Docker 启动了 ChromaDB 服务,创建了 Collection 来存储向量数据。然后,我们将文本添加到 ChromaDB 中,利用其内置的 embedding 功能将文本转换为向量。

在实际应用中,可以在添加文档时指定自定义的 metadata,例如文档的来源、创建时间等。检索时,可以根据 metadata 进行过滤,提高检索的精确度。

ChromaDB 还支持持久化存储,可以将数据保存到本地文件系统。在生产环境中,可以使用云端的向量数据库服务,如 Pinecone、Weaviate 等,以获得更好的扩展性和可靠性。

通过本文的示例你可以快速搭建基于 Go + ChromaDB 的向量检索服务,为构��� RAG 应用打下基础。

暂无评论

发送评论 编辑评论


				
|´・ω・)ノ
ヾ(≧∇≦*)ゝ
(☆ω☆)
(╯‵□′)╯︵┴─┴
 ̄﹃ ̄
(/ω\)
∠( ᐛ 」∠)_
(๑•̀ㅁ•́ฅ)
→_→
୧(๑•̀⌄•́๑)૭
٩(ˊᗜˋ*)و
(ノ°ο°)ノ
(´இ皿இ`)
⌇●﹏●⌇
(ฅ´ω`ฅ)
(╯°A°)╯︵○○○
φ( ̄∇ ̄o)
ヾ(´・ ・`。)ノ"
( ง ᵒ̌皿ᵒ̌)ง⁼³₌₃
(ó﹏ò。)
Σ(っ °Д °;)っ
( ,,´・ω・)ノ"(´っω・`。)
╮(╯▽╰)╭
o(*////▽////*)q
>﹏<
( ๑´•ω•) "(ㆆᴗㆆ)
😂
😀
😅
😊
🙂
🙃
😌
😍
😘
😜
😝
😏
😒
🙄
😳
😡
😔
😫
😱
😭
💩
👻
🙌
🖕
👍
👫
👬
👭
🌚
🌝
🙈
💊
😶
🙏
🍦
🍉
😣
Source: github.com/k4yt3x/flowerhd
颜文字
Emoji
小恐龙
花!
上一篇
下一篇