Refactor structured output API improving flexibility, support native structured output for OpenAI and Google. (#443)

Author: EugeneTheDev

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

@@ -13,8 +13,11 @@ 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.json.JsonSchemaGenerator
+import ai.koog.prompt.structure.StructureFixingParser
+import ai.koog.prompt.structure.StructuredOutput
+import ai.koog.prompt.structure.StructuredOutputConfig
 import ai.koog.prompt.structure.json.JsonStructuredData
+import ai.koog.prompt.structure.json.generator.StandardJsonSchemaGenerator
 import io.github.oshai.kotlinlogging.KotlinLogging
 import kotlinx.serialization.Serializable
 import kotlin.reflect.KType
@@ -99,11 +102,15 @@ public open class AIAgentSubgraph<Input, Output>(
             }
 
             val selectedTools = this.requestLLMStructured(
-                structure = JsonStructuredData.createJsonStructure<SelectedTools>(
-                    schemaFormat = JsonSchemaGenerator.SchemaFormat.JsonSchema,
-                    examples = listOf(SelectedTools(listOf()), SelectedTools(tools.map { it.name }.take(3))),
-                ),
-                retries = toolSelectionStrategy.maxRetries,
+                config = StructuredOutputConfig(
+                    default = StructuredOutput.Manual(
+                        JsonStructuredData.createJsonStructure<SelectedTools>(
+                            schemaGenerator = StandardJsonSchemaGenerator,
+                            examples = listOf(SelectedTools(listOf()), SelectedTools(tools.map { it.name }.take(3))),
+                        ),
+                    ),
+                    fixingParser = toolSelectionStrategy.fixingParser,
+                )
             ).getOrThrow()
 
             prompt = initialPrompt
@@ -262,8 +269,9 @@ 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 maxRetries: Int = 3) : ToolSelectionStrategy
+    public data class AutoSelectForTask(val subtaskDescription: String, val fixingParser: StructureFixingParser? = null) : ToolSelectionStrategy
 
     /**
      * Represents a subset of tools to be utilized within a subgraph or task.

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

@@ -6,16 +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.clients.openai.OpenAIModels
 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.Message
 import ai.koog.prompt.params.LLMParams
-import ai.koog.prompt.structure.StructuredData
+import ai.koog.prompt.structure.StructureFixingParser
+import ai.koog.prompt.structure.StructuredOutputConfig
 import ai.koog.prompt.structure.StructuredResponse
 import ai.koog.prompt.structure.executeStructured
-import ai.koog.prompt.structure.executeStructuredOneShot
+import kotlinx.serialization.KSerializer
+import kotlinx.serialization.serializer
 
 /**
  * Represents a session for an AI agent that interacts with an LLM (Language Learning Model).
@@ -49,9 +50,7 @@ public sealed class AIAgentLLMSession(
      * Typical usage includes providing input to LLM requests, such as:
      * - [requestLLMWithoutTools]
      * - [requestLLM]
-     * - [requestLLMMultiple]
-     * - [requestLLMStructured]
-     * - [requestLLMStructuredOneShot]
+     * etc.
      */
     public open val prompt: Prompt by ActiveProperty(prompt) { isActive }
 
@@ -238,32 +237,81 @@ public sealed class AIAgentLLMSession(
     }
 
     /**
-     * Coerce LLM to provide a structured output.
+     * Sends a request to LLM and gets a structured response.
+     *
+     * @param config A configuration defining structures and behavior.
      *
      * @see [executeStructured]
      */
     public open suspend fun <T> requestLLMStructured(
-        structure: StructuredData<T>,
-        retries: Int = 1,
-        fixingModel: LLModel = OpenAIModels.Chat.GPT4o
+        config: StructuredOutputConfig<T>,
     ): Result<StructuredResponse<T>> {
         validateSession()
+
         val preparedPrompt = preparePrompt(prompt, tools = emptyList())
-        return executor.executeStructured(preparedPrompt, model, structure, retries, fixingModel)
+
+        return executor.executeStructured(
+            prompt = preparedPrompt,
+            model = model,
+            config = config,
+        )
     }
 
     /**
-     * Expect LLM to reply in a structured format and try to parse it.
-     * For more robust version with model coercion and correction see [requestLLMStructured]
+     * 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.
      *
-     * @see [executeStructuredOneShot]
+     * @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.
      */
