fix: buffer completion text with linear accumulation

Castiron-Internal-PR: openai/openai-java-internal#99
Castiron-Source-SHA: 5f7f4c7ae92197dc18b6c03d32ee58b218844535
Castiron-Public-Base-SHA: 5ac7585

Author: dpiet-oai

.castiron.stats.yml

@@ -1,8 +1,8 @@
 schema_version: 1
-generation_id: 41062568-96ed-43ad-acf7-4e906cec816a
-openapi_spec_hash: c7a2088dc0fb8d92f9a934ba4dca2d80
-openapi_transformed_spec_hash: 7281a983306f72fc39720d7057f11ef7
-config_hash: 60435bd668a5724c530fe7e9d6869e9a
-codegen_sha: d8cde0e2bacc733acbb03dc3a3a0d1529f04d9c7
-codegen_hash: 6d6c043d25782a4d8090ccdfcc621f32341ba1116513d40b535b89085b9704bd
-public_codegen_sha: ff3aa2b162a124130056f2f79c2846068f315842
+generation_id: 7667f385-fe72-41e4-bbad-28974f953856
+openapi_spec_hash: 3bbe70558d1cb860c1efbfd901285796
+openapi_transformed_spec_hash: 2daccabc873f61eec5ef07ab28b20a1a
+config_hash: c234968a5221983d6e4315cda9d36d23
+codegen_sha: b704e89df9b355295f0a210ded78f0abb7a851ac
+codegen_hash: 376441665ba8246e8237451b8fd77dd3d047a2db02eeba4b3373041d6e173bcf
+public_codegen_sha: c95741280f105f8280e004ce564afdf944adaff7

openai-java-core/src/main/kotlin/com/openai/helpers/ChatCompletionAccumulator.kt

@@ -48,13 +48,13 @@ class ChatCompletionAccumulator private constructor() {
      * The accumulated content for each message. The keys correspond to the indexes in
      * [messageBuilders].
      */
-    private val messageContents = mutableMapOf<Long, String>()
+    private val messageContents = mutableMapOf<Long, StringBuilder>()
 
     /**
      * The accumulated refusal for each message. The keys correspond to the indexes in
      * [messageBuilders].
      */
-    private val messageRefusals = mutableMapOf<Long, String>()
+    private val messageRefusals = mutableMapOf<Long, StringBuilder>()
 
     /**
      * The builders for the [ChatCompletion.Choice.Logprobs] of each choice. These are only
@@ -262,8 +262,8 @@ class ChatCompletionAccumulator private constructor() {
     internal fun accumulateMessage(index: Long, delta: ChatCompletionChunk.Choice.Delta) {
         val messageBuilder = messageBuilders.getOrPut(index) { ChatCompletionMessage.builder() }
 
-        delta.content().ifPresent { messageContents[index] = (messageContents[index] ?: "") + it }
-        delta.refusal().ifPresent { messageRefusals[index] = (messageRefusals[index] ?: "") + it }
+        delta.content().ifPresent { messageContents.getOrPut(index) { StringBuilder() }.append(it) }
+        delta.refusal().ifPresent { messageRefusals.getOrPut(index) { StringBuilder() }.append(it) }
         // The `role` defaults to "assistant", so if no other `role` is set on the delta, there is
         // no need to set it explicitly to anything else.
         delta.role().ifPresent { messageBuilder.role(JsonValue.from(it.asString())) }
@@ -325,8 +325,8 @@ class ChatCompletionAccumulator private constructor() {
             .getOrElse(index) {
                 throw OpenAIInvalidDataException("Missing message for index $index.")
             }
-            .content(messageContents[index])
-            .refusal(messageRefusals[index])
+            .content(messageContents[index]?.toString())
+            .refusal(messageRefusals[index]?.toString())
             .toolCalls(buildToolCalls(index))
             .build()
 

openai-java-core/src/test/kotlin/com/openai/helpers/FragmentedCompletionTextTest.kt

@@ -0,0 +1,47 @@
+package com.openai.helpers
+
+import com.openai.models.chat.completions.ChatCompletionChunk
+import org.assertj.core.api.Assertions.assertThat
+import org.junit.jupiter.api.Test
+
+internal class FragmentedCompletionTextTest {
+    @Test
+    fun preservesContentAndRefusalAcrossManyFragments() {
+        val accumulator = ChatCompletionAccumulator.create()
+        val delta = ChatCompletionChunk.Choice.Delta.builder().content("aλ😀").refusal("no").build()
+        val fragment = chunk(delta, false)
+        repeat(10_000) { accumulator.accumulate(fragment) }
+        accumulator.accumulate(chunk(ChatCompletionChunk.Choice.Delta.builder().build(), true))
+        val message = accumulator.chatCompletion().choices().single().message()
+        assertThat(message.content()).hasValue("aλ😀".repeat(10_000))
+        assertThat(message.refusal()).hasValue("no".repeat(10_000))
+    }
+
+    @Test
+    fun distinguishesEmptyContentFromMissingRefusal() {
+        val accumulator = ChatCompletionAccumulator.create()
+        accumulator.accumulate(
+            chunk(ChatCompletionChunk.Choice.Delta.builder().content("").build(), false)
+        )
+        accumulator.accumulate(chunk(ChatCompletionChunk.Choice.Delta.builder().build(), true))
+        val message = accumulator.chatCompletion().choices().single().message()
+        assertThat(message.content()).hasValue("")
+        assertThat(message.refusal()).isEmpty()
+    }
+
+    private fun chunk(delta: ChatCompletionChunk.Choice.Delta, finished: Boolean) =
+        ChatCompletionChunk.builder()
+            .id("test")
+            .created(1)
+            .model("test")
+            .addChoice(
+                ChatCompletionChunk.Choice.builder()
+                    .index(0)
+                    .delta(delta)
+                    .finishReason(
+                        if (finished) ChatCompletionChunk.Choice.FinishReason.STOP else null
+                    )
+                    .build()
+            )
+            .build()
+}