diff --git a/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/http/AzureHttpResponseAdapter.java b/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/http/AzureHttpResponseAdapter.java index e299a62a1d48..b2a5e6c60e81 100644 --- a/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/http/AzureHttpResponseAdapter.java +++ b/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/http/AzureHttpResponseAdapter.java @@ -5,7 +5,6 @@ import com.azure.core.http.HttpHeader; import com.azure.core.http.HttpHeaders; -import com.azure.core.util.logging.ClientLogger; import com.openai.core.http.Headers; import com.openai.core.http.HttpResponse; @@ -17,8 +16,6 @@ */ final class AzureHttpResponseAdapter implements HttpResponse { - private static final ClientLogger LOGGER = new ClientLogger(AzureHttpResponseAdapter.class); - private final com.azure.core.http.HttpResponse azureResponse; /** @@ -42,7 +39,9 @@ public Headers headers() { @Override public InputStream body() { - return azureResponse.getBodyAsBinaryData().toStream(); + // replace with azureResponse.bodyStream() and delete FluxInputStream class from this package + // when new version of azure-core is released. + return new FluxInputStream(azureResponse.getBody()); } @Override diff --git a/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/http/FluxInputStream.java b/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/http/FluxInputStream.java new file mode 100644 index 000000000000..98b2657b4248 --- /dev/null +++ b/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/http/FluxInputStream.java @@ -0,0 +1,246 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.ai.agents.implementation.http; + +import com.azure.core.util.FluxUtil; +import com.azure.core.util.logging.ClientLogger; +import org.reactivestreams.Subscription; +import reactor.core.publisher.Flux; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.Buffer; +import java.nio.ByteBuffer; +import java.util.concurrent.locks.Condition; +import java.util.concurrent.locks.Lock; +import java.util.concurrent.locks.ReentrantLock; + +/** + * An InputStream that subscribes to a Flux. + */ +public class FluxInputStream extends InputStream { + + private static final ClientLogger LOGGER = new ClientLogger(FluxInputStream.class); + + // The data to subscribe to. + private final Flux data; + + // Subscription to request more data from as needed + private Subscription subscription; + + private ByteArrayInputStream buffer; + + private volatile boolean subscribed; + private volatile boolean fluxComplete; + private volatile boolean waitingForData; + + /* The following lock and condition variable is to synchronize access between the reader and the + reactor thread asynchronously reading data from the Flux. If no data is available, the reader + acquires the lock and waits on the dataAvailable condition variable. Once data is available + (or an error or completion event occurs) the reactor thread acquires the lock and signals that + data is available. */ + private final Lock lock; + private final Condition dataAvailable; + + private IOException lastError; + + /** + * Creates a new FluxInputStream + * + * @param data The data to subscribe to and read from. + */ + public FluxInputStream(Flux data) { + this.subscribed = false; + this.fluxComplete = false; + this.waitingForData = false; + this.data = data; + this.lock = new ReentrantLock(); + this.dataAvailable = lock.newCondition(); + } + + @Override + public int read() throws IOException { + byte[] ret = new byte[1]; + int count = read(ret, 0, 1); + return count == -1 ? -1 : (ret[0] & 0xFF); + } + + @Override + public int read(byte[] b, int off, int len) throws IOException { + validateParameters(b, off, len); + + /* If len is 0, then no bytes are read and 0 is returned. */ + if (len == 0) { + return 0; + } + /* Attempt to read at least one byte. If no byte is available because the stream is at end of file, + the value -1 is returned; otherwise, at least one byte is read and stored into b. */ + + /* Not subscribed? subscribe and block for data */ + if (!subscribed) { + blockForData(); + } + /* Now, we have subscribed. */ + /* At this point, buffer should not be null. If it is, that indicates either an error or completion event + was emitted by the Flux. */ + if (this.buffer == null) { // Only executed on first subscription. + if (this.lastError != null) { + throw LOGGER.logThrowableAsError(this.lastError); + } + if (this.fluxComplete) { + return -1; + } + throw LOGGER.logExceptionAsError(new IllegalStateException("An unexpected error occurred. No data was " + + "read from the stream but the stream did not indicate completion.")); + } + + /* Now we are guaranteed that buffer is SOMETHING. */ + /* No data is available in the buffer. */ + if (this.buffer.available() == 0) { + /* If the flux completed, there is no more data available to be read from the stream. Return -1. */ + if (this.fluxComplete) { + return -1; + } + /* Block current thread until data is available. */ + blockForData(); + } + + /* Data available in buffer, read the buffer. */ + if (this.buffer.available() > 0) { + return this.buffer.read(b, off, len); + } + + /* If the flux completed, there is no more data available to be read from the stream. Return -1. */ + if (this.fluxComplete) { + return -1; + } else { + throw LOGGER.logExceptionAsError(new IllegalStateException("An unexpected error occurred. No data was " + + "read from the stream but the stream did not indicate completion.")); + } + } + + @Override + public void close() throws IOException { + if (subscription != null) { + subscription.cancel(); + } + + if (this.buffer != null) { + this.buffer.close(); + } + super.close(); + if (this.lastError != null) { + throw LOGGER.logThrowableAsError(this.lastError); + } + } + + /** + * Request more data and wait on data to become available. + */ + private void blockForData() { + lock.lock(); + try { + waitingForData = true; + if (!subscribed) { + subscribeToData(); + } else { + subscription.request(1); + } + // Block current thread until data is available. + while (waitingForData) { + if (fluxComplete) { + break; + } else { + try { + dataAvailable.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw LOGGER.logExceptionAsError(new RuntimeException(e)); + } + } + } + } finally { + lock.unlock(); + } + } + + /** + * Subscribes to the data with a special subscriber. + */ + @SuppressWarnings("deprecation") + private void subscribeToData() { + this.data.filter(Buffer::hasRemaining) /* Filter to make sure only non empty byte buffers are emitted. */ + .onBackpressureBuffer() + .subscribe( + // ByteBuffer consumer + byteBuffer -> { + this.buffer = new ByteArrayInputStream(FluxUtil.byteBufferToArray(byteBuffer)); + lock.lock(); + try { + this.waitingForData = false; + // Signal the consumer when data is available. + dataAvailable.signal(); + } finally { + lock.unlock(); + } + }, + // Error consumer + throwable -> { + // Signal the consumer in case an error occurs (indicates we completed without data). + if (throwable instanceof IOException) { + this.lastError = (IOException) throwable; + } else { + this.lastError = new IOException(throwable); + } + signalOnCompleteOrError(); + }, + // Complete consumer + // Signal the consumer in case we completed without data. + this::signalOnCompleteOrError, + // Subscription consumer + subscription -> { + this.subscription = subscription; + this.subscribed = true; + this.subscription.request(1); + }); + } + + /** + * Signals to the subscriber when the flux completes without data (onCompletion or onError) + */ + private void signalOnCompleteOrError() { + this.fluxComplete = true; + lock.lock(); + try { + this.waitingForData = false; + dataAvailable.signal(); + } finally { + lock.unlock(); + } + } + + /** + * Validates parameters according to {@link InputStream#read(byte[], int, int)} spec. + * + * @param bytes the buffer into which the data is read. + * @param offset the start offset in array bytes at which the data is written. + * @param length the maximum number of bytes to read. + */ + private void validateParameters(byte[] bytes, int offset, int length) { + if (bytes == null) { + throw LOGGER.logExceptionAsError(new NullPointerException("'bytes' cannot be null")); + } + if (offset < 0) { + throw LOGGER.logExceptionAsError(new IndexOutOfBoundsException("'offset' cannot be less than 0")); + } + if (length < 0) { + throw LOGGER.logExceptionAsError(new IndexOutOfBoundsException("'length' cannot be less than 0")); + } + if (length > (bytes.length - offset)) { + throw LOGGER.logExceptionAsError( + new IndexOutOfBoundsException("'length' cannot be greater than 'bytes'.length - 'offset'")); + } + } +} diff --git a/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/http/HttpClientHelper.java b/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/http/HttpClientHelper.java index b0060f6301a8..67c01ba40d62 100644 --- a/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/http/HttpClientHelper.java +++ b/sdk/ai/azure-ai-agents/src/main/java/com/azure/ai/agents/implementation/http/HttpClientHelper.java @@ -21,6 +21,7 @@ import com.openai.core.http.HttpRequestBody; import com.openai.core.http.HttpResponse; import com.openai.errors.BadRequestException; +import reactor.core.scheduler.Schedulers; import com.openai.errors.InternalServerException; import com.openai.errors.NotFoundException; import com.openai.errors.OpenAIException; @@ -110,6 +111,10 @@ public CompletableFuture executeAsync(HttpRequest request, Request .flatMap(azureRequest -> this.httpPipeline.send(azureRequest, buildRequestContext(requestOptions))) .map(response -> (HttpResponse) new AzureHttpResponseAdapter(response)) .onErrorMap(HttpClientWrapper::mapAzureExceptionToOpenAI) + // publishOn moves the CompletableFuture completion (and all OpenAI SDK continuations that + // run synchronously on it) off the Netty/OkHttp I/O thread and onto a thread pool that + // is safe to block. + .publishOn(Schedulers.boundedElastic()) .toFuture(); } @@ -240,7 +245,7 @@ private static HttpHeaders toAzureHeaders(Headers sourceHeaders) { * @return Azure request {@link Context} */ private static Context buildRequestContext(RequestOptions requestOptions) { - Context context = new Context("azure-eagerly-read-response", true); + Context context = Context.NONE; Timeout timeout = requestOptions.getTimeout(); // we use "read" as it's the closest thing to the "response timeout" if (timeout != null && !timeout.read().isZero() && !timeout.read().isNegative()) { diff --git a/sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/implementation/http/FluxInputStreamTests.java b/sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/implementation/http/FluxInputStreamTests.java new file mode 100644 index 000000000000..4be3a3afe4eb --- /dev/null +++ b/sdk/ai/azure-ai-agents/src/test/java/com/azure/ai/agents/implementation/http/FluxInputStreamTests.java @@ -0,0 +1,143 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.ai.agents.implementation.http; + +import com.azure.core.exception.HttpResponseException; +import com.azure.core.http.HttpHeaders; +import com.azure.core.http.HttpResponse; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +import java.io.IOException; +import java.io.InputStream; +import java.io.UncheckedIOException; +import java.nio.ByteBuffer; +import java.nio.charset.Charset; +import java.util.ArrayList; +import java.util.List; +import java.util.stream.Stream; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +public class FluxInputStreamTests { + private static final int KB = 1024; + private static final int MB = KB * KB; + + /* Network tests to be performed by implementors of the FluxInputStream. */ + private Flux generateData(int num) { + List buffers = new ArrayList<>(); + for (int i = 0; i < num; i++) { + buffers.add(ByteBuffer.wrap(new byte[] { (byte) i })); + } + return Flux.fromIterable(buffers); + } + + @ParameterizedTest + @ValueSource(ints = { 1, 10, 100, KB, MB }) + public void fluxInputStreamMin(int num) throws IOException { + + try (InputStream is = new FluxInputStream(generateData(num))) { + byte[] bytes = new byte[num]; + int totalRead = 0; + int bytesRead = 0; + + while (bytesRead != -1 && totalRead < num) { + bytesRead = is.read(bytes, totalRead, num); + if (bytesRead != -1) { + totalRead += bytesRead; + num -= bytesRead; + } + } + + for (int i = 0; i < num; i++) { + assertEquals((byte) i, bytes[i]); + } + } + } + + @Test + public void fluxInputStreamWithEmptyByteBuffers() throws IOException { + int num = KB; + List buffers = new ArrayList<>(num * 2); + for (int i = 0; i < num; i++) { + buffers.add(ByteBuffer.wrap(new byte[] { (byte) i })); + buffers.add(ByteBuffer.wrap(new byte[0])); + } + + try (InputStream is = new FluxInputStream(Flux.fromIterable(buffers))) { + byte[] bytes = new byte[num]; + int totalRead = 0; + int bytesRead = 0; + + while (bytesRead != -1 && totalRead < num) { + bytesRead = is.read(bytes, totalRead, num); + if (bytesRead != -1) { + totalRead += bytesRead; + num -= bytesRead; + } + } + + for (int i = 0; i < num; i++) { + assertEquals((byte) i, bytes[i]); + } + } + } + + @ParameterizedTest + @MethodSource("fluxInputStreamErrorSupplier") + public void fluxInputStreamError(RuntimeException exception) { + assertThrows(IOException.class, () -> { + InputStream is = new FluxInputStream(Flux.error(exception)); + is.read(); + is.close(); + }); + } + + @SuppressWarnings("deprecation") + private static Stream fluxInputStreamErrorSupplier() { + HttpResponse httpResponse = new HttpResponse(null) { + @Override + public int getStatusCode() { + return 404; + } + + @Override + public String getHeaderValue(String name) { + return ""; + } + + @Override + public HttpHeaders getHeaders() { + return null; + } + + @Override + public Flux getBody() { + return null; + } + + @Override + public Mono getBodyAsByteArray() { + return null; + } + + @Override + public Mono getBodyAsString() { + return null; + } + + @Override + public Mono getBodyAsString(Charset charset) { + return null; + } + }; + return Stream.of(new IllegalArgumentException("Mock illegal argument exception."), + new HttpResponseException("Mock exception", httpResponse, null), + new UncheckedIOException(new IOException("Mock IO Exception."))); + } +} diff --git a/sdk/ai/azure-ai-projects/src/main/java/com/azure/ai/projects/implementation/http/AzureHttpResponseAdapter.java b/sdk/ai/azure-ai-projects/src/main/java/com/azure/ai/projects/implementation/http/AzureHttpResponseAdapter.java index 84db64020502..7af5eb323033 100644 --- a/sdk/ai/azure-ai-projects/src/main/java/com/azure/ai/projects/implementation/http/AzureHttpResponseAdapter.java +++ b/sdk/ai/azure-ai-projects/src/main/java/com/azure/ai/projects/implementation/http/AzureHttpResponseAdapter.java @@ -5,7 +5,6 @@ import com.azure.core.http.HttpHeader; import com.azure.core.http.HttpHeaders; -import com.azure.core.util.logging.ClientLogger; import com.openai.core.http.Headers; import com.openai.core.http.HttpResponse; @@ -17,8 +16,6 @@ */ final class AzureHttpResponseAdapter implements HttpResponse { - private static final ClientLogger LOGGER = new ClientLogger(AzureHttpResponseAdapter.class); - private final com.azure.core.http.HttpResponse azureResponse; /** @@ -42,7 +39,9 @@ public Headers headers() { @Override public InputStream body() { - return azureResponse.getBodyAsBinaryData().toStream(); + // replace with azureResponse.bodyStream() and delete FluxInputStream class from this package + // when new version of azure-core is released. + return new FluxInputStream(azureResponse.getBody()); } @Override diff --git a/sdk/ai/azure-ai-projects/src/main/java/com/azure/ai/projects/implementation/http/FluxInputStream.java b/sdk/ai/azure-ai-projects/src/main/java/com/azure/ai/projects/implementation/http/FluxInputStream.java new file mode 100644 index 000000000000..dfe88b19e290 --- /dev/null +++ b/sdk/ai/azure-ai-projects/src/main/java/com/azure/ai/projects/implementation/http/FluxInputStream.java @@ -0,0 +1,246 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.ai.projects.implementation.http; + +import com.azure.core.util.FluxUtil; +import com.azure.core.util.logging.ClientLogger; +import org.reactivestreams.Subscription; +import reactor.core.publisher.Flux; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.Buffer; +import java.nio.ByteBuffer; +import java.util.concurrent.locks.Condition; +import java.util.concurrent.locks.Lock; +import java.util.concurrent.locks.ReentrantLock; + +/** + * An InputStream that subscribes to a Flux. + */ +public class FluxInputStream extends InputStream { + + private static final ClientLogger LOGGER = new ClientLogger(FluxInputStream.class); + + // The data to subscribe to. + private final Flux data; + + // Subscription to request more data from as needed + private Subscription subscription; + + private ByteArrayInputStream buffer; + + private volatile boolean subscribed; + private volatile boolean fluxComplete; + private volatile boolean waitingForData; + + /* The following lock and condition variable is to synchronize access between the reader and the + reactor thread asynchronously reading data from the Flux. If no data is available, the reader + acquires the lock and waits on the dataAvailable condition variable. Once data is available + (or an error or completion event occurs) the reactor thread acquires the lock and signals that + data is available. */ + private final Lock lock; + private final Condition dataAvailable; + + private IOException lastError; + + /** + * Creates a new FluxInputStream + * + * @param data The data to subscribe to and read from. + */ + public FluxInputStream(Flux data) { + this.subscribed = false; + this.fluxComplete = false; + this.waitingForData = false; + this.data = data; + this.lock = new ReentrantLock(); + this.dataAvailable = lock.newCondition(); + } + + @Override + public int read() throws IOException { + byte[] ret = new byte[1]; + int count = read(ret, 0, 1); + return count == -1 ? -1 : (ret[0] & 0xFF); + } + + @Override + public int read(byte[] b, int off, int len) throws IOException { + validateParameters(b, off, len); + + /* If len is 0, then no bytes are read and 0 is returned. */ + if (len == 0) { + return 0; + } + /* Attempt to read at least one byte. If no byte is available because the stream is at end of file, + the value -1 is returned; otherwise, at least one byte is read and stored into b. */ + + /* Not subscribed? subscribe and block for data */ + if (!subscribed) { + blockForData(); + } + /* Now, we have subscribed. */ + /* At this point, buffer should not be null. If it is, that indicates either an error or completion event + was emitted by the Flux. */ + if (this.buffer == null) { // Only executed on first subscription. + if (this.lastError != null) { + throw LOGGER.logThrowableAsError(this.lastError); + } + if (this.fluxComplete) { + return -1; + } + throw LOGGER.logExceptionAsError(new IllegalStateException("An unexpected error occurred. No data was " + + "read from the stream but the stream did not indicate completion.")); + } + + /* Now we are guaranteed that buffer is SOMETHING. */ + /* No data is available in the buffer. */ + if (this.buffer.available() == 0) { + /* If the flux completed, there is no more data available to be read from the stream. Return -1. */ + if (this.fluxComplete) { + return -1; + } + /* Block current thread until data is available. */ + blockForData(); + } + + /* Data available in buffer, read the buffer. */ + if (this.buffer.available() > 0) { + return this.buffer.read(b, off, len); + } + + /* If the flux completed, there is no more data available to be read from the stream. Return -1. */ + if (this.fluxComplete) { + return -1; + } else { + throw LOGGER.logExceptionAsError(new IllegalStateException("An unexpected error occurred. No data was " + + "read from the stream but the stream did not indicate completion.")); + } + } + + @Override + public void close() throws IOException { + if (subscription != null) { + subscription.cancel(); + } + + if (this.buffer != null) { + this.buffer.close(); + } + super.close(); + if (this.lastError != null) { + throw LOGGER.logThrowableAsError(this.lastError); + } + } + + /** + * Request more data and wait on data to become available. + */ + private void blockForData() { + lock.lock(); + try { + waitingForData = true; + if (!subscribed) { + subscribeToData(); + } else { + subscription.request(1); + } + // Block current thread until data is available. + while (waitingForData) { + if (fluxComplete) { + break; + } else { + try { + dataAvailable.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw LOGGER.logExceptionAsError(new RuntimeException(e)); + } + } + } + } finally { + lock.unlock(); + } + } + + /** + * Subscribes to the data with a special subscriber. + */ + @SuppressWarnings("deprecation") + private void subscribeToData() { + this.data.filter(Buffer::hasRemaining) /* Filter to make sure only non empty byte buffers are emitted. */ + .onBackpressureBuffer() + .subscribe( + // ByteBuffer consumer + byteBuffer -> { + this.buffer = new ByteArrayInputStream(FluxUtil.byteBufferToArray(byteBuffer)); + lock.lock(); + try { + this.waitingForData = false; + // Signal the consumer when data is available. + dataAvailable.signal(); + } finally { + lock.unlock(); + } + }, + // Error consumer + throwable -> { + // Signal the consumer in case an error occurs (indicates we completed without data). + if (throwable instanceof IOException) { + this.lastError = (IOException) throwable; + } else { + this.lastError = new IOException(throwable); + } + signalOnCompleteOrError(); + }, + // Complete consumer + // Signal the consumer in case we completed without data. + this::signalOnCompleteOrError, + // Subscription consumer + subscription -> { + this.subscription = subscription; + this.subscribed = true; + this.subscription.request(1); + }); + } + + /** + * Signals to the subscriber when the flux completes without data (onCompletion or onError) + */ + private void signalOnCompleteOrError() { + this.fluxComplete = true; + lock.lock(); + try { + this.waitingForData = false; + dataAvailable.signal(); + } finally { + lock.unlock(); + } + } + + /** + * Validates parameters according to {@link InputStream#read(byte[], int, int)} spec. + * + * @param bytes the buffer into which the data is read. + * @param offset the start offset in array bytes at which the data is written. + * @param length the maximum number of bytes to read. + */ + private void validateParameters(byte[] bytes, int offset, int length) { + if (bytes == null) { + throw LOGGER.logExceptionAsError(new NullPointerException("'bytes' cannot be null")); + } + if (offset < 0) { + throw LOGGER.logExceptionAsError(new IndexOutOfBoundsException("'offset' cannot be less than 0")); + } + if (length < 0) { + throw LOGGER.logExceptionAsError(new IndexOutOfBoundsException("'length' cannot be less than 0")); + } + if (length > (bytes.length - offset)) { + throw LOGGER.logExceptionAsError( + new IndexOutOfBoundsException("'length' cannot be greater than 'bytes'.length - 'offset'")); + } + } +} diff --git a/sdk/ai/azure-ai-projects/src/main/java/com/azure/ai/projects/implementation/http/HttpClientHelper.java b/sdk/ai/azure-ai-projects/src/main/java/com/azure/ai/projects/implementation/http/HttpClientHelper.java index 013e7d770eba..278288d73fb9 100644 --- a/sdk/ai/azure-ai-projects/src/main/java/com/azure/ai/projects/implementation/http/HttpClientHelper.java +++ b/sdk/ai/azure-ai-projects/src/main/java/com/azure/ai/projects/implementation/http/HttpClientHelper.java @@ -30,6 +30,7 @@ import com.openai.errors.UnexpectedStatusCodeException; import com.openai.errors.UnprocessableEntityException; import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; import java.io.ByteArrayOutputStream; import java.net.MalformedURLException; @@ -110,6 +111,10 @@ public CompletableFuture executeAsync(HttpRequest request, Request .flatMap(azureRequest -> this.httpPipeline.send(azureRequest, buildRequestContext(requestOptions))) .map(response -> (HttpResponse) new AzureHttpResponseAdapter(response)) .onErrorMap(HttpClientWrapper::mapAzureExceptionToOpenAI) + // publishOn moves the CompletableFuture completion (and all OpenAI SDK continuations that + // run synchronously on it) off the Netty/OkHttp I/O thread and onto a thread pool that + // is safe to block. + .publishOn(Schedulers.boundedElastic()) .toFuture(); } @@ -240,7 +245,7 @@ private static HttpHeaders toAzureHeaders(Headers sourceHeaders) { * @return Azure request {@link Context} */ private static Context buildRequestContext(RequestOptions requestOptions) { - Context context = new Context("azure-eagerly-read-response", true); + Context context = Context.NONE; Timeout timeout = requestOptions.getTimeout(); // we use "read" as it's the closest thing to the "response timeout" if (timeout != null && !timeout.read().isZero() && !timeout.read().isNegative()) { diff --git a/sdk/ai/azure-ai-projects/src/test/java/com/azure/ai/projects/implementation/http/FluxInputStreamTests.java b/sdk/ai/azure-ai-projects/src/test/java/com/azure/ai/projects/implementation/http/FluxInputStreamTests.java new file mode 100644 index 000000000000..919376c028d0 --- /dev/null +++ b/sdk/ai/azure-ai-projects/src/test/java/com/azure/ai/projects/implementation/http/FluxInputStreamTests.java @@ -0,0 +1,143 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.ai.projects.implementation.http; + +import com.azure.core.exception.HttpResponseException; +import com.azure.core.http.HttpHeaders; +import com.azure.core.http.HttpResponse; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +import java.io.IOException; +import java.io.InputStream; +import java.io.UncheckedIOException; +import java.nio.ByteBuffer; +import java.nio.charset.Charset; +import java.util.ArrayList; +import java.util.List; +import java.util.stream.Stream; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +public class FluxInputStreamTests { + private static final int KB = 1024; + private static final int MB = KB * KB; + + /* Network tests to be performed by implementors of the FluxInputStream. */ + private Flux generateData(int num) { + List buffers = new ArrayList<>(); + for (int i = 0; i < num; i++) { + buffers.add(ByteBuffer.wrap(new byte[] { (byte) i })); + } + return Flux.fromIterable(buffers); + } + + @ParameterizedTest + @ValueSource(ints = { 1, 10, 100, KB, MB }) + public void fluxInputStreamMin(int num) throws IOException { + + try (InputStream is = new FluxInputStream(generateData(num))) { + byte[] bytes = new byte[num]; + int totalRead = 0; + int bytesRead = 0; + + while (bytesRead != -1 && totalRead < num) { + bytesRead = is.read(bytes, totalRead, num); + if (bytesRead != -1) { + totalRead += bytesRead; + num -= bytesRead; + } + } + + for (int i = 0; i < num; i++) { + assertEquals((byte) i, bytes[i]); + } + } + } + + @Test + public void fluxInputStreamWithEmptyByteBuffers() throws IOException { + int num = KB; + List buffers = new ArrayList<>(num * 2); + for (int i = 0; i < num; i++) { + buffers.add(ByteBuffer.wrap(new byte[] { (byte) i })); + buffers.add(ByteBuffer.wrap(new byte[0])); + } + + try (InputStream is = new FluxInputStream(Flux.fromIterable(buffers))) { + byte[] bytes = new byte[num]; + int totalRead = 0; + int bytesRead = 0; + + while (bytesRead != -1 && totalRead < num) { + bytesRead = is.read(bytes, totalRead, num); + if (bytesRead != -1) { + totalRead += bytesRead; + num -= bytesRead; + } + } + + for (int i = 0; i < num; i++) { + assertEquals((byte) i, bytes[i]); + } + } + } + + @ParameterizedTest + @MethodSource("fluxInputStreamErrorSupplier") + public void fluxInputStreamError(RuntimeException exception) { + assertThrows(IOException.class, () -> { + InputStream is = new FluxInputStream(Flux.error(exception)); + is.read(); + is.close(); + }); + } + + @SuppressWarnings("deprecation") + private static Stream fluxInputStreamErrorSupplier() { + HttpResponse httpResponse = new HttpResponse(null) { + @Override + public int getStatusCode() { + return 404; + } + + @Override + public String getHeaderValue(String name) { + return ""; + } + + @Override + public HttpHeaders getHeaders() { + return null; + } + + @Override + public Flux getBody() { + return null; + } + + @Override + public Mono getBodyAsByteArray() { + return null; + } + + @Override + public Mono getBodyAsString() { + return null; + } + + @Override + public Mono getBodyAsString(Charset charset) { + return null; + } + }; + return Stream.of(new IllegalArgumentException("Mock illegal argument exception."), + new HttpResponseException("Mock exception", httpResponse, null), + new UncheckedIOException(new IOException("Mock IO Exception."))); + } +}