feat(chat): add tool call index (#256)

aallam · Nov 8, 2023 · 71ea728 · 71ea728
1 parent 39762ac
commit 71ea728
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 10 deletions.
diff --git a/...src/commonMain/kotlin/com.aallam.openai.client/extension/internal/ChatMessageAssembler.kt b/...src/commonMain/kotlin/com.aallam.openai.client/extension/internal/ChatMessageAssembler.kt
@@ -6,12 +6,12 @@ import com.aallam.openai.api.chat.internal.ToolType
 /**
  * A class to help assemble chat messages from chat chunks.
  */
-internal class ChatMessageAssembler() {
+internal class ChatMessageAssembler {
     private val chatFuncName = StringBuilder()
     private val chatFuncArgs = StringBuilder()
     private val chatContent = StringBuilder()
     private var chatRole: ChatRole? = null
-    private var toolCallsAssemblers = mutableListOf<ToolCallAssembler>()
+    private val toolCallsAssemblers = mutableMapOf<Int, ToolCallAssembler>()
 
     /**
      * Merges a chat chunk into the chat message being assembled.
@@ -24,12 +24,9 @@ internal class ChatMessageAssembler() {
                 call.nameOrNull?.let { chatFuncName.append(it) }
                 call.argumentsOrNull?.let { chatFuncArgs.append(it) }
             }
-            toolCalls?.first()?.let { toolCall -> // TBC: tool calls come one by one
-                if (toolCall.idOrNull != null) {
-                    val toolCallAssembler = ToolCallAssembler()
-                    toolCallsAssemblers.add(toolCallAssembler)
-                }
-                toolCallsAssemblers.last().merge(toolCall)
+            toolCalls?.onEach { toolCall ->
+                val assembler = toolCallsAssemblers.getOrPut(toolCall.index) { ToolCallAssembler() }
+                assembler.merge(toolCall)
             }
         }
         return this
@@ -46,18 +43,20 @@ internal class ChatMessageAssembler() {
             this.name = chatFuncName.toString()
         }
         if (toolCallsAssemblers.isNotEmpty()) {
-            this.toolCalls = toolCallsAssemblers.map { it.build() }.toList()
+            this.toolCalls = toolCallsAssemblers.map { (_, value) -> value.build() }.toList()
         }
     }
 }
 
 internal class ToolCallAssembler {
+    private var toolIndex: Int? = null
     private var toolId: ToolId? = null
     private var toolType: ToolType? = null
     private var funcName: String? = null
     private val funcArgs = StringBuilder()
 
     fun merge(toolCall: ToolCall): ToolCallAssembler {
+        toolCall.indexOrNull?.let { toolIndex = it }
         toolCall.idOrNull?.let { toolId = it }
         toolCall.typeOrNull?.let { toolType = it }
         toolCall.functionOrNull?.let { call ->
@@ -71,6 +70,7 @@ internal class ToolCallAssembler {
      * Builds and returns the assembled chat message.
      */
     fun build(): ToolCall = toolCall {
+        this.index = toolIndex
         this.id = toolId
         this.type = toolType
         if (funcName?.isNotEmpty() == true || funcArgs.isNotEmpty()) {

diff --git a/openai-core/src/commonMain/kotlin/com.aallam.openai.api/chat/ToolCall.kt b/openai-core/src/commonMain/kotlin/com.aallam.openai.api/chat/ToolCall.kt
@@ -8,10 +8,12 @@ import kotlinx.serialization.Serializable
 /**
  * The tool call generated by the model, such as function call.
  *
- * In case of streaming variant of the chat API, the object parameters can be `null`.
+ * In case of streaming variant of the chat API, [index] is required, other object parameters can be `null`.
  */
 @Serializable
 public data class ToolCall(
+    /** Tool call index. Required in the case of chat stream variant **/
+    @SerialName("index") val indexOrNull: Int? = null,
     /** The ID of the tool call. **/
     @SerialName("id") val idOrNull: ToolId? = null,
     /** The type of the tool. **/
@@ -20,6 +22,9 @@ public data class ToolCall(
     @SerialName("function") val functionOrNull: FunctionCall? = null,
 ) {
 
+    val index: Int
+        get() = requireNotNull(indexOrNull)
+
     /** The ID of the tool call. **/
     val id: ToolId
         get() = requireNotNull(idOrNull)
@@ -40,6 +45,10 @@ public fun toolCall(block: ToolCallBuilder.() -> Unit): ToolCall = ToolCallBuild
 
 @OpenAIDsl
 public class ToolCallBuilder {
+
+    /** Tool call index. Required in the case of chat stream variant **/
+    public var index: Int? = null
+
     /** The ID of the tool call. **/
     public var id: ToolId? = null
 
@@ -53,6 +62,7 @@ public class ToolCallBuilder {
      * Create [ToolCall] instance.
      */
     public fun build(): ToolCall = ToolCall(
+        indexOrNull = index,
         idOrNull = id,
         typeOrNull = type,
         functionOrNull = function,