1use crate::domain::error::A2AError;
2
3pub 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 #[inline]
15 pub fn text(content: String) -> Self {
16 Self {
17 content: Some(part::Content::Text(content)),
18 ..Default::default()
19 }
20 }
21
22 #[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 #[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 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 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 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 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 pub fn text_builder(content: String) -> PartBuilder {
88 PartBuilder {
89 part: Self::text(content),
90 }
91 }
92
93 pub fn data_builder(data: ::buffa_types::google::protobuf::Value) -> PartBuilder {
95 PartBuilder {
96 part: Self::data(data),
97 }
98 }
99
100 pub fn file_builder() -> FilePartBuilder {
102 FilePartBuilder::new()
103 }
104}
105
106pub struct PartBuilder {
108 part: Part,
109}
110
111impl PartBuilder {
112 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 pub fn build(self) -> Part {
120 self.part
121 }
122}
123
124pub 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 pub fn name(mut self, name: String) -> Self {
152 self.name = Some(name);
153 self
154 }
155
156 pub fn mime_type(mut self, mime_type: String) -> Self {
158 self.mime_type = Some(mime_type);
159 self
160 }
161
162 pub fn bytes(mut self, bytes: Vec<u8>) -> Self {
164 self.bytes = Some(bytes);
165 self.uri = None; self
167 }
168
169 pub fn uri(mut self, uri: String) -> Self {
171 self.uri = Some(uri);
172 self.bytes = None; self
174 }
175
176 pub fn with_metadata(mut self, metadata: ::buffa_types::google::protobuf::Struct) -> Self {
178 self.metadata = Some(metadata);
179 self
180 }
181
182 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
209pub 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 pub fn builder() -> MessageBuilder {
299 MessageBuilder::new()
300 }
301
302 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 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 pub fn add_part(&mut self, part: Part) {
324 self.parts.push(part);
325 }
326
327 pub fn add_part_validated(&mut self, part: Part) -> Result<(), A2AError> {
329 part.validate()?;
330 self.parts.push(part);
331 Ok(())
332 }
333
334 pub fn validate(&self) -> Result<(), A2AError> {
336 for part in &self.parts {
337 part.validate()?;
338 }
339 Ok(())
340 }
341}