From 7f68fcacf6d790c2c23dc177203fb6566f818650 Mon Sep 17 00:00:00 2001 From: "Eric (Codex)" Date: Fri, 7 Aug 2026 13:15:53 +0800 Subject: [PATCH] feat(thinking): add tag selection API with include/exclude filtering (#1790) Add ThinkingTagSelection with include/exclude filtering for XML-style thinking blocks. Non-TAG blocks (prefix/untagged) are always retained. Missing tags trigger a WARN log to make silent hasThinking()==false diagnosable. thinking(tag) convenience shorthand for single-tag extraction. --- .../embabel/agent/api/common/PromptRunner.kt | 12 ++ .../DelegatingStreamingPromptRunner.kt | 10 ++ .../api/common/support/DelegatingThinking.kt | 10 +- .../support/OperationContextDelegate.kt | 22 +++- .../common/support/PromptExecutionDelegate.kt | 13 +- .../agent/core/support/LlmInteraction.kt | 47 +++++++ .../spi/support/ToolLoopLlmOperations.kt | 23 +++- .../springai/ChatClientLlmOperations.kt | 15 ++- .../agent/test/unit/FakePromptRunner.kt | 9 +- .../DelegatingStreamingPromptRunnerTest.kt | 32 +++++ .../common/support/DelegatingThinkingTest.kt | 71 +++++++++-- .../LlmInteractionSerializationTest.kt | 115 ++++++++++++++++++ .../ChatClientLlmOperationsThinkingTest.kt | 35 ++++++ .../asciidoc/reference/thinking/page.adoc | 59 ++++++++- 14 files changed, 438 insertions(+), 35 deletions(-) diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/PromptRunner.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/PromptRunner.kt index 2ec012430..da12eb82a 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/PromptRunner.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/PromptRunner.kt @@ -487,6 +487,18 @@ interface PromptRunner : LlmUse, PromptRunnerOperations, ToolChaining): Thinking = thinking(include, emptySet()) + + // Non-empty include/exclude throws UnsupportedOperationException by default — implementations must override. + fun thinking(include: Set, exclude: Set): Thinking = + if (include.isEmpty() && exclude.isEmpty()) thinking() + else throw UnsupportedOperationException( + "Thinking tag selection is not supported by this PromptRunner implementation" + ) + + // Default impl ignores the tag — implementations that support tag selection must override. + fun thinking(tag: String): Thinking = thinking() + override fun respond( messages: List, ): AssistantMessage = diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/DelegatingStreamingPromptRunner.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/DelegatingStreamingPromptRunner.kt index 651e5b87b..590d9e29b 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/DelegatingStreamingPromptRunner.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/DelegatingStreamingPromptRunner.kt @@ -30,6 +30,7 @@ import com.embabel.agent.api.tool.callback.ToolLoopTransformer import com.embabel.agent.api.validation.guardrails.GuardRail import com.embabel.agent.core.ToolGroup import com.embabel.agent.core.ToolGroupRequirement +import com.embabel.agent.core.support.ThinkingTagSelection import com.embabel.agent.experimental.primitive.Determination import com.embabel.agent.spi.loop.ToolInjectionStrategy import com.embabel.agent.spi.loop.ToolNotFoundPolicy @@ -235,4 +236,13 @@ internal data class DelegatingStreamingPromptRunner( DelegatingThinking( delegate = delegate, ) + + override fun thinking(include: Set, exclude: Set): PromptRunner.Thinking = + DelegatingThinking( + delegate = delegate, + selection = ThinkingTagSelection(include = include, exclude = exclude), + ) + + override fun thinking(tag: String): PromptRunner.Thinking = + thinking(include = setOf(tag)) } diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/DelegatingThinking.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/DelegatingThinking.kt index 017d0e43d..536f4baec 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/DelegatingThinking.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/DelegatingThinking.kt @@ -16,6 +16,7 @@ package com.embabel.agent.api.common.support import com.embabel.agent.api.common.PromptRunner +import com.embabel.agent.core.support.ThinkingTagSelection import com.embabel.chat.AssistantMessage import com.embabel.chat.Message import com.embabel.common.core.thinking.ThinkingResponse @@ -26,27 +27,28 @@ import com.embabel.common.core.types.ZeroToOne */ internal data class DelegatingThinking( private val delegate: PromptExecutionDelegate, + private val selection: ThinkingTagSelection = ThinkingTagSelection(), ) : PromptRunner.Thinking { override fun createObjectIfPossible( messages: List, outputClass: Class ): ThinkingResponse = - delegate.createObjectIfPossibleWithThinking(messages, outputClass) + delegate.createObjectIfPossibleWithThinking(messages, outputClass, selection) override fun createObject( messages: List, outputClass: Class ): ThinkingResponse = - delegate.createObjectWithThinking(messages, outputClass) + delegate.createObjectWithThinking(messages, outputClass, selection) override fun respond(messages: List): ThinkingResponse = - delegate.respondWithThinking(messages) + delegate.respondWithThinking(messages, selection) override fun evaluateCondition( condition: String, context: String, confidenceThreshold: ZeroToOne ): ThinkingResponse = - delegate.evaluateConditionWithThinking(condition, context, confidenceThreshold) + delegate.evaluateConditionWithThinking(condition, context, confidenceThreshold, selection) } diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/OperationContextDelegate.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/OperationContextDelegate.kt index 0b5f9299f..bbb2ed16a 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/OperationContextDelegate.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/OperationContextDelegate.kt @@ -34,6 +34,7 @@ import com.embabel.agent.core.ToolGroup import com.embabel.agent.core.ToolGroupRequirement import com.embabel.agent.core.internal.LlmOperations import com.embabel.agent.core.support.LlmInteraction +import com.embabel.agent.core.support.ThinkingTagSelection import com.embabel.agent.core.support.safelyGetTools import com.embabel.agent.experimental.primitive.Determination import com.embabel.agent.core.internal.streaming.StreamingLlmOperationsFactory @@ -420,11 +421,13 @@ internal data class OperationContextDelegate( // Patterned after createObject() - uses ProcessContext flow override fun createObjectWithThinking( messages: List, - outputClass: Class + outputClass: Class, + thinkingTags: ThinkingTagSelection, ): ThinkingResponse { val combinedMessages = combineImagesWithMessages(this.messages + messages) val interaction = thinkingInteraction( toolGroups = this.toolGroups + toolGroups, + thinkingTags = thinkingTags, ) return context.processContext.createObjectWithThinking( messages = combinedMessages, @@ -438,11 +441,13 @@ internal data class OperationContextDelegate( // Patterned after createObjectWithThinking() - uses ProcessContext flow override fun createObjectIfPossibleWithThinking( messages: List, - outputClass: Class + outputClass: Class, + thinkingTags: ThinkingTagSelection, ): ThinkingResponse { val combinedMessages = combineImagesWithMessages(this.messages + messages) val interaction = thinkingInteraction( toolGroups = this.toolGroups + toolGroups, + thinkingTags = thinkingTags, ) val result = context.processContext.createObjectIfPossibleWithThinking( messages = combinedMessages, @@ -478,14 +483,18 @@ internal data class OperationContextDelegate( } } - override fun respondWithThinking(messages: List): ThinkingResponse { - return createObjectWithThinking(messages, AssistantMessage::class.java) + override fun respondWithThinking( + messages: List, + thinkingTags: ThinkingTagSelection, + ): ThinkingResponse { + return createObjectWithThinking(messages, AssistantMessage::class.java, thinkingTags) } override fun evaluateConditionWithThinking( condition: String, context: String, - confidenceThreshold: ZeroToOne + confidenceThreshold: ZeroToOne, + thinkingTags: ThinkingTagSelection, ): ThinkingResponse { val prompt = """ @@ -503,6 +512,7 @@ internal data class OperationContextDelegate( val response = createObjectWithThinking( messages = listOf(UserMessage(prompt)), outputClass = Determination::class.java, + thinkingTags = thinkingTags, ) val result = response.result?.let { @@ -517,6 +527,7 @@ internal data class OperationContextDelegate( private fun thinkingInteraction( toolGroups: Set = this.toolGroups, + thinkingTags: ThinkingTagSelection = ThinkingTagSelection(), ): LlmInteraction { val thinkingEnabledLlm = llm.withThinking(Thinking.withExtraction()) val toolConfig = resolveToolConfig() @@ -537,6 +548,7 @@ internal data class OperationContextDelegate( toolLoopTransformers = toolLoopTransformers, toolCallInspectors = toolCallInspectors, toolCallContext = toolCallContext, + thinkingTags = thinkingTags, ) } diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/PromptExecutionDelegate.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/PromptExecutionDelegate.kt index a62738c21..3e14009ea 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/PromptExecutionDelegate.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/support/PromptExecutionDelegate.kt @@ -31,6 +31,7 @@ import com.embabel.agent.core.ToolGroup import com.embabel.agent.core.ToolGroupRequirement import com.embabel.agent.core.internal.LlmOperations import com.embabel.agent.core.support.LlmUse +import com.embabel.agent.core.support.ThinkingTagSelection import com.embabel.agent.spi.loop.ToolInjectionStrategy import com.embabel.agent.spi.loop.ToolNotFoundPolicy import com.embabel.chat.AssistantMessage @@ -148,19 +149,25 @@ internal interface PromptExecutionDelegate : LlmUse { fun createObjectIfPossibleWithThinking( messages: List, outputClass: Class, + thinkingTags: ThinkingTagSelection = ThinkingTagSelection(), ): ThinkingResponse fun createObjectWithThinking( messages: List, - outputClass: Class + outputClass: Class, + thinkingTags: ThinkingTagSelection = ThinkingTagSelection(), ): ThinkingResponse - fun respondWithThinking(messages: List): ThinkingResponse + fun respondWithThinking( + messages: List, + thinkingTags: ThinkingTagSelection = ThinkingTagSelection(), + ): ThinkingResponse fun evaluateConditionWithThinking( condition: String, context: String, - confidenceThreshold: ZeroToOne + confidenceThreshold: ZeroToOne, + thinkingTags: ThinkingTagSelection = ThinkingTagSelection(), ): ThinkingResponse } diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/core/support/LlmInteraction.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/core/support/LlmInteraction.kt index 0ae595a2e..1ac68f3a8 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/core/support/LlmInteraction.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/core/support/LlmInteraction.kt @@ -31,12 +31,15 @@ import com.embabel.common.ai.model.LlmOptions import com.embabel.common.ai.prompt.PromptContributor import com.embabel.common.ai.prompt.PromptContributorConsumer import com.embabel.common.core.MobyNameGenerator +import com.embabel.common.core.thinking.ThinkingBlock +import com.embabel.common.core.thinking.ThinkingTagType import com.embabel.common.core.types.HasInfoString import com.embabel.common.util.indent import com.fasterxml.jackson.annotation.JsonIgnore import jakarta.validation.ConstraintViolation import java.lang.reflect.Field import java.util.function.Predicate +import org.slf4j.Logger /** * Spec for calling an LLM. Optional LlmOptions, @@ -99,6 +102,49 @@ private data class LlmCallImpl( override val validation: Boolean = true, ) : LlmCall +data class ThinkingTagSelection( + val include: Set = emptySet(), + val exclude: Set = emptySet(), +) { + init { + require((include + exclude).all(TAG_NAME_REGEX::matches)) { + "Thinking tags must be valid XML-style tag names matching [a-zA-Z][a-zA-Z0-9_-]*" + } + require(include.intersect(exclude).isEmpty()) { + "Thinking tags cannot be both included and excluded" + } + } + + // Non-TAG blocks (prefix/untagged) are always retained — only XML TAG blocks are filtered. + fun filter(blocks: List): List = + blocks.filter { block -> + block.tagType != ThinkingTagType.TAG || + ((include.isEmpty() || block.tagValue in include) && block.tagValue !in exclude) + } + + fun missingFrom(systemPrompt: String): Set = + include.filter { tag -> !systemPrompt.contains(TAG_OPEN_REGEX(tag)) }.toSet() + + // Warns when a declared tag is absent from the system prompt — without it, hasThinking() silently returns false. + fun warnIfMissing(systemPrompt: String, logger: Logger) { + val missing = missingFrom(systemPrompt) + if (missing.isNotEmpty()) { + logger.warn( + "Thinking tag(s) {} not found in the system prompt — the LLM may not generate matching blocks", + missing, + ) + } + } + + companion object { + val TAG_NAME_REGEX = Regex("[a-zA-Z][a-zA-Z0-9_-]*") + + // Boundary after the tag name: '>' (open), whitespace (attributes), or '/' (self-closing). + // Without it, "think" would match "". + private fun TAG_OPEN_REGEX(tag: String): Regex = Regex("<$tag(?:>|\\s|/)") + } +} + /** * Encapsulates an interaction with an LLM. * An LlmInteraction is a specific instance of an LlmCall. @@ -133,6 +179,7 @@ data class LlmInteraction( val toolCallInspectors: List = emptyList(), val toolCallContext: ToolCallContext = ToolCallContext.EMPTY, val toolNotFoundPolicy: ToolNotFoundPolicy? = null, + val thinkingTags: ThinkingTagSelection = ThinkingTagSelection(), ) : LlmCall { override val name: String = id.value diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/ToolLoopLlmOperations.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/ToolLoopLlmOperations.kt index c7feed805..7efeffe1d 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/ToolLoopLlmOperations.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/ToolLoopLlmOperations.kt @@ -36,6 +36,7 @@ import com.embabel.agent.core.ReplanRequestedException import com.embabel.agent.core.Usage import com.embabel.agent.core.support.LlmCall import com.embabel.agent.core.support.LlmInteraction +import com.embabel.agent.core.support.ThinkingTagSelection import com.embabel.agent.spi.AutoLlmSelectionCriteriaResolver import com.embabel.agent.spi.LlmService import com.embabel.agent.spi.ToolDecorator @@ -361,7 +362,7 @@ open class ToolLoopLlmOperations( // For String output: return raw text (with thinking tags preserved) // For other types: converter chain handles thinking suppression for JSON parsing val outputParser: (String) -> ThinkingResponse = { text -> - val thinkingBlocks = extractAllThinkingBlocks(text) + val thinkingBlocks = interaction.thinkingTags.filter(extractAllThinkingBlocks(text)) val result = if (outputClass == String::class.java) { @Suppress("UNCHECKED_CAST") text as O // Raw text, not sanitized - thinking blocks preserved in response @@ -401,6 +402,8 @@ open class ToolLoopLlmOperations( val initialMessages = buildInitialMessages(promptContributions, messages, schemaFormat) + interaction.thinkingTags.warnIfMissing(systemPromptOf(initialMessages), logger) + emitCallEvent(llmRequestEvent, promptContributions, messages, schemaFormat) // Guardrails: Pre-validation of user input @@ -426,7 +429,7 @@ open class ToolLoopLlmOperations( // Filter by role to catch both AssistantMessage and AssistantMessageWithToolCalls val allThinkingBlocks = result.conversationHistory .filter { it.role == com.embabel.chat.Role.ASSISTANT } - .flatMap { extractAllThinkingBlocks(it.content) } + .flatMap { interaction.thinkingTags.filter(extractAllThinkingBlocks(it.content)) } // Merge accumulated thinking blocks with the final result val thinkingResponse = ThinkingResponse( @@ -459,7 +462,7 @@ open class ToolLoopLlmOperations( // Output parser: extract thinking blocks FIRST, then parse MaybeReturn val outputParser: (String) -> Result> = { text -> - val thinkingBlocks = extractAllThinkingBlocks(text) + val thinkingBlocks = interaction.thinkingTags.filter(extractAllThinkingBlocks(text)) try { val maybeResult = if (text.isNotBlank()) { converter.convert(text)!! @@ -530,6 +533,8 @@ open class ToolLoopLlmOperations( schemaFormat, ) + interaction.thinkingTags.warnIfMissing(systemPromptOf(initialMessages), logger) + emitCallEvent(llmRequestEvent, promptContributions, messages, schemaFormat) // Guardrails: Pre-validation of user input @@ -553,7 +558,7 @@ open class ToolLoopLlmOperations( // Accumulate thinking blocks from ALL assistant messages across all iterations // Filter by role to catch both AssistantMessage and AssistantMessageWithToolCalls - val allThinkingBlocks = accumulateThinkingBlocks(result.conversationHistory) + val allThinkingBlocks = accumulateThinkingBlocks(interaction.thinkingTags, result.conversationHistory) // Merge accumulated thinking blocks with the final result (success or failure path) val thinkingResult = mergeThinkingBlocksWithResult(finalIterationResult, allThinkingBlocks) @@ -694,6 +699,9 @@ open class ToolLoopLlmOperations( llm: LlmService<*>, ): String = buildPromptContributionsString(interaction.promptContributors, llm.promptContributors) + private fun systemPromptOf(messages: List): String = + messages.filterIsInstance().joinToString("\n\n") { it.content } + /** * Build initial messages for the tool loop, including system prompt contributions and schema. * All system content is consolidated into a single system message at the beginning @@ -945,10 +953,13 @@ open class ToolLoopLlmOperations( * Filters by ASSISTANT role to catch both AssistantMessage and AssistantMessageWithToolCalls. */ @OptIn(InternalThinkingApi::class) - private fun accumulateThinkingBlocks(conversationHistory: List): List { + private fun accumulateThinkingBlocks( + thinkingTags: ThinkingTagSelection, + conversationHistory: List, + ): List { return conversationHistory .filter { it.role == com.embabel.chat.Role.ASSISTANT } - .flatMap { extractAllThinkingBlocks(it.content) } + .flatMap { thinkingTags.filter(extractAllThinkingBlocks(it.content)) } } /** diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/springai/ChatClientLlmOperations.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/springai/ChatClientLlmOperations.kt index a33a52044..c1f953891 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/springai/ChatClientLlmOperations.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/springai/ChatClientLlmOperations.kt @@ -343,6 +343,8 @@ internal class ChatClientLlmOperations( buildBasicPrompt(promptContributions, messages) } + interaction.thinkingTags.warnIfMissing(systemPromptOf(springAiPrompt), logger) + // Guardrails: Pre-validation of user input val userMessages = messages.filterIsInstance() validateUserInput(userMessages, interaction, llmRequestEvent?.agentProcess?.blackboard) @@ -382,7 +384,7 @@ internal class ChatClientLlmOperations( recordUsage(llm, chatResponse, llmRequestEvent) val rawText = chatResponse.result.output.text as String - val thinkingBlocks = extractAllThinkingBlocks(rawText) + val thinkingBlocks = interaction.thinkingTags.filter(extractAllThinkingBlocks(rawText)) logger.debug("Extracted {} thinking blocks for String response", thinkingBlocks.size) val thinkingResponse = ThinkingResponse( @@ -400,7 +402,7 @@ internal class ChatClientLlmOperations( recordUsage(llm, chatResponse, llmRequestEvent) val rawText = chatResponse.result.output.text ?: "" - val thinkingBlocks = extractAllThinkingBlocks(rawText) + val thinkingBlocks = interaction.thinkingTags.filter(extractAllThinkingBlocks(rawText)) logger.debug( "Extracted {} thinking blocks for {} response", thinkingBlocks.size, @@ -490,6 +492,8 @@ internal class ChatClientLlmOperations( schemaFormat ) + interaction.thinkingTags.warnIfMissing(systemPromptOf(springAiPrompt), logger) + // Guardrails: Pre-validation of user input val userMessages = messages.filterIsInstance() validateUserInput(userMessages, interaction, llmRequestEvent?.agentProcess?.blackboard) @@ -524,7 +528,7 @@ internal class ChatClientLlmOperations( val chatResponse = requireChatResponse(callResponse, interaction) recordUsage(llm, chatResponse, llmRequestEvent) val rawText = chatResponse.result.output.text ?: "" - val thinkingBlocks = extractAllThinkingBlocks(rawText) + val thinkingBlocks = interaction.thinkingTags.filter(extractAllThinkingBlocks(rawText)) // Execute converter chain manually instead of using responseEntity try { @@ -672,6 +676,11 @@ internal class ChatClientLlmOperations( // SPRING AI PROMPT BUILDERS // ==================================== + private fun systemPromptOf(springAiPrompt: Prompt): String = + springAiPrompt.instructions + .filterIsInstance() + .joinToString("\n\n") { it.text ?: "" } + /** * Base prompt builder - consolidates all system messages at the beginning. * Extracts SystemMessages from the input messages and merges with promptContributions diff --git a/embabel-agent-api/src/main/kotlin/com/embabel/agent/test/unit/FakePromptRunner.kt b/embabel-agent-api/src/main/kotlin/com/embabel/agent/test/unit/FakePromptRunner.kt index 79539594b..1de363f5a 100644 --- a/embabel-agent-api/src/main/kotlin/com/embabel/agent/test/unit/FakePromptRunner.kt +++ b/embabel-agent-api/src/main/kotlin/com/embabel/agent/test/unit/FakePromptRunner.kt @@ -32,6 +32,7 @@ import com.embabel.agent.core.ToolGroup import com.embabel.agent.core.ToolGroupRequirement import com.embabel.agent.core.internal.LlmOperations import com.embabel.agent.core.support.LlmInteraction +import com.embabel.agent.core.support.ThinkingTagSelection import com.embabel.agent.core.support.safelyGetTools import com.embabel.agent.spi.loop.ToolInjectionStrategy import com.embabel.agent.spi.loop.ToolNotFoundPolicy @@ -268,6 +269,7 @@ data class FakePromptRunner( override fun createObjectIfPossibleWithThinking( messages: List, outputClass: Class, + thinkingTags: ThinkingTagSelection, ): ThinkingResponse { TODO("Not yet implemented") } @@ -275,11 +277,15 @@ data class FakePromptRunner( override fun createObjectWithThinking( messages: List, outputClass: Class, + thinkingTags: ThinkingTagSelection, ): ThinkingResponse { TODO("Not yet implemented") } - override fun respondWithThinking(messages: List): ThinkingResponse { + override fun respondWithThinking( + messages: List, + thinkingTags: ThinkingTagSelection, + ): ThinkingResponse { TODO("Not yet implemented") } @@ -287,6 +293,7 @@ data class FakePromptRunner( condition: String, context: String, confidenceThreshold: ZeroToOne, + thinkingTags: ThinkingTagSelection, ): ThinkingResponse { TODO("Not yet implemented") } diff --git a/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/DelegatingStreamingPromptRunnerTest.kt b/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/DelegatingStreamingPromptRunnerTest.kt index ca5913fd3..2be41b76b 100644 --- a/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/DelegatingStreamingPromptRunnerTest.kt +++ b/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/DelegatingStreamingPromptRunnerTest.kt @@ -24,9 +24,11 @@ import com.embabel.agent.api.tool.ToolObject import com.embabel.agent.core.ToolGroup import com.embabel.agent.core.ToolGroupRequirement import com.embabel.agent.experimental.primitive.Determination +import com.embabel.agent.core.support.ThinkingTagSelection import com.embabel.chat.UserMessage import com.embabel.common.ai.model.LlmOptions import com.embabel.common.ai.prompt.PromptContributor +import com.embabel.common.core.thinking.ThinkingResponse import com.embabel.common.textio.template.TemplateRenderer import com.fasterxml.jackson.databind.ObjectMapper import io.mockk.every @@ -371,6 +373,36 @@ class DelegatingStreamingPromptRunnerTest { verify { mockDelegate.supportsStreaming() } } + @Test + fun `thinking should apply tag selection`() { + val messages = listOf(UserMessage("test")) + val outputClass = String::class.java + val expectedResponse = ThinkingResponse( + result = "test", + thinkingBlocks = emptyList(), + ) + + every { + mockDelegate.createObjectIfPossibleWithThinking( + messages, + outputClass, + thinkingTags = ThinkingTagSelection( + include = setOf("reasoning"), + exclude = setOf("div"), + ), + ) + } returns expectedResponse + + val runner = createPromptRunner() + val result = runner.thinking( + include = setOf("reasoning"), + exclude = setOf("div"), + ) + + val actualResponse = result.createObjectIfPossible(messages, outputClass) + assertEquals(expectedResponse, actualResponse) + } + @Test fun `thinking should not return null`() { val runner = createPromptRunner() diff --git a/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/DelegatingThinkingTest.kt b/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/DelegatingThinkingTest.kt index d6aa866fe..8e41391e7 100644 --- a/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/DelegatingThinkingTest.kt +++ b/embabel-agent-api/src/test/kotlin/com/embabel/agent/api/common/support/DelegatingThinkingTest.kt @@ -15,6 +15,7 @@ */ package com.embabel.agent.api.common.support +import com.embabel.agent.core.support.ThinkingTagSelection import com.embabel.chat.AssistantMessage import com.embabel.chat.UserMessage import com.embabel.common.core.thinking.ThinkingBlock @@ -57,13 +58,13 @@ class DelegatingThinkingTest { ) every { - mockDelegate.createObjectIfPossibleWithThinking(messages, outputClass) + mockDelegate.createObjectIfPossibleWithThinking(messages, outputClass, any()) } returns expectedResponse val operations = createThinkingOperations() val result = operations.createObjectIfPossible(messages, outputClass) - verify { mockDelegate.createObjectIfPossibleWithThinking(messages, outputClass) } + verify { mockDelegate.createObjectIfPossibleWithThinking(messages, outputClass, any()) } assertEquals(expectedResponse, result) assertEquals("test result", result.result) assertEquals(1, result.thinkingBlocks.size) @@ -87,13 +88,13 @@ class DelegatingThinkingTest { ) every { - mockDelegate.createObjectIfPossibleWithThinking(messages, outputClass) + mockDelegate.createObjectIfPossibleWithThinking(messages, outputClass, any()) } returns expectedResponse val operations = createThinkingOperations() val result = operations.createObjectIfPossible(messages, outputClass) - verify { mockDelegate.createObjectIfPossibleWithThinking(messages, outputClass) } + verify { mockDelegate.createObjectIfPossibleWithThinking(messages, outputClass, any()) } assertEquals(expectedResponse, result) assertEquals(null, result.result) } @@ -119,13 +120,13 @@ class DelegatingThinkingTest { ) every { - mockDelegate.createObjectWithThinking(messages, outputClass) + mockDelegate.createObjectWithThinking(messages, outputClass, any()) } returns expectedResponse val operations = createThinkingOperations() val result = operations.createObject(messages, outputClass) - verify { mockDelegate.createObjectWithThinking(messages, outputClass) } + verify { mockDelegate.createObjectWithThinking(messages, outputClass, any()) } assertEquals(expectedResponse, result) assertEquals("test", result.result?.name) assertEquals(42, result.result?.value) @@ -150,12 +151,12 @@ class DelegatingThinkingTest { thinkingBlocks = thinkingBlocks ) - every { mockDelegate.respondWithThinking(messages) } returns expectedResponse + every { mockDelegate.respondWithThinking(messages, any()) } returns expectedResponse val operations = createThinkingOperations() val result = operations.respond(messages) - verify { mockDelegate.respondWithThinking(messages) } + verify { mockDelegate.respondWithThinking(messages, any()) } assertEquals(expectedResponse, result) assertEquals("4", result.result?.content) assertEquals(1, result.thinkingBlocks.size) @@ -183,13 +184,13 @@ class DelegatingThinkingTest { ) every { - mockDelegate.evaluateConditionWithThinking(condition, context, any()) + mockDelegate.evaluateConditionWithThinking(condition, context, any(), any()) } returns expectedResponse val operations = createThinkingOperations() val result = operations.evaluateCondition(condition, context, threshold) - verify { mockDelegate.evaluateConditionWithThinking(condition, context, any()) } + verify { mockDelegate.evaluateConditionWithThinking(condition, context, any(), any()) } assertEquals(expectedResponse, result) assertEquals(true, result.result) } @@ -211,16 +212,62 @@ class DelegatingThinkingTest { ) every { - mockDelegate.evaluateConditionWithThinking(condition, context, any()) + mockDelegate.evaluateConditionWithThinking(condition, context, any(), any()) } returns expectedResponse val operations = createThinkingOperations() val result = operations.evaluateCondition(condition, context) - verify { mockDelegate.evaluateConditionWithThinking(condition, context, any()) } + verify { mockDelegate.evaluateConditionWithThinking(condition, context, any(), any()) } assertEquals(expectedResponse, result) } } data class TestItem(val name: String, val value: Int) + + @Nested + inner class TagSelectionDelegationTest { + + @Test + fun `should pass include and exclude tags to delegate`() { + val messages = listOf(UserMessage("test prompt")) + val outputClass = String::class.java + val expectedResponse = ThinkingResponse( + result = "test result", + thinkingBlocks = emptyList() + ) + + every { + mockDelegate.createObjectIfPossibleWithThinking( + messages, + outputClass, + thinkingTags = ThinkingTagSelection( + include = setOf("decision_reasoning"), + exclude = setOf("div"), + ), + ) + } returns expectedResponse + + val operations = DelegatingThinking( + delegate = mockDelegate, + selection = ThinkingTagSelection( + include = setOf("decision_reasoning"), + exclude = setOf("div"), + ), + ) + val result = operations.createObjectIfPossible(messages, outputClass) + + verify { + mockDelegate.createObjectIfPossibleWithThinking( + messages, + outputClass, + thinkingTags = ThinkingTagSelection( + include = setOf("decision_reasoning"), + exclude = setOf("div"), + ), + ) + } + assertEquals(expectedResponse, result) + } + } } diff --git a/embabel-agent-api/src/test/kotlin/com/embabel/agent/core/support/LlmInteractionSerializationTest.kt b/embabel-agent-api/src/test/kotlin/com/embabel/agent/core/support/LlmInteractionSerializationTest.kt index f1c058abb..5bf288344 100644 --- a/embabel-agent-api/src/test/kotlin/com/embabel/agent/core/support/LlmInteractionSerializationTest.kt +++ b/embabel-agent-api/src/test/kotlin/com/embabel/agent/core/support/LlmInteractionSerializationTest.kt @@ -16,9 +16,20 @@ package com.embabel.agent.core.support import com.embabel.agent.api.common.InteractionId +import com.embabel.common.core.thinking.ThinkingBlock +import com.embabel.common.core.thinking.ThinkingTagType import com.fasterxml.jackson.databind.ObjectMapper import com.fasterxml.jackson.module.kotlin.registerKotlinModule +import ch.qos.logback.classic.Level +import ch.qos.logback.classic.Logger +import ch.qos.logback.classic.spi.ILoggingEvent +import ch.qos.logback.core.read.ListAppender +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.Nested import org.junit.jupiter.api.Test +import org.junit.jupiter.api.assertThrows +import org.slf4j.LoggerFactory +import kotlin.test.assertEquals /** * Test for GitHub issue #1309: Serialization issue with LlmInteraction. @@ -35,6 +46,110 @@ class LlmInteractionSerializationTest { private val objectMapper = ObjectMapper().registerKotlinModule() + @Nested + inner class ThinkingTagSelectionTest { + + private val blocks = listOf( + ThinkingBlock("selected", ThinkingTagType.TAG, "reasoning"), + ThinkingBlock("html", ThinkingTagType.TAG, "div"), + ThinkingBlock("prefix", ThinkingTagType.PREFIX, "legacy_prefix"), + ) + + @Test + fun `include retains selected tags and non-tag blocks`() { + val selection = ThinkingTagSelection(include = setOf("reasoning")) + assertEquals(listOf(blocks[0], blocks[2]), selection.filter(blocks)) + } + + @Test + fun `exclude removes matching tags`() { + val selection = ThinkingTagSelection(exclude = setOf("div")) + assertEquals(listOf(blocks[0], blocks[2]), selection.filter(blocks)) + } + + @Test + fun `overlapping include and exclude tags are rejected`() { + assertThrows { + ThinkingTagSelection(include = setOf("reasoning"), exclude = setOf("reasoning")) + } + } + + @Test + fun `invalid tag names are rejected`() { + assertThrows { + ThinkingTagSelection(include = setOf("invalid tag")) + } + assertThrows { + ThinkingTagSelection(exclude = setOf("a> tags.")) + } + + @Test + fun `missingFrom does not treat longer tag names as occurrences`() { + val selection = ThinkingTagSelection(include = setOf("think")) + assertEquals(setOf("think"), selection.missingFrom("Reason inside tags.")) + } + + @Test + fun `missingFrom treats self-closing tags as present`() { + val selection = ThinkingTagSelection(include = setOf("reasoning")) + assertEquals(emptySet(), selection.missingFrom("Reason inside tags.")) + } + + @Test + fun `missingFrom is empty when include is empty`() { + assertEquals(emptySet(), ThinkingTagSelection().missingFrom("anything")) + } + + @Test + fun `exclude does not affect missingFrom`() { + val selection = ThinkingTagSelection(include = setOf("reasoning"), exclude = setOf("div")) + assertEquals(emptySet(), selection.missingFrom("Reason inside tags.")) + } + + @Test + fun `warnIfMissing logs a warning when declared tags are absent`() { + val appender = ListAppender().apply { start() } + val logger = LoggerFactory.getLogger(ThinkingTagSelection::class.java) as Logger + logger.addAppender(appender) + try { + ThinkingTagSelection(include = setOf("missing_tag")).warnIfMissing("no tags here", logger) + assertTrue( + appender.list.any { it.level == Level.WARN && it.formattedMessage.contains("missing_tag") }, + ) + } finally { + logger.detachAppender(appender) + } + } + + @Test + fun `warnIfMissing does not log when all declared tags are present`() { + val appender = ListAppender().apply { start() } + val logger = LoggerFactory.getLogger(ThinkingTagSelection::class.java) as Logger + logger.addAppender(appender) + try { + ThinkingTagSelection(include = setOf("reasoning")).warnIfMissing("Reason inside ", logger) + assertTrue( + appender.list.isEmpty(), + ) + } finally { + logger.detachAppender(appender) + } + } + } @Test fun `LlmInteraction can be serialized to JSON without conflicting getter error`() { val interaction = LlmInteraction( diff --git a/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/support/ChatClientLlmOperationsThinkingTest.kt b/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/support/ChatClientLlmOperationsThinkingTest.kt index 4f65368e3..d9ca6a01a 100644 --- a/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/support/ChatClientLlmOperationsThinkingTest.kt +++ b/embabel-agent-api/src/test/kotlin/com/embabel/agent/spi/support/ChatClientLlmOperationsThinkingTest.kt @@ -29,6 +29,7 @@ import com.embabel.agent.core.ProcessContext import com.embabel.agent.core.support.InvalidLlmReturnFormatException import com.embabel.agent.core.support.LlmCall import com.embabel.agent.core.support.LlmInteraction +import com.embabel.agent.core.support.ThinkingTagSelection import com.embabel.agent.spi.support.springai.ChatClientLlmOperations import com.embabel.agent.spi.support.springai.SpringAiLlmService import com.embabel.agent.spi.validation.DefaultValidationPromptGenerator @@ -311,6 +312,40 @@ class ChatClientLlmOperationsThinkingTest { assertTrue(result.thinkingBlocks[0].content.contains("process this request carefully")) } + @Test + fun `doTransformWithThinking should apply selected thinking tags`() { + // Given: LLM response with a selected tag, an unselected tag, and valid JSON + val rawLlmResponse = """ + Keep this block. +
Discard this markup.
+ + { + "status": "selected", + "value": 101 + } + """.trimIndent() + + val setup = createChatClientLlmOperations(FakeChatModel(rawLlmResponse)) + val interaction = LlmInteraction( + id = InteractionId("selected-thinking"), + thinkingTags = ThinkingTagSelection(include = setOf("reasoning")), + ) + + // When: Use doTransformWithThinking with tag selection + val result = setup.llmOperations.doTransformWithThinkingSpringAi( + messages = listOf(UserMessage("Select reasoning")), + interaction = interaction, + outputClass = SimpleResult::class.java, + llmRequestEvent = null, + agentProcess = null, + action = null, + ) + + // Then: Should extract only the selected thinking tags + assertEquals("selected", result.result!!.status) + assertEquals(listOf("reasoning"), result.thinkingBlocks.map { it.tagValue }) + } + @Test fun `ChatResponseWithThinkingException should preserve message and thinking blocks`() { // Test the actual constructor and properties (new code) diff --git a/embabel-agent-docs/src/main/asciidoc/reference/thinking/page.adoc b/embabel-agent-docs/src/main/asciidoc/reference/thinking/page.adoc index c8264655d..d0bc588c3 100644 --- a/embabel-agent-docs/src/main/asciidoc/reference/thinking/page.adoc +++ b/embabel-agent-docs/src/main/asciidoc/reference/thinking/page.adoc @@ -21,7 +21,7 @@ An even more important use case arises when the LLM cannot fulfill a request—f - `ThinkingTagType` — An enum defining the types of reasoning markers: `TAG` (XML-style tags like ``), `PREFIX` (line prefixes like `//THINKING:`), and `NO_PREFIX` (untagged reasoning text before JSON output). - `ThinkingResponse` — A response wrapper that holds both the result object and a list of `ThinkingBlock` instances. - `ThinkingException` — An exception that preserves thinking blocks when object instantiation fails, enabling debugging even in error scenarios. -- `thinking()` — The core `PromptRunner` API method that enables thinking extraction. +- `thinking()` — The core `PromptRunner` API method that enables thinking extraction. Use `thinking(tag)` to extract reasoning blocks from a single custom tag, or `thinking(include, exclude)` to select or suppress XML-style thinking tags without changing the prompt sent to the LLM. ==== Example: Handling Objects and Thinking Blocks @@ -93,6 +93,62 @@ for (block in thinkingBlocks) { ---- ==== +==== Selecting Thinking Tags + +XML-style tags are discovered dynamically, so ordinary markup such as `
...
` can otherwise appear as a thinking block. Select the tags relevant to an interaction: + +[tabs] +==== +Java:: ++ +[source,java] +---- +ThinkingResponse response = runner + .thinking(Set.of("reasoning"), Set.of("div")) + .createObject(prompt, MonthItem.class); +---- + +Kotlin:: ++ +[source,kotlin] +---- +val response = runner + .thinking(include = setOf("reasoning"), exclude = setOf("div")) + .createObject(prompt, MonthItem::class.java) +---- +==== + +When `include` is non-empty, only matching XML-style `TAG` blocks are retained. Tags in `exclude` are removed. Prefix-based and untagged reasoning blocks are unaffected. A tag cannot appear in both sets. Calling `thinking()` without arguments preserves the default extraction behavior. + +When tag selection is used, the framework checks whether each declared tag appears as an XML open tag (` response = runner + .thinking("decision_reasoning") + .createObject(prompt, MonthItem.class); +---- + +Kotlin:: ++ +[source,kotlin] +---- +val response = runner + .thinking("decision_reasoning") + .createObject(prompt, MonthItem::class.java) +---- +==== + +The tag must be a valid XML name matching `[a-zA-Z][a-zA-Z0-9_-]*`. Implementations that support tag selection reject invalid tags with an `IllegalArgumentException`. + +The tag selection methods do not modify the prompt — `thinking()`, `thinking(tag)`, and `thinking(include, exclude)` all leave the prompt unchanged. The LLM will only produce reasoning inside `...` if your prompt asks for it. For example, include in your system prompt: "You MUST provide reasoning inside `...` tags before providing your final answer." + ==== Example: Handling Failures Gracefully Use `createObjectIfPossible` when the LLM might not be able to produce a valid result: @@ -143,6 +199,7 @@ if (result != null) { Embabel exposes thinking through a provider-neutral API: - `PromptRunner.thinking()` +- `PromptRunner.thinking(tag)` - `LlmOptions.thinking` This remains the primary public API for enabling reasoning/thinking mode.