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::tool::{ToolT, to_llm_tool};
7use autoagents_llm::LLMProvider;
8use autoagents_llm::chat::{ChatMessage, Tool};
9use autoagents_protocol::Event;
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    serialized_tools: Option<Arc<Vec<Tool>>>,
32    config: AgentConfig,
33    state: Arc<Mutex<AgentState>>,
34    tx: Option<mpsc::Sender<Event>>,
35    stream: bool,
36}
37
38#[derive(Clone, Debug, thiserror::Error)]
39pub enum ContextError {
40    #[error("Tx value is None, Tx is only set for Actor agents")]
41    EmptyTx,
42    /// Error when sending events
43    #[error("Failed to send event: {0}")]
44    EventSendError(String),
45}
46
47impl Context {
48    pub fn new(llm: Arc<dyn LLMProvider>, tx: Option<mpsc::Sender<Event>>) -> Self {
49        Self {
50            llm,
51            messages: vec![],
52            memory: None,
53            tools: vec![],
54            serialized_tools: None,
55            config: AgentConfig::default(),
56            state: Arc::new(Mutex::new(AgentState::new())),
57            stream: false,
58            tx,
59        }
60    }
61
62    #[cfg(not(target_arch = "wasm32"))]
63    pub async fn publish<M: ActorMessage>(
64        &self,
65        topic: Topic<M>,
66        message: M,
67    ) -> Result<(), ContextError> {
68        self.tx
69            .as_ref()
70            .ok_or(ContextError::EmptyTx)?
71            .send(Event::PublishMessage {
72                topic_name: topic.name().to_string(),
73                message: Arc::new(message) as Arc<dyn Any + Send + Sync>,
74                topic_type: topic.type_id(),
75            })
76            .await
77            .map_err(|e| ContextError::EventSendError(e.to_string()))
78    }
79
80    pub fn with_memory(mut self, memory: Option<Arc<Mutex<Box<dyn MemoryProvider>>>>) -> Self {
81        self.memory = memory;
82        self
83    }
84
85    pub fn with_tools(mut self, tools: Vec<Box<dyn ToolT>>) -> Self {
86        if tools.is_empty() {
87            self.serialized_tools = None;
88        } else if self.serialized_tools.is_none() {
89            let serialized = tools.iter().map(to_llm_tool).collect::<Vec<_>>();
90            self.serialized_tools = Some(Arc::new(serialized));
91        }
92        self.tools = tools;
93        self
94    }
95
96    pub fn with_serialized_tools(mut self, tools: Option<Arc<Vec<Tool>>>) -> Self {
97        self.serialized_tools = tools;
98        self
99    }
100
101    pub fn with_config(mut self, config: AgentConfig) -> Self {
102        self.config = config;
103        self
104    }
105
106    pub fn with_messages(mut self, messages: Vec<ChatMessage>) -> Self {
107        self.messages = messages;
108        self
109    }
110
111    pub fn with_stream(mut self, stream: bool) -> Self {
112        self.stream = stream;
113        self
114    }
115
116    // Getters
117    pub fn llm(&self) -> &Arc<dyn LLMProvider> {
118        &self.llm
119    }
120
121    pub fn messages(&self) -> &[ChatMessage] {
122        &self.messages
123    }
124
125    pub fn memory(&self) -> Option<Arc<Mutex<Box<dyn MemoryProvider>>>> {
126        self.memory.clone()
127    }
128
129    pub fn tools(&self) -> &[Box<dyn ToolT>] {
130        &self.tools
131    }
132
133    pub fn serialized_tools(&self) -> Option<Arc<Vec<Tool>>> {
134        self.serialized_tools.clone()
135    }
136
137    pub fn config(&self) -> &AgentConfig {
138        &self.config
139    }
140
141    pub fn state(&self) -> Arc<Mutex<AgentState>> {
142        self.state.clone()
143    }
144
145    /// Get a clone of the event transmitter (only present for Actor agents).
146    pub fn tx(&self) -> Result<mpsc::Sender<Event>, ContextError> {
147        Ok(self.tx.as_ref().ok_or(ContextError::EmptyTx)?.clone())
148    }
149
150    pub fn stream(&self) -> bool {
151        self.stream
152    }
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158    use crate::agent::memory::SlidingWindowMemory;
159    use crate::tests::MockLLMProvider;
160    use autoagents_llm::chat::{ChatMessage, ChatMessageBuilder, ChatRole};
161    use std::sync::Arc;
162
163    #[test]
164    fn test_context_with_memory() {
165        let llm = Arc::new(MockLLMProvider);
166        let memory = Box::new(SlidingWindowMemory::new(5));
167        let message = ChatMessage::user().content("Hello".to_string()).build();
168        let system_message = ChatMessageBuilder::new(ChatRole::System)
169            .content("System prompt".to_string())
170            .build();
171
172        let context = Context::new(llm, None)
173            .with_memory(Some(Arc::new(Mutex::new(memory))))
174            .with_messages(vec![message, system_message])
175            .with_stream(true);
176
177        assert!(context.memory().is_some());
178        assert_eq!(context.messages().len(), 2);
179        assert_eq!(context.messages()[0].role, ChatRole::User);
180        assert_eq!(context.messages()[0].content, "Hello");
181        assert!(context.stream());
182    }
183
184    #[test]
185    fn test_context_tx_missing_returns_error() {
186        let llm = Arc::new(MockLLMProvider);
187        let context = Context::new(llm, None);
188        let err = context.tx().unwrap_err();
189        assert!(matches!(err, ContextError::EmptyTx));
190    }
191}
autoagents_core/agent/context.rs

autoagents_core/agent/
context.rs
