ai_xiaopei
9 days ago 1ba86b102ecc01de3191c3861592d9d4d22dfd81
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
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
}