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:
parent
ddde66dc08
commit
fac4df6b1a
4 changed files with 1584 additions and 0 deletions
126
Sources/MCP/Base/UnitInterval.swift
Normal file
126
Sources/MCP/Base/UnitInterval.swift
Normal 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))!
|
||||
}
|
||||
}
|
||||
|
|
@ -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 {
|
||||
|
|
|
|||
196
Sources/MCP/Server/Sampling.swift
Normal file
196
Sources/MCP/Server/Sampling.swift
Normal 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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue