Skip to main content

autoagents_core/agent/
base.rs

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/// Core trait that defines agent metadata and behavior
32/// This trait is implemented via the #[agent] macro
33#[async_trait]
34pub trait AgentDeriveT: Send + Sync + 'static + Debug {
35    /// The output type this agent produces
36    type Output: AgentOutputT;
37
38    /// Get the agent's description
39    fn description(&self) -> &str;
40
41    // If you provide None then its taken as String output
42    fn output_schema(&self) -> Option<Value>;
43
44    /// Get the agent's name
45    fn name(&self) -> &str;
46
47    /// Get the tools available to this agent
48    fn tools(&self) -> Vec<Box<dyn ToolT>>;
49}
50
51pub trait AgentType: 'static + Send + Sync {
52    fn type_name() -> &'static str;
53}
54
55/// Base agent type that wraps an AgentDeriveT implementation with additional runtime components
56#[derive(Clone)]
57pub struct BaseAgent<T: AgentDeriveT + AgentExecutor + AgentHooks + Send + Sync, A: AgentType> {
58    /// The inner agent implementation (from macro)
59    pub(crate) inner: Arc<T>,
60    /// LLM provider for this agent
61    pub(crate) llm: Arc<dyn LLMProvider>,
62    /// Agent ID
63    pub id: ActorID,
64    /// Optional memory provider
65    pub(crate) memory: Option<Arc<Mutex<Box<dyn MemoryProvider>>>>,
66    /// Cached serialized tool definitions
67    pub(crate) serialized_tools: Option<Arc<Vec<Tool>>>,
68    /// Tx sender
69    pub(crate) tx: Option<Sender<Event>>,
70    //Stream
71    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    /// Create a new BaseAgent wrapping an AgentDeriveT implementation
83    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        //Run Hook
110        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    /// Get the agent's name
120    pub fn name(&self) -> &str {
121        self.inner.name()
122    }
123
124    /// Get the agent's description
125    pub fn description(&self) -> &str {
126        self.inner.description()
127    }
128
129    /// Get the tools as Arc-wrapped references
130    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    /// Get the LLM provider
170    pub fn llm(&self) -> Arc<dyn LLMProvider> {
171        self.llm.clone()
172    }
173
174    /// Get the memory provider if available
175    pub fn memory(&self) -> Option<Arc<Mutex<Box<dyn MemoryProvider>>>> {
176        self.memory.clone()
177    }
178
179    /// Clone handle-style fields without requiring `T: Clone`.
180    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    /// Emit `TaskComplete`, run `on_run_complete`, and return the agent output.
194    ///
195    /// Uses `Value: From<AgentExecutor::Output>` so `TaskComplete` serialization matches
196    /// the executor payload regardless of `AgentDeriveT::Output`.
197    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}