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

package rag

import (
	"context"
	"log"

	"agent-system-backend/internal/settings"
	"agent-system-backend/internal/store"
	core "agent-system-core"
)

// Pipeline 实现设计 4.12 检索管线：
//
//	索引：文档 → 分块 → 全文入库 →（embedding 配置时）向量入库
//	查询：Embedding → 向量检索 + 全文检索 → RRF 融合 →（rerank 配置时）重排 → TopK
//
// 未配置 embedding 时自动退化为纯全文档（默认行为不变）。
type Pipeline struct {
	st       *store.Store
	set      *settings.Service
	chunker  FixedChunker
	embedder core.EmbeddingProvider
	reranker core.Reranker
}

func NewPipeline(st *store.Store, set *settings.Service, embedder core.EmbeddingProvider, reranker core.Reranker) *Pipeline {
	return &Pipeline{st: st, set: set, chunker: FixedChunker{}, embedder: embedder, reranker: reranker}
}

// SetEmbedder/SetReranker 允许组合根后注入（embedder 的路由解析依赖本管线自身）。
func (p *Pipeline) SetEmbedder(e core.EmbeddingProvider) { p.embedder = e }
func (p *Pipeline) SetReranker(r core.Reranker)          { p.reranker = r }

// EmbeddingRoute 是 embedding 模型路由解析（供 HTTP 客户端闭包使用）。
func (p *Pipeline) EmbeddingRoute(ctx context.Context) Route { return p.embeddingRoute(ctx) }

// RerankRoute 是 rerank 模型路由解析。
func (p *Pipeline) RerankRoute(ctx context.Context) Route { return p.rerankRoute(ctx) }

func (p *Pipeline) embeddingRoute(ctx context.Context) Route {
	eff := p.set.Effective(store.OwnerDefault, "")
	return Route{
		Base:  eff.GetString("models.embedding.baseUrl", ""),
		Key:   eff.GetString("models.embedding.apiKey", ""),
		Model: eff.GetString("models.embedding.model", ""),
	}
}

func (p *Pipeline) rerankRoute(ctx context.Context) Route {
	eff := p.set.Effective(store.OwnerDefault, "")
	return Route{
		Base:  eff.GetString("models.rerank.baseUrl", ""),
		Key:   eff.GetString("models.rerank.apiKey", ""),
		Model: eff.GetString("models.rerank.model", ""),
	}
}

// IndexDocument 检索索引入口（core.DocumentIndexer）：分块 + 全文 + 可选向量。
// embedding 失败降级为纯全文（记日志，不阻断上传/工件链路）。
func (p *Pipeline) IndexDocument(owner, projectID core.ID, source, title, ref, plainText string) error {
	if plainText == "" {
		return nil
	}
	doc := core.ParsedDocument{Blocks: []core.DocBlock{{Type: core.BlockPara, Text: plainText}}}
	chunks := p.chunker.Chunk(doc)
	if len(chunks) == 0 {
		return nil
	}
	docs := make([]core.KnowledgeDoc, len(chunks))
	for i, c := range chunks {
		docs[i] = core.KnowledgeDoc{Owner: owner, ProjectID: projectID, Source: source, Title: title, Ref: ref, Chunk: c.Text}
	}
	if err := p.st.AddKnowledge(docs...); err != nil {
		return err
	}
	if p.embedder == nil {
		return nil
	}
	texts := make([]string, len(chunks))
	for i, c := range chunks {
		texts[i] = c.Text
	}
	vecs, err := p.embedder.Embed(context.Background(), texts)
	if err != nil {
		log.Printf("[rag] embedding 降级为纯全文（ref=%s）: %v", ref, err)
		return nil
	}
	for i, v := range vecs {
		if len(v) == 0 {
			continue
		}
		if err := p.st.UpsertKnowledgeVector(context.Background(), owner, docs[i].ID, ref, v); err != nil {
			return err
		}
	}
	return nil
}

// Retrieve 统一检索（core.Retriever）。
func (p *Pipeline) Retrieve(ctx context.Context, owner core.ID, q core.RetrievalQuery) ([]core.RetrievalHit, error) {
	topK := q.TopK
	if topK <= 0 {
		topK = 5
	}
	if p.embedder == nil {
		return p.st.SearchKnowledge(owner, q)
	}
	route := p.embeddingRoute(ctx)
	if !route.OK() {
		return p.st.SearchKnowledge(owner, q)
	}
	qvecs, err := p.embedder.Embed(ctx, []string{q.Text})
	if err != nil || len(qvecs) == 0 || len(qvecs[0]) == 0 {
		// 查询向量失败：退化为全文（不因可选能力故障阻断主链路）。
		if err != nil {
			log.Printf("[rag] query embedding fallback to fulltext: %v", err)
		}
		return p.st.SearchKnowledge(owner, q)
	}
	q.TopK = topK * 2 // 每路多取，融合后再裁剪
	vhits, err := p.st.SearchKnowledgeVectors(ctx, owner, qvecs[0], topK*2, q)
	if err != nil {
		return nil, err
	}
	fhits, err := p.st.SearchKnowledge(owner, q)
	if err != nil {
		return nil, err
	}
	merged := rrfMerge(vhits, fhits)
	if p.reranker != nil && p.rerankRoute(ctx).OK() && len(merged) > 0 {
		docs := make([]core.KnowledgeDoc, len(merged))
		for i, h := range merged {
			docs[i] = h.Doc
		}
		reordered, err := p.reranker.Rerank(ctx, q.Text, docs, topK)
		if err != nil {
			log.Printf("[rag] rerank fallback to rrf order: %v", err)
		} else {
			byID := map[core.ID]core.RetrievalHit{}
			for _, h := range merged {
				byID[h.Doc.ID] = h
			}
			out := make([]core.RetrievalHit, 0, len(reordered))
			for _, d := range reordered {
				if h, ok := byID[d.ID]; ok {
					out = append(out, h)
				}
			}
			if len(out) >= topK {
				return out[:topK], nil
			}
			return out, nil
		}
	}
	if len(merged) > topK {
		merged = merged[:topK]
	}
	return merged, nil
}

// rrfMerge 倒数排名融合（k=60）：同文档多路命中分数相加；向量路保留 Kind=vector。
func rrfMerge(lists ...[]core.RetrievalHit) []core.RetrievalHit {
	const k = 60
	type acc struct {
		hit   core.RetrievalHit
		score float64
	}
	order := []core.ID{}
	byID := map[core.ID]*acc{}
	for _, list := range lists {
		for rank, h := range list {
			a, ok := byID[h.Doc.ID]
			if !ok {
				a = &acc{hit: h}
				byID[h.Doc.ID] = a
				order = append(order, h.Doc.ID)
			}
			a.score += 1.0 / float64(k+rank+1)
			if h.Kind == "vector" {
				a.hit.Kind = "vector" // 向量命中优先标注
			}
		}
	}
	out := make([]core.RetrievalHit, 0, len(order))
	for _, id := range order {
		h := byID[id].hit
		h.Score = byID[id].score
		out = append(out, h)
	}
	// 稳定排序（分数降序）。
	for i := 1; i < len(out); i++ {
		for j := i; j > 0 && out[j].Score > out[j-1].Score; j-- {
			out[j], out[j-1] = out[j-1], out[j]
		}
	}
	return out
}