use std::sync::Arc;
use bamboo_llm::LLMProvider;
#[derive(Clone, Default)]
pub struct RoleModel {
pub name: Option<String>,
pub provider_override: Option<Arc<dyn LLMProvider>>,
}
impl RoleModel {
pub fn from_parts(
name: Option<String>,
provider_override: Option<Arc<dyn LLMProvider>>,
) -> Option<Self> {
if name.is_none() && provider_override.is_none() {
None
} else {
Some(Self {
name,
provider_override,
})
}
}
}
#[derive(Clone, Default)]
pub struct ModelRoster {
pub model: Option<String>,
pub provider_name: Option<String>,
pub provider_type: Option<String>,
pub fast: Option<RoleModel>,
pub background: Option<RoleModel>,
pub summarization: Option<RoleModel>,
}
impl ModelRoster {
pub fn empty() -> Self {
Self::default()
}
pub fn from_areas(
model: Option<String>,
provider_name: Option<String>,
provider_type: Option<String>,
areas: crate::model_areas::GlobalAreaModels,
) -> Self {
Self {
model,
provider_name,
provider_type,
fast: RoleModel::from_parts(
areas.fast.as_ref().map(|m| m.model_name.clone()),
areas.fast.map(|m| m.provider),
),
background: RoleModel::from_parts(
areas.background.as_ref().map(|m| m.model_name.clone()),
areas.background.map(|m| m.provider),
),
summarization: RoleModel::from_parts(
areas.summarization.as_ref().map(|m| m.model_name.clone()),
areas.summarization.map(|m| m.provider),
),
}
}
pub fn fast_model(&self) -> Option<String> {
self.fast.as_ref().and_then(|r| r.name.clone())
}
pub fn fast_model_provider(&self) -> Option<Arc<dyn LLMProvider>> {
self.fast.as_ref().and_then(|r| r.provider_override.clone())
}
pub fn background_model(&self) -> Option<String> {
self.background.as_ref().and_then(|r| r.name.clone())
}
pub fn background_model_provider(&self) -> Option<Arc<dyn LLMProvider>> {
self.background
.as_ref()
.and_then(|r| r.provider_override.clone())
}
pub fn summarization_model(&self) -> Option<String> {
self.summarization.as_ref().and_then(|r| r.name.clone())
}
pub fn summarization_model_provider(&self) -> Option<Arc<dyn LLMProvider>> {
self.summarization
.as_ref()
.and_then(|r| r.provider_override.clone())
}
pub fn set_fast(&mut self, name: Option<String>, provider: Option<Arc<dyn LLMProvider>>) {
self.fast = RoleModel::from_parts(name, provider);
}
pub fn set_background(&mut self, name: Option<String>, provider: Option<Arc<dyn LLMProvider>>) {
self.background = RoleModel::from_parts(name, provider);
}
pub fn set_summarization(
&mut self,
name: Option<String>,
provider: Option<Arc<dyn LLMProvider>>,
) {
self.summarization = RoleModel::from_parts(name, provider);
}
}