Skip to content

Commit 99b8819

Browse files
EugeneTheDevEugeneTheDev
committed
Refactor structured output API to make it more flexible, support native structured output for OpenAI and Google
1 parent c730582 commit 99b8819

File tree

54 files changed

+3510
-1620
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

54 files changed

+3510
-1620
lines changed

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

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,10 @@ import ai.koog.agents.core.tools.ToolDescriptor
1313
import ai.koog.agents.core.tools.annotations.LLMDescription
1414
import ai.koog.prompt.llm.LLModel
1515
import ai.koog.prompt.params.LLMParams
16-
import ai.koog.prompt.structure.json.JsonSchemaGenerator
16+
import ai.koog.prompt.structure.StructureFixingParser
17+
import ai.koog.prompt.structure.StructuredOutput
18+
import ai.koog.prompt.structure.StructuredOutputConfig
19+
import ai.koog.prompt.structure.json.generator.FullJsonSchemaGenerator
1720
import ai.koog.prompt.structure.json.JsonStructuredData
1821
import io.github.oshai.kotlinlogging.KotlinLogging
1922
import kotlinx.serialization.Serializable
@@ -99,11 +102,15 @@ public open class AIAgentSubgraph<Input, Output>(
99102
}
100103

101104
val selectedTools = this.requestLLMStructured(
102-
structure = JsonStructuredData.createJsonStructure<SelectedTools>(
103-
schemaFormat = JsonSchemaGenerator.SchemaFormat.JsonSchema,
104-
examples = listOf(SelectedTools(listOf()), SelectedTools(tools.map { it.name }.take(3))),
105-
),
106-
retries = toolSelectionStrategy.maxRetries,
105+
config = StructuredOutputConfig(
106+
default = StructuredOutput.Manual(
107+
JsonStructuredData.createJsonStructure<SelectedTools>(
108+
schemaGenerator = FullJsonSchemaGenerator,
109+
examples = listOf(SelectedTools(listOf()), SelectedTools(tools.map { it.name }.take(3))),
110+
),
111+
),
112+
fixingParser = toolSelectionStrategy.fixingParser,
113+
)
107114
).getOrThrow()
108115

