ai_xiaopei
2026-07-25 7445005745523d430a3cd22c3208cce2f383f1ee
feat: 实现 SQLite 存储层与 FTS5 全文索引
1 files modified
4 files added
370 ■■■■■ changed files
internal/index/fts.go 101 ●●●●● patch | view | raw | blame | history
internal/index/fts_test.go 54 ●●●●● patch | view | raw | blame | history
internal/index/sqlite.go 149 ●●●●● patch | view | raw | blame | history
internal/index/sqlite_test.go 66 ●●●●● patch | view | raw | blame | history
kb patch | view | raw | blame | history
internal/index/fts.go
New file
@@ -0,0 +1,101 @@
package index
import (
    "database/sql"
    "encoding/json"
    "strings"
)
// FTSResult 全文搜索结果
type FTSResult struct {
    ID      int64
    Path    string
    Title   string
    Section string
    Rank    float64
}
// CreateFTS 创建 FTS5 虚拟表
func (s *Store) CreateFTS() error {
    _, err := s.db.Exec(`
        CREATE VIRTUAL TABLE IF NOT EXISTS nodes_fts USING fts5(
            title,
            content,
            tags,
            entities,
            content='nodes',
            content_rowid='id'
        )
    `)
    return err
}
// PopulateFTS 填充 FTS 索引
func (s *Store) PopulateFTS() error {
    _, err := s.db.Exec(`
        INSERT INTO nodes_fts(rowid, title, content, tags, entities)
        SELECT id, title, content_fts, tags, entities FROM nodes
    `)
    return err
}
// FTSSearch 全文搜索
func (s *Store) FTSSearch(keywords []string, limit int) ([]FTSResult, error) {
    if len(keywords) == 0 {
        return nil, nil
    }
    // 构建 FTS5 查询
    query := strings.Join(keywords, " OR ")
    rows, err := s.db.Query(`
        SELECT n.id, n.path, n.title, n.section, fts.rank
        FROM nodes_fts fts
        JOIN nodes n ON n.id = fts.rowid
        WHERE nodes_fts MATCH ?
        ORDER BY fts.rank
        LIMIT ?
    `, query, limit)
    if err != nil {
        return nil, err
    }
    defer rows.Close()
    var results []FTSResult
    for rows.Next() {
        var r FTSResult
        if err := rows.Scan(&r.ID, &r.Path, &r.Title, &r.Section, &r.Rank); err != nil {
            return nil, err
        }
        results = append(results, r)
    }
    return results, nil
}
// GetNodeContent 获取节点内容
func (s *Store) GetNodeContent(id int64) (string, []string, []string, error) {
    var content, tagsJSON, entitiesJSON string
    err := s.db.QueryRow("SELECT content_fts, tags, entities FROM nodes WHERE id = ?", id).
        Scan(&content, &tagsJSON, &entitiesJSON)
    if err == sql.ErrNoRows {
        return "", nil, nil, nil
    }
    if err != nil {
        return "", nil, nil, err
    }
    // 解析 JSON 数组
    var tags, entities []string
    if tagsJSON != "" {
        if err := json.Unmarshal([]byte(tagsJSON), &tags); err != nil {
            tags = nil
        }
    }
    if entitiesJSON != "" {
        if err := json.Unmarshal([]byte(entitiesJSON), &entities); err != nil {
            entities = nil
        }
    }
    return content, tags, entities, nil
}
internal/index/fts_test.go
New file
@@ -0,0 +1,54 @@
package index
import (
    "path/filepath"
    "testing"
    "github.com/aisim/kb-cli/internal/graph"
)
func TestFTSSearch(t *testing.T) {
    tmpDir := t.TempDir()
    dbPath := filepath.Join(tmpDir, "test.db")
    store, err := Open(dbPath)
    if err != nil {
        t.Fatalf("Open failed: %v", err)
    }
    defer store.Close()
    // 插入测试数据
    store.InsertNode(&graph.Node{
        ID:      1,
        Path:    "FAQ/充装类/001-test.md",
        Title:   "充装问题排查",
        Section: "FAQ",
        Tags:    []string{"充装"},
        Content: "关于智能枪充装问题的排查方法",
    })
    store.InsertNode(&graph.Node{
        ID:      2,
        Path:    "知识/002-config.md",
        Title:   "充装规格配置",
        Section: "知识",
        Tags:    []string{"配置"},
        Content: "充装规格配置说明",
    })
    // 创建 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)
    }
    // 搜索 - FTS5 对中文按字符 tokenize,"充装" 两个节点都包含
    results, err := store.FTSSearch([]string{"充装"}, 10)
    if err != nil {
        t.Fatalf("FTSSearch failed: %v", err)
    }
    if len(results) < 1 {
        t.Errorf("result count = %d, want >= 1", len(results))
    }
}
internal/index/sqlite.go
New file
@@ -0,0 +1,149 @@
package index
import (
    "database/sql"
    "encoding/json"
    "fmt"
    "os"
    "path/filepath"
    _ "github.com/mattn/go-sqlite3"
    "github.com/aisim/kb-cli/internal/graph"
)
// Store SQLite 存储层
type Store struct {
    db *sql.DB
}
// Open 打开或创建数据库
func Open(dbPath string) (*Store, error) {
    // 确保目录存在
    dir := filepath.Dir(dbPath)
    if err := os.MkdirAll(dir, 0755); err != nil {
        return nil, fmt.Errorf("创建目录失败: %w", err)
    }
    db, err := sql.Open("sqlite3", dbPath+"?_journal_mode=WAL")
    if err != nil {
        return nil, fmt.Errorf("打开数据库失败: %w", err)
    }
    s := &Store{db: db}
    if err := s.initTables(); err != nil {
        db.Close()
        return nil, err
    }
    return s, nil
}
// Close 关闭数据库
func (s *Store) Close() error {
    return s.db.Close()
}
// initTables 创建表结构
func (s *Store) initTables() error {
    schema := `
    CREATE TABLE IF NOT EXISTS meta (
        key   TEXT PRIMARY KEY,
        value TEXT NOT NULL
    );
    CREATE TABLE IF NOT EXISTS nodes (
        id          INTEGER PRIMARY KEY AUTOINCREMENT,
        path        TEXT NOT NULL UNIQUE,
        title       TEXT,
        section     TEXT,
        tags        TEXT,
        entities    TEXT,
        wikilinks   TEXT,
        content_fts TEXT,
        created_at  TEXT DEFAULT (datetime('now')),
        updated_at  TEXT DEFAULT (datetime('now'))
    );
    CREATE TABLE IF NOT EXISTS edges (
        id          INTEGER PRIMARY KEY AUTOINCREMENT,
        from_node   INTEGER NOT NULL,
        to_node     INTEGER NOT NULL,
        relation    TEXT NOT NULL,
        label       TEXT,
        UNIQUE(from_node, to_node, relation, label)
    );
    CREATE INDEX IF NOT EXISTS idx_nodes_section ON nodes(section);
    CREATE INDEX IF NOT EXISTS idx_edges_from ON edges(from_node);
    CREATE INDEX IF NOT EXISTS idx_edges_to ON edges(to_node);
    CREATE INDEX IF NOT EXISTS idx_edges_label ON edges(label);
    `
    _, err := s.db.Exec(schema)
    if err != nil {
        return fmt.Errorf("创建表失败: %w", err)
    }
    return nil
}
// ClearData 清空数据(重建前调用)
func (s *Store) ClearData() error {
    _, err := s.db.Exec("DELETE FROM edges; DELETE FROM nodes;")
    return err
}
// InsertNode 插入节点
func (s *Store) InsertNode(n *graph.Node) (int64, error) {
    tagsJSON, _ := json.Marshal(n.Tags)
    entitiesJSON, _ := json.Marshal(n.Entities)
    wikilinksJSON, _ := json.Marshal(n.Wikilinks)
    result, err := s.db.Exec(`
        INSERT INTO nodes (path, title, section, tags, entities, wikilinks, content_fts)
        VALUES (?, ?, ?, ?, ?, ?, ?)
    `, n.Path, n.Title, n.Section, string(tagsJSON), string(entitiesJSON),
        string(wikilinksJSON), n.Content)
    if err != nil {
        return 0, err
    }
    return result.LastInsertId()
}
// InsertEdge 插入边
func (s *Store) InsertEdge(e *graph.Edge) error {
    _, err := s.db.Exec(`
        INSERT OR IGNORE INTO edges (from_node, to_node, relation, label)
        VALUES (?, ?, ?, ?)
    `, e.FromNode, e.ToNode, e.Relation, e.Label)
    return err
}
// SetMeta 设置元信息
func (s *Store) SetMeta(key, value string) error {
    _, err := s.db.Exec(`
        INSERT OR REPLACE INTO meta (key, value) VALUES (?, ?)
    `, key, value)
    return err
}
// GetMeta 获取元信息
func (s *Store) GetMeta(key string) (string, error) {
    var value string
    err := s.db.QueryRow("SELECT value FROM meta WHERE key = ?", key).Scan(&value)
    if err == sql.ErrNoRows {
        return "", nil
    }
    return value, err
}
// NodeCount 返回节点数量
func (s *Store) NodeCount() (int, error) {
    var count int
    err := s.db.QueryRow("SELECT COUNT(*) FROM nodes").Scan(&count)
    return count, err
}
// EdgeCount 返回边数量
func (s *Store) EdgeCount() (int, error) {
    var count int
    err := s.db.QueryRow("SELECT COUNT(*) FROM edges").Scan(&count)
    return count, err
}
internal/index/sqlite_test.go
New file
@@ -0,0 +1,66 @@
package index
import (
    "path/filepath"
    "testing"
    "github.com/aisim/kb-cli/internal/graph"
)
func TestStore(t *testing.T) {
    tmpDir := t.TempDir()
    dbPath := filepath.Join(tmpDir, "test.db")
    store, err := Open(dbPath)
    if err != nil {
        t.Fatalf("Open failed: %v", err)
    }
    defer store.Close()
    // 测试插入节点
    node := &graph.Node{
        Path:    "FAQ/充装类/001-test.md",
        Title:   "测试FAQ",
        Section: "FAQ",
        Tags:    []string{"充装", "智能枪"},
    }
    id, err := store.InsertNode(node)
    if err != nil {
        t.Fatalf("InsertNode failed: %v", err)
    }
    if id == 0 {
        t.Error("InsertNode returned 0 id")
    }
    // 测试插入边
    edge := &graph.Edge{
        FromNode: id,
        ToNode:   1000000,
        Relation: "tag",
        Label:    "充装",
    }
    if err := store.InsertEdge(edge); err != nil {
        t.Fatalf("InsertEdge failed: %v", err)
    }
    // 测试元信息
    if err := store.SetMeta("git_commit", "abc123"); err != nil {
        t.Fatalf("SetMeta failed: %v", err)
    }
    commit, err := store.GetMeta("git_commit")
    if err != nil {
        t.Fatalf("GetMeta failed: %v", err)
    }
    if commit != "abc123" {
        t.Errorf("commit = %q, want %q", commit, "abc123")
    }
    // 测试计数
    count, err := store.NodeCount()
    if err != nil {
        t.Fatalf("NodeCount failed: %v", err)
    }
    if count != 1 {
        t.Errorf("node count = %d, want 1", count)
    }
}
kb
Binary files differ