-
Notifications
You must be signed in to change notification settings - Fork 5
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(client): add async streaming methods (#6)
- Loading branch information
1 parent
2a1c09d
commit c077dca
Showing
10 changed files
with
459 additions
and
4 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
97 changes: 97 additions & 0 deletions
97
openai-java-core/src/main/kotlin/com/openai/core/http/AsyncStreamResponse.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,97 @@ | ||
| package com.openai.core.http | ||
|
|
||
| import com.openai.core.http.AsyncStreamResponse.Handler | ||
| import java.util.Optional | ||
| import java.util.concurrent.CompletableFuture | ||
| import java.util.concurrent.Executor | ||
| import java.util.concurrent.atomic.AtomicReference | ||
|
|
||
| interface AsyncStreamResponse<T> { | ||
|
|
||
| fun subscribe(handler: Handler<T>): AsyncStreamResponse<T> | ||
|
|
||
| fun subscribe(handler: Handler<T>, executor: Executor): AsyncStreamResponse<T> | ||
|
|
||
| /** | ||
| * Closes this resource, relinquishing any underlying resources. | ||
| * | ||
| * This is purposefully not inherited from [AutoCloseable] because this response should not be | ||
| * synchronously closed via try-with-resources. | ||
| */ | ||
| fun close() | ||
|
|
||
| fun interface Handler<in T> { | ||
|
|
||
| fun onNext(value: T) | ||
|
|
||
| fun onComplete(error: Optional<Throwable>) {} | ||
| } | ||
| } | ||
|
|
||
| @JvmSynthetic | ||
| internal fun <T> CompletableFuture<StreamResponse<T>>.toAsync(streamHandlerExecutor: Executor) = | ||
| PhantomReachableClosingAsyncStreamResponse( | ||
| object : AsyncStreamResponse<T> { | ||
|
|
||
| private val state = AtomicReference(State.NEW) | ||
|
|
||
| override fun subscribe(handler: Handler<T>): AsyncStreamResponse<T> = | ||
| subscribe(handler, streamHandlerExecutor) | ||
|
|
||
| override fun subscribe( | ||
| handler: Handler<T>, | ||
| executor: Executor | ||
| ): AsyncStreamResponse<T> = apply { | ||
| // TODO(JDK): Use `compareAndExchange` once targeting JDK 9. | ||
| check(state.compareAndSet(State.NEW, State.SUBSCRIBED)) { | ||
| if (state.get() == State.SUBSCRIBED) "Cannot subscribe more than once" | ||
| else "Cannot subscribe after the response is closed" | ||
| } | ||
|
|
||
| this@toAsync.whenCompleteAsync( | ||
| { streamResponse, futureError -> | ||
| if (state.get() == State.CLOSED) { | ||
| // Avoid doing any work if `close` was called before the future | ||
| // completed. | ||
| return@whenCompleteAsync | ||
| } | ||
|
|
||
| if (futureError != null) { | ||
| // An error occurred before we started passing chunks to the handler. | ||
| handler.onComplete(Optional.of(futureError)) | ||
| return@whenCompleteAsync | ||
| } | ||
|
|
||
| var streamError: Throwable? = null | ||
| try { | ||
| streamResponse.stream().forEach(handler::onNext) | ||
| } catch (e: Throwable) { | ||
| streamError = e | ||
| } | ||
|
|
||
| try { | ||
| handler.onComplete(Optional.ofNullable(streamError)) | ||
| } finally { | ||
| close() | ||
| } | ||
| }, | ||
| executor | ||
| ) | ||
| } | ||
|
|
||
| override fun close() { | ||
| val previousState = state.getAndSet(State.CLOSED) | ||
| if (previousState == State.CLOSED) { | ||
| return | ||
| } | ||
|
|
||
| this@toAsync.whenComplete { streamResponse, _ -> streamResponse?.close() } | ||
| } | ||
| } | ||
| ) | ||
|
|
||
| private enum class State { | ||
| NEW, | ||
| SUBSCRIBED, | ||
| CLOSED | ||
| } |
24 changes: 24 additions & 0 deletions
24
...a-core/src/main/kotlin/com/openai/core/http/PhantomReachableClosingAsyncStreamResponse.kt
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,24 @@ | ||
| package com.openai.core.http | ||
|
|
||
| import com.openai.core.closeWhenPhantomReachable | ||
| import com.openai.core.http.AsyncStreamResponse.Handler | ||
| import java.util.concurrent.Executor | ||
|
|
||
| internal class PhantomReachableClosingAsyncStreamResponse<T>( | ||
| private val asyncStreamResponse: AsyncStreamResponse<T> | ||
| ) : AsyncStreamResponse<T> { | ||
| init { | ||
| closeWhenPhantomReachable(this, asyncStreamResponse::close) | ||
| } | ||
|
|
||
| override fun subscribe(handler: Handler<T>): AsyncStreamResponse<T> = apply { | ||
| asyncStreamResponse.subscribe(handler) | ||
| } | ||
|
|
||
| override fun subscribe(handler: Handler<T>, executor: Executor): AsyncStreamResponse<T> = | ||
| apply { | ||
| asyncStreamResponse.subscribe(handler, executor) | ||
| } | ||
|
|
||
| override fun close() = asyncStreamResponse.close() | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.