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
10 changes: 9 additions & 1 deletion src/agent/agent_loop.rs
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,9 @@ enum AgenticLoopResult {
pub struct AgentDeps {
pub store: Option<Arc<dyn Database>>,
pub llm: Arc<dyn LlmProvider>,
/// Cheap/fast LLM for lightweight tasks (heartbeat, routing, evaluation).
/// Falls back to the main `llm` if None.
pub cheap_llm: Option<Arc<dyn LlmProvider>>,
pub safety: Arc<SafetyLayer>,
pub tools: Arc<ToolRegistry>,
pub workspace: Option<Arc<Workspace>>,
Expand Down Expand Up @@ -138,6 +141,11 @@ impl Agent {
&self.deps.llm
}

/// Get the cheap/fast LLM provider, falling back to the main one.
fn cheap_llm(&self) -> &Arc<dyn LlmProvider> {
self.deps.cheap_llm.as_ref().unwrap_or(&self.deps.llm)
}

fn safety(&self) -> &Arc<SafetyLayer> {
&self.deps.safety
}
Expand Down Expand Up @@ -301,7 +309,7 @@ impl Agent {
Some(spawn_heartbeat(
config,
workspace.clone(),
self.llm().clone(),
self.cheap_llm().clone(),
Some(notify_tx),
))
} else {
Expand Down
4 changes: 4 additions & 0 deletions src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -388,6 +388,9 @@ impl std::str::FromStr for NearAiApiMode {
pub struct NearAiConfig {
/// Model to use (e.g., "claude-3-5-sonnet-20241022", "gpt-4o")
pub model: String,
/// Cheap/fast model for lightweight tasks (heartbeat, routing, evaluation).
/// Falls back to the main model if not set.
pub cheap_model: Option<String>,
/// Base URL for the NEAR AI API (default: https://api.near.ai)
pub base_url: String,
/// Base URL for auth/refresh endpoints (default: https://private.near.ai)
Expand Down Expand Up @@ -454,6 +457,7 @@ impl LlmConfig {
"fireworks::accounts/fireworks/models/llama4-maverick-instruct-basic"
.to_string()
}),
cheap_model: optional_env("NEARAI_CHEAP_MODEL")?,
base_url: optional_env("NEARAI_BASE_URL")?
.unwrap_or_else(|| "https://cloud-api.near.ai".to_string()),
auth_base_url: optional_env("NEARAI_AUTH_URL")?
Expand Down
103 changes: 103 additions & 0 deletions src/llm/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -183,3 +183,106 @@ fn create_openai_compatible_provider(config: &LlmConfig) -> Result<Arc<dyn LlmPr
);
Ok(Arc::new(RigAdapter::new(model, &compat.model)))
}

/// 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 backends (Responses and ChatCompletions modes).
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 {
return Ok(None);
};

if config.backend != LlmBackend::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);
}

let mut cheap_config = config.nearai.clone();
cheap_config.model = cheap_model.clone();

tracing::info!("Cheap LLM provider: {}", cheap_model);

match cheap_config.api_mode {
NearAiApiMode::Responses => Ok(Some(Arc::new(NearAiProvider::new(cheap_config, session)))),
NearAiApiMode::ChatCompletions => {
Ok(Some(Arc::new(NearAiChatProvider::new(cheap_config)?)))
}
}
Comment on lines +213 to +218
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The logic for creating the provider based on api_mode is identical to create_llm_provider_with_config. You can refactor this to use the existing helper function to reduce duplication and ensure consistent provider initialization.

}

#[cfg(test)]
mod tests {
use super::*;
use crate::config::{LlmBackend, NearAiApiMode, NearAiConfig};
use std::path::PathBuf;

fn test_nearai_config() -> NearAiConfig {
NearAiConfig {
model: "test-model".to_string(),
cheap_model: None,
base_url: "https://api.near.ai".to_string(),
auth_base_url: "https://private.near.ai".to_string(),
session_path: PathBuf::from("/tmp/test-session.json"),
api_mode: NearAiApiMode::Responses,
api_key: None,
fallback_model: None,
max_retries: 3,
}
}

fn test_llm_config() -> LlmConfig {
LlmConfig {
backend: LlmBackend::NearAi,
nearai: test_nearai_config(),
openai: None,
anthropic: None,
ollama: None,
openai_compatible: None,
}
}

#[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);
assert!(result.is_ok());
assert!(result.unwrap().is_none());
}

#[test]
fn test_create_cheap_llm_provider_creates_provider_when_configured() {
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);

