From 021c5b728cf499a313cfeb13888fe6218d4e3105 Mon Sep 17 00:00:00 2001 From: Arnaud Mengus Date: Tue, 4 Jun 2024 15:45:31 +0000 Subject: [PATCH] Begin Server Side Event --- .../ragapi/OpenAiCompatibleV1ApiController.java | 15 +++++++++++---- .../api/ragapi/SearchRagApiController.java | 2 +- 2 files changed, 12 insertions(+), 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 index 540555b2..a6327813 100644 --- 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 index 73592956..34ffe52c 100644 --- 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 searchDocumentsInSilo( for (EmbeddingMatch embeddingMatch : embeddingSearchResult.matches()) { searchResults.add(EmbeddingMatchResponse.builder() .score(embeddingMatch.score()) - //.metadata(embeddingMatch.embedded().metadata()) + // .metadata(embeddingMatch.embedded().metadata()) .text(embeddingMatch.embedded().text()) .build()); }