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,
41 context::AgentContext,
42 model::{AgentOutput, FunctionDefinition, Resource},
43 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 /// Initializes the tool with the given context.
95 /// It will be called once when building the engine.
96 fn init(&self, _ctx: C) -> impl Future<Output = Result<(), BoxError>> + Send {
97 futures::future::ready(Ok(()))
98 }
99
100 /// Returns a list of tool dependencies required by the agent.
101 /// The tool dependencies are checked when building the engine.
102 fn tool_dependencies(&self) -> Vec<String> {
103 Vec::new()
104 }
105
106 /// Executes the agent's main logic with given context and inputs
107 ///
108 /// # Arguments
109 /// - `ctx`: The execution context implementing `AgentContext`
110 /// - `prompt`: The input prompt or message for the agent
111 /// - `resources`: Optional additional resources
112 ///
113 /// # Returns
114 /// - A future resolving to `Result<AgentOutput, BoxError>`
115 fn run(
116 &self,
117 ctx: C,
118 prompt: String,
119 resources: Option<Vec<Resource>>,
120 ) -> impl Future<Output = Result<AgentOutput, BoxError>> + Send;
121}
122
123/// Dynamic dispatch version of Agent trait for runtime flexibility
124///
125/// This trait allows for runtime polymorphism of agents, enabling dynamic agent selection
126/// and execution without knowing the concrete type at compile time.
127pub trait AgentDyn<C>: Send + Sync
128where
129 C: AgentContext + Send + Sync,
130{
131 fn name(&self) -> String;
132
133 fn definition(&self) -> FunctionDefinition;
134
135 fn tool_dependencies(&self) -> Vec<String>;
136
137 fn init(&self, ctx: C) -> BoxPinFut<Result<(), BoxError>>;
138
139 fn run(
140 &self,
141 ctx: C,
142 prompt: String,
143 resources: Option<Vec<Resource>>,
144 ) -> BoxPinFut<Result<AgentOutput, BoxError>>;
145}
146
147/// Adapter for converting static Agent to dynamic dispatch
148struct AgentWrapper<T, C>(Arc<T>, PhantomData<C>)
149where
150 T: Agent<C> + 'static,
151 C: AgentContext + Send + Sync + 'static;
152
153impl<T, C> AgentDyn<C> for AgentWrapper<T, C>
154where
155 T: Agent<C> + 'static,
156 C: AgentContext + Send + Sync + 'static,
157{
158 fn name(&self) -> String {
159 self.0.name()
160 }
161
162 fn definition(&self) -> FunctionDefinition {
163 self.0.definition()
164 }
165
166 fn tool_dependencies(&self) -> Vec<String> {
167 self.0.tool_dependencies()
168 }
169
170 fn init(&self, ctx: C) -> BoxPinFut<Result<(), BoxError>> {
171 let agent = self.0.clone();
172 Box::pin(async move { agent.init(ctx).await })
173 }
174
175 fn run(
176 &self,
177 ctx: C,
178 prompt: String,
179 resources: Option<Vec<Resource>>,
180 ) -> BoxPinFut<Result<AgentOutput, BoxError>> {
181 let agent = self.0.clone();
182 Box::pin(async move { agent.run(ctx, prompt, resources).await })
183 }
184}
185
186/// Collection of registered agents with lookup and execution capabilities
187///
188/// # Type Parameters
189/// - `C`: The context type that implements `AgentContext`
190#[derive(Default)]
191pub struct AgentSet<C: AgentContext> {
192 pub set: BTreeMap<String, Box<dyn AgentDyn<C>>>,
193}
194
195impl<C> AgentSet<C>
196where
197 C: AgentContext + Send + Sync + 'static,
198{
199 /// Creates a new empty AgentSet
200 pub fn new() -> Self {
201 Self {
202 set: BTreeMap::new(),
203 }
204 }
205
206 /// Checks if an agent with given name exists
207 pub fn contains(&self, name: &str) -> bool {
208 self.set.contains_key(&name.to_ascii_lowercase())
209 }
210
211 /// Retrieves definition for a specific agent
212 pub fn definition(&self, name: &str) -> Option<FunctionDefinition> {
213 self.set
214 .get(&name.to_ascii_lowercase())
215 .map(|agent| agent.definition())
216 }
217
218 /// Returns definitions for all or specified agents
219 ///
220 /// # Arguments
221 /// - `names`: Optional slice of agent names to filter by
222 ///
223 /// # Returns
224 /// - `Vec<FunctionDefinition>`: Vector of agent definitions
225 pub fn definitions(&self, names: Option<&[&str]>) -> Vec<FunctionDefinition> {
226 let names: Option<Vec<String>> =
227 names.map(|names| names.iter().map(|n| n.to_ascii_lowercase()).collect());
228 self.set
229 .iter()
230 .filter_map(|(name, agent)| match &names {
231 Some(names) => {
232 if names.contains(name) {
233 Some(agent.definition())
234 } else {
235 None
236 }
237 }
238 None => Some(agent.definition()),
239 })
240 .collect()
241 }
242
243 /// Registers a new agent in the set
244 ///
245 /// # Arguments
246 /// - `agent`: The agent to register, must implement `Agent<C>`
247 pub fn add<T>(&mut self, agent: T) -> Result<(), BoxError>
248 where
249 T: Agent<C> + Send + Sync + 'static,
250 {
251 let name = agent.name().to_ascii_lowercase();
252 if self.set.contains_key(&name) {
253 return Err(format!("agent {} already exists", name).into());
254 }
255
256 validate_function_name(&name)?;
257 let agent_dyn = AgentWrapper(Arc::new(agent), PhantomData);
258 self.set.insert(name, Box::new(agent_dyn));
259 Ok(())
260 }
261
262 /// Retrieves an agent by name
263 pub fn get(&self, name: &str) -> Option<&Box<dyn AgentDyn<C>>> {
264 self.set.get(name)
265 }
266}