Add extension point for custom tool schemas (#1158)

<!--
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.
-->

## Motivation and Context

The aim of this PR is to introduce extension point for clients to give
ability to provide custom tool schema / add fields for existing schemas.
It might be useful for example in cases when you have almost OpenAI
compatible API, but with slight differences in tool schema.

## Breaking Changes
<!-- Will users need to update their code or configurations? -->
No
---

#### Type of the changes
- [x] New feature (non-breaking change which adds functionality)
- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Tests improvement
- [ ] Refactoring

#### Checklist
- [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
- [ ] 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
- [x] Docs have been added / updated

Author: Rizzen

agents/agents-features/agents-features-event-handler/src/jvmTest/kotlin/ai/koog/agents/features/eventHandler/feature/EventHandlerTest.kt

@@ -11,6 +11,7 @@ import ai.koog.agents.core.dsl.extension.nodeLLMSendToolResult
 import ai.koog.agents.core.dsl.extension.onAssistantMessage
 import ai.koog.agents.core.dsl.extension.onToolCall
 import ai.koog.agents.core.feature.handler.subgraph.SubgraphExecutionEventContext
+import ai.koog.agents.core.tools.ToolDescriptor
 import ai.koog.agents.core.tools.ToolRegistry
 import ai.koog.agents.features.eventHandler.eventString
 import ai.koog.agents.testing.tools.DummyTool
@@ -577,13 +578,13 @@ class EventHandlerTest {
             override suspend fun execute(
                 prompt: Prompt,
                 model: ai.koog.prompt.llm.LLModel,
-                tools: List<ai.koog.agents.core.tools.ToolDescriptor>
+                tools: List<ToolDescriptor>
             ): List<Message.Response> = emptyList()
 
             override fun executeStreaming(
                 prompt: Prompt,
                 model: ai.koog.prompt.llm.LLModel,
-                tools: List<ai.koog.agents.core.tools.ToolDescriptor>
+                tools: List<ToolDescriptor>
             ): Flow<StreamFrame> = flow {
                 throw IllegalStateException(testStreamingErrorMessage)
             }

agents/agents-features/agents-features-trace/src/jvmTest/kotlin/ai/koog/agents/features/tracing/writer/TraceFeatureMessageTestWriterTest.kt

@@ -19,6 +19,7 @@ import ai.koog.agents.core.feature.model.events.SubgraphExecutionFailedEvent
 import ai.koog.agents.core.feature.model.events.SubgraphExecutionStartingEvent
 import ai.koog.agents.core.feature.model.events.ToolCallCompletedEvent
 import ai.koog.agents.core.feature.model.events.ToolCallStartingEvent
+import ai.koog.agents.core.tools.ToolDescriptor
 import ai.koog.agents.core.tools.ToolRegistry
 import ai.koog.agents.core.utils.SerializationUtils
 import ai.koog.agents.features.tracing.feature.Tracing
@@ -443,13 +444,13 @@ class TraceFeatureMessageTestWriterTest {
             override suspend fun execute(
                 prompt: Prompt,
                 model: ai.koog.prompt.llm.LLModel,
-                tools: List<ai.koog.agents.core.tools.ToolDescriptor>
+                tools: List<ToolDescriptor>
             ): List<Message.Response> = emptyList()
 
             override fun executeStreaming(
                 prompt: Prompt,
                 model: ai.koog.prompt.llm.LLModel,
-                tools: List<ai.koog.agents.core.tools.ToolDescriptor>
+                tools: List<ToolDescriptor>
             ): Flow<StreamFrame> = flow {
                 val testException = IllegalStateException(testStreamingErrorMessage)
                 testStreamingStackTrace = testException.stackTraceToString()

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

@@ -0,0 +1,68 @@
+package ai.koog.agents.core.tools
+
+/**
+ * Represents a descriptor for a tool that contains information about the tool's name, description, required parameters,
+ * and optional parameters.
+ *
+ * This class is annotated with @Serializable to support serialization/deserialization using kotlinx.serialization.
+ *
+ * @property name The name of the tool.
+ * @property description The description of the tool.
+ * @property requiredParameters A list of ToolParameterDescriptor representing the required parameters for the tool.
+ * @property optionalParameters A list of ToolParameterDescriptor representing the optional parameters for the tool.
+ */
+public open class ToolDescriptor(
+    public val name: String,
+    public val description: String,
+    public val requiredParameters: List<ToolParameterDescriptor> = emptyList(),
+    public val optionalParameters: List<ToolParameterDescriptor> = emptyList(),
+) {
+    /**
+     * Creates a copy of the current ToolDescriptor with the option to modify specific attributes.
+     *
+     * @param name The name of the tool. Defaults to the current tool's name if not provided.
+     * @param description The description of the tool. Defaults to the current tool's description if not provided.
+     * @param requiredParameters A list of ToolParameterDescriptor representing the required parameters for the tool.
+     * Defaults to the current required parameters if not provided.
+     * @param optionalParameters A list of ToolParameterDescriptor representing the optional parameters for the tool.
+     * Defaults to the current optional parameters if not provided.
+     * @return A new instance of ToolDescriptor with the updated attributes.
+     */
+    public fun copy(
+        name: String = this.name,
+        description: String = this.description,
+        requiredParameters: List<ToolParameterDescriptor> = this.requiredParameters.toList(),
+        optionalParameters: List<ToolParameterDescriptor> = this.optionalParameters.toList(),
+    ): ToolDescriptor {
+        return ToolDescriptor(
+            name = name,
+            description = description,
+            requiredParameters = requiredParameters,
+            optionalParameters = optionalParameters,
+        )
+    }
+
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (other !is ToolDescriptor) return false
+
+        if (name != other.name) return false
+        if (description != other.description) return false
+        if (requiredParameters != other.requiredParameters) return false
+        if (optionalParameters != other.optionalParameters) return false
+
+        return true
+    }
+
+    override fun toString(): String {
+        return "ToolDescriptor(name=$name, description=$description, requiredParameters=$requiredParameters, optionalParameters=$optionalParameters)"
+    }
+
+    override fun hashCode(): Int {
+        var result = name.hashCode()
+        result = 31 * result + description.hashCode()
+        result = 31 * result + requiredParameters.hashCode()
+        result = 31 * result + optionalParameters.hashCode()
+        return result
+    }
+}

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

@@ -2,24 +2,6 @@ package ai.koog.agents.core.tools
 
 import kotlin.enums.EnumEntries
 
-/**
- * Represents a descriptor for a tool that contains information about the tool's name, description, required parameters,
- * and optional parameters.
- *
- * This class is annotated with @Serializable to support serialization/deserialization using kotlinx.serialization.
- *
- * @property name The name of the tool.
- * @property description The description of the tool.
- * @property requiredParameters A list of ToolParameterDescriptor representing the required parameters for the tool.
- * @property optionalParameters A list of ToolParameterDescriptor representing the optional parameters for the tool.
- */
-public data class ToolDescriptor(
-    val name: String,
-    val description: String,
-    val requiredParameters: List<ToolParameterDescriptor> = emptyList(),
-    val optionalParameters: List<ToolParameterDescriptor> = emptyList(),
-)
-
 /**
  * Represents a descriptor for a tool parameter.
  * A tool parameter descriptor contains information about a specific tool parameter, such as its name, description,

agents/agents-tools/src/commonMain/kotlin/ai/koog/agents/core/tools/serialization/ToolDescriptorSchemaGenerator.kt

@@ -0,0 +1,17 @@
+package ai.koog.agents.core.tools.serialization
+
+import ai.koog.agents.core.tools.ToolDescriptor
+import kotlinx.serialization.json.JsonObject
+
+/**
+ * Interface for converting a list of ToolDescriptor objects to a JSON schema representation.
+ */
+public interface ToolDescriptorSchemaGenerator {
+    /**
+     * Converts a list of ToolDescriptor objects into a JSON object representation.
+     *
+     * @param toolDescriptor The ToolDescriptor to convert.
+     * @return A JsonObject containing the JSON representation of the provided list of ToolDescriptor objects.
+     */
+    public fun generate(toolDescriptor: ToolDescriptor): JsonObject
+}

agents/agents-tools/src/commonTest/kotlin/ai/koog/agents/core/tools/SerialToToolDescriptionTest.kt

@@ -272,10 +272,10 @@ class SerialToToolDescriptionTest {
     fun verify_optional_description_applies() {
         val tripPlanDescriptor = serializer<TripPlan>().descriptor.asToolDescriptor(
             toolName = "provideTripPlan",
-            toolDescription = "Custom tool ,call me!"
+            toolDescription = "Custom tool, call me!"
         )
 
-        assertEquals("Custom tool ,call me!", tripPlanDescriptor.description)
-        assertEquals(expectedTripPlanToolDescriptor.copy(description = "Custom tool ,call me!"), tripPlanDescriptor)
+        assertEquals("Custom tool, call me!", tripPlanDescriptor.description)
+        assertEquals(expectedTripPlanToolDescriptor.copy(description = "Custom tool, call me!"), tripPlanDescriptor)
     }
 }

docs/docs/tools/tool-descriptor-schemer.md

@@ -0,0 +1,176 @@
+# ToolDescriptorSchemer
+
+## What is it
+
+`ToolDescriptorSchemer` is a extension point that converts a `ToolDescriptor` into a JSON Schema object compatible with specific LLM providers.
+
+Key points:
+
+- Location: `ai.koog.agents.core.tools.serialization.ToolDescriptorSchemer`
+- Contract: a single function `scheme(toolDescriptor: ToolDescriptor): JsonObject`
+- Implementations provided:
+  - `OpenAICompatibleToolDescriptorSchemer` — generates schemas compatible with OpenAI‑style function/tool definitions.
+  - `OllamaToolDescriptorSchemer` — generates schemas compatible with Ollama tool JSON.
+
+
+<!--- INCLUDE
+import ai.koog.agents.core.tools.ToolDescriptor
+import ai.koog.agents.core.tools.ToolParameterDescriptor
+import ai.koog.agents.core.tools.ToolParameterType
+import kotlinx.serialization.json.JsonObject
+-->
+```kotlin
+// Interface
+interface ToolDescriptorSchemaGenerator {
+  fun generate(toolDescriptor: ToolDescriptor): JsonObject
+}
+```
+<!--- KNIT example-tool-descriptor-schemer-01.kt -->
+
+## Why to use it?
+If you want to provide custom scheme for existing or new LLM providers, implement this interface to convert Koog’s `ToolDescriptor` into the expected JSON Schema format.
+
+## Implementation example
+
+Below is a minimal custom implementation that renders only a subset of parameter types to illustrate how to plug into the SPI. Real implementations should cover all `ToolParameterType`s (String, Integer, Float, Boolean, Null, Enum, List, Object, AnyOf).
+
+<!--- INCLUDE
+import ai.koog.agents.core.tools.ToolDescriptor
+import ai.koog.agents.core.tools.ToolParameterDescriptor
+import ai.koog.agents.core.tools.ToolParameterType
+import ai.koog.agents.core.tools.serialization.ToolDescriptorSchemaGenerator
+import kotlinx.serialization.json.JsonPrimitive
+import kotlinx.serialization.json.JsonObject
+import kotlinx.serialization.json.buildJsonObject
+import kotlinx.serialization.json.put
+import kotlinx.serialization.json.putJsonArray
+import kotlinx.serialization.json.putJsonObject
+-->
+```kotlin
+
+class MinimalSchemer : ToolDescriptorSchemaGenerator {
+    override fun generate(toolDescriptor: ToolDescriptor): JsonObject = buildJsonObject {
+        put("type", "object")
+        putJsonObject("properties") {
+            (toolDescriptor.requiredParameters + toolDescriptor.optionalParameters).forEach { p ->
+                put(p.name, buildJsonObject {
+                    put("description", p.description)
+                    when (val t = p.type) {
+                        ToolParameterType.String -> put("type", "string")
+                        ToolParameterType.Integer -> put("type", "integer")
+                        is ToolParameterType.Enum -> {
+                            put("type", "string")
+                            putJsonArray("enum") { t.entries.forEach { add(JsonPrimitive(it)) } }
+                        }
+                        else -> put("type", "string") // fallback for brevity
+                    }
+                })
+            }
+        }
+        putJsonArray("required") { toolDescriptor.requiredParameters.forEach { add(JsonPrimitive(it.name)) } }
+    }
+}
+```
+<!--- KNIT example-tool-descriptor-schemer-02.kt -->
+
+## Example of usage with client
+
+Typically you do not need to call a schemer directly. Koog clients accept a list of `ToolDescriptor` objects and apply the correct schemer internally when serializing requests for the provider.
+
+The example below defines a simple tool and passes it to the OpenAI client. The client will use `OpenAICompatibleToolDescriptorSchemer` under the hood to build the JSON schema.
+
+<!--- INCLUDE
+import ai.koog.prompt.executor.clients.openai.OpenAILLMClient
+import ai.koog.agents.core.tools.ToolDescriptor
+import ai.koog.agents.core.tools.ToolParameterDescriptor
+import ai.koog.agents.core.tools.ToolParameterType
+import ai.koog.prompt.dsl.Prompt
+import ai.koog.prompt.executor.clients.openai.OpenAIModels
+import ai.koog.prompt.executor.clients.openai.base.OpenAICompatibleToolDescriptorSchemaGenerator
+import kotlinx.serialization.json.JsonPrimitive
+import kotlinx.serialization.json.JsonObject
+import kotlinx.serialization.json.buildJsonObject
+import kotlinx.serialization.json.put
+import kotlinx.serialization.json.putJsonArray
+import kotlinx.serialization.json.putJsonObject
+import kotlinx.coroutines.runBlocking
+
+class MinimalSchemer : OpenAICompatibleToolDescriptorSchemaGenerator() {
+    override fun generate(toolDescriptor: ToolDescriptor): JsonObject = buildJsonObject {
+        put("type", "object")
+        putJsonObject("properties") {
+            (toolDescriptor.requiredParameters + toolDescriptor.optionalParameters).forEach { p ->
+                put(p.name, buildJsonObject {
+                    put("description", p.description)
+                    when (val t = p.type) {
+                        ToolParameterType.String -> put("type", "string")
+                        ToolParameterType.Integer -> put("type", "integer")
+                        is ToolParameterType.Enum -> {
+                            put("type", "string")
+                            putJsonArray("enum") { t.entries.forEach { add(JsonPrimitive(it)) } }
+                        }
+                        else -> put("type", "string") // fallback for brevity
+                    }
+                })
+            }
+        }
+        putJsonArray("required") { toolDescriptor.requiredParameters.forEach { add(JsonPrimitive(it.name)) } }
+    }
+}
+
+-->
+```kotlin
+val client = OpenAILLMClient(apiKey = System.getenv("OPENAI_API_KEY"), toolsConverter = MinimalSchemer())
+
+val getUserTool = ToolDescriptor(
+    name = "get_user",
+    description = "Returns user profile by id",
+    requiredParameters = listOf(
+        ToolParameterDescriptor(
+            name = "id",
+            description = "User id",
+            type = ToolParameterType.String
+        )
+    )
+)
+
+val prompt = Prompt.build(id = "p1") { user("Hello") }
+val responses = runBlocking {
+    client.execute(
+        prompt = prompt,
+        model = OpenAIModels.Chat.GPT4o,
+        tools = listOf(getUserTool)
+    )
+}
+```
+<!--- KNIT example-tool-descriptor-schemer-03.kt -->
+
+If you need direct access to the produced schema (for debugging or for a custom transport), you can instantiate the provider‑specific schemer and serialize the JSON yourself:
+
+<!--- INCLUDE
+import kotlinx.serialization.json.Json
+import ai.koog.prompt.executor.clients.openai.base.OpenAICompatibleToolDescriptorSchemaGenerator
+import ai.koog.agents.core.tools.ToolDescriptor
+import ai.koog.agents.core.tools.ToolParameterDescriptor
+import ai.koog.agents.core.tools.ToolParameterType
+
+fun getUserTool(): ToolDescriptor {
+    return ToolDescriptor(
+        name = "get_user",
+        description = "Returns user profile by id",
+        requiredParameters = listOf(
+            ToolParameterDescriptor(
+                name = "id",
+                description = "User id",
+                type = ToolParameterType.String
+            )
+        )
+    )
+}
+-->
+
+```kotlin
+val json = Json { prettyPrint = true }
+val schema = OpenAICompatibleToolDescriptorSchemaGenerator().generate(getUserTool())
+```
+<!--- KNIT example-tool-descriptor-schemer-04.kt -->

docs/mkdocs.yml

@@ -19,6 +19,7 @@ nav:
           - Built-in tools: built-in-tools.md
           - Annotation-based tools: annotation-based-tools.md
           - Class-based tools: class-based-tools.md
+          - Custom tools schema: tools/tool-descriptor-schemer.md
       - Events: agent-events.md
       - Strategies:
           - Pre-defined nodes and components: nodes-and-components.md