cortiq-gateway 0.2.38

Universal LLM gateway with intelligent routing and an embedded multilingual admin console
//! Model selection: `(task_label, complexity_tier) → ordered list of model_ids`.
//! The list serves as both a priority order and a fallback chain. See docs/ROUTING.md.

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>>,
    /// Per-task rules — a matching rule's models lead the candidate chain.
    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(),
        }
    }

    /// Candidates for `(task, tier)`, in preference order, deduplicated:
    /// the task rule's models first (when one exists), then the tier chain —
    /// so a task pin adds preference without ever weakening failover.
    /// Returns `(candidates, task_rule_matched)`.
    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)
    }

    /// Candidates for the given complexity tier, in preference order.
    /// If the tier is not defined, the sole candidate is `default`.
    pub fn candidates(&self, tier: &str) -> Vec<String> {
        self.candidates_for("", tier).0
    }

    /// Default model (used when the router is unavailable).
    pub fn default_model(&self) -> &str {
        &self.default
    }
}

// TODO(v0.2): cost_aware mode and circuit breaker — pick the cheapest available
// model no lower than min_class[tier], respecting max_cost_usd_per_request,
// and skip models whose circuit breaker is open.

#[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"]);
        // dedup: rule model already in the tier chain appears once, first
        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"]);
    }
}