feat:add AIChat endpoint

Author: msslulu

base/src/main/java/com/tinyengine/it/controller/AiChatController.java

@@ -1,13 +1,12 @@
 /**
  * Copyright (c) 2023 - present TinyEngine Authors.
  * Copyright (c) 2023 - present Huawei Cloud Computing Technologies Co., Ltd.
- *
+ * <p>
  * Use of this source code is governed by an MIT-style license.
- *
+ * <p>
  * THE OPEN SOURCE SOFTWARE IN THIS PRODUCT IS DISTRIBUTED IN THE HOPE THAT IT WILL BE USEFUL,
  * BUT WITHOUT ANY WARRANTY, WITHOUT EVEN THE IMPLIED WARRANTY OF MERCHANTABILITY OR FITNESS FOR
  * A PARTICULAR PURPOSE. SEE THE APPLICABLE LICENSES FOR MORE DETAILS.
- *
  */
 
 package com.tinyengine.it.controller;
@@ -18,6 +17,9 @@
 import com.tinyengine.it.model.dto.AiToken;
 import com.tinyengine.it.model.dto.ChatRequest;
 
+import com.tinyengine.it.rag.entity.EmbeddingMatchDto;
+import com.tinyengine.it.rag.entity.SearchRequest;
+import com.tinyengine.it.rag.service.StorageService;
 import com.tinyengine.it.service.app.v1.AiChatV1Service;
 import io.swagger.v3.oas.annotations.Operation;
 import io.swagger.v3.oas.annotations.Parameter;
@@ -37,6 +39,8 @@
 import org.springframework.web.bind.annotation.RestController;
 import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;
 
+import java.util.List;
+
 /**
  * The type Ai chat controller.
  *
@@ -52,6 +56,11 @@ public class AiChatController {
      */
     @Autowired
     private AiChatV1Service aiChatV1Service;
