diff --git a/llm-router/.gitignore b/llm-router/.gitignore new file mode 100644 index 000000000..2c96eb1b6 --- /dev/null +++ b/llm-router/.gitignore @@ -0,0 +1,2 @@ +target/ +Cargo.lock diff --git a/llm-router/Cargo.toml b/llm-router/Cargo.toml new file mode 100644 index 000000000..ee69c428b --- /dev/null +++ b/llm-router/Cargo.toml @@ -0,0 +1,25 @@ +[workspace] + +[package] +name = "iii-llm-router" +version = "0.1.0" +edition = "2021" +publish = false + +[[bin]] +name = "iii-llm-router" +path = "src/main.rs" + +[dependencies] +iii-sdk = "=0.11.2" +tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "signal"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +serde_yaml = "0.9" +anyhow = "1" +thiserror = "2" +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] } +clap = { version = "4", features = ["derive"] } +uuid = { version = "1", features = ["v4"] } +rand = "0.8" diff --git a/llm-router/README.md b/llm-router/README.md new file mode 100644 index 000000000..c9dfaefcb --- /dev/null +++ b/llm-router/README.md @@ -0,0 +1,127 @@ +# iii-llm-router + +Policy-based LLM routing brain. **Unopinionated** — ships with zero built-in model names, zero hardcoded pricing, zero provider assumptions. Wraps any gateway (LiteLLM, Bifrost, OpenRouter, a local vLLM, your own proxy) by sitting *in front* of it: gateway asks `router::decide` before every call, router returns a model ID, gateway forwards. + +## Why unopinionated matters + +Every existing LLM router (RouteLLM, Portkey, LiteLLM's routing block) bakes a specific catalog of models and a specific rank ordering into the library. That catalog is wrong the day you read it — new models ship weekly, pricing moves, quality tiers shift. This worker doesn't know what "Opus" or "GPT" is. You register what you actually use at runtime. The only thing the router enforces is its own logic: match → classify → budget → health → fallback. + +## Functions (18) + +| id | shape | +|----|-------| +| `router::decide` | hot path — returns `{model, reason, policy_id?, ab_test_id?, fallback?, confidence, request_id}` | +| `router::policy_create` / `update` / `delete` / `list` / `test` | CRUD + dry-run | +| `router::classify` | run the prompt heuristic only; returns `{complexity, confidence, suggested_model}` (suggested_model respects your classifier map) | +| `router::classifier_config` | register `{id, thresholds: {simple/moderate/complex/expert → }}` | +| `router::ab_create` / `ab_record` / `ab_report` / `ab_conclude` | A/B tests with weighted variants + quality/latency/cost aggregation | +| `router::health_update` / `health_list` | per-model availability + error rate; feeds fallback path | +| `router::model_register` / `model_unregister` / `model_list` | you tell the router what models exist; used only by the budget-downgrade path and stats | +| `router::stats` | usage by model, by policy, over a day window | + +## HTTP triggers (18) + +``` +POST /api/router/decide +POST /api/router/policy/{create,update,delete,test} +GET /api/router/policy/list +POST /api/router/classify +POST /api/router/classifier +POST /api/router/ab/{create,record,report,conclude} +POST /api/router/health/update +GET /api/router/health/list +POST /api/router/model/{register,unregister} +GET /api/router/model/list +GET /api/router/stats +``` + +## Decide logic + +``` +match policies (by tenant, feature, tags) and pick highest priority + if matching A/B test is running → sample a variant → return + else if policy.action.model == "auto" → classify → look up user mapping + if chosen model is unhealthy → use policy.fallback + if policy.max_cost_per_request > budget_remaining → search registered + models for a cheaper one meeting min_quality (if none: return original, + flag reason) + return {model, reason, policy_id, fallback, confidence} +no policy matched: + if classifier exists → classify → map + else → return empty model + reason (caller should handle) +``` + +Router **never** invents a model name. If you ask it to pick "auto" without a classifier registered, it tells you so in `reason` and returns the policy's fallback (or empty). + +## State (engine-managed) + +All stored in `state_scope: "llm-router"` (configurable). + +``` +policies: — policy definitions +ab_tests: — A/B test definitions +ab_events::… — recorded outcomes +routing_log:: — decision audit trail +model_health: — availability + latency + error_rate +classifier: — category → model mapping +models: — registered models (quality, pricing, provider) +``` + +## Example + +```bash +# 1. register two models you actually use +curl -X POST localhost:3111/api/router/model/register -d '{ + "model": "gw/cheap-fast", "quality": "low", + "input_per_1m": 0.1, "output_per_1m": 0.4 +}' +curl -X POST localhost:3111/api/router/model/register -d '{ + "model": "gw/strong", "quality": "high", + "input_per_1m": 15, "output_per_1m": 75 +}' + +# 2. configure the classifier (category → model is YOUR choice) +curl -X POST localhost:3111/api/router/classifier -d '{ + "id": "default", + "thresholds": { + "simple": "gw/cheap-fast", + "moderate": "gw/cheap-fast", + "complex": "gw/strong", + "expert": "gw/strong" + } +}' + +# 3. write a policy +curl -X POST localhost:3111/api/router/policy/create -d '{ + "name": "support-auto", + "match": { "feature": "support-chat" }, + "action": { "model": "auto", "fallback": "gw/cheap-fast" }, + "priority": 100 +}' + +# 4. ask before every call +curl -X POST localhost:3111/api/router/decide -d '{ + "feature": "support-chat", + "prompt": "How do I reset my password?" +}' +# → {"model":"gw/cheap-fast", "reason":"policy: support-auto + classifier: simple", ...} +``` + +Your gateway (LiteLLM/Bifrost/OpenRouter/your-own) takes `model` and forwards. The router doesn't make any LLM call itself. + +## What this is NOT + +- Not a gateway — no LLM traffic passes through it, no API keys stored. +- Not an observability platform — `routing_log` is for audit, use iii's OTel for real telemetry. +- Not a training-based classifier — the shipped classifier is a cheap prompt heuristic (length, code markers, math markers). Swap it by calling `router::classifier_config` with your own mapping, or wrap a stronger classifier as a separate worker and call `router::decide` after you've called it. + +## SDK + stack + +- `iii-sdk 0.11.0` stable +- State via `state::get`/`set`/`delete`/`list` against scope `llm-router` +- `rand` for A/B variant weighted sampling +- `serde_json` everywhere — all state blobs are JSON + +## Tests + +17 passing — policy matching, priority ordering, A/B weighted sampling, classifier mapping, auto-without-classifier, unhealthy-fallback, budget-downgrade with and without registered models, health skip thresholds, heuristic category classification. diff --git a/llm-router/build.rs b/llm-router/build.rs new file mode 100644 index 000000000..33143a5e0 --- /dev/null +++ b/llm-router/build.rs @@ -0,0 +1,6 @@ +fn main() { + println!( + "cargo:rustc-env=TARGET={}", + std::env::var("TARGET").unwrap_or_default() + ); +} diff --git a/llm-router/config.yaml b/llm-router/config.yaml new file mode 100644 index 000000000..b83a8da88 --- /dev/null +++ b/llm-router/config.yaml @@ -0,0 +1,4 @@ +state_scope: "llm-router" +classifier_default_id: "default" +stats_default_days: 7 +health_skip_threshold_error_rate: 0.3 diff --git a/llm-router/src/config.rs b/llm-router/src/config.rs new file mode 100644 index 000000000..151a21d0d --- /dev/null +++ b/llm-router/src/config.rs @@ -0,0 +1,107 @@ +use anyhow::{Context, Result}; +use serde::{Deserialize, Serialize}; +use std::fs; + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct RouterConfig { + #[serde(default = "default_state_scope")] + pub state_scope: String, + + #[serde(default = "default_classifier_id")] + pub classifier_default_id: String, + + #[serde(default = "default_stats_days")] + pub stats_default_days: u32, + + #[serde(default = "default_health_skip_error_rate")] + pub health_skip_threshold_error_rate: f64, +} + +fn default_state_scope() -> String { + "llm-router".to_string() +} +fn default_classifier_id() -> String { + "default".to_string() +} +fn default_stats_days() -> u32 { + 7 +} +fn default_health_skip_error_rate() -> f64 { + 0.3 +} + +impl Default for RouterConfig { + fn default() -> Self { + Self { + state_scope: default_state_scope(), + classifier_default_id: default_classifier_id(), + stats_default_days: default_stats_days(), + health_skip_threshold_error_rate: default_health_skip_error_rate(), + } + } +} + +pub fn load_config(path: &str) -> Result { + let content = fs::read_to_string(path).with_context(|| format!("read {}", path))?; + let cfg: RouterConfig = + serde_yaml::from_str(&content).with_context(|| format!("parse {}", path))?; + validate(&cfg)?; + Ok(cfg) +} + +fn validate(cfg: &RouterConfig) -> Result<()> { + if cfg.state_scope.trim().is_empty() { + anyhow::bail!("config: state_scope must be non-empty"); + } + if cfg.classifier_default_id.trim().is_empty() { + anyhow::bail!("config: classifier_default_id must be non-empty"); + } + if cfg.stats_default_days == 0 { + anyhow::bail!("config: stats_default_days must be >= 1"); + } + let rate = cfg.health_skip_threshold_error_rate; + if !(0.0..=1.0).contains(&rate) || rate.is_nan() { + anyhow::bail!( + "config: health_skip_threshold_error_rate must be within 0.0..=1.0 (got {})", + rate + ); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_defaults() { + let c = RouterConfig::default(); + assert_eq!(c.state_scope, "llm-router"); + assert_eq!(c.classifier_default_id, "default"); + assert_eq!(c.stats_default_days, 7); + } + + #[test] + fn validate_rejects_out_of_range_error_rate() { + let mut c = RouterConfig::default(); + c.health_skip_threshold_error_rate = 2.0; + assert!(validate(&c).is_err()); + c.health_skip_threshold_error_rate = -0.1; + assert!(validate(&c).is_err()); + } + + #[test] + fn validate_rejects_empty_strings() { + let mut c = RouterConfig::default(); + c.state_scope = "".into(); + assert!(validate(&c).is_err()); + } + + #[test] + fn validate_rejects_zero_stats_days() { + let mut c = RouterConfig::default(); + c.stats_default_days = 0; + assert!(validate(&c).is_err()); + } +} diff --git a/llm-router/src/functions/ab.rs b/llm-router/src/functions/ab.rs new file mode 100644 index 000000000..37d25f1e4 --- /dev/null +++ b/llm-router/src/functions/ab.rs @@ -0,0 +1,278 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use iii_sdk::{IIIError, III}; +use serde_json::{json, Value}; +use uuid::Uuid; + +use crate::config::RouterConfig; +use crate::state; +use crate::types::{AbEvent, AbTest}; + +fn key_test(id: &str) -> String { + format!("ab_tests:{}", id) +} +fn key_event(test_id: &str, timestamp_ms: u64, id: &str) -> String { + format!("ab_events:{}:{:020}:{}", test_id, timestamp_ms, id) +} + +pub fn create_handler( + iii: III, + cfg: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + let cfg = cfg.clone(); + Box::pin(async move { + let mut v = payload; + if v.get("id").is_none() { + if let Value::Object(ref mut m) = v { + m.insert("id".into(), Value::String(format!("ab-{}", Uuid::new_v4()))); + } + } + let mut t: AbTest = serde_json::from_value(v) + .map_err(|e| IIIError::Handler(format!("parse ab-test: {}", e)))?; + t.created_at_ms = crate::functions::decide::now_ms(); + state::state_set( + &iii, + &cfg.state_scope, + &key_test(&t.id), + serde_json::to_value(&t).unwrap(), + ) + .await?; + Ok(json!({ "test_id": t.id, "created": true })) + }) + } +} + +pub fn record_handler( + iii: III, + cfg: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + let cfg = cfg.clone(); + Box::pin(async move { + let test_id = payload + .get("test_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| IIIError::Handler("missing 'test_id'".into()))? + .to_string(); + let variant = payload + .get("variant_model") + .and_then(|v| v.as_str()) + .ok_or_else(|| IIIError::Handler("missing 'variant_model'".into()))? + .trim() + .to_string(); + if variant.is_empty() { + return Err(IIIError::Handler("empty 'variant_model'".into())); + } + + // Verify the variant belongs to the test before persisting. + let test_val = state::state_get(&iii, &cfg.state_scope, &key_test(&test_id)) + .await? + .ok_or_else(|| { + IIIError::Handler(format!("no such ab-test: {}", test_id)) + })?; + let test: AbTest = serde_json::from_value(test_val).map_err(|e| { + IIIError::Handler(format!("parse test {}: {}", test_id, e)) + })?; + if !test.variants.iter().any(|v| v.model == variant) { + return Err(IIIError::Handler(format!( + "variant_model '{}' is not registered on ab-test '{}'", + variant, test_id + ))); + } + + let quality_score = payload + .get("quality_score") + .and_then(|v| v.as_f64()) + .unwrap_or(0.0); + if !(0.0..=1.0).contains(&quality_score) || quality_score.is_nan() { + return Err(IIIError::Handler(format!( + "quality_score must be within 0.0..=1.0 (got {})", + quality_score + ))); + } + let latency_ms = payload + .get("latency_ms") + .and_then(|v| v.as_i64()) + .unwrap_or(0); + if latency_ms < 0 { + return Err(IIIError::Handler("latency_ms must be >= 0".into())); + } + let cost_usd = payload.get("cost_usd").and_then(|v| v.as_f64()).unwrap_or(0.0); + if cost_usd < 0.0 || cost_usd.is_nan() { + return Err(IIIError::Handler("cost_usd must be >= 0".into())); + } + + let ev = AbEvent { + test_id: test_id.clone(), + variant_model: variant, + quality_score, + latency_ms: latency_ms as u64, + cost_usd, + recorded_at_ms: crate::functions::decide::now_ms(), + }; + let evt_id = format!("evt-{}", Uuid::new_v4()); + state::state_set( + &iii, + &cfg.state_scope, + &key_event(&test_id, ev.recorded_at_ms, &evt_id), + serde_json::to_value(&ev).unwrap(), + ) + .await?; + Ok(json!({ "recorded": true, "event_id": evt_id })) + }) + } +} + +pub fn report_handler( + iii: III, + cfg: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + let cfg = cfg.clone(); + Box::pin(async move { + let test_id = payload + .get("test_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| IIIError::Handler("missing 'test_id'".into()))? + .to_string(); + + let test_val = state::state_get(&iii, &cfg.state_scope, &key_test(&test_id)) + .await? + .ok_or_else(|| IIIError::Handler(format!("no such ab-test: {}", test_id)))?; + let test: AbTest = serde_json::from_value(test_val) + .map_err(|e| IIIError::Handler(format!("parse test: {}", e)))?; + + let items = state::state_list( + &iii, + &cfg.state_scope, + &format!("ab_events:{}:", test_id), + ) + .await?; + let events: Vec = items + .into_iter() + .filter_map(|it| state::parse_item::(&it)) + .collect(); + + let mut summary: std::collections::HashMap = + std::collections::HashMap::new(); + for e in &events { + let row = summary.entry(e.variant_model.clone()).or_insert((0, 0.0, 0.0, 0.0)); + row.0 += 1; + row.1 += e.quality_score; + row.2 += e.latency_ms as f64; + row.3 += e.cost_usd; + } + + let variants_out: Vec = test + .variants + .iter() + .map(|v| { + let (n, q, l, c) = summary + .get(&v.model) + .copied() + .unwrap_or((0, 0.0, 0.0, 0.0)); + let n_f = (n as f64).max(1.0); + json!({ + "model": v.model, + "weight": v.weight, + "samples": n, + "avg_quality": q / n_f, + "avg_latency_ms": l / n_f, + "avg_cost_usd": c / n_f, + }) + }) + .collect(); + + let total_samples: u64 = summary.values().map(|(n, _, _, _)| *n).sum(); + let status = if test.status == "running" && total_samples < test.min_samples as u64 { + "insufficient_data" + } else if test.status == "concluded" { + "concluded" + } else { + "running" + }; + + Ok(json!({ + "test_id": test.id, + "name": test.name, + "status": status, + "total_samples": total_samples, + "variants": variants_out, + })) + }) + } +} + +pub fn conclude_handler( + iii: III, + cfg: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + let cfg = cfg.clone(); + Box::pin(async move { + let test_id = payload + .get("test_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| IIIError::Handler("missing 'test_id'".into()))? + .to_string(); + let winner = payload + .get("winner_model") + .and_then(|v| v.as_str()) + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) + .ok_or_else(|| IIIError::Handler("missing 'winner_model'".into()))?; + + let test_val = state::state_get(&iii, &cfg.state_scope, &key_test(&test_id)) + .await? + .ok_or_else(|| IIIError::Handler(format!("no such ab-test: {}", test_id)))?; + let mut test: AbTest = serde_json::from_value(test_val) + .map_err(|e| IIIError::Handler(format!("parse test: {}", e)))?; + if !test.variants.iter().any(|v| v.model == winner) { + return Err(IIIError::Handler(format!( + "winner_model '{}' is not one of the variants on ab-test '{}'", + winner, test_id + ))); + } + + test.status = "concluded".into(); + state::state_set( + &iii, + &cfg.state_scope, + &key_test(&test_id), + serde_json::to_value(&test).unwrap(), + ) + .await?; + + // Rollout is intentionally not automatic here: a policy can target + // the test via match-rules without naming the winner, so a write + // would need policy IDs we don't have. Callers drive the rollout + // explicitly via router::policy_update. + Ok(json!({ + "concluded": true, + "test_id": test_id, + "winner_model": winner, + "rollout_applied": false, + "note": "call router::policy_update to roll the winner into the active policy", + })) + }) + } +} diff --git a/llm-router/src/functions/classify.rs b/llm-router/src/functions/classify.rs new file mode 100644 index 000000000..0c5ee2ccb --- /dev/null +++ b/llm-router/src/functions/classify.rs @@ -0,0 +1,92 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use iii_sdk::{IIIError, III}; +use serde_json::{json, Value}; + +use crate::config::RouterConfig; +use crate::router::heuristic_complexity; +use crate::state; +use crate::types::ClassifierConfig; + +pub fn classify_handler( + iii: III, + cfg: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + let cfg = cfg.clone(); + Box::pin(async move { + let prompt = payload + .get("prompt") + .and_then(|v| v.as_str()) + .ok_or_else(|| IIIError::Handler("missing 'prompt'".into()))? + .trim() + .to_string(); + if prompt.is_empty() { + return Err(IIIError::Handler("empty 'prompt'".into())); + } + let classifier_id = payload + .get("classifier_id") + .and_then(|v| v.as_str()) + .map(String::from) + .unwrap_or_else(|| cfg.classifier_default_id.clone()); + + let (category, confidence) = heuristic_complexity(&prompt); + + let classifier = state::state_get( + &iii, + &cfg.state_scope, + &format!("classifier:{}", classifier_id), + ) + .await?; + let mapped_model = classifier + .as_ref() + .and_then(|v| serde_json::from_value::(v.clone()).ok()) + .and_then(|c| c.thresholds.get(category).cloned()); + + Ok(json!({ + "classifier_id": classifier_id, + "complexity": category, + "confidence": confidence, + "suggested_model": mapped_model, + })) + }) + } +} + +pub fn config_handler( + iii: III, + cfg: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + let cfg = cfg.clone(); + Box::pin(async move { + let id = payload + .get("id") + .and_then(|v| v.as_str()) + .map(String::from) + .unwrap_or_else(|| cfg.classifier_default_id.clone()); + let mut c: ClassifierConfig = serde_json::from_value(payload) + .map_err(|e| IIIError::Handler(format!("parse classifier: {}", e)))?; + c.id = id.clone(); + c.created_at_ms = crate::functions::decide::now_ms(); + state::state_set( + &iii, + &cfg.state_scope, + &format!("classifier:{}", id), + serde_json::to_value(&c).unwrap_or(Value::Null), + ) + .await?; + Ok(json!({ "configured": true, "id": id })) + }) + } +} diff --git a/llm-router/src/functions/decide.rs b/llm-router/src/functions/decide.rs new file mode 100644 index 000000000..89af538cb --- /dev/null +++ b/llm-router/src/functions/decide.rs @@ -0,0 +1,172 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; + +use iii_sdk::{IIIError, III}; +use rand::SeedableRng; +use serde_json::{json, Value}; +use std::any::type_name; +use uuid::Uuid; + +use crate::config::RouterConfig; +use crate::router::{decide, DecideContext}; +use crate::state; +use crate::types::{ + AbTest, ClassifierConfig, ModelHealth, ModelRegistration, Policy, RoutingLogEntry, + RoutingRequest, +}; + +pub fn build_handler( + iii: III, + cfg: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + let cfg = cfg.clone(); + Box::pin(async move { handle(iii, cfg, payload).await }) + } +} + +async fn handle(iii: III, cfg: Arc, payload: Value) -> Result { + let mut req: RoutingRequest = serde_json::from_value(payload) + .map_err(|e| IIIError::Handler(format!("parse request: {}", e)))?; + req.prompt = req.prompt.trim().to_string(); + if req.prompt.is_empty() { + return Err(IIIError::Handler("missing or empty 'prompt'".to_string())); + } + + let classifier_id = req + .classifier_id + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(String::from) + .unwrap_or_else(|| cfg.classifier_default_id.clone()); + + let (policies, ab_tests, health, classifier, models) = tokio::join!( + load_policies(&iii, &cfg), + load_ab_tests(&iii, &cfg), + load_health(&iii, &cfg), + load_classifier(&iii, &cfg, &classifier_id), + load_models(&iii, &cfg), + ); + let policies = policies?; + let ab_tests = ab_tests?; + let health = health?; + let classifier = classifier?; + let models = models?; + + // Entropy-backed RNG — millisecond-timestamp seeding collided on burst + // requests, biasing A/B variant picks. + let mut rng = rand::rngs::StdRng::from_entropy(); + let ctx = DecideContext { + policies: &policies, + ab_tests: &ab_tests, + health: &health, + classifier: classifier.as_ref(), + models: &models, + }; + let decision = decide(&req, ctx, &cfg, &mut rng); + + let request_id = format!("req-{}", Uuid::new_v4()); + let log = RoutingLogEntry { + timestamp_ms: now_ms(), + request_id: request_id.clone(), + tenant: req.tenant.clone(), + feature: req.feature.clone(), + model_selected: decision.model.clone(), + policy_matched: decision.policy_id.clone(), + ab_test_id: decision.ab_test_id.clone(), + reason: decision.reason.clone(), + cost_usd: None, + }; + if let Err(e) = state::state_set( + &iii, + &cfg.state_scope, + &format!("routing_log:{:020}:{}", log.timestamp_ms, request_id), + serde_json::to_value(&log).unwrap_or(Value::Null), + ) + .await + { + tracing::warn!(error = %e, "failed to write routing log"); + } + + let mut out = serde_json::to_value(&decision).unwrap_or(Value::Null); + if let Value::Object(ref mut m) = out { + m.insert("request_id".into(), json!(request_id)); + } + Ok(out) +} + +async fn load_policies(iii: &III, cfg: &RouterConfig) -> Result, IIIError> { + load_typed(iii, cfg, "policies:").await +} + +async fn load_ab_tests(iii: &III, cfg: &RouterConfig) -> Result, IIIError> { + load_typed(iii, cfg, "ab_tests:").await +} + +async fn load_health(iii: &III, cfg: &RouterConfig) -> Result, IIIError> { + load_typed(iii, cfg, "model_health:").await +} + +async fn load_models(iii: &III, cfg: &RouterConfig) -> Result, IIIError> { + load_typed(iii, cfg, "models:").await +} + +async fn load_classifier( + iii: &III, + cfg: &RouterConfig, + id: &str, +) -> Result, IIIError> { + let key = format!("classifier:{}", id); + match state::state_get(iii, &cfg.state_scope, &key).await? { + Some(v) => match serde_json::from_value(v) { + Ok(c) => Ok(Some(c)), + Err(e) => { + tracing::warn!(error = %e, "failed to parse classifier config"); + Ok(None) + } + }, + None => Ok(None), + } +} + +async fn load_typed( + iii: &III, + cfg: &RouterConfig, + prefix: &str, +) -> Result, IIIError> { + let items = state::state_list(iii, &cfg.state_scope, prefix).await?; + let mut out = Vec::with_capacity(items.len()); + for it in items { + let key_hint = it + .as_object() + .and_then(|o| o.get("key")) + .and_then(|k| k.as_str()); + match state::parse_item::(&it) { + Some(parsed) => out.push(parsed), + None => { + tracing::warn!( + scope = %cfg.state_scope, + prefix = %prefix, + key = %key_hint.unwrap_or(""), + target_type = %type_name::(), + "skipping malformed state entry" + ); + } + } + } + Ok(out) +} + +pub fn now_ms() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0) +} diff --git a/llm-router/src/functions/health.rs b/llm-router/src/functions/health.rs new file mode 100644 index 000000000..eb1b308b8 --- /dev/null +++ b/llm-router/src/functions/health.rs @@ -0,0 +1,79 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use iii_sdk::{IIIError, III}; +use serde_json::{json, Value}; + +use crate::config::RouterConfig; +use crate::state; +use crate::types::ModelHealth; + +fn key(model: &str) -> String { + format!("model_health:{}", model) +} + +pub fn update_handler( + iii: III, + cfg: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + let cfg = cfg.clone(); + Box::pin(async move { + let model = payload + .get("model") + .and_then(|v| v.as_str()) + .ok_or_else(|| IIIError::Handler("missing 'model'".into()))? + .trim() + .to_string(); + if model.is_empty() { + return Err(IIIError::Handler("empty 'model'".into())); + } + let mut h: ModelHealth = serde_json::from_value(payload) + .map_err(|e| IIIError::Handler(format!("parse health: {}", e)))?; + if let Some(rate) = h.error_rate { + if !(0.0..=1.0).contains(&rate) || rate.is_nan() { + return Err(IIIError::Handler(format!( + "error_rate must be within 0.0..=1.0 (got {})", + rate + ))); + } + } + h.model = model.clone(); + h.last_checked_ms = crate::functions::decide::now_ms(); + state::state_set( + &iii, + &cfg.state_scope, + &key(&model), + serde_json::to_value(&h).unwrap(), + ) + .await?; + Ok(json!({ "updated": true, "model": model })) + }) + } +} + +pub fn list_handler( + iii: III, + cfg: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |_payload: Value| { + let iii = iii.clone(); + let cfg = cfg.clone(); + Box::pin(async move { + let items = state::state_list(&iii, &cfg.state_scope, "model_health:").await?; + let out: Vec = items + .into_iter() + .filter_map(|it| state::parse_item::(&it)) + .collect(); + Ok(json!({ "models": out, "count": out.len() })) + }) + } +} diff --git a/llm-router/src/functions/mod.rs b/llm-router/src/functions/mod.rs new file mode 100644 index 000000000..26ac0a991 --- /dev/null +++ b/llm-router/src/functions/mod.rs @@ -0,0 +1,7 @@ +pub mod ab; +pub mod classify; +pub mod decide; +pub mod health; +pub mod model; +pub mod policy; +pub mod stats; diff --git a/llm-router/src/functions/model.rs b/llm-router/src/functions/model.rs new file mode 100644 index 000000000..ba3db7798 --- /dev/null +++ b/llm-router/src/functions/model.rs @@ -0,0 +1,104 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use iii_sdk::{IIIError, III}; +use serde_json::{json, Value}; + +use crate::config::RouterConfig; +use crate::state; +use crate::types::ModelRegistration; + +fn key(name: &str) -> String { + format!("models:{}", name) +} + +pub fn register_handler( + iii: III, + cfg: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + let cfg = cfg.clone(); + Box::pin(async move { + let mut m: ModelRegistration = serde_json::from_value(payload) + .map_err(|e| IIIError::Handler(format!("parse model: {}", e)))?; + m.model = m.model.trim().to_string(); + if m.model.is_empty() { + return Err(IIIError::Handler("missing or empty 'model'".into())); + } + for (field, value) in [ + ("input_per_1m", m.input_per_1m), + ("output_per_1m", m.output_per_1m), + ] { + if let Some(v) = value { + if v < 0.0 || v.is_nan() { + return Err(IIIError::Handler(format!( + "{} must be >= 0 (got {})", + field, v + ))); + } + } + } + m.registered_at_ms = crate::functions::decide::now_ms(); + state::state_set( + &iii, + &cfg.state_scope, + &key(&m.model), + serde_json::to_value(&m).unwrap(), + ) + .await?; + Ok(json!({ "registered": true, "model": m.model })) + }) + } +} + +pub fn unregister_handler( + iii: III, + cfg: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + let cfg = cfg.clone(); + Box::pin(async move { + let model = payload + .get("model") + .and_then(|v| v.as_str()) + .ok_or_else(|| IIIError::Handler("missing 'model'".into()))? + .trim() + .to_string(); + if model.is_empty() { + return Err(IIIError::Handler("empty 'model'".into())); + } + state::state_delete(&iii, &cfg.state_scope, &key(&model)).await?; + Ok(json!({ "unregistered": true, "model": model })) + }) + } +} + +pub fn list_handler( + iii: III, + cfg: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |_payload: Value| { + let iii = iii.clone(); + let cfg = cfg.clone(); + Box::pin(async move { + let items = state::state_list(&iii, &cfg.state_scope, "models:").await?; + let out: Vec = items + .into_iter() + .filter_map(|it| state::parse_item::(&it)) + .collect(); + Ok(json!({ "models": out, "count": out.len() })) + }) + } +} diff --git a/llm-router/src/functions/policy.rs b/llm-router/src/functions/policy.rs new file mode 100644 index 000000000..ef29c8c20 --- /dev/null +++ b/llm-router/src/functions/policy.rs @@ -0,0 +1,270 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use iii_sdk::{IIIError, III}; +use rand::SeedableRng; +use serde_json::{json, Value}; +use uuid::Uuid; + +use crate::config::RouterConfig; +use crate::router::{decide, match_policy, DecideContext}; +use crate::state; +use crate::types::{ + AbTest, ClassifierConfig, ModelHealth, ModelRegistration, Policy, RoutingRequest, +}; + +fn key_for(id: &str) -> String { + format!("policies:{}", id) +} + +pub fn create_handler( + iii: III, + cfg: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + let cfg = cfg.clone(); + Box::pin(async move { + let mut p: Policy = parse_policy(payload)?; + if p.id.is_empty() { + p.id = format!("pol-{}", Uuid::new_v4()); + } + validate_policy_semantics(&p)?; + p.created_at_ms = crate::functions::decide::now_ms(); + state::state_set( + &iii, + &cfg.state_scope, + &key_for(&p.id), + serde_json::to_value(&p).unwrap(), + ) + .await?; + Ok(json!({ "policy_id": p.id, "created": true })) + }) + } +} + +pub fn update_handler( + iii: III, + cfg: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + let cfg = cfg.clone(); + Box::pin(async move { + let id = payload + .get("policy_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| IIIError::Handler("missing 'policy_id'".into()))? + .to_string(); + let existing = state::state_get(&iii, &cfg.state_scope, &key_for(&id)) + .await? + .ok_or_else(|| IIIError::Handler(format!("policy not found: {}", id)))?; + let mut p: Policy = serde_json::from_value(existing) + .map_err(|e| IIIError::Handler(format!("parse stored policy: {}", e)))?; + merge_policy(&mut p, &payload)?; + validate_policy_semantics(&p)?; + state::state_set( + &iii, + &cfg.state_scope, + &key_for(&id), + serde_json::to_value(&p).unwrap(), + ) + .await?; + Ok(serde_json::to_value(&p).unwrap_or(Value::Null)) + }) + } +} + +pub fn delete_handler( + iii: III, + cfg: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + let cfg = cfg.clone(); + Box::pin(async move { + let id = payload + .get("policy_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| IIIError::Handler("missing 'policy_id'".into()))? + .to_string(); + state::state_delete(&iii, &cfg.state_scope, &key_for(&id)).await?; + Ok(json!({ "deleted": true, "policy_id": id })) + }) + } +} + +pub fn list_handler( + iii: III, + cfg: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + let cfg = cfg.clone(); + Box::pin(async move { + let tenant = payload + .get("tenant") + .and_then(|v| v.as_str()) + .map(String::from); + let enabled_only = payload + .get("enabled") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + let items = state::state_list(&iii, &cfg.state_scope, "policies:").await?; + let mut out: Vec = items + .into_iter() + .filter_map(|it| state::parse_item::(&it)) + .collect(); + if let Some(t) = &tenant { + out.retain(|p| p.match_rule.tenant.as_deref() == Some(t.as_str())); + } + if enabled_only { + out.retain(|p| p.enabled); + } + Ok(json!({ "policies": out, "count": out.len() })) + }) + } +} + +pub fn test_handler( + iii: III, + cfg: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + let cfg = cfg.clone(); + Box::pin(async move { + let req: RoutingRequest = serde_json::from_value(payload) + .map_err(|e| IIIError::Handler(format!("parse request: {}", e)))?; + + // Mirror the production decide path — load every scope the real + // handler consults so dry-runs can't silently diverge. + let classifier_id = req + .classifier_id + .clone() + .unwrap_or_else(|| cfg.classifier_default_id.clone()); + let policies = load_policies(&iii, &cfg).await?; + let ab_tests = load_list::(&iii, &cfg, "ab_tests:").await?; + let health = load_list::(&iii, &cfg, "model_health:").await?; + let models = load_list::(&iii, &cfg, "models:").await?; + let classifier = match state::state_get( + &iii, + &cfg.state_scope, + &format!("classifier:{}", classifier_id), + ) + .await? + { + Some(v) => serde_json::from_value::(v).ok(), + None => None, + }; + + let matched: Vec<_> = policies + .iter() + .filter(|p| match_policy(&req, p)) + .cloned() + .collect(); + + let mut rng = rand::rngs::StdRng::seed_from_u64(0); + let ctx = DecideContext { + policies: &policies, + ab_tests: &ab_tests, + health: &health, + classifier: classifier.as_ref(), + models: &models, + }; + let decision = decide(&req, ctx, &cfg, &mut rng); + Ok(json!({ + "matched_policies": matched, + "decision": decision, + })) + }) + } +} + +async fn load_policies(iii: &III, cfg: &RouterConfig) -> Result, IIIError> { + load_list::(iii, cfg, "policies:").await +} + +async fn load_list( + iii: &III, + cfg: &RouterConfig, + prefix: &str, +) -> Result, IIIError> { + let items = state::state_list(iii, &cfg.state_scope, prefix).await?; + Ok(items + .into_iter() + .filter_map(|it| state::parse_item::(&it)) + .collect()) +} + +fn parse_policy(payload: Value) -> Result { + let mut v = payload; + if v.get("id").is_none() { + if let Value::Object(ref mut m) = v { + m.insert("id".into(), Value::String(String::new())); + } + } + serde_json::from_value::(v) + .map_err(|e| IIIError::Handler(format!("parse policy: {}", e))) +} + +fn validate_policy_semantics(p: &Policy) -> Result<(), IIIError> { + if p.action.model.trim().is_empty() { + return Err(IIIError::Handler("policy.action.model must be non-empty".into())); + } + if let Some(max) = p.action.max_cost_per_request_usd { + if max < 0.0 || max.is_nan() { + return Err(IIIError::Handler(format!( + "policy.action.max_cost_per_request_usd must be >= 0 (got {})", + max + ))); + } + } + Ok(()) +} + +fn merge_policy(target: &mut Policy, patch: &Value) -> Result<(), IIIError> { + if let Some(n) = patch.get("name").and_then(|v| v.as_str()) { + target.name = n.to_string(); + } + if let Some(m) = patch.get("match") { + target.match_rule = serde_json::from_value(m.clone()) + .map_err(|e| IIIError::Handler(format!("invalid 'match' in patch: {}", e)))?; + } + if let Some(a) = patch.get("action") { + target.action = serde_json::from_value(a.clone()) + .map_err(|e| IIIError::Handler(format!("invalid 'action' in patch: {}", e)))?; + } + if let Some(raw) = patch.get("priority") { + let p = raw + .as_i64() + .ok_or_else(|| IIIError::Handler("invalid 'priority': not an integer".into()))?; + if p < i32::MIN as i64 || p > i32::MAX as i64 { + return Err(IIIError::Handler(format!( + "'priority' out of range for i32: {}", + p + ))); + } + target.priority = p as i32; + } + if let Some(e) = patch.get("enabled").and_then(|v| v.as_bool()) { + target.enabled = e; + } + Ok(()) +} diff --git a/llm-router/src/functions/stats.rs b/llm-router/src/functions/stats.rs new file mode 100644 index 000000000..6e23084df --- /dev/null +++ b/llm-router/src/functions/stats.rs @@ -0,0 +1,163 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use iii_sdk::{IIIError, III}; +use serde_json::{json, Value}; + +use crate::config::RouterConfig; +use crate::state; +use crate::types::RoutingLogEntry; + +// Routing log keys are `routing_log::`. The +// 20-digit zero-padded timestamp makes lexicographic order match chronological +// order, so we can build a bucket-prefixed list per day in the window instead +// of loading every log entry into memory. +// +// Hard cap on total entries per stats call — a stalled consumer shouldn't be +// able to pull megabytes of log into process memory. +const SCAN_HARD_CAP: usize = 50_000; + +pub fn handler( + iii: III, + cfg: Arc, +) -> impl Fn(Value) -> Pin> + Send>> + + Send + + Sync + + 'static { + move |payload: Value| { + let iii = iii.clone(); + let cfg = cfg.clone(); + Box::pin(async move { + let tenant = payload + .get("tenant") + .and_then(|v| v.as_str()) + .map(String::from); + let feature = payload + .get("feature") + .and_then(|v| v.as_str()) + .map(String::from); + let days = payload + .get("days") + .and_then(|v| v.as_u64()) + .unwrap_or(cfg.stats_default_days as u64); + + let now = crate::functions::decide::now_ms(); + let horizon = now.saturating_sub(days.saturating_mul(86_400_000)); + + // Narrow the prefix to the shared leading digits of horizon..now, + // so state::list returns only entries that might match instead of + // the full audit log. + // + // Known limitation: iii-sdk 0.11.2 exposes no cursor/paginated + // state::list, so the scan still materializes a Vec before the + // hard cap applies. Track SDK-side pagination in iii-hq/iii and + // switch to a streaming scan once available. Until then the + // prefix narrowing + hard cap is the best we can do. + let prefix = scan_prefix(horizon, now); + let items = state::state_list(&iii, &cfg.state_scope, &prefix).await?; + + let mut total = 0u64; + let mut scanned = 0usize; + let mut truncated = false; + let mut by_model: std::collections::HashMap = + std::collections::HashMap::new(); + let mut by_policy: std::collections::HashMap = + std::collections::HashMap::new(); + + if items.len() > SCAN_HARD_CAP { + tracing::warn!( + returned = items.len(), + cap = SCAN_HARD_CAP, + "routing_log scan returned more than SCAN_HARD_CAP; aggregate is truncated" + ); + } + + for it in items { + scanned += 1; + if scanned > SCAN_HARD_CAP { + truncated = true; + break; + } + let Some(e) = state::parse_item::(&it) else { + tracing::warn!("skipping malformed routing_log entry"); + continue; + }; + if e.timestamp_ms < horizon { + continue; + } + if let Some(t) = &tenant { + if e.tenant.as_deref() != Some(t.as_str()) { + continue; + } + } + if let Some(f) = &feature { + if e.feature.as_deref() != Some(f.as_str()) { + continue; + } + } + total += 1; + *by_model.entry(e.model_selected.clone()).or_insert(0) += 1; + if let Some(p) = e.policy_matched { + *by_policy.entry(p).or_insert(0) += 1; + } + } + + Ok(json!({ + "total_requests": total, + "days": days, + "by_model": by_model, + "by_policy": by_policy, + "scanned": scanned, + "truncated": truncated, + "scan_hard_cap": SCAN_HARD_CAP, + })) + }) + } +} + +// Build the narrowest log key prefix that still covers [horizon, now]. +// E.g. now=1713696000000, horizon=1713091200000 share the first 4 digits, +// so prefix = "routing_log:1713". +fn scan_prefix(horizon_ms: u64, now_ms: u64) -> String { + let lo = format!("{:020}", horizon_ms); + let hi = format!("{:020}", now_ms); + let mut shared = 0; + for (a, b) in lo.chars().zip(hi.chars()) { + if a == b { + shared += 1; + } else { + break; + } + } + format!("routing_log:{}", &lo[..shared]) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn scan_prefix_narrows_when_window_is_small() { + let now: u64 = 1_713_696_000_000; + let hour_ago = now - 3_600_000; + let prefix = scan_prefix(hour_ago, now); + // Last few digits differ, leading 16 are identical. + assert!(prefix.starts_with("routing_log:")); + assert!(prefix.len() > "routing_log:".len() + 10); + } + + #[test] + fn scan_prefix_falls_back_when_window_is_large() { + // Window of months / years — the first differing 13th-ish digit means + // prefix won't narrow much beyond the zero-padding at the front. We + // only care that we don't over-narrow. + let now: u64 = 2_000_000_000_000; + let very_old: u64 = 1_000_000_000_000; + let prefix = scan_prefix(very_old, now); + assert!(prefix.starts_with("routing_log:")); + // Shared prefix is at most the zero-padding (13 digits + leading zeros). + let suffix = &prefix["routing_log:".len()..]; + assert!(suffix.len() < 13, "over-narrowed: {}", prefix); + } +} diff --git a/llm-router/src/main.rs b/llm-router/src/main.rs new file mode 100644 index 000000000..42a7cb265 --- /dev/null +++ b/llm-router/src/main.rs @@ -0,0 +1,193 @@ +use anyhow::Result; +use clap::Parser; +use iii_sdk::{ + register_worker, InitOptions, OtelConfig, RegisterFunctionMessage, RegisterTriggerInput, +}; +use serde_json::json; +use std::sync::Arc; + +mod config; +mod functions; +mod manifest; +mod router; +mod state; +mod types; + +#[derive(Parser, Debug)] +#[command(name = "iii-llm-router", about = "Policy-based LLM routing brain for iii")] +struct Cli { + #[arg(long, default_value = "./config.yaml")] + config: String, + + #[arg(long, default_value = "ws://127.0.0.1:49134")] + url: String, + + #[arg(long)] + manifest: bool, +} + +#[tokio::main] +async fn main() -> Result<()> { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")), + ) + .init(); + + let cli = Cli::parse(); + + if cli.manifest { + let m = manifest::build_manifest(); + println!("{}", serde_json::to_string_pretty(&m).unwrap()); + return Ok(()); + } + + let router_config = match config::load_config(&cli.config) { + Ok(c) => { + tracing::info!( + scope = %c.state_scope, + classifier = %c.classifier_default_id, + stats_days = c.stats_default_days, + "loaded config" + ); + c + } + Err(e) => { + // Distinguish a missing config (acceptable — use defaults) from a + // present-but-invalid one (abort: running with wrong routing + // settings is worse than not starting). + if is_missing_file(&e) { + tracing::warn!(path = %cli.config, "config file not found, using defaults"); + config::RouterConfig::default() + } else { + tracing::error!(error = %e, path = %cli.config, "invalid config, aborting startup"); + return Err(e); + } + } + }; + let cfg = Arc::new(router_config); + + tracing::info!(url = %cli.url, "connecting to III engine"); + let iii = register_worker( + &cli.url, + InitOptions { + otel: Some(OtelConfig::default()), + ..Default::default() + }, + ); + + register_functions(&iii, cfg.clone()); + register_triggers(&iii)?; + + tracing::info!( + "iii-llm-router registered {} functions and HTTP triggers, ready", + manifest::FUNCTIONS.len() + ); + + tokio::signal::ctrl_c().await?; + tracing::info!("iii-llm-router shutting down"); + iii.shutdown_async().await; + Ok(()) +} + +fn register_functions(iii: &iii_sdk::III, cfg: Arc) { + // iii_sdk::III::register_function_with returns an infallible FunctionRef, + // so no error path to aggregate here — register_triggers below is the one + // that can fail. If the SDK ever makes this fallible we'll want to + // collect errors and abort startup. + let desc_for = |id: &str| -> &'static str { + manifest::FUNCTIONS + .iter() + .find(|(fid, _)| *fid == id) + .map(|(_, d)| *d) + .unwrap_or("") + }; + + macro_rules! reg { + ($id:expr, $handler:expr) => {{ + let msg = RegisterFunctionMessage { + id: $id.to_string(), + description: Some(desc_for($id).to_string()), + request_format: None, + response_format: None, + metadata: None, + invocation: None, + }; + iii.register_function_with(msg, $handler); + }}; + } + + reg!("router::decide", functions::decide::build_handler(iii.clone(), cfg.clone())); + reg!("router::policy_create", functions::policy::create_handler(iii.clone(), cfg.clone())); + reg!("router::policy_update", functions::policy::update_handler(iii.clone(), cfg.clone())); + reg!("router::policy_delete", functions::policy::delete_handler(iii.clone(), cfg.clone())); + reg!("router::policy_list", functions::policy::list_handler(iii.clone(), cfg.clone())); + reg!("router::policy_test", functions::policy::test_handler(iii.clone(), cfg.clone())); + reg!("router::classify", functions::classify::classify_handler(iii.clone(), cfg.clone())); + reg!("router::classifier_config", functions::classify::config_handler(iii.clone(), cfg.clone())); + reg!("router::ab_create", functions::ab::create_handler(iii.clone(), cfg.clone())); + reg!("router::ab_record", functions::ab::record_handler(iii.clone(), cfg.clone())); + reg!("router::ab_report", functions::ab::report_handler(iii.clone(), cfg.clone())); + reg!("router::ab_conclude", functions::ab::conclude_handler(iii.clone(), cfg.clone())); + reg!("router::health_update", functions::health::update_handler(iii.clone(), cfg.clone())); + reg!("router::health_list", functions::health::list_handler(iii.clone(), cfg.clone())); + reg!("router::model_register", functions::model::register_handler(iii.clone(), cfg.clone())); + reg!("router::model_unregister", functions::model::unregister_handler(iii.clone(), cfg.clone())); + reg!("router::model_list", functions::model::list_handler(iii.clone(), cfg.clone())); + reg!("router::stats", functions::stats::handler(iii.clone(), cfg.clone())); +} + +fn register_triggers(iii: &iii_sdk::III) -> Result<()> { + let mut errors: Vec<(String, iii_sdk::IIIError)> = Vec::new(); + for (fn_id, path, method) in [ + ("router::decide", "router/decide", "POST"), + ("router::policy_create", "router/policy/create", "POST"), + ("router::policy_update", "router/policy/update", "POST"), + ("router::policy_delete", "router/policy/delete", "POST"), + ("router::policy_list", "router/policy/list", "GET"), + ("router::policy_test", "router/policy/test", "POST"), + ("router::classify", "router/classify", "POST"), + ("router::classifier_config", "router/classifier", "POST"), + ("router::ab_create", "router/ab/create", "POST"), + ("router::ab_record", "router/ab/record", "POST"), + ("router::ab_report", "router/ab/report", "POST"), + ("router::ab_conclude", "router/ab/conclude", "POST"), + ("router::health_update", "router/health/update", "POST"), + ("router::health_list", "router/health/list", "GET"), + ("router::model_register", "router/model/register", "POST"), + ("router::model_unregister", "router/model/unregister", "POST"), + ("router::model_list", "router/model/list", "GET"), + ("router::stats", "router/stats", "GET"), + ] { + if let Err(e) = iii.register_trigger(RegisterTriggerInput { + trigger_type: "http".to_string(), + function_id: fn_id.to_string(), + config: json!({ "api_path": path, "http_method": method }), + metadata: None, + }) { + tracing::error!(error = %e, "failed to register trigger for {}", fn_id); + errors.push((fn_id.to_string(), e)); + } + } + if !errors.is_empty() { + anyhow::bail!( + "iii-llm-router startup aborted: {} trigger registration(s) failed", + errors.len() + ); + } + Ok(()) +} + +fn is_missing_file(e: &anyhow::Error) -> bool { + // Walk the source chain to avoid brittle string matching on the outer + // with_context() wrapper. + for cause in e.chain() { + if let Some(io_err) = cause.downcast_ref::() { + if io_err.kind() == std::io::ErrorKind::NotFound { + return true; + } + } + } + false +} diff --git a/llm-router/src/manifest.rs b/llm-router/src/manifest.rs new file mode 100644 index 000000000..aefeac4c9 --- /dev/null +++ b/llm-router/src/manifest.rs @@ -0,0 +1,67 @@ +use serde_json::{json, Value}; + +// Canonical list of router functions and their descriptions. Both +// register_functions() in main.rs and build_manifest() below derive from +// this, so the published manifest and the registered handlers can never drift. +pub const FUNCTIONS: &[(&str, &str)] = &[ + ("router::decide", "Pick a model for a request (hot path)"), + ("router::policy_create", "Register a routing policy"), + ("router::policy_update", "Patch a policy"), + ("router::policy_delete", "Remove a policy"), + ("router::policy_list", "List all policies"), + ("router::policy_test", "Dry-run router::decide without logging"), + ("router::classify", "Run prompt-complexity classifier only"), + ("router::classifier_config", "Configure the category→model mapping"), + ("router::ab_create", "Create an A/B test"), + ("router::ab_record", "Record a quality/latency/cost outcome"), + ("router::ab_report", "Aggregate A/B samples"), + ("router::ab_conclude", "Mark an A/B test concluded"), + ("router::health_update", "Update per-model health + latency"), + ("router::health_list", "List health for all models"), + ( + "router::model_register", + "Register a model (name, quality, pricing)", + ), + ("router::model_unregister", "Remove a model registration"), + ("router::model_list", "List registered models"), + ("router::stats", "Usage stats over a window"), +]; + +pub fn build_manifest() -> Value { + let fns: Vec = FUNCTIONS + .iter() + .map(|(id, desc)| json!({ "id": id, "description": desc })) + .collect(); + json!({ + "name": "iii-llm-router", + "version": env!("CARGO_PKG_VERSION"), + "description": "Unopinionated LLM routing brain. Wraps any gateway (LiteLLM/Bifrost/OpenRouter). Models, classifiers, policies, A/B tests, health — all registered at runtime via state.", + "functions": fns, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_manifest_has_required_fields() { + let m = build_manifest(); + assert!(m.get("name").is_some()); + assert!(m.get("version").is_some()); + let fns = m.get("functions").unwrap().as_array().unwrap(); + assert_eq!(fns.len(), FUNCTIONS.len()); + } + + #[test] + fn test_manifest_json_output() { + let s = serde_json::to_string(&build_manifest()).unwrap(); + assert!(s.contains("router::decide")); + assert!(s.contains("router::model_register")); + } + + #[test] + fn functions_const_has_18_entries() { + assert_eq!(FUNCTIONS.len(), 18); + } +} diff --git a/llm-router/src/router.rs b/llm-router/src/router.rs new file mode 100644 index 000000000..ac0bfedba --- /dev/null +++ b/llm-router/src/router.rs @@ -0,0 +1,595 @@ +use crate::config::RouterConfig; +use crate::types::{ + AbTest, AbVariant, ClassifierConfig, ModelHealth, ModelRegistration, Policy, RoutingDecision, + RoutingRequest, +}; +use rand::Rng; + +// Router is intentionally UNOPINIONATED about model names. It only matches +// user-registered policies, classifiers, models, and health records stored in +// engine state. No hardcoded model catalog. + +pub fn match_policy(req: &RoutingRequest, p: &Policy) -> bool { + if !p.enabled { + return false; + } + if let Some(t) = &p.match_rule.tenant { + if req.tenant.as_deref() != Some(t.as_str()) { + return false; + } + } + if let Some(f) = &p.match_rule.feature { + if req.feature.as_deref() != Some(f.as_str()) { + return false; + } + } + if let Some(want_tags) = &p.match_rule.tags { + match &req.tags { + Some(have) => { + if !want_tags.iter().any(|t| have.iter().any(|h| h == t)) { + return false; + } + } + None => return false, + } + } + true +} + +pub fn match_ab(req: &RoutingRequest, t: &AbTest) -> bool { + if t.status != "running" { + return false; + } + if let Some(tn) = &t.match_rule.tenant { + if req.tenant.as_deref() != Some(tn.as_str()) { + return false; + } + } + if let Some(f) = &t.match_rule.feature { + if req.feature.as_deref() != Some(f.as_str()) { + return false; + } + } + true +} + +pub fn pick_ab_variant(variants: &[AbVariant], rng: &mut R) -> Option { + // Accumulate as u64 so N variants with max-u32 weights can't overflow. + let total: u64 = variants.iter().map(|v| v.weight as u64).sum(); + if total == 0 { + return None; + } + let mut pick = rng.gen_range(0..total); + for v in variants { + let w = v.weight as u64; + if pick < w { + return Some(v.model.clone()); + } + pick -= w; + } + None +} + +pub fn policy_specificity(p: &Policy) -> usize { + let mut score = 0usize; + if p.match_rule.tenant.is_some() { + score += 1; + } + if p.match_rule.feature.is_some() { + score += 1; + } + if let Some(tags) = &p.match_rule.tags { + score += tags.len(); + } + score +} + +pub fn skip_unavailable(model: &str, health: &[ModelHealth], error_rate_skip: f64) -> bool { + if let Some(h) = health.iter().find(|h| h.model == model) { + if !h.available { + return true; + } + if let Some(r) = h.error_rate { + if r >= error_rate_skip { + return true; + } + } + } + false +} + +/// Classify a prompt into a category label. Returns (category, confidence). +/// The category is abstract — it does NOT name a model. The user's classifier +/// config is what maps category → model. +pub fn heuristic_complexity(prompt: &str) -> (&'static str, f64) { + let len = prompt.chars().count(); + let has_code = prompt.contains("```") || prompt.contains("fn ") || prompt.contains("def "); + let has_math = prompt.contains('$') || prompt.contains("prove") || prompt.contains("derive"); + let multi_step = + prompt.contains("first") || prompt.contains("then") || prompt.contains("after that"); + + if len < 80 && !has_code && !has_math { + ("simple", 0.85) + } else if len < 300 && !multi_step && !has_math { + ("moderate", 0.75) + } else if len < 1200 && !has_math { + ("complex", 0.8) + } else { + ("expert", 0.9) + } +} + +#[derive(Default)] +pub struct DecideContext<'a> { + pub policies: &'a [Policy], + pub ab_tests: &'a [AbTest], + pub health: &'a [ModelHealth], + pub classifier: Option<&'a ClassifierConfig>, + pub models: &'a [ModelRegistration], +} + +pub fn decide( + req: &RoutingRequest, + ctx: DecideContext<'_>, + cfg: &RouterConfig, + rng: &mut impl Rng, +) -> RoutingDecision { + let mut matched: Vec<&Policy> = ctx.policies.iter().filter(|p| match_policy(req, p)).collect(); + // Deterministic ordering: priority desc, then specificity desc (more + // match-rule fields = more specific), then policy id asc as a stable + // tie-breaker so the same inputs always pick the same policy. + matched.sort_by(|a, b| { + b.priority + .cmp(&a.priority) + .then_with(|| policy_specificity(b).cmp(&policy_specificity(a))) + .then_with(|| a.id.cmp(&b.id)) + }); + + // Path 1: policy match → resolve → health → budget. + if let Some(policy) = matched.first().copied() { + let mut chosen = policy.action.model.clone(); + let mut confidence = 0.9; + let mut reason = format!("policy: {}", policy.name); + + if chosen == "auto" { + let (cls, conf) = heuristic_complexity(&req.prompt); + match ctx.classifier.and_then(|c| c.thresholds.get(cls)) { + Some(m) => { + chosen = m.clone(); + confidence = conf; + reason = format!("policy: {} + classifier: {}", policy.name, cls); + } + None => { + return RoutingDecision { + model: policy.action.fallback.clone().unwrap_or_default(), + reason: format!( + "policy: {} asks auto but no classifier mapping for '{}'", + policy.name, cls + ), + policy_id: Some(policy.id.clone()), + ab_test_id: None, + fallback: None, + confidence: 0.3, + }; + } + } + } + + if skip_unavailable(&chosen, ctx.health, cfg.health_skip_threshold_error_rate) { + if let Some(fb) = &policy.action.fallback { + if skip_unavailable(fb, ctx.health, cfg.health_skip_threshold_error_rate) { + // Both primary and fallback are unhealthy. Return the + // primary with a warning so the caller sees the + // degradation instead of masking it with a bad fallback. + reason = format!( + "{} (primary + fallback both unhealthy — returning primary)", + reason + ); + confidence *= 0.4; + } else { + return RoutingDecision { + model: fb.clone(), + reason: format!("{} (primary unhealthy → fallback)", reason), + policy_id: Some(policy.id.clone()), + ab_test_id: None, + fallback: None, + confidence: confidence * 0.8, + }; + } + } + } + + if let Some(remaining) = req.budget_remaining_usd { + if let Some(max_per_req) = policy.action.max_cost_per_request_usd { + if max_per_req > remaining && remaining > 0.0 { + if let Some(downgraded) = downgrade_to_fit(remaining, req, ctx.models) { + let degraded = skip_unavailable( + &downgraded, + ctx.health, + cfg.health_skip_threshold_error_rate, + ); + let mut dreason = format!("budget constraint: downgraded from {}", chosen); + let mut dconf = confidence * 0.7; + if degraded { + dreason = format!("{} (downgrade target unhealthy)", dreason); + dconf *= 0.5; + } + return RoutingDecision { + model: downgraded, + reason: dreason, + policy_id: Some(policy.id.clone()), + ab_test_id: None, + fallback: policy.action.fallback.clone(), + confidence: dconf, + }; + } + reason = format!( + "{} (over budget but no registered model fits — using original)", + reason + ); + } + } + } + + return RoutingDecision { + model: chosen, + reason, + policy_id: Some(policy.id.clone()), + ab_test_id: None, + fallback: policy.action.fallback.clone(), + confidence, + }; + } + + // Path 2: no policy — try classifier, still subject to health check. + if let Some(classifier) = ctx.classifier { + let (cls, conf) = heuristic_complexity(&req.prompt); + if let Some(m) = classifier.thresholds.get(cls) { + let mut confidence = conf; + let mut reason = format!("no policy, classifier: {}", cls); + if skip_unavailable(m, ctx.health, cfg.health_skip_threshold_error_rate) { + reason = format!("{} (model unhealthy, no policy fallback)", reason); + confidence *= 0.5; + } + return RoutingDecision { + model: m.clone(), + reason, + policy_id: None, + ab_test_id: None, + fallback: None, + confidence, + }; + } + } + + // Path 3: no policy, no classifier — AB test is the last resort. Still + // subject to health check. + if let Some(ab) = ctx.ab_tests.iter().find(|t| match_ab(req, t)) { + if let Some(model) = pick_ab_variant(&ab.variants, rng) { + let mut confidence = 1.0; + let mut reason = format!("ab-test: {}", ab.name); + if skip_unavailable(&model, ctx.health, cfg.health_skip_threshold_error_rate) { + reason = format!("{} (variant unhealthy, no policy fallback)", reason); + confidence *= 0.5; + } + return RoutingDecision { + model, + reason, + policy_id: None, + ab_test_id: Some(ab.id.clone()), + fallback: None, + confidence, + }; + } + } + + RoutingDecision { + model: String::new(), + reason: "no policy matched and no classifier configured".to_string(), + policy_id: None, + ab_test_id: None, + fallback: None, + confidence: 0.0, + } +} + +fn downgrade_to_fit( + remaining_usd: f64, + req: &RoutingRequest, + models: &[ModelRegistration], +) -> Option { + if models.is_empty() { + return None; + } + let mut candidates: Vec<(&ModelRegistration, f64)> = models + .iter() + .filter_map(|m| { + let est = match (m.input_per_1m, m.output_per_1m) { + (Some(i), Some(o)) => (i + o) / 2_000_000.0 * 1_000.0, + _ => return None, + }; + if est > remaining_usd { + return None; + } + if let Some(min_q) = &req.min_quality { + if m.quality.as_deref() != Some(min_q.as_str()) + && !matches_higher_or_equal(&m.quality, min_q) + { + return None; + } + } + Some((m, est)) + }) + .collect(); + candidates.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap()); + candidates.first().map(|(m, _)| m.model.clone()) +} + +fn matches_higher_or_equal(have: &Option, want: &str) -> bool { + const ORDER: &[&str] = &["low", "medium", "high", "flagship"]; + let rank = |s: &str| ORDER.iter().position(|x| *x == s); + match (have.as_deref().and_then(rank), rank(want)) { + (Some(h), Some(w)) => h >= w, + _ => false, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::{PolicyAction, PolicyMatch}; + use rand::SeedableRng; + use std::collections::HashMap; + + fn mk_policy(id: &str, tenant: Option<&str>, feature: Option<&str>, model: &str, priority: i32) -> Policy { + Policy { + id: id.into(), + name: id.into(), + match_rule: PolicyMatch { + tenant: tenant.map(String::from), + feature: feature.map(String::from), + tags: None, + }, + action: PolicyAction { + model: model.into(), + fallback: None, + max_cost_per_request_usd: None, + }, + priority, + enabled: true, + created_at_ms: 0, + } + } + + fn mk_req(tenant: Option<&str>, feature: Option<&str>, prompt: &str) -> RoutingRequest { + RoutingRequest { + tenant: tenant.map(String::from), + feature: feature.map(String::from), + user: None, + prompt: prompt.into(), + tags: None, + budget_remaining_usd: None, + latency_slo_ms: None, + min_quality: None, + classifier_id: None, + } + } + + fn empty_ctx<'a>() -> DecideContext<'a> { + DecideContext::default() + } + + #[test] + fn test_match_policy_tenant_and_feature() { + let p = mk_policy("p1", Some("acme"), Some("support"), "model-a", 100); + assert!(match_policy(&mk_req(Some("acme"), Some("support"), "hi"), &p)); + assert!(!match_policy(&mk_req(Some("other"), Some("support"), "hi"), &p)); + } + + #[test] + fn test_match_policy_disabled_rejected() { + let mut p = mk_policy("p", None, None, "m", 1); + p.enabled = false; + assert!(!match_policy(&mk_req(None, None, "hi"), &p)); + } + + #[test] + fn test_priority_highest_wins() { + let lo = mk_policy("lo", Some("acme"), None, "lo-model", 10); + let hi = mk_policy("hi", Some("acme"), None, "hi-model", 100); + let cfg = RouterConfig::default(); + let mut rng = rand::rngs::StdRng::seed_from_u64(1); + let ctx = DecideContext { + policies: &[lo, hi], + ..empty_ctx() + }; + let d = decide(&mk_req(Some("acme"), None, "hi"), ctx, &cfg, &mut rng); + assert_eq!(d.model, "hi-model"); + assert_eq!(d.policy_id.as_deref(), Some("hi")); + } + + #[test] + fn test_unhealthy_primary_uses_fallback() { + let mut p = mk_policy("p", None, None, "primary-m", 10); + p.action.fallback = Some("fallback-m".into()); + let health = vec![ModelHealth { + model: "primary-m".into(), + available: false, + latency_p99_ms: None, + error_rate: None, + last_checked_ms: 0, + }]; + let cfg = RouterConfig::default(); + let mut rng = rand::rngs::StdRng::seed_from_u64(1); + let ctx = DecideContext { + policies: &[p], + health: &health, + ..empty_ctx() + }; + let d = decide(&mk_req(None, None, "hi"), ctx, &cfg, &mut rng); + assert_eq!(d.model, "fallback-m"); + assert!(d.reason.contains("fallback")); + } + + #[test] + fn test_auto_needs_classifier_mapping() { + let p = mk_policy("auto-p", None, None, "auto", 10); + let cfg = RouterConfig::default(); + let mut thresholds = HashMap::new(); + thresholds.insert("simple".to_string(), "cheap-model".to_string()); + thresholds.insert("moderate".to_string(), "mid-model".to_string()); + thresholds.insert("complex".to_string(), "strong-model".to_string()); + thresholds.insert("expert".to_string(), "frontier-model".to_string()); + let classifier = ClassifierConfig { + id: "default".to_string(), + thresholds, + created_at_ms: 0, + }; + let mut rng = rand::rngs::StdRng::seed_from_u64(1); + let ctx = DecideContext { + policies: &[p], + classifier: Some(&classifier), + ..empty_ctx() + }; + let d = decide(&mk_req(None, None, "hi there"), ctx, &cfg, &mut rng); + assert_eq!(d.model, "cheap-model"); + assert!(d.reason.contains("classifier: simple")); + } + + #[test] + fn test_auto_without_classifier_returns_empty_and_reason() { + let p = mk_policy("auto-p", None, None, "auto", 10); + let cfg = RouterConfig::default(); + let mut rng = rand::rngs::StdRng::seed_from_u64(1); + let ctx = DecideContext { + policies: &[p], + ..empty_ctx() + }; + let d = decide(&mk_req(None, None, "hi"), ctx, &cfg, &mut rng); + assert!(d.reason.contains("no classifier mapping")); + assert_eq!(d.confidence, 0.3); + } + + #[test] + fn test_no_policy_no_classifier_empty_model() { + let cfg = RouterConfig::default(); + let mut rng = rand::rngs::StdRng::seed_from_u64(1); + let d = decide(&mk_req(None, None, "hi"), empty_ctx(), &cfg, &mut rng); + assert!(d.model.is_empty()); + assert!(d.reason.contains("no policy")); + } + + #[test] + fn test_no_policy_with_classifier_falls_through() { + let cfg = RouterConfig::default(); + let mut thresholds = HashMap::new(); + thresholds.insert("simple".to_string(), "cheap".to_string()); + let classifier = ClassifierConfig { + id: "default".to_string(), + thresholds, + created_at_ms: 0, + }; + let mut rng = rand::rngs::StdRng::seed_from_u64(1); + let ctx = DecideContext { + classifier: Some(&classifier), + ..empty_ctx() + }; + let d = decide(&mk_req(None, None, "hi"), ctx, &cfg, &mut rng); + assert_eq!(d.model, "cheap"); + } + + #[test] + fn test_heuristic_returns_category_only() { + assert_eq!(heuristic_complexity("hi").0, "simple"); + let long = "prove that ".repeat(200); + assert_eq!(heuristic_complexity(&long).0, "expert"); + } + + #[test] + fn test_ab_pick_variant() { + let variants = vec![ + AbVariant { model: "a".into(), weight: 50 }, + AbVariant { model: "b".into(), weight: 50 }, + ]; + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + let picked = pick_ab_variant(&variants, &mut rng).unwrap(); + assert!(picked == "a" || picked == "b"); + } + + #[test] + fn test_ab_zero_weight_none() { + let variants = vec![AbVariant { model: "a".into(), weight: 0 }]; + let mut rng = rand::rngs::StdRng::seed_from_u64(1); + assert!(pick_ab_variant(&variants, &mut rng).is_none()); + } + + #[test] + fn test_downgrade_needs_registered_models() { + let mut p = mk_policy("p", None, None, "expensive", 10); + p.action.max_cost_per_request_usd = Some(5.0); + let cfg = RouterConfig::default(); + let mut req = mk_req(None, None, "hi"); + req.budget_remaining_usd = Some(0.01); + let mut rng = rand::rngs::StdRng::seed_from_u64(1); + let ctx = DecideContext { + policies: &[p], + ..empty_ctx() + }; + let d = decide(&req, ctx, &cfg, &mut rng); + assert_eq!(d.model, "expensive"); + assert!(d.reason.contains("no registered model fits")); + } + + #[test] + fn test_downgrade_with_registered_model_picks_cheapest() { + let mut p = mk_policy("p", None, None, "expensive", 10); + p.action.max_cost_per_request_usd = Some(5.0); + let cfg = RouterConfig::default(); + let mut req = mk_req(None, None, "hi"); + req.budget_remaining_usd = Some(0.01); + let models = vec![ + ModelRegistration { + model: "cheap-mini".into(), + quality: Some("low".into()), + input_per_1m: Some(0.5), + output_per_1m: Some(1.0), + provider: None, + max_tokens: None, + metadata: None, + registered_at_ms: 0, + }, + ModelRegistration { + model: "cheap-nano".into(), + quality: Some("low".into()), + input_per_1m: Some(0.1), + output_per_1m: Some(0.4), + provider: None, + max_tokens: None, + metadata: None, + registered_at_ms: 0, + }, + ]; + let mut rng = rand::rngs::StdRng::seed_from_u64(1); + let ctx = DecideContext { + policies: &[p], + models: &models, + ..empty_ctx() + }; + let d = decide(&req, ctx, &cfg, &mut rng); + assert_eq!(d.model, "cheap-nano"); + assert!(d.reason.contains("downgraded")); + } + + #[test] + fn test_skip_unavailable_respects_error_rate() { + let h = vec![ModelHealth { + model: "m".into(), + available: true, + latency_p99_ms: None, + error_rate: Some(0.5), + last_checked_ms: 0, + }]; + assert!(skip_unavailable("m", &h, 0.3)); + assert!(!skip_unavailable("m", &h, 0.8)); + } +} diff --git a/llm-router/src/state.rs b/llm-router/src/state.rs new file mode 100644 index 000000000..c109f66d6 --- /dev/null +++ b/llm-router/src/state.rs @@ -0,0 +1,188 @@ +use iii_sdk::{IIIError, TriggerRequest, III}; +use serde_json::{json, Value}; + +// Per-request timeout for all state helpers on the hot path. Routing has a +// sub-2s SLO so a stalled backend must error fast, not hang. +const STATE_TIMEOUT_MS: u64 = 1_500; + +pub async fn state_get(iii: &III, scope: &str, key: &str) -> Result, IIIError> { + let result = iii + .trigger(TriggerRequest { + function_id: "state::get".to_string(), + payload: json!({ "scope": scope, "key": key }), + action: None, + timeout_ms: Some(STATE_TIMEOUT_MS), + }) + .await; + match result { + Ok(val) => extract_value(&val, scope, key), + Err(e) => { + let msg = e.to_string().to_lowercase(); + if msg.contains("not found") || msg.contains("no such") { + Ok(None) + } else { + Err(e) + } + } + } +} + +pub async fn state_set(iii: &III, scope: &str, key: &str, value: Value) -> Result<(), IIIError> { + iii.trigger(TriggerRequest { + function_id: "state::set".to_string(), + payload: json!({ "scope": scope, "key": key, "value": value }), + action: None, + timeout_ms: Some(STATE_TIMEOUT_MS), + }) + .await?; + Ok(()) +} + +pub async fn state_delete(iii: &III, scope: &str, key: &str) -> Result<(), IIIError> { + iii.trigger(TriggerRequest { + function_id: "state::delete".to_string(), + payload: json!({ "scope": scope, "key": key }), + action: None, + timeout_ms: Some(STATE_TIMEOUT_MS), + }) + .await?; + Ok(()) +} + +pub async fn state_list(iii: &III, scope: &str, prefix: &str) -> Result, IIIError> { + let result = iii + .trigger(TriggerRequest { + function_id: "state::list".to_string(), + payload: json!({ "scope": scope, "prefix": prefix }), + action: None, + timeout_ms: Some(STATE_TIMEOUT_MS), + }) + .await?; + extract_items(&result, scope, prefix) +} + +// state::get envelope. iii-sdk may return { "value": ... } or the value +// directly depending on engine version. Treat null as absent; reject malformed +// responses instead of silently returning None. +fn extract_value(val: &Value, scope: &str, key: &str) -> Result, IIIError> { + if val.is_null() { + return Ok(None); + } + match val { + Value::Object(m) => match m.get("value") { + Some(Value::Null) => Ok(None), + Some(v) => Ok(Some(v.clone())), + None => { + // An object without a "value" key is the raw stored value. + // Preserve prior behavior, don't error. + if m.is_empty() { + Ok(None) + } else { + Ok(Some(Value::Object(m.clone()))) + } + } + }, + other => Ok(Some(other.clone())), + // Suppress unreachable arm — above already covers all variants. + #[allow(unreachable_patterns)] + _ => Err(IIIError::Handler(format!( + "malformed state::get response for scope={} key={}", + scope, key + ))), + } +} + +// Deserialize a state::list item into T. Handles both envelope +// { key, value } and bare-value shapes so callers don't re-implement the +// fallback. Returns None on shape mismatch or deserialize error, both kinds +// of failure are expected in mixed-shape responses. +pub fn parse_item(item: &Value) -> Option { + if let Some(obj) = item.as_object() { + if let Some(v) = obj.get("value") { + if let Ok(parsed) = serde_json::from_value::(v.clone()) { + return Some(parsed); + } + } + } + serde_json::from_value::(item.clone()).ok() +} + +// state::list envelope. Handle three shapes seen across engine versions: +// { "items": [...] } (0.11.0) +// [ ... ] (0.11.2 bare array) +// null (empty) +// Anything else is a hard error, not a silent empty. +fn extract_items(val: &Value, scope: &str, prefix: &str) -> Result, IIIError> { + if val.is_null() { + return Ok(Vec::new()); + } + if let Some(arr) = val.as_array() { + return Ok(arr.clone()); + } + if let Some(items) = val.get("items") { + if let Some(arr) = items.as_array() { + return Ok(arr.clone()); + } + return Err(IIIError::Handler(format!( + "malformed state::list response: 'items' not an array (scope={} prefix={})", + scope, prefix + ))); + } + Err(IIIError::Handler(format!( + "malformed state::list response: missing 'items' and not an array (scope={} prefix={})", + scope, prefix + ))) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn extract_value_unwraps_envelope() { + let v = json!({"value": {"a": 1}}); + assert_eq!( + extract_value(&v, "s", "k").unwrap(), + Some(json!({"a": 1})) + ); + } + + #[test] + fn extract_value_treats_null_value_as_absent() { + let v = json!({"value": null}); + assert_eq!(extract_value(&v, "s", "k").unwrap(), None); + } + + #[test] + fn extract_value_accepts_bare_object() { + let v = json!({"a": 1}); + assert_eq!( + extract_value(&v, "s", "k").unwrap(), + Some(json!({"a": 1})) + ); + } + + #[test] + fn extract_items_accepts_wrapped_array() { + let v = json!({"items": [{"a": 1}]}); + assert_eq!(extract_items(&v, "s", "p").unwrap(), vec![json!({"a": 1})]); + } + + #[test] + fn extract_items_accepts_bare_array() { + let v = json!([{"a": 1}]); + assert_eq!(extract_items(&v, "s", "p").unwrap(), vec![json!({"a": 1})]); + } + + #[test] + fn extract_items_rejects_bad_shape() { + let v = json!({"items": "oops"}); + assert!(extract_items(&v, "s", "p").is_err()); + } + + #[test] + fn extract_items_rejects_missing() { + let v = json!({"other": 1}); + assert!(extract_items(&v, "s", "p").is_err()); + } +} diff --git a/llm-router/src/types.rs b/llm-router/src/types.rs new file mode 100644 index 000000000..a17433151 --- /dev/null +++ b/llm-router/src/types.rs @@ -0,0 +1,198 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +#[serde(deny_unknown_fields)] +pub struct PolicyMatch { + #[serde(default)] + pub tenant: Option, + #[serde(default)] + pub feature: Option, + #[serde(default)] + pub tags: Option>, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PolicyAction { + /// Opaque model identifier. Pass "auto" to defer to the classifier. + /// Router never interprets this — the downstream gateway does. + pub model: String, + #[serde(default)] + pub fallback: Option, + #[serde(default)] + pub max_cost_per_request_usd: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Policy { + pub id: String, + pub name: String, + #[serde(default, rename = "match", alias = "match_rule")] + pub match_rule: PolicyMatch, + pub action: PolicyAction, + #[serde(default = "default_priority")] + pub priority: i32, + #[serde(default = "default_enabled")] + pub enabled: bool, + #[serde(default)] + pub created_at_ms: u64, +} + +fn default_priority() -> i32 { + 100 +} +fn default_enabled() -> bool { + true +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct RoutingRequest { + #[serde(default)] + pub tenant: Option, + #[serde(default)] + pub feature: Option, + #[serde(default)] + pub user: Option, + pub prompt: String, + #[serde(default)] + pub tags: Option>, + #[serde(default)] + pub budget_remaining_usd: Option, + #[serde(default)] + pub latency_slo_ms: Option, + #[serde(default)] + pub min_quality: Option, + #[serde(default)] + pub classifier_id: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RoutingDecision { + pub model: String, + pub reason: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub policy_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub ab_test_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub fallback: Option, + pub confidence: f64, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ModelHealth { + pub model: String, + #[serde(default = "default_available")] + pub available: bool, + #[serde(default)] + pub latency_p99_ms: Option, + #[serde(default)] + pub error_rate: Option, + #[serde(default)] + pub last_checked_ms: u64, +} + +fn default_available() -> bool { + true +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AbVariant { + pub model: String, + pub weight: u32, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct AbTest { + pub id: String, + pub name: String, + #[serde(default, rename = "match", alias = "match_rule")] + pub match_rule: PolicyMatch, + pub variants: Vec, + #[serde(default = "default_metric")] + pub metric: String, + #[serde(default = "default_min_samples")] + pub min_samples: u32, + #[serde(default = "default_max_days")] + pub max_duration_days: u32, + #[serde(default = "default_status")] + pub status: String, + #[serde(default)] + pub created_at_ms: u64, +} + +fn default_metric() -> String { + "quality_score".to_string() +} +fn default_min_samples() -> u32 { + 100 +} +fn default_max_days() -> u32 { + 14 +} +fn default_status() -> String { + "running".to_string() +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AbEvent { + pub test_id: String, + pub variant_model: String, + pub quality_score: f64, + pub latency_ms: u64, + pub cost_usd: f64, + pub recorded_at_ms: u64, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RoutingLogEntry { + pub timestamp_ms: u64, + pub request_id: String, + pub tenant: Option, + pub feature: Option, + pub model_selected: String, + pub policy_matched: Option, + pub ab_test_id: Option, + pub reason: String, + pub cost_usd: Option, +} + +/// Classifier that maps prompt-complexity categories to user-chosen model IDs. +/// Router ships with a simple prompt heuristic — the `thresholds` map is what +/// the user controls to keep the router unopinionated. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ClassifierConfig { + pub id: String, + /// `simple`, `moderate`, `complex`, `expert` → any opaque model ID the gateway understands. + pub thresholds: std::collections::HashMap, + #[serde(default)] + pub created_at_ms: u64, +} + +/// A model registration. The router does NOT know any model names out-of-the-box. +/// Users register whatever model IDs their gateway supports, with optional +/// quality/pricing attributes used only for the downgrade and stats paths. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ModelRegistration { + pub model: String, + /// Any label the user wants: `low`, `medium`, `high`, `flagship`, etc. + /// Used only when the request specifies `min_quality`. + #[serde(default)] + pub quality: Option, + #[serde(default)] + pub input_per_1m: Option, + #[serde(default)] + pub output_per_1m: Option, + #[serde(default)] + pub provider: Option, + #[serde(default)] + pub max_tokens: Option, + #[serde(default)] + pub metadata: Option, + #[serde(default)] + pub registered_at_ms: u64, +}