Skip to main content

a2a_rs/domain/core/
message.rs

1use crate::domain::error::A2AError;
2
3// Re-export the generated types so downstream code gets them from `domain::core::message`
4pub use crate::domain::generated::{Artifact, Message, Part, Role, part};
5
6#[allow(non_upper_case_globals)]
7impl Role {
8    pub const User: Self = Self::ROLE_USER;
9    pub const Agent: Self = Self::ROLE_AGENT;
10}
11
12impl Part {
13    /// Create a text part
14    #[inline]
15    pub fn text(content: String) -> Self {
16        Self {
17            content: Some(part::Content::Text(content)),
18            ..Default::default()
19        }
20    }
21
22    /// Create a text part with metadata
23    #[inline]
24    pub fn text_with_metadata(
25        content: String,
26        metadata: ::buffa_types::google::protobuf::Struct,
27    ) -> Self {
28        Self {
29            content: Some(part::Content::Text(content)),
30            metadata: ::buffa::MessageField::some(metadata),
31            ..Default::default()
32        }
33    }
34
35    /// Create a data part
36    #[inline]
37    pub fn data(data: ::buffa_types::google::protobuf::Value) -> Self {
38        Self {
39            content: Some(part::Content::Data(Box::new(data))),
40            ..Default::default()
41        }
42    }
43
44    /// Create a file part from base64 encoded data (or bytes)
45    pub fn file_from_bytes(
46        bytes: Vec<u8>,
47        name: Option<String>,
48        mime_type: Option<String>,
49    ) -> Self {
50        Self {
51            content: Some(part::Content::Raw(bytes)),
52            filename: name.unwrap_or_default(),
53            media_type: mime_type.unwrap_or_default(),
54            ..Default::default()
55        }
56    }
57
58    /// Create a file part from a URI
59    pub fn file_from_uri(uri: String, name: Option<String>, mime_type: Option<String>) -> Self {
60        Self {
61            content: Some(part::Content::Url(uri)),
62            filename: name.unwrap_or_default(),
63            media_type: mime_type.unwrap_or_default(),
64            ..Default::default()
65        }
66    }
67
68    /// Helper method to get the text content if this is a Text part
69    pub fn get_text(&self) -> Option<&str> {
70        match &self.content {
71            Some(part::Content::Text(text)) => Some(text.as_str()),
72            _ => None,
73        }
74    }
75
76    /// Validate the part content
77    pub fn validate(&self) -> Result<(), A2AError> {
78        match &self.content {
79            Some(_) => Ok(()),
80            None => Err(A2AError::InvalidParams(
81                "Part must contain content (text, raw, url, or data)".to_string(),
82            )),
83        }
84    }
85
86    /// Create a builder-style text part that can be chained
87    pub fn text_builder(content: String) -> PartBuilder {
88        PartBuilder {
89            part: Self::text(content),
90        }
91    }
92
93    /// Create a builder-style data part that can be chained
94    pub fn data_builder(data: ::buffa_types::google::protobuf::Value) -> PartBuilder {
95        PartBuilder {
96            part: Self::data(data),
97        }
98    }
99
100    /// Create a builder-style file part that can be chained
101    pub fn file_builder() -> FilePartBuilder {
102        FilePartBuilder::new()
103    }
104}
105
106/// Builder for Part instances
107pub struct PartBuilder {
108    part: Part,
109}
110
111impl PartBuilder {
112    /// Add metadata to any part type
113    pub fn with_metadata(mut self, metadata: ::buffa_types::google::protobuf::Struct) -> Self {
114        self.part.metadata = ::buffa::MessageField::some(metadata);
115        self
116    }
117
118    /// Build the final Part
119    pub fn build(self) -> Part {
120        self.part
121    }
122}
123
124/// Builder for file parts with validation
125pub struct FilePartBuilder {
126    name: Option<String>,
127    mime_type: Option<String>,
128    bytes: Option<Vec<u8>>,
129    uri: Option<String>,
130    metadata: Option<::buffa_types::google::protobuf::Struct>,
131}
132
133impl Default for FilePartBuilder {
134    fn default() -> Self {
135        Self::new()
136    }
137}
138
139impl FilePartBuilder {
140    pub fn new() -> Self {
141        Self {
142            name: None,
143            mime_type: None,
144            bytes: None,
145            uri: None,
146            metadata: None,
147        }
148    }
149
150    /// Set the file name
151    pub fn name(mut self, name: String) -> Self {
152        self.name = Some(name);
153        self
154    }
155
156    /// Set the MIME type
157    pub fn mime_type(mut self, mime_type: String) -> Self {
158        self.mime_type = Some(mime_type);
159        self
160    }
161
162    /// Set file content as bytes
163    pub fn bytes(mut self, bytes: Vec<u8>) -> Self {
164        self.bytes = Some(bytes);
165        self.uri = None; // Clear URI if setting bytes
166        self
167    }
168
169    /// Set file URI
170    pub fn uri(mut self, uri: String) -> Self {
171        self.uri = Some(uri);
172        self.bytes = None; // Clear bytes if setting URI
173        self
174    }
175
176    /// Add metadata
177    pub fn with_metadata(mut self, metadata: ::buffa_types::google::protobuf::Struct) -> Self {
178        self.metadata = Some(metadata);
179        self
180    }
181
182    /// Build the file part with validation
183    pub fn build(self) -> Result<Part, A2AError> {
184        let content = match (self.bytes, self.uri) {
185            (Some(b), None) => part::Content::Raw(b),
186            (None, Some(u)) => part::Content::Url(u),
187            (Some(_), Some(_)) => {
188                return Err(A2AError::InvalidParams(
189                    "Cannot provide both bytes and uri".to_string(),
190                ));
191            }
192            (None, None) => {
193                return Err(A2AError::InvalidParams(
194                    "Must provide either bytes or uri".to_string(),
195                ));
196            }
197        };
198
199        Ok(Part {
200            content: Some(content),
201            filename: self.name.unwrap_or_default(),
202            media_type: self.mime_type.unwrap_or_default(),
203            metadata: self.metadata.into(),
204            ..Default::default()
205        })
206    }
207}
208
209/// Builder for Message instances to keep compatibility
210pub struct MessageBuilder {
211    message_id: String,
212    context_id: String,
213    task_id: String,
214    role: Role,
215    parts: Vec<Part>,
216    metadata: Option<::buffa_types::google::protobuf::Struct>,
217    extensions: Vec<String>,
218    reference_task_ids: Vec<String>,
219}
220
221impl Default for MessageBuilder {
222    fn default() -> Self {
223        Self::new()
224    }
225}
226
227impl MessageBuilder {
228    pub fn new() -> Self {
229        Self {
230            message_id: String::new(),
231            context_id: String::new(),
232            task_id: String::new(),
233            role: Role::ROLE_UNSPECIFIED,
234            parts: Vec::new(),
235            metadata: None,
236            extensions: Vec::new(),
237            reference_task_ids: Vec::new(),
238        }
239    }
240
241    pub fn message_id(mut self, message_id: String) -> Self {
242        self.message_id = message_id;
243        self
244    }
245
246    pub fn context_id(mut self, context_id: String) -> Self {
247        self.context_id = context_id;
248        self
249    }
250
251    pub fn task_id(mut self, task_id: String) -> Self {
252        self.task_id = task_id;
253        self
254    }
255
256    pub fn role(mut self, role: Role) -> Self {
257        self.role = role;
258        self
259    }
260
261    pub fn parts(mut self, parts: Vec<Part>) -> Self {
262        self.parts = parts;
263        self
264    }
265
266    pub fn metadata(mut self, metadata: ::buffa_types::google::protobuf::Struct) -> Self {
267        self.metadata = Some(metadata);
268        self
269    }
270
271    pub fn extensions(mut self, extensions: Vec<String>) -> Self {
272        self.extensions = extensions;
273        self
274    }
275
276    pub fn reference_task_ids(mut self, reference_task_ids: Vec<String>) -> Self {
277        self.reference_task_ids = reference_task_ids;
278        self
279    }
280
281    pub fn build(self) -> Message {
282        Message {
283            message_id: self.message_id,
284            context_id: self.context_id,
285            task_id: self.task_id,
286            role: ::buffa::EnumValue::from(self.role),
287            parts: self.parts,
288            metadata: self.metadata.into(),
289            extensions: self.extensions,
290            reference_task_ids: self.reference_task_ids,
291            ..Default::default()
292        }
293    }
294}
295
296impl Message {
297    /// Create a new Message builder
298    pub fn builder() -> MessageBuilder {
299        MessageBuilder::new()
300    }
301
302    /// Create a new user message with a single text part
303    pub fn user_text(text: String, message_id: String) -> Self {
304        Self {
305            role: ::buffa::EnumValue::from(Role::ROLE_USER),
306            parts: vec![Part::text(text)],
307            message_id,
308            ..Default::default()
309        }
310    }
311
312    /// Create a new agent message with a single text part
313    pub fn agent_text(text: String, message_id: String) -> Self {
314        Self {
315            role: ::buffa::EnumValue::from(Role::ROLE_AGENT),
316            parts: vec![Part::text(text)],
317            message_id,
318            ..Default::default()
319        }
320    }
321
322    /// Add a part to this message
323    pub fn add_part(&mut self, part: Part) {
324        self.parts.push(part);
325    }
326
327    /// Add a part to this message, validating and returning Result
328    pub fn add_part_validated(&mut self, part: Part) -> Result<(), A2AError> {
329        part.validate()?;
330        self.parts.push(part);
331        Ok(())
332    }
333
334    /// Validate a message
335    pub fn validate(&self) -> Result<(), A2AError> {
336        for part in &self.parts {
337            part.validate()?;
338        }
339        Ok(())
340    }
341}