571 lines
21 KiB
Java
571 lines
21 KiB
Java
package com.superbiz.agent.service;
|
||
|
||
import io.milvus.client.MilvusServiceClient;
|
||
import io.milvus.grpc.MutationResult;
|
||
import io.milvus.param.R;
|
||
import io.milvus.param.RpcStatus;
|
||
import io.milvus.param.collection.LoadCollectionParam;
|
||
import io.milvus.param.dml.DeleteParam;
|
||
import io.milvus.param.dml.InsertParam;
|
||
import lombok.Getter;
|
||
import lombok.Setter;
|
||
import com.superbiz.agent.constant.MilvusConstants;
|
||
import com.superbiz.agent.dto.DocumentChunk;
|
||
import com.superbiz.agent.dto.Frontmatter;
|
||
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 java.io.File;
|
||
import java.nio.file.Files;
|
||
import java.nio.file.Path;
|
||
import java.nio.file.Paths;
|
||
import java.time.LocalDateTime;
|
||
import java.util.*;
|
||
|
||
/**
|
||
* 向量索引服务
|
||
* 负责读取文件、生成向量、存储到 Milvus
|
||
*/
|
||
@Service
|
||
public class VectorIndexService {
|
||
|
||
private static final Logger logger = LoggerFactory.getLogger(VectorIndexService.class);
|
||
|
||
@Autowired
|
||
private MilvusServiceClient milvusClient;
|
||
|
||
@Autowired
|
||
private VectorEmbeddingService embeddingService;
|
||
|
||
@Autowired
|
||
private DocumentChunkService chunkService;
|
||
|
||
@Value("${file.upload.path}")
|
||
private String uploadPath;
|
||
|
||
/**
|
||
* 索引指定目录下的所有文件
|
||
*
|
||
* @param directoryPath 目录路径（可选，默认使用配置的上传目录）
|
||
* @return 索引结果 这里可以优化：定时重建目录下所有文件的索引
|
||
*/
|
||
public IndexingResult indexDirectory(String directoryPath) {
|
||
IndexingResult result = new IndexingResult();
|
||
result.setStartTime(LocalDateTime.now());
|
||
|
||
try {
|
||
// 使用指定目录或默认上传目录
|
||
String targetPath = (directoryPath != null && !directoryPath.trim().isEmpty())
|
||
? directoryPath : uploadPath;
|
||
|
||
Path dirPath = Paths.get(targetPath).normalize();
|
||
File directory = dirPath.toFile();
|
||
|
||
if (!directory.exists() || !directory.isDirectory()) {
|
||
throw new IllegalArgumentException("目录不存在或不是有效目录: " + targetPath);
|
||
}
|
||
|
||
result.setDirectoryPath(directory.getAbsolutePath());
|
||
|
||
// 获取所有支持的文件
|
||
File[] files = directory.listFiles((dir, name) ->
|
||
name.endsWith(".txt") || name.endsWith(".md")
|
||
);
|
||
|
||
if (files == null || files.length == 0) {
|
||
logger.warn("目录中没有找到支持的文件: {}", targetPath);
|
||
result.setTotalFiles(0);
|
||
result.setSuccess(true);
|
||
result.setEndTime(LocalDateTime.now());
|
||
return result;
|
||
}
|
||
|
||
result.setTotalFiles(files.length);
|
||
logger.info("开始索引目录: {}, 找到 {} 个文件", targetPath, files.length);
|
||
|
||
// 遍历并索引每个文件
|
||
for (File file : files) {
|
||
try {
|
||
indexSingleFile(file.getAbsolutePath());
|
||
result.incrementSuccessCount();
|
||
logger.info("✓ 文件索引成功: {}", file.getName());
|
||
} catch (Exception e) {
|
||
result.incrementFailCount();
|
||
result.addFailedFile(file.getAbsolutePath(), e.getMessage());
|
||
logger.error("✗ 文件索引失败: {}", file.getName(), e);
|
||
}
|
||
}
|
||
|
||
result.setSuccess(result.getFailCount() == 0);
|
||
result.setEndTime(LocalDateTime.now());
|
||
|
||
logger.info("目录索引完成: 总数={}, 成功={}, 失败={}",
|
||
result.getTotalFiles(), result.getSuccessCount(), result.getFailCount());
|
||
|
||
return result;
|
||
|
||
} catch (Exception e) {
|
||
logger.error("索引目录失败", e);
|
||
result.setSuccess(false);
|
||
result.setErrorMessage(e.getMessage());
|
||
result.setEndTime(LocalDateTime.now());
|
||
return result;
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 索引单个文件
|
||
*
|
||
* @param filePath 文件路径
|
||
* @throws Exception 索引失败时抛出异常
|
||
*/
|
||
public void indexSingleFile(String filePath) throws Exception {
|
||
Path path = Paths.get(filePath).normalize();
|
||
File file = path.toFile();
|
||
|
||
if (!file.exists() || !file.isFile()) {
|
||
throw new IllegalArgumentException("文件不存在: " + filePath);
|
||
}
|
||
|
||
logger.info("开始索引文件: {}", path);
|
||
|
||
// 1. 读取文件内容
|
||
String content = Files.readString(path);
|
||
logger.info("读取文件: {}, 内容长度: {} 字符", path, content.length());
|
||
|
||
// 2. 删除该文件的旧数据（如果存在）
|
||
deleteExistingData(path.toString());
|
||
|
||
// 3. 文档分片
|
||
List<DocumentChunk> chunks = chunkService.chunkDocument(content, path.toString());
|
||
logger.info("文档分片完成: {} -> {} 个分片", filePath, chunks.size());
|
||
|
||
// 4. 为每个分片生成向量并插入 Milvus
|
||
for (int i = 0; i < chunks.size(); i++) {
|
||
DocumentChunk chunk = chunks.get(i);
|
||
|
||
try {
|
||
// 生成向量
|
||
List<Float> vector = embeddingService.generateEmbedding(buildEmbeddingText(chunk));
|
||
|
||
// 构建元数据（包含文件信息）
|
||
Map<String, Object> metadata = buildMetadata(path.toString(), chunk, chunks.size());
|
||
|
||
// 插入到 Milvus
|
||
insertToMilvus(chunk.getContent(), vector, metadata, chunk.getChunkIndex());
|
||
|
||
logger.info("✓ 分片 {}/{} 索引成功", i + 1, chunks.size());
|
||
|
||
} catch (Exception e) {
|
||
logger.error("✗ 分片 {}/{} 索引失败", i + 1, chunks.size(), e);
|
||
throw new RuntimeException("分片索引失败: " + e.getMessage(), e);
|
||
}
|
||
}
|
||
|
||
logger.info("文件索引完成: {}, 共 {} 个分片", filePath, chunks.size());
|
||
}
|
||
|
||
/**
|
||
* 索引文档分块（用于上传文档的向量化）
|
||
*
|
||
* @param docId 文档ID
|
||
* @param chunks 文档分块列表
|
||
* @param category 文档类别（api、domain、troubleshoot 等）
|
||
* @throws Exception 索引失败时抛出异常
|
||
*/
|
||
public void indexDocumentChunks(String docId, List<DocumentChunk> chunks, String category) throws Exception {
|
||
indexDocumentChunks(docId, chunks, category, null);
|
||
}
|
||
|
||
public void indexDocumentChunks(String docId, List<DocumentChunk> chunks, String category, Frontmatter frontmatter) throws Exception {
|
||
if (chunks == null || chunks.isEmpty()) {
|
||
throw new IllegalArgumentException("文档分块列表为空");
|
||
}
|
||
|
||
logger.info("开始索引文档分块，docId: {}, 分块数: {}, 类别: {}", docId, chunks.size(), category);
|
||
|
||
// 1. 删除该文档的旧数据（如果存在）
|
||
deleteDocumentChunks(docId);
|
||
|
||
// 2. 为每个分块生成向量并插入 Milvus
|
||
for (int i = 0; i < chunks.size(); i++) {
|
||
DocumentChunk chunk = chunks.get(i);
|
||
|
||
try {
|
||
// 生成向量
|
||
List<Float> vector = embeddingService.generateEmbedding(buildEmbeddingText(chunk));
|
||
|
||
// 构建元数据（使用 docId 和 category）
|
||
Map<String, Object> metadata = buildDocumentMetadata(docId, chunk, chunks.size(), category, frontmatter);
|
||
|
||
// 插入到 Milvus
|
||
insertToMilvus(chunk.getContent(), vector, metadata, chunk.getChunkIndex());
|
||
|
||
logger.info("✓ 文档分块 {}/{} 索引成功，docId: {}", i + 1, chunks.size(), docId);
|
||
|
||
} catch (Exception e) {
|
||
logger.error("✗ 文档分块 {}/{} 索引失败，docId: {}", i + 1, chunks.size(), docId, e);
|
||
throw new RuntimeException("文档分块索引失败: " + e.getMessage(), e);
|
||
}
|
||
}
|
||
|
||
logger.info("文档索引完成，docId: {}, 共 {} 个分块，类别: {}", docId, chunks.size(), category);
|
||
}
|
||
|
||
/**
|
||
* 删除文档的所有分块（根据 docId）
|
||
*/
|
||
public void deleteDocumentChunks(String docId) {
|
||
try {
|
||
// 构建删除表达式：metadata["docId"] == "xxx"
|
||
String expr = String.format("metadata[\"docId\"] == \"%s\"", docId);
|
||
|
||
logger.info("准备删除文档旧数据，docId: {}, 表达式: {}", docId, expr);
|
||
|
||
// 确保 collection 已加载
|
||
R<RpcStatus> loadResponse = milvusClient.loadCollection(
|
||
LoadCollectionParam.newBuilder()
|
||
.withCollectionName(MilvusConstants.MILVUS_COLLECTION_NAME)
|
||
.build()
|
||
);
|
||
|
||
if (loadResponse.getStatus() != 0 && loadResponse.getStatus() != 65535) {
|
||
logger.warn("加载 collection 失败: {}", loadResponse.getMessage());
|
||
return;
|
||
}
|
||
|
||
DeleteParam deleteParam = DeleteParam.newBuilder()
|
||
.withCollectionName(MilvusConstants.MILVUS_COLLECTION_NAME)
|
||
.withExpr(expr)
|
||
.build();
|
||
|
||
R<MutationResult> deleteResponse = milvusClient.delete(deleteParam);
|
||
|
||
if (deleteResponse.getStatus() == 0) {
|
||
logger.info("删除文档旧数据成功，docId: {}", docId);
|
||
} else {
|
||
logger.warn("删除文档旧数据失败，docId: {}, 原因: {}", docId, deleteResponse.getMessage());
|
||
}
|
||
|
||
} catch (Exception e) {
|
||
logger.warn("删除文档旧数据异常，docId: {}", docId, e);
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 构建文档元数据（用于上传文档）
|
||
*/
|
||
static Map<String, Object> buildDocumentMetadata(String docId, DocumentChunk chunk, int totalChunks, String category) {
|
||
return buildDocumentMetadata(docId, chunk, totalChunks, category, null);
|
||
}
|
||
|
||
static Map<String, Object> buildDocumentMetadata(String docId,
|
||
DocumentChunk chunk,
|
||
int totalChunks,
|
||
String category,
|
||
Frontmatter frontmatter) {
|
||
Map<String, Object> metadata = new HashMap<>();
|
||
|
||
// 文档标识
|
||
String source = firstNonBlank(frontmatter != null ? frontmatter.getSource() : null, "upload:" + docId);
|
||
metadata.put("docId", docId);
|
||
metadata.put("_source", source); // 区分文件索引和上传文档
|
||
metadata.put("source", source);
|
||
|
||
// 分片信息
|
||
metadata.put("chunkIndex", chunk.getChunkIndex());
|
||
metadata.put("totalChunks", totalChunks);
|
||
|
||
// 标题信息
|
||
String title = firstNonBlank(chunk.getTitle(), frontmatter != null ? frontmatter.getTitle() : null);
|
||
if (title != null) {
|
||
metadata.put("title", title);
|
||
}
|
||
|
||
// 面包屑导航（完整标题层级路径）
|
||
String breadcrumb = firstNonBlank(frontmatter != null ? frontmatter.getBreadcrumb() : null, chunk.getBreadcrumb());
|
||
if (breadcrumb != null) {
|
||
metadata.put("breadcrumb", breadcrumb);
|
||
}
|
||
|
||
// 文档类别
|
||
metadata.put("category", category != null && !category.isBlank() ? category : "upload");
|
||
|
||
String kbScope = trimToNull(frontmatter != null ? frontmatter.getKbScope() : null);
|
||
if (kbScope != null) {
|
||
metadata.put("kb_scope", kbScope);
|
||
}
|
||
|
||
return metadata;
|
||
}
|
||
|
||
static String buildEmbeddingText(DocumentChunk chunk) {
|
||
String content = trimToEmpty(chunk.getContent());
|
||
String title = trimToEmpty(chunk.getTitle());
|
||
String breadcrumb = trimToEmpty(chunk.getBreadcrumb());
|
||
|
||
if (title.isEmpty() && breadcrumb.isEmpty()) {
|
||
return content;
|
||
}
|
||
|
||
StringBuilder text = new StringBuilder();
|
||
if (!title.isEmpty()) {
|
||
text.append("Title: ").append(title).append("\n");
|
||
}
|
||
if (!breadcrumb.isEmpty()) {
|
||
text.append("Path: ").append(breadcrumb).append("\n");
|
||
}
|
||
text.append("Content:\n").append(content);
|
||
return text.toString();
|
||
}
|
||
|
||
private static String trimToEmpty(String value) {
|
||
return value == null ? "" : value.trim();
|
||
}
|
||
|
||
private static String trimToNull(String value) {
|
||
if (value == null || value.isBlank()) {
|
||
return null;
|
||
}
|
||
return value.trim();
|
||
}
|
||
|
||
private static String firstNonBlank(String... values) {
|
||
for (String value : values) {
|
||
String trimmed = trimToNull(value);
|
||
if (trimmed != null) {
|
||
return trimmed;
|
||
}
|
||
}
|
||
return null;
|
||
}
|
||
|
||
/**
|
||
* 删除文件的旧数据（根据 metadata._source）
|
||
*/
|
||
private void deleteExistingData(String filePath) {
|
||
try {
|
||
// 使用统一的路径分隔符（正斜杠）用于Milvus存储，避免表达式解析错误
|
||
// 将系统路径转换为统一格式
|
||
Path path = Paths.get(filePath).normalize();
|
||
String normalizedPath = path.toString().replace(File.separator, "/");
|
||
|
||
// 构建删除表达式：metadata["_source"] == "xxx"
|
||
String expr = String.format("metadata[\"_source\"] == \"%s\"", normalizedPath);
|
||
|
||
logger.info("准备删除旧数据，路径: {}, 表达式: {}", normalizedPath, expr);
|
||
|
||
// 确保 collection 已加载（删除操作需要集合已加载）
|
||
R<RpcStatus> loadResponse = milvusClient.loadCollection(
|
||
LoadCollectionParam.newBuilder()
|
||
.withCollectionName(MilvusConstants.MILVUS_COLLECTION_NAME)
|
||
.build()
|
||
);
|
||
|
||
// 状态码 65535 表示集合已经加载，这不是错误
|
||
if (loadResponse.getStatus() != 0 && loadResponse.getStatus() != 65535) {
|
||
logger.warn("加载 collection 失败: {}", loadResponse.getMessage());
|
||
return;
|
||
}
|
||
|
||
DeleteParam deleteParam = DeleteParam.newBuilder()
|
||
.withCollectionName(MilvusConstants.MILVUS_COLLECTION_NAME)
|
||
.withExpr(expr)
|
||
.build();
|
||
|
||
R<MutationResult> response = milvusClient.delete(deleteParam);
|
||
|
||
if (response.getStatus() != 0) {
|
||
logger.warn("删除旧数据时出现警告: {}", response.getMessage());
|
||
} else {
|
||
long deletedCount = response.getData().getDeleteCnt();
|
||
logger.info("✓ 已删除文件的旧数据: {}, 删除记录数: {}", normalizedPath, deletedCount);
|
||
}
|
||
|
||
} catch (Exception e) {
|
||
logger.warn("删除旧数据失败（可能是首次索引）: {}", e.getMessage());
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 构建元数据（包含文件信息）
|
||
*/
|
||
private Map<String, Object> buildMetadata(String filePath, DocumentChunk chunk, int totalChunks) {
|
||
Map<String, Object> metadata = new HashMap<>();
|
||
|
||
// 标准化路径：使用统一的路径分隔符（正斜杠）用于存储，确保跨平台一致性
|
||
Path path = Paths.get(filePath).normalize();
|
||
String normalizedPath = path.toString().replace(File.separator, "/");
|
||
|
||
// 文件信息
|
||
Path fileName = path.getFileName();
|
||
String fileNameStr = fileName != null ? fileName.toString() : "";
|
||
String extension = "";
|
||
int dotIndex = fileNameStr.lastIndexOf('.');
|
||
if (dotIndex > 0) {
|
||
extension = fileNameStr.substring(dotIndex);
|
||
}
|
||
|
||
metadata.put("_source", normalizedPath);
|
||
metadata.put("_extension", extension);
|
||
metadata.put("_file_name", fileNameStr);
|
||
|
||
// 提取类别（从文件路径中提取目录名）
|
||
String category = extractCategory(normalizedPath);
|
||
if (category != null && !category.isEmpty()) {
|
||
metadata.put("category", category);
|
||
}
|
||
|
||
// 分片信息
|
||
metadata.put("chunkIndex", chunk.getChunkIndex());
|
||
metadata.put("totalChunks", totalChunks);
|
||
|
||
// 标题信息
|
||
if (chunk.getTitle() != null && !chunk.getTitle().isEmpty()) {
|
||
metadata.put("title", chunk.getTitle());
|
||
}
|
||
|
||
// 面包屑导航（完整标题层级路径）
|
||
if (chunk.getBreadcrumb() != null && !chunk.getBreadcrumb().isEmpty()) {
|
||
metadata.put("breadcrumb", chunk.getBreadcrumb());
|
||
}
|
||
|
||
return metadata;
|
||
}
|
||
|
||
/**
|
||
* 从文件路径中提取类别
|
||
* 例如：aiops-docs/api/redis-api.md → "api"
|
||
*/
|
||
private String extractCategory(String filePath) {
|
||
try {
|
||
// 标准化路径分隔符
|
||
String normalized = filePath.replace("\\", "/");
|
||
|
||
// 查找 aiops-docs/ 后的第一级目录
|
||
int docsIndex = normalized.indexOf("aiops-docs/");
|
||
if (docsIndex >= 0) {
|
||
String afterDocs = normalized.substring(docsIndex + "aiops-docs/".length());
|
||
int slashIndex = afterDocs.indexOf("/");
|
||
if (slashIndex > 0) {
|
||
return afterDocs.substring(0, slashIndex);
|
||
}
|
||
}
|
||
|
||
// 如果没有 aiops-docs，返回第一级目录
|
||
int firstSlash = normalized.indexOf("/");
|
||
if (firstSlash > 0) {
|
||
return normalized.substring(0, firstSlash);
|
||
}
|
||
|
||
return null;
|
||
} catch (Exception e) {
|
||
logger.warn("提取类别失败，路径: {}", filePath, e);
|
||
return null;
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 插入向量到 Milvus
|
||
*/
|
||
private void insertToMilvus(String content, List<Float> vector,
|
||
Map<String, Object> metadata, int chunkIndex) throws Exception {
|
||
try {
|
||
// 确保 collection 已加载
|
||
R<RpcStatus> loadResponse = milvusClient.loadCollection(
|
||
LoadCollectionParam.newBuilder()
|
||
.withCollectionName(MilvusConstants.MILVUS_COLLECTION_NAME)
|
||
.build()
|
||
);
|
||
|
||
if (loadResponse.getStatus() != 0 && loadResponse.getStatus() != 65535) {
|
||
throw new RuntimeException("加载 collection 失败: " + loadResponse.getMessage());
|
||
}
|
||
|
||
// 生成唯一 ID（使用 _source + 分片索引）
|
||
String source = (String) metadata.get("_source");
|
||
String id = UUID.nameUUIDFromBytes((source + "_" + chunkIndex).getBytes()).toString();
|
||
|
||
// 构建字段数据
|
||
List<InsertParam.Field> fields = new ArrayList<>();
|
||
|
||
// ID 字段
|
||
fields.add(new InsertParam.Field("id", Collections.singletonList(id)));
|
||
|
||
// content 字段
|
||
fields.add(new InsertParam.Field("content", Collections.singletonList(content)));
|
||
|
||
// vector 字段
|
||
fields.add(new InsertParam.Field("vector", Collections.singletonList(vector)));
|
||
|
||
// metadata 字段（JSON 对象）
|
||
com.google.gson.Gson gson = new com.google.gson.Gson();
|
||
com.google.gson.JsonObject metadataJson = gson.toJsonTree(metadata).getAsJsonObject();
|
||
fields.add(new InsertParam.Field("metadata", Collections.singletonList(metadataJson)));
|
||
|
||
// 构建插入参数
|
||
InsertParam insertParam = InsertParam.newBuilder()
|
||
.withCollectionName(MilvusConstants.MILVUS_COLLECTION_NAME)
|
||
.withFields(fields)
|
||
.build();
|
||
|
||
// 执行插入
|
||
R<MutationResult> insertResponse = milvusClient.insert(insertParam);
|
||
|
||
if (insertResponse.getStatus() != 0) {
|
||
throw new RuntimeException("插入向量失败: " + insertResponse.getMessage());
|
||
}
|
||
|
||
logger.debug("向量插入成功: id={}, source={}, chunk={}", id, source, chunkIndex);
|
||
|
||
} catch (Exception e) {
|
||
logger.error("插入向量到 Milvus 失败", e);
|
||
throw e;
|
||
}
|
||
}
|
||
|
||
/**
|
||
* 索引结果类
|
||
*/
|
||
@Getter
|
||
public static class IndexingResult {
|
||
@Setter
|
||
private boolean success;
|
||
@Setter
|
||
private String directoryPath;
|
||
@Setter
|
||
private int totalFiles;
|
||
private int successCount;
|
||
private int failCount;
|
||
@Setter
|
||
private LocalDateTime startTime;
|
||
@Setter
|
||
private LocalDateTime endTime;
|
||
@Setter
|
||
private String errorMessage;
|
||
private Map<String, String> failedFiles = new HashMap<>();
|
||
|
||
public void incrementSuccessCount() {
|
||
this.successCount++;
|
||
}
|
||
|
||
public void incrementFailCount() {
|
||
this.failCount++;
|
||
}
|
||
|
||
public long getDurationMs() {
|
||
if (startTime != null && endTime != null) {
|
||
return java.time.Duration.between(startTime, endTime).toMillis();
|
||
}
|
||
return 0;
|
||
}
|
||
|
||
public void addFailedFile(String filePath, String error) {
|
||
this.failedFiles.put(filePath, error);
|
||
}
|
||
}
|
||
}
|