Add id and encrypted fields to Message.Reasoning and integrate across clients

devcrocod · devcrocod · commit d704f74a5897 · 2025-10-22T21:21:52.000+02:00
diff --git a/prompt/prompt-executor/prompt-executor-clients/prompt-executor-anthropic-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/anthropic/AnthropicLLMClient.kt b/prompt/prompt-executor/prompt-executor-clients/prompt-executor-anthropic-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/anthropic/AnthropicLLMClient.kt
@@ -328,9 +328,9 @@ public open class AnthropicLLMClient(
                     messages.add(
                         AnthropicMessage.Assistant(
                             content = listOf(
-                                message.original as? AnthropicContent.Thinking ?: AnthropicContent.Thinking(
-                                    signature = Uuid.random().toString(),
-                                    thinking = message.content // TODO: Use signature from original message if available
+                                AnthropicContent.Thinking(
+                                    signature = message.encrypted!!,
+                                    thinking = message.content
                                 )
                             )
                         )
@@ -507,7 +507,7 @@ public open class AnthropicLLMClient(
                 }
 
                 is AnthropicContent.Thinking -> {
-                    Message.Reasoning(original = content, content = content.thinking, metaInfo = metaInfo)
+                    Message.Reasoning(encrypted = content.signature, content = content.thinking, metaInfo = metaInfo)
                 }
 
                 is AnthropicContent.ToolUse -> {
diff --git a/prompt/prompt-executor/prompt-executor-clients/prompt-executor-bedrock-client/src/jvmMain/kotlin/ai/koog/prompt/executor/clients/bedrock/modelfamilies/anthropic/BedrockAnthropicClaudeSerialization.kt b/prompt/prompt-executor/prompt-executor-clients/prompt-executor-bedrock-client/src/jvmMain/kotlin/ai/koog/prompt/executor/clients/bedrock/modelfamilies/anthropic/BedrockAnthropicClaudeSerialization.kt
@@ -22,8 +22,6 @@ import kotlinx.serialization.json.buildJsonArray
 import kotlinx.serialization.json.buildJsonObject
 import kotlinx.serialization.json.encodeToJsonElement
 import kotlinx.serialization.json.put
-import kotlin.uuid.ExperimentalUuidApi
-import kotlin.uuid.Uuid
 
 internal object BedrockAnthropicClaudeSerialization {
 
@@ -35,7 +33,6 @@ internal object BedrockAnthropicClaudeSerialization {
         explicitNulls = false
     }
 
-    @OptIn(ExperimentalUuidApi::class)
     private fun buildMessagesHistory(prompt: Prompt): MutableList<BedrockAnthropicInvokeModelMessage> {
         val messages = mutableListOf<BedrockAnthropicInvokeModelMessage>()
         prompt.messages.forEach { msg ->
@@ -68,11 +65,10 @@ internal object BedrockAnthropicClaudeSerialization {
                         messages.add(
                             BedrockAnthropicInvokeModelMessage.Assistant(
                                 content = listOf(
-                                    msg.original as? BedrockAnthropicInvokeModelContent.Thinking
-                                        ?: BedrockAnthropicInvokeModelContent.Thinking(
-                                            signature = Uuid.random().toString(),
-                                            thinking = msg.content
-                                        )
+                                    BedrockAnthropicInvokeModelContent.Thinking(
+                                        signature = msg.encrypted!!,
+                                        thinking = msg.content
+                                    )
                                 )
                             )
                         )
@@ -181,7 +177,6 @@ internal object BedrockAnthropicClaudeSerialization {
         )
     }
 
-    @OptIn(ExperimentalUuidApi::class)
     internal fun parseAnthropicResponse(responseBody: String, clock: Clock = Clock.System): List<Message.Response> {
         val response = json.decodeFromString<BedrockAnthropicResponse>(responseBody)
 
@@ -204,7 +199,7 @@ internal object BedrockAnthropicClaudeSerialization {
                 )
 
                 is AnthropicContent.Thinking -> Message.Reasoning(
-                    original = content,
+                    encrypted = content.signature,
                     content = content.thinking,
                     metaInfo = metaInfo
                 )
@@ -255,6 +250,7 @@ internal object BedrockAnthropicClaudeSerialization {
                                 name = content.name,
                                 content = content.input.toString()
                             )
+
                         else -> throw IllegalArgumentException(
                             "Unsupported AnthropicContent type in message_delta. Content: $content"
                         )
diff --git a/prompt/prompt-executor/prompt-executor-clients/prompt-executor-google-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/google/GoogleLLMClient.kt b/prompt/prompt-executor/prompt-executor-clients/prompt-executor-google-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/google/GoogleLLMClient.kt
@@ -333,9 +333,10 @@ public open class GoogleLLMClient(
                         GoogleContent(
                             role = "assistant",
                             parts = listOf(
-                                message.original as? GooglePart ?: GooglePart.Text(
+                                GooglePart.Text(
                                     text = message.content,
-                                    thought = true
+                                    thoughtSignature = message.encrypted,
+                                    thought = true,
                                 )
                             )
                         )
@@ -606,7 +607,7 @@ public open class GoogleLLMClient(
                 is GooglePart.Text -> {
                     if (part.thought ?: false) {
                         Message.Reasoning(
-                            original = part,
+                            encrypted = part.thoughtSignature,
                             content = part.text,
                             metaInfo = metaInfo
                         )
diff --git a/prompt/prompt-executor/prompt-executor-clients/prompt-executor-openai-client-base/src/commonMain/kotlin/ai/koog/prompt/executor/clients/openai/base/AbstractOpenAILLMClient.kt b/prompt/prompt-executor/prompt-executor-clients/prompt-executor-openai-client-base/src/commonMain/kotlin/ai/koog/prompt/executor/clients/openai/base/AbstractOpenAILLMClient.kt
@@ -283,7 +283,7 @@ public abstract class AbstractOpenAILLMClient<TResponse : OpenAIBaseLLMResponse,
 
                 is Message.Reasoning -> {
                     flushPendingCalls()
-                    messages += message.original as? OpenAIMessage.Assistant ?: OpenAIMessage.Assistant(
+                    messages += OpenAIMessage.Assistant(
                         content = Content.Text(message.content),
                         reasoningContent = message.content
                     )
@@ -471,7 +471,6 @@ public abstract class AbstractOpenAILLMClient<TResponse : OpenAIBaseLLMResponse,
 
             this is OpenAIMessage.Assistant && this.reasoningContent != null && this.content != null -> listOf(
                 Message.Reasoning(
-                    original = this,
                     content = this.reasoningContent,
                     metaInfo = metaInfo
                 ),
diff --git a/prompt/prompt-executor/prompt-executor-clients/prompt-executor-openai-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/openai/OpenAILLMClient.kt b/prompt/prompt-executor/prompt-executor-clients/prompt-executor-openai-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/openai/OpenAILLMClient.kt
@@ -602,9 +602,10 @@ public open class OpenAILLMClient(
                     is Message.Reasoning -> {
                         flushPendingCalls()
                         add(
-                            message.original as? Item.Reasoning ?: Item.Reasoning(
-                                Uuid.random().toString(),
-                                listOf(Item.Reasoning.Summary(message.content))
+                            Item.Reasoning(
+                                id = message.id ?: Uuid.random().toString(),
+                                encryptedContent = message.encrypted,
+                                summary = listOf(Item.Reasoning.Summary(message.content))
                             )
                         )
                     }
@@ -709,7 +710,8 @@ public open class OpenAILLMClient(
                     )
 
                     is Item.Reasoning -> Message.Reasoning(
-                        original = output,
+                        id = output.id,
+                        encrypted = output.encryptedContent,
                         content = output.summary.joinToString(separator = "\n") { it.text },
                         metaInfo = metaInfo
                     )
@@ -735,13 +737,15 @@ public open class OpenAILLMClient(
             )
             params
         }
+
         params is OpenAIChatParams -> {
             model.requireCapability(
                 LLMCapability.OpenAIEndpoint.Completions,
                 message = "Must be supported to use OpenAI chat params."
             )
             params
         }
+
         model.supports(LLMCapability.OpenAIEndpoint.Completions) -> params.toOpenAIChatParams()
         model.supports(LLMCapability.OpenAIEndpoint.Responses) -> params.toOpenAIResponsesParams()
         else -> error("Cannot determine proper LLM params for OpenAI model: ${model.id}")
diff --git a/prompt/prompt-model/src/commonMain/kotlin/ai/koog/prompt/message/Message.kt b/prompt/prompt-model/src/commonMain/kotlin/ai/koog/prompt/message/Message.kt
@@ -3,7 +3,6 @@ package ai.koog.prompt.message
 import kotlinx.datetime.Clock
 import kotlinx.datetime.Instant
 import kotlinx.serialization.Serializable
-import kotlinx.serialization.Transient
 import kotlinx.serialization.json.Json
 import kotlinx.serialization.json.JsonObject
 import kotlinx.serialization.json.jsonObject
@@ -143,9 +142,8 @@ public sealed interface Message {
      * Represents a reasoning message exchanged in a chat system, encapsulating the content,
      * role, and associated metadata, with an optional reference to the original thinking process.
      *
-     * @property original An optional reference to the original [Thinking] process that generated this reasoning,
-     *                    providing context or tracing capabilities during operations. This property is transient and
-     *                    will not be serialized.
+     * @property id An optional identifier for the reasoning process.
+     * @property encrypted The encrypted content of the reasoning message.
      * @property content The content of the message as a string.
      * @property role The [Role] of the message, indicating its source or function in the chat (e.g., assistant, user).
      *                Defaults to [Role.Assistant].
@@ -154,8 +152,8 @@ public sealed interface Message {
      */
     @Serializable
     public data class Reasoning(
-        @Transient
-        public val original: Thinking? = null,
+        public val id: String? = null,
+        public val encrypted: String? = null,
         override val content: String,
         override val role: Role = Role.Assistant,
         override val metaInfo: ResponseMetaInfo

Original file line number	Diff line number	Diff line change
`@@ -328,9 +328,9 @@ public open class AnthropicLLMClient(`
`328`	`328`	`messages.add(`
`329`	`329`	`AnthropicMessage.Assistant(`
`330`	`330`	`content = listOf(`
`331`		`- message.original as? AnthropicContent.Thinking ?: AnthropicContent.Thinking(`
`332`		`- signature = Uuid.random().toString(),`
`333`		`- thinking = message.content // TODO: Use signature from original message if available`
	`331`	`+ AnthropicContent.Thinking(`
	`332`	`+ signature = message.encrypted!!,`
	`333`	`+ thinking = message.content`
`334`	`334`	`)`
`335`	`335`	`)`
`336`	`336`	`)`
`@@ -507,7 +507,7 @@ public open class AnthropicLLMClient(`
`507`	`507`	`}`
`508`	`508`
`509`	`509`	`is AnthropicContent.Thinking -> {`
`510`		`- Message.Reasoning(original = content, content = content.thinking, metaInfo = metaInfo)`
	`510`	`+ Message.Reasoning(encrypted = content.signature, content = content.thinking, metaInfo = metaInfo)`
`511`	`511`	`}`
`512`	`512`
`513`	`513`	`is AnthropicContent.ToolUse -> {`
Original file line number	Diff line number	Diff line change
`@@ -22,8 +22,6 @@ import kotlinx.serialization.json.buildJsonArray`
`22`	`22`	`import kotlinx.serialization.json.buildJsonObject`
`23`	`23`	`import kotlinx.serialization.json.encodeToJsonElement`
`24`	`24`	`import kotlinx.serialization.json.put`
`25`		`-import kotlin.uuid.ExperimentalUuidApi`
`26`		`-import kotlin.uuid.Uuid`
`27`	`25`
`28`	`26`	`internal object BedrockAnthropicClaudeSerialization {`
`29`	`27`
`@@ -35,7 +33,6 @@ internal object BedrockAnthropicClaudeSerialization {`
`35`	`33`	`explicitNulls = false`
`36`	`34`	`}`
`37`	`35`
`38`		`- @OptIn(ExperimentalUuidApi::class)`
`39`	`36`	`private fun buildMessagesHistory(prompt: Prompt): MutableList<BedrockAnthropicInvokeModelMessage> {`
`40`	`37`	`val messages = mutableListOf<BedrockAnthropicInvokeModelMessage>()`
`41`	`38`	`prompt.messages.forEach { msg ->`
`@@ -68,11 +65,10 @@ internal object BedrockAnthropicClaudeSerialization {`
`68`	`65`	`messages.add(`
`69`	`66`	`BedrockAnthropicInvokeModelMessage.Assistant(`
`70`	`67`	`content = listOf(`
`71`		`- msg.original as? BedrockAnthropicInvokeModelContent.Thinking`
`72`		`- ?: BedrockAnthropicInvokeModelContent.Thinking(`
`73`		`- signature = Uuid.random().toString(),`
`74`		`- thinking = msg.content`
`75`		`- )`
	`68`	`+ BedrockAnthropicInvokeModelContent.Thinking(`
	`69`	`+ signature = msg.encrypted!!,`
	`70`	`+ thinking = msg.content`
	`71`	`+ )`
`76`	`72`	`)`
`77`	`73`	`)`
`78`	`74`	`)`
`@@ -181,7 +177,6 @@ internal object BedrockAnthropicClaudeSerialization {`
`181`	`177`	`)`
`182`	`178`	`}`
`183`	`179`
`184`		`- @OptIn(ExperimentalUuidApi::class)`
`185`	`180`	`internal fun parseAnthropicResponse(responseBody: String, clock: Clock = Clock.System): List<Message.Response> {`
`186`	`181`	`val response = json.decodeFromString<BedrockAnthropicResponse>(responseBody)`
`187`	`182`
`@@ -204,7 +199,7 @@ internal object BedrockAnthropicClaudeSerialization {`
`204`	`199`	`)`
`205`	`200`
`206`	`201`	`is AnthropicContent.Thinking -> Message.Reasoning(`
`207`		`- original = content,`
	`202`	`+ encrypted = content.signature,`
`208`	`203`	`content = content.thinking,`
`209`	`204`	`metaInfo = metaInfo`
`210`	`205`	`)`
`@@ -255,6 +250,7 @@ internal object BedrockAnthropicClaudeSerialization {`
`255`	`250`	`name = content.name,`
`256`	`251`	`content = content.input.toString()`
`257`	`252`	`)`
	`253`	`+`
`258`	`254`	`else -> throw IllegalArgumentException(`
`259`	`255`	`"Unsupported AnthropicContent type in message_delta. Content: $content"`
`260`	`256`	`)`
Original file line number	Diff line number	Diff line change
`@@ -333,9 +333,10 @@ public open class GoogleLLMClient(`
`333`	`333`	`GoogleContent(`
`334`	`334`	`role = "assistant",`
`335`	`335`	`parts = listOf(`
`336`		`- message.original as? GooglePart ?: GooglePart.Text(`
	`336`	`+ GooglePart.Text(`
`337`	`337`	`text = message.content,`
`338`		`- thought = true`
	`338`	`+ thoughtSignature = message.encrypted,`
	`339`	`+ thought = true,`
`339`	`340`	`)`
`340`	`341`	`)`
`341`	`342`	`)`
`@@ -606,7 +607,7 @@ public open class GoogleLLMClient(`
`606`	`607`	`is GooglePart.Text -> {`
`607`	`608`	`if (part.thought ?: false) {`
`608`	`609`	`Message.Reasoning(`
`609`		`- original = part,`
	`610`	`+ encrypted = part.thoughtSignature,`
`610`	`611`	`content = part.text,`
`611`	`612`	`metaInfo = metaInfo`
`612`	`613`	`)`
Original file line number	Diff line number	Diff line change
`@@ -602,9 +602,10 @@ public open class OpenAILLMClient(`
`602`	`602`	`is Message.Reasoning -> {`
`603`	`603`	`flushPendingCalls()`
`604`	`604`	`add(`
`605`		`- message.original as? Item.Reasoning ?: Item.Reasoning(`
`606`		`- Uuid.random().toString(),`
`607`		`- listOf(Item.Reasoning.Summary(message.content))`
	`605`	`+ Item.Reasoning(`
	`606`	`+ id = message.id ?: Uuid.random().toString(),`
	`607`	`+ encryptedContent = message.encrypted,`
	`608`	`+ summary = listOf(Item.Reasoning.Summary(message.content))`
`608`	`609`	`)`
`609`	`610`	`)`
`610`	`611`	`}`
`@@ -709,7 +710,8 @@ public open class OpenAILLMClient(`
`709`	`710`	`)`
`710`	`711`
`711`	`712`	`is Item.Reasoning -> Message.Reasoning(`
`712`		`- original = output,`
	`713`	`+ id = output.id,`
	`714`	`+ encrypted = output.encryptedContent,`
`713`	`715`	`content = output.summary.joinToString(separator = "\n") { it.text },`
`714`	`716`	`metaInfo = metaInfo`
`715`	`717`	`)`
`@@ -735,13 +737,15 @@ public open class OpenAILLMClient(`
`735`	`737`	`)`
`736`	`738`	`params`
`737`	`739`	`}`
	`740`	`+`
`738`	`741`	`params is OpenAIChatParams -> {`
`739`	`742`	`model.requireCapability(`
`740`	`743`	`LLMCapability.OpenAIEndpoint.Completions,`
`741`	`744`	`message = "Must be supported to use OpenAI chat params."`
`742`	`745`	`)`
`743`	`746`	`params`
`744`	`747`	`}`
	`748`	`+`
`745`	`749`	`model.supports(LLMCapability.OpenAIEndpoint.Completions) -> params.toOpenAIChatParams()`
`746`	`750`	`model.supports(LLMCapability.OpenAIEndpoint.Responses) -> params.toOpenAIResponsesParams()`
`747`	`751`	`else -> error("Cannot determine proper LLM params for OpenAI model: ${model.id}")`