feat(chat): add stream options (#363)

Author: aallam

CHANGELOG.md

@@ -1,3 +1,8 @@
+# Unreleased
+
+### Added
+- **chat**: add stream options (#363)
+
 # 3.8.1
 > Published 28 Jun 2024
 

openai-client/src/commonTest/kotlin/com/aallam/openai/client/TestChatCompletions.kt

@@ -203,4 +203,40 @@ class TestChatCompletions : TestOpenAI() {
 
         assertTrue(job.isCancelled, "Job should be cancelled")
     }
+
+    @Test
+    fun streamOptions() = test {
+        val request = chatCompletionRequest {
+            model = ModelId("gpt-3.5-turbo")
+            messages {
+                message {
+                    role = ChatRole.System
+                    content = "You are a helpful assistant.!"
+                }
+                message {
+                    role = ChatRole.User
+                    content = "Who won the world series in 2020?"
+                }
+                message {
+                    role = ChatRole.Assistant
+                    content = "The Los Angeles Dodgers won the World Series in 2020."
+                }
+                message {
+                    role = ChatRole.User
+                    content = "Where was it played?"
+                }
+            }
+            streamOptions = streamOptions {
+                includeUsage = true
+            }
+        }
+
+        val results = mutableListOf<ChatCompletionChunk>()
+        openAI.chatCompletions(request).onEach { results += it }.launchIn(this).join()
+
+        assertNotNull(results.last().usage)
+        assertNotNull(results.last().usage?.promptTokens)
+        assertNotNull(results.last().usage?.completionTokens)
+        assertNotNull(results.last().usage?.totalTokens)
+    }
 }

openai-core/src/commonMain/kotlin/com.aallam.openai.api/chat/ChatCompletionRequest.kt

@@ -163,6 +163,11 @@ public class ChatCompletionRequest(
      * A unique identifier representing the Multi LORA reserved instance.
      */
     @SerialName("instance_id") public val instanceId: String? = null,
+
+    /**
+     * Options for streaming response. Only used when in streaming mode.
+     */
+    @SerialName("stream_options") public val streamOptions: StreamOptions? = null
 )
 
 /**
@@ -338,6 +343,11 @@ public class ChatCompletionRequestBuilder {
         functions = FunctionsBuilder().apply(block).functions
     }
 
+    /**
+     * Options for streaming response. Only used when in streaming mode.
+     */
+    public var streamOptions: StreamOptions? = null
+
     /**
      * Builder of [ChatCompletionRequest] instances.
      */
@@ -361,6 +371,7 @@ public class ChatCompletionRequestBuilder {
         logprobs = logprobs,
         topLogprobs = topLogprobs,
         instanceId = instanceId,
+        streamOptions = streamOptions,
     )
 }
 

openai-core/src/commonMain/kotlin/com.aallam.openai.api/chat/StreamOptions.kt

@@ -0,0 +1,44 @@
+package com.aallam.openai.api.chat
+
+import kotlinx.serialization.SerialName
+import kotlinx.serialization.Serializable
+
+/**
+ * Options for streaming response.
+ */
+@Serializable
+public data class StreamOptions(
+    /**
+     * If set, an additional chunk will be streamed before the data: `[DONE]` message.
+     * The usage field on this chunk shows the token usage statistics for the entire request, and the choices field will
+     * always be an empty array. All other chunks will also include a usage field, but with a null value.
+     */
+    @SerialName("include_usage") public val includeUsage: Boolean? = null,
+)
+
+/**
+ * Create a new [StreamOptions] instance.
+ */
+public fun streamOptions(block: StreamOptionsBuilder.() -> Unit): StreamOptions {
+    return StreamOptionsBuilder().apply(block).build()
+}
+
+/**
+ * Builder for [StreamOptions].
+ */
+public class StreamOptionsBuilder {
+
+    /**
+     * If set, an additional chunk will be streamed before the data: `[DONE]` message.
+     * The usage field on this chunk shows the token usage statistics for the entire request, and the choices field will
+     * always be an empty array. All other chunks will also include a usage field, but with a null value.
+     */
+    public var includeUsage: Boolean? = null
+
+    /**
+     * Build the [StreamOptions] instance.
+     */
+    public fun build(): StreamOptions = StreamOptions(
+        includeUsage = includeUsage,
+    )
+}