claude_agent/types/
message.rs

1//! Message types for the Claude API.
2
3use serde::{Deserialize, Serialize};
4
5use super::ContentBlock;
6use super::document::DocumentBlock;
7use super::search::SearchResultBlock;
8
9/// Role of a message participant
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
11#[serde(rename_all = "lowercase")]
12pub enum Role {
13    /// User message
14    User,
15    /// Assistant (Claude) message
16    Assistant,
17}
18
19/// A message in a conversation
20#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct Message {
22    /// Role of the message sender
23    pub role: Role,
24    /// Content of the message
25    pub content: Vec<ContentBlock>,
26}
27
28impl Message {
29    pub fn user(text: impl Into<String>) -> Self {
30        Self {
31            role: Role::User,
32            content: vec![ContentBlock::text(text)],
33        }
34    }
35
36    pub fn assistant(text: impl Into<String>) -> Self {
37        Self {
38            role: Role::Assistant,
39            content: vec![ContentBlock::text(text)],
40        }
41    }
42
43    pub fn tool_results(results: Vec<super::ToolResultBlock>) -> Self {
44        Self {
45            role: Role::User,
46            content: results.into_iter().map(ContentBlock::ToolResult).collect(),
47        }
48    }
49
50    pub fn user_with_content(content: Vec<ContentBlock>) -> Self {
51        Self {
52            role: Role::User,
53            content,
54        }
55    }
56
57    pub fn user_with_document(text: impl Into<String>, doc: DocumentBlock) -> Self {
58        Self {
59            role: Role::User,
60            content: vec![ContentBlock::Document(doc), ContentBlock::text(text)],
61        }
62    }
63
64    pub fn user_with_documents(text: impl Into<String>, docs: Vec<DocumentBlock>) -> Self {
65        let mut content: Vec<ContentBlock> = docs.into_iter().map(ContentBlock::Document).collect();
66        content.push(ContentBlock::text(text));
67        Self {
68            role: Role::User,
69            content,
70        }
71    }
72
73    pub fn user_with_search_results(
74        text: impl Into<String>,
75        results: Vec<SearchResultBlock>,
76    ) -> Self {
77        let mut content: Vec<ContentBlock> = results
78            .into_iter()
79            .map(ContentBlock::SearchResult)
80            .collect();
81        content.push(ContentBlock::text(text));
82        Self {
83            role: Role::User,
84            content,
85        }
86    }
87
88    pub fn text(&self) -> String {
89        self.content
90            .iter()
91            .filter_map(|block| block.as_text())
92            .collect::<Vec<_>>()
93            .join("")
94    }
95
96    pub fn has_tool_use(&self) -> bool {
97        self.content
98            .iter()
99            .any(|block| matches!(block, ContentBlock::ToolUse { .. }))
100    }
101
102    pub fn tool_uses(&self) -> Vec<&super::ToolUseBlock> {
103        self.content
104            .iter()
105            .filter_map(|block| match block {
106                ContentBlock::ToolUse(tool_use) => Some(tool_use),
107                _ => None,
108            })
109            .collect()
110    }
111
112    pub fn documents(&self) -> Vec<&DocumentBlock> {
113        self.content
114            .iter()
115            .filter_map(|block| block.as_document())
116            .collect()
117    }
118
119    pub fn search_results(&self) -> Vec<&SearchResultBlock> {
120        self.content
121            .iter()
122            .filter_map(|block| block.as_search_result())
123            .collect()
124    }
125
126    pub fn with_cache_on_last_block(mut self) -> Self {
127        if let Some(last) = self.content.pop() {
128            self.content
129                .push(last.with_cache_control(CacheControl::ephemeral()));
130        }
131        self
132    }
133
134    pub fn set_cache_on_last_block(&mut self, cache: CacheControl) {
135        if let Some(last) = self.content.last_mut() {
136            last.set_cache_control(Some(cache));
137        }
138    }
139
140    pub fn clear_cache_control(&mut self) {
141        for block in &mut self.content {
142            block.set_cache_control(None);
143        }
144    }
145
146    pub fn has_cache_control(&self) -> bool {
147        self.content.iter().any(|b| b.is_cached())
148    }
149}
150
151/// System prompt configuration
152#[derive(Debug, Clone, Serialize, Deserialize)]
153#[serde(untagged)]
154pub enum SystemPrompt {
155    /// Simple text system prompt
156    Text(String),
157    /// Structured system prompt with cache control
158    Blocks(Vec<SystemBlock>),
159}
160
161impl Default for SystemPrompt {
162    fn default() -> Self {
163        Self::Text(String::new())
164    }
165}
166
167impl SystemPrompt {
168    pub fn is_empty(&self) -> bool {
169        match self {
170            Self::Text(s) => s.is_empty(),
171            Self::Blocks(b) => b.is_empty(),
172        }
173    }
174
175    pub fn as_text(&self) -> String {
176        match self {
177            Self::Text(s) => s.clone(),
178            Self::Blocks(b) => b
179                .iter()
180                .map(|block| block.text.as_str())
181                .collect::<Vec<_>>()
182                .join("\n\n"),
183        }
184    }
185}
186
187impl std::fmt::Display for SystemPrompt {
188    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
189        write!(f, "{}", self.as_text())
190    }
191}
192
193/// A block in a structured system prompt.
194#[derive(Debug, Clone, Serialize, Deserialize)]
195pub struct SystemBlock {
196    /// Type of the block (always "text" for now).
197    #[serde(rename = "type")]
198    pub block_type: String,
199    /// Text content.
200    pub text: String,
201    /// Optional cache control.
202    #[serde(skip_serializing_if = "Option::is_none")]
203    pub cache_control: Option<CacheControl>,
204}
205
206impl SystemBlock {
207    /// Create a new system block with caching enabled.
208    pub fn cached(text: impl Into<String>) -> Self {
209        Self {
210            block_type: "text".to_string(),
211            text: text.into(),
212            cache_control: Some(CacheControl::ephemeral()),
213        }
214    }
215
216    /// Create a new system block without caching.
217    pub fn uncached(text: impl Into<String>) -> Self {
218        Self {
219            block_type: "text".to_string(),
220            text: text.into(),
221            cache_control: None,
222        }
223    }
224}
225
226/// Cache control for prompt caching.
227#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
228pub struct CacheControl {
229    #[serde(rename = "type")]
230    pub cache_type: CacheType,
231    #[serde(skip_serializing_if = "Option::is_none")]
232    pub ttl: Option<CacheTtl>,
233}
234
235#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
236#[serde(rename_all = "snake_case")]
237pub enum CacheType {
238    Ephemeral,
239}
240
241#[derive(Debug, Clone, Copy, PartialEq, Eq)]
242pub enum CacheTtl {
243    FiveMinutes,
244    OneHour,
245}
246
247impl Serialize for CacheTtl {
248    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
249    where
250        S: serde::Serializer,
251    {
252        match self {
253            CacheTtl::FiveMinutes => serializer.serialize_str("5m"),
254            CacheTtl::OneHour => serializer.serialize_str("1h"),
255        }
256    }
257}
258
259impl<'de> Deserialize<'de> for CacheTtl {
260    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
261    where
262        D: serde::Deserializer<'de>,
263    {
264        let s = String::deserialize(deserializer)?;
265        match s.as_str() {
266            "5m" => Ok(CacheTtl::FiveMinutes),
267            "1h" => Ok(CacheTtl::OneHour),
268            _ => Err(serde::de::Error::custom(format!("unknown TTL: {}", s))),
269        }
270    }
271}
272
273impl CacheControl {
274    pub fn ephemeral() -> Self {
275        Self {
276            cache_type: CacheType::Ephemeral,
277            ttl: None,
278        }
279    }
280
281    pub fn ephemeral_5m() -> Self {
282        Self {
283            cache_type: CacheType::Ephemeral,
284            ttl: Some(CacheTtl::FiveMinutes),
285        }
286    }
287
288    pub fn ephemeral_1h() -> Self {
289        Self {
290            cache_type: CacheType::Ephemeral,
291            ttl: Some(CacheTtl::OneHour),
292        }
293    }
294
295    pub fn with_ttl(mut self, ttl: CacheTtl) -> Self {
296        self.ttl = Some(ttl);
297        self
298    }
299}
300
301impl SystemPrompt {
302    /// Create a simple text system prompt
303    pub fn text(prompt: impl Into<String>) -> Self {
304        Self::Text(prompt.into())
305    }
306
307    /// Create a system prompt with caching enabled
308    pub fn cached(prompt: impl Into<String>) -> Self {
309        Self::Blocks(vec![SystemBlock {
310            block_type: "text".to_string(),
311            text: prompt.into(),
312            cache_control: Some(CacheControl {
313                cache_type: CacheType::Ephemeral,
314                ttl: None,
315            }),
316        }])
317    }
318}
319
320#[cfg(test)]
321mod tests {
322    use super::*;
323
324    #[test]
325    fn test_user_message() {
326        let msg = Message::user("Hello");
327        assert_eq!(msg.role, Role::User);
328        assert_eq!(msg.text(), "Hello");
329    }
330
331    #[test]
332    fn test_assistant_message() {
333        let msg = Message::assistant("Hi there!");
334        assert_eq!(msg.role, Role::Assistant);
335        assert_eq!(msg.text(), "Hi there!");
336    }
337}
claude_agent/types/message.rs

claude_agent/types/
message.rs
