pj0235-eai_agentplatform/eai_agentplatform/backend-go/internal/ai/retrieve.go

package ai

import (
	"math"
	"sort"
	"strings"

	"eai_agentplatform/backend/internal/config"
	"eai_agentplatform/backend/internal/model"
	"eai_agentplatform/backend/internal/store"
)

// Retrieve 混合检索：向量 brute-force 余弦 + 关键词兜底，去重合并取 topK
func Retrieve(cfg *config.Config, query string, topK int) []string {
	if topK <= 0 {
		topK = 5
	}
	var vectorResults []string
	if route, err := config.GetRoute("embed_gen"); err == nil {
		if v, err := vectorRetrieve(route, query, topK); err == nil {
			vectorResults = v
		}
	}
	keywordResults := keywordRetrieve(query, topK*2)

	seen := map[string]bool{}
	out := make([]string, 0, topK)
	for _, s := range append(vectorResults, keywordResults...) {
		s = strings.TrimSpace(s)
		if s == "" || seen[s] {
			continue
		}
		seen[s] = true
		out = append(out, s)
		if len(out) >= topK {
			break
		}
	}
	return out
}

// vectorRetrieve 向量检索：query + 所有 chunk 一次批量 embedding，brute-force 余弦 topK
func vectorRetrieve(route *config.RouteConfig, query string, topK int) ([]string, error) {
	var chunks []model.KnowledgeChunk
	store.DB.Order("id ASC").Find(&chunks)
	if len(chunks) == 0 {
		return nil, nil
	}
	client := NewClient(route)
	inputs := make([]string, 0, len(chunks)+1)
	inputs = append(inputs, query)
	for _, ch := range chunks {
		inputs = append(inputs, ch.Content)
	}
	vecs, err := client.Embed(inputs)
	if err != nil || len(vecs) != len(inputs) {
		if err != nil {
			return nil, err
		}
		return nil, nil
	}
	qv := vecs[0]
	type scored struct {
		idx int
		sim float64
	}
	ss := make([]scored, 0, len(chunks))
	for i := 1; i < len(vecs); i++ {
		ss = append(ss, scored{i - 1, cosine(qv, vecs[i])})
	}
	sort.Slice(ss, func(a, b int) bool { return ss[a].sim > ss[b].sim })

	seen := map[string]bool{}
	out := make([]string, 0, topK)
	for _, s := range ss {
		content := chunks[s.idx].Content
		if seen[content] {
			continue
		}
		seen[content] = true
		out = append(out, content)
		if len(out) >= topK {
			break
		}
	}
	return out, nil
}

// keywordRetrieve 关键词兜底：term 命中数排序
func keywordRetrieve(query string, topK int) []string {
	terms := splitTerms(query)
	var chunks []model.KnowledgeChunk
	store.DB.Order("id ASC").Find(&chunks)

	type scored struct {
		content string
		score   int
	}
	var ss []scored
	for _, ch := range chunks {
		s := 0
		for _, t := range terms {
			if strings.Contains(ch.Content, t) {
				s++
			}
		}
		if s > 0 {
			ss = append(ss, scored{ch.Content, s})
		}
	}
	sort.Slice(ss, func(a, b int) bool { return ss[a].score > ss[b].score })

	seen := map[string]bool{}
	out := make([]string, 0, topK)
	for _, s := range ss {
		if seen[s.content] {
			continue
		}
		seen[s.content] = true
		out = append(out, s.content)
		if len(out) >= topK {
			break
		}
	}
	return out
}

func splitTerms(q string) []string {
	f := func(r rune) bool {
		switch r {
		case ' ', '，', '。', '？', '！', '、', ',', '.', '?', '!', ':', '：', ';', '；':
			return true
		}
		return false
	}
	terms := strings.FieldsFunc(q, f)
	var out []string
	for _, t := range terms {
		if len([]rune(t)) >= 2 {
			out = append(out, t)
		}
	}
	if len(out) == 0 {
		out = []string{q}
	}
	return out
}

func cosine(a, b []float64) float64 {
	if len(a) == 0 || len(a) != len(b) {
		return 0
	}
	var dot, na, nb float64
	for i := range a {
		dot += a[i] * b[i]
		na += a[i] * a[i]
		nb += b[i] * b[i]
	}
	if na == 0 || nb == 0 {
		return 0
	}
	return dot / (math.Sqrt(na) * math.Sqrt(nb))
}