claude_agent/client/messages/
request.rs

1//! Message request and response types.
2
3use serde::{Deserialize, Serialize};
4
5use super::config::{EffortLevel, OutputConfig, OutputFormat, ThinkingConfig, ToolChoice};
6use super::context::ContextManagement;
7use super::types::{ApiTool, RequestMetadata};
8use crate::types::{Message, SystemPrompt, ToolDefinition, WebFetchTool, WebSearchTool};
9
10#[derive(Debug, Clone, Serialize)]
11pub struct CreateMessageRequest {
12    pub model: String,
13    pub max_tokens: u32,
14    pub messages: Vec<Message>,
15    #[serde(skip_serializing_if = "Option::is_none")]
16    pub system: Option<SystemPrompt>,
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub tools: Option<Vec<ApiTool>>,
19    #[serde(skip_serializing_if = "Option::is_none")]
20    pub tool_choice: Option<ToolChoice>,
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub stream: Option<bool>,
23    #[serde(skip_serializing_if = "Option::is_none")]
24    pub stop_sequences: Option<Vec<String>>,
25    #[serde(skip_serializing_if = "Option::is_none")]
26    pub temperature: Option<f32>,
27    #[serde(skip_serializing_if = "Option::is_none")]
28    pub top_p: Option<f32>,
29    #[serde(skip_serializing_if = "Option::is_none")]
30    pub top_k: Option<u32>,
31    #[serde(skip_serializing_if = "Option::is_none")]
32    pub metadata: Option<RequestMetadata>,
33    #[serde(skip_serializing_if = "Option::is_none")]
34    pub thinking: Option<ThinkingConfig>,
35    #[serde(skip_serializing_if = "Option::is_none")]
36    pub output_format: Option<OutputFormat>,
37    #[serde(skip_serializing_if = "Option::is_none")]
38    pub context_management: Option<ContextManagement>,
39    #[serde(skip_serializing_if = "Option::is_none")]
40    pub output_config: Option<OutputConfig>,
41}
42
43impl CreateMessageRequest {
44    pub fn new(model: impl Into<String>, messages: Vec<Message>) -> Self {
45        Self {
46            model: model.into(),
47            max_tokens: 8192,
48            messages,
49            system: None,
50            tools: None,
51            tool_choice: None,
52            stream: None,
53            stop_sequences: None,
54            temperature: None,
55            top_p: None,
56            top_k: None,
57            metadata: None,
58            thinking: None,
59            output_format: None,
60            context_management: None,
61            output_config: None,
62        }
63    }
64
65    pub fn with_metadata(mut self, metadata: RequestMetadata) -> Self {
66        self.metadata = Some(metadata);
67        self
68    }
69
70    pub fn with_system(mut self, system: impl Into<SystemPrompt>) -> Self {
71        self.system = Some(system.into());
72        self
73    }
74
75    pub fn with_tools(mut self, tools: Vec<ToolDefinition>) -> Self {
76        let api_tools: Vec<ApiTool> = tools.into_iter().map(ApiTool::Custom).collect();
77        self.tools = Some(api_tools);
78        self
79    }
80
81    pub fn with_web_search(mut self, config: WebSearchTool) -> Self {
82        let mut tools = self.tools.unwrap_or_default();
83        tools.push(ApiTool::WebSearch(config));
84        self.tools = Some(tools);
85        self
86    }
87
88    pub fn with_web_fetch(mut self, config: WebFetchTool) -> Self {
89        let mut tools = self.tools.unwrap_or_default();
90        tools.push(ApiTool::WebFetch(config));
91        self.tools = Some(tools);
92        self
93    }
94
95    pub fn with_api_tools(mut self, tools: Vec<ApiTool>) -> Self {
96        self.tools = Some(tools);
97        self
98    }
99
100    pub fn with_tool_choice(mut self, choice: ToolChoice) -> Self {
101        self.tool_choice = Some(choice);
102        self
103    }
104
105    pub fn with_tool_choice_auto(mut self) -> Self {
106        self.tool_choice = Some(ToolChoice::Auto);
107        self
108    }
109
110    pub fn with_tool_choice_any(mut self) -> Self {
111        self.tool_choice = Some(ToolChoice::Any);
112        self
113    }
114
115    pub fn with_tool_choice_none(mut self) -> Self {
116        self.tool_choice = Some(ToolChoice::None);
117        self
118    }
119
120    pub fn with_required_tool(mut self, name: impl Into<String>) -> Self {
121        self.tool_choice = Some(ToolChoice::tool(name));
122        self
123    }
124
125    pub fn with_stream(mut self) -> Self {
126        self.stream = Some(true);
127        self
128    }
129
130    pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
131        self.max_tokens = max_tokens;
132        self
133    }
134
135    pub fn with_model(mut self, model: impl Into<String>) -> Self {
136        self.model = model.into();
137        self
138    }
139
140    pub fn with_temperature(mut self, temperature: f32) -> Self {
141        self.temperature = Some(temperature);
142        self
143    }
144
145    pub fn with_top_p(mut self, top_p: f32) -> Self {
146        self.top_p = Some(top_p);
147        self
148    }
149
150    pub fn with_top_k(mut self, top_k: u32) -> Self {
151        self.top_k = Some(top_k);
152        self
153    }
154
155    pub fn with_stop_sequences(mut self, sequences: Vec<String>) -> Self {
156        self.stop_sequences = Some(sequences);
157        self
158    }
159
160    pub fn with_thinking(mut self, config: ThinkingConfig) -> Self {
161        self.thinking = Some(config);
162        self
163    }
164
165    pub fn with_extended_thinking(mut self, budget_tokens: u32) -> Self {
166        self.thinking = Some(ThinkingConfig::enabled(budget_tokens));
167        self
168    }
169
170    pub fn with_output_format(mut self, format: OutputFormat) -> Self {
171        self.output_format = Some(format);
172        self
173    }
174
175    pub fn with_json_schema(mut self, schema: serde_json::Value) -> Self {
176        self.output_format = Some(OutputFormat::json_schema(schema));
177        self
178    }
179
180    pub fn with_context_management(mut self, management: ContextManagement) -> Self {
181        self.context_management = Some(management);
182        self
183    }
184
185    pub fn with_effort(mut self, level: EffortLevel) -> Self {
186        self.output_config = Some(OutputConfig::with_effort(level));
187        self
188    }
189
190    pub fn with_output_config(mut self, config: OutputConfig) -> Self {
191        self.output_config = Some(config);
192        self
193    }
194}
195
196impl From<String> for SystemPrompt {
197    fn from(s: String) -> Self {
198        SystemPrompt::Text(s)
199    }
200}
201
202impl From<&str> for SystemPrompt {
203    fn from(s: &str) -> Self {
204        SystemPrompt::Text(s.to_string())
205    }
206}
207
208#[derive(Debug, Clone, Serialize)]
209pub struct CountTokensRequest {
210    pub model: String,
211    pub messages: Vec<Message>,
212    #[serde(skip_serializing_if = "Option::is_none")]
213    pub system: Option<SystemPrompt>,
214    #[serde(skip_serializing_if = "Option::is_none")]
215    pub tools: Option<Vec<ApiTool>>,
216    #[serde(skip_serializing_if = "Option::is_none")]
217    pub thinking: Option<ThinkingConfig>,
218}
219
220impl CountTokensRequest {
221    pub fn new(model: impl Into<String>, messages: Vec<Message>) -> Self {
222        Self {
223            model: model.into(),
224            messages,
225            system: None,
226            tools: None,
227            thinking: None,
228        }
229    }
230
231    pub fn with_system(mut self, system: impl Into<SystemPrompt>) -> Self {
232        self.system = Some(system.into());
233        self
234    }
235
236    pub fn with_tools(mut self, tools: Vec<ToolDefinition>) -> Self {
237        self.tools = Some(tools.into_iter().map(ApiTool::Custom).collect());
238        self
239    }
240
241    pub fn with_api_tools(mut self, tools: Vec<ApiTool>) -> Self {
242        self.tools = Some(tools);
243        self
244    }
245
246    pub fn with_thinking(mut self, config: ThinkingConfig) -> Self {
247        self.thinking = Some(config);
248        self
249    }
250
251    pub fn from_message_request(request: &CreateMessageRequest) -> Self {
252        Self {
253            model: request.model.clone(),
254            messages: request.messages.clone(),
255            system: request.system.clone(),
256            tools: request.tools.clone(),
257            thinking: request.thinking.clone(),
258        }
259    }
260}
261
262#[derive(Debug, Clone, Deserialize)]
263pub struct CountTokensResponse {
264    pub input_tokens: u32,
265    #[serde(default, skip_serializing_if = "Option::is_none")]
266    pub context_management: Option<CountTokensContextManagement>,
267}
268
269#[derive(Debug, Clone, Default, Deserialize)]
270pub struct CountTokensContextManagement {
271    #[serde(default)]
272    pub original_input_tokens: Option<u32>,
273}
274
275#[cfg(test)]
276mod tests {
277    use super::*;
278
279    #[test]
280    fn test_create_request() {
281        let request = CreateMessageRequest::new("claude-sonnet-4-5", vec![Message::user("Hello")])
282            .with_max_tokens(1000)
283            .with_temperature(0.7);
284
285        assert_eq!(request.model, "claude-sonnet-4-5");
286        assert_eq!(request.max_tokens, 1000);
287        assert_eq!(request.temperature, Some(0.7));
288    }
289
290    #[test]
291    fn test_count_tokens_request() {
292        let request = CountTokensRequest::new("claude-sonnet-4-5", vec![Message::user("Hello")])
293            .with_system("You are a helpful assistant");
294
295        assert_eq!(request.model, "claude-sonnet-4-5");
296        assert!(request.system.is_some());
297    }
298
299    #[test]
300    fn test_count_tokens_from_message_request() {
301        let msg_request = CreateMessageRequest::new("claude-sonnet-4-5", vec![Message::user("Hi")])
302            .with_system("System prompt")
303            .with_extended_thinking(10000);
304
305        let count_request = CountTokensRequest::from_message_request(&msg_request);
306
307        assert_eq!(count_request.model, msg_request.model);
308        assert_eq!(count_request.messages.len(), msg_request.messages.len());
309        assert!(count_request.system.is_some());
310        assert!(count_request.thinking.is_some());
311    }
312
313    #[test]
314    fn test_request_with_effort() {
315        let request = CreateMessageRequest::new("claude-opus-4-5", vec![Message::user("Hi")])
316            .with_effort(EffortLevel::Medium);
317        assert!(request.output_config.is_some());
318        assert_eq!(
319            request.output_config.unwrap().effort,
320            Some(EffortLevel::Medium)
321        );
322    }
323
324    #[test]
325    fn test_request_with_context_management() {
326        let mgmt = ContextManagement::new().with_edit(ContextManagement::clear_thinking(2));
327        let request = CreateMessageRequest::new("claude-sonnet-4-5", vec![Message::user("Hi")])
328            .with_context_management(mgmt);
329        assert!(request.context_management.is_some());
330    }
331
332    #[test]
333    fn test_request_with_tool_choice() {
334        let request = CreateMessageRequest::new("claude-sonnet-4-5", vec![Message::user("Hi")])
335            .with_tool_choice_any();
336        assert_eq!(request.tool_choice, Some(ToolChoice::Any));
337
338        let request = CreateMessageRequest::new("claude-sonnet-4-5", vec![Message::user("Hi")])
339            .with_required_tool("Grep");
340        assert_eq!(request.tool_choice, Some(ToolChoice::tool("Grep")));
341    }
342}