| | |
| | | "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, |
| | | } |
| | | |
| | |
| | | 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 { |
| | |
| | | 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 |
| | | } |