109116
prompt = initialPrompt
@@ -262,8 +269,9 @@ public sealed interface ToolSelectionStrategy {
262269
* This ensures that unnecessary tools are excluded, optimizing the toolset for the specific use case.
263270
*
264271
* @property subtaskDescription A description of the subtask for which the relevant tools should be selected.
272+
* @property fixingParser Optional [StructureFixingParser] to attempt fixes when malformed structured response with tool list is received.
265273
*/
266-
public data class AutoSelectForTask(val subtaskDescription: String, val maxRetries: Int = 3) : ToolSelectionStrategy
274+
public data class AutoSelectForTask(val subtaskDescription: String, val fixingParser: StructureFixingParser? = null) : ToolSelectionStrategy
267275

268276
/**
269277
* 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

Lines changed: 11 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -6,16 +6,14 @@ import ai.koog.agents.core.tools.ToolDescriptor
66
import ai.koog.agents.core.utils.ActiveProperty
77
import ai.koog.prompt.dsl.ModerationResult
88
import ai.koog.prompt.dsl.Prompt
9-
import ai.koog.prompt.executor.clients.openai.OpenAIModels
109
import ai.koog.prompt.executor.model.LLMChoice
1110
import ai.koog.prompt.executor.model.PromptExecutor
1211
import ai.koog.prompt.llm.LLModel
1312
import ai.koog.prompt.message.Message
1413
import ai.koog.prompt.params.LLMParams
15-
import ai.koog.prompt.structure.StructuredData
14+
import ai.koog.prompt.structure.StructuredOutputConfig
1615
import ai.koog.prompt.structure.StructuredResponse
1716
import ai.koog.prompt.structure.executeStructured
18-
import ai.koog.prompt.structure.executeStructuredOneShot
1917

2018
/**
2119
* Represents a session for an AI agent that interacts with an LLM (Language Learning Model).
@@ -240,30 +238,24 @@ public sealed class AIAgentLLMSession(
240238
}
241239

242240
/**
243-
* Coerce LLM to provide a structured output.
241+
* Sends a request to LLM and gets a structured response.
242+
*
243+
* @param config A configuration defining structures and behavior.
244244
*
245245
* @see [executeStructured]
246246
*/
247247
public open suspend fun <T> requestLLMStructured(
248-
structure: StructuredData<T>,
249-
retries: Int = 1,
250-
fixingModel: LLModel = OpenAIModels.Chat.GPT4o
248+
config: StructuredOutputConfig<T>,
251249
): Result<StructuredResponse<T>> {
252250
validateSession()
253-
val preparedPrompt = preparePrompt(prompt, tools = emptyList())
254-
return executor.executeStructured(preparedPrompt, model, structure, retries, fixingModel)
255-
}
256251

257-
/**
258-
* Expect LLM to reply in a structured format and try to parse it.
259-
* For more robust version with model coercion and correction see [requestLLMStructured]
260-
*
261-
* @see [executeStructuredOneShot]
262-
*/
263-
public open suspend fun <T> requestLLMStructuredOneShot(structure: StructuredData<T>): StructuredResponse<T> {
264-
validateSession()
265252
val preparedPrompt = preparePrompt(prompt, tools = emptyList())
266-
return executor.executeStructuredOneShot(preparedPrompt, model, structure)
253+
254+
return executor.executeStructured(
255+
prompt = preparedPrompt,
256+
model = model,
257+
config = config,
258+
)
267259
}
268260

269261
/**

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

Lines changed: 10 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,7 @@ package ai.koog.agents.core.agent.session
33
import ai.koog.agents.core.agent.config.AIAgentConfigBase
44
import ai.koog.agents.core.environment.AIAgentEnvironment
55
import ai.koog.agents.core.environment.SafeTool
6-
import ai.koog.agents.core.tools.Tool
7-
import ai.koog.agents.core.tools.ToolArgs
8-
import ai.koog.agents.core.tools.ToolDescriptor
9-
import ai.koog.agents.core.tools.ToolRegistry
10-
import ai.koog.agents.core.tools.ToolResult
6+
import ai.koog.agents.core.tools.*
117
import ai.koog.agents.core.utils.ActiveProperty
128
import ai.koog.prompt.dsl.Prompt
139
import ai.koog.prompt.dsl.PromptBuilder
@@ -16,9 +12,10 @@ import ai.koog.prompt.executor.model.PromptExecutor
1612
import ai.koog.prompt.llm.LLModel
1713
import ai.koog.prompt.message.Message
1814
import ai.koog.prompt.params.LLMParams
19-
import ai.koog.prompt.structure.StructuredData
2015
import ai.koog.prompt.structure.StructuredDataDefinition
16+
import ai.koog.prompt.structure.StructuredOutputConfig
2117
import ai.koog.prompt.structure.StructuredResponse
18+
import ai.koog.prompt.structure.executeStructured
2219
import kotlinx.coroutines.flow.Flow
2320
import kotlinx.coroutines.flow.flatMapMerge
2421
import kotlinx.coroutines.flow.flow
@@ -414,23 +411,19 @@ public class AIAgentLLMWriteSession internal constructor(
414411
}
415412

416413
/**
417-
* Requests an LLM (Language Model) to generate a structured output based on the provided structure.
418-
* The response is post-processed to update the prompt with the raw response.
414+
* Sends a request to LLM and gets a structured response.
419415
*
420-
* @param structure The structured data definition specifying the expected structured output format, schema, and parsing logic.
421-
* @param retries The number of retry attempts to allow in case of generation failures.
422-
* @param fixingModel The language model to use for re-parsing or error correction during retries.
423-
* @return A structured response containing both the parsed structure and the raw response text.
416+
* @param config A configuration defining structures and behavior.
417+
*
418+
* @see [executeStructured]
424419
*/
425420
override suspend fun <T> requestLLMStructured(
426-
structure: StructuredData<T>,
427-
retries: Int,
428-
fixingModel: LLModel
421+
config: StructuredOutputConfig<T>,
429422
): Result<StructuredResponse<T>> {
430-
return super.requestLLMStructured(structure, retries, fixingModel).also {
423+
return super.requestLLMStructured(config).also {
431424
it.onSuccess { response ->
432425
updatePrompt {
433-
assistant(response.raw)
426+
message(response.message)
434427
}
435428
}
436429
}
@@ -455,19 +448,4 @@ public class AIAgentLLMWriteSession internal constructor(
455448

456449
return executor.executeStreaming(prompt, model)
457450
}
458-
459-
/**
460-
* Sends a request to the LLM using the given structured data and expects a structured response in one attempt.
461-
* Updates the prompt with the raw response received from the LLM.
462-
*
463-
* @param structure The structured data defining the schema, examples, and parsing logic for the response.
464-
* @return A structured response containing both the parsed data and the raw response text from the LLM.
465-
*/
466-
override suspend fun <T> requestLLMStructuredOneShot(structure: StructuredData<T>): StructuredResponse<T> {
467-
return super.requestLLMStructuredOneShot(structure).also { response ->
468-
updatePrompt {
469-
assistant(response.raw)
470-
}
471-
}
472-
}
473451
}

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

Lines changed: 5 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,8 @@ import ai.koog.prompt.dsl.PromptBuilder
1717
import ai.koog.prompt.dsl.prompt
1818
import ai.koog.prompt.llm.LLModel
1919
import ai.koog.prompt.message.Message
20-
import ai.koog.prompt.structure.StructuredData
2120
import ai.koog.prompt.structure.StructuredDataDefinition
21+
import ai.koog.prompt.structure.StructuredOutputConfig
2222
import ai.koog.prompt.structure.StructuredResponse
2323
import kotlinx.coroutines.flow.Flow
2424

@@ -167,31 +167,23 @@ public fun AIAgentSubgraphBuilderBase<*, *>.nodeLLMModerateMessage(
167167
}
168168

169169
/**
170-
* A node that appends a user message to the LLM prompt and requests structured data from the LLM with error correction capabilities.
170+
* A node that appends a user message to the LLM prompt and requests structured data from the LLM with optional error correction capabilities.
171171
*
172172
* @param name Optional node name.
173-
* @param structure Definition of expected output format and parsing logic.
174-
* @param retries Number of retry attempts for failed generations.
175-
* @param fixingModel LLM used for error correction.
173+
* @param config A configuration defining structures and behavior.
176174
*/
177175
@AIAgentBuilderDslMarker
178176
public inline fun <reified T> AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestStructured(
179177
name: String? = null,
180-
structure: StructuredData<T>,
181-
retries: Int,
182-
fixingModel: LLModel
178+
config: StructuredOutputConfig<T>,
183179
): AIAgentNodeDelegate<String, Result<StructuredResponse<T>>> =
184180
node(name) { message ->
185181
llm.writeSession {
186182
updatePrompt {
187183
user(message)
188184
}
189185

190-
requestLLMStructured(
191-
structure,
192-
retries,
193-
fixingModel
194-
)
186+
requestLLMStructured(config)
195187
}
196188
}
197189

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

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@ import ai.koog.agents.memory.providers.NoMemory
2121
import ai.koog.prompt.dsl.Prompt
2222
import ai.koog.prompt.llm.LLModel
2323
import ai.koog.prompt.message.Message
24+
import ai.koog.prompt.structure.StructuredOutput
25+
import ai.koog.prompt.structure.StructuredOutputConfig
2426
import ai.koog.prompt.structure.json.JsonStructuredData
2527
import io.github.oshai.kotlinlogging.KotlinLogging
2628
import kotlinx.serialization.Serializable
@@ -537,12 +539,16 @@ internal suspend fun AIAgentLLMWriteSession.retrieveFactsFromHistory(
537539

538540
val facts = when (concept.factType) {
539541
FactType.SINGLE -> {
540-
val response = requestLLMStructured(JsonStructuredData.createJsonStructure<FactStructure>())
542+
val response = requestLLMStructured(
543+
config = StructuredOutputConfig(default = StructuredOutput.Manual(JsonStructuredData.createJsonStructure<FactStructure>()))
544+
)
541545
SingleFact(concept = concept, value = response.getOrNull()?.structure?.fact ?: "No facts extracted", timestamp = timestamp)
542546
}
543547

544548
FactType.MULTIPLE -> {
545-
val response = requestLLMStructured(JsonStructuredData.createJsonStructure<FactListStructure>())
549+
val response = requestLLMStructured(
550+
config = StructuredOutputConfig(default = StructuredOutput.Manual(JsonStructuredData.createJsonStructure<FactListStructure>()))
551+
)
546552
val factsList = response.getOrNull()?.structure?.facts ?: emptyList()
547553
MultipleFacts(concept = concept, values = factsList.map { it.fact }, timestamp = timestamp)
548554
}

‎examples/build.gradle.kts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,8 @@ registerRunExampleTask("runExampleEssay", "ai.koog.agents.example.essay.EssayWri
7272
registerRunExampleTask("runExampleFleetProjectTemplateGeneration", "ai.koog.agents.example.templategen.FleetProjectTemplateGenerationKt")
7373
registerRunExampleTask("runExampleTemplate", "ai.koog.agents.example.template.TemplateKt")
7474
registerRunExampleTask("runProjectAnalyzer", "ai.koog.agents.example.ProjectAnalyzerAgentKt")
75-
registerRunExampleTask("runExampleStructuredOutput", "ai.koog.agents.example.structureddata.StructuredDataExampleKt")
75+
registerRunExampleTask("runExampleSimpleStructuredData", "ai.koog.agents.example.structureddata.SimpleStructuredDataExampleKt")
76+
registerRunExampleTask("runExampleFullStructuredData", "ai.koog.agents.example.structureddata.FullStructuredDataExampleKt")
7677
registerRunExampleTask("runExampleMarkdownStreaming", "ai.koog.agents.example.structureddata.MarkdownStreamingDataExampleKt")
7778
registerRunExampleTask("runExampleMarkdownStreamingWithTool", "ai.koog.agents.example.structureddata.MarkdownStreamingWithToolsExampleKt")
7879
registerRunExampleTask("runExampleRiderProjectTemplate", "ai.koog.agents.example.rider.project.template.RiderProjectTemplateKt")

‎examples/src/main/kotlin/ai/koog/agents/example/banking/routing/RoutingViaGraph.kt

Lines changed: 24 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,14 @@ import ai.koog.agents.example.banking.tools.transactionAnalysisPrompt
1515
import ai.koog.agents.ext.agent.ProvideStringSubgraphResult
1616
import ai.koog.agents.ext.agent.subgraphWithTask
1717
import ai.koog.agents.ext.tool.AskUser
18-
import ai.koog.prompt.structure.json.JsonSchemaGenerator
19-
import ai.koog.prompt.structure.json.JsonStructuredData
2018
import ai.koog.prompt.dsl.prompt
2119
import ai.koog.prompt.executor.clients.openai.OpenAIModels
2220
import ai.koog.prompt.executor.llms.all.simpleOpenAIExecutor
21+
import ai.koog.prompt.structure.StructureFixingParser
22+
import ai.koog.prompt.structure.StructuredOutput
23+
import ai.koog.prompt.structure.StructuredOutputConfig
24+
import ai.koog.prompt.structure.json.JsonStructuredData
25+
import ai.koog.prompt.structure.json.generator.FullJsonSchemaGenerator
2326
import kotlinx.coroutines.runBlocking
2427

2528
fun main() = runBlocking {
@@ -38,21 +41,27 @@ fun main() = runBlocking {
3841
tools = listOf(AskUser)
3942
) {
4043
val requestClassification by nodeLLMRequestStructured(
41-
structure = JsonStructuredData.createJsonStructure<ClassifiedBankRequest>(
42-
schemaFormat = JsonSchemaGenerator.SchemaFormat.JsonSchema,
43-
examples = listOf(
44-
ClassifiedBankRequest(
45-
requestType = RequestType.Transfer,
46-
userRequest = "Send 25 euros to Daniel for dinner at the restaurant."
44+
config = StructuredOutputConfig(
45+
default = StructuredOutput.Manual(
46+
structure = JsonStructuredData.createJsonStructure<ClassifiedBankRequest>(
47+
schemaGenerator = FullJsonSchemaGenerator,
48+
examples = listOf(
49+
ClassifiedBankRequest(
50+
requestType = RequestType.Transfer,
51+
userRequest = "Send 25 euros to Daniel for dinner at the restaurant."
52+
),
53+
ClassifiedBankRequest(
54+
requestType = RequestType.Analytics,
55+
userRequest = "Provide transaction overview for the last month"
56+
)
57+
)
4758
),
48-
ClassifiedBankRequest(
49-
requestType = RequestType.Analytics,
50-
userRequest = "Provide transaction overview for the last month"
51-
)
52-
)
59+
),
60+
fixingParser = StructureFixingParser(
61+
fixingModel = OpenAIModels.CostOptimized.GPT4oMini,
62+
retries = 2,
63+
),
5364
),
54-
retries = 2,
55-
fixingModel = OpenAIModels.CostOptimized.GPT4oMini
5665
)
5766

5867
val callLLM by nodeLLMRequest()

‎examples/src/main/kotlin/ai/koog/agents/example/parallelexecution/BestJokeAgent.kt

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@ import ai.koog.prompt.executor.clients.openai.OpenAILLMClient
1414
import ai.koog.prompt.executor.clients.openai.OpenAIModels
1515
import ai.koog.prompt.executor.llms.MultiLLMPromptExecutor
1616
import ai.koog.prompt.llm.LLMProvider
17+
import ai.koog.prompt.structure.StructuredOutput
18+
import ai.koog.prompt.structure.StructuredOutputConfig
1719
import ai.koog.prompt.structure.json.JsonStructuredData
1820
import io.opentelemetry.exporter.logging.LoggingSpanExporter
1921
import kotlinx.coroutines.runBlocking
@@ -94,7 +96,13 @@ fun main(args: Array<String>) = runBlocking {
9496
}
9597
}
9698

97-
val response = requestLLMStructured(JsonStructuredData.createJsonStructure<JokeWinner>())
99+
val response = requestLLMStructured(
100+
config = StructuredOutputConfig(
101+
default = StructuredOutput.Manual(
102+
structure = JsonStructuredData.createJsonStructure<JokeWinner>()
103+
)
104+
)
105+
)
98106
val bestJoke = response.getOrNull()!!.structure
99107
bestJoke.index
100108
}

0 commit comments

Comments
 (0)