use std::sync::Arc;
use crate::{
agent::multi_agent::subagents::SubagentTool,
agent::{create_agent, AgentError, UnifiedAgent},
tools::Tool,
};
pub struct SubagentInfo {
pub agent: Arc<UnifiedAgent>,
pub name: String,
pub description: String,
}
impl SubagentInfo {
pub fn new(agent: Arc<UnifiedAgent>, name: String, description: String) -> Self {
Self {
agent,
name,
description,
}
}
}
pub struct SubagentsBuilder {
model: Option<String>,
system_prompt: Option<String>,
tools: Vec<Arc<dyn Tool>>,
subagents: Vec<SubagentInfo>,
}
impl SubagentsBuilder {
pub fn new() -> Self {
Self {
model: None,
system_prompt: None,
tools: Vec::new(),
subagents: Vec::new(),
}
}
pub fn with_model<S: Into<String>>(mut self, model: S) -> Self {
self.model = Some(model.into());
self
}
pub fn with_system_prompt<S: Into<String>>(mut self, system_prompt: S) -> Self {
self.system_prompt = Some(system_prompt.into());
self
}
pub fn with_tools(mut self, tools: &[Arc<dyn Tool>]) -> Self {
self.tools.extend_from_slice(tools);
self
}
pub fn with_subagent(mut self, subagent: SubagentInfo) -> Self {
self.subagents.push(subagent);
self
}
pub fn with_subagents(mut self, subagents: Vec<SubagentInfo>) -> Self {
self.subagents.extend(subagents);
self
}
pub fn build(self) -> Result<UnifiedAgent, AgentError> {
let model = self
.model
.ok_or_else(|| AgentError::MissingObject("model".to_string()))?;
let mut all_tools: Vec<Arc<dyn Tool>> = self.tools;
for subagent_info in self.subagents {
let subagent_tool = Arc::new(SubagentTool::new(
subagent_info.agent,
subagent_info.name,
subagent_info.description,
));
all_tools.push(subagent_tool);
}
let system_prompt = self.system_prompt.unwrap_or_else(|| {
"You are a helpful assistant that coordinates specialized subagents to help users."
.to_string()
});
create_agent(&model, &all_tools, Some(&system_prompt), None)
}
}
impl Default for SubagentsBuilder {
fn default() -> Self {
Self::new()
}
}