feat(memory/vector): 向量检索引擎(SQLite BLOB 内积默认 + PG pgvector 真语义增强,配置后真生效)
This commit is contained in:
1 parent
ef79edd96a
commit
fb880572a4
6 files changed
+1009
No files matched your search
@@ -0,0 +1,92 @@
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNewPG_InvalidURL_异常(t *testing.T) {
|
||||
t.Parallel()
|
||||
if _, err := NewPG(context.Background(), " ", 4, testLogger()); !errors.Is(err, ErrInvalidPGURL) {
|
||||
t.Fatalf("NewPG(blank url) err = %v, want ErrInvalidPGURL", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewPG_InvalidDim_异常(t *testing.T) {
|
||||
t.Parallel()
|
||||
// dim is validated before any connection is opened, so no network I/O here.
|
||||
if _, err := NewPG(context.Background(), "postgres://host/db", 0, testLogger()); !errors.Is(err, ErrInvalidDim) {
|
||||
t.Fatalf("NewPG(dim 0) err = %v, want ErrInvalidDim", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPGVectorLiteral(t *testing.T) {
|
||||
t.Parallel()
|
||||
tests := []struct {
|
||||
name string
|
||||
v Vector
|
||||
want string
|
||||
}{
|
||||
{name: "single", v: Vector{1}, want: "[1]"},
|
||||
{name: "three", v: Vector{1.5, 2, -0.25}, want: "[1.5,2,-0.25]"},
|
||||
{name: "empty", v: Vector{}, want: "[]"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
if got := pgVectorLiteral(tt.v); got != tt.want {
|
||||
t.Fatalf("pgVectorLiteral(%v) = %q, want %q", tt.v, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestPG_RequiresEnv runs a real PostgreSQL+pgvector round trip when DB_PG_URL
|
||||
// is set, and skips otherwise. There is no PG server in this environment by
|
||||
// default, so the pgvector-backed Add/Search path is validated end-to-end only
|
||||
// when the operator supplies a live database with the pgvector extension. The
|
||||
// SQLite path (the default) is fully covered by the sqlite_test.go suite.
|
||||
func TestPG_RequiresEnv(t *testing.T) {
|
||||
url := os.Getenv("DB_PG_URL")
|
||||
if url == "" {
|
||||
t.Skip("DB_PG_URL not set; skipping real PostgreSQL+pgvector integration test")
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
idx, err := NewPG(ctx, url, 4, testLogger())
|
||||
if err != nil {
|
||||
t.Fatalf("NewPG: %v", err)
|
||||
}
|
||||
idx.withEmbedder(&mockEmbedder{dim: 4})
|
||||
defer func() { _ = idx.Close() }()
|
||||
|
||||
if err := idx.Add(ctx, "a", Vector{1, 0, 0, 0}); err != nil {
|
||||
t.Fatalf("Add a: %v", err)
|
||||
}
|
||||
if err := idx.Add(ctx, "b", Vector{0, 1, 0, 0}); err != nil {
|
||||
t.Fatalf("Add b: %v", err)
|
||||
}
|
||||
|
||||
hits, err := idx.Search(ctx, Vector{1, 0, 0, 0}, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("Search: %v", err)
|
||||
}
|
||||
if len(hits) != 1 || hits[0].ID != "a" {
|
||||
t.Fatalf("pg top hit = %+v, want single a", hits)
|
||||
}
|
||||
if hits[0].Score < 0.99 {
|
||||
t.Fatalf("pg a score = %v, want ~1.0", hits[0].Score)
|
||||
}
|
||||
|
||||
// The embedder is wired, so Index.Embed is usable.
|
||||
vs, err := idx.Embed(ctx, []string{"hello"})
|
||||
if err != nil {
|
||||
t.Fatalf("Embed: %v", err)
|
||||
}
|
||||
if len(vs) != 1 || len(vs[0]) != 4 {
|
||||
t.Fatalf("Embed => %d vectors, want one 4-dim", len(vs))
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user