feat: add graph command for knowledge graph queries

ai_xiaopei

9 days ago 6b58a8fd1d10191d5720b037af1137d2a747cfeb

 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"`
}