Skip to content

Commit 33329bf

Browse files
tzolovtzolov
authored
feat(mcp): Add builder for CreateMessageRequest (#60)
- Implements a builder pattern for CreateMessageRequest - Updates corresponding tests to use the new builder syntax - Add ModelPreferences builder and ModelHint helper - Use builder pattern for CreateMessageRequest in integration tests Signed-off-by: Christian Tzolov <[email protected]>
1 parent 3e2139f commit 33329bf

File tree

5 files changed

+163
-32
lines changed

5 files changed

+163
-32
lines changed

Diff for: ‎mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxSseIntegrationTests.java

+12-7
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest;
2525
import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult;
2626
import io.modelcontextprotocol.spec.McpSchema.InitializeResult;
27+
import io.modelcontextprotocol.spec.McpSchema.ModelPreferences;
2728
import io.modelcontextprotocol.spec.McpSchema.Role;
2829
import io.modelcontextprotocol.spec.McpSchema.Root;
2930
import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities;
@@ -45,6 +46,7 @@
4546

4647
import static org.assertj.core.api.Assertions.assertThat;
4748
import static org.awaitility.Awaitility.await;
49+
import static org.junit.Assert.assertThat;
4850
import static org.mockito.Mockito.mock;
4951

5052
public class WebFluxSseIntegrationTests {
@@ -142,13 +144,16 @@ void testCreateMessageSuccess(String clientType) throws InterruptedException {
142144
McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification(
143145
new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> {
144146

145-
var messages = List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER,
146-
new McpSchema.TextContent("Test message")));
147-
var modelPrefs = new McpSchema.ModelPreferences(List.of(), 1.0, 1.0, 1.0);
148-
149-
var craeteMessageRequest = new McpSchema.CreateMessageRequest(messages, modelPrefs, null,
150-
McpSchema.CreateMessageRequest.ContextInclusionStrategy.NONE, null, 100, List.of(),
151-
Map.of());
147+
var craeteMessageRequest = McpSchema.CreateMessageRequest.builder()
148+
.messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER,
149+
new McpSchema.TextContent("Test message"))))
150+
.modelPreferences(ModelPreferences.builder()
151+
.hints(List.of())
152+
.costPriority(1.0)
153+
.speedPriority(1.0)
154+
.intelligencePriority(1.0)
155+
.build())
156+
.build();
152157

