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