1use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9
10use super::{content::Content, core::Role};
11
12use super::tools::Tool;
13
14#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
16#[serde(rename_all = "camelCase")]
17pub enum IncludeContext {
18 None,
20 ThisServer,
22 AllServers,
24}
25
26#[derive(Debug, Clone, Serialize, Deserialize)]
28pub struct SamplingMessage {
29 pub role: Role,
31 pub content: Content,
33 #[serde(skip_serializing_if = "Option::is_none")]
35 pub metadata: Option<HashMap<String, serde_json::Value>>,
36}
37
38#[derive(Debug, Clone, Serialize, Deserialize)]
40pub struct CreateMessageRequest {
41 pub messages: Vec<SamplingMessage>,
43 #[serde(skip_serializing_if = "Option::is_none")]
45 pub model_preferences: Option<ModelPreferences>,
46 #[serde(rename = "systemPrompt", skip_serializing_if = "Option::is_none")]
48 pub system_prompt: Option<String>,
49 #[serde(rename = "includeContext", skip_serializing_if = "Option::is_none")]
51 pub include_context: Option<IncludeContext>,
52 #[serde(skip_serializing_if = "Option::is_none")]
54 pub temperature: Option<f64>,
55 #[serde(rename = "maxTokens")]
57 pub max_tokens: u32,
58 #[serde(rename = "stopSequences", skip_serializing_if = "Option::is_none")]
60 pub stop_sequences: Option<Vec<String>>,
61 #[serde(skip_serializing_if = "Option::is_none")]
65 pub tools: Option<Vec<Tool>>,
66 #[serde(rename = "toolChoice", skip_serializing_if = "Option::is_none")]
71 pub tool_choice: Option<ToolChoice>,
72 #[serde(skip_serializing_if = "Option::is_none")]
78 pub task: Option<crate::types::tasks::TaskMetadata>,
79 #[serde(skip_serializing_if = "Option::is_none")]
81 pub _meta: Option<serde_json::Value>,
82}
83
84#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
86pub struct ModelHint {
87 #[serde(skip_serializing_if = "Option::is_none")]
90 pub name: Option<String>,
91}
92
93impl ModelHint {
94 pub fn new(name: impl Into<String>) -> Self {
96 Self {
97 name: Some(name.into()),
98 }
99 }
100}
101
102#[derive(Debug, Clone, Serialize, Deserialize)]
107pub struct ModelPreferences {
108 #[serde(skip_serializing_if = "Option::is_none")]
110 pub hints: Option<Vec<ModelHint>>,
111
112 #[serde(rename = "costPriority", skip_serializing_if = "Option::is_none")]
114 pub cost_priority: Option<f64>,
115
116 #[serde(rename = "speedPriority", skip_serializing_if = "Option::is_none")]
118 pub speed_priority: Option<f64>,
119
120 #[serde(
122 rename = "intelligencePriority",
123 skip_serializing_if = "Option::is_none"
124 )]
125 pub intelligence_priority: Option<f64>,
126}
127
128#[derive(Debug, Clone, Serialize, Deserialize)]
130pub struct CreateMessageResult {
131 pub role: super::core::Role,
133 pub content: Content,
135 pub model: String,
137 #[serde(rename = "stopReason", skip_serializing_if = "Option::is_none")]
141 pub stop_reason: Option<StopReason>,
142 #[serde(skip_serializing_if = "Option::is_none")]
144 pub _meta: Option<serde_json::Value>,
145}
146
147#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
151#[serde(rename_all = "camelCase")]
152pub enum StopReason {
153 EndTurn,
155 MaxTokens,
157 StopSequence,
159 ContentFilter,
161 ToolUse,
163}
164
165#[derive(Debug, Clone, Serialize, Deserialize)]
167pub struct UsageStats {
168 #[serde(rename = "inputTokens", skip_serializing_if = "Option::is_none")]
170 pub input_tokens: Option<u32>,
171 #[serde(rename = "outputTokens", skip_serializing_if = "Option::is_none")]
173 pub output_tokens: Option<u32>,
174 #[serde(rename = "totalTokens", skip_serializing_if = "Option::is_none")]
176 pub total_tokens: Option<u32>,
177}
178
179#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash, Default)]
181#[serde(rename_all = "lowercase")]
182pub enum ToolChoiceMode {
183 #[default]
185 Auto,
186 Required,
188 None,
190}
191
192#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
194pub struct ToolChoice {
195 #[serde(skip_serializing_if = "Option::is_none")]
200 pub mode: Option<ToolChoiceMode>,
201}
202
203impl ToolChoice {
204 pub fn auto() -> Self {
206 Self {
207 mode: Some(ToolChoiceMode::Auto),
208 }
209 }
210
211 pub fn required() -> Self {
213 Self {
214 mode: Some(ToolChoiceMode::Required),
215 }
216 }
217
218 pub fn none() -> Self {
220 Self {
221 mode: Some(ToolChoiceMode::None),
222 }
223 }
224}
225
226impl Default for ToolChoice {
227 fn default() -> Self {
228 Self::auto()
229 }
230}
231
232#[cfg(test)]
233mod tests {
234 use super::*;
235
236 #[test]
237 fn test_tool_choice_mode_serialization() {
238 assert_eq!(
239 serde_json::to_string(&ToolChoiceMode::Auto).unwrap(),
240 "\"auto\""
241 );
242 assert_eq!(
243 serde_json::to_string(&ToolChoiceMode::Required).unwrap(),
244 "\"required\""
245 );
246 assert_eq!(
247 serde_json::to_string(&ToolChoiceMode::None).unwrap(),
248 "\"none\""
249 );
250 }
251
252 #[test]
253 fn test_tool_choice_constructors() {
254 let auto = ToolChoice::auto();
255 assert_eq!(auto.mode, Some(ToolChoiceMode::Auto));
256
257 let required = ToolChoice::required();
258 assert_eq!(required.mode, Some(ToolChoiceMode::Required));
259
260 let none = ToolChoice::none();
261 assert_eq!(none.mode, Some(ToolChoiceMode::None));
262 }
263
264 #[test]
265 fn test_tool_choice_default() {
266 let default = ToolChoice::default();
267 assert_eq!(default.mode, Some(ToolChoiceMode::Auto));
268 }
269
270 #[test]
271 fn test_tool_choice_serialization() {
272 let choice = ToolChoice::required();
273 let json = serde_json::to_string(&choice).unwrap();
274 assert!(json.contains("\"mode\":\"required\""));
275 }
276}