Skip to content

Commit 6e6f1a0

Browse files
sualehsualeh
committed
Register all functions with MCP server
1 parent e20775d commit 6e6f1a0

File tree

4 files changed

+94
-91
lines changed

4 files changed

+94
-91
lines changed

‎schemacrawler-ai-mcp/src/main/java/schemacrawler/tools/command/aichat/mcp/CommonService.java‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
package schemacrawler.tools.command.aichat.mcp;
2+
3+
import org.springframework.ai.tool.annotation.Tool;
4+
import org.springframework.ai.tool.annotation.ToolParam;
5+
import org.springframework.stereotype.Service;
6+
import schemacrawler.Version;
7+
8+
@Service
9+
public class CommonService {
10+
11+
@Tool(name = "get-schemacrawler-version", description = "Gets the version of SchemaCrawler", returnDirect = true)
12+
public String getSchemaCrawlerVersion(
13+
@ToolParam(
14+
description =
15+
"""
16+
Current date, as an ISO 8601 local date.
17+
""",
18+
required = false) final String date) {
19+
System.out.printf("get-schemacrawler-version called with %s", date);
20+
return Version.about();
21+
}
22+
}

‎schemacrawler-ai-mcp/src/main/java/schemacrawler/tools/command/aichat/mcp/SchemaCrawlerMCPServer.java‎

Lines changed: 12 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,14 @@
11
package schemacrawler.tools.command.aichat.mcp;
22

3-
import java.util.List;
43
import org.springframework.ai.tool.ToolCallback;
54
import org.springframework.ai.tool.ToolCallbackProvider;
65
import org.springframework.ai.tool.method.MethodToolCallbackProvider;
76
import org.springframework.boot.SpringApplication;
87
import org.springframework.boot.autoconfigure.SpringBootApplication;
98
import org.springframework.context.annotation.Bean;
109

