rexis_rag/agent/
builder.rs

1//! Agent builder pattern
2
3use super::memory::{AgentMemoryManager, MemoryConfig};
4use super::{Agent, AgentConfig, ConversationMode, ToolExecutor};
5use crate::error::RragResult;
6
7#[cfg(feature = "rexis-llm-client")]
8use rexis_llm::tools::{Tool, ToolRegistry};
9#[cfg(feature = "rexis-llm-client")]
10use rexis_llm::Client;
11
12/// Builder for creating agents
13pub struct AgentBuilder {
14    llm_client: Option<Client>,
15    tools: Vec<Box<dyn Tool>>,
16    config: AgentConfig,
17    memory_config: Option<MemoryConfig>,
18}
19
20impl AgentBuilder {
21    /// Create a new agent builder
22    pub fn new() -> Self {
23        Self {
24            llm_client: None,
25            tools: Vec::new(),
26            config: AgentConfig::default(),
27            memory_config: None,
28        }
29    }
30
31    /// Set the LLM client
32    pub fn with_llm(mut self, client: Client) -> Self {
33        self.llm_client = Some(client);
34        self
35    }
36
37    /// Add a single tool
38    pub fn with_tool(mut self, tool: Box<dyn Tool>) -> Self {
39        self.tools.push(tool);
40        self
41    }
42
43    /// Add multiple tools
44    pub fn with_tools(mut self, tools: Vec<Box<dyn Tool>>) -> Self {
45        self.tools.extend(tools);
46        self
47    }
48
49    /// Set system prompt
50    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
51        self.config.system_prompt = prompt.into();
52        self
53    }
54
55    /// Set max iterations
56    pub fn with_max_iterations(mut self, max: usize) -> Self {
57        self.config.max_iterations = max;
58        self
59    }
60
61    /// Enable verbose output
62    pub fn verbose(mut self, verbose: bool) -> Self {
63        self.config.verbose = verbose;
64        self
65    }
66
67    /// Set conversation mode
68    pub fn with_conversation_mode(mut self, mode: ConversationMode) -> Self {
69        self.config.conversation_mode = mode;
70        self
71    }
72
73    /// Enable stateful mode (maintains conversation)
74    pub fn stateful(mut self) -> Self {
75        self.config.conversation_mode = ConversationMode::Stateful;
76        self
77    }
78
79    /// Enable stateless mode (fresh each call)
80    pub fn stateless(mut self) -> Self {
81        self.config.conversation_mode = ConversationMode::Stateless;
82        self
83    }
84
85    /// Set max conversation length
86    pub fn with_max_conversation_length(mut self, length: usize) -> Self {
87        self.config.max_conversation_length = length;
88        self
89    }
90
91    /// Set memory configuration (enables persistent memory)
92    pub fn with_memory(mut self, memory_config: MemoryConfig) -> Self {
93        self.memory_config = Some(memory_config);
94        self
95    }
96
97    /// Build the agent
98    pub fn build(self) -> RragResult<Agent> {
99        let llm_client = self
100            .llm_client
101            .ok_or_else(|| crate::error::RragError::Agent {
102                agent_id: "builder".to_string(),
103                message: "LLM client is required. Use with_llm()".to_string(),
104                source: None,
105            })?;
106
107        // Create tool registry
108        let mut registry = ToolRegistry::new();
109        for tool in self.tools {
110            registry
111                .register(tool)
112                .map_err(|e| crate::error::RragError::Agent {
113                    agent_id: "builder".to_string(),
114                    message: format!("Failed to register tool: {}", e),
115                    source: Some(Box::new(e)),
116                })?;
117        }
118
119        let tool_executor = ToolExecutor::new(registry);
120
121        // Build with or without persistent memory
122        if let Some(memory_config) = self.memory_config {
123            let memory_manager = AgentMemoryManager::new(memory_config);
124            Agent::new_with_memory(llm_client, tool_executor, memory_manager, self.config)
125        } else {
126            Agent::new(llm_client, tool_executor, self.config)
127        }
128    }
129}
130
131impl Default for AgentBuilder {
132    fn default() -> Self {
133        Self::new()
134    }
135}