systemprompt_agent/services/a2a_server/processing/task_builder/
mod.rs1mod builders;
2pub mod helpers;
3
4pub use builders::{
5 build_canceled_task, build_completed_task, build_mock_task, build_multiturn_task,
6 build_submitted_task,
7};
8
9use crate::models::a2a::{
10 Artifact, Message, MessageRole, Part, Task, TaskState, TaskStatus, TextPart,
11};
12use systemprompt_identifiers::{ContextId, MessageId, TaskId};
13use systemprompt_models::a2a::TaskMetadata;
14
15#[derive(Debug)]
16pub struct TaskBuilder {
17 task_id: TaskId,
18 context_id: ContextId,
19 state: TaskState,
20 response_text: String,
21 id: MessageId,
22 user_message: Option<Message>,
23 artifacts: Vec<Artifact>,
24 metadata: Option<TaskMetadata>,
25}
26
27impl TaskBuilder {
28 pub fn new(context_id: ContextId) -> Self {
29 Self {
30 task_id: TaskId::generate(),
31 context_id,
32 state: TaskState::Completed,
33 response_text: String::new(),
34 id: MessageId::generate(),
35 user_message: None,
36 artifacts: Vec::new(),
37 metadata: None,
38 }
39 }
40
41 pub fn with_task_id(mut self, task_id: TaskId) -> Self {
42 self.task_id = task_id;
43 self
44 }
45
46 pub const fn with_state(mut self, state: TaskState) -> Self {
47 self.state = state;
48 self
49 }
50
51 pub fn with_response_text(mut self, text: String) -> Self {
52 self.response_text = text;
53 self
54 }
55
56 pub fn with_message_id(mut self, id: impl Into<MessageId>) -> Self {
57 self.id = id.into();
58 self
59 }
60
61 pub fn with_user_message(mut self, message: Message) -> Self {
62 self.user_message = Some(message);
63 self
64 }
65
66 pub fn with_artifacts(mut self, artifacts: Vec<Artifact>) -> Self {
67 self.artifacts = artifacts;
68 self
69 }
70
71 pub fn with_metadata(mut self, metadata: TaskMetadata) -> Self {
72 self.metadata = Some(metadata);
73 self
74 }
75
76 pub fn build(self) -> Task {
77 let agent_message = Message {
78 role: MessageRole::Agent,
79 parts: vec![Part::Text(TextPart {
80 text: self.response_text.clone(),
81 })],
82 message_id: self.id.clone(),
83 task_id: Some(self.task_id.clone()),
84 context_id: self.context_id.clone(),
85 metadata: None,
86 extensions: None,
87 reference_task_ids: None,
88 };
89
90 let history = if let Some(user_msg) = self.user_message {
91 Some(vec![
92 user_msg,
93 Message {
94 role: MessageRole::Agent,
95 parts: vec![Part::Text(TextPart {
96 text: self.response_text.clone(),
97 })],
98 message_id: MessageId::generate(),
99 task_id: Some(self.task_id.clone()),
100 context_id: self.context_id.clone(),
101 metadata: None,
102 extensions: None,
103 reference_task_ids: None,
104 },
105 ])
106 } else {
107 None
108 };
109
110 Task {
111 id: self.task_id.clone(),
112 context_id: self.context_id.clone(),
113 status: TaskStatus {
114 state: self.state,
115 message: Some(agent_message),
116 timestamp: Some(chrono::Utc::now()),
117 },
118 history,
119 artifacts: if self.artifacts.is_empty() {
120 None
121 } else {
122 Some(self.artifacts)
123 },
124 metadata: self.metadata,
125 created_at: Some(chrono::Utc::now()),
126 last_modified: Some(chrono::Utc::now()),
127 }
128 }
129}