agenterra_rmcp/model/
annotated.rs

1use std::ops::{Deref, DerefMut};
2
3use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5
6use super::{
7    RawAudioContent, RawContent, RawEmbeddedResource, RawImageContent, RawResource,
8    RawResourceTemplate, RawTextContent, Role,
9};
10
11#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
12#[serde(rename_all = "camelCase")]
13#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
14pub struct Annotations {
15    #[serde(skip_serializing_if = "Option::is_none")]
16    pub audience: Option<Vec<Role>>,
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub priority: Option<f32>,
19    #[serde(skip_serializing_if = "Option::is_none")]
20    pub timestamp: Option<DateTime<Utc>>,
21}
22
23impl Annotations {
24    /// Creates a new Annotations instance specifically for resources
25    /// optional priority, and a timestamp (defaults to now if None)
26    pub fn for_resource(priority: f32, timestamp: DateTime<Utc>) -> Self {
27        assert!(
28            (0.0..=1.0).contains(&priority),
29            "Priority {priority} must be between 0.0 and 1.0"
30        );
31        Annotations {
32            priority: Some(priority),
33            timestamp: Some(timestamp),
34            audience: None,
35        }
36    }
37}
38
39#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
40#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
41pub struct Annotated<T: AnnotateAble> {
42    #[serde(flatten)]
43    pub raw: T,
44    #[serde(skip_serializing_if = "Option::is_none")]
45    pub annotations: Option<Annotations>,
46}
47
48impl<T: AnnotateAble> Deref for Annotated<T> {
49    type Target = T;
50    fn deref(&self) -> &Self::Target {
51        &self.raw
52    }
53}
54
55impl<T: AnnotateAble> DerefMut for Annotated<T> {
56    fn deref_mut(&mut self) -> &mut Self::Target {
57        &mut self.raw
58    }
59}
60
61impl<T: AnnotateAble> Annotated<T> {
62    pub fn new(raw: T, annotations: Option<Annotations>) -> Self {
63        Self { raw, annotations }
64    }
65    pub fn remove_annotation(&mut self) -> Option<Annotations> {
66        self.annotations.take()
67    }
68    pub fn audience(&self) -> Option<&Vec<Role>> {
69        self.annotations.as_ref().and_then(|a| a.audience.as_ref())
70    }
71    pub fn priority(&self) -> Option<f32> {
72        self.annotations.as_ref().and_then(|a| a.priority)
73    }
74    pub fn timestamp(&self) -> Option<DateTime<Utc>> {
75        self.annotations.as_ref().and_then(|a| a.timestamp)
76    }
77    pub fn with_audience(self, audience: Vec<Role>) -> Annotated<T>
78    where
79        Self: Sized,
80    {
81        if let Some(annotations) = self.annotations {
82            Annotated {
83                raw: self.raw,
84                annotations: Some(Annotations {
85                    audience: Some(audience),
86                    ..annotations
87                }),
88            }
89        } else {
90            Annotated {
91                raw: self.raw,
92                annotations: Some(Annotations {
93                    audience: Some(audience),
94                    priority: None,
95                    timestamp: None,
96                }),
97            }
98        }
99    }
100    pub fn with_priority(self, priority: f32) -> Annotated<T>
101    where
102        Self: Sized,
103    {
104        if let Some(annotations) = self.annotations {
105            Annotated {
106                raw: self.raw,
107                annotations: Some(Annotations {
108                    priority: Some(priority),
109                    ..annotations
110                }),
111            }
112        } else {
113            Annotated {
114                raw: self.raw,
115                annotations: Some(Annotations {
116                    priority: Some(priority),
117                    timestamp: None,
118                    audience: None,
119                }),
120            }
121        }
122    }
123    pub fn with_timestamp(self, timestamp: DateTime<Utc>) -> Annotated<T>
124    where
125        Self: Sized,
126    {
127        if let Some(annotations) = self.annotations {
128            Annotated {
129                raw: self.raw,
130                annotations: Some(Annotations {
131                    timestamp: Some(timestamp),
132                    ..annotations
133                }),
134            }
135        } else {
136            Annotated {
137                raw: self.raw,
138                annotations: Some(Annotations {
139                    timestamp: Some(timestamp),
140                    priority: None,
141                    audience: None,
142                }),
143            }
144        }
145    }
146    pub fn with_timestamp_now(self) -> Annotated<T>
147    where
148        Self: Sized,
149    {
150        self.with_timestamp(Utc::now())
151    }
152}
153
154mod sealed {
155    pub trait Sealed {}
156}
157macro_rules! annotate {
158    ($T: ident) => {
159        impl sealed::Sealed for $T {}
160        impl AnnotateAble for $T {}
161    };
162}
163
164annotate!(RawContent);
165annotate!(RawTextContent);
166annotate!(RawImageContent);
167annotate!(RawAudioContent);
168annotate!(RawEmbeddedResource);
169annotate!(RawResource);
170annotate!(RawResourceTemplate);
171
172pub trait AnnotateAble: sealed::Sealed {
173    fn optional_annotate(self, annotations: Option<Annotations>) -> Annotated<Self>
174    where
175        Self: Sized,
176    {
177        Annotated::new(self, annotations)
178    }
179    fn annotate(self, annotations: Annotations) -> Annotated<Self>
180    where
181        Self: Sized,
182    {
183        Annotated::new(self, Some(annotations))
184    }
185    fn no_annotation(self) -> Annotated<Self>
186    where
187        Self: Sized,
188    {
189        Annotated::new(self, None)
190    }
191    fn with_audience(self, audience: Vec<Role>) -> Annotated<Self>
192    where
193        Self: Sized,
194    {
195        self.annotate(Annotations {
196            audience: Some(audience),
197            ..Default::default()
198        })
199    }
200    fn with_priority(self, priority: f32) -> Annotated<Self>
201    where
202        Self: Sized,
203    {
204        self.annotate(Annotations {
205            priority: Some(priority),
206            ..Default::default()
207        })
208    }
209    fn with_timestamp(self, timestamp: DateTime<Utc>) -> Annotated<Self>
210    where
211        Self: Sized,
212    {
213        self.annotate(Annotations {
214            timestamp: Some(timestamp),
215            ..Default::default()
216        })
217    }
218    fn with_timestamp_now(self) -> Annotated<Self>
219    where
220        Self: Sized,
221    {
222        self.with_timestamp(Utc::now())
223    }
224}
agenterra_rmcp/model/annotated.rs

agenterra_rmcp/model/
annotated.rs
