Skip to content

Commit 504ea71

Browse files
tiginamariageorge.lemeshko
tiginamaria
and
george.lemeshko
authored
Support MCP enum arg types and object additionalParameters (#214)
Co-authored-by: george.lemeshko <[email protected]>
1 parent 295f0a9 commit 504ea71

File tree

8 files changed

+463
-40
lines changed

8 files changed

+463
-40
lines changed

‎agents/agents-mcp/Module.md‎

Lines changed: 136 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,139 @@
11
# Module agents-mcp
22

3-
Provides facilities to integrate agents with Model Context Protocol (MCP) servers via Tools API.
3+
A module provides integration with [Model Context Protocol (MCP)](https://modelcontextprotocol.io) servers.
4+
The main components of the MCP integration in Koog are:
5+
- [**McpToolRegistryProvider**](src/jvmMain/kotlin/ai/koog/agents/mcp/McpToolRegistryProvider.kt): Creates tool registries that connect to MCP servers
6+
- [**McpTool**](src/jvmMain/kotlin/ai/koog/agents/mcp/McpTool.kt): A bridge between the Koog agent framework's Tool interface and the MCP SDK
7+
- [**McpToolDescriptorParser**](src/jvmMain/kotlin/ai/koog/agents/mcp/McpToolDefinitionParser.kt): Parses tool definitions from the MCP SDK to the Koog tool descriptor format
48

5-
<!-- TODO -->
9+
10+
## Overview
11+
12+
### What is MCP?
13+
14+
The Model Context Protocol (MCP) is a standardized protocol that enables AI agents to interact with external tools and services through a consistent interface.
15+
MCP works by exposing tools and prompts as API endpoints that can be called by AI agents.
16+
Each tool has a defined name and input schema that describes its inputs and outputs in JSON SHEMA format.
17+
To read more about MCP visit [https://modelcontextprotocol.io](https://modelcontextprotocol.io)
18+
19+
### How to use MCP servers?
20+
You can find ready-to-use mcp servers in the [MCP Marketplace](https://mcp.so/) or [MCP DockerHub](https://hub.docker.com/u/mcp).
21+
MCP servers support stdio transport and optionally sse transport protocols to communicate with the agent.
22+
23+
### How MCP is integrated with Koog?
24+
25+
The Koog framework integrates with MCP using the [MCP SDK](https://github.com/modelcontextprotocol/kotlin-sdk) with the additional api extensions presented in module `agent-mcp`.
26+
This integration allows Koog agents to:
27+
28+
1. Connect to MCP servers through various transport mechanisms (stdio, SSE)
29+
2. Retrieve available tools from the MCP server
30+
3. Transform MCP tools into the Koog agent framework's Tool interface
31+
4. Register the transformed tools in a ToolRegistry
32+
5. Call MCP tools with arguments provided by the LLM
33+
34+
### How to Use MCP with Koog?
35+
36+
#### Setting Up an MCP Connection
37+
38+
To use MCP with Koog, you need to:
39+
40+
1. Start an MCP server (either as a process, Docker container, or web service)
41+
2. Create a transport to communicate with the server
42+
3. Create a ToolRegistry with tools from the MCP server
43+
4. Use the tools in an AI agent
44+
45+
Here's a basic example of setting up an MCP connection:
46+
47+
```kotlin
48+
// Start the MCP server (e.g., as a process)
49+
val process = ProcessBuilder("path/to/mcp/server").start()
50+
51+
// Create a ToolRegistry with tools from the MCP server
52+
val toolRegistry = McpToolRegistryProvider.fromTransport(
53+
transport = McpToolRegistryProvider.defaultStdioTransport(process)
54+
)
55+
56+
// Use the tools in an AI agent
57+
val agent = AIAgent(
58+
promptExecutor = executor,
59+
strategy = strategy,
60+
agentConfig = agentConfig,
61+
toolRegistry = toolRegistry
62+
)
63+
64+
// Run the agent
65+
agent.runAndGetResult("Your task here")
66+
```
67+
68+
#### Transport Types
69+
70+
MCP supports different transport mechanisms for communication:
71+
72+
##### Standard Input/Output (stdio)
73+
74+
Use stdio transport when the MCP server is running as a separate process:
75+
76+
```kotlin
77+
val process = ProcessBuilder("path/to/mcp/server").start()
78+
val transport = McpToolRegistryProvider.defaultStdioTransport(process)
79+
```
80+
81+
##### Server-Sent Events (SSE)
82+
83+
Use SSE transport when the MCP server is running as a web service:
84+
85+
```kotlin
86+
val transport = McpToolRegistryProvider.defaultSseTransport("http://localhost:8931")
87+
```
88+
89+
### Examples
90+
91+
#### Google Maps MCP Integration
92+
93+
This example demonstrates using MCP to connect to a [Google Maps](https://mcp.so/server/google-maps/modelcontextprotocol) server for geographic data:
94+
95+
```kotlin
96+
// Start the Docker container with the Google Maps MCP server
97+
val process = ProcessBuilder(
98+
"docker", "run", "-i",
99+
"-e", "GOOGLE_MAPS_API_KEY=$googleMapsApiKey",
100+
"mcp/google-maps"
101+
).start()
102+
103+
// Create the ToolRegistry with tools from the MCP server
104+
val toolRegistry = McpToolRegistryProvider.fromTransport(
105+
transport = McpToolRegistryProvider.defaultStdioTransport(process)
106+
)
107+
108+
// Create and run the agent
109+
val agent = simpleSingleRunAgent(
110+
executor = simpleOpenAIExecutor(openAIApiToken),
111+
llmModel = OpenAIModels.Chat.GPT4o,
112+
toolRegistry = toolRegistry,
113+
)
114+
agent.run("Get elevation of the Jetbrains Office in Munich, Germany?")
115+
```
116+
117+
#### Playwright MCP Integration
118+
119+
This example demonstrates using MCP to connect to a [Playwright](https://mcp.so/server/playwright-mcp/microsoft) server for web automation:
120+
121+
```kotlin
122+
// Start the Playwright MCP server
123+
val process = ProcessBuilder(
124+
"npx", "@playwright/mcp@latest", "--port", "8931"
125+
).start()
126+
127+
// Create the ToolRegistry with tools from the MCP server
128+
val toolRegistry = McpToolRegistryProvider.fromTransport(
129+
transport = McpToolRegistryProvider.defaultSseTransport("http://localhost:8931")
130+
)
131+
132+
// Create and run the agent
133+
val agent = simpleSingleRunAgent(
134+
executor = simpleOpenAIExecutor(openAIApiToken),
135+
llmModel = OpenAIModels.Chat.GPT4o,
136+
toolRegistry = toolRegistry,
137+
)
138+
agent.run("Open a browser, navigate to jetbrains.com, accept all cookies, click AI in toolbar")
139+
```

‎agents/agents-mcp/src/jvmMain/kotlin/ai/koog/agents/mcp/McpToolDefinitionParser.kt‎

Lines changed: 43 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,7 @@ package ai.koog.agents.mcp
33
import ai.koog.agents.core.tools.ToolDescriptor
44
import ai.koog.agents.core.tools.ToolParameterDescriptor
55
import ai.koog.agents.core.tools.ToolParameterType
6-
import kotlinx.serialization.json.JsonObject
7-
import kotlinx.serialization.json.jsonObject
8-
import kotlinx.serialization.json.jsonPrimitive
9-
import kotlin.collections.component1
10-
import kotlin.collections.component2
11-
import kotlin.collections.contains
6+
import kotlinx.serialization.json.*
127
import io.modelcontextprotocol.kotlin.sdk.Tool as SDKTool
138

149
/**
@@ -49,11 +44,8 @@ public object DefaultMcpToolDescriptorParser : McpToolDescriptorParser {
4944

5045
private fun parseParameterType(element: JsonObject): ToolParameterType {
5146
// Extract the type string from the JSON object
52-
val typeStr = if ("type" in element) {
53-
element.getValue("type").jsonPrimitive.content
54-
} else {
55-
throw IllegalArgumentException("Parameter type must have type property")
56-
}
47+
val typeStr = element["type"]?.jsonPrimitive?.content
48+
?: throw IllegalArgumentException("Parameter type must have type property")
5749

5850
// Convert the type string to a ToolParameterType
5951
return when (typeStr.lowercase()) {
@@ -62,31 +54,59 @@ public object DefaultMcpToolDescriptorParser : McpToolDescriptorParser {
6254
"integer" -> ToolParameterType.Integer
6355
"number" -> ToolParameterType.Float
6456
"boolean" -> ToolParameterType.Boolean
57+
"enum" -> ToolParameterType.Enum(
58+
element.getValue("enum").jsonArray.map { it.jsonPrimitive.content }.toTypedArray()
59+
)
6560

6661
// Array type
6762
"array" -> {
68-
val items = if ("items" in element) {
69-
element.getValue("items").jsonObject
70-
} else {
71-
throw IllegalArgumentException("Array type parameters must have items property")
72-
}
63+
val items = element["items"]?.jsonObject
64+
?: throw IllegalArgumentException("Array type parameters must have items property")
65+
7366
val itemType = parseParameterType(items)
7467

7568
ToolParameterType.List(itemsType = itemType)
7669
}
7770

7871
// Object type
7972
"object" -> {
80-
val properties = if ("properties" in element) {
81-
element.getValue("properties").jsonObject
73+
val properties = element["properties"]?.let { properties ->
74+
val rawProperties = properties.jsonObject
75+
rawProperties.map { (name, property) ->
76+
// Description is optional
77+
val description = element["description"]?.jsonPrimitive?.content.orEmpty()
78+
ToolParameterDescriptor(name, description, parseParameterType(property.jsonObject))
79+
}
80+
} ?: emptyList()
81+
82+
val required = element["required"]?.jsonArray?.map { it.jsonPrimitive.content } ?: emptyList()
83+
84+
85+
val additionalProperties = if ("additionalProperties" in element) {
86+
when (element.getValue("additionalProperties")) {
87+
is JsonPrimitive -> element.getValue("additionalProperties").jsonPrimitive.boolean
88+
is JsonObject -> true
89+
else -> null
90+
}
91+
} else {
92+
null
93+
}
94+
95+
val additionalPropertiesType = if ("additionalProperties" in element) {
96+
when (element.getValue("additionalProperties")) {
97+
is JsonObject -> parseParameterType(element.getValue("additionalProperties").jsonObject)
98+
else -> null
99+
}
82100
} else {
83-
throw IllegalArgumentException("Object type parameters must have properties property")
101+
null
84102
}
85103

86-
ToolParameterType.Object(properties.map { (name, property) ->
87-
val description = element["description"]?.jsonPrimitive?.content.orEmpty()
88-
ToolParameterDescriptor(name, description, parseParameterType(property.jsonObject))
89-
})
104+
ToolParameterType.Object(
105+
properties = properties,
106+
requiredProperties = required,
107+
additionalPropertiesType = additionalPropertiesType,
108+
additionalProperties = additionalProperties
109+
)
90110
}
91111

92112
// Unsupported type

‎agents/agents-mcp/src/jvmMain/kotlin/ai/koog/agents/mcp/McpToolRegistryProvider.kt‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package ai.koog.agents.mcp
22

33
import ai.koog.agents.core.tools.ToolRegistry
4+
import io.github.oshai.kotlinlogging.KotlinLogging
45
import io.ktor.client.*
56
import io.ktor.client.plugins.sse.*
67
import io.modelcontextprotocol.kotlin.sdk.Implementation
@@ -22,6 +23,8 @@ import kotlinx.io.buffered
2223
* 4. Registering the transformed tools in a ToolRegistry
2324
*/
2425
public object McpToolRegistryProvider {
26+
private val logger = KotlinLogging.logger (McpToolRegistryProvider::class.qualifiedName!!)
27+
2528
/**
2629
* Default name for the MCP client when connecting to an MCP server.
2730
*/
@@ -78,8 +81,12 @@ public object McpToolRegistryProvider {
7881
val sdkTools = mcpClient.listTools()?.tools.orEmpty()
7982
return ToolRegistry {
8083
sdkTools.forEach { sdkTool ->
81-
val toolDescriptor = mcpToolParser.parse(sdkTool)
82-
tool(McpTool(mcpClient, toolDescriptor))
84+
try {
85+
val toolDescriptor = mcpToolParser.parse(sdkTool)
86+
tool(McpTool(mcpClient, toolDescriptor))
87+
} catch (e: Throwable) {
88+
logger.error(e) { "Failed to parse descriptor parameters for tool: ${sdkTool.name}" }
89+
}
8390
}
8491
}
8592
}

0 commit comments

Comments
 (0)