t
18c9354095
- 注入数据库、运行时配置、统一响应、文件存储和平台 AI 能力，补充业务读写工具与客户快捷操作契约。 - 移除模块内重复的组织、客户、工单、标签、技能、旧工作流、MCP 和迁移实现，将身份权限与业务主体交由宿主管理。 - 使用 libSQL 重构向量存储，并完善图片消息、访客身份、排队调度、企业微信和支持聊天页面。 - 统一 HTTP、DTO 与 WebSocket 的 snake_case 协议，补齐模块初始化、业务动作和公共载荷等回归测试。
95 lines
3.0 KiB
Go
95 lines
3.0 KiB
Go
package vectordb
|
|
|
|
import (
|
|
"context"
|
|
"path/filepath"
|
|
"testing"
|
|
|
|
"code.tczkiot.com/wlw/ai-agent/internal/pkg/config"
|
|
)
|
|
|
|
func TestLibSQLProviderVectorLifecycle(t *testing.T) {
|
|
databaseDir := t.TempDir()
|
|
provider, err := NewLibSQLProvider(&config.VectorDBConfig{
|
|
Path: filepath.Join(databaseDir, "vectors.db"),
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewLibSQLProvider() error = %v", err)
|
|
}
|
|
t.Cleanup(func() { _ = provider.Close() })
|
|
|
|
ctx := context.Background()
|
|
const collection = "knowledge_chunks"
|
|
if err := provider.CreateCollection(ctx, collection, 3); err != nil {
|
|
t.Fatalf("CreateCollection() error = %v", err)
|
|
}
|
|
vectors := []Vector{
|
|
{ID: "a", Vector: []float32{1, 0, 0}, Payload: ChunkPayload{KnowledgeBaseID: 1, DocumentID: 10, Content: "alpha"}},
|
|
{ID: "b", Vector: []float32{0, 1, 0}, Payload: ChunkPayload{KnowledgeBaseID: 2, DocumentID: 20, Content: "beta"}},
|
|
{ID: "c", Vector: []float32{0.9, 0.1, 0}, Payload: ChunkPayload{KnowledgeBaseID: 1, DocumentID: 11, Content: "gamma"}},
|
|
}
|
|
if err := provider.UpsertVectors(ctx, collection, vectors); err != nil {
|
|
t.Fatalf("UpsertVectors() error = %v", err)
|
|
}
|
|
|
|
info, err := provider.GetCollection(ctx, collection)
|
|
if err != nil {
|
|
t.Fatalf("GetCollection() error = %v", err)
|
|
}
|
|
if info.Dimension != 3 || info.PointCount != 3 || info.Status != "ready" {
|
|
t.Fatalf("GetCollection() = %+v", info)
|
|
}
|
|
|
|
results, err := provider.Search(ctx, &SearchRequest{
|
|
CollectionName: collection,
|
|
Vector: []float32{1, 0, 0},
|
|
TopK: 2,
|
|
ScoreThreshold: 0,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Search() error = %v", err)
|
|
}
|
|
if len(results) != 2 || results[0].ID != "a" {
|
|
t.Fatalf("Search() = %+v, want a first", results)
|
|
}
|
|
|
|
filtered, err := provider.Search(ctx, &SearchRequest{
|
|
CollectionName: collection,
|
|
Vector: []float32{1, 0, 0},
|
|
TopK: 10,
|
|
ScoreThreshold: 0,
|
|
Filter: &SearchFilter{KnowledgeBaseIDs: []int64{2}},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("filtered Search() error = %v", err)
|
|
}
|
|
if len(filtered) != 1 || filtered[0].ID != "b" || filtered[0].Payload.Content != "beta" {
|
|
t.Fatalf("filtered Search() = %+v", filtered)
|
|
}
|
|
if err := provider.Close(); err != nil {
|
|
t.Fatalf("Close() error = %v", err)
|
|
}
|
|
provider, err = NewLibSQLProvider(&config.VectorDBConfig{Path: filepath.Join(databaseDir, "vectors.db")})
|
|
if err != nil {
|
|
t.Fatalf("reopen NewLibSQLProvider() error = %v", err)
|
|
}
|
|
info, err = provider.GetCollection(ctx, collection)
|
|
if err != nil || info.PointCount != 3 {
|
|
t.Fatalf("reopened GetCollection() = %+v, %v", info, err)
|
|
}
|
|
|
|
if err := provider.DeleteVectors(ctx, collection, []string{"a"}); err != nil {
|
|
t.Fatalf("DeleteVectors() error = %v", err)
|
|
}
|
|
if err := provider.DeleteCollection(ctx, collection); err != nil {
|
|
t.Fatalf("DeleteCollection() error = %v", err)
|
|
}
|
|
collections, err := provider.ListCollections(ctx)
|
|
if err != nil {
|
|
t.Fatalf("ListCollections() error = %v", err)
|
|
}
|
|
if len(collections) != 0 {
|
|
t.Fatalf("ListCollections() = %v, want empty", collections)
|
|
}
|
|
}
|