第9章: Spring AI 做 RAG：从检索到回答的整条链路

传统 AI：
用户："公司2023年营收是多少？"
AI："抱歉，我不知道..."  ❌

RAG：
1. 从知识库检索相关文档
2. 找到："2023年营收10亿元"
3. AI 基于文档回答："根据资料，2023年营收10亿元"  ✅

package com.example.chatbot.util;

/**
 * 向量距离计算工具类
 */
public class VectorDistanceUtils {

    /**
     * 计算余弦相似度，值越接近1越相似
     */
    public static double cosineSimilarity(float[] vector1, float[] vector2) {
        double dotProduct = 0.0;
        double norm1 = 0.0;
        double norm2 = 0.0;

        for (int i = 0; i < vector1.length; i++) {
            dotProduct += vector1[i] * vector2[i];
            norm1 += vector1[i] * vector1[i];
            norm2 += vector2[i] * vector2[i];
        }

        return dotProduct / (Math.sqrt(norm1) * Math.sqrt(norm2));
    }

    /**
     * 计算欧氏距离，值越小越相似
     */
    public static double euclideanDistance(float[] vector1, float[] vector2) {
        double sum = 0.0;
        for (int i = 0; i < vector1.length; i++) {
            double diff = vector1[i] - vector2[i];
            sum += diff * diff;
        }
        return Math.sqrt(sum);
    }
}

mvn compile
mvn test-compile

package com.example.chatbot.knowledge;

import com.example.chatbot.util.VectorDistanceUtils;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;

import java.util.*;
import java.util.stream.Collectors;

import static org.junit.jupiter.api.Assertions.*;

@SpringBootTest(classes = com.example.chatbot.AiChatbotApplication.class)
@DisplayName("知识点9：RAG 完整教程")
public class Test09_RAGTutorial {

    @Autowired
    private EmbeddingModel embeddingModel;

    @Autowired
    private ChatModel chatModel;

    // 简单的内存向量库
    static class SimpleVectorStore {
        private final List<VectorDocument> documents = new ArrayList<>();

        static class VectorDocument {
            String id;
            String content;
            float[] vector;
            Map<String, Object> metadata = new HashMap<>();

            VectorDocument(String id, String content, float[] vector) {
                this.id = id;
                this.content = content;
                this.vector = vector;
            }
        }

        void add(String id, String content, float[] vector) {
            documents.add(new VectorDocument(id, content, vector));
        }

        List<VectorDocument> search(float[] queryVector, int topK) {
            return documents.stream()
                .map(doc -> {
                    double similarity = VectorDistanceUtils.cosineSimilarity(
                        queryVector, doc.vector
                    );
                    doc.metadata.put("similarity", similarity);
                    return doc;
                })
                .sorted((a, b) -> Double.compare(
                    (Double) b.metadata.get("similarity"),
                    (Double) a.metadata.get("similarity")
                ))
                .limit(topK)
                .collect(Collectors.toList());
        }

        int size() {
            return documents.size();
        }
    }
}

/**
 * 测试 9.1：验证 EmbeddingModel 注入
 */
@Test
@DisplayName("9.1 EmbeddingModel 注入")
public void testEmbeddingModelInjection() {
    System.out.println("\n========== 测试：EmbeddingModel 注入 ==========");

    assertNotNull(embeddingModel, "EmbeddingModel 应该被注入");

    System.out.println("✅ EmbeddingModel 注入成功");
    System.out.println("📝 类型: " + embeddingModel.getClass().getSimpleName());
    System.out.println("✅ 测试通过");
}

mvn test -Dtest=Test09_RAGTutorial#testEmbeddingModelInjection

/**
 * 测试 9.2：文本向量化
 */
@Test
@DisplayName("9.2 文本向量化")
public void testEmbedding() {
    System.out.println("\n=== 测试向量化 ===");

    String text = "Spring AI 是一个强大的框架";
    System.out.println("📝 原始文本: " + text);

    System.out.println("\n🔄 向量化中...");
    var response = embeddingModel.embedForResponse(List.of(text));
    float[] vector = response.getResults().get(0).getOutput();

    System.out.println("✅ 向量化完成");
    System.out.println("📊 向量维度: " + vector.length);
    System.out.println("📊 前5个值: " + java.util.Arrays.toString(
        java.util.Arrays.copyOfRange(vector, 0, Math.min(5, vector.length))
    ));

    assertTrue(vector.length > 0, "向量不应为空");
    assertEquals(1536, vector.length, "通义千问的向量维度应该是1536");

    System.out.println("✅ 向量化测试通过");
}

mvn test -Dtest=Test09_RAGTutorial#testEmbedding

/**
 * 测试 9.3：相似度计算
 */
