Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import com.embabel.common.ai.autoconfig.NativeSupport
import com.embabel.common.ai.model.LlmMetadata
import com.embabel.common.util.loggerFor
import com.embabel.agent.spi.support.nativeoutput.shouldUseNativeStructuredOutput
import org.springframework.ai.chat.messages.AssistantMessage
import org.springframework.ai.chat.model.ChatModel
import org.springframework.ai.chat.model.ChatResponse
import org.springframework.ai.chat.prompt.ChatOptions
Expand Down Expand Up @@ -96,12 +97,15 @@ internal class SpringAiLlmMessageSender(

logger.debug("Prompt: {}\nResponse: {}", prompt, response)

// Convert response to Embabel message
// Note: Some providers (e.g., Bedrock) may return multiple generations where
// the first is empty and the second contains tool calls. We need to find the
// generation with tool calls, or fall back to the first one if none have them.
// See: https://github.com/embabel/embabel-agent/issues/1350
val assistantMessage = findGenerationWithToolCalls(response) ?: response.result!!.output
// Convert response to Embabel message.

// Providers may return multiple generations in one ChatResponse:
// - Bedrock: empty first generation, tool calls on a later one (#1350)
// - Google GenAI includeThoughts: thought parts first (isThought=true), answer later.

// Using only ChatResponse.result (first generation) discards later answer text
// and breaks structured output / createObject after thought-signature support.
val assistantMessage = resolveAssistantMessage(response)
val embabelMessage = assistantMessage.toEmbabelMessage()

// Extract usage information
Expand All @@ -115,60 +119,129 @@ internal class SpringAiLlmMessageSender(
}

/**
* Find the best generation to use from the response.
* Resolve a Spring AI response into the assistant message Embabel should store and inspect.
*
* Some providers (e.g., Bedrock) may return multiple generations where
* the first is empty and a subsequent one contains tool calls.
* Spring AI exposes provider response parts as generations. For providers that split a
* single answer across generations, this method preserves the pieces Embabel needs while
* avoiding unsafe concatenation of alternative structured answers.
*
* Strategy:
* 1. Collect all tool calls from all generations
* 2. Collect all text content from all generations
* 3. If there are tool calls, create a merged AssistantMessage with all tool calls and combined text
* 4. If no tool calls, return null to fall back to first generation
*
* This ensures we don't lose valuable content (text or tool calls) from any generation.
*
* @return A merged AssistantMessage with all tool calls and text, or null if no tool calls found
* 1. Collect tool calls and metadata from every generation.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

GEN AI-specific logic?

* 2. Prefer non-thought text (Google GenAI sets metadata `isThought=true` on thought parts).
* Fall back to the first non-blank generation text when `isThought` is absent.
* 3. If tool calls exist, return a merged message with all tool calls and selected text.
* 4. If only text exists, return selected answer text with merged metadata.
*/
private fun findGenerationWithToolCalls(response: ChatResponse): org.springframework.ai.chat.messages.AssistantMessage? {
private fun resolveAssistantMessage(
response: ChatResponse,
): AssistantMessage {
val allOutputs = response.results.map { it.output }
require(allOutputs.isNotEmpty()) { "ChatResponse contained no generations" }

// Collect all tool calls from all generations
val allToolCalls = allOutputs
.flatMap { it.toolCalls ?: emptyList() }

if (allToolCalls.isEmpty()) {
return null // No tool calls found, let caller use first generation
}

// Collect all metadata from all generations
val allToolCalls = allOutputs.flatMap { it.toolCalls ?: emptyList() }
val allMetaData: Map<String, Any> = allOutputs
.mapNotNull { it.metadata }
.fold(emptyMap()) { acc, metadata -> acc + metadata }

// Collect all non-empty text from all generations
val allText = allOutputs
.mapNotNull { it.text?.takeIf { text -> text.isNotBlank() } }
.joinToString("\n")

// Log if we're merging content from multiple generations
val answerText = if (allToolCalls.isNotEmpty()) {
selectToolCallText(allOutputs)
} else {
selectAnswerText(allOutputs)
}
val generationsWithToolCalls = allOutputs.count { !it.toolCalls.isNullOrEmpty() }
val generationsWithText = allOutputs.count { !it.text.isNullOrBlank() }
if (generationsWithToolCalls > 1 || generationsWithText > 1) {
logger.debug(
"Merging content from multiple generations: {} with tool calls, {} with text",
"Resolving multi-generation ChatResponse: {} with tool calls, {} with text, selected answer length={}",
generationsWithToolCalls,
generationsWithText
generationsWithText,
answerText.length,
)
}

return org.springframework.ai.chat.messages.AssistantMessage.builder()
.content(allText)
.toolCalls(allToolCalls)
if (allToolCalls.isNotEmpty()) {
return AssistantMessage.builder()
.content(answerText)
.toolCalls(allToolCalls)
.properties(allMetaData)
.build()
}

return AssistantMessage.builder()
.content(answerText)
.properties(allMetaData)
.build()
}

/**
* Select text that should be treated as the model answer for structured conversion.
*
* Google GenAI with includeThoughts emits one generation per part; thought parts are
* marked with metadata isThought=true and must not be used alone as the JSON payload.
* When multiple non-thought texts are present, the first one is selected because those
* generations may be alternative candidates rather than chunks of one JSON document.
*/
private fun selectAnswerText(
allOutputs: List<AssistantMessage>,
): String {
val nonThoughtTexts = allOutputs
.filterNot { isThoughtGeneration(it) }
.mapNotNull { it.text?.takeIf { text -> text.isNotBlank() } }
if (nonThoughtTexts.isNotEmpty()) {
return nonThoughtTexts.first()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The previous code returned text from all generations. Here you are just returning the first. I understand the reason from the function comment, however is this specific to the Google GenAI implementation? How does this affect responses from other providers?

}
// No non-thought text (or provider does not mark thoughts): use first non-blank content
return allOutputs
.mapNotNull { it.text?.takeIf { text -> text.isNotBlank() } }
.firstOrNull()
?: ""
}

/**
* Select text for responses that include tool calls.
*
* Tool-call responses need to keep tool calls from all generations. Text handling is more
* conservative: if no generation is marked as thought, all non-blank text is joined to
* preserve Bedrock-style split responses. If thought markers are present, thought text is
* removed so structured answer content and tool continuation metadata stay coherent.
*/
private fun selectToolCallText(
allOutputs: List<AssistantMessage>,
): String {
val textOutputs = allOutputs.mapNotNull { it.text?.takeIf { text -> text.isNotBlank() } }
if (allOutputs.none { isThoughtGeneration(it) }) {
return textOutputs.joinToString("\n")
}
return allOutputs
.filterNot { isThoughtGeneration(it) }
.mapNotNull { it.text?.takeIf { text -> text.isNotBlank() } }
.joinToString("\n")
}

/**
* Return true when a generation is provider-marked as model thinking rather than final
* assistant answer text.
*
* Spring AI's Google GenAI adapter uses Boolean `true`; trimmed string values are accepted
* so metadata copied through less strongly typed paths is still filtered correctly.
*/
private fun isThoughtGeneration(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would be concerned about putting provider-specific metadata handling like isThought directly into SpringAiLlmMessageSender. SpringAiLlmMessageSender is a shared Spring AI integration layer, so once we start coding Google-specific response semantics in here we are on a slippery slope. It would be better if we could make response resolution provider-pluggable in some manner. The Spring team have yet to deal with this themselves as outlined here: spring-projects/spring-ai#4269

message: AssistantMessage,
): Boolean = when (val isThought = message.metadata?.get(IS_THOUGHT_METADATA_KEY)) {
true -> true
is String -> isThought.trim().equals("true", ignoreCase = true)
else -> false
}

companion object {
/**
* Metadata key set by Spring AI Google GenAI on thought parts
* (`GoogleGenAiChatModel.responseCandidateToGeneration`).
*/
const val IS_THOUGHT_METADATA_KEY: String = "isThought"
}

/**
* Build ChatOptions with tool definitions.
* Tools are passed to the LLM so it knows what's available,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -147,11 +147,13 @@ fun SpringAiAssistantMessage.toEmbabelMessage(): Message {
val toolCalls = this.toolCalls
val content = this.text ?: ""
val metadata = this.metadata ?: emptyMap()
val hasProviderMetadata = metadata.keys.any { it != "messageType" }

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could do with a comment to explain how this resolves to having provider metadata.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

{ it != "messageType" } ==> specific to GEN AI provider?

return if (toolCalls.isNullOrEmpty()) {
// AssistantMessage requires non-empty content (TextPart validation).
// For empty content, use AssistantMessageWithToolCalls which handles empty content gracefully.
if (content.isEmpty()) {
AssistantMessageWithToolCalls(content = "", toolCalls = emptyList(), metadata = metadata)

// AssistantMessage requires non-empty content and does not carry provider metadata. Use
// AssistantMessageWithToolCalls with an empty tool-call list when metadata must survive.
if (content.isEmpty() || hasProviderMetadata) {
AssistantMessageWithToolCalls(content = content, toolCalls = emptyList(), metadata = metadata)
} else {
AssistantMessage(content = content)
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -349,6 +349,95 @@ class MessageConversionTest {
assertThat(signatures[1] as ByteArray).containsExactly(30, 40)
}

@Test
fun `preserves thoughtSignatures metadata when tool calls are present with non-empty content`() {
// Prepare
val thoughtSignatures = listOf(byteArrayOf(1, 2, 3))
val toolCalls = listOf(
SpringAiAssistantMessage.ToolCall("call-1", "function", "lookup", "{}"),
)
val springMessage = SpringAiAssistantMessage.builder()
.content("calling tool")
.toolCalls(toolCalls)
.properties(mapOf("thoughtSignatures" to thoughtSignatures, "isThought" to false))
.build()

// Execute
val embabelMessage = springMessage.toEmbabelMessage()

// Verify
assertThat(embabelMessage).isInstanceOf(AssistantMessageWithToolCalls::class.java)
val messageWithCalls = embabelMessage as AssistantMessageWithToolCalls
assertThat(messageWithCalls.toolCalls).hasSize(1)

val signatures = messageWithCalls.metadata["thoughtSignatures"] as? List<*>
assertThat(signatures).isNotNull
assertThat(signatures!![0] as ByteArray).containsExactly(1, 2, 3)
assertThat(messageWithCalls.metadata["isThought"]).isEqualTo(false)
}

@Test
fun `preserves custom metadata for non-empty assistant content without tool calls`() {
// Prepare
val thoughtSignatures = listOf(byteArrayOf(7, 8))
val springMessage = SpringAiAssistantMessage.builder()
.content("""{"name":"July"}""")
.properties(mapOf("thoughtSignatures" to thoughtSignatures, "isThought" to false))
.build()

// Execute
val embabelMessage = springMessage.toEmbabelMessage()

// Verify
assertThat(embabelMessage).isInstanceOf(AssistantMessageWithToolCalls::class.java)
assertThat(embabelMessage.content).isEqualTo("""{"name":"July"}""")

val messageWithCalls = embabelMessage as AssistantMessageWithToolCalls
assertThat(messageWithCalls.toolCalls).isEmpty()

val signatures = messageWithCalls.metadata["thoughtSignatures"] as? List<*>
assertThat(signatures).isNotNull
assertThat(signatures!![0] as ByteArray).containsExactly(7, 8)
assertThat(messageWithCalls.metadata["isThought"]).isEqualTo(false)
}

@Test
fun `default messageType metadata alone keeps non-empty content as plain assistant message`() {
// Prepare
val springMessage = SpringAiAssistantMessage.builder()
.content("plain answer")
.properties(mapOf("messageType" to "ASSISTANT"))
.build()

// Execute
val embabelMessage = springMessage.toEmbabelMessage()

// Verify
assertThat(embabelMessage).isInstanceOf(AssistantMessage::class.java)
assertThat(embabelMessage).isNotInstanceOf(AssistantMessageWithToolCalls::class.java)
assertThat(embabelMessage.content).isEqualTo("plain answer")
}

@Test
fun `preserves empty thoughtSignatures metadata for non-empty assistant content without tool calls`() {
// Prepare
val springMessage = SpringAiAssistantMessage.builder()
.content("""{"name":"July"}""")
.properties(mapOf("thoughtSignatures" to emptyList<ByteArray>(), "isThought" to false))
.build()

// Execute
val embabelMessage = springMessage.toEmbabelMessage()

// Verify
assertThat(embabelMessage).isInstanceOf(AssistantMessageWithToolCalls::class.java)

val messageWithCalls = embabelMessage as AssistantMessageWithToolCalls
assertThat(messageWithCalls.toolCalls).isEmpty()
assertThat(messageWithCalls.metadata["thoughtSignatures"] as? List<*>).isEmpty()
assertThat(messageWithCalls.metadata["isThought"]).isEqualTo(false)
}

@Test
fun `converts Spring AI AssistantMessage with tool calls`() {
val toolCalls = listOf(
Expand Down
Loading