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}