Add stream to Ask

Author: johnoliver

app/backend/src/main/java/com/microsoft/openai/samples/rag/ask/controller/AskController.java

@@ -10,11 +10,14 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.http.HttpStatus;
+import org.springframework.http.MediaType;
 import org.springframework.http.ResponseEntity;
 import org.springframework.util.StringUtils;
 import org.springframework.web.bind.annotation.PostMapping;
 import org.springframework.web.bind.annotation.RequestBody;
 import org.springframework.web.bind.annotation.RestController;
+import org.springframework.web.server.ResponseStatusException;
+import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;
 
 @RestController
 public class AskController {
@@ -26,8 +29,64 @@ public class AskController {
         this.ragApproachFactory = ragApproachFactory;
     }
 
+    @PostMapping(
+            value = "/api/ask",
+            produces = MediaType.APPLICATION_NDJSON_VALUE
+    )
+    public ResponseEntity openAIAskStream(
+            @RequestBody ChatAppRequest askRequest
+    ) {
+        if (!askRequest.stream()) {
+            LOGGER.warn("Requested a content-type of application/ndjson however did not requested streaming. Please use a content-type of application/json");
+            throw new ResponseStatusException(HttpStatus.BAD_REQUEST, "Requested a content-type of application/ndjson however did not requested streaming. Please use a content-type of application/json");
+        }
+
+        String question = askRequest.messages().get(askRequest.messages().size() - 1).content();
+        LOGGER.info("Received request for ask api with question [{}] and approach[{}]", question, askRequest.approach());
+
+        if (!StringUtils.hasText(askRequest.approach())) {
+            LOGGER.warn("approach cannot be null in ASK request");
+            return ResponseEntity.status(HttpStatus.BAD_REQUEST).body(null);
+        }
+
+        if (!StringUtils.hasText(question)) {
+            LOGGER.warn("question cannot be null in ASK request");
+            return ResponseEntity.status(HttpStatus.BAD_REQUEST).body(null);
+        }
+
+        var ragOptions = new RAGOptions.Builder()
+                .retrievialMode(askRequest.context().overrides().retrieval_mode().name())
+                .semanticKernelMode(askRequest.context().overrides().semantic_kernel_mode())
+                .semanticRanker(askRequest.context().overrides().semantic_ranker())
+                .semanticCaptions(askRequest.context().overrides().semantic_captions())
+                .excludeCategory(askRequest.context().overrides().exclude_category())
+                .promptTemplate(askRequest.context().overrides().prompt_template())
+                .top(askRequest.context().overrides().top())
+                .build();
+
+        RAGApproach<String, RAGResponse> ragApproach = ragApproachFactory.createApproach(askRequest.approach(), RAGType.ASK, ragOptions);
+
+        StreamingResponseBody response = output -> {
+            try {
+                ragApproach.runStreaming(question, ragOptions, output);
+            } finally {
+                output.flush();
+                output.close();
+            }
+        };
+
+        return ResponseEntity.ok()
+                .contentType(MediaType.APPLICATION_NDJSON)
+                .body(response);
+    }
+
     @PostMapping("/api/ask")
     public ResponseEntity<ChatResponse> openAIAsk(@RequestBody ChatAppRequest askRequest) {
+        if (askRequest.stream()) {
+            LOGGER.warn("Requested a content-type of application/json however also requested streaming. Please use a content-type of application/ndjson");
+            throw new ResponseStatusException(HttpStatus.BAD_REQUEST, "Requested a content-type of application/json however also requested streaming. Please use a content-type of application/ndjson");
+        }
+
         String question = askRequest.messages().get(askRequest.messages().size() - 1).content();
         LOGGER.info("Received request for ask api with question [{}] and approach[{}]", question, askRequest.approach());
 

app/backend/src/main/java/com/microsoft/openai/samples/rag/chat/controller/ChatController.java

@@ -19,8 +19,8 @@
 import org.springframework.web.bind.annotation.PostMapping;
 import org.springframework.web.bind.annotation.RequestBody;
 import org.springframework.web.bind.annotation.RestController;
+import org.springframework.web.server.ResponseStatusException;
 import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;
-import reactor.core.publisher.Flux;
 
 import java.util.ArrayList;
 import java.util.Collection;
@@ -43,6 +43,11 @@ public ChatController(RAGApproachFactory<ChatGPTConversation, RAGResponse> ragAp
     public ResponseEntity<StreamingResponseBody> openAIAskStream(
             @RequestBody ChatAppRequest chatRequest
     ) {
+        if (!chatRequest.stream()) {
+            LOGGER.warn("Requested a content-type of application/ndjson however did not requested streaming. Please use a content-type of application/json");
+            throw new ResponseStatusException(HttpStatus.BAD_REQUEST, "Requested a content-type of application/ndjson however did not requested streaming. Please use a content-type of application/json");
+        }
+
         LOGGER.info("Received request for chat api with approach[{}]", chatRequest.approach());
 
         if (!StringUtils.hasText(chatRequest.approach())) {
@@ -69,9 +74,6 @@ public ResponseEntity<StreamingResponseBody> openAIAskStream(
 
         ChatGPTConversation chatGPTConversation = convertToChatGPT(chatRequest.messages());
 
-
-        Flux<Integer> counter = Flux.range(0, Integer.MAX_VALUE);
-
         StreamingResponseBody response = output -> {
             try {
                 ragApproach.runStreaming(chatGPTConversation, ragOptions, output);
@@ -91,6 +93,11 @@ public ResponseEntity<StreamingResponseBody> openAIAskStream(
             produces = MediaType.APPLICATION_JSON_VALUE
     )
     public ResponseEntity<ChatResponse> openAIAsk(@RequestBody ChatAppRequest chatRequest) {
+        if (chatRequest.stream()) {
+            LOGGER.warn("Requested a content-type of application/json however also requested streaming. Please use a content-type of application/ndjson");
+            throw new ResponseStatusException(HttpStatus.BAD_REQUEST, "Requested a content-type of application/json however also requested streaming. Please use a content-type of application/ndjson");
+        }
+
         LOGGER.info("Received request for chat api with approach[{}]", chatRequest.approach());
 
         if (!StringUtils.hasText(chatRequest.approach())) {