向量数据库：相似度查询

pom.xml

@@ -105,6 +105,21 @@
             <artifactId>spring-ai-markdown-document-reader</artifactId>
             <version>1.0.0-M6</version>
         </dependency>
+        <!-- pg数据库 -->
+        <dependency>
+            <groupId>org.springframework.boot</groupId>
+            <artifactId>spring-boot-starter-jdbc</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.postgresql</groupId>
+            <artifactId>postgresql</artifactId>
+            <scope>runtime</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.springframework.ai</groupId>
+            <artifactId>spring-ai-pgvector-store</artifactId>
+            <version>1.0.0-M6</version>
+        </dependency>
 
     </dependencies>
 

src/main/java/com/huangge1199/aiagent/LongAiAgentApplication.java

@@ -1,9 +1,13 @@
 package com.huangge1199.aiagent;
 
+import org.springframework.ai.autoconfigure.vectorstore.pgvector.PgVectorStoreAutoConfiguration;
 import org.springframework.boot.SpringApplication;
 import org.springframework.boot.autoconfigure.SpringBootApplication;
 
-@SpringBootApplication
+/**
+ * @author hyy
+ */
+@SpringBootApplication(exclude = PgVectorStoreAutoConfiguration.class)
 public class LongAiAgentApplication {
 
     public static void main(String[] args) {

src/main/java/com/huangge1199/aiagent/Service/DBService.java

@@ -0,0 +1,15 @@
+package com.huangge1199.aiagent.Service;
+
+import org.springframework.ai.document.Document;
+
+import java.util.List;
+
+/**
+ * DBService
+ *
+ * @author huangge1199
+ * @since 2025/5/27 10:46:40
+ */
+public interface DBService {
+    List<Document> similaritySearch();
+}

src/main/java/com/huangge1199/aiagent/Service/impl/DBServiceImpl.java

@@ -0,0 +1,37 @@
+package com.huangge1199.aiagent.Service.impl;
+
+import com.huangge1199.aiagent.Service.DBService;
+import jakarta.annotation.Resource;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.vectorstore.SearchRequest;
+import org.springframework.ai.vectorstore.VectorStore;
+import org.springframework.stereotype.Service;
+
+import java.util.List;
+import java.util.Map;
+
+/**
+ * DBServiceImpl
+ *
+ * @author huangge1199
+ * @since 2025/5/27 10:46:56
+ */
+@Service
+public class DBServiceImpl implements DBService {
+
+    @Resource
+    VectorStore pgVectorVectorStore;
+
+
+    @Override
+    public List<Document> similaritySearch() {
+        List<Document> documents = List.of(
+                new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", Map.of("meta1", "meta1")),
+                new Document("The World is Big and Salvation Lurks Around the Corner"),
+                new Document("You walk forward facing the past and you turn back toward the future.", Map.of("meta2", "meta2")));
+        // 添加文档
+        pgVectorVectorStore.add(documents);
+        // 相似度查询
+        return pgVectorVectorStore.similaritySearch(SearchRequest.builder().query("Spring").topK(5).build());
+    }
+}

src/main/java/com/huangge1199/aiagent/config/PgVectorVectorStoreConfig.java

@@ -0,0 +1,55 @@
+package com.huangge1199.aiagent.config;
+
+import com.huangge1199.aiagent.rag.DocumentLoaderUtils;
+import jakarta.annotation.Resource;
+import org.springframework.ai.embedding.EmbeddingModel;
+import org.springframework.ai.vectorstore.VectorStore;
+import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
+import org.springframework.beans.factory.annotation.Qualifier;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.jdbc.core.JdbcTemplate;
+
+import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgDistanceType.COSINE_DISTANCE;
+import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType.HNSW;
+
+/**
+ * pgConfi
+ *
+ * @author huangge1199
+ * @since 2025/5/27 9:52:30
+ */
+@Configuration
+public class PgVectorVectorStoreConfig {
+
+    @Resource
+    private DocumentLoaderUtils documentLoaderUtils;
+
+    @Bean
+    public VectorStore vectorStore(JdbcTemplate jdbcTemplate, @Qualifier("ollamaEmbeddingModel") EmbeddingModel embeddingModel) {
+        // 设置向量维度，默认为模型维度或1536
+        // 设置距离类型，默认为 COSINE_DISTANCE
+        // 设置索引类型，默认为 HNSW
+        // 是否初始化模式，默认为 false
+        // 设置模式名称，默认为 "public"
+        // 设置向量表名称，默认为 "vector_store"
+        // 设置最大文档批处理大小，默认为 10000
+        return PgVectorStore.builder(jdbcTemplate, embeddingModel)
+                // 设置向量维度，默认为模型维度或1536
+                .dimensions(1024)
+                // 设置距离类型，默认为 COSINE_DISTANCE
+                .distanceType(COSINE_DISTANCE)
+                // 设置索引类型，默认为 HNSW
+                .indexType(HNSW)
+                // 是否初始化模式，默认为 false
+                .initializeSchema(true)
+                // 设置模式名称，默认为 "public"
+                .schemaName("public")
+                // 设置向量表名称，默认为 "vector_store"
+                .vectorTableName("vector_store")
+                // 设置最大文档批处理大小，默认为 10000
+                .maxDocumentBatchSize(10000)
+                .build();
+    }
+}
+

src/main/java/com/huangge1199/aiagent/controller/DBController.java

@@ -0,0 +1,35 @@
+package com.huangge1199.aiagent.controller;
+
+import com.huangge1199.aiagent.Service.DBService;
+import com.huangge1199.aiagent.common.R;
+import io.swagger.v3.oas.annotations.Operation;
+import io.swagger.v3.oas.annotations.tags.Tag;
+import jakarta.annotation.Resource;
+import org.springframework.ai.document.Document;
+import org.springframework.web.bind.annotation.PostMapping;
+import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RestController;
+
+import java.util.List;
+
+/**
+ * DBController
+ *
+ * @author huangge1199
+ * @since 2025/5/27 10:37:25
+ */
+@RestController
+@RequestMapping("/db")
+@Tag(name = "数据库相关")
+public class DBController {
+
+    @Resource
+    private DBService dbService;
+
+    @PostMapping("/similaritySearch")
+    @Operation(summary = "相似度查询")
+    public R<List<Document>> similaritySearch() {
+        List<Document> results = dbService.similaritySearch();
+        return R.ok(results);
+    }
+}

src/main/java/com/huangge1199/aiagent/rag/RagConfig.java

@@ -2,16 +2,9 @@ package com.huangge1199.aiagent.rag;
 
 import jakarta.annotation.Resource;
 import org.springframework.ai.chat.client.ChatClient;
-import org.springframework.ai.document.Document;
-import org.springframework.ai.embedding.EmbeddingModel;
-import org.springframework.ai.vectorstore.SimpleVectorStore;
-import org.springframework.ai.vectorstore.VectorStore;
-import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 
-import java.util.List;
-
 /**
  * RagConfig
  *
@@ -35,17 +28,4 @@ public class RagConfig {
         return builder.defaultSystem("你将作为一名恋爱大师，对于用户的问题作出解答")
                 .build();
     }
-
-    @Bean
-    VectorStore vectorStore(@Qualifier("ollamaEmbeddingModel") EmbeddingModel embeddingModel) {
-        SimpleVectorStore simpleVectorStore = SimpleVectorStore.builder(embeddingModel).build();
-//        // 加载文档
-//        List<Document> documentList = documentLoaderUtils.loadMarkdowns();
-//        // 自主切分文档
-//        List<Document> splitDocuments = myTokenTextSplitter.splitCustomized(documentList);
-//        // 自动补充关键词元信息
-//        List<Document> enrichedDocuments = myKeywordEnricher.enrichDocuments(documentList);
-//        simpleVectorStore.add(enrichedDocuments);
-        return simpleVectorStore;
-    }
 }

src/main/resources/application.yml

@@ -12,6 +12,12 @@ spring:
       base-url: http://192.168.188.2:11435
       chat:
         model: gemma3:1b
+    vectorstore:
+      pgvector:
+        index-type: HNSW
+        distance-type: COSINE_DISTANCE
+        dimensions: 1536
+        max-document-batch-size: 10000
 server:
   port: 8123
   servlet: