helios_engine/
agent.rs

1//! # Agent Module
2//!
3//! This module defines the `Agent` struct, the core of the Helios Engine. Agents are
4//! autonomous entities that can interact with users, use tools, and manage their
5//! own chat history. The `AgentBuilder` provides a convenient way to construct
6//! and configure agents.
7
8#![allow(dead_code)]
9#![allow(unused_variables)]
10use crate::chat::{ChatMessage, ChatSession};
11use crate::config::Config;
12use crate::error::{HeliosError, Result};
13use crate::llm::{LLMClient, LLMProviderType};
14use crate::tools::{ToolRegistry, ToolResult};
15use serde_json::Value;
16
17/// Prefix for agent-specific keys in the chat session metadata.
18const AGENT_MEMORY_PREFIX: &str = "agent:";
19
20/// Represents an LLM-powered agent that can chat, use tools, and manage a conversation.
21pub struct Agent {
22    /// The name of the agent.
23    name: String,
24    /// The client for interacting with the Large Language Model.
25    llm_client: LLMClient,
26    /// The registry of tools available to the agent.
27    tool_registry: ToolRegistry,
28    /// The chat session, which stores the conversation history.
29    chat_session: ChatSession,
30    /// The maximum number of iterations for tool execution in a single turn.
31    max_iterations: usize,
32}
33
34impl Agent {
35    /// Creates a new agent with the given name and configuration.
36    ///
37    /// # Arguments
38    ///
39    /// * `name` - The name of the agent.
40    /// * `config` - The configuration for the agent.
41    ///
42    /// # Returns
43    ///
44    /// A `Result` containing the new `Agent` instance.
45    async fn new(name: impl Into<String>, config: Config) -> Result<Self> {
46        #[cfg(feature = "local")]
47        let provider_type = if let Some(local_config) = config.local {
48            LLMProviderType::Local(local_config)
49        } else {
50            LLMProviderType::Remote(config.llm)
51        };
52
53        #[cfg(not(feature = "local"))]
54        let provider_type = LLMProviderType::Remote(config.llm);
55
56        let llm_client = LLMClient::new(provider_type).await?;
57
58        Ok(Self {
59            name: name.into(),
60            llm_client,
61            tool_registry: ToolRegistry::new(),
62            chat_session: ChatSession::new(),
63            max_iterations: 10,
64        })
65    }
66
67    /// Returns a new `AgentBuilder` for constructing an agent.
68    ///
69    /// # Arguments
70    ///
71    /// * `name` - The name of the agent.
72    pub fn builder(name: impl Into<String>) -> AgentBuilder {
73        AgentBuilder::new(name)
74    }
75
76    /// Returns the name of the agent.
77    pub fn name(&self) -> &str {
78        &self.name
79    }
80
81    /// Sets the system prompt for the agent.
82    ///
83    /// # Arguments
84    ///
85    /// * `prompt` - The system prompt to set.
86    pub fn set_system_prompt(&mut self, prompt: impl Into<String>) {
87        self.chat_session = self.chat_session.clone().with_system_prompt(prompt);
88    }
89
90    /// Registers a tool with the agent.
91    ///
92    /// # Arguments
93    ///
94    /// * `tool` - The tool to register.
95    pub fn register_tool(&mut self, tool: Box<dyn crate::tools::Tool>) {
96        self.tool_registry.register(tool);
97    }
98
99    /// Returns a reference to the agent's tool registry.
100    pub fn tool_registry(&self) -> &ToolRegistry {
101        &self.tool_registry
102    }
103
104    /// Returns a mutable reference to the agent's tool registry.
105    pub fn tool_registry_mut(&mut self) -> &mut ToolRegistry {
106        &mut self.tool_registry
107    }
108
109    /// Returns a reference to the agent's chat session.
110    pub fn chat_session(&self) -> &ChatSession {
111        &self.chat_session
112    }
113
114    /// Returns a mutable reference to the agent's chat session.
115    pub fn chat_session_mut(&mut self) -> &mut ChatSession {
116        &mut self.chat_session
117    }
118
119    /// Clears the agent's chat history.
120    pub fn clear_history(&mut self) {
121        self.chat_session.clear();
122    }
123
124    /// Sends a message to the agent and gets a response.
125    ///
126    /// # Arguments
127    ///
128    /// * `message` - The message to send.
129    ///
130    /// # Returns
131    ///
132    /// A `Result` containing the agent's response.
133    pub async fn send_message(&mut self, message: impl Into<String>) -> Result<String> {
134        let user_message = message.into();
135        self.chat_session.add_user_message(user_message.clone());
136
137        // Execute agent loop with tool calling
138        let response = self.execute_with_tools().await?;
139
140        Ok(response)
141    }
142
143    /// Executes the agent's main loop, including tool calls.
144    async fn execute_with_tools(&mut self) -> Result<String> {
145        self.execute_with_tools_streaming().await
146    }
147
148    /// Executes the agent's main loop with streaming, including tool calls.
149    async fn execute_with_tools_streaming(&mut self) -> Result<String> {
150        self.execute_with_tools_streaming_with_params(None, None, None)
151            .await
152    }
153
154    /// Executes the agent's main loop with parameters, including tool calls.
155    async fn execute_with_tools_with_params(
156        &mut self,
157        temperature: Option<f32>,
158        max_tokens: Option<u32>,
159        stop: Option<Vec<String>>,
160    ) -> Result<String> {
161        let mut iterations = 0;
162        let tool_definitions = self.tool_registry.get_definitions();
163
164        loop {
165            if iterations >= self.max_iterations {
166                return Err(HeliosError::AgentError(
167                    "Maximum iterations reached".to_string(),
168                ));
169            }
170
171            let messages = self.chat_session.get_messages();
172            let tools_option = if tool_definitions.is_empty() {
173                None
174            } else {
175                Some(tool_definitions.clone())
176            };
177
178            let response = self
179                .llm_client
180                .chat(
181                    messages,
182                    tools_option,
183                    temperature,
184                    max_tokens,
185                    stop.clone(),
186                )
187                .await?;
188
189            // Check if the response includes tool calls
190            if let Some(ref tool_calls) = response.tool_calls {
191                // Add assistant message with tool calls
192                self.chat_session.add_message(response.clone());
193
194                // Execute each tool call
195                for tool_call in tool_calls {
196                    let tool_name = &tool_call.function.name;
197                    let tool_args: Value = serde_json::from_str(&tool_call.function.arguments)
198                        .unwrap_or(Value::Object(serde_json::Map::new()));
199
200                    let tool_result = self
201                        .tool_registry
202                        .execute(tool_name, tool_args)
203                        .await
204                        .unwrap_or_else(|e| {
205                            ToolResult::error(format!("Tool execution failed: {}", e))
206                        });
207
208                    // Add tool result message
209                    let tool_message = ChatMessage::tool(tool_result.output, tool_call.id.clone());
210                    self.chat_session.add_message(tool_message);
211                }
212
213                iterations += 1;
214                continue;
215            }
216
217            // No tool calls, we have the final response
218            self.chat_session.add_message(response.clone());
219            return Ok(response.content);
220        }
221    }
222
223    /// Executes the agent's main loop with parameters and streaming, including tool calls.
224    async fn execute_with_tools_streaming_with_params(
225        &mut self,
226        temperature: Option<f32>,
227        max_tokens: Option<u32>,
228        stop: Option<Vec<String>>,
229    ) -> Result<String> {
230        let mut iterations = 0;
231        let tool_definitions = self.tool_registry.get_definitions();
232
233        loop {
234            if iterations >= self.max_iterations {
235                return Err(HeliosError::AgentError(
236                    "Maximum iterations reached".to_string(),
237                ));
238            }
239
240            let messages = self.chat_session.get_messages();
241            let tools_option = if tool_definitions.is_empty() {
242                None
243            } else {
244                Some(tool_definitions.clone())
245            };
246
247            let mut streamed_content = String::new();
248
249            let stream_result = self
250                .llm_client
251                .chat_stream(
252                    messages,
253                    tools_option, // Enable tools for streaming
254                    temperature,
255                    max_tokens,
256                    stop.clone(),
257                    |chunk| {
258                        // Print chunk to stdout for visible streaming
259                        print!("{}", chunk);
260                        let _ = std::io::Write::flush(&mut std::io::stdout());
261                        streamed_content.push_str(chunk);
262                    },
263                )
264                .await;
265
266            let response = stream_result?;
267
268            // Print newline after streaming completes
269            println!();
270
271            // Check if the response includes tool calls
272            if let Some(ref tool_calls) = response.tool_calls {
273                // Add assistant message with tool calls
274                let mut msg_with_content = response.clone();
275                msg_with_content.content = streamed_content.clone();
276                self.chat_session.add_message(msg_with_content);
277
278                // Execute each tool call
279                for tool_call in tool_calls {
280                    let tool_name = &tool_call.function.name;
281                    let tool_args: Value = serde_json::from_str(&tool_call.function.arguments)
282                        .unwrap_or(Value::Object(serde_json::Map::new()));
283
284                    let tool_result = self
285                        .tool_registry
286                        .execute(tool_name, tool_args)
287                        .await
288                        .unwrap_or_else(|e| {
289                            ToolResult::error(format!("Tool execution failed: {}", e))
290                        });
291
292                    // Add tool result message
293                    let tool_message = ChatMessage::tool(tool_result.output, tool_call.id.clone());
294                    self.chat_session.add_message(tool_message);
295                }
296
297                iterations += 1;
298                continue;
299            }
300
301            // No tool calls, we have the final response with streamed content
302            let mut final_msg = response;
303            final_msg.content = streamed_content.clone();
304            self.chat_session.add_message(final_msg);
305            return Ok(streamed_content);
306        }
307    }
308
309    /// A convenience method for sending a message to the agent.
310    pub async fn chat(&mut self, message: impl Into<String>) -> Result<String> {
311        self.send_message(message).await
312    }
313
314    /// Sets the maximum number of iterations for tool execution.
315    ///
316    /// # Arguments
317    ///
318    /// * `max` - The maximum number of iterations.
319    pub fn set_max_iterations(&mut self, max: usize) {
320        self.max_iterations = max;
321    }
322
323    /// Returns a summary of the current chat session.
324    pub fn get_session_summary(&self) -> String {
325        self.chat_session.get_summary()
326    }
327
328    /// Clears the agent's memory (agent-scoped metadata).
329    pub fn clear_memory(&mut self) {
330        // Only clear agent-scoped memory keys to avoid wiping general session metadata
331        self.chat_session
332            .metadata
333            .retain(|k, _| !k.starts_with(AGENT_MEMORY_PREFIX));
334    }
335
336    /// Prefixes a key with the agent memory prefix.
337    #[inline]
338    fn prefixed_key(key: &str) -> String {
339        format!("{}{}", AGENT_MEMORY_PREFIX, key)
340    }
341
342    // Agent-scoped memory API (namespaced under "agent:")
343    /// Sets a value in the agent's memory.
344    pub fn set_memory(&mut self, key: impl Into<String>, value: impl Into<String>) {
345        let key = key.into();
346        self.chat_session
347            .set_metadata(Self::prefixed_key(&key), value);
348    }
349
350    /// Gets a value from the agent's memory.
351    pub fn get_memory(&self, key: &str) -> Option<&String> {
352        self.chat_session.get_metadata(&Self::prefixed_key(key))
353    }
354
355    /// Removes a value from the agent's memory.
356    pub fn remove_memory(&mut self, key: &str) -> Option<String> {
357        self.chat_session.remove_metadata(&Self::prefixed_key(key))
358    }
359
360    // Convenience helpers to reduce duplication in examples
361    /// Increments a counter in the agent's memory.
362    pub fn increment_counter(&mut self, key: &str) -> u32 {
363        let current = self
364            .get_memory(key)
365            .and_then(|v| v.parse::<u32>().ok())
366            .unwrap_or(0);
367        let next = current + 1;
368        self.set_memory(key, next.to_string());
369        next
370    }
371
372    /// Increments the "tasks_completed" counter in the agent's memory.
373    pub fn increment_tasks_completed(&mut self) -> u32 {
374        self.increment_counter("tasks_completed")
375    }
376
377    /// Executes a stateless conversation with the provided message history.
378    ///
379    /// This method creates a temporary chat session with the provided messages
380    /// and executes the agent logic without modifying the agent's persistent session.
381    /// This is useful for OpenAI API compatibility where each request contains
382    /// the full conversation history.
383    ///
384    /// # Arguments
385    ///
386    /// * `messages` - The complete conversation history for this request
387    /// * `temperature` - Optional temperature parameter for generation
388    /// * `max_tokens` - Optional maximum tokens parameter for generation
389    /// * `stop` - Optional stop sequences for generation
390    ///
391    /// # Returns
392    ///
393    /// A `Result` containing the assistant's response content.
394    pub async fn chat_with_history(
395        &mut self,
396        messages: Vec<ChatMessage>,
397        temperature: Option<f32>,
398        max_tokens: Option<u32>,
399        stop: Option<Vec<String>>,
400    ) -> Result<String> {
401        // Create a temporary session with the provided messages
402        let mut temp_session = ChatSession::new();
403
404        // Add all messages to the temporary session
405        for message in messages {
406            temp_session.add_message(message);
407        }
408
409        // Execute agent loop with tool calling using the temporary session
410        self.execute_with_tools_temp_session(temp_session, temperature, max_tokens, stop)
411            .await
412    }
413
414    /// Executes the agent's main loop with a temporary session, including tool calls.
415    async fn execute_with_tools_temp_session(
416        &mut self,
417        mut temp_session: ChatSession,
418        temperature: Option<f32>,
419        max_tokens: Option<u32>,
420        stop: Option<Vec<String>>,
421    ) -> Result<String> {
422        let mut iterations = 0;
423        let tool_definitions = self.tool_registry.get_definitions();
424
425        loop {
426            if iterations >= self.max_iterations {
427                return Err(HeliosError::AgentError(
428                    "Maximum iterations reached".to_string(),
429                ));
430            }
431
432            let messages = temp_session.get_messages();
433            let tools_option = if tool_definitions.is_empty() {
434                None
435            } else {
436                Some(tool_definitions.clone())
437            };
438
439            let response = self
440                .llm_client
441                .chat(
442                    messages,
443                    tools_option,
444                    temperature,
445                    max_tokens,
446                    stop.clone(),
447                )
448                .await?;
449
450            // Check if the response includes tool calls
451            if let Some(ref tool_calls) = response.tool_calls {
452                // Add assistant message with tool calls to temp session
453                temp_session.add_message(response.clone());
454
455                // Execute each tool call
456                for tool_call in tool_calls {
457                    let tool_name = &tool_call.function.name;
458                    let tool_args: Value = serde_json::from_str(&tool_call.function.arguments)
459                        .unwrap_or(Value::Object(serde_json::Map::new()));
460
461                    let tool_result = self
462                        .tool_registry
463                        .execute(tool_name, tool_args)
464                        .await
465                        .unwrap_or_else(|e| {
466                            ToolResult::error(format!("Tool execution failed: {}", e))
467                        });
468
469                    // Add tool result message to temp session
470                    let tool_message = ChatMessage::tool(tool_result.output, tool_call.id.clone());
471                    temp_session.add_message(tool_message);
472                }
473
474                iterations += 1;
475                continue;
476            }
477
478            // No tool calls, we have the final response
479            return Ok(response.content);
480        }
481    }
482
483    /// Executes a stateless conversation with the provided message history and streams the response.
484    ///
485    /// This method creates a temporary chat session with the provided messages
486    /// and streams the agent's response in real-time as tokens are generated.
487    /// Note: Tool calls are not supported in streaming mode yet - they will be
488    /// executed after the initial response is complete.
489    ///
490    /// # Arguments
491    ///
492    /// * `messages` - The complete conversation history for this request
493    /// * `temperature` - Optional temperature parameter for generation
494    /// * `max_tokens` - Optional maximum tokens parameter for generation
495    /// * `stop` - Optional stop sequences for generation
496    /// * `on_chunk` - Callback function called for each chunk of generated text
497    ///
498    /// # Returns
499    ///
500    /// A `Result` containing the final assistant message after streaming is complete.
501    pub async fn chat_stream_with_history<F>(
502        &mut self,
503        messages: Vec<ChatMessage>,
504        temperature: Option<f32>,
505        max_tokens: Option<u32>,
506        stop: Option<Vec<String>>,
507        on_chunk: F,
508    ) -> Result<ChatMessage>
509    where
510        F: FnMut(&str) + Send,
511    {
512        // Create a temporary session with the provided messages
513        let mut temp_session = ChatSession::new();
514
515        // Add all messages to the temporary session
516        for message in messages {
517            temp_session.add_message(message);
518        }
519
520        // For now, use streaming for the initial response only
521        // Tool calls will be handled after the stream completes
522        self.execute_streaming_with_tools_temp_session(
523            temp_session,
524            temperature,
525            max_tokens,
526            stop,
527            on_chunk,
528        )
529        .await
530    }
531
532    /// Executes the agent's main loop with streaming and a temporary session.
533    async fn execute_streaming_with_tools_temp_session<F>(
534        &mut self,
535        mut temp_session: ChatSession,
536        temperature: Option<f32>,
537        max_tokens: Option<u32>,
538        stop: Option<Vec<String>>,
539        mut on_chunk: F,
540    ) -> Result<ChatMessage>
541    where
542        F: FnMut(&str) + Send,
543    {
544        let mut iterations = 0;
545        let tool_definitions = self.tool_registry.get_definitions();
546
547        loop {
548            if iterations >= self.max_iterations {
549                return Err(HeliosError::AgentError(
550                    "Maximum iterations reached".to_string(),
551                ));
552            }
553
554            let messages = temp_session.get_messages();
555            let tools_option = if tool_definitions.is_empty() {
556                None
557            } else {
558                Some(tool_definitions.clone())
559            };
560
561            // Use streaming for all iterations
562            let mut streamed_content = String::new();
563
564            let stream_result = self
565                .llm_client
566                .chat_stream(
567                    messages,
568                    tools_option,
569                    temperature,
570                    max_tokens,
571                    stop.clone(),
572                    |chunk| {
573                        on_chunk(chunk);
574                        streamed_content.push_str(chunk);
575                    },
576                )
577                .await;
578
579            match stream_result {
580                Ok(response) => {
581                    // Check if the response includes tool calls
582                    if let Some(ref tool_calls) = response.tool_calls {
583                        // Add assistant message with tool calls to temp session
584                        let mut msg_with_content = response.clone();
585                        msg_with_content.content = streamed_content.clone();
586                        temp_session.add_message(msg_with_content);
587
588                        // Execute each tool call
589                        for tool_call in tool_calls {
590                            let tool_name = &tool_call.function.name;
591                            let tool_args: Value =
592                                serde_json::from_str(&tool_call.function.arguments)
593                                    .unwrap_or(Value::Object(serde_json::Map::new()));
594
595                            let tool_result = self
596                                .tool_registry
597                                .execute(tool_name, tool_args)
598                                .await
599                                .unwrap_or_else(|e| {
600                                    ToolResult::error(format!("Tool execution failed: {}", e))
601                                });
602
603                            // Add tool result message to temp session
604                            let tool_message =
605                                ChatMessage::tool(tool_result.output, tool_call.id.clone());
606                            temp_session.add_message(tool_message);
607                        }
608
609                        iterations += 1;
610                        continue; // Continue the loop for another iteration
611                    } else {
612                        // No tool calls, return the final response with streamed content
613                        let mut final_msg = response;
614                        final_msg.content = streamed_content;
615                        return Ok(final_msg);
616                    }
617                }
618                Err(e) => return Err(e),
619            }
620        }
621    }
622}
623
624#[cfg(test)]
625mod tests {
626    use super::*;
627    use crate::config::Config;
628    use crate::tools::{CalculatorTool, Tool, ToolParameter, ToolResult};
629    use serde_json::Value;
630    use std::collections::HashMap;
631
632    /// Tests that an agent can be created using the builder.
633    #[tokio::test]
634    async fn test_agent_creation_via_builder() {
635        let config = Config::new_default();
636        let agent = Agent::builder("test_agent").config(config).build().await;
637        assert!(agent.is_ok());
638    }
639
640    /// Tests the namespacing of agent memory.
641    #[tokio::test]
642    async fn test_agent_memory_namespacing_set_get_remove() {
643        let config = Config::new_default();
644        let mut agent = Agent::builder("test_agent")
645            .config(config)
646            .build()
647            .await
648            .unwrap();
649
650        // Set and get namespaced memory
651        agent.set_memory("working_directory", "/tmp");
652        assert_eq!(
653            agent.get_memory("working_directory"),
654            Some(&"/tmp".to_string())
655        );
656
657        // Ensure underlying chat session stored the prefixed key
658        assert_eq!(
659            agent.chat_session().get_metadata("agent:working_directory"),
660            Some(&"/tmp".to_string())
661        );
662        // Non-prefixed key should not exist
663        assert!(agent
664            .chat_session()
665            .get_metadata("working_directory")
666            .is_none());
667
668        // Remove should also be namespaced
669        let removed = agent.remove_memory("working_directory");
670        assert_eq!(removed.as_deref(), Some("/tmp"));
671        assert!(agent.get_memory("working_directory").is_none());
672    }
673
674    /// Tests that clearing agent memory only affects agent-scoped data.
675    #[tokio::test]
676    async fn test_agent_clear_memory_scoped() {
677        let config = Config::new_default();
678        let mut agent = Agent::builder("test_agent")
679            .config(config)
680            .build()
681            .await
682            .unwrap();
683
684        // Set an agent memory and a general (non-agent) session metadata key
685        agent.set_memory("tasks_completed", "3");
686        agent
687            .chat_session_mut()
688            .set_metadata("session_start", "now");
689
690        // Clear only agent-scoped memory
691        agent.clear_memory();
692
693        // Agent memory removed
694        assert!(agent.get_memory("tasks_completed").is_none());
695        // General session metadata preserved
696        assert_eq!(
697            agent.chat_session().get_metadata("session_start"),
698            Some(&"now".to_string())
699        );
700    }
701
702    /// Tests the increment helper methods for agent memory.
703    #[tokio::test]
704    async fn test_agent_increment_helpers() {
705        let config = Config::new_default();
706        let mut agent = Agent::builder("test_agent")
707            .config(config)
708            .build()
709            .await
710            .unwrap();
711
712        // tasks_completed increments from 0
713        let n1 = agent.increment_tasks_completed();
714        assert_eq!(n1, 1);
715        assert_eq!(agent.get_memory("tasks_completed"), Some(&"1".to_string()));
716
717        let n2 = agent.increment_tasks_completed();
718        assert_eq!(n2, 2);
719        assert_eq!(agent.get_memory("tasks_completed"), Some(&"2".to_string()));
720
721        // generic counter
722        let f1 = agent.increment_counter("files_accessed");
723        assert_eq!(f1, 1);
724        let f2 = agent.increment_counter("files_accessed");
725        assert_eq!(f2, 2);
726        assert_eq!(agent.get_memory("files_accessed"), Some(&"2".to_string()));
727    }
728
729    /// Tests the full functionality of the agent builder.
730    #[tokio::test]
731    async fn test_agent_builder() {
732        let config = Config::new_default();
733        let agent = Agent::builder("test_agent")
734            .config(config)
735            .system_prompt("You are a helpful assistant")
736            .max_iterations(5)
737            .tool(Box::new(CalculatorTool))
738            .build()
739            .await
740            .unwrap();
741
742        assert_eq!(agent.name(), "test_agent");
743        assert_eq!(agent.max_iterations, 5);
744        assert_eq!(
745            agent.tool_registry().list_tools(),
746            vec!["calculator".to_string()]
747        );
748    }
749
750    /// Tests setting the system prompt for an agent.
751    #[tokio::test]
752    async fn test_agent_system_prompt() {
753        let config = Config::new_default();
754        let mut agent = Agent::builder("test_agent")
755            .config(config)
756            .build()
757            .await
758            .unwrap();
759        agent.set_system_prompt("You are a test agent");
760
761        // Check that the system prompt is set in chat session
762        let session = agent.chat_session();
763        assert_eq!(
764            session.system_prompt,
765            Some("You are a test agent".to_string())
766        );
767    }
768
769    /// Tests the tool registry functionality of an agent.
770    #[tokio::test]
771    async fn test_agent_tool_registry() {
772        let config = Config::new_default();
773        let mut agent = Agent::builder("test_agent")
774            .config(config)
775            .build()
776            .await
777            .unwrap();
778
779        // Initially no tools
780        assert!(agent.tool_registry().list_tools().is_empty());
781
782        // Register a tool
783        agent.register_tool(Box::new(CalculatorTool));
784        assert_eq!(
785            agent.tool_registry().list_tools(),
786            vec!["calculator".to_string()]
787        );
788    }
789
790    /// Tests clearing the chat history of an agent.
791    #[tokio::test]
792    async fn test_agent_clear_history() {
793        let config = Config::new_default();
794        let mut agent = Agent::builder("test_agent")
795            .config(config)
796            .build()
797            .await
798            .unwrap();
799
800        // Add a message to the chat session
801        agent.chat_session_mut().add_user_message("Hello");
802        assert!(!agent.chat_session().messages.is_empty());
803
804        // Clear history
805        agent.clear_history();
806        assert!(agent.chat_session().messages.is_empty());
807    }
808
809    // Mock tool for testing
810    struct MockTool;
811
812    #[async_trait::async_trait]
813    impl Tool for MockTool {
814        fn name(&self) -> &str {
815            "mock_tool"
816        }
817
818        fn description(&self) -> &str {
819            "A mock tool for testing"
820        }
821
822        fn parameters(&self) -> HashMap<String, ToolParameter> {
823            let mut params = HashMap::new();
824            params.insert(
825                "input".to_string(),
826                ToolParameter {
827                    param_type: "string".to_string(),
828                    description: "Input parameter".to_string(),
829                    required: Some(true),
830                },
831            );
832            params
833        }
834
835        async fn execute(&self, args: Value) -> crate::Result<ToolResult> {
836            let input = args
837                .get("input")
838                .and_then(|v| v.as_str())
839                .unwrap_or("default");
840            Ok(ToolResult::success(format!("Mock tool output: {}", input)))
841        }
842    }
843}
844
845pub struct AgentBuilder {
846    name: String,
847    config: Option<Config>,
848    system_prompt: Option<String>,
849    tools: Vec<Box<dyn crate::tools::Tool>>,
850    max_iterations: usize,
851}
852
853impl AgentBuilder {
854    pub fn new(name: impl Into<String>) -> Self {
855        Self {
856            name: name.into(),
857            config: None,
858            system_prompt: None,
859            tools: Vec::new(),
860            max_iterations: 10,
861        }
862    }
863
864    pub fn config(mut self, config: Config) -> Self {
865        self.config = Some(config);
866        self
867    }
868
869    pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
870        self.system_prompt = Some(prompt.into());
871        self
872    }
873
874    pub fn tool(mut self, tool: Box<dyn crate::tools::Tool>) -> Self {
875        self.tools.push(tool);
876        self
877    }
878
879    pub fn max_iterations(mut self, max: usize) -> Self {
880        self.max_iterations = max;
881        self
882    }
883
884    pub async fn build(self) -> Result<Agent> {
885        let config = self
886            .config
887            .ok_or_else(|| HeliosError::AgentError("Config is required".to_string()))?;
888
889        let mut agent = Agent::new(self.name, config).await?;
890
891        if let Some(prompt) = self.system_prompt {
892            agent.set_system_prompt(prompt);
893        }
894
895        for tool in self.tools {
896            agent.register_tool(tool);
897        }
898
899        agent.set_max_iterations(self.max_iterations);
900
901        Ok(agent)
902    }
903}