feat(completion): implement completion API support (#141)
Add completion API support to the MCP protocol implementation: - Add CompleteRequest and CompleteResult schema classes - Implement completion capabilities in ServerCapabilities - Add completion handlers in McpAsyncServer and McpServer - Add completion client methods in McpAsyncClient and McpSyncClient - Add CompletionRefKey and completion specifications in McpServerFeatures - Add integration test for completion functionality - Fix isPresent() check to use isEmpty() in WebMvcSseServerTransportProvider - Replace McpServerFeatures.CompletionRefKey by McpSchemaCompleteReference Co-authored-by: Christian Tzolov <christian.tzolov@broadcom.com> Signed-off-by: jitokim <pigberger70@gmail.com>
This commit is contained in:
parent
84adde16a5
commit
7853efefda
2 changed files with 340 additions and 33 deletions
|
|
@ -10,6 +10,7 @@ import java.util.Map;
|
|||
import java.util.concurrent.ConcurrentHashMap;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.function.BiFunction;
|
||||
import java.util.function.Function;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
|
|
@ -20,19 +21,12 @@ import io.modelcontextprotocol.client.transport.WebFluxSseClientTransport;
|
|||
import io.modelcontextprotocol.server.McpServer;
|
||||
import io.modelcontextprotocol.server.McpServerFeatures;
|
||||
import io.modelcontextprotocol.server.TestUtil;
|
||||
import io.modelcontextprotocol.server.McpSyncServerExchange;
|
||||
import io.modelcontextprotocol.server.transport.WebFluxSseServerTransportProvider;
|
||||
import io.modelcontextprotocol.spec.McpError;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpSchema.CallToolResult;
|
||||
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
|
||||
import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest;
|
||||
import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult;
|
||||
import io.modelcontextprotocol.spec.McpSchema.InitializeResult;
|
||||
import io.modelcontextprotocol.spec.McpSchema.ModelPreferences;
|
||||
import io.modelcontextprotocol.spec.McpSchema.Role;
|
||||
import io.modelcontextprotocol.spec.McpSchema.Root;
|
||||
import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities;
|
||||
import io.modelcontextprotocol.spec.McpSchema.Tool;
|
||||
import io.modelcontextprotocol.spec.McpSchema.*;
|
||||
import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities.CompletionCapabilities;
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.params.ParameterizedTest;
|
||||
|
|
@ -759,4 +753,53 @@ class WebFluxSseIntegrationTests {
|
|||
mcpServer.close();
|
||||
}
|
||||
|
||||
}
|
||||
// ---------------------------------------
|
||||
// Completion Tests
|
||||
// ---------------------------------------
|
||||
@ParameterizedTest(name = "{0} : Completion call")
|
||||
@ValueSource(strings = { "httpclient", "webflux" })
|
||||
void testCompletionShouldReturnExpectedSuggestions(String clientType) {
|
||||
var clientBuilder = clientBuilders.get(clientType);
|
||||
|
||||
var expectedValues = List.of("python", "pytorch", "pyside");
|
||||
var completionResponse = new McpSchema.CompleteResult(new CompleteResult.CompleteCompletion(expectedValues, 10, // total
|
||||
true // hasMore
|
||||
));
|
||||
|
||||
AtomicReference<CompleteRequest> samplingRequest = new AtomicReference<>();
|
||||
BiFunction<McpSyncServerExchange, CompleteRequest, CompleteResult> completionHandler = (mcpSyncServerExchange,
|
||||
request) -> {
|
||||
samplingRequest.set(request);
|
||||
return completionResponse;
|
||||
};
|
||||
|
||||
var mcpServer = McpServer.sync(mcpServerTransportProvider)
|
||||
.capabilities(ServerCapabilities.builder().completions(new CompletionCapabilities()).build())
|
||||
.prompts(new McpServerFeatures.SyncPromptSpecification(
|
||||
new Prompt("code_review", "this is code review prompt", List.of()),
|
||||
(mcpSyncServerExchange, getPromptRequest) -> null))
|
||||
.completions(new McpServerFeatures.SyncCompletionSpecification(
|
||||
new McpSchema.PromptReference("ref/prompt", "code_review"), completionHandler))
|
||||
.build();
|
||||
|
||||
try (var mcpClient = clientBuilder.build()) {
|
||||
|
||||
InitializeResult initResult = mcpClient.initialize();
|
||||
assertThat(initResult).isNotNull();
|
||||
|
||||
CompleteRequest request = new CompleteRequest(new PromptReference("ref/prompt", "code_review"),
|
||||
new CompleteRequest.CompleteArgument("language", "py"));
|
||||
|
||||
CompleteResult result = mcpClient.completeCompletion(request);
|
||||
|
||||
assertThat(result).isNotNull();
|
||||
|
||||
assertThat(samplingRequest.get().argument().name()).isEqualTo("language");
|
||||
assertThat(samplingRequest.get().argument().value()).isEqualTo("py");
|
||||
assertThat(samplingRequest.get().ref().type()).isEqualTo("ref/prompt");
|
||||
}
|
||||
|
||||
mcpServer.close();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -300,7 +300,7 @@ public class WebMvcSseServerTransportProvider implements McpServerTransportProvi
|
|||
return ServerResponse.status(HttpStatus.SERVICE_UNAVAILABLE).body("Server is shutting down");
|
||||
}
|
||||
|
||||
if (!request.param("sessionId").isPresent()) {
|
||||
if (request.param("sessionId").isEmpty()) {
|
||||
return ServerResponse.badRequest().body(new McpError("Session ID missing in message endpoint"));
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue