diff --git a/awscurl/src/main/java/ai/djl/awscurl/HttpClient.java b/awscurl/src/main/java/ai/djl/awscurl/HttpClient.java index 2e374662d..a5a8015d1 100644 --- 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 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 list = new ArrayList<>();