From 33442eb9e666c449833f9aba064a62eaf8088656 Mon Sep 17 00:00:00 2001 From: Maria Tigina Date: Fri, 14 Nov 2025 11:38:37 +0100 Subject: [PATCH] Move structured output to client TODO: enable fixing parser in executor --- .../core/agent/entity/AIAgentSubgraph.kt | 30 +- .../core/agent/session/AIAgentLLMSession.kt | 13 +- .../agent/session/AIAgentLLMWriteSession.kt | 5 +- .../extension/AIAgentFunctionalContextExt.kt | 2 +- .../agents/core/dsl/extension/AIAgentNodes.kt | 8 +- .../agents/ext/agent/AIAgentStrategies.kt | 6 +- .../ai/koog/agents/ext/agent/LLMAsAJudge.kt | 2 +- .../banking/routing/RoutingViaGraph.kt | 2 +- .../AdvancedWithBasicSchema.kt | 15 +- .../AdvancedWithStandardSchema.kt | 14 +- .../AdvancedWithStandardSchemaAndTools.kt | 15 +- .../example/structuredoutput/SimpleExample.kt | 2 +- .../utils/structuredOutput/WeatherReport.kt | 2 +- .../prompt-executor-clients/build.gradle.kts | 1 + .../clients/dashscope/DashscopeLLMClient.kt | 5 - .../clients/deepseek/DeepSeekLLMClient.kt | 5 - .../clients/google/GoogleLLMClient.kt | 24 +- .../clients/mistralai/MistralAILLMClient.kt | 5 - .../openai/base/AbstractOpenAILLMClient.kt | 33 +- .../clients/openai/OpenAILLMClient.kt | 5 - .../clients/openrouter/OpenRouterLLMClient.kt | 5 - .../koog/prompt/executor/clients/LLMClient.kt | 18 + .../clients/LLMClientStructuredOutput.kt | 94 +++++ .../executor/llms/MultiLLMPromptExecutor.kt | 27 ++ .../executor/llms/SingleLLMPromptExecutor.kt | 21 ++ .../prompt-executor-model/build.gradle.kts | 9 + .../model}/LLMStructuredParsingError.kt | 4 +- .../prompt/executor/model/PromptExecutor.kt | 43 +++ .../model/PromptExecutorStructuredOutput.kt | 144 ++++++++ .../executor/model}/StructureFixingParser.kt | 4 +- .../model}/StructureFixingParserTest.kt | 8 +- prompt/prompt-structure/build.gradle.kts | 3 +- .../structure/PromptExecutorExtensions.kt | 334 ------------------ .../koog/prompt/structure/StructuredOutput.kt | 163 +++++++++ .../prompt/structure/StructuredPrompts.kt | 39 +- 35 files changed, 664 insertions(+), 446 deletions(-) create mode 100644 prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/LLMClientStructuredOutput.kt rename prompt/{prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure => prompt-executor/prompt-executor-model/src/commonMain/kotlin/ai/koog/prompt/executor/model}/LLMStructuredParsingError.kt (82%) create mode 100644 prompt/prompt-executor/prompt-executor-model/src/commonMain/kotlin/ai/koog/prompt/executor/model/PromptExecutorStructuredOutput.kt rename prompt/{prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure => prompt-executor/prompt-executor-model/src/commonMain/kotlin/ai/koog/prompt/executor/model}/StructureFixingParser.kt (98%) rename prompt/{prompt-structure/src/commonTest/kotlin/ai/koog/prompt/structure => prompt-executor/prompt-executor-model/src/commonTest/kotlin/ai/koog/prompt/executor/model}/StructureFixingParserTest.kt (87%) delete mode 100644 prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/PromptExecutorExtensions.kt create mode 100644 prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/StructuredOutput.kt diff --git a/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/entity/AIAgentSubgraph.kt b/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/entity/AIAgentSubgraph.kt index fa24c6a382..28d0e7cff7 100644 --- a/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/entity/AIAgentSubgraph.kt +++ b/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/entity/AIAgentSubgraph.kt @@ -14,9 +14,9 @@ import ai.koog.agents.core.dsl.extension.replaceHistoryWithTLDR import ai.koog.agents.core.prompt.Prompts.selectRelevantTools import ai.koog.agents.core.tools.ToolDescriptor import ai.koog.agents.core.tools.annotations.LLMDescription +import ai.koog.prompt.executor.model.StructureFixingParser import ai.koog.prompt.llm.LLModel import ai.koog.prompt.params.LLMParams -import ai.koog.prompt.structure.StructureFixingParser import ai.koog.prompt.structure.StructuredRequest import ai.koog.prompt.structure.StructuredRequestConfig import ai.koog.prompt.structure.json.JsonStructure @@ -137,8 +137,8 @@ public open class AIAgentSubgraph( examples = listOf(SelectedTools(listOf()), SelectedTools(tools.map { it.name }.take(3))), ), ), - fixingParser = toolSelectionStrategy.fixingParser, - ) + ), + fixingParser = toolSelectionStrategy.fixingParser, ).getOrThrow() prompt = initialPrompt @@ -157,7 +157,15 @@ public open class AIAgentSubgraph( */ @OptIn(InternalAgentsApi::class, DetachedPromptExecutorAPI::class, ExperimentalUuidApi::class) override suspend fun execute(context: AIAgentGraphContextBase, input: TInput): TOutput? = - withContext(NodeInfoContextElement(Uuid.random().toString(), getNodeInfoElement()?.id, name, input, inputType)) { + withContext( + NodeInfoContextElement( + Uuid.random().toString(), + getNodeInfoElement()?.id, + name, + input, + inputType + ) + ) { val newTools = selectTools(context) // Copy inner context with new tools, model and LLM params. @@ -201,7 +209,14 @@ public open class AIAgentSubgraph( } runIfNonRootContext(context) { - pipeline.onSubgraphExecutionCompleted(this@AIAgentSubgraph, innerContext, input, inputType, result, outputType) + pipeline.onSubgraphExecutionCompleted( + this@AIAgentSubgraph, + innerContext, + input, + inputType, + result, + outputType + ) } result @@ -284,7 +299,10 @@ public open class AIAgentSubgraph( * effectively skipping execution for root contexts. */ @OptIn(InternalAgentsApi::class) - private suspend fun runIfNonRootContext(context: AIAgentGraphContextBase, action: suspend AIAgentGraphContextBase.() -> Unit) { + private suspend fun runIfNonRootContext( + context: AIAgentGraphContextBase, + action: suspend AIAgentGraphContextBase.() -> Unit + ) { if (context.parentContext == null) return action(context) } diff --git a/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/session/AIAgentLLMSession.kt b/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/session/AIAgentLLMSession.kt index e4dd3f400c..1e78d00dbc 100644 --- a/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/session/AIAgentLLMSession.kt +++ b/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/session/AIAgentLLMSession.kt @@ -7,16 +7,16 @@ 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.PromptExecutor +import ai.koog.prompt.executor.model.StructureFixingParser +import ai.koog.prompt.executor.model.executeStructured +import ai.koog.prompt.executor.model.parseResponseToStructuredResponse 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.StructuredRequestConfig import ai.koog.prompt.structure.StructuredResponse -import ai.koog.prompt.structure.executeStructured -import ai.koog.prompt.structure.parseResponseToStructuredResponse import kotlinx.coroutines.flow.Flow import kotlinx.serialization.KSerializer import kotlinx.serialization.serializer @@ -264,6 +264,7 @@ public sealed class AIAgentLLMSession( */ public open suspend fun requestLLMStructured( config: StructuredRequestConfig, + fixingParser: StructureFixingParser? = null ): Result> { validateSession() @@ -273,6 +274,7 @@ public sealed class AIAgentLLMSession( prompt = preparedPrompt, model = model, config = config, + fixingParser = fixingParser ) } @@ -346,8 +348,9 @@ public sealed class AIAgentLLMSession( */ public suspend fun parseResponseToStructuredResponse( response: Message.Assistant, - config: StructuredRequestConfig - ): StructuredResponse = executor.parseResponseToStructuredResponse(response, config, model) + config: StructuredRequestConfig, + fixingParser: StructureFixingParser? = null + ): StructuredResponse = executor.parseResponseToStructuredResponse(response, config, model, fixingParser) /** * Sends a request to the language model, potentially receiving multiple choices, diff --git a/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/session/AIAgentLLMWriteSession.kt b/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/session/AIAgentLLMWriteSession.kt index ff1daffc8d..34abd8c8da 100644 --- a/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/session/AIAgentLLMWriteSession.kt +++ b/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/session/AIAgentLLMWriteSession.kt @@ -11,12 +11,12 @@ import ai.koog.prompt.dsl.Prompt import ai.koog.prompt.dsl.PromptBuilder import ai.koog.prompt.dsl.prompt import ai.koog.prompt.executor.model.PromptExecutor +import ai.koog.prompt.executor.model.StructureFixingParser 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.StructureDefinition -import ai.koog.prompt.structure.StructureFixingParser import ai.koog.prompt.structure.StructuredRequestConfig import ai.koog.prompt.structure.StructuredResponse import kotlinx.coroutines.flow.Flow @@ -444,8 +444,9 @@ public class AIAgentLLMWriteSession internal constructor( */ override suspend fun requestLLMStructured( config: StructuredRequestConfig, + fixingParser: StructureFixingParser? ): Result> { - return super.requestLLMStructured(config).also { + return super.requestLLMStructured(config, fixingParser).also { it.onSuccess { response -> appendPrompt { message(response.message) diff --git a/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/extension/AIAgentFunctionalContextExt.kt b/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/extension/AIAgentFunctionalContextExt.kt index d19caab204..f80bf9c76e 100644 --- a/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/extension/AIAgentFunctionalContextExt.kt +++ b/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/extension/AIAgentFunctionalContextExt.kt @@ -9,10 +9,10 @@ import ai.koog.agents.core.tools.Tool import ai.koog.agents.core.tools.ToolArgs import ai.koog.agents.core.tools.ToolDescriptor import ai.koog.agents.core.tools.ToolResult +import ai.koog.prompt.executor.model.StructureFixingParser import ai.koog.prompt.message.Message import ai.koog.prompt.streaming.StreamFrame import ai.koog.prompt.structure.StructureDefinition -import ai.koog.prompt.structure.StructureFixingParser import ai.koog.prompt.structure.StructuredResponse import kotlinx.coroutines.flow.Flow import kotlinx.serialization.serializer diff --git a/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/extension/AIAgentNodes.kt b/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/extension/AIAgentNodes.kt index b0deaa1d63..bce258de4c 100644 --- a/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/extension/AIAgentNodes.kt +++ b/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/extension/AIAgentNodes.kt @@ -13,12 +13,13 @@ import ai.koog.agents.core.tools.ToolDescriptor import ai.koog.prompt.dsl.ModerationResult import ai.koog.prompt.dsl.PromptBuilder import ai.koog.prompt.dsl.prompt +import ai.koog.prompt.executor.model.StructureFixingParser 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.StructureDefinition -import ai.koog.prompt.structure.StructureFixingParser +import ai.koog.prompt.structure.StructuredOutputPrompts.appendStructuredOutputInstructions import ai.koog.prompt.structure.StructuredRequestConfig import ai.koog.prompt.structure.StructuredResponse import kotlinx.coroutines.flow.Flow @@ -199,6 +200,7 @@ public fun AIAgentSubgraphBuilderBase<*, *>.nodeLLMModerateMessage( public inline fun AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestStructured( name: String? = null, config: StructuredRequestConfig, + fixingParser: StructureFixingParser? = null, ): AIAgentNodeDelegate>> = node(name) { message -> llm.writeSession { @@ -206,7 +208,7 @@ public inline fun AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestStr user(message) } - requestLLMStructured(config) + requestLLMStructured(config, fixingParser) } } @@ -540,7 +542,7 @@ public inline fun AIAgentSubgraphBuilderBase<*, *>.nodeSetSt ): AIAgentNodeDelegate = node(name) { message -> llm.writeSession { - prompt = config.updatePrompt(model, prompt) + prompt = appendStructuredOutputInstructions(prompt, config.structuredRequest(model)) message } } diff --git a/agents/agents-ext/src/commonMain/kotlin/ai/koog/agents/ext/agent/AIAgentStrategies.kt b/agents/agents-ext/src/commonMain/kotlin/ai/koog/agents/ext/agent/AIAgentStrategies.kt index d78d1c7a7e..223925defd 100644 --- a/agents/agents-ext/src/commonMain/kotlin/ai/koog/agents/ext/agent/AIAgentStrategies.kt +++ b/agents/agents-ext/src/commonMain/kotlin/ai/koog/agents/ext/agent/AIAgentStrategies.kt @@ -18,6 +18,7 @@ import ai.koog.agents.core.dsl.extension.onMultipleToolCalls import ai.koog.agents.core.dsl.extension.onToolCall import ai.koog.agents.core.environment.ReceivedToolResult import ai.koog.agents.core.environment.result +import ai.koog.prompt.executor.model.StructureFixingParser import ai.koog.prompt.message.Message import ai.koog.prompt.structure.StructuredRequestConfig @@ -184,9 +185,11 @@ public fun reActStrategy( */ public inline fun structuredOutputWithToolsStrategy( config: StructuredRequestConfig, + fixingParser: StructureFixingParser? = null, parallelTools: Boolean = false ): AIAgentGraphStrategy = structuredOutputWithToolsStrategy( config, + fixingParser, parallelTools ) { it } @@ -204,6 +207,7 @@ public inline fun structuredOutputWithToolsStrategy( */ public inline fun structuredOutputWithToolsStrategy( config: StructuredRequestConfig, + fixingParser: StructureFixingParser? = null, parallelTools: Boolean = false, noinline transform: suspend AIAgentGraphContextBase.(input: Input) -> String ): AIAgentGraphStrategy = strategy("structured_output_with_tools_strategy") { @@ -214,7 +218,7 @@ public inline fun structuredOutputWithToolsStrat val sendToolResult by nodeLLMSendMultipleToolResults() val transformToStructuredOutput by node { response -> llm.writeSession { - parseResponseToStructuredResponse(response, config).data + parseResponseToStructuredResponse(response, config, fixingParser).data } } diff --git a/agents/agents-ext/src/commonMain/kotlin/ai/koog/agents/ext/agent/LLMAsAJudge.kt b/agents/agents-ext/src/commonMain/kotlin/ai/koog/agents/ext/agent/LLMAsAJudge.kt index 960dddcc36..68b5cd8337 100644 --- a/agents/agents-ext/src/commonMain/kotlin/ai/koog/agents/ext/agent/LLMAsAJudge.kt +++ b/agents/agents-ext/src/commonMain/kotlin/ai/koog/agents/ext/agent/LLMAsAJudge.kt @@ -6,9 +6,9 @@ import ai.koog.agents.core.dsl.builder.AIAgentSubgraphBuilderBase import ai.koog.agents.core.tools.annotations.LLMDescription import ai.koog.prompt.dsl.prompt import ai.koog.prompt.executor.clients.openai.OpenAIModels +import ai.koog.prompt.executor.model.StructureFixingParser import ai.koog.prompt.llm.LLModel import ai.koog.prompt.message.Message -import ai.koog.prompt.structure.StructureFixingParser import kotlinx.serialization.Serializable /** diff --git a/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/banking/routing/RoutingViaGraph.kt b/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/banking/routing/RoutingViaGraph.kt index afb1a79a03..b1bac6a648 100644 --- a/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/banking/routing/RoutingViaGraph.kt +++ b/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/banking/routing/RoutingViaGraph.kt @@ -21,7 +21,7 @@ import ai.koog.agents.ext.tool.AskUser import ai.koog.prompt.dsl.prompt import ai.koog.prompt.executor.clients.openai.OpenAIModels import ai.koog.prompt.executor.llms.all.simpleOpenAIExecutor -import ai.koog.prompt.structure.StructureFixingParser +import ai.koog.prompt.executor.model.StructureFixingParser import kotlinx.coroutines.runBlocking suspend fun main() { diff --git a/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/structuredoutput/AdvancedWithBasicSchema.kt b/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/structuredoutput/AdvancedWithBasicSchema.kt index cc5f322eb8..90a6e88c27 100644 --- a/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/structuredoutput/AdvancedWithBasicSchema.kt +++ b/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/structuredoutput/AdvancedWithBasicSchema.kt @@ -19,8 +19,8 @@ import ai.koog.prompt.executor.clients.google.structure.GoogleBasicJsonSchemaGen import ai.koog.prompt.executor.clients.openai.OpenAILLMClient import ai.koog.prompt.executor.clients.openai.base.structure.OpenAIBasicJsonSchemaGenerator import ai.koog.prompt.executor.llms.MultiLLMPromptExecutor +import ai.koog.prompt.executor.model.StructureFixingParser import ai.koog.prompt.llm.LLMProvider -import ai.koog.prompt.structure.StructureFixingParser import ai.koog.prompt.structure.StructuredRequest import ai.koog.prompt.structure.StructuredRequestConfig import ai.koog.prompt.structure.json.JsonStructure @@ -212,13 +212,12 @@ suspend fun main() { // Fallback manual structured output mode, via explicit prompting with additional message, not native model support default = StructuredRequest.Manual(genericWeatherStructure), - - // Helper parser to attempt a fix if a malformed output is produced. - fixingParser = StructureFixingParser( - model = AnthropicModels.Haiku_3_5, - retries = 2, - ), - ) + ), + // Helper parser to attempt a fix if a malformed output is produced. + fixingParser = StructureFixingParser( + model = AnthropicModels.Haiku_3_5, + retries = 2, + ), ) nodeStart then prepareRequest then getStructuredForecast diff --git a/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/structuredoutput/AdvancedWithStandardSchema.kt b/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/structuredoutput/AdvancedWithStandardSchema.kt index b2b7d25a9a..3a33b9330e 100644 --- a/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/structuredoutput/AdvancedWithStandardSchema.kt +++ b/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/structuredoutput/AdvancedWithStandardSchema.kt @@ -18,8 +18,8 @@ import ai.koog.prompt.executor.clients.google.structure.GoogleStandardJsonSchema import ai.koog.prompt.executor.clients.openai.OpenAILLMClient import ai.koog.prompt.executor.clients.openai.base.structure.OpenAIStandardJsonSchemaGenerator import ai.koog.prompt.executor.llms.MultiLLMPromptExecutor +import ai.koog.prompt.executor.model.StructureFixingParser import ai.koog.prompt.llm.LLMProvider -import ai.koog.prompt.structure.StructureFixingParser import ai.koog.prompt.structure.StructuredRequest import ai.koog.prompt.structure.StructuredRequestConfig import ai.koog.prompt.structure.json.JsonStructure @@ -81,13 +81,13 @@ suspend fun main() { // Fallback manual structured output mode, via explicit prompting with additional message, not native model support default = StructuredRequest.Manual(genericWeatherStructure), + ), - // Helper parser to attempt a fix if a malformed output is produced. - fixingParser = StructureFixingParser( - model = AnthropicModels.Haiku_3_5, - retries = 2, - ), - ) + // Helper parser to attempt a fix if a malformed output is produced. + fixingParser = StructureFixingParser( + model = AnthropicModels.Haiku_3_5, + retries = 2, + ), ) nodeStart then prepareRequest then getStructuredForecast diff --git a/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/structuredoutput/AdvancedWithStandardSchemaAndTools.kt b/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/structuredoutput/AdvancedWithStandardSchemaAndTools.kt index a6a6a92f7e..f56c0efa6e 100644 --- a/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/structuredoutput/AdvancedWithStandardSchemaAndTools.kt +++ b/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/structuredoutput/AdvancedWithStandardSchemaAndTools.kt @@ -19,8 +19,8 @@ import ai.koog.prompt.executor.clients.openai.OpenAILLMClient import ai.koog.prompt.executor.clients.openai.OpenAIModels import ai.koog.prompt.executor.clients.openai.base.structure.OpenAIStandardJsonSchemaGenerator import ai.koog.prompt.executor.llms.MultiLLMPromptExecutor +import ai.koog.prompt.executor.model.StructureFixingParser import ai.koog.prompt.llm.LLMProvider -import ai.koog.prompt.structure.StructureFixingParser import ai.koog.prompt.structure.StructuredRequest import ai.koog.prompt.structure.StructuredRequestConfig import ai.koog.prompt.structure.json.JsonStructure @@ -71,16 +71,17 @@ suspend fun main() { // Fallback manual structured output mode, via explicit prompting with additional message, not native model support default = StructuredRequest.Manual(genericWeatherStructure), + ) - // Helper parser to attempt a fix if a malformed output is produced. - fixingParser = StructureFixingParser( - model = AnthropicModels.Haiku_3_5, - retries = 2, - ), + // Helper parser to attempt a fix if a malformed output is produced. + val fixingParser = StructureFixingParser( + model = AnthropicModels.Haiku_3_5, + retries = 2, ) val agentStrategy = structuredOutputWithToolsStrategy( - config + config, + fixingParser ) { request -> text { +"Requesting forecast for" diff --git a/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/structuredoutput/SimpleExample.kt b/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/structuredoutput/SimpleExample.kt index 41c5b3c885..9c009168c8 100644 --- a/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/structuredoutput/SimpleExample.kt +++ b/examples/simple-examples/src/main/kotlin/ai/koog/agents/example/structuredoutput/SimpleExample.kt @@ -15,8 +15,8 @@ import ai.koog.prompt.executor.clients.google.GoogleModels import ai.koog.prompt.executor.clients.openai.OpenAILLMClient import ai.koog.prompt.executor.clients.openai.OpenAIModels import ai.koog.prompt.executor.llms.MultiLLMPromptExecutor +import ai.koog.prompt.executor.model.StructureFixingParser import ai.koog.prompt.llm.LLMProvider -import ai.koog.prompt.structure.StructureFixingParser import ai.koog.prompt.text.text import kotlinx.serialization.SerialName import kotlinx.serialization.Serializable diff --git a/integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/utils/structuredOutput/WeatherReport.kt b/integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/utils/structuredOutput/WeatherReport.kt index b05e0dc155..517b9493e8 100644 --- a/integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/utils/structuredOutput/WeatherReport.kt +++ b/integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/utils/structuredOutput/WeatherReport.kt @@ -2,9 +2,9 @@ package ai.koog.integration.tests.utils.structuredOutput import ai.koog.agents.core.tools.annotations.LLMDescription import ai.koog.prompt.dsl.Prompt +import ai.koog.prompt.executor.model.StructureFixingParser import ai.koog.prompt.llm.LLModel import ai.koog.prompt.structure.RegisteredStandardJsonSchemaGenerators -import ai.koog.prompt.structure.StructureFixingParser import ai.koog.prompt.structure.StructuredRequest import ai.koog.prompt.structure.StructuredRequestConfig import ai.koog.prompt.structure.StructuredResponse diff --git a/prompt/prompt-executor/prompt-executor-clients/build.gradle.kts b/prompt/prompt-executor/prompt-executor-clients/build.gradle.kts index e74a6de399..c4a2d5906a 100644 --- a/prompt/prompt-executor/prompt-executor-clients/build.gradle.kts +++ b/prompt/prompt-executor/prompt-executor-clients/build.gradle.kts @@ -13,6 +13,7 @@ kotlin { dependencies { api(project(":prompt:prompt-model")) api(project(":agents:agents-tools")) + api(project(":prompt:prompt-structure")) api(project(":prompt:prompt-executor:prompt-executor-model")) api(libs.kotlinx.coroutines.core) } diff --git a/prompt/prompt-executor/prompt-executor-clients/prompt-executor-dashscope-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/dashscope/DashscopeLLMClient.kt b/prompt/prompt-executor/prompt-executor-clients/prompt-executor-dashscope-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/dashscope/DashscopeLLMClient.kt index a72f657cc0..fdcac02bba 100644 --- a/prompt/prompt-executor/prompt-executor-clients/prompt-executor-dashscope-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/dashscope/DashscopeLLMClient.kt +++ b/prompt/prompt-executor/prompt-executor-clients/prompt-executor-dashscope-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/dashscope/DashscopeLLMClient.kt @@ -60,11 +60,6 @@ public class DashscopeLLMClient( private companion object { private val staticLogger = KotlinLogging.logger { } - - init { - // On class load register custom OpenAI JSON schema generators for structured output. - registerOpenAIJsonSchemaGenerators(LLMProvider.Alibaba) - } } override fun llmProvider(): LLMProvider = LLMProvider.Alibaba diff --git a/prompt/prompt-executor/prompt-executor-clients/prompt-executor-deepseek-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/deepseek/DeepSeekLLMClient.kt b/prompt/prompt-executor/prompt-executor-clients/prompt-executor-deepseek-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/deepseek/DeepSeekLLMClient.kt index 1b222ff735..b5cbe27452 100644 --- a/prompt/prompt-executor/prompt-executor-clients/prompt-executor-deepseek-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/deepseek/DeepSeekLLMClient.kt +++ b/prompt/prompt-executor/prompt-executor-clients/prompt-executor-deepseek-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/deepseek/DeepSeekLLMClient.kt @@ -61,11 +61,6 @@ public class DeepSeekLLMClient( private companion object { private val staticLogger = KotlinLogging.logger { } - - init { - // On class load register custom OpenAI JSON schema generators for structured output. - registerOpenAIJsonSchemaGenerators(LLMProvider.DeepSeek) - } } /** diff --git a/prompt/prompt-executor/prompt-executor-clients/prompt-executor-google-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/google/GoogleLLMClient.kt b/prompt/prompt-executor/prompt-executor-clients/prompt-executor-google-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/google/GoogleLLMClient.kt index 7a762d7e8b..bd692c9136 100644 --- a/prompt/prompt-executor/prompt-executor-clients/prompt-executor-google-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/google/GoogleLLMClient.kt +++ b/prompt/prompt-executor/prompt-executor-clients/prompt-executor-google-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/google/GoogleLLMClient.kt @@ -36,9 +36,9 @@ import ai.koog.prompt.streaming.emitAppend import ai.koog.prompt.streaming.emitEnd import ai.koog.prompt.streaming.emitToolCall import ai.koog.prompt.streaming.streamFrameFlow -import ai.koog.prompt.structure.RegisteredBasicJsonSchemaGenerators -import ai.koog.prompt.structure.RegisteredStandardJsonSchemaGenerators import ai.koog.prompt.structure.annotations.InternalStructuredOutputApi +import ai.koog.prompt.structure.json.generator.BasicJsonSchemaGenerator +import ai.koog.prompt.structure.json.generator.StandardJsonSchemaGenerator import ai.koog.utils.io.SuitableForIO import io.github.oshai.kotlinlogging.KotlinLogging import io.ktor.client.HttpClient @@ -113,12 +113,6 @@ public open class GoogleLLMClient( private const val DEFAULT_PATH = "v1beta/models" private const val DEFAULT_METHOD_GENERATE_CONTENT = "generateContent" private const val DEFAULT_METHOD_STREAM_GENERATE_CONTENT = "streamGenerateContent" - - init { - // On class load register custom Google JSON schema generators for structured output. - RegisteredBasicJsonSchemaGenerators[LLMProvider.Google] = GoogleBasicJsonSchemaGenerator - RegisteredStandardJsonSchemaGenerators[LLMProvider.Google] = GoogleStandardJsonSchemaGenerator - } } private val json = Json { @@ -252,6 +246,20 @@ public open class GoogleLLMClient( return processGoogleResponse(getGoogleResponse(prompt, model, tools)) } + /** + * Standard JSON schema generator supported by all models provided by the Google LLMClient. + */ + override fun getStandardJsonSchemaGenerator(model: LLModel): StandardJsonSchemaGenerator { + return GoogleStandardJsonSchemaGenerator + } + + /** + * Basic JSON schema generator supported by all models provided by the Google LLMClient. + */ + override fun getBasicJsonSchemaGenerator(model: LLModel): BasicJsonSchemaGenerator { + return GoogleBasicJsonSchemaGenerator + } + /** * Gets a response from the Google AI API. * diff --git a/prompt/prompt-executor/prompt-executor-clients/prompt-executor-mistralai-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/mistralai/MistralAILLMClient.kt b/prompt/prompt-executor/prompt-executor-clients/prompt-executor-mistralai-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/mistralai/MistralAILLMClient.kt index dd3f74de2b..955016f0e9 100644 --- a/prompt/prompt-executor/prompt-executor-clients/prompt-executor-mistralai-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/mistralai/MistralAILLMClient.kt +++ b/prompt/prompt-executor/prompt-executor-clients/prompt-executor-mistralai-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/mistralai/MistralAILLMClient.kt @@ -76,11 +76,6 @@ public open class MistralAILLMClient( private companion object { private val staticLogger = KotlinLogging.logger { } - - init { - // On class load register custom OpenAI JSON schema generators for structured output. - registerOpenAIJsonSchemaGenerators(LLMProvider.DeepSeek) - } } /** diff --git a/prompt/prompt-executor/prompt-executor-clients/prompt-executor-openai-client-base/src/commonMain/kotlin/ai/koog/prompt/executor/clients/openai/base/AbstractOpenAILLMClient.kt b/prompt/prompt-executor/prompt-executor-clients/prompt-executor-openai-client-base/src/commonMain/kotlin/ai/koog/prompt/executor/clients/openai/base/AbstractOpenAILLMClient.kt index 3a367aaf8b..87e6cfba7a 100644 --- a/prompt/prompt-executor/prompt-executor-clients/prompt-executor-openai-client-base/src/commonMain/kotlin/ai/koog/prompt/executor/clients/openai/base/AbstractOpenAILLMClient.kt +++ b/prompt/prompt-executor/prompt-executor-clients/prompt-executor-openai-client-base/src/commonMain/kotlin/ai/koog/prompt/executor/clients/openai/base/AbstractOpenAILLMClient.kt @@ -24,7 +24,6 @@ import ai.koog.prompt.executor.clients.openai.base.models.OpenAIUsage import ai.koog.prompt.executor.clients.openai.base.structure.OpenAIBasicJsonSchemaGenerator import ai.koog.prompt.executor.clients.openai.base.structure.OpenAIStandardJsonSchemaGenerator import ai.koog.prompt.llm.LLMCapability -import ai.koog.prompt.llm.LLMProvider import ai.koog.prompt.llm.LLModel import ai.koog.prompt.message.AttachmentContent import ai.koog.prompt.message.ContentPart @@ -35,9 +34,8 @@ import ai.koog.prompt.params.LLMParams import ai.koog.prompt.streaming.StreamFrame import ai.koog.prompt.streaming.StreamFrameFlowBuilder import ai.koog.prompt.streaming.buildStreamFrameFlow -import ai.koog.prompt.structure.RegisteredBasicJsonSchemaGenerators -import ai.koog.prompt.structure.RegisteredStandardJsonSchemaGenerators -import ai.koog.prompt.structure.annotations.InternalStructuredOutputApi +import ai.koog.prompt.structure.json.generator.BasicJsonSchemaGenerator +import ai.koog.prompt.structure.json.generator.StandardJsonSchemaGenerator import io.github.oshai.kotlinlogging.KLogger import io.ktor.client.HttpClient import io.ktor.client.plugins.HttpTimeout @@ -94,19 +92,6 @@ public abstract class AbstractOpenAILLMClient { diff --git a/prompt/prompt-executor/prompt-executor-clients/prompt-executor-openrouter-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/openrouter/OpenRouterLLMClient.kt b/prompt/prompt-executor/prompt-executor-clients/prompt-executor-openrouter-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/openrouter/OpenRouterLLMClient.kt index daff3db38d..eafb749ce9 100644 --- a/prompt/prompt-executor/prompt-executor-clients/prompt-executor-openrouter-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/openrouter/OpenRouterLLMClient.kt +++ b/prompt/prompt-executor/prompt-executor-clients/prompt-executor-openrouter-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/openrouter/OpenRouterLLMClient.kt @@ -63,11 +63,6 @@ public class OpenRouterLLMClient( private companion object { private val staticLogger = KotlinLogging.logger { } - - init { - // On class load register custom OpenAI JSON schema generators for structured output. - registerOpenAIJsonSchemaGenerators(LLMProvider.OpenRouter) - } } /** diff --git a/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/LLMClient.kt b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/LLMClient.kt index 1795f77261..9187232f35 100644 --- a/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/LLMClient.kt +++ b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/LLMClient.kt @@ -8,6 +8,8 @@ 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.json.generator.BasicJsonSchemaGenerator +import ai.koog.prompt.structure.json.generator.StandardJsonSchemaGenerator import kotlinx.coroutines.flow.Flow /** @@ -84,6 +86,22 @@ public interface LLMClient : AutoCloseable { * @return The LLMProvider instance used for executing prompts and managing LLM operations. */ public fun llmProvider(): LLMProvider + + /** + * Standard JSON schema generator supported by the LLMClient. + * Return [StandardJsonSchemaGenerator] by default. + */ + public fun getStandardJsonSchemaGenerator(model: LLModel): StandardJsonSchemaGenerator { + return StandardJsonSchemaGenerator + } + + /** + * Basic JSON schema generator supported by the LLMClient. + * Return [BasicJsonSchemaGenerator] by default. + */ + public fun getBasicJsonSchemaGenerator(model: LLModel): BasicJsonSchemaGenerator { + return BasicJsonSchemaGenerator + } } /** diff --git a/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/LLMClientStructuredOutput.kt b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/LLMClientStructuredOutput.kt new file mode 100644 index 0000000000..6f71cef048 --- /dev/null +++ b/prompt/prompt-executor/prompt-executor-clients/src/commonMain/kotlin/ai/koog/prompt/executor/clients/LLMClientStructuredOutput.kt @@ -0,0 +1,94 @@ +import ai.koog.prompt.dsl.Prompt +import ai.koog.prompt.executor.clients.LLMClient +import ai.koog.prompt.llm.LLModel +import ai.koog.prompt.message.Message +import ai.koog.prompt.structure.StructuredOutputPrompts.appendStructuredOutputInstructions +import ai.koog.prompt.structure.StructuredRequest +import ai.koog.prompt.structure.StructuredResponse +import ai.koog.prompt.structure.buildStructuredRequest +import kotlinx.serialization.KSerializer +import kotlinx.serialization.serializer + +/** + * Executes a prompt with structured output, enhancing it with schema instructions or native structured output + * parameter, and parses the response into the defined structure. + * + * This is a simple version of the full `executeStructured`. 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. + * + * For example, it chooses which JSON schema to use (simple or full) and with which mode (native or manual). + * + * @param T The structure to request. + * @param prompt The prompt to be executed. + * @param model LLM to execute requests. + * @param examples Optional list of examples in case manual mode will be used. These examples might help the model to + * understand the format better. + * + * @return [kotlin.Result] with parsed [StructuredResponse] or error. + */ +public suspend inline fun LLMClient.executeStructured( + prompt: Prompt, + model: LLModel, + examples: List = emptyList(), +): Result> { + return executeStructured( + prompt = prompt, + model = model, + serializer = serializer(), + examples = examples, + ) +} + +public suspend fun LLMClient.executeStructured( + prompt: Prompt, + model: LLModel, + serializer: KSerializer, + examples: List = emptyList(), +): Result> { + val structuredOutput = buildStructuredRequest( + model = model, + serializer = serializer, + examples = examples, + standardJsonSchemaGenerator = getStandardJsonSchemaGenerator(model), + basicJsonSchemaGenerator = getBasicJsonSchemaGenerator(model), + ) + return executeStructured( + prompt = prompt, + model = model, + structuredRequest = structuredOutput, + ) +} + +/** + * Executes a prompt with structured output, enhancing it with schema instructions or native structured output + * parameter, and parses the response into the defined structure. + * + * **Note**: While many language models advertise support for structured output via JSON schema, + * the actual level of support varies between models and even between versions + * of the same model. Some models may produce malformed outputs or deviate from + * the schema in subtle ways, especially with complex structures like polymorphic types. + * + * @param prompt The prompt to be executed. + * @param model LLM to execute requests. + * @param structuredRequest The structured output to request. + * + * @return [kotlin.Result] with parsed [StructuredResponse] or error. + */ +public suspend fun LLMClient.executeStructured( + prompt: Prompt, + model: LLModel, + structuredRequest: StructuredRequest, +): Result> { + val updatedPrompt = appendStructuredOutputInstructions(prompt, structuredRequest) + val message = this.execute(prompt = updatedPrompt, model = model).single() + + return runCatching { + require(message is Message.Assistant) { "Response for structured output must be an assistant message, got ${message::class.simpleName} instead" } + StructuredResponse( + data = structuredRequest.structure.parse(message.content), + structure = structuredRequest.structure, + message = message + ) + } +} diff --git a/prompt/prompt-executor/prompt-executor-llms/src/commonMain/kotlin/ai/koog/prompt/executor/llms/MultiLLMPromptExecutor.kt b/prompt/prompt-executor/prompt-executor-llms/src/commonMain/kotlin/ai/koog/prompt/executor/llms/MultiLLMPromptExecutor.kt index 489a20091a..c10b92672c 100644 --- a/prompt/prompt-executor/prompt-executor-llms/src/commonMain/kotlin/ai/koog/prompt/executor/llms/MultiLLMPromptExecutor.kt +++ b/prompt/prompt-executor/prompt-executor-llms/src/commonMain/kotlin/ai/koog/prompt/executor/llms/MultiLLMPromptExecutor.kt @@ -10,8 +10,12 @@ 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.StructuredRequest +import ai.koog.prompt.structure.StructuredResponse +import executeStructured import io.github.oshai.kotlinlogging.KotlinLogging import kotlinx.coroutines.flow.Flow +import kotlinx.serialization.KSerializer /** * MultiLLMPromptExecutor is a class responsible for executing prompts @@ -203,6 +207,29 @@ public open class MultiLLMPromptExecutor( return choices } + override suspend fun executeStructured( + prompt: Prompt, + model: LLModel, + serializer: KSerializer, + examples: List + ): Result> { + val provider = model.provider + val client = llmClients[provider] ?: throw IllegalArgumentException("No client found for provider: $provider") + + return client.executeStructured(prompt, model, serializer, examples) + } + + override suspend fun executeStructured( + prompt: Prompt, + model: LLModel, + structuredRequest: StructuredRequest + ): Result> { + val provider = model.provider + val client = llmClients[provider] ?: throw IllegalArgumentException("No client found for provider: $provider") + + return client.executeStructured(prompt, model, structuredRequest) + } + /** * Moderates the provided multi-modal content using the specified model. * diff --git a/prompt/prompt-executor/prompt-executor-llms/src/commonMain/kotlin/ai/koog/prompt/executor/llms/SingleLLMPromptExecutor.kt b/prompt/prompt-executor/prompt-executor-llms/src/commonMain/kotlin/ai/koog/prompt/executor/llms/SingleLLMPromptExecutor.kt index d065c439aa..9fb53710fa 100644 --- a/prompt/prompt-executor/prompt-executor-llms/src/commonMain/kotlin/ai/koog/prompt/executor/llms/SingleLLMPromptExecutor.kt +++ b/prompt/prompt-executor/prompt-executor-llms/src/commonMain/kotlin/ai/koog/prompt/executor/llms/SingleLLMPromptExecutor.kt @@ -9,8 +9,12 @@ 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.StructuredRequest +import ai.koog.prompt.structure.StructuredResponse +import executeStructured import io.github.oshai.kotlinlogging.KotlinLogging import kotlinx.coroutines.flow.Flow +import kotlinx.serialization.KSerializer /** * Executes prompts using a direct client for communication with large language model (LLM) providers. @@ -58,6 +62,23 @@ public open class SingleLLMPromptExecutor( return choices } + override suspend fun executeStructured( + prompt: Prompt, + model: LLModel, + serializer: KSerializer, + examples: List + ): Result> { + return llmClient.executeStructured(prompt, model, serializer, examples) + } + + override suspend fun executeStructured( + prompt: Prompt, + model: LLModel, + structuredRequest: StructuredRequest + ): Result> { + return llmClient.executeStructured(prompt, model, structuredRequest) + } + override suspend fun moderate(prompt: Prompt, model: LLModel): ModerationResult = llmClient.moderate(prompt, model) override suspend fun models(): List = llmClient.models() diff --git a/prompt/prompt-executor/prompt-executor-model/build.gradle.kts b/prompt/prompt-executor/prompt-executor-model/build.gradle.kts index b8d953f1e4..9822d88408 100644 --- a/prompt/prompt-executor/prompt-executor-model/build.gradle.kts +++ b/prompt/prompt-executor/prompt-executor-model/build.gradle.kts @@ -13,6 +13,7 @@ kotlin { dependencies { api(project(":agents:agents-tools")) api(project(":prompt:prompt-model")) + api(project(":prompt:prompt-structure")) api(libs.kotlinx.coroutines.core) api(libs.oshai.kotlin.logging) } @@ -23,6 +24,14 @@ kotlin { api(libs.kotlinx.coroutines.jdk8) } } + + commonTest { + dependencies { + implementation(kotlin("test")) + implementation(libs.kotlinx.coroutines.test) + implementation(project(":agents:agents-test")) + } + } } explicitApi() diff --git a/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/LLMStructuredParsingError.kt b/prompt/prompt-executor/prompt-executor-model/src/commonMain/kotlin/ai/koog/prompt/executor/model/LLMStructuredParsingError.kt similarity index 82% rename from prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/LLMStructuredParsingError.kt rename to prompt/prompt-executor/prompt-executor-model/src/commonMain/kotlin/ai/koog/prompt/executor/model/LLMStructuredParsingError.kt index bccb2f1833..dc20565b5b 100644 --- a/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/LLMStructuredParsingError.kt +++ b/prompt/prompt-executor/prompt-executor-model/src/commonMain/kotlin/ai/koog/prompt/executor/model/LLMStructuredParsingError.kt @@ -1,4 +1,4 @@ -package ai.koog.prompt.structure +package ai.koog.prompt.executor.model import kotlinx.serialization.SerializationException @@ -10,6 +10,6 @@ import kotlinx.serialization.SerializationException * * @constructor Creates a new instance of [LLMStructuredParsingError]. * @param message A detailed message describing the cause of the parsing error. - * @param cause [SerializationException] that caused parsing exception. + * @param cause [kotlinx.serialization.SerializationException] that caused parsing exception. */ public class LLMStructuredParsingError(message: String, cause: SerializationException?) : Exception(message, cause) diff --git a/prompt/prompt-executor/prompt-executor-model/src/commonMain/kotlin/ai/koog/prompt/executor/model/PromptExecutor.kt b/prompt/prompt-executor/prompt-executor-model/src/commonMain/kotlin/ai/koog/prompt/executor/model/PromptExecutor.kt index 2ea8acfdd4..4d73730565 100644 --- a/prompt/prompt-executor/prompt-executor-model/src/commonMain/kotlin/ai/koog/prompt/executor/model/PromptExecutor.kt +++ b/prompt/prompt-executor/prompt-executor-model/src/commonMain/kotlin/ai/koog/prompt/executor/model/PromptExecutor.kt @@ -7,7 +7,10 @@ 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.StructuredRequest +import ai.koog.prompt.structure.StructuredResponse import kotlinx.coroutines.flow.Flow +import kotlinx.serialization.KSerializer /** * An interface representing an executor for processing language model prompts. @@ -64,6 +67,46 @@ public interface PromptExecutor : AutoCloseable { ): List = listOf(execute(prompt, model, tools)) + /** + * Executes structured output generation for a given prompt and model. + * + * This method generates structured output based on the provided prompt and model. + * It uses the specified structured output configuration to define the expected output format. + * + * @param prompt The prompt containing input messages and parameters to guide the language model execution. + * @param model The language model to be used for processing the prompt. + * + * @return The generated structured output of type T. + */ + public suspend fun executeStructured( + prompt: Prompt, + model: LLModel, + structuredRequest: StructuredRequest + ): Result> { + throw UnsupportedOperationException("Not implemented for this executor") + } + + /** + * Executes structured output generation for a given prompt and model. + * + * This method generates structured output based on the provided prompt and model. + * It uses the specified structured output configuration to define the expected output format. + * + * @param prompt The prompt containing input messages and parameters to guide the language model execution. + * @param model The language model to be used for processing the prompt. + * @param serializer The serializer for converting the structured output to the expected type. + * + * @return The generated structured output of type T. + */ + public suspend fun executeStructured( + prompt: Prompt, + model: LLModel, + serializer: KSerializer, + examples: List = emptyList(), + ): Result> { + throw UnsupportedOperationException("Not implemented for this executor") + } + /** * Moderates the content of a given message with attachments using a specified language model. * diff --git a/prompt/prompt-executor/prompt-executor-model/src/commonMain/kotlin/ai/koog/prompt/executor/model/PromptExecutorStructuredOutput.kt b/prompt/prompt-executor/prompt-executor-model/src/commonMain/kotlin/ai/koog/prompt/executor/model/PromptExecutorStructuredOutput.kt new file mode 100644 index 0000000000..e9cd08e6f5 --- /dev/null +++ b/prompt/prompt-executor/prompt-executor-model/src/commonMain/kotlin/ai/koog/prompt/executor/model/PromptExecutorStructuredOutput.kt @@ -0,0 +1,144 @@ +package ai.koog.prompt.executor.model + +import ai.koog.prompt.dsl.Prompt +import ai.koog.prompt.llm.LLModel +import ai.koog.prompt.message.Message +import ai.koog.prompt.structure.StructuredRequestConfig +import ai.koog.prompt.structure.StructuredResponse +import ai.koog.prompt.structure.annotations.InternalStructuredOutputApi +import kotlinx.serialization.KSerializer +import kotlinx.serialization.serializer + +/** + * Executes a prompt with structured output, enhancing it with schema instructions or native structured output + * parameter, and parses the response into the defined structure. + * + * **Note**: While many language models advertise support for structured output via JSON schema, + * the actual level of support varies between models and even between versions + * of the same model. Some models may produce malformed outputs or deviate from + * the schema in subtle ways, especially with complex structures like polymorphic types. + * + * @param prompt The prompt to be executed. + * @param model LLM to execute requests. + * @param config A configuration defining structures and behavior. + * + * @return [kotlin.Result] with parsed [StructuredResponse] or error. + */ +public suspend fun PromptExecutor.executeStructured( + prompt: Prompt, + model: LLModel, + config: StructuredRequestConfig, + fixingParser: StructureFixingParser? = null, +): Result> { + TODO("Implement fixing parser yet implemented") + return executeStructured( + prompt = prompt, + model = model, + structuredRequest = config.structuredRequest(model), + ) +} + +/** + * Executes a prompt with structured output, enhancing it with schema instructions or native structured output + * parameter, and parses the response into the defined structure. + * + * This is a simple version of the full `executeStructured`. 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. + * + * For example, it chooses which JSON schema to use (simple or full) and with which mode (native or manual). + * + * @param T The structure to request. + * @param prompt The prompt to be executed. + * @param model LLM to execute requests. + * @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. + * + * @return [kotlin.Result] with parsed [StructuredResponse] or error. + */ +public suspend inline fun PromptExecutor.executeStructured( + prompt: Prompt, + model: LLModel, + examples: List = emptyList(), + fixingParser: StructureFixingParser? = null, +): Result> { + return executeStructured( + prompt = prompt, + model = model, + serializer = serializer(), + examples = examples, + fixingParser = fixingParser + ) +} + +/** + * Executes a prompt with structured output, enhancing it with schema instructions or native structured output + * parameter, and parses the response into the defined structure. + * + * This is a simple version of the full `executeStructured`. 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. + * + * For example, it chooses which JSON schema to use (simple or full) and with which mode (native or manual). + * + * @param prompt The prompt to be executed. + * @param model LLM to execute requests. + * @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. + * + * @return [kotlin.Result] with parsed [StructuredResponse] or error. + */ +@OptIn(InternalStructuredOutputApi::class) +public suspend fun PromptExecutor.executeStructured( + prompt: Prompt, + model: LLModel, + serializer: KSerializer, + examples: List = emptyList(), + fixingParser: StructureFixingParser? = null, +): Result> { + TODO("Implement fixing parser yet implemented") + return executeStructured( + prompt = prompt, + model = model, + serializer = serializer, + examples = examples + ) +} + +/** + * Parses a structured response from the assistant message using the provided structured output configuration + * and language model. If a fixing parser is specified in the configuration, it will be used; otherwise, + * the structure will be parsed directly. + * + * @param T The type of the structured output. + * @param response The assistant's response message to be parsed. + * @param config The structured output configuration defining how the response should be parsed. + * @param model The language model to be used for parsing the structured output. + * @return A `StructuredResponse` containing the parsed structure and the original assistant message. + */ +public suspend fun PromptExecutor.parseResponseToStructuredResponse( + response: Message.Assistant, + config: StructuredRequestConfig, + model: LLModel, + fixingParser: StructureFixingParser? = null +): StructuredResponse { + // Use fixingParser if provided, otherwise parse directly + TODO("Not yet implemented") +// val structure = config.structure(model) +// val structureResponse = config.fixingParser +// ?.parse(this, structure, response.content) +// ?: structure.parse(response.content) +// +// return StructuredResponse( +// data = structureResponse, +// structure = structure, +// message = response +// ) +} diff --git a/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/StructureFixingParser.kt b/prompt/prompt-executor/prompt-executor-model/src/commonMain/kotlin/ai/koog/prompt/executor/model/StructureFixingParser.kt similarity index 98% rename from prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/StructureFixingParser.kt rename to prompt/prompt-executor/prompt-executor-model/src/commonMain/kotlin/ai/koog/prompt/executor/model/StructureFixingParser.kt index 5b32ef2b95..ee17b48b5e 100644 --- a/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/StructureFixingParser.kt +++ b/prompt/prompt-executor/prompt-executor-model/src/commonMain/kotlin/ai/koog/prompt/executor/model/StructureFixingParser.kt @@ -1,12 +1,12 @@ -package ai.koog.prompt.structure +package ai.koog.prompt.executor.model import ai.koog.prompt.dsl.PromptBuilder import ai.koog.prompt.dsl.prompt -import ai.koog.prompt.executor.model.PromptExecutor import ai.koog.prompt.llm.LLModel import ai.koog.prompt.markdown.markdown import ai.koog.prompt.message.Message import ai.koog.prompt.params.LLMParams +import ai.koog.prompt.structure.Structure import io.github.oshai.kotlinlogging.KotlinLogging import kotlinx.serialization.SerializationException diff --git a/prompt/prompt-structure/src/commonTest/kotlin/ai/koog/prompt/structure/StructureFixingParserTest.kt b/prompt/prompt-executor/prompt-executor-model/src/commonTest/kotlin/ai/koog/prompt/executor/model/StructureFixingParserTest.kt similarity index 87% rename from prompt/prompt-structure/src/commonTest/kotlin/ai/koog/prompt/structure/StructureFixingParserTest.kt rename to prompt/prompt-executor/prompt-executor-model/src/commonTest/kotlin/ai/koog/prompt/executor/model/StructureFixingParserTest.kt index 21bd45bf41..58600bd823 100644 --- a/prompt/prompt-structure/src/commonTest/kotlin/ai/koog/prompt/structure/StructureFixingParserTest.kt +++ b/prompt/prompt-executor/prompt-executor-model/src/commonTest/kotlin/ai/koog/prompt/executor/model/StructureFixingParserTest.kt @@ -1,7 +1,9 @@ -package ai.koog.prompt.structure +package ai.koog.prompt.executor.model.ai.koog.prompt.executor.model import ai.koog.agents.testing.tools.getMockExecutor import ai.koog.prompt.executor.clients.openai.OpenAIModels +import ai.koog.prompt.executor.model.LLMStructuredParsingError +import ai.koog.prompt.executor.model.StructureFixingParser import ai.koog.prompt.structure.json.JsonStructure import kotlinx.coroutines.test.runTest import kotlinx.serialization.Serializable @@ -18,8 +20,8 @@ class StructureFixingParserTest { ) private val testData = TestData("test", 42) - private val testDataJson = Json.encodeToString(testData) - private val testStructure = JsonStructure.create() + private val testDataJson = Json.Default.encodeToString(testData) + private val testStructure = JsonStructure.Companion.create() @Test fun testParseValidContentWithoutFixing() = runTest { diff --git a/prompt/prompt-structure/build.gradle.kts b/prompt/prompt-structure/build.gradle.kts index 0be9a8b8a5..56e83dda54 100644 --- a/prompt/prompt-structure/build.gradle.kts +++ b/prompt/prompt-structure/build.gradle.kts @@ -12,9 +12,8 @@ kotlin { sourceSets { commonMain { dependencies { - api(project(":prompt:prompt-executor:prompt-executor-model")) - api(project(":prompt:prompt-markdown")) + api(project(":agents:agents-tools")) api(libs.kotlinx.coroutines.core) api(libs.kotlinx.serialization.json) implementation(libs.oshai.kotlin.logging) diff --git a/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/PromptExecutorExtensions.kt b/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/PromptExecutorExtensions.kt deleted file mode 100644 index 9f4b0ec689..0000000000 --- a/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/PromptExecutorExtensions.kt +++ /dev/null @@ -1,334 +0,0 @@ -package ai.koog.prompt.structure - -import ai.koog.prompt.dsl.Prompt -import ai.koog.prompt.dsl.prompt -import ai.koog.prompt.executor.model.PromptExecutor -import ai.koog.prompt.llm.LLMCapability -import ai.koog.prompt.llm.LLMProvider -import ai.koog.prompt.llm.LLModel -import ai.koog.prompt.markdown.markdown -import ai.koog.prompt.message.Message -import ai.koog.prompt.structure.annotations.InternalStructuredOutputApi -import ai.koog.prompt.structure.json.JsonStructure -import ai.koog.prompt.structure.json.generator.BasicJsonSchemaGenerator -import ai.koog.prompt.structure.json.generator.JsonSchemaGenerator -import ai.koog.prompt.structure.json.generator.StandardJsonSchemaGenerator -import kotlinx.serialization.KSerializer -import kotlinx.serialization.serializer -import kotlin.reflect.KType - -/** - * Represents a container for structured data parsed from a response message. - * - * This class is designed to encapsulate both the parsed structured output and the original raw - * text as returned from a processing step, such as a language model execution. - * - * @param T The type of the structured data contained within this response. - * @property data The parsed structured data corresponding to the specific schema. - * @property structure The structure used for the response. - * @property message The original assistant message from which the structure was parsed. - */ -public data class StructuredResponse( - val data: T, - val structure: Structure, - val message: Message.Assistant -) - -/** - * Configures structured output behavior. - * Defines which structures in which modes should be used for each provider when requesting a structured output. - * - * @property default Fallback [StructuredRequest] to be used when there's no suitable structure found in [byProvider] - * for a requested [LLMProvider]. Defaults to `null`, meaning structured output would fail with error in such a case. - * - * @property byProvider A map matching [LLMProvider] to compatible [StructuredRequest] definitions. Each provider may - * require different schema formats. E.g. for [JsonStructure] this means you have to use the appropriate - * [JsonSchemaGenerator] implementation for each provider for [StructuredRequest.Native], or fallback to [StructuredRequest.Manual] - * - * @property 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 data class StructuredRequestConfig( - public val default: StructuredRequest? = null, - public val byProvider: Map> = emptyMap(), - public val fixingParser: StructureFixingParser? = null -) { - /** - * Updates a given prompt to configure structured output using the specified large language model (LLM). - * Depending on the model's support for structured outputs, the prompt is updated either manually or natively. - * - * @param model The large language model (LLModel) used to determine the structured output configuration. - * @param prompt The original prompt to be updated with the structured output configuration. - * @return A new prompt reflecting the updated structured output configuration. - */ - public fun updatePrompt(model: LLModel, prompt: Prompt): Prompt { - return when (val mode = structuredRequest(model)) { - // Don't set schema parameter in prompt and coerce the model manually with user message to provide a structured response. - is StructuredRequest.Manual -> { - prompt(prompt) { - user( - markdown { - StructuredOutputPrompts.outputInstructionPrompt(this, mode.structure) - } - ) - } - } - - // Rely on built-in model capabilities to provide structured response. - is StructuredRequest.Native -> { - prompt.withUpdatedParams { schema = mode.structure.schema } - } - } - } - - /** - * Retrieves the structured data configuration for a specific large language model (LLM). - * - * The method determines the appropriate structured data setup based on the given LLM - * instance, ensuring compatibility with the model's provider and capabilities. - * - * @param model The large language model (LLM) instance used to identify the structured data configuration. - * @return The structured data configuration represented as a `StructuredData` instance. - */ - public fun structure(model: LLModel): Structure { - return structuredRequest(model).structure - } - - /** - * Retrieves the structured output configuration for a specific large language model (LLM). - * - * The method determines the appropriate structured output instance based on the model's provider. - * If no specific configuration is found for the provider, it falls back to the default configuration. - * Throws an exception if no default configuration is available. - * - * @param model The large language model (LLM) used to identify the structured output configuration. - * @return An instance of `StructuredOutput` that represents the structured output configuration for the model. - * @throws IllegalArgumentException if no configuration is found for the provider and no default configuration is set. - */ - private fun structuredRequest(model: LLModel): StructuredRequest { - return byProvider[model.provider] - ?: default - ?: throw IllegalArgumentException("No structure found for provider ${model.provider}") - } -} - -/** - * Defines how structured outputs should be generated. - * - * Can be [StructuredRequest.Manual] or [StructuredRequest.Native] - * - * @param T The type of structured data. - */ -public sealed interface StructuredRequest { - /** - * The definition of a structure. - */ - public val structure: Structure - - /** - * Instructs the model to produce structured output through explicit prompting. - * - * Uses an additional user message containing [Structure.definition] to guide - * the model in generating correctly formatted responses. - * - * @property structure The structure definition to be used in output generation. - */ - public data class Manual(override val structure: Structure) : StructuredRequest - - /** - * Leverages a model's built-in structured output capabilities. - * - * Uses [Structure.schema] to define the expected response format through the model's - * native structured output functionality. - * - * Note: [Structure.examples] are not used with this mode, only the schema is sent via parameters. - * - * @property structure The structure definition to be used in output generation. - */ - public data class Native(override val structure: Structure) : StructuredRequest -} - -/** - * Executes a prompt with structured output, enhancing it with schema instructions or native structured output - * parameter, and parses the response into the defined structure. - * - * **Note**: While many language models advertise support for structured output via JSON schema, - * the actual level of support varies between models and even between versions - * of the same model. Some models may produce malformed outputs or deviate from - * the schema in subtle ways, especially with complex structures like polymorphic types. - * - * @param prompt The prompt to be executed. - * @param model LLM to execute requests. - * @param config A configuration defining structures and behavior. - * - * @return [kotlin.Result] with parsed [StructuredResponse] or error. - */ -public suspend fun PromptExecutor.executeStructured( - prompt: Prompt, - model: LLModel, - config: StructuredRequestConfig, -): Result> { - val updatedPrompt = config.updatePrompt(model, prompt) - val response = this.execute(prompt = updatedPrompt, model = model).single() - - return runCatching { - require(response is Message.Assistant) { "Response for structured output must be an assistant message, got ${response::class.simpleName} instead" } - - parseResponseToStructuredResponse(response, config, model) - } -} - -/** - * Registered mapping of providers to their respective known simple JSON schema format generators. - * The registration is supposed to be done by the LLM clients when they are loaded, to communicate their custom formats. - * - * Used to attempt to get a proper generator implicitly in the simple version of [executeStructured] (that does not accept [StructuredRequest] explicitly) - * to attempt to generate an appropriate schema for the passed [KType]. - */ -@InternalStructuredOutputApi -public val RegisteredBasicJsonSchemaGenerators: MutableMap = mutableMapOf() - -/** - * Registered mapping of providers to their respective known full JSON schema format generators. - * The registration is supposed to be done by the LLM clients on their initialization, to communicate their custom formats. - * - * Used to attempt to get a proper generator implicitly in the simple version of [executeStructured] (that does not accept [StructuredRequest] explicitly) - * to attempt to generate an appropriate schema for the passed [KType]. - */ -@InternalStructuredOutputApi -public val RegisteredStandardJsonSchemaGenerators: MutableMap = mutableMapOf() - -/** - * Executes a prompt with structured output, enhancing it with schema instructions or native structured output - * parameter, and parses the response into the defined structure. - * - * This is a simple version of the full `executeStructured`. 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. - * - * For example, it chooses which JSON schema to use (simple or full) and with which mode (native or manual). - * - * @param prompt The prompt to be executed. - * @param model LLM to execute requests. - * @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. - * - * @return [kotlin.Result] with parsed [StructuredResponse] or error. - */ -@OptIn(InternalStructuredOutputApi::class) -public suspend fun PromptExecutor.executeStructured( - prompt: Prompt, - model: LLModel, - serializer: KSerializer, - examples: List = emptyList(), - fixingParser: StructureFixingParser? = null, -): Result> { - @Suppress("UNCHECKED_CAST") - val id = serializer.descriptor.serialName.substringAfterLast(".") - - val structuredRequest = when { - LLMCapability.Schema.JSON.Standard in model.capabilities -> StructuredRequest.Native( - JsonStructure.create( - id = id, - serializer = serializer, - schemaGenerator = RegisteredStandardJsonSchemaGenerators[model.provider] ?: StandardJsonSchemaGenerator - ) - ) - - LLMCapability.Schema.JSON.Basic in model.capabilities -> StructuredRequest.Native( - JsonStructure.create( - id = id, - serializer = serializer, - schemaGenerator = RegisteredBasicJsonSchemaGenerators[model.provider] ?: BasicJsonSchemaGenerator - ) - ) - - else -> StructuredRequest.Manual( - JsonStructure.create( - id = id, - serializer = serializer, - schemaGenerator = StandardJsonSchemaGenerator, - examples = examples, - ) - ) - } - - return executeStructured( - prompt = prompt, - model = model, - config = StructuredRequestConfig( - default = structuredRequest, - fixingParser = fixingParser, - ) - ) -} - -/** - * Executes a prompt with structured output, enhancing it with schema instructions or native structured output - * parameter, and parses the response into the defined structure. - * - * This is a simple version of the full `executeStructured`. 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. - * - * For example, it chooses which JSON schema to use (simple or full) and with which mode (native or manual). - * - * @param T The structure to request. - * @param prompt The prompt to be executed. - * @param model LLM to execute requests. - * @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. - * - * @return [kotlin.Result] with parsed [StructuredResponse] or error. - */ -public suspend inline fun PromptExecutor.executeStructured( - prompt: Prompt, - model: LLModel, - examples: List = emptyList(), - fixingParser: StructureFixingParser? = null, -): Result> { - return executeStructured( - prompt = prompt, - model = model, - serializer = serializer(), - examples = examples, - fixingParser = fixingParser, - ) -} - -/** - * Parses a structured response from the assistant message using the provided structured output configuration - * and language model. If a fixing parser is specified in the configuration, it will be used; otherwise, - * the structure will be parsed directly. - * - * @param T The type of the structured output. - * @param response The assistant's response message to be parsed. - * @param config The structured output configuration defining how the response should be parsed. - * @param model The language model to be used for parsing the structured output. - * @return A `StructuredResponse` containing the parsed structure and the original assistant message. - */ -public suspend fun PromptExecutor.parseResponseToStructuredResponse( - response: Message.Assistant, - config: StructuredRequestConfig, - model: LLModel -): StructuredResponse { - // Use fixingParser if provided, otherwise parse directly - val structure = config.structure(model) - val structureResponse = config.fixingParser - ?.parse(this, structure, response.content) - ?: structure.parse(response.content) - - return StructuredResponse( - data = structureResponse, - structure = structure, - message = response - ) -} diff --git a/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/StructuredOutput.kt b/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/StructuredOutput.kt new file mode 100644 index 0000000000..c6898a2aa9 --- /dev/null +++ b/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/StructuredOutput.kt @@ -0,0 +1,163 @@ +package ai.koog.prompt.structure + +import ai.koog.prompt.llm.LLMCapability +import ai.koog.prompt.llm.LLMProvider +import ai.koog.prompt.llm.LLModel +import ai.koog.prompt.message.Message +import ai.koog.prompt.structure.json.JsonStructure +import ai.koog.prompt.structure.json.generator.BasicJsonSchemaGenerator +import ai.koog.prompt.structure.json.generator.JsonSchemaGenerator +import ai.koog.prompt.structure.json.generator.StandardJsonSchemaGenerator +import kotlinx.serialization.KSerializer + +/** + * Represents a container for structured data parsed from a response message. + * + * This class is designed to encapsulate both the parsed structured output and the original raw + * text as returned from a processing step, such as a language model execution. + * + * @param T The type of the structured data contained within this response. + * @property data The parsed structured data corresponding to the specific schema. + * @property structure The structure used for the response. + * @property message The original assistant message from which the structure was parsed. + */ +public data class StructuredResponse( + val data: T, + val structure: Structure, + val message: Message.Assistant +) + +/** + * Configures structured output behavior. + * Defines which structures in which modes should be used for each provider when requesting a structured output. + * + * @property default Fallback [StructuredRequest] to be used when there's no suitable structure found in [byProvider] + * for a requested [LLMProvider]. Defaults to `null`, meaning structured output would fail with error in such a case. + * + * @property byProvider A map matching [LLMProvider] to compatible [StructuredRequest] definitions. Each provider may + * require different schema formats. E.g. for [JsonStructure] this means you have to use the appropriate + * [JsonSchemaGenerator] implementation for each provider for [StructuredRequest.Native], or fallback to [StructuredRequest.Manual] + * + * intelligently fix parsing errors. When specified, parsing errors trigger additional + * LLM calls with error context to attempt correction of the structure format. + */ +public data class StructuredRequestConfig( + public val default: StructuredRequest? = null, + public val byProvider: Map> = emptyMap(), +) { + /** + * Retrieves the structured data configuration for a specific large language model (LLM). + * + * The method determines the appropriate structured data setup based on the given LLM + * instance, ensuring compatibility with the model's provider and capabilities. + * + * @param model The large language model (LLM) instance used to identify the structured data configuration. + * @return The structured data configuration represented as a `StructuredData` instance. + */ + public fun structure(model: LLModel): Structure { + return structuredRequest(model).structure + } + + /** + * Retrieves the structured output configuration for a specific large language model (LLM). + * + * The method determines the appropriate structured output instance based on the model's provider. + * If no specific configuration is found for the provider, it falls back to the default configuration. + * Throws an exception if no default configuration is available. + * + * @param model The large language model (LLM) used to identify the structured output configuration. + * @return An instance of `StructuredOutput` that represents the structured output configuration for the model. + * @throws IllegalArgumentException if no configuration is found for the provider and no default configuration is set. + */ + public fun structuredRequest(model: LLModel): StructuredRequest { + return byProvider[model.provider] + ?: default + ?: throw IllegalArgumentException("No structure found for provider ${model.provider}") + } +} + +/** + * Defines how structured outputs should be generated. + * + * Can be [StructuredRequest.Manual] or [StructuredRequest.Native] + * + * @param T The type of structured data. + */ +public sealed interface StructuredRequest { + /** + * The definition of a structure. + */ + public val structure: Structure + + /** + * Instructs the model to produce structured output through explicit prompting. + * + * Uses an additional user message containing [Structure.definition] to guide + * the model in generating correctly formatted responses. + * + * @property structure The structure definition to be used in output generation. + */ + public data class Manual(override val structure: Structure) : StructuredRequest + + /** + * Leverages a model's built-in structured output capabilities. + * + * Uses [Structure.schema] to define the expected response format through the model's + * native structured output functionality. + * + * Note: [Structure.examples] are not used with this mode, only the schema is sent via parameters. + * + * @property structure The structure definition to be used in output generation. + */ + public data class Native(override val structure: Structure) : StructuredRequest +} + +/** + * Defines how structured outputs should be generated. + * + * @param model The model to be used for generating structured outputs. + * @param serializer The serializer for the structured data type. + * @param examples Optional list of examples to be used for schema generation. + * @param standardJsonSchemaGenerator The generator for standard JSON schemas. + * @param basicJsonSchemaGenerator The generator for basic JSON schemas. + * @return A StructuredRequest instance representing the selected structured output mode. + */ +public fun buildStructuredRequest( + model: LLModel, + serializer: KSerializer, + examples: List = emptyList(), + standardJsonSchemaGenerator: StandardJsonSchemaGenerator, + basicJsonSchemaGenerator: BasicJsonSchemaGenerator, +): StructuredRequest { + @Suppress("UNCHECKED_CAST") + val id = serializer.descriptor.serialName.substringAfterLast(".") + + val structuredRequest = when { + LLMCapability.Schema.JSON.Standard in model.capabilities -> StructuredRequest.Native( + JsonStructure.create( + id = id, + serializer = serializer, + schemaGenerator = standardJsonSchemaGenerator + ) + ) + + LLMCapability.Schema.JSON.Basic in model.capabilities -> StructuredRequest.Native( + JsonStructure.create( + id = id, + serializer = serializer, + schemaGenerator = basicJsonSchemaGenerator + ) + ) + + else -> StructuredRequest.Manual( + JsonStructure.create( + id = id, + serializer = serializer, + schemaGenerator = StandardJsonSchemaGenerator, + examples = examples, + ) + ) + } + + return structuredRequest +} diff --git a/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/StructuredPrompts.kt b/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/StructuredPrompts.kt index 1bb5f837ab..28e96db4f4 100644 --- a/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/StructuredPrompts.kt +++ b/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/StructuredPrompts.kt @@ -1,24 +1,45 @@ package ai.koog.prompt.structure +import ai.koog.prompt.dsl.Prompt +import ai.koog.prompt.dsl.prompt import ai.koog.prompt.markdown.markdown -import ai.koog.prompt.text.TextContentBuilderBase /** * An object that provides utilities for formatting structured output prompts. */ public object StructuredOutputPrompts { + /** - * Formats and appends the structured data output to the provided MarkdownContentBuilder. + * Updates a given prompt to configure structured output using the specified large language model (LLM). + * Depending on the model's support for structured outputs, the prompt is updated either manually or natively. * - * @param structure The StructuredData instance containing the format ID and definition for the output. + * @param prompt The original prompt to be updated with the structured output configuration. + * @param structuredRequest The structured output configuration to be applied to the prompt. + * @return A new prompt reflecting the updated structured output configuration. */ - public fun outputInstructionPrompt(builder: TextContentBuilderBase<*>, structure: Structure<*, *>): TextContentBuilderBase<*> = builder.apply { - markdown { - h2("NEXT MESSAGE OUTPUT FORMAT") - +"The output in the next message MUST ADHERE TO ${structure.id} format." - br() + public fun appendStructuredOutputInstructions(prompt: Prompt, structuredRequest: StructuredRequest<*>): Prompt { + return when (structuredRequest) { + // Don't set schema parameter in prompt and coerce the model manually with user message to provide a structured response. + is StructuredRequest.Manual -> { + prompt(prompt) { + user( + markdown { + markdown { + h2("NEXT MESSAGE OUTPUT FORMAT") + +"The output in the next message MUST ADHERE TO ${structuredRequest.structure.id} format." + br() + + structuredRequest.structure.definition(this) + } + } + ) + } + } - structure.definition(this) + // Rely on built-in model capabilities to provide structured response. + is StructuredRequest.Native -> { + prompt.withUpdatedParams { schema = structuredRequest.structure.schema } + } } } }