使用 Embedding 模型和向量数据库的 Spring AI RAG

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-openai-spring-boot-starter</artifactId>
    <version>1.0.0-M1</version>
</dependency>

spring.ai.openai.api-key=${OPENAI_API_KEY}

# 你可以使用以下属性覆盖上述用于 embedding 的通用 api-key
spring.ai.openai.embedding.api-key=${OPENAI_API_KEY}

@Component
class MyComponent {
    private final EmbeddingModel embeddingModel;
    
    public MyComponent(EmbeddingModel embeddingModel) {
        this.embeddingModel = embeddingModel;
    }
    
    public void convertTextToEmbedding() {
        // 例 1：将文本转换为 Embedding
        List<Double> embeddings1 = embeddingModel.embed("I like Spring Boot");
        
        // 例 2：将文档转换为 Embedding
        List<Double> embeddings2 = embeddingModel.embed(new Document("I like Spring Boot"));
        
        // 例 3：使用选项将文本转换为 Embedding
        EmbeddingRequest embeddingRequest =
                new EmbeddingRequest(List.of("I like Spring Boot"),
                        OpenAiEmbeddingOptions.builder()
                                .withModel("text-davinci-003")
                                .build());
        EmbeddingResponse embeddingResponse = embeddingModel.call(embeddingRequest);
        List<Double> embeddings3 = embeddingResponse.getResult().getOutput();
    }
}

@Configuration
class AppConfig {
    @Bean
    VectorStore vectorStore(EmbeddingModel embeddingModel) {
        return new SimpleVectorStore(embeddingModel);
    }
}

@Component
class MyComponent {
    private final VectorStore vectorStore;
    
    public MyComponent(VectorStore vectorStore) {
        this.vectorStore = vectorStore;
    }
    
    public void storeAndRetrieveEmbeddings() {
        // 存储 Embedding
        List<Document> documents = 
                List.of(new Document("I like Spring Boot"),
                        new Document("I love Java programming language"));
        vectorStore.add(documents);
        
        // 检索 Embedding
        SearchRequest query = SearchRequest.query("Spring Boot").withTopK(2);
        List<Document> similarDocuments = vectorStore.similaritySearch(query);
        String relevantData = similarDocuments.stream()
                            .map(Document::getContent)
                            .collect(Collectors.joining(System.lineSeparator()));
    }
}

@Component
class MyComponent {
    private final VectorStore vectorStore;
    
    @Value("classpath:myfile.txt")
    private Resource resource;
    
    public MyComponent(VectorStore vectorStore) {
        this.vectorStore = vectorStore;
    }
    
    public void storeEmbeddingsFromTextFile() {
        var textReader = new TextReader(resource);
        textReader.setCharset(Charset.defaultCharset());
        List<Document> documents = textReader.get();

        vectorStore.add(documents);
    }
}

@RestController
class RAGController {
    private final ChatClient chatClient;
    private final VectorStore vectorStore;

    RAGController(ChatClient.Builder chatClientBuilder, VectorStore vectorStore) {
        this.chatClient = chatClientBuilder.build();
        this.vectorStore = vectorStore;
    }
    
    // 假设我们已经从包含人员信息的文件中读取了文件  并将其存储在 VectorStore 中，如上一节所述。
    
    @GetMapping("/ai/rag/people")
    Person chatWithRag(@RequestParam String name) {
        // 使用自然语言查询 VectorStore，查找有关个人的信息。
        List<Document> similarDocuments = 
                vectorStore.similaritySearch(SearchRequest.query(name).withTopK(2));
        String information = similarDocuments.stream()
                .map(Document::getContent)
                .collect(Collectors.joining(System.lineSeparator()));
        
        // 构建系统消息（systemMessage），指示 AI 模型使用传递的信息来回答问题。
        var systemPromptTemplate = new SystemPromptTemplate("""
              You are a helpful assistant.
              
              Use the following information to answer the question:
              {information}
              """);
        var systemMessage = systemPromptTemplate.createMessage(
                Map.of("information", information));

        // 使用 BeanOutputConverter 将响应解析为 Person 的实例。
        var outputConverter = new BeanOutputConverter<>(Person.class);
        
        // 构建用户信息（userMessage），要求 AI 模型介绍这个人。
        PromptTemplate userMessagePromptTemplate = new PromptTemplate("""
        Tell me about {name} as if current date is {current_date}.

        {format}
        """);
        Map<String,Object> model = Map.of("name", name,
                "current_date", LocalDate.now(),
                "format", outputConverter.getFormat());
        var userMessage = new UserMessage(userMessagePromptTemplate.create(model).getContents());

        var prompt = new Prompt(List.of(systemMessage, userMessage));

        var response = chatClient.prompt(prompt).call().content();

        return outputConverter.convert(response);
    }
}

record Person(String name,
              String dateOfBirth,
              int experienceInYears,
              List<String> books) {
}