cortiq-gateway 0.2.40

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> {
    t.models()
}

/// A task rule, normalized: per-tier model lists ("which model serves this
/// task at each complexity") plus any-tier models that lead regardless.
#[derive(Clone, Default)]
pub struct TaskRule {
    pub any: Vec<String>,
    pub per: std::collections::HashMap<String, Vec<String>>,
}

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, TaskRule>,
    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)| {
                    let (any, per) = targets.task_rule();
                    (task.clone(), TaskRule { any, per })
                })
                .filter(|(_, r)| !r.any.is_empty() || r.per.values().any(|v| !v.is_empty()))
                .collect(),
            default: cfg.routing.default.clone(),
        }
    }

    /// Candidate chain for `(task, tier)` and whether a task rule shaped it.
    ///
    /// Order of preference:
    ///   1. the task's models FOR THIS TIER (the user said exactly which model
    ///      serves this task at this complexity);
    ///   2. at tier=high — the global high chain (escalation: any-tier task
    ///      models must not undercut it), then the task's any-tier models;
    ///      otherwise the task's any-tier models, then the global tier chain;
    ///   3. `default` as the last resort. Deduplicated, empties dropped.
    pub fn candidates_for(&self, task: &str, tier: &str) -> (Vec<String>, bool) {
        fn push(out: &mut Vec<String>, id: &str) {
            if !id.is_empty() && !out.iter().any(|x| x == id) {
                out.push(id.to_string());
            }
        }
        let mut out: Vec<String> = Vec::new();
        let rule = self.tasks.get(task);
        let tier_chain = self
            .tiers
            .get(tier)
            .cloned()
            .unwrap_or_else(|| vec![self.default.clone()]);
        if let Some(r) = rule {
            if let Some(list) = r.per.get(tier) {
                for id in list {
                    push(&mut out, id);
                }
            }
        }
        let escalate = tier == "high" && self.tiers.contains_key("high");
        if escalate {
            // high verdicts escalate over any-tier task pins
            for id in &tier_chain {
                push(&mut out, id);
            }
            if let Some(r) = rule {
                for id in &r.any {
                    push(&mut out, id);
                }
            }
        } else {
            if let Some(r) = rule {
                for id in &r.any {
                    push(&mut out, id);
                }
            }
            for id in &tier_chain {
                push(&mut out, id);
            }
        }
        let d = self.default.clone();
        push(&mut out, &d);
        // "task-rule" must mean "the rule chose the head of this chain" — a
        // rule that exists but lost to the tier chain (math has a high pick,
        // the request came in low) must not claim the credit.
        let matched = match (rule, out.first()) {
            (Some(r), Some(head)) => {
                r.per.get(tier).map(|v| v.contains(head)).unwrap_or(false)
                    || r.any.first() == Some(head)
            }
            _ => false,
        };
        (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(),
                        TaskRule {
                            any: v.iter().map(|s| s.to_string()).collect(),
                            per: Default::default(),
                        },
                    )
                })
                .collect(),
            default: default.to_string(),
        }
    }

    fn with_task_tier_models(
        mut t: RoutingTable,
        task: &str,
        tier: &str,
        models: &[&str],
    ) -> RoutingTable {
        t.tasks.entry(task.to_string()).or_default().per.insert(
            tier.to_string(),
            models.iter().map(|s| s.to_string()).collect(),
        );
        t
    }

    #[test]
    fn task_rule_leads_tier_chain_follows() {
        let t = table(
            &[("translation", &["local"])],
            &[("low", &["cheap", "local"]), ("high", &["big"])],
            "cheap",
        );
        // dedup: rule model already in the tier chain appears once, first
        let (c, matched) = t.candidates_for("translation", "low");
        assert!(matched);
        assert_eq!(c, vec!["local", "cheap"]);
    }

    #[test]
    fn high_tier_escalates_over_task_rule() {
        let t = table(
            &[("code", &["cheap-coder"])],
            &[("low", &["cheap"]), ("high", &["big"])],
            "cheap",
        );
        // medium (no chain defined → default): task rule still leads
        let (c, m) = t.candidates_for("code", "medium");
        assert!(m);
        assert_eq!(c[0], "cheap-coder");
        // high: tier chain leads, task rule follows as fallback
        let (c, m) = t.candidates_for("code", "high");
        // escalation put the global high chain first → the rule did not pick
        // the head, so no task-rule credit
        assert!(!m);
        // default rides along as the very last fallback
        assert_eq!(c, vec!["big", "cheap-coder", "cheap"]);
    }

    #[test]
    fn per_tier_models_lead_their_tier() {
        let base = table(
            &[("code", &["generic-coder"])],
            &[("low", &["cheap"]), ("high", &["big"])],
            "cheap",
        );
        let t = with_task_tier_models(base, "code", "high", &["best-coder"]);
        // at high: the task's HIGH pick leads, escalation chain next,
        // any-tier task models after it, default last
        let (c, m) = t.candidates_for("code", "high");
        assert!(m);
        assert_eq!(c, vec!["best-coder", "big", "generic-coder", "cheap"]);
        // at low: any-tier task models lead as before (and the badge is
        // honest — the head came from the rule)
        let (c, m) = t.candidates_for("code", "low");
        assert!(m);
        assert_eq!(c, vec!["generic-coder", "cheap"]);
        // a tier with an explicit pick but no any-models
        let t2 = with_task_tier_models(
            table(&[], &[("low", &["cheap"]), ("high", &["big"])], "cheap"),
            "legal",
            "high",
            &["strong"],
        );
        let (c, m) = t2.candidates_for("legal", "high");
        assert!(m);
        assert_eq!(c, vec!["strong", "big", "cheap"]);
        // the same rule at LOW did not shape the head → no task-rule credit
        let (c, m) = t2.candidates_for("legal", "low");
        assert!(!m);
        assert_eq!(c, vec!["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", "d"]);
        let (c, matched) = t.candidates_for("qa", "unknown-tier");
        assert!(!matched);
        assert_eq!(c, vec!["d"]);
    }
}