feat: 添加 classify 命令和实体提取功能

ai_xiaopei

2026-07-26 1ba86b102ecc01de3191c3861592d9d4d22dfd81

feat: 添加 classify 命令和实体提取功能

 cmd/classify.go

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

import (
    "fmt"

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

var classifyCmd = &cobra.Command{
    Use:   "classify",
    Short: "使用 LLM 对知识库文档进行分类",
    Long: `使用本地 LLM 分析知识库文档，提取平台、设备、内容类型等信息。

输出：
  - entities.json: 实体定义（从 实体/ 目录提取）
  - relations.json: 实体关系
  - classification.json: 文档分类结果`,
    RunE: runClassify,
}

var (
    classifyBatchSize int
    classifyDryRun    bool
)

func init() {
    rootCmd.AddCommand(classifyCmd)
    classifyCmd.Flags().IntVar(&classifyBatchSize, "batch-size", 10, "每批处理的文档数量")
    classifyCmd.Flags().BoolVar(&classifyDryRun, "dry-run", false, "仅显示将要执行的操作")
}

func runClassify(cmd *cobra.Command, args []string) error {
    vaultPath := cmd.Flag("vault").Value.String()

    fmt.Println("步骤 1/3: 提取实体信息...")
    entities, relations, err := classify.ExtractEntities(vaultPath)
    if err != nil {
        return fmt.Errorf("提取实体失败: %w", err)
    }

    if !classifyDryRun {
        if err := classify.SaveEntities(vaultPath, entities, relations); err != nil {
            return fmt.Errorf("保存实体失败: %w", err)
        }
        fmt.Printf("✓ 已生成 entities.json 和 relations.json\n")
    } else {
        fmt.Printf("[DRY-RUN] 将生成 entities.json 和 relations.json\n")
    }

    fmt.Println("\n步骤 2/3: 文档分类（待实现）...")
    fmt.Println("步骤 3/3: 生成 classification.json（待实现）...")

    return nil
}

 internal/classify/entities.go

New file
@@ -0,0 +1,164 @@
package classify

import (
    "encoding/json"
    "fmt"
    "os"
    "path/filepath"
)

// Entity 实体定义
type Entity struct {
    ID          string   `json:"id"`
    Name        string   `json:"name"`
    Aliases     []string `json:"aliases"`
    Platform    string   `json:"platform,omitempty"`
    Description string   `json:"description,omitempty"`
    Users       []string `json:"users,omitempty"`
    URL         string   `json:"url,omitempty"`
}

// Entities 实体集合
type Entities struct {
    Platforms []Entity `json:"platforms"`
    Apps      []Entity `json:"apps"`
    Devices   []Entity `json:"devices"`
}

// Relation 实体关系
type Relation struct {
    From        string `json:"from"`
    To          string `json:"to"`
    Type        string `json:"type"`
    Description string `json:"description"`
}

// Relations 关系集合
type Relations struct {
    Relations []Relation `json:"relations"`
}

// ExtractEntities 从实体目录提取实体信息
func ExtractEntities(vaultPath string) (*Entities, *Relations, error) {
    entityDir := filepath.Join(vaultPath, "实体")

    // 检查目录是否存在
    if _, err := os.Stat(entityDir); os.IsNotExist(err) {
        return nil, nil, fmt.Errorf("实体目录不存在: %s", entityDir)
    }

    // 预定义实体（从设计文档）
    entities := &Entities{
        Platforms: []Entity{
            {
                ID:          "elc",
                Name:        "电子秤平台",
                Aliases:     []string{"电子秤后台", "elc平台", "elc.zhiheiot.com"},
                URL:         "https://elc.zhiheiot.com",
                Description: "气站充装作业管理系统",
                Users:       []string{"气站站长", "充装员", "开票员"},
            },
            {
                ID:          "ops",
                Name:        "运营管理平台",
                Aliases:     []string{"运营平台", "运营后台", "rb.zhiheiot.com"},
                URL:         "https://rb.zhiheiot.com",
                Description: "气站运营管理系统",
                Users:       []string{"老板", "管理层", "客服"},
            },
        },
        Apps: []Entity{
            {
                ID:          "safety",
                Name:        "安全用气App",
                Aliases:     []string{"配送App", "电子秤配送App"},
                Platform:    "elc",
                Description: "二维码配送",
            },
            {
                ID:          "delivery",
                Name:        "易配送App",
                Aliases:     []string{"配送App", "运营配送"},
                Platform:    "ops",
                Description: "NFC识别芯片配送",
            },
            {
                ID:          "assistant",
                Name:        "艾信助手App",
                Aliases:     []string{"建档App"},
                Platform:    "both",
                Description: "融合两平台建档",
            },
        },
        Devices: []Entity{
            {
                ID:          "scale",
                Name:        "电子秤",
                Aliases:     []string{"充装电子秤"},
                Platform:    "elc",
                Description: "充装作业设备",
            },
            {
                ID:          "gun",
                Name:        "智能枪",
                Aliases:     []string{"充装枪", "智能充装枪"},
                Platform:    "both",
                Description: "共有硬件",
            },
            {
                ID:          "valve",
                Name:        "智能阀",
                Aliases:     []string{"智能角阀"},
                Platform:    "ops",
                Description: "气瓶阀门",
            },
            {
                ID:          "box",
                Name:        "艾信盒子",
                Aliases:     []string{"4G通信设备"},
                Platform:    "both",
                Description: "共有硬件",
            },
        },
    }

    // 预定义关系
    relations := &Relations{
        Relations: []Relation{
            {From: "gun", To: "scale", Type: "connects", Description: "智能枪连接电子秤"},
            {From: "gun", To: "box", Type: "syncs", Description: "充装数据同步到艾信盒子"},
            {From: "box", To: "elc", Type: "uploads", Description: "上传充装记录"},
            {From: "box", To: "ops", Type: "uploads", Description: "上传充装记录"},
            {From: "valve", To: "delivery", Type: "identified_by", Description: "NFC识别"},
            {From: "safety", To: "elc", Type: "belongs_to", Description: "属于电子秤平台"},
            {From: "delivery", To: "ops", Type: "belongs_to", Description: "属于运营管理平台"},
        },
    }

    return entities, relations, nil
}

// SaveEntities 保存实体和关系到 JSON 文件
func SaveEntities(vaultPath string, entities *Entities, relations *Relations) error {
    // 保存 entities.json
    entitiesPath := filepath.Join(vaultPath, "entities.json")
    entitiesData, err := json.MarshalIndent(entities, "", "  ")
    if err != nil {
        return fmt.Errorf("序列化 entities 失败: %w", err)
    }
    if err := os.WriteFile(entitiesPath, entitiesData, 0644); err != nil {
        return fmt.Errorf("写入 entities.json 失败: %w", err)
    }

    // 保存 relations.json
    relationsPath := filepath.Join(vaultPath, "relations.json")
    relationsData, err := json.MarshalIndent(relations, "", "  ")
    if err != nil {
        return fmt.Errorf("序列化 relations 失败: %w", err)
    }
    if err := os.WriteFile(relationsPath, relationsData, 0644); err != nil {
        return fmt.Errorf("写入 relations.json 失败: %w", err)
    }

    return nil
}