Add support for sampling (#119)

* Add UnitInterval type

* Add test coverage for UnitInterval

* Add support for sampling

* Add integration tests for sampling
This commit is contained in:
Mattt 2025-05-26 04:57:24 -07:00 • committed by GitHub
commit fac4df6b1a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 1584 additions and 0 deletions

View file

@ -0,0 +1,126 @@
/// A value constrained to the range 0.0 to 1.0, inclusive.
///
/// `UnitInterval` represents a normalized value that is guaranteed to be within
/// the unit interval [0, 1]. This type is commonly used for representing
/// priorities in sampling request model preferences.
///
/// The type provides safe initialization that returns `nil` for values outside
/// the valid range, ensuring that all instances contain valid unit interval values.
///
/// - Example:
/// ```swift
/// let zero: UnitInterval = 0 // 0.0
/// let half = UnitInterval(0.5)! // 0.5
/// let one: UnitInterval = 1.0 // 1.0
/// let invalid = UnitInterval(1.5) // nil
/// ```
public struct UnitInterval: Hashable, Sendable {
private let value: Double
/// Creates a unit interval value from a `Double`.
///
/// - Parameter value: A double value that must be in the range 0.0...1.0
/// - Returns: A `UnitInterval` instance if the value is valid, `nil` otherwise
///
/// - Example:
/// ```swift
/// let valid = UnitInterval(0.75) // Optional(0.75)
/// let invalid = UnitInterval(-0.1) // nil
/// let boundary = UnitInterval(1.0) // Optional(1.0)
/// ```
public init?(_ value: Double) {
guard (0...1).contains(value) else { return nil }
self.value = value
}
/// The underlying double value.
///
/// This property provides access to the raw double value that is guaranteed
/// to be within the range [0, 1].
///
/// - Returns: The double value between 0.0 and 1.0, inclusive
public var doubleValue: Double { value }
}
// MARK: - Comparable
extension UnitInterval: Comparable {
public static func < (lhs: UnitInterval, rhs: UnitInterval) -> Bool {
lhs.value < rhs.value
}
}
// MARK: - CustomStringConvertible
extension UnitInterval: CustomStringConvertible {
public var description: String { "\(value)" }
}
// MARK: - Codable
extension UnitInterval: Codable {
public init(from decoder: Decoder) throws {
let container = try decoder.singleValueContainer()
let doubleValue = try container.decode(Double.self)
guard let interval = UnitInterval(doubleValue) else {
throw DecodingError.dataCorrupted(
DecodingError.Context(
codingPath: decoder.codingPath,
debugDescription: "Value \(doubleValue) is not in range 0...1")
)
}
self = interval
}
public func encode(to encoder: Encoder) throws {
var container = encoder.singleValueContainer()
try container.encode(value)
}
}
// MARK: - ExpressibleByFloatLiteral
extension UnitInterval: ExpressibleByFloatLiteral {
/// Creates a unit interval from a floating-point literal.
///
/// This initializer allows you to create `UnitInterval` instances using
/// floating-point literals. The literal value must be in the range [0, 1]
/// or a runtime error will occur.
///
/// - Parameter value: A floating-point literal between 0.0 and 1.0
///
/// - Warning: This initializer will crash if the literal is outside the valid range.
/// Use the failable initializer `init(_:)` for runtime validation.
///
/// - Example:
/// ```swift
/// let quarter: UnitInterval = 0.25
/// let half: UnitInterval = 0.5
/// ```
public init(floatLiteral value: Double) {
self.init(value)!
}
}
// MARK: - ExpressibleByIntegerLiteral
extension UnitInterval: ExpressibleByIntegerLiteral {
/// Creates a unit interval from an integer literal.
///
/// This initializer allows you to create `UnitInterval` instances using
/// integer literals. Only the values 0 and 1 are valid.
///
/// - Parameter value: An integer literal, either 0 or 1
///
/// - Warning: This initializer will crash if the literal is outside the valid range.
/// Use the failable initializer `init(_:)` for runtime validation.
///
/// - Example:
/// ```swift
/// let zero: UnitInterval = 0
/// let one: UnitInterval = 1
/// ```
public init(integerLiteral value: Int) {
self.init(Double(value))!
}
}

View file

@ -583,6 +583,45 @@ public actor Client {
return (content: result.content, isError: result.isError)
}
// MARK: - Sampling
/// Register a handler for sampling requests from servers
///
/// Sampling allows servers to request LLM completions through the client,
/// enabling sophisticated agentic behaviors while maintaining human-in-the-loop control.
///
/// The sampling flow follows these steps:
/// 1. Server sends a `sampling/createMessage` request to the client
/// 2. Client reviews the request and can modify it (via this handler)
/// 3. Client samples from an LLM (via this handler)
/// 4. Client reviews the completion (via this handler)
/// 5. Client returns the result to the server
///
/// - Parameter handler: A closure that processes sampling requests and returns completions
/// - Returns: Self for method chaining
/// - SeeAlso: https://modelcontextprotocol.io/docs/concepts/sampling#how-sampling-works
@discardableResult
public func withSamplingHandler(
_ handler: @escaping @Sendable (CreateSamplingMessage.Parameters) async throws ->
CreateSamplingMessage.Result
) -> Self {
// Note: This would require extending the client architecture to handle incoming requests from servers.
// The current MCP Swift SDK architecture assumes clients only send requests to servers,
// but sampling requires bidirectional communication where servers can send requests to clients.
//
// A full implementation would need:
// 1. Request handlers in the client (similar to how servers handle requests)
// 2. Bidirectional transport support
// 3. Request/response correlation for server-to-client requests
//
// For now, this serves as the correct API design for when bidirectional support is added.
// This would register the handler similar to how servers register method handlers:
// methodHandlers[CreateSamplingMessage.name] = TypedRequestHandler(handler)
return self
}
// MARK: -
private func handleResponse(_ response: Response<AnyMethod>) async {

View file

@ -0,0 +1,196 @@
import Foundation
/// The Model Context Protocol (MCP) allows servers to request LLM completions
/// through the client, enabling sophisticated agentic behaviors while maintaining
/// security and privacy.
///
/// - SeeAlso: https://modelcontextprotocol.io/docs/concepts/sampling#how-sampling-works
public enum Sampling {
/// A message in the conversation history.
public struct Message: Hashable, Codable, Sendable {
/// The message role
public enum Role: String, Hashable, Codable, Sendable {
/// A user message
case user
/// An assistant message
case assistant
}
/// The message role
public let role: Role
/// The message content
public let content: Content
public init(role: Role, content: Content) {
self.role = role
self.content = content
}
/// Content types for sampling messages
public enum Content: Hashable, Codable, Sendable {
/// Text content
case text(String)
/// Image content
case image(data: String, mimeType: String)
private enum CodingKeys: String, CodingKey {
case type, text, data, mimeType
}
public init(from decoder: Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
let type = try container.decode(String.self, forKey: .type)
switch type {
case "text":
let text = try container.decode(String.self, forKey: .text)
self = .text(text)
case "image":
let data = try container.decode(String.self, forKey: .data)
let mimeType = try container.decode(String.self, forKey: .mimeType)
self = .image(data: data, mimeType: mimeType)
default:
throw DecodingError.dataCorruptedError(
forKey: .type, in: container,
debugDescription: "Unknown sampling message content type")
}
}
public func encode(to encoder: Encoder) throws {
var container = encoder.container(keyedBy: CodingKeys.self)
switch self {
case .text(let text):
try container.encode("text", forKey: .type)
try container.encode(text, forKey: .text)
case .image(let data, let mimeType):
try container.encode("image", forKey: .type)
try container.encode(data, forKey: .data)
try container.encode(mimeType, forKey: .mimeType)
}
}
}
}
/// Model preferences for sampling requests
public struct ModelPreferences: Hashable, Codable, Sendable {
/// Model hints for selection
public struct Hint: Hashable, Codable, Sendable {
/// Suggested model name/family
public let name: String?
public init(name: String? = nil) {
self.name = name
}
}
/// Array of model name suggestions that clients can use to select an appropriate model
public let hints: [Hint]?
/// Importance of minimizing costs (0-1 normalized)
public let costPriority: UnitInterval?
/// Importance of low latency response (0-1 normalized)
public let speedPriority: UnitInterval?
/// Importance of advanced model capabilities (0-1 normalized)
public let intelligencePriority: UnitInterval?
public init(
hints: [Hint]? = nil,
costPriority: UnitInterval? = nil,
speedPriority: UnitInterval? = nil,
intelligencePriority: UnitInterval? = nil
) {
self.hints = hints
self.costPriority = costPriority
self.speedPriority = speedPriority
self.intelligencePriority = intelligencePriority
}
}
/// Context inclusion options for sampling requests
public enum ContextInclusion: String, Hashable, Codable, Sendable {
/// No additional context
case none
/// Include context from the requesting server
case thisServer
/// Include context from all connected MCP servers
case allServers
}
/// Stop reason for sampling completion
public enum StopReason: String, Hashable, Codable, Sendable {
/// Natural end of turn
case endTurn
/// Hit a stop sequence
case stopSequence
/// Reached maximum tokens
case maxTokens
}
}
/// To request sampling from a client, servers send a `sampling/createMessage` request.
/// - SeeAlso: https://modelcontextprotocol.io/docs/concepts/sampling#how-sampling-works
public enum CreateSamplingMessage: Method {
public static let name = "sampling/createMessage"
public struct Parameters: Hashable, Codable, Sendable {
/// The conversation history to send to the LLM
public let messages: [Sampling.Message]
/// Model selection preferences
public let modelPreferences: Sampling.ModelPreferences?
/// Optional system prompt
public let systemPrompt: String?
/// What MCP context to include
public let includeContext: Sampling.ContextInclusion?
/// Controls randomness (0.0 to 1.0)
public let temperature: Double?
/// Maximum tokens to generate
public let maxTokens: Int
/// Array of sequences that stop generation
public let stopSequences: [String]?
/// Additional provider-specific parameters
public let metadata: [String: Value]?
public init(
messages: [Sampling.Message],
modelPreferences: Sampling.ModelPreferences? = nil,
systemPrompt: String? = nil,
includeContext: Sampling.ContextInclusion? = nil,
temperature: Double? = nil,
maxTokens: Int,
stopSequences: [String]? = nil,
metadata: [String: Value]? = nil
) {
self.messages = messages
self.modelPreferences = modelPreferences
self.systemPrompt = systemPrompt
self.includeContext = includeContext
self.temperature = temperature
self.maxTokens = maxTokens
self.stopSequences = stopSequences
self.metadata = metadata
}
}
public struct Result: Hashable, Codable, Sendable {
/// Name of the model used
public let model: String
/// Why sampling stopped
public let stopReason: Sampling.StopReason?
/// The role of the completion
public let role: Sampling.Message.Role
/// The completion content
public let content: Sampling.Message.Content
public init(
model: String,
stopReason: Sampling.StopReason? = nil,
role: Sampling.Message.Role,
content: Sampling.Message.Content
) {
self.model = model
self.stopReason = stopReason
self.role = role
self.content = content
}
}
}

View file

@ -82,12 +82,19 @@ public actor Server {
public init() {}
}
/// Sampling capabilities
public struct Sampling: Hashable, Codable, Sendable {
public init() {}
}
/// Logging capabilities
public var logging: Logging?
/// Prompts capabilities
public var prompts: Prompts?
/// Resources capabilities
public var resources: Resources?
/// Sampling capabilities
public var sampling: Sampling?
/// Tools capabilities
public var tools: Tools?
@ -95,11 +102,13 @@ public actor Server {
logging: Logging? = nil,
prompts: Prompts? = nil,
resources: Resources? = nil,
sampling: Sampling? = nil,
tools: Tools? = nil
) {
self.logging = logging
self.prompts = prompts
self.resources = resources
self.sampling = sampling
self.tools = tools
}
}
@ -290,6 +299,69 @@ public actor Server {
try await connection.send(notificationData)
}
// MARK: - Sampling
/// Request sampling from the connected client
///
/// Sampling allows servers to request LLM completions through the client,
/// enabling sophisticated agentic behaviors while maintaining human-in-the-loop control.
///
/// The sampling flow follows these steps:
/// 1. Server sends a `sampling/createMessage` request to the client
/// 2. Client reviews the request and can modify it
/// 3. Client samples from an LLM
/// 4. Client reviews the completion
/// 5. Client returns the result to the server
///
/// - Parameters:
/// - messages: The conversation history to send to the LLM
/// - modelPreferences: Model selection preferences
/// - systemPrompt: Optional system prompt
/// - includeContext: What MCP context to include
/// - temperature: Controls randomness (0.0 to 1.0)
/// - maxTokens: Maximum tokens to generate
/// - stopSequences: Array of sequences that stop generation
/// - metadata: Additional provider-specific parameters
/// - Returns: The sampling result containing the model used, stop reason, role, and content
/// - Throws: MCPError if the request fails
/// - SeeAlso: https://modelcontextprotocol.io/docs/concepts/sampling#how-sampling-works
public func requestSampling(
messages: [Sampling.Message],
modelPreferences: Sampling.ModelPreferences? = nil,
systemPrompt: String? = nil,
includeContext: Sampling.ContextInclusion? = nil,
temperature: Double? = nil,
maxTokens: Int,
stopSequences: [String]? = nil,
metadata: [String: Value]? = nil
) async throws -> CreateSamplingMessage.Result {
guard connection != nil else {
throw MCPError.internalError("Server connection not initialized")
}
// Note: This is a conceptual implementation. The actual implementation would require
// bidirectional communication support in the transport layer, allowing servers to
// send requests to clients and receive responses.
_ = CreateSamplingMessage.request(
.init(
messages: messages,
modelPreferences: modelPreferences,
systemPrompt: systemPrompt,
includeContext: includeContext,
temperature: temperature,
maxTokens: maxTokens,
stopSequences: stopSequences,
metadata: metadata
)
)
// This would need to be implemented with proper request/response handling
// similar to how the client sends requests to servers
throw MCPError.internalError(
"Bidirectional sampling requests not yet implemented in transport layer")
}
/// A JSON-RPC batch containing multiple requests and/or notifications
struct Batch: Sendable {
/// An item in a JSON-RPC batch