Skip to content

Commit 87dbad4

Browse files
committed
refactoring
1 parent e2c0180 commit 87dbad4

File tree

6 files changed

+135
-149
lines changed

6 files changed

+135
-149
lines changed

activators/llm/src/main/kotlin/com/justai/jaicf/activator/llm/LLMProps.kt

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,6 @@ package com.justai.jaicf.activator.llm
33
import com.justai.jaicf.activator.llm.builder.JsonSchemaBuilder
44
import com.justai.jaicf.activator.llm.mcp.McpService
55
import com.justai.jaicf.activator.llm.mcp.McpServiceResponseBuilder
6-
import com.justai.jaicf.activator.llm.mcp.asTools
7-
import com.justai.jaicf.activator.llm.mcp.getTool
86
import com.justai.jaicf.activator.llm.tool.*
97
import com.justai.jaicf.activator.llm.vectorstore.LLMVectorStore
108
import com.justai.jaicf.activator.llm.vectorstore.LLMVectorStoreResponseBuilder
@@ -146,7 +144,7 @@ data class LLMProps(
146144
toolName: String,
147145
description: String? = null,
148146
noinline responseBuilder: McpServiceResponseBuilder = { it }
149-
) = tool(mcp.getTool(toolName, description, responseBuilder))
147+
) = tool(mcp.tool(toolName, description, responseBuilder))
150148

151149
fun vectorStore(
152150
store: LLMVectorStore,
@@ -159,7 +157,7 @@ data class LLMProps(
159157
service: McpService,
160158
tools: List<String> = emptyList(),
161159
responseBuilder: McpServiceResponseBuilder = { it }
162-
) = service.asTools(tools, responseBuilder).map { mcpTool -> tool(mcpTool) }
160+
) = service.tools(tools, responseBuilder).map { mcpTool -> tool(mcpTool) }
163161

164162
fun build() = LLMProps(
165163
model,

activators/llm/src/main/kotlin/com/justai/jaicf/activator/llm/builder/JsonSchemaBuilder.kt

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,3 +97,64 @@ class JsonSchemaBuilder {
9797
}
9898
}.toMap()
9999
}
100+
101+
fun JsonSchemaBuilder.inputSchemaToBuilder(inputSchema: Map<String, Any>?) {
102+
inputSchema ?: return
103+
104+
val properties = inputSchema["properties"] as? Map<String, Any> ?: return
105+
val required = inputSchema["required"] as? List<String> ?: emptyList()
106+
107+
properties.forEach { (propName, propDef) ->
108+
val propMap = when (propDef) {
109+
is kotlinx.serialization.json.JsonObject -> {
110+
propDef.entries.associate { (key, value) ->
111+
key to when (value) {
112+
is kotlinx.serialization.json.JsonPrimitive -> {
113+
when {
114+
value.isString -> value.content
115+
value.content == "true" || value.content == "false" -> value.content.toBoolean()
116+
value.content.toIntOrNull() != null -> value.content.toInt()
117+
value.content.toDoubleOrNull() != null -> value.content.toDouble()
118+
else -> value.content
119+
}
120+
}
121+
122+
is kotlinx.serialization.json.JsonArray -> {
123+
value.map {
124+
if (it is kotlinx.serialization.json.JsonPrimitive) it.content else it.toString()
125+
}
126+
}
127+
128+
else -> value.toString()
129+
}
130+
}
131+
}
132+
133+
is Map<*, *> -> propDef as? Map<String, Any>
134+
else -> null
135+
} ?: return@forEach
136+
137+
val type = propMap["type"] as? String ?: return@forEach
138+
val description = propMap["description"] as? String
139+
val isRequired = propName in required
140+
141+
when (type) {
142+
"string" -> {
143+
val enumValues = (propMap["enum"] as? List<*>)?.mapNotNull { it as? String }
144+
str(propName, description, isRequired, enumValues)
145+
}
146+
147+
"integer" -> {
148+
int(propName, description, isRequired)
149+
}
150+
151+
"number" -> {
152+
num(propName, description, isRequired)
153+
}
154+
155+
"boolean" -> {
156+
bool(propName, description, isRequired)
157+
}
158+
}
159+
}
160+
}
Lines changed: 52 additions & 113 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
package com.justai.jaicf.activator.llm.mcp
22

