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)
|
||
|
|
}
|
||
|
|
}
|