1use crate::{InvocationContext, Result, event::Event};
2use async_trait::async_trait;
3use futures::stream::Stream;
4use std::pin::Pin;
5use std::sync::Arc;
6
7pub type EventStream = Pin<Box<dyn Stream<Item = Result<Event>> + Send>>;
9
10#[async_trait]
17pub trait Agent: Send + Sync {
18 fn name(&self) -> &str;
20 fn description(&self) -> &str;
22 fn sub_agents(&self) -> &[Arc<dyn Agent>];
24
25 async fn run(&self, ctx: Arc<dyn InvocationContext>) -> Result<EventStream>;
27}
28
29#[derive(Clone)]
40pub struct ResolvedContext {
41 pub system_instruction: String,
43 pub active_tools: Vec<Arc<dyn crate::Tool>>,
45}
46
47impl std::fmt::Debug for ResolvedContext {
48 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
49 f.debug_struct("ResolvedContext")
50 .field("system_instruction_len", &self.system_instruction.len())
51 .field("active_tools_count", &self.active_tools.len())
52 .finish()
53 }
54}
55
56#[cfg(test)]
57mod tests {
58 use super::*;
59 use crate::{Content, ReadonlyContext, RunConfig};
60 use async_stream::stream;
61
62 struct TestAgent {
63 name: String,
64 }
65
66 use crate::{CallbackContext, Session, State};
67 use std::collections::HashMap;
68
69 struct MockState;
70 impl State for MockState {
71 fn get(&self, _key: &str) -> Option<serde_json::Value> {
72 None
73 }
74 fn set(&mut self, _key: String, _value: serde_json::Value) {}
75 fn all(&self) -> HashMap<String, serde_json::Value> {
76 HashMap::new()
77 }
78 }
79
80 struct MockSession;
81 impl Session for MockSession {
82 fn id(&self) -> &str {
83 "session"
84 }
85 fn app_name(&self) -> &str {
86 "app"
87 }
88 fn user_id(&self) -> &str {
89 "user"
90 }
91 fn state(&self) -> &dyn State {
92 &MockState
93 }
94 fn conversation_history(&self) -> Vec<Content> {
95 Vec::new()
96 }
97 }
98
99 #[allow(dead_code)]
100 struct TestContext {
101 content: Content,
102 config: RunConfig,
103 session: MockSession,
104 }
105
106 #[allow(dead_code)]
107 impl TestContext {
108 fn new() -> Self {
109 Self {
110 content: Content::new("user"),
111 config: RunConfig::default(),
112 session: MockSession,
113 }
114 }
115 }
116
117 #[async_trait]
118 impl ReadonlyContext for TestContext {
119 fn invocation_id(&self) -> &str {
120 "test"
121 }
122 fn agent_name(&self) -> &str {
123 "test"
124 }
125 fn user_id(&self) -> &str {
126 "user"
127 }
128 fn app_name(&self) -> &str {
129 "app"
130 }
131 fn session_id(&self) -> &str {
132 "session"
133 }
134 fn branch(&self) -> &str {
135 ""
136 }
137 fn user_content(&self) -> &Content {
138 &self.content
139 }
140 }
141
142 #[async_trait]
143 impl CallbackContext for TestContext {
144 fn artifacts(&self) -> Option<Arc<dyn crate::Artifacts>> {
145 None
146 }
147 }
148
149 #[async_trait]
150 impl InvocationContext for TestContext {
151 fn agent(&self) -> Arc<dyn Agent> {
152 unimplemented!()
153 }
154 fn memory(&self) -> Option<Arc<dyn crate::Memory>> {
155 None
156 }
157 fn session(&self) -> &dyn Session {
158 &self.session
159 }
160 fn run_config(&self) -> &RunConfig {
161 &self.config
162 }
163 fn end_invocation(&self) {}
164 fn ended(&self) -> bool {
165 false
166 }
167 }
168
169 #[async_trait]
170 impl Agent for TestAgent {
171 fn name(&self) -> &str {
172 &self.name
173 }
174
175 fn description(&self) -> &str {
176 "test agent"
177 }
178
179 fn sub_agents(&self) -> &[Arc<dyn Agent>] {
180 &[]
181 }
182
183 async fn run(&self, _ctx: Arc<dyn InvocationContext>) -> Result<EventStream> {
184 let s = stream! {
185 yield Ok(Event::new("test"));
186 };
187 Ok(Box::pin(s))
188 }
189 }
190
191 #[test]
192 fn test_agent_trait() {
193 let agent = TestAgent { name: "test".to_string() };
194 assert_eq!(agent.name(), "test");
195 assert_eq!(agent.description(), "test agent");
196 }
197}