3-
import com.justai.jaicf.activator.llm.builder.JsonSchemaBuilder
3+
import com.justai.jaicf.activator.llm.builder.inputSchemaToBuilder
44
import com.justai.jaicf.activator.llm.tool.LLMTool
55
import com.justai.jaicf.activator.llm.tool.LLMToolCallContext
66
import com.justai.jaicf.activator.llm.tool.llmTool
@@ -26,7 +26,7 @@ class McpService() : AutoCloseable {
2626
private val mcp: Client = Client(clientInfo = Implementation(name = "mcp-client-jaicf", version = "1.0.0"))
2727
private var process: Process? = null
2828

29-
suspend fun connectStdio(command: List<String>) {
29+
private suspend fun connectStdio(command: List<String>) {
3030
process = runCatching { ProcessBuilder(command).start() }
3131
.getOrElse { throw IllegalStateException("Failed to start MCP process with command: $command", it) }
3232

@@ -39,7 +39,7 @@ class McpService() : AutoCloseable {
3939
}
4040
}
4141

42-
suspend fun connectSse(
42+
private suspend fun connectSse(
4343
urlString: String,
4444
client: HttpClient,
4545
reconnectionTime: Duration?,
@@ -49,7 +49,7 @@ class McpService() : AutoCloseable {
4949
mcp.connect(transport)
5050
}
5151

52-
suspend fun connectWebSocket(
52+
private suspend fun connectWebSocket(
5353
urlString: String,
5454
client: HttpClient,
5555
requestBuilder: HttpRequestBuilder.() -> Unit
@@ -58,7 +58,7 @@ class McpService() : AutoCloseable {
5858
mcp.connect(transport)
5959
}
6060

61-
suspend fun connectStreamableHttp(
61+
private suspend fun connectStreamableHttp(
6262
urlString: String,
6363
client: HttpClient,
6464
reconnectionTime: Duration?,
@@ -68,7 +68,51 @@ class McpService() : AutoCloseable {
6868
mcp.connect(transport)
6969
}
7070

71-
suspend fun getTools(): List<Tool> = mcp.listTools()?.tools ?: emptyList()
71+
private suspend fun getTools(): List<Tool> = mcp.listTools()?.tools ?: emptyList()
72+
73+
private fun asTool(
74+
tool: Tool,
75+
description: String?,
76+
responseBuilder: McpServiceResponseBuilder
77+
): LLMTool<Map<String, Any>> {
78+
return llmTool<Map<String, Any>>(
79+
name = tool.name,
80+
description = description ?: tool.description ?: "",
81+
parameters = {
82+
val inputSchemaMap = mapOf(
83+
"type" to tool.inputSchema.type,
84+
"properties" to tool.inputSchema.properties,
85+
"required" to (tool.inputSchema.required ?: emptyList())
86+
)
87+
inputSchemaToBuilder(inputSchemaMap)
88+
}
89+
) {
90+
responseBuilder(callTool(tool.name, call.arguments))
91+
}
92+
}
93+
94+
fun tools(
95+
toolNames: List<String> = emptyList(),
96+
responseBuilder: McpServiceResponseBuilder = { it }
97+
): List<LLMTool<Map<String, Any>>> {
98+
val availableTools = runBlocking { getTools() }
99+
return availableTools
100+
.filter { toolNames.isEmpty() || it.name in toolNames }
101+
.map { asTool(it, null, responseBuilder) }
102+
}
103+
104+
fun tool(
105+
toolName: String,
106+
description: String? = null,
107+
responseBuilder: McpServiceResponseBuilder = { it }
108+
): LLMTool<Map<String, Any>> {
109+
110+
val availableTools = runBlocking { getTools() }
111+
val tool = availableTools.find { it.name == toolName }
112+
?: throw IllegalArgumentException("Unknown tool $toolName")
113+
114+
return asTool(tool, description, responseBuilder)
115+
}
72116

73117
suspend fun callTool(name: String, arguments: Map<String, Any?>): CallToolResultBase {
74118
return mcp.callTool(name, arguments)
@@ -86,11 +130,11 @@ class McpService() : AutoCloseable {
86130
) = McpService().apply { runBlocking { connectStdio(command) } }
87131

88132
fun sse(
89-
urlString: String,
133+
url: String,
90134
client: HttpClient = HttpClient { install(SSE) },
91135
reconnectionTime: Duration? = null,
92136
requestBuilder: HttpRequestBuilder.() -> Unit = {}
93-
) = McpService().apply { runBlocking { connectSse(urlString, client, reconnectionTime) { requestBuilder() } } }
137+
) = McpService().apply { runBlocking { connectSse(url, client, reconnectionTime) { requestBuilder() } } }
94138

95139
fun websocket(
96140
urlString: String,
@@ -115,108 +159,3 @@ class McpService() : AutoCloseable {
115159
}
116160
}
117161
}
118-
119-
fun McpService.asTools(
120-
tools: List<String> = emptyList(),
121-
responseBuilder: McpServiceResponseBuilder = { it }
122-
): List<LLMTool<Map<String, Any>>> {
123-
val availableTools = runBlocking { getTools() }
124-
return availableTools
125-
.filter { tools.isEmpty() || it.name in tools }
126-
.map { getLlmTool(it, null, responseBuilder) }
127-
}
128-
129-
fun McpService.getTool(
130-
toolName: String,
131-
description: String? = null,
132-
responseBuilder: McpServiceResponseBuilder = { it }
133-
): LLMTool<Map<String, Any>> {
134-
135-
val availableTools = runBlocking { getTools() }
136-
val tool = availableTools.find { it.name == toolName }
137-
?: throw IllegalArgumentException("Unknown tool $toolName")
138-
139-
return getLlmTool(tool, description, responseBuilder)
140-
}
141-
142-
private fun McpService.getLlmTool(
143-
tool: Tool,
144-
description: String?,
145-
responseBuilder: McpServiceResponseBuilder
146-
): LLMTool<Map<String, Any>> {
147-
return llmTool<Map<String, Any>>(
148-
name = tool.name,
149-
description = description ?: tool.description ?: "",
150-
parameters = {
151-
val inputSchemaMap = mapOf(
152-
"type" to tool.inputSchema.type,
153-
"properties" to tool.inputSchema.properties,
154-
"required" to (tool.inputSchema.required ?: emptyList())
155-
)
156-
this.convertInputSchemaToBuilder(inputSchemaMap)
157-
}
158-
) {
159-
responseBuilder(callTool(tool.name, call.arguments))
160-
}
161-
}
162-
163-
private fun JsonSchemaBuilder.convertInputSchemaToBuilder(inputSchema: Map<String, Any>?) {
164-
inputSchema ?: return
165-
166-
val properties = inputSchema["properties"] as? Map<String, Any> ?: return
167-
val required = inputSchema["required"] as? List<String> ?: emptyList()
168-
169-
properties.forEach { (propName, propDef) ->
170-
val propMap = when (propDef) {
171-
is kotlinx.serialization.json.JsonObject -> {
172-
propDef.entries.associate { (key, value) ->
173-
key to when (value) {
174-
is kotlinx.serialization.json.JsonPrimitive -> {
175-
when {
176-
value.isString -> value.content
177-
value.content == "true" || value.content == "false" -> value.content.toBoolean()
178-
value.content.toIntOrNull() != null -> value.content.toInt()
179-
value.content.toDoubleOrNull() != null -> value.content.toDouble()
180-
else -> value.content
181-
}
182-
}
183-
184-
is kotlinx.serialization.json.JsonArray -> {
185-
value.map {
186-
if (it is kotlinx.serialization.json.JsonPrimitive) it.content else it.toString()
187-
}
188-
}
189-
190-
else -> value.toString()
191-
}
192-
}
193-
}
194-
195-
is Map<*, *> -> propDef as? Map<String, Any>
196-
else -> null
197-
} ?: return@forEach
198-
199-
val type = propMap["type"] as? String ?: return@forEach
200-
val description = propMap["description"] as? String
201-
val isRequired = propName in required
202-
203-
when (type) {
204-
"string" -> {
205-
val enumValues = (propMap["enum"] as? List<*>)?.mapNotNull { it as? String }
206-
str(propName, description, isRequired, enumValues)
207-
}
208-
209-
"integer" -> {
210-
int(propName, description, isRequired)
211-
}
212-
213-
"number" -> {
214-
num(propName, description, isRequired)
215-
}
216-
217-
"boolean" -> {
218-
bool(propName, description, isRequired)
219-
}
220-
}
221-
}
222-
}
Lines changed: 12 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,7 @@
11
package com.justai.jaicf.examples.llm
22

3-
import com.justai.jaicf.BotEngine
3+
import com.justai.jaicf.activator.llm.agent.LLMAgent
44
import com.justai.jaicf.activator.llm.mcp.McpService
5-
import com.justai.jaicf.activator.llm.mcp.getTool
6-
import com.justai.jaicf.activator.llm.scenario.llmState
7-
import com.justai.jaicf.builder.Scenario
85
import com.justai.jaicf.examples.llm.channel.ConsoleChannel
96

107
/**
@@ -17,7 +14,7 @@ import com.justai.jaicf.examples.llm.channel.ConsoleChannel
1714
* - **Streaming** (`--transport=streaming`) - HTTP with streaming support, use with `McpService.streamableHttp()`
1815
*
1916
* ## Setup:
20-
* 1. See docker-compose.yml at: activators/llm/src/main/kotlin/com/justai/jaicf/activator/llm/mcp/gateway/docker-compose.yml
17+
* 1. See docker-compose.yml at: activators/llm/mcp-gateway/docker-compose.yml
2118
* 2. Start the gateway: `docker-compose up`
2219
* 3. Connect to the gateway using one of the transports below
2320
*
@@ -28,11 +25,11 @@ import com.justai.jaicf.examples.llm.channel.ConsoleChannel
2825
* EXA_API_KEY: your-api-key-here
2926
* ```
3027
*
31-
* For STDIO transport examples, see ScenarioWithMcpServiceStdio.kt
28+
* For STDIO transport examples, see AgentWithMcpServiceStdio.kt
3229
*/
3330

34-
val mcpServiceSse = McpService.sse(
35-
urlString = "http://localhost:8081/sse"
31+
val searchMCPService = McpService.sse(
32+
url = "http://localhost:8081/sse"
3633
)
3734

3835
/**
@@ -41,12 +38,13 @@ val mcpServiceSse = McpService.sse(
4138
*
4239
* IMPORTANT! Set up your OPENAI_API_KEY and OPENAI_BASE_URL env before running
4340
*/
44-
private val scenario = Scenario {
45-
llmState("main", {
41+
private val agent = LLMAgent(
42+
name = "mcp-agent",
43+
props = {
4644
model = "gpt-4.1-nano"
4745

4846
// The most simple way to add a tool. Provide a tool name via `toolName` parameter.
49-
tool(mcpServiceSse.getTool("search"))
47+
tool(searchMCPService.tool("search"))
5048

5149
/**
5250
* You can also use other options for your scenario:
@@ -84,10 +82,8 @@ private val scenario = Scenario {
8482
* }
8583
* ```
8684
*/
87-
})
88-
}
89-
85+
}
86+
)
9087
fun main() {
91-
ConsoleChannel(BotEngine(scenario))
92-
.run("Hi")
88+
ConsoleChannel(agent.asBot).run("What tools you have?")
9389
}

0 commit comments

Comments
 (0)