use crate::{
builder::LLMBuilder,
error::LLMError,
memory::{ChatWithMemory, ChatWithMemoryConfig, MemoryProvider, MessageCondition},
LLMProvider,
};
use std::sync::Arc;
use tokio::sync::RwLock;
pub struct AgentBuilder {
llm_builder: LLMBuilder,
role: Option<String>,
role_triggers: Vec<(String, MessageCondition)>,
max_cycles: Option<u32>,
single_reply_per_turn: bool,
debounce_ms: Option<u64>,
stt_builder: Option<LLMBuilder>,
tts_builder: Option<LLMBuilder>,
memory: Option<Box<dyn MemoryProvider>>,
}
impl AgentBuilder {
pub fn new() -> Self {
Self {
llm_builder: LLMBuilder::new(),
role: None,
role_triggers: Vec::new(),
max_cycles: None,
single_reply_per_turn: false,
debounce_ms: None,
stt_builder: None,
tts_builder: None,
memory: None,
}
}
pub fn role(mut self, role: impl Into<String>) -> Self {
self.role = Some(role.into());
self
}
pub fn on(mut self, role: impl Into<String>, condition: MessageCondition) -> Self {
self.role_triggers.push((role.into(), condition));
self
}
pub fn max_cycles(mut self, max: u32) -> Self {
self.max_cycles = Some(max);
self
}
pub fn single_reply_per_turn(mut self, enabled: bool) -> Self {
self.single_reply_per_turn = enabled;
self
}
pub fn debounce(mut self, ms: u64) -> Self {
self.debounce_ms = Some(ms);
self
}
pub fn llm(mut self, llm_builder: LLMBuilder) -> Self {
self.llm_builder = llm_builder;
self
}
pub fn stt(mut self, stt_builder: LLMBuilder) -> Self {
self.stt_builder = Some(stt_builder);
self
}
pub fn tts(mut self, tts_builder: LLMBuilder) -> Self {
self.tts_builder = Some(tts_builder);
self
}
pub fn memory(mut self, memory: impl MemoryProvider + 'static) -> Self {
self.memory = Some(Box::new(memory));
self
}
pub fn build(self) -> Result<Box<dyn LLMProvider>, LLMError> {
let base_provider = self.llm_builder.build()?;
let stt_provider = match self.stt_builder {
Some(builder) => Some(Arc::from(builder.build()?)),
None => None,
};
if let Some(memory) = self.memory {
let memory_arc = Arc::new(RwLock::new(memory));
let provider_arc = Arc::from(base_provider);
let mut config = ChatWithMemoryConfig::new(provider_arc, memory_arc)
.role_triggers(self.role_triggers)
.stt_provider(stt_provider);
if let Some(role) = self.role {
config = config.role(role);
}
if let Some(max) = self.max_cycles {
config = config.max_cycles(max);
}
let agent_provider = ChatWithMemory::with_config(config);
Ok(Box::new(agent_provider))
} else {
Ok(base_provider)
}
}
}
impl Default for AgentBuilder {
fn default() -> Self {
Self::new()
}
}