diff --git a/openai-java-core/src/main/kotlin/com/openai/core/handlers/StreamHandler.kt b/openai-java-core/src/main/kotlin/com/openai/core/handlers/StreamHandler.kt index 13a1949a6..809a4e57c 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/handlers/StreamHandler.kt +++ b/openai-java-core/src/main/kotlin/com/openai/core/handlers/StreamHandler.kt @@ -8,6 +8,7 @@ import com.openai.core.http.PhantomReachableClosingStreamResponse import com.openai.core.http.StreamResponse import com.openai.errors.OpenAIIoException import java.io.IOException +import java.util.Optional import java.util.stream.Stream import kotlin.streams.asStream @@ -18,6 +19,7 @@ internal fun streamHandler( object : Handler> { override fun handle(response: HttpResponse): StreamResponse { + val requestId = response.requestId() val reader = response.body().bufferedReader() val sequence = // Wrap in a `CloseableSequence` to avoid performing a read on the `reader` @@ -40,6 +42,8 @@ internal fun streamHandler( return PhantomReachableClosingStreamResponse( object : StreamResponse { + override fun requestId(): Optional = requestId + override fun stream(): Stream = sequence.asStream() override fun close() { diff --git a/openai-java-core/src/main/kotlin/com/openai/core/http/AsyncStreamResponse.kt b/openai-java-core/src/main/kotlin/com/openai/core/http/AsyncStreamResponse.kt index 019dc6aff..47e608c17 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/http/AsyncStreamResponse.kt +++ b/openai-java-core/src/main/kotlin/com/openai/core/http/AsyncStreamResponse.kt @@ -1,8 +1,11 @@ package com.openai.core.http import com.openai.core.http.AsyncStreamResponse.Handler +import com.openai.errors.OpenAIServiceException import java.util.Optional +import java.util.concurrent.CancellationException import java.util.concurrent.CompletableFuture +import java.util.concurrent.CompletionException import java.util.concurrent.Executor import java.util.concurrent.atomic.AtomicReference @@ -13,6 +16,14 @@ import java.util.concurrent.atomic.AtomicReference */ interface AsyncStreamResponse { + /** + * Returns the value of the `x-request-id` header, or an empty [Optional] if it's unavailable. + * + * This method does not wait for response headers. Its result is empty until those headers are + * available, and may remain empty if the response has no `x-request-id` header. + */ + fun requestId(): Optional = Optional.empty() + /** * Registers [handler] to be called for events of this stream. * @@ -78,6 +89,19 @@ internal fun CompletableFuture>.toAsync(streamHandlerExecu } } + override fun requestId(): Optional { + val streamResponse = + try { + this@toAsync.getNow(null) + } catch (error: CompletionException) { + return requestIdFromError(error) + } catch (_: CancellationException) { + return Optional.empty() + } + + return streamResponse?.requestId() ?: Optional.empty() + } + override fun subscribe(handler: Handler): AsyncStreamResponse = subscribe(handler, streamHandlerExecutor) @@ -150,6 +174,14 @@ internal fun CompletableFuture>.toAsync(streamHandlerExecu } ) +private tailrec fun requestIdFromError(error: Throwable): Optional = + when { + error is CompletionException && error.cause != null -> requestIdFromError(error.cause!!) + error is OpenAIServiceException -> + Optional.ofNullable(error.headers().values("x-request-id").firstOrNull()) + else -> Optional.empty() + } + private enum class State { NEW, SUBSCRIBED, diff --git a/openai-java-core/src/main/kotlin/com/openai/core/http/PhantomReachableClosingAsyncStreamResponse.kt b/openai-java-core/src/main/kotlin/com/openai/core/http/PhantomReachableClosingAsyncStreamResponse.kt index 8c32247d8..f41128e55 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/http/PhantomReachableClosingAsyncStreamResponse.kt +++ b/openai-java-core/src/main/kotlin/com/openai/core/http/PhantomReachableClosingAsyncStreamResponse.kt @@ -25,6 +25,8 @@ internal class PhantomReachableClosingAsyncStreamResponse( closeWhenPhantomReachable(reachabilityTracker, asyncStreamResponse::close) } + override fun requestId(): Optional = asyncStreamResponse.requestId() + override fun subscribe(handler: Handler): AsyncStreamResponse = apply { asyncStreamResponse.subscribe(TrackedHandler(handler, reachabilityTracker)) } diff --git a/openai-java-core/src/main/kotlin/com/openai/core/http/PhantomReachableClosingStreamResponse.kt b/openai-java-core/src/main/kotlin/com/openai/core/http/PhantomReachableClosingStreamResponse.kt index 4321edb49..68da7809b 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/http/PhantomReachableClosingStreamResponse.kt +++ b/openai-java-core/src/main/kotlin/com/openai/core/http/PhantomReachableClosingStreamResponse.kt @@ -1,6 +1,7 @@ package com.openai.core.http import com.openai.core.closeWhenPhantomReachable +import java.util.Optional import java.util.stream.Stream /** @@ -15,6 +16,8 @@ internal class PhantomReachableClosingStreamResponse( closeWhenPhantomReachable(this, streamResponse) } + override fun requestId(): Optional = streamResponse.requestId() + override fun stream(): Stream = streamResponse.stream() override fun close() = streamResponse.close() diff --git a/openai-java-core/src/main/kotlin/com/openai/core/http/StreamResponse.kt b/openai-java-core/src/main/kotlin/com/openai/core/http/StreamResponse.kt index 821f21db1..72284c7f0 100644 --- a/openai-java-core/src/main/kotlin/com/openai/core/http/StreamResponse.kt +++ b/openai-java-core/src/main/kotlin/com/openai/core/http/StreamResponse.kt @@ -1,9 +1,15 @@ package com.openai.core.http +import java.util.Optional import java.util.stream.Stream interface StreamResponse : AutoCloseable { + /** + * Returns the value of the `x-request-id` header, or an empty [Optional] if it's unavailable. + */ + fun requestId(): Optional = Optional.empty() + fun stream(): Stream /** Overridden from [AutoCloseable] to not have a checked exception in its signature. */ @@ -13,6 +19,8 @@ interface StreamResponse : AutoCloseable { @JvmSynthetic internal fun StreamResponse.map(transform: (T) -> R): StreamResponse = object : StreamResponse { + override fun requestId(): Optional = this@map.requestId() + override fun stream(): Stream = this@map.stream().map(transform) override fun close() = this@map.close() diff --git a/openai-java-core/src/test/kotlin/com/openai/core/handlers/StreamHandlerTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/handlers/StreamHandlerTest.kt index 1ebcb16d7..34000cfad 100644 --- a/openai-java-core/src/test/kotlin/com/openai/core/handlers/StreamHandlerTest.kt +++ b/openai-java-core/src/test/kotlin/com/openai/core/handlers/StreamHandlerTest.kt @@ -22,6 +22,22 @@ internal class StreamHandlerTest { assertThat(lines).containsExactly("a", "bb", "ccc", "dddd") } + @Test + fun streamHandler_exposesRequestIdAfterClose() { + val handler = streamHandler { _, lines -> yieldAll(lines) } + val streamResponse = + handler.handle( + httpResponse( + "a\n".byteInputStream(), + Headers.builder().put("x-request-id", "req_123").build(), + ) + ) + + assertThat(streamResponse.requestId()).contains("req_123") + streamResponse.close() + assertThat(streamResponse.requestId()).contains("req_123") + } + @Test fun streamHandler_whenClosedEarly_stopsYielding() { val handler = streamHandler { _, lines -> yieldAll(lines) } @@ -68,12 +84,15 @@ internal class StreamHandlerTest { assertThat(e).isSameAs(ioException) } - private fun httpResponse(body: InputStream): HttpResponse = + private fun httpResponse( + body: InputStream, + headers: Headers = Headers.builder().build(), + ): HttpResponse = object : HttpResponse { override fun statusCode(): Int = 0 - override fun headers(): Headers = Headers.builder().build() + override fun headers(): Headers = headers override fun body(): InputStream = body diff --git a/openai-java-core/src/test/kotlin/com/openai/core/http/AsyncStreamResponseTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/http/AsyncStreamResponseTest.kt index 976bcd6da..2673c70a9 100644 --- a/openai-java-core/src/test/kotlin/com/openai/core/http/AsyncStreamResponseTest.kt +++ b/openai-java-core/src/test/kotlin/com/openai/core/http/AsyncStreamResponseTest.kt @@ -1,7 +1,9 @@ package com.openai.core.http +import com.openai.errors.BadRequestException import java.util.* import java.util.concurrent.CompletableFuture +import java.util.concurrent.CompletionException import java.util.concurrent.Executor import java.util.stream.Stream import kotlin.streams.asStream @@ -45,6 +47,78 @@ internal class AsyncStreamResponseTest { } private val handler = mock>() + @Test + fun requestId_whenStreamFutureIsIncomplete_isEmpty() { + val asyncStreamResponse = CompletableFuture>().toAsync(executor) + + assertThat(asyncStreamResponse.requestId()).isEmpty + } + + @Test + fun requestId_whenStreamFutureCompletes_isAvailable() { + val future = CompletableFuture>() + val asyncStreamResponse = future.toAsync(executor) + doReturn(Optional.of("req_123")).whenever(streamResponse).requestId() + + future.complete(streamResponse) + + assertThat(asyncStreamResponse.requestId()).contains("req_123") + } + + @Test + fun requestId_whenHandlerRuns_isAvailable() { + val future = CompletableFuture>() + val asyncStreamResponse = future.toAsync(executor) + val requestIds = mutableListOf>() + doReturn(Optional.of("req_123")).whenever(streamResponse).requestId() + asyncStreamResponse.subscribe { requestIds.add(asyncStreamResponse.requestId()) } + + future.complete(streamResponse) + + assertThat(requestIds).containsOnly(Optional.of("req_123")) + } + + @Test + fun requestId_whenStreamFutureFailsWithServiceException_isAvailableInExceptionally() { + val future = CompletableFuture>() + val asyncStreamResponse = future.toAsync(executor) + val requestIds = mutableListOf>() + val error = + BadRequestException.builder() + .headers(Headers.builder().put("x-request-id", "req_error").build()) + .build() + val assertion = + asyncStreamResponse.onCompleteFuture().exceptionally { + requestIds.add(asyncStreamResponse.requestId()) + null + } + + future.completeExceptionally(CompletionException(error)) + + assertion.join() + assertThat(requestIds).containsExactly(Optional.of("req_error")) + } + + @Test + fun requestId_whenStreamFutureFailsBeforeResponse_isEmpty() { + val future = CompletableFuture>() + val asyncStreamResponse = future.toAsync(executor) + + future.completeExceptionally(ERROR) + + assertThat(asyncStreamResponse.requestId()).isEmpty + } + + @Test + fun requestId_whenStreamFutureIsCancelled_isEmpty() { + val future = CompletableFuture>() + val asyncStreamResponse = future.toAsync(executor) + + future.cancel(false) + + assertThat(asyncStreamResponse.requestId()).isEmpty + } + @Test fun subscribe_whenAlreadySubscribed_throws() { val asyncStreamResponse = CompletableFuture>().toAsync(executor) diff --git a/openai-java-core/src/test/kotlin/com/openai/core/http/StreamResponseTest.kt b/openai-java-core/src/test/kotlin/com/openai/core/http/StreamResponseTest.kt new file mode 100644 index 000000000..d03a5579d --- /dev/null +++ b/openai-java-core/src/test/kotlin/com/openai/core/http/StreamResponseTest.kt @@ -0,0 +1,36 @@ +package com.openai.core.http + +import java.util.Optional +import java.util.stream.Stream +import org.assertj.core.api.Assertions.assertThat +import org.junit.jupiter.api.Test + +internal class StreamResponseTest { + + @Test + fun requestId_byDefault_isEmpty() { + val response = streamResponse() + + assertThat(response.requestId()).isEmpty + } + + @Test + fun map_preservesRequestId() { + val response = streamResponse(Optional.of("req_123")) + + val mappedResponse = response.map(String::length).map(Int::toString) + + assertThat(mappedResponse.requestId()).contains("req_123") + } + + private fun streamResponse( + requestId: Optional = Optional.empty() + ): StreamResponse = + object : StreamResponse { + override fun requestId(): Optional = requestId + + override fun stream(): Stream = Stream.empty() + + override fun close() {} + } +} diff --git a/openai-java-core/src/test/kotlin/com/openai/services/StreamingRequestIdTest.kt b/openai-java-core/src/test/kotlin/com/openai/services/StreamingRequestIdTest.kt new file mode 100644 index 000000000..cf2a66b99 --- /dev/null +++ b/openai-java-core/src/test/kotlin/com/openai/services/StreamingRequestIdTest.kt @@ -0,0 +1,109 @@ +package com.openai.services + +import com.github.tomakehurst.wiremock.client.WireMock.anyUrl +import com.github.tomakehurst.wiremock.client.WireMock.ok +import com.github.tomakehurst.wiremock.client.WireMock.post +import com.github.tomakehurst.wiremock.client.WireMock.status +import com.github.tomakehurst.wiremock.client.WireMock.stubFor +import com.github.tomakehurst.wiremock.junit5.WireMockRuntimeInfo +import com.github.tomakehurst.wiremock.junit5.WireMockTest +import com.openai.client.OpenAIClient +import com.openai.client.OpenAIClientAsync +import com.openai.client.okhttp.OpenAIOkHttpClient +import com.openai.client.okhttp.OpenAIOkHttpClientAsync +import com.openai.models.completions.CompletionCreateParams +import org.assertj.core.api.Assertions.assertThat +import org.junit.jupiter.api.BeforeEach +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.parallel.ResourceLock + +@WireMockTest +@ResourceLock("https://github.com/wiremock/wiremock/issues/169") +internal class StreamingRequestIdTest { + + private lateinit var blockingClient: OpenAIClient + private lateinit var asyncClient: OpenAIClientAsync + + @BeforeEach + fun beforeEach(wmRuntimeInfo: WireMockRuntimeInfo) { + blockingClient = + OpenAIOkHttpClient.builder() + .baseUrl(wmRuntimeInfo.httpBaseUrl) + .apiKey("My API Key") + .build() + asyncClient = + OpenAIOkHttpClientAsync.builder() + .baseUrl(wmRuntimeInfo.httpBaseUrl) + .apiKey("My API Key") + .build() + } + + @Test + fun requestId_isAvailableOnBlockingStream() { + stubStreamingResponse("req_blocking") + + val response = blockingClient.completions().createStreaming(params()) + + response.use { + assertThat(response.requestId()).contains("req_blocking") + assertThat(response.stream().count()).isEqualTo(1L) + } + } + + @Test + fun requestId_isAvailableInAsyncHandler() { + stubStreamingResponse("req_async") + val response = asyncClient.completions().createStreaming(params()) + val requestIds = mutableListOf() + + response.subscribe { response.requestId().ifPresent(requestIds::add) } + response.onCompleteFuture().join() + + assertThat(requestIds).containsExactly("req_async") + assertThat(response.requestId()).contains("req_async") + } + + @Test + fun requestId_isAvailableForHttpErrors() { + stubFor( + post(anyUrl()) + .willReturn( + status(400) + .withHeader("Content-Type", "application/json") + .withHeader("x-request-id", "req_error") + .withBody("""{"error":{"message":"bad request"}}""") + ) + ) + val response = asyncClient.completions().createStreaming(params()) + val requestIds = mutableListOf() + val assertion = + response.onCompleteFuture().exceptionally { + response.requestId().ifPresent(requestIds::add) + null + } + + assertion.join() + assertThat(requestIds).containsExactly("req_error") + } + + private fun stubStreamingResponse(requestId: String) { + stubFor( + post(anyUrl()) + .willReturn( + ok() + .withHeader("Content-Type", "text/event-stream") + .withHeader("x-request-id", requestId) + .withBody( + "data: {\"id\":\"id\",\"choices\":[],\"created\":0,\"model\":\"model\"}\n\n" + + "data: [DONE]\n\n" + ) + ) + ) + } + + private fun params(): CompletionCreateParams = + CompletionCreateParams.builder() + .model(CompletionCreateParams.Model.GPT_3_5_TURBO_INSTRUCT) + .prompt("This is a test.") + .build() +}