use crate::config::{Config, TierTargets};
fn to_list(t: &TierTargets) -> Vec<String> {
match t {
TierTargets::List(v) => v.clone(),
TierTargets::One(s) => vec![s.clone()],
}
}
pub struct RoutingTable {
tiers: std::collections::HashMap<String, Vec<String>>,
tasks: std::collections::HashMap<String, Vec<String>>,
default: String,
}
impl RoutingTable {
pub fn from_config(cfg: &Config) -> Self {
Self {
tiers: cfg
.routing
.tiers
.iter()
.map(|(tier, targets)| (tier.clone(), to_list(targets)))
.collect(),
tasks: cfg
.routing
.tasks
.iter()
.map(|(task, targets)| (task.clone(), to_list(targets)))
.filter(|(_, v)| !v.is_empty())
.collect(),
default: cfg.routing.default.clone(),
}
}
pub fn candidates_for(&self, task: &str, tier: &str) -> (Vec<String>, bool) {
let mut out: Vec<String> = Vec::new();
let matched = if let Some(rule) = self.tasks.get(task) {
out.extend(rule.iter().cloned());
true
} else {
false
};
let tier_chain = self
.tiers
.get(tier)
.cloned()
.unwrap_or_else(|| vec![self.default.clone()]);
for id in tier_chain {
if !out.contains(&id) {
out.push(id);
}
}
out.retain(|id| !id.is_empty());
(out, matched)
}
pub fn candidates(&self, tier: &str) -> Vec<String> {
self.candidates_for("", tier).0
}
pub fn default_model(&self) -> &str {
&self.default
}
}
#[cfg(test)]
mod tests {
use super::*;
fn table(tasks: &[(&str, &[&str])], tiers: &[(&str, &[&str])], default: &str) -> RoutingTable {
RoutingTable {
tiers: tiers
.iter()
.map(|(k, v)| (k.to_string(), v.iter().map(|s| s.to_string()).collect()))
.collect(),
tasks: tasks
.iter()
.map(|(k, v)| (k.to_string(), v.iter().map(|s| s.to_string()).collect()))
.collect(),
default: default.to_string(),
}
}
#[test]
fn task_rule_leads_tier_chain_follows() {
let t = table(
&[("translation", &["local"])],
&[("low", &["cheap", "local"]), ("high", &["big"])],
"cheap",
);
let (c, matched) = t.candidates_for("translation", "high");
assert!(matched);
assert_eq!(c, vec!["local", "big"]);
let (c, _) = t.candidates_for("translation", "low");
assert_eq!(c, vec!["local", "cheap"]);
}
#[test]
fn no_rule_falls_back_to_tier_then_default() {
let t = table(&[], &[("low", &["a"])], "d");
assert_eq!(t.candidates_for("qa", "low").0, vec!["a"]);
let (c, matched) = t.candidates_for("qa", "unknown-tier");
assert!(!matched);
assert_eq!(c, vec!["d"]);
}
}