153158
StepVerifier.create(exchange.createMessage(craeteMessageRequest)).consumeNextWith(result -> {
154159
assertThat(result).isNotNull();

Diff for: ‎mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseIntegrationTests.java

+12-7
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest;
2121
import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult;
2222
import io.modelcontextprotocol.spec.McpSchema.InitializeResult;
23+
import io.modelcontextprotocol.spec.McpSchema.ModelPreferences;
2324
import io.modelcontextprotocol.spec.McpSchema.Role;
2425
import io.modelcontextprotocol.spec.McpSchema.Root;
2526
import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities;
@@ -45,6 +46,7 @@
4546

4647
import static org.assertj.core.api.Assertions.assertThat;
4748
import static org.awaitility.Awaitility.await;
49+
import static org.junit.Assert.assertThat;
4850
import static org.mockito.Mockito.mock;
4951

5052
public class WebMvcSseIntegrationTests {
@@ -199,13 +201,16 @@ void testCreateMessageSuccess() throws InterruptedException {
199201
McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification(
200202
new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> {
201203

202-
var messages = List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER,
203-
new McpSchema.TextContent("Test message")));
204-
var modelPrefs = new McpSchema.ModelPreferences(List.of(), 1.0, 1.0, 1.0);
205-
206-
var craeteMessageRequest = new McpSchema.CreateMessageRequest(messages, modelPrefs, null,
207-
McpSchema.CreateMessageRequest.ContextInclusionStrategy.NONE, null, 100, List.of(),
208-
Map.of());
204+
var craeteMessageRequest = McpSchema.CreateMessageRequest.builder()
205+
.messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER,
206+
new McpSchema.TextContent("Test message"))))
207+
.modelPreferences(ModelPreferences.builder()
208+
.hints(List.of())
209+
.costPriority(1.0)
210+
.speedPriority(1.0)
211+
.intelligencePriority(1.0)
212+
.build())
213+
.build();
209214

210215
StepVerifier.create(exchange.createMessage(craeteMessageRequest)).consumeNextWith(result -> {
211216
assertThat(result).isNotNull();

Diff for: ‎mcp/src/main/java/io/modelcontextprotocol/spec/McpSchema.java

+112-5
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
package io.modelcontextprotocol.spec;
66

77
import java.io.IOException;
8+
import java.util.ArrayList;
89
import java.util.HashMap;
910
import java.util.List;
1011
import java.util.Map;
@@ -763,15 +764,61 @@ public record CallToolResult( // @formatter:off
763764
@JsonInclude(JsonInclude.Include.NON_ABSENT)
764765
@JsonIgnoreProperties(ignoreUnknown = true)
765766
public record ModelPreferences(// @formatter:off
766-
@JsonProperty("hints") List<ModelHint> hints,
767-
@JsonProperty("costPriority") Double costPriority,
768-
@JsonProperty("speedPriority") Double speedPriority,
769-
@JsonProperty("intelligencePriority") Double intelligencePriority) {
770-
} // @formatter:on
767+
@JsonProperty("hints") List<ModelHint> hints,
768+
@JsonProperty("costPriority") Double costPriority,
769+
@JsonProperty("speedPriority") Double speedPriority,
770+
@JsonProperty("intelligencePriority") Double intelligencePriority) {
771+
772+
public static Builder builder() {
773+
return new Builder();
774+
}
775+
776+
public static class Builder {
777+
private List<ModelHint> hints;
778+
private Double costPriority;
779+
private Double speedPriority;
780+
private Double intelligencePriority;
781+
782+
public Builder hints(List<ModelHint> hints) {
783+
this.hints = hints;
784+
return this;
785+
}
786+
787+
public Builder addHint(String name) {
788+
if (this.hints == null) {
789+
this.hints = new ArrayList<>();
790+
}
791+
this.hints.add(new ModelHint(name));
792+
return this;
793+
}
794+
795+
public Builder costPriority(Double costPriority) {
796+
this.costPriority = costPriority;
797+
return this;
798+
}
799+
800+
public Builder speedPriority(Double speedPriority) {
801+
this.speedPriority = speedPriority;
802+
return this;
803+
}
804+
805+
public Builder intelligencePriority(Double intelligencePriority) {
806+
this.intelligencePriority = intelligencePriority;
807+
return this;
808+
}
809+
810+
public ModelPreferences build() {
811+
return new ModelPreferences(hints, costPriority, speedPriority, intelligencePriority);
812+
}
813+
}
814+
} // @formatter:on
771815

772816
@JsonInclude(JsonInclude.Include.NON_ABSENT)
773817
@JsonIgnoreProperties(ignoreUnknown = true)
774818
public record ModelHint(@JsonProperty("name") String name) {
819+
public static ModelHint of(String name) {
820+
return new ModelHint(name);
821+
}
775822
}
776823

777824
@JsonInclude(JsonInclude.Include.NON_ABSENT)
@@ -799,6 +846,66 @@ public enum ContextInclusionStrategy {
799846
@JsonProperty("thisServer") THIS_SERVER,
800847
@JsonProperty("allServers") ALL_SERVERS
801848
}
849+
850+
public static Builder builder() {
851+
return new Builder();
852+
}
853+
854+
public static class Builder {
855+
private List<SamplingMessage> messages;
856+
private ModelPreferences modelPreferences;
857+
private String systemPrompt;
858+
private ContextInclusionStrategy includeContext;
859+
private Double temperature;
860+
private int maxTokens;
861+
private List<String> stopSequences;
862+
private Map<String, Object> metadata;
863+
864+
public Builder messages(List<SamplingMessage> messages) {
865+
this.messages = messages;
866+
return this;
867+
}
868+
869+
public Builder modelPreferences(ModelPreferences modelPreferences) {
870+
this.modelPreferences = modelPreferences;
871+
return this;
872+
}
873+
874+
public Builder systemPrompt(String systemPrompt) {
875+
this.systemPrompt = systemPrompt;
876+
return this;
877+
}
878+
879+
public Builder includeContext(ContextInclusionStrategy includeContext) {
880+
this.includeContext = includeContext;
881+
return this;
882+
}
883+
884+
public Builder temperature(Double temperature) {
885+
this.temperature = temperature;
886+
return this;
887+
}
888+
889+
public Builder maxTokens(int maxTokens) {
890+
this.maxTokens = maxTokens;
891+
return this;
892+
}
893+
894+
public Builder stopSequences(List<String> stopSequences) {
895+
this.stopSequences = stopSequences;
896+
return this;
897+
}
898+
899+
public Builder metadata(Map<String, Object> metadata) {
900+
this.metadata = metadata;
901+
return this;
902+
}
903+
904+
public CreateMessageRequest build() {
905+
return new CreateMessageRequest(messages, modelPreferences, systemPrompt,
906+
includeContext, temperature, maxTokens, stopSequences, metadata);
907+
}
908+
}
802909
}// @formatter:on
803910

804911
@JsonInclude(JsonInclude.Include.NON_ABSENT)

Diff for: ‎mcp/src/test/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProviderIntegrationTests.java

+11-7
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest;
2222
import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult;
2323
import io.modelcontextprotocol.spec.McpSchema.InitializeResult;
24+
import io.modelcontextprotocol.spec.McpSchema.ModelPreferences;
2425
import io.modelcontextprotocol.spec.McpSchema.Role;
2526
import io.modelcontextprotocol.spec.McpSchema.Root;
2627
import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities;
@@ -162,13 +163,16 @@ void testCreateMessageSuccess() throws InterruptedException {
162163
McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification(
163164
new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> {
164165

165-
var messages = List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER,
166-
new McpSchema.TextContent("Test message")));
167-
var modelPrefs = new McpSchema.ModelPreferences(List.of(), 1.0, 1.0, 1.0);
168-
169-
var craeteMessageRequest = new McpSchema.CreateMessageRequest(messages, modelPrefs, null,
170-
McpSchema.CreateMessageRequest.ContextInclusionStrategy.NONE, null, 100, List.of(),
171-
Map.of());
166+
var craeteMessageRequest = McpSchema.CreateMessageRequest.builder()
167+
.messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER,
168+
new McpSchema.TextContent("Test message"))))
169+
.modelPreferences(ModelPreferences.builder()
170+
.hints(List.of())
171+
.costPriority(1.0)
172+
.speedPriority(1.0)
173+
.intelligencePriority(1.0)
174+
.build())
175+
.build();
172176

