feat: add graph command for knowledge graph queries

ai_xiaopei

9 days ago 6b58a8fd1d10191d5720b037af1137d2a747cfeb

feat: add graph command for knowledge graph queries

- graph stats: show node/edge counts and relation distribution
- graph query: find nodes by keyword and show their edges
- graph related: find related nodes through shared tags/entities
- support --json output for all graph commands
- add GetRelationStats, FindNodesByKeyword, GetNodeEdges, FindRelatedNodes methods to Store

 cmd/graph.go

New file
@@ -0,0 +1,240 @@
package cmd

import (
    "encoding/json"
    "fmt"

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

var graphCmd = &cobra.Command{
    Use:   "graph",
    Short: "知识图谱查询与管理",
}

var graphStatsCmd = &cobra.Command{
    Use:   "stats",
    Short: "查看知识图谱统计信息",
    Long:  `kb-cli graph stats [--json] # 查看知识图谱统计信息（节点数、边数、关系类型分布）`,
    RunE:  runGraphStats,
}

var graphQueryCmd = &cobra.Command{
    Use:   "query <关键词>",
    Short: "查询节点的关联关系",
    Long:  `kb-cli graph query <关键词> [--relation=<类型>] [--depth=<深度>] [--json] # 查询节点的关联关系`,
    Args:  cobra.ExactArgs(1),
    RunE:  runGraphQuery,
}

var graphRelatedCmd = &cobra.Command{
    Use:   "related <关键词>",
    Short: "查找与关键词相关的所有节点",
    Long:  `kb-cli graph related <关键词> [--top N] [--json] # 查找与关键词相关的所有节点`,
    Args:  cobra.ExactArgs(1),
    RunE:  runGraphRelated,
}

var (
    graphRelation string
    graphDepth    int
    graphJSON     bool
    graphTopN     int
)

func init() {
    rootCmd.AddCommand(graphCmd)
    graphCmd.AddCommand(graphStatsCmd)
    graphCmd.AddCommand(graphQueryCmd)
    graphCmd.AddCommand(graphRelatedCmd)

    // graph stats flags
    graphStatsCmd.Flags().BoolVar(&graphJSON, "json", false, "JSON 格式输出")

    // graph query flags
    graphQueryCmd.Flags().StringVar(&graphRelation, "relation", "", "关系类型（tag/entity/wikilink）")
    graphQueryCmd.Flags().IntVar(&graphDepth, "depth", 1, "查询深度（1-3）")
    graphQueryCmd.Flags().BoolVar(&graphJSON, "json", false, "JSON 格式输出")

    // graph related flags
    graphRelatedCmd.Flags().IntVar(&graphTopN, "top", 10, "返回前 N 条结果")
    graphRelatedCmd.Flags().BoolVar(&graphJSON, "json", false, "JSON 格式输出")
}

func runGraphStats(cmd *cobra.Command, args []string) error {
    // 打开索引
    store, err := index.Open(dbPath)
    if err != nil {
        return fmt.Errorf("打开索引失败: %w", err)
    }
    defer store.Close()

    // 获取节点数和边数
    nodeCount, err := store.NodeCount()
    if err != nil {
        return fmt.Errorf("获取节点数失败: %w", err)
    }

    edgeCount, err := store.EdgeCount()
    if err != nil {
        return fmt.Errorf("获取边数失败: %w", err)
    }

    // 获取关系类型分布
    relationStats, err := store.GetRelationStats()
    if err != nil {
        return fmt.Errorf("获取关系统计失败: %w", err)
    }

    if graphJSON {
        stats := map[string]interface{}{
            "node_count": nodeCount,
            "edge_count": edgeCount,
            "relations":  relationStats,
        }
        data, err := json.MarshalIndent(stats, "", "  ")
        if err != nil {
            return err
        }
        fmt.Println(string(data))
    } else {
        fmt.Printf("知识图谱统计:\n")
        fmt.Printf("  节点数: %d\n", nodeCount)
        fmt.Printf("  边数: %d\n", edgeCount)
        fmt.Printf("\n关系类型分布:\n")
        for rel, count := range relationStats {
            fmt.Printf("  %-12s %d\n", rel, count)
        }
    }

    return nil
}

func runGraphQuery(cmd *cobra.Command, args []string) error {
    keyword := args[0]

    // 打开索引
    store, err := index.Open(dbPath)
    if err != nil {
        return fmt.Errorf("打开索引失败: %w", err)
    }
    defer store.Close()

    // 查找匹配关键词的节点
    nodes, err := store.FindNodesByKeyword(keyword)
    if err != nil {
        return fmt.Errorf("查找节点失败: %w", err)
    }

    if len(nodes) == 0 {
        fmt.Println("未找到匹配的节点")
        return nil
    }

    // 查询每个节点的关联关系
    var results []QueryResult
    for _, node := range nodes {
        edges, err := store.GetNodeEdges(node.ID, graphRelation, graphDepth)
        if err != nil {
            return fmt.Errorf("查询关联失败: %w", err)
        }

        results = append(results, QueryResult{
            Node:  node,
            Edges: edges,
        })
    }

    if graphJSON {
        data, err := json.MarshalIndent(results, "", "  ")
        if err != nil {
            return err
        }
        fmt.Println(string(data))
    } else {
        for _, r := range results {
            fmt.Printf("\n节点: %s (%s)\n", r.Node.Title, r.Node.Path)
            if len(r.Edges) == 0 {
                fmt.Println("  无关联关系")
                continue
            }

            // 按关系类型分组
            grouped := make(map[string][]string)
            for _, edge := range r.Edges {
                grouped[edge.Relation] = append(grouped[edge.Relation], edge.Label)
            }

            for rel, labels := range grouped {
                fmt.Printf("  [%s]\n", rel)
                for _, label := range labels {
                    fmt.Printf("    - %s\n", label)
                }
            }
        }
    }

    return nil
}

func runGraphRelated(cmd *cobra.Command, args []string) error {
    keyword := args[0]

    // 打开索引
    store, err := index.Open(dbPath)
    if err != nil {
        return fmt.Errorf("打开索引失败: %w", err)
    }
    defer store.Close()

    // 查找相关节点
    related, err := store.FindRelatedNodes(keyword, graphTopN)
    if err != nil {
        return fmt.Errorf("查找相关节点失败: %w", err)
    }

    if len(related) == 0 {
        fmt.Println("未找到相关节点")
        return nil
    }

    if graphJSON {
        data, err := json.MarshalIndent(related, "", "  ")
        if err != nil {
            return err
        }
        fmt.Println(string(data))
    } else {
        fmt.Printf("与 '%s' 相关的节点:\n\n", keyword)
        for i, r := range related {
            fmt.Printf("%d. %s\n", i+1, r.Title)
            fmt.Printf("   路径: %s\n", r.Path)
            fmt.Printf("   板块: %s\n", r.Section)
            fmt.Printf("   关联度: %d\n", r.Relevance)
            if len(r.Tags) > 0 {
                fmt.Printf("   标签: %v\n", r.Tags)
            }
            fmt.Println()
        }
    }

    return nil
}

// QueryResult 查询结果
type QueryResult struct {
    Node  *graph.Node  `json:"node"`
    Edges []*graph.Edge `json:"edges"`
}

// RelatedNode 相关节点
type RelatedNode struct {
    ID        int64    `json:"id"`
    Path      string   `json:"path"`
    Title     string   `json:"title"`
    Section   string   `json:"section"`
    Tags      []string `json:"tags"`
    Relevance int      `json:"relevance"`
}

 internal/index/sqlite.go

@@ -6,6 +6,7 @@
    "fmt"
    "os"
    "path/filepath"
    "sort"
    "strings"

    _ "github.com/mattn/go-sqlite3"
@@ -258,3 +259,283 @@
    affected, _ := result.RowsAffected()
    return int(affected), nil
}

