From 045587b9100054c85480ea88977a8bbb19b57af3 Mon Sep 17 00:00:00 2001 From: Damjan Polugic Date: Thu, 29 May 2025 11:50:13 +0200 Subject: [PATCH 1/4] Add assertion for schema description in StructuredDataTest.kt --- .../ai/koog/prompt/structure/StructuredDataTest.kt | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/prompt/prompt-structure/src/commonTest/kotlin/ai/koog/prompt/structure/StructuredDataTest.kt b/prompt/prompt-structure/src/commonTest/kotlin/ai/koog/prompt/structure/StructuredDataTest.kt index 4cc975238a..83d3735063 100644 --- a/prompt/prompt-structure/src/commonTest/kotlin/ai/koog/prompt/structure/StructuredDataTest.kt +++ b/prompt/prompt-structure/src/commonTest/kotlin/ai/koog/prompt/structure/StructuredDataTest.kt @@ -1,5 +1,6 @@ package ai.koog.prompt.structure +import ai.koog.agents.core.tools.annotations.LLMDescription import ai.koog.prompt.structure.json.JsonStructuredData import ai.koog.prompt.text.TextContentBuilder import kotlinx.serialization.SerialName @@ -13,7 +14,11 @@ import kotlin.test.assertTrue class JsonStructuredDataTest { // Simple data structure for basic tests @Serializable - data class SimpleData(val value: String) + @LLMDescription("SimpleData description") + data class SimpleData( + @property:LLMDescription("SimpleData.value description") + val value: String + ) // Array data structure @Serializable @@ -124,6 +129,8 @@ class JsonStructuredDataTest { val content = builder.build() assertTrue(content.contains("DEFINITION OF simple")) + assertTrue(content.contains("SimpleData description")) + assertTrue(content.contains("SimpleData.value description")) assertTrue(content.contains("is defined only and solely with JSON, without any additional characters, backticks or anything similar.")) } From 91a7b3ba9598751bb7fd1141b9bd1bce41d534d3 Mon Sep 17 00:00:00 2001 From: Damjan Polugic Date: Thu, 29 May 2025 11:55:45 +0200 Subject: [PATCH 2/4] Add @SerialInfo annotation to LLMDescription --- .../ai/koog/agents/core/tools/annotations/LLMDescription.kt | 3 +++ 1 file changed, 3 insertions(+) diff --git a/agents/agents-tools/src/commonMain/kotlin/ai/koog/agents/core/tools/annotations/LLMDescription.kt b/agents/agents-tools/src/commonMain/kotlin/ai/koog/agents/core/tools/annotations/LLMDescription.kt index b29cc621cb..3e82d22ffd 100644 --- a/agents/agents-tools/src/commonMain/kotlin/ai/koog/agents/core/tools/annotations/LLMDescription.kt +++ b/agents/agents-tools/src/commonMain/kotlin/ai/koog/agents/core/tools/annotations/LLMDescription.kt @@ -1,11 +1,14 @@ package ai.koog.agents.core.tools.annotations +import kotlinx.serialization.SerialInfo + /** * Description for an entity that can be provided to LLMs. * You may use it to annotate properties, functions, parameters, classes, return types, etc. * * @property description The description of the entity. */ +@SerialInfo @Target( AnnotationTarget.PROPERTY, AnnotationTarget.CLASS, From 23f41669f38afb29f681bdb132ec1b89d06d6085 Mon Sep 17 00:00:00 2001 From: Andrey Bragin Date: Fri, 30 May 2025 00:26:35 +0200 Subject: [PATCH 3/4] [prompt] Fix JSON structured output descriptions Properly handle description annotation during schema generation --- .gitignore | 1 + .../core/tools/annotations/LLMDescription.kt | 1 - .../structureddata/StructuredDataExample.kt | 34 +++--- .../prompt/structure/DescriptionMetadata.kt | 44 -------- .../structure/json/JsonSchemaGenerator.kt | 25 ++++- .../structure/json/JsonStructuredData.kt | 101 +++--------------- .../structure/json/JsonSchemaGeneratorTest.kt | 42 +++++--- 7 files changed, 76 insertions(+), 172 deletions(-) delete mode 100644 prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/DescriptionMetadata.kt diff --git a/.gitignore b/.gitignore index 22de8b8e5e..f16ac09ebb 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,4 @@ .kotlin build **/.claude/settings.local.json +env.properties diff --git a/agents/agents-tools/src/commonMain/kotlin/ai/koog/agents/core/tools/annotations/LLMDescription.kt b/agents/agents-tools/src/commonMain/kotlin/ai/koog/agents/core/tools/annotations/LLMDescription.kt index 3e82d22ffd..0b6d7370ea 100644 --- a/agents/agents-tools/src/commonMain/kotlin/ai/koog/agents/core/tools/annotations/LLMDescription.kt +++ b/agents/agents-tools/src/commonMain/kotlin/ai/koog/agents/core/tools/annotations/LLMDescription.kt @@ -12,7 +12,6 @@ import kotlinx.serialization.SerialInfo @Target( AnnotationTarget.PROPERTY, AnnotationTarget.CLASS, - AnnotationTarget.PROPERTY, AnnotationTarget.TYPE, AnnotationTarget.VALUE_PARAMETER, AnnotationTarget.FUNCTION diff --git a/examples/src/main/kotlin/ai/koog/agents/example/structureddata/StructuredDataExample.kt b/examples/src/main/kotlin/ai/koog/agents/example/structureddata/StructuredDataExample.kt index 9395ff9cb3..c1fa4d984a 100644 --- a/examples/src/main/kotlin/ai/koog/agents/example/structureddata/StructuredDataExample.kt +++ b/examples/src/main/kotlin/ai/koog/agents/example/structureddata/StructuredDataExample.kt @@ -5,12 +5,11 @@ import ai.koog.agents.core.agent.config.AIAgentConfig import ai.koog.agents.core.dsl.builder.forwardTo import ai.koog.agents.core.dsl.builder.strategy import ai.koog.agents.core.dsl.extension.nodeLLMRequest +import ai.koog.agents.core.dsl.extension.nodeLLMRequestStructured import ai.koog.agents.core.tools.ToolRegistry import ai.koog.agents.core.tools.annotations.LLMDescription import ai.koog.agents.example.ApiKeyService import ai.koog.agents.features.eventHandler.feature.handleEvents -import ai.koog.prompt.structure.json.JsonSchemaGenerator -import ai.koog.prompt.structure.json.JsonStructuredData import ai.koog.prompt.dsl.prompt import ai.koog.prompt.executor.clients.anthropic.AnthropicLLMClient import ai.koog.prompt.executor.clients.anthropic.AnthropicModels @@ -18,7 +17,8 @@ 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.llm.LLMProvider -import ai.koog.prompt.message.Message +import ai.koog.prompt.structure.json.JsonSchemaGenerator +import ai.koog.prompt.structure.json.JsonStructuredData import kotlinx.coroutines.runBlocking import kotlinx.serialization.SerialName import kotlinx.serialization.Serializable @@ -210,25 +210,16 @@ fun main(): Unit = runBlocking { val agentStrategy = strategy("weather-forecast") { val setup by nodeLLMRequest() - - val getStructuredForecast by node { _ -> - val structuredResponse = llm.writeSession { - this.requestLLMStructured( - structure = weatherForecastStructure, - // the model that would handle coercion if the output does not conform to the requested structure - fixingModel = OpenAIModels.Reasoning.GPT4oMini, - ).getOrThrow() - } - - """ - Response structure: - ${structuredResponse.structure} - """.trimIndent() - } + val getStructuredForecast by nodeLLMRequestStructured( + structure = weatherForecastStructure, + retries = 1, + // the model that would handle coercion if the output does not conform to the requested structure + fixingModel = OpenAIModels.Reasoning.GPT4oMini, + ) edge(nodeStart forwardTo setup) - edge(setup forwardTo getStructuredForecast) - edge(getStructuredForecast forwardTo nodeFinish) + edge(setup forwardTo getStructuredForecast transformed { it.content }) + edge(getStructuredForecast forwardTo nodeFinish transformed { "Structured response:\n${it.getOrThrow().structure}" }) } val agentConfig = AIAgentConfig( @@ -269,8 +260,7 @@ fun main(): Unit = runBlocking { === Weather Forecast Example === This example demonstrates how to use StructuredData and sendStructuredAndUpdatePrompt to get properly structured output from the LLM. - - """.trimIndent() + """.trimIndent() ) runner.run("Get weather forecast for New York") diff --git a/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/DescriptionMetadata.kt b/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/DescriptionMetadata.kt deleted file mode 100644 index 40d6e27554..0000000000 --- a/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/DescriptionMetadata.kt +++ /dev/null @@ -1,44 +0,0 @@ -package ai.koog.prompt.structure - -/** - * Represents a metadata definition for annotated descriptions of a class and its fields. - * This interface facilitates managing descriptions for classes and their corresponding fields. - */ -public interface DescriptionMetadata { - /** - * Represents the name of a class associated with a description or metadata entry. - * - * This property is intended to uniquely identify a class in a descriptive metadata context. - * It can be used to associate additional information, such as a description or field details, - * with the class it represents. - */ - public val className: String - /** - * Represents a detailed description of a class provided in a metadata context. - * - * This property may contain additional information or insights related to the class it belongs to, - * offering explanations, context, or clarifications about its purpose or behavior. - * - * Can be `null` if no description is provided or applicable. - */ - public val classDescription: String? - /** - * A map containing metadata descriptions for fields associated with a specific class. - * - * The keys in the map represent the names of fields, and the corresponding values provide human-readable - * descriptions or details about each field. This is commonly used for dynamically managing and accessing - * metadata about object fields. - */ - public val fieldDescriptions: Map - - /** - * Aggregates all the available descriptions into a single map. This includes field descriptions - * and, if available, the class description associated with the class name. - * - * @return A map where the keys are field or class names and the values are their corresponding descriptions. - */ - public fun allDescriptions(): Map = buildMap { - putAll(fieldDescriptions) - classDescription?.let { put(className, it) } - } -} \ No newline at end of file diff --git a/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/JsonSchemaGenerator.kt b/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/JsonSchemaGenerator.kt index b10f9d2c54..7566ceae72 100644 --- a/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/JsonSchemaGenerator.kt +++ b/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/JsonSchemaGenerator.kt @@ -1,5 +1,6 @@ package ai.koog.prompt.structure.json +import ai.koog.agents.core.tools.annotations.LLMDescription import kotlinx.serialization.KSerializer import kotlinx.serialization.descriptors.* import kotlinx.serialization.json.* @@ -65,7 +66,8 @@ public class JsonSchemaGenerator( * Generate a JSON schema for a serializable class. * * @param serializer The serializer for the class - * @param descriptions Optional map of serial class names and property names to descriptions + * @param descriptions Optional map of serial class names and property names to descriptions. + * If a property/type is already described with [LLMDescription] annotation, value from the map will override this description. * @return A JsonObject representing the JSON schema */ public fun generate( @@ -202,8 +204,18 @@ public class JsonSchemaGenerator( for (i in 0 until descriptor.elementsCount) { val propertyName = descriptor.getElementName(i) val propertyDescriptor = descriptor.getElementDescriptor(i) + val propertyAnnotations = descriptor.getElementAnnotations(i) + + // Description for a property val lookupKey = "${descriptor.serialName}.$propertyName" - val propertyDescription = descriptions[lookupKey] + val propertyDescriptionMap = descriptions[lookupKey] + val propertyDescriptionAnnotation = propertyAnnotations + .filterIsInstance() + .firstOrNull() + ?.description + + // Look at the explicit map first, then at the annotation + val propertyDescription = propertyDescriptionMap ?: propertyDescriptionAnnotation put( propertyName, @@ -243,7 +255,14 @@ public class JsonSchemaGenerator( } // Description for a whole type (definition) - val typeDescription = descriptions[descriptor.serialName] + val typeDescriptionMap = descriptions[descriptor.serialName] + val typeDescriptionAnnotation = descriptor.annotations + .filterIsInstance() + .firstOrNull() + ?.description + + // Look at the explicit map first, then at the annotation + val typeDescription = typeDescriptionMap ?: typeDescriptionAnnotation // Build type definition val typeDefinition = buildJsonObject { diff --git a/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/JsonStructuredData.kt b/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/JsonStructuredData.kt index 53261d3da8..3cdf144719 100644 --- a/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/JsonStructuredData.kt +++ b/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/JsonStructuredData.kt @@ -1,10 +1,9 @@ package ai.koog.prompt.structure.json import ai.koog.agents.core.tools.annotations.LLMDescription -import ai.koog.prompt.structure.DescriptionMetadata +import ai.koog.prompt.params.LLMParams import ai.koog.prompt.structure.StructuredData import ai.koog.prompt.structure.structure -import ai.koog.prompt.params.LLMParams import ai.koog.prompt.text.TextContentBuilder import kotlinx.serialization.KSerializer import kotlinx.serialization.json.Json @@ -67,6 +66,16 @@ public class JsonStructuredData( // TODO: Class.simpleName is the only reason to make the function inline, perhaps we can hide most of the implementation /** * Factory method to create JSON structure with auto-generated JSON schema. + * + * @param id Unique identifier for the structure. + * @param serializer Serializer used for converting the data to and from JSON. + * @param json JSON configuration instance used for serialization. + * @param schemaFormat Format of the generated schema, can be simple or detailed. + * @param maxDepth Maximum recursion depth when generating schema to prevent infinite recursion for circular references. + * @param descriptions Optional map of serial class names and property names to descriptions. + * If a property/type is already described with [LLMDescription] annotation, value from the map will override this description. + * @param examples List of example data items that conform to the structure, used for demonstrating valid formats. + * @param schemaType Type of JSON schema to generate, determines the level of detail in the schema. */ public inline fun createJsonStructure( id: String = T::class.simpleName ?: error("Class name is required for JSON structure"), @@ -74,19 +83,12 @@ public class JsonStructuredData( json: Json = JsonStructureLanguage.defaultJson, schemaFormat: JsonSchemaGenerator.SchemaFormat = JsonSchemaGenerator.SchemaFormat.Simple, maxDepth: Int = 20, - propertyDescriptionOverrides: Map = emptyMap(), + descriptions: Map = emptyMap(), examples: List = emptyList(), schemaType: JsonSchemaType = JsonSchemaType.SIMPLE ): StructuredData { val structureLanguage = JsonStructureLanguage(json) - val metadata = getDescriptionMetadata(serializer) - - // Use platform-specific implementations to get property descriptions - val propertyDescriptions = metadata?.allDescriptions().orEmpty() - .merge(propertyDescriptionOverrides) { _, _, override -> override } - - val schema = - JsonSchemaGenerator(json, schemaFormat, maxDepth).generate(id, serializer, propertyDescriptions) + val schema = JsonSchemaGenerator(json, schemaFormat, maxDepth).generate(id, serializer, descriptions) return JsonStructuredData( id = id, @@ -99,82 +101,5 @@ public class JsonStructuredData( } ) } - - /** - * Retrieves description metadata for a given serializer. The metadata includes - * a description of the class (if annotated) and descriptions of its fields - * based on the presence of the `LLMDescription` annotation. - * - * @param T The type of the serializer. - * @param serializer The serializer for the type T, used to extract metadata. - * @return A `DescriptionMetadata` object containing the class and field descriptions, - * or `null` if no descriptions are found. - */ - @PublishedApi - internal fun getDescriptionMetadata(serializer: KSerializer): DescriptionMetadata? { - // Try to find the class in the registry - val className = serializer.descriptor.serialName - - // Check if the class has LLMDescription annotation - val classDescription = serializer.descriptor.annotations - .filterIsInstance() - .firstOrNull() - ?.description - - // Collect field descriptions - val fieldDescriptions = mutableMapOf() - val descriptor = serializer.descriptor - - for (i in 0 until descriptor.elementsCount) { - val propertyName = descriptor.getElementName(i) - val propertyAnnotations = descriptor.getElementAnnotations(i) - - val description = propertyAnnotations - .filterIsInstance() - .firstOrNull() - ?.description - - if (description != null) { - // Use the format expected by JsonSchemaGenerator: "${descriptor.serialName}.$propertyName" - fieldDescriptions["$className.$propertyName"] = description - } - } - - // If no class description and no field descriptions, return null - if (classDescription == null && fieldDescriptions.isEmpty()) { - return null - } - - // Create a new DescriptionMetadata object with the class description and field descriptions - return object : DescriptionMetadata { - override val className: String = className - override val classDescription: String? = classDescription - override val fieldDescriptions: Map = fieldDescriptions - } - } - - /** - * Merges [other] into the current map. - * If the key already exists, it calls [merger] and puts its result, - * otherwise it simply adds a new pair. - */ - @PublishedApi - internal inline fun MutableMap.mergeInPlace( - other: Map, - merger: (key: K, first: V, second: V) -> V - ): MutableMap = apply { - other.forEach { (k, v) -> - this[k] = get(k)?.let { merger(k, it, v) } ?: v - } - } - - /** - * Non-blocking merge: returns a new map, leaving the original collections unchanged. - */ - @PublishedApi - internal inline fun Map.merge( - other: Map, - merger: (key: K, first: V, second: V) -> V - ): Map = toMutableMap().mergeInPlace(other, merger) } } diff --git a/prompt/prompt-structure/src/commonTest/kotlin/ai/koog/prompt/structure/json/JsonSchemaGeneratorTest.kt b/prompt/prompt-structure/src/commonTest/kotlin/ai/koog/prompt/structure/json/JsonSchemaGeneratorTest.kt index 0c081f5693..477877caa2 100644 --- a/prompt/prompt-structure/src/commonTest/kotlin/ai/koog/prompt/structure/json/JsonSchemaGeneratorTest.kt +++ b/prompt/prompt-structure/src/commonTest/kotlin/ai/koog/prompt/structure/json/JsonSchemaGeneratorTest.kt @@ -1,5 +1,6 @@ package ai.koog.prompt.structure.json +import ai.koog.agents.core.tools.annotations.LLMDescription import kotlinx.serialization.SerialName import kotlinx.serialization.Serializable import kotlinx.serialization.encodeToString @@ -34,7 +35,9 @@ class JsonSchemaGeneratorTest { @Serializable @SerialName("TestClass") + @LLMDescription("A test class") data class TestClass( + @property:LLMDescription("A string property") val stringProperty: String, val intProperty: Int, val booleanProperty: Boolean, @@ -45,7 +48,9 @@ class JsonSchemaGeneratorTest { @Serializable @SerialName("NestedTestClass") + @LLMDescription("Nested test class") data class NestedTestClass( + @LLMDescription("The name") val name: String, val nested: NestedProperty, val nestedList: List = emptyList(), @@ -53,7 +58,9 @@ class JsonSchemaGeneratorTest { ) { @Serializable @SerialName("NestedProperty") + @LLMDescription("Nested property class") data class NestedProperty( + @property:LLMDescription("Nested foo property") val foo: String, val bar: Int ) @@ -117,9 +124,11 @@ class JsonSchemaGeneratorTest { "${"$"}defs": { "TestClass": { "type": "object", + "description": "A test class", "properties": { "stringProperty": { - "type": "string" + "type": "string", + "description": "A string property" }, "intProperty": { "type": "integer" @@ -166,9 +175,11 @@ class JsonSchemaGeneratorTest { val expectedSchema = """ { "type": "object", + "description": "A test class", "properties": { "stringProperty": { - "type": "string" + "type": "string", + "description": "A string property" }, "intProperty": { "type": "integer" @@ -207,8 +218,7 @@ class JsonSchemaGeneratorTest { @Test fun testGenerateJsonSchemaWithDescriptions() { val descriptions = mapOf( - "TestClass" to "A test class", - "TestClass.stringProperty" to "A string property", + "TestClass" to "A test class (override)", "TestClass.intProperty" to "An integer property" ) @@ -221,7 +231,7 @@ class JsonSchemaGeneratorTest { "${"$"}defs": { "TestClass": { "type": "object", - "description": "A test class", + "description": "A test class (override)", "properties": { "stringProperty": { "type": "string", @@ -269,8 +279,7 @@ class JsonSchemaGeneratorTest { @Test fun testGenerateSimpleSchemaWithDescriptions() { val descriptions = mapOf( - "TestClass" to "A test class", - "TestClass.stringProperty" to "A string property", + "TestClass" to "A test class (override)", "TestClass.intProperty" to "An integer property" ) @@ -279,7 +288,7 @@ class JsonSchemaGeneratorTest { val expectedSchema = """ { "type": "object", - "description": "A test class", + "description": "A test class (override)", "properties": { "stringProperty": { "type": "string", @@ -323,11 +332,10 @@ class JsonSchemaGeneratorTest { @Test fun testJsonSchemaNestedDescriptions() { val descriptions = mapOf( - "NestedTestClass.name" to "The name", + "NestedTestClass.name" to "The name (override)", "NestedTestClass.nestedList" to "List of nested properties", "NestedTestClass.nestedMap" to "Map of nested properties", - "NestedProperty.foo" to "Nested foo property", "NestedProperty.bar" to "Nested bar property", ) @@ -340,6 +348,7 @@ class JsonSchemaGeneratorTest { "${"$"}defs": { "NestedProperty": { "type": "object", + "description": "Nested property class", "properties": { "foo": { "type": "string", @@ -357,10 +366,11 @@ class JsonSchemaGeneratorTest { }, "NestedTestClass": { "type": "object", + "description": "Nested test class", "properties": { "name": { "type": "string", - "description": "The name" + "description": "The name (override)" }, "nested": { "${"$"}ref": "#/defs/NestedProperty" @@ -397,26 +407,28 @@ class JsonSchemaGeneratorTest { @Test fun testSimpleSchemaNestedDescriptions() { val descriptions = mapOf( - "NestedTestClass.name" to "The name", + "NestedTestClass.name" to "The name (override)", "NestedTestClass.nestedList" to "List of nested properties", "NestedTestClass.nestedMap" to "Map of nested properties", - "NestedProperty.foo" to "Nested foo property", "NestedProperty.bar" to "Nested bar property", ) + val schema = json.encodeToString(simpleSchemaGenerator.generate("NestedTestClass", serializer(), descriptions)) val expectedDotSchema = """ { "type": "object", + "description": "Nested test class", "properties": { "name": { "type": "string", - "description": "The name" + "description": "The name (override)" }, "nested": { "type": "object", + "description": "Nested property class", "properties": { "foo": { "type": "string", @@ -436,6 +448,7 @@ class JsonSchemaGeneratorTest { "type": "array", "items": { "type": "object", + "description": "Nested property class", "properties": { "foo": { "type": "string", @@ -457,6 +470,7 @@ class JsonSchemaGeneratorTest { "type": "object", "additionalProperties": { "type": "object", + "description": "Nested property class", "properties": { "foo": { "type": "string", From cb74609492f372754b5c9a8f6b8df79bcf27f48b Mon Sep 17 00:00:00 2001 From: Andrey Bragin Date: Fri, 30 May 2025 13:25:40 +0200 Subject: [PATCH 4/4] [prompt] Rename to descriptionOverrides, add example usage in KDocs --- .../structureddata/StructuredDataExample.kt | 2 +- .../structure/json/JsonSchemaGenerator.kt | 24 ++++---- .../structure/json/JsonStructuredData.kt | 58 ++++++++++++++++++- 3 files changed, 68 insertions(+), 16 deletions(-) diff --git a/examples/src/main/kotlin/ai/koog/agents/example/structureddata/StructuredDataExample.kt b/examples/src/main/kotlin/ai/koog/agents/example/structureddata/StructuredDataExample.kt index c1fa4d984a..83711081a3 100644 --- a/examples/src/main/kotlin/ai/koog/agents/example/structureddata/StructuredDataExample.kt +++ b/examples/src/main/kotlin/ai/koog/agents/example/structureddata/StructuredDataExample.kt @@ -205,7 +205,7 @@ fun main(): Unit = runBlocking { // some models don't work well with json schema, so you may try simple, but it has more limitations (no polymorphism!) schemaFormat = JsonSchemaGenerator.SchemaFormat.JsonSchema, examples = exampleForecasts, - schemaType = JsonStructuredData.JsonSchemaType.SIMPLE + schemaType = JsonStructuredData.JsonSchemaType.FULL ) val agentStrategy = strategy("weather-forecast") { diff --git a/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/JsonSchemaGenerator.kt b/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/JsonSchemaGenerator.kt index 7566ceae72..6e52da3241 100644 --- a/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/JsonSchemaGenerator.kt +++ b/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/JsonSchemaGenerator.kt @@ -66,21 +66,21 @@ public class JsonSchemaGenerator( * Generate a JSON schema for a serializable class. * * @param serializer The serializer for the class - * @param descriptions Optional map of serial class names and property names to descriptions. + * @param descriptionOverrides Optional map of serial class names and property names to descriptions. * If a property/type is already described with [LLMDescription] annotation, value from the map will override this description. * @return A JsonObject representing the JSON schema */ public fun generate( id: String, serializer: KSerializer<*>, - descriptions: Map = emptyMap() + descriptionOverrides: Map = emptyMap() ): JsonObject { val rootSchema: JsonObject val definitions = buildJsonObject { rootSchema = generatePropertySchema( rootDefsBuilder = this, processedDefs = emptySet(), - descriptions = descriptions, + descriptionOverrides = descriptionOverrides, descriptor = serializer.descriptor, currentDepth = 0, ) @@ -118,7 +118,7 @@ public class JsonSchemaGenerator( private fun generatePropertySchema( rootDefsBuilder: JsonObjectBuilder, processedDefs: Set, - descriptions: Map, + descriptionOverrides: Map, descriptor: SerialDescriptor, currentDepth: Int, isPolymorphicSubtype: Boolean = false, @@ -155,7 +155,7 @@ public class JsonSchemaGenerator( generatePropertySchema( rootDefsBuilder = rootDefsBuilder, processedDefs = processedDefs, - descriptions = descriptions, + descriptionOverrides = descriptionOverrides, descriptor = itemDescriptor, currentDepth = currentDepth + 1, ) @@ -177,7 +177,7 @@ public class JsonSchemaGenerator( generatePropertySchema( rootDefsBuilder = rootDefsBuilder, processedDefs = processedDefs, - descriptions = descriptions, + descriptionOverrides = descriptionOverrides, descriptor = valueDescriptor, currentDepth = currentDepth + 1, ) @@ -208,14 +208,14 @@ public class JsonSchemaGenerator( // Description for a property val lookupKey = "${descriptor.serialName}.$propertyName" - val propertyDescriptionMap = descriptions[lookupKey] + val propertyDescriptionOverride = descriptionOverrides[lookupKey] val propertyDescriptionAnnotation = propertyAnnotations .filterIsInstance() .firstOrNull() ?.description // Look at the explicit map first, then at the annotation - val propertyDescription = propertyDescriptionMap ?: propertyDescriptionAnnotation + val propertyDescription = propertyDescriptionOverride ?: propertyDescriptionAnnotation put( propertyName, @@ -223,7 +223,7 @@ public class JsonSchemaGenerator( generatePropertySchema( rootDefsBuilder = rootDefsBuilder, processedDefs = updatedProcessedDefs, - descriptions = descriptions, + descriptionOverrides = descriptionOverrides, descriptor = propertyDescriptor, currentDepth = currentDepth + 1, ).let { propertySchema -> @@ -255,14 +255,14 @@ public class JsonSchemaGenerator( } // Description for a whole type (definition) - val typeDescriptionMap = descriptions[descriptor.serialName] + val typeDescriptionOverride = descriptionOverrides[descriptor.serialName] val typeDescriptionAnnotation = descriptor.annotations .filterIsInstance() .firstOrNull() ?.description // Look at the explicit map first, then at the annotation - val typeDescription = typeDescriptionMap ?: typeDescriptionAnnotation + val typeDescription = typeDescriptionOverride ?: typeDescriptionAnnotation // Build type definition val typeDefinition = buildJsonObject { @@ -295,7 +295,7 @@ public class JsonSchemaGenerator( generatePropertySchema( rootDefsBuilder = rootDefsBuilder, processedDefs = processedDefs, - descriptions = descriptions, + descriptionOverrides = descriptionOverrides, descriptor = polymorphicDescriptor, currentDepth = currentDepth + 1, isPolymorphicSubtype = true, diff --git a/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/JsonStructuredData.kt b/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/JsonStructuredData.kt index 3cdf144719..8b1005b2c3 100644 --- a/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/JsonStructuredData.kt +++ b/prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/JsonStructuredData.kt @@ -67,12 +67,64 @@ public class JsonStructuredData( /** * Factory method to create JSON structure with auto-generated JSON schema. * + * Example usage: + * ```kotlin + * @Serializable + * @SerialName("LatLon") + * @LLMDescription("Coordinates of the location in latitude and longitude format") + * data class LatLon( + * @LLMDescription("Latitude of the location") + * val lat: Double, + * @LLMDescription("Longitude of the location") + * val lon: Double + * ) + * + * @Serializable + * @SerialName("WeatherDatapoint") + * @LLMDescription("Weather datapoint for a given timestamp in the given location") + * data class WeatherDatapoint( + * @LLMDescription("Forecast timestamp") + * val timestampt: Long, + * @LLMDescription("Forecast temperature in Celsius") + * val temperature: Double, + * @LLMDescription("Precipitation in mm/h") + * val precipitation: Double, + * ) + * + * @Serializable + * @SerialName("Weather") + * data class Weather( + * @LLMDescription("Country code of the location") + * val countryCode: String, + * @LLMDescription("City name of the location") + * val cityName: String, + * @LLMDescription("Coordinates of the location") + * val latLon: LatLon, + * val forecast: List, + * ) + * + * val weatherStructure = JsonStructuredData.createJsonStructure( + * // some models don't work well with full json schema, so you may try simple, but it has more limitations (no polymorphism!) + * schemaFormat = JsonSchemaGenerator.SchemaFormat.JsonSchema, + * schemaType = JsonStructuredData.JsonSchemaType.FULL, + * descriptionOverrides = mapOf( + * // type descriptions + * "Weather" to "Weather forecast for a given location", // the class doesn't have description annotation, this will add description + * "WeatherDatapoint" to "Weather data at a given time", // the class has description annotation, this will override description + * + * // property descriptions + * "Weather.forecast" to "List of forecasted weather conditions for a given location", // the property doesn't have description annotation, this will add description + * "Weather.countryCode" to "Country code of the location in the ISO2 format", // the property has description annotation, this will override description + * ) + * ) + * ``` + * * @param id Unique identifier for the structure. * @param serializer Serializer used for converting the data to and from JSON. * @param json JSON configuration instance used for serialization. * @param schemaFormat Format of the generated schema, can be simple or detailed. * @param maxDepth Maximum recursion depth when generating schema to prevent infinite recursion for circular references. - * @param descriptions Optional map of serial class names and property names to descriptions. + * @param descriptionOverrides Optional map of serial class names and property names to descriptions. * If a property/type is already described with [LLMDescription] annotation, value from the map will override this description. * @param examples List of example data items that conform to the structure, used for demonstrating valid formats. * @param schemaType Type of JSON schema to generate, determines the level of detail in the schema. @@ -83,12 +135,12 @@ public class JsonStructuredData( json: Json = JsonStructureLanguage.defaultJson, schemaFormat: JsonSchemaGenerator.SchemaFormat = JsonSchemaGenerator.SchemaFormat.Simple, maxDepth: Int = 20, - descriptions: Map = emptyMap(), + descriptionOverrides: Map = emptyMap(), examples: List = emptyList(), schemaType: JsonSchemaType = JsonSchemaType.SIMPLE ): StructuredData { val structureLanguage = JsonStructureLanguage(json) - val schema = JsonSchemaGenerator(json, schemaFormat, maxDepth).generate(id, serializer, descriptions) + val schema = JsonSchemaGenerator(json, schemaFormat, maxDepth).generate(id, serializer, descriptionOverrides) return JsonStructuredData( id = id,