|
21 | 21 | import io.modelcontextprotocol.server.McpSyncServerExchange;
|
22 | 22 | import io.modelcontextprotocol.spec.McpSchema;
|
23 | 23 | import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult;
|
| 24 | +import io.modelcontextprotocol.spec.McpSchema.LoggingLevel; |
| 25 | +import io.modelcontextprotocol.spec.McpSchema.LoggingMessageNotification; |
24 | 26 | import io.modelcontextprotocol.spec.McpSchema.ModelPreferences;
|
25 | 27 | import org.slf4j.Logger;
|
26 | 28 |
|
27 | 29 | import org.springframework.ai.chat.model.ToolContext;
|
| 30 | +import org.springframework.ai.mcp.McpToolUtils; |
28 | 31 | import org.springframework.ai.model.ModelOptionsUtils;
|
29 | 32 | import org.springframework.ai.tool.annotation.Tool;
|
30 | 33 | import org.springframework.ai.tool.annotation.ToolParam;
|
@@ -72,43 +75,54 @@ public String getTemperature(@ToolParam(description = "The location latitude") d
|
72 | 75 |
|
73 | 76 | public String callMcpSampling(ToolContext toolContext, WeatherResponse weatherResponse) {
|
74 | 77 |
|
75 |
| - String openAiWeatherPoem = "<no OpenAI poem>"; |
76 |
| - String anthropicWeatherPoem = "<no Anthropic poem>"; |
| 78 | + StringBuilder openAiWeatherPoem = new StringBuilder(); |
| 79 | + StringBuilder anthropicWeatherPoem = new StringBuilder(); |
77 | 80 |
|
78 |
| - if (toolContext != null && toolContext.getContext().containsKey("exchange")) { |
| 81 | + McpToolUtils.getMcpExchange(toolContext) |
| 82 | + .ifPresent(exchange -> { |
79 | 83 |
|
80 |
| - // Spring AI MCP Auto-configuration injects the McpSyncServerExchange into the ToolContext under the key "exchange" |
81 |
| - McpSyncServerExchange exchange = (McpSyncServerExchange) toolContext.getContext().get("exchange"); |
82 |
| - if (exchange.getClientCapabilities().sampling() != null) { |
83 |
| - var messageRequestBuilder = McpSchema.CreateMessageRequest.builder() |
84 |
| - .systemPrompt("You are a poet!") |
85 |
| - .messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, |
86 |
| - new McpSchema.TextContent( |
87 |
| - "Please write a poem about thius weather forecast (temperature is in Celsious). Use markdown format :\n " |
88 |
| - + ModelOptionsUtils.toJsonStringPrettyPrinter(weatherResponse))))); |
| 84 | + exchange.loggingNotification(LoggingMessageNotification.builder() |
| 85 | + .level(LoggingLevel.INFO) |
| 86 | + .data("Start sampling") |
| 87 | + .build()); |
89 | 88 |
|
90 |
| - var opeAiLlmMessageRequest = messageRequestBuilder |
91 |
| - .modelPreferences(ModelPreferences.builder().addHint("openai").build()) |
92 |
| - .build(); |
93 |
| - CreateMessageResult openAiLlmResponse = exchange.createMessage(opeAiLlmMessageRequest); |
| 89 | + if (exchange.getClientCapabilities().sampling() != null) { |
| 90 | + var messageRequestBuilder = McpSchema.CreateMessageRequest.builder() |
| 91 | + .systemPrompt("You are a poet!") |
| 92 | + .messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, |
| 93 | + new McpSchema.TextContent( |
| 94 | + "Please write a poem about thius weather forecast (temperature is in Celsious). Use markdown format :\n " |
| 95 | + + ModelOptionsUtils |
| 96 | + .toJsonStringPrettyPrinter(weatherResponse))))); |
94 | 97 |
|
95 |
| - openAiWeatherPoem = ((McpSchema.TextContent) openAiLlmResponse.content()).text(); |
| 98 | + var opeAiLlmMessageRequest = messageRequestBuilder |
| 99 | + .modelPreferences(ModelPreferences.builder().addHint("openai").build()) |
| 100 | + .build(); |
| 101 | + CreateMessageResult openAiLlmResponse = exchange.createMessage(opeAiLlmMessageRequest); |
96 | 102 |
|
97 |
| - var anthropicLlmMessageRequest = messageRequestBuilder |
98 |
| - .modelPreferences(ModelPreferences.builder().addHint("anthropic").build()) |
99 |
| - .build(); |
100 |
| - CreateMessageResult anthropicAiLlmResponse = exchange.createMessage(anthropicLlmMessageRequest); |
| 103 | + openAiWeatherPoem.append(((McpSchema.TextContent) openAiLlmResponse.content()).text()); |
101 | 104 |
|
102 |
| - anthropicWeatherPoem = ((McpSchema.TextContent) anthropicAiLlmResponse.content()).text(); |
| 105 | + var anthropicLlmMessageRequest = messageRequestBuilder |
| 106 | + .modelPreferences(ModelPreferences.builder().addHint("anthropic").build()) |
| 107 | + .build(); |
| 108 | + CreateMessageResult anthropicAiLlmResponse = exchange.createMessage(anthropicLlmMessageRequest); |
103 | 109 |
|
104 |
| - } |
105 |
| - } |
| 110 | + anthropicWeatherPoem.append(((McpSchema.TextContent) anthropicAiLlmResponse.content()).text()); |
| 111 | + |
| 112 | + } |
| 113 | + |
| 114 | + exchange.loggingNotification(LoggingMessageNotification.builder() |
| 115 | + .level(LoggingLevel.INFO) |
| 116 | + .data("Finish Sampling") |
| 117 | + .build()); |
| 118 | + |
| 119 | + }); |
106 | 120 |
|
107 |
| - String responseWithPoems = "OpenAI poem about the weather: " + openAiWeatherPoem + "\n\n" + |
108 |
| - "Anthropic poem about the weather: " + anthropicWeatherPoem + "\n" |
| 121 | + String responseWithPoems = "OpenAI poem about the weather: " + openAiWeatherPoem.toString() + "\n\n" + |
| 122 | + "Anthropic poem about the weather: " + anthropicWeatherPoem.toString() + "\n" |
109 | 123 | + ModelOptionsUtils.toJsonStringPrettyPrinter(weatherResponse);
|
110 | 124 |
|
111 |
| - logger.info(anthropicWeatherPoem, responseWithPoems); |
| 125 | + logger.info(anthropicWeatherPoem.toString(), responseWithPoems.toString()); |
112 | 126 |
|
113 | 127 | return responseWithPoems;
|
114 | 128 |
|
|
0 commit comments