RAG问答系统：Spring Boot + ChromaDB 知识库检索实战

一、系统架构设计

二、核心组件实现

1. 依赖配置

<!-- pom.xml -->
<dependencies>
    <!-- ChromaDB Java客户端 -->
    <dependency>
        <groupId>io.chroma</groupId>
        <artifactId>chromadb-client</artifactId>
        <version>0.4.0</version>
    </dependency>
    
    <!-- 嵌入模型 -->
    <dependency>
        <groupId>ai.djl</groupId>
        <artifactId>model-zoo</artifactId>
        <version>0.22.0</version>
    </dependency>
    
    <!-- OpenAI集成 -->
    <dependency>
        <groupId>com.theokanning.openai-java</groupId>
        <artifactId>service</artifactId>
        <version>0.14.0</version>
    </dependency>
</dependencies>

2. ChromaDB配置

@Configuration
public class ChromaConfig {
    
    @Value("${chromadb.host:localhost}")
    private String host;
    
    @Value("${chromadb.port:8000}")
    private int port;
    
    @Bean
    public ChromaClient chromaClient() {
        return new ChromaClient(host, port);
    }
    
    @Bean
    public Collection knowledgeCollection(ChromaClient client) {
        return client.getOrCreateCollection("knowledge-base", 
            CollectionSpec.builder()
                .dimension(384) // 嵌入维度
                .build());
    }
}

三、知识库处理流水线

1. 文档切分服务

@Service
public class DocumentChunker {
    
    private static final int CHUNK_SIZE = 512;
    private static final int OVERLAP = 50;
    
    public List<TextChunk> chunkDocument(String content) {
        List<TextChunk> chunks = new ArrayList<>();
        int start = 0;
        
        while (start < content.length()) {
            int end = Math.min(start + CHUNK_SIZE, content.length());
            String chunkText = content.substring(start, end);
            
            chunks.add(new TextChunk(
                chunkText,
                start,
                end
            ));
            
            start = end - OVERLAP;
        }
        
        return chunks;
    }
    
    @Data
    @AllArgsConstructor
    public static class TextChunk {
        private String content;
        private int start;
        private int end;
    }
}

2. 向量化服务

@Service
public class EmbeddingService {
    
    private final ZooModel<String, float[]> embeddingModel;
    
    public EmbeddingService() throws ModelException, IOException {
        this.embeddingModel = ModelZoo.loadModel(
            new Criteria.Builder()
                .setTypes(String.class, float[].class)
                .optEngine("PyTorch")
                .optModelUrls("djl://ai.djl.huggingface.pytorch/sentence-transformers/all-MiniLM-L6-v2")
                .build()
        );
    }
    
    public float[] embed(String text) {
        try (Predictor<String, float[]> predictor = embeddingModel.newPredictor()) {
            return predictor.predict(text);
        }
    }
    
    public List<float[]> batchEmbed(List<String> texts) {
        return texts.stream()
            .parallel()
            .map(this::embed)
            .collect(Collectors.toList());
    }
}

3. 知识库索引服务

@Service
@RequiredArgsConstructor
public class KnowledgeIndexer {
    
    private final DocumentChunker chunker;
    private final EmbeddingService embeddingService;
    private final Collection collection;
    
    public void indexDocument(String docId, String content) {
        List<TextChunk> chunks = chunker.chunkDocument(content);
        List<String> texts = chunks.stream()
            .map(TextChunk::getContent)
            .collect(Collectors.toList());
        
        List<float[]> embeddings = embeddingService.batchEmbed(texts);
        
        List<String> ids = chunks.stream()
            .map(chunk -> docId + "_" + chunk.getStart())
            .collect(Collectors.toList());
        
        collection.add(
            ids,
            embeddings,
            texts.stream().map(text -> 
                Metadata.of("doc_id", docId)
            ).collect(Collectors.toList()),
            texts
        );
    }
}

四、检索增强生成核心

1. 检索服务

@Service
@RequiredArgsConstructor
public class RetrieverService {
    
    private final Collection collection;
    private final EmbeddingService embeddingService;
    
    public List<RetrievalResult> retrieve(String query, int topK) {
        float[] queryEmbedding = embeddingService.embed(query);
        
        QueryResult result = collection.query()
            .queryEmbeddings(List.of(queryEmbedding))
            .nResults(topK)
            .execute();
        
        return IntStream.range(0, result.getIds().get(0).size())
            .mapToObj(i -> new RetrievalResult(
                result.getIds().get(0).get(i),
                result.getDistances().get(0).get(i),
                result.getDocuments().get(0).get(i),
                result.getMetadatas().get(0).get(i)
            ))
            .collect(Collectors.toList());
    }
    
