1use super::role::Role;
5use crate::resource::ResourceContents;
6use chrono::{DateTime, Utc};
7use serde::{Deserialize, Serialize};
8
9#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
10#[serde(rename_all = "camelCase")]
11pub struct Annotations {
12 #[serde(skip_serializing_if = "Option::is_none")]
13 pub audience: Option<Vec<Role>>,
14 #[serde(skip_serializing_if = "Option::is_none")]
15 pub priority: Option<f32>,
16 #[serde(skip_serializing_if = "Option::is_none")]
17 pub timestamp: Option<DateTime<Utc>>,
18}
19
20impl Annotations {
21 pub fn for_resource(priority: f32, timestamp: DateTime<Utc>) -> Self {
24 assert!(
25 (0.0..=1.0).contains(&priority),
26 "Priority {priority} must be between 0.0 and 1.0"
27 );
28 Annotations {
29 priority: Some(priority),
30 timestamp: Some(timestamp),
31 audience: None,
32 }
33 }
34}
35
36#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
37#[serde(rename_all = "camelCase")]
38pub struct TextContent {
39 pub text: String,
40 #[serde(skip_serializing_if = "Option::is_none")]
41 pub annotations: Option<Annotations>,
42}
43
44#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
45#[serde(rename_all = "camelCase")]
46pub struct ImageContent {
47 pub data: String,
48 pub mime_type: String,
49 #[serde(skip_serializing_if = "Option::is_none")]
50 pub annotations: Option<Annotations>,
51}
52
53#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
54#[serde(rename_all = "camelCase")]
55pub struct EmbeddedResource {
56 pub resource: ResourceContents,
57 #[serde(skip_serializing_if = "Option::is_none")]
58 pub annotations: Option<Annotations>,
59}
60
61impl EmbeddedResource {
62 pub fn get_text(&self) -> String {
63 match &self.resource {
64 ResourceContents::TextResourceContents { text, .. } => text.clone(),
65 _ => String::new(),
66 }
67 }
68}
69
70#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
71#[serde(tag = "type", rename_all = "camelCase")]
72pub enum Content {
73 Text(TextContent),
74 Image(ImageContent),
75 Resource(EmbeddedResource),
76}
77
78impl Content {
79 pub fn text<S: Into<String>>(text: S) -> Self {
80 Content::Text(TextContent {
81 text: text.into(),
82 annotations: None,
83 })
84 }
85
86 pub fn image<S: Into<String>, T: Into<String>>(data: S, mime_type: T) -> Self {
87 Content::Image(ImageContent {
88 data: data.into(),
89 mime_type: mime_type.into(),
90 annotations: None,
91 })
92 }
93
94 pub fn resource(resource: ResourceContents) -> Self {
95 Content::Resource(EmbeddedResource {
96 resource,
97 annotations: None,
98 })
99 }
100
101 pub fn embedded_text<S: Into<String>, T: Into<String>>(uri: S, content: T) -> Self {
102 Content::Resource(EmbeddedResource {
103 resource: ResourceContents::TextResourceContents {
104 uri: uri.into(),
105 mime_type: Some("text".to_string()),
106 text: content.into(),
107 },
108 annotations: None,
109 })
110 }
111
112 pub fn as_text(&self) -> Option<&str> {
114 match self {
115 Content::Text(text) => Some(&text.text),
116 _ => None,
117 }
118 }
119
120 pub fn as_image(&self) -> Option<(&str, &str)> {
122 match self {
123 Content::Image(image) => Some((&image.data, &image.mime_type)),
124 _ => None,
125 }
126 }
127
128 pub fn with_audience(mut self, audience: Vec<Role>) -> Self {
130 let annotations = match &mut self {
131 Content::Text(text) => &mut text.annotations,
132 Content::Image(image) => &mut image.annotations,
133 Content::Resource(resource) => &mut resource.annotations,
134 };
135 *annotations = Some(match annotations.take() {
136 Some(mut a) => {
137 a.audience = Some(audience);
138 a
139 }
140 None => Annotations {
141 audience: Some(audience),
142 priority: None,
143 timestamp: None,
144 },
145 });
146 self
147 }
148
149 pub fn with_priority(mut self, priority: f32) -> Self {
153 if !(0.0..=1.0).contains(&priority) {
154 panic!("Priority must be between 0.0 and 1.0");
155 }
156 let annotations = match &mut self {
157 Content::Text(text) => &mut text.annotations,
158 Content::Image(image) => &mut image.annotations,
159 Content::Resource(resource) => &mut resource.annotations,
160 };
161 *annotations = Some(match annotations.take() {
162 Some(mut a) => {
163 a.priority = Some(priority);
164 a
165 }
166 None => Annotations {
167 audience: None,
168 priority: Some(priority),
169 timestamp: None,
170 },
171 });
172 self
173 }
174
175 pub fn audience(&self) -> Option<&Vec<Role>> {
177 match self {
178 Content::Text(text) => text.annotations.as_ref().and_then(|a| a.audience.as_ref()),
179 Content::Image(image) => image.annotations.as_ref().and_then(|a| a.audience.as_ref()),
180 Content::Resource(resource) => resource
181 .annotations
182 .as_ref()
183 .and_then(|a| a.audience.as_ref()),
184 }
185 }
186
187 pub fn priority(&self) -> Option<f32> {
189 match self {
190 Content::Text(text) => text.annotations.as_ref().and_then(|a| a.priority),
191 Content::Image(image) => image.annotations.as_ref().and_then(|a| a.priority),
192 Content::Resource(resource) => resource.annotations.as_ref().and_then(|a| a.priority),
193 }
194 }
195
196 pub fn unannotated(&self) -> Self {
197 match self {
198 Content::Text(text) => Content::text(text.text.clone()),
199 Content::Image(image) => Content::image(image.data.clone(), image.mime_type.clone()),
200 Content::Resource(resource) => Content::resource(resource.resource.clone()),
201 }
202 }
203}
204
205#[cfg(test)]
206mod tests {
207 use super::*;
208
209 #[test]
210 fn test_content_text() {
211 let content = Content::text("hello");
212 assert_eq!(content.as_text(), Some("hello"));
213 assert_eq!(content.as_image(), None);
214 }
215
216 #[test]
217 fn test_content_image() {
218 let content = Content::image("data", "image/png");
219 assert_eq!(content.as_text(), None);
220 assert_eq!(content.as_image(), Some(("data", "image/png")));
221 }
222
223 #[test]
224 fn test_content_annotations_basic() {
225 let content = Content::text("hello")
226 .with_audience(vec![Role::User])
227 .with_priority(0.5);
228 assert_eq!(content.audience(), Some(&vec![Role::User]));
229 assert_eq!(content.priority(), Some(0.5));
230 }
231
232 #[test]
233 fn test_content_annotations_order_independence() {
234 let content1 = Content::text("hello")
235 .with_audience(vec![Role::User])
236 .with_priority(0.5);
237 let content2 = Content::text("hello")
238 .with_priority(0.5)
239 .with_audience(vec![Role::User]);
240
241 assert_eq!(content1.audience(), content2.audience());
242 assert_eq!(content1.priority(), content2.priority());
243 }
244
245 #[test]
246 fn test_content_annotations_overwrite() {
247 let content = Content::text("hello")
248 .with_audience(vec![Role::User])
249 .with_priority(0.5)
250 .with_audience(vec![Role::Assistant])
251 .with_priority(0.8);
252
253 assert_eq!(content.audience(), Some(&vec![Role::Assistant]));
254 assert_eq!(content.priority(), Some(0.8));
255 }
256
257 #[test]
258 fn test_content_annotations_image() {
259 let content = Content::image("data", "image/png")
260 .with_audience(vec![Role::User])
261 .with_priority(0.5);
262
263 assert_eq!(content.audience(), Some(&vec![Role::User]));
264 assert_eq!(content.priority(), Some(0.5));
265 }
266
267 #[test]
268 fn test_content_annotations_preservation() {
269 let text_content = Content::text("hello")
270 .with_audience(vec![Role::User])
271 .with_priority(0.5);
272
273 match &text_content {
274 Content::Text(TextContent { annotations, .. }) => {
275 assert!(annotations.is_some());
276 let ann = annotations.as_ref().unwrap();
277 assert_eq!(ann.audience, Some(vec![Role::User]));
278 assert_eq!(ann.priority, Some(0.5));
279 }
280 _ => panic!("Expected Text content"),
281 }
282 }
283
284 #[test]
285 #[should_panic(expected = "Priority must be between 0.0 and 1.0")]
286 fn test_invalid_priority() {
287 Content::text("hello").with_priority(1.5);
288 }
289
290 #[test]
291 fn test_unannotated() {
292 let content = Content::text("hello")
293 .with_audience(vec![Role::User])
294 .with_priority(0.5);
295 let unannotated = content.unannotated();
296 assert_eq!(unannotated.audience(), None);
297 assert_eq!(unannotated.priority(), None);
298 }
299
300 #[test]
301 fn test_partial_annotations() {
302 let content = Content::text("hello").with_priority(0.5);
303 assert_eq!(content.audience(), None);
304 assert_eq!(content.priority(), Some(0.5));
305
306 let content = Content::text("hello").with_audience(vec![Role::User]);
307 assert_eq!(content.audience(), Some(&vec![Role::User]));
308 assert_eq!(content.priority(), None);
309 }
310}