package cmd
|
|
import (
|
"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: `# classify - LLM 分类
|
kb-cli classify [--batch-size=<数量>] [--dry-run] # 使用 LLM 分析文档，提取平台、设备、内容类型`,
|
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()
|
|
// 创建 LLM 客户端（从配置文件加载）
|
llmClient := llm.NewClient()
|
|
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: 扫描文档...")
|
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
|
}
|
