Begin Server Side Event

dRAGon-Okinawa · Jun 4, 2024 · 021c5b7 · 021c5b7
1 parent d4ffb0d
commit 021c5b7
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 5 deletions.
diff --git a/backend/src/main/java/ai/dragon/controller/api/ragapi/OpenAiCompatibleV1ApiController.java b/backend/src/main/java/ai/dragon/controller/api/ragapi/OpenAiCompatibleV1ApiController.java
@@ -8,15 +8,16 @@
 import org.springframework.web.bind.annotation.RequestBody;
 import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RestController;
+import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
 
 import ai.dragon.dto.openai.completion.OpenAiChatCompletionChoice;
 import ai.dragon.dto.openai.completion.OpenAiChatCompletionRequest;
 import ai.dragon.dto.openai.completion.OpenAiChatCompletionResponse;
 import ai.dragon.dto.openai.completion.OpenAiCompletionChoice;
+import ai.dragon.dto.openai.completion.OpenAiCompletionMessage;
 import ai.dragon.dto.openai.completion.OpenAiCompletionRequest;
 import ai.dragon.dto.openai.completion.OpenAiCompletionResponse;
 import ai.dragon.dto.openai.completion.OpenAiCompletionUsage;
-import ai.dragon.dto.openai.completion.OpenAiCompletionMessage;
 import io.swagger.v3.oas.annotations.Operation;
 import io.swagger.v3.oas.annotations.tags.Tag;
 import jakarta.validation.Valid;
@@ -27,7 +28,8 @@
 public class OpenAiCompatibleV1ApiController {
     @PostMapping("/completions")
     @Operation(summary = "Creates a completion", description = "Creates a completion for the provided prompt and parameters.")
-    public OpenAiCompletionResponse completions(@Valid @RequestBody OpenAiCompletionRequest request) throws Exception {
+    public OpenAiCompletionResponse completions(@Valid @RequestBody OpenAiCompletionRequest request)
+            throws Exception {
         OpenAiCompletionResponse response = new OpenAiCompletionResponse();
 
         response.setId(UUID.randomUUID().toString());
@@ -56,8 +58,7 @@ public OpenAiCompletionResponse completions(@Valid @RequestBody OpenAiCompletion
 
     @PostMapping("/chat/completions")
     @Operation(summary = "Creates a chat completion", description = "Creates a chat completion for the provided prompt and parameters.")
-    public OpenAiChatCompletionResponse chatCompletions(@Valid @RequestBody OpenAiChatCompletionRequest request)
-            throws Exception {
+    public Object chatCompletions(@Valid @RequestBody OpenAiChatCompletionRequest request) throws Exception {
         OpenAiChatCompletionResponse response = new OpenAiChatCompletionResponse();
 
         response.setId(UUID.randomUUID().toString());
@@ -85,6 +86,12 @@ public OpenAiChatCompletionResponse chatCompletions(@Valid @RequestBody OpenAiCh
 
         response.setChoices(choices);
 
+        if (request.getStream() != null && request.getStream()) {
+            SseEmitter emitter = new SseEmitter(Long.MAX_VALUE);
+            emitter.send("okok");
+            return emitter;
+        }
+
         return response;
     }
 }
diff --git a/backend/src/main/java/ai/dragon/controller/api/ragapi/SearchRagApiController.java b/backend/src/main/java/ai/dragon/controller/api/ragapi/SearchRagApiController.java
@@ -41,7 +41,7 @@ public List<EmbeddingMatchResponse> searchDocumentsInSilo(
         for (EmbeddingMatch<TextSegment> embeddingMatch : embeddingSearchResult.matches()) {
             searchResults.add(EmbeddingMatchResponse.builder()
                     .score(embeddingMatch.score())
-                    //.metadata(embeddingMatch.embedded().metadata())
+                    // .metadata(embeddingMatch.embedded().metadata())
                     .text(embeddingMatch.embedded().text())
                     .build());
         }