RAG检索增强系统实战指南：从向量数据库到语义搜索

┌─────────────────────────────────────────────────────────┐
│                     应用层                               │
│  ┌─────────────┐ ┌─────────────┐ ┌─────────────┐        │
│  │  知识库管理  │ │  语义搜索   │ │  AI问答     │        │
│  └──────┬──────┘ └──────┬──────┘ └──────┬──────┘        │
└─────────┼───────────────┼───────────────┼───────────────┘
          │               │               │
          └───────────────┴───────────────┘
                          │
          ┌───────────────▼───────────────┐
│         │      RAG Core Service         │
│  ┌──────┴───────────────────────┬──────┐ │
│  │      Embedding Service       │      │ │
│  │  ┌────────────────────────┐  │      │ │
│  │  │  Text Splitter         │  │      │ │
│  │  │  Embedding Model       │  │      │ │
│  │  └────────────────────────┘  │      │ │
│  └──────────────────────────────┘      │ │
│  ┌──────────────────────────────────┐  │ │
│  │      Vector Database             │  │ │
│  │  ┌──────────┐    ┌──────────┐    │  │ │
│  │  │ Collection│    │  Index   │    │  │ │
│  │  └──────────┘    └──────────┘    │  │ │
│  └──────────────────────────────────┘  │ │
│  ┌──────────────────────────────────┐  │ │
│  │      Retrieval & Ranking         │  │ │
│  │  ┌──────────┐    ┌──────────┐    │  │ │
│  │  │ Similarity│    │ Rerank   │    │  │ │
│  │  └──────────┘    └──────────┘    │  │ │
│  └──────────────────────────────────┘  │ │
└────────────────────────────────────────┘

@Component
public class MilvusSchemaBuilder {
    
    public CollectionSchemaParam buildDocumentSchema() {
        return CollectionSchemaParam.newBuilder()
            // 主键字段 - 使用业务ID（fixId）
            .addFieldType(FieldType.newBuilder()
                .withName("id")
                .withDataType(DataType.VarChar)
                .withPrimaryKey(true)
                .withMaxLength(64)
                .withAutoID(false)
                .build())
            
            // 向量字段（1536维，对应Gemini embedding-001模型）
            .addFieldType(FieldType.newBuilder()
                .withName("embedding")
                .withDataType(DataType.FloatVector)
                .withDimension(1536)
                .build())
            
            // 设备类型（用于分区优化）
            .addFieldType(FieldType.newBuilder()
                .withName("device_type")
                .withDataType(DataType.VarChar)
                .withMaxLength(50)
                .build())
            
            // 设备名称（标题）
            .addFieldType(FieldType.newBuilder()
                .withName("title")
                .withDataType(DataType.VarChar)
                .withMaxLength(500)
                .build())
            
            // 故障描述
            .addFieldType(FieldType.newBuilder()
                .withName("troubleshooting")
                .withDataType(DataType.VarChar)
                .withMaxLength(65535)
                .build())
            
            // 来源URL
            .addFieldType(FieldType.newBuilder()
                .withName("url")
                .withDataType(DataType.VarChar)
                .withMaxLength(500)
                .build())
            
            // 最后更新时间
            .addFieldType(FieldType.newBuilder()
                .withName("last_update_time")
                .withDataType(DataType.Int64)
                .build())
            
            .build();
    }
}

public interface TextSplitter {
    List<String> split(String text);
}

/**
 * 固定长度分块
 */
@Component
public class FixedLengthSplitter implements TextSplitter {
    
    private final int chunkSize = 500;  // 每块500字符
    private final int overlap = 50;     // 重叠50字符
    
    @Override
    public List<String> split(String text) {
        List<String> chunks = new ArrayList<>();
        
        int start = 0;
        while (start < text.length()) {
            int end = Math.min(start + chunkSize, text.length());
            chunks.add(text.substring(start, end));
            start = end - overlap;  // 重叠部分
        }
        
        return chunks;
    }
}

