[awscurl] Fix output for server sent event content-type (#1784)

deepjavalibrary · Apr 16, 2024 · 9719228 · 9719228
1 parent 472e117
commit 9719228
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/awscurl/src/main/java/ai/djl/awscurl/HttpClient.java b/awscurl/src/main/java/ai/djl/awscurl/HttpClient.java
@@ -172,7 +172,7 @@ public static HttpResponse sendRequest(
                             ret = new BasicHttpResponse(status);
                         }
                     } else if ("text/event-stream".equals(contentType)) {
-                        handleServerSentEvent(is, requestTime, begin, names, tokens, request);
+                        handleServerSentEvent(is, requestTime, begin, names, tokens, request, ps);
                     } else if ("application/vnd.amazon.eventstream".equals(contentType)) {
                         Header header = resp.getFirstHeader("X-Amzn-SageMaker-Content-Type");
                         String realContentType = header == null ? null : header.getValue();
@@ -210,7 +210,8 @@ private static void handleServerSentEvent(
             long begin,
             String[] names,
             AtomicInteger tokens,
-            SignableRequest request)
+            SignableRequest request,
+            OutputStream ps)
             throws IOException {
         List<String> list = new ArrayList<>();
         try (BufferedReader reader =
@@ -225,10 +226,14 @@ private static void handleServerSentEvent(
                     requestTime[1] = System.nanoTime() - begin;
                 }
                 line = line.substring(5);
+                ps.write(line.getBytes(StandardCharsets.UTF_8));
+                ps.write(new byte[] {'\n'});
                 JsonElement element = JsonUtils.GSON.fromJson(line, JsonElement.class);
                 JsonUtils.getJsonList(element, list, names);
             }
-            updateTokenCount(list, tokens, request);
+            if (tokens != null) {
+                updateTokenCount(list, tokens, request);
+            }
         }
     }
 
@@ -271,7 +276,7 @@ private static void handleEventStream(
         byte[] bytes = bos.toByteArray();
         InputStream bis = new ByteArrayInputStream(bytes);
         if ("text/event-stream".equalsIgnoreCase(realContentType)) {
-            handleServerSentEvent(bis, requestTime, begin, names, tokens, request);
+            handleServerSentEvent(bis, requestTime, begin, names, tokens, request, ps);
             return;
         }
         List<StringBuilder> list = new ArrayList<>();