    @Data
    @AllArgsConstructor
    public static class RetrievalResult {
        private String id;
        private float score;
        private String content;
        private Map<String, String> metadata;
    }
}

2. 提示工程

public class PromptBuilder {
    
    public static String buildRAGPrompt(String question, List<String> contexts) {
        StringBuilder sb = new StringBuilder();
        sb.append("基于以下上下文信息回答问题。如果上下文不包含答案，请回答'我不知道'。\n\n");
        
        sb.append("上下文：\n");
        for (int i = 0; i < contexts.size(); i++) {
            sb.append(String.format("[片段%d]: %s\n\n", i+1, contexts.get(i)));
        }
        
        sb.append("\n问题：").append(question).append("\n");
        sb.append("答案：");
        
        return sb.toString();
    }
}

3. 生成服务

@Service
public class GenerationService {
    
    private final OpenAiService openAiService;
    
    public GenerationService(@Value("${openai.api-key}") String apiKey) {
        this.openAiService = new OpenAiService(apiKey, Duration.ofSeconds(30));
    }
    
    public String generateAnswer(String prompt) {
        ChatCompletionRequest request = ChatCompletionRequest.builder()
            .model("gpt-3.5-turbo")
            .messages(List.of(
                new ChatMessage("user", prompt)
            ))
            .maxTokens(500)
            .temperature(0.3)
            .build();
        
        return openAiService.createChatCompletion(request)
            .getChoices().get(0).getMessage().getContent();
    }
}

五、REST API设计

1. 问答端点

@RestController
@RequestMapping("/api/rag")
@RequiredArgsConstructor
public class RAGController {
    
    private final RetrieverService retrieverService;
    private final GenerationService generationService;
    
    @PostMapping("/ask")
    public ResponseEntity<RAGResponse> askQuestion(@RequestBody QuestionRequest request) {
        // 检索相关上下文
        List<RetrieverService.RetrievalResult> results = 
            retrieverService.retrieve(request.getQuestion(), 3);
        
        // 构建提示
        List<String> contexts = results.stream()
            .map(RetrieverService.RetrievalResult::getContent)
            .collect(Collectors.toList());
        
        String prompt = PromptBuilder.buildRAGPrompt(request.getQuestion(), contexts);
        
        // 生成答案
        String answer = generationService.generateAnswer(prompt);
        
        // 构建响应
        RAGResponse response = new RAGResponse();
        response.setAnswer(answer);
        response.setContexts(contexts);
        response.setSources(results.stream()
            .map(r -> r.getMetadata().get("doc_id"))
            .distinct()
            .collect(Collectors.toList()));
        
        return ResponseEntity.ok(response);
    }
    
    @Data
    public static class QuestionRequest {
        @NotBlank
        private String question;
    }
    
    @Data
    public static class RAGResponse {
        private String answer;
        private List<String> contexts;
        private List<String> sources;
    }
}

六、高级检索策略

1. 混合检索

public List<RetrievalResult> hybridRetrieve(String query, int topK) {
    // 向量检索
    List<RetrievalResult> vectorResults = retrieverService.retrieve(query, topK * 2);
    
    // 关键词检索
    List<RetrievalResult> keywordResults = keywordSearch(query, topK * 2);
    
    // 结果融合与重排序
    return fuseResults(vectorResults, keywordResults, topK);
}

private List<RetrievalResult> fuseResults(
    List<RetrievalResult> list1, 
    List<RetrievalResult> list2, 
    int topK
) {
    // 使用RRF(Reciprocal Rank Fusion)算法
    Map<String, RetrievalResult> fused = new HashMap<>();
    
    // 合并结果
    fuseList(list1, fused, 1);
    fuseList(list2, fused, 1);
    
    // 按分数排序
    return fused.values().stream()
        .sorted(Comparator.comparingDouble(RetrievalResult::getScore).reversed())
        .limit(topK)
        .collect(Collectors.toList());
}

private void fuseList(List<RetrievalResult> list, Map<String, RetrievalResult> fused, int k) {
    for (int i = 0; i < list.size(); i++) {
        RetrievalResult result = list.get(i);
        double rrfScore = 1.0 / (k + i);
        
        RetrievalResult existing = fused.get(result.getId());
        if (existing != null) {
            existing.setScore(existing.getScore() + rrfScore);
        } else {
            result.setScore(rrfScore);
            fused.put(result.getId(), result);
        }
    }
}

