Agent-System/internal/kb/retrieve.go

package kb

import (
	"context"
	"database/sql"
	"fmt"
	"log"
	"sort"
	"strconv"
	"strings"
)

// Retriever 检索调度器
type Retriever struct {
	Store    *Store
	Embedder *Embedder
	Reranker *Reranker
}

// NewRetriever 创建检索调度器
func NewRetriever(store *Store, embedder *Embedder, reranker *Reranker) *Retriever {
	return &Retriever{
		Store:    store,
		Embedder: embedder,
		Reranker: reranker,
	}
}

// Search 执行跨领域语义检索与重排
func (r *Retriever) Search(ctx context.Context, query string, opts SearchOptions) ([]HitItem, error) {
	q := strings.TrimSpace(query)
	if q == "" {
		return []HitItem{}, nil
	}
	if opts.Limit <= 0 {
		opts.Limit = 5
	}
	if opts.CandidateLimit <= 0 {
		opts.CandidateLimit = 20
	}
	if len(opts.Domains) == 0 {
		opts.Domains = []KnowledgeDomain{DomainRegulatory}
	}

	var allCandidates []HitItem

	for _, domain := range opts.Domains {
		switch domain {
		case DomainRegulatory:
			hits, err := r.searchRegulatory(ctx, q, opts)
			if err != nil {
				log.Printf("[kb-retriever] searchRegulatory err: %v", err)
			} else {
				allCandidates = append(allCandidates, hits...)
			}
		case DomainProjectAsset:
			if opts.ProjectID <= 0 || opts.RequesterID <= 0 {
				log.Printf("[kb-retriever] searchProjectAsset skipped: project_id and requester_id required")
				continue
			}
			hits, err := r.searchProjectAsset(ctx, q, opts)
			if err != nil {
				log.Printf("[kb-retriever] searchProjectAsset err: %v", err)
			} else {
				allCandidates = append(allCandidates, hits...)
			}
		}
	}

	if len(allCandidates) == 0 {
		return []HitItem{}, nil
	}

	// 排序候选集
	sort.Slice(allCandidates, func(i, j int) bool {
		return allCandidates[i].Score > allCandidates[j].Score
	})

	// 二阶段语义重排 (Cross-Encoder)
	if r.Reranker != nil && r.Reranker.Endpoint != "" && len(allCandidates) > 1 {
		rerankLimit := opts.CandidateLimit
		if len(allCandidates) < rerankLimit {
			rerankLimit = len(allCandidates)
		}
		topList := allCandidates[:rerankLimit]
		docs := make([]string, len(topList))
		for i, h := range topList {
			docs[i] = firstNRunes(h.DocCode+" "+h.Title+"\n"+h.HeadingPath+"\n"+h.Content, 720)
		}
		scores, err := r.Reranker.Rerank(ctx, q, docs)
		if err != nil {
			log.Printf("[kb-retriever] reranker failed (keeping RRF scores): %v", err)
		} else if len(scores) == len(topList) {
			for i := range topList {
				topList[i].Score = scores[i]
			}
			sort.Slice(topList, func(i, j int) bool {
				return topList[i].Score > topList[j].Score
			})
			allCandidates = topList
		}
	}

	// 过滤分数与截断
	var results []HitItem
	for _, it := range allCandidates {
		if opts.MinScore > 0 && it.Score < opts.MinScore {
			continue
		}
		results = append(results, it)
		if len(results) >= opts.Limit {
			break
		}
	}

	return results, nil
}