10+
import java.util.List;
11+
1112
/**
1213
* Spring Boot application for the SchemaCrawler AI MCP server. This class enables the Spring AI MCP
1314
* server capabilities.
@@ -22,31 +23,27 @@ public static void main(final String[] args) {
2223
@Bean
2324
public ToolCallbackProvider schemaCrawlerTools() {
2425
final List<ToolCallback> tools = SpringAIUtility.toolCallbacks(SpringAIUtility.tools());
25-
final List<ToolCallback> tools1 =
26-
tools.stream()
27-
.filter(toolCallback -> toolCallback.getToolDefinition().name().equals("lint"))
28-
.collect(java.util.stream.Collectors.toList());
29-
final ToolCallbackProvider toolCallbackProvider = ToolCallbackProvider.from(tools1);
26+
final ToolCallbackProvider toolCallbackProvider = ToolCallbackProvider.from(tools);
3027
printTools(toolCallbackProvider);
3128
return toolCallbackProvider;
3229
}
3330

3431
@Bean
35-
public ToolCallbackProvider weatherTools(final WeatherService weatherService) {
32+
public ToolCallbackProvider weatherTools(final CommonService weatherService) {
3633
final MethodToolCallbackProvider toolCallbackProvider =
37-
MethodToolCallbackProvider.builder().toolObjects(weatherService).build();
34+
MethodToolCallbackProvider.builder().toolObjects(weatherService).build();
3835
printTools(toolCallbackProvider);
3936
return toolCallbackProvider;
4037
}
4138

4239
private void printTools(final ToolCallbackProvider toolCallbackProvider) {
4340
List.of(toolCallbackProvider.getToolCallbacks())
44-
.forEach(
45-
toolCallback -> {
46-
System.out.println(toolCallback.getToolDefinition().name());
47-
System.out.println(toolCallback.getToolDefinition().description());
48-
System.out.println(toolCallback.getToolDefinition().inputSchema());
49-
System.out.println("----------------------------------------");
50-
});
41+
.forEach(
42+
toolCallback -> {
43+
System.out.println(toolCallback.getToolDefinition().name());
44+
System.out.println(toolCallback.getToolDefinition().description());
45+
System.out.println(toolCallback.getToolDefinition().inputSchema());
46+
System.out.println("----------------------------------------");
47+
});
5148
}
5249
}

‎schemacrawler-ai-mcp/src/main/java/schemacrawler/tools/command/aichat/mcp/SpringAIUtility.java‎

Lines changed: 60 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -28,64 +28,37 @@
2828

2929
package schemacrawler.tools.command.aichat.mcp;
3030

31-
import java.util.ArrayList;
32-
import java.util.HashMap;
33-
import java.util.HashSet;
34-
import java.util.List;
35-
import java.util.Map;
36-
import java.util.Map.Entry;
37-
import java.util.Objects;
38-
import java.util.Set;
39-
import java.util.logging.Level;
40-
import java.util.logging.Logger;
41-
import org.springframework.ai.chat.model.ToolContext;
42-
import org.springframework.ai.tool.ToolCallback;
43-
import org.springframework.ai.tool.definition.ToolDefinition;
44-
import org.springframework.ai.util.json.JsonParser;
45-
import org.springframework.lang.Nullable;
4631
import com.fasterxml.jackson.databind.JsonNode;
4732
import com.fasterxml.jackson.databind.ObjectMapper;
33+
import com.fasterxml.jackson.databind.node.ArrayNode;
4834
import com.fasterxml.jackson.databind.node.ObjectNode;
4935
import com.fasterxml.jackson.module.jsonSchema.JsonSchema;
5036
import com.fasterxml.jackson.module.jsonSchema.JsonSchemaGenerator;
5137
import com.github.victools.jsonschema.generator.SchemaVersion;
38+
import org.springframework.ai.chat.model.ToolContext;
39+
import org.springframework.ai.tool.ToolCallback;
40+
import org.springframework.ai.tool.definition.ToolDefinition;
41+
import org.springframework.ai.util.json.JsonParser;
42+
import org.springframework.lang.Nullable;
5243
import schemacrawler.tools.command.aichat.FunctionDefinition;
5344
import schemacrawler.tools.command.aichat.FunctionDefinition.FunctionType;
5445
import schemacrawler.tools.command.aichat.functions.FunctionDefinitionRegistry;
5546
import us.fatehi.utility.UtilityMarker;
5647

57-
@UtilityMarker
58-
public class SpringAIUtility {
59-
60-
public record SpringAIToolCallback(ToolDefinition toolDefinition) implements ToolCallback {
61-
62-
public SpringAIToolCallback {
63-
Objects.requireNonNull(toolDefinition, "Tool definition must not be null");
64-
}
48+
import java.util.*;
49+
import java.util.Map.Entry;
50+
import java.util.logging.Level;
51+
import java.util.logging.Logger;
6552

66-
@Override
67-
public ToolDefinition getToolDefinition() {
68-
return toolDefinition;
69-
}
53+
@UtilityMarker
54+
public final class SpringAIUtility {
7055

71-
@Override
72-
public String call(final String toolInput) {
73-
final String callMessage =
74-
String.format(
75-
"Call to <%s>%n%s%nTool was successfully executed with no return value.",
76-
toolDefinition.name(), toolInput);
77-
System.out.println(callMessage);
78-
return callMessage;
79-
}
56+
private static final Logger LOGGER = Logger.getLogger(SpringAIUtility.class.getCanonicalName());
8057

81-
@Override
82-
public String call(final String toolInput, @Nullable final ToolContext tooContext) {
83-
return call(toolInput);
84-
}
58+
private SpringAIUtility() {
59+
// Prevent instantiation
8560
}
8661

87-
private static final Logger LOGGER = Logger.getLogger(SpringAIUtility.class.getCanonicalName());
88-
8962
public static List<ToolCallback> toolCallbacks(final List<ToolDefinition> tools) {
9063
Objects.requireNonNull(tools, "Tools must not be null");
9164
final List<ToolCallback> toolCallbacks = new ArrayList<>();
@@ -99,22 +72,22 @@ public static List<ToolDefinition> tools() {
9972

10073
final List<ToolDefinition> toolDefinitions = new ArrayList<>();
10174
for (final FunctionDefinition<?> functionDefinition :
102-
FunctionDefinitionRegistry.getFunctionDefinitionRegistry().getFunctionDefinitions()) {
75+
FunctionDefinitionRegistry.getFunctionDefinitionRegistry().getFunctionDefinitions()) {
10376
if (functionDefinition.getFunctionType() != FunctionType.USER) {
10477
continue;
10578
}
10679

10780
try {
10881
final ToolDefinition toolDefinition =
109-
ToolDefinition.builder()
110-
.name(functionDefinition.getName())
111-
.description(functionDefinition.getDescription())
112-
.inputSchema(generateToolInput(functionDefinition.getParametersClass()))
113-
.build();
82+
ToolDefinition.builder()
83+
.name(functionDefinition.getName())
84+
.description(functionDefinition.getDescription())
85+
.inputSchema(generateToolInput(functionDefinition.getParametersClass()))
86+
.build();
11487
toolDefinitions.add(toolDefinition);
11588
} catch (final Exception e) {
11689
LOGGER.log(
117-
Level.WARNING, String.format("Could not load <%s>", functionDefinition.getName()), e);
90+
Level.WARNING, String.format("Could not load <%s>", functionDefinition.getName()), e);
11891
}
11992
}
12093

@@ -132,10 +105,21 @@ private static String generateToolInput(final Class<?> parametersClass) throws E
132105
schema.put("$schema", SchemaVersion.DRAFT_2020_12.getIdentifier());
133106
schema.put("type", "object");
134107

108+
final List<String> required = new ArrayList<>();
135109
final ObjectNode properties = schema.putObject("properties");
136110
for (final Entry<String, JsonNode> parameter : parametersJsonSchema.entrySet()) {
137-
properties.set(parameter.getKey(), parameter.getValue());
111+
final String parameterName = parameter.getKey();
112+
final JsonNode parameterSchema = parameter.getValue();
113+
if (parameterSchema.has("required") && parameterSchema.get("required").asBoolean()) {
114+
((ObjectNode) parameterSchema).remove("required");
115+
required.add(parameterName);
116+
}
117+
properties.set(parameterName, parameterSchema);
138118
}
119+
final ArrayNode requiredArray = schema.putArray("required");
120+
required.forEach(requiredArray::add);
121+
122+
schema.put("additionalProperties", false);
139123

140124
return schema.toPrettyString();
141125
}
@@ -159,7 +143,31 @@ private static Map<String, JsonNode> jsonSchema(final Class<?> parametersClass)
159143
return propertiesMap;
160144
}
161145

162-
private SpringAIUtility() {
163-
// Prevent instantiation
146+
public record SpringAIToolCallback(
147+
ToolDefinition toolDefinition) implements ToolCallback {
148+
149+
public SpringAIToolCallback {
150+
Objects.requireNonNull(toolDefinition, "Tool definition must not be null");
151+
}
152+
153+
@Override
154+
public ToolDefinition getToolDefinition() {
155+
return toolDefinition;
156+
}
157+
158+
@Override
159+
public String call(final String toolInput) {
160+
final String callMessage =
161+
String.format(
162+
"Call to <%s>%n%s%nTool was successfully executed with no return value.",
163+
toolDefinition.name(), toolInput);
164+
System.out.println(callMessage);
165+
return callMessage;
166+
}
167+
168+
@Override
169+
public String call(final String toolInput, @Nullable final ToolContext tooContext) {
170+
return call(toolInput);
171+
}
164172
}
165173
}

‎schemacrawler-ai-mcp/src/main/java/schemacrawler/tools/command/aichat/mcp/WeatherService.java‎

Lines changed: 0 additions & 24 deletions
This file was deleted.

0 commit comments

Comments
 (0)