go整合向量数据库ChromaDB | 卡卡罗特

卡卡罗特

2026-04-28 · via 卡卡罗特

package main

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

// --- 配置和全局变量 ---
var (
	baseURL       string
	client        *http.Client
	defaultTenant = "default_tenant"
	defaultDB     = "default_database"
)

// --- 数据结构 ---
type CreateCollectionV2Req struct {
	Name     string                 `json:"name"`
	Metadata map[string]interface{} `json:"metadata,omitempty"`
}

type CreateCollectionV2Resp struct {
	ID       string `json:"id"`
	Name     string `json:"name"`
	Tenant   string `json:"tenant"`
	Database string `json:"database"`
}

type AddRecordsV2Req struct {
	IDs        []string                 `json:"ids"`
	Embeddings [][]float32              `json:"embeddings"`
	Documents  []string                 `json:"documents,omitempty"`
	Metadatas  []map[string]interface{} `json:"metadatas,omitempty"`
}

type QueryV2Req struct {
	QueryEmbeddings [][]float32 `json:"query_embeddings"`
	NRResults       int         `json:"n_results"`
	Include         []string    `json:"include,omitempty"`
}

type QueryV2Resp struct {
	IDs       [][]string                 `json:"ids"`
	Documents [][]string                 `json:"documents"`
	Distances [][]float32                `json:"distances"`
	Metadatas [][]map[string]interface{} `json:"metadatas,omitempty"`
}

// 存储集合信息
var currentCollection struct {
	ID       string
	Name     string
	Tenant   string
	Database string
}

func main() {
	baseURL = "http://127.0.0.1:8000"
	client = &http.Client{Timeout: 10 * time.Second}
	fmt.Printf("🔗 连接 ChromaDB v2: %s\n", baseURL)

	// 1. 心跳检测
	if !heartbeat() {
		fmt.Println("❌ 无法连接到 ChromaDB")
		return
	}
	fmt.Println("✅ 服务连接正常")

	// 2. 获取或创建集合
	collectionName := "my-go-collection"
	if err := getOrCreateCollection(collectionName); err != nil {
		fmt.Printf("❌ 获取/创建集合失败: %v\n", err)
		return
	}
	fmt.Printf("✅ 集合已就绪，ID: %s, 名称: %s\n", currentCollection.ID, currentCollection.Name)

	// 3. 添加数据（使用完整的路径）
	if err := addData(); err != nil {
		fmt.Printf("❌ 添加数据失败: %v\n", err)
		return
	}
	fmt.Println("✅ 数据添加成功")

	// 4. 查询数据
	queryData()
}

// 心跳检测
func heartbeat() bool {
	resp, err := client.Get(baseURL + "/api/v2/heartbeat")
	if err != nil {
		return false
	}
	defer resp.Body.Close()
	return resp.StatusCode == http.StatusOK
}

// 获取或创建集合
func getOrCreateCollection(name string) error {
	// 先尝试获取已存在的集合
	err := getCollection(name)
	if err == nil {
		fmt.Printf("📂 找到已存在的集合: %s (ID: %s)\n", name, currentCollection.ID)
		return nil
	}

	// 集合不存在，创建新集合
	fmt.Printf("📂 集合 '%s' 不存在，正在创建...\n", name)
	return createCollection(name)
}

// 获取已存在的集合
func getCollection(name string) error {
	url := fmt.Sprintf("%s/api/v2/tenants/%s/databases/%s/collections",
		baseURL, defaultTenant, defaultDB)

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

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body))
	}

	// 解析集合列表
	var collections []CreateCollectionV2Resp
	if err := json.Unmarshal(body, &collections); err != nil {
		return err
	}

	// 查找指定名称的集合
	for _, col := range collections {
		if col.Name == name {
			currentCollection.ID = col.ID
			currentCollection.Name = col.Name
			currentCollection.Tenant = col.Tenant
			currentCollection.Database = col.Database
			return nil
		}
	}

	return fmt.Errorf("集合 '%s' 不存在", name)
}

// 创建新集合
func createCollection(name string) error {
	url := fmt.Sprintf("%s/api/v2/tenants/%s/databases/%s/collections",
		baseURL, defaultTenant, defaultDB)

	reqBody := CreateCollectionV2Req{
		Name: name,
		Metadata: map[string]interface{}{
			"description": "Go 客户端测试集合",
			"created_at":  time.Now().Format(time.RFC3339),
		},
	}

	jsonData, _ := json.Marshal(reqBody)
	fmt.Printf("📤 创建集合请求: %s\n", string(jsonData))

	resp, err := client.Post(url, "application/json", bytes.NewBuffer(jsonData))
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(resp.Body)

	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated {
		return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body))
	}

	// 解析响应获取集合信息
	var result CreateCollectionV2Resp
	if err := json.Unmarshal(body, &result); err != nil {
		return err
	}

	currentCollection.ID = result.ID
	currentCollection.Name = result.Name
	currentCollection.Tenant = result.Tenant
	currentCollection.Database = result.Database

	fmt.Printf("📥 创建成功: ID=%s, Name=%s, Tenant=%s, Database=%s\n",
		result.ID, result.Name, result.Tenant, result.Database)
	return nil
}

