ai_xiaopei
6 days ago 1196f409d86bc61e7596eb274840244a62ce84ba
cmd/classify.go
@@ -4,18 +4,15 @@
   "fmt"
   "github.com/aisim/kb-cli/internal/classify"
   "github.com/aisim/kb-cli/internal/llm"
   "github.com/spf13/cobra"
)
var classifyCmd = &cobra.Command{
   Use:   "classify",
   Short: "使用 LLM 对知识库文档进行分类",
   Long: `使用本地 LLM 分析知识库文档,提取平台、设备、内容类型等信息。
输出:
  - entities.json: 实体定义(从 实体/ 目录提取)
  - relations.json: 实体关系
  - classification.json: 文档分类结果`,
   Long:  `# classify - LLM 分类
kb-cli classify [--batch-size <数量>] [--dry-run] # 使用 LLM 分析文档,提取平台、设备、内容类型`,
   RunE: runClassify,
}
@@ -33,6 +30,9 @@
func runClassify(cmd *cobra.Command, args []string) error {
   vaultPath := cmd.Flag("vault").Value.String()
   // 创建 LLM 客户端(从配置文件加载)
   llmClient := llm.NewClient()
   fmt.Println("步骤 1/3: 提取实体信息...")
   entities, relations, err := classify.ExtractEntities(vaultPath)
   if err != nil {
@@ -48,8 +48,52 @@
      fmt.Printf("[DRY-RUN] 将生成 entities.json 和 relations.json\n")
   }
   fmt.Println("\n步骤 2/3: 文档分类(待实现)...")
   fmt.Println("步骤 3/3: 生成 classification.json(待实现)...")
   fmt.Println("\n步骤 2/3: 扫描文档...")
   docs, err := classify.ScanDocuments(vaultPath)
   if err != nil {
      return fmt.Errorf("扫描文档失败: %w", err)
   }
   fmt.Printf("✓ 发现 %d 个文档\n", len(docs))
   if len(docs) == 0 {
      fmt.Println("警告: 未发现任何文档")
      return nil
   }
   fmt.Println("\n步骤 3/3: 使用 LLM 分类...")
   if classifyDryRun {
      fmt.Printf("[DRY-RUN] 将对 %d 个文档进行分类(批次大小: %d)\n", len(docs), classifyBatchSize)
      return nil
   }
   // 分批处理
   var allResults []classify.Classification
   for i := 0; i < len(docs); i += classifyBatchSize {
      end := i + classifyBatchSize
      if end > len(docs) {
         end = len(docs)
      }
      batch := docs[i:end]
      fmt.Printf("处理批次 %d-%d / %d...\n", i+1, end, len(docs))
      results, err := classify.ClassifyBatch(batch, llmClient)
      if err != nil {
         fmt.Printf("警告: 批次 %d 分类失败: %v\n", i/classifyBatchSize+1, err)
         continue
      }
      allResults = append(allResults, results...)
   }
   if len(allResults) == 0 {
      return fmt.Errorf("分类失败: 未获得任何结果")
   }
   if err := classify.SaveClassification(vaultPath, allResults); err != nil {
      return fmt.Errorf("保存分类结果失败: %w", err)
   }
   fmt.Printf("✓ 已生成 classification.json(%d 个文档)\n", len(allResults))
   return nil
}