From 853b6800863ec82ad2dba0149cfa97883c0e4db6 Mon Sep 17 00:00:00 2001
From: Siddharth Venkatesan <siddhave@amazon.com>
Date: Thu, 6 Feb 2025 13:32:45 -0800
Subject: [PATCH] [fix][lmi] fix issue where streaming response that finishes
 with empty token is not completed

---
 .../python/src/main/java/ai/djl/python/engine/RollingBatch.java | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/engines/python/src/main/java/ai/djl/python/engine/RollingBatch.java b/engines/python/src/main/java/ai/djl/python/engine/RollingBatch.java
index 5eb191119..a7330813a 100644
--- a/engines/python/src/main/java/ai/djl/python/engine/RollingBatch.java
+++ b/engines/python/src/main/java/ai/djl/python/engine/RollingBatch.java
@@ -342,7 +342,7 @@ void addResponse(byte[] json, Map<String, String> properties) {
                         break;
                 }
             }
-            if ((nextToken == null || nextToken.isEmpty()) && code == null) {
+            if ((nextToken == null || nextToken.isEmpty()) && !last) {
                 // in non-streaming cases, we do not return content until generation is finished
                 return;
             }