2. 查询扩展

public String expandQuery(String originalQuery) {
    // 使用LLM扩展查询
    String prompt = "生成3个与以下问题相关的查询：\n" + originalQuery;
    
    String expansion = generationService.generateAnswer(prompt);
    
    // 解析生成的查询
    List<String> queries = parseExpansion(expansion);
    queries.add(0, originalQuery);
    
    return String.join(" ", queries);
}

private List<String> parseExpansion(String expansion) {
    // 解析格式如：1. 查询1\n2. 查询2\n3. 查询3
    return Arrays.stream(expansion.split("\n"))
        .map(line -> line.replaceAll("^\\d+\\.\\s*", ""))
        .collect(Collectors.toList());
}

七、性能优化方案

1. 缓存策略

@Cacheable(value = "retrievalCache", key = "#query.hashCode()")
public List<RetrievalResult> cachedRetrieve(String query, int topK) {
    return retrieverService.retrieve(query, topK);
}

@Cacheable(value = "generationCache", key = "{#prompt.hashCode()}")
public String cachedGenerate(String prompt) {
    return generationService.generateAnswer(prompt);
}

2. 异步处理

@Async
public CompletableFuture<List<RetrievalResult>> retrieveAsync(String query, int topK) {
    return CompletableFuture.completedFuture(retrieverService.retrieve(query, topK));
}

@Async
public CompletableFuture<String> generateAsync(String prompt) {
    return CompletableFuture.completedFuture(generationService.generateAnswer(prompt));
}

// 控制器中使用
@PostMapping("/ask-async")
public CompletableFuture<ResponseEntity<RAGResponse>> askQuestionAsync(@RequestBody QuestionRequest request) {
    return retrieveAsync(request.getQuestion(), 3)
        .thenCompose(results -> {
            List<String> contexts = results.stream()
                .map(RetrieverService.RetrievalResult::getContent)
                .collect(Collectors.toList());
            
            String prompt = PromptBuilder.buildRAGPrompt(request.getQuestion(), contexts);
            return generateAsync(prompt)
                .thenApply(answer -> {
                    RAGResponse response = new RAGResponse();
                    response.setAnswer(answer);
                    response.setContexts(contexts);
                    return ResponseEntity.ok(response);
                });
        });
}

八、生产部署方案

1. Docker Compose部署

version: '3.8'

services:
  chromadb:
    image: chromadb/chroma
    ports:
      - "8000:8000"
    volumes:
      - chroma-data:/chroma/chroma
    
  rag-service:
    build: .
    ports:
      - "8080:8080"
    environment:
      - CHROMADB_HOST=chromadb
      - OPENAI_API_KEY=${OPENAI_API_KEY}
    depends_on:
      - chromadb

volumes:
  chroma-data:

2. Kubernetes部署

apiVersion: apps/v1
kind: Deployment
metadata:
  name: chromadb
spec:
  replicas: 1
  selector:
    matchLabels:
      app: chromadb
  template:
    metadata:
      labels:
        app: chromadb
    spec:
      containers:
      - name: chromadb
        image: chromadb/chroma
        ports:
        - containerPort: 8000
        volumeMounts:
        - name: chroma-data
          mountPath: /chroma/chroma
      volumes:
      - name: chroma-data
        persistentVolumeClaim:
          claimName: chroma-pvc

---
apiVersion: apps/v1
kind: Deployment
metadata:
  name: rag-service
spec:
  replicas: 3
  selector:
    matchLabels:
      app: rag-service
  template:
    metadata:
      labels:
        app: rag-service
    spec:
      containers:
      - name: rag-service
        image: rag-service:1.0
        ports:
        - containerPort: 8080
        env:
        - name: CHROMADB_HOST
          value: "chromadb"
        - name: OPENAI_API_KEY
          valueFrom:
            secretKeyRef:
              name: openai-secret
              key: api-key

九、监控与评估

1. 评估指标

public class EvaluationService {
    
    public RAGEvaluation evaluate(List<QAExample> examples) {
        RAGEvaluation evaluation = new RAGEvaluation();
        
        for (QAExample example : examples) {
            RAGResponse response = askQuestion(example.getQuestion());
            
            // 计算答案相似度
            double similarity = calculateSimilarity(example.getExpectedAnswer(), response.getAnswer());
            
            // 计算检索召回率
            double recall = calculateRecall(example.getExpectedContexts(), response.getContexts());
            
            evaluation.addResult(similarity, recall);
        }
        
        return evaluation;
    }
    