/**
 * 递归字符分块（推荐）
 * 优先按段落、句子分割，保持语义完整
 */
@Component
public class RecursiveCharacterSplitter implements TextSplitter {
    
    private final List<String> separators = Arrays.asList(
        "\n\n",  // 段落
        "\n",    // 换行
        "。",    // 句号
        "！",    // 感叹号
        "？",    // 问号
        " "      // 空格
    );
    
    @Override
    public List<String> split(String text) {
        return splitRecursive(text, 0);
    }
    
    private List<String> splitRecursive(String text, int separatorIndex) {
        List<String> chunks = new ArrayList<>();
        
        if (separatorIndex >= separators.size()) {
            chunks.add(text);
            return chunks;
        }
        
        String separator = separators.get(separatorIndex);
        String[] parts = text.split(separator);
        
        for (String part : parts) {
            if (part.length() > 500) {
                // 仍太长，继续细分
                chunks.addAll(splitRecursive(part, separatorIndex + 1));
            } else {
                chunks.add(part);
            }
        }
        
        return chunks;
    }
}

@Service
public class DocumentIngestionService {
    
    @Autowired
    private TextSplitter textSplitter;
    
    @Autowired
    private AIService aiService;
    
    @Autowired
    private MilvusClient milvusClient;
    
    /**
     * 文档入库流程
     */
    public void ingestDocument(Document document) {
        // 1. 文本预处理
        String cleanedText = preprocess(document.getContent());
        
        // 2. 分块
        List<String> chunks = textSplitter.split(cleanedText);
        
        // 3. 批量向量化（优化：并行处理）
        List<List<Float>> embeddings = chunks.parallelStream()
            .map(chunk -> aiService.embed(EmbeddingRequest.builder()
                .text(chunk)
            .model("gemini-embedding-001")
                .build()))
            .map(response -> response.getEmbedding())
            .collect(Collectors.toList());
        
        // 4. 构建插入数据
        List<InsertParam.Field> fields = Arrays.asList(
            new InsertParam.Field("id", 
                chunks.stream().map(c -> document.getId() + "_" + c.hashCode()).collect(Collectors.toList())),
            new InsertParam.Field("embedding", embeddings),
            new InsertParam.Field("device_type", 
                chunks.stream().map(c -> document.getDeviceType()).collect(Collectors.toList())),
            new InsertParam.Field("title", 
                chunks.stream().map(c -> document.getTitle()).collect(Collectors.toList())),
            new InsertParam.Field("troubleshooting", chunks),
            new InsertParam.Field("url", 
                chunks.stream().map(c -> document.getUrl()).collect(Collectors.toList())),
            new InsertParam.Field("last_update_time", 
                chunks.stream().map(c -> System.currentTimeMillis()).collect(Collectors.toList()))
        );
        
        // 5. 批量插入Milvus
        InsertParam insertParam = InsertParam.newBuilder()
            .withCollectionName("ifixit_vectors")
            .withFields(fields)
            .build();
        
        R<MutationResult> result = milvusClient.insert(insertParam);
        
        if (result.getStatus() != R.Status.Success.getCode()) {
            throw new RuntimeException("插入失败: " + result.getMessage());
        }
        
        log.info("文档入库成功，docId={}, chunks={}", document.getId(), chunks.size());
    }
    
    /**
     * 批量入库（大数据量优化）
     */
    public void batchIngest(List<Document> documents) {
        // 分批处理，每批1000条
        List<List<Document>> batches = Lists.partition(documents, 1000);
        
        for (List<Document> batch : batches) {
            batch.parallelStream().forEach(this::ingestDocument);
            
            // 防止内存溢出，主动GC
            System.gc();
        }
    }
}

@Service
public class LargeScaleVectorizationService {
    
    @Autowired
    private ThreadPoolExecutor taskExecutor;
    