// GetRelationStats 获取关系类型分布统计
func (s *Store) GetRelationStats() (map[string]int, error) {
    rows, err := s.db.Query(`
        SELECT relation, COUNT(*) as count
        FROM edges
        GROUP BY relation
        ORDER BY count DESC
    `)
    if err != nil {
        return nil, fmt.Errorf("查询关系统计失败: %w", err)
    }
    defer rows.Close()

    stats := make(map[string]int)
    for rows.Next() {
        var relation string
        var count int
        if err := rows.Scan(&relation, &count); err != nil {
            return nil, fmt.Errorf("扫描关系统计失败: %w", err)
        }
        stats[relation] = count
    }

    if err := rows.Err(); err != nil {
        return nil, fmt.Errorf("遍历关系统计失败: %w", err)
    }

    return stats, nil
}

// FindNodesByKeyword 根据关键词查找节点（模糊匹配标题、路径、标签、实体）
func (s *Store) FindNodesByKeyword(keyword string) ([]*graph.Node, error) {
    keyword = "%" + keyword + "%"
    query := `
        SELECT id, path, title, section, tags, entities, wikilinks, content_fts
        FROM nodes
        WHERE title LIKE ? OR path LIKE ? OR tags LIKE ? OR entities LIKE ?
        LIMIT 20
    `
    rows, err := s.db.Query(query, keyword, keyword, keyword, keyword)
    if err != nil {
        return nil, fmt.Errorf("查询节点失败: %w", err)
    }
    defer rows.Close()

    var nodes []*graph.Node
    for rows.Next() {
        var n graph.Node
        var tagsJSON, entitiesJSON, wikilinksJSON string
        var content sql.NullString

        err := rows.Scan(&n.ID, &n.Path, &n.Title, &n.Section,
            &tagsJSON, &entitiesJSON, &wikilinksJSON, &content)
        if err != nil {
            return nil, fmt.Errorf("扫描节点失败: %w", err)
        }

        json.Unmarshal([]byte(tagsJSON), &n.Tags)
        json.Unmarshal([]byte(entitiesJSON), &n.Entities)
        json.Unmarshal([]byte(wikilinksJSON), &n.Wikilinks)
        if content.Valid {
            n.Content = content.String
        }

        nodes = append(nodes, &n)
    }

    if err := rows.Err(); err != nil {
        return nil, fmt.Errorf("遍历节点失败: %w", err)
    }

    return nodes, nil
}

