claude_agent/client/messages/
request.rs1use 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}