Spring AI中使用嵌入模型和向量数据库实现RAG应用

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

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

# You can override the above common api-key for embedding using the following property
spring.ai.openai.embedding.api-key=${OPENAI_API_KEY}

@Component
class MyComponent {
    private final EmbeddingClient embeddingClient;
    
    public MyComponent(EmbeddingClient embeddingClient) {
        this.embeddingClient = embeddingClient;
    }
    
    public void convertTextToEmbedding() {
        //Example 1: Convert text to embeddings
        List<Double> embeddings1 = embeddingClient.embed("I like Spring Boot");
        
        //Example 2: Convert document to embeddings
        List<Double> embeddings2 = embeddingClient.embed(new Document("I like Spring Boot"));
        
        //Example 3: Convert text to embeddings using options
        EmbeddingRequest embeddingRequest =
                new EmbeddingRequest(List.of("I like Spring Boot"),
                        OpenAiEmbeddingOptions.builder()
                                .withModel("text-davinci-003")
                                .build());
        EmbeddingResponse embeddingResponse = embeddingClient.call(embeddingRequest);
        List<Double> embeddings3 = embeddingResponse.getResult().getOutput();
    }
}

@Configuration
class AppConfig {
    @Bean
    VectorStore vectorStore(EmbeddingClient embeddingClient) {
        return new SimpleVectorStore(embeddingClient);
    }
}

@Component
class MyComponent {
    private final VectorStore vectorStore;
    
    public MyComponent(VectorStore vectorStore) {
        this.vectorStore = vectorStore;
    }
    
    public void storeAndRetrieveEmbeddings() {
        // Store embeddings
        List<Document> documents = 
                List.of(new Document("I like Spring Boot"),
                        new Document("I love Java programming language"));
        vectorStore.add(documents);
        
        // Retrieve embeddings
        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 chatClient, VectorStore vectorStore) {
        this.chatClient = chatClient;
        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 以指示人工智能模型使用传递的信息
        // 回答问题。
        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));

        // 使用 BeanOutputParser 将响应解析为 Person 的实例。
        var outputParser = new BeanOutputParser<>(Person.class);
        
        // 构建用户信息（userMessage），要求人工智能模型介绍这个人。
        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", outputParser.getFormat());
        var userMessage = new UserMessage(userMessagePromptTemplate.create(model).getContents());

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

        var response = chatClient.call(prompt).getResult().getOutput().getContent();

        return outputParser.parse(response);
    }
}

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