claude_agent/types/
document.rs

1//! Document content block types for citations.
2
3use serde::{Deserialize, Serialize};
4
5use super::citations::CitationsConfig;
6use super::message::CacheControl;
7
8#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
9#[serde(tag = "type", rename_all = "snake_case")]
10pub enum DocumentSource {
11    Text { media_type: String, data: String },
12    Base64 { media_type: String, data: String },
13    Content { content: Vec<DocumentContentBlock> },
14    File { file_id: String },
15    Url { url: String },
16}
17
18#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
19#[serde(tag = "type", rename_all = "snake_case")]
20pub enum DocumentContentBlock {
21    Text { text: String },
22}
23
24impl DocumentContentBlock {
25    pub fn text(content: impl Into<String>) -> Self {
26        Self::Text {
27            text: content.into(),
28        }
29    }
30}
31
32#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
33pub struct DocumentBlock {
34    pub source: DocumentSource,
35    #[serde(skip_serializing_if = "Option::is_none")]
36    pub title: Option<String>,
37    #[serde(skip_serializing_if = "Option::is_none")]
38    pub context: Option<String>,
39    #[serde(skip_serializing_if = "Option::is_none")]
40    pub citations: Option<CitationsConfig>,
41    #[serde(skip_serializing_if = "Option::is_none")]
42    pub cache_control: Option<CacheControl>,
43}
44
45impl DocumentBlock {
46    pub fn text(content: impl Into<String>) -> Self {
47        Self {
48            source: DocumentSource::Text {
49                media_type: "text/plain".to_string(),
50                data: content.into(),
51            },
52            title: None,
53            context: None,
54            citations: None,
55            cache_control: None,
56        }
57    }
58
59    pub fn html(content: impl Into<String>) -> Self {
60        Self {
61            source: DocumentSource::Text {
62                media_type: "text/html".to_string(),
63                data: content.into(),
64            },
65            title: None,
66            context: None,
67            citations: None,
68            cache_control: None,
69        }
70    }
71
72    pub fn markdown(content: impl Into<String>) -> Self {
73        Self {
74            source: DocumentSource::Text {
75                media_type: "text/markdown".to_string(),
76                data: content.into(),
77            },
78            title: None,
79            context: None,
80            citations: None,
81            cache_control: None,
82        }
83    }
84
85    pub fn pdf_base64(data: impl Into<String>) -> Self {
86        Self {
87            source: DocumentSource::Base64 {
88                media_type: "application/pdf".to_string(),
89                data: data.into(),
90            },
91            title: None,
92            context: None,
93            citations: None,
94            cache_control: None,
95        }
96    }
97
98    pub fn from_url(url: impl Into<String>) -> Self {
99        Self {
100            source: DocumentSource::Url { url: url.into() },
101            title: None,
102            context: None,
103            citations: None,
104            cache_control: None,
105        }
106    }
107
108    pub fn from_file(file_id: impl Into<String>) -> Self {
109        Self {
110            source: DocumentSource::File {
111                file_id: file_id.into(),
112            },
113            title: None,
114            context: None,
115            citations: None,
116            cache_control: None,
117        }
118    }
119
120    pub fn structured(blocks: Vec<DocumentContentBlock>) -> Self {
121        Self {
122            source: DocumentSource::Content { content: blocks },
123            title: None,
124            context: None,
125            citations: None,
126            cache_control: None,
127        }
128    }
129
130    pub fn with_title(mut self, title: impl Into<String>) -> Self {
131        self.title = Some(title.into());
132        self
133    }
134
135    pub fn with_context(mut self, context: impl Into<String>) -> Self {
136        self.context = Some(context.into());
137        self
138    }
139
140    pub fn with_citations(mut self, enabled: bool) -> Self {
141        self.citations = Some(if enabled {
142            CitationsConfig::enabled()
143        } else {
144            CitationsConfig::disabled()
145        });
146        self
147    }
148
149    pub fn without_citations(mut self) -> Self {
150        self.citations = Some(CitationsConfig::disabled());
151        self
152    }
153
154    pub fn with_cache_control(mut self, cache_control: CacheControl) -> Self {
155        self.cache_control = Some(cache_control);
156        self
157    }
158
159    pub fn cached(mut self) -> Self {
160        self.cache_control = Some(CacheControl::ephemeral());
161        self
162    }
163}
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168
169    #[test]
170    fn test_document_text() {
171        let doc = DocumentBlock::text("Hello world")
172            .with_title("Test Doc")
173            .cached();
174
175        assert!(doc.title.is_some());
176        assert!(doc.cache_control.is_some());
177        assert!(matches!(doc.source, DocumentSource::Text { .. }));
178    }
179
180    #[test]
181    fn test_document_structured() {
182        let doc = DocumentBlock::structured(vec![
183            DocumentContentBlock::text("Block 1"),
184            DocumentContentBlock::text("Block 2"),
185        ]);
186
187        if let DocumentSource::Content { content } = &doc.source {
188            assert_eq!(content.len(), 2);
189        } else {
190            panic!("Expected Content source");
191        }
192    }
193
194    #[test]
195    fn test_document_serialization() {
196        let doc = DocumentBlock::text("test content").with_title("Title");
197        let json = serde_json::to_string(&doc).unwrap();
198
199        assert!(json.contains("text/plain"));
200        assert!(json.contains("test content"));
201    }
202}
claude_agent/types/document.rs

claude_agent/types/
document.rs