    /**
     * 流式处理，避免OOM
     */
    public void streamIngest(Stream<Document> documentStream) {
        documentStream
            .parallel()
            .map(this::processDocument)
            .filter(Objects::nonNull)
            .collect(Collectors.groupingByConcurrent(
                chunk -> chunk.hashCode() % 100,  // 分100组
                Collectors.toList()
            ))
            .values()
            .forEach(this::batchInsert);
    }
    
    /**
     * 异步批量向量化
     */
    @Async("taskExecutor")
    public CompletableFuture<Void> asyncVectorize(List<Document> documents) {
        return CompletableFuture.runAsync(() -> {
            documents.forEach(doc -> {
                try {
                    ingestDocument(doc);
                } catch (Exception e) {
                    log.error("文档向量化失败，docId={}", doc.getId(), e);
                    // 记录失败，后续补偿
                    saveFailedDoc(doc);
                }
            });
        }, taskExecutor);
    }
    
    /**
     * 定时补偿失败任务
     */
    @Scheduled(fixedDelay = 300000) // 每5分钟
    public void compensateFailedDocs() {
        List<Document> failedDocs = getFailedDocs();
        if (!failedDocs.isEmpty()) {
            log.info("补偿处理{}条失败文档", failedDocs.size());
            failedDocs.forEach(this::ingestDocument);
        }
    }
}

@Service
public class VectorSearchService {
    
    @Autowired
    private AIService aiService;
    
    @Autowired
    private MilvusClient milvusClient;
    
    /**
     * 语义搜索
     */
    public List<SearchResult> search(String query, int topK) {
        // 1. Query向量化
        EmbeddingResponse embedding = aiService.embed(EmbeddingRequest.builder()
            .text(query)
            .model("gemini-embedding-001")
            .build());
        
        // 2. 向量检索
        SearchParam searchParam = SearchParam.newBuilder()
            .withCollectionName("ifixit_vectors")
            .withVectors(Collections.singletonList(embedding.getEmbedding()))
            .withVectorFieldName("embedding")
            .withTopK(topK)
            .withMetricType(MetricType.COSINE)  // 余弦相似度
            .withParams("{\"nprobe\": 10}")      // IVF_FLAT搜索参数，探测10个簇
            .build();
        
        R<SearchResults> result = milvusClient.search(searchParam);
        
        if (result.getStatus() != R.Status.Success.getCode()) {
            throw new RuntimeException("搜索失败: " + result.getMessage());
        }
        
        // 3. 解析结果
        return parseSearchResults(result.getData());
    }
}

@Service
public class HybridSearchService {
    
    /**
     * 混合检索
     */
    public List<SearchResult> hybridSearch(String query, int topK) {
        // 1. 向量检索
        List<SearchResult> vectorResults = vectorSearchService.search(query, topK * 2);
        
        // 2. 关键词检索（BM25）
        List<SearchResult> keywordResults = keywordSearchService.search(query, topK * 2);
        
        // 3. 结果融合（RRF - Reciprocal Rank Fusion）
        return reciprocalRankFusion(vectorResults, keywordResults, topK);
    }
    
    /**
     * RRF算法融合结果
     */
    private List<SearchResult> reciprocalRankFusion(
            List<SearchResult> vectorResults, 
            List<SearchResult> keywordResults, 
            int topK) {
        
        Map<String, Double> scores = new HashMap<>();
        int k = 60; // RRF常数
        
        // 向量检索结果打分
        for (int i = 0; i < vectorResults.size(); i++) {
            String docId = vectorResults.get(i).getDocId();
            double score = 1.0 / (k + i + 1);
            scores.merge(docId, score, Double::sum);
        }
        
        // 关键词检索结果打分
        for (int i = 0; i < keywordResults.size(); i++) {
            String docId = keywordResults.get(i).getDocId();
            double score = 1.0 / (k + i + 1);
            scores.merge(docId, score, Double::sum);
        }
        
        // 排序取TopK
        return scores.entrySet().stream()
            .sorted(Map.Entry.<String, Double>comparingByValue().reversed())
            .limit(topK)
            .map(entry -> {
                SearchResult result = new SearchResult();
                result.setDocId(entry.getKey());
                result.setScore(entry.getValue());
                return result;
            })
            .collect(Collectors.toList());
    }
}

