diff --git a/src/app.rs b/src/app.rs index 22982e5a..48407eaf 100644 --- a/src/app.rs +++ b/src/app.rs @@ -289,94 +289,14 @@ impl AppBuilder { /// Phase 3: Initialize LLM provider chain. /// - /// Creates the primary provider, then wraps with failover, circuit - /// breaker, and response cache as configured. + /// Delegates to `build_provider_chain` which applies all decorators + /// (retry, smart routing, failover, circuit breaker, response cache). #[allow(clippy::type_complexity)] pub fn init_llm( &self, ) -> Result<(Arc, Option>), anyhow::Error> { - use crate::llm::{ - CachedProvider, CircuitBreakerConfig, CircuitBreakerProvider, CooldownConfig, - FailoverProvider, ResponseCacheConfig, create_cheap_llm_provider, create_llm_provider, - create_llm_provider_with_config, - }; - - let llm = create_llm_provider(&self.config.llm, self.session.clone())?; - tracing::info!("LLM provider initialized: {}", llm.model_name()); - - // Wrap in failover if a fallback model is configured - let llm: Arc = if let Some(fallback_model) = - self.config.llm.nearai.fallback_model.as_ref() - { - if fallback_model == &self.config.llm.nearai.model { - tracing::warn!( - "fallback_model is the same as primary model, failover may not be effective" - ); - } - let mut fallback_config = self.config.llm.nearai.clone(); - fallback_config.model = fallback_model.clone(); - let fallback = create_llm_provider_with_config(&fallback_config, self.session.clone())?; - tracing::info!( - primary = %llm.model_name(), - fallback = %fallback.model_name(), - "LLM failover enabled" - ); - let cooldown_config = CooldownConfig { - cooldown_duration: std::time::Duration::from_secs( - self.config.llm.nearai.failover_cooldown_secs, - ), - failure_threshold: self.config.llm.nearai.failover_cooldown_threshold, - }; - Arc::new(FailoverProvider::with_cooldown( - vec![llm, fallback], - cooldown_config, - )?) - } else { - llm - }; - - // Wrap in circuit breaker if configured - let llm: Arc = - if let Some(threshold) = self.config.llm.nearai.circuit_breaker_threshold { - let cb_config = CircuitBreakerConfig { - failure_threshold: threshold, - recovery_timeout: std::time::Duration::from_secs( - self.config.llm.nearai.circuit_breaker_recovery_secs, - ), - ..CircuitBreakerConfig::default() - }; - tracing::info!( - threshold, - recovery_secs = self.config.llm.nearai.circuit_breaker_recovery_secs, - "LLM circuit breaker enabled" - ); - Arc::new(CircuitBreakerProvider::new(llm, cb_config)) - } else { - llm - }; - - // Wrap in response cache if configured - let llm: Arc = if self.config.llm.nearai.response_cache_enabled { - let rc_config = ResponseCacheConfig { - ttl: std::time::Duration::from_secs(self.config.llm.nearai.response_cache_ttl_secs), - max_entries: self.config.llm.nearai.response_cache_max_entries, - }; - tracing::info!( - ttl_secs = self.config.llm.nearai.response_cache_ttl_secs, - max_entries = self.config.llm.nearai.response_cache_max_entries, - "LLM response cache enabled" - ); - Arc::new(CachedProvider::new(llm, rc_config)) - } else { - llm - }; - - // Cheap LLM for lightweight tasks - let cheap_llm = create_cheap_llm_provider(&self.config.llm, self.session.clone())?; - if let Some(ref cheap) = cheap_llm { - tracing::info!("Cheap LLM provider initialized: {}", cheap.model_name()); - } - + let (llm, cheap_llm) = + crate::llm::build_provider_chain(&self.config.llm, self.session.clone())?; Ok((llm, cheap_llm)) } diff --git a/src/config/llm.rs b/src/config/llm.rs index a134284e..35e4ae48 100644 --- a/src/config/llm.rs +++ b/src/config/llm.rs @@ -167,6 +167,10 @@ pub struct NearAiConfig { /// Number of consecutive retryable failures before a provider enters /// cooldown (default: 3). pub failover_cooldown_threshold: u32, + /// Enable cascade mode for smart routing: when a moderate-complexity task + /// gets an uncertain response from the cheap model, re-send to primary. + /// Default: true. + pub smart_routing_cascade: bool, } impl LlmConfig { @@ -232,6 +236,7 @@ impl LlmConfig { response_cache_max_entries: parse_optional_env("RESPONSE_CACHE_MAX_ENTRIES", 1000)?, failover_cooldown_secs: parse_optional_env("LLM_FAILOVER_COOLDOWN_SECS", 300)?, failover_cooldown_threshold: parse_optional_env("LLM_FAILOVER_THRESHOLD", 3)?, + smart_routing_cascade: parse_optional_env("SMART_ROUTING_CASCADE", true)?, }; // Resolve provider-specific configs based on backend diff --git a/src/llm/mod.rs b/src/llm/mod.rs index e03a0482..efd550e2 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -17,6 +17,7 @@ pub mod response_cache; pub mod retry; mod rig_adapter; pub mod session; +pub mod smart_routing; pub use circuit_breaker::{CircuitBreakerConfig, CircuitBreakerProvider}; pub use failover::{CooldownConfig, FailoverProvider}; @@ -33,6 +34,7 @@ pub use response_cache::{CachedProvider, ResponseCacheConfig}; pub use retry::{RetryConfig, RetryProvider}; pub use rig_adapter::RigAdapter; pub use session::{SessionConfig, SessionManager, create_session_manager}; +pub use smart_routing::{SmartRoutingConfig, SmartRoutingProvider, TaskComplexity}; use std::sync::Arc; @@ -273,6 +275,147 @@ pub fn create_cheap_llm_provider( )?))) } +/// Build the full LLM provider chain with all configured wrappers. +/// +/// Applies decorators in this order: +/// 1. Raw provider (from config) +/// 2. RetryProvider (per-provider retry with exponential backoff) +/// 3. SmartRoutingProvider (cheap/primary split when cheap model is configured) +/// 4. FailoverProvider (fallback model when primary fails) +/// 5. CircuitBreakerProvider (fast-fail when backend is degraded) +/// 6. CachedProvider (in-memory response cache) +/// +/// Also returns a separate cheap LLM provider for heartbeat/evaluation (not +/// part of the chain — it's a standalone provider for explicitly cheap tasks). +/// +/// This is the single source of truth for provider chain construction, +/// called by both `main.rs` and `app.rs`. +#[allow(clippy::type_complexity)] +pub fn build_provider_chain( + config: &LlmConfig, + session: Arc, +) -> Result<(Arc, Option>), LlmError> { + let llm = create_llm_provider(config, session.clone())?; + tracing::info!("LLM provider initialized: {}", llm.model_name()); + + // 1. Retry + let retry_config = RetryConfig { + max_retries: config.nearai.max_retries, + }; + let llm: Arc = if retry_config.max_retries > 0 { + tracing::info!( + max_retries = retry_config.max_retries, + "LLM retry wrapper enabled" + ); + Arc::new(RetryProvider::new(llm, retry_config.clone())) + } else { + llm + }; + + // 2. Smart routing (cheap/primary split) + let llm: Arc = if let Some(ref cheap_model) = config.nearai.cheap_model { + let mut cheap_config = config.nearai.clone(); + cheap_config.model = cheap_model.clone(); + let cheap = create_llm_provider_with_config(&cheap_config, session.clone())?; + let cheap: Arc = if retry_config.max_retries > 0 { + Arc::new(RetryProvider::new(cheap, retry_config.clone())) + } else { + cheap + }; + tracing::info!( + primary = %llm.model_name(), + cheap = %cheap.model_name(), + "Smart routing enabled" + ); + Arc::new(SmartRoutingProvider::new( + llm, + cheap, + SmartRoutingConfig { + cascade_enabled: config.nearai.smart_routing_cascade, + ..SmartRoutingConfig::default() + }, + )) + } else { + llm + }; + + // 3. Failover + let llm: Arc = if let Some(ref fallback_model) = config.nearai.fallback_model { + if fallback_model == &config.nearai.model { + tracing::warn!( + "fallback_model is the same as primary model, failover may not be effective" + ); + } + let mut fallback_config = config.nearai.clone(); + fallback_config.model = fallback_model.clone(); + let fallback = create_llm_provider_with_config(&fallback_config, session.clone())?; + tracing::info!( + primary = %llm.model_name(), + fallback = %fallback.model_name(), + "LLM failover enabled" + ); + let fallback: Arc = if retry_config.max_retries > 0 { + Arc::new(RetryProvider::new(fallback, retry_config.clone())) + } else { + fallback + }; + let cooldown_config = CooldownConfig { + cooldown_duration: std::time::Duration::from_secs(config.nearai.failover_cooldown_secs), + failure_threshold: config.nearai.failover_cooldown_threshold, + }; + Arc::new(FailoverProvider::with_cooldown( + vec![llm, fallback], + cooldown_config, + )?) + } else { + llm + }; + + // 4. Circuit breaker + let llm: Arc = if let Some(threshold) = config.nearai.circuit_breaker_threshold + { + let cb_config = CircuitBreakerConfig { + failure_threshold: threshold, + recovery_timeout: std::time::Duration::from_secs( + config.nearai.circuit_breaker_recovery_secs, + ), + ..CircuitBreakerConfig::default() + }; + tracing::info!( + threshold, + recovery_secs = config.nearai.circuit_breaker_recovery_secs, + "LLM circuit breaker enabled" + ); + Arc::new(CircuitBreakerProvider::new(llm, cb_config)) + } else { + llm + }; + + // 5. Response cache + let llm: Arc = if config.nearai.response_cache_enabled { + let rc_config = ResponseCacheConfig { + ttl: std::time::Duration::from_secs(config.nearai.response_cache_ttl_secs), + max_entries: config.nearai.response_cache_max_entries, + }; + tracing::info!( + ttl_secs = config.nearai.response_cache_ttl_secs, + max_entries = config.nearai.response_cache_max_entries, + "LLM response cache enabled" + ); + Arc::new(CachedProvider::new(llm, rc_config)) + } else { + llm + }; + + // Standalone cheap LLM for heartbeat/evaluation (not part of the chain) + let cheap_llm = create_cheap_llm_provider(config, session)?; + if let Some(ref cheap) = cheap_llm { + tracing::info!("Cheap LLM provider initialized: {}", cheap.model_name()); + } + + Ok((llm, cheap_llm)) +} + #[cfg(test)] mod tests { use super::*; @@ -296,6 +439,7 @@ mod tests { response_cache_max_entries: 1000, failover_cooldown_secs: 300, failover_cooldown_threshold: 3, + smart_routing_cascade: true, } } diff --git a/src/llm/nearai_chat.rs b/src/llm/nearai_chat.rs index ccc996c7..8d87dc37 100644 --- a/src/llm/nearai_chat.rs +++ b/src/llm/nearai_chat.rs @@ -766,6 +766,7 @@ mod tests { response_cache_max_entries: 1000, failover_cooldown_secs: 300, failover_cooldown_threshold: 3, + smart_routing_cascade: true, } } diff --git a/src/llm/smart_routing.rs b/src/llm/smart_routing.rs new file mode 100644 index 00000000..b8aa24ce --- /dev/null +++ b/src/llm/smart_routing.rs @@ -0,0 +1,700 @@ +//! Smart routing provider that routes requests to cheap or primary models based on task complexity. +//! +//! Inspired by RelayPlane's cost-reduction approach: simple tasks (status checks, greetings, +//! short questions) go to a cheap model (e.g. Haiku), while complex tasks (code generation, +//! analysis, multi-step reasoning) go to the primary model (e.g. Sonnet/Opus). +//! +//! This is a decorator that wraps two `LlmProvider`s and implements `LlmProvider` itself, +//! following the same pattern as `RetryProvider`, `CachedProvider`, and `CircuitBreakerProvider`. + +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use async_trait::async_trait; +use rust_decimal::Decimal; + +use crate::error::LlmError; +use crate::llm::provider::{ + CompletionRequest, CompletionResponse, LlmProvider, ModelMetadata, Role, ToolCompletionRequest, + ToolCompletionResponse, +}; + +/// 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 + Simple, + /// Ambiguous complexity -> cheap model first, cascade to primary if uncertain + Moderate, + /// Code generation, analysis, multi-step reasoning -> primary model + Complex, +} + +/// 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, + /// Message length threshold below which a message may be classified as Simple (default: 200). + pub simple_max_chars: usize, + /// Message length threshold above which a message is classified as Complex (default: 1000). + pub complex_min_chars: usize, +} + +impl Default for SmartRoutingConfig { + fn default() -> Self { + Self { + cascade_enabled: true, + simple_max_chars: 200, + complex_min_chars: 1000, + } + } +} + +/// 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()` — classifies and routes to cheap or primary model +/// - `complete_with_tools()` — always routes to primary (tool use requires reliable structured output) +pub struct SmartRoutingProvider { + primary: Arc, + cheap: Arc, + config: SmartRoutingConfig, + stats: SmartRoutingStats, +} + +impl SmartRoutingProvider { + /// Create a new smart routing provider wrapping a primary and cheap provider. + pub fn new( + primary: Arc, + cheap: Arc, + config: SmartRoutingConfig, + ) -> Self { + Self { + primary, + cheap, + config, + 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. + 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(""); + + classify_message(last_user_msg, &self.config) + } + + /// 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)) + } +} + +/// Classify a message's complexity based on content patterns and length. +/// +/// Exposed as a free function for testability. +fn classify_message(msg: &str, config: &SmartRoutingConfig) -> TaskComplexity { + let trimmed = msg.trim(); + let len = trimmed.len(); + + // Empty or very short -> Simple + if len == 0 { + return TaskComplexity::Simple; + } + + // Check for code blocks (triple backticks) -> Complex + if trimmed.contains("```") { + return TaskComplexity::Complex; + } + + let lower = trimmed.to_lowercase(); + + // Complex keywords/patterns -> Complex regardless of length + const COMPLEX_KEYWORDS: &[&str] = &[ + "implement", + "refactor", + "analyze", + "debug", + "create a", + "build a", + "design", + "fix the", + "fix this", + "write a", + "write the", + "explain how", + "explain why", + "explain the", + "compare", + "optimize", + "review", + "rewrite", + "migrate", + "architect", + "integrate", + ]; + + if COMPLEX_KEYWORDS.iter().any(|k| lower.contains(k)) { + return TaskComplexity::Complex; + } + + // Long messages -> Complex + if len >= config.complex_min_chars { + return TaskComplexity::Complex; + } + + // Simple keywords/patterns for short messages + const SIMPLE_KEYWORDS: &[&str] = &[ + "list", + "show", + "what is", + "what's", + "status", + "help", + "yes", + "no", + "ok", + "thanks", + "thank you", + "hello", + "hi", + "hey", + "ping", + "version", + "how many", + "when", + "where is", + "who", + ]; + + if len <= config.simple_max_chars && SIMPLE_KEYWORDS.iter().any(|k| lower.contains(k)) { + return TaskComplexity::Simple; + } + + // Short confirmations / single words -> Simple + if len <= 10 { + return TaskComplexity::Simple; + } + + // Everything else -> Moderate + TaskComplexity::Moderate +} + +#[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() + } + + async fn complete(&self, request: CompletionRequest) -> Result { + self.stats.total_requests.fetch_add(1, Ordering::Relaxed); + + let complexity = self.classify(&request); + + match complexity { + TaskComplexity::Simple => { + tracing::debug!( + 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::debug!( + 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::debug!( + 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::debug!( + 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 { + self.stats.total_requests.fetch_add(1, Ordering::Relaxed); + self.stats.primary_requests.fetch_add(1, Ordering::Relaxed); + tracing::debug!( + 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, LlmError> { + self.primary.list_models().await + } + + async fn model_metadata(&self) -> Result { + self.primary.model_metadata().await + } + + 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() + } + + // -- Classification tests -- + + #[test] + fn classify_empty_message_as_simple() { + assert_eq!( + classify_message("", &default_config()), + TaskComplexity::Simple + ); + } + + #[test] + fn classify_greeting_as_simple() { + assert_eq!( + classify_message("hello", &default_config()), + TaskComplexity::Simple + ); + assert_eq!( + classify_message("hi there", &default_config()), + TaskComplexity::Simple + ); + } + + #[test] + fn classify_short_question_with_simple_keyword() { + assert_eq!( + classify_message("what is the status?", &default_config()), + TaskComplexity::Simple + ); + assert_eq!( + classify_message("show me the list", &default_config()), + TaskComplexity::Simple + ); + assert_eq!( + classify_message("help", &default_config()), + TaskComplexity::Simple + ); + } + + #[test] + fn classify_yes_no_as_simple() { + assert_eq!( + classify_message("yes", &default_config()), + TaskComplexity::Simple + ); + assert_eq!( + classify_message("no", &default_config()), + TaskComplexity::Simple + ); + assert_eq!( + classify_message("ok", &default_config()), + TaskComplexity::Simple + ); + } + + #[test] + fn classify_code_generation_as_complex() { + assert_eq!( + classify_message("implement a binary search function", &default_config()), + TaskComplexity::Complex + ); + assert_eq!( + classify_message("refactor the auth module", &default_config()), + TaskComplexity::Complex + ); + assert_eq!( + classify_message("debug this error", &default_config()), + TaskComplexity::Complex + ); + } + + #[test] + fn classify_code_blocks_as_complex() { + let msg = "What does this do?\n```rust\nfn main() {}\n```"; + assert_eq!( + classify_message(msg, &default_config()), + TaskComplexity::Complex + ); + } + + #[test] + fn classify_long_message_as_complex() { + let long_msg = "a ".repeat(600); // 1200 chars + assert_eq!( + classify_message(&long_msg, &default_config()), + TaskComplexity::Complex + ); + } + + #[test] + fn classify_medium_message_without_keywords_as_moderate() { + // > 10 chars, < 1000 chars, no simple or complex keywords + let msg = "Tell me about the weather patterns in the Pacific Ocean during summer months"; + assert_eq!( + classify_message(msg, &default_config()), + TaskComplexity::Moderate + ); + } + + #[test] + fn classify_very_short_unknown_as_simple() { + // <= 10 chars, no keywords + assert_eq!( + classify_message("foo", &default_config()), + TaskComplexity::Simple + ); + } + + // -- Uncertainty detection tests -- + + #[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, + }; + 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, + }; + 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, + }; + 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, + }; + assert!(!SmartRoutingProvider::response_is_uncertain(&response)); + } + + // -- 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()); + + let resp = router + .complete(make_request("implement a binary search")) + .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 + router.complete(make_request("hello")).await.unwrap(); + // Complex -> primary + router + .complete(make_request("implement a search")) + .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() { + // Cheap model returns an 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 moderate task (no simple/complex keywords, medium length) + let resp = router + .complete(make_request( + "Tell me about the weather patterns in the Pacific Ocean during summer months", + )) + .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( + "The Pacific Ocean weather patterns during summer are characterized by trade winds.", + ) + .with_model_name("cheap"), + ); + + let router = SmartRoutingProvider::new( + primary.clone(), + cheap.clone(), + SmartRoutingConfig { + cascade_enabled: true, + ..default_config() + }, + ); + + let resp = router + .complete(make_request( + "Tell me about the weather patterns in the Pacific Ocean during summer months", + )) + .await + .unwrap(); + + // Should NOT have escalated + assert!(resp.content.contains("Pacific Ocean")); + 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"); + } +} diff --git a/src/main.rs b/src/main.rs index c2ed03bf..b4a5831e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -24,12 +24,7 @@ use ironclaw::{ context::ContextManager, extensions::ExtensionManager, hooks::{HookRegistry, bootstrap_hooks}, - llm::{ - CachedProvider, CircuitBreakerConfig, CircuitBreakerProvider, CooldownConfig, - FailoverProvider, LlmProvider, ResponseCacheConfig, RetryConfig, RetryProvider, - SessionConfig, create_cheap_llm_provider, create_llm_provider, - create_llm_provider_with_config, create_session_manager, - }, + llm::{SessionConfig, build_provider_chain, create_session_manager}, orchestrator::{ ContainerJobConfig, ContainerJobManager, OrchestratorApi, TokenStore, api::OrchestratorState, @@ -626,103 +621,8 @@ async fn main() -> anyhow::Result<()> { None }; - // Initialize LLM provider (clone session so we can reuse it for embeddings) - let llm = create_llm_provider(&config.llm, session.clone())?; - tracing::info!("LLM provider initialized: {}", llm.model_name()); - - // Wrap each provider with RetryProvider for automatic retries on transient errors. - // RetryProvider sits inside FailoverProvider so each provider in the failover chain - // gets its own retry attempts before the failover moves to the next provider. - let retry_config = RetryConfig { - max_retries: config.llm.nearai.max_retries, - }; - let llm: Arc = if retry_config.max_retries > 0 { - tracing::info!( - max_retries = retry_config.max_retries, - "LLM retry wrapper enabled" - ); - Arc::new(RetryProvider::new(llm, retry_config.clone())) - } else { - llm - }; - - // Wrap in failover if a fallback model is configured - let llm: Arc = - if let Some(fallback_model) = config.llm.nearai.fallback_model.as_ref() { - if fallback_model == &config.llm.nearai.model { - tracing::warn!( - "fallback_model is the same as primary model, failover may not be effective" - ); - } - let mut fallback_config = config.llm.nearai.clone(); - fallback_config.model = fallback_model.clone(); - let fallback = create_llm_provider_with_config(&fallback_config, session.clone())?; - tracing::info!( - primary = %llm.model_name(), - fallback = %fallback.model_name(), - "LLM failover enabled" - ); - // Wrap fallback with retry too - let fallback: Arc = if retry_config.max_retries > 0 { - Arc::new(RetryProvider::new(fallback, retry_config.clone())) - } else { - fallback - }; - let cooldown_config = CooldownConfig { - cooldown_duration: std::time::Duration::from_secs( - config.llm.nearai.failover_cooldown_secs, - ), - failure_threshold: config.llm.nearai.failover_cooldown_threshold, - }; - Arc::new(FailoverProvider::with_cooldown( - vec![llm, fallback], - cooldown_config, - )?) - } else { - llm - }; - - // Wrap in circuit breaker if configured - let llm: Arc = - if let Some(threshold) = config.llm.nearai.circuit_breaker_threshold { - let cb_config = CircuitBreakerConfig { - failure_threshold: threshold, - recovery_timeout: std::time::Duration::from_secs( - config.llm.nearai.circuit_breaker_recovery_secs, - ), - ..CircuitBreakerConfig::default() - }; - tracing::info!( - threshold, - recovery_secs = config.llm.nearai.circuit_breaker_recovery_secs, - "LLM circuit breaker enabled" - ); - Arc::new(CircuitBreakerProvider::new(llm, cb_config)) - } else { - llm - }; - - // Wrap in response cache if configured - let llm: Arc = if config.llm.nearai.response_cache_enabled { - let rc_config = ResponseCacheConfig { - ttl: std::time::Duration::from_secs(config.llm.nearai.response_cache_ttl_secs), - max_entries: config.llm.nearai.response_cache_max_entries, - }; - tracing::info!( - ttl_secs = config.llm.nearai.response_cache_ttl_secs, - max_entries = config.llm.nearai.response_cache_max_entries, - "LLM response cache enabled" - ); - Arc::new(CachedProvider::new(llm, rc_config)) - } else { - llm - }; - - // Initialize cheap LLM provider for lightweight tasks (heartbeat, evaluation) - let cheap_llm = create_cheap_llm_provider(&config.llm, session.clone())?; - if let Some(ref cheap) = cheap_llm { - tracing::info!("Cheap LLM provider initialized: {}", cheap.model_name()); - } + // Build the full LLM provider chain (retry → smart routing → failover → circuit breaker → cache) + let (llm, cheap_llm) = build_provider_chain(&config.llm, session.clone())?; // Initialize safety layer let safety = Arc::new(SafetyLayer::new(&config.safety)); diff --git a/src/setup/wizard.rs b/src/setup/wizard.rs index f314348d..b7258407 100644 --- a/src/setup/wizard.rs +++ b/src/setup/wizard.rs @@ -1176,6 +1176,7 @@ impl SetupWizard { response_cache_max_entries: 1000, failover_cooldown_secs: 300, failover_cooldown_threshold: 3, + smart_routing_cascade: true, }, openai: None, anthropic: None,