// GetNodeEdges 获取节点的关联边（支持关系类型过滤和深度查询）
func (s *Store) GetNodeEdges(nodeID int64, relation string, depth int) ([]*graph.Edge, error) {
    if depth < 1 {
        depth = 1
    }
    if depth > 3 {
        depth = 3
    }

    var edges []*graph.Edge
    visited := make(map[int64]bool)
    currentLevel := []int64{nodeID}

    for d := 0; d < depth; d++ {
        var nextLevel []int64

        for _, id := range currentLevel {
            if visited[id] && d > 0 {
                continue
            }
            visited[id] = true

            query := `
                SELECT id, from_node, to_node, relation, label
                FROM edges
                WHERE from_node = ?
            `
            args := []interface{}{id}

            if relation != "" {
                query += " AND relation = ?"
                args = append(args, relation)
            }

            rows, err := s.db.Query(query, args...)
            if err != nil {
                return nil, fmt.Errorf("查询边失败: %w", err)
            }

            for rows.Next() {
                var e graph.Edge
                var edgeID int64
                if err := rows.Scan(&edgeID, &e.FromNode, &e.ToNode, &e.Relation, &e.Label); err != nil {
                    rows.Close()
                    return nil, fmt.Errorf("扫描边失败: %w", err)
                }
                edges = append(edges, &e)
                nextLevel = append(nextLevel, e.ToNode)
            }
            rows.Close()
        }

        currentLevel = nextLevel
    }

    return edges, nil
}

