Skip to content

Commit d82e620

Browse files
adwsinghadwsingh
committed
Move away from using model to using Schema
1 parent a33d1c1 commit d82e620

File tree

6 files changed

+51
-166
lines changed

6 files changed

+51
-166
lines changed

‎examples/mcp-server/src/main/java/software/amazon/smithy/java/example/server/mcp/ProxyMCPExample.java

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55
import software.amazon.smithy.java.example.server.mcp.operations.GetEmployeeDetails;
66
import software.amazon.smithy.java.example.server.mcp.service.EmployeeService;
77
import software.amazon.smithy.java.mcp.server.McpServer;
8-
import software.amazon.smithy.java.mcp.server.PromptLoader;
98
import software.amazon.smithy.java.server.ProxyService;
109
import software.amazon.smithy.java.server.Server;
1110
import software.amazon.smithy.model.Model;
@@ -43,7 +42,6 @@ public static void main(String[] args) {
4342
.stdio()
4443
.name("smithy-mcp-server")
4544
.addService(mcpService)
46-
.addModel(model)
4745
.build();
4846
mcpServer.start();
4947

‎mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServer.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ public final class McpServer implements Server {
8080

8181
McpServer(McpServerBuilder builder) {
8282
this.tools = createTools(builder.serviceList);
83-
this.prompts = PromptLoader.loadPrompts(builder.modelList);
83+
this.prompts = PromptLoader.loadPrompts(builder.serviceList);
8484
this.promptProcessor = new PromptProcessor();
8585
this.is = builder.is;
8686
this.os = builder.os;

‎mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServerBuilder.java

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -69,16 +69,6 @@ public McpServerBuilder addMcpService(McpServerProxy proxy) {
6969
return this;
7070
}
7171

72-
public McpServerBuilder addModel(Model... model) {
73-
modelList.addAll(Arrays.asList(model));
74-
return this;
75-
}
76-
77-
public McpServerBuilder addModels(List<Model> models) {
78-
modelList.addAll(models);
79-
return this;
80-
}
81-
8272
private void validate() {
8373
Objects.requireNonNull(is, "MCP server input stream is required");
8474
Objects.requireNonNull(os, "MCP server output stream is required");

‎mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/PromptLoader.java

Lines changed: 49 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -6,97 +6,98 @@
66
package software.amazon.smithy.java.mcp.server;
77

88
import java.util.ArrayList;
9+
import java.util.HashMap;
910
import java.util.LinkedHashMap;
1011
import java.util.List;
1112
import java.util.Map;
12-
import java.util.Set;
1313
import software.amazon.smithy.ai.PromptTemplateDefinition;
1414
import software.amazon.smithy.ai.PromptsTrait;
15+
import software.amazon.smithy.java.core.schema.Schema;
16+
import software.amazon.smithy.java.core.schema.TraitKey;
1517
import software.amazon.smithy.java.mcp.model.PromptArgument;
1618
import software.amazon.smithy.java.mcp.model.PromptInfo;
17-
import software.amazon.smithy.model.Model;
18-
import software.amazon.smithy.model.shapes.Shape;
19-
import software.amazon.smithy.model.shapes.ShapeId;
20-
import software.amazon.smithy.model.shapes.StructureShape;
21-
import software.amazon.smithy.model.traits.DocumentationTrait;
22-
import software.amazon.smithy.model.traits.RequiredTrait;
19+
import software.amazon.smithy.java.server.Service;
2320
import software.amazon.smithy.utils.SmithyUnstableApi;
2421

2522
/**
2623
* Handles loading and parsing of prompts from Smithy models.
2724
*/
2825
@SmithyUnstableApi
29-
public final class PromptLoader {
26+
final class PromptLoader {
27+
28+
private static final TraitKey<PromptsTrait> PROMPTS_TRAIT_KEY = TraitKey.get(PromptsTrait.class);
3029

3130
public static final String TOOL_PREFERENCE_PREFIX = ".Tool preference: ";
3231

3332
/**
3433
* Loads prompts from the provided Smithy models.
3534
*
36-
* @param models List of Smithy models to extract prompts from
3735
* @return Map of prompt names to PromptInfo objects
3836
*/
39-
public static Map<String, PromptInfo> loadPrompts(List<Model> models) {
37+
public static Map<String, PromptInfo> loadPrompts(List<Service> services) {
4038
Map<String, PromptInfo> promptInfos = new LinkedHashMap<>();
4139

42-
for (Model model : models) {
43-
Set<Shape> promptShapes = model.getShapesWithTrait(PromptsTrait.ID);
44-
for (Shape prompt : promptShapes) {
45-
46-
Map<String, PromptTemplateDefinition> promptDefinitions =
47-
prompt.expectTrait(PromptsTrait.class).getValues();
48-
for (Map.Entry<String, PromptTemplateDefinition> entry : promptDefinitions.entrySet()) {
49-
var promptName = entry.getKey().toLowerCase();
50-
var promptTemplateDefinition = entry.getValue();
51-
var templateString = promptTemplateDefinition.getTemplate();
52-
53-
promptInfos.put(
54-
promptName,
55-
PromptInfo
56-
.builder()
57-
.name(promptName)
58-
.description(promptTemplateDefinition.getDescription())
59-
.template(
60-
promptTemplateDefinition.getPreferWhen().isPresent()
61-
? templateString + TOOL_PREFERENCE_PREFIX
62-
+ promptTemplateDefinition.getPreferWhen().get()
63-
: templateString)
64-
.arguments(promptTemplateDefinition.getArguments().isPresent()
65-
? convertArgumentShapeToPromptArgument(model,
66-
promptTemplateDefinition.getArguments().get())
67-
: List.of())
68-
.build());
40+
for (var service : services) {
41+
Map<String, PromptTemplateDefinition> promptDefinitions = new HashMap<>();
42+
var servicePromptTrait = service.schema().getTrait(PROMPTS_TRAIT_KEY);
43+
if (servicePromptTrait != null) {
44+
promptDefinitions.putAll(servicePromptTrait.getValues());
45+
}
46+
service.getAllOperations().forEach(operation -> {
47+
var operationPromptsTrait = operation.getApiOperation().schema().getTrait(PROMPTS_TRAIT_KEY);
48+
if (operationPromptsTrait != null) {
49+
promptDefinitions.putAll(operationPromptsTrait.getValues());
6950
}
51+
52+
});
53+
for (Map.Entry<String, PromptTemplateDefinition> entry : promptDefinitions.entrySet()) {
54+
var promptName = entry.getKey().toLowerCase();
55+
var promptTemplateDefinition = entry.getValue();
56+
var templateString = promptTemplateDefinition.getTemplate();
57+
58+
promptInfos.put(
59+
promptName,
60+
PromptInfo
61+
.builder()
62+
.name(promptName)
63+
.description(promptTemplateDefinition.getDescription())
64+
.template(
65+
promptTemplateDefinition.getPreferWhen().isPresent()
66+
? templateString + TOOL_PREFERENCE_PREFIX
67+
+ promptTemplateDefinition.getPreferWhen().get()
68+
: templateString)
69+
.arguments(promptTemplateDefinition.getArguments().isPresent()
70+
? convertArgumentShapeToPromptArgument(
71+
service.schemaIndex()
72+
.getSchema(promptTemplateDefinition.getArguments().get()))
73+
: List.of())
74+
.build());
7075
}
7176
}
72-
7377
return promptInfos;
7478
}
7579

7680
/**
7781
* Converts a Smithy structure shape to a list of PromptArgument objects.
7882
*
79-
* @param model The Smithy model containing the shape
8083
* @param argumentShapeId The ShapeId of the structure to convert
8184
* @return List of PromptArgument objects representing the structure members
8285
*/
83-
public static List<PromptArgument> convertArgumentShapeToPromptArgument(Model model, ShapeId argumentShapeId) {
84-
StructureShape argument = model.expectShape(argumentShapeId, StructureShape.class);
86+
public static List<PromptArgument> convertArgumentShapeToPromptArgument(Schema argument) {
8587
List<PromptArgument> promptArguments = new ArrayList<>();
8688

87-
for (var member : argument.getAllMembers().entrySet()) {
88-
String memberName = member.getKey();
89-
var memberShape = member.getValue();
89+
for (var member : argument.members()) {
90+
String memberName = member.memberName();
9091

9192
// Get description from documentation trait, use empty string if not present
9293
String description = "";
93-
var documentationTrait = memberShape.getTrait(DocumentationTrait.class);
94-
if (documentationTrait.isPresent()) {
95-
description = documentationTrait.get().getValue();
94+
var documentationTrait = member.getTrait(TraitKey.DOCUMENTATION_TRAIT);
95+
if (documentationTrait != null) {
96+
description = documentationTrait.getValue();
9697
}
9798

9899
// Check if member is required
99-
boolean isRequired = memberShape.hasTrait(RequiredTrait.class);
100+
boolean isRequired = member.getTrait(TraitKey.REQUIRED_TRAIT) != null;
100101

101102
// Build the PromptArgument
102103
PromptArgument promptArgument = PromptArgument.builder()

‎mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/McpServerTest.java

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -368,7 +368,6 @@ void testPromptsList() {
368368
.proxyEndpoint("http://localhost")
369369
.model(MODEL)
370370
.build())
371-
.addModel(MODEL)
372371
.build();
373372

374373
server.start();
@@ -377,6 +376,7 @@ void testPromptsList() {
377376
var response = read();
378377
var prompts = response.getResult().asStringMap().get("prompts").asList();
379378

379+
prompts.forEach(prompt -> System.out.println(prompt.asStringMap()));
380380
assertEquals(2, prompts.size());
381381

382382
// Check the prompt (service and operation have same name, so only one is returned)
@@ -401,7 +401,6 @@ void testPromptsGetWithValidPrompt() {
401401
.proxyEndpoint("http://localhost")
402402
.model(MODEL)
403403
.build())
404-
.addModel(MODEL)
405404
.build();
406405

407406
server.start();
@@ -434,7 +433,6 @@ void testPromptsGetWithInvalidPrompt() {
434433
.proxyEndpoint("http://localhost")
435434
.model(MODEL)
436435
.build())
437-
.addModel(MODEL)
438436
.build();
439437

440438
server.start();
@@ -464,7 +462,6 @@ void testPromptsGetWithTemplateArguments() {
464462
.proxyEndpoint("http://localhost")
465463
.model(modelWithArgs)
466464
.build())
467-
.addModel(modelWithArgs)
468465
.build();
469466

470467
server.start();
@@ -505,7 +502,6 @@ void testPromptsGetWithMissingRequiredArguments() {
505502
.proxyEndpoint("http://localhost")
506503
.model(modelWithArgs)
507504
.build())
508-
.addModel(modelWithArgs)
509505
.build();
510506

511507
server.start();
@@ -542,7 +538,6 @@ void testApplyTemplateArgumentsEdgeCases() {
542538
.proxyEndpoint("http://localhost")
543539
.model(modelEdgeCases)
544540
.build())
545-
.addModel(modelEdgeCases)
546541
.build();
547542

548543
server.start();

‎mcp/mcp-server/src/test/java/software/amazon/smithy/java/mcp/server/PromptLoaderTest.java

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

0 commit comments

Comments
 (0)