Просмотр исходного кода

feat:实现RAG功能;
使用RetrievalAugmentationAdvisor进行RAG业务的背景知识添加处理;
使用VectorStoreDocumentRetriever进行相关性知识召回;

yangyi 7 месяцев назад
Родитель
Сommit
dfee1f391d

+ 69 - 0
src/main/java/space/anyi/springAiAlibabaLearn/controller/RAGController.java

@@ -0,0 +1,69 @@
+package space.anyi.springAiAlibabaLearn.controller;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.ai.chat.client.ChatClient;
+import org.springframework.ai.chat.model.ChatModel;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
+import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
+import org.springframework.ai.vectorstore.VectorStore;
+import org.springframework.web.bind.annotation.*;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@RestController
+@RequestMapping("/rag")
+public class RAGController {
+    private static final Logger log = LoggerFactory.getLogger(RAGController.class);
+    private final  ChatClient chatClient;
+    private final VectorStore vectorStore;
+
+    public RAGController(ChatModel chatModel,VectorStore vectorStore) {
+        //文档检索增强器
+        VectorStoreDocumentRetriever vectorStoreDocumentRetriever = VectorStoreDocumentRetriever.builder().vectorStore(vectorStore).topK(3).similarityThreshold(0.5).build();
+        //RAG Advisor
+        RetrievalAugmentationAdvisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
+                .documentRetriever(vectorStoreDocumentRetriever)
+                .build();
+        //构建chatClient使用RetrievalAugmentationAdvisor进行RAG业务的处理
+        this.chatClient = ChatClient.builder(chatModel)
+                .defaultAdvisors(retrievalAugmentationAdvisor)
+                .build();
+        this.vectorStore = vectorStore;
+    }
+    @GetMapping("/ask")
+    public String rag(@RequestParam("message")String message){
+        log.debug("message:{}",message);
+        return chatClient.prompt(message).call().content();
+    }
+    @PostMapping("/import")
+    public String importData(@RequestParam("data")String data){
+        log.debug("data:{}",data);
+        List<Document> documents = new ArrayList<>();
+        String[] lines = data.split("\n");
+        //数据预处理
+        for (int i = 1; i < lines.length; i++) {
+            String line = lines[i];
+            int index = line.indexOf(',');
+            String text = new StringBuilder("问题:\n")
+                    .append(line.substring(0, index))
+                    .append("\n回答:\n")
+                    .append(line.substring(index + 1))
+                    .append('\n').toString();
+            log.debug("text:{}",text);
+            Document document = Document.builder().text(text).build();
+            log.debug("document:{}",document);
+            documents.add(document);
+            log.debug("documents.size:{}",documents.size());
+            //向量化模型一次最多只能处理25个document
+            if (documents.size()%25==24){
+                //文本数据向量化,并插入到postgres数据库中
+                vectorStore.add(documents);
+                documents.clear();
+            }
+        }
+        return "success";
+    }
+}