claude_agent_sdk/orchestration/
agent.rs

1//! # Agent trait and core types
2//!
3//! This module defines the core Agent trait and associated types for the
4//! multi-agent orchestration framework.
5
6use async_trait::async_trait;
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9
10/// Error type for agent operations
11#[derive(Debug, thiserror::Error)]
12pub enum AgentError {
13    #[error("Execution failed: {0}")]
14    ExecutionFailed(String),
15
16    #[error("Invalid input: {0}")]
17    InvalidInput(String),
18
19    #[error("Timeout")]
20    Timeout,
21
22    #[error("Other error: {0}")]
23    Other(#[from] anyhow::Error),
24}
25
26/// Result type for agent operations
27pub type Result<T> = std::result::Result<T, AgentError>;
28
29/// Input to an agent
30#[derive(Debug, Clone, Serialize, Deserialize)]
31pub struct AgentInput {
32    /// Main content/prompt for the agent
33    pub content: String,
34
35    /// Additional context data (JSON-serializable)
36    #[serde(default)]
37    pub context: serde_json::Value,
38
39    /// Metadata key-value pairs
40    #[serde(default)]
41    pub metadata: HashMap<String, String>,
42}
43
44impl AgentInput {
45    /// Create a new agent input with content
46    pub fn new(content: impl Into<String>) -> Self {
47        Self {
48            content: content.into(),
49            context: serde_json::json!({}),
50            metadata: HashMap::new(),
51        }
52    }
53
54    /// Add context data
55    pub fn with_context(mut self, context: serde_json::Value) -> Self {
56        self.context = context;
57        self
58    }
59
60    /// Add metadata
61    pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
62        self.metadata.insert(key.into(), value.into());
63        self
64    }
65}
66
67/// Output from an agent
68#[derive(Debug, Clone, Serialize, Deserialize)]
69pub struct AgentOutput {
70    /// Main response content
71    pub content: String,
72
73    /// Additional data (JSON-serializable)
74    #[serde(default)]
75    pub data: serde_json::Value,
76
77    /// Confidence score (0.0 - 1.0)
78    pub confidence: f64,
79
80    /// Metadata key-value pairs
81    #[serde(default)]
82    pub metadata: HashMap<String, String>,
83}
84
85impl AgentOutput {
86    /// Create a new agent output with content
87    pub fn new(content: impl Into<String>) -> Self {
88        Self {
89            content: content.into(),
90            data: serde_json::json!({}),
91            confidence: 1.0,
92            metadata: HashMap::new(),
93        }
94    }
95
96    /// Set confidence score
97    pub fn with_confidence(mut self, confidence: f64) -> Self {
98        self.confidence = confidence.clamp(0.0, 1.0);
99        self
100    }
101
102    /// Add data
103    pub fn with_data(mut self, data: serde_json::Value) -> Self {
104        self.data = data;
105        self
106    }
107
108    /// Add metadata
109    pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
110        self.metadata.insert(key.into(), value.into());
111        self
112    }
113
114    /// Check if output is successful (confidence > 0.5)
115    pub fn is_successful(&self) -> bool {
116        self.confidence > 0.5
117    }
118}
119
120/// Core Agent trait
121///
122/// Agents implement this trait to participate in orchestration.
123/// Each agent has a name, description, and execution logic.
124#[async_trait]
125pub trait Agent: Send + Sync {
126    /// Agent name (must be unique)
127    fn name(&self) -> &str;
128
129    /// Agent description (what it does)
130    fn description(&self) -> &str;
131
132    /// Execute the agent's logic
133    ///
134    /// Takes input and produces output asynchronously.
135    /// Returns an error if execution fails.
136    async fn execute(&self, input: AgentInput) -> Result<AgentOutput>;
137}
138
139/// Simple wrapper agent for easy creation
140pub struct SimpleAgent<F>
141where
142    F: Fn(AgentInput) -> Result<AgentOutput> + Send + Sync,
143{
144    name: String,
145    description: String,
146    handler: F,
147}
148
149impl<F> SimpleAgent<F>
150where
151    F: Fn(AgentInput) -> Result<AgentOutput> + Send + Sync,
152{
153    /// Create a new simple agent
154    pub fn new(name: impl Into<String>, description: impl Into<String>, handler: F) -> Self {
155        Self {
156            name: name.into(),
157            description: description.into(),
158            handler,
159        }
160    }
161}
162
163#[async_trait]
164impl<F> Agent for SimpleAgent<F>
165where
166    F: Fn(AgentInput) -> Result<AgentOutput> + Send + Sync,
167{
168    fn name(&self) -> &str {
169        &self.name
170    }
171
172    fn description(&self) -> &str {
173        &self.description
174    }
175
176    async fn execute(&self, input: AgentInput) -> Result<AgentOutput> {
177        // Call the synchronous handler
178        (self.handler)(input)
179    }
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185
186    #[test]
187    fn test_agent_input_creation() {
188        let input = AgentInput::new("test content")
189            .with_context(serde_json::json!({"key": "value"}))
190            .with_metadata("meta1", "value1");
191
192        assert_eq!(input.content, "test content");
193        assert_eq!(input.context["key"], "value");
194        assert_eq!(input.metadata["meta1"], "value1");
195    }
196
197    #[test]
198    fn test_agent_output_creation() {
199        let output = AgentOutput::new("test response")
200            .with_confidence(0.8)
201            .with_data(serde_json::json!({"result": 42}))
202            .with_metadata("meta1", "value1");
203
204        assert_eq!(output.content, "test response");
205        assert_eq!(output.confidence, 0.8);
206        assert_eq!(output.data["result"], 42);
207        assert_eq!(output.metadata["meta1"], "value1");
208        assert!(output.is_successful());
209    }
210
211    #[test]
212    fn test_simple_agent() {
213        let agent = SimpleAgent::new("TestAgent", "A test agent", |input| {
214            Ok(AgentOutput::new(format!("Processed: {}", input.content)))
215        });
216
217        assert_eq!(agent.name(), "TestAgent");
218        assert_eq!(agent.description(), "A test agent");
219    }
220
221    #[tokio::test]
222    async fn test_simple_agent_execute() {
223        let agent = SimpleAgent::new("TestAgent", "A test agent", |input| {
224            Ok(AgentOutput::new(format!("Echo: {}", input.content)).with_confidence(0.9))
225        });
226
227        let input = AgentInput::new("Hello");
228        let output = agent.execute(input).await.unwrap();
229
230        assert_eq!(output.content, "Echo: Hello");
231        assert_eq!(output.confidence, 0.9);
232        assert!(output.is_successful());
233    }
234}
claude_agent_sdk/orchestration/agent.rs

claude_agent_sdk/orchestration/
agent.rs
