add stream support for tools

Author: CJCrafter

src/main/kotlin/com/cjcrafter/openai/OpenAI.kt

@@ -5,8 +5,6 @@ import com.cjcrafter.openai.chat.tool.ToolChoice
 import com.cjcrafter.openai.completions.CompletionRequest
 import com.cjcrafter.openai.completions.CompletionResponse
 import com.cjcrafter.openai.completions.CompletionResponseChunk
-import com.cjcrafter.openai.jackson.ChatChoiceChunkDeserializer
-import com.cjcrafter.openai.jackson.ChatChoiceChunkSerializer
 import com.cjcrafter.openai.jackson.ToolChoiceDeserializer
 import com.cjcrafter.openai.jackson.ToolChoiceSerializer
 import com.fasterxml.jackson.annotation.JsonInclude
@@ -145,8 +143,6 @@ interface OpenAI {
 
             // Register modules with custom serializers/deserializers
             val module = SimpleModule().apply {
-                addSerializer(ChatChoiceChunk::class.java, ChatChoiceChunkSerializer())
-                addDeserializer(ChatChoiceChunk::class.java, ChatChoiceChunkDeserializer())
                 addSerializer(ToolChoice::class.java, ToolChoiceSerializer())
                 addDeserializer(ToolChoice::class.java, ToolChoiceDeserializer())
             }

src/main/kotlin/com/cjcrafter/openai/OpenAIImpl.kt

@@ -9,45 +9,11 @@ import com.fasterxml.jackson.databind.node.ObjectNode
 import okhttp3.*
 import okhttp3.MediaType.Companion.toMediaType
 import okhttp3.RequestBody.Companion.toRequestBody
+import org.jetbrains.annotations.ApiStatus
 import java.io.BufferedReader
 import java.io.IOException
 
-/**
- * The `OpenAI` class contains all the API calls to OpenAI's endpoint. Whether
- * you are working with images, chat, or completions, you need to have an
- * `OpenAI` instance to make the API requests.
- *
- * To get your API key:
- * 1. Log in to your account: Go to [https://www.openai.com/](openai.com) and
- * log in.
- * 2. Access the API dashboard: After logging in, click on the "API" tab.
- * 3. Choose a subscription plan: Select a suitable plan based on your needs
- * and complete the payment process.
- * 4. Obtain your API key: After subscribing to a plan, you will be redirected
- * to the API dashboard, where you can find your unique API key. Copy and store it securely.
- *
- * All API methods in this class have a non-blocking option which will enqueues
- * the HTTPS request on a different thread. These method names have `Async
- * appended to the end of their names.
- *
- * Completions API:
- * * [createCompletion]
- * * [streamCompletion]
- * * [createCompletionAsync]
- * * [streamCompletionAsync]
- *
- * Chat API:
- * * [createChatCompletion]
- * * [streamChatCompletion]
- * * [createChatCompletionAsync]
- * * [streamChatCompletionAsync]
- *
- * @property apiKey Your OpenAI API key. It starts with `"sk-"` (without the quotes).
- * @property organization If you belong to multiple organizations, specify which one to use (else `null`).
- * @property client Controls proxies, timeouts, etc.
- * @constructor Create a ChatBot for responding to requests.
- */
-open class OpenAIImpl @JvmOverloads constructor(
+open class OpenAIImpl @ApiStatus.Internal constructor(
     protected val apiKey: String,
     protected val organization: String? = null,
     private val client: OkHttpClient = OkHttpClient()
@@ -98,7 +64,9 @@ open class OpenAIImpl @JvmOverloads constructor(
         }
 
         val json = httpResponse.body?.byteStream()?.bufferedReader() ?: throw IOException("Response body is null")
-        return objectMapper.readValue(json, ChatResponse::class.java)
+        val str = json.readText()
+        println(str)
+        return objectMapper.readValue(str, ChatResponse::class.java)
     }
 
     override fun streamChatCompletion(request: ChatRequest): Iterable<ChatResponseChunk> {
@@ -143,10 +111,11 @@ open class OpenAIImpl @JvmOverloads constructor(
 
                     override fun next(): ChatResponseChunk {
                         val currentLine = nextLine ?: throw NoSuchElementException("No more lines")
-                        //println("         $currentLine")
+                        //println("    $currentLine")
                         chunk = chunk?.apply { update(objectMapper.readTree(currentLine) as ObjectNode) } ?: objectMapper.readValue(currentLine, ChatResponseChunk::class.java)
                         nextLine = readNextLine(reader) // Prepare the next line
                         return chunk!!
+                        //return ChatResponseChunk("1", 1, listOf())
                     }
                 }
             }

src/main/kotlin/com/cjcrafter/openai/chat/ChatChoiceChunk.kt

@@ -1,11 +1,12 @@
 package com.cjcrafter.openai.chat
 
 import com.cjcrafter.openai.FinishReason
-import com.cjcrafter.openai.jackson.ChatChoiceChunkDeserializer
-import com.cjcrafter.openai.jackson.ChatChoiceChunkSerializer
+import com.cjcrafter.openai.OpenAI
+import com.cjcrafter.openai.chat.tool.ChatMessageDelta
+import com.cjcrafter.openai.chat.tool.ToolCall
 import com.fasterxml.jackson.annotation.JsonProperty
 import com.fasterxml.jackson.databind.JsonNode
-import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
+import java.lang.IllegalArgumentException
 
 /**
  *
@@ -29,24 +30,38 @@ import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
  */
 data class ChatChoiceChunk(
     val index: Int,
-    val message: ChatMessage,
-    var delta: String,
+    var delta: ChatMessageDelta? = null,
     @JsonProperty("finish_reason") var finishReason: FinishReason?
 ) {
+    val message: ChatMessage = ChatMessage(delta?.role!!, delta?.content, delta?.toolCalls?.map { it.toToolCall() })
 
-    internal fun update(json: String) {
-        val node: JsonNode = jacksonObjectMapper().readTree(json)
-        val deltaNode = node.get("delta")
-        delta = if (deltaNode?.has("content") == true && !deltaNode.get("content").isNull) {
-            deltaNode.get("content").asText()
-        } else {
-            ""
-        }
+    // Reads the json from the OpenAI API, and sets delta. message accumulates changes
+    internal fun update(json: JsonNode) {
+        val deltaNode = json.get("delta") ?: throw IllegalArgumentException("Passed a json without delta")
+        val delta = OpenAI.createObjectMapper().treeToValue(deltaNode, ChatMessageDelta::class.java)
+
+        // The "bread and butter" of streaming. You can start showing the user
+        // generated content usually *within a second*. However, for Tool Calls,
+        // this will always be null.
+        if (message.content == null)
+            message.content = delta?.content // Always null for tool calls
+        else if (delta.content != null)
+            message.content += delta.content
 
-        message.content += delta
-        finishReason = node.get("finish_reason")?.takeIf { !it.isNull }?.asText()?.let {
-            FinishReason.valueOf(it.uppercase())
+        // Handle updating the tool call
+        if (message.toolCalls != null && delta.toolCalls != null) {
+            for (deltaToolCall in delta.toolCalls) {
+                val toolCall = message.toolCalls!![deltaToolCall.index]
+                toolCall.update(deltaToolCall)
+            }
         }
+
+        // The reason the bot stopped generating tokens. This is null until done.
+        finishReason = json.get("finish_reason")?.let { if (it.isNull) null else FinishReason.valueOf(it.asText().uppercase()) }
+
+        // People can manually check changes instead of using this API's
+        // accumulative changes.
+        this.delta = delta
     }
 
     /**
@@ -55,30 +70,4 @@ data class ChatChoiceChunk(
      * complete message.
      */
     fun isFinished() = finishReason != null
-
-    companion object {
-        @JvmStatic
-        fun serializer() = ChatChoiceChunkSerializer()
-
-        @JvmStatic
-        fun deserializer() = ChatChoiceChunkDeserializer()
-    }
 }
-
-/*
-Below is a potential Steam response from OpenAI. You can see that the first
-message contains 0 generated content, and the last message (before "[DONE]")
-adds the finish_reason.
-
-data: {"id":"chatcmpl-6xUB4Vi8jEG8u4hMBTMeO8KXgA87z","object":"chat.completion.chunk","created":1679635374,"model":"gpt-3.5-turbo-0301","choices":[{"delta":{"role":"assistant"},"index":0,"finish_reason":null}]}
-
-data: {"id":"chatcmpl-6xUB4Vi8jEG8u4hMBTMeO8KXgA87z","object":"chat.completion.chunk","created":1679635374,"model":"gpt-3.5-turbo-0301","choices":[{"delta":{"content":"Hello"},"index":0,"finish_reason":null}]}
-
-data: {"id":"chatcmpl-6xUB4Vi8jEG8u4hMBTMeO8KXgA87z","object":"chat.completion.chunk","created":1679635374,"model":"gpt-3.5-turbo-0301","choices":[{"delta":{"content":" World"},"index":0,"finish_reason":null}]}
-
-data: {"id":"chatcmpl-6xUB4Vi8jEG8u4hMBTMeO8KXgA87z","object":"chat.completion.chunk","created":1679635374,"model":"gpt-3.5-turbo-0301","choices":[{"delta":{"content":"."},"index":0,"finish_reason":null}]}
-
-data: {"id":"chatcmpl-6xUB4Vi8jEG8u4hMBTMeO8KXgA87z","object":"chat.completion.chunk","created":1679635374,"model":"gpt-3.5-turbo-0301","choices":[{"delta":{},"index":0,"finish_reason":"stop"}]}
-
-data: [DONE]
- */

src/main/kotlin/com/cjcrafter/openai/chat/ChatResponseChunk.kt

@@ -32,7 +32,7 @@ data class ChatResponseChunk(
     internal fun update(json: ObjectNode) {
         val choicesArray = json.get("choices") as? ArrayNode
         choicesArray?.forEachIndexed { index, jsonNode ->
-            choices[index].update(jsonNode.toString())
+            choices[index].update(jsonNode)
         }
     }
 

src/main/kotlin/com/cjcrafter/openai/chat/tool/ChatMessageDelta.kt

@@ -0,0 +1,11 @@
+package com.cjcrafter.openai.chat.tool
+
+import com.cjcrafter.openai.chat.ChatUser
+import com.fasterxml.jackson.annotation.JsonProperty
+
+data class ChatMessageDelta(
+    val role: ChatUser? = null,
+    val content: String? = null,
+    @JsonProperty("tool_calls") val toolCalls: List<ToolCallDelta>? = null,
+) {
+}

src/main/kotlin/com/cjcrafter/openai/chat/tool/FunctionCall.kt

@@ -5,6 +5,7 @@ import com.fasterxml.jackson.core.JsonProcessingException
 import com.fasterxml.jackson.databind.JsonNode
 import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
 import com.fasterxml.jackson.module.kotlin.readValue
+import org.jetbrains.annotations.ApiStatus
 
 /**
  * Represents a function call by ChatGPT. When ChatGPT calls a function, you
@@ -23,6 +24,12 @@ data class FunctionCall(
     var arguments: String,
 ) {
 
+    @ApiStatus.Internal
+    fun update(delta: FunctionCallDelta) {
+        // The only field that updates is arguments
+        arguments += delta.arguments
+    }
+
     /**
      * Attempts to parse the arguments passed to the function.
      *

src/main/kotlin/com/cjcrafter/openai/chat/tool/FunctionCallDelta.kt

@@ -0,0 +1,20 @@
+package com.cjcrafter.openai.chat.tool
+
+import org.jetbrains.annotations.ApiStatus
+
+data class FunctionCallDelta(
+    val name: String?,
+    val arguments: String,
+) {
+
+    /**
+     * Returns an **incomplete** function call.
+     */
+    @ApiStatus.Internal
+    fun toFunctionCall(): FunctionCall {
+        return FunctionCall(
+            name ?: throw IllegalStateException("name must be set"),
+            arguments
+        )
+    }
+}

src/main/kotlin/com/cjcrafter/openai/chat/tool/ToolCall.kt

@@ -1,8 +1,16 @@
 package com.cjcrafter.openai.chat.tool
 
+import org.jetbrains.annotations.ApiStatus
+
 data class ToolCall(
     var id: String,
     var type: ToolType,
     var function: FunctionCall,
 ) {
+    @ApiStatus.Internal
+    internal fun update(delta: ToolCallDelta) {
+        // The only field that updates is function
+        if (delta.function != null)
+            function.update(delta.function)
+    }
 }

src/main/kotlin/com/cjcrafter/openai/chat/tool/ToolCallDelta.kt

@@ -0,0 +1,17 @@
+package com.cjcrafter.openai.chat.tool
+
+import org.jetbrains.annotations.ApiStatus
+
+data class ToolCallDelta(
+    val index: Int,
+    val id: String? = null,
+    val type: ToolType? = null,
+    val function: FunctionCallDelta? = null,
+) {
+    @ApiStatus.Internal
+    fun toToolCall() = ToolCall(
+        id = id ?: throw IllegalStateException("id must be set"),
+        type = type ?: throw IllegalStateException("type must be set"),
+        function = function?.toFunctionCall() ?: throw IllegalStateException("function must be set"),
+    )
+}