turul_mcp_protocol_2025_06_18/
sampling.rs

1//! MCP Sampling Protocol Types
2//!
3//! This module defines types for sampling requests in MCP.
4
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7use crate::prompts::ContentBlock;
8
9/// Sampling request parameters
10#[derive(Debug, Clone, Serialize, Deserialize)]
11#[serde(rename_all = "camelCase")]
12pub struct SamplingRequest {
13    /// The sampling method to use
14    pub method: String,
15    /// Parameters for the sampling method
16    #[serde(skip_serializing_if = "Option::is_none")]
17    pub params: Option<Value>,
18}
19
20/// Sampling response
21#[derive(Debug, Clone, Serialize, Deserialize)]
22#[serde(rename_all = "camelCase")]
23pub struct SamplingResult {
24    /// The sampled result
25    pub result: Value,
26}
27
28impl SamplingResult {
29    pub fn new(result: Value) -> Self {
30        Self { result }
31    }
32}
33
34/// Role enum for messages (per MCP spec)
35#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
36#[serde(rename_all = "lowercase")]
37pub enum Role {
38    User,
39    Assistant,
40    System,
41}
42
43/// Model hint enum (per MCP spec)
44#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
45#[serde(rename_all = "kebab-case")]
46pub enum ModelHint {
47    #[serde(rename = "claude-3-5-sonnet-20241022")]
48    Claude35Sonnet20241022,
49    #[serde(rename = "claude-3-5-haiku-20241022")]
50    Claude35Haiku20241022,
51    #[serde(rename = "gpt-4o")]
52    Gpt4o,
53    #[serde(rename = "gpt-4o-mini")]
54    Gpt4oMini,
55    #[serde(rename = "o1-preview")]
56    O1Preview,
57    #[serde(rename = "o1-mini")]
58    O1Mini,
59}
60
61/// Model preferences (per MCP spec)
62#[derive(Debug, Clone, Serialize, Deserialize)]
63#[serde(rename_all = "camelCase")]
64pub struct ModelPreferences {
65    /// Optional hints about which models to use
66    #[serde(skip_serializing_if = "Option::is_none")]
67    pub hints: Option<Vec<ModelHint>>,
68    /// Optional cost priority (0.0-1.0)
69    #[serde(skip_serializing_if = "Option::is_none")]
70    pub cost_priority: Option<f64>,
71    /// Optional speed priority (0.0-1.0)
72    #[serde(skip_serializing_if = "Option::is_none")]
73    pub speed_priority: Option<f64>,
74    /// Optional intelligence priority (0.0-1.0)
75    #[serde(skip_serializing_if = "Option::is_none")]
76    pub intelligence_priority: Option<f64>,
77}
78
79/// Sampling message (per MCP spec)
80#[derive(Debug, Clone, Serialize, Deserialize)]
81#[serde(rename_all = "camelCase")]
82pub struct SamplingMessage {
83    /// Role of the message
84    pub role: Role,
85    /// Content of the message
86    pub content: ContentBlock,
87}
88
89/// Parameters for sampling/createMessage request (per MCP spec)
90#[derive(Debug, Clone, Serialize, Deserialize)]
91#[serde(rename_all = "camelCase")]
92pub struct CreateMessageParams {
93    /// Messages for context
94    pub messages: Vec<SamplingMessage>,
95    /// Optional model preferences
96    #[serde(skip_serializing_if = "Option::is_none")]
97    pub model_preferences: Option<ModelPreferences>,
98    /// Optional system prompt
99    #[serde(skip_serializing_if = "Option::is_none")]
100    pub system_prompt: Option<String>,
101    /// Optional include context
102    #[serde(skip_serializing_if = "Option::is_none")]
103    pub include_context: Option<String>,
104    /// Optional temperature
105    #[serde(skip_serializing_if = "Option::is_none")]
106    pub temperature: Option<f64>,
107    /// Maximum tokens (required field)
108    pub max_tokens: u32,
109    /// Optional stop sequences
110    #[serde(skip_serializing_if = "Option::is_none")]
111    pub stop_sequences: Option<Vec<String>>,
112    /// Optional metadata
113    #[serde(skip_serializing_if = "Option::is_none")]
114    pub metadata: Option<Value>,
115    /// Meta information (optional _meta field inside params)
116    #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
117    pub meta: Option<std::collections::HashMap<String, Value>>,
118}
119
120/// Complete sampling/createMessage request (matches TypeScript CreateMessageRequest interface)
121#[derive(Debug, Clone, Serialize, Deserialize)]
122#[serde(rename_all = "camelCase")]
123pub struct CreateMessageRequest {
124    /// Method name (always "sampling/createMessage")
125    pub method: String,
126    /// Request parameters
127    pub params: CreateMessageParams,
128}
129
130/// Result for sampling/createMessage (per MCP spec)
131#[derive(Debug, Clone, Serialize, Deserialize)]
132#[serde(rename_all = "camelCase")]
133pub struct CreateMessageResult {
134    /// The generated message
135    pub message: SamplingMessage,
136    /// Model used for generation
137    pub model: String,
138    /// Stop reason
139    #[serde(skip_serializing_if = "Option::is_none")]
140    pub stop_reason: Option<String>,
141    /// Meta information (follows MCP Result interface)
142    #[serde(
143        default,
144        skip_serializing_if = "Option::is_none",
145        alias = "_meta",
146        rename = "_meta"
147    )]
148    pub meta: Option<std::collections::HashMap<String, Value>>,
149}
150
151impl CreateMessageParams {
152    pub fn new(messages: Vec<SamplingMessage>, max_tokens: u32) -> Self {
153        Self {
154            messages,
155            model_preferences: None,
156            system_prompt: None,
157            include_context: None,
158            temperature: None,
159            max_tokens,
160            stop_sequences: None,
161            metadata: None,
162            meta: None,
163        }
164    }
165
166    pub fn with_model_preferences(mut self, preferences: ModelPreferences) -> Self {
167        self.model_preferences = Some(preferences);
168        self
169    }
170
171    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
172        self.system_prompt = Some(prompt.into());
173        self
174    }
175
176    pub fn with_temperature(mut self, temperature: f64) -> Self {
177        self.temperature = Some(temperature);
178        self
179    }
180
181    pub fn with_stop_sequences(mut self, sequences: Vec<String>) -> Self {
182        self.stop_sequences = Some(sequences);
183        self
184    }
185
186    pub fn with_meta(mut self, meta: std::collections::HashMap<String, Value>) -> Self {
187        self.meta = Some(meta);
188        self
189    }
190}
191
192impl CreateMessageRequest {
193    pub fn new(messages: Vec<SamplingMessage>, max_tokens: u32) -> Self {
194        Self {
195            method: "sampling/createMessage".to_string(),
196            params: CreateMessageParams::new(messages, max_tokens),
197        }
198    }
199
200    pub fn with_model_preferences(mut self, preferences: ModelPreferences) -> Self {
201        self.params = self.params.with_model_preferences(preferences);
202        self
203    }
204
205    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
206        self.params = self.params.with_system_prompt(prompt);
207        self
208    }
209
210    pub fn with_temperature(mut self, temperature: f64) -> Self {
211        self.params = self.params.with_temperature(temperature);
212        self
213    }
214
215    pub fn with_stop_sequences(mut self, sequences: Vec<String>) -> Self {
216        self.params = self.params.with_stop_sequences(sequences);
217        self
218    }
219
220    pub fn with_meta(mut self, meta: std::collections::HashMap<String, Value>) -> Self {
221        self.params = self.params.with_meta(meta);
222        self
223    }
224}
225
226impl CreateMessageResult {
227    pub fn new(message: SamplingMessage, model: impl Into<String>) -> Self {
228        Self {
229            message,
230            model: model.into(),
231            stop_reason: None,
232            meta: None,
233        }
234    }
235
236    pub fn with_stop_reason(mut self, reason: impl Into<String>) -> Self {
237        self.stop_reason = Some(reason.into());
238        self
239    }
240
241    pub fn with_meta(mut self, meta: std::collections::HashMap<String, Value>) -> Self {
242        self.meta = Some(meta);
243        self
244    }
245}
246
247// Trait implementations for sampling
248
249use crate::traits::*;
250use std::collections::HashMap;
251
252// Trait implementations for CreateMessageParams
253impl Params for CreateMessageParams {}
254
255impl HasCreateMessageParams for CreateMessageParams {
256    fn messages(&self) -> &Vec<SamplingMessage> {
257        &self.messages
258    }
259    
260    fn model_preferences(&self) -> Option<&ModelPreferences> {
261        self.model_preferences.as_ref()
262    }
263    
264    fn system_prompt(&self) -> Option<&String> {
265        self.system_prompt.as_ref()
266    }
267    
268    fn include_context(&self) -> Option<&String> {
269        self.include_context.as_ref()
270    }
271    
272    fn temperature(&self) -> Option<&f64> {
273        self.temperature.as_ref()
274    }
275    
276    fn max_tokens(&self) -> u32 {
277        self.max_tokens
278    }
279    
280    fn stop_sequences(&self) -> Option<&Vec<String>> {
281        self.stop_sequences.as_ref()
282    }
283    
284    fn metadata(&self) -> Option<&Value> {
285        self.metadata.as_ref()
286    }
287}
288
289impl HasMetaParam for CreateMessageParams {
290    fn meta(&self) -> Option<&std::collections::HashMap<String, Value>> {
291        self.meta.as_ref()
292    }
293}
294
295// Trait implementations for CreateMessageRequest
296impl HasMethod for CreateMessageRequest {
297    fn method(&self) -> &str {
298        &self.method
299    }
300}
301
302impl HasParams for CreateMessageRequest {
303    fn params(&self) -> Option<&dyn Params> {
304        Some(&self.params)
305    }
306}
307
308// Trait implementations for CreateMessageResult
309impl HasData for CreateMessageResult {
310    fn data(&self) -> HashMap<String, Value> {
311        let mut data = HashMap::new();
312        data.insert("role".to_string(), serde_json::to_value(&self.message.role).unwrap_or(Value::String("user".to_string())));
313        data.insert("content".to_string(), serde_json::to_value(&self.message.content).unwrap_or(Value::Null));
314        data.insert("model".to_string(), Value::String(self.model.clone()));
315        if let Some(ref stop_reason) = self.stop_reason {
316            data.insert("stopReason".to_string(), Value::String(stop_reason.clone()));
317        }
318        data
319    }
320}
321
322impl HasMeta for CreateMessageResult {
323    fn meta(&self) -> Option<HashMap<String, Value>> {
324        self.meta.clone()
325    }
326}
327
328impl RpcResult for CreateMessageResult {}
329
330impl crate::traits::CreateMessageResult for CreateMessageResult {
331    fn role(&self) -> &Role {
332        &self.message.role
333    }
334    
335    fn content(&self) -> &ContentBlock {
336        &self.message.content
337    }
338    
339    fn model(&self) -> &String {
340        &self.model
341    }
342    
343    fn stop_reason(&self) -> Option<&String> {
344        self.stop_reason.as_ref()
345    }
346}
347
348// ===========================================
349// === Fine-Grained Sampling Traits ===
350// ===========================================
351
352/// Trait for sampling message metadata (role, content from MCP spec)
353pub trait HasSamplingMessageMetadata {
354    /// Role of the message (from spec)
355    fn role(&self) -> &Role;
356    
357    /// Content of the message (from spec)
358    fn content(&self) -> &ContentBlock;
359}
360
361/// Trait for sampling configuration (from CreateMessageRequest spec)
362pub trait HasSamplingConfig {
363    /// Maximum tokens to generate (required field from spec)
364    fn max_tokens(&self) -> u32;
365    
366    /// Temperature for sampling (optional from spec)
367    fn temperature(&self) -> Option<f64> {
368        None
369    }
370    
371    /// Stop sequences (optional from spec)
372    fn stop_sequences(&self) -> Option<&Vec<String>> {
373        None
374    }
375}
376
377/// Trait for sampling context (from CreateMessageRequest spec)
378pub trait HasSamplingContext {
379    /// Messages for context (required from spec)
380    fn messages(&self) -> &[SamplingMessage];
381    
382    /// System prompt (optional from spec)
383    fn system_prompt(&self) -> Option<&str> {
384        None
385    }
386    
387    /// Include context setting (optional from spec)
388    fn include_context(&self) -> Option<&str> {
389        None
390    }
391}
392
393/// Trait for model preferences (from CreateMessageRequest spec)
394pub trait HasModelPreferences {
395    /// Model preferences (optional from spec)
396    fn model_preferences(&self) -> Option<&ModelPreferences> {
397        None
398    }
399    
400    /// Metadata (optional from spec)
401    fn metadata(&self) -> Option<&Value> {
402        None
403    }
404}
405
406/// Composed sampling definition trait (automatically implemented via blanket impl)
407pub trait SamplingDefinition: 
408    HasSamplingConfig + 
409    HasSamplingContext + 
410    HasModelPreferences 
411{
412    /// Convert to CreateMessageParams
413    fn to_create_params(&self) -> CreateMessageParams {
414        CreateMessageParams {
415            messages: self.messages().to_vec(),
416            model_preferences: self.model_preferences().cloned(),
417            system_prompt: self.system_prompt().map(|s| s.to_string()),
418            include_context: self.include_context().map(|s| s.to_string()),
419            temperature: self.temperature(),
420            max_tokens: self.max_tokens(),
421            stop_sequences: self.stop_sequences().cloned(),
422            metadata: self.metadata().cloned(),
423            meta: None,
424        }
425    }
426}
427
428// Blanket implementation: any type implementing the fine-grained traits automatically gets SamplingDefinition
429impl<T> SamplingDefinition for T 
430where 
431    T: HasSamplingConfig + HasSamplingContext + HasModelPreferences 
432{}
433
434// ================== TRAIT IMPLEMENTATIONS FOR CONCRETE TYPES ==================
435
436impl HasSamplingMessageMetadata for SamplingMessage {
437    fn role(&self) -> &Role { &self.role }
438    fn content(&self) -> &ContentBlock { &self.content }
439}
440
441impl HasSamplingConfig for CreateMessageParams {
442    fn max_tokens(&self) -> u32 { self.max_tokens }
443    fn temperature(&self) -> Option<f64> { self.temperature }
444    fn stop_sequences(&self) -> Option<&Vec<String>> { self.stop_sequences.as_ref() }
445}
446
447impl HasSamplingContext for CreateMessageParams {
448    fn messages(&self) -> &[SamplingMessage] { &self.messages }
449    fn system_prompt(&self) -> Option<&str> { self.system_prompt.as_deref() }
450    fn include_context(&self) -> Option<&str> { self.include_context.as_deref() }
451}
452
453impl HasModelPreferences for CreateMessageParams {
454    fn model_preferences(&self) -> Option<&ModelPreferences> { self.model_preferences.as_ref() }
455    fn metadata(&self) -> Option<&Value> { self.metadata.as_ref() }
456}
457
458// CreateMessageParams automatically implements SamplingDefinition via trait composition!
459
460// ================== CONVENIENCE CONSTRUCTORS ==================
461
462impl ModelPreferences {
463    pub fn new() -> Self {
464        Self {
465            hints: None,
466            cost_priority: None,
467            speed_priority: None,
468            intelligence_priority: None,
469        }
470    }
471
472    pub fn with_hints(mut self, hints: Vec<ModelHint>) -> Self {
473        self.hints = Some(hints);
474        self
475    }
476
477    pub fn with_cost_priority(mut self, priority: f64) -> Self {
478        self.cost_priority = Some(priority);
479        self
480    }
481
482    pub fn with_speed_priority(mut self, priority: f64) -> Self {
483        self.speed_priority = Some(priority);
484        self
485    }
486
487    pub fn with_intelligence_priority(mut self, priority: f64) -> Self {
488        self.intelligence_priority = Some(priority);
489        self
490    }
491}
492
493impl Default for ModelPreferences {
494    fn default() -> Self {
495        Self::new()
496    }
497}
498
499impl SamplingMessage {
500    pub fn new(role: Role, content: ContentBlock) -> Self {
501        Self { role, content }
502    }
503
504    pub fn user_text(text: impl Into<String>) -> Self {
505        Self::new(Role::User, ContentBlock::Text { text: text.into() })
506    }
507
508    pub fn assistant_text(text: impl Into<String>) -> Self {
509        Self::new(Role::Assistant, ContentBlock::Text { text: text.into() })
510    }
511
512    pub fn system_text(text: impl Into<String>) -> Self {
513        Self::new(Role::System, ContentBlock::Text { text: text.into() })
514    }
515}