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

package rag

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"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}
}

// SecretLookup 由组合根注入。service 是 models.embedding / models.rerank。
var SecretLookup func(service, fallback string) string

func lookupSecret(service, fallback string) string {
	if SecretLookup != nil {
		return SecretLookup(service, fallback)
	}
	return fallback
}

// 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:   lookupSecret("models.embedding", 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:   lookupSecret("models.rerank", 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}
	}
	return p.IndexChunks(context.Background(), owner, docs)
}

// IndexChunks 写入已构造好的知识分块并建向量（EXT-1 KnowledgeWrite embed=true 的路径）。
// AddKnowledge 原地回填 docs[i].ID；embedding 失败降级为纯全文（记日志，不阻断写入）。
func (p *Pipeline) IndexChunks(ctx context.Context, owner core.ID, docs []core.KnowledgeDoc) error {
	if len(docs) == 0 {
		return nil
	}
	if err := p.st.AddKnowledge(docs...); err != nil {
		return err
	}
	if p.embedder == nil {
		return nil
	}
	texts := make([]string, len(docs))
	for i, d := range docs {
		texts[i] = d.Chunk
	}
	vecs, err := p.embedder.Embed(ctx, texts)
	if err != nil {
		log.Printf("[rag] embedding 降级为纯全文（ref=%s）: %v", firstRef(docs), err)
		return nil
	}
	model := embedderModel(p.embedder)
	for i, v := range vecs {
		if len(v) == 0 {
			continue
		}
		if err := p.st.UpsertKnowledgeVectorExt(ctx, owner, docs[i].ID, docs[i].Ref, model, contentHash(docs[i].Chunk), v); err != nil {
			return err
		}
	}
	return nil
}

// embedderModel 模型身份（可选接口：EmbeddingProvider 契约未含，实现侧
// 声明；未声明=空=legacy 无身份行，向后兼容）。
func embedderModel(e core.EmbeddingProvider) string {
	if m, ok := e.(interface{ Model() string }); ok {
		return m.Model()
	}
	return ""
}

// contentHash 分块内容摘要（写入幂等核对用；sha256 前 16 字节十六进制）。
func contentHash(chunk string) string {
	sum := sha256.Sum256([]byte(chunk))
	return hex.EncodeToString(sum[:])[:16]
}

func firstRef(docs []core.KnowledgeDoc) string {
	for _, d := range docs {
		if d.Ref != "" {
			return d.Ref
		}
	}
	return ""
}

// RetrieveShared EXT-8 余项：本人两路如常 + 追加共享两路（向量 Shared + 全文
// Shared，allow=适配器判定），RRF 融合后裁剪 topK。embedder 不可用时共享仅全文
// （与切片 2 行为一致）。
func (p *Pipeline) RetrieveShared(ctx context.Context, owner core.ID, q core.RetrievalQuery, allow func(docOwner, projectID core.ID) bool) ([]core.RetrievalHit, error) {
	own, err := p.Retrieve(ctx, owner, q)
	if err != nil {
		return nil, err
	}
	shared, err := p.st.SearchKnowledgeShared(string(owner), q, allow)
	if err != nil {
		return nil, err
	}
	// 向量共享路：embedder 可用时并入。
	if p.embedder != nil && p.embeddingRoute(ctx).OK() {
		if qvecs, e := p.embedder.Embed(ctx, []string{q.Text}); e == nil && len(qvecs) > 0 && len(qvecs[0]) > 0 {
			topK := q.TopK
			if topK <= 0 {
				topK = 5
			}
			vs, e2 := p.st.SearchKnowledgeVectorsSharedModel(ctx, string(owner), embedderModel(p.embedder), qvecs[0], topK, q, allow)
			if e2 == nil {
				shared = rrfMerge(vs, shared)
			}
		}
	}
	merged := rrfMerge(own, shared)
	if limit := q.TopK; limit > 0 && len(merged) > limit {
		merged = merged[:limit]
	}
	return merged, 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)
	offline, _ := p.embedder.(interface{ Offline() bool })
	if !route.OK() && (offline == nil || !offline.Offline()) {
		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.SearchKnowledgeVectorsModel(ctx, owner, embedderModel(p.embedder), 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
}