package com.openai.helpers

import com.openai.core.JsonNull

import com.openai.core.JsonValue

import com.openai.errors.OpenAIInvalidDataException

import com.openai.models.chat.completions.*

import java.util.Optional

import kotlin.jvm.optionals.getOrNull

class ChatCompletionAccumulator private constructor() {

/**

private var chatCompletion: ChatCompletion? = null

/**

private var chatCompletionBuilder: ChatCompletion.Builder? = null

/**

private val choiceBuilders = mutableMapOf<Long, ChatCompletion.Choice.Builder>()

/**

private val messageBuilders = mutableMapOf<Long, ChatCompletionMessage.Builder>()

/**

private val messageContents = mutableMapOf<Long, StringBuilder>()

/**

private val messageRefusals = mutableMapOf<Long, StringBuilder>()

private val messageAudio = mutableMapOf<Long, AudioState>()

private class AudioState {

var id: String? = null

var expiresAt: Long? = null

var data: StringBuilder? = null

var transcript: StringBuilder? = null

val additionalProperties = mutableMapOf<String, JsonValue>()

fun isComplete() = id != null && expiresAt != null && data != null && transcript != null

fun build(): ChatCompletionAudio {

if (!isComplete()) throw OpenAIInvalidDataException("Incomplete streamed audio.")

return ChatCompletionAudio.builder()

.id(id!!)

.data(data.toString())

.transcript(transcript.toString())

.expiresAt(expiresAt!!)

.putAllAdditionalProperties(additionalProperties)

.build()

}

}

/**

private val logprobsBuilders = mutableMapOf<Long, ChatCompletion.Choice.Logprobs.Builder>()

/**

private val toolCallBuilders =

mutableMapOf<Long, MutableMap<Long, ChatCompletionMessageFunctionToolCall.Builder>>()

/**

private val toolCallFunctionBuilders =

mutableMapOf<

Long,

MutableMap<Long, ChatCompletionMessageFunctionToolCall.Function.Builder>,

>()

/**

private val toolCallFunctionArgs = mutableMapOf<Long, MutableMap<Long, StringBuilder>>()

/**

private val isFinished = mutableMapOf<Long, Boolean>()

companion object {

@JvmStatic fun create() = ChatCompletionAccumulator()

@JvmSynthetic

internal fun convertFunctionCall(

chunkFunctionCall: ChatCompletionChunk.Choice.Delta.FunctionCall

) =

ChatCompletionMessage.FunctionCall.builder()

.name(chunkFunctionCall._name())

.arguments(chunkFunctionCall._arguments())

.additionalProperties(chunkFunctionCall._additionalProperties())

.build()

}

/**

fun chatCompletion(): ChatCompletion =

checkNotNull(chatCompletion) { "Final chat completion chunk(s) not yet received." }

/**

fun <T : Any> chatCompletion(responseType: Class<T>) =

StructuredChatCompletion(responseType, chatCompletion())

/**

fun accumulate(chunk: ChatCompletionChunk): ChatCompletionChunk {

val localChatCompletion = chatCompletion

check(

localChatCompletion == null ||

(chunk.usage().getOrNull() != null &&

localChatCompletion.usage().getOrNull() == null)

) {

"Already accumulated the final chunk(s)."

}

val chatCompletionBuilder = ensureChatCompletionBuilder()

chunk.usage().getOrNull()?.let {

chatCompletionBuilder.usage(it)

chatCompletion = chatCompletionBuilder.build()

// This final `usage` chunk will have no `choices`, so stop here.

return@accumulate chunk

}

// Each chunk in the stream will repeat these values. They could be copied just once, but it

// is simpler copy them each time they are encountered.

chatCompletionBuilder

// These are the mandatory fields that must be set to something.

.id(chunk.id())

.created(chunk.created())

.model(chunk.model())

// The other fields are optional and `object` has an appropriate default value.

.systemFingerprint(chunk._systemFingerprint())

.apply {

chunk.serviceTier().ifPresent {

serviceTier(ChatCompletion.ServiceTier.of(it.asString()))

}

}

.putAllAdditionalProperties(chunk._additionalProperties())

// Each chunk contains a `choices` list (the size depends on `n` in the request). Each

// chunked `choice` has an `index`, that identifies the corresponding `Choice` within the

// accumulated list of choices; and a `Delta`, that contains additional content to be

// accumulated in that `Choice`, mostly within its `message`.

chunk.choices().forEach { choice ->

val index = choice.index()

val choiceBuilder =

choiceBuilders.getOrPut(index) { ChatCompletion.Choice.builder().index(index) }

accumulateLogprobs(index, choice.logprobs())

accumulateMessage(index, choice.delta())

choiceBuilder.additionalProperties(choice._additionalProperties())

if (choice.finishReason().isPresent) {

choiceBuilder.finishReason(

ChatCompletion.Choice.FinishReason.of(choice.finishReason().get().asString())

)

// A non-null `finishReason` signals that this is the final `Chunk` for its index.

// For each of the `n` messages completed, there will be one chunk carrying only a

// `finishReason` (its `Choice` will have no `content`). These `finalReason` chunks

// will all arrive after the non-empty chunks have been received for all `n`

// choices. Therefore, the value of `n` can be inferred from the number of distinct

// indexes observed thus far when building choices. When all `n` `finishReason`

// chunks have been received, presumptively build the `ChatCompletion`. (One further

// usage chunk _might_ be notified.)

isFinished[index] = true

} else if (isTerminalAudioUpdate(choice) && messageAudio[index]?.isComplete() == true) {

if (isFinished[index] != true) {

choiceBuilder.finishReason(ChatCompletion.Choice.FinishReason.STOP)

}

isFinished[index] = true

} else if (isFinished[index] != true) {

choiceBuilder.finishReason(JsonNull.of())

}

}

if (

chunk.choices().any { it.finishReason().isPresent || isTerminalAudioUpdate(it) } &&

choiceBuilders.keys.all { isFinished[it] == true } &&

messageAudio.values.all { it.isComplete() }

) {

chatCompletion = chatCompletionBuilder.choices(buildChoices()).build()

// Release mutable storage only after the whole chunk and final build succeed.

toolCallFunctionArgs.clear()

messageContents.clear()

messageRefusals.clear()

messageAudio.clear()

}

return chunk

}

private fun ensureChatCompletionBuilder(): ChatCompletion.Builder =

chatCompletionBuilder ?: ChatCompletion.builder().also { chatCompletionBuilder = it }

@JvmSynthetic

internal fun accumulateLogprobs(

index: Long,

logprobs: Optional<ChatCompletionChunk.Choice.Logprobs>,

) {

val chunkLogprobs = logprobs.getOrNull() ?: return

val logprobsBuilder =

logprobsBuilders.getOrPut(index) {

ChatCompletion.Choice.Logprobs.builder()

// These must be set explicitly to avoid an error when `build()` is called in

// the event that no `content` or `refusal` is added later.

.content(mutableListOf())

.refusal(mutableListOf())

}

chunkLogprobs.content().ifPresent { chunkContentList ->

chunkContentList.forEach(logprobsBuilder::addContent)

}

chunkLogprobs.refusal().ifPresent { chunkRefusalsList ->

chunkRefusalsList.forEach(logprobsBuilder::addRefusal)

}

logprobsBuilder.putAllAdditionalProperties(chunkLogprobs._additionalProperties())

}

@JvmSynthetic

internal fun accumulateMessage(index: Long, delta: ChatCompletionChunk.Choice.Delta) {

val messageBuilder = messageBuilders.getOrPut(index) { ChatCompletionMessage.builder() }

delta.audio().ifPresent { audio ->

val state = messageAudio.getOrPut(index) { AudioState() }

audio.id().ifPresent { state.id = it }

audio.expiresAt().ifPresent { state.expiresAt = it }

audio.data().ifPresent {

val buffer = state.data ?: StringBuilder().also { state.data = it }

buffer.append(it)

}

audio.transcript().ifPresent {

val buffer = state.transcript ?: StringBuilder().also { state.transcript = it }

buffer.append(it)

}

state.additionalProperties.putAll(audio._additionalProperties())

}

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())) }

delta.functionCall().ifPresent { messageBuilder.functionCall(convertFunctionCall(it)) }

delta.toolCalls().ifPresent {

it.map { deltaToolCall ->

// The first chunk delta will carry the tool call ID and the function name. Later

// deltas will carry only fragments of the function arguments, but the tool call

// index will identify the function to which those argument fragments belong.

val messageToolCallBuilders = toolCallBuilders.getOrPut(index) { mutableMapOf() }

messageToolCallBuilders.getOrPut(deltaToolCall.index()) {

ChatCompletionMessageFunctionToolCall.builder()

.id(deltaToolCall._id())

.additionalProperties(deltaToolCall._additionalProperties())

// Must wait until the `function` is accumulated and built before adding it to

// the tool call later when `buildChoices` is called.

}

val messageToolCallFunctionBuilders =

toolCallFunctionBuilders.getOrPut(index) { mutableMapOf() }

messageToolCallFunctionBuilders.getOrPut(deltaToolCall.index()) {

ChatCompletionMessageFunctionToolCall.Function.builder()

.name(ensureFunction(deltaToolCall.function())._name())

.additionalProperties(deltaToolCall._additionalProperties())

}

val messageToolCallFunctionArgs =

toolCallFunctionArgs.getOrPut(index) { mutableMapOf() }

val fragment =

ensureFunction(deltaToolCall.function()).arguments().getOrNull() ?: ""

messageToolCallFunctionArgs

.getOrPut(deltaToolCall.index()) { StringBuilder() }

.append(fragment)

}

}

messageBuilder.putAllAdditionalProperties(delta._additionalProperties())

}

private fun ensureFunction(

function: Optional<ChatCompletionChunk.Choice.Delta.ToolCall.Function>

): ChatCompletionChunk.Choice.Delta.ToolCall.Function =

function.orElseThrow { OpenAIInvalidDataException("Tool call chunk missing function.") }

private fun isTerminalAudioUpdate(choice: ChatCompletionChunk.Choice): Boolean =

choice.delta().audio().getOrNull()?.let {

it.expiresAt().isPresent &&

!it.id().isPresent &&

!it.data().isPresent &&

!it.transcript().isPresent

} ?: false

private fun buildChoices() =

choiceBuilders.entries

.sortedBy { it.key }

.map {

it.value

.message(buildMessage(it.key))

.logprobs(logprobsBuilders[it.key]?.build())

.build()

}

private fun buildMessage(index: Long) =

messageBuilders

.getOrElse(index) {

throw OpenAIInvalidDataException("Missing message for index $index.")

}

.content(messageContents[index]?.toString())

.refusal(messageRefusals[index]?.toString())

.toolCalls(buildToolCalls(index))

.apply { messageAudio[index]?.let { audio(it.build()) } }

.build()

private fun buildToolCalls(index: Long): List<ChatCompletionMessageToolCall> =

// It is OK for a message not to have any tool calls; most will not and an empty list will

// be returned. An entry (if it exists) will be a collection of tool call builders and each

// has a function that needs to be set.

toolCallBuilders[index]

?.entries

?.sortedBy { it.key }

?.map { (toolCallIndex, toolCallBuilder) ->

ChatCompletionMessageToolCall.ofFunction(

toolCallBuilder.function(buildFunction(index, toolCallIndex)).build()

)

} ?: listOf()

private fun buildFunction(index: Long, toolCallIndex: Long) =

// Every tool call is expected to have a function with arguments.

toolCallFunctionBuilders[index]

?.get(toolCallIndex)

?.arguments(

toolCallFunctionArgs[index]?.get(toolCallIndex)?.toString()

?: throw OpenAIInvalidDataException(

"Missing function arguments for index $index.$toolCallIndex."

)

)

?.build()

?: throw OpenAIInvalidDataException(

"Missing function builder for index $index.$toolCallIndex."

)

}