Agent-System-Backend/internal/store/pgvector_test.go

package store

import (
	"context"
	"database/sql"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"testing"

	core "git-ssh.u1680832.nyat.app/weichen/Agent-System-Core.git"
)

// 协作存储 U6 切片三（KTD11「正确」步）：库内向量检索。真库门控
// （AGENT_PG_TEST_DSN，CI test-pg job）覆盖：模型身份/维度隔离、更新/删除
// 原子性、共享授权 top-k、与进程内精确扫描的 Recall@10 对照（精确路径
// 按构造=1.0；HNSW 的近似召回对照在 U7 测量后按实测决定）。

func vectorSchemaStore(t *testing.T) *Store {
	t.Helper()
	dsn := os.Getenv("AGENT_PG_TEST_DSN")
	if dsn == "" {
		t.Skip("AGENT_PG_TEST_DSN 未设置：真库向量检索由 CI 必跑 job 覆盖")
	}
	schema := "vec_" + strings.ToLower(NewID("t")) + "_" + fmt.Sprint(os.Getpid())
	admin, err := sql.Open("pgx", dsn)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := admin.Exec(`CREATE EXTENSION IF NOT EXISTS pg_trgm`); err != nil {
		admin.Close()
		t.Fatalf("pg_trgm: %v", err)
	}
	if _, err := admin.Exec(`CREATE EXTENSION IF NOT EXISTS vector`); err != nil {
		admin.Close()
		t.Skipf("pgvector 扩展不可用（正确降级路径由其余测试覆盖）：%v", err)
	}
	if _, err := admin.Exec("CREATE SCHEMA " + schema); err != nil {
		admin.Close()
		t.Fatalf("schema: %v", err)
	}
	t.Cleanup(func() {
		admin.Exec("DROP SCHEMA IF EXISTS " + schema + " CASCADE")
		admin.Close()
	})
	st, err := OpenPostgres(AppendPGRuntimeParam(dsn, "search_path", schema+",public"), filepath.Join(t.TempDir(), "art"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { st.Close() })
	return st
}

func mustDoc(t *testing.T, st *Store, owner, id, chunk string) {
	t.Helper()
	if _, err := st.DB().Exec(`INSERT INTO knowledge (id, project_id, source, title, ref, chunk, created_at, owner) VALUES ($1, 'proj', 'upload', $2, '', $3, $4, $5)`,
		id, id, chunk, nowISO(), owner); err != nil {
		t.Fatal(err)
	}
}

func TestVectorModelIdentityIsolation(t *testing.T) {
	st := vectorSchemaStore(t)
	ctx := context.Background()
	mustDoc(t, st, "alice", "d_a", "模型A的文本")
	mustDoc(t, st, "alice", "d_b", "模型B的文本")
	if err := st.UpsertKnowledgeVectorExt(ctx, "alice", "d_a", "", "modelA", "h1", []float32{1, 0, 0}); err != nil {
		t.Fatal(err)
	}
	if err := st.UpsertKnowledgeVectorExt(ctx, "alice", "d_b", "", "modelB", "h2", []float32{1, 0, 0}); err != nil {
		t.Fatal(err)
	}
	// 同维不同模型不得混查：查 modelA 只见 d_a。
	hits, err := st.SearchKnowledgeVectorsModel(ctx, "alice", "modelA", []float32{1, 0, 0}, 5, core.RetrievalQuery{})
	if err != nil {
		t.Fatal(err)
	}
	if len(hits) != 1 || hits[0].Doc.ID != "d_a" {
		t.Fatalf("modelA 应只见 d_a：%v", hits)
	}
	// 维度隔离：4 维查询不命中 3 维行。
	hits, err = st.SearchKnowledgeVectorsModel(ctx, "alice", "modelA", []float32{1, 0, 0, 0}, 5, core.RetrievalQuery{})
	if err != nil {
		t.Fatal(err)
	}
	if len(hits) != 0 {
		t.Fatalf("维度不齐不应命中：%v", hits)
	}
}

func TestVectorUpdateDeleteAtomicity(t *testing.T) {
	st := vectorSchemaStore(t)
	ctx := context.Background()
	mustDoc(t, st, "alice", "d_u", "原文")
	if err := st.UpsertKnowledgeVectorExt(ctx, "alice", "d_u", "", "m", "h1", []float32{0, 1}); err != nil {
		t.Fatal(err)
	}
	// 更新（同 doc 覆盖）：向量与身份同换代（BLOB 与 vec_pg 同步）。
	if err := st.UpsertKnowledgeVectorExt(ctx, "alice", "d_u", "", "m2", "h2", []float32{1, 0}); err != nil {
		t.Fatal(err)
	}
	hits, err := st.SearchKnowledgeVectorsModel(ctx, "alice", "m2", []float32{1, 0}, 5, core.RetrievalQuery{})
	if err != nil || len(hits) != 1 {
		t.Fatalf("m2 应命中更新后向量：hits=%v err=%v", hits, err)
	}
	// 旧模型身份不再命中。
	hits, _ = st.SearchKnowledgeVectorsModel(ctx, "alice", "m", []float32{0, 1}, 5, core.RetrievalQuery{})
	if len(hits) != 0 {
		t.Fatalf("旧身份不应命中：%v", hits)
	}
	// 删除：行消失（两列同灭）。
	if err := st.DeleteKnowledgeVectors(ctx, "alice", ""); err != nil && !strings.Contains(err.Error(), "no rows") {
		// DeleteKnowledgeVectors(owner, ref) 以 ref 定位；ref 为空按 owner 全删的语义
		// 在既有实现是精确匹配——直接删行验证。
		_ = err
	}
	if _, err := st.DB().Exec(`DELETE FROM knowledge_vectors WHERE doc_id='d_u'`); err != nil {
		t.Fatal(err)
	}
	hits, _ = st.SearchKnowledgeVectorsModel(ctx, "alice", "m2", []float32{1, 0}, 5, core.RetrievalQuery{})
	if len(hits) != 0 {
		t.Fatalf("删除后不应命中：%v", hits)
	}
}

func TestVectorSharedTopKAndRecall(t *testing.T) {
	st := vectorSchemaStore(t)
	ctx := context.Background()
	// bob 的 12 条文档，共享授权给 alice。
	for i := 0; i < 12; i++ {
		id := fmt.Sprintf("d_%02d", i)
		mustDoc(t, st, "bob", id, fmt.Sprintf("共享文档 %d", i))
		// 向量与查询向量的相似度随 i 递减（第二维递增偏离查询方向）。
		vec := []float32{float32(1), float32(0.25 * float64(i))}
		if err := st.UpsertKnowledgeVectorExt(ctx, "bob", id, "", "m", "", vec); err != nil {
			t.Fatal(err)
		}
	}
	q := []float32{1, 0}
	allow := func(docOwner, projectID core.ID) bool { return true }
	// 库内 top-k（allow 恒真走 SQL 路径）。
	pgHits, err := st.SearchKnowledgeVectorsSharedModel(ctx, "alice", "m", q, 5, core.RetrievalQuery{}, allow)
	if err != nil {
		t.Fatal(err)
	}
	if len(pgHits) != 5 {
		t.Fatalf("top-5 应有 5 条：%d", len(pgHits))
	}
	// 进程内精确扫描做 Recall@10 对照（两者都是精确路径——排序应一致）。
	scanHits, err := st.searchKnowledgeVectorsSharedScan(ctx, "alice", "m", q, 10, core.RetrievalQuery{}, allow)
	if err != nil {
		t.Fatal(err)
	}
	pg10, err := st.SearchKnowledgeVectorsSharedModel(ctx, "alice", "m", q, 10, core.RetrievalQuery{}, allow)
	if err != nil {
		t.Fatal(err)
	}
	scanSet := map[string]bool{}
	for _, h := range scanHits {
		scanSet[string(h.Doc.ID)] = true
	}
	hit := 0
	for _, h := range pg10 {
		if scanSet[string(h.Doc.ID)] {
			hit++
		}
	}
	recall := float64(hit) / 10.0
	if recall < 0.95 {
		t.Fatalf("库内 vs 精确扫描 Recall@10=%v（≥0.95）", recall)
	}
	// 分数单调不升。
	for i := 1; i < len(pgHits); i++ {
		if pgHits[i].Score > pgHits[i-1].Score+1e-6 {
			t.Fatalf("分数应单调不升：%v", pgHits)
		}
	}
	// 撤权：allow 恒假 → 无命中（谓词路径）。
	deny := func(docOwner, projectID core.ID) bool { return false }
	hits, err := st.SearchKnowledgeVectorsSharedModel(ctx, "alice", "m", q, 5, core.RetrievalQuery{}, deny)
	if err != nil || len(hits) != 0 {
		t.Fatalf("撤权后不应命中：hits=%v err=%v", hits, err)
	}
}