添加向量搜索命中测试 接口

ask-data-ai/ask-data-ai-biz/src/main/java/com/ask/controller/KnowledgeBaseController.java

@@ -10,16 +10,27 @@ import io.swagger.v3.oas.annotations.Operation;
 import io.swagger.v3.oas.annotations.Parameter;
 import io.swagger.v3.oas.annotations.tags.Tag;
 import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.vectorstore.SearchRequest;
+import org.springframework.ai.vectorstore.VectorStore;
+import org.springframework.ai.vectorstore.filter.Filter;
+import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
 import org.springframework.http.HttpHeaders;
 import org.springframework.util.StringUtils;
 import org.springframework.web.bind.annotation.*;
 
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
 /**
  * 知识库管理
  *
  * @author ai
  * @date 2024/12/19
  */
+@Slf4j
 @RestController
 @RequiredArgsConstructor
 @RequestMapping("/knowledge/base")
@@ -27,6 +38,7 @@ import org.springframework.web.bind.annotation.*;
 public class KnowledgeBaseController {
 
     private final KnowledgeBaseService knowledgeBaseService;
+    private final VectorStore vectorStore;
 
     /**
      * 分页查询
@@ -36,7 +48,7 @@ public class KnowledgeBaseController {
      */
     @Operation(summary = "分页查询", description = "分页查询")
     @GetMapping("/page")
-    public R<IPage<KnowledgeBase>> getPage(@Parameter(description = "分页对象") Page page, 
+    public R<IPage<KnowledgeBase>> getPage(@Parameter(description = "分页对象") Page page,
                                           @Parameter(description = "查询条件") KnowledgeBase knowledgeBase) {
         return R.ok(knowledgeBaseService.page(page, Wrappers.query(knowledgeBase)));
     }
@@ -64,12 +76,12 @@ public class KnowledgeBaseController {
         if (!StringUtils.hasText(knowledgeBase.getName())) {
             return R.failed("知识库名称不能为空");
         }
-        
+
         // 校验知识库名称是否重复
         if (knowledgeBaseService.checkNameExists(knowledgeBase.getName(), null)) {
             return R.failed("知识库名称已存在，请修改后重试");
         }
-        
+
         return R.ok(knowledgeBaseService.save(knowledgeBase));
     }
 
@@ -85,12 +97,12 @@ public class KnowledgeBaseController {
         if (!StringUtils.hasText(knowledgeBase.getName())) {
             return R.failed("知识库名称不能为空");
         }
-        
+
         // 校验知识库名称是否重复
         if (knowledgeBaseService.checkNameExists(knowledgeBase.getName(), knowledgeBase.getId())) {
             return R.failed("知识库名称已存在，请修改后重试");
         }
-        
+
         return R.ok(knowledgeBaseService.updateById(knowledgeBase));
     }
 
@@ -113,17 +125,81 @@ public class KnowledgeBaseController {
      */
     @Operation(summary = "校验知识库名称", description = "校验知识库名称是否重复")
     @GetMapping("/checkName")
-    public R<Boolean> checkName(@Parameter(description = "知识库名称") @RequestParam String name, 
+    public R<Boolean> checkName(@Parameter(description = "知识库名称") @RequestParam String name,
                                @Parameter(description = "知识库ID（可选，修改时传入）") @RequestParam(required = false) Long id) {
         if (!StringUtils.hasText(name)) {
             return R.ok(false);
         }
-        
+
         boolean exists = knowledgeBaseService.checkNameExists(name, id);
         if (exists) {
             return R.ok(false);
         }
-        
+
         return R.ok(true);
     }
+
+    /**
+     * 向量搜索命中测试
+     * @param knowledgeBaseId 知识库ID
+     * @param content 搜索内容
+     * @param topK 返回最相似的K个结果（可选，默认5）
+     * @param similarityThreshold 相似度阈值（可选，默认0.7）
+     * @return 搜索结果
+     */
+    @Operation(summary = "向量搜索命中测试", description = "根据知识库ID和内容进行向量相似度搜索测试")
+    @GetMapping("/search/test")
+    public R<List<Document>> searchTest(
+            @Parameter(description = "知识库ID") @RequestParam Long knowledgeBaseId,
+            @Parameter(description = "搜索内容") @RequestParam String content,
+            @Parameter(description = "返回最相似的K个结果") @RequestParam(defaultValue = "5") Integer topK,
+            @Parameter(description = "相似度阈值") @RequestParam(defaultValue = "0.7") Double similarityThreshold) {
+
+        try {
+            // 参数校验
+            if (knowledgeBaseId == null) {
+                return R.failed("知识库ID不能为空");
+            }
+            if (!StringUtils.hasText(content)) {
+                return R.failed("搜索内容不能为空");
+            }
+            if (topK == null || topK <= 0) {
+                topK = 5;
+            }
+            if (similarityThreshold == null || similarityThreshold < 0 || similarityThreshold > 1) {
+                similarityThreshold = 0.7;
+            }
+
+            log.info("开始向量搜索测试 - 知识库ID: {}, 内容: {}, topK: {}, 阈值: {}",
+                    knowledgeBaseId, content, topK, similarityThreshold);
+
+            // 构建搜索过滤器，只搜索指定知识库下的向量
+            // 多个条件 - AND逻辑
+            Filter.Expression filterExpression = new Filter.Expression(
+                    Filter.ExpressionType.AND,
+                    new Filter.Expression(
+                            Filter.ExpressionType.EQ,
+                            new Filter.Key("knowledgeBaseId"),
+                            new Filter.Value(knowledgeBaseId)
+                    ),
+                    new Filter.Expression(
+                            Filter.ExpressionType.EQ,
+                            new Filter.Key("isEnabled"),
+                            new Filter.Value(1)
+                    )
+            );
+            // 执行向量搜索
+            List<Document> searchResults = vectorStore.similaritySearch(SearchRequest.builder().filterExpression(filterExpression).similarityThreshold(similarityThreshold).topK(topK).build());
+
+
+            log.info("向量搜索测试完成 - 知识库ID: {}, 找到 {} 个相似结果", knowledgeBaseId, searchResults.size());
+
+            return R.ok(searchResults);
+
+        } catch (Exception e) {
+            log.error("向量搜索测试失败 - 知识库ID: {}, 内容: {}, 错误: {}",
+                    knowledgeBaseId, content, e.getMessage(), e);
+            return R.failed("搜索失败: " + e.getMessage());
+        }
+    }
 } 
\ No newline at end of file