Agent-System-Backend/internal/rag/rag_test.go

package rag

import (
	"context"
	"errors"
	"path/filepath"
	"strings"
	"testing"

	"git-ssh.u1680832.nyat.app/weichen/Agent-System-Backend.git/internal/settings"
	"git-ssh.u1680832.nyat.app/weichen/Agent-System-Backend.git/internal/store"
	core "git-ssh.u1680832.nyat.app/weichen/Agent-System-Core.git"
)

// fakeEmbedder 确定性向量：按首字符映射，验证混合检索排序无需真实模型。
type fakeEmbedder struct{ fail bool }

func (f fakeEmbedder) Embed(_ context.Context, texts []string) ([][]float32, error) {
	if f.fail {
		return nil, errors.New("embed down")
	}
	out := make([][]float32, len(texts))
	for i, t := range texts {
		if strings.HasPrefix(t, "A") {
			out[i] = []float32{1, 0}
		} else {
			out[i] = []float32{0, 1}
		}
	}
	return out, nil
}

type fakeReranker struct{ called *bool }

func (f fakeReranker) Rerank(_ context.Context, _ string, docs []core.KnowledgeDoc, topK int) ([]core.KnowledgeDoc, error) {
	*f.called = true
	// 反转顺序以验证 rerank 生效。
	out := make([]core.KnowledgeDoc, 0, len(docs))
	for i := len(docs) - 1; i >= 0; i-- {
		out = append(out, docs[i])
	}
	if len(out) > topK {
		out = out[:topK]
	}
	return out, nil
}