+    /**
+     *  vector storage service
+     */
+    @Autowired
+    private StorageService vectorStorageService;
 
     /**
      * AI api
@@ -60,17 +69,17 @@ public class AiChatController {
      * @return ai回答信息 result
      */
     @Operation(summary = "获取ai回答信息", description = "获取ai回答信息",
-        parameters = {
-            @Parameter(name = "ChatRequest", description = "入参对象")
-        }, responses = {
+            parameters = {
+                    @Parameter(name = "ChatRequest", description = "入参对象")
+            }, responses = {
             @ApiResponse(responseCode = "200", description = "返回信息",
-                content = @Content(mediaType = "application/json", schema = @Schema())),
+                    content = @Content(mediaType = "application/json", schema = @Schema())),
             @ApiResponse(responseCode = "400", description = "请求失败")
     })
     @SystemControllerLog(description = "AI chat")
     @PostMapping("/ai/chat")
     public ResponseEntity<?> aiChat(@RequestBody ChatRequest request,
-        @RequestHeader(value = "Authorization", required = false) String authorization) throws Exception {
+                                    @RequestHeader(value = "Authorization", required = false) String authorization) throws Exception {
 
         if (authorization != null && authorization.startsWith("Bearer ")) {
             String token = authorization.replace("Bearer ", "");
@@ -81,16 +90,36 @@ public ResponseEntity<?> aiChat(@RequestBody ChatRequest request,
 
         if (request.isStream()) {
             return ResponseEntity.ok()
-                .contentType(MediaType.TEXT_EVENT_STREAM)
-                .header("Cache-Control", "no-cache")
-                .header("X-Accel-Buffering", "no") // 禁用Nginx缓冲
-                .body((StreamingResponseBody) response);
+                    .contentType(MediaType.TEXT_EVENT_STREAM)
+                    .header("Cache-Control", "no-cache")
+                    .header("X-Accel-Buffering", "no") // 禁用Nginx缓冲
+                    .body((StreamingResponseBody) response);
         } else {
             return ResponseEntity.ok(response);
         }
 
     }
 
+    /**
+     * search in collection
+     *
+     * @param searchDto the searchDto
+     * @return result
+     */
+    @Operation(summary = "在指定集合中搜索", description = "在指定集合中搜索",
+            parameters = {
+                    @Parameter(name = "searchDto", description = "搜索请求参数体"),
+            }, responses = {
+            @ApiResponse(responseCode = "200", description = "返回信息",
+                    content = @Content(mediaType = "application/json", schema = @Schema())),
+            @ApiResponse(responseCode = "400", description = "请求失败")
+    })
+    @SystemControllerLog(description = "AI search in collection")
+    @PostMapping("/ai/search")
+    public Result<List<EmbeddingMatchDto>> searchInCollection(@RequestBody SearchRequest searchDto) {
+        List<EmbeddingMatchDto> results = vectorStorageService.search(searchDto);
+        return Result.success(results);
+    }
 
     /**
      * AI api v1
@@ -99,17 +128,17 @@ public ResponseEntity<?> aiChat(@RequestBody ChatRequest request,
      * @return ai回答信息 result
      */
     @Operation(summary = "获取ai回答信息", description = "获取ai回答信息",
-        parameters = {
-            @Parameter(name = "ChatRequest", description = "入参对象")
-        }, responses = {
+            parameters = {
+                    @Parameter(name = "ChatRequest", description = "入参对象")
+            }, responses = {
             @ApiResponse(responseCode = "200", description = "返回信息",
-                content = @Content(mediaType = "application/json", schema = @Schema())),
+                    content = @Content(mediaType = "application/json", schema = @Schema())),
             @ApiResponse(responseCode = "400", description = "请求失败")
     })
     @SystemControllerLog(description = "AI completions")
     @PostMapping("/chat/completions")
     public ResponseEntity<?> completions(@RequestBody ChatRequest request,
-        @RequestHeader(value = "Authorization", required = false) String authorization) throws Exception {
+                                         @RequestHeader(value = "Authorization", required = false) String authorization) throws Exception {
         if (authorization != null && authorization.startsWith("Bearer ")) {
             String token = authorization.replace("Bearer ", "");
             request.setApiKey(token);
@@ -119,33 +148,34 @@ public ResponseEntity<?> completions(@RequestBody ChatRequest request,
 
         if (request.isStream()) {
             return ResponseEntity.ok()
-                .contentType(MediaType.TEXT_EVENT_STREAM)
-                .header("Cache-Control", "no-cache")
-                .header("X-Accel-Buffering", "no") // 禁用Nginx缓冲
-                .body((StreamingResponseBody) response);
+                    .contentType(MediaType.TEXT_EVENT_STREAM)
+                    .header("Cache-Control", "no-cache")
+                    .header("X-Accel-Buffering", "no") // 禁用Nginx缓冲
+                    .body((StreamingResponseBody) response);
         } else {
             return ResponseEntity.ok(response);
         }
     }
+
     /**
      * get token
      *
      * @param request the request
      * @return ai回答信息 result
      */
     @Operation(summary = "获取加密key信息", description = "获取加密key信息",
-        parameters = {
-            @Parameter(name = "request", description = "入参对象")
-        }, responses = {
+            parameters = {
+                    @Parameter(name = "request", description = "入参对象")
+            }, responses = {
             @ApiResponse(responseCode = "200", description = "返回信息",
-                content = @Content(mediaType = "application/json", schema = @Schema())),
+                    content = @Content(mediaType = "application/json", schema = @Schema())),
             @ApiResponse(responseCode = "400", description = "请求失败")
     })
     @SystemControllerLog(description = "get token")
     @PostMapping("/encrypt-key")
     public Result<AiToken> getToken(@RequestBody ChatRequest request) throws Exception {
         String apiKey = request.getApiKey();
-        if(apiKey == null || apiKey.isEmpty()) {
+        if (apiKey == null || apiKey.isEmpty()) {
             return Result.failed(ExceptionEnum.CM320);
         }
         String token = aiChatV1Service.getToken(apiKey);