1use crate::agent::config::AgentConfig;
2use crate::agent::executor::event_helper::EventHelper;
3use crate::agent::memory::MemoryProvider;
4use crate::agent::task::Task;
5use crate::agent::{AgentExecutor, Context, output::AgentOutputT};
6use crate::tool::{ToolT, to_llm_tool};
7use async_trait::async_trait;
8use autoagents_llm::LLMProvider;
9use autoagents_llm::chat::Tool;
10use autoagents_protocol::{ActorID, Event, SubmissionId};
11
12use serde_json::Value;
13use std::marker::PhantomData;
14use std::{fmt::Debug, sync::Arc};
15
16#[cfg(target_arch = "wasm32")]
17pub use futures::lock::Mutex;
18#[cfg(not(target_arch = "wasm32"))]
19pub use tokio::sync::Mutex;
20
21#[cfg(target_arch = "wasm32")]
22use futures::channel::mpsc::Sender;
23
24#[cfg(not(target_arch = "wasm32"))]
25use tokio::sync::mpsc::Sender;
26
27use crate::agent::error::RunnableAgentError;
28use crate::agent::hooks::AgentHooks;
29use uuid::Uuid;
30
31#[async_trait]
34pub trait AgentDeriveT: Send + Sync + 'static + Debug {
35 type Output: AgentOutputT;
37
38 fn description(&self) -> &str;
40
41 fn output_schema(&self) -> Option<Value>;
43
44 fn name(&self) -> &str;
46
47 fn tools(&self) -> Vec<Box<dyn ToolT>>;
49}
50
51pub trait AgentType: 'static + Send + Sync {
52 fn type_name() -> &'static str;
53}
54
55#[derive(Clone)]
57pub struct BaseAgent<T: AgentDeriveT + AgentExecutor + AgentHooks + Send + Sync, A: AgentType> {
58 pub(crate) inner: Arc<T>,
60 pub(crate) llm: Arc<dyn LLMProvider>,
62 pub id: ActorID,
64 pub(crate) memory: Option<Arc<Mutex<Box<dyn MemoryProvider>>>>,
66 pub(crate) serialized_tools: Option<Arc<Vec<Tool>>>,
68 pub(crate) tx: Option<Sender<Event>>,
70 pub(crate) stream: bool,
72 pub(crate) marker: PhantomData<A>,
73}
74
75impl<T: AgentDeriveT + AgentExecutor + AgentHooks, A: AgentType> Debug for BaseAgent<T, A> {
76 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77 f.write_str(format!("A: {} - T: {}", self.inner().name(), A::type_name()).as_str())
78 }
79}
80
81impl<T: AgentDeriveT + AgentExecutor + AgentHooks, A: AgentType> BaseAgent<T, A> {
82 pub async fn new(
84 inner: T,
85 llm: Arc<dyn LLMProvider>,
86 memory: Option<Box<dyn MemoryProvider>>,
87 tx: Sender<Event>,
88 stream: bool,
89 ) -> Result<Self, RunnableAgentError> {
90 let tool_defs = inner.tools();
91 let serialized_tools = if tool_defs.is_empty() {
92 None
93 } else {
94 Some(Arc::new(
95 tool_defs.iter().map(to_llm_tool).collect::<Vec<_>>(),
96 ))
97 };
98 let agent = Self {
99 inner: Arc::new(inner),
100 id: Uuid::new_v4(),
101 llm,
102 tx: Some(tx),
103 memory: memory.map(|m| Arc::new(Mutex::new(m))),
104 serialized_tools,
105 stream,
106 marker: PhantomData,
107 };
108
109 agent.inner().on_agent_create().await;
111
112 Ok(agent)
113 }
114
115 pub fn inner(&self) -> Arc<T> {
116 self.inner.clone()
117 }
118
119 pub fn name(&self) -> &str {
121 self.inner.name()
122 }
123
124 pub fn description(&self) -> &str {
126 self.inner.description()
127 }
128
129 pub fn tools(&self) -> Vec<Box<dyn ToolT>> {
131 self.inner.tools()
132 }
133
134 pub fn serialized_tools(&self) -> Option<Arc<Vec<Tool>>> {
135 self.serialized_tools.clone()
136 }
137
138 pub fn stream(&self) -> bool {
139 self.stream
140 }
141
142 pub(crate) fn create_context(&self) -> Arc<Context> {
143 let tools = self.tools();
144 let cached_tools = self
145 .serialized_tools()
146 .filter(|cached| tools_match_cached(&tools, cached));
147 Arc::new(
148 Context::new(self.llm(), self.tx.clone())
149 .with_memory(self.memory())
150 .with_serialized_tools(cached_tools)
151 .with_tools(tools)
152 .with_config(self.agent_config())
153 .with_stream(self.stream()),
154 )
155 }
156
157 pub fn agent_config(&self) -> AgentConfig {
158 let output_schema = self.inner().output_schema();
159 let structured_schema =
160 output_schema.and_then(|schema| serde_json::from_value(schema).ok());
161 AgentConfig {
162 name: self.name().into(),
163 description: self.description().into(),
164 id: self.id,
165 output_schema: structured_schema,
166 }
167 }
168
169 pub fn llm(&self) -> Arc<dyn LLMProvider> {
171 self.llm.clone()
172 }
173
174 pub fn memory(&self) -> Option<Arc<Mutex<Box<dyn MemoryProvider>>>> {
176 self.memory.clone()
177 }
178
179 pub(crate) fn clone_shallow(&self) -> Self {
181 Self {
182 inner: self.inner.clone(),
183 llm: self.llm.clone(),
184 id: self.id,
185 memory: self.memory.clone(),
186 serialized_tools: self.serialized_tools.clone(),
187 tx: self.tx.clone(),
188 stream: self.stream,
189 marker: PhantomData,
190 }
191 }
192
193 pub(crate) async fn finish_executor_run(
198 &self,
199 task: &Task,
200 context: &Context,
201 submission_id: SubmissionId,
202 executor_out: <T as AgentExecutor>::Output,
203 ) -> Result<<T as AgentDeriveT>::Output, RunnableAgentError>
204 where
205 Value: From<<T as AgentExecutor>::Output>,
206 <T as AgentDeriveT>::Output: From<<T as AgentExecutor>::Output>,
207 <T as AgentExecutor>::Output: Clone,
208 {
209 let tx_event = self.tx.clone();
210 let value: Value = executor_out.clone().into();
211 #[cfg(not(target_arch = "wasm32"))]
212 if let Err(e) = EventHelper::send_task_completed_value(
213 &tx_event,
214 submission_id,
215 self.id,
216 self.name().to_string(),
217 &value,
218 )
219 .await
220 {
221 let err = RunnableAgentError::ExecutorError(e.to_string());
222 EventHelper::send_task_error(&tx_event, submission_id, self.id, err.to_string()).await;
223 return Err(err);
224 }
225
226 let agent_out: <T as AgentDeriveT>::Output = executor_out.into();
227 self.inner.on_run_complete(task, &agent_out, context).await;
228 Ok(agent_out)
229 }
230}
231
232fn tools_match_cached(tools: &[Box<dyn ToolT>], cached: &[Tool]) -> bool {
233 if tools.len() != cached.len() {
234 return false;
235 }
236
237 tools.iter().zip(cached.iter()).all(|(tool, cached_tool)| {
238 cached_tool.tool_type == "function"
239 && cached_tool.function.name == tool.name()
240 && cached_tool.function.description == tool.description()
241 && cached_tool.function.parameters == tool.args_schema()
242 })
243}
244
245#[cfg(test)]
246mod tests {
247 use super::*;
248 use crate::agent::memory::SlidingWindowMemory;
249 use crate::agent::{AgentConfig, DirectAgent};
250 use crate::tests::{MockAgentImpl, MockLLMProvider};
251 use autoagents_llm::chat::StructuredOutputFormat;
252 use std::sync::Arc;
253 use tokio::sync::mpsc::{Receiver, channel};
254 use uuid::Uuid;
255
256 #[test]
257 fn test_agent_config_with_schema() {
258 let schema = StructuredOutputFormat {
259 name: "TestSchema".to_string(),
260 description: Some("Test schema".to_string()),
261 schema: Some(serde_json::json!({"type": "object"})),
262 strict: Some(true),
263 };
264
265 let config = AgentConfig {
266 name: "test_agent".to_string(),
267 id: Uuid::new_v4(),
268 description: "A test agent".to_string(),
269 output_schema: Some(schema.clone()),
270 };
271
272 assert_eq!(config.name, "test_agent");
273 assert_eq!(config.description, "A test agent");
274 assert!(config.output_schema.is_some());
275 assert_eq!(config.output_schema.unwrap().name, "TestSchema");
276 }
277
278 #[tokio::test]
279 async fn test_base_agent_creation_with_memory_and_stream() {
280 let mock_agent = MockAgentImpl::new("test", "test description");
281 let llm = Arc::new(MockLLMProvider);
282 let memory = Box::new(SlidingWindowMemory::new(5));
283 let (tx, _): (Sender<Event>, Receiver<Event>) = channel(32);
284 let base_agent = BaseAgent::<_, DirectAgent>::new(mock_agent, llm, Some(memory), tx, true)
285 .await
286 .unwrap();
287
288 assert_eq!(base_agent.name(), "test");
289 assert_eq!(base_agent.description(), "test description");
290 assert!(base_agent.memory().is_some());
291 assert!(base_agent.stream);
292 }
293
294 #[tokio::test]
295 async fn test_base_agent_create_context_populates_config() {
296 let mock_agent = MockAgentImpl::new("ctx_agent", "context agent");
297 let llm = Arc::new(MockLLMProvider);
298 let (tx, _): (Sender<Event>, Receiver<Event>) = channel(32);
299 let base_agent = BaseAgent::<_, DirectAgent>::new(mock_agent, llm, None, tx, false)
300 .await
301 .unwrap();
302
303 let context = base_agent.create_context();
304 let config = context.config();
305 assert_eq!(config.name, "ctx_agent");
306 assert_eq!(config.description, "context agent");
307 }
308}