Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions src/config/llm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,8 @@ impl LlmConfig {
provider: None,
bedrock: None,
request_timeout_secs: 120,
cheap_model: None,
smart_routing_cascade: false,
}
}

Expand Down Expand Up @@ -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()
Expand All @@ -184,6 +194,8 @@ impl LlmConfig {
provider,
bedrock,
request_timeout_secs,
cheap_model,
smart_routing_cascade,
})
}

Expand Down
24 changes: 24 additions & 0 deletions src/llm/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
149 changes: 121 additions & 28 deletions src/llm/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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));
}

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

let mut cheap_config = config.nearai.clone();
cheap_config.model = cheap_model.clone();
// 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
),
})?;

Ok(Some(Arc::new(NearAiChatProvider::new(
cheap_config,
session,
)?)))
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.
Expand Down Expand Up @@ -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 {
Expand All @@ -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()
},
))
Expand Down Expand Up @@ -560,6 +590,8 @@ mod tests {
provider: None,
bedrock: None,
request_timeout_secs: 120,
cheap_model: None,
smart_routing_cascade: true,
}
}

Expand All @@ -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());

Expand All @@ -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());
Expand All @@ -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);
}
}
2 changes: 2 additions & 0 deletions src/setup/wizard.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
}

Expand Down
Loading