Files
optimclaw/src/llm/smart_routing.rs
T
716629809c fix: eliminate panic paths in production code (#1184)
* fix: eliminate panic paths in production code and document infallible operations

PolicyRule::new() now returns Result instead of panicking on invalid
caller-supplied regex. CreateJobTool returns ToolError when job_manager
is unconfigured instead of panicking. Remaining infallible unwrap/expect
calls (hardcoded regexes, compile-time constants, guarded accesses)
are annotated with SAFETY comments. Where possible, unwraps are replaced
with safer patterns: split_last(), if-let, match-destructure, and
reusing peek() values.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: use inline lowercase safety comments to match CI pattern

The no-panics CI check greps for '// safety:' (lowercase, inline)
to suppress false positives. Switch from block SAFETY comments to
inline safety comments on the .unwrap() lines.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* test: add regression tests for panic-path fixes

- PolicyRule::new returns Err on invalid regex (not panic)
- CreateJobTool::execute_sandbox returns ToolError when job_manager is None

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: add inline // safety: comments on all infallible unwrap/expect lines

The CI no-panics check requires '// safety:' on the same line as
unwrap()/expect() to suppress false positives. Move safety annotations
from block comments to inline comments on every infallible production
unwrap/expect across all touched files.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* chore: trigger CI with skip-regression-check label

[skip-regression-check]

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* refactor: remove redundant block-level SAFETY comments

Each unwrap/expect now carries its own inline // safety: annotation,
making the standalone block comments above them redundant.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-15 03:17:03 +00:00

1751 lines
61 KiB
Rust

//! Smart routing provider that routes requests to cheap or primary models based on task complexity.
//!
//! Uses a 13-dimension complexity scorer (from PR #208 by @onlyamicrowave) to analyze prompts
//! across reasoning, code, multi-step, domain-specific, creativity, precision, safety, and other
//! dimensions. Pattern overrides provide fast-path routing for obvious cases (greetings → cheap,
//! security audits → primary).
//!
//! This is a decorator that wraps two `LlmProvider`s and implements `LlmProvider` itself,
//! following the same pattern as `RetryProvider`, `CachedProvider`, and `CircuitBreakerProvider`.
//!
//! # Complexity Tiers
//!
//! The scorer produces a 0-100 score mapped to four tiers:
//! - **Flash** (0-15): Greetings, quick lookups → cheap model
//! - **Standard** (16-40): Writing, comparisons → cheap model
//! - **Pro** (41-65): Multi-step analysis, code review → cheap with cascade, or primary
//! - **Frontier** (66+): Security audits, critical decisions → primary model
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use async_trait::async_trait;
use regex::Regex;
use rust_decimal::Decimal;
use crate::llm::error::LlmError;
use crate::llm::provider::{
CompletionRequest, CompletionResponse, LlmProvider, ModelMetadata, Role, ToolCompletionRequest,
ToolCompletionResponse,
};
// ---------------------------------------------------------------------------
// Complexity tiers & scoring
// ---------------------------------------------------------------------------
/// Complexity tier produced by the 13-dimension scorer.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Tier {
/// Simple requests: greetings, quick lookups (score 0-15).
Flash,
/// Standard tasks: writing, comparisons (score 16-40).
Standard,
/// Complex work: multi-step analysis, code review (score 41-65).
Pro,
/// Critical tasks: security audits, high-stakes decisions (score 66+).
Frontier,
}
impl Tier {
/// Convert a complexity score to a tier.
pub fn from_score(score: u32) -> Self {
match score {
0..=15 => Tier::Flash,
16..=40 => Tier::Standard,
41..=65 => Tier::Pro,
_ => Tier::Frontier,
}
}
/// Get a representative score for this tier (used when score is not computed).
pub fn to_score(self) -> u32 {
match self {
Tier::Flash => 8,
Tier::Standard => 28,
Tier::Pro => 52,
Tier::Frontier => 80,
}
}
/// Tier name as string.
pub fn as_str(&self) -> &'static str {
match self {
Tier::Flash => "flash",
Tier::Standard => "standard",
Tier::Pro => "pro",
Tier::Frontier => "frontier",
}
}
}
impl std::fmt::Display for Tier {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.as_str())
}
}
/// Weights for each of the 13 scoring dimensions.
#[derive(Debug, Clone)]
pub struct ScorerWeights {
pub reasoning_words: f32,
pub token_estimate: f32,
pub code_indicators: f32,
pub multi_step: f32,
pub domain_specific: f32,
pub ambiguity: f32,
pub creativity: f32,
pub precision: f32,
pub context_dependency: f32,
pub tool_likelihood: f32,
pub safety_sensitivity: f32,
pub question_complexity: f32,
pub sentence_complexity: f32,
}
impl Default for ScorerWeights {
fn default() -> Self {
Self {
reasoning_words: 0.14,
token_estimate: 0.12,
code_indicators: 0.10,
multi_step: 0.10,
domain_specific: 0.10,
ambiguity: 0.05,
creativity: 0.07,
precision: 0.06,
context_dependency: 0.05,
tool_likelihood: 0.05,
safety_sensitivity: 0.04,
question_complexity: 0.07,
sentence_complexity: 0.05,
}
}
}
/// Default domain-specific keywords for complexity scoring.
pub const DEFAULT_DOMAIN_KEYWORDS: &[&str] = &[
// Infrastructure
"kubernetes",
"k8s",
"docker",
"terraform",
"nginx",
"apache",
"linux",
"unix",
"bash",
"shell",
// Languages & frameworks
"solidity",
"rust",
"typescript",
"react",
"nextjs",
"vue",
"angular",
"svelte",
// Databases
"postgresql",
"postgres",
"mysql",
"mongodb",
"redis",
// APIs & protocols
"graphql",
"grpc",
"protobuf",
"websocket",
"oauth",
"jwt",
"cors",
"csrf",
"xss",
"sql.?injection",
"api",
"rest",
"http",
"https",
"tcp",
"udp",
"dns",
"cdn",
// Cloud & deployment
"aws",
"gcp",
"azure",
"vercel",
"netlify",
"cloudflare",
"ci/cd",
"devops",
// Version control
"git",
"github",
"gitlab",
// Web3 general
"blockchain",
"web3",
"defi",
"nft",
"smart.?contract",
// Ethereum
"ethereum",
"evm",
"anchor",
// NEAR ecosystem
"near",
"near.?sdk",
"near.?api",
"testnet",
"mainnet",
"meteor",
"ledger",
"cold.?wallet",
"rpc",
"indexer",
"relayer",
"cross.?chain",
"intents",
// Fogo/SVM
"fogo",
"svm",
"firedancer",
"paymaster",
"gasless",
"sessions.?sdk",
// Rust/NEAR tooling
"cargo.?near",
"workspaces",
"sandbox",
// Project-specific
"lobo",
"trezu",
"multisig",
"treasury",
"openclaw",
"ironclaw",
];
/// Configuration for the complexity scorer.
#[derive(Debug, Clone, Default)]
pub struct ScorerConfig {
/// Weights for each scoring dimension.
pub weights: ScorerWeights,
/// Custom domain-specific keywords (overrides defaults if provided).
/// Each entry is a word or regex pattern fragment.
pub domain_keywords: Option<Vec<String>>,
}
/// Build a domain regex from a keyword list, with fallback on invalid patterns.
///
/// An empty keyword list falls back to the default keywords so scoring
/// doesn't break when `domain_keywords: Some(vec![])` is configured.
fn build_domain_regex(keywords: &[&str]) -> Regex {
if keywords.is_empty() {
return RE_DOMAIN_DEFAULT.clone();
}
let pattern = format!(r"(?i)\b({})\b", keywords.join("|"));
Regex::new(&pattern).unwrap_or_else(|e| {
tracing::warn!(error = %e, "Invalid domain keywords pattern, using minimal fallback");
Regex::new(r"(?i)\b(api|code|deploy)\b").expect("fallback regex is valid") // safety: hardcoded literal
})
}
/// Breakdown of complexity score by dimension.
#[derive(Debug, Clone)]
pub struct ScoreBreakdown {
/// Total complexity score (0-100).
pub total: u32,
/// Computed tier.
pub tier: Tier,
/// Per-dimension scores (0-100 each).
pub components: HashMap<String, u32>,
/// Human-readable hints about why this score.
pub hints: Vec<String>,
}
// ---------------------------------------------------------------------------
// Static regex patterns (compiled once via LazyLock)
// ---------------------------------------------------------------------------
use std::sync::LazyLock;
static RE_REASONING: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(
r"(?i)\b(why|how|explain|analyze|analyse|compare|contrast|evaluate|assess|reason|think|consider|implications?|consequences?|trade-?offs?|pros?\s*(and|&)\s*cons?|advantages?|disadvantages?|benefits?|drawbacks?|differs?|difference|versus|vs\.?|better|worse|optimal|best|worst)\b"
).expect("RE_REASONING is a valid regex") // safety: hardcoded literal
});
static RE_MULTI_STEP: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(
r"(?i)\b(first|then|next|after|before|finally|step|steps|phase|stages?|process|workflow|sequence|procedure|pipeline|chain|series|order|followed by)\b"
).expect("RE_MULTI_STEP is a valid regex") // safety: hardcoded literal
});
static RE_CREATIVITY: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(
r"(?i)\b(write|create|generate|compose|design|imagine|brainstorm|ideate|draft|invent|story|poem|essay|article|blog|content|narrative|script|summarize|summarise|rewrite|paraphrase|translate|adapt|tweet|post|thread|outline|structure|format|style|tone|voice)\b"
).expect("RE_CREATIVITY is a valid regex") // safety: hardcoded literal
});
static RE_PRECISION: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(
r"(?i)\b(\d{4}|\d+\.\d+|exactly|precisely|specific|accurate|correct|verify|confirm|date|time|number|calculate|compute|measure|count)\b"
).expect("RE_PRECISION is a valid regex") // safety: hardcoded literal
});
static RE_CODE: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(
r"(?i)(`{1,3}|```|function|const|let|var|import|export|class|def |async|await|=>|\.ts|\.js|\.py|\.rs|\.go|\.sol|\(\)|\[\]|\{\}|<[A-Z][a-z]+>|useState|useEffect|npm|yarn|pnpm|cargo|pip|implement|rebase|merge|commit|branch|PR|pull.?request|columns?|migrations?|module|refactor|debug|fix|bug|error|schema|database|query)"
).expect("RE_CODE is a valid regex") // safety: hardcoded literal
});
static RE_TOOL: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(
r"(?i)\b(file|read|write|search|fetch|run|execute|check|look up|find|open|save|send|post|get|download|upload|install|deploy|build|compile|test|add|update|remove|delete|modify|change|edit|create|resolve|push|pull|clone)\b"
).expect("RE_TOOL is a valid regex") // safety: hardcoded literal
});
static RE_SAFETY: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(
r"(?i)\b(password|secret|private|confidential|medical|legal|financial|personal|sensitive|ssn|credit.?card|auth|token|key|encrypt|decrypt|hash|vulnerability|exploit|attack|breach)\b"
).expect("RE_SAFETY is a valid regex") // safety: hardcoded literal
});
static RE_CONTEXT: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(
r"(?i)\b(previous|earlier|above|before|last|that|those|it|they|we discussed|you said|mentioned|remember|recall|as I said|like I mentioned)\b"
).expect("RE_CONTEXT is a valid regex") // safety: hardcoded literal
});
static RE_VAGUE: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(r"(?i)\b(it|this|that|something|stuff|thing|things)\b")
.expect("RE_VAGUE is a valid regex") // safety: hardcoded literal
});
static RE_OPEN_ENDED: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(r"(?i)\b(why|how|what if|explain|describe|elaborate|discuss)\b")
.expect("RE_OPEN_ENDED is a valid regex") // safety: hardcoded literal
});
static RE_CONJUNCTIONS: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(
r"(?i)\b(and|but|or|however|therefore|because|although|while|whereas|moreover|furthermore)\b",
)
.expect("RE_CONJUNCTIONS is a valid regex") // safety: hardcoded literal
});
static RE_TIER_HINT: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(r"(?i)\[tier:(flash|standard|pro|frontier)\]")
.expect("RE_TIER_HINT is a valid regex") // safety: hardcoded literal
});
/// Default domain regex, compiled once from `DEFAULT_DOMAIN_KEYWORDS`.
static RE_DOMAIN_DEFAULT: LazyLock<Regex> =
LazyLock::new(|| build_domain_regex(DEFAULT_DOMAIN_KEYWORDS));
// ---------------------------------------------------------------------------
// Pattern overrides (fast-path before scoring)
// ---------------------------------------------------------------------------
/// A compiled pattern override entry.
struct PatternOverride {
regex: Regex,
tier: Tier,
}
/// Default pattern overrides, compiled once.
static DEFAULT_OVERRIDES: LazyLock<Vec<PatternOverride>> = LazyLock::new(|| {
vec![
// Flash tier: greetings and acknowledgments
PatternOverride {
regex: Regex::new(
r"(?i)^(hi|hello|hey|thanks|ok|sure|yes|no|yep|nope|cool|nice|great|got it)$",
)
.expect("greeting pattern is valid"), // safety: hardcoded literal
tier: Tier::Flash,
},
// Flash tier: quick lookups (end-anchored to avoid matching complex questions
// like "What time complexity is merge sort?")
PatternOverride {
regex: Regex::new(
r"(?i)^what(?:'s|\s+is)?\s+(?:the\s+)?(time|date|day|weather)\b(?:\s+(?:is\s+it|today|now|in\s+\S+))?[?.!]*$",
)
.expect("lookup pattern is valid"), // safety: hardcoded literal
tier: Tier::Flash,
},
// Frontier tier: security audits
PatternOverride {
regex: Regex::new(r"(?i)security.*(audit|review|scan)")
.expect("security audit pattern is valid"), // safety: hardcoded literal
tier: Tier::Frontier,
},
PatternOverride {
regex: Regex::new(r"(?i)vulnerabilit(y|ies).*(review|scan|check|audit)")
.expect("vulnerability pattern is valid"), // safety: hardcoded literal
tier: Tier::Frontier,
},
// Pro tier: production deployments
PatternOverride {
regex: Regex::new(r"(?i)deploy.*(mainnet|production)")
.expect("deploy pattern is valid"), // safety: hardcoded literal
tier: Tier::Pro,
},
PatternOverride {
regex: Regex::new(r"(?i)production.*(deploy|release|push)")
.expect("production pattern is valid"), // safety: hardcoded literal
tier: Tier::Pro,
},
]
});
// ---------------------------------------------------------------------------
// Scoring functions
// ---------------------------------------------------------------------------
/// Count regex matches in text.
fn count_matches(re: &Regex, text: &str) -> usize {
re.find_iter(text).count()
}
/// Score a prompt's complexity across 13 dimensions.
///
/// Returns a `ScoreBreakdown` with a total score (0-100) and per-dimension breakdown.
pub fn score_complexity(prompt: &str) -> ScoreBreakdown {
score_complexity_with_config(prompt, &ScorerConfig::default())
}
/// Score with custom configuration (weights + domain keywords).
///
/// If you will call this repeatedly with the same config, prefer
/// [`score_complexity_with_regex`] and pre-build the domain regex once.
pub fn score_complexity_with_config(prompt: &str, config: &ScorerConfig) -> ScoreBreakdown {
let domain_regex = match &config.domain_keywords {
Some(custom) => {
let refs: Vec<&str> = custom.iter().map(|s| s.as_str()).collect();
build_domain_regex(&refs)
}
None => RE_DOMAIN_DEFAULT.clone(),
};
score_complexity_internal(prompt, &config.weights, &domain_regex)
}
/// Score with a pre-compiled domain regex (avoids rebuilding per call).
pub fn score_complexity_with_regex(
prompt: &str,
weights: &ScorerWeights,
domain_regex: &Regex,
) -> ScoreBreakdown {
score_complexity_internal(prompt, weights, domain_regex)
}
/// Internal scoring implementation.
fn score_complexity_internal(
prompt: &str,
weights: &ScorerWeights,
domain_regex: &Regex,
) -> ScoreBreakdown {
let mut hints = Vec::new();
let mut components = HashMap::new();
// Check for explicit tier hint (e.g. "[tier:flash]")
if let Some(caps) = RE_TIER_HINT.captures(prompt) {
let tier_str = caps.get(1).expect("capture group 1 exists").as_str(); // safety: RE_TIER_HINT has group 1
let tier = match tier_str.to_lowercase().as_str() {
"flash" => Tier::Flash,
"standard" => Tier::Standard,
"pro" => Tier::Pro,
"frontier" => Tier::Frontier,
// The regex only captures valid tiers, so this is defensive.
other => {
tracing::error!(tier = %other, "Unexpected tier in hint despite regex constraint");
Tier::Standard
}
};
hints.push(format!("Explicit tier hint: {tier}"));
return ScoreBreakdown {
total: tier.to_score(),
tier,
components,
hints,
};
}
// Token estimate (based on char count): <20 chars = 0, >=520 chars = 100
let char_count = prompt.len();
let token_score = ((char_count as i32 - 20).max(0) as f32 / 5.0).min(100.0) as u32;
components.insert("token_estimate".to_string(), token_score);
if char_count > 200 {
hints.push(format!("Long prompt ({char_count} chars)"));
}
// Reasoning words
let reasoning_count = count_matches(&RE_REASONING, prompt);
let reasoning_score = (reasoning_count * 50).min(100) as u32;
components.insert("reasoning_words".to_string(), reasoning_score);
if reasoning_count >= 2 {
hints.push(format!("reasoning_words: {reasoning_count} matches"));
}
// Multi-step
let multi_step_count = count_matches(&RE_MULTI_STEP, prompt);
let multi_step_score = (multi_step_count * 50).min(100) as u32;
components.insert("multi_step".to_string(), multi_step_score);
if multi_step_count >= 2 {
hints.push(format!("multi_step: {multi_step_count} matches"));
}
// Creativity
let creativity_count = count_matches(&RE_CREATIVITY, prompt);
let creativity_score = (creativity_count * 50).min(100) as u32;
components.insert("creativity".to_string(), creativity_score);
if creativity_count >= 2 {
hints.push(format!("creativity: {creativity_count} matches"));
}
// Precision
let precision_count = count_matches(&RE_PRECISION, prompt);
let precision_score = (precision_count * 50).min(100) as u32;
components.insert("precision".to_string(), precision_score);
// Code indicators
let code_count = count_matches(&RE_CODE, prompt);
let code_score = (code_count * 50).min(100) as u32;
components.insert("code_indicators".to_string(), code_score);
if code_count >= 2 {
hints.push(format!("code_indicators: {code_count} matches"));
}
// Tool likelihood
let tool_count = count_matches(&RE_TOOL, prompt);
let tool_score = (tool_count * 50).min(100) as u32;
components.insert("tool_likelihood".to_string(), tool_score);
// Safety sensitivity
let safety_count = count_matches(&RE_SAFETY, prompt);
let safety_score = (safety_count * 50).min(100) as u32;
components.insert("safety_sensitivity".to_string(), safety_score);
if safety_count >= 1 {
hints.push(format!("safety_sensitivity: {safety_count} matches"));
}
// Context dependency
let context_count = count_matches(&RE_CONTEXT, prompt);
let context_score = (context_count * 50).min(100) as u32;
components.insert("context_dependency".to_string(), context_score);
// Domain specific
let domain_count = count_matches(domain_regex, prompt);
let domain_score = (domain_count * 50).min(100) as u32;
components.insert("domain_specific".to_string(), domain_score);
if domain_count >= 2 {
hints.push(format!("domain_specific: {domain_count} matches"));
}
// Ambiguity (vague pronouns)
let vague_count = count_matches(&RE_VAGUE, prompt);
let ambiguity_score = (vague_count * 25).min(100) as u32;
components.insert("ambiguity".to_string(), ambiguity_score);
// Question complexity
let question_marks = prompt.matches('?').count();
let open_ended_count = count_matches(&RE_OPEN_ENDED, prompt);
let question_score = ((question_marks * 20) + (open_ended_count * 25)).min(100) as u32;
components.insert("question_complexity".to_string(), question_score);
if question_marks >= 2 {
hints.push(format!("Multiple questions: {question_marks}"));
}
// Sentence complexity (commas, semicolons, conjunctions)
let commas = prompt.matches(',').count();
let semicolons = prompt.matches(';').count();
let conjunctions = count_matches(&RE_CONJUNCTIONS, prompt);
let clauses = commas + (semicolons * 2) + conjunctions;
let sentence_score = (clauses * 12).min(100) as u32;
components.insert("sentence_complexity".to_string(), sentence_score);
if clauses >= 5 {
hints.push(format!("Complex structure: {clauses} clauses"));
}
// Calculate weighted total using data-driven iteration
let total: f32 = [
("reasoning_words", weights.reasoning_words),
("token_estimate", weights.token_estimate),
("code_indicators", weights.code_indicators),
("multi_step", weights.multi_step),
("domain_specific", weights.domain_specific),
("ambiguity", weights.ambiguity),
("creativity", weights.creativity),
("precision", weights.precision),
("context_dependency", weights.context_dependency),
("tool_likelihood", weights.tool_likelihood),
("safety_sensitivity", weights.safety_sensitivity),
("question_complexity", weights.question_complexity),
("sentence_complexity", weights.sentence_complexity),
]
.iter()
.map(|(name, weight)| components.get(*name).copied().unwrap_or(0) as f32 * weight)
.sum();
// Multi-dimensional boost: +30% when 3+ dimensions fire above threshold
let triggered_dimensions = components.values().filter(|&&v| v > 20).count();
let total = if triggered_dimensions >= 3 {
hints.push(format!(
"Multi-dimensional ({triggered_dimensions} triggers)"
));
total * 1.3
} else if triggered_dimensions >= 2 {
total * 1.15
} else {
total
};
// Clamp to 0-100
let total = (total as u32).clamp(0, 100);
let tier = Tier::from_score(total);
ScoreBreakdown {
total,
tier,
components,
hints,
}
}
// ---------------------------------------------------------------------------
// TaskComplexity (provider-level classification)
// ---------------------------------------------------------------------------
/// Classification of a request's complexity, determining which model handles it.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TaskComplexity {
/// Short, simple queries -> cheap model (Flash + Standard tiers)
Simple,
/// Ambiguous complexity -> cheap model first, cascade to primary if uncertain (Pro tier)
Moderate,
/// Code generation, analysis, multi-step reasoning -> primary model (Frontier tier)
Complex,
}
impl From<Tier> for TaskComplexity {
fn from(tier: Tier) -> Self {
match tier {
Tier::Flash | Tier::Standard => TaskComplexity::Simple,
Tier::Pro => TaskComplexity::Moderate,
Tier::Frontier => TaskComplexity::Complex,
}
}
}
// ---------------------------------------------------------------------------
// SmartRoutingConfig & Provider
// ---------------------------------------------------------------------------
/// Configuration for the smart routing provider.
#[derive(Debug, Clone)]
pub struct SmartRoutingConfig {
/// Enable cascade mode: retry with primary if cheap model response seems uncertain.
pub cascade_enabled: bool,
/// Custom domain keywords for the scorer (None uses defaults).
pub domain_keywords: Option<Vec<String>>,
}
impl Default for SmartRoutingConfig {
fn default() -> Self {
Self {
cascade_enabled: true,
domain_keywords: None,
}
}
}
/// Atomic counters for routing observability.
struct SmartRoutingStats {
total_requests: AtomicU64,
cheap_requests: AtomicU64,
primary_requests: AtomicU64,
cascade_escalations: AtomicU64,
}
impl SmartRoutingStats {
fn new() -> Self {
Self {
total_requests: AtomicU64::new(0),
cheap_requests: AtomicU64::new(0),
primary_requests: AtomicU64::new(0),
cascade_escalations: AtomicU64::new(0),
}
}
}
/// Snapshot of routing statistics for external consumption.
#[derive(Debug, Clone)]
pub struct SmartRoutingSnapshot {
pub total_requests: u64,
pub cheap_requests: u64,
pub primary_requests: u64,
pub cascade_escalations: u64,
}
/// Smart routing provider that classifies task complexity and routes to the appropriate model.
///
/// - `complete()` — scores complexity across 13 dimensions, checks pattern overrides, then
/// routes to cheap or primary model. Moderate tasks use cascade (try cheap, escalate if uncertain).
/// - `complete_with_tools()` — always routes to primary (tool use requires reliable structured output)
pub struct SmartRoutingProvider {
primary: Arc<dyn LlmProvider>,
cheap: Arc<dyn LlmProvider>,
config: SmartRoutingConfig,
scorer_config: ScorerConfig,
/// Pre-compiled domain regex (built once at construction time).
domain_regex: Regex,
stats: SmartRoutingStats,
}
impl SmartRoutingProvider {
/// Create a new smart routing provider wrapping a primary and cheap provider.
pub fn new(
primary: Arc<dyn LlmProvider>,
cheap: Arc<dyn LlmProvider>,
config: SmartRoutingConfig,
) -> Self {
let scorer_config = ScorerConfig {
weights: ScorerWeights::default(),
domain_keywords: config.domain_keywords.clone(),
};
let domain_regex = match &scorer_config.domain_keywords {
Some(custom) => {
let refs: Vec<&str> = custom.iter().map(|s| s.as_str()).collect();
build_domain_regex(&refs)
}
None => RE_DOMAIN_DEFAULT.clone(),
};
Self {
primary,
cheap,
config,
scorer_config,
domain_regex,
stats: SmartRoutingStats::new(),
}
}
/// Get a snapshot of routing statistics.
pub fn stats(&self) -> SmartRoutingSnapshot {
SmartRoutingSnapshot {
total_requests: self.stats.total_requests.load(Ordering::Relaxed),
cheap_requests: self.stats.cheap_requests.load(Ordering::Relaxed),
primary_requests: self.stats.primary_requests.load(Ordering::Relaxed),
cascade_escalations: self.stats.cascade_escalations.load(Ordering::Relaxed),
}
}
/// Classify the complexity of a request based on its last user message.
///
/// Priority: explicit tier hints > pattern overrides > 13-dimension scorer.
fn classify(&self, request: &CompletionRequest) -> TaskComplexity {
let last_user_msg = request
.messages
.iter()
.rev()
.find(|m| m.role == Role::User)
.map(|m| m.content.as_str())
.unwrap_or("");
// Normalize: trim whitespace so anchored regexes and token scoring are consistent.
let last_user_msg = last_user_msg.trim();
// Highest priority: explicit tier hints (e.g. "[tier:flash]")
if let Some(caps) = RE_TIER_HINT.captures(last_user_msg) {
// SAFETY: RE_TIER_HINT has exactly one capture group; get(1) is guaranteed Some after match.
let tier_str = caps.get(1).expect("capture group 1 exists").as_str(); // safety: RE_TIER_HINT has group 1
let tier = match tier_str.to_lowercase().as_str() {
"flash" => Tier::Flash,
"standard" => Tier::Standard,
"pro" => Tier::Pro,
"frontier" => Tier::Frontier,
other => {
tracing::error!(tier = %other, "Unexpected tier in hint despite regex constraint");
Tier::Standard
}
};
let complexity = TaskComplexity::from(tier);
tracing::trace!(
%tier,
?complexity,
"Smart routing: explicit tier hint"
);
return complexity;
}
// Fast-path: check pattern overrides
for po in DEFAULT_OVERRIDES.iter() {
if po.regex.is_match(last_user_msg) {
let complexity = TaskComplexity::from(po.tier);
tracing::trace!(
tier = %po.tier,
?complexity,
"Smart routing: pattern override matched"
);
return complexity;
}
}
// Full 13-dimension scoring (uses pre-compiled domain regex)
let breakdown = score_complexity_with_regex(
last_user_msg,
&self.scorer_config.weights,
&self.domain_regex,
);
let complexity = TaskComplexity::from(breakdown.tier);
tracing::trace!(
score = breakdown.total,
tier = %breakdown.tier,
?complexity,
hints = ?breakdown.hints,
"Smart routing: scored complexity"
);
complexity
}
/// Check if a response from the cheap model shows uncertainty, warranting escalation.
fn response_is_uncertain(response: &CompletionResponse) -> bool {
let content = response.content.trim();
// Empty response is always uncertain
if content.is_empty() {
return true;
}
let lower = content.to_lowercase();
// Uncertainty signals
let uncertainty_patterns = [
"i'm not sure",
"i am not sure",
"i don't know",
"i do not know",
"i'm unable to",
"i am unable to",
"i cannot",
"i can't",
"beyond my capabilities",
"beyond my ability",
"i'm not able to",
"i am not able to",
"i don't have enough",
"i do not have enough",
"i need more context",
"i need more information",
"could you clarify",
"could you provide more",
"i'm not confident",
"i am not confident",
];
uncertainty_patterns.iter().any(|p| lower.contains(p))
}
}
#[async_trait]
impl LlmProvider for SmartRoutingProvider {
fn model_name(&self) -> &str {
self.primary.model_name()
}
fn cost_per_token(&self) -> (Decimal, Decimal) {
self.primary.cost_per_token()
}
fn cache_write_multiplier(&self) -> Decimal {
self.primary.cache_write_multiplier()
}
fn cache_read_discount(&self) -> Decimal {
self.primary.cache_read_discount()
}
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
self.stats.total_requests.fetch_add(1, Ordering::Relaxed);
let complexity = self.classify(&request);
match complexity {
TaskComplexity::Simple => {
tracing::trace!(
model = %self.cheap.model_name(),
"Smart routing: Simple task -> cheap model"
);
self.stats.cheap_requests.fetch_add(1, Ordering::Relaxed);
self.cheap.complete(request).await
}
TaskComplexity::Complex => {
tracing::trace!(
model = %self.primary.model_name(),
"Smart routing: Complex task -> primary model"
);
self.stats.primary_requests.fetch_add(1, Ordering::Relaxed);
self.primary.complete(request).await
}
TaskComplexity::Moderate => {
if self.config.cascade_enabled {
tracing::trace!(
model = %self.cheap.model_name(),
"Smart routing: Moderate task -> cheap model (cascade enabled)"
);
self.stats.cheap_requests.fetch_add(1, Ordering::Relaxed);
let response = self.cheap.complete(request.clone()).await?;
if Self::response_is_uncertain(&response) {
tracing::info!(
cheap_model = %self.cheap.model_name(),
primary_model = %self.primary.model_name(),
"Smart routing: Escalating to primary (cheap model response uncertain)"
);
self.stats
.cascade_escalations
.fetch_add(1, Ordering::Relaxed);
self.stats.primary_requests.fetch_add(1, Ordering::Relaxed);
self.primary.complete(request).await
} else {
Ok(response)
}
} else {
// Without cascade, moderate tasks go to cheap model
tracing::trace!(
model = %self.cheap.model_name(),
"Smart routing: Moderate task -> cheap model (cascade disabled)"
);
self.stats.cheap_requests.fetch_add(1, Ordering::Relaxed);
self.cheap.complete(request).await
}
}
}
}
/// Tool use always goes to the primary model for reliable structured output.
async fn complete_with_tools(
&self,
request: ToolCompletionRequest,
) -> Result<ToolCompletionResponse, LlmError> {
self.stats.total_requests.fetch_add(1, Ordering::Relaxed);
self.stats.primary_requests.fetch_add(1, Ordering::Relaxed);
tracing::trace!(
model = %self.primary.model_name(),
"Smart routing: Tool use -> primary model (always)"
);
self.primary.complete_with_tools(request).await
}
async fn list_models(&self) -> Result<Vec<String>, LlmError> {
self.primary.list_models().await
}
async fn model_metadata(&self) -> Result<ModelMetadata, LlmError> {
self.primary.model_metadata().await
}
fn effective_model_name(&self, requested_model: Option<&str>) -> String {
self.primary.effective_model_name(requested_model)
}
fn active_model_name(&self) -> String {
self.primary.active_model_name()
}
fn set_model(&self, model: &str) -> Result<(), LlmError> {
self.primary.set_model(model)
}
fn calculate_cost(&self, input_tokens: u32, output_tokens: u32) -> Decimal {
self.primary.calculate_cost(input_tokens, output_tokens)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::llm::ChatMessage;
use crate::testing::StubLlm;
fn default_config() -> SmartRoutingConfig {
SmartRoutingConfig::default()
}
// -----------------------------------------------------------------------
// Score complexity: tier boundaries
// -----------------------------------------------------------------------
#[test]
fn score_empty_prompt_is_flash() {
let result = score_complexity("");
assert_eq!(result.tier, Tier::Flash);
assert!(result.total <= 15);
}
#[test]
fn score_simple_greeting_is_flash() {
let result = score_complexity("Hi");
assert_eq!(result.tier, Tier::Flash);
assert!(result.total <= 15);
}
#[test]
fn score_quick_question_is_flash_or_standard() {
let result = score_complexity("What time is it?");
assert!(
result.tier == Tier::Flash || result.tier == Tier::Standard,
"Expected Flash or Standard, got {:?} (score {})",
result.tier,
result.total
);
}
#[test]
fn score_code_task_is_standard_or_higher() {
let result = score_complexity("Implement a function to sort an array in TypeScript");
assert!(
result.tier == Tier::Standard || result.tier == Tier::Pro,
"Expected Standard or Pro, got {:?} (score {})",
result.tier,
result.total
);
}
#[test]
fn score_complex_analysis_is_at_least_standard() {
let result = score_complexity(
"Explain why React uses a virtual DOM and compare it to Svelte's approach. \
Consider the trade-offs for performance and developer experience.",
);
assert!(
result.total >= 20,
"Expected score >= 20, got {}",
result.total
);
assert!(
result.tier == Tier::Standard || result.tier == Tier::Pro,
"Expected Standard or Pro, got {:?}",
result.tier
);
}
#[test]
fn score_security_audit_prompt_is_at_least_standard() {
let result = score_complexity(
"Analyze this Solidity contract for reentrancy vulnerabilities, \
check for authentication bypass, and provide a security audit report.",
);
assert!(
result.total >= 16,
"Expected score >= 16, got {}",
result.total
);
}
// -----------------------------------------------------------------------
// Score complexity: individual dimensions
// -----------------------------------------------------------------------
#[test]
fn score_reasoning_dimension() {
let result = score_complexity("Why is this better? Explain the trade-offs and compare");
let reasoning = result
.components
.get("reasoning_words")
.copied()
.unwrap_or(0);
assert!(
reasoning >= 100,
"Expected reasoning >= 100, got {reasoning}"
);
}
#[test]
fn score_multi_step_dimension() {
let result = score_complexity(
"First, read the file at src/auth.ts. Then analyze it for security issues. \
After that, write a detailed report.",
);
let multi_step = result.components.get("multi_step").copied().unwrap_or(0);
assert!(
multi_step >= 100,
"Expected multi_step >= 100, got {multi_step}"
);
assert!(result.hints.iter().any(|h| h.contains("multi_step")));
}
#[test]
fn score_code_dimension() {
let result = score_complexity("Fix the bug in the async function, refactor the module");
let code = result
.components
.get("code_indicators")
.copied()
.unwrap_or(0);
assert!(code >= 50, "Expected code_indicators >= 50, got {code}");
}
#[test]
fn score_safety_dimension() {
let result = score_complexity("Store the password and encrypt the auth token");
let safety = result
.components
.get("safety_sensitivity")
.copied()
.unwrap_or(0);
assert!(safety >= 100, "Expected safety >= 100, got {safety}");
}
#[test]
fn score_domain_dimension() {
let result = score_complexity("Deploy the kubernetes cluster on aws with terraform");
let domain = result
.components
.get("domain_specific")
.copied()
.unwrap_or(0);
assert!(
domain >= 100,
"Expected domain_specific >= 100, got {domain}"
);
}
#[test]
fn score_creativity_dimension() {
let result = score_complexity("Write a blog post about design patterns, then summarize");
let creativity = result.components.get("creativity").copied().unwrap_or(0);
assert!(
creativity >= 100,
"Expected creativity >= 100, got {creativity}"
);
}
#[test]
fn score_question_complexity_dimension() {
let result = score_complexity("Why does this fail? How can I fix it? What if I try X?");
let qc = result
.components
.get("question_complexity")
.copied()
.unwrap_or(0);
assert!(qc >= 60, "Expected question_complexity >= 60, got {qc}");
assert!(
result
.hints
.iter()
.any(|h| h.contains("Multiple questions"))
);
}
#[test]
fn score_sentence_complexity_dimension() {
let result = score_complexity(
"This is complex, because it has commas, and conjunctions, \
however it also has semicolons; moreover, it keeps going, and going",
);
let sc = result
.components
.get("sentence_complexity")
.copied()
.unwrap_or(0);
assert!(sc >= 60, "Expected sentence_complexity >= 60, got {sc}");
}
#[test]
fn score_token_estimate_for_long_prompt() {
let long_prompt = "a ".repeat(300); // 600 chars
let result = score_complexity(&long_prompt);
let token = result
.components
.get("token_estimate")
.copied()
.unwrap_or(0);
assert!(token >= 80, "Expected token_estimate >= 80, got {token}");
}
#[test]
fn score_token_estimate_for_short_prompt() {
let result = score_complexity("hi");
let token = result
.components
.get("token_estimate")
.copied()
.unwrap_or(0);
assert_eq!(token, 0, "Expected token_estimate == 0, got {token}");
}
// -----------------------------------------------------------------------
// Score complexity: multi-dimensional boost
// -----------------------------------------------------------------------
#[test]
fn score_multi_dimensional_boost() {
// This triggers reasoning, multi-step, code, domain, creativity, safety
let result = score_complexity(
"First, explain why the kubernetes deployment fails. \
Then refactor the auth module to fix the vulnerability. \
After that, write a security report comparing the approaches.",
);
assert!(
result.hints.iter().any(|h| h.contains("Multi-dimensional")),
"Expected multi-dimensional boost, hints: {:?}",
result.hints
);
}
// -----------------------------------------------------------------------
// Score complexity: explicit tier hint
// -----------------------------------------------------------------------
#[test]
fn score_explicit_tier_hint_flash() {
let result = score_complexity("[tier:flash] This looks complex but override to flash");
assert_eq!(result.tier, Tier::Flash);
assert!(
result
.hints
.iter()
.any(|h| h.contains("Explicit tier hint"))
);
}
#[test]
fn score_explicit_tier_hint_frontier() {
let result = score_complexity("[tier:frontier] Simple question but I want the best");
assert_eq!(result.tier, Tier::Frontier);
}
#[test]
fn score_explicit_tier_hint_case_insensitive() {
let result = score_complexity("[tier:PRO] some message");
assert_eq!(result.tier, Tier::Pro);
}
// -----------------------------------------------------------------------
// Score complexity: custom domain keywords
// -----------------------------------------------------------------------
#[test]
fn score_custom_domain_keywords_override_defaults() {
// Default keywords should match "kubernetes"
let default_result = score_complexity("How do I deploy kubernetes?");
let default_domain = default_result
.components
.get("domain_specific")
.copied()
.unwrap_or(0);
assert!(
default_domain > 0,
"Default keywords should match 'kubernetes'"
);
// Custom keywords that DON'T include kubernetes
let config = ScorerConfig {
weights: ScorerWeights::default(),
domain_keywords: Some(vec!["mycompany".to_string(), "myproduct".to_string()]),
};
let custom_result = score_complexity_with_config("How do I deploy kubernetes?", &config);
let custom_domain = custom_result
.components
.get("domain_specific")
.copied()
.unwrap_or(0);
assert_eq!(
custom_domain, 0,
"Custom keywords shouldn't match 'kubernetes'"
);
// Custom keywords should match their own terms
let custom_result2 =
score_complexity_with_config("Tell me about myproduct features", &config);
let custom_domain2 = custom_result2
.components
.get("domain_specific")
.copied()
.unwrap_or(0);
assert!(
custom_domain2 > 0,
"Custom keywords should match 'myproduct'"
);
}
// -----------------------------------------------------------------------
// Score complexity: edge cases
// -----------------------------------------------------------------------
#[test]
fn score_whitespace_only_is_flash() {
let result = score_complexity(" \n\t ");
assert_eq!(result.tier, Tier::Flash);
}
#[test]
fn score_single_word_no_keywords() {
let result = score_complexity("banana");
assert!(
result.tier == Tier::Flash || result.tier == Tier::Standard,
"Single non-keyword word should be Flash or Standard, got {:?}",
result.tier
);
}
#[test]
fn score_very_long_prompt_is_at_least_standard() {
let long = "Tell me about ".to_string() + &"things ".repeat(200);
let result = score_complexity(&long);
assert!(
result.total >= 16,
"Very long prompt should score at least Standard, got {}",
result.total
);
}
#[test]
fn score_all_dimensions_have_entries() {
let result = score_complexity(
"First, explain why the function fails. Then write a fix and deploy it.",
);
let expected_keys = [
"reasoning_words",
"token_estimate",
"code_indicators",
"multi_step",
"domain_specific",
"ambiguity",
"creativity",
"precision",
"context_dependency",
"tool_likelihood",
"safety_sensitivity",
"question_complexity",
"sentence_complexity",
];
for key in &expected_keys {
assert!(
result.components.contains_key(*key),
"Missing component: {key}"
);
}
}
#[test]
fn score_is_clamped_to_100() {
// Trigger every dimension hard
let prompt = "First, explain why the kubernetes docker terraform deployment on aws fails. \
Then analyze the security vulnerability and compare the trade-offs. \
After that, write a detailed blog post report with code examples: \
```rust\nfn main() {}\n``` \
Calculate exactly how many steps are needed? Why? How? \
Deploy to production mainnet. Review the authentication token password.";
let result = score_complexity(prompt);
assert!(
result.total <= 100,
"Score should be clamped to 100, got {}",
result.total
);
}
// -----------------------------------------------------------------------
// Pattern overrides
// -----------------------------------------------------------------------
#[test]
fn pattern_override_greeting_is_simple() {
let primary = Arc::new(StubLlm::new("p").with_model_name("primary"));
let cheap = Arc::new(StubLlm::new("c").with_model_name("cheap"));
let provider = SmartRoutingProvider::new(primary, cheap, default_config());
let req = CompletionRequest::new(vec![ChatMessage::user("Hi")]);
let complexity = provider.classify(&req);
assert_eq!(complexity, TaskComplexity::Simple);
}
#[test]
fn pattern_override_security_audit_is_complex() {
let primary = Arc::new(StubLlm::new("p").with_model_name("primary"));
let cheap = Arc::new(StubLlm::new("c").with_model_name("cheap"));
let provider = SmartRoutingProvider::new(primary, cheap, default_config());
let req = CompletionRequest::new(vec![ChatMessage::user(
"Please do a security audit of this contract",
)]);
let complexity = provider.classify(&req);
assert_eq!(complexity, TaskComplexity::Complex);
}
#[test]
fn pattern_override_production_deploy_is_moderate() {
let primary = Arc::new(StubLlm::new("p").with_model_name("primary"));
let cheap = Arc::new(StubLlm::new("c").with_model_name("cheap"));
let provider = SmartRoutingProvider::new(primary, cheap, default_config());
let req = CompletionRequest::new(vec![ChatMessage::user("Deploy this to production")]);
let complexity = provider.classify(&req);
assert_eq!(complexity, TaskComplexity::Moderate);
}
#[test]
fn pattern_override_time_question_is_simple() {
let primary = Arc::new(StubLlm::new("p").with_model_name("primary"));
let cheap = Arc::new(StubLlm::new("c").with_model_name("cheap"));
let provider = SmartRoutingProvider::new(primary, cheap, default_config());
let req = CompletionRequest::new(vec![ChatMessage::user("What time is it?")]);
let complexity = provider.classify(&req);
assert_eq!(complexity, TaskComplexity::Simple);
}
#[test]
fn pattern_override_time_does_not_match_complex_questions() {
// The quick-lookup override regex should NOT match "What time complexity..."
// because it's end-anchored. Verify the regex itself doesn't fire.
let overrides = &*DEFAULT_OVERRIDES;
let lookup_override = overrides
.iter()
.find(|po| po.tier == Tier::Flash && po.regex.as_str().contains("time"))
.expect("time lookup override exists");
assert!(
!lookup_override
.regex
.is_match("What time complexity is merge sort?"),
"Time override should not match 'What time complexity is merge sort?'"
);
// But it should still match actual time lookups
assert!(lookup_override.regex.is_match("What time is it?"));
assert!(lookup_override.regex.is_match("what's the date today?"));
}
#[test]
fn empty_domain_keywords_uses_defaults() {
// An empty custom keywords list should fall back to defaults, not produce
// a broken regex that matches empty strings everywhere.
let config = ScorerConfig {
domain_keywords: Some(vec![]),
..ScorerConfig::default()
};
let result = score_complexity_with_config("deploy kubernetes to mainnet", &config);
// Should still detect domain keywords via the default fallback
assert!(
result
.components
.get("domain_specific")
.copied()
.unwrap_or(0)
> 0,
"Empty custom keywords should fall back to defaults"
);
}
// -----------------------------------------------------------------------
// Tier → TaskComplexity mapping
// -----------------------------------------------------------------------
#[test]
fn tier_to_task_complexity_mapping() {
assert_eq!(TaskComplexity::from(Tier::Flash), TaskComplexity::Simple);
assert_eq!(TaskComplexity::from(Tier::Standard), TaskComplexity::Simple);
assert_eq!(TaskComplexity::from(Tier::Pro), TaskComplexity::Moderate);
assert_eq!(
TaskComplexity::from(Tier::Frontier),
TaskComplexity::Complex
);
}
#[test]
fn tier_from_score_boundaries() {
assert_eq!(Tier::from_score(0), Tier::Flash);
assert_eq!(Tier::from_score(15), Tier::Flash);
assert_eq!(Tier::from_score(16), Tier::Standard);
assert_eq!(Tier::from_score(40), Tier::Standard);
assert_eq!(Tier::from_score(41), Tier::Pro);
assert_eq!(Tier::from_score(65), Tier::Pro);
assert_eq!(Tier::from_score(66), Tier::Frontier);
assert_eq!(Tier::from_score(100), Tier::Frontier);
}
#[test]
fn tier_display() {
assert_eq!(Tier::Flash.as_str(), "flash");
assert_eq!(Tier::Frontier.to_string(), "frontier");
}
// -----------------------------------------------------------------------
// Uncertainty detection
// -----------------------------------------------------------------------
#[test]
fn detects_uncertain_short_response() {
let response = CompletionResponse {
content: "I'm not sure.".to_string(),
input_tokens: 10,
output_tokens: 5,
finish_reason: crate::llm::FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
};
assert!(SmartRoutingProvider::response_is_uncertain(&response));
}
#[test]
fn detects_empty_response_as_uncertain() {
let response = CompletionResponse {
content: "".to_string(),
input_tokens: 10,
output_tokens: 0,
finish_reason: crate::llm::FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
};
assert!(SmartRoutingProvider::response_is_uncertain(&response));
}
#[test]
fn short_confident_response_is_not_uncertain() {
let response = CompletionResponse {
content: "Yes.".to_string(),
input_tokens: 10,
output_tokens: 1,
finish_reason: crate::llm::FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
};
assert!(!SmartRoutingProvider::response_is_uncertain(&response));
}
#[test]
fn confident_response_is_not_uncertain() {
let response = CompletionResponse {
content: "The answer is 42. This is a well-known constant from the Hitchhiker's Guide."
.to_string(),
input_tokens: 10,
output_tokens: 20,
finish_reason: crate::llm::FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
};
assert!(!SmartRoutingProvider::response_is_uncertain(&response));
}
// -----------------------------------------------------------------------
// Provider routing tests
// -----------------------------------------------------------------------
fn make_request(content: &str) -> CompletionRequest {
CompletionRequest::new(vec![ChatMessage::user(content)])
}
fn make_tool_request() -> ToolCompletionRequest {
ToolCompletionRequest::new(vec![ChatMessage::user("implement a search")], vec![])
}
#[tokio::test]
async fn simple_task_routes_to_cheap() {
let primary = Arc::new(StubLlm::new("primary-response").with_model_name("primary"));
let cheap = Arc::new(StubLlm::new("cheap-response").with_model_name("cheap"));
let router = SmartRoutingProvider::new(
primary.clone(),
cheap.clone(),
SmartRoutingConfig {
cascade_enabled: false,
..default_config()
},
);
let resp = router.complete(make_request("hello")).await.unwrap();
assert_eq!(resp.content, "cheap-response");
assert_eq!(cheap.calls(), 1);
assert_eq!(primary.calls(), 0);
}
#[tokio::test]
async fn complex_task_routes_to_primary() {
let primary = Arc::new(StubLlm::new("primary-response").with_model_name("primary"));
let cheap = Arc::new(StubLlm::new("cheap-response").with_model_name("cheap"));
let router = SmartRoutingProvider::new(primary.clone(), cheap.clone(), default_config());
// Security audit triggers Frontier via pattern override → Complex → primary
let resp = router
.complete(make_request(
"Please do a security audit of this smart contract",
))
.await
.unwrap();
assert_eq!(resp.content, "primary-response");
assert_eq!(primary.calls(), 1);
assert_eq!(cheap.calls(), 0);
}
#[tokio::test]
async fn tool_use_always_routes_to_primary() {
let primary = Arc::new(StubLlm::new("primary-response").with_model_name("primary"));
let cheap = Arc::new(StubLlm::new("cheap-response").with_model_name("cheap"));
let router = SmartRoutingProvider::new(primary.clone(), cheap.clone(), default_config());
let resp = router
.complete_with_tools(make_tool_request())
.await
.unwrap();
assert_eq!(resp.content, Some("primary-response".to_string()));
assert_eq!(primary.calls(), 1);
assert_eq!(cheap.calls(), 0);
}
#[tokio::test]
async fn stats_increment_correctly() {
let primary = Arc::new(StubLlm::new("primary").with_model_name("primary"));
let cheap = Arc::new(StubLlm::new("cheap").with_model_name("cheap"));
let router = SmartRoutingProvider::new(
primary,
cheap,
SmartRoutingConfig {
cascade_enabled: false,
..default_config()
},
);
// Simple → cheap (greeting pattern override)
router.complete(make_request("hello")).await.unwrap();
// Complex → primary (security audit pattern override → Frontier)
router
.complete(make_request("security audit review"))
.await
.unwrap();
// Tool use → primary
router
.complete_with_tools(make_tool_request())
.await
.unwrap();
let stats = router.stats();
assert_eq!(stats.total_requests, 3);
assert_eq!(stats.cheap_requests, 1);
assert_eq!(stats.primary_requests, 2);
assert_eq!(stats.cascade_escalations, 0);
}
#[tokio::test]
async fn cascade_escalates_on_uncertain_response() {
let primary = Arc::new(StubLlm::new("primary-response").with_model_name("primary"));
let cheap = Arc::new(StubLlm::new("I'm not sure about that.").with_model_name("cheap"));
let router = SmartRoutingProvider::new(
primary.clone(),
cheap.clone(),
SmartRoutingConfig {
cascade_enabled: true,
..default_config()
},
);
// A Pro-tier task (triggers Moderate → cascade)
let resp = router
.complete(make_request("Deploy this to production"))
.await
.unwrap();
// Should have escalated to primary
assert_eq!(resp.content, "primary-response");
assert_eq!(cheap.calls(), 1);
assert_eq!(primary.calls(), 1);
let stats = router.stats();
assert_eq!(stats.cascade_escalations, 1);
}
#[tokio::test]
async fn cascade_does_not_escalate_on_confident_response() {
let primary = Arc::new(StubLlm::new("primary-response").with_model_name("primary"));
let cheap = Arc::new(
StubLlm::new("Deployed successfully to production mainnet.").with_model_name("cheap"),
);
let router = SmartRoutingProvider::new(
primary.clone(),
cheap.clone(),
SmartRoutingConfig {
cascade_enabled: true,
..default_config()
},
);
let resp = router
.complete(make_request("Deploy this to production"))
.await
.unwrap();
// Should NOT have escalated
assert!(resp.content.contains("Deployed successfully"));
assert_eq!(cheap.calls(), 1);
assert_eq!(primary.calls(), 0);
let stats = router.stats();
assert_eq!(stats.cascade_escalations, 0);
}
#[tokio::test]
async fn model_name_returns_primary() {
let primary = Arc::new(StubLlm::new("ok").with_model_name("sonnet"));
let cheap = Arc::new(StubLlm::new("ok").with_model_name("haiku"));
let router = SmartRoutingProvider::new(primary, cheap, default_config());
assert_eq!(router.model_name(), "sonnet");
assert_eq!(router.active_model_name(), "sonnet");
}
#[tokio::test]
async fn tier_hint_overrides_pattern_override() {
// "[tier:flash] security audit review" has both a Flash tier hint and
// a Frontier pattern override. Tier hints should win.
let primary = Arc::new(StubLlm::new("primary").with_model_name("primary"));
let cheap = Arc::new(StubLlm::new("cheap").with_model_name("cheap"));
let router = SmartRoutingProvider::new(
primary.clone(),
cheap.clone(),
SmartRoutingConfig {
cascade_enabled: false,
..default_config()
},
);
router
.complete(make_request("[tier:flash] security audit review"))
.await
.unwrap();
// Tier hint → Flash → Simple → cheap model
assert_eq!(cheap.calls(), 1);
assert_eq!(primary.calls(), 0);
}
#[tokio::test]
async fn trimmed_greeting_matches_override() {
// Trailing whitespace should not prevent the greeting override from matching.
let primary = Arc::new(StubLlm::new("primary").with_model_name("primary"));
let cheap = Arc::new(StubLlm::new("cheap").with_model_name("cheap"));
let router = SmartRoutingProvider::new(
primary.clone(),
cheap.clone(),
SmartRoutingConfig {
cascade_enabled: false,
..default_config()
},
);
router.complete(make_request(" hello \n")).await.unwrap();
// Should match greeting override → Flash → Simple → cheap model
assert_eq!(cheap.calls(), 1);
assert_eq!(primary.calls(), 0);
}
}