Move structured output from executor to client

Author: tiginamaria

agents/agents-core/build.gradle.kts

@@ -17,7 +17,6 @@ kotlin {
                 api(project(":utils"))
                 api(project(":prompt:prompt-executor:prompt-executor-model"))
                 api(project(":prompt:prompt-llm"))
-                api(project(":prompt:prompt-structure"))
 
                 api(project(":prompt:prompt-executor:prompt-executor-clients:prompt-executor-openai-client"))
                 api(project(":prompt:prompt-markdown"))

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

@@ -14,9 +14,8 @@ import ai.koog.agents.core.tools.ToolDescriptor
 import ai.koog.agents.core.tools.annotations.LLMDescription
 import ai.koog.prompt.llm.LLModel
 import ai.koog.prompt.params.LLMParams
-import ai.koog.prompt.structure.StructureFixingParser
 import ai.koog.prompt.structure.StructuredOutput
-import ai.koog.prompt.structure.StructuredOutputConfig
+import ai.koog.prompt.structure.StructuredOutputFixingConfig
 import ai.koog.prompt.structure.json.JsonStructuredData
 import ai.koog.prompt.structure.json.generator.StandardJsonSchemaGenerator
 import io.github.oshai.kotlinlogging.KotlinLogging
@@ -124,15 +123,13 @@ public open class AIAgentSubgraph<TInput, TOutput>(
             }
 
             val selectedTools = this.requestLLMStructured(
-                config = StructuredOutputConfig(
-                    default = StructuredOutput.Manual(
-                        JsonStructuredData.createJsonStructure<SelectedTools>(
-                            schemaGenerator = StandardJsonSchemaGenerator,
-                            examples = listOf(SelectedTools(listOf()), SelectedTools(tools.map { it.name }.take(3))),
-                        ),
+                structuredOutput = StructuredOutput.Manual(
+                    JsonStructuredData.createJsonStructure<SelectedTools>(
+                        schemaGenerator = StandardJsonSchemaGenerator,
+                        examples = listOf(SelectedTools(listOf()), SelectedTools(tools.map { it.name }.take(3))),
                     ),
-                    fixingParser = toolSelectionStrategy.fixingParser,
-                )
+                ),
+                fixingConfig = toolSelectionStrategy.fixingConfig,
             ).getOrThrow()
 
             prompt = initialPrompt
@@ -294,11 +291,10 @@ public sealed interface ToolSelectionStrategy {
      * This ensures that unnecessary tools are excluded, optimizing the toolset for the specific use case.
      *
      * @property subtaskDescription A description of the subtask for which the relevant tools should be selected.
-     * @property fixingParser Optional [StructureFixingParser] to attempt fixes when malformed structured response with tool list is received.
      */
     public data class AutoSelectForTask(
         val subtaskDescription: String,
-        val fixingParser: StructureFixingParser? = null
+        val fixingConfig: StructuredOutputFixingConfig? = null
     ) : ToolSelectionStrategy
 
     /**

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

@@ -6,17 +6,17 @@ import ai.koog.agents.core.tools.ToolDescriptor
 import ai.koog.agents.core.utils.ActiveProperty
 import ai.koog.prompt.dsl.ModerationResult
 import ai.koog.prompt.dsl.Prompt
-import ai.koog.prompt.executor.model.LLMChoice
 import ai.koog.prompt.executor.model.PromptExecutor
 import ai.koog.prompt.llm.LLModel
+import ai.koog.prompt.message.LLMChoice
 import ai.koog.prompt.message.Message
 import ai.koog.prompt.params.LLMParams
 import ai.koog.prompt.streaming.StreamFrame
-import ai.koog.prompt.structure.StructureFixingParser
+import ai.koog.prompt.structure.StructuredOutput
 import ai.koog.prompt.structure.StructuredOutputConfig
+import ai.koog.prompt.structure.StructuredOutputFixingConfig
 import ai.koog.prompt.structure.StructuredResponse
-import ai.koog.prompt.structure.executeStructured
-import ai.koog.prompt.structure.parseResponseToStructuredResponse
+import ai.koog.prompt.structure.annotations.InternalStructuredOutputApi
 import kotlinx.coroutines.flow.Flow
 import kotlinx.serialization.KSerializer
 import kotlinx.serialization.serializer
@@ -260,10 +260,11 @@ public sealed class AIAgentLLMSession(
      *
      * @param config A configuration defining structures and behavior.
      *
-     * @see [executeStructured]
+     * @see [PromptExecutor.executeStructured]
      */
     public open suspend fun <T> requestLLMStructured(
-        config: StructuredOutputConfig<T>,
+        structuredOutput: StructuredOutput<T>,
+        fixingConfig: StructuredOutputFixingConfig? = null,
     ): Result<StructuredResponse<T>> {
         validateSession()
 
@@ -272,10 +273,39 @@ public sealed class AIAgentLLMSession(
         return executor.executeStructured(
             prompt = preparedPrompt,
             model = model,
-            config = config,
+            structuredOutput = structuredOutput,
+            fixingConfig = fixingConfig,
         )
     }
 
+    public open suspend fun <T> requestLLMStructured(
+        config: StructuredOutputConfig<T>
+    ): Result<StructuredResponse<T>> = requestLLMStructured(
+        structuredOutput = config.structuredOutput(model),
+        fixingConfig = config.fixingConfig,
+    )
+
+    /**
+     * Sends a request to LLM and gets a structured response.
+     *
+     * This is a simple version of the full `requestLLMStructured`. Unlike the full version, it does not require specifying
+     * struct definitions and structured output modes manually. It attempts to find the best approach to provide a structured
+     * output based on the defined [model] capabilities.
+     *
+     * @param T The structure to request.
+     * @param examples Optional list of examples in case manual mode will be used. These examples might help the model to
+     * understand the format better.
+     * @param fixingConfig Optional configuration for fixing the structured response.
+     */
+    public suspend inline fun <reified T> requestLLMStructured(
+        examples: List<T> = emptyList(),
+        fixingConfig: StructuredOutputFixingConfig? = null
+    ): Result<StructuredResponse<T>> = requestLLMStructured(
+        serializer = serializer<T>(),
+        examples = examples,
+        fixingConfig = fixingConfig,
+    )
+
     /**
      * Sends a request to LLM and gets a structured response.
      *
@@ -286,14 +316,13 @@ public sealed class AIAgentLLMSession(
      * @param serializer Serializer for the requested structure type.
      * @param examples Optional list of examples in case manual mode will be used. These examples might help the model to
      * understand the format better.
-     * @param fixingParser Optional parser that handles malformed responses by using an auxiliary LLM to
-     * intelligently fix parsing errors. When specified, parsing errors trigger additional
-     * LLM calls with error context to attempt correction of the structure format.
+     * @param fixingConfig Optional configuration for fixing the structured response.
      */
+    @OptIn(InternalStructuredOutputApi::class)
     public open suspend fun <T> requestLLMStructured(
         serializer: KSerializer<T>,
         examples: List<T> = emptyList(),
-        fixingParser: StructureFixingParser? = null
+        fixingConfig: StructuredOutputFixingConfig? = null,
     ): Result<StructuredResponse<T>> {
         validateSession()
 
@@ -304,51 +333,10 @@ public sealed class AIAgentLLMSession(
             model = model,
             serializer = serializer,
             examples = examples,
-            fixingParser = fixingParser,
+            fixingConfig = fixingConfig,
         )
     }
 
-    /**
-     * Sends a request to LLM and gets a structured response.
-     *
-     * This is a simple version of the full `requestLLMStructured`. Unlike the full version, it does not require specifying
-     * struct definitions and structured output modes manually. It attempts to find the best approach to provide a structured
-     * output based on the defined [model] capabilities.
-     *
-     * @param T The structure to request.
-     * @param examples Optional list of examples in case manual mode will be used. These examples might help the model to
-     * understand the format better.
-     * @param fixingParser Optional parser that handles malformed responses by using an auxiliary LLM to
-     * intelligently fix parsing errors. When specified, parsing errors trigger additional
-     * LLM calls with error context to attempt correction of the structure format.
-     */
-    public suspend inline fun <reified T> requestLLMStructured(
-        examples: List<T> = emptyList(),
-        fixingParser: StructureFixingParser? = null
-    ): Result<StructuredResponse<T>> = requestLLMStructured(
-        serializer = serializer<T>(),
-        examples = examples,
-        fixingParser = fixingParser,
-    )
-
-    /**
-     * Parses a structured response from the language model using the specified configuration.
-     *
-     * This function takes a response message and a structured output configuration,
-     * parses the response content based on the defined structure, and returns
-     * a structured response containing the parsed data and the original message.
-     *
-     * @param response The response message from the language model that contains the content to be parsed.
-     * The message is expected to match the defined structured output.
-     * @param config The configuration defining the expected structure and additional parsing behavior.
-     * It includes options such as structure definitions and optional parsers for error handling.
-     * @return A structured response containing the parsed data of type `T` along with the original message.
-     */
-    public suspend fun <T> parseResponseToStructuredResponse(
-        response: Message.Assistant,
-        config: StructuredOutputConfig<T>
-    ): StructuredResponse<T> = executor.parseResponseToStructuredResponse(response, config, model)
-
     /**
      * Sends a request to the language model, potentially receiving multiple choices,
      * and returns a list of choices from the model.

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

@@ -15,9 +15,9 @@ import ai.koog.prompt.llm.LLModel
 import ai.koog.prompt.message.Message
 import ai.koog.prompt.params.LLMParams
 import ai.koog.prompt.streaming.StreamFrame
-import ai.koog.prompt.structure.StructureFixingParser
 import ai.koog.prompt.structure.StructuredDataDefinition
 import ai.koog.prompt.structure.StructuredOutputConfig
+import ai.koog.prompt.structure.StructuredOutputFixingConfig
 import ai.koog.prompt.structure.StructuredResponse
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.flatMapMerge
@@ -442,7 +442,7 @@ public class AIAgentLLMWriteSession internal constructor(
      *
      * @param config A configuration defining structures and behavior.
      *
-     * @see [executeStructured]
+     * @see [PromptExecutor.executeStructured]
      */
     override suspend fun <T> requestLLMStructured(
         config: StructuredOutputConfig<T>,
@@ -466,16 +466,14 @@ public class AIAgentLLMWriteSession internal constructor(
      * @param serializer Serializer for the requested structure type.
      * @param examples Optional list of examples in case manual mode will be used. These examples might help the model to
      * understand the format better.
-     * @param fixingParser Optional parser that handles malformed responses by using an auxiliary LLM to
-     * intelligently fix parsing errors. When specified, parsing errors trigger additional
-     * LLM calls with error context to attempt correction of the structure format.
+     * @param fixingConfig Optional configuration for fixing the structured response.
      */
     override suspend fun <T> requestLLMStructured(
         serializer: KSerializer<T>,
         examples: List<T>,
-        fixingParser: StructureFixingParser?
+        fixingConfig: StructuredOutputFixingConfig?
     ): Result<StructuredResponse<T>> {
-        return super.requestLLMStructured(serializer, examples, fixingParser).also {
+        return super.requestLLMStructured(serializer, examples, fixingConfig).also {
             it.onSuccess { response ->
                 appendPrompt {
                     message(response.message)

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

@@ -11,8 +11,8 @@ import ai.koog.agents.core.tools.ToolDescriptor
 import ai.koog.agents.core.tools.ToolResult
 import ai.koog.prompt.message.Message
 import ai.koog.prompt.streaming.StreamFrame
-import ai.koog.prompt.structure.StructureFixingParser
 import ai.koog.prompt.structure.StructuredDataDefinition
+import ai.koog.prompt.structure.StructuredOutputFixingConfig
 import ai.koog.prompt.structure.StructuredResponse
 import kotlinx.coroutines.flow.Flow
 import kotlinx.serialization.serializer
@@ -150,7 +150,7 @@ public suspend fun AIAgentFunctionalContext.latestTokenUsage(): Int {
 public suspend inline fun <reified T> AIAgentFunctionalContext.requestLLMStructured(
     message: String,
     examples: List<T> = emptyList(),
-    fixingParser: StructureFixingParser? = null
+    fixingConfig: StructuredOutputFixingConfig? = null,
 ): Result<StructuredResponse<T>> {
     return llm.writeSession {
         appendPrompt {
@@ -160,7 +160,7 @@ public suspend inline fun <reified T> AIAgentFunctionalContext.requestLLMStructu
         requestLLMStructured(
             serializer<T>(),
             examples,
-            fixingParser
+            fixingConfig
         )
     }
 }

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

@@ -17,9 +17,10 @@ import ai.koog.prompt.llm.LLModel
 import ai.koog.prompt.message.Message
 import ai.koog.prompt.streaming.StreamFrame
 import ai.koog.prompt.streaming.toMessageResponses
-import ai.koog.prompt.structure.StructureFixingParser
 import ai.koog.prompt.structure.StructuredDataDefinition
 import ai.koog.prompt.structure.StructuredOutputConfig
+import ai.koog.prompt.structure.StructuredOutputFixingConfig
+import ai.koog.prompt.structure.StructuredOutputPrompts.getStructuredOutputPrompt
 import ai.koog.prompt.structure.StructuredResponse
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.toList
@@ -229,7 +230,7 @@ public inline fun <reified T> AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestStr
 public inline fun <reified T> AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestStructured(
     name: String? = null,
     examples: List<T> = emptyList(),
-    fixingParser: StructureFixingParser? = null
+    fixingConfig: StructuredOutputFixingConfig
 ): AIAgentNodeDelegate<String, Result<StructuredResponse<T>>> =
     node(name) { message ->
         llm.writeSession {
@@ -239,7 +240,7 @@ public inline fun <reified T> AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestStr
 
             requestLLMStructured<T>(
                 examples = examples,
-                fixingParser = fixingParser
+                fixingConfig = fixingConfig
             )
         }
     }
@@ -540,7 +541,7 @@ public inline fun <reified TInput, T> AIAgentSubgraphBuilderBase<*, *>.nodeSetSt
 ): AIAgentNodeDelegate<TInput, TInput> =
     node(name) { message ->
         llm.writeSession {
-            prompt = config.updatePrompt(model, prompt)
+            prompt = getStructuredOutputPrompt(prompt, config.structuredOutput(model))
             message
         }
     }

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/feature/PromptExecutorProxy.kt

@@ -4,17 +4,21 @@ import ai.koog.agents.core.feature.pipeline.AIAgentPipeline
 import ai.koog.agents.core.tools.ToolDescriptor
 import ai.koog.prompt.dsl.ModerationResult
 import ai.koog.prompt.dsl.Prompt
-import ai.koog.prompt.executor.model.LLMChoice
 import ai.koog.prompt.executor.model.PromptExecutor
 import ai.koog.prompt.llm.LLModel
+import ai.koog.prompt.message.LLMChoice
 import ai.koog.prompt.message.Message
 import ai.koog.prompt.streaming.StreamFrame
+import ai.koog.prompt.structure.StructuredOutput
+import ai.koog.prompt.structure.StructuredOutputFixingConfig
+import ai.koog.prompt.structure.StructuredResponse
 import io.github.oshai.kotlinlogging.KotlinLogging
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.catch
 import kotlinx.coroutines.flow.onCompletion
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.onStart
+import kotlinx.serialization.KSerializer
 import kotlin.uuid.ExperimentalUuidApi
 import kotlin.uuid.Uuid
 
@@ -115,6 +119,27 @@ public class PromptExecutorProxy(
         return responses
     }
 
+    override suspend fun <T> executeStructured(
+        prompt: Prompt,
+        model: LLModel,
+        structuredOutput: StructuredOutput<T>,
+        fixingConfig: StructuredOutputFixingConfig?
+    ): Result<StructuredResponse<T>> {
+        logger.debug { "Executing structured LLM request (prompt: $prompt)" }
+        return executor.executeStructured(prompt, model, structuredOutput, fixingConfig)
+    }
+
+    override suspend fun <T> executeStructured(
+        prompt: Prompt,
+        model: LLModel,
+        serializer: KSerializer<T>,
+        examples: List<T>,
+        fixingConfig: StructuredOutputFixingConfig?
+    ): Result<StructuredResponse<T>> {
+        logger.debug { "Executing structured LLM request (prompt: $prompt)" }
+        return executor.executeStructured(prompt, model, serializer, examples, fixingConfig)
+    }
+
     override suspend fun moderate(
         prompt: Prompt,
         model: LLModel

agents/agents-core/src/commonTest/kotlin/ai/koog/agents/core/agent/session/AIAgentLLMSessionStructuredOutputTest.kt

@@ -10,7 +10,7 @@ import ai.koog.prompt.executor.clients.openai.OpenAIModels
 import ai.koog.prompt.message.Message
 import ai.koog.prompt.message.ResponseMetaInfo
 import ai.koog.prompt.structure.StructuredOutput
-import ai.koog.prompt.structure.StructuredOutputConfig
+import ai.koog.prompt.structure.StructuredOutputFixingConfig
 import ai.koog.prompt.structure.json.JsonStructuredData
 import kotlinx.coroutines.test.runTest
 import kotlinx.serialization.Serializable
@@ -33,7 +33,7 @@ class AIAgentLLMSessionStructuredOutputTest : AgentTestBase() {
     @Test
     fun testParseResponseToStructuredResponse() = runTest {
         val structure = JsonStructuredData.createJsonStructure<TestStructure>()
-        val config = StructuredOutputConfig(
+        val config = StructuredOutputFixingConfig(
             default = StructuredOutput.Manual(structure)
         )
 
@@ -86,7 +86,7 @@ class AIAgentLLMSessionStructuredOutputTest : AgentTestBase() {
     @Test
     fun testParseResponseToStructuredResponseWithNullableField() = runTest {
         val structure = JsonStructuredData.createJsonStructure<TestStructure>()
-        val config = StructuredOutputConfig(
+        val config = StructuredOutputFixingConfig(
             default = StructuredOutput.Manual(structure)
         )
 
@@ -156,7 +156,7 @@ class AIAgentLLMSessionStructuredOutputTest : AgentTestBase() {
         )
 
         val structure = JsonStructuredData.createJsonStructure<ComplexStructure>()
-        val config = StructuredOutputConfig(
+        val config = StructuredOutputFixingConfig(
             default = StructuredOutput.Manual(structure)
         )