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 crate::content::ContentBlock;
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
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(
313            "role".to_string(),
314            serde_json::to_value(&self.message.role).unwrap_or(Value::String("user".to_string())),
315        );
316        data.insert(
317            "content".to_string(),
318            serde_json::to_value(&self.message.content).unwrap_or(Value::Null),
319        );
320        data.insert("model".to_string(), Value::String(self.model.clone()));
321        if let Some(ref stop_reason) = self.stop_reason {
322            data.insert("stopReason".to_string(), Value::String(stop_reason.clone()));
323        }
324        data
325    }
326}
327
328impl HasMeta for CreateMessageResult {
329    fn meta(&self) -> Option<HashMap<String, Value>> {
330        self.meta.clone()
331    }
332}
333
334impl RpcResult for CreateMessageResult {}
335
336impl crate::traits::CreateMessageResult for CreateMessageResult {
337    fn role(&self) -> &Role {
338        &self.message.role
339    }
340
341    fn content(&self) -> &ContentBlock {
342        &self.message.content
343    }
344
345    fn model(&self) -> &String {
346        &self.model
347    }
348
349    fn stop_reason(&self) -> Option<&String> {
350        self.stop_reason.as_ref()
351    }
352}
353
354// ===========================================
355// === Fine-Grained Sampling Traits ===
356// ===========================================
357
358// ================== CONVENIENCE CONSTRUCTORS ==================
359
360impl ModelPreferences {
361    pub fn new() -> Self {
362        Self {
363            hints: None,
364            cost_priority: None,
365            speed_priority: None,
366            intelligence_priority: None,
367        }
368    }
369
370    pub fn with_hints(mut self, hints: Vec<ModelHint>) -> Self {
371        self.hints = Some(hints);
372        self
373    }
374
375    pub fn with_cost_priority(mut self, priority: f64) -> Self {
376        self.cost_priority = Some(priority);
377        self
378    }
379
380    pub fn with_speed_priority(mut self, priority: f64) -> Self {
381        self.speed_priority = Some(priority);
382        self
383    }
384
385    pub fn with_intelligence_priority(mut self, priority: f64) -> Self {
386        self.intelligence_priority = Some(priority);
387        self
388    }
389}
390
391impl Default for ModelPreferences {
392    fn default() -> Self {
393        Self::new()
394    }
395}
396
397impl SamplingMessage {
398    pub fn new(role: Role, content: ContentBlock) -> Self {
399        Self { role, content }
400    }
401
402    pub fn user_text(text: impl Into<String>) -> Self {
403        Self::new(Role::User, ContentBlock::text(text))
404    }
405
406    pub fn assistant_text(text: impl Into<String>) -> Self {
407        Self::new(Role::Assistant, ContentBlock::text(text))
408    }
409
410    pub fn system_text(text: impl Into<String>) -> Self {
411        Self::new(Role::System, ContentBlock::text(text))
412    }
413}