use std::sync::Arc;
use std::path::PathBuf;
use tokio::sync::RwLock;
use bamboo_agent_core::tools::{Tool, ToolExecutor};
use bamboo_engine::AgentBuilder as EngineAgentBuilder;
use bamboo_llm::{create_provider_with_dir, Config, LLMProvider};
use bamboo_metrics::{MetricsCollector, SqliteMetricsStorage};
use bamboo_skills::{SkillManager, SkillStoreConfig};
use bamboo_storage::{LockedSessionStore, SessionStoreV2};
use bamboo_tools::ToolRegistry;
use super::Agent;
const DEFAULT_METRICS_RETENTION_DAYS: u32 = 90;
pub struct AgentBuilder {
inner: EngineAgentBuilder,
system_prompt: Option<String>,
tools: Vec<Arc<dyn Tool>>,
model: Option<String>,
api_key: Option<String>,
}
impl AgentBuilder {
pub fn new() -> Self {
Self {
inner: EngineAgentBuilder::new(),
system_prompt: None,
tools: Vec::new(),
model: None,
api_key: None,
}
}
pub fn model(mut self, model: impl Into<String>) -> Self {
self.model = Some(model.into());
self
}
pub fn instruction(mut self, instruction: impl Into<String>) -> Self {
self.system_prompt = Some(instruction.into());
self
}
pub fn tools<I>(mut self, tools: I) -> Self
where
I: IntoIterator<Item = Arc<dyn Tool>>,
{
self.tools = tools.into_iter().collect();
self
}
pub fn tool<T: Tool + 'static>(mut self, tool: T) -> Self {
self.tools.push(Arc::new(tool));
self
}
pub fn tool_shared(mut self, tool: Arc<dyn Tool>) -> Self {
self.tools.push(tool);
self
}
pub fn api_key(mut self, api_key: impl Into<String>) -> Self {
self.api_key = Some(api_key.into());
self
}
pub fn provider(mut self, provider: Arc<dyn LLMProvider>) -> Self {
self.inner = self.inner.provider(provider);
self
}
pub fn default_tools(mut self, tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor>) -> Self {
self.inner = self.inner.default_tools(tools);
self
}
pub fn config(mut self, config: Arc<RwLock<Config>>) -> Self {
self.inner = self.inner.config(config);
self
}
pub async fn with_defaults_for_data_dir(mut self, data_dir: PathBuf) -> Result<Self, String> {
let mut config = Config::from_data_dir(Some(data_dir.clone()));
if let Some(api_key) = self.api_key.clone() {
apply_api_key(&mut config, &api_key);
}
let provider = create_provider_with_dir(&config, data_dir.clone())
.await
.map_err(|e| format!("failed to create provider: {e}"))?;
let config = Arc::new(RwLock::new(config));
let default_tools: Arc<dyn bamboo_agent_core::tools::ToolExecutor> = Arc::new(
bamboo_tools::BuiltinToolExecutor::new_with_config(config.clone()),
);
let store = Arc::new(
SessionStoreV2::new(data_dir.clone())
.await
.map_err(|e| format!("failed to initialize session store: {e}"))?,
);
let persistence = Arc::new(LockedSessionStore::new(store.clone()));
let skill_manager = Arc::new(SkillManager::with_config(SkillStoreConfig {
skills_dir: data_dir.join("skills"),
project_dir: std::env::current_dir().ok(),
active_mode: None,
}));
skill_manager
.initialize()
.await
.map_err(|e| format!("failed to initialize skill manager: {e}"))?;
let metrics_storage: Arc<dyn bamboo_metrics::storage::MetricsStorage> =
Arc::new(SqliteMetricsStorage::new(data_dir.join("metrics.db")));
let metrics_collector =
MetricsCollector::spawn(metrics_storage, DEFAULT_METRICS_RETENTION_DAYS);
self.inner = self
.inner
.storage(store.clone())
.persistence(persistence)
.attachment_reader(store)
.skill_manager(skill_manager)
.metrics_collector(metrics_collector)
.config(config)
.provider(provider)
.default_tools(default_tools);
Ok(self)
}
pub fn build(mut self) -> Result<Agent, String> {
if !self.tools.is_empty() {
let registry = ToolRegistry::new();
for tool in &self.tools {
let _ = registry.register_shared(tool.clone());
}
let executor: Arc<dyn ToolExecutor> =
Arc::new(bamboo_tools::BuiltinToolExecutor::with_registry(registry));
self.inner = self.inner.default_tools(executor);
}
let runtime = self.inner.build().map_err(|e| e.to_string())?;
Ok(Agent::from_runtime_with_config(
runtime,
self.system_prompt,
self.model,
))
}
}
impl Default for AgentBuilder {
fn default() -> Self {
Self::new()
}
}
fn apply_api_key(config: &mut Config, api_key: &str) {
let stanza = || serde_json::json!({ "api_key": api_key });
let applied = match config.provider.as_str() {
"openai" => match config.providers.openai.as_mut() {
Some(c) => {
c.api_key = api_key.to_string();
true
}
None => {
config.providers.openai = serde_json::from_value(stanza()).ok();
config.providers.openai.is_some()
}
},
"anthropic" => match config.providers.anthropic.as_mut() {
Some(c) => {
c.api_key = api_key.to_string();
true
}
None => {
config.providers.anthropic = serde_json::from_value(stanza()).ok();
config.providers.anthropic.is_some()
}
},
"gemini" => match config.providers.gemini.as_mut() {
Some(c) => {
c.api_key = api_key.to_string();
true
}
None => {
config.providers.gemini = serde_json::from_value(stanza()).ok();
config.providers.gemini.is_some()
}
},
_ => false,
};
if !applied {
tracing::warn!(
provider = %config.provider,
"AgentBuilder::api_key: key not applied — the active provider either \
takes no plain api_key (e.g. copilot uses cached OAuth) or its config \
could not be built from a key alone"
);
}
}
#[cfg(test)]
mod tests {
use super::apply_api_key;
use bamboo_llm::Config;
#[test]
fn api_key_fabricates_stanza_for_keyed_providers() {
for provider in ["openai", "anthropic", "gemini"] {
let mut config = Config::default();
config.provider = provider.to_string();
config.providers.openai = None;
config.providers.anthropic = None;
config.providers.gemini = None;
apply_api_key(&mut config, "sk-test-123");
let key = match provider {
"openai" => config.providers.openai.as_ref().map(|c| c.api_key.as_str()),
"anthropic" => config
.providers
.anthropic
.as_ref()
.map(|c| c.api_key.as_str()),
"gemini" => config.providers.gemini.as_ref().map(|c| c.api_key.as_str()),
_ => unreachable!(),
};
assert_eq!(
key,
Some("sk-test-123"),
"expected a fabricated {provider} stanza carrying the api_key"
);
}
}
}