feat: 实现搜索引擎与评分算法

ai_xiaopei

2026-07-25 beb1530ffd6ef2f505acc2bc8809599fbae30c0a

feat: 实现搜索引擎与评分算法

 internal/search/engine.go

New file
@@ -0,0 +1,106 @@
package search

import (
    "sort"
    "strings"

    "github.com/aisim/kb-cli/internal/index"
)

// SearchOptions 搜索选项
type SearchOptions struct {
    Expanded []string // 扩展词
    Symptom  []string // 症状词
    TopN     int      // 返回前 N 条
}

// SearchResult 搜索结果
type SearchResult struct {
    ID      int64
    Path    string
    Title   string
    Section string
    Score   int
}

// Search 执行搜索
func Search(store *index.Store, keywords []string, opts SearchOptions) ([]SearchResult, error) {
    // 合并所有关键词
    allKeywords := append(keywords, opts.Expanded...)
    allKeywords = append(allKeywords, opts.Symptom...)

    // FTS5 搜索
    ftsResults, err := store.FTSSearch(allKeywords, 100)
    if err != nil {
        return nil, err
    }

    // 评分
    scoreMap := make(map[int64]int)
    for _, r := range ftsResults {
        score := 0

        // 关键词匹配评分
        for _, kw := range keywords {
            score += scoreResult(r, kw, ScoreNormal)
        }
        for _, kw := range opts.Expanded {
            score += scoreResult(r, kw, ScoreExpanded)
        }
        for _, kw := range opts.Symptom {
            score += scoreResult(r, kw, ScoreSymptom)
        }

        scoreMap[r.ID] = score
    }

    // 转换为结果列表
    var results []SearchResult
    for _, r := range ftsResults {
        results = append(results, SearchResult{
            ID:      r.ID,
            Path:    r.Path,
            Title:   r.Title,
            Section: r.Section,
            Score:   scoreMap[r.ID],
        })
    }

    // 按分数排序
    sort.Slice(results, func(i, j int) bool {
        return results[i].Score > results[j].Score
    })

    // 限制返回数量
    if opts.TopN > 0 && len(results) > opts.TopN {
        results = results[:opts.TopN]
    }

    return results, nil
}

// scoreResult 计算单个结果的得分
func scoreResult(r index.FTSResult, keyword string, scoreType ScoreType) int {
    score := 0
    kw := strings.ToLower(keyword)

    // 路径匹配
    if strings.Contains(strings.ToLower(r.Path), kw) {
        score += CalcScore(keyword, "path", scoreType)
    }

    // 标题匹配
    if strings.Contains(strings.ToLower(r.Title), kw) {
        score += CalcScore(keyword, "title", scoreType)
    }

    // 板块匹配
    if strings.Contains(strings.ToLower(r.Section), kw) {
        score += CalcScore(keyword, "tag", scoreType)
    }

    // 内容匹配（FTS 已经匹配，给基础分）
    score += CalcScore(keyword, "content", scoreType)

    return score
}

 internal/search/engine_test.go

New file
@@ -0,0 +1,98 @@
package search

import (
    "path/filepath"
    "strings"
    "testing"

    "github.com/aisim/kb-cli/internal/graph"
    "github.com/aisim/kb-cli/internal/index"
)

func TestSearch(t *testing.T) {
    // 创建临时数据库
    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{
        {
            ID:      1,
            Path:    "FAQ/充装类/001-test.md",
            Title:   "充装问题排查",
            Section: "FAQ",
            Tags:    []string{"充装"},
            Content: "关于智能枪充装问题的排查方法",
        },
        {
            ID:      2,
            Path:    "知识/002-config.md",
            Title:   "充装规格配置",
            Section: "知识",
            Tags:    []string{"配置"},
            Content: "充装规格配置说明",
        },
    }

    for _, n := range nodes {
        _, err := store.InsertNode(n)
        if err != nil {
            t.Fatalf("InsertNode failed: %v", err)
        }
    }

    // 创建 FTS 索引
    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) == 0 {
        t.Error("Expected results, got none")
    }

    // 验证结果包含关键词
    found := false
    for _, r := range results {
        if strings.Contains(r.Title, "充装") || strings.Contains(r.Path, "充装") {
            found = true
            break
        }
    }
    if !found {
        t.Error("Expected results to contain '充装'")
    }
}

func TestIsGenericWord(t *testing.T) {
    tests := []struct {
        word     string
        expected bool
    }{
        {"问题", true},
        {"故障", true},
        {"充装", false},
        {"智能枪", false},
    }

    for _, tt := range tests {
        result := IsGenericWord(tt.word)
        if result != tt.expected {
            t.Errorf("IsGenericWord(%q) = %v, want %v", tt.word, result, tt.expected)
        }
    }
}

 internal/search/scorer.go

New file
@@ -0,0 +1,66 @@
package search

import (
    "strings"
)

// GenericWords 通用词列表（降低权重避免匹配所有文件）
var GenericWords = []string{
    "问题", "故障", "报错", "异常", "无法", "不能",
    "怎么", "如何", "为什么", "是什么", "哪里", "什么",
    "可以", "需要", "应该", "是否", "有没有",
}

// IsGenericWord 判断是否为通用词
func IsGenericWord(word string) bool {
    word = strings.ToLower(word)
    for _, g := range GenericWords {
        if word == g {
            return true
        }
    }
    return false
}

// ScoreType 评分类型
type ScoreType int

const (
    ScoreNormal  ScoreType = iota // 普通词
    ScoreExpanded                 // 扩展词
    ScoreSymptom                  // 症状词
)

// ScoreTable 评分表
var ScoreTable = map[ScoreType]map[string]int{
    ScoreNormal: {
        "path":    2,
        "title":   3,
        "tag":     2,
        "content": 1,
    },
    ScoreExpanded: {
        "path":    40,
        "title":   30,
        "tag":     20,
        "content": 5,
    },
    ScoreSymptom: {
        "path":    60,
        "title":   50,
        "tag":     35,
        "content": 15,
    },
}

// CalcScore 计算单个关键词在某个位置的得分
func CalcScore(keyword string, position string, scoreType ScoreType) int {
    base := ScoreTable[scoreType][position]
    if base == 0 {
        return 0
    }
    if IsGenericWord(keyword) {
        return base / 10 // 通用词降权
    }
    return base
}