// searchRegulatory 法规领域检索实现 (Vector + Keyword + RRF)
func (r *Retriever) searchRegulatory(ctx context.Context, q string, opts SearchOptions) ([]HitItem, error) {
	if r.Store == nil || r.Store.Pool == nil {
		return nil, fmt.Errorf("database not connected")
	}

	var vecHits []HitItem
	if r.Embedder != nil && r.Embedder.Endpoint != "" {
		vec, err := r.Embedder.EmbedOne(ctx, q)
		if err != nil {
			log.Printf("[kb-retriever] embed failed: %v", err)
		} else if len(vec) == EmbeddingDim {
			vecHits, err = r.regVectorHits(ctx, vec, opts.IncludeRepealed, 40)
			if err != nil {
				log.Printf("[kb-retriever] regVectorHits failed: %v", err)
			}
		}
	}

	kwHits, err := r.regKeywordHits(ctx, q, opts.IncludeRepealed, 30)
	if err != nil {
		log.Printf("[kb-retriever] regKeywordHits failed: %v", err)
	}

	// Reciprocal Rank Fusion (k=60.0)
	const rrfK = 60.0
	type rrfItem struct {
		hit   HitItem
		score float64
	}
	fused := make(map[string]*rrfItem)

	for rank, h := range vecHits {
		score := 1.0 / (rrfK + float64(rank+1))
		key := fmt.Sprintf("%s|%s", h.DocCode, h.HeadingPath)
		if existing, ok := fused[key]; ok {
			existing.score += score
		} else {
			fused[key] = &rrfItem{hit: h, score: score}
		}
	}

	for rank, h := range kwHits {
		score := 1.0 / (rrfK + float64(rank+1))
		key := fmt.Sprintf("%s|%s", h.DocCode, h.HeadingPath)
		if existing, ok := fused[key]; ok {
			existing.score += score
		} else {
			fused[key] = &rrfItem{hit: h, score: score}
		}
	}

	out := make([]HitItem, 0, len(fused))
	for _, it := range fused {
		it.hit.Score = it.score
		out = append(out, it.hit)
	}
	return out, nil
}

