From 106a37ab3ce082a834d9fbfc46df0c7fe619b146 Mon Sep 17 00:00:00 2001 From: SMKRV Date: Fri, 13 Mar 2026 00:32:10 +0300 Subject: [PATCH 1/2] feat: add LLM_CHEAP_MODEL for generic smart routing across all backends Add generic cheap model support that works with any LLM backend, not just NearAI. New env vars: LLM_CHEAP_MODEL (cheap model for any backend) and SMART_ROUTING_CASCADE (top-level cascade flag). Resolution order: LLM_CHEAP_MODEL > NEARAI_CHEAP_MODEL (backward compat). Registry-based providers (OpenAI, Anthropic, Groq, etc.) clone their RegistryProviderConfig with the cheap model swapped in. Bedrock returns an explicit error (not yet supported). All error paths use ok_or_else with proper LlmError variants -- no unwrap/expect in production code. --- src/config/llm.rs | 12 ++++ src/llm/config.rs | 7 ++ src/llm/mod.rs | 161 ++++++++++++++++++++++++++++++++++---------- src/setup/wizard.rs | 2 + 4 files changed, 145 insertions(+), 37 deletions(-) diff --git a/src/config/llm.rs b/src/config/llm.rs index dd2c95632c2..58510473fa7 100644 --- a/src/config/llm.rs +++ b/src/config/llm.rs @@ -39,6 +39,8 @@ impl LlmConfig { provider: None, bedrock: None, request_timeout_secs: 120, + cheap_model: None, + smart_routing_cascade: false, } } @@ -169,6 +171,14 @@ impl LlmConfig { let request_timeout_secs = parse_optional_env("LLM_REQUEST_TIMEOUT_SECS", 120)?; + // Generic cheap model (works with any backend). + // Falls back to NearAI-specific cheap_model in provider chain logic. + let cheap_model = optional_env("LLM_CHEAP_MODEL")?; + + // Generic smart routing cascade flag. + // Defaults to true. Overrides NearAI-specific smart_routing_cascade. + let smart_routing_cascade = parse_optional_env("SMART_ROUTING_CASCADE", true)?; + Ok(Self { backend: if is_nearai { "nearai".to_string() @@ -184,6 +194,8 @@ impl LlmConfig { provider, bedrock, request_timeout_secs, + cheap_model, + smart_routing_cascade, }) } diff --git a/src/llm/config.rs b/src/llm/config.rs index 1902f128b32..b0016ff8426 100644 --- a/src/llm/config.rs +++ b/src/llm/config.rs @@ -129,6 +129,13 @@ pub struct LlmConfig { /// Default: 120. Increase for local LLMs (Ollama, vLLM, LM Studio) that /// need more time for prompt evaluation on consumer hardware. pub request_timeout_secs: u64, + /// Generic cheap/fast model for lightweight tasks (heartbeat, routing, evaluation). + /// Works with any backend. Set via `LLM_CHEAP_MODEL` env var. + /// When set, takes priority over the NearAI-specific `NEARAI_CHEAP_MODEL`. + pub cheap_model: Option, + /// Enable cascade mode for smart routing (retry with primary if cheap model + /// response seems uncertain). Default: true. Set via `SMART_ROUTING_CASCADE`. + pub smart_routing_cascade: bool, } /// NEAR AI configuration. diff --git a/src/llm/mod.rs b/src/llm/mod.rs index b49e4974a1f..96e7be446b3 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -336,32 +336,70 @@ fn create_ollama_from_registry( /// Create a cheap/fast LLM provider for lightweight tasks (heartbeat, routing, evaluation). /// -/// Uses `NEARAI_CHEAP_MODEL` if set, otherwise falls back to the main provider. -/// Currently only supports NEAR AI backend. -pub fn create_cheap_llm_provider( +/// Resolution order: +/// 1. `LLM_CHEAP_MODEL` (generic, works with any backend) +/// 2. `NEARAI_CHEAP_MODEL` (NearAI-only, backward compatibility) +/// +/// Returns `None` if no cheap model is configured. +pub async fn create_cheap_llm_provider( config: &LlmConfig, session: Arc, ) -> Result>, LlmError> { - let Some(ref cheap_model) = config.nearai.cheap_model else { + // Resolve cheap model: generic LLM_CHEAP_MODEL first, then NearAI-specific + let cheap_model = config.cheap_model.as_deref().or_else(|| { + if config.backend == "nearai" { + config.nearai.cheap_model.as_deref() + } else { + None + } + }); + + let Some(cheap_model) = cheap_model else { return Ok(None); }; - if config.backend != "nearai" { - tracing::warn!( - "NEARAI_CHEAP_MODEL is set but LLM_BACKEND is '{}', not nearai. \ - Cheap model setting will be ignored.", - config.backend - ); - return Ok(None); + create_cheap_provider_for_backend(config, session, cheap_model).await +} + +/// Create a cheap provider for a specific backend. +/// +/// Handles backend-specific provider construction: +/// - `nearai` — clones NearAiConfig, swaps model, uses `create_llm_provider_with_config` +/// - `bedrock` — returns error (smart routing not yet supported) +/// - All others — clones `RegistryProviderConfig`, swaps model, uses `create_registry_provider` +async fn create_cheap_provider_for_backend( + config: &LlmConfig, + session: Arc, + cheap_model: &str, +) -> Result>, LlmError> { + if config.backend == "nearai" { + let mut cheap_config = config.nearai.clone(); + cheap_config.model = cheap_model.to_string(); + let provider = + create_llm_provider_with_config(&cheap_config, session, config.request_timeout_secs)?; + return Ok(Some(provider)); } - let mut cheap_config = config.nearai.clone(); - cheap_config.model = cheap_model.clone(); + if config.backend == "bedrock" { + return Err(LlmError::RequestFailed { + provider: "bedrock".to_string(), + reason: "Smart routing with cheap model is not supported for Bedrock yet".to_string(), + }); + } - Ok(Some(Arc::new(NearAiChatProvider::new( - cheap_config, - session, - )?))) + // Registry-based provider: clone config and swap model + let reg_config = config.provider.as_ref().ok_or_else(|| LlmError::RequestFailed { + provider: config.backend.clone(), + reason: format!( + "Cannot create cheap provider for backend '{}': no registry provider config available", + config.backend + ), + })?; + + let mut cheap_reg_config = reg_config.clone(); + cheap_reg_config.model = cheap_model.to_string(); + let provider = create_registry_provider(&cheap_reg_config)?; + Ok(Some(provider)) } /// Build the full LLM provider chain with all configured wrappers. @@ -409,14 +447,24 @@ pub async fn build_provider_chain( }; // 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(), - config.request_timeout_secs, - )?; + // Resolve cheap model: generic LLM_CHEAP_MODEL first, then NearAI-specific + let effective_cheap_model = config.cheap_model.as_deref().or_else(|| { + if config.backend == "nearai" { + config.nearai.cheap_model.as_deref() + } else { + None + } + }); + let llm: Arc = if let Some(cheap_model) = effective_cheap_model { + let cheap = create_cheap_provider_for_backend(config, session.clone(), cheap_model) + .await? + .ok_or_else(|| LlmError::RequestFailed { + provider: config.backend.clone(), + reason: format!( + "Failed to create cheap provider for model '{cheap_model}' on backend '{}'", + config.backend + ), + })?; let cheap: Arc = if retry_config.max_retries > 0 { Arc::new(RetryProvider::new(cheap, retry_config.clone())) } else { @@ -431,7 +479,7 @@ pub async fn build_provider_chain( llm, cheap, SmartRoutingConfig { - cascade_enabled: config.nearai.smart_routing_cascade, + cascade_enabled: config.smart_routing_cascade, ..SmartRoutingConfig::default() }, )) @@ -520,7 +568,7 @@ pub async fn build_provider_chain( }; // Standalone cheap LLM for heartbeat/evaluation (not part of the chain) - let cheap_llm = create_cheap_llm_provider(config, session)?; + let cheap_llm = create_cheap_llm_provider(config, session).await?; if let Some(ref cheap) = cheap_llm { tracing::debug!("Cheap LLM provider initialized: {}", cheap.model_name()); } @@ -560,26 +608,28 @@ mod tests { provider: None, bedrock: None, request_timeout_secs: 120, + cheap_model: None, + smart_routing_cascade: true, } } - #[test] - fn test_create_cheap_llm_provider_returns_none_when_not_configured() { + #[tokio::test] + async fn test_create_cheap_llm_provider_returns_none_when_not_configured() { let config = test_llm_config(); let session = Arc::new(SessionManager::new(SessionConfig::default())); - let result = create_cheap_llm_provider(&config, session); + let result = create_cheap_llm_provider(&config, session).await; assert!(result.is_ok()); assert!(result.unwrap().is_none()); } - #[test] - fn test_create_cheap_llm_provider_creates_provider_when_configured() { + #[tokio::test] + async fn test_create_cheap_llm_provider_creates_provider_with_nearai_cheap_model() { let mut config = test_llm_config(); config.nearai.cheap_model = Some("cheap-test-model".to_string()); let session = Arc::new(SessionManager::new(SessionConfig::default())); - let result = create_cheap_llm_provider(&config, session); + let result = create_cheap_llm_provider(&config, session).await; assert!(result.is_ok()); let provider = result.unwrap(); @@ -587,16 +637,53 @@ mod tests { assert_eq!(provider.unwrap().model_name(), "cheap-test-model"); } - #[test] - fn test_create_cheap_llm_provider_ignored_for_non_nearai_backend() { + #[tokio::test] + async fn test_create_cheap_llm_provider_generic_overrides_nearai() { + let mut config = test_llm_config(); + config.nearai.cheap_model = Some("nearai-cheap".to_string()); + config.cheap_model = Some("generic-cheap".to_string()); + + let session = Arc::new(SessionManager::new(SessionConfig::default())); + let result = create_cheap_llm_provider(&config, session).await; + + assert!(result.is_ok()); + let provider = result.unwrap(); + assert!(provider.is_some()); + assert_eq!( + provider.unwrap().model_name(), + "generic-cheap", + "LLM_CHEAP_MODEL should take priority over NEARAI_CHEAP_MODEL" + ); + } + + #[tokio::test] + async fn test_create_cheap_llm_provider_nearai_cheap_ignored_for_non_nearai_backend() { let mut config = test_llm_config(); config.backend = "openai".to_string(); config.nearai.cheap_model = Some("cheap-test-model".to_string()); let session = Arc::new(SessionManager::new(SessionConfig::default())); - let result = create_cheap_llm_provider(&config, session); + let result = create_cheap_llm_provider(&config, session).await; assert!(result.is_ok()); - assert!(result.unwrap().is_none()); + assert!( + result.unwrap().is_none(), + "NEARAI_CHEAP_MODEL should be ignored when backend is not nearai" + ); + } + + #[tokio::test] + async fn test_create_cheap_llm_provider_bedrock_returns_error() { + let mut config = test_llm_config(); + config.backend = "bedrock".to_string(); + config.cheap_model = Some("cheap-model".to_string()); + + let session = Arc::new(SessionManager::new(SessionConfig::default())); + let result = create_cheap_llm_provider(&config, session).await; + + assert!( + result.is_err(), + "Bedrock should return an error for cheap model" + ); } } diff --git a/src/setup/wizard.rs b/src/setup/wizard.rs index 6c7d03cbadc..516151154f3 100644 --- a/src/setup/wizard.rs +++ b/src/setup/wizard.rs @@ -3436,6 +3436,8 @@ fn build_nearai_model_fetch_config() -> crate::config::LlmConfig { provider: None, bedrock: None, request_timeout_secs: 120, + cheap_model: None, + smart_routing_cascade: true, } } From c70b49f654d99750afaf22eb1d77394f3c1e1b9d Mon Sep 17 00:00:00 2001 From: SMKRV Date: Fri, 13 Mar 2026 01:04:12 +0300 Subject: [PATCH 2/2] =?UTF-8?q?refactor:=20address=20Gemini=20review=20?= =?UTF-8?q?=E2=80=94=20remove=20unnecessary=20async,=20extract=20cheap=5Fm?= =?UTF-8?q?odel=5Fname()?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Remove async from create_cheap_provider_for_backend() and create_cheap_llm_provider() — neither contains .await calls - Extract duplicated cheap model resolution logic into LlmConfig::cheap_model_name() helper method (DRY) - Revert tests from tokio::test async back to sync #[test] - Add test_cheap_model_name_resolution() unit test for the helper --- src/llm/config.rs | 17 ++++++++++ src/llm/mod.rs | 86 +++++++++++++++++++++++++---------------------- 2 files changed, 63 insertions(+), 40 deletions(-) diff --git a/src/llm/config.rs b/src/llm/config.rs index b0016ff8426..9bf1b79bc96 100644 --- a/src/llm/config.rs +++ b/src/llm/config.rs @@ -138,6 +138,23 @@ pub struct LlmConfig { pub smart_routing_cascade: bool, } +impl LlmConfig { + /// Resolve the effective cheap model name. + /// + /// Resolution order: + /// 1. `LLM_CHEAP_MODEL` (generic, works with any backend) + /// 2. `NEARAI_CHEAP_MODEL` (NearAI-only, backward compatibility) + pub fn cheap_model_name(&self) -> Option<&str> { + self.cheap_model.as_deref().or_else(|| { + if self.backend == "nearai" { + self.nearai.cheap_model.as_deref() + } else { + None + } + }) + } +} + /// NEAR AI configuration. #[derive(Debug, Clone)] pub struct NearAiConfig { diff --git a/src/llm/mod.rs b/src/llm/mod.rs index 96e7be446b3..11e1ad71c77 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -341,24 +341,15 @@ fn create_ollama_from_registry( /// 2. `NEARAI_CHEAP_MODEL` (NearAI-only, backward compatibility) /// /// Returns `None` if no cheap model is configured. -pub async fn create_cheap_llm_provider( +pub fn create_cheap_llm_provider( config: &LlmConfig, session: Arc, ) -> Result>, LlmError> { - // Resolve cheap model: generic LLM_CHEAP_MODEL first, then NearAI-specific - let cheap_model = config.cheap_model.as_deref().or_else(|| { - if config.backend == "nearai" { - config.nearai.cheap_model.as_deref() - } else { - None - } - }); - - let Some(cheap_model) = cheap_model else { + let Some(cheap_model) = config.cheap_model_name() else { return Ok(None); }; - create_cheap_provider_for_backend(config, session, cheap_model).await + create_cheap_provider_for_backend(config, session, cheap_model) } /// Create a cheap provider for a specific backend. @@ -367,7 +358,7 @@ pub async fn create_cheap_llm_provider( /// - `nearai` — clones NearAiConfig, swaps model, uses `create_llm_provider_with_config` /// - `bedrock` — returns error (smart routing not yet supported) /// - All others — clones `RegistryProviderConfig`, swaps model, uses `create_registry_provider` -async fn create_cheap_provider_for_backend( +fn create_cheap_provider_for_backend( config: &LlmConfig, session: Arc, cheap_model: &str, @@ -447,17 +438,8 @@ pub async fn build_provider_chain( }; // 2. Smart routing (cheap/primary split) - // Resolve cheap model: generic LLM_CHEAP_MODEL first, then NearAI-specific - let effective_cheap_model = config.cheap_model.as_deref().or_else(|| { - if config.backend == "nearai" { - config.nearai.cheap_model.as_deref() - } else { - None - } - }); - let llm: Arc = if let Some(cheap_model) = effective_cheap_model { - let cheap = create_cheap_provider_for_backend(config, session.clone(), cheap_model) - .await? + let llm: Arc = if let Some(cheap_model) = config.cheap_model_name() { + let cheap = create_cheap_provider_for_backend(config, session.clone(), cheap_model)? .ok_or_else(|| LlmError::RequestFailed { provider: config.backend.clone(), reason: format!( @@ -568,7 +550,7 @@ pub async fn build_provider_chain( }; // Standalone cheap LLM for heartbeat/evaluation (not part of the chain) - let cheap_llm = create_cheap_llm_provider(config, session).await?; + let cheap_llm = create_cheap_llm_provider(config, session)?; if let Some(ref cheap) = cheap_llm { tracing::debug!("Cheap LLM provider initialized: {}", cheap.model_name()); } @@ -613,23 +595,23 @@ mod tests { } } - #[tokio::test] - async fn test_create_cheap_llm_provider_returns_none_when_not_configured() { + #[test] + fn test_create_cheap_llm_provider_returns_none_when_not_configured() { let config = test_llm_config(); let session = Arc::new(SessionManager::new(SessionConfig::default())); - let result = create_cheap_llm_provider(&config, session).await; + let result = create_cheap_llm_provider(&config, session); assert!(result.is_ok()); assert!(result.unwrap().is_none()); } - #[tokio::test] - async fn test_create_cheap_llm_provider_creates_provider_with_nearai_cheap_model() { + #[test] + fn test_create_cheap_llm_provider_creates_provider_with_nearai_cheap_model() { let mut config = test_llm_config(); config.nearai.cheap_model = Some("cheap-test-model".to_string()); let session = Arc::new(SessionManager::new(SessionConfig::default())); - let result = create_cheap_llm_provider(&config, session).await; + let result = create_cheap_llm_provider(&config, session); assert!(result.is_ok()); let provider = result.unwrap(); @@ -637,14 +619,14 @@ mod tests { assert_eq!(provider.unwrap().model_name(), "cheap-test-model"); } - #[tokio::test] - async fn test_create_cheap_llm_provider_generic_overrides_nearai() { + #[test] + fn test_create_cheap_llm_provider_generic_overrides_nearai() { let mut config = test_llm_config(); config.nearai.cheap_model = Some("nearai-cheap".to_string()); config.cheap_model = Some("generic-cheap".to_string()); let session = Arc::new(SessionManager::new(SessionConfig::default())); - let result = create_cheap_llm_provider(&config, session).await; + let result = create_cheap_llm_provider(&config, session); assert!(result.is_ok()); let provider = result.unwrap(); @@ -656,14 +638,14 @@ mod tests { ); } - #[tokio::test] - async fn test_create_cheap_llm_provider_nearai_cheap_ignored_for_non_nearai_backend() { + #[test] + fn test_create_cheap_llm_provider_nearai_cheap_ignored_for_non_nearai_backend() { let mut config = test_llm_config(); config.backend = "openai".to_string(); config.nearai.cheap_model = Some("cheap-test-model".to_string()); let session = Arc::new(SessionManager::new(SessionConfig::default())); - let result = create_cheap_llm_provider(&config, session).await; + let result = create_cheap_llm_provider(&config, session); assert!(result.is_ok()); assert!( @@ -672,18 +654,42 @@ mod tests { ); } - #[tokio::test] - async fn test_create_cheap_llm_provider_bedrock_returns_error() { + #[test] + fn test_create_cheap_llm_provider_bedrock_returns_error() { let mut config = test_llm_config(); config.backend = "bedrock".to_string(); config.cheap_model = Some("cheap-model".to_string()); let session = Arc::new(SessionManager::new(SessionConfig::default())); - let result = create_cheap_llm_provider(&config, session).await; + let result = create_cheap_llm_provider(&config, session); assert!( result.is_err(), "Bedrock should return an error for cheap model" ); } + + #[test] + fn test_cheap_model_name_resolution() { + // Generic takes priority + let mut config = test_llm_config(); + config.cheap_model = Some("generic".to_string()); + config.nearai.cheap_model = Some("nearai".to_string()); + assert_eq!(config.cheap_model_name(), Some("generic")); + + // NearAI fallback when backend is nearai + let mut config = test_llm_config(); + config.nearai.cheap_model = Some("nearai".to_string()); + assert_eq!(config.cheap_model_name(), Some("nearai")); + + // NearAI ignored for non-nearai backend + let mut config = test_llm_config(); + config.backend = "openai".to_string(); + config.nearai.cheap_model = Some("nearai".to_string()); + assert_eq!(config.cheap_model_name(), None); + + // None when nothing configured + let config = test_llm_config(); + assert_eq!(config.cheap_model_name(), None); + } }