Files
ai-gateway-go/internal/workbench/vector_integration_test.go
T
superidou b536672000 feat(m8): P2 pgvector + Ollama 向量化与语义检索
- PostgreSQL 切换 pgvector/pgvector:pg17 镜像;迁移 000024 建 vector 扩展、
  knowledge_chunks.embedding vector(1024) + HNSW 余弦索引,retrieval_mode 放宽三态
- OllamaEmbedder 本地 bge-m3 批量嵌入,404 惰性 pull 重试,维度/超时校验,可整体关闭
- SemanticRetriever/HybridRetriever + NewRetriever 按 retrieval_mode 分发,缺 embedder 回退 FTS
- 文档入库同步批量向量化;Ollama 故障降级入库 + embedding_failed 事件
- 修复 pgx CopyFrom 对 vector 列二进制编码误读:COPY 基础列后同事务 unnest 批量回填
- 修复降级路径 embeddings=nil 索引越界 panic(Add 与 Reprocess)
- 知识库列表 vectorized_chunk_count + 前端三态检索模式选择与向量化覆盖率
- 单测 embedder/retrievers + 集成 TestKnowledgeVectorLifecycle 全绿

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-12 15:16:32 +08:00

91 lines
3.7 KiB
Go

package workbench
import (
"context"
"os"
"strings"
"testing"
"time"
"aigateway.local/core/internal/platform/config"
"aigateway.local/core/internal/platform/database"
)
// TestKnowledgeVectorLifecycle 跑真实 PostgreSQL(pgvector)+ Ollama:
// vector 模式知识库导入文档 → embedding 非空 → SemanticRetriever 语义命中;
// 并验证 embedder 失败时文档照常入库且发出 embedding_failed 事件。
// 需要 WORKBENCH_TEST_DATABASE_URL 与 WORKBENCH_TEST_OLLAMA_URL。
func TestKnowledgeVectorLifecycle(t *testing.T) {
databaseURL := os.Getenv("WORKBENCH_TEST_DATABASE_URL")
ollamaURL := os.Getenv("WORKBENCH_TEST_OLLAMA_URL")
if databaseURL == "" || ollamaURL == "" {
t.Skip("WORKBENCH_TEST_DATABASE_URL and WORKBENCH_TEST_OLLAMA_URL are not set")
}
ctx := context.Background()
pool, err := database.Open(ctx, config.Database{URL: databaseURL, MaxConns: 8, MinConns: 0})
if err != nil {
t.Fatal(err)
}
defer pool.Close()
actorID := "55555555-5555-4555-8555-555555555555"
_, err = pool.Exec(ctx, `INSERT INTO gateway.admin_accounts(id,username,password_hash,role) VALUES($1,'m8-vector-test','test','superadmin') ON CONFLICT(id) DO NOTHING`, actorID)
if err != nil {
t.Fatal(err)
}
name := "m8-vector-integration-kb"
cleanup := func() {
_, _ = pool.Exec(ctx, `DELETE FROM gateway.knowledge_bases WHERE name=$1 OR name=$2`, name, name+"-fts")
}
cleanup()
defer cleanup()
assets := NewService(pool)
assets.SetEmbedder(NewOllamaEmbedder(OllamaEmbedderConfig{BaseURL: ollamaURL, Model: "bge-m3", Dim: 1024, BatchSize: 8, Timeout: 60 * time.Second}))
kb, err := assets.SaveKnowledgeBase(ctx, KnowledgeBase{Name: name, Description: "vector itest", RetrievalMode: "vector", ChunkSize: 300, ChunkOverlap: 40, DepartmentIDs: []string{}, Enabled: true}, actorID, true)
if err != nil {
t.Fatal(err)
}
doc, err := assets.AddKnowledgeDocument(ctx, kb.ID, "向量检索测试", "text", "", "pgvector 语义检索依赖 bge-m3 向量。Ollama 本地生成嵌入。", actorID)
if err != nil {
t.Fatal(err)
}
var embedded int
if err = pool.QueryRow(ctx, `SELECT count(*) FROM gateway.knowledge_chunks WHERE document_id=$1 AND embedding IS NOT NULL`, doc.ID).Scan(&embedded); err != nil {
t.Fatal(err)
}
if embedded == 0 || embedded != doc.ChunkCount {
t.Fatalf("expected all chunks embedded, got %d/%d", embedded, doc.ChunkCount)
}
// 语义检索:查询词与正文无字面重合也应命中(余弦相似度)。
hitChunks, err := (&SemanticRetriever{pool: pool, embedder: assets.Embedder()}).Search(ctx, kb.ID, "语义相似度匹配", 4)
if err != nil {
t.Fatalf("semantic search: %v", err)
}
if len(hitChunks) == 0 {
t.Fatal("expected semantic hit")
}
if !strings.Contains(hitChunks[0].Content, "pgvector") {
t.Fatalf("expected pgvector content in top hit, got %q", hitChunks[0].Content)
}
// 降级路径:embedder 失败 → 文档照常入库 + embedding_failed 事件。
assets.SetEmbedder(failingEmbedder{})
kbFts, err := assets.SaveKnowledgeBase(ctx, KnowledgeBase{Name: name + "-fts", Description: "degraded", RetrievalMode: "vector", ChunkSize: 300, ChunkOverlap: 40, DepartmentIDs: []string{}, Enabled: true}, actorID, true)
if err != nil {
t.Fatal(err)
}
degradedDoc, err := assets.AddKnowledgeDocument(ctx, kbFts.ID, "降级测试", "text", "", "没有向量也能入库。", actorID)
if err != nil {
t.Fatalf("degraded add should succeed, got %v", err)
}
var failed bool
if err = pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM gateway.outbox_events WHERE event_type='knowledge_document.embedding_failed' AND aggregate_id=$1)`, degradedDoc.ID).Scan(&failed); err != nil {
t.Fatal(err)
}
if !failed {
t.Fatal("expected embedding_failed event when embedder is down")
}
}