@Service
public class RerankService {
    
    @Autowired
    private AIService aiService;
    
    /**
     * 重排序
     */
    public List<SearchResult> rerank(String query, List<SearchResult> candidates) {
        // 构建交叉编码器输入
        List<RerankRequest> requests = candidates.stream()
            .map(result -> RerankRequest.builder()
                .query(query)
                .document(result.getContent())
                .build())
            .collect(Collectors.toList());
        
        // 调用重排序模型（如BGE-Reranker）
        List<RerankResponse> responses = aiService.rerankBatch(requests);
        
        // 按重排序分数排序
        return IntStream.range(0, candidates.size())
            .mapToObj(i -> {
                SearchResult result = candidates.get(i);
                result.setRerankScore(responses.get(i).getScore());
                return result;
            })
            .sorted(Comparator.comparing(SearchResult::getRerankScore).reversed())
            .collect(Collectors.toList());
    }
}

@Component
public class RAGPromptBuilder {
    
    private final String template = """
        你是一个专业的AI助手。请基于以下参考资料回答用户问题。
        
        参考资料：
        {% for doc in documents %}
        [{{ loop.index }}] {{ doc.content }}
        来源：{{ doc.source }}
        {% endfor %}
        
        用户问题：{{ query }}
        
        要求：
        1. 基于参考资料回答，不要编造信息
        2. 如果参考资料不足以回答问题，请明确说明
        3. 回答要简洁明了，突出重点
        4. 如果可能，请指出信息来源的编号
        
        请回答：
        """;
    
    public String buildPrompt(String query, List<Document> documents) {
        Map<String, Object> variables = new HashMap<>();
        variables.put("query", query);
        variables.put("documents", documents);
        
        return renderTemplate(template, variables);
    }
}

@Service
public class RAGService {
    
    @Autowired
    private HybridSearchService searchService;
    
    @Autowired
    private RerankService rerankService;
    
    @Autowired
    private RAGPromptBuilder promptBuilder;
    
    @Autowired
    private AIService aiService;
    
    /**
     * 完整的RAG流程
     */
    public RAGResponse answer(String query) {
        // 1. 混合检索（召回）
        List<SearchResult> candidates = searchService.hybridSearch(query, 20);
        
        // 2. 重排序（精排）
        List<SearchResult> rankedResults = rerankService.rerank(query, candidates);
        
        // 3. 取Top5作为上下文
        List<Document> contextDocs = rankedResults.stream()
            .limit(5)
            .map(this::getDocumentById)
            .collect(Collectors.toList());
        
        // 4. 构建Prompt
        String prompt = promptBuilder.buildPrompt(query, contextDocs);
        
        // 5. 调用LLM生成回答
        ChatResponse response = aiService.chat(ChatRequest.builder()
            .message(ChatMessage.user(prompt))
            .temperature(0.3)  // 降低随机性
            .maxTokens(2000)
            .build());
        
        // 6. 返回结果
        return RAGResponse.builder()
            .answer(response.getContent())
            .sources(contextDocs)
            .confidence(calculateConfidence(rankedResults))
            .build();
    }
}

/**
 * 创建IVF_FLAT索引（当前使用）
 * 适合有明显聚类特征的数据，内存友好
 */
