anda_core/agent.rs
1//! Module providing core agent functionality for AI systems.
2//!
3//! This module defines the core traits and structures for creating and managing AI agents. It provides:
4//! - The [`Agent`] trait for defining custom agents with specific capabilities.
5//! - Dynamic dispatch capabilities through [`AgentDyn`] trait.
6//! - An [`AgentSet`] collection for managing multiple agents.
7//!
8//! # Key Features
9//! - Type-safe agent definitions with clear interfaces.
10//! - Asynchronous execution model.
11//! - Dynamic dispatch support for runtime agent selection.
12//! - Agent registration and management system.
13//! - Tool dependency management.
14//!
15//! # Architecture Overview
16//! The module follows a dual-trait pattern:
17//! 1. [`Agent`] - Static trait for defining concrete agent implementations.
18//! 2. [`AgentDyn`] - Dynamic trait for runtime polymorphism.
19//!
20//! The [`AgentSet`] acts as a registry and execution manager for agents, providing:
21//! - Agent registration and lookup.
22//! - Bulk definition retrieval.
23//! - Execution routing.
24//!
25//! # Usage
26//!
27//! ## Reference Implementations
28//! 1. [`Extractor`](https://github.com/ldclabs/anda/blob/main/anda_engine/src/extension/extractor.rs) -
29//! An agent for structured data extraction using LLMs
30//! 2. [`DocumentSegmenter`](https://github.com/ldclabs/anda/blob/main/anda_engine/src/extension/segmenter.rs) -
31//! A document segmentation tool using LLMs
32//! 3. [`CharacterAgent`](https://github.com/ldclabs/anda/blob/main/anda_engine/src/extension/character.rs) -
33//! A role-playing AI agent, also serving as the core agent for [`anda_bot`](https://github.com/ldclabs/anda/blob/main/agents/anda_bot/README.md)
34
35use serde::{Deserialize, Serialize};
36use serde_json::json;
37use std::{collections::BTreeMap, future::Future, marker::PhantomData, sync::Arc};
38
39use crate::{
40 BoxError, BoxPinFut, Function,
41 context::AgentContext,
42 model::{AgentOutput, FunctionDefinition, Resource},
43 select_resources, validate_function_name,
44};
45
46/// Arguments for an AI agent.
47#[derive(Debug, Clone, Deserialize, Serialize)]
48pub struct AgentArgs {
49 /// optimized prompt or message.
50 pub prompt: String,
51}
52
53/// Core trait defining an AI agent's behavior.
54///
55/// # Type Parameters
56/// - `C`: The context type that implements [`AgentContext`], must be thread-safe and have a static lifetime.
57pub trait Agent<C>: Send + Sync
58where
59 C: AgentContext + Send + Sync,
60{
61 /// Returns the agent's name as a String
62 /// The unique name of the agent, case-insensitive, must follow these rules in lowercase:
63 ///
64 /// # Rules
65 /// - Must not be empty;
66 /// - Must not exceed 64 characters;
67 /// - Must start with a lowercase letter;
68 /// - Can only contain: lowercase letters (a-z), digits (0-9), and underscores (_);
69 /// - Unique within the engine in lowercase.
70 fn name(&self) -> String;
71
72 /// Returns the agent's capabilities description in a short string.
73 fn description(&self) -> String;
74
75 /// Returns the agent's function definition for API integration.
76 ///
77 /// # Returns
78 /// - `FunctionDefinition`: The structured definition of the agent's capabilities.
79 fn definition(&self) -> FunctionDefinition {
80 FunctionDefinition {
81 name: self.name().to_ascii_lowercase(),
82 description: self.description(),
83 parameters: json!({
84 "type": "object",
85 "properties": {
86 "prompt": {"type": "string", "description": "optimized prompt or message."},
87 },
88 "required": ["prompt"],
89 }),
90 strict: None,
91 }
92 }
93
94 /// It is used to select resources based on the provided tags.
95 /// If the agent requires specific resources, it can filter them based on the tags.
96 /// By default, it returns an empty list.
97 ///
98 /// # Returns
99 /// - A list of resource tags from the tags provided that supported by the agent.
100 fn supported_resource_tags(&self) -> Vec<String> {
101 Vec::new()
102 }
103
104 /// Selects resources based on the agent's supported tags.
105 fn select_resources(&self, resources: &mut Vec<Resource>) -> Vec<Resource> {
106 let supported_tags = self.supported_resource_tags();
107 select_resources(resources, &supported_tags)
108 }
109
110 /// Initializes the tool with the given context.
111 /// It will be called once when building the Anda engine.
112 fn init(&self, _ctx: C) -> impl Future<Output = Result<(), BoxError>> + Send {
113 futures::future::ready(Ok(()))
114 }
115
116 /// Returns a list of tool dependencies required by the agent.
117 /// The tool dependencies are checked when building the engine.
118 fn tool_dependencies(&self) -> Vec<String> {
119 Vec::new()
120 }
121
122 /// Executes the agent's main logic with given context and inputs.
123 ///
124 /// # Arguments
125 /// - `ctx`: The execution context implementing [`AgentContext`].
126 /// - `prompt`: The input prompt or message for the agent.
127 /// - `resources`: Optional additional resources, If resources don’t meet the agent’s expectations, ignore them.
128 ///
129 /// # Returns
130 /// - A future resolving to Result<[`AgentOutput`], BoxError>.
131 fn run(
132 &self,
133 ctx: C,
134 prompt: String,
135 resources: Vec<Resource>,
136 ) -> impl Future<Output = Result<AgentOutput, BoxError>> + Send;
137}
138
139/// Dynamic dispatch version of Agent trait for runtime flexibility.
140///
141/// This trait allows for runtime polymorphism of agents, enabling dynamic agent selection
142/// and execution without knowing the concrete type at compile time.
143pub trait AgentDyn<C>: Send + Sync
144where
145 C: AgentContext + Send + Sync,
146{
147 fn label(&self) -> &str;
148
149 fn name(&self) -> String;
150
151 fn definition(&self) -> FunctionDefinition;
152
153 fn tool_dependencies(&self) -> Vec<String>;
154
155 fn supported_resource_tags(&self) -> Vec<String>;
156
157 fn init(&self, ctx: C) -> BoxPinFut<Result<(), BoxError>>;
158
159 fn run(
160 &self,
161 ctx: C,
162 prompt: String,
163 resources: Vec<Resource>,
164 ) -> BoxPinFut<Result<AgentOutput, BoxError>>;
165}
166
167/// Adapter for converting static Agent to dynamic dispatch.
168struct AgentWrapper<T, C>
169where
170 T: Agent<C> + 'static,
171 C: AgentContext + Send + Sync + 'static,
172{
173 inner: Arc<T>,
174 label: String,
175 _phantom: PhantomData<C>,
176}
177
178impl<T, C> AgentDyn<C> for AgentWrapper<T, C>
179where
180 T: Agent<C> + 'static,
181 C: AgentContext + Send + Sync + 'static,
182{
183 fn label(&self) -> &str {
184 &self.label
185 }
186
187 fn name(&self) -> String {
188 self.inner.name()
189 }
190
191 fn definition(&self) -> FunctionDefinition {
192 self.inner.definition()
193 }
194
195 fn tool_dependencies(&self) -> Vec<String> {
196 self.inner.tool_dependencies()
197 }
198
199 fn supported_resource_tags(&self) -> Vec<String> {
200 self.inner.supported_resource_tags()
201 }
202
203 fn init(&self, ctx: C) -> BoxPinFut<Result<(), BoxError>> {
204 let agent = self.inner.clone();
205 Box::pin(async move { agent.init(ctx).await })
206 }
207
208 fn run(
209 &self,
210 ctx: C,
211 prompt: String,
212 resources: Vec<Resource>,
213 ) -> BoxPinFut<Result<AgentOutput, BoxError>> {
214 let agent = self.inner.clone();
215 Box::pin(async move { agent.run(ctx, prompt, resources).await })
216 }
217}
218
219/// Collection of agents with lookup capabilities.
220///
221/// # Type Parameters
222/// - `C`: The context type that implements [`AgentContext`].
223#[derive(Default)]
224pub struct AgentSet<C: AgentContext> {
225 pub set: BTreeMap<String, Box<dyn AgentDyn<C>>>,
226}
227
228impl<C> AgentSet<C>
229where
230 C: AgentContext + Send + Sync + 'static,
231{
232 /// Creates a new empty AgentSet.
233 pub fn new() -> Self {
234 Self {
235 set: BTreeMap::new(),
236 }
237 }
238
239 /// Checks if an agent with given name exists.
240 pub fn contains(&self, name: &str) -> bool {
241 self.set.contains_key(&name.to_ascii_lowercase())
242 }
243
244 /// Returns the names of all agents in the set.
245 pub fn names(&self) -> Vec<String> {
246 self.set.keys().cloned().collect()
247 }
248
249 /// Retrieves definition for a specific agent.
250 pub fn definition(&self, name: &str) -> Option<FunctionDefinition> {
251 self.set
252 .get(&name.to_ascii_lowercase())
253 .map(|agent| agent.definition())
254 }
255
256 /// Returns definitions for all or specified agents.
257 ///
258 /// # Arguments
259 /// - `names`: Optional slice of agent names to filter by.
260 ///
261 /// # Returns
262 /// - Vec<[`FunctionDefinition`]>: Vector of agent definitions.
263 pub fn definitions(&self, names: Option<&[&str]>) -> Vec<FunctionDefinition> {
264 let names: Option<Vec<String>> =
265 names.map(|names| names.iter().map(|n| n.to_ascii_lowercase()).collect());
266 self.set
267 .iter()
268 .filter_map(|(name, agent)| match &names {
269 Some(names) => {
270 if names.contains(name) {
271 Some(agent.definition())
272 } else {
273 None
274 }
275 }
276 None => Some(agent.definition()),
277 })
278 .collect()
279 }
280
281 /// Returns a list of functions for all or specified agents.
282 ///
283 /// # Arguments
284 /// - `names`: Optional slice of agent names to filter by.
285 ///
286 /// # Returns
287 /// - Vec<[`Function`]>: Vector of agent functions.
288 pub fn functions(&self, names: Option<&[&str]>) -> Vec<Function> {
289 let names: Option<Vec<String>> =
290 names.map(|names| names.iter().map(|n| n.to_ascii_lowercase()).collect());
291 self.set
292 .iter()
293 .filter_map(|(name, agent)| match &names {
294 Some(names) => {
295 if names.contains(name) {
296 Some(Function {
297 definition: agent.definition(),
298 supported_resource_tags: agent.supported_resource_tags(),
299 })
300 } else {
301 None
302 }
303 }
304 None => Some(Function {
305 definition: agent.definition(),
306 supported_resource_tags: agent.supported_resource_tags(),
307 }),
308 })
309 .collect()
310 }
311
312 /// Extracts resources from the provided list based on the agent's supported tags.
313 pub fn select_resources(&self, name: &str, resources: &mut Vec<Resource>) -> Vec<Resource> {
314 if resources.is_empty() {
315 return Vec::new();
316 }
317
318 self.set
319 .get(&name.to_ascii_lowercase())
320 .map(|agent| {
321 let supported_tags = agent.supported_resource_tags();
322 select_resources(resources, &supported_tags)
323 })
324 .unwrap_or_default()
325 }
326
327 /// Registers a new agent in the set.
328 ///
329 /// # Arguments
330 /// - `agent`: The agent to register, must implement [`Agent`] trait.
331 pub fn add<T>(&mut self, agent: T, label: Option<String>) -> Result<(), BoxError>
332 where
333 T: Agent<C> + Send + Sync + 'static,
334 {
335 let name = agent.name().to_ascii_lowercase();
336 if self.set.contains_key(&name) {
337 return Err(format!("agent {} already exists", name).into());
338 }
339
340 validate_function_name(&name)?;
341 let agent_dyn = AgentWrapper {
342 inner: Arc::new(agent),
343 label: label.unwrap_or_else(|| name.clone()),
344 _phantom: PhantomData,
345 };
346 self.set.insert(name, Box::new(agent_dyn));
347 Ok(())
348 }
349
350 /// Retrieves an agent by name.
351 pub fn get(&self, name: &str) -> Option<&dyn AgentDyn<C>> {
352 self.set.get(&name.to_ascii_lowercase()).map(|v| &**v)
353 }
354}