Repository: deepjavalibrary/djl
Pull request #3727 (2025-07-08 02:24:08)

Adds jsonlines streaming support for RpcEngine
## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here



Patch:
diff --git a/api/src/main/java/ai/djl/engine/rpc/RpcClient.java b/api/src/main/java/ai/djl/engine/rpc/RpcClient.java
index a86478dae..455dd4220 100644
--- a/api/src/main/java/ai/djl/engine/rpc/RpcClient.java
+++ b/api/src/main/java/ai/djl/engine/rpc/RpcClient.java
@@ -26,6 +26,7 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.io.OutputStream;
+import java.io.Reader;
 import java.net.HttpURLConnection;
 import java.net.MalformedURLException;
 import java.net.URL;
@@ -123,7 +124,7 @@ public final class RpcClient {
             throw new IOException("Offline mode is enabled.");
         }
         HttpURLConnection conn = (HttpURLConnection) url.openConnection();
-        boolean isSse = false;
+        boolean isStream = false;
         try {
             conn.setRequestMethod(method);
             if ("POST".equals(method) || "PUT".equals(method)) {
@@ -153,43 +154,18 @@ public final class RpcClient {
                 String value = entry.getValue().get(0);
                 if (key != null) {
                     if ("content-type".equalsIgnoreCase(key)
-                            && "text/event-stream".equalsIgnoreCase(value)) {
-                        isSse = true;
+                            && ("text/event-stream".equalsIgnoreCase(value)
+                                    || "application/jsonlines".equalsIgnoreCase(value))) {
+                        isStream = true;
                     }
                     out.addProperty(key, value);
                 }
             }
             if (code == 200) {
-                if (isSse) {
+                if (isStream) {
                     ChunkedBytesSupplier cbs = new ChunkedBytesSupplier();
                     out.add(cbs);
-                    CompletableFuture.<Void>supplyAsync(
-                            () -> {
-                                try (BufferedReader reader =
-                                        new BufferedReader(
-                                                new InputStreamReader(
-                                                        conn.getInputStream(),
-                                                        StandardCharsets.UTF_8))) {
-                                    String line;
-                                    while ((line = reader.readLine()) != null) {
-                                        if (line.startsWith("data: ")) {
-                                            line = line.substring(6);
-                                        }
-                                        if (!line.isEmpty()) {
-                                            cbs.appendContent(BytesSupplier.wrap(line), false);
-                                        }
-                                    }
-                                } catch (IOException e) {
-                                    logger.warn("Failed run inference.", e);
-                                    cbs.appendContent(
-                                            BytesSupplier.wrap("connection abort exceptionally"),
-                                            false);
-                                } finally {
-                                    cbs.appendContent(new byte[0], true);
-                                    conn.disconnect();
-                                }
-                                return null;
-                            });
+                    CompletableFuture.supplyAsync(() -> handleStream(conn, cbs));
                 } else {
                     try (InputStream is = conn.getInputStream()) {
                         out.add(Utils.toByteArray(is));
@@ -208,7 +184,7 @@ public final class RpcClient {
             }
             return out;
         } finally {
-            if (!isSse) {
+            if (!isStream) {
                 conn.disconnect();
             }
         }
@@ -223,6 +199,35 @@ public final class RpcClient {
         return def;
     }
 
+    private static Void handleStream(HttpURLConnection conn, ChunkedBytesSupplier cbs) {
+        try (Reader r = new InputStreamReader(conn.getInputStream(), StandardCharsets.UTF_8);
+                BufferedReader reader = new BufferedReader(r)) {
+            String line;
+            StringBuilder sb = new StringBuilder();
+            while ((line = reader.readLine()) != null) {
+                if (line.startsWith("data: ")) {
+                    if (sb.length() > 0) {
+                        sb.append('\n');
+                    }
+                    sb.append(line.substring(6));
+                } else if (!line.isEmpty()) {
+                    // jsonlines
+                    cbs.appendContent(BytesSupplier.wrap(line), false);
+                } else if (sb.length() > 0) {
+                    cbs.appendContent(BytesSupplier.wrap(sb.toString()), false);
+                    sb.setLength(0);
+                }
+            }
+        } catch (IOException e) {
+            logger.warn("Failed run inference.", e);
+            cbs.appendContent(BytesSupplier.wrap("connection abort exceptionally"), false);
+        } finally {
+            cbs.appendContent(new byte[0], true);
+            conn.disconnect();
+        }
+        return null;
+    }
+
     static final class CaseInsensitiveKey {
         String key;
 


Test patch:
diff --git a/api/src/test/java/ai/djl/engine/rpc/RpcTranslatorFactoryTest.java b/api/src/test/java/ai/djl/engine/rpc/RpcTranslatorFactoryTest.java
index 0087760cd..308f1e321 100644
--- a/api/src/test/java/ai/djl/engine/rpc/RpcTranslatorFactoryTest.java
+++ b/api/src/test/java/ai/djl/engine/rpc/RpcTranslatorFactoryTest.java
@@ -26,6 +26,7 @@ import ai.djl.util.JsonUtils;
 import ai.djl.util.Pair;
 
 import com.google.gson.JsonElement;
+import com.google.gson.JsonObject;
 
 import org.testng.Assert;
 import org.testng.annotations.Test;
@@ -129,11 +130,21 @@ public class RpcTranslatorFactoryTest {
                 Assert.assertEquals(ret.data, "echo");
 
                 // test SSE response
-                server.setContent("data: line1\n\ndata:  line2\n\n");
+                server.setContent(
+                        "data: {\"value\":\n"
+                                + "data: \"line1\"}\n\n"
+                                + "data: {\"value\"\n"
+                                + "data: :\" line2\"}\n\n");
                 server.setContentType("text/event-stream");
                 ret = predictor.predict(testInput);
                 Assert.assertEquals(ret.data, "line1 line2");
 
+                // test jsonlines response
+                server.setContent("{\"value\":\"line1\"}\n{\"value\":\" line2\"}\n\n");
+                server.setContentType("application/jsonlines");
+                ret = predictor.predict(testInput);
+                Assert.assertEquals(ret.data, "line1 line2");
+
                 server.setCode(400);
                 Assert.assertThrows(() -> predictor.predict(testInput));
             }
@@ -156,7 +167,13 @@ public class RpcTranslatorFactoryTest {
 
         public static TestData fromJson(Iterator<String> it) {
             StringBuilder sb = new StringBuilder();
-            it.forEachRemaining(sb::append);
+            while (it.hasNext()) {
+                String line = it.next();
+                if (!line.isEmpty()) {
+                    JsonObject json = JsonUtils.GSON.fromJson(line, JsonObject.class);
+                    sb.append(json.get("value").getAsString());
+                }
+            }
             return new TestData(sb.toString());
         }