public void createIVFIndex() {
    CreateIndexParam indexParam = CreateIndexParam.newBuilder()
        .withCollectionName("ifixit_vectors")
        .withFieldName("embedding")
        .withIndexType(IndexType.IVF_FLAT)      // 倒排文件索引
        .withMetricType(MetricType.COSINE)      // 余弦相似度
        .withExtraParam("{\"nlist\": 2816}")     // 聚类中心数，根据数据量调整
        .build();
    
    milvusClient.createIndex(indexParam);
}

/**
 * 创建HNSW索引（未来升级选项）
 * 适合高维向量，查询速度更快但内存占用更高
 */
public void createHNSWIndex() {
    CreateIndexParam indexParam = CreateIndexParam.newBuilder()
        .withCollectionName("ifixit_vectors")
        .withFieldName("embedding")
        .withIndexType(IndexType.HNSW)
        .withMetricType(MetricType.COSINE)
        .withExtraParam(Map.of(
            "M", 16,        // 每个节点的邻居数
            "efConstruction", 200  // 构建时的搜索深度
        ))
        .build();
    
    milvusClient.createIndex(indexParam);
}

/**
 * 按文档类型分区（可选优化）
 * 实际项目中可以根据设备类型创建分区，如手机、电脑、家电等
 */
public void createPartitions() {
    String[] deviceTypes = {"smartphone", "laptop", "appliance", "wearable"};
    
    for (String type : deviceTypes) {
        CreatePartitionParam partitionParam = CreatePartitionParam.newBuilder()
            .withCollectionName("ifixit_vectors")
            .withPartitionName("type_" + type)
            .build();
        
        milvusClient.createPartition(partitionParam);
    }
}

/**
 * 搜索指定分区
 */
public List<SearchResult> searchInPartition(String query, String deviceType) {
    SearchParam searchParam = SearchParam.newBuilder()
        .withCollectionName("ifixit_vectors")
        .withPartitionNames(Collections.singletonList("type_" + deviceType))
        // ... 其他参数
        .build();
    
    return milvusClient.search(searchParam);
}

@Component
public class RAGEvaluator {
    
    /**
     * 评估指标
     */
    public EvaluationMetrics evaluate(List<TestCase> testCases) {
        int hitCount = 0;
        double totalPrecision = 0;
        double totalRecall = 0;
        
        for (TestCase testCase : testCases) {
            // 执行检索
            List<SearchResult> results = ragService.search(testCase.getQuery(), 10);
            Set<String> retrievedIds = results.stream()
                .map(SearchResult::getDocId)
                .collect(Collectors.toSet());
            
            Set<String> relevantIds = testCase.getRelevantDocIds();
            
            // 计算Precision@K
            int relevantRetrieved = (int) retrievedIds.stream()
                .filter(relevantIds::contains)
                .count();
            double precision = (double) relevantRetrieved / retrievedIds.size();
            
            // 计算Recall
            double recall = (double) relevantRetrieved / relevantIds.size();
            
            // Hit@1
            if (!results.isEmpty() && relevantIds.contains(results.get(0).getDocId())) {
                hitCount++;
            }
            
            totalPrecision += precision;
            totalRecall += recall;
        }
        
        return EvaluationMetrics.builder()
            .averagePrecision(totalPrecision / testCases.size())
            .averageRecall(totalRecall / testCases.size())
            .hitAt1((double) hitCount / testCases.size())
            .build();
    }
}

@Component
public class RAGMetrics {
    
    private final MeterRegistry meterRegistry;
    
    public void recordSearch(String query, int resultCount, long latency) {
        // 记录检索延迟
        meterRegistry.timer("rag.search.latency")
            .record(latency, TimeUnit.MILLISECONDS);
        
        // 记录结果数量
        meterRegistry.distributionSummary("rag.search.result.count")
            .record(resultCount);
        
        // 记录Query长度分布
        meterRegistry.distributionSummary("rag.query.length")
            .record(query.length());
    }
    
    public void recordAnswerQuality(String query, double score) {
        meterRegistry.distributionSummary("rag.answer.quality.score")
            .record(score);
    }
}