Fix reasoning messages in nodes

- Add `onReasoningMessage` edge transformation for reasoning message filtering
- Reorder `onMultipleToolCalls` edge to prioritize execution flow
- Expand `EdgeTransformationDslMarker` target scope to include classes and annotate `AIAgentEdgeBuilderIntermediate`
- Add `skipReasoningMessage` support and `requestLLMMultipleWithoutTools` implementation
- Document `skipReasoningMessage` behavior in `nodeLLMRequest` API function

Author: devcrocod

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/AIAgentSimpleStrategies.kt

@@ -52,13 +52,13 @@ private fun singleRunWithParallelAbility(parallelTools: Boolean) = strategy("sin
 
     edge(nodeExecuteTool forwardTo nodeSendToolResult)
 
+    edge(nodeSendToolResult forwardTo nodeExecuteTool onMultipleToolCalls { true })
+
     edge(
         nodeSendToolResult forwardTo nodeFinish
             onMultipleAssistantMessages { true }
             transformed { it.joinToString("\n") { message -> message.content } }
     )
-
-    edge(nodeSendToolResult forwardTo nodeExecuteTool onMultipleToolCalls { true })
 }
 
 private fun singleRunModeStrategy() = strategy("single_run") {

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/session/AIAgentLLMSession.kt

@@ -122,6 +122,21 @@ public sealed class AIAgentLLMSession(
     protected suspend fun executeSingle(prompt: Prompt, tools: List<ToolDescriptor>): Message.Response =
         executeMultiple(prompt, tools).first()
 
+    /**
+     * Sends a request to the language model without utilizing any tools and returns multiple responses.
+     *
+     * @return A list of response messages from the language model.
+     */
+    public open suspend fun requestLLMMultipleWithoutTools(): List<Message.Response> {
+        validateSession()
+
+        val promptWithDisabledTools = prompt
+            .withUpdatedParams { toolChoice = null }
+            .let { preparePrompt(it, emptyList()) }
+
+        return executeMultiple(promptWithDisabledTools, emptyList())
+    }
+
     /**
      * Sends a request to the language model without utilizing any tools and returns the response.
      *
@@ -130,10 +145,14 @@ public sealed class AIAgentLLMSession(
      * to ensure compatibility with the underlying LLM client's behavior. It then executes the request
      * and retrieves the response from the LLM.
      *
+     * When `skipReasoningMessage` is true, filters out reasoning messages and returns the first
+     * non-reasoning response.
+     *
+     * @param skipReasoningMessage If true, skips reasoning messages and returns first non-reasoning response (default: true).
      * @return The response message from the language model after executing the request, represented
      *         as a [Message.Response] instance.
      */
-    public open suspend fun requestLLMWithoutTools(): Message.Response {
+    public open suspend fun requestLLMWithoutTools(skipReasoningMessage: Boolean = true): Message.Response {
         validateSession()
         /*
             Not all LLM providers support tool list when tool choice is set to "none", so we are rewriting all tool messages to regular messages,
@@ -143,7 +162,11 @@ public sealed class AIAgentLLMSession(
             .withUpdatedParams { toolChoice = null }
             .let { preparePrompt(it, emptyList()) }
 
-        return executeSingle(promptWithDisabledTools, emptyList())
+        return if (skipReasoningMessage) {
+            executeMultiple(promptWithDisabledTools, emptyList()).first { it !is Message.Reasoning }
+        } else {
+            executeSingle(promptWithDisabledTools, emptyList())
+        }
     }
 
     /**
@@ -205,11 +228,19 @@ public sealed class AIAgentLLMSession(
      * Sends a request to the underlying LLM and returns the first response.
      * This method ensures the session is active before executing the request.
      *
+     * When `skipReasoningMessage` is true, filters out reasoning messages and returns the first
+     * non-reasoning response.
+     *
+     * @param skipReasoningMessage If true, skips reasoning messages and returns first non-reasoning response (default: true).
      * @return The first response message from the LLM after executing the request.
      */
-    public open suspend fun requestLLM(): Message.Response {
+    public open suspend fun requestLLM(skipReasoningMessage: Boolean = true): Message.Response {
         validateSession()
-        return executeSingle(prompt, tools)
+        return if (skipReasoningMessage) {
+            executeMultiple(prompt, tools).first { it !is Message.Reasoning }
+        } else {
+            executeSingle(prompt, tools)
+        }
     }
 
     /**

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/session/AIAgentLLMWriteSession.kt

@@ -365,16 +365,40 @@ public class AIAgentLLMWriteSession internal constructor(
         prompt.withParams(newParams)
     }
 
+    /**
+     * Sends a request to the language model without utilizing any tools, returns multiple responses,
+     * and updates the prompt with the received messages.
+     *
+     * @return A list of response messages from the language model.
+     */
+    override suspend fun requestLLMMultipleWithoutTools(): List<Message.Response> {
+        return super.requestLLMMultipleWithoutTools().also { responses ->
+            appendPrompt {
+                responses.forEach { message(it) }
+            }
+        }
+    }
+
     /**
      * Sends a request to the Language Model (LLM) without including any tools, processes the response,
      * and updates the prompt with the returned message.
      *
      * LLM might answer only with a textual assistant message.
      *
+     * When `skipReasoningMessage` is true, all responses (including reasoning) are added to the prompt,
+     * but only the first non-reasoning response is returned.
+     *
+     * @param skipReasoningMessage If true, filters out reasoning messages from the return value (default: true).
      * @return the response from the LLM after processing the request, as a [Message.Response].
      */
-    override suspend fun requestLLMWithoutTools(): Message.Response {
-        return super.requestLLMWithoutTools().also { response -> appendPrompt { message(response) } }
+    override suspend fun requestLLMWithoutTools(skipReasoningMessage: Boolean): Message.Response {
+        return if (skipReasoningMessage) {
+            requestLLMMultipleWithoutTools().first { it !is Message.Reasoning }
+        } else {
+            super.requestLLMWithoutTools(skipReasoningMessage).also { response ->
+                appendPrompt { message(response) }
+            }
+        }
     }
 
     /**
@@ -412,11 +436,19 @@ public class AIAgentLLMWriteSession internal constructor(
      * Makes an asynchronous request to a Large Language Model (LLM) and updates the current prompt
      * with the response received from the LLM.
      *
+     * When `skipReasoningMessage` is true, all responses (including reasoning) are added to the prompt,
+     * but only the first non-reasoning response is returned.
+     *
+     * @param skipReasoningMessage If true, filters out reasoning messages from the return value (default: true).
      * @return A [Message.Response] object containing the response from the LLM.
      */
-    override suspend fun requestLLM(): Message.Response {
-        return super.requestLLM().also { response ->
-            appendPrompt { message(response) }
+    override suspend fun requestLLM(skipReasoningMessage: Boolean): Message.Response {
+        return if (skipReasoningMessage) {
+            requestLLMMultiple().first { it !is Message.Reasoning }
+        } else {
+            super.requestLLM(skipReasoningMessage).also { response ->
+                appendPrompt { message(response) }
+            }
         }
     }
 

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/builder/AIAgentEdgeBuilder.kt

@@ -10,7 +10,7 @@ import ai.koog.agents.core.utils.Option
  * to ensure its proper highlighting.
  */
 @DslMarker
-@Target(AnnotationTarget.FUNCTION)
+@Target(AnnotationTarget.CLASS, AnnotationTarget.FUNCTION)
 public annotation class EdgeTransformationDslMarker
 
 /**
@@ -54,6 +54,7 @@ public class AIAgentEdgeBuilder<IncomingOutput, OutgoingInput, CompatibleOutput
  * the originating node's output into an intermediate representation
  * or filtering the flow based on specific conditions.
  */
+@EdgeTransformationDslMarker
 public class AIAgentEdgeBuilderIntermediate<IncomingOutput, IntermediateOutput, OutgoingInput> internal constructor(
     internal val fromNode: AIAgentNodeBase<*, IncomingOutput>,
     internal val toNode: AIAgentNodeBase<OutgoingInput, *>,

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/extension/AIAgentEdges.kt

@@ -195,6 +195,19 @@ public infix fun <IncomingOutput, IntermediateOutput, OutgoingInput> AIAgentEdge
         .transformed { it.content }
 }
 
+/**
+ * Creates an edge that filters a reasoning message based on a custom condition
+ *
+ * @param block A function that evaluates whether to accept a reasoning message
+ */
+@EdgeTransformationDslMarker
+public infix fun <IncomingOutput, IntermediateOutput, OutgoingInput> AIAgentEdgeBuilderIntermediate<IncomingOutput, IntermediateOutput, OutgoingInput>.onReasoningMessage(
+    block: suspend (Message.Reasoning) -> Boolean
+): AIAgentEdgeBuilderIntermediate<IncomingOutput, Message.Reasoning, OutgoingInput> {
+    return onIsInstance(Message.Reasoning::class)
+        .onCondition { signature -> block(signature) }
+}
+
 /**
  * Creates an edge that filters assistant messages based on a custom condition and extracts their content.
  *
@@ -206,7 +219,21 @@ public infix fun <IncomingOutput, OutgoingInput> AIAgentEdgeBuilderIntermediate<
 ): AIAgentEdgeBuilderIntermediate<IncomingOutput, List<Message.Assistant>, OutgoingInput> {
     return onIsInstance(List::class)
         .transformed { it.filterIsInstance<Message.Assistant>() }
-        .onCondition { toolResults -> block(toolResults) }
+        .onCondition { messages -> block(messages) }
+}
+
+/**
+ * Creates an edge that filters lists of reasoning messages based on a custom condition.
+ *
+ * @param block A function that evaluates whether to accept a list of reasoning messages
+ */
+public infix fun <IncomingOutput, OutgoingInput> AIAgentEdgeBuilderIntermediate<IncomingOutput, List<Message.Response>, OutgoingInput>.onMultipleReasoningMessages(
+    block: suspend (List<Message.Reasoning>) -> Boolean
+): AIAgentEdgeBuilderIntermediate<IncomingOutput, List<Message.Reasoning>, OutgoingInput> {
+    return onIsInstance(List::class)
+        .transformed { it.filterIsInstance<Message.Reasoning>() }
+        .onCondition { it.isNotEmpty() }
+        .onCondition { messages -> block(messages) }
 }
 
 /**

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/extension/AIAgentNodes.kt

@@ -132,13 +132,18 @@ public fun AIAgentSubgraphBuilderBase<*, *>.nodeLLMSendMessageForceOneTool(
 /**
  * A node that appends a user message to the LLM prompt and gets a response with optional tool usage.
  *
+ * When `skipReasoningMessage` is enabled, the node requests multiple responses and returns the first
+ * non-reasoning message. Otherwise, it returns a single response that may include reasoning.
+ *
  * @param name Optional node name.
  * @param allowToolCalls Controls whether LLM can use tools (default: true).
+ * @param skipReasoningMessage If true, filters out reasoning messages and returns the first non-reasoning response (default: true).
  */
 @AIAgentBuilderDslMarker
 public fun AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequest(
     name: String? = null,
-    allowToolCalls: Boolean = true
+    allowToolCalls: Boolean = true,
+    skipReasoningMessage: Boolean = true,
 ): AIAgentNodeDelegate<String, Message.Response> =
     node(name) { message ->
         llm.writeSession {
@@ -147,9 +152,9 @@ public fun AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequest(
             }
 
             if (allowToolCalls) {
-                requestLLM()
+                requestLLM(skipReasoningMessage)
             } else {
-                requestLLMWithoutTools()
+                requestLLMWithoutTools(skipReasoningMessage)
             }
         }
     }

prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/PromptExecutorExtensions.kt

@@ -170,7 +170,7 @@ public suspend fun <T> PromptExecutor.executeStructured(
     config: StructuredRequestConfig<T>,
 ): Result<StructuredResponse<T>> {
     val updatedPrompt = config.updatePrompt(model, prompt)
-    val response = this.execute(prompt = updatedPrompt, model = model).single()
+    val response = this.execute(prompt = updatedPrompt, model = model).filterNot { it is Message.Reasoning }.single()
 
     return runCatching {
         require(response is Message.Assistant) { "Response for structured output must be an assistant message, got ${response::class.simpleName} instead" }