Prechádzať zdrojové kódy

feat:实现文本向量化并插入的postgres数据库;实现文本语义相似度的查询;

yangyi 7 mesiacov pred
rodič
commit
851fe727a5

+ 53 - 0
src/main/java/space/anyi/springAiAlibabaLearn/controller/VectorTestController.java

@@ -0,0 +1,53 @@
+package space.anyi.springAiAlibabaLearn.controller;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.vectorstore.SearchRequest;
+import org.springframework.ai.vectorstore.VectorStore;
+import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RequestParam;
+import org.springframework.web.bind.annotation.RestController;
+
+import java.util.List;
+
+@RestController
+@RequestMapping("/vector")
+public class VectorTestController {
+    private final Logger log = LoggerFactory.getLogger(VectorTestController.class);
+    public final VectorStore vectorStore;
+
+    public VectorTestController(VectorStore vectorStore) {
+        this.vectorStore = vectorStore;
+    }
+
+    /**
+     * 向向量数据库添加数据;
+     * @param message
+     * @return
+     */
+    @GetMapping("/add")
+    public String add(@RequestParam("message") String message){
+        log.debug("message:{}",message);
+        //构建一个document对象
+        Document document = Document.builder().text(message).build();
+        log.debug("document:{}",document);
+        //将数据向量化,然后插入postgres数据库中
+        vectorStore.add(List.of(document));
+        return "success";
+    }
+    @GetMapping("/search")
+    public List<Document> search(@RequestParam("message") String message){
+        //构建向量查询请求对象
+        SearchRequest searchRequest = SearchRequest.builder()
+                .query(message)
+                //返回相识度最高的五条记录
+                .topK(5)
+                .build();
+        log.debug("searchRequest:{}",searchRequest);
+        List<Document> documents = vectorStore.similaritySearch(searchRequest);
+        documents.stream().forEach(document->log.debug("document:{}",document));
+        return documents;
+    }
+}