    private double calculateSimilarity(String expected, String actual) {
        // 使用BERTScore或ROUGE
        return bertScore.score(expected, actual);
    }
    
    private double calculateRecall(List<String> expected, List<String> actual) {
        Set<String> expectedSet = new HashSet<>(expected);
        Set<String> actualSet = new HashSet<>(actual);
        
        Set<String> intersection = new HashSet<>(expectedSet);
        intersection.retainAll(actualSet);
        
        return (double) intersection.size() / expectedSet.size();
    }
}

2. Prometheus监控

@Bean
MeterRegistryCustomizer<MeterRegistry> metrics() {
    return registry -> {
        Timer.builder("rag.retrieval.time")
            .register(registry);
        
        Timer.builder("rag.generation.time")
            .register(registry);
        
        Counter.builder("rag.requests")
            .register(registry);
    };
}

@Aspect
@Component
public class MonitoringAspect {
    
    @Around("execution(* RetrieverService.retrieve(..))")
    public Object timeRetrieval(ProceedingJoinPoint pjp) throws Throwable {
        Timer.Sample sample = Timer.start();
        Object result = pjp.proceed();
        sample.stop(Metrics.timer("rag.retrieval.time"));
        return result;
    }
    
    @Around("execution(* GenerationService.generateAnswer(..))")
    public Object timeGeneration(ProceedingJoinPoint pjp) throws Throwable {
        Timer.Sample sample = Timer.start();
        Object result = pjp.proceed();
        sample.stop(Metrics.timer("rag.generation.time"));
        return result;
    }
}

十、安全增强措施

1. 输入过滤

@Aspect
@Component
public class InputValidationAspect {
    
    @Before("execution(* RAGController.askQuestion(..)) && args(request)")
    public void validateInput(QuestionRequest request) {
        if (containsMaliciousContent(request.getQuestion())) {
            throw new SecurityException("检测到恶意输入");
        }
    }
    
    private boolean containsMaliciousContent(String text) {
        // 检测注入攻击、恶意指令等
        return text.contains("DROP TABLE") || 
               text.contains("<script>") ||
               text.contains("sudo rm -rf");
    }
}

2. 内容审核

public String safeGenerate(String prompt) {
    String answer = generationService.generateAnswer(prompt);
    
    if (isUnsafeContent(answer)) {
        return "抱歉，我无法回答这个问题";
    }
    
    return answer;
}

private boolean isUnsafeContent(String text) {
    // 使用敏感词库或AI内容审核
    return moderationService.moderate(text).isFlagged();
}

十一、应用场景扩展

1. 多语言支持

public String translateToEnglish(String query) {
    if (isEnglish(query)) return query;
    
    String prompt = "将以下文本翻译为英文：" + query;
    return generationService.generateAnswer(prompt);
}

private boolean isEnglish(String text) {
    // 简单检测
    return text.matches(".*[a-zA-Z].*");
}

// 在检索前翻译
public List<RetrievalResult> multilingualRetrieve(String query, int topK) {
    String englishQuery = translateToEnglish(query);
    return retrieverService.retrieve(englishQuery, topK);
}

2. 领域适配

public void configureForDomain(String domain) {
    // 加载领域特定知识库
    collection = chromaClient.getCollection("knowledge_" + domain);
    
    // 加载领域特定嵌入模型
    embeddingModel = loadDomainEmbeddingModel(domain);
    
    // 配置领域特定提示模板
    promptTemplate = loadPromptTemplate(domain);
}

十二、性能压测数据

测试环境

性能指标

总结：RAG系统优势

1. 知识实时更新：无需重新训练模型，更新知识库即可
2. 来源可追溯：提供答案来源文档片段
3. 减少幻觉：基于事实知识生成答案
4. 领域适应性强：快速适配不同行业知识库
5. 成本效益：比微调大模型成本低90%
   典型应用场景：

• 企业知识问答系统
• 智能客服助手
• 教育领域智能辅导
• 医疗诊断辅助
• 法律条文查询

最佳实践建议：
1. 知识库文档需预处理（清洗、结构化）
2. 关键业务问题设置人工审核流程
3. 定期评估和优化检索效果
4. 敏感领域增加本地LLM支持