func newTestPipeline(t *testing.T, embedder core.EmbeddingProvider, reranker core.Reranker) (*Pipeline, *store.Store) {
	t.Helper()
	st, err := store.OpenSQLite(filepath.Join(t.TempDir(), "test.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { st.Close() })
	set := settings.NewService(st)
	return NewPipeline(st, set, embedder, reranker), st
}

func TestChunkerPacksParagraphs(t *testing.T) {
	text := strings.Repeat("段", 300) + "\n\n" + strings.Repeat("落", 300) + "\n\n" + strings.Repeat("字", 300)
	chunks := chunkText(text, 500, 80)
	if len(chunks) < 2 {
		t.Fatalf("expected packing into >=2 chunks, got %d", len(chunks))
	}
	total := 0
	for i, c := range chunks {
		if c.Index != i {
			t.Fatalf("chunk index %d != %d", c.Index, i)
		}
		total += len([]rune(c.Text))
		if len([]rune(c.Text)) > 500+1 { // 允许连接换行
			t.Fatalf("chunk %d oversized: %d", i, len([]rune(c.Text)))
		}
	}
	if total == 0 {
		t.Fatal("no content")
	}
}

func TestChunkerHardSplitsLongParagraphWithOverlap(t *testing.T) {
	text := strings.Repeat("长", 1200)
	chunks := chunkText(text, 500, 100)
	if len(chunks) < 3 {
		t.Fatalf("expected >=3 window chunks, got %d", len(chunks))
	}
	// 相邻窗口有重叠上下文。
	a := []rune(chunks[0].Text)
	b := []rune(chunks[1].Text)
	if len(a) < 500 || len(b) == 0 {
		t.Fatalf("window sizes: %d %d", len(a), len(b))
	}
	tail := string(a[len(a)-100:])
	head := string(b[:min(100, len(b))])
	if tail != head {
		t.Fatalf("overlap mismatch:\n%q\n%q", tail, head)
	}
}

func TestChunkerEmpty(t *testing.T) {
	if got := chunkText("  \n\n  ", 500, 80); len(got) != 0 {
		t.Fatalf("empty input should produce no chunks, got %d", len(got))
	}
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

// 索引：分块入全文库；embedding 配置时向量入库；失败降级不阻断。
func TestIndexDocumentChunksAndVectors(t *testing.T) {
	pipe, st := newTestPipeline(t, fakeEmbedder{}, nil)
	text := strings.Repeat("Alpha 文档内容。", 120) // 约 1200 字 → 多块
	if err := pipe.IndexDocument("o", "p", "upload", "a.md", "art_1", text); err != nil {
		t.Fatal(err)
	}
	hits, err := st.SearchKnowledge("o", core.RetrievalQuery{Text: "文档"})
	if err != nil {
		t.Fatal(err)
	}
	if len(hits) < 2 {
		t.Fatalf("chunked fulltext expected >=2 hits, got %d", len(hits))
	}
	// 向量已入库（fakeEmbedder 首字符 'A' → [1,0]）。
	vecHits, err := st.SearchKnowledgeVectors(context.Background(), "o", []float32{1, 0}, 5, core.RetrievalQuery{})
	if err != nil {
		t.Fatal(err)
	}
	if len(vecHits) == 0 || vecHits[0].Kind != "vector" {
		t.Fatalf("vector hits: %+v", vecHits)
	}
}

func TestIndexDocumentEmbeddingFailureDegrades(t *testing.T) {
	pipe, st := newTestPipeline(t, fakeEmbedder{fail: true}, nil)
	if err := pipe.IndexDocument("o", "p", "upload", "a.md", "art_1", "Beta 内容"); err != nil {
		t.Fatalf("embedding failure must degrade, got %v", err)
	}
	vecHits, _ := st.SearchKnowledgeVectors(context.Background(), "o", []float32{1, 0}, 5, core.RetrievalQuery{})
	if len(vecHits) != 0 {
		t.Fatalf("no vectors expected on embedding failure, got %d", len(vecHits))
	}
	fhits, _ := st.SearchKnowledge("o", core.RetrievalQuery{Text: "Beta"})
	if len(fhits) == 0 {
		t.Fatal("fulltext must still have the doc")
	}
}

// 检索：未配置 embedding → 纯全文；配置 → 向量+全文 RRF；rerank 生效。
func TestRetrieveFallbackAndHybrid(t *testing.T) {
	// 无 embedder：纯全文路径。
	pipe, _ := newTestPipeline(t, nil, nil)
	pipe.IndexDocument("o", "p", "upload", "d", "r1", "Alpha 目标文档")
	hits, err := pipe.Retrieve(context.Background(), "o", core.RetrievalQuery{Text: "目标"})
	if err != nil {
		t.Fatal(err)
	}
	if len(hits) == 0 || hits[0].Kind != "fulltext" {
		t.Fatalf("fulltext fallback: %+v", hits)
	}

	// 有 embedder 且路由已配置：查询向量 A 类 → A 文档排前（RRF 向量路主导）。
	// P2 受管：models.* 是系统受管项，配置写入系统层行（owner=OwnerSystemRow）。
	pipe2, st2 := newTestPipeline(t, fakeEmbedder{}, nil)
	if err := st2.MergeSettingsScope(settings.OwnerSystemRow, settings.ScopeSystem, map[string]any{
		"models": map[string]any{"embedding": map[string]any{"baseUrl": "http://fake.internal", "model": "e-test"}},
	}); err != nil {
		t.Fatal(err)
	}
	pipe2.IndexDocument("o", "", "upload", "d1", "rA", "Alpha 语义相关文档")
	pipe2.IndexDocument("o", "", "upload", "d2", "rB", "Beta 另一个主题文档")
	hits2, err := pipe2.Retrieve(context.Background(), "o", core.RetrievalQuery{Text: "A 语义查询", TopK: 2})
	if err != nil {
		t.Fatal(err)
	}
	if len(hits2) == 0 {
		t.Fatal("hybrid hits empty")
	}
	if hits2[0].Doc.Ref != "rA" {
		t.Fatalf("vector path should rank rA first, got ref=%s", hits2[0].Doc.Ref)
	}
	if hits2[0].Kind != "vector" {
		t.Fatalf("top hit kind: %s", hits2[0].Kind)
	}
}

func TestRetrieveRerankReorders(t *testing.T) {
	called := false
	pipe, _ := newTestPipeline(t, fakeEmbedder{}, fakeReranker{called: &called})
	pipe.IndexDocument("o", "", "upload", "d1", "rA", "Alpha 文档一")
	pipe.IndexDocument("o", "", "upload", "d2", "rB", "Beta 文档二")
	// rerank 未配置路由（settings 无 models.rerank）→ 不调用；此处直接验证 Reranker
	// 组件行为（管线内调用受路由开关，见 Retrieve 实现）。
	var docs []core.KnowledgeDoc
	docs = append(docs, core.KnowledgeDoc{ID: "a", Chunk: "A"}, core.KnowledgeDoc{ID: "b", Chunk: "B"})
	out, err := fakeReranker{called: &called}.Rerank(context.Background(), "q", docs, 2)
	if err != nil || !called {
		t.Fatalf("reranker: %v called=%v", err, called)
	}
	if out[0].ID != "b" || out[1].ID != "a" {
		t.Fatalf("rerank order: %v", out)
	}
}

// 租户边界：向量检索不越权。
func TestVectorTenantIsolation(t *testing.T) {
	pipe, st := newTestPipeline(t, fakeEmbedder{}, nil)
	pipe.IndexDocument("alice", "", "upload", "d", "rA", "Alice 私有文档")
	vecHits, err := st.SearchKnowledgeVectors(context.Background(), "bob", []float32{1, 0}, 5, core.RetrievalQuery{})
	if err != nil {
		t.Fatal(err)
	}
	if len(vecHits) != 0 {
		t.Fatalf("cross-tenant vector leak: %+v", vecHits)
	}
}

// 维度不匹配的向量不参与检索（换 embedding 模型后旧向量自然失效）。
func TestVectorDimMismatchSkipped(t *testing.T) {
	pipe, st := newTestPipeline(t, fakeEmbedder{}, nil)
	pipe.IndexDocument("o", "", "upload", "d", "r", "Alpha 文档")
	hits, err := st.SearchKnowledgeVectors(context.Background(), "o", []float32{1, 0, 0}, 5, core.RetrievalQuery{})
	if err != nil {
		t.Fatal(err)
	}
	if len(hits) != 0 {
		t.Fatalf("dim mismatch must skip, got %+v", hits)
	}
}

// ---- EXT-1：IndexChunks（KnowledgeWrite embed=true 的底层路径） ----

func TestIndexChunksWritesFulltextAndVectors(t *testing.T) {
	p, st := newTestPipeline(t, fakeEmbedder{}, nil)
	docs := []core.KnowledgeDoc{
		{Owner: "alice", Source: "plugin:rdm", Title: "规格", Ref: "spec-1", Chunk: "Alpha 内容甲"},
		{Owner: "alice", Source: "plugin:rdm", Title: "规格", Ref: "spec-1", Chunk: "beta 内容乙"},
	}
	if err := p.IndexChunks(context.Background(), "alice", docs); err != nil {
		t.Fatal(err)
	}
	// ID 已被 AddKnowledge 回填，向量按回填 ID 关联。
	for _, d := range docs {
		if d.ID == "" {
			t.Fatal("doc ID not backfilled")
		}
	}
	// 全文行落库。
	hits, err := st.SearchKnowledge("alice", core.RetrievalQuery{Text: "内容甲"})
	if err != nil {
		t.Fatal(err)
	}
	if len(hits) == 0 {
		t.Fatal("fulltext row missing")
	}
	// 向量行落库（fakeEmbedder：非 A 前缀 → [0,1]）。
	vhits, err := st.SearchKnowledgeVectors(context.Background(), "alice", []float32{0, 1}, 5, core.RetrievalQuery{})
	if err != nil {
		t.Fatal(err)
	}
	if len(vhits) == 0 {
		t.Fatal("vector rows missing")
	}
	// owner 作用域：bob 检索不到 alice 的向量。
	bhits, err := st.SearchKnowledgeVectors(context.Background(), "bob", []float32{0, 1}, 5, core.RetrievalQuery{})
	if err != nil {
		t.Fatal(err)
	}
	if len(bhits) != 0 {
		t.Fatalf("owner leak: %+v", bhits)
	}
}

func TestIndexChunksEmbedFailureDegrades(t *testing.T) {
	p, st := newTestPipeline(t, fakeEmbedder{fail: true}, nil)
	docs := []core.KnowledgeDoc{{Owner: "alice", Source: "plugin:rdm", Ref: "spec-2", Chunk: "降级内容"}}
	if err := p.IndexChunks(context.Background(), "alice", docs); err != nil {
		t.Fatalf("embed failure must degrade to fulltext, got %v", err)
	}
	hits, err := st.SearchKnowledge("alice", core.RetrievalQuery{Text: "降级"})
	if err != nil {
		t.Fatal(err)
	}
	if len(hits) == 0 {
		t.Fatal("fulltext row missing after degrade")
	}
	vhits, err := st.SearchKnowledgeVectors(context.Background(), "alice", []float32{0, 1}, 5, core.RetrievalQuery{})
	if err != nil {
		t.Fatal(err)
	}
	if len(vhits) != 0 {
		t.Fatalf("no vectors expected after degrade, got %+v", vhits)
	}
}