@Test
@DisplayName("9.3 相似度计算")
public void testSimilarity() {
    System.out.println("\n=== 测试相似度计算 ===");

    String text1 = "我喜欢编程";
    String text2 = "我热爱写代码";
    String text3 = "今天天气真好";

    System.out.println("📝 文本1: " + text1);
    System.out.println("📝 文本2: " + text2);
    System.out.println("📝 文本3: " + text3);

    System.out.println("\n🔄 向量化中...");
    java.util.List<float[]> vectors = embeddingModel.embed(java.util.List.of(text1, text2, text3));

    System.out.println("\n📊 相似度计算:");
    double sim12 = VectorDistanceUtils.cosineSimilarity(vectors.get(0), vectors.get(1));
    double sim13 = VectorDistanceUtils.cosineSimilarity(vectors.get(0), vectors.get(2));
    double sim23 = VectorDistanceUtils.cosineSimilarity(vectors.get(1), vectors.get(2));

    System.out.println("  文本1 vs 文本2: " + String.format("%.4f", sim12));
    System.out.println("  文本1 vs 文本3: " + String.format("%.4f", sim13));
    System.out.println("  文本2 vs 文本3: " + String.format("%.4f", sim23));

    System.out.println("\n💡 分析:");
    assertTrue(sim12 > sim13, "相似的文本（1和2）相似度应该更高");
    assertTrue(sim12 > sim23, "相似的文本（1和2）相似度应该更高");
    System.out.println("  ✅ 编程相关的文本确实更相似");

    System.out.println("✅ 相似度测试通过");
}

/**
 * 测试 9.4：构建简单向量库
 */
@Test
@DisplayName("9.4 构建简单向量库")
public void testVectorStore() {
    System.out.println("\n=== 测试向量库 ===");

    java.util.List<String> documents = java.util.List.of(
        "Spring AI 是 Spring 官方的 AI 集成框架",
        "Spring Boot 简化了 Spring 应用的开发",
        "Java 是一种面向对象的编程语言",
        "Python 在数据科学领域很流行",
        "RAG 是检索增强生成技术"
    );

    System.out.println("📚 准备知识库文档...");
    documents.forEach(doc -> System.out.println("  - " + doc));

    System.out.println("\n🔄 向量化文档...");
    java.util.List<float[]> vectors = embeddingModel.embed(documents);
    System.out.println("✅ 向量化完成，共 " + vectors.size() + " 个向量");

    System.out.println("\n💾 存入向量库...");
    SimpleVectorStore vectorStore = new SimpleVectorStore();
    for (int i = 0; i < documents.size(); i++) {
        vectorStore.add("doc-" + i, documents.get(i), vectors.get(i));
    }
    System.out.println("✅ 存储完成，向量库大小: " + vectorStore.size());

    System.out.println("\n🔍 测试检索...");
    String query = "什么是 Spring AI？";
    System.out.println("  查询: " + query);

    float[] queryVector = embeddingModel.embed(query);
    java.util.List<SimpleVectorStore.VectorDocument> results = vectorStore.search(queryVector, 3);

    System.out.println("\n📊 検索结果（Top 3）:");
    for (int i = 0; i < results.size(); i++) {
        var doc = results.get(i);
        System.out.println("  " + (i + 1) + ". [相似度: " +
            String.format("%.4f", doc.metadata.get("similarity")) + "] " +
            doc.content);
    }

    assertTrue(results.get(0).content.contains("Spring AI"),
        "最相关的文档应该包含 Spring AI");

    System.out.println("\n✅ 向量库测试通过");
}

/**
 * 测试 9.5：完整 RAG 流程
 */
@Test
@DisplayName("9.5 完整 RAG 流程")
public void testCompleteRAG() {
    System.out.println("\n=== 完整 RAG 流程 ===");

    System.out.println("【步骤1：构建知识库】");
    java.util.List<String> knowledge = java.util.List.of(
        "公司成立于2020年，总部位于北京",
        "2023年公司营收达到10亿元，同比增长50%",
        "公司主要产品是AI助手，服务超过100万用户",
        "公司有员工500人，其中研发人员占60%"
    );
    knowledge.forEach(doc -> System.out.println("  📄 " + doc));

    System.out.println("\n【步骤2：向量化并存储】");
    SimpleVectorStore vectorStore = new SimpleVectorStore();
    java.util.List<float[]> vectors = embeddingModel.embed(knowledge);
    for (int i = 0; i < knowledge.size(); i++) {
        vectorStore.add("doc-" + i, knowledge.get(i), vectors.get(i));
    }
    System.out.println("  ✅ 存储了 " + vectorStore.size() + " 个文档");

    System.out.println("\n【步骤3：用户提问】");
    String question = "公司2023年的营收是多少？";
    System.out.println("  👤 问题: " + question);

    System.out.println("\n【步骤4：检索相关文档】");
    float[] questionVector = embeddingModel.embed(question);
    java.util.List<SimpleVectorStore.VectorDocument> relevantDocs =
        vectorStore.search(questionVector, 2);
    System.out.println("  🔍 找到 " + relevantDocs.size() + " 个相关文档:");
    relevantDocs.forEach(doc ->
        System.out.println("    - " + doc.content)
    );

    System.out.println("\n【步骤5：构造增强 Prompt】");
    String context = relevantDocs.stream()
        .map(doc -> doc.content)
        .collect(java.util.stream.Collectors.joining("\n"));
    String enhancedPrompt = String.format(
        "基于以下信息回答问题：\n\n%s\n\n问题：%s",
        context, question
    );
    System.out.println("  📝 增强后的 Prompt:\n  " + enhancedPrompt.replace("\n", "\n  "));

    System.out.println("\n【步骤6：AI 生成回答】");
    var response = chatModel.call(new org.springframework.ai.chat.prompt.Prompt(enhancedPrompt));
    String answer = response.getResult().getOutput().getText();
    System.out.println("  🤖 AI 回答: " + answer);

    assertTrue(answer.contains("10亿") || answer.contains("10 亿"),
        "回答应该包含营收数据");

    System.out.println("\n✅ 完整 RAG 流程测试通过");
}

