feat: add LLM_CHEAP_MODEL for generic smart routing across all backends (#1081)

* 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.

* refactor: address Gemini review — remove unnecessary async, extract cheap_model_name()

- 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

---------

Co-authored-by: SMKRV <[email protected]>

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,
         })
     }
 

src/llm/config.rs

@@ -129,6 +129,30 @@ 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<String>,
+    /// 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,
+}
+
+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.

src/llm/mod.rs

@@ -336,32 +336,61 @@ 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.
+/// 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 fn create_cheap_llm_provider(
     config: &LlmConfig,
     session: Arc<SessionManager>,
 ) -> Result<Option<Arc<dyn LlmProvider>>, LlmError> {
-    let Some(ref cheap_model) = config.nearai.cheap_model else {
+    let Some(cheap_model) = config.cheap_model_name() 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)
+}
+
+/// 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`
+fn create_cheap_provider_for_backend(
+    config: &LlmConfig,
+    session: Arc<SessionManager>,
+    cheap_model: &str,
+) -> Result<Option<Arc<dyn LlmProvider>>, 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 +438,15 @@ pub async fn build_provider_chain(
     };
 
     // 2. Smart routing (cheap/primary split)
-    let llm: Arc<dyn LlmProvider> = 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,
-        )?;
+    let llm: Arc<dyn LlmProvider> = 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!(
+                    "Failed to create cheap provider for model '{cheap_model}' on backend '{}'",
+                    config.backend
+                ),
+            })?;
         let cheap: Arc<dyn LlmProvider> = if retry_config.max_retries > 0 {
             Arc::new(RetryProvider::new(cheap, retry_config.clone()))
         } else {
@@ -431,7 +461,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()
             },
         ))
@@ -560,6 +590,8 @@ mod tests {
             provider: None,
             bedrock: None,
             request_timeout_secs: 120,
+            cheap_model: None,
+            smart_routing_cascade: true,
         }
     }
 
@@ -574,7 +606,7 @@ mod tests {
     }
 
     #[test]
-    fn test_create_cheap_llm_provider_creates_provider_when_configured() {
+    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());
 
@@ -588,7 +620,26 @@ mod tests {
     }
 
     #[test]
-    fn test_create_cheap_llm_provider_ignored_for_non_nearai_backend() {
+    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);
+
+        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"
+        );
+    }
+
+    #[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());
@@ -597,6 +648,48 @@ mod tests {
         let result = create_cheap_llm_provider(&config, session);
 
         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"
+        );
+    }
+
+    #[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);
+
+        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);
     }
 }

src/setup/wizard.rs

@@ -3429,6 +3429,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,
     }
 }