feat: RWR 图结构排序 + CJK LIKE 混合检索 + 双信号加权 + aliases/status 参与排序
3 files added
5 files modified
| | |
| | | |
| | | "github.com/aisim/kb-cli/internal/index" |
| | | "github.com/aisim/kb-cli/internal/search" |
| | | "github.com/aisim/kb-cli/internal/llm" |
| | | "github.com/aisim/kb-cli/internal/output" |
| | | "github.com/spf13/cobra" |
| | | ) |
| | |
| | | } |
| | | |
| | | // 执行搜索 |
| | | cfg := llm.LoadConfig() |
| | | textWeight := cfg.Search.TextWeight |
| | | if textWeight <= 0 || textWeight > 1 { |
| | | textWeight = 0.5 // 缺省 0.5 |
| | | } |
| | | opts := search.SearchOptions{ |
| | | Expanded: expanded, |
| | | Symptom: symptom, |
| | | TopN: topN, |
| | | WithContent: withContent, |
| | | WithLinks: withLinks, |
| | | TextWeight: textWeight, |
| | | } |
| | | |
| | | results, err := search.Search(store, allKeywords, opts) |
| New file |
| | |
| | | package graph |
| | | |
| | | // AdjEdge 带权邻接边(导出:index 包 LoadRWRGraph 需引用该类型) |
| | | type AdjEdge struct { |
| | | To int64 |
| | | Weight float64 |
| | | } |
| | | |
| | | // RWR Random-Walk-with-Restart(个性化 PageRank): |
| | | // 从种子集合出发在无向带权图上做 power iteration,restart 概率 alpha 回到种子。 |
| | | // 返回归一化到 [0,1] 的游走质量。确定性算法,无随机数。 |
| | | func RWR(seedIDs []int64, adj map[int64][]AdjEdge, alpha float64) map[int64]float64 { |
| | | if len(seedIDs) == 0 { |
| | | return map[int64]float64{} |
| | | } |
| | | // 收集参与节点:种子 + 邻接可达 |
| | | nodes := make(map[int64]bool) |
| | | for _, s := range seedIDs { |
| | | nodes[s] = true |
| | | } |
| | | for from := range adj { |
| | | nodes[from] = true |
| | | for _, e := range adj[from] { |
| | | nodes[e.To] = true |
| | | } |
| | | } |
| | | n := len(nodes) |
| | | if n == 0 { |
| | | return map[int64]float64{} |
| | | } |
| | | idx := make(map[int64]int, n) |
| | | ids := make([]int64, 0, n) |
| | | for id := range nodes { |
| | | idx[id] = len(ids) |
| | | ids = append(ids, id) |
| | | } |
| | | |
| | | // 重启向量:种子均匀 |
| | | r := make([]float64, n) |
| | | for _, s := range seedIDs { |
| | | if i, ok := idx[s]; ok { |
| | | r[i] = 1.0 / float64(len(seedIDs)) |
| | | } |
| | | } |
| | | |
| | | const maxIter = 50 |
| | | const eps = 1e-6 |
| | | newR := make([]float64, n) |
| | | for iter := 0; iter < maxIter; iter++ { |
| | | for i := range newR { |
| | | newR[i] = 0 |
| | | } |
| | | // 游走传播:out[i] = sum(w_ij) |
| | | for i, id := range ids { |
| | | var total float64 |
| | | for _, e := range adj[id] { |
| | | total += e.Weight |
| | | } |
| | | if total == 0 { |
| | | continue |
| | | } |
| | | for _, e := range adj[id] { |
| | | j, ok := idx[e.To] |
| | | if !ok { |
| | | continue |
| | | } |
| | | newR[j] += (1 - alpha) * r[i] * e.Weight / total |
| | | } |
| | | } |
| | | // restart |
| | | for i := range newR { |
| | | newR[i] += alpha * r[i] |
| | | } |
| | | // 收敛判断 |
| | | diff := 0.0 |
| | | for i := range newR { |
| | | d := newR[i] - r[i] |
| | | if d < 0 { |
| | | d = -d |
| | | } |
| | | diff += d |
| | | } |
| | | r, newR = newR, r |
| | | if diff < eps { |
| | | break |
| | | } |
| | | } |
| | | // 归一化 [0,1] |
| | | maxV := 0.0 |
| | | for _, v := range r { |
| | | if v > maxV { |
| | | maxV = v |
| | | } |
| | | } |
| | | out := make(map[int64]float64, n) |
| | | if maxV == 0 { |
| | | for _, id := range ids { |
| | | out[id] = 0 |
| | | } |
| | | return out |
| | | } |
| | | for i, id := range ids { |
| | | out[id] = r[i] / maxV |
| | | } |
| | | return out |
| | | } |
| New file |
| | |
| | | package graph |
| | | |
| | | import "testing" |
| | | |
| | | func TestRWRConvergenceAndSeed(t *testing.T) { |
| | | // 图: A-B-C 链 + D 孤立 |
| | | adj := map[int64][]AdjEdge{ |
| | | 1: {{2, 1}}, 2: {{1, 1}, {3, 1}}, 3: {{2, 1}}, 4: {{}}, |
| | | } |
| | | r := RWR([]int64{1}, adj, 0.25) |
| | | if r[1] <= 0 { |
| | | t.Fatal("种子节点质量必须 > 0") |
| | | } |
| | | // 与种子连通的质量应高于孤立节点 |
| | | if r[4] >= r[3] { |
| | | t.Errorf("孤立节点质量不应高于连通节点: r4=%f r3=%f", r[4], r[3]) |
| | | } |
| | | // 归一化 [0,1] |
| | | for _, v := range r { |
| | | if v < 0 || v > 1 { |
| | | t.Errorf("质量未归一化: %f", v) |
| | | } |
| | | } |
| | | } |
| | | |
| | | func TestRWRWeightedByProvenance(t *testing.T) { |
| | | // A -exact-> B, A -fuzzy-> C:B 的质量应高于 C |
| | | adj := map[int64][]AdjEdge{ |
| | | 1: {{2, 1.0}, {3, 0.5}}, |
| | | 2: {{1, 1.0}}, |
| | | 3: {{1, 0.5}}, |
| | | } |
| | | r := RWR([]int64{1}, adj, 0.25) |
| | | if r[2] <= r[3] { |
| | | t.Errorf("exact 边节点质量应高于 fuzzy: r2=%f r3=%f", r[2], r[3]) |
| | | } |
| | | } |
| | |
| | | Title string |
| | | Section string |
| | | Rank float64 |
| | | Aliases []string // 别名(JSON 解析) |
| | | Status string // 状态(草稿/待确认/跟进中 等参与降权) |
| | | } |
| | | |
| | | // CreateFTS 创建 FTS5 虚拟表(external-content 模式,由触发器增量维护) |
| | |
| | | |
| | | for _, kw := range keywords { |
| | | rows, err := s.db.Query(` |
| | | SELECT n.id, n.path, n.title, n.section, fts.rank |
| | | SELECT n.id, n.path, n.title, n.section, fts.rank, n.aliases, n.status |
| | | FROM nodes_fts fts |
| | | JOIN nodes n ON n.id = fts.rowid |
| | | WHERE nodes_fts MATCH ? |
| | |
| | | |
| | | for rows.Next() { |
| | | var r FTSResult |
| | | if err := rows.Scan(&r.ID, &r.Path, &r.Title, &r.Section, &r.Rank); err != nil { |
| | | var aliasesJSON string |
| | | if err := rows.Scan(&r.ID, &r.Path, &r.Title, &r.Section, &r.Rank, &aliasesJSON, &r.Status); err != nil { |
| | | rows.Close() |
| | | continue |
| | | } |
| | | if aliasesJSON != "" { |
| | | _ = json.Unmarshal([]byte(aliasesJSON), &r.Aliases) |
| | | } |
| | | // 去重 |
| | | if !seen[r.ID] { |
| | | seen[r.ID] = true |
| | |
| | | return allResults, nil |
| | | } |
| | | |
| | | // containsCJK 是否含 CJK 统一表意文字 |
| | | func containsCJK(s string) bool { |
| | | for _, r := range s { |
| | | if r >= 0x4E00 && r <= 0x9FFF { |
| | | return true |
| | | } |
| | | } |
| | | return false |
| | | } |
| | | |
| | | // KeywordSearch 双通道关键词检索:ASCII 词走 FTS5 MATCH,CJK 词走 LIKE(title/aliases/content/tags)。 |
| | | // FTS5 unicode61 把连续中文当整串单 token,多字符 CJK 词 MATCH 匹配不到,必须走 LIKE。 |
| | | // 结果按"首命中顺序"排列(FTS 通道按 rank,LIKE 通道补在尾部),engine 再按双信号重排。 |
| | | func (s *Store) KeywordSearch(keywords []string, limit int) ([]FTSResult, error) { |
| | | if len(keywords) == 0 { |
| | | return nil, nil |
| | | } |
| | | seen := make(map[int64]*FTSResult) |
| | | var order []int64 |
| | | |
| | | // 通道 1: FTS(ASCII 词) |
| | | var ascii []string |
| | | for _, kw := range keywords { |
| | | if !containsCJK(kw) { |
| | | ascii = append(ascii, kw) |
| | | } |
| | | } |
| | | for _, kw := range ascii { |
| | | rows, err := s.db.Query(` |
| | | SELECT n.id, n.path, n.title, n.section, fts.rank, n.aliases, n.status |
| | | FROM nodes_fts fts JOIN nodes n ON n.id = fts.rowid |
| | | WHERE nodes_fts MATCH ? ORDER BY fts.rank LIMIT ?`, kw, limit) |
| | | if err != nil { |
| | | continue // 语法异常词跳过 |
| | | } |
| | | for rows.Next() { |
| | | var r FTSResult |
| | | var aliasesJSON string |
| | | if err := rows.Scan(&r.ID, &r.Path, &r.Title, &r.Section, &r.Rank, &aliasesJSON, &r.Status); err != nil { |
| | | continue |
| | | } |
| | | if aliasesJSON != "" { |
| | | _ = json.Unmarshal([]byte(aliasesJSON), &r.Aliases) |
| | | } |
| | | if _, ok := seen[r.ID]; !ok { |
| | | seen[r.ID] = &r |
| | | order = append(order, r.ID) |
| | | } |
| | | } |
| | | rows.Close() |
| | | } |
| | | |
| | | // 通道 2: LIKE(CJK 词) |
| | | for _, kw := range keywords { |
| | | if !containsCJK(kw) { |
| | | continue |
| | | } |
| | | pat := "%" + kw + "%" |
| | | rows, err := s.db.Query(` |
| | | SELECT id, path, title, section, aliases, status FROM nodes |
| | | WHERE title LIKE ? OR aliases LIKE ? OR content_fts LIKE ? OR tags LIKE ?`, |
| | | pat, pat, pat, pat) |
| | | if err != nil { |
| | | continue |
| | | } |
| | | for rows.Next() { |
| | | var r FTSResult |
| | | var aliasesJSON string |
| | | if err := rows.Scan(&r.ID, &r.Path, &r.Title, &r.Section, &aliasesJSON, &r.Status); err != nil { |
| | | continue |
| | | } |
| | | if aliasesJSON != "" { |
| | | _ = json.Unmarshal([]byte(aliasesJSON), &r.Aliases) |
| | | } |
| | | r.Rank = 0 // LIKE 无 rank |
| | | if _, ok := seen[r.ID]; !ok { |
| | | seen[r.ID] = &r |
| | | order = append(order, r.ID) |
| | | } |
| | | } |
| | | rows.Close() |
| | | } |
| | | |
| | | results := make([]FTSResult, 0, len(order)) |
| | | for _, id := range order { |
| | | results = append(results, *seen[id]) |
| | | } |
| | | if limit > 0 && len(results) > limit { |
| | | results = results[:limit] |
| | | } |
| | | return results, nil |
| | | } |
| | | |
| | | // GetNodeContent 获取节点内容 |
| | | func (s *Store) GetNodeContent(id int64) (string, []string, []string, error) { |
| | | var content, tagsJSON, entitiesJSON string |
| New file |
| | |
| | | package index |
| | | |
| | | import ( |
| | | "fmt" |
| | | |
| | | "github.com/aisim/kb-cli/internal/graph" |
| | | ) |
| | | |
| | | // LoadRWRGraph 加载 RWR 邻接(无向、带权)。 |
| | | // 权重:wikilink exact=1.0 / fuzzy=0.5 / entity=1.0 / tag=0.5(tag 扇出大降权)。 |
| | | // 只加载与种子同连通域的边不可行(SQLite 无图查询),全量加载后 RWR 内部按种子收敛—— |
| | | // 458 节点 / ~2k 边规模下全量加载 <5ms,可接受。 |
| | | func (s *Store) LoadRWRGraph() (map[int64][]graph.AdjEdge, error) { |
| | | rows, err := s.db.Query(` |
| | | SELECT from_node, to_node, relation, provenance FROM edges |
| | | WHERE relation IN ('wikilink', 'entity', 'tag')`) |
| | | if err != nil { |
| | | return nil, fmt.Errorf("加载边失败: %w", err) |
| | | } |
| | | defer rows.Close() |
| | | adj := make(map[int64][]graph.AdjEdge) |
| | | add := func(from, to int64, w float64) { |
| | | adj[from] = append(adj[from], graph.AdjEdge{To: to, Weight: w}) |
| | | adj[to] = append(adj[to], graph.AdjEdge{To: from, Weight: w}) |
| | | } |
| | | for rows.Next() { |
| | | var from, to int64 |
| | | var relation, prov string |
| | | if err := rows.Scan(&from, &to, &relation, &prov); err != nil { |
| | | return nil, err |
| | | } |
| | | var w float64 |
| | | switch relation { |
| | | case "wikilink": |
| | | if prov == "exact" { |
| | | w = 1.0 |
| | | } else { |
| | | w = 0.5 |
| | | } |
| | | case "entity": |
| | | w = 1.0 |
| | | case "tag": |
| | | w = 0.5 |
| | | default: |
| | | continue |
| | | } |
| | | add(from, to, w) |
| | | } |
| | | return adj, rows.Err() |
| | | } |
| | |
| | | ReviewDir string `yaml:"review_dir"` |
| | | DefaultType string `yaml:"default_type"` |
| | | } `yaml:"draft"` |
| | | // Search 搜索配置(双信号加权权重) |
| | | Search struct { |
| | | TextWeight float64 `yaml:"text_weight"` // 文本分权重(0~1,缺省 0.5) |
| | | } `yaml:"search"` |
| | | // Explore 探索预算配置(Task 7 预留) |
| | | Explore struct { |
| | | DefaultBudget int `yaml:"default_budget"` |
| | | HardBudget int `yaml:"hard_budget"` |
| | | TopN int `yaml:"top_n"` |
| | | } `yaml:"explore"` |
| | | } |
| | | |
| | | // Client LLM 客户端 |
| | |
| | | "sort" |
| | | "strings" |
| | | |
| | | "github.com/aisim/kb-cli/internal/graph" |
| | | "github.com/aisim/kb-cli/internal/index" |
| | | ) |
| | | |
| | |
| | | TopN int // 返回前 N 条 |
| | | WithContent bool // 返回完整文件内容 |
| | | WithLinks bool // 返回关联文档链接 |
| | | TextWeight float64 // 文本分权重(0~1,缺省 0.5;RWR 图质量权重 = 1 - TextWeight) |
| | | } |
| | | |
| | | // SearchResult 搜索结果 |
| | |
| | | Title string `json:"title"` |
| | | Section string `json:"section"` |
| | | Score int `json:"score"` |
| | | Status string `json:"status,omitempty"` |
| | | Content string `json:"content,omitempty"` // 文件内容(WithContent=true 时填充) |
| | | Links []string `json:"links,omitempty"` // 关联文档路径(WithLinks=true 时填充) |
| | | } |
| | | |
| | | // Search 执行搜索 |
| | | // Search 执行搜索(双信号加权:文本位置分 × TextWeight + RWR 图质量 × (1-TextWeight), |
| | | // 草稿态/待审阅降权 0.6) |
| | | func Search(store *index.Store, keywords []string, opts SearchOptions) ([]SearchResult, error) { |
| | | // 合并所有关键词 |
| | | allKeywords := append(keywords, opts.Expanded...) |
| | | allKeywords = append(allKeywords, opts.Symptom...) |
| | | allKeywords := append(append(keywords, opts.Expanded...), opts.Symptom...) |
| | | |
| | | // FTS5 搜索 |
| | | ftsResults, err := store.FTSSearch(allKeywords, 100) |
| | | // 1. 双通道检索(ASCII 走 FTS,CJK 走 LIKE) |
| | | candidates, err := store.KeywordSearch(allKeywords, 100) |
| | | if err != nil { |
| | | return nil, err |
| | | } |
| | | if len(candidates) == 0 { |
| | | return nil, nil |
| | | } |
| | | |
| | | // 评分 |
| | | scoreMap := make(map[int64]int) |
| | | for _, r := range ftsResults { |
| | | score := 0 |
| | | |
| | | // 关键词匹配评分 |
| | | // 2. 文本分(沿用位置加权,aliases 按 title 档计权) |
| | | textScore := make(map[int64]int, len(candidates)) |
| | | for _, r := range candidates { |
| | | var score int |
| | | for _, kw := range keywords { |
| | | score += scoreResult(r, kw, ScoreNormal) |
| | | } |
| | |
| | | for _, kw := range opts.Symptom { |
| | | score += scoreResult(r, kw, ScoreSymptom) |
| | | } |
| | | |
| | | scoreMap[r.ID] = score |
| | | textScore[r.ID] = score |
| | | } |
| | | |
| | | // 转换为结果列表 |
| | | // 3. RWR 图质量(种子 = 候选前 20,全量加载邻接后按种子收敛) |
| | | seedIDs := make([]int64, 0, len(candidates)) |
| | | for _, r := range candidates { |
| | | if len(seedIDs) >= 20 { |
| | | break |
| | | } |
| | | seedIDs = append(seedIDs, r.ID) |
| | | } |
| | | rwrMass := map[int64]float64{} |
| | | if len(seedIDs) > 0 { |
| | | adj, err := store.LoadRWRGraph() |
| | | if err == nil { |
| | | rwrMass = graph.RWR(seedIDs, adj, 0.25) |
| | | } |
| | | } |
| | | |
| | | // 4. 双信号加权:finalScore = norm(textScore)*tw + rwrMass*(1-tw) |
| | | tw := opts.TextWeight |
| | | if tw <= 0 || tw > 1 { |
| | | tw = 0.5 |
| | | } |
| | | // 先求文本分真实 min/max,再按 (s-minT)/(maxT-minT) 归一化 |
| | | minT, maxT := 0, 0 |
| | | for _, s := range textScore { |
| | | if minT == 0 && maxT == 0 { |
| | | minT, maxT = s, s |
| | | } else { |
| | | if s < minT { |
| | | minT = s |
| | | } |
| | | if s > maxT { |
| | | maxT = s |
| | | } |
| | | } |
| | | } |
| | | // status 降权:草稿态或待审阅板块 ×0.6 |
| | | statusFactor := make(map[int64]float64, len(candidates)) |
| | | for _, r := range candidates { |
| | | if r.Section == "待审阅" || isDraftStatus(r.Status) { |
| | | statusFactor[r.ID] = 0.6 |
| | | } else { |
| | | statusFactor[r.ID] = 1.0 |
| | | } |
| | | } |
| | | |
| | | var results []SearchResult |
| | | for _, r := range ftsResults { |
| | | for _, r := range candidates { |
| | | ts := float64(textScore[r.ID]) |
| | | if maxT > minT { |
| | | ts = (ts - float64(minT)) / float64(maxT - minT) |
| | | } else if maxT > 0 { |
| | | ts = 1 |
| | | } |
| | | final := ts*tw + rwrMass[r.ID]*(1-tw) |
| | | final *= statusFactor[r.ID] |
| | | results = append(results, SearchResult{ |
| | | ID: r.ID, |
| | | Path: r.Path, |
| | | Title: r.Title, |
| | | Section: r.Section, |
| | | Score: scoreMap[r.ID], |
| | | Score: int(final * 100), |
| | | Status: r.Status, |
| | | }) |
| | | } |
| | | |
| | |
| | | return results, nil |
| | | } |
| | | |
| | | // isDraftStatus 草稿态降权判断(草稿/待确认/跟进中 降权 0.6) |
| | | func isDraftStatus(status string) bool { |
| | | switch status { |
| | | case "草稿", "待确认", "跟进中": |
| | | return true |
| | | } |
| | | return false |
| | | } |
| | | |
| | | // scoreResult 计算单个结果的得分 |
| | | func scoreResult(r index.FTSResult, keyword string, scoreType ScoreType) int { |
| | | score := 0 |
| | |
| | | score += CalcScore(keyword, "tag", scoreType) |
| | | } |
| | | |
| | | // 别名匹配(命中按 title 档计权,一次命中即计) |
| | | if len(r.Aliases) > 0 { |
| | | for _, a := range r.Aliases { |
| | | if strings.Contains(strings.ToLower(a), kw) { |
| | | score += CalcScore(keyword, "title", scoreType) |
| | | break |
| | | } |
| | | } |
| | | } |
| | | |
| | | // 内容匹配(FTS 已经匹配,给基础分) |
| | | score += CalcScore(keyword, "content", scoreType) |
| | | |
| | |
| | | } |
| | | } |
| | | |
| | | func TestCJKLikeChannel(t *testing.T) { |
| | | // CJK LIKE 通道:多字符中文词 FTS MATCH 匹配不到(unicode61 整串 token), |
| | | // 必须走 LIKE 才能命中 content/title |
| | | tmpDir := t.TempDir() |
| | | dbPath := filepath.Join(tmpDir, "test.db") |
| | | store, err := index.Open(dbPath) |
| | | if err != nil { |
| | | t.Fatalf("Open failed: %v", err) |
| | | } |
| | | defer store.Close() |
| | | |
| | | nodes := []*graph.Node{ |
| | | { // content 含"补气",title 不含(靠 LIKE content 通道命中) |
| | | Path: "FAQ/001-a.md", |
| | | Title: "设备故障排查", |
| | | Section: "FAQ", |
| | | Content: "电子秤补气失败时的排查步骤", |
| | | }, |
| | | { // title 含"补气"(title LIKE 命中,应排前) |
| | | Path: "FAQ/002-b.md", |
| | | Title: "电子秤补气失败", |
| | | Section: "FAQ", |
| | | Content: "补气失败的处理方法", |
| | | }, |
| | | { // 不含"补气"(不应出现在结果中) |
| | | Path: "FAQ/003-c.md", |
| | | Title: "阀门漏气处理", |
| | | Section: "FAQ", |
| | | Content: "阀门漏气的原因和处理", |
| | | }, |
| | | } |
| | | for _, n := range nodes { |
| | | if err := store.UpsertNode(n, 100, 1, "x"); err != nil { |
| | | t.Fatalf("UpsertNode failed: %v", err) |
| | | } |
| | | } |
| | | if err := store.CreateFTS(); err != nil { |
| | | t.Fatalf("CreateFTS failed: %v", err) |
| | | } |
| | | if err := store.PopulateFTS(); err != nil { |
| | | t.Fatalf("PopulateFTS failed: %v", err) |
| | | } |
| | | |
| | | results, err := Search(store, []string{"补气"}, SearchOptions{}) |
| | | if err != nil { |
| | | t.Fatalf("Search failed: %v", err) |
| | | } |
| | | if len(results) < 2 { |
| | | t.Fatalf("Expected 2 results (content 命中 + title 命中), got %d", len(results)) |
| | | } |
| | | // 两个含"补气"的节点都应在结果中 |
| | | if results[0].Path != "FAQ/002-b.md" && results[1].Path != "FAQ/002-b.md" { |
| | | t.Error("title 含'补气'的节点应出现在结果中") |
| | | } |
| | | if results[0].Path != "FAQ/001-a.md" && results[1].Path != "FAQ/001-a.md" { |
| | | t.Error("content 含'补气'的节点应出现在结果中") |
| | | } |
| | | // 不含"补气"的节点不应出现 |
| | | for _, r := range results { |
| | | if r.Path == "FAQ/003-c.md" { |
| | | t.Error("不含关键词的节点不应出现在结果中") |
| | | } |
| | | } |
| | | // title 命中者应排前(title 计权高于 content) |
| | | if results[0].Path != "FAQ/002-b.md" { |
| | | t.Errorf("title 命中者应排第一, got %q", results[0].Path) |
| | | } |
| | | } |
| | | |
| | | func TestIsGenericWord(t *testing.T) { |
| | | tests := []struct { |
| | | word string |