173177
StepVerifier.create(exchange.createMessage(craeteMessageRequest)).consumeNextWith(result -> {
174178
assertThat(result).isNotNull();

Diff for: ‎mcp/src/test/java/io/modelcontextprotocol/spec/McpSchemaTests.java

+16-6
Original file line numberDiff line numberDiff line change
@@ -524,10 +524,16 @@ void testCreateMessageRequest() throws Exception {
524524
Map<String, Object> metadata = new HashMap<>();
525525
metadata.put("session", "test-session");
526526

527-
McpSchema.CreateMessageRequest request = new McpSchema.CreateMessageRequest(Collections.singletonList(message),
528-
preferences, "You are a helpful assistant",
529-
McpSchema.CreateMessageRequest.ContextInclusionStrategy.THIS_SERVER, 0.7, 1000,
530-
Arrays.asList("STOP", "END"), metadata);
527+
McpSchema.CreateMessageRequest request = McpSchema.CreateMessageRequest.builder()
528+
.messages(Collections.singletonList(message))
529+
.modelPreferences(preferences)
530+
.systemPrompt("You are a helpful assistant")
531+
.includeContext(McpSchema.CreateMessageRequest.ContextInclusionStrategy.THIS_SERVER)
532+
.temperature(0.7)
533+
.maxTokens(1000)
534+
.stopSequences(Arrays.asList("STOP", "END"))
535+
.metadata(metadata)
536+
.build();
531537

532538
String value = mapper.writeValueAsString(request);
533539

@@ -543,8 +549,12 @@ void testCreateMessageRequest() throws Exception {
543549
void testCreateMessageResult() throws Exception {
544550
McpSchema.TextContent content = new McpSchema.TextContent("Assistant response");
545551

546-
McpSchema.CreateMessageResult result = new McpSchema.CreateMessageResult(McpSchema.Role.ASSISTANT, content,
547-
"gpt-4", McpSchema.CreateMessageResult.StopReason.END_TURN);
552+
McpSchema.CreateMessageResult result = McpSchema.CreateMessageResult.builder()
553+
.role(McpSchema.Role.ASSISTANT)
554+
.content(content)
555+
.model("gpt-4")
556+
.stopReason(McpSchema.CreateMessageResult.StopReason.END_TURN)
557+
.build();
548558

549559
String value = mapper.writeValueAsString(result);
550560

0 commit comments

Comments
 (0)