(#178) Fix structured output JSON descriptions

.gitignore

@@ -3,3 +3,4 @@
 .kotlin
 build
 **/.claude/settings.local.json
+env.properties

agents/agents-tools/src/commonMain/kotlin/ai/koog/agents/core/tools/annotations/LLMDescription.kt

@@ -1,15 +1,17 @@
 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,
-    AnnotationTarget.PROPERTY,
     AnnotationTarget.TYPE,
     AnnotationTarget.VALUE_PARAMETER,
     AnnotationTarget.FUNCTION

examples/src/main/kotlin/ai/koog/agents/example/structureddata/StructuredDataExample.kt

@@ -5,20 +5,20 @@ 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
 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
@@ -205,30 +205,21 @@ 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") {
         val setup by nodeLLMRequest()
-
-        val getStructuredForecast by node<Message.Response, String> { _ ->
-            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")

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,20 +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<String, String> = emptyMap()
+        descriptionOverrides: Map<String, String> = emptyMap()
     ): JsonObject {
         val rootSchema: JsonObject
         val definitions = buildJsonObject {
             rootSchema = generatePropertySchema(
                 rootDefsBuilder = this,
                 processedDefs = emptySet(),
-                descriptions = descriptions,
+                descriptionOverrides = descriptionOverrides,
                 descriptor = serializer.descriptor,
                 currentDepth = 0,
             )
@@ -116,7 +118,7 @@ public class JsonSchemaGenerator(
     private fun generatePropertySchema(
         rootDefsBuilder: JsonObjectBuilder,
         processedDefs: Set<String>,
-        descriptions: Map<String, String>,
+        descriptionOverrides: Map<String, String>,
         descriptor: SerialDescriptor,
         currentDepth: Int,
         isPolymorphicSubtype: Boolean = false,
@@ -153,7 +155,7 @@ public class JsonSchemaGenerator(
                         generatePropertySchema(
                             rootDefsBuilder = rootDefsBuilder,
                             processedDefs = processedDefs,
-                            descriptions = descriptions,
+                            descriptionOverrides = descriptionOverrides,
                             descriptor = itemDescriptor,
                             currentDepth = currentDepth + 1,
                         )
@@ -175,7 +177,7 @@ public class JsonSchemaGenerator(
                         generatePropertySchema(
                             rootDefsBuilder = rootDefsBuilder,
                             processedDefs = processedDefs,
-                            descriptions = descriptions,
+                            descriptionOverrides = descriptionOverrides,
                             descriptor = valueDescriptor,
                             currentDepth = currentDepth + 1,
                         )
@@ -202,16 +204,26 @@ 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 propertyDescriptionOverride = descriptionOverrides[lookupKey]
+                                val propertyDescriptionAnnotation = propertyAnnotations
+                                    .filterIsInstance<LLMDescription>()
+                                    .firstOrNull()
+                                    ?.description
+
+                                // Look at the explicit map first, then at the annotation
+                                val propertyDescription = propertyDescriptionOverride ?: propertyDescriptionAnnotation
 
                                 put(
                                     propertyName,
                                     JsonObject(
                                         generatePropertySchema(
                                             rootDefsBuilder = rootDefsBuilder,
                                             processedDefs = updatedProcessedDefs,
-                                            descriptions = descriptions,
+                                            descriptionOverrides = descriptionOverrides,
                                             descriptor = propertyDescriptor,
                                             currentDepth = currentDepth + 1,
                                         ).let { propertySchema ->
@@ -243,7 +255,14 @@ public class JsonSchemaGenerator(
                         }
 
                         // Description for a whole type (definition)
-                        val typeDescription = descriptions[descriptor.serialName]
+                        val typeDescriptionOverride = descriptionOverrides[descriptor.serialName]
+                        val typeDescriptionAnnotation = descriptor.annotations
+                            .filterIsInstance<LLMDescription>()
+                            .firstOrNull()
+                            ?.description
+
+                        // Look at the explicit map first, then at the annotation
+                        val typeDescription = typeDescriptionOverride ?: typeDescriptionAnnotation
 
                         // Build type definition
                         val typeDefinition = buildJsonObject {
@@ -276,7 +295,7 @@ public class JsonSchemaGenerator(
                                 generatePropertySchema(
                                     rootDefsBuilder = rootDefsBuilder,
                                     processedDefs = processedDefs,
-                                    descriptions = descriptions,
+                                    descriptionOverrides = descriptionOverrides,
                                     descriptor = polymorphicDescriptor,
                                     currentDepth = currentDepth + 1,
                                     isPolymorphicSubtype = true,