use crate::error::DispatchError;
use crate::ir::WarbleIr;
use std::collections::BTreeSet;
pub const BINDING_SPEC_VERSION: &str = "1.0";
const STRONG_TIER: &str = "strong";
const CHEAP_TIER: &str = "cheap";
const ORCHESTRATOR_TIER: &str = "orchestrator";
pub const ANTHROPIC_PROVIDER: &str = "anthropic";
pub const OPENAI_COMPAT_PROVIDER: &str = "openai_compat";
pub type Provider = String;
#[derive(Debug, Clone)]
pub struct TierBinding {
pub provider: Provider,
pub endpoint: Option<String>,
pub model: String,
}
impl TierBinding {
fn anthropic(model: String) -> Self {
TierBinding {
provider: ANTHROPIC_PROVIDER.to_string(),
endpoint: None,
model,
}
}
}
#[derive(Debug, Clone)]
pub struct ModelConfig {
tiers: Vec<(String, TierBinding)>,
}
impl Default for ModelConfig {
fn default() -> Self {
ModelConfig {
tiers: vec![
(
STRONG_TIER.to_string(),
TierBinding::anthropic("opus".into()),
),
(
CHEAP_TIER.to_string(),
TierBinding::anthropic("haiku".into()),
),
(
ORCHESTRATOR_TIER.to_string(),
TierBinding::anthropic("sonnet".into()),
),
],
}
}
}
impl ModelConfig {
pub fn from_flags(strong: String, cheap: String, orchestrator: String) -> Self {
ModelConfig {
tiers: vec![
(STRONG_TIER.to_string(), TierBinding::anthropic(strong)),
(CHEAP_TIER.to_string(), TierBinding::anthropic(cheap)),
(
ORCHESTRATOR_TIER.to_string(),
TierBinding::anthropic(orchestrator),
),
],
}
}
pub fn from_yaml(text: &str) -> Result<Self, DispatchError> {
#[derive(serde::Deserialize)]
struct Raw {
#[serde(default)]
tiers: serde_yaml::Mapping,
}
let raw: Raw = serde_yaml::from_str(text)
.map_err(|e| DispatchError(format!("invalid models config: {e}")))?;
let mut tiers = Vec::new();
for (key, value) in raw.tiers {
let name = key
.as_str()
.ok_or_else(|| DispatchError::new("models config: tier name must be a string"))?
.to_string();
let binding = Self::parse_tier_value(&name, &value)?;
tiers.push((name, binding));
}
if tiers.is_empty() {
return Err(DispatchError::new(
"models config: `tiers` must not be empty",
));
}
Ok(ModelConfig { tiers })
}
fn parse_tier_value(
name: &str,
value: &serde_yaml::Value,
) -> Result<TierBinding, DispatchError> {
if let Some(model) = value.as_str() {
return Ok(TierBinding::anthropic(model.to_string()));
}
let map = value.as_mapping().ok_or_else(|| {
DispatchError(format!(
"models config: tier '{name}' must be a model-alias string or a \
{{provider, endpoint?, model}} map"
))
})?;
let get = |k: &str| map.get(serde_yaml::Value::from(k)).and_then(|v| v.as_str());
let model = get("model")
.ok_or_else(|| {
DispatchError(format!(
"models config: tier '{name}' map is missing `model`"
))
})?
.to_string();
let provider = get("provider")
.map(str::to_string)
.unwrap_or_else(|| ANTHROPIC_PROVIDER.to_string());
let endpoint = get("endpoint").map(str::to_string);
if provider == OPENAI_COMPAT_PROVIDER && endpoint.is_none() {
return Err(DispatchError(format!(
"models config: tier '{name}' uses provider openai_compat but has no `endpoint`"
)));
}
Ok(TierBinding {
provider,
endpoint,
model,
})
}
pub fn orchestrator(&self) -> Result<&str, DispatchError> {
self.require(ORCHESTRATOR_TIER)
}
fn binding_for(&self, tier: &str) -> Option<&TierBinding> {
self.tiers
.iter()
.find(|(name, _)| name == tier)
.map(|(_, binding)| binding)
}
fn model_for(&self, tier: &str) -> Option<&str> {
self.binding_for(tier).map(|b| b.model.as_str())
}
fn rank(&self, tier: &str) -> usize {
self.tiers
.iter()
.position(|(name, _)| name == tier)
.unwrap_or(usize::MAX)
}
fn tier_names(&self) -> String {
self.tiers
.iter()
.map(|(name, _)| name.as_str())
.collect::<Vec<_>>()
.join(", ")
}
pub fn require(&self, tier: &str) -> Result<&str, DispatchError> {
self.model_for(tier).ok_or_else(|| {
DispatchError(format!(
"tier '{tier}' has no model binding — define it in --models-config or via \
--strong/--cheap (known tiers: {})",
self.tier_names()
))
})
}
pub fn binding(&self, tier: &str) -> Result<&TierBinding, DispatchError> {
self.binding_for(tier).ok_or_else(|| {
DispatchError(format!(
"tier '{tier}' has no model binding — define it in --models-config or via \
--strong/--cheap (known tiers: {})",
self.tier_names()
))
})
}
pub fn collapsed_model<'a>(
&'a self,
calls: &[crate::ir::LlmCall],
) -> Result<&'a str, DispatchError> {
let strongest = calls
.iter()
.min_by_key(|c| self.rank(&c.tier))
.ok_or_else(|| {
DispatchError::new("component has no llm_calls; cannot select a model")
})?;
self.require(&strongest.tier)
}
pub fn validate(&self, ir: &WarbleIr) -> Result<(), DispatchError> {
let mut checked = BTreeSet::new();
for node in &ir.components {
for call in &node.llm_calls {
if checked.insert(call.tier.as_str()) {
self.require(&call.tier)?;
}
}
}
Ok(())
}
}