// FindRelatedNodes 查找与关键词相关的所有节点（通过共同标签、实体、wikilink）
func (s *Store) FindRelatedNodes(keyword string, topN int) ([]RelatedNode, error) {
    // 先找到匹配关键词的节点
    keyword = "%" + keyword + "%"
    query := `
        SELECT id, path, title, section, tags, entities
        FROM nodes
        WHERE title LIKE ? OR path LIKE ? OR tags LIKE ? OR entities LIKE ?
        LIMIT 5
    `
    rows, err := s.db.Query(query, keyword, keyword, keyword, keyword)
    if err != nil {
        return nil, fmt.Errorf("查询节点失败: %w", err)
    }
    defer rows.Close()

    var seedNodes []*graph.Node
    for rows.Next() {
        var n graph.Node
        var tagsJSON, entitiesJSON string
        err := rows.Scan(&n.ID, &n.Path, &n.Title, &n.Section, &tagsJSON, &entitiesJSON)
        if err != nil {
            return nil, fmt.Errorf("扫描节点失败: %w", err)
        }
        json.Unmarshal([]byte(tagsJSON), &n.Tags)
        json.Unmarshal([]byte(entitiesJSON), &n.Entities)
        seedNodes = append(seedNodes, &n)
    }

    if len(seedNodes) == 0 {
        return []RelatedNode{}, nil
    }

    // 收集种子节点的标签和实体
    tagSet := make(map[string]bool)
    entitySet := make(map[string]bool)
    for _, n := range seedNodes {
        for _, tag := range n.Tags {
            tagSet[tag] = true
        }
        for _, entity := range n.Entities {
            entitySet[entity] = true
        }
    }

    // 查找共享标签或实体的节点
    relevanceMap := make(map[int64]int)
    pathMap := make(map[int64]string)
    titleMap := make(map[int64]string)
    sectionMap := make(map[int64]string)
    tagsMap := make(map[int64][]string)

    // 通过标签查找
    for tag := range tagSet {
        tagPattern := "%" + tag + "%"
        rows, err := s.db.Query(`
            SELECT id, path, title, section, tags
            FROM nodes
            WHERE tags LIKE ?
        `, tagPattern)
        if err != nil {
            continue
        }

        for rows.Next() {
            var id int64
            var path, title, section, tagsJSON string
            if err := rows.Scan(&id, &path, &title, &section, &tagsJSON); err != nil {
                rows.Close()
                continue
            }
            relevanceMap[id]++
            pathMap[id] = path
            titleMap[id] = title
            sectionMap[id] = section
            var tags []string
            json.Unmarshal([]byte(tagsJSON), &tags)
            tagsMap[id] = tags
        }
        rows.Close()
    }

    // 通过实体查找
    for entity := range entitySet {
        entityPattern := "%" + entity + "%"
        rows, err := s.db.Query(`
            SELECT id, path, title, section, tags
            FROM nodes
            WHERE entities LIKE ?
        `, entityPattern)
        if err != nil {
            continue
        }

        for rows.Next() {
            var id int64
            var path, title, section, tagsJSON string
            if err := rows.Scan(&id, &path, &title, &section, &tagsJSON); err != nil {
                rows.Close()
                continue
            }
            relevanceMap[id]++
            pathMap[id] = path
            titleMap[id] = title
            sectionMap[id] = section
            var tags []string
            json.Unmarshal([]byte(tagsJSON), &tags)
            tagsMap[id] = tags
        }
        rows.Close()
    }

    // 转换为结果列表
    var results []RelatedNode
    for id, relevance := range relevanceMap {
        results = append(results, RelatedNode{
            ID:        id,
            Path:      pathMap[id],
            Title:     titleMap[id],
            Section:   sectionMap[id],
            Tags:      tagsMap[id],
            Relevance: relevance,
        })
    }

    // 按相关度排序
    sort.Slice(results, func(i, j int) bool {
        return results[i].Relevance > results[j].Relevance
    })

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

    return results, nil
}

// RelatedNode 相关节点
type RelatedNode struct {
    ID        int64    `json:"id"`
    Path      string   `json:"path"`
    Title     string   `json:"title"`
    Section   string   `json:"section"`
    Tags      []string `json:"tags"`
    Relevance int      `json:"relevance"`
}