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
2 changed files with 1584 additions and 0 deletions

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