From 1ba86b102ecc01de3191c3861592d9d4d22dfd81 Mon Sep 17 00:00:00 2001
From: ai_xiaopei <xiaopei@aisim.cn>
Date: Sun, 26 Jul 2026 15:17:48 +0800
Subject: [PATCH] feat: 添加 classify 命令和实体提取功能
---
internal/classify/entities.go | 164 +++++++++++++++++++++++++++++++++++++++++
cmd/classify.go | 55 +++++++++++++
2 files changed, 219 insertions(+), 0 deletions(-)
diff --git a/cmd/classify.go b/cmd/classify.go
new file mode 100644
index 0000000..d698dd0
--- /dev/null
+++ b/cmd/classify.go
@@ -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
+}
diff --git a/internal/classify/entities.go b/internal/classify/entities.go
new file mode 100644
index 0000000..14e4c7c
--- /dev/null
+++ b/internal/classify/entities.go
@@ -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
+}
--
Gitblit v1.9.1