From 8c6dfd53c8333a89116bfe8c4e51efd678cb53bf Mon Sep 17 00:00:00 2001 From: Peder Holdgaard Pedersen <127606677+PederHP@users.noreply.github.com> Date: Sun, 5 Jan 2025 20:05:32 +0100 Subject: [PATCH] Sampling capability (#7) --- src/mcpdotnet.Tests/IntegrationTests.cs | 65 +++++++++++++++ src/mcpdotnet/Client/IMcpClient.cs | 5 ++ src/mcpdotnet/Client/McpClient.cs | 64 ++++++++++++++- .../Protocol/Messages/JsonRpcResponse.cs | 4 +- .../Protocol/Transport/StdioTransport.cs | 1 + src/mcpdotnet/Protocol/Types/Content.cs | 10 ++- .../Protocol/Types/ContextInclusion.cs | 11 +++ .../Types/CreateMessageRequestParams.cs | 52 ++++++++++++ .../Protocol/Types/CreateMessageResult.cs | 35 ++++++++ src/mcpdotnet/Protocol/Types/ModelHint.cs | 22 +++++ .../Protocol/Types/ModelPreferences.cs | 82 +++++++++++++++++++ .../Protocol/Types/ResourceContents.cs | 8 +- .../Protocol/Types/SamplingMessage.cs | 18 ++++ src/mcpdotnet/mcpdotnet.csproj | 2 +- 14 files changed, 370 insertions(+), 9 deletions(-) create mode 100644 src/mcpdotnet/Protocol/Types/ContextInclusion.cs create mode 100644 src/mcpdotnet/Protocol/Types/CreateMessageRequestParams.cs create mode 100644 src/mcpdotnet/Protocol/Types/CreateMessageResult.cs create mode 100644 src/mcpdotnet/Protocol/Types/ModelHint.cs create mode 100644 src/mcpdotnet/Protocol/Types/ModelPreferences.cs create mode 100644 src/mcpdotnet/Protocol/Types/SamplingMessage.cs diff --git a/src/mcpdotnet.Tests/IntegrationTests.cs b/src/mcpdotnet.Tests/IntegrationTests.cs index 8245bdb..ef3c24b 100644 --- a/src/mcpdotnet.Tests/IntegrationTests.cs +++ b/src/mcpdotnet.Tests/IntegrationTests.cs @@ -341,4 +341,69 @@ public class IntegrationTests Assert.Single(result.Contents); Assert.NotNull(result.Contents[0].Blob); } + + [Fact] + public async Task Sampling_Stdio_EverythingServer() + { + // arrange + var config = new McpServerConfig + { + Id = "everything", + Name = "Everything", + TransportType = "stdio", + TransportOptions = new Dictionary + { + ["command"] = "npx", + ["arguments"] = "-y @modelcontextprotocol/server-everything" + } + }; + + var options = new McpClientOptions + { + ClientInfo = new() { Name = "IntegrationTestClient", Version = "1.0.0" }, + Capabilities = new ClientCapabilities + { + Sampling = new() + } + }; + + var factory = new McpClientFactory( + new[] { config }, + options + ); + var client = await factory.GetClientAsync("everything"); + + // Set up the sampling handler + int samplingHandlerCalls = 0; + client.SamplingHandler = async (_) => + { + samplingHandlerCalls++; + return new CreateMessageResult + { + Model = "test-model", + Role = "assistant", + Content = new Content + { + Type = "text", + Text = "Test response" + } + }; + }; + + // Call the server's sampleLLM tool which should trigger our sampling handler + var result = await client.CallToolAsync( + "sampleLLM", + new Dictionary + { + ["prompt"] = "Test prompt", + ["maxTokens"] = 100 + } + ); + + // assert + Assert.NotNull(result); + var textContent = Assert.Single(result.Content); + Assert.Equal("text", textContent.Type); + Assert.Contains("Test response", textContent.Text); + } } diff --git a/src/mcpdotnet/Client/IMcpClient.cs b/src/mcpdotnet/Client/IMcpClient.cs index 9d5e29d..39783b2 100644 --- a/src/mcpdotnet/Client/IMcpClient.cs +++ b/src/mcpdotnet/Client/IMcpClient.cs @@ -25,6 +25,11 @@ public interface IMcpClient : IAsyncDisposable /// Implementation? ServerInfo { get; } + /// + /// Gets or sets the handler for the server's sampling requests. + /// + Func>? SamplingHandler { get; set; } + /// /// Establishes a connection to the server. /// diff --git a/src/mcpdotnet/Client/McpClient.cs b/src/mcpdotnet/Client/McpClient.cs index d209fa0..5f02409 100644 --- a/src/mcpdotnet/Client/McpClient.cs +++ b/src/mcpdotnet/Client/McpClient.cs @@ -7,6 +7,8 @@ using McpDotNet.Protocol.Messages; using System.Diagnostics; using System.Text.Json; using McpDotNet.Utils.Json; +using System.Text; +using System.Threading; /// internal class McpClient : IMcpClient @@ -15,6 +17,7 @@ internal class McpClient : IMcpClient private readonly McpClientOptions _options; private readonly ConcurrentDictionary> _pendingRequests; private readonly ConcurrentDictionary>> _notificationHandlers; + private readonly Dictionary>> _requestHandlers = new(); private int _nextRequestId; private bool _isInitialized; private Task? _messageProcessingTask; @@ -42,6 +45,17 @@ internal class McpClient : IMcpClient _pendingRequests = new(); _notificationHandlers = new(); _nextRequestId = 1; + + if (options.Capabilities?.Sampling != null) + { + SetRequestHandler("sampling/createMessage", + async (request) => { + if (SamplingHandler == null) + throw new McpClientException("Sampling handler not configured"); + + return await SamplingHandler(request); + }); + } } /// @@ -271,7 +285,7 @@ internal class McpClient : IMcpClient { await foreach (var message in _transport.MessageReader.ReadAllAsync(cancellationToken).ConfigureAwait(false)) { - await HandleMessageAsync(message).ConfigureAwait(false); + await HandleMessageAsync(message, cancellationToken).ConfigureAwait(false); } } catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) @@ -284,10 +298,39 @@ internal class McpClient : IMcpClient } } - private async Task HandleMessageAsync(IJsonRpcMessage message) + private async Task HandleMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken) { switch (message) { + case JsonRpcRequest request: + if (_requestHandlers.TryGetValue(request.Method, out var handler)) + { + try + { + var result = await handler(request); + await _transport.SendMessageAsync(new JsonRpcResponse + { + Id = request.Id, + JsonRpc = "2.0", + Result = result + }, cancellationToken); + } + catch (Exception ex) + { + // Send error response + await _transport.SendMessageAsync(new JsonRpcError + { + Id = request.Id, + JsonRpc = "2.0", + Error = new JsonRpcErrorDetail + { + Code = -32000, // Implementation defined error + Message = ex.Message + } + }, cancellationToken); + } + } + break; case IJsonRpcMessageWithId messageWithId: if (_pendingRequests.TryRemove(messageWithId.Id, out var tcs)) { @@ -298,11 +341,11 @@ internal class McpClient : IMcpClient case JsonRpcNotification notification: if (_notificationHandlers.TryGetValue(notification.Method, out var handlers)) { - foreach (var handler in handlers) + foreach (var notificationHandler in handlers) { try { - await handler(notification).ConfigureAwait(false); + await notificationHandler(notification).ConfigureAwait(false); } catch (Exception ex) { @@ -383,6 +426,18 @@ internal class McpClient : IMcpClient } } + public Func>? SamplingHandler { get; set; } + + private void SetRequestHandler(string method, Func> handler) + { + _requestHandlers[method] = async (request) => { + // Convert the params JsonElement to our type using the same options + var jsonString = JsonSerializer.Serialize(request.Params); + var typedRequest = JsonSerializer.Deserialize(jsonString, _jsonOptions); + return await handler(typedRequest); + }; + } + private async Task CleanupAsync() { _cts?.Cancel(); @@ -419,3 +474,4 @@ internal class McpClient : IMcpClient GC.SuppressFinalize(this); } } + diff --git a/src/mcpdotnet/Protocol/Messages/JsonRpcResponse.cs b/src/mcpdotnet/Protocol/Messages/JsonRpcResponse.cs index dde1693..eca1ca7 100644 --- a/src/mcpdotnet/Protocol/Messages/JsonRpcResponse.cs +++ b/src/mcpdotnet/Protocol/Messages/JsonRpcResponse.cs @@ -1,7 +1,7 @@ -namespace McpDotNet.Protocol.Messages; - + using System.Text.Json.Serialization; +namespace McpDotNet.Protocol.Messages; /// /// A successful response message in the JSON-RPC protocol. /// diff --git a/src/mcpdotnet/Protocol/Transport/StdioTransport.cs b/src/mcpdotnet/Protocol/Transport/StdioTransport.cs index 9469dcb..963fbfa 100644 --- a/src/mcpdotnet/Protocol/Transport/StdioTransport.cs +++ b/src/mcpdotnet/Protocol/Transport/StdioTransport.cs @@ -142,6 +142,7 @@ public sealed class StdioTransport : TransportBase try { var json = JsonSerializer.Serialize(message, _jsonOptions); + Debug.WriteLine($"Sending message: {json}"); // Write the message followed by a newline await _process!.StandardInput.WriteLineAsync(json.AsMemory(), cancellationToken); diff --git a/src/mcpdotnet/Protocol/Types/Content.cs b/src/mcpdotnet/Protocol/Types/Content.cs index 313d421..8bfbe87 100644 --- a/src/mcpdotnet/Protocol/Types/Content.cs +++ b/src/mcpdotnet/Protocol/Types/Content.cs @@ -1,4 +1,6 @@ -namespace McpDotNet.Protocol.Types; +using System.Text.Json.Serialization; + +namespace McpDotNet.Protocol.Types; /// /// Represents the content of a tool response. @@ -10,26 +12,32 @@ public class Content /// /// The type of content. This determines the structure of the content object. Can be "image", "text", "resource". /// + + [JsonPropertyName("type")] public string Type { get; set; } = string.Empty; /// /// The text content of the message. /// + [JsonPropertyName("text")] public string? Text { get; set; } /// /// The base64-encoded image data. /// + [JsonPropertyName("data")] public string? Data { get; set; } /// /// The MIME type of the image. /// + [JsonPropertyName("mime_type")] public string? MimeType { get; set; } /// /// The resource content of the message (if embedded). /// + [JsonPropertyName("resource")] public ResourceContents? Resource { get; set; } /// diff --git a/src/mcpdotnet/Protocol/Types/ContextInclusion.cs b/src/mcpdotnet/Protocol/Types/ContextInclusion.cs new file mode 100644 index 0000000..0aa1996 --- /dev/null +++ b/src/mcpdotnet/Protocol/Types/ContextInclusion.cs @@ -0,0 +1,11 @@ +namespace McpDotNet.Protocol.Types; + +/// +/// A request to include context from one or more MCP servers (including the caller), to be attached to the prompt. +/// +public enum ContextInclusion +{ + None, + ThisServer, + AllServers +} diff --git a/src/mcpdotnet/Protocol/Types/CreateMessageRequestParams.cs b/src/mcpdotnet/Protocol/Types/CreateMessageRequestParams.cs new file mode 100644 index 0000000..66cd3ea --- /dev/null +++ b/src/mcpdotnet/Protocol/Types/CreateMessageRequestParams.cs @@ -0,0 +1,52 @@ +namespace McpDotNet.Protocol.Types; + +/// +/// A request from the server to sample an LLM via the client. +/// The client has full discretion over which model to select. +/// The client should also inform the user before beginning sampling, to allow them to inspect the request (human in the loop) and decide whether to approve it. +/// +/// While these align with the protocol specification, +/// clients have full discretion over model selection and should inform users before sampling. +/// +public class CreateMessageRequestParams +{ + /// + /// A request to include context from one or more MCP servers (including the caller), to be attached to the prompt. The client MAY ignore this request. + /// + public ContextInclusion? IncludeContext { get; init; } + + /// + /// The maximum number of tokens to sample, as requested by the server. The client MAY choose to sample fewer tokens than requested. + /// + public int MaxTokens { get; init; } + + /// + /// Messages requested by the server to be included in the prompt. + /// + public required IReadOnlyList Messages { get; init; } + + /// + /// Optional metadata to pass through to the LLM provider. The format of this metadata is provider-specific. + /// + public object? Metadata { get; init; } + + /// + /// The server's preferences for which model to select. The client MAY ignore these preferences. + /// + public ModelPreferences? ModelPreferences { get; init; } + + /// + /// Optional stop sequences that the server wants to use for sampling. + /// + public IReadOnlyList? StopSequences { get; init; } + + /// + /// An optional system prompt the server wants to use for sampling. The client MAY modify or omit this prompt. + /// + public string? SystemPrompt { get; init; } + + /// + /// The temperature to use for sampling, as requested by the server. + /// + public float? Temperature { get; init; } +} diff --git a/src/mcpdotnet/Protocol/Types/CreateMessageResult.cs b/src/mcpdotnet/Protocol/Types/CreateMessageResult.cs new file mode 100644 index 0000000..48352ce --- /dev/null +++ b/src/mcpdotnet/Protocol/Types/CreateMessageResult.cs @@ -0,0 +1,35 @@ +using System.Text.Json.Serialization; + +namespace McpDotNet.Protocol.Types; + +/// +/// The client's response to a sampling/create_message request from the server. +/// The client should inform the user before returning the sampled message, to allow them to inspect the response (human in the loop) +/// and decide whether to allow the server to see it. +/// +public class CreateMessageResult +{ + /// + /// Text or image content of the message. + /// + [JsonPropertyName("content")] + public required Content Content { get; init; } + + /// + /// The name of the model that generated the message. + /// + [JsonPropertyName("model")] + public required string Model { get; init; } + + /// + /// The reason why sampling stopped, if known. + /// + [JsonPropertyName("stop_reason")] + public string? StopReason { get; init; } + + /// + /// The role of the user who generated the message. + /// + [JsonPropertyName("role")] + public required string Role { get; init; } +} diff --git a/src/mcpdotnet/Protocol/Types/ModelHint.cs b/src/mcpdotnet/Protocol/Types/ModelHint.cs new file mode 100644 index 0000000..01026af --- /dev/null +++ b/src/mcpdotnet/Protocol/Types/ModelHint.cs @@ -0,0 +1,22 @@ +namespace McpDotNet.Protocol.Types; + +/// +/// Hints to use for model selection. +/// Keys not declared here are currently left unspecified by the spec and are up +/// to the client to interpret. +/// +public class ModelHint +{ + /// + /// A hint for a model name. + /// + /// The client SHOULD treat this as a substring of a model name; for example: + /// - `claude-3-5-sonnet` should match `claude-3-5-sonnet-20241022` + /// - `sonnet` should match `claude-3-5-sonnet-20241022`, `claude-3-sonnet-20240229`, etc. + /// - `claude` should match any Claude model + /// + /// The client MAY also map the string to a different provider's model name or a different model family, as long as it fills a similar niche; for example: + /// - `gemini-1.5-flash` could match `claude-3-haiku-20240307` + /// + public string? Name { get; init; } +} \ No newline at end of file diff --git a/src/mcpdotnet/Protocol/Types/ModelPreferences.cs b/src/mcpdotnet/Protocol/Types/ModelPreferences.cs new file mode 100644 index 0000000..6fa469f --- /dev/null +++ b/src/mcpdotnet/Protocol/Types/ModelPreferences.cs @@ -0,0 +1,82 @@ +namespace McpDotNet.Protocol.Types; + +/// +/// The server's preferences for model selection, requested of the client during sampling. +/// Because LLMs can vary along multiple dimensions, choosing the \"best\" model is +/// rarely straightforward. Different models excel in different areas—some are +/// faster but less capable, others are more capable but more expensive, and so +/// on. This interface allows servers to express their priorities across multiple +/// dimensions to help clients make an appropriate selection for their use case. +/// +/// These preferences are always advisory. The client MAY ignore them. It is also +/// up to the client to decide how to interpret these preferences and how to +/// balance them against other considerations. +/// +public class ModelPreferences +{ + /// + /// How much to prioritize cost when selecting a model. A value of 0 means cost\nis not important, while a value of 1 means cost is the most important + /// factor. + /// + public float? CostPriority { get; init; } + + /// + /// Optional hints to use for model selection. + /// + /// If multiple hints are specified, the client MUST evaluate them in order + /// (such that the first match is taken). + /// + /// The client SHOULD prioritize these hints over the numeric priorities, but + /// MAY still use the priorities to select from ambiguous matches. + /// + public IReadOnlyList? Hints { get; init; } + + /// + /// How much to prioritize intelligence and capabilities when selecting a + /// model. A value of 0 means intelligence is not important, while a value of 1 + /// means intelligence is the most important factor. + /// + public float? SpeedPriority { get; init; } + + /// + /// How much to prioritize sampling speed (latency) when selecting a model. A + /// value of 0 means speed is not important, while a value of 1 means speed is + /// the most important factor. + /// + public float? IntelligencePriority { get; init; } + + /// + /// Validates the model preferences. + /// + /// Error message if object isn't valid + /// True if valid, false if invalid + public bool Validate(out string errorMessage) + { + bool valid = true; + List errors = new(); + if (CostPriority < 0 || CostPriority > 1) + { + errors.Add("CostPriority must be between 0 and 1"); + valid = false; + } + if (SpeedPriority < 0 || SpeedPriority > 1) + { + errors.Add("SpeedPriority must be between 0 and 1"); + valid = false; + } + if (IntelligencePriority < 0 || IntelligencePriority > 1) + { + errors.Add("IntelligencePriority must be between 0 and 1"); + valid = false; + } + if (!valid) + { + errorMessage = string.Join(", ", errors); + } + else + { + errorMessage = ""; + } + return valid; + } +} diff --git a/src/mcpdotnet/Protocol/Types/ResourceContents.cs b/src/mcpdotnet/Protocol/Types/ResourceContents.cs index 2950f95..06bace9 100644 --- a/src/mcpdotnet/Protocol/Types/ResourceContents.cs +++ b/src/mcpdotnet/Protocol/Types/ResourceContents.cs @@ -1,4 +1,6 @@ -namespace McpDotNet.Protocol.Types; +using System.Text.Json.Serialization; + +namespace McpDotNet.Protocol.Types; /// /// Represents the content of a resource. @@ -8,21 +10,25 @@ public class ResourceContents /// /// The URI of the resource. /// + [JsonPropertyName("uri")] public string Uri { get; set; } = string.Empty; /// /// The type of content. /// + [JsonPropertyName("mime_type")] public string? MimeType { get; set; } /// /// The text content of the resource. /// + [JsonPropertyName("text")] public string? Text { get; set; } /// /// The base64-encoded binary content of the resource. /// + [JsonPropertyName("blob")] public string? Blob { get; set; } } \ No newline at end of file diff --git a/src/mcpdotnet/Protocol/Types/SamplingMessage.cs b/src/mcpdotnet/Protocol/Types/SamplingMessage.cs new file mode 100644 index 0000000..1b36e6d --- /dev/null +++ b/src/mcpdotnet/Protocol/Types/SamplingMessage.cs @@ -0,0 +1,18 @@ +namespace McpDotNet.Protocol.Types; + +/// +/// Describes a message issued to or received from an LLM API. +/// +public class SamplingMessage +{ + /// + /// Text or image content of the message. + /// + public required Content Content { get; init; } + + /// + /// The role of the message ("user" or "assistant"). + /// + + public required Role Role { get; init; } // "user" or "assistant" +} diff --git a/src/mcpdotnet/mcpdotnet.csproj b/src/mcpdotnet/mcpdotnet.csproj index 9e5afee..00c8ceb 100644 --- a/src/mcpdotnet/mcpdotnet.csproj +++ b/src/mcpdotnet/mcpdotnet.csproj @@ -9,7 +9,7 @@ mcpdotnet - 0.2.0-alpha1 + 0.3.0.1 PederHP .NET client library for the Model Context Protocol (MCP) https://github.com/PederHP/mcpdotnet