Skip to content

Commit 4602d15

Browse files
tiginamariatiginamaria
authored
Improve naming in structured output module (#1107)
<!-- Thank you for opening a pull request! Please add a brief description of the proposed change here. Also, please tick the appropriate points in the checklist below. --> Before structural refactoring of the whole structured output pipeline, it's important to make some renaming ## Motivation and Context <!-- Why is this change needed? What problem does it solve? --> Renaming done: * `StructuredOutput` -> `StructuredRequest` (Because it can be messaged up with `StructuredResponce`, but it should be opposite to) * `StructuredData`-> `Structure` (Because there is no data in there, this entity just defines the structure and parsing) * `JsonStructuredData` -> `JsonStructure`, `MarkdownStructuredDataDefinition` -> `MarkdownStructureDefinition` as prev * Add `structure: Structure` into `StructuredResponse` to be able to turn in back to string ## Breaking Changes <!-- Will users need to update their code or configurations? --> --- #### Type of the changes - [ ] New feature (non-breaking change which adds functionality) - [ ] Bug fix (non-breaking change which fixes an issue) - [x] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Tests improvement - [x] Refactoring #### Checklist - [ ] The pull request has a description of the proposed change - [ ] I read the [Contributing Guidelines](https://github.com/JetBrains/koog/blob/main/CONTRIBUTING.md) before opening the pull request - [ ] The pull request uses **`develop`** as the base branch - [ ] Tests for the changes have been added - [ ] All new and existing tests passed ##### Additional steps for pull requests adding a new feature - [ ] An issue describing the proposed change exists - [ ] The pull request includes a link to the issue - [ ] The change was discussed and approved in the issue - [ ] Docs have been added / updated
1 parent bb6a6c5 commit 4602d15

File tree

38 files changed

+331
-330
lines changed

38 files changed

+331
-330
lines changed

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

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,9 @@ import ai.koog.agents.core.tools.annotations.LLMDescription
1717
import ai.koog.prompt.llm.LLModel
1818
import ai.koog.prompt.params.LLMParams
1919
import ai.koog.prompt.structure.StructureFixingParser
20-
import ai.koog.prompt.structure.StructuredOutput
21-
import ai.koog.prompt.structure.StructuredOutputConfig
22-
import ai.koog.prompt.structure.json.JsonStructuredData
20+
import ai.koog.prompt.structure.StructuredRequest
21+
import ai.koog.prompt.structure.StructuredRequestConfig
22+
import ai.koog.prompt.structure.json.JsonStructure
2323
import ai.koog.prompt.structure.json.generator.StandardJsonSchemaGenerator
2424
import io.github.oshai.kotlinlogging.KotlinLogging
2525
import kotlinx.coroutines.CancellationException
@@ -130,9 +130,9 @@ public open class AIAgentSubgraph<TInput, TOutput>(
130130
}
131131

132132
val selectedTools = this.requestLLMStructured(
133-
config = StructuredOutputConfig(
134-
default = StructuredOutput.Manual(
135-
JsonStructuredData.createJsonStructure<SelectedTools>(
133+
config = StructuredRequestConfig(
134+
default = StructuredRequest.Manual(
135+
JsonStructure.create<SelectedTools>(
136136
schemaGenerator = StandardJsonSchemaGenerator,
137137
examples = listOf(SelectedTools(listOf()), SelectedTools(tools.map { it.name }.take(3))),
138138
),
@@ -143,7 +143,7 @@ public open class AIAgentSubgraph<TInput, TOutput>(
143143

144144
prompt = initialPrompt
145145

146-
tools.filter { it.name in selectedTools.structure.tools.toSet() }
146+
tools.filter { it.name in selectedTools.data.tools.toSet() }
147147
}
148148
}
149149

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import ai.koog.prompt.message.Message
1313
import ai.koog.prompt.params.LLMParams
1414
import ai.koog.prompt.streaming.StreamFrame
1515
import ai.koog.prompt.structure.StructureFixingParser
16-
import ai.koog.prompt.structure.StructuredOutputConfig
16+
import ai.koog.prompt.structure.StructuredRequestConfig
1717
import ai.koog.prompt.structure.StructuredResponse
1818
import ai.koog.prompt.structure.executeStructured
1919
import ai.koog.prompt.structure.parseResponseToStructuredResponse
@@ -263,7 +263,7 @@ public sealed class AIAgentLLMSession(
263263
* @see [executeStructured]
264264
*/
265265
public open suspend fun <T> requestLLMStructured(
266-
config: StructuredOutputConfig<T>,
266+
config: StructuredRequestConfig<T>,
267267
): Result<StructuredResponse<T>> {
268268
validateSession()
269269

@@ -346,7 +346,7 @@ public sealed class AIAgentLLMSession(
346346
*/
347347
public suspend fun <T> parseResponseToStructuredResponse(
348348
response: Message.Assistant,
349-
config: StructuredOutputConfig<T>
349+
config: StructuredRequestConfig<T>
350350
): StructuredResponse<T> = executor.parseResponseToStructuredResponse(response, config, model)
351351

352352
/**

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

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,9 @@ import ai.koog.prompt.llm.LLModel
1515
import ai.koog.prompt.message.Message
1616
import ai.koog.prompt.params.LLMParams
1717
import ai.koog.prompt.streaming.StreamFrame
18+
import ai.koog.prompt.structure.StructureDefinition
1819
import ai.koog.prompt.structure.StructureFixingParser
19-
import ai.koog.prompt.structure.StructuredDataDefinition
20-
import ai.koog.prompt.structure.StructuredOutputConfig
20+
import ai.koog.prompt.structure.StructuredRequestConfig
2121
import ai.koog.prompt.structure.StructuredResponse
2222
import kotlinx.coroutines.flow.Flow
2323
import kotlinx.coroutines.flow.flatMapMerge
@@ -441,11 +441,9 @@ public class AIAgentLLMWriteSession internal constructor(
441441
* Sends a request to LLM and gets a structured response.
442442
*
443443
* @param config A configuration defining structures and behavior.
444-
*
445-
* @see [executeStructured]
446444
*/
447445
override suspend fun <T> requestLLMStructured(
448-
config: StructuredOutputConfig<T>,
446+
config: StructuredRequestConfig<T>,
449447
): Result<StructuredResponse<T>> {
450448
return super.requestLLMStructured(config).also {
451449
it.onSuccess { response ->
@@ -491,7 +489,7 @@ public class AIAgentLLMWriteSession internal constructor(
491489
* in constructing the prompt for the language model request.
492490
* @return a flow of `StreamingFrame` objects that streams the responses from the language model.
493491
*/
494-
public fun requestLLMStreaming(definition: StructuredDataDefinition? = null): Flow<StreamFrame> {
492+
public fun requestLLMStreaming(definition: StructureDefinition? = null): Flow<StreamFrame> {
495493
if (definition != null) {
496494
val prompt = prompt(prompt, clock) {
497495
user {

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

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,8 @@ import ai.koog.agents.core.tools.ToolDescriptor
1111
import ai.koog.agents.core.tools.ToolResult
1212
import ai.koog.prompt.message.Message
1313
import ai.koog.prompt.streaming.StreamFrame
14+
import ai.koog.prompt.structure.StructureDefinition
1415
import ai.koog.prompt.structure.StructureFixingParser
15-
import ai.koog.prompt.structure.StructuredDataDefinition
1616
import ai.koog.prompt.structure.StructuredResponse
1717
import kotlinx.coroutines.flow.Flow
1818
import kotlinx.serialization.serializer
@@ -142,9 +142,8 @@ public suspend fun AIAgentFunctionalContext.latestTokenUsage(): Int {
142142
* The message becomes part of the current prompt, and the LLM's response is processed to extract structured data.
143143
*
144144
* @param message The content of the message to be sent to the LLM.
145-
* @param structure Definition of expected output format and parsing logic.
146-
* @param retries Number of retry attempts for failed generations.
147-
* @param fixingModel LLM used for error correction.
145+
* @param examples Examples of the structured output.
146+
* @param fixingParser Optional parser to fix generated structured data.
148147
* @return Result containing the structured response if successful, or an error if parsing failed.
149148
*/
150149
public suspend inline fun <reified T> AIAgentFunctionalContext.requestLLMStructured(
@@ -175,7 +174,7 @@ public suspend inline fun <reified T> AIAgentFunctionalContext.requestLLMStructu
175174
*/
176175
public suspend fun AIAgentFunctionalContext.requestLLMStreaming(
177176
message: String,
178-
structureDefinition: StructuredDataDefinition? = null
177+
structureDefinition: StructureDefinition? = null
179178
): Flow<StreamFrame> {
180179
return llm.writeSession {
181180
appendPrompt {

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

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,9 @@ import ai.koog.prompt.llm.LLModel
1717
import ai.koog.prompt.message.Message
1818
import ai.koog.prompt.streaming.StreamFrame
1919
import ai.koog.prompt.streaming.toMessageResponses
20+
import ai.koog.prompt.structure.StructureDefinition
2021
import ai.koog.prompt.structure.StructureFixingParser
21-
import ai.koog.prompt.structure.StructuredDataDefinition
22-
import ai.koog.prompt.structure.StructuredOutputConfig
22+
import ai.koog.prompt.structure.StructuredRequestConfig
2323
import ai.koog.prompt.structure.StructuredResponse
2424
import kotlinx.coroutines.flow.Flow
2525
import kotlinx.coroutines.flow.toList
@@ -198,7 +198,7 @@ public fun AIAgentSubgraphBuilderBase<*, *>.nodeLLMModerateMessage(
198198
@AIAgentBuilderDslMarker
199199
public inline fun <reified T> AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestStructured(
200200
name: String? = null,
201-
config: StructuredOutputConfig<T>,
201+
config: StructuredRequestConfig<T>,
202202
): AIAgentNodeDelegate<String, Result<StructuredResponse<T>>> =
203203
node(name) { message ->
204204
llm.writeSession {
@@ -254,7 +254,7 @@ public inline fun <reified T> AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestStr
254254
@AIAgentBuilderDslMarker
255255
public fun <T> AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestStreaming(
256256
name: String? = null,
257-
structureDefinition: StructuredDataDefinition? = null,
257+
structureDefinition: StructureDefinition? = null,
258258
transformStreamData: suspend (Flow<StreamFrame>) -> Flow<T>
259259
): AIAgentNodeDelegate<String, Flow<T>> =
260260
node(name) { message ->
@@ -278,7 +278,7 @@ public fun <T> AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestStreaming(
278278
@AIAgentBuilderDslMarker
279279
public fun AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestStreaming(
280280
name: String? = null,
281-
structureDefinition: StructuredDataDefinition? = null,
281+
structureDefinition: StructureDefinition? = null,
282282
): AIAgentNodeDelegate<String, Flow<StreamFrame>> = nodeLLMRequestStreaming(name, structureDefinition) { it }
283283

284284
/**
@@ -357,7 +357,7 @@ public inline fun <reified T> AIAgentSubgraphBuilderBase<*, *>.nodeLLMCompressHi
357357
@AIAgentBuilderDslMarker
358358
public inline fun <reified T> AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestStreamingAndSendResults(
359359
name: String? = null,
360-
structureDefinition: StructuredDataDefinition? = null
360+
structureDefinition: StructureDefinition? = null
361361
): AIAgentNodeDelegate<T, List<Message.Response>> = node(name) { input ->
362362
llm.writeSession {
363363
requestLLMStreaming(structureDefinition)
@@ -536,7 +536,7 @@ public inline fun <reified ToolArg, reified TResult> AIAgentSubgraphBuilderBase<
536536
@AIAgentBuilderDslMarker
537537
public inline fun <reified TInput, T> AIAgentSubgraphBuilderBase<*, *>.nodeSetStructuredOutput(
538538
name: String? = null,
539-
config: StructuredOutputConfig<T>
539+
config: StructuredRequestConfig<T>
540540
): AIAgentNodeDelegate<TInput, TInput> =
541541
node(name) { message ->
542542
llm.writeSession {

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

Lines changed: 25 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,9 @@ import ai.koog.prompt.dsl.prompt
99
import ai.koog.prompt.executor.clients.openai.OpenAIModels
1010
import ai.koog.prompt.message.Message
1111
import ai.koog.prompt.message.ResponseMetaInfo
12-
import ai.koog.prompt.structure.StructuredOutput
13-
import ai.koog.prompt.structure.StructuredOutputConfig
14-
import ai.koog.prompt.structure.json.JsonStructuredData
12+
import ai.koog.prompt.structure.StructuredRequest
13+
import ai.koog.prompt.structure.StructuredRequestConfig
14+
import ai.koog.prompt.structure.json.JsonStructure
1515
import kotlinx.coroutines.test.runTest
1616
import kotlinx.serialization.Serializable
1717
import kotlin.test.Test
@@ -32,9 +32,9 @@ class AIAgentLLMSessionStructuredOutputTest : AgentTestBase() {
3232

3333
@Test
3434
fun testParseResponseToStructuredResponse() = runTest {
35-
val structure = JsonStructuredData.createJsonStructure<TestStructure>()
36-
val config = StructuredOutputConfig(
37-
default = StructuredOutput.Manual(structure)
35+
val structure = JsonStructure.create<TestStructure>()
36+
val config = StructuredRequestConfig(
37+
default = StructuredRequest.Manual(structure)
3838
)
3939

4040
val mockExecutor = getMockExecutor {
@@ -77,17 +77,17 @@ class AIAgentLLMSessionStructuredOutputTest : AgentTestBase() {
7777
}
7878

7979
assertNotNull(result)
80-
assertEquals("John Doe", result.structure.name)
81-
assertEquals(30, result.structure.age)
82-
assertEquals("A test person", result.structure.description)
80+
assertEquals("John Doe", result.data.name)
81+
assertEquals(30, result.data.age)
82+
assertEquals("A test person", result.data.description)
8383
assertEquals(assistantMessage, result.message)
8484
}
8585

8686
@Test
8787
fun testParseResponseToStructuredResponseWithNullableField() = runTest {
88-
val structure = JsonStructuredData.createJsonStructure<TestStructure>()
89-
val config = StructuredOutputConfig(
90-
default = StructuredOutput.Manual(structure)
88+
val structure = JsonStructure.create<TestStructure>()
89+
val config = StructuredRequestConfig(
90+
default = StructuredRequest.Manual(structure)
9191
)
9292

9393
val mockExecutor = getMockExecutor {
@@ -129,9 +129,9 @@ class AIAgentLLMSessionStructuredOutputTest : AgentTestBase() {
129129
}
130130

131131
assertNotNull(result)
132-
assertEquals("Jane Doe", result.structure.name)
133-
assertEquals(25, result.structure.age)
134-
assertEquals(null, result.structure.description)
132+
assertEquals("Jane Doe", result.data.name)
133+
assertEquals(25, result.data.age)
134+
assertEquals(null, result.data.description)
135135
assertEquals(assistantMessage, result.message)
136136
}
137137

@@ -155,9 +155,9 @@ class AIAgentLLMSessionStructuredOutputTest : AgentTestBase() {
155155
val tags: Set<String>
156156
)
157157

158-
val structure = JsonStructuredData.createJsonStructure<ComplexStructure>()
159-
val config = StructuredOutputConfig(
160-
default = StructuredOutput.Manual(structure)
158+
val structure = JsonStructure.create<ComplexStructure>()
159+
val config = StructuredRequestConfig(
160+
default = StructuredRequest.Manual(structure)
161161
)
162162

163163
val mockExecutor = getMockExecutor {
@@ -209,13 +209,13 @@ class AIAgentLLMSessionStructuredOutputTest : AgentTestBase() {
209209
}
210210

211211
assertNotNull(result)
212-
assertEquals(123, result.structure.id)
213-
assertEquals(2, result.structure.addresses.size)
214-
assertEquals("123 Main St", result.structure.addresses[0].street)
215-
assertEquals("New York", result.structure.addresses[0].city)
216-
assertEquals("456 Oak Ave", result.structure.addresses[1].street)
217-
assertEquals("Los Angeles", result.structure.addresses[1].city)
218-
assertEquals(setOf("vip", "premium", "verified"), result.structure.tags)
212+
assertEquals(123, result.data.id)
213+
assertEquals(2, result.data.addresses.size)
214+
assertEquals("123 Main St", result.data.addresses[0].street)
215+
assertEquals("New York", result.data.addresses[0].city)
216+
assertEquals("456 Oak Ave", result.data.addresses[1].street)
217+
assertEquals("Los Angeles", result.data.addresses[1].city)
218+
assertEquals(setOf("vip", "premium", "verified"), result.data.tags)
219219
assertEquals(assistantMessage, result.message)
220220
}
221221
}

‎agents/agents-core/src/commonTest/kotlin/ai/koog/agents/core/dsl/extension/AIAgentNodesTest.kt‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,9 @@ import ai.koog.prompt.dsl.Prompt
1212
import ai.koog.prompt.dsl.prompt
1313
import ai.koog.prompt.executor.clients.openai.OpenAIModels
1414
import ai.koog.prompt.llm.OllamaModels
15-
import ai.koog.prompt.structure.StructuredOutput
16-
import ai.koog.prompt.structure.StructuredOutputConfig
17-
import ai.koog.prompt.structure.json.JsonStructuredData
15+
import ai.koog.prompt.structure.StructuredRequest
16+
import ai.koog.prompt.structure.StructuredRequestConfig
17+
import ai.koog.prompt.structure.json.JsonStructure
1818
import ai.koog.utils.io.use
1919
import kotlinx.coroutines.test.runTest
2020
import kotlinx.serialization.Serializable
@@ -154,9 +154,9 @@ class AIAgentNodesTest {
154154
)
155155

156156
// Test Manual mode
157-
val manualStructure = JsonStructuredData.createJsonStructure<TestOutput>()
158-
val manualConfig = StructuredOutputConfig(
159-
default = StructuredOutput.Manual(manualStructure)
157+
val manualStructure = JsonStructure.create<TestOutput>()
158+
val manualConfig = StructuredRequestConfig(
159+
default = StructuredRequest.Manual(manualStructure)
160160
)
161161

162162
var capturedPrompt: Prompt? = null
@@ -201,9 +201,9 @@ class AIAgentNodesTest {
201201
)
202202

203203
// Test Native mode
204-
val nativeStructure = JsonStructuredData.createJsonStructure<TestOutput>()
205-
val nativeConfig = StructuredOutputConfig(
206-
default = StructuredOutput.Native(nativeStructure)
204+
val nativeStructure = JsonStructure.create<TestOutput>()
205+
val nativeConfig = StructuredRequestConfig(
206+
default = StructuredRequest.Native(nativeStructure)
207207
)
208208

209209
val nativeStrategy = strategy<String, String>("test-native") {

‎agents/agents-ext/src/commonMain/kotlin/ai/koog/agents/ext/agent/AIAgentStrategies.kt‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ import ai.koog.agents.core.dsl.extension.onToolCall
1919
import ai.koog.agents.core.environment.ReceivedToolResult
2020
import ai.koog.agents.core.environment.result
2121
import ai.koog.prompt.message.Message
22-
import ai.koog.prompt.structure.StructuredOutputConfig
22+
import ai.koog.prompt.structure.StructuredRequestConfig
2323

2424
// FIXME improve this strategy to use Message.Assistant to chat, it works better than tools
2525

@@ -183,7 +183,7 @@ public fun reActStrategy(
183183
* @param config The configuration for structured output processing, specifying schema, providers, and optional error handling mechanisms.
184184
*/
185185
public inline fun <reified Output> structuredOutputWithToolsStrategy(
186-
config: StructuredOutputConfig<Output>,
186+
config: StructuredRequestConfig<Output>,
187187
parallelTools: Boolean = false
188188
): AIAgentGraphStrategy<String, Output> = structuredOutputWithToolsStrategy(
189189
config,
@@ -203,7 +203,7 @@ public inline fun <reified Output> structuredOutputWithToolsStrategy(
203203
* that serves as the input for further processing in the structured output pipeline.
204204
*/
205205
public inline fun <reified Input, reified Output> structuredOutputWithToolsStrategy(
206-
config: StructuredOutputConfig<Output>,
206+
config: StructuredRequestConfig<Output>,
207207
parallelTools: Boolean = false,
208208
noinline transform: suspend AIAgentGraphContextBase.(input: Input) -> String
209209
): AIAgentGraphStrategy<Input, Output> = strategy<Input, Output>("structured_output_with_tools_strategy") {
@@ -214,7 +214,7 @@ public inline fun <reified Input, reified Output> structuredOutputWithToolsStrat
214214
val sendToolResult by nodeLLMSendMultipleToolResults()
215215
val transformToStructuredOutput by node<Message.Assistant, Output> { response ->
216216
llm.writeSession {
217-
parseResponseToStructuredResponse(response, config).structure
217+
parseResponseToStructuredResponse(response, config).data
218218
}
219219
}
220220

‎agents/agents-ext/src/commonMain/kotlin/ai/koog/agents/ext/agent/LLMAsAJudge.kt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -109,10 +109,10 @@ public inline fun <reified T> AIAgentSubgraphBuilderBase<*, *>.llmAsAJudge(
109109
),
110110
// optional field -- recommented for reliability of the format
111111
fixingParser = StructureFixingParser(
112-
fixingModel = OpenAIModels.CostOptimized.GPT4oMini,
112+
model = OpenAIModels.CostOptimized.GPT4oMini,
113113
retries = 3,
114114
)
115-
).getOrThrow().structure
115+
).getOrThrow().data
116116

117117
prompt = initialPrompt
118118
model = initialModel

0 commit comments

Comments
 (0)