-    public open suspend fun <T> requestLLMStructuredOneShot(structure: StructuredData<T>): StructuredResponse<T> {
+    public open suspend fun <T> requestLLMStructured(
+        serializer: KSerializer<T>,
+        examples: List<T> = emptyList(),
+        fixingParser: StructureFixingParser? = null
+    ): Result<StructuredResponse<T>> {
         validateSession()
+
         val preparedPrompt = preparePrompt(prompt, tools = emptyList())
-        return executor.executeStructuredOneShot(preparedPrompt, model, structure)
+
+        return executor.executeStructured(
+            prompt = preparedPrompt,
+            model = model,
+            serializer = serializer,
+            examples = examples,
+            fixingParser = fixingParser,
+        )
     }
 
+    /**
+     * 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,
+    )
+
     /**
      * 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

@@ -16,13 +16,15 @@ import ai.koog.prompt.executor.model.PromptExecutor
 import ai.koog.prompt.llm.LLModel
 import ai.koog.prompt.message.Message
 import ai.koog.prompt.params.LLMParams
-import ai.koog.prompt.structure.StructuredData
+import ai.koog.prompt.structure.StructureFixingParser
 import ai.koog.prompt.structure.StructuredDataDefinition
+import ai.koog.prompt.structure.StructuredOutputConfig
 import ai.koog.prompt.structure.StructuredResponse
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.flatMapMerge
 import kotlinx.coroutines.flow.flow
 import kotlinx.datetime.Clock
+import kotlinx.serialization.KSerializer
 import kotlin.reflect.KClass
 
 /**
@@ -424,23 +426,47 @@ public class AIAgentLLMWriteSession internal constructor(
     }
 
     /**
-     * Requests an LLM (Language Model) to generate a structured output based on the provided structure.
-     * The response is post-processed to update the prompt with the raw response.
+     * Sends a request to LLM and gets a structured response.
      *
-     * @param structure The structured data definition specifying the expected structured output format, schema, and parsing logic.
-     * @param retries The number of retry attempts to allow in case of generation failures.
-     * @param fixingModel The language model to use for re-parsing or error correction during retries.
-     * @return A structured response containing both the parsed structure and the raw response text.
+     * @param config A configuration defining structures and behavior.
+     *
+     * @see [executeStructured]
+     */
+    override suspend fun <T> requestLLMStructured(
+        config: StructuredOutputConfig<T>,
+    ): Result<StructuredResponse<T>> {
+        return super.requestLLMStructured(config).also {
+            it.onSuccess { response ->
+                updatePrompt {
+                    message(response.message)
+                }
+            }
+        }
+    }
+
+    /**
+     * 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 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.
      */
     override suspend fun <T> requestLLMStructured(
-        structure: StructuredData<T>,
-        retries: Int,
-        fixingModel: LLModel
+        serializer: KSerializer<T>,
+        examples: List<T>,
+        fixingParser: StructureFixingParser?
     ): Result<StructuredResponse<T>> {
-        return super.requestLLMStructured(structure, retries, fixingModel).also {
+        return super.requestLLMStructured(serializer, examples, fixingParser).also {
             it.onSuccess { response ->
                 updatePrompt {
-                    assistant(response.raw)
+                    message(response.message)
                 }
             }
         }
@@ -465,19 +491,4 @@ public class AIAgentLLMWriteSession internal constructor(
 
         return executor.executeStreaming(prompt, model)
     }
-
-    /**
-     * Sends a request to the LLM using the given structured data and expects a structured response in one attempt.
-     * Updates the prompt with the raw response received from the LLM.
-     *
-     * @param structure The structured data defining the schema, examples, and parsing logic for the response.
-     * @return A structured response containing both the parsed data and the raw response text from the LLM.
-     */
-    override suspend fun <T> requestLLMStructuredOneShot(structure: StructuredData<T>): StructuredResponse<T> {
-        return super.requestLLMStructuredOneShot(structure).also { response ->
-            updatePrompt {
-                assistant(response.raw)
-            }
-        }
-    }
 }

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

@@ -17,8 +17,9 @@ import ai.koog.prompt.dsl.PromptBuilder
 import ai.koog.prompt.dsl.prompt
 import ai.koog.prompt.llm.LLModel
 import ai.koog.prompt.message.Message
-import ai.koog.prompt.structure.StructuredData
+import ai.koog.prompt.structure.StructureFixingParser
 import ai.koog.prompt.structure.StructuredDataDefinition
+import ai.koog.prompt.structure.StructuredOutputConfig
 import ai.koog.prompt.structure.StructuredResponse
 import kotlinx.coroutines.flow.Flow
 
@@ -173,30 +174,57 @@ public fun AIAgentSubgraphBuilderBase<*, *>.nodeLLMModerateMessage(
     }
 
 /**
- * A node that appends a user message to the LLM prompt and requests structured data from the LLM with error correction capabilities.
+ * A node that appends a user message to the LLM prompt and requests structured data from the LLM with optional error
+ * correction capabilities.
  *
  * @param name Optional node name.
- * @param structure Definition of expected output format and parsing logic.
- * @param retries Number of retry attempts for failed generations.
- * @param fixingModel LLM used for error correction.
+ * @param config A configuration defining structures and behavior.
  */
 @AIAgentBuilderDslMarker
 public inline fun <reified T> AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestStructured(
     name: String? = null,
-    structure: StructuredData<T>,
-    retries: Int,
-    fixingModel: LLModel
+    config: StructuredOutputConfig<T>,
 ): AIAgentNodeDelegate<String, Result<StructuredResponse<T>>> =
     node(name) { message ->
         llm.writeSession {
             updatePrompt {
                 user(message)
             }
 
-            requestLLMStructured(
-                structure,
-                retries,
-                fixingModel
+            requestLLMStructured(config)
+        }
+    }
+
+/**
+ * A node that appends a user message to the LLM prompt and requests structured data from the LLM with optional error
+ * correction capabilities.
+ *
+ * This is a simple version of the full `nodeLLMRequestStructured`. 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 name Optional node name.
+ * @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.
+ */
+@AIAgentBuilderDslMarker
+public inline fun <reified T> AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestStructured(
+    name: String? = null,
+    examples: List<T> = emptyList(),
+    fixingParser: StructureFixingParser? = null
+): AIAgentNodeDelegate<String, Result<StructuredResponse<T>>> =
+    node(name) { message ->
+        llm.writeSession {
+            updatePrompt {
+                user(message)
+            }
+
+            requestLLMStructured<T>(
+                examples = examples,
+                fixingParser = fixingParser
             )
         }
     }

agents/agents-features/agents-features-memory/src/commonMain/kotlin/ai/koog/agents/memory/feature/AgentMemory.kt

@@ -28,6 +28,8 @@ import ai.koog.agents.memory.providers.NoMemory
 import ai.koog.prompt.dsl.Prompt
 import ai.koog.prompt.llm.LLModel
 import ai.koog.prompt.message.Message
+import ai.koog.prompt.structure.StructuredOutput
+import ai.koog.prompt.structure.StructuredOutputConfig
 import ai.koog.prompt.structure.json.JsonStructuredData
 import io.github.oshai.kotlinlogging.KotlinLogging
 import kotlinx.serialization.Serializable
@@ -554,7 +556,10 @@ internal suspend fun AIAgentLLMWriteSession.retrieveFactsFromHistory(
 
     val facts = when (concept.factType) {
         FactType.SINGLE -> {
-            val response = requestLLMStructured(JsonStructuredData.createJsonStructure<FactStructure>())
+            val response = requestLLMStructured(
+                config = StructuredOutputConfig(default = StructuredOutput.Manual(JsonStructuredData.createJsonStructure<FactStructure>()))
+            )
+
             SingleFact(
                 concept = concept,
                 value = response.getOrNull()?.structure?.fact ?: "No facts extracted",
@@ -563,7 +568,9 @@ internal suspend fun AIAgentLLMWriteSession.retrieveFactsFromHistory(
         }
 
         FactType.MULTIPLE -> {
-            val response = requestLLMStructured(JsonStructuredData.createJsonStructure<FactListStructure>())
+            val response = requestLLMStructured(
+                config = StructuredOutputConfig(default = StructuredOutput.Manual(JsonStructuredData.createJsonStructure<FactListStructure>()))
+            )
             val factsList = response.getOrNull()?.structure?.facts ?: emptyList()
             MultipleFacts(concept = concept, values = factsList.map { it.fact }, timestamp = timestamp)
         }

agents/agents-features/agents-features-memory/src/jvmTest/kotlin/ai/koog/agents/memory/AIAgentMemoryTest.kt

@@ -25,6 +25,7 @@ import ai.koog.prompt.dsl.prompt
 import ai.koog.prompt.executor.clients.anthropic.AnthropicModels
 import ai.koog.prompt.executor.clients.openai.OpenAIModels
 import ai.koog.prompt.executor.model.PromptExecutor
+import ai.koog.prompt.llm.LLMProvider
 import ai.koog.prompt.llm.LLModel
 import ai.koog.prompt.message.Message
 import io.mockk.coEvery
@@ -60,6 +61,7 @@ class AIAgentMemoryTest {
 
     private val testModel = mockk<LLModel> {
         every { id } returns "test-model"
+        every { provider } returns mockk<LLMProvider>()
     }
 
     private val testClock: Clock = object : Clock {