Java构建RAG系统全攻略

1. 知识库准备与索引

• 数据源： 你需要一个结构化的知识库，可以是文本文件、数据库记录、维基百科 dump 等。
• 文档处理：
  • 加载： 使用 Java IO/NIO 或数据库连接库读取文档。
  • 清洗： 移除无关字符、标准化格式。
  • 分块： 将大文档分割成更小的、语义上有意义的片段（chunks）。例如，按段落或固定长度分割。这有助于后续的精确检索。
  • 向量化： 这是关键步骤。为每个文本块生成一个数值向量（嵌入向量），该向量能捕捉其语义信息。你需要使用一个嵌入模型。
    • Java 嵌入模型库： 可以使用如 Deep Java Library 加载预训练的 Sentence Transformers 模型（如 all-MiniLM-L6-v2）来计算文本嵌入向量。

    // 示例伪代码 - 使用DJL计算嵌入向量
    import ai.djl.*;
    import ai.djl.inference.*;
    import ai.djl.modality.*;
    import ai.djl.translate.*;
    
    Translator<String, float[]> translator = ... // 定义字符串到浮点数数组的转换器
    try (Predictor<String, float[]> predictor = model.newPredictor(translator)) {
        float[] embedding = predictor.predict("This is a text chunk.");
        // 保存 embedding 和对应的文本块
    }
    

• 向量索引： 将所有文档块的嵌入向量存储在一个高效的向量数据库中，以便进行快速的近似最近邻搜索。
  • Java 向量数据库选项：
    • Lucene (近似向量搜索)： Apache Lucene 支持通过 KnnVectorField 进行近似 KNN 搜索。你需要将向量存储为这种类型的字段。

    // 示例伪代码 - 使用Lucene索引向量
    Document doc = new Document();
    doc.add(new KnnVectorField("vector", embedding, VectorSimilarityFunction.COSINE));
    doc.add(new StoredField("text", chunkText));
    indexWriter.addDocument(doc);
    

    • 专用向量数据库客户端： 如 Milvus、Qdrant、Pinecone 等通常提供 Java SDK。

    // 示例伪代码 - 使用Milvus Java SDK插入向量
    milvusClient.insert(
        InsertParam.newBuilder()
            .withCollectionName("knowledge_base")
            .addFieldValue("vector", embedding)
            .addFieldValue("text", chunkText)
            .build()
    );
    

2. 检索阶段

• 用户查询： 接收用户的输入问题/提示。
• 查询向量化： 使用与知识库文档块相同的嵌入模型，将用户查询也转换为一个嵌入向量。

  float[] queryEmbedding = predictor.predict(userQuery);
  

• 相似度搜索： 在向量索引/数据库中，执行近似最近邻搜索，查找与查询向量最相似的 K 个文档块向量。
  • Lucene：

    // 示例伪代码 - Lucene KNN搜索
    KnnVectorQuery query = new KnnVectorQuery("vector", queryEmbedding, k);
    TopDocs topDocs = searcher.search(query, k);
    for (ScoreDoc scoreDoc : topDocs.scoreDocs) {
        Document hitDoc = searcher.doc(scoreDoc.doc);
        String retrievedText = hitDoc.get("text");
        // 收集检索到的文本块
    }
    

  • 专用向量数据库：

    // 示例伪代码 - Milvus搜索
    SearchParam searchParam = SearchParam.newBuilder()
        .withCollectionName("knowledge_base")
        .withFloatVectorFieldName("vector")
        .withFloatVectors(Arrays.asList(queryEmbedding))
        .withTopK(k)
        .build();
    SearchResults results = milvusClient.search(searchParam);
    for (QueryResult result : results.getResult()) {
        String retrievedText = result.getFieldValue("text", String.class);
        // 收集检索到的文本块
    }
    

• 结果聚合： 将检索到的 Top K 个文本块（上下文）组合起来，准备输入给 LLM。

3. 生成阶段

• 构造提示： 将用户的原始查询（query）和检索到的相关上下文（context）按照特定的模板组合成一个新的提示（prompt），输入给 LLM。例如：

  Answer the question based on the following context:
  Context: {retrieved_context_1} ... {retrieved_context_k}
  Question: {user_query}
  Answer:
  

  String promptTemplate = "Answer the question based on the following context:\nContext: %s\nQuestion: %s\nAnswer:";
  String context = String.join("\n", retrievedTexts); // 合并检索到的文本块
  String finalPrompt = String.format(promptTemplate, context, userQuery);
  

• 调用 LLM： 使用 Java 调用 LLM 的 API 或本地运行的模型来生成文本。
  • API 调用 (如 OpenAI)：

    // 示例伪代码 - 调用OpenAI API (使用HttpClient)
    HttpRequest request = HttpRequest.newBuilder()
        .uri(URI.create("https://api.openai.com/v1/chat/completions"))
        .header("Authorization", "Bearer YOUR_API_KEY")
        .header("Content-Type", "application/json")
        .POST(HttpRequest.BodyPublishers.ofString(
            String.format("{\"model\": \"gpt-3.5-turbo\", \"messages\": [{\"role\": \"user\", \"content\": \"%s\"}]}", finalPrompt)
        ))
        .build();
    HttpResponse<String> response = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
    // 解析response获取生成的文本 (response.body())
    

  • 本地模型 (使用 DJL)： 类似嵌入模型调用，但使用文本生成模型（如 GPT、Llama）的 Predictor 和 Translator。

    try (Predictor<String, String> generator = generatorModel.newPredictor(translator)) {
        String generatedText = generator.predict(finalPrompt);
    }
    

• 输出： 将 LLM 生成的文本作为最终答案返回给用户。

4. 整体流程整合

将上述步骤串联起来：

// 1. 用户输入
String userQuery = "Java中如何实现多线程？";

// 2. 检索阶段
float[] queryEmbedding = embeddingPredictor.predict(userQuery);
List<String> retrievedContexts = vectorIndex.searchTopK(queryEmbedding, 5); // 假设搜索Top 5

// 3. 构造提示
String context = String.join("\n\n", retrievedContexts);
String prompt = String.format("基于以下上下文回答问题：\n%s\n\n问题：%s\n回答：", context, userQuery);

// 4. 生成阶段
String answer = llmGenerator.generate(prompt); // 调用LLM生成

// 5. 输出结果
System.out.println("回答: " + answer);

依赖库概览

• 文本处理/IO： Java 标准库。
• 向量计算/嵌入模型： Deep Java Library (DJL)。
• 向量索引：
  • 内置选项：Apache Lucene。
  • 外部服务：Milvus, Qdrant, Pinecone 等的 Java SDK。
• LLM 调用：
  • API 调用：Java HttpClient (标准库或 Apache HttpClient, OkHttp)。
  • 本地模型：Deep Java Library (DJL)。