185 lines
6.9 KiB
Go
185 lines
6.9 KiB
Go
package kb
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"unicode/utf8"
|
||
|
||
"tcm-agent/internal/dao"
|
||
)
|
||
|
||
// ========================================================================
|
||
// Searcher：知识库检索（V1 仅全文检索；V2 预留向量分支）
|
||
// ========================================================================
|
||
// V1 路径：
|
||
// 1. 从 xk_system_config 读 search_mode（默认 fulltext）
|
||
// 2. 走 KBFullTextSearch（MySQL ngram FULLTEXT）
|
||
// 先 BOOLEAN MODE（短关键词=精确子串匹配），零命中或语法报错时
|
||
// 自动回退 NATURAL LANGUAGE MODE（整句自动 2-gram 分词 + 相关度排序），
|
||
// 让 Agent 链路传整段病历上下文时也能检索到内容
|
||
// 3. 返回带 score 的结果
|
||
//
|
||
// V2 路径（设计已留好）：
|
||
// - search_mode=vector：调 Embedder 把 query 转向量 → 拉库内已向量化 chunk → 算余弦
|
||
// - search_mode=blend：fulltext + vector 加权融合
|
||
// - 都靠 KBHasVectorized() 判断库是否已经向量化，未向量化自动降级到 fulltext
|
||
// ========================================================================
|
||
|
||
// maxQueryRunes 检索词长度上限（rune 计）
|
||
//
|
||
// Agent 链路会把整段病历上下文当 query 传进来（可能几千字），
|
||
// ngram 自然语言模式下 query 每个 2-gram 都参与匹配，过长会拖慢查询；
|
||
// 截断到与 chunk 尺寸同量级（500）足够表达检索意图
|
||
const maxQueryRunes = 500
|
||
|
||
// naturalFallbackMinRunes 触发自然语言回退的最小查询长度
|
||
//
|
||
// ≤2 字的查询本身就是单个 2-gram，BOOLEAN 短语匹配和自然语言模式
|
||
// 结果完全等价，零命中时回退纯属浪费一次查询
|
||
const naturalFallbackMinRunes = 3
|
||
|
||
// SearchResult 单条检索结果
|
||
type SearchResult struct {
|
||
ChunkID uint `json:"chunk_id"`
|
||
LibraryID uint `json:"library_id"`
|
||
DocID uint `json:"doc_id"`
|
||
Title string `json:"title"`
|
||
Content string `json:"content"`
|
||
Score float64 `json:"score"` // 相关度（V1 是 FULLTEXT BM25 分数）
|
||
SourceType string `json:"source_type"` // 命中类型：fulltext / vector / blend
|
||
}
|
||
|
||
// SearchOptions 检索参数
|
||
type SearchOptions struct {
|
||
LibraryID uint // 必填：在哪个库里检索
|
||
Query string // 必填：检索词
|
||
TopK int // 默认 5
|
||
Mode string // fulltext / vector / blend（V1 强制 fulltext）
|
||
}
|
||
|
||
// Searcher 检索器
|
||
//
|
||
// 通过 NewSearcher 构造，依赖 Embedder（V1 是 NoopEmbedder）
|
||
type Searcher struct {
|
||
embedder Embedder
|
||
}
|
||
|
||
// NewSearcher 构造检索器
|
||
//
|
||
// 参数：
|
||
// - embedder：向量化器（V1 传 NoopEmbedder，V2 传真实实现）
|
||
func NewSearcher(embedder Embedder) *Searcher {
|
||
return &Searcher{embedder: embedder}
|
||
}
|
||
|
||
// Search 执行检索
|
||
//
|
||
// 决策流程：
|
||
// 1. mode=fulltext 或库未向量化 → 走 KBFullTextSearch
|
||
// 2. mode=vector 且 embedder 可用且库已向量化 → 走 V2 向量分支
|
||
// 3. mode=blend → 两路都跑（V2）
|
||
//
|
||
// V1 实现只走分支 1，分支 2/3 留 TODO 注释
|
||
func (s *Searcher) Search(ctx context.Context, opt SearchOptions) ([]*SearchResult, error) {
|
||
if opt.LibraryID == 0 {
|
||
return nil, fmt.Errorf("kb: library_id 不能为空")
|
||
}
|
||
if opt.Query == "" {
|
||
return []*SearchResult{}, nil
|
||
}
|
||
if opt.TopK <= 0 {
|
||
opt.TopK = 5
|
||
}
|
||
// 超长查询截断（Agent 链路传整段病历上下文的场景）
|
||
opt.Query = truncateRunes(opt.Query, maxQueryRunes)
|
||
|
||
mode := opt.Mode
|
||
if mode == "" {
|
||
mode = "fulltext"
|
||
}
|
||
|
||
// 判断库是否已向量化（V2 用，V1 始终 false）
|
||
vectorized, _ := dao.KBHasVectorized(opt.LibraryID)
|
||
|
||
// 分支选择
|
||
switch {
|
||
case mode == "fulltext":
|
||
return s.searchFullText(opt)
|
||
case mode == "vector" && s.embedder.Available() && vectorized:
|
||
// V2 实现：return s.searchVector(ctx, opt)
|
||
return nil, fmt.Errorf("kb: 向量检索 V2 未实现，请配置 ai_kb_search_mode=fulltext")
|
||
case mode == "blend" && s.embedder.Available() && vectorized:
|
||
// V2 实现：fulltext + vector 加权融合
|
||
return nil, fmt.Errorf("kb: 混合检索 V2 未实现，请配置 ai_kb_search_mode=fulltext")
|
||
default:
|
||
// 自动降级到全文检索（最稳妥）
|
||
return s.searchFullText(opt)
|
||
}
|
||
}
|
||
|
||
// searchFullText 走 MySQL FULLTEXT + ngram 检索（V1 主路径）
|
||
//
|
||
// 两段式策略（查询侧「自动分词」的实现方式）：
|
||
// 1. 先 BOOLEAN MODE：ngram 把整个 query 转成连续 2-gram 短语，等价子串匹配，
|
||
// 短关键词（麻黄汤 / 风寒感冒）和带空格的多词查询（OR 语义）都精确高效
|
||
// 2. 零命中或报错、且 query ≥ 3 字时，回退 NATURAL LANGUAGE MODE：
|
||
// MySQL 自动把 query 拆成 2-gram 做 OR 匹配 + 相关度排序，
|
||
// 自然语言长句（如整段病历上下文）也能召回相关分段
|
||
//
|
||
// 为什么 BOOLEAN 报错也走回退：自然文本里的半角括号/引号/加减号
|
||
// 会被 BOOLEAN MODE 当运算符解析，畸形组合会直接 SQL 报错（如括号不配对），
|
||
// 回退到自然语言模式对这类输入天然免疫
|
||
func (s *Searcher) searchFullText(opt SearchOptions) ([]*SearchResult, error) {
|
||
rows, err := dao.KBFullTextSearch(opt.LibraryID, opt.Query, opt.TopK, false)
|
||
matchType := "fulltext"
|
||
if (err != nil || len(rows) == 0) && utf8.RuneCountInString(opt.Query) >= naturalFallbackMinRunes {
|
||
var nErr error
|
||
rows, nErr = dao.KBFullTextSearch(opt.LibraryID, opt.Query, opt.TopK, true)
|
||
if nErr != nil {
|
||
// 两种模式都失败：报自然语言模式的错（同一根因，如 DB 不可用）
|
||
return nil, fmt.Errorf("kb: FULLTEXT 检索失败: %w", nErr)
|
||
}
|
||
matchType = "fulltext-natural"
|
||
} else if err != nil {
|
||
return nil, fmt.Errorf("kb: FULLTEXT 检索失败: %w", err)
|
||
}
|
||
out := make([]*SearchResult, 0, len(rows))
|
||
for _, r := range rows {
|
||
out = append(out, &SearchResult{
|
||
ChunkID: r.ID,
|
||
LibraryID: r.LibraryID,
|
||
DocID: r.DocID,
|
||
Title: r.Title,
|
||
Content: r.Content,
|
||
Score: r.Score,
|
||
SourceType: matchType,
|
||
})
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
// truncateRunes 按 rune 截断字符串（不破坏 UTF-8 字符边界）
|
||
func truncateRunes(s string, max int) string {
|
||
if max <= 0 || utf8.RuneCountInString(s) <= max {
|
||
return s
|
||
}
|
||
runes := []rune(s)
|
||
return string(runes[:max])
|
||
}
|
||
|
||
// ========================================================================
|
||
// V2 预留：向量检索（不实现，仅占位）
|
||
// ========================================================================
|
||
// func (s *Searcher) searchVector(ctx context.Context, opt SearchOptions) ([]*SearchResult, error) {
|
||
// // 1. 把 query 转向量
|
||
// qv, err := s.embedder.Embed(ctx, []string{opt.Query})
|
||
// if err != nil || len(qv) == 0 {
|
||
// return nil, err
|
||
// }
|
||
// // 2. 拉库内已向量化 chunk
|
||
// chunks, _ := dao.KBListChunksByLibrary(opt.LibraryID)
|
||
// // 3. 算余弦相似度，排序取 TopK
|
||
// // （这一步 V2 用 flat 暴力搜索；如果 chunk 量大需要换 HNSW/faiss）
|
||
// // ...
|
||
// }
|