/**
 * 测试 9.6：对比有无 RAG
 */
@Test
@DisplayName("9.6 对比有无 RAG")
public void testWithAndWithoutRAG() {
    System.out.println("\n=== 对比有无 RAG ===");

    java.util.List<String> knowledge = java.util.List.of(
        "XYZ 公司成立于2020年",
        "XYZ 公司2023年营收10亿元"
    );

    SimpleVectorStore vectorStore = new SimpleVectorStore();
    java.util.List<float[]> vectors = embeddingModel.embed(knowledge);
    for (int i = 0; i < knowledge.size(); i++) {
        vectorStore.add("doc-" + i, knowledge.get(i), vectors.get(i));
    }

    String question = "XYZ 公司2023年营收多少？";

    System.out.println("\n【测试1：没有 RAG】");
    System.out.println("👤 问题: " + question);
    var response1 = chatModel.call(new org.springframework.ai.chat.prompt.Prompt(question));
    String answer1 = response1.getResult().getOutput().getText();
    System.out.println("🤖 回答: " + answer1);
    System.out.println("❌ AI 可能不知道答案");

    System.out.println("\n【测试2：使用 RAG】");
    float[] qVector = embeddingModel.embed(question);
    java.util.List<SimpleVectorStore.VectorDocument> docs = vectorStore.search(qVector, 1);
    String context = docs.get(0).content;
    String enhancedPrompt = String.format(
        "根据以下信息回答：%s\n\n问题：%s", context, question
    );
    var response2 = chatModel.call(new org.springframework.ai.chat.prompt.Prompt(enhancedPrompt));
    String answer2 = response2.getResult().getOutput().getText();
    System.out.println("🤖 回答: " + answer2);
    System.out.println("✅ AI 基于知识库准确回答");

    assertTrue(answer2.contains("10亿") || answer2.contains("10 亿"));

    System.out.println("\n✅ 对比测试通过");
}

<dependency>
  <groupId>org.springframework.ai</groupId>
  <artifactId>spring-ai-chroma-store-spring-boot-starter</artifactId>
</dependency>

spring:
  ai:
    vectorstore:
      chroma:
        client:
          host: http://localhost
          port: 8000
        collection-name: SpringAiCollection
        initialize-schema: true

<!-- Spring AI ChromaDB Store -->
<dependency>
  <groupId>org.springframework.ai</groupId>
  <artifactId>spring-ai-chroma-store-spring-boot-starter</artifactId>
  <version>1.0.0-M6</version>
</dependency>

<!-- 当前已验证：ChromaDB 0.4.24 + Spring AI 1.0.0-M6 -->

import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;

@Autowired
private VectorStore vectorStore;

@Test
@DisplayName("9.7 ChromaDB 向量库（可选）")
public void testChromaDbVectorStore() {
    System.out.println("\n=== 测试 9.7：ChromaDB VectorStore ===");

    java.util.List<Document> docs = java.util.List.of(
        new Document("Spring AI rocks!!"),
        new Document("RAG = 检索 + 增强 Prompt + 生成"),
        new Document("向量数据库用于存储 embedding 并做相似度检索")
    );

    vectorStore.add(docs);

    java.util.List<Document> results = vectorStore.similaritySearch(
        SearchRequest.builder().query("Spring AI").topK(2).build()
    );

    System.out.println("🔍 TopK 结果数量: " + results.size());
    results.forEach(d -> System.out.println("  - " + d.getText()));
}

// 运行
// mvn test -Dtest=Test09_RAGTutorial#testChromaDbVectorStore

mvn test -Dtest=Test09_RAGTutorial

Tests run: 6, Failures: 0, Errors: 0, Skipped: 0

✅ 所有测试通过！