chore(internal): refactor streaming implementation (#129)

Author: stainless-app[bot]

openai-java-core/src/main/kotlin/com/openai/core/handlers/SseHandler.kt

@@ -6,69 +6,40 @@ import com.fasterxml.jackson.databind.json.JsonMapper
 import com.openai.core.JsonValue
 import com.openai.core.http.HttpResponse
 import com.openai.core.http.HttpResponse.Handler
-import com.openai.core.http.PhantomReachableClosingStreamResponse
 import com.openai.core.http.SseMessage
 import com.openai.core.http.StreamResponse
 import com.openai.errors.OpenAIException
 import java.util.stream.Stream
 import kotlin.jvm.optionals.getOrNull
-import kotlin.streams.asStream
 
 @JvmSynthetic
 internal fun sseHandler(jsonMapper: JsonMapper): Handler<StreamResponse<SseMessage>> =
-    object : Handler<StreamResponse<SseMessage>> {
-
-        override fun handle(response: HttpResponse): StreamResponse<SseMessage> {
-            val reader = response.body().bufferedReader()
-            val sequence =
-                sequence {
-                        reader.useLines { lines ->
-                            val state = SseState(jsonMapper)
-                            var done = false
-                            for (line in lines) {
-                                // Stop emitting messages, but iterate through the full stream.
-                                if (done) {
-                                    continue
-                                }
-                                val message = state.decode(line) ?: continue
-
-                                if (message.data.startsWith("[DONE]")) {
-                                    // In this case we don't break because we still want to iterate
-                                    // through the full stream.
-                                    done = true
-                                    continue
-                                }
-
-                                val error =
-                                    message.json<JsonValue>().asObject().getOrNull()?.get("error")
-                                if (error != null) {
-                                    val errorMessage =
-                                        error.asString().getOrNull()
-                                            ?: error
-                                                .asObject()
-                                                .getOrNull()
-                                                ?.get("message")
-                                                ?.asString()
-                                                ?.getOrNull()
-                                            ?: "An error occurred during streaming"
-                                    throw OpenAIException(errorMessage)
-                                }
-                                yield(message)
-                            }
-                        }
-                    }
-                    .constrainOnce()
-
-            return PhantomReachableClosingStreamResponse(
-                object : StreamResponse<SseMessage> {
-                    override fun stream(): Stream<SseMessage> = sequence.asStream()
-
-                    override fun close() {
-                        reader.close()
-                        response.close()
-                    }
-                }
-            )
+    streamHandler { lines ->
+        val state = SseState(jsonMapper)
+        var done = false
+        for (line in lines) {
+            // Stop emitting messages, but iterate through the full stream.
+            if (done) {
+                continue
+            }
+            val message = state.decode(line) ?: continue
+
+            if (message.data.startsWith("[DONE]")) {
+                // In this case we don't break because we still want to iterate through the full
+                // stream.
+                done = true
+                continue
+            }
+
+            val error = message.json<JsonValue>().asObject().getOrNull()?.get("error")
+            if (error != null) {
+                val errorMessage =
+                    error.asString().getOrNull()
+                        ?: error.asObject().getOrNull()?.get("message")?.asString()?.getOrNull()
+                        ?: "An error occurred during streaming"
+                throw OpenAIException(errorMessage)
+            }
+            yield(message)
         }
     }
 

openai-java-core/src/main/kotlin/com/openai/core/handlers/StreamHandler.kt

@@ -0,0 +1,32 @@
+@file:JvmName("StreamHandler")
+
+package com.openai.core.handlers
+
+import com.openai.core.http.HttpResponse
+import com.openai.core.http.HttpResponse.Handler
+import com.openai.core.http.PhantomReachableClosingStreamResponse
+import com.openai.core.http.StreamResponse
+import java.util.stream.Stream
+import kotlin.streams.asStream
+
+@JvmSynthetic
+internal fun <T> streamHandler(
+    block: suspend SequenceScope<T>.(lines: Sequence<String>) -> Unit
+): Handler<StreamResponse<T>> =
+    object : Handler<StreamResponse<T>> {
+        override fun handle(response: HttpResponse): StreamResponse<T> {
+            val reader = response.body().bufferedReader()
+            val sequence = sequence { reader.useLines { block(it) } }.constrainOnce()
+
+            return PhantomReachableClosingStreamResponse(
+                object : StreamResponse<T> {
+                    override fun stream(): Stream<T> = sequence.asStream()
+
+                    override fun close() {
+                        reader.close()
+                        response.close()
+                    }
+                }
+            )
+        }
+    }