-
Notifications
You must be signed in to change notification settings - Fork 1.4k
feat: direct api key and cheap model #116
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
@@ -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)?))) | ||||||||||||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||
| #[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()); | ||||||||||||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||
| } | |
| } | |
| #[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"); | |
| } |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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, | ||
|
|
@@ -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
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The condition for skipping session-based authentication should ideally check the |
||
|
|
||
|
|
@@ -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
|
||
|
|
||
| // Initialize safety layer | ||
| let safety = Arc::new(SafetyLayer::new(&config.safety)); | ||
| tracing::info!("Safety layer initialized"); | ||
|
|
@@ -1185,6 +1194,7 @@ async fn main() -> anyhow::Result<()> { | |
| let deps = AgentDeps { | ||
| store: db, | ||
| llm, | ||
| cheap_llm, | ||
| safety, | ||
| tools, | ||
| workspace, | ||
|
|
@@ -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 | ||
| } | ||
|
|
||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The logic for creating the provider based on
api_modeis identical tocreate_llm_provider_with_config. You can refactor this to use the existing helper function to reduce duplication and ensure consistent provider initialization.