// 添加数据 - 使用完整的 API 路径
func addData() error {
	// 关键修正：使用完整的路径包含 tenant 和 database
	url := fmt.Sprintf("%s/api/v2/tenants/%s/databases/%s/collections/%s/add",
		baseURL,
		currentCollection.Tenant,
		currentCollection.Database,
		currentCollection.ID)

	fmt.Printf("\n📤 添加数据到集合: %s\n", currentCollection.Name)
	fmt.Printf("🔗 完整 URL: %s\n", url)

	reqBody := AddRecordsV2Req{
		IDs: []string{"doc1", "doc2"},
		Embeddings: [][]float32{
			{0.1, 0.2, 0.3, 0.4, 0.5},
			{0.6, 0.7, 0.8, 0.9, 1.0},
		},
		Documents: []string{
			"Go 语言是并发的王者",
			"ChromaDB 是一个向量数据库",
		},
		Metadatas: []map[string]interface{}{
			{"source": "go-demo", "language": "go", "type": "programming"},
			{"source": "go-demo", "topic": "database", "type": "technology"},
		},
	}

	jsonData, err := json.Marshal(reqBody)
	if err != nil {
		return err
	}

	fmt.Printf("📦 请求体: %s\n", string(jsonData))

	resp, err := client.Post(url, "application/json", bytes.NewBuffer(jsonData))
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	fmt.Printf("📥 响应状态: %d\n", resp.StatusCode)

	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated {
		return fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body))
	}

	if len(body) > 0 {
		fmt.Printf("📥 响应内容: %s\n", string(body))
	}
	return nil
}

// 查询数据 - 使用完整的 API 路径
func queryData() {
	url := fmt.Sprintf("%s/api/v2/tenants/%s/databases/%s/collections/%s/query",
		baseURL,
		currentCollection.Tenant,
		currentCollection.Database,
		currentCollection.ID)

	fmt.Printf("\n🔍 查询 URL: %s\n", url)

	reqBody := QueryV2Req{
		QueryEmbeddings: [][]float32{{0.1, 0.2, 0.3, 0.4, 0.5}},
		NRResults:       3,
		Include:         []string{"documents", "distances", "metadatas"},
	}

	jsonData, _ := json.Marshal(reqBody)
	fmt.Printf("🔍 查询请求: %s\n", string(jsonData))

	resp, err := client.Post(url, "application/json", bytes.NewBuffer(jsonData))
	if err != nil {
		fmt.Printf("❌ 查询失败: %v\n", err)
		return
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(resp.Body)

	if resp.StatusCode != http.StatusOK {
		fmt.Printf("❌ 查询失败: HTTP %d - %s\n", resp.StatusCode, string(body))
		return
	}

	var result QueryV2Resp
	if err := json.Unmarshal(body, &result); err != nil {
		fmt.Printf("❌ 解析失败: %v\n原始响应: %s\n", err, string(body))
		return
	}

	// 显示结果
	if len(result.Documents) > 0 && len(result.Documents[0]) > 0 {
		fmt.Println("\n🎉 查询成功！找到以下结果：")
		for i, doc := range result.Documents[0] {
			distance := float32(0)
			if len(result.Distances) > 0 && len(result.Distances[0]) > i {
				distance = result.Distances[0][i]
			}

			// 显示元数据（如果有）
			metadata := ""
			if len(result.Metadatas) > 0 && len(result.Metadatas[0]) > i {
				if lang, ok := result.Metadatas[0][i]["language"]; ok {
					metadata = fmt.Sprintf(" [语言: %v]", lang)
				}
			}

			fmt.Printf("  %d.%s %s (距离: %.4f)\n", i+1, metadata, doc, distance)
		}
	} else {
		fmt.Println("\n⚠️ 未找到相关结果")
	}
}

text

🔗 连接 ChromaDB v2: http://127.0.0.1:8000
✅ 服务连接正常
📂 集合 'my-go-collection' 不存在，正在创建...
📤 创建集合请求: {"name":"my-go-collection","metadata":{"created_at":"2026-04-05T16:21:22+08:00","description":"Go 客户端测试集合"}}
📥 创建成功: ID=29578383-f2a0-4e6c-a02a-5490554493d8, Name=my-go-collection, Tenant=default_tenant, Database=default_database
✅ 集合已就绪，ID: 29578383-f2a0-4e6c-a02a-5490554493d8, 名称: my-go-collection

📤 添加数据到集合: my-go-collection
🔗 完整 URL: http://127.0.0.1:8000/api/v2/tenants/default_tenant/databases/default_database/collections/29578383-f2a0-4e6c-a02a-5490554493d8/add
📦 请求体: {"ids":["doc1","doc2"],"embeddings":[[0.1,0.2,0.3,0.4,0.5],[0.6,0.7,0.8,0.9,1]],"documents":["Go 语言是并发的王者","ChromaDB 是一个向量数据库"],"metadatas":[{"language":"go","source":"go-demo","type":"programming"},{"source":"go-demo","topic":"database","type":"technology"}]}
📥 响应状态: 201
📥 响应内容: {}
✅ 数据添加成功

🔍 查询 URL: http://127.0.0.1:8000/api/v2/tenants/default_tenant/databases/default_database/collections/29578383-f2a0-4e6c-a02a-5490554493d8/query
🔍 查询请求: {"query_embeddings":[[0.1,0.2,0.3,0.4,0.5]],"n_results":3,"include":["documents","distances","metadatas"]}

🎉 查询成功！找到以下结果：
  1. [语言: go] Go 语言是并发的王者 (距离: 0.0000)
  2. ChromaDB 是一个向量数据库 (距离: 1.2500)

此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。原文来自 — 版权归原作者所有。

推荐订阅源

卡卡罗特