assert!(result.is_ok());
let provider = result.unwrap();
assert!(provider.is_some());
assert_eq!(provider.unwrap().model_name(), "cheap-test-model");
}

#[test]
fn test_create_cheap_llm_provider_ignored_for_non_nearai_backend() {
let mut config = test_llm_config();
config.backend = LlmBackend::OpenAi;
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);

assert!(result.is_ok());
assert!(result.unwrap().is_none());
}
Copy link

Copilot AI Feb 17, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider adding a test case for creating a cheap LLM provider with ChatCompletions API mode. The existing tests only cover the Responses mode, but the code supports both modes (lines 213-218). This would improve test coverage for this feature.

Suggested change
}
}
#[test]
fn test_create_cheap_llm_provider_with_chat_completions_mode() {
let mut config = test_llm_config();
config.nearai.cheap_model = Some("cheap-test-model".to_string());
config.nearai.api_mode = NearAiApiMode::ChatCompletions;
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(), "cheap-test-model");
}

Copilot uses AI. Check for mistakes.
}
30 changes: 26 additions & 4 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,8 @@ use ironclaw::{
context::ContextManager,
extensions::ExtensionManager,
llm::{
FailoverProvider, LlmProvider, SessionConfig, create_llm_provider,
create_llm_provider_with_config, create_session_manager,
FailoverProvider, LlmProvider, SessionConfig, create_cheap_llm_provider,
create_llm_provider, create_llm_provider_with_config, create_session_manager,
},
orchestrator::{
ContainerJobConfig, ContainerJobManager, OrchestratorApi, TokenStore,
Expand Down Expand Up @@ -307,8 +307,11 @@ async fn main() -> anyhow::Result<()> {
};
let session = create_session_manager(session_config).await;

// Ensure we're authenticated before proceeding (only needed for NEAR AI backend)
if config.llm.backend == ironclaw::config::LlmBackend::NearAi {
// Session-based auth is only needed for NEAR AI backend without an API key.
// ChatCompletions mode with an API key skips session auth entirely.
if config.llm.backend == ironclaw::config::LlmBackend::NearAi
&& config.llm.nearai.api_key.is_none()
{
session.ensure_authenticated().await?;
}
Comment on lines +310 to 316
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The condition for skipping session-based authentication should ideally check the api_mode rather than just the presence of an API key. While api_key.is_some() usually implies ChatCompletions mode, a user could explicitly configure Responses mode while having an API key set. In Responses mode, the NearAiProvider always requires a session token, so ensure_authenticated() must be called.


Expand Down Expand Up @@ -534,6 +537,12 @@ async fn main() -> anyhow::Result<()> {
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());
}
Comment on lines +540 to +544
Copy link

Copilot AI Feb 17, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider adding a warning when the cheap model is the same as the main model, similar to the fallback model check above (lines 522-526). This would help users avoid misconfiguration where they set NEARAI_CHEAP_MODEL to the same value as NEARAI_MODEL, defeating the purpose of having a cheaper model for lightweight tasks.

Copilot uses AI. Check for mistakes.

// Initialize safety layer
let safety = Arc::new(SafetyLayer::new(&config.safety));
tracing::info!("Safety layer initialized");
Expand Down Expand Up @@ -1185,6 +1194,7 @@ async fn main() -> anyhow::Result<()> {
let deps = AgentDeps {
store: db,
llm,
cheap_llm,
safety,
tools,
workspace,
Expand Down Expand Up @@ -1229,6 +1239,18 @@ fn check_onboard_needed() -> Option<&'static str> {
return Some("Database not configured");
}

// First run (onboarding never completed and no session).
// Reads NEARAI_API_KEY env var directly because this function runs
// before Config is loaded -- Config::from_env() may fail without a
// database URL, which is what triggers onboarding in the first place.
if std::env::var("NEARAI_API_KEY").is_err() {
let settings = ironclaw::settings::Settings::load();
let session_path = ironclaw::llm::session::default_session_path();
if !settings.onboard_completed && !session_path.exists() {
return Some("First run");
}
}

None
}

Expand Down
1 change: 1 addition & 0 deletions src/setup/wizard.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1014,6 +1014,7 @@ impl SetupWizard {
backend: crate::config::LlmBackend::NearAi,
nearai: crate::config::NearAiConfig {
model: "dummy".to_string(),
cheap_model: None,
base_url,
auth_base_url,
session_path: crate::llm::session::default_session_path(),
Expand Down
Loading