[Streaming] Fix issues when not using try-with-resource

foundation-models/openai/src/main/java/com/sap/ai/sdk/foundationmodels/openai/IterableStreamConverter.java

@@ -0,0 +1,120 @@
+package com.sap.ai.sdk.foundationmodels.openai;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static java.util.Spliterator.NONNULL;
+import static java.util.Spliterator.ORDERED;
+
+import io.vavr.control.Try;
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.util.Iterator;
+import java.util.NoSuchElementException;
+import java.util.Spliterators;
+import java.util.concurrent.Callable;
+import java.util.function.Function;
+import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+import lombok.AccessLevel;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.hc.core5.http.HttpEntity;
+
+/**
+ * Internal utility class to convert from a reading handler to {@link Iterable} and {@link Stream}.
+ *
+ * <p><strong>Note:</strong> All operations are sequential in nature. Thread safety is not
+ * guaranteed.
+ *
+ * @param <T> Iterated item type.
+ */
+@Slf4j
+@RequiredArgsConstructor(access = AccessLevel.PRIVATE)
+class IterableStreamConverter<T> implements Iterator<T> {
+  /** see DEFAULT_CHAR_BUFFER_SIZE in {@link BufferedReader} * */
+  static final int BUFFER_SIZE = 8192;
+
+  /** Read next entry for Stream or {@code null} when no further entry can be read. */
+  private final Callable<T> readHandler;
+
+  /** Close handler to be called when Stream terminated. */
+  private final Runnable stopHandler;
+
+  /** Error handler to be called when Stream is interrupted. */
+  private final Function<Exception, RuntimeException> errorHandler;
+
+  private boolean isDone = false;
+  private boolean isNextFetched = false;
+  private T next = null;
+
+  @SuppressWarnings("checkstyle:IllegalCatch")
+  @Override
+  public boolean hasNext() {
+    if (isDone) {
+      return false;
+    }
+    if (isNextFetched) {
+      return true;
+    }
+    try {
+      next = readHandler.call();
+      isNextFetched = true;
+      if (next == null) {
+        isDone = true;
+        stopHandler.run();
+      }
+    } catch (final Exception e) {
+      isDone = true;
+      stopHandler.run();
+      log.debug("Error while reading next element.", e);
+      throw errorHandler.apply(e);
+    }
+    return !isDone;
+  }
+
+  @Override
+  public T next() {
+    if (next == null && !hasNext()) {
+      throw new NoSuchElementException(); // normally not reached with Stream API
+    }
+    isNextFetched = false;
+    return next;
+  }
+
+  /**
+   * Create a sequential Stream of lines from an HTTP response string (UTF-8). The underlying {@link
+   * InputStream} is closed, when the resulting Stream is closed (e.g. via try-with-resources) or
+   * when an exception occurred.
+   *
+   * @param entity The HTTP entity object.
+   * @return A sequential Stream object.
+   * @throws OpenAiClientException if the provided HTTP entity object is {@code null} or empty.
+   */
+  @SuppressWarnings("PMD.CloseResource") // Stream is closed automatically when consumed
+  @Nonnull
+  static Stream<String> lines(@Nullable final HttpEntity entity) throws OpenAiClientException {
+    if (entity == null) {
+      throw new OpenAiClientException("OpenAI response was empty.");
+    }
+
+    final InputStream inputStream;
+    try {
+      inputStream = entity.getContent();
+    } catch (final IOException e) {
+      throw new OpenAiClientException("Failed to read response content.", e);
+    }
+
+    final var reader = new BufferedReader(new InputStreamReader(inputStream, UTF_8), BUFFER_SIZE);
+    final Runnable closeHandler =
+        () -> Try.run(reader::close).onFailure(e -> log.error("Could not close input stream", e));
+    final Function<Exception, RuntimeException> errHandler =
+        e -> new OpenAiClientException("Parsing response content was interrupted.", e);
+
+    final var iterator = new IterableStreamConverter<>(reader::readLine, closeHandler, errHandler);
+    final var spliterator = Spliterators.spliteratorUnknownSize(iterator, ORDERED | NONNULL);
+    return StreamSupport.stream(spliterator, /* NOT PARALLEL */ false).onClose(closeHandler);
+  }
+}

foundation-models/openai/src/main/java/com/sap/ai/sdk/foundationmodels/openai/OpenAiStreamingHandler.java

@@ -5,57 +5,31 @@
 import static com.sap.ai.sdk.foundationmodels.openai.OpenAiResponseHandler.parseErrorAndThrow;
 
 import com.sap.ai.sdk.foundationmodels.openai.model.StreamedDelta;
-import io.vavr.control.Try;
-import java.io.BufferedReader;
 import java.io.IOException;
-import java.io.InputStream;
-import java.io.InputStreamReader;
-import java.nio.charset.StandardCharsets;
 import java.util.stream.Stream;
 import javax.annotation.Nonnull;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.hc.core5.http.ClassicHttpResponse;
