package com.superbiz.agent.service; import com.superbiz.agent.domain.entity.ApiDocument; import com.superbiz.agent.domain.enums.FaultCategory; import com.superbiz.agent.repository.ApiDocumentRepository; import com.superbiz.agent.dto.KnowledgeEntry; import com.superbiz.agent.dto.Frontmatter; import com.superbiz.agent.dto.DocumentChunk; import lombok.Data; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Value; import org.springframework.stereotype.Service; import org.springframework.transaction.annotation.Transactional; import java.io.IOException; import java.nio.file.*; import java.nio.file.attribute.BasicFileAttributes; import java.time.LocalDateTime; import java.util.*; import java.util.stream.Collectors; import java.util.stream.Collectors; /** * 知识库初始化服务 * 负责批量导入 knowledge_base 目录下的文档到数据库和 Milvus */ @Service public class KnowledgeBaseInitService { private static final Logger logger = LoggerFactory.getLogger(KnowledgeBaseInitService.class); @Value("${knowledge.base-path:knowledge_base}") private String knowledgeBasePath; @Autowired private ApiDocumentRepository apiDocumentRepository; @Autowired private FrontmatterParser frontmatterParser; @Autowired private DocumentChunkService documentChunkService; @Autowired private VectorIndexService vectorIndexService; @Autowired private VectorEmbeddingService vectorEmbeddingService; @Autowired private KnowledgeIndexService knowledgeIndexService; /** * 初始化知识库 * * @param force 是否强制重新导入(跳过去重检查) * @return 初始化结果 */ @Transactional(rollbackFor = Exception.class) public InitResult initializeKnowledgeBase(boolean force) { logger.info("开始初始化知识库: basePath={}, force={}", knowledgeBasePath, force); InitResult result = new InitResult(); Path baseDir = Paths.get(knowledgeBasePath); if (!Files.exists(baseDir)) { logger.error("知识库目录不存在: {}", knowledgeBasePath); throw new RuntimeException("知识库目录不存在: " + knowledgeBasePath); } // 1. 扫描所有 Markdown 文件 List markdownFiles = scanMarkdownFiles(baseDir); result.setScanned(markdownFiles.size()); logger.info("扫描到 {} 个 Markdown 文件", markdownFiles.size()); // 2. 如果非强制模式,获取已存在的文档(用于去重) Set existingFilePaths = new HashSet<>(); if (!force) { existingFilePaths = apiDocumentRepository.findAll().stream() .map(ApiDocument::getFilePath) .collect(Collectors.toSet()); logger.info("已存在 个文档记录", existingFilePaths.size()); } // 3. 逐个处理文档 for (Path file : markdownFiles) { String relativePath = baseDir.relativize(file).toString().replace("\\", "/"); try { // 去重检查 if (!force && existingFilePaths.contains(relativePath)) { logger.debug("跳过已存在的文档: {}", relativePath); result.incrementSkipped(); result.addDetail(relativePath, "已存在,跳过"); continue; } // 解析文档 String content = Files.readString(file); Frontmatter frontmatter = frontmatterParser.parse(content); if (frontmatter == null) { logger.warn("文档格式无效: {}, frontmatter 解析失败", relativePath); result.incrementFailed(); result.addDetail(relativePath, "格式无效: frontmatter 解析失败"); continue; } // 提取字段 String title = frontmatter.getTitle(); String summary = frontmatter.getSummary(); String category = frontmatter.getCategory() != null ? frontmatter.getCategory() : "general"; List keywords = frontmatter.getKeywords(); if (title == null || title.isBlank()) { logger.warn("文档缺少标题: {}", relativePath); result.incrementFailed(); result.addDetail(relativePath, "缺少标题"); continue; } // 保存到数据库 ApiDocument document = saveToDatabase(relativePath, title, summary, category, content, keywords); // 提取文档正文(去除 frontmatter) String body = extractBody(content); // 文档分块 List chunks = documentChunkService.chunkDocument(body, relativePath); logger.debug("文档分块完成: {} -> {} 个 chunk", relativePath, chunks.size()); // 上传到 Milvus try { vectorIndexService.indexDocumentChunks(document.getDocId(), chunks, category); document.setStatus("INDEXED"); document.setChunkCount(chunks.size()); document.setIndexedAt(LocalDateTime.now()); apiDocumentRepository.save(document); logger.info("文档已索引到 Milvus: {} (docId={}, chunks={})", title, document.getDocId(), chunks.size()); } catch (Exception e) { logger.error("上传到 Milvus 失败: {}", relativePath, e); document.setStatus("FAILED"); document.setErrorMessage(e.getMessage()); apiDocumentRepository.save(document); result.incrementFailed(); result.addDetail(relativePath, "Milvus 索引失败: " + e.getMessage()); continue; // 跳过该文档,继续处理下一个 } // 添加到 L0 内存索引 KnowledgeEntry entry = KnowledgeEntry.builder() .filePath(relativePath) .title(title) .keywords(keywords) .summary(summary) .category(category) .build(); knowledgeIndexService.addToIndex(entry); result.incrementInserted(); result.addDetail(relativePath, "导入成功(L0+L1)"); logger.info("文档导入成功: {} -> {} (L0+L1 索引已更新)", relativePath, title); } catch (Exception e) { logger.error("处理文档失败: {}", relativePath, e); result.incrementFailed(); result.addDetail(relativePath, "处理失败: " + e.getMessage()); } } logger.info("知识库初始化完成: 扫描={}, 跳过={}, 新增={}, 失败={}", result.getScanned(), result.getSkipped(), result.getInserted(), result.getFailed()); return result; } /** * 获取知识库统计信息 */ public Stats getStats() { Stats stats = new Stats(); // 数据库中的文档数量 long totalDocuments = apiDocumentRepository.count(); stats.setTotalDocuments(totalDocuments); // L0 索引中的文档数量 int indexSize = knowledgeIndexService.getIndexSize(); logger.debug("L0 索引大小: {}", indexSize); // 按分类统计(从 fault_category 字段读取) Map categoryCount = apiDocumentRepository.findAll().stream() .collect(Collectors.groupingBy( doc -> doc.getFaultCategory() != null ? doc.getFaultCategory().name() : "GENERAL", Collectors.counting() )); stats.setCategoryCount(categoryCount); // Milvus 中的向量数量(需要实现) // TODO: 查询 Milvus collection 的实体数量 stats.setTotalVectors(0L); return stats; } /** * 扫描目录下所有 Markdown 文件 */ private List scanMarkdownFiles(Path baseDir) { List files = new ArrayList<>(); try { Files.walkFileTree(baseDir, new SimpleFileVisitor() { @Override public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) { if (file.toString().endsWith(".md")) { files.add(file); } return FileVisitResult.CONTINUE; } @Override public FileVisitResult visitFileFailed(Path file, IOException exc) { logger.warn("访问文件失败: {}", file, exc); return FileVisitResult.CONTINUE; } }); } catch (IOException e) { logger.error("扫描目录失败: {}", baseDir, e); throw new RuntimeException("扫描目录失败", e); } return files; } /** * 保存文档到数据库 */ private ApiDocument saveToDatabase(String filePath, String title, String summary, String category, String content, List keywords) { ApiDocument document = new ApiDocument(); document.setDocId(UUID.randomUUID().toString()); document.setFileName(Paths.get(filePath).getFileName().toString()); document.setFilePath(filePath); document.setApiName(title); // 使用 title 作为 apiName document.setStatus("PENDING"); // 初始状态为 PENDING,索引成功后更新为 INDEXED // 映射 category 到 FaultCategory 枚举 FaultCategory faultCategory = FaultCategory.fromString(category); document.setFaultCategory(faultCategory); // 将 frontmatter 信息保存到 metadata(JSON 格式) String metadataJson = String.format( "{\"title\":\"%s\",\"summary\":\"%s\",\"category\":\"%s\",\"keywords\":%s}", escapeJson(title), escapeJson(summary), escapeJson(category), "[\"" + String.join("\",\"", keywords.stream().map(this::escapeJson).toArray(String[]::new)) + "\"]" ); document.setMetadata(metadataJson); document.setFileSize((long) content.length()); return apiDocumentRepository.save(document); } /** * JSON 转义 */ private String escapeJson(String str) { if (str == null) { return ""; } return str.replace("\\", "\\\\") .replace("\"", "\\\"") .replace("\n", "\\n") .replace("\r", "\\r"); } /** * 提取文档正文(去除 frontmatter) */ private String extractBody(String content) { if (!content.trim().startsWith("---")) { return content; } int firstEnd = content.indexOf("---", 3); if (firstEnd == -1) { return content; } int secondEnd = content.indexOf("---", firstEnd + 3); if (secondEnd == -1) { return content.substring(firstEnd + 3).trim(); } return content.substring(secondEnd + 3).trim(); } // ==================== 数据模型 ==================== /** * 初始化结果 */ @Data public static class InitResult { private int scanned; // 扫描到的文件数量 private int skipped; // 跳过的文件数量(已存在) private int inserted; // 成功导入的文件数量 private int failed; // 失败的文件数量 private Map details = new LinkedHashMap<>(); // 详细信息 public void incrementSkipped() { this.skipped++; } public void incrementInserted() { this.inserted++; } public void incrementFailed() { this.failed++; } public void addDetail(String filePath, String message) { this.details.put(filePath, message); } } /** * 统计信息 */ @Data public static class Stats { private long totalDocuments; // 数据库中的文档总数 private long totalVectors; // Milvus 中的向量总数 private Map categoryCount; // 按分类统计 } }