weichen
4fdbfbc9a2
Details (original message): Add the vector RAG pipeline: paragraph chunker, OpenAI-compatible embeddings and rerank clients, dual-tier vector store with cosine search, and RRF hybrid retrieval
231 lines
7.5 KiB
Go
231 lines
7.5 KiB
Go
package rag
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
|
|
core "agent-system-core"
|
|
"agent-system-backend/internal/settings"
|
|
"agent-system-backend/internal/store"
|
|
)
|
|
|
|
// 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 向量路主导）。
|
|
pipe2, st2 := newTestPipeline(t, fakeEmbedder{}, nil)
|
|
if err := st2.MergeSettingsScope(store.OwnerDefault, settings.ScopeUser, 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)
|
|
}
|
|
}
|