-import org.apache.hc.core5.http.HttpEntity;
 
 @Slf4j
 @RequiredArgsConstructor
 class OpenAiStreamingHandler<D extends StreamedDelta> {
 
   @Nonnull private final Class<D> deltaType;
 
-  @Nonnull
-  Stream<D> handleResponse(@Nonnull final ClassicHttpResponse response)
-      throws OpenAiClientException {
-    if (response.getCode() >= 300) {
-      buildExceptionAndThrow(response);
-    }
-    return parseResponse(response);
-  }
-
   /**
    * @param response The response to process
    * @return A {@link Stream} of a model class instantiated from the response
-   * @author stippi
    */
-  // The stream is closed by the user of the Stream
-  @SuppressWarnings("PMD.CloseResource")
-  private Stream<D> parseResponse(@Nonnull final ClassicHttpResponse response)
+  @SuppressWarnings("PMD.CloseResource") // Stream is closed automatically when consumed
+  @Nonnull
+  Stream<D> handleResponse(@Nonnull final ClassicHttpResponse response)
       throws OpenAiClientException {
-    final HttpEntity responseEntity = response.getEntity();
-    if (responseEntity == null) {
-      throw new OpenAiClientException("Response from OpenAI model was empty.");
-    }
-    final InputStream inputStream;
-    try {
-      inputStream = responseEntity.getContent();
-    } catch (IOException e) {
-      throw new OpenAiClientException("Failed to read response content.", e);
+    if (response.getCode() >= 300) {
+      buildExceptionAndThrow(response);
     }
-    final var br = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8));
-
-    // https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#event_stream_format
-    return br.lines()
+    return IterableStreamConverter.lines(response.getEntity())
         // half of the lines are empty newlines, the last line is "data: [DONE]"
         .filter(line -> !line.isEmpty() && !"data: [DONE]".equals(line.trim()))
         .peek(
@@ -74,10 +48,6 @@ private Stream<D> parseResponse(@Nonnull final ClassicHttpResponse response)
                 log.error("Failed to parse the following response from OpenAI model: {}", line);
                 throw new OpenAiClientException("Failed to parse delta message: " + line, e);
               }
-            })
-        .onClose(
-            () ->
-                Try.run(inputStream::close)
-                    .onFailure(e -> log.error("Could not close HTTP input stream", e)));
+            });
   }
 }

foundation-models/openai/src/test/java/com/sap/ai/sdk/foundationmodels/openai/IterableStreamConverterTest.java

@@ -0,0 +1,106 @@
+package com.sap.ai.sdk.foundationmodels.openai;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
+import java.util.concurrent.atomic.AtomicInteger;
+import lombok.SneakyThrows;
+import org.apache.hc.core5.http.ContentType;
+import org.apache.hc.core5.http.io.entity.InputStreamEntity;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+
+public class IterableStreamConverterTest {
+  @SneakyThrows
+  @Test
+  @DisplayName("Stream is fully consumed")
+  void testLines() {
+    final var TEMPLATE = "THIS\nIS\nA\nTEST\n";
+    final var input = TEMPLATE.repeat(IterableStreamConverter.BUFFER_SIZE);
+    final var inputStream = spy(new ByteArrayInputStream(input.getBytes(StandardCharsets.UTF_8)));
+    final var entity = new InputStreamEntity(inputStream, ContentType.TEXT_PLAIN);
+
+    final var sut = IterableStreamConverter.lines(entity);
+    verify(inputStream, never()).read();
+    verify(inputStream, never()).read(any());
+    verify(inputStream, never()).read(any(), anyInt(), anyInt());
+
+    final var streamCounter = new AtomicInteger(0);
+    sut.peek(s -> streamCounter.incrementAndGet())
+        .forEach(
+            s ->
+                assertThat(s)
+                    .containsAnyOf("THIS", "IS", "A", "TEST")
+                    .doesNotContainAnyWhitespaces());
+
+    assertThat(streamCounter).hasValue(IterableStreamConverter.BUFFER_SIZE * 4);
+    verify(inputStream, times(TEMPLATE.length() + 1))
+        .read(any(), anyInt(), eq(IterableStreamConverter.BUFFER_SIZE));
+    verify(inputStream, times(1)).close();
+  }
+
+  @SneakyThrows
+  @Test
+  @DisplayName("Stream may only read first entry without closing")
+  void testLinesFindFirst() {
+    final var TEMPLATE = "Foo Bar\n";
+    final var inputStream = mock(InputStream.class);
+    when(inputStream.read(any(), anyInt(), anyInt()))
+        .thenAnswer(
+            arg -> {
+              byte[] ar = arg.getArgument(0, byte[].class);
+              byte[] bytes = TEMPLATE.getBytes(StandardCharsets.UTF_8);
+              for (int i = 0; i < ar.length; i++) ar[i] = bytes[i % bytes.length];
+              return ar.length;
+            });
+
+    final var entity = new InputStreamEntity(inputStream, ContentType.TEXT_PLAIN);
+
+    final var sut = IterableStreamConverter.lines(entity);
+    assertThat(sut.findFirst()).contains("Foo Bar");
+    verify(inputStream, times(1)).read(any(), anyInt(), anyInt());
+    verify(inputStream, never()).close();
+  }
+
+  @SneakyThrows
+  @Test
+  @DisplayName("Stream may close unexpectedly")
+  void testLinesThrows() {
+    final var TEMPLATE = "Foo Bar\n";
+    final var inputStream = mock(InputStream.class);
+    when(inputStream.read(any(), anyInt(), anyInt()))
+        .thenAnswer(
+            arg -> {
+              byte[] ar = arg.getArgument(0, byte[].class);
+              byte[] bytes = TEMPLATE.getBytes(StandardCharsets.UTF_8);
+              for (int i = 0; i < ar.length; i++) ar[i] = bytes[i % bytes.length];
+              return ar.length;
+            })
+        .thenThrow(new IOException("Ups!"));
+
+    final var entity = new InputStreamEntity(inputStream, ContentType.TEXT_PLAIN);
+
+    final var sut = IterableStreamConverter.lines(entity);
+    assertThatThrownBy(sut::count)
+        .isInstanceOf(OpenAiClientException.class)
+        .hasMessage("Parsing response content was interrupted.")
+        .cause()
+        .isInstanceOf(IOException.class)
+        .hasMessage("Ups!");
+    verify(inputStream, times(2)).read(any(), anyInt(), anyInt());
+    verify(inputStream, times(1)).close();
+  }
+}