Skip to main content

agentrs_agents/
builder.rs

1use std::{marker::PhantomData, sync::Arc};
2
3use agentrs_core::{AgentError, LlmProvider, Memory, Result, Tool};
4use agentrs_memory::InMemoryMemory;
5use agentrs_tools::ToolRegistry;
6
7use crate::runner::{AgentConfig, AgentRunner, LoopStrategy};
8
9/// Typestate marker indicating no LLM has been configured.
10#[derive(Debug, Clone, Copy, Default)]
11pub struct NoLlm;
12
13/// Typestate marker indicating an LLM has been configured.
14#[derive(Debug, Clone, Copy, Default)]
15pub struct WithLlm;
16
17/// User-facing entrypoint for agent construction.
18pub struct Agent;
19
20impl Agent {
21    /// Starts building an agent.
22    pub fn builder() -> AgentBuilder<NoLlm, InMemoryMemory> {
23        AgentBuilder::new()
24    }
25}
26
27/// Typestate builder for [`AgentRunner`].
28pub struct AgentBuilder<State, M = InMemoryMemory> {
29    state: PhantomData<State>,
30    llm: Option<Arc<dyn LlmProvider>>,
31    memory: M,
32    tools: ToolRegistry,
33    system_prompt: Option<String>,
34    config: AgentConfig,
35}
36
37impl AgentBuilder<NoLlm, InMemoryMemory> {
38    /// Creates a builder with in-memory defaults.
39    pub fn new() -> Self {
40        Self {
41            state: PhantomData,
42            llm: None,
43            memory: InMemoryMemory::new(),
44            tools: ToolRegistry::new(),
45            system_prompt: None,
46            config: AgentConfig::default(),
47        }
48    }
49}
50
51impl<State, M> AgentBuilder<State, M> {
52    /// Sets the system prompt.
53    pub fn system(mut self, prompt: impl Into<String>) -> Self {
54        self.system_prompt = Some(prompt.into());
55        self
56    }
57
58    /// Replaces the tool registry.
59    pub fn tools(mut self, tools: ToolRegistry) -> Self {
60        self.tools = tools;
61        self
62    }
63
64    /// Registers a single tool.
65    pub fn tool(mut self, tool: impl Tool + 'static) -> Self {
66        self.tools = self.tools.register(tool);
67        self
68    }
69
70    /// Replaces the memory backend.
71    pub fn memory<NewM>(self, memory: NewM) -> AgentBuilder<State, NewM> {
72        AgentBuilder {
73            state: PhantomData,
74            llm: self.llm,
75            memory,
76            tools: self.tools,
77            system_prompt: self.system_prompt,
78            config: self.config,
79        }
80    }
81
82    /// Sets the default model.
83    pub fn model(mut self, model: impl Into<String>) -> Self {
84        self.config.model = model.into();
85        self
86    }
87
88    /// Sets the sampling temperature.
89    pub fn temperature(mut self, temperature: f32) -> Self {
90        self.config.temperature = Some(temperature);
91        self
92    }
93
94    /// Sets the maximum output tokens.
95    pub fn max_tokens(mut self, max_tokens: u32) -> Self {
96        self.config.max_tokens = Some(max_tokens);
97        self
98    }
99
100    /// Sets the loop strategy.
101    pub fn loop_strategy(mut self, loop_strategy: LoopStrategy) -> Self {
102        self.config.max_steps = loop_strategy.max_steps_hint(self.config.max_steps);
103        self.config.loop_strategy = loop_strategy;
104        self
105    }
106
107    /// Sets the maximum tool/action steps.
108    pub fn max_steps(mut self, max_steps: usize) -> Self {
109        self.config.max_steps = max_steps;
110        self
111    }
112}
113
114impl<M> AgentBuilder<NoLlm, M> {
115    /// Configures the LLM provider and advances the typestate.
116    pub fn llm(self, llm: impl LlmProvider + 'static) -> AgentBuilder<WithLlm, M> {
117        AgentBuilder {
118            state: PhantomData,
119            llm: Some(Arc::new(llm)),
120            memory: self.memory,
121            tools: self.tools,
122            system_prompt: self.system_prompt,
123            config: self.config,
124        }
125    }
126
127    /// Configures a shared LLM provider and advances the typestate.
128    pub fn llm_arc(self, llm: Arc<dyn LlmProvider>) -> AgentBuilder<WithLlm, M> {
129        AgentBuilder {
130            state: PhantomData,
131            llm: Some(llm),
132            memory: self.memory,
133            tools: self.tools,
134            system_prompt: self.system_prompt,
135            config: self.config,
136        }
137    }
138}
139
140impl<M> AgentBuilder<WithLlm, M>
141where
142    M: Memory,
143{
144    /// Builds the configured agent.
145    pub fn build(self) -> Result<AgentRunner<M>> {
146        let llm = self.llm.ok_or(AgentError::NoLlmProvider)?;
147        Ok(AgentRunner::new(
148            llm,
149            self.memory,
150            self.tools,
151            self.system_prompt,
152            self.config,
153        ))
154    }
155}