autoagents_core/agent/
context.rs

1#[cfg(not(target_arch = "wasm32"))]
2use crate::actor::{ActorMessage, Topic};
3use crate::agent::AgentConfig;
4use crate::agent::memory::MemoryProvider;
5use crate::agent::state::AgentState;
6use crate::protocol::Event;
7use crate::tool::ToolT;
8use autoagents_llm::LLMProvider;
9use autoagents_llm::chat::ChatMessage;
10use std::any::Any;
11use std::sync::Arc;
12#[cfg(not(target_arch = "wasm32"))]
13use tokio::sync::{Mutex, mpsc};
14
15#[cfg(target_arch = "wasm32")]
16use futures::channel::mpsc;
17#[cfg(target_arch = "wasm32")]
18use futures::lock::Mutex;
19
20/// Execution context shared across an agent run.
21///
22/// Holds the configured LLM provider, accumulated chat messages, optional
23/// memory and tools, agent configuration, ephemeral execution state, and
24/// an optional event transmitter used by actor-based agents to emit protocol
25/// `Event`s. Also carries a `stream` flag to indicate streaming mode.
26pub struct Context {
27    llm: Arc<dyn LLMProvider>,
28    messages: Vec<ChatMessage>,
29    memory: Option<Arc<Mutex<Box<dyn MemoryProvider>>>>,
30    tools: Vec<Box<dyn ToolT>>,
31    config: AgentConfig,
32    state: Arc<Mutex<AgentState>>,
33    tx: Option<mpsc::Sender<Event>>,
34    stream: bool,
35}
36
37#[derive(Clone, Debug, thiserror::Error)]
38pub enum ContextError {
39    #[error("Tx value is None, Tx is only set for Actor agents")]
40    EmptyTx,
41    /// Error when sending events
42    #[error("Failed to send event: {0}")]
43    EventSendError(String),
44}
45
46impl Context {
47    pub fn new(llm: Arc<dyn LLMProvider>, tx: Option<mpsc::Sender<Event>>) -> Self {
48        Self {
49            llm,
50            messages: vec![],
51            memory: None,
52            tools: vec![],
53            config: AgentConfig::default(),
54            state: Arc::new(Mutex::new(AgentState::new())),
55            stream: false,
56            tx,
57        }
58    }
59
60    #[cfg(not(target_arch = "wasm32"))]
61    pub async fn publish<M: ActorMessage>(
62        &self,
63        topic: Topic<M>,
64        message: M,
65    ) -> Result<(), ContextError> {
66        self.tx
67            .as_ref()
68            .ok_or(ContextError::EmptyTx)?
69            .send(Event::PublishMessage {
70                topic_name: topic.name().to_string(),
71                message: Arc::new(message) as Arc<dyn Any + Send + Sync>,
72                topic_type: topic.type_id(),
73            })
74            .await
75            .map_err(|e| ContextError::EventSendError(e.to_string()))
76    }
77
78    pub fn with_memory(mut self, memory: Option<Arc<Mutex<Box<dyn MemoryProvider>>>>) -> Self {
79        self.memory = memory;
80        self
81    }
82
83    pub fn with_tools(mut self, tools: Vec<Box<dyn ToolT>>) -> Self {
84        self.tools = tools;
85        self
86    }
87
88    pub fn with_config(mut self, config: AgentConfig) -> Self {
89        self.config = config;
90        self
91    }
92
93    pub fn with_messages(mut self, messages: Vec<ChatMessage>) -> Self {
94        self.messages = messages;
95        self
96    }
97
98    pub fn with_stream(mut self, stream: bool) -> Self {
99        self.stream = stream;
100        self
101    }
102
103    // Getters
104    pub fn llm(&self) -> &Arc<dyn LLMProvider> {
105        &self.llm
106    }
107
108    pub fn messages(&self) -> &[ChatMessage] {
109        &self.messages
110    }
111
112    pub fn memory(&self) -> Option<Arc<Mutex<Box<dyn MemoryProvider>>>> {
113        self.memory.clone()
114    }
115
116    pub fn tools(&self) -> &[Box<dyn ToolT>] {
117        &self.tools
118    }
119
120    pub fn config(&self) -> &AgentConfig {
121        &self.config
122    }
123
124    pub fn state(&self) -> Arc<Mutex<AgentState>> {
125        self.state.clone()
126    }
127
128    /// Get a clone of the event transmitter (only present for Actor agents).
129    pub fn tx(&self) -> Result<mpsc::Sender<Event>, ContextError> {
130        Ok(self.tx.as_ref().ok_or(ContextError::EmptyTx)?.clone())
131    }
132
133    pub fn stream(&self) -> bool {
134        self.stream
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141    use crate::agent::memory::SlidingWindowMemory;
142    use autoagents_llm::chat::{ChatMessage, ChatMessageBuilder, ChatRole};
143    use autoagents_test_utils::llm::MockLLMProvider;
144    use std::sync::Arc;
145
146    #[test]
147    fn test_context_creation() {
148        let llm = Arc::new(MockLLMProvider);
149        let context = Context::new(llm, None);
150
151        assert!(context.messages.is_empty());
152        assert!(context.memory.is_none());
153        assert!(context.tools.is_empty());
154        assert!(!context.stream);
155    }
156
157    #[test]
158    fn test_context_with_llm_provider() {
159        let llm = Arc::new(MockLLMProvider);
160        let context = Context::new(llm.clone(), None);
161
162        // Verify the LLM provider is set correctly
163        let context_llm = context.llm();
164        assert!(Arc::strong_count(context_llm) > 0);
165    }
166
167    #[test]
168    fn test_context_with_memory() {
169        let llm = Arc::new(MockLLMProvider);
170        let memory = Box::new(SlidingWindowMemory::new(5));
171        let context = Context::new(llm, None).with_memory(Some(Arc::new(Mutex::new(memory))));
172
173        assert!(context.memory().is_some());
174    }
175
176    #[test]
177    fn test_context_with_messages() {
178        let llm = Arc::new(MockLLMProvider);
179        let message = ChatMessage::user().content("Hello".to_string()).build();
180        let context = Context::new(llm, None).with_messages(vec![message]);
181
182        assert_eq!(context.messages().len(), 1);
183        assert_eq!(context.messages()[0].role, ChatRole::User);
184        assert_eq!(context.messages()[0].content, "Hello");
185    }
186
187    #[test]
188    fn test_context_streaming_flag() {
189        let llm = Arc::new(MockLLMProvider);
190        let context = Context::new(llm, None).with_stream(true);
191        assert!(context.stream());
192    }
193
194    #[test]
195    fn test_context_fluent_interface() {
196        let llm = Arc::new(MockLLMProvider);
197        let memory = Box::new(SlidingWindowMemory::new(3));
198        let message = ChatMessageBuilder::new(ChatRole::System)
199            .content("System prompt".to_string())
200            .build();
201
202        let context = Context::new(llm, None)
203            .with_memory(Some(Arc::new(Mutex::new(memory))))
204            .with_messages(vec![message])
205            .with_stream(true);
206
207        assert!(context.memory().is_some());
208        assert_eq!(context.messages().len(), 1);
209        assert!(context.stream());
210    }
211}
autoagents_core/agent/context.rs

autoagents_core/agent/
context.rs
