Add support for excluding properties in JSON schema generation (JetBrains#638)

Fixes JetBrains#377 
Changes are the same as JetBrains#378, it's just the updated version after the
refactor of JetBrains#443

Proposed changes:
Add an `excludedProperties` parameter like this:
```kotlin
simpleSchemaGenerator.generate(
    "TestClass",
    serializer<TestClass>(),
    descriptionOverrides = emptyMap(),
    excludedProperties = setOf("TestClass.someProperty")
)
```

---

#### Type of the change
- [x] New feature
- [ ] Bug fix
- [ ] Documentation fix

#### Checklist for all pull requests
- [x] The pull request has a description of the proposed change
- [x] I read the [Contributing
Guidelines](https://github.com/JetBrains/koog/blob/main/CONTRIBUTING.md)
before opening the pull request
- [x] The pull request uses **`develop`** as the base branch
- [x] Tests for the changes have been added
- [x] All new and existing tests passed

##### Additional steps for pull requests adding a new feature
- [x] An issue describing the proposed change exists
- [x] The pull request includes a link to the issue
- [ ] The change was discussed and approved in the issue
- [ ] Docs have been added / updated

Author: nathanfallet

prompt/prompt-executor/prompt-executor-clients/prompt-executor-openai-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/openai/structure/OpenAIStandardJsonSchemaGenerator.kt

@@ -36,9 +36,10 @@ public open class OpenAIStandardJsonSchemaGenerator : StandardJsonSchemaGenerato
         json: Json,
         name: String,
         serializer: KSerializer<*>,
-        descriptionOverrides: Map<String, String>
+        descriptionOverrides: Map<String, String>,
+        excludedProperties: Set<String>,
     ): LLMParams.Schema.JSON.Standard {
-        val param = super.generate(json, name, serializer, descriptionOverrides)
+        val param = super.generate(json, name, serializer, descriptionOverrides, excludedProperties)
         val schema = param.schema.toMutableMap()
 
         // OpenAI doesn't accept "$ref" at the root of the schema, so copying this definition explicitly.

prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/JsonStructuredData.kt

@@ -167,6 +167,7 @@ public class JsonStructuredData<TStruct>(
          * @param schemaGenerator JSON schema generator
          * @param descriptionOverrides Optional map of serial class names and property names to descriptions.
          * If a property/type is already described with [ai.koog.agents.core.tools.annotations.LLMDescription] annotation, value from the map will override this description.
+         * @param excludedProperties Optional set of property names to exclude from the schema generation.
          * @param examples List of example data items that conform to the structure, used for demonstrating valid formats.
          * @param definitionPrompt Prompt with definition, explaining the structure to the LLM when the manual mode for
          * structured output is used. Default is [JsonStructuredData.defaultDefinitionPrompt]
@@ -177,6 +178,7 @@ public class JsonStructuredData<TStruct>(
             json: Json = defaultJson,
             schemaGenerator: JsonSchemaGenerator = StandardJsonSchemaGenerator.Default,
             descriptionOverrides: Map<String, String> = emptyMap(),
+            excludedProperties: Set<String> = emptySet(),
             examples: List<TStruct> = emptyList(),
             definitionPrompt: (
                 builder: TextContentBuilderBase<*>,
@@ -185,7 +187,7 @@ public class JsonStructuredData<TStruct>(
         ): JsonStructuredData<TStruct> {
             return JsonStructuredData(
                 id = id,
-                schema = schemaGenerator.generate(json, id, serializer, descriptionOverrides),
+                schema = schemaGenerator.generate(json, id, serializer, descriptionOverrides, excludedProperties),
                 examples = examples,
                 serializer = serializer,
                 json = json,

prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/generator/BasicJsonSchemaGenerator.kt

@@ -32,7 +32,8 @@ public open class BasicJsonSchemaGenerator : GenericJsonSchemaGenerator() {
         json: Json,
         name: String,
         serializer: KSerializer<*>,
-        descriptionOverrides: Map<String, String>
+        descriptionOverrides: Map<String, String>,
+        excludedProperties: Set<String>,
     ): LLMParams.Schema.JSON.Basic {
         val descriptorKind = serializer.descriptor.kind
 
@@ -46,6 +47,7 @@ public open class BasicJsonSchemaGenerator : GenericJsonSchemaGenerator() {
             processedTypeDefs = mutableMapOf(),
             currentDefPath = listOf(),
             descriptionOverrides = descriptionOverrides,
+            excludedProperties = excludedProperties,
             currentDescription = null,
         )
 

prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/generator/GenericJsonSchemaGenerator.kt

@@ -150,6 +150,15 @@ public abstract class GenericJsonSchemaGenerator : JsonSchemaGenerator() {
                     val propertyName = context.descriptor.getElementName(i)
                     val propertyDescriptor = context.descriptor.getElementDescriptor(i)
 
+                    // Check if the property is excluded
+                    val lookupKey = "${context.descriptor.serialName}.$propertyName"
+                    if (context.excludedProperties.contains(lookupKey)) {
+                        if (!context.descriptor.isElementOptional(i)) {
+                            throw IllegalArgumentException("Property '$lookupKey' is marked as excluded, but it is required in the schema.")
+                        }
+                        continue
+                    }
+
                     put(
                         propertyName,
                         process(

prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/generator/JsonSchemaGenerator.kt

@@ -31,6 +31,7 @@ public abstract class JsonSchemaGenerator {
      * during schema generation processing.
      * @property descriptionOverrides A map of user-defined properties and types descriptions to override [LLMDescription]
      * from the provided class.
+     * @property excludedProperties A set of property names to exclude from the schema generation.
      * @property currentDescription Description for the current element
      */
     public data class GenerationContext(
@@ -39,6 +40,7 @@ public abstract class JsonSchemaGenerator {
         public val processedTypeDefs: MutableMap<SerialDescriptor, JsonObject>,
         public val currentDefPath: List<SerialDescriptor>,
         public val descriptionOverrides: Map<String, String>,
+        public val excludedProperties: Set<String>,
         public val currentDescription: String?,
     ) {
         /**
@@ -89,12 +91,14 @@ public abstract class JsonSchemaGenerator {
      * @param name The name of the schema.
      * @param serializer The serializer for the type for which the schema has to be generated.
      * @param descriptionOverrides A map containing overrides for [LLMDescription].
+     * @param excludedProperties A set of property names to exclude from the schema generation.
      */
     public abstract fun generate(
         json: Json,
         name: String,
         serializer: KSerializer<*>,
         descriptionOverrides: Map<String, String>,
+        excludedProperties: Set<String> = emptySet(),
     ): LLMParams.Schema.JSON
 
     /**

prompt/prompt-structure/src/commonMain/kotlin/ai/koog/prompt/structure/json/generator/StandardJsonSchemaGenerator.kt

@@ -39,7 +39,8 @@ public open class StandardJsonSchemaGenerator : GenericJsonSchemaGenerator() {
         json: Json,
         name: String,
         serializer: KSerializer<*>,
-        descriptionOverrides: Map<String, String>
+        descriptionOverrides: Map<String, String>,
+        excludedProperties: Set<String>,
     ): LLMParams.Schema.JSON.Standard {
         val descriptorKind = serializer.descriptor.kind
 
@@ -56,6 +57,7 @@ public open class StandardJsonSchemaGenerator : GenericJsonSchemaGenerator() {
             processedTypeDefs = mutableMapOf(),
             currentDefPath = listOf(),
             descriptionOverrides = descriptionOverrides,
+            excludedProperties = excludedProperties,
             currentDescription = null
         )
 

prompt/prompt-structure/src/commonTest/kotlin/ai/koog/prompt/structure/json/generator/JsonSchemaGeneratorTest.kt

@@ -267,6 +267,64 @@ class JsonSchemaGeneratorTest {
         assertEquals(expectedSchema, schema)
     }
 
+    @Test
+    fun testGenerateBasicSchemaExcludingProperties() {
+        val result = basicGenerator.generate(json, "TestClass", serializer<TestClass>(), emptyMap(), setOf("TestClass.nullableProperty"))
+        val schema = json.encodeToString(result.schema)
+
+        val expectedSchema = """
+            {
+              "type": "object",
+              "properties": {
+                "stringProperty": {
+                  "type": "string",
+                  "description": "A string property"
+                },
+                "intProperty": {
+                  "type": "integer"
+                },
+                "booleanProperty": {
+                  "type": "boolean"
+                },
+                "listProperty": {
+                  "type": "array",
+                  "items": {
+                    "type": "string"
+                  }
+                },
+                "mapProperty": {
+                  "type": "object",
+                  "additionalProperties": {
+                    "type": "integer"
+                  }
+                }
+              },
+              "required": [
+                "stringProperty",
+                "intProperty",
+                "booleanProperty"
+              ],
+              "additionalProperties": false
+            }
+        """.trimIndent()
+
+        assertEquals(expectedSchema, schema)
+    }
+
+    @Test
+    fun testGenerateBasicSchemaExcludingRequiredProperties() {
+        val exception = assertFailsWith<IllegalArgumentException> {
+            basicGenerator.generate(
+                json,
+                "TestClass",
+                serializer<TestClass>(),
+                emptyMap(),
+                setOf("TestClass.stringProperty")
+            )
+        }
+        assertEquals("Property 'TestClass.stringProperty' is marked as excluded, but it is required in the schema.", exception.message)
+    }
+
     @Test
     fun testGenerateStandardSchemaWithDescriptions() {
         val descriptions = mapOf(