mofa_foundation/agent/
base.rs1use mofa_kernel::agent::{
6 AgentCapabilities, AgentContext, AgentError, AgentOutput, AgentResult, AgentState, AgentStats,
7 InterruptResult, MoFAAgent,
8};
9
10use async_trait::async_trait;
11
12pub struct BaseAgent {
16 pub id: String,
18 pub name: String,
20 pub description: Option<String>,
22 pub version: Option<String>,
24 pub capabilities: AgentCapabilities,
26 pub state: AgentState,
28 stats: AgentStats,
30}
31
32impl BaseAgent {
33 pub fn new(id: impl Into<String>, name: impl Into<String>) -> Self {
35 Self {
36 id: id.into(),
37 name: name.into(),
38 description: None,
39 version: None,
40 capabilities: AgentCapabilities::default(),
41 state: AgentState::Created,
42 stats: AgentStats::default(),
43 }
44 }
45
46 pub fn with_description(mut self, description: impl Into<String>) -> Self {
48 self.description = Some(description.into());
49 self
50 }
51
52 pub fn with_version(mut self, version: impl Into<String>) -> Self {
54 self.version = Some(version.into());
55 self
56 }
57
58 pub fn with_capabilities(mut self, capabilities: AgentCapabilities) -> Self {
60 self.capabilities = capabilities;
61 self
62 }
63
64 pub fn transition_to(&mut self, new_state: AgentState) -> AgentResult<()> {
66 if self.state.can_transition_to(&new_state) {
67 self.state = new_state;
68 Ok(())
69 } else {
70 Err(AgentError::invalid_state_transition(
71 &self.state,
72 &new_state,
73 ))
74 }
75 }
76
77 pub fn record_success(&mut self, duration_ms: u64, tokens: u64, tool_calls: u64) {
79 self.stats.total_executions += 1;
80 self.stats.successful_executions += 1;
81 self.stats.total_tokens_used += tokens;
82 self.stats.total_tool_calls += tool_calls;
83
84 let n = self.stats.total_executions as f64;
86 self.stats.avg_execution_time_ms =
87 (self.stats.avg_execution_time_ms * (n - 1.0) + duration_ms as f64) / n;
88 }
89
90 pub fn record_failure(&mut self) {
92 self.stats.total_executions += 1;
93 self.stats.failed_executions += 1;
94 }
95
96 pub fn stats(&self) -> &AgentStats {
98 &self.stats
99 }
100
101 pub fn id(&self) -> &str {
103 &self.id
104 }
105
106 pub fn name(&self) -> &str {
108 &self.name
109 }
110
111 pub fn capabilities(&self) -> &AgentCapabilities {
113 &self.capabilities
114 }
115
116 pub fn state(&self) -> AgentState {
118 self.state.clone()
119 }
120
121 pub async fn initialize(&mut self, _ctx: &AgentContext) -> AgentResult<()> {
123 self.transition_to(AgentState::Initializing)?;
124 self.transition_to(AgentState::Ready)?;
125 Ok(())
126 }
127
128 pub async fn interrupt(&mut self) -> AgentResult<InterruptResult> {
130 Ok(InterruptResult::Acknowledged)
131 }
132
133 pub async fn shutdown(&mut self) -> AgentResult<()> {
135 self.transition_to(AgentState::ShuttingDown)?;
136 self.transition_to(AgentState::Shutdown)?;
137 Ok(())
138 }
139}
140
141#[async_trait]
142impl MoFAAgent for BaseAgent {
143 fn id(&self) -> &str {
144 &self.id
145 }
146
147 fn name(&self) -> &str {
148 &self.name
149 }
150
151 fn capabilities(&self) -> &AgentCapabilities {
152 &self.capabilities
153 }
154
155 async fn initialize(
156 &mut self,
157 _ctx: &mofa_kernel::agent::context::AgentContext,
158 ) -> AgentResult<()> {
159 self.transition_to(AgentState::Initializing)?;
160 self.transition_to(AgentState::Ready)?;
161 Ok(())
162 }
163
164 async fn execute(
165 &mut self,
166 _input: mofa_kernel::agent::AgentInput,
167 _ctx: &mofa_kernel::agent::context::AgentContext,
168 ) -> AgentResult<AgentOutput> {
169 Ok(AgentOutput::text("BaseAgent execute"))
170 }
171
172 async fn shutdown(&mut self) -> AgentResult<()> {
173 self.transition_to(AgentState::ShuttingDown)?;
174 self.transition_to(AgentState::Shutdown)?;
175 Ok(())
176 }
177
178 fn state(&self) -> AgentState {
179 self.state.clone()
180 }
181}