func (r *Retriever) regVectorHits(ctx context.Context, vec []float32, includeRepealed bool, limit int) ([]HitItem, error) {
	lit := VectorLiteral(vec)
	sqlStr := `
		SELECT c.id, COALESCE(d.doc_no,''), COALESCE(d.title,''), COALESCE(d.kind,'unknown'), COALESCE(d.status,'unknown'),
		       COALESCE(c.content,''), COALESCE(c.heading_path,''), COALESCE(c.page,0), COALESCE(c.chunk_no,''),
		       1 - (c.embedding <=> $1::vector) AS score
		FROM reg_chunks c
		JOIN reg_documents d ON d.id=c.document_id
		WHERE c.embedding IS NOT NULL` + regStatusFilter(includeRepealed) + `
		ORDER BY c.embedding <=> $1::vector
		LIMIT $2`

	rows, err := r.Store.Pool.Query(ctx, sqlStr, lit, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	return scanRegHits(rows)
}

func (r *Retriever) regKeywordHits(ctx context.Context, q string, includeRepealed bool, limit int) ([]HitItem, error) {
	like1 := "%" + escapeLike(q) + "%"
	compactQ := strings.ReplaceAll(q, " ", "")
	like2 := "%" + escapeLike(compactQ) + "%"

	sqlStr := `
		SELECT c.id, COALESCE(d.doc_no,''), COALESCE(d.title,''), COALESCE(d.kind,'unknown'), COALESCE(d.status,'unknown'),
		       COALESCE(c.content,''), COALESCE(c.heading_path,''), COALESCE(c.page,0), COALESCE(c.chunk_no,''),
		       CASE WHEN c.tsv @@ plainto_tsquery('simple', $1) THEN ts_rank_cd(c.tsv, plainto_tsquery('simple', $1)) ELSE 0.1 END AS score
		FROM reg_chunks c
		JOIN reg_documents d ON d.id=c.document_id
		WHERE (
		    (c.tsv IS NOT NULL AND c.tsv @@ plainto_tsquery('simple', $1)) OR
		    COALESCE(d.doc_no,'') ILIKE $2 ESCAPE '\' OR COALESCE(d.doc_no,'') ILIKE $3 ESCAPE '\' OR
		    COALESCE(c.heading_path,'') ILIKE $2 ESCAPE '\' OR COALESCE(c.heading_path,'') ILIKE $3 ESCAPE '\' OR
		    COALESCE(c.content,'') ILIKE $2 ESCAPE '\' OR COALESCE(c.content,'') ILIKE $3 ESCAPE '\' OR
		    COALESCE(c.cite_key,'') ILIKE $2 ESCAPE '\' OR COALESCE(c.cite_key,'') ILIKE $3 ESCAPE '\'
		)` + regStatusFilter(includeRepealed) + `
		ORDER BY score DESC
		LIMIT $4`

	rows, err := r.Store.Pool.Query(ctx, sqlStr, q, like1, like2, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	return scanRegHits(rows)
}

func (r *Retriever) searchProjectAsset(ctx context.Context, q string, opts SearchOptions) ([]HitItem, error) {
	sqlStr := `
		SELECT c.id, p.project_code, p.name, d.doc_type, 'active',
		       c.content, c.chunk_title, 0, '',
		       CASE WHEN c.tsv @@ plainto_tsquery('simple', $1) THEN ts_rank_cd(c.tsv, plainto_tsquery('simple', $1)) ELSE 0.1 END AS score
		FROM asset_chunks c
		JOIN asset_projects p ON p.id=c.project_id
		JOIN asset_documents d ON d.id=c.document_id
		WHERE c.tsv IS NOT NULL AND c.tsv @@ plainto_tsquery('simple', $1)
		  AND p.id = $3
		ORDER BY score DESC LIMIT $2`

	rows, err := r.Store.Pool.Query(ctx, sqlStr, q, opts.Limit, opts.ProjectID)
	if err != nil {
		return nil, nil // 预留表空查询容错
	}
	defer rows.Close()

	var out []HitItem
	for rows.Next() {
		var id int64
		var code, title, category, status, content, heading, chunkNo string
		var page int
		var score float64
		if err := rows.Scan(&id, &code, &title, &category, &status, &content, &heading, &page, &chunkNo, &score); err != nil {
			continue
		}
		out = append(out, HitItem{
			Domain:      DomainProjectAsset,
			Category:    category,
			DocCode:     code,
			Title:       title,
			HeadingPath: heading,
			Content:     content,
			Score:       score,
			CiteKey:     fmt.Sprintf("〔项目:%s §%s〕", code, heading),
			Source:      "项目数字资产库",
		})
	}
	return out, rows.Err()
}

func regStatusFilter(includeRepealed bool) string {
	if includeRepealed {
		return ""
	}
	return ` AND LOWER(d.status) IN ('in_force','unknown','active','published')`
}

func scanRegHits(rows interface {
	Next() bool
	Scan(dest ...any) error
	Err() error
}) ([]HitItem, error) {
	var out []HitItem
	for rows.Next() {
		var id int64
		var docNo, title, kind, status, content, heading, chunkNo sql.NullString
		var page sql.NullInt64
		var score sql.NullFloat64

		if err := rows.Scan(&id, &docNo, &title, &kind, &status, &content, &heading, &page, &chunkNo, &score); err != nil {
			return nil, err
		}

		docCode := docNo.String
		if docCode == "" {
			docCode = title.String
		}

		cite := formatCiteKey(docCode, heading.String, chunkNo.String, int(page.Int64), status.String)

		out = append(out, HitItem{
			Domain:      DomainRegulatory,
			Category:    kind.String,
			DocCode:     docCode,
			Title:       title.String,
			HeadingPath: heading.String,
			Content:     content.String,
			Score:       score.Float64,
			CiteKey:     cite,
			Source:      "国家医疗器械法规与技术标准库",
		})
	}
	return out, rows.Err()
}

func formatCiteKey(docCode, heading, chunkNo string, page int, status string) string {
	clause := chunkNo
	if clause == "" {
		clause = heading
	}
	pageStr := ""
	if page > 0 {
		pageStr = " p." + strconv.Itoa(page)
	}
	clauseStr := ""
	if clause != "" {
		clauseStr = " §" + clause
	}
	core := "〔" + docCode + clauseStr + pageStr + "〕"
	if status == "repealed" || status == "superseded" {
		core = "〔已废止〕" + core
	}
	return core
}

func escapeLike(s string) string {
	s = strings.ReplaceAll(s, `\`, `\\`)
	s = strings.ReplaceAll(s, `%`, `\%`)
	s = strings.ReplaceAll(s, `_`, `\_`)
	return s
}

func firstNRunes(s string, n int) string {
	if n <= 0 || len(s) == 0 {
		return ""
	}
	runes := []rune(s)
	if len(runes) <= n {
		return s
	}
	return string(runes[:n])
}