use std::sync::Arc;
use lellm_core::{ReasoningConfig, ToolChoice};
use lellm_provider::ResolvedModel;
use super::config::{ToolUseConfig, ToolUseDeps};
use super::context::ContextBudget;
use super::fallback::FallbackStrategy;
use super::request_opts::RequestOptions;
use super::retry::RetryPolicy;
use super::runtime::ToolUseLoop;
use super::tools::{CompositeCatalog, StaticCatalog, ToolCatalog, ToolExecutor, ToolRegistration};
pub struct AgentBuilder {
model: ResolvedModel,
static_tools: Vec<ToolRegistration>,
catalogs: Vec<Arc<dyn ToolCatalog>>,
config: ToolUseConfig,
deps: ToolUseDeps,
}
impl AgentBuilder {
pub fn new(model: ResolvedModel) -> Self {
Self {
model,
static_tools: Vec::new(),
catalogs: Vec::new(),
config: ToolUseConfig::default(),
deps: ToolUseDeps::default(),
}
}
pub fn tool(mut self, reg: ToolRegistration) -> Self {
self.static_tools.push(reg);
self
}
pub fn tools(mut self, registrations: impl IntoIterator<Item = ToolRegistration>) -> Self {
self.static_tools.extend(registrations);
self
}
pub fn catalog(mut self, catalog: Arc<dyn ToolCatalog>) -> Self {
self.catalogs.push(catalog);
self
}
pub fn max_iterations(mut self, max: usize) -> Self {
self.config.max_iterations = max;
self
}
pub fn max_output_tokens(mut self, max: u32) -> Self {
self.config.max_output_tokens = max;
self
}
pub fn max_total_output_tokens(mut self, max: u32) -> Self {
self.config.max_total_output_tokens = Some(max);
self
}
pub fn system_prompt(mut self, prompt: String) -> Self {
self.config.system_prompt = Some(prompt);
self
}
pub fn request_options(mut self, opts: RequestOptions) -> Self {
self.config.request_options = opts;
self
}
pub fn temperature(mut self, t: f64) -> Self {
self.config.request_options.temperature = Some(t);
self
}
pub fn top_p(mut self, p: f64) -> Self {
self.config.request_options.top_p = Some(p);
self
}
pub fn seed(mut self, s: u64) -> Self {
self.config.request_options.seed = Some(s);
self
}
pub fn tool_choice(mut self, choice: ToolChoice) -> Self {
self.config.request_options.tool_choice = Some(choice);
self
}
pub fn stop_sequences(mut self, seqs: Vec<String>) -> Self {
self.config.request_options.stop_sequences = Some(seqs);
self
}
pub fn prefill(mut self, text: String) -> Self {
self.config.request_options.prefill = Some(text);
self
}
pub fn reasoning(mut self, r: ReasoningConfig) -> Self {
self.config.request_options.reasoning = Some(r);
self
}
pub fn stream_thinking(mut self, enable: bool) -> Self {
self.config.stream_thinking = enable;
self
}
pub fn reasoning_budget(mut self, max: u32) -> Self {
self.config.request_options.max_reasoning_tokens = Some(max);
self
}
pub fn max_total_reasoning_tokens(mut self, max: u32) -> Self {
self.config.max_total_reasoning_tokens = Some(max);
self
}
pub fn fallback(mut self, fallback: Arc<dyn FallbackStrategy>) -> Self {
self.deps.fallback = fallback;
self
}
pub fn retry_policy(mut self, policy: RetryPolicy) -> Self {
self.config.retry_policy = policy;
self
}
pub fn context_budget(mut self, budget: ContextBudget) -> Self {
self.config.context_budget = budget;
self
}
pub fn build(self) -> ToolUseLoop {
let mut sources: Vec<Arc<dyn ToolCatalog>> = Vec::new();
if !self.static_tools.is_empty() {
sources.push(Arc::new(StaticCatalog::from_tools(self.static_tools)));
}
sources.extend(self.catalogs);
let final_catalog: Arc<dyn ToolCatalog> = match sources.len() {
0 => Arc::new(StaticCatalog::empty()),
1 => sources.remove(0),
_ => Arc::new(CompositeCatalog::new(sources)),
};
let executor =
ToolExecutor::with_retry_policy(final_catalog, self.config.retry_policy.clone());
ToolUseLoop::new(self.model, executor, self.config, self.deps)
}
}