feat(thinking): add tag selection API with include/exclude filtering (#1790)

embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/PromptRunner.kt

@@ -487,6 +487,18 @@ interface PromptRunner : LlmUse, PromptRunnerOperations, ToolChaining<PromptRunn
         error("Implementation error: supportsThinking() returned true but withThinking() not overridden")
     }
 
+    fun thinking(include: Set<String>): Thinking = thinking(include, emptySet())
+
+    // Non-empty include/exclude throws UnsupportedOperationException by default — implementations must override.
+    fun thinking(include: Set<String>, exclude: Set<String>): 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<Message>,
     ): AssistantMessage =

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<String>, exclude: Set<String>): PromptRunner.Thinking =
+        DelegatingThinking(
+            delegate = delegate,
+            selection = ThinkingTagSelection(include = include, exclude = exclude),
+        )
+
+    override fun thinking(tag: String): PromptRunner.Thinking =
+        thinking(include = setOf(tag))
 }

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 <T> createObjectIfPossible(
         messages: List<Message>,
         outputClass: Class<T>
     ): ThinkingResponse<T?> =
-        delegate.createObjectIfPossibleWithThinking(messages, outputClass)
+        delegate.createObjectIfPossibleWithThinking(messages, outputClass, selection)
 
     override fun <T> createObject(
         messages: List<Message>,
         outputClass: Class<T>
     ): ThinkingResponse<T> =
-        delegate.createObjectWithThinking(messages, outputClass)
+        delegate.createObjectWithThinking(messages, outputClass, selection)
 
     override fun respond(messages: List<Message>): ThinkingResponse<AssistantMessage> =
-        delegate.respondWithThinking(messages)
+        delegate.respondWithThinking(messages, selection)
 
     override fun evaluateCondition(
         condition: String,
         context: String,
         confidenceThreshold: ZeroToOne
     ): ThinkingResponse<Boolean> =
-        delegate.evaluateConditionWithThinking(condition, context, confidenceThreshold)
+        delegate.evaluateConditionWithThinking(condition, context, confidenceThreshold, selection)
 }

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 <T> createObjectWithThinking(
         messages: List<Message>,
-        outputClass: Class<T>
+        outputClass: Class<T>,
+        thinkingTags: ThinkingTagSelection,
     ): ThinkingResponse<T> {
         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 <T> createObjectIfPossibleWithThinking(
         messages: List<Message>,
-        outputClass: Class<T>
+        outputClass: Class<T>,
+        thinkingTags: ThinkingTagSelection,
     ): ThinkingResponse<T?> {
         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<Message>): ThinkingResponse<AssistantMessage> {
-        return createObjectWithThinking(messages, AssistantMessage::class.java)
+    override fun respondWithThinking(
+        messages: List<Message>,
+        thinkingTags: ThinkingTagSelection,
+    ): ThinkingResponse<AssistantMessage> {
+        return createObjectWithThinking(messages, AssistantMessage::class.java, thinkingTags)
     }
 
     override fun evaluateConditionWithThinking(
         condition: String,
         context: String,
-        confidenceThreshold: ZeroToOne
+        confidenceThreshold: ZeroToOne,
+        thinkingTags: ThinkingTagSelection,
     ): ThinkingResponse<Boolean> {
         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<ToolGroupRequirement> = 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,
         )
     }
 

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 <T> createObjectIfPossibleWithThinking(
         messages: List<Message>,
         outputClass: Class<T>,
+        thinkingTags: ThinkingTagSelection = ThinkingTagSelection(),
     ): ThinkingResponse<T?>
 
     fun <T> createObjectWithThinking(
         messages: List<Message>,
-        outputClass: Class<T>
+        outputClass: Class<T>,
+        thinkingTags: ThinkingTagSelection = ThinkingTagSelection(),
     ): ThinkingResponse<T>
 
-    fun respondWithThinking(messages: List<Message>): ThinkingResponse<AssistantMessage>
+    fun respondWithThinking(
+        messages: List<Message>,
+        thinkingTags: ThinkingTagSelection = ThinkingTagSelection(),
+    ): ThinkingResponse<AssistantMessage>
 
     fun evaluateConditionWithThinking(
         condition: String,
         context: String,
-        confidenceThreshold: ZeroToOne
+        confidenceThreshold: ZeroToOne,
+        thinkingTags: ThinkingTagSelection = ThinkingTagSelection(),
     ): ThinkingResponse<Boolean>
 
 }

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<String> = emptySet(),
+    val exclude: Set<String> = 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<ThinkingBlock>): List<ThinkingBlock> =
+        blocks.filter { block ->
+            block.tagType != ThinkingTagType.TAG ||
+                    ((include.isEmpty() || block.tagValue in include) && block.tagValue !in exclude)
+        }
+
+    fun missingFrom(systemPrompt: String): Set<String> =
+        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 "<thinking>".
+        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<ToolCallInspector> = emptyList(),
     val toolCallContext: ToolCallContext = ToolCallContext.EMPTY,
     val toolNotFoundPolicy: ToolNotFoundPolicy? = null,
+    val thinkingTags: ThinkingTagSelection = ThinkingTagSelection(),
 ) : LlmCall {
 
     override val name: String = id.value

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<O> = { 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<ThinkingResponse<O>> = { 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<Message>): String =
+        messages.filterIsInstance<SystemMessage>().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<Message>): List<ThinkingBlock> {
+    private fun accumulateThinkingBlocks(
+        thinkingTags: ThinkingTagSelection,
+        conversationHistory: List<Message>,
+    ): List<ThinkingBlock> {
         return conversationHistory
             .filter { it.role == com.embabel.chat.Role.ASSISTANT }
-            .flatMap { extractAllThinkingBlocks(it.content) }
+            .flatMap { thinkingTags.filter(extractAllThinkingBlocks(it.content)) }
     }
 
     /**