From d230c6eb0fba2b22d7dd045784854ea4e89829f1 Mon Sep 17 00:00:00 2001 From: Bradley Axen Date: Thu, 23 Jan 2025 10:49:18 -0800 Subject: [PATCH] feat: general purpose config (#701) Co-authored-by: Salman Mohammed --- crates/goose-cli/src/commands/configure.rs | 335 ++++++------- crates/goose-cli/src/commands/session.rs | 217 +++----- crates/goose-cli/src/config.rs | 203 -------- crates/goose-cli/src/main.rs | 129 +---- crates/goose-server/src/routes/agent.rs | 13 +- crates/goose-server/src/routes/extension.rs | 9 +- crates/goose-server/src/routes/reply.rs | 14 +- crates/goose-server/src/routes/secrets.rs | 11 +- crates/goose/Cargo.toml | 13 +- crates/goose/examples/agent.rs | 2 +- crates/goose/examples/databricks_oauth.rs | 2 +- crates/goose/examples/image_tool.rs | 4 +- crates/goose/src/agents/capabilities.rs | 16 +- crates/goose/src/agents/truncate.rs | 32 +- crates/goose/src/config/base.rs | 462 ++++++++++++++++++ crates/goose/src/config/extensions.rs | 106 ++++ crates/goose/src/config/mod.rs | 6 + crates/goose/src/key_manager.rs | 192 -------- crates/goose/src/lib.rs | 3 +- .../src/{providers/configs.rs => model.rs} | 0 crates/goose/src/providers/anthropic.rs | 52 +- crates/goose/src/providers/base.rs | 74 ++- crates/goose/src/providers/databricks.rs | 58 ++- crates/goose/src/providers/factory.rs | 38 +- .../goose/src/providers/formats/anthropic.rs | 2 +- crates/goose/src/providers/formats/google.rs | 2 +- crates/goose/src/providers/formats/openai.rs | 2 +- crates/goose/src/providers/google.rs | 47 +- crates/goose/src/providers/groq.rs | 47 +- crates/goose/src/providers/mod.rs | 5 +- crates/goose/src/providers/ollama.rs | 48 +- crates/goose/src/providers/openai.rs | 67 ++- crates/goose/src/providers/openrouter.rs | 52 +- crates/goose/src/token_counter.rs | 2 +- crates/goose/tests/providers.rs | 44 +- download_cli.sh | 24 +- .../settings/providers/ProviderGrid.tsx | 3 +- 37 files changed, 1231 insertions(+), 1105 deletions(-) delete mode 100644 crates/goose-cli/src/config.rs create mode 100644 crates/goose/src/config/base.rs create mode 100644 crates/goose/src/config/extensions.rs create mode 100644 crates/goose/src/config/mod.rs delete mode 100644 crates/goose/src/key_manager.rs rename crates/goose/src/{providers/configs.rs => model.rs} (100%) diff --git a/crates/goose-cli/src/commands/configure.rs b/crates/goose-cli/src/commands/configure.rs index 3bb0c2c6dc..afc8ced40a 100644 --- a/crates/goose-cli/src/commands/configure.rs +++ b/crates/goose-cli/src/commands/configure.rs @@ -1,29 +1,17 @@ use cliclack::spinner; use console::style; use goose::agents::{extension::Envs, ExtensionConfig}; -use goose::key_manager::{get_keyring_secret, save_to_keyring, KeyRetrievalStrategy}; +use goose::config::{Config, ExtensionEntry, ExtensionManager}; use goose::message::Message; -use goose::providers::anthropic::ANTHROPIC_DEFAULT_MODEL; -use goose::providers::databricks::DATABRICKS_DEFAULT_MODEL; -use goose::providers::factory; -use goose::providers::google::GOOGLE_DEFAULT_MODEL; -use goose::providers::groq::GROQ_DEFAULT_MODEL; -use goose::providers::ollama::OLLAMA_MODEL; -use goose::providers::openai::OPEN_AI_DEFAULT_MODEL; -use goose::providers::openrouter::OPENROUTER_DEFAULT_MODEL; +use goose::providers::{create, providers}; +use serde_json::Value; use std::collections::HashMap; use std::error::Error; -use crate::config::{Config, ExtensionEntry}; +pub async fn handle_configure() -> Result<(), Box> { + let config = Config::global(); -pub async fn handle_configure( - provided_provider: Option, - provided_model: Option, -) -> Result<(), Box> { - // Load existing config or create new one - let config_exists = Config::config_path()?.exists(); - - if !config_exists { + if !config.exists() { // First time setup flow println!(); println!( @@ -36,12 +24,24 @@ pub async fn handle_configure( ); println!(); cliclack::intro(style(" goose-configure ").on_cyan().black())?; - configure_provider_dialog(provided_provider, provided_model).await?; + configure_provider_dialog().await?; println!( "\n {}: Run '{}' again to adjust your config or add extensions", style("Tip").green().italic(), style("goose configure").cyan() ); + + // Since we are setting up for the first time, we'll also enable the developer system + ExtensionManager::set( + "developer", + ExtensionEntry { + enabled: true, + config: ExtensionConfig::Builtin { + name: "developer".to_string(), + }, + }, + )?; + Ok(()) } else { println!(); @@ -52,7 +52,7 @@ pub async fn handle_configure( println!( "{} {}", style(" if you prefer, you can edit it directly at").dim(), - Config::config_path()?.display() + config.path() ); println!(); @@ -74,100 +74,126 @@ pub async fn handle_configure( match action { "toggle" => toggle_extensions_dialog(), "add" => configure_extensions_dialog(), - "providers" => configure_provider_dialog(provided_provider, provided_model).await, + "providers" => configure_provider_dialog().await, _ => unreachable!(), } } } /// Dialog for configuring the AI provider and model -pub async fn configure_provider_dialog( - provided_provider: Option, - provided_model: Option, -) -> Result<(), Box> { - // Load existing config or create new one - let mut config = Config::load().unwrap_or_default(); +pub async fn configure_provider_dialog() -> Result<(), Box> { + // Get global config instance + let config = Config::global(); - // TODO offload to separate provider - let provider_name = if let Some(provider) = provided_provider { - provider - } else { - let providers = [ - "openai".to_string(), - "databricks".to_string(), - "ollama".to_string(), - "anthropic".to_string(), - "google".to_string(), - "groq".to_string(), - "openrouter".to_string(), - ]; - let provider = cliclack::select("Which model provider should we use?") - .initial_value(&config.default_provider) - .items(&[ - (&providers[0], "OpenAI", "GPT-4 etc"), - (&providers[1], "Databricks", "Models on AI Gateway"), - (&providers[2], "Ollama", "Local open source models"), - (&providers[3], "Anthropic", "Claude models"), - (&providers[4], "Google Gemini", "Gemini models"), - (&providers[5], "Groq", "Fast inference"), - (&providers[6], "OpenRouter", "Router for many models"), - ]) - .interact()?; - provider.to_string() - }; + // Get all available providers and their metadata + let available_providers = providers(); - // Configure provider keys - for key in get_required_keys(&provider_name).iter() { - // If the key is in the keyring, ask if we want to overwrite - if get_keyring_secret(key, KeyRetrievalStrategy::KeyringOnly).is_ok() { - let _ = cliclack::log::info(format!("{} is already available in the keyring", key)); - if cliclack::confirm("Would you like to overwrite this value?").interact()? { - let value = cliclack::password(format!("Enter the value for {}", key)) + // Create selection items from provider metadata + let provider_items: Vec<(&String, &str, &str)> = available_providers + .iter() + .map(|p| (&p.name, p.display_name.as_str(), p.description.as_str())) + .collect(); + + // Get current default provider if it exists + let current_provider: Option = config.get("GOOSE_PROVIDER").ok(); + let default_provider = current_provider.unwrap_or_default(); + + // Select provider + let provider_name = cliclack::select("Which model provider should we use?") + .initial_value(&default_provider) + .items(&provider_items) + .interact()?; + + // Get the selected provider's metadata + let provider_meta = available_providers + .iter() + .find(|p| &p.name == provider_name) + .expect("Selected provider must exist in metadata"); + + // Configure required provider keys + for key in &provider_meta.config_keys { + if !key.required { + continue; + } + + // Try to get existing value from environment or config + let existing: Result = if key.secret { + config.get_secret(&key.name) + } else { + config.get(&key.name) + }; + + match existing { + Ok(_) => { + let _ = cliclack::log::info(format!("{} is already configured", key.name)); + if cliclack::confirm("Would you like to update this value?").interact()? { + let new_value: String = if key.secret { + cliclack::password(format!("Enter new value for {}", key.name)) + .mask('▪') + .interact()? + } else { + cliclack::input(format!("Enter new value for {}", key.name)).interact()? + }; + + if key.secret { + config.set_secret(&key.name, Value::String(new_value))?; + } else { + config.set(&key.name, Value::String(new_value))?; + } + } + } + Err(_) => { + let value: String = if key.secret { + cliclack::password(format!( + "Provider {} requires {}, please enter a value", + provider_meta.display_name, key.name + )) .mask('▪') - .interact()?; + .interact()? + } else { + cliclack::input(format!( + "Provider {} requires {}, please enter a value", + provider_meta.display_name, key.name + )) + .interact()? + }; - save_to_keyring(key, &value)?; + if key.secret { + config.set_secret(&key.name, Value::String(value))?; + } else { + config.set(&key.name, Value::String(value))?; + } } } - // If the key is in the env, ask if we want to save to keyring - else if let Ok(value) = get_keyring_secret(key, KeyRetrievalStrategy::EnvironmentOnly) { - let _ = cliclack::log::info(format!("Detected {} in env, we can use this from your environment.\nIt will need to continue to be set in future goose usage.", key)); - if cliclack::confirm("Would you like to save it to your keyring?").interact()? { - save_to_keyring(key, &value)?; - } - } - // We don't have a value, so we prompt for one - else { - let value = cliclack::password(format!( - "Provider {} requires {}, please enter a value. (Will be saved to your keyring)", - provider_name, key - )) - .mask('▪') - .interact()?; - - save_to_keyring(key, &value)?; - } } - let model = if let Some(model) = provided_model { - model - } else { - let recommended_model = get_recommended_model(&provider_name); - cliclack::input("Enter a model from that provider:") - .default_input(recommended_model) - .interact()? - }; + // Select model, defaulting to the provider's recommended model + let model: String = cliclack::input("Enter a model from that provider:") + .default_input(&provider_meta.default_model) + .interact()?; // Update config with new values - config.default_provider = provider_name.clone(); - config.default_model = model.clone(); + config.set("GOOSE_PROVIDER", Value::String(provider_name.to_string()))?; + config.set("GOOSE_MODEL", Value::String(model.clone()))?; // Test the configuration let spin = spinner(); spin.start("Checking your configuration..."); - let provider = factory::get_provider(&provider_name).unwrap(); - let message = Message::user().with_text("Please give a nice welcome messsage (one sentence) and let them know they are all set to use this agent"); - let result = provider.complete("You are an AI agent called Goose. You use tools of connected extensions to solve problems.", &[message], &[]).await; + + let model_config = goose::model::ModelConfig::new(model.clone()); + let provider = create(provider_name, model_config)?; + + let message = Message::user().with_text( + "Please give a nice welcome message (one sentence) and let them know they are all set to use this agent" + ); + + let result = provider + .complete( + "You are an AI agent called Goose. You use tools of connected extensions to solve problems.", + &[message], + &[] + ) + .await; match result { Ok((message, _usage)) => { @@ -181,13 +207,7 @@ pub async fn configure_provider_dialog( spin.stop("No response content available"); } - let _ = match config.save() { - Ok(()) => { - let msg = format!("Configuration saved to: {:?}", Config::config_path()?); - cliclack::outro(msg) - } - Err(e) => cliclack::outro(format!("Failed to save configuration: {}", e)), - }; + cliclack::outro("Configuration saved successfully")?; } Err(e) => { println!("{:?}", e); @@ -199,39 +219,12 @@ pub async fn configure_provider_dialog( Ok(()) } -pub fn get_recommended_model(provider_name: &str) -> &str { - match provider_name { - "openai" => OPEN_AI_DEFAULT_MODEL, - "databricks" => DATABRICKS_DEFAULT_MODEL, - "ollama" => OLLAMA_MODEL, - "anthropic" => ANTHROPIC_DEFAULT_MODEL, - "google" => GOOGLE_DEFAULT_MODEL, - "groq" => GROQ_DEFAULT_MODEL, - "openrouter" => OPENROUTER_DEFAULT_MODEL, - _ => panic!("Invalid provider name"), - } -} - -pub fn get_required_keys(provider_name: &str) -> Vec<&'static str> { - match provider_name { - "openai" => vec!["OPENAI_API_KEY"], - "databricks" => vec!["DATABRICKS_HOST"], - "ollama" => vec!["OLLAMA_HOST"], - "anthropic" => vec!["ANTHROPIC_API_KEY"], - "google" => vec!["GOOGLE_API_KEY"], - "groq" => vec!["GROQ_API_KEY"], - "openrouter" => vec!["OPENROUTER_API_KEY"], - _ => panic!("Invalid provider name"), - } -} - /// Configure extensions that can be used with goose /// Dialog for toggling which extensions are enabled/disabled pub fn toggle_extensions_dialog() -> Result<(), Box> { - // Load existing config - let mut config = Config::load().unwrap_or_default(); + let extensions = ExtensionManager::get_all()?; - if config.extensions.is_empty() { + if extensions.is_empty() { cliclack::outro( "No extensions configured yet. Run configure and add some extensions first.", )?; @@ -239,10 +232,10 @@ pub fn toggle_extensions_dialog() -> Result<(), Box> { } // Create a list of extension names and their enabled status - let mut extension_status: Vec<(String, bool)> = Vec::new(); - for (name, entry) in config.extensions.iter() { - extension_status.push((name.clone(), entry.enabled)); - } + let extension_status: Vec<(String, bool)> = extensions + .iter() + .map(|(name, entry)| (name.clone(), entry.enabled)) + .collect(); // Get currently enabled extensions for the selection let enabled_extensions: Vec<&String> = extension_status @@ -265,35 +258,16 @@ pub fn toggle_extensions_dialog() -> Result<(), Box> { .initial_values(enabled_extensions) .interact()?; - // Update the config with new enabled/disabled status - for (name, _) in extension_status.iter() { - if let Some(entry) = config.extensions.get_mut(name) { - entry.enabled = selected.contains(&name); - } + // Update enabled status for each extension + for name in extension_status.iter().map(|(name, _)| name) { + ExtensionManager::set_enabled(name, selected.iter().any(|s| s.as_str() == name))?; } - config.save()?; cliclack::outro("Extension settings updated successfully")?; Ok(()) } pub fn configure_extensions_dialog() -> Result<(), Box> { - println!(); - println!( - "{}", - style("Configure will help you add extensions that goose can use").dim() - ); - println!( - "{}", - style(" extensions provide tools and capabilities to the AI agent").dim() - ); - println!(); - - cliclack::intro(style(" goose-configure-extensions ").on_cyan().black())?; - - // Load existing config or create new one - let mut config = Config::load().unwrap_or_default(); - let extension_type = cliclack::select("What type of extension would you like to add?") .item( "built-in", @@ -313,6 +287,7 @@ pub fn configure_extensions_dialog() -> Result<(), Box> { .interact()?; match extension_type { + // TODO we'll want a place to collect all these options, maybe just an enum in goose-mcp "built-in" => { let extension = cliclack::select("Which built-in extension would you like to enable?") .item( @@ -325,23 +300,34 @@ pub fn configure_extensions_dialog() -> Result<(), Box> { "Non Developer", "AI driven scripting for non developers", ) + .item( + "google_drive", + "Google Drive", + "Search and read content from google drive - additional config required", + ) + .item( + "memory", + "Memory", + "Tools to save and retrieve durable memories", + ) .item("jetbrains", "JetBrains", "Connect to jetbrains IDEs") - .interact()?; + .interact()? + .to_string(); - config.extensions.insert( - extension.to_string(), + ExtensionManager::set( + &extension, ExtensionEntry { enabled: true, config: ExtensionConfig::Builtin { - name: extension.to_string(), + name: extension.clone(), }, }, - ); + )?; cliclack::outro(format!("Enabled {} extension", style(extension).green()))?; } "stdio" => { - let extensions = config.extensions.clone(); + let extensions = ExtensionManager::get_all()?; let name: String = cliclack::input("What would you like to call this extension?") .placeholder("my-extension") .validate(move |input: &String| { @@ -377,11 +363,11 @@ pub fn configure_extensions_dialog() -> Result<(), Box> { let mut envs = HashMap::new(); if add_env { loop { - let key = cliclack::input("Environment variable name:") + let key: String = cliclack::input("Environment variable name:") .placeholder("API_KEY") .interact()?; - let value = cliclack::password("Environment variable value:") + let value: String = cliclack::password("Environment variable value:") .mask('▪') .interact()?; @@ -393,8 +379,8 @@ pub fn configure_extensions_dialog() -> Result<(), Box> { } } - config.extensions.insert( - name.clone(), + ExtensionManager::set( + &name, ExtensionEntry { enabled: true, config: ExtensionConfig::Stdio { @@ -403,12 +389,12 @@ pub fn configure_extensions_dialog() -> Result<(), Box> { envs: Envs::new(envs), }, }, - ); + )?; cliclack::outro(format!("Added {} extension", style(name).green()))?; } "sse" => { - let extensions = config.extensions.clone(); + let extensions = ExtensionManager::get_all()?; let name: String = cliclack::input("What would you like to call this extension?") .placeholder("my-remote-extension") .validate(move |input: &String| { @@ -422,7 +408,7 @@ pub fn configure_extensions_dialog() -> Result<(), Box> { }) .interact()?; - let uri = cliclack::input("What is the SSE endpoint URI?") + let uri: String = cliclack::input("What is the SSE endpoint URI?") .placeholder("http://localhost:8000/events") .validate(|input: &String| { if input.is_empty() { @@ -441,11 +427,11 @@ pub fn configure_extensions_dialog() -> Result<(), Box> { let mut envs = HashMap::new(); if add_env { loop { - let key = cliclack::input("Environment variable name:") + let key: String = cliclack::input("Environment variable name:") .placeholder("API_KEY") .interact()?; - let value = cliclack::password("Environment variable value:") + let value: String = cliclack::password("Environment variable value:") .mask('▪') .interact()?; @@ -457,8 +443,8 @@ pub fn configure_extensions_dialog() -> Result<(), Box> { } } - config.extensions.insert( - name.clone(), + ExtensionManager::set( + &name, ExtensionEntry { enabled: true, config: ExtensionConfig::Sse { @@ -466,13 +452,12 @@ pub fn configure_extensions_dialog() -> Result<(), Box> { envs: Envs::new(envs), }, }, - ); + )?; cliclack::outro(format!("Added {} extension", style(name).green()))?; } _ => unreachable!(), }; - config.save()?; Ok(()) } diff --git a/crates/goose-cli/src/commands/session.rs b/crates/goose-cli/src/commands/session.rs index b6d39638e7..0365e168ea 100644 --- a/crates/goose-cli/src/commands/session.rs +++ b/crates/goose-cli/src/commands/session.rs @@ -1,97 +1,43 @@ -use console::style; -use goose::agents::AgentFactory; -use goose::providers::factory; use rand::{distributions::Alphanumeric, Rng}; -use std::path::{Path, PathBuf}; use std::process; -use crate::config::Config; use crate::prompt::rustyline::RustylinePrompt; -use crate::prompt::Prompt; use crate::session::{ensure_session_dir, get_most_recent_session, Session}; use goose::agents::extension::ExtensionError; +use goose::agents::AgentFactory; +use goose::config::{Config, ExtensionManager}; +use goose::providers::create; + use mcp_client::transport::Error as McpClientError; -/// Get the provider and model to use, following priority: -/// 1. CLI arguments -/// 2. Environment variables -/// 3. Config file -fn get_provider_and_model( - cli_provider: Option, - cli_model: Option, - config: &Config, -) -> (String, String) { - let provider = cli_provider - .or_else(|| std::env::var("GOOSE_PROVIDER").ok()) - .unwrap_or_else(|| config.default_provider.clone()); +pub async fn build_session(name: Option, resume: bool) -> Session<'static> { + // Load config and get provider/model + let config = Config::global(); - let model = cli_model - .or_else(|| std::env::var("GOOSE_MODEL").ok()) - .unwrap_or_else(|| config.default_model.clone()); - - (provider, model) -} - -pub async fn build_session<'a>( - session: Option, - provider: Option, - model: Option, - agent_version: Option, - resume: bool, -) -> Box> { + let provider_name: String = config + .get("GOOSE_PROVIDER") + .expect("No provider configured. Run 'goose configure' first"); let session_dir = ensure_session_dir().expect("Failed to create session directory"); - let session_file = if resume && session.is_none() { - // When resuming without a specific session name, use the most recent session - get_most_recent_session().expect("Failed to get most recent session") - } else { - session_path(session.clone(), &session_dir, session.is_none() && !resume) - }; - // Guard against resuming a non-existent session - if resume && !session_file.exists() { - panic!( - "Cannot resume session: file {} does not exist", - session_file.display() - ); + let model = config + .get("GOOSE_MODEL") + .expect("No model configured. Run 'goose configure' first"); + let model_config = goose::model::ModelConfig::new(model); + let provider = create(&provider_name, model_config).expect("Failed to create provider"); + + // Create the agent + let agent_version: Option = config.get("GOOSE_AGENT").ok(); + let mut agent = match agent_version { + Some(version) => AgentFactory::create(&version, provider), + None => AgentFactory::create(AgentFactory::default_version(), provider), } + .expect("Failed to create agent"); - // Guard against running a new session with a file that already exists - if !resume && session_file.exists() { - panic!( - "Session file {} already exists. Use --resume to continue an existing session", - session_file.display() - ); - } - - let config_path = Config::config_path().expect("should identify default config path"); - - if !config_path.exists() { - println!("No configuration found. Please run 'goose configure' first."); - process::exit(1); - } - - let config = Config::load().unwrap_or_else(|_| { - println!("The loaded configuration from {} was invalid", config_path.display()); - println!(" please edit the file to make it valid or consider deleting and recreating it via `goose configure`"); - process::exit(1); - }); - - let (provider_name, model_name) = get_provider_and_model(provider, model, &config); - let provider = factory::get_provider(&provider_name).unwrap(); - - let mut agent = AgentFactory::create( - agent_version - .as_deref() - .unwrap_or(AgentFactory::default_version()), - provider, - ) - .unwrap(); - - // Add configured extensions - for (name, _) in config.extensions.iter() { - if let Some(extension_config) = config.get_extension_config(name) { + // Setup extensions for the agent + for (name, extension) in ExtensionManager::get_all().expect("should load extensions") { + if extension.enabled { agent - .add_extension(extension_config.clone()) + .add_extension(extension.config.clone()) .await .unwrap_or_else(|e| { let err = match e { @@ -107,82 +53,43 @@ pub async fn build_session<'a>( } } - let prompt = match std::env::var("GOOSE_INPUT") { - Ok(val) => match val.as_str() { - "rustyline" => Box::new(RustylinePrompt::new()) as Box, - _ => Box::new(RustylinePrompt::new()) as Box, - }, - Err(_) => Box::new(RustylinePrompt::new()), - }; - - display_session_info(resume, provider_name, model_name, session_file.as_path()); - Box::new(Session::new(agent, prompt, session_file)) -} - -fn session_path( - provided_session_name: Option, - session_dir: &Path, - retry_on_conflict: bool, -) -> PathBuf { - let session_name = provided_session_name.unwrap_or(random_session_name()); - let session_file = session_dir.join(format!("{}.jsonl", session_name)); - - if session_file.exists() && retry_on_conflict { - generate_new_session_path(session_dir) - } else { - session_file - } -} - -fn random_session_name() -> String { - rand::thread_rng() - .sample_iter(&Alphanumeric) - .take(8) - .map(char::from) - .collect::() - .to_lowercase() -} - -// For auto-generated names, try up to 5 times to get a unique name -fn generate_new_session_path(session_dir: &Path) -> PathBuf { - let mut attempts = 0; - let max_attempts = 5; - - loop { - let generated_name = random_session_name(); - let generated_file = session_dir.join(format!("{}.jsonl", generated_name)); - - if !generated_file.exists() { - break generated_file; - } - - attempts += 1; - if attempts >= max_attempts { - panic!( - "Failed to generate unique session name after {} attempts", - max_attempts - ); + // If resuming, try to find the session + if resume { + if let Some(ref session_name) = name { + // Try to resume specific session + let session_file = session_dir.join(format!("{}.jsonl", session_name)); + if session_file.exists() { + let prompt = Box::new(RustylinePrompt::new()); + return Session::new(agent, prompt, session_file); + } else { + eprintln!("Session '{}' not found, starting new session", session_name); + } + } else { + // Try to resume most recent session + if let Ok(session_file) = get_most_recent_session() { + let prompt = Box::new(RustylinePrompt::new()); + return Session::new(agent, prompt, session_file); + } else { + eprintln!("No previous sessions found, starting new session"); + } } } -} -fn display_session_info(resume: bool, provider: String, model: String, session_file: &Path) { - let start_session_msg = if resume { - "resuming session |" - } else { - "starting session |" - }; - println!( - "{} {} {} {} {}", - style(start_session_msg).dim(), - style("provider:").dim(), - style(provider).cyan().dim(), - style("model:").dim(), - style(model).cyan().dim(), - ); - println!( - " {} {}", - style("logging to").dim(), - style(session_file.display()).dim().cyan(), - ); + // Generate session name if not provided + let name = name.unwrap_or_else(|| { + rand::thread_rng() + .sample_iter(&Alphanumeric) + .take(8) + .map(char::from) + .collect() + }); + + let session_file = session_dir.join(format!("{}.jsonl", name)); + if session_file.exists() { + eprintln!("Session '{}' already exists", name); + process::exit(1); + } + + let prompt = Box::new(RustylinePrompt::new()); + Session::new(agent, prompt, session_file) } diff --git a/crates/goose-cli/src/config.rs b/crates/goose-cli/src/config.rs deleted file mode 100644 index 272f0acacc..0000000000 --- a/crates/goose-cli/src/config.rs +++ /dev/null @@ -1,203 +0,0 @@ -use std::collections::HashMap; -use std::fs; -use std::path::PathBuf; - -use anyhow::{Context, Result}; -use serde::{Deserialize, Serialize}; - -use goose::agents::ExtensionConfig; - -const DEFAULT_EXTENSION: &str = "developer"; - -/// Core configuration for Goose CLI -#[derive(Debug, Deserialize, Serialize)] -pub struct Config { - pub default_provider: String, - pub default_model: String, - pub extensions: HashMap, -} - -/// An extension configuration entry with an enabled flag and configuration -#[derive(Debug, Deserialize, Serialize, Clone)] -pub struct ExtensionEntry { - pub enabled: bool, - #[serde(flatten)] - pub config: ExtensionConfig, -} - -impl Config { - /// Get the configuration file path - pub fn config_path() -> Result { - let home_dir = dirs::home_dir().context("Could not determine home directory")?; - let config_dir = home_dir.join(".config").join("goose"); - if !config_dir.exists() { - fs::create_dir_all(&config_dir)?; - } - Ok(config_dir.join("config.yaml")) - } - - /// Load the configuration from disk - pub fn load() -> Result { - let path = Self::config_path()?; - if !path.exists() { - return Err(anyhow::anyhow!("Config has not yet been created")); - } - let content = fs::read_to_string(path)?; - Ok(serde_yaml::from_str(&content)?) - } - - /// Save the configuration to disk - pub fn save(&self) -> Result<()> { - let path = Self::config_path()?; - let content = serde_yaml::to_string(self)?; - fs::write(path, content)?; - Ok(()) - } - - /// Get the extension configuration if enabled - pub fn get_extension_config(&self, name: &str) -> Option { - let entry = self.extensions.get(name)?; - if entry.enabled { - Some(entry.config.clone()) - } else { - None - } - } -} - -impl Default for Config { - fn default() -> Self { - Self { - default_provider: "".to_string(), - default_model: "".to_string(), - extensions: HashMap::from([( - DEFAULT_EXTENSION.to_string(), - ExtensionEntry { - enabled: true, - config: ExtensionConfig::Builtin { - name: DEFAULT_EXTENSION.to_string(), - }, - }, - )]), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - /// This test provides a comprehensive example of all possible configuration options - /// and validates that they are correctly parsed - #[test] - fn test_comprehensive_config() { - let yaml = r#" -# Core settings for the default provider and model -default_provider: openai -default_model: gpt-4 - -# Extension configurations showing all possible variants -extensions: - # Built-in extension that just needs to be enabled - developer: - enabled: true - type: builtin - name: developer - - # Built-in extension that is disabled - unused: - enabled: false - type: builtin - name: unused - - # Full stdio extension configuration with all options - python: - enabled: true - type: stdio - cmd: python3 - args: - - "-m" - - "goose.extensions.python" - envs: - PYTHONPATH: /path/to/python - DEBUG: "true" - - # Full SSE extension configuration - remote: - enabled: true - type: sse - uri: http://localhost:8000/events - envs: - API_KEY: secret - DEBUG: "true" - - # Disabled full extension configuration - disabled_extension: - enabled: false - type: stdio - cmd: test - args: [] - envs: {} -"#; - let config: Config = serde_yaml::from_str(yaml).unwrap(); - - // Check core settings - assert_eq!(config.default_provider, "openai"); - assert_eq!(config.default_model, "gpt-4"); - - // Check builtin enabled extension - match &config.extensions.get("developer").unwrap().config { - ExtensionConfig::Builtin { name } => assert_eq!(name, "developer"), - _ => panic!("Expected builtin extension config"), - } - assert!(config.extensions.get("developer").unwrap().enabled); - - // Check builtin disabled extension - match &config.extensions.get("unused").unwrap().config { - ExtensionConfig::Builtin { name } => assert_eq!(name, "unused"), - _ => panic!("Expected builtin extension config"), - } - assert!(!config.extensions.get("unused").unwrap().enabled); - - // Check full stdio extension - let python = config.extensions.get("python").unwrap(); - assert!(python.enabled); - match &python.config { - ExtensionConfig::Stdio { cmd, args, envs } => { - assert_eq!(cmd, "python3"); - assert_eq!( - args, - &vec!["-m".to_string(), "goose.extensions.python".to_string()] - ); - let env = envs.get_env(); - assert_eq!(env.get("PYTHONPATH").unwrap(), "/path/to/python"); - assert_eq!(env.get("DEBUG").unwrap(), "true"); - } - _ => panic!("Expected stdio extension config"), - } - - // Check full SSE extension - let remote = config.extensions.get("remote").unwrap(); - assert!(remote.enabled); - match &remote.config { - ExtensionConfig::Sse { uri, envs } => { - assert_eq!(uri, "http://localhost:8000/events"); - let env = envs.get_env(); - assert_eq!(env.get("API_KEY").unwrap(), "secret"); - assert_eq!(env.get("DEBUG").unwrap(), "true"); - } - _ => panic!("Expected sse extension config"), - } - - // Check disabled full extension - assert!(!config.extensions.get("disabled_extension").unwrap().enabled); - - // Test the get_extension_config helper - assert!(config.get_extension_config("developer").is_some()); - assert!(config.get_extension_config("unused").is_none()); - assert!(config.get_extension_config("python").is_some()); - assert!(config.get_extension_config("remote").is_some()); - assert!(config.get_extension_config("disabled_extension").is_none()); - assert!(config.get_extension_config("nonexistent").is_none()); - } -} diff --git a/crates/goose-cli/src/main.rs b/crates/goose-cli/src/main.rs index 2820fa2fd9..897da1ab18 100644 --- a/crates/goose-cli/src/main.rs +++ b/crates/goose-cli/src/main.rs @@ -1,9 +1,7 @@ use anyhow::Result; use clap::{CommandFactory, Parser, Subcommand}; -use goose::agents::AgentFactory; mod commands; -mod config; mod log_usage; mod logging; mod prompt; @@ -14,8 +12,8 @@ use commands::configure::handle_configure; use commands::mcp::run_server; use commands::session::build_session; use commands::version::print_version; -use config::Config; use console::style; +use goose::config::Config; use logging::setup_logging; use std::io::{self, Read}; @@ -36,25 +34,7 @@ struct Cli { enum Command { /// Configure Goose settings #[command(about = "Configure Goose settings")] - Configure { - /// AI Provider to use - #[arg( - short, - long, - help = "AI Provider to use (e.g., 'openai', 'databricks', 'ollama')", - long_help = "Specify AI Provider to use (e.g., 'openai', 'databricks', 'ollama')." - )] - provider: Option, - - /// Model to use - #[arg( - short, - long, - help = "Model to use (e.g., 'gpt-4', 'llama2')", - long_help = "Specify which model to use." - )] - model: Option, - }, + Configure {}, /// Manage system prompts and behaviors #[command(about = "Run one of the mcp servers bundled with goose")] @@ -73,33 +53,6 @@ enum Command { )] name: Option, - /// Provider to use (overrides config) - #[arg( - short, - long, - help = "Provider to use (e.g., 'openai', 'anthropic')", - long_help = "Override the default provider from config" - )] - provider: Option, - - /// Model to use (overrides config) - #[arg( - short, - long, - help = "Model to use (e.g., 'gpt-4', 'claude-3')", - long_help = "Override the default model from config" - )] - model: Option, - - /// Agent version to use (e.g., 'default', 'v1') - #[arg( - short, - long, - help = "Agent version to use (e.g., 'default', 'v1'), defaults to 'default'", - long_help = "Specify which agent version to use for this session." - )] - agent: Option, - /// Resume a previous session #[arg( short, @@ -134,24 +87,6 @@ enum Command { )] input_text: Option, - /// Provider to use (overrides config) - #[arg( - short, - long, - help = "Provider to use (e.g., 'openai', 'anthropic')", - long_help = "Override the default provider from config" - )] - provider: Option, - - /// Model to use (overrides config) - #[arg( - short, - long, - help = "Model to use (e.g., 'gpt-4', 'claude-3')", - long_help = "Override the default model from config" - )] - model: Option, - /// Name for this run session #[arg( short, @@ -162,15 +97,6 @@ enum Command { )] name: Option, - /// Agent version to use (e.g., 'default', 'v1') - #[arg( - short, - long, - help = "Agent version to use (e.g., 'default', 'v1')", - long_help = "Specify which agent version to use for this session." - )] - agent: Option, - /// Resume a previous run #[arg( short, @@ -203,36 +129,15 @@ async fn main() -> Result<()> { } match cli.command { - Some(Command::Configure { provider, model }) => { - let _ = handle_configure(provider, model).await; + Some(Command::Configure {}) => { + let _ = handle_configure().await; return Ok(()); } Some(Command::Mcp { name }) => { let _ = run_server(&name).await; } - Some(Command::Session { - name, - provider, - model, - agent, - resume, - }) => { - if let Some(agent_version) = agent.clone() { - if !AgentFactory::available_versions().contains(&agent_version.as_str()) { - eprintln!("Error: Invalid agent version '{}'", agent_version); - eprintln!("Available versions:"); - for version in AgentFactory::available_versions() { - if version == AgentFactory::default_version() { - eprintln!("* {} (default)", version); - } else { - eprintln!(" {}", version); - } - } - std::process::exit(1); - } - } - - let mut session = build_session(name, provider, model, agent, resume).await; + Some(Command::Session { name, resume }) => { + let mut session = build_session(name, resume).await; setup_logging(session.session_file().file_stem().and_then(|s| s.to_str()))?; let _ = session.start().await; @@ -241,10 +146,7 @@ async fn main() -> Result<()> { Some(Command::Run { instructions, input_text, - provider, - model, name, - agent, resume, }) => { // Validate that we have some input source @@ -253,21 +155,6 @@ async fn main() -> Result<()> { std::process::exit(1); } - if let Some(agent_version) = agent.clone() { - if !AgentFactory::available_versions().contains(&agent_version.as_str()) { - eprintln!("Error: Invalid agent version '{}'", agent_version); - eprintln!("Available versions:"); - for version in AgentFactory::available_versions() { - if version == AgentFactory::default_version() { - eprintln!("* {} (default)", version); - } else { - eprintln!(" {}", version); - } - } - std::process::exit(1); - } - } - let contents = if let Some(file_name) = instructions { let file_path = std::path::Path::new(&file_name); std::fs::read_to_string(file_path).expect("Failed to read the instruction file") @@ -280,7 +167,7 @@ async fn main() -> Result<()> { .expect("Failed to read from stdin"); stdin }; - let mut session = build_session(name, provider, model, agent, resume).await; + let mut session = build_session(name, resume).await; let _ = session.headless_start(contents.clone()).await; return Ok(()); } @@ -291,7 +178,7 @@ async fn main() -> Result<()> { None => { Cli::command().print_help()?; println!(); - if Config::load().is_err() { + if !Config::global().exists() { println!( "\n {}: Run '{}' to setup goose for the first time", style("Tip").green().italic(), diff --git a/crates/goose-server/src/routes/agent.rs b/crates/goose-server/src/routes/agent.rs index 9b651bd1b6..dfd683075e 100644 --- a/crates/goose-server/src/routes/agent.rs +++ b/crates/goose-server/src/routes/agent.rs @@ -5,7 +5,8 @@ use axum::{ routing::{get, post}, Json, Router, }; -use goose::{agents::AgentFactory, providers::factory}; +use goose::config::Config; +use goose::{agents::AgentFactory, model::ModelConfig, providers}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::env; @@ -82,7 +83,15 @@ async fn create_agent( println!("Set environment variable: {}={}", env_var_key, model); } - let provider = factory::get_provider(&payload.provider).expect("Failed to create provider"); + let config = Config::global(); + let model = payload.model.unwrap_or_else(|| { + config + .get("GOOSE_MODEL") + .expect("Did not find a model on payload or in env") + }); + let model_config = ModelConfig::new(model); + let provider = + providers::create(&payload.provider, model_config).expect("Failed to create provider"); let version = payload .version diff --git a/crates/goose-server/src/routes/extension.rs b/crates/goose-server/src/routes/extension.rs index 2287379b83..d0f04bb106 100644 --- a/crates/goose-server/src/routes/extension.rs +++ b/crates/goose-server/src/routes/extension.rs @@ -4,7 +4,7 @@ use crate::state::AppState; use axum::{extract::State, routing::post, Json, Router}; use goose::{ agents::{extension::Envs, ExtensionConfig}, - key_manager::{get_keyring_secret, KeyRetrievalStrategy}, + config::Config, }; use http::{HeaderMap, StatusCode}; use serde::{Deserialize, Serialize}; @@ -65,6 +65,9 @@ async fn add_extension( return Err(StatusCode::UNAUTHORIZED); } + // Load the configuration + let config = Config::global(); + // Initialize a vector to collect any missing keys. let mut missing_keys = Vec::new(); @@ -73,7 +76,7 @@ async fn add_extension( ExtensionConfigRequest::Sse { uri, env_keys } => { let mut env_map = HashMap::new(); for key in env_keys { - match get_keyring_secret(&key, KeyRetrievalStrategy::KeyringOnly) { + match config.get_secret(&key) { Ok(value) => { env_map.insert(key, value); } @@ -105,7 +108,7 @@ async fn add_extension( } => { let mut env_map = HashMap::new(); for key in env_keys { - match get_keyring_secret(&key, KeyRetrievalStrategy::KeyringOnly) { + match config.get_secret(&key) { Ok(value) => { env_map.insert(key, value); } diff --git a/crates/goose-server/src/routes/reply.rs b/crates/goose-server/src/routes/reply.rs index b10833288f..803a7a0c26 100644 --- a/crates/goose-server/src/routes/reply.rs +++ b/crates/goose-server/src/routes/reply.rs @@ -440,8 +440,8 @@ mod tests { use super::*; use goose::{ agents::AgentFactory, + model::ModelConfig, providers::base::{Provider, ProviderUsage, Usage}, - providers::configs::ModelConfig, }; use mcp_core::tool::Tool; @@ -453,8 +453,12 @@ mod tests { #[async_trait::async_trait] impl Provider for MockProvider { - fn get_model_config(&self) -> &ModelConfig { - &self.model_config + fn metadata() -> goose::providers::base::ProviderMetadata { + goose::providers::base::ProviderMetadata::empty() + } + + fn get_model_config(&self) -> ModelConfig { + self.model_config.clone() } async fn complete( @@ -468,10 +472,6 @@ mod tests { ProviderUsage::new("mock".to_string(), Usage::default()), )) } - - fn get_usage(&self, _data: &Value) -> anyhow::Result { - Ok(Usage::new(None, None, None)) - } } #[test] diff --git a/crates/goose-server/src/routes/secrets.rs b/crates/goose-server/src/routes/secrets.rs index e2303c4a8e..03278e1713 100644 --- a/crates/goose-server/src/routes/secrets.rs +++ b/crates/goose-server/src/routes/secrets.rs @@ -1,11 +1,10 @@ use crate::state::AppState; use axum::{extract::State, routing::delete, routing::post, Json, Router}; -use goose::key_manager::{ - delete_from_keyring, get_keyring_secret, save_to_keyring, KeyRetrievalStrategy, -}; +use goose::config::Config; use http::{HeaderMap, StatusCode}; use once_cell::sync::Lazy; use serde::{Deserialize, Serialize}; +use serde_json::Value; use std::collections::HashMap; #[derive(Serialize)] @@ -34,7 +33,7 @@ async fn store_secret( return Err(StatusCode::UNAUTHORIZED); } - match save_to_keyring(&request.key, &request.value) { + match Config::global().set_secret(&request.key, Value::String(request.value)) { Ok(_) => Ok(Json(SecretResponse { error: false })), Err(_) => Ok(Json(SecretResponse { error: true })), } @@ -76,7 +75,7 @@ static PROVIDER_ENV_REQUIREMENTS: Lazy> = Lazy:: fn check_key_status(key: &str) -> (bool, Option) { if let Ok(_value) = std::env::var(key) { (true, Some("env".to_string())) - } else if get_keyring_secret(key, KeyRetrievalStrategy::KeyringOnly).is_ok() { + } else if Config::global().get_secret::(key).is_ok() { (true, Some("keyring".to_string())) } else { (false, None) @@ -151,7 +150,7 @@ async fn delete_secret( } // Attempt to delete the key - match delete_from_keyring(&request.key) { + match Config::global().delete_secret(&request.key) { Ok(_) => Ok(StatusCode::NO_CONTENT), Err(_) => Err(StatusCode::NOT_FOUND), } diff --git a/crates/goose/Cargo.toml b/crates/goose/Cargo.toml index 2cd19b908f..e194007544 100644 --- a/crates/goose/Cargo.toml +++ b/crates/goose/Cargo.toml @@ -53,16 +53,17 @@ lazy_static = "1.5" tracing = "0.1" tracing-subscriber = "0.3" wiremock = "0.6.0" -keyring = { version = "3.6.1", features = [ - "apple-native", - "windows-native", - "sync-secret-service", -] } +keyring = { version = "3.6.1", features = ["apple-native", "windows-native", "sync-secret-service"] } ctor = "0.2.7" paste = "1.0" +serde_yaml = "0.9.34" +once_cell = "1.20.2" +dirs = "6.0.0" [dev-dependencies] criterion = "0.5" +tempfile = "3.15.0" +serial_test = "3.2.0" [[example]] name = "agent" @@ -74,4 +75,4 @@ path = "examples/databricks_oauth.rs" [[bench]] name = "tokenization_benchmark" -harness = false \ No newline at end of file +harness = false diff --git a/crates/goose/examples/agent.rs b/crates/goose/examples/agent.rs index 02816e622a..a6c9cc27c8 100644 --- a/crates/goose/examples/agent.rs +++ b/crates/goose/examples/agent.rs @@ -9,7 +9,7 @@ async fn main() { // Setup a model provider from env vars let _ = dotenv(); - let provider = Box::new(DatabricksProvider::from_env().expect("should create provider")); + let provider = Box::new(DatabricksProvider::default()); // Setup an agent with the developer extension let mut agent = AgentFactory::create("reference", provider).expect("default should exist"); diff --git a/crates/goose/examples/databricks_oauth.rs b/crates/goose/examples/databricks_oauth.rs index 568fd131b4..2b49cc03a8 100644 --- a/crates/goose/examples/databricks_oauth.rs +++ b/crates/goose/examples/databricks_oauth.rs @@ -14,7 +14,7 @@ async fn main() -> Result<()> { std::env::remove_var("DATABRICKS_TOKEN"); // Create the provider - let provider = DatabricksProvider::from_env()?; + let provider = DatabricksProvider::default(); // Create a simple message let message = Message::user().with_text("Tell me a short joke about programming."); diff --git a/crates/goose/examples/image_tool.rs b/crates/goose/examples/image_tool.rs index 9a833d42bd..f3b5d94a52 100644 --- a/crates/goose/examples/image_tool.rs +++ b/crates/goose/examples/image_tool.rs @@ -19,8 +19,8 @@ async fn main() -> Result<()> { // Create providers let providers: Vec> = vec![ - Box::new(DatabricksProvider::from_env()?), - Box::new(OpenAiProvider::from_env()?), + Box::new(DatabricksProvider::default()), + Box::new(OpenAiProvider::default()), ]; for provider in providers { diff --git a/crates/goose/src/agents/capabilities.rs b/crates/goose/src/agents/capabilities.rs index f6b888b57b..4daa8d9867 100644 --- a/crates/goose/src/agents/capabilities.rs +++ b/crates/goose/src/agents/capabilities.rs @@ -517,8 +517,8 @@ impl Capabilities { mod tests { use super::*; use crate::message::Message; - use crate::providers::base::{Provider, ProviderUsage, Usage}; - use crate::providers::configs::ModelConfig; + use crate::model::ModelConfig; + use crate::providers::base::{Provider, ProviderMetadata, ProviderUsage, Usage}; use mcp_client::client::Error; use mcp_client::client::McpClientTrait; use mcp_core::protocol::{ @@ -534,8 +534,12 @@ mod tests { #[async_trait::async_trait] impl Provider for MockProvider { - fn get_model_config(&self) -> &ModelConfig { - &self.model_config + fn metadata() -> ProviderMetadata { + ProviderMetadata::empty() + } + + fn get_model_config(&self) -> ModelConfig { + self.model_config.clone() } async fn complete( @@ -549,10 +553,6 @@ mod tests { ProviderUsage::new("mock".to_string(), Usage::default()), )) } - - fn get_usage(&self, _data: &serde_json::Value) -> anyhow::Result { - Ok(Usage::new(None, None, None)) - } } struct MockClient {} diff --git a/crates/goose/src/agents/truncate.rs b/crates/goose/src/agents/truncate.rs index 9fb891f3e6..94218d820f 100644 --- a/crates/goose/src/agents/truncate.rs +++ b/crates/goose/src/agents/truncate.rs @@ -258,8 +258,8 @@ register_agent!("truncate", TruncateAgent); mod tests { use crate::agents::truncate::TruncateAgent; use crate::message::Message; - use crate::providers::base::{Provider, ProviderUsage, Usage}; - use crate::providers::configs::ModelConfig; + use crate::model::ModelConfig; + use crate::providers::base::{Provider, ProviderMetadata, ProviderUsage, Usage}; use mcp_core::{Content, Tool}; use std::iter; @@ -271,8 +271,12 @@ mod tests { #[async_trait::async_trait] impl Provider for MockProvider { - fn get_model_config(&self) -> &ModelConfig { - &self.model_config + fn metadata() -> ProviderMetadata { + ProviderMetadata::empty() + } + + fn get_model_config(&self) -> ModelConfig { + self.model_config.clone() } async fn complete( @@ -286,10 +290,6 @@ mod tests { ProviderUsage::new("mock".to_string(), Usage::default()), )) } - - fn get_usage(&self, _data: &serde_json::Value) -> anyhow::Result { - Ok(Usage::new(None, None, None)) - } } const SMALL_MESSAGE: &str = "This is a test, this is just a test, this is only a test.\n"; @@ -365,9 +365,9 @@ mod tests { #[tokio::test] async fn test_truncation_when_conversation_history_too_big() -> anyhow::Result<()> { let conversation = create_basic_valid_conversation(5000, false); - let messages = call_enforce_ctx_limit(&*conversation).await?; - assert_eq!(conversation.len() > messages.len(), true); - assert_eq!(messages.len() > 0, true); + let messages = call_enforce_ctx_limit(&conversation).await?; + assert!(conversation.len() > messages.len()); + assert!(!messages.is_empty()); Ok(()) } @@ -380,18 +380,18 @@ mod tests { let mut conversation = create_basic_valid_conversation(3, false); conversation.push(Message::user().with_text(oversized_message)); - let messages = call_enforce_ctx_limit(&*conversation).await; + let messages = call_enforce_ctx_limit(&conversation).await; - assert!(matches!(messages, Err(_, ..))); + assert!(messages.is_err()); Ok(()) } #[tokio::test] async fn test_truncation_when_tool_response_set_too_big() -> anyhow::Result<()> { let conversation = create_basic_valid_conversation(5000, true); - let messages = call_enforce_ctx_limit(&*conversation).await?; - assert_eq!(conversation.len() > messages.len(), true); - assert_eq!(messages.len() > 0, true); + let messages = call_enforce_ctx_limit(&conversation).await?; + assert!(conversation.len() > messages.len()); + assert!(!messages.is_empty()); Ok(()) } } diff --git a/crates/goose/src/config/base.rs b/crates/goose/src/config/base.rs new file mode 100644 index 0000000000..d8416b2642 --- /dev/null +++ b/crates/goose/src/config/base.rs @@ -0,0 +1,462 @@ +use keyring::Entry; +use once_cell::sync::OnceCell; +use serde::Deserialize; +use serde_json::Value; +use std::collections::HashMap; +use std::env; +use std::path::{Path, PathBuf}; +use thiserror::Error; + +const KEYRING_SERVICE: &str = "goose"; +const KEYRING_USERNAME: &str = "secrets"; + +#[cfg(test)] +const TEST_KEYRING_SERVICE: &str = "goose-test"; + +#[derive(Error, Debug)] +pub enum ConfigError { + #[error("Configuration value not found: {0}")] + NotFound(String), + #[error("Failed to deserialize value: {0}")] + DeserializeError(String), + #[error("Failed to read config file: {0}")] + FileError(#[from] std::io::Error), + #[error("Failed to create config directory: {0}")] + DirectoryError(String), + #[error("Failed to access keyring: {0}")] + KeyringError(String), +} + +impl From for ConfigError { + fn from(err: serde_json::Error) -> Self { + ConfigError::DeserializeError(err.to_string()) + } +} + +impl From for ConfigError { + fn from(err: serde_yaml::Error) -> Self { + ConfigError::DeserializeError(err.to_string()) + } +} + +impl From for ConfigError { + fn from(err: keyring::Error) -> Self { + ConfigError::KeyringError(err.to_string()) + } +} + +/// Configuration management for Goose. +/// +/// This module provides a flexible configuration system that supports: +/// - Dynamic configuration keys +/// - Multiple value types through serde deserialization +/// - Environment variable overrides +/// - YAML-based configuration file storage +/// - Hot reloading of configuration changes +/// - Secure secret storage in system keyring +/// +/// Configuration values are loaded with the following precedence: +/// 1. Environment variables (exact key match) +/// 2. Configuration file (~/.config/goose/config.yaml by default) +/// +/// Secrets are loaded with the following precedence: +/// 1. Environment variables (exact key match) +/// 2. System keyring +/// +/// # Examples +/// +/// ```no_run +/// use goose::config::Config; +/// use serde::Deserialize; +/// +/// // Get a string value +/// let config = Config::global(); +/// let api_key: String = config.get("OPENAI_API_KEY").unwrap(); +/// +/// // Get a complex type +/// #[derive(Deserialize)] +/// struct ServerConfig { +/// host: String, +/// port: u16, +/// } +/// +/// let server_config: ServerConfig = config.get("server").unwrap(); +/// ``` +/// +/// # Naming Convention +/// we recommend snake_case for keys, and will convert to UPPERCASE when +/// checking for environment overrides. e.g. openai_api_key will check for an +/// environment variable OPENAI_API_KEY +/// +/// For Goose-specific configuration, consider prefixing with "goose_" to avoid conflicts. +pub struct Config { + config_path: PathBuf, + keyring_service: String, +} + +// Global instance +static GLOBAL_CONFIG: OnceCell = OnceCell::new(); + +impl Default for Config { + fn default() -> Self { + let config_dir = dirs::home_dir() + .expect("goose requires a home dir") + .join(".config") + .join("goose"); + std::fs::create_dir_all(&config_dir).expect("Failed to create config directory"); + + let config_path = config_dir.join("config.yaml"); + Config { + config_path, + keyring_service: KEYRING_SERVICE.to_string(), + } + } +} + +impl Config { + /// Get the global configuration instance. + /// + /// This will initialize the configuration with the default path (~/.config/goose/config.yaml) + /// if it hasn't been initialized yet. + pub fn global() -> &'static Config { + GLOBAL_CONFIG.get_or_init(Config::default) + } + + /// Create a new configuration instance with custom paths + /// + /// This is primarily useful for testing or for applications that need + /// to manage multiple configuration files. + pub fn new>(config_path: P, service: &str) -> Result { + Ok(Config { + config_path: config_path.as_ref().to_path_buf(), + keyring_service: service.to_string(), + }) + } + + /// Check if this config already exists + pub fn exists(&self) -> bool { + self.config_path.exists() + } + + /// Get the path to the configuration file + pub fn path(&self) -> String { + self.config_path.to_string_lossy().to_string() + } + + // Load current values from the config file + fn load_values(&self) -> Result, ConfigError> { + if self.config_path.exists() { + let file_content = std::fs::read_to_string(&self.config_path)?; + // Parse YAML into JSON Value for consistent internal representation + let yaml_value: serde_yaml::Value = serde_yaml::from_str(&file_content)?; + let json_value: Value = serde_json::to_value(yaml_value)?; + + match json_value { + Value::Object(map) => Ok(map.into_iter().collect()), + _ => Ok(HashMap::new()), + } + } else { + Ok(HashMap::new()) + } + } + + // Load current secrets from the keyring + fn load_secrets(&self) -> Result, ConfigError> { + let entry = Entry::new(&self.keyring_service, KEYRING_USERNAME)?; + + match entry.get_password() { + Ok(content) => { + let values: HashMap = serde_json::from_str(&content)?; + Ok(values) + } + Err(keyring::Error::NoEntry) => Ok(HashMap::new()), + Err(e) => Err(ConfigError::KeyringError(e.to_string())), + } + } + + /// Get a configuration value. + /// + /// This will attempt to get the value from: + /// 1. Environment variable with the exact key name + /// 2. Configuration file + /// + /// The value will be deserialized into the requested type. This works with + /// both simple types (String, i32, etc.) and complex types that implement + /// serde::Deserialize. + /// + /// # Errors + /// + /// Returns a ConfigError if: + /// - The key doesn't exist in either environment or config file + /// - The value cannot be deserialized into the requested type + /// - There is an error reading the config file + pub fn get Deserialize<'de>>(&self, key: &str) -> Result { + // First check environment variables (convert to uppercase) + let env_key = key.to_uppercase(); + if let Ok(val) = env::var(&env_key) { + // Parse the environment variable value into a serde_json::Value + let value: Value = serde_json::from_str(&val).unwrap_or(Value::String(val)); + return Ok(serde_json::from_value(value)?); + } + + // Load current values from file + let values = self.load_values()?; + + // Then check our stored values + values + .get(key) + .ok_or_else(|| ConfigError::NotFound(key.to_string())) + .and_then(|v| Ok(serde_json::from_value(v.clone())?)) + } + + /// Set a configuration value in the config file. + /// + /// This will immediately write the value to the config file. The value + /// can be any type that can be serialized to JSON/YAML. + /// + /// Note that this does not affect environment variables - those can only + /// be set through the system environment. + /// + /// # Errors + /// + /// Returns a ConfigError if: + /// - There is an error reading or writing the config file + /// - There is an error serializing the value + pub fn set(&self, key: &str, value: Value) -> Result<(), ConfigError> { + let mut values = self.load_values()?; + values.insert(key.to_string(), value); + + // Convert to YAML for storage + let yaml_value = serde_yaml::to_string(&values)?; + + // Ensure the directory exists + if let Some(parent) = self.config_path.parent() { + std::fs::create_dir_all(parent) + .map_err(|e| ConfigError::DirectoryError(e.to_string()))?; + } + + std::fs::write(&self.config_path, yaml_value)?; + Ok(()) + } + + /// Get a secret value. + /// + /// This will attempt to get the value from: + /// 1. Environment variable with the exact key name + /// 2. System keyring + /// + /// The value will be deserialized into the requested type. This works with + /// both simple types (String, i32, etc.) and complex types that implement + /// serde::Deserialize. + /// + /// # Errors + /// + /// Returns a ConfigError if: + /// - The key doesn't exist in either environment or keyring + /// - The value cannot be deserialized into the requested type + /// - There is an error accessing the keyring + pub fn get_secret Deserialize<'de>>(&self, key: &str) -> Result { + // First check environment variables (convert to uppercase) + let env_key = key.to_uppercase(); + if let Ok(val) = env::var(&env_key) { + let value: Value = serde_json::from_str(&val).unwrap_or(Value::String(val)); + return Ok(serde_json::from_value(value)?); + } + + // Then check keyring + let values = self.load_secrets()?; + values + .get(key) + .ok_or_else(|| ConfigError::NotFound(key.to_string())) + .and_then(|v| Ok(serde_json::from_value(v.clone())?)) + } + + /// Set a secret value in the system keyring. + /// + /// This will store the value in a single JSON object in the system keyring, + /// alongside any other secrets. The value can be any type that can be + /// serialized to JSON. + /// + /// Note that this does not affect environment variables - those can only + /// be set through the system environment. + /// + /// # Errors + /// + /// Returns a ConfigError if: + /// - There is an error accessing the keyring + /// - There is an error serializing the value + pub fn set_secret(&self, key: &str, value: Value) -> Result<(), ConfigError> { + let mut values = self.load_secrets()?; + values.insert(key.to_string(), value); + + let json_value = serde_json::to_string(&values)?; + let entry = Entry::new(&self.keyring_service, KEYRING_USERNAME)?; + entry.set_password(&json_value)?; + Ok(()) + } + + /// Delete a secret from the system keyring. + /// + /// This will remove the specified key from the JSON object in the system keyring. + /// Other secrets will remain unchanged. + /// + /// # Errors + /// + /// Returns a ConfigError if: + /// - There is an error accessing the keyring + /// - There is an error serializing the remaining values + pub fn delete_secret(&self, key: &str) -> Result<(), ConfigError> { + let mut values = self.load_secrets()?; + values.remove(key); + + let json_value = serde_json::to_string(&values)?; + let entry = Entry::new(&self.keyring_service, KEYRING_USERNAME)?; + entry.set_password(&json_value)?; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serial_test::serial; + use tempfile::NamedTempFile; + + fn cleanup_keyring() -> Result<(), ConfigError> { + let entry = Entry::new(TEST_KEYRING_SERVICE, KEYRING_USERNAME)?; + match entry.delete_credential() { + Ok(_) => Ok(()), + Err(keyring::Error::NoEntry) => Ok(()), + Err(e) => Err(ConfigError::KeyringError(e.to_string())), + } + } + + #[test] + fn test_basic_config() -> Result<(), ConfigError> { + let temp_file = NamedTempFile::new().unwrap(); + let config = Config::new(temp_file.path(), TEST_KEYRING_SERVICE)?; + + // Set a simple string value + config.set("test_key", Value::String("test_value".to_string()))?; + + // Test simple string retrieval + let value: String = config.get("test_key")?; + assert_eq!(value, "test_value"); + + // Test with environment variable override + std::env::set_var("TEST_KEY", "env_value"); + let value: String = config.get("test_key")?; + assert_eq!(value, "env_value"); + + Ok(()) + } + + #[test] + fn test_complex_type() -> Result<(), ConfigError> { + #[derive(Deserialize, Debug, PartialEq)] + struct TestStruct { + field1: String, + field2: i32, + } + + let temp_file = NamedTempFile::new().unwrap(); + let config = Config::new(temp_file.path(), TEST_KEYRING_SERVICE)?; + + // Set a complex value + config.set( + "complex_key", + serde_json::json!({ + "field1": "hello", + "field2": 42 + }), + )?; + + let value: TestStruct = config.get("complex_key")?; + assert_eq!(value.field1, "hello"); + assert_eq!(value.field2, 42); + + Ok(()) + } + + #[test] + fn test_missing_value() { + let temp_file = NamedTempFile::new().unwrap(); + let config = Config::new(temp_file.path(), TEST_KEYRING_SERVICE).unwrap(); + + let result: Result = config.get("nonexistent_key"); + assert!(matches!(result, Err(ConfigError::NotFound(_)))); + } + + #[test] + fn test_yaml_formatting() -> Result<(), ConfigError> { + let temp_file = NamedTempFile::new().unwrap(); + let config = Config::new(temp_file.path(), TEST_KEYRING_SERVICE)?; + + config.set("key1", Value::String("value1".to_string()))?; + config.set("key2", Value::Number(42.into()))?; + + // Read the file directly to check YAML formatting + let content = std::fs::read_to_string(temp_file.path())?; + assert!(content.contains("key1: value1")); + assert!(content.contains("key2: 42")); + + Ok(()) + } + + #[test] + #[serial] + fn test_secret_management() -> Result<(), ConfigError> { + cleanup_keyring()?; + let temp_file = NamedTempFile::new().unwrap(); + let config = Config::new(temp_file.path(), TEST_KEYRING_SERVICE)?; + + // Test setting and getting a simple secret + config.set_secret("api_key", Value::String("secret123".to_string()))?; + let value: String = config.get_secret("api_key")?; + assert_eq!(value, "secret123"); + + // Test environment variable override + std::env::set_var("API_KEY", "env_secret"); + let value: String = config.get_secret("api_key")?; + assert_eq!(value, "env_secret"); + std::env::remove_var("API_KEY"); + + // Test deleting a secret + config.delete_secret("api_key")?; + let result: Result = config.get_secret("api_key"); + assert!(matches!(result, Err(ConfigError::NotFound(_)))); + + cleanup_keyring()?; + Ok(()) + } + + #[test] + #[serial] + fn test_multiple_secrets() -> Result<(), ConfigError> { + cleanup_keyring()?; + let temp_file = NamedTempFile::new().unwrap(); + let config = Config::new(temp_file.path(), TEST_KEYRING_SERVICE)?; + + // Set multiple secrets + config.set_secret("key1", Value::String("secret1".to_string()))?; + config.set_secret("key2", Value::String("secret2".to_string()))?; + + // Verify both exist + let value1: String = config.get_secret("key1")?; + let value2: String = config.get_secret("key2")?; + assert_eq!(value1, "secret1"); + assert_eq!(value2, "secret2"); + + // Delete one secret + config.delete_secret("key1")?; + + // Verify key1 is gone but key2 remains + let result1: Result = config.get_secret("key1"); + let value2: String = config.get_secret("key2")?; + assert!(matches!(result1, Err(ConfigError::NotFound(_)))); + assert_eq!(value2, "secret2"); + + cleanup_keyring()?; + Ok(()) + } +} diff --git a/crates/goose/src/config/extensions.rs b/crates/goose/src/config/extensions.rs new file mode 100644 index 0000000000..50e0181774 --- /dev/null +++ b/crates/goose/src/config/extensions.rs @@ -0,0 +1,106 @@ +use anyhow::Result; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; + +use super::base::Config; +use crate::agents::ExtensionConfig; + +const DEFAULT_EXTENSION: &str = "developer"; + +#[derive(Debug, Deserialize, Serialize, Clone)] +pub struct ExtensionEntry { + pub enabled: bool, + #[serde(flatten)] + pub config: ExtensionConfig, +} + +/// Extension configuration management +pub struct ExtensionManager; + +impl ExtensionManager { + /// Get the extension configuration if enabled + pub fn get_config(name: &str) -> Result> { + let config = Config::global(); + + // Try to get the extension entry + let extensions: HashMap = match config.get("extensions") { + Ok(exts) => exts, + Err(super::ConfigError::NotFound(_)) => { + // Initialize with default developer extension + let defaults = HashMap::from([( + DEFAULT_EXTENSION.to_string(), + ExtensionEntry { + enabled: true, + config: ExtensionConfig::Builtin { + name: DEFAULT_EXTENSION.to_string(), + }, + }, + )]); + config.set("extensions", serde_json::to_value(&defaults)?)?; + defaults + } + Err(e) => return Err(e.into()), + }; + + Ok(extensions.get(name).and_then(|entry| { + if entry.enabled { + Some(entry.config.clone()) + } else { + None + } + })) + } + + /// Set or update an extension configuration + pub fn set(name: &str, entry: ExtensionEntry) -> Result<()> { + let config = Config::global(); + + let mut extensions: HashMap = + config.get("extensions").unwrap_or_else(|_| HashMap::new()); + + extensions.insert(name.to_string(), entry); + config.set("extensions", serde_json::to_value(extensions)?)?; + Ok(()) + } + + /// Remove an extension configuration + pub fn remove(name: &str) -> Result<()> { + let config = Config::global(); + + let mut extensions: HashMap = + config.get("extensions").unwrap_or_else(|_| HashMap::new()); + + extensions.remove(name); + config.set("extensions", serde_json::to_value(extensions)?)?; + Ok(()) + } + + /// Enable or disable an extension + pub fn set_enabled(name: &str, enabled: bool) -> Result<()> { + let config = Config::global(); + + let mut extensions: HashMap = + config.get("extensions").unwrap_or_else(|_| HashMap::new()); + + if let Some(entry) = extensions.get_mut(name) { + entry.enabled = enabled; + config.set("extensions", serde_json::to_value(extensions)?)?; + } + Ok(()) + } + + /// Get all extensions and their configurations + pub fn get_all() -> Result> { + let config = Config::global(); + Ok(config.get("extensions").unwrap_or_else(|_| HashMap::new())) + } + + /// Check if an extension is enabled + pub fn is_enabled(name: &str) -> Result { + let config = Config::global(); + let extensions: HashMap = + config.get("extensions").unwrap_or_else(|_| HashMap::new()); + + Ok(extensions.get(name).map(|e| e.enabled).unwrap_or(false)) + } +} diff --git a/crates/goose/src/config/mod.rs b/crates/goose/src/config/mod.rs new file mode 100644 index 0000000000..fdaa8a864d --- /dev/null +++ b/crates/goose/src/config/mod.rs @@ -0,0 +1,6 @@ +mod base; +mod extensions; + +pub use crate::agents::ExtensionConfig; +pub use base::{Config, ConfigError}; +pub use extensions::{ExtensionEntry, ExtensionManager}; diff --git a/crates/goose/src/key_manager.rs b/crates/goose/src/key_manager.rs deleted file mode 100644 index 9d6355b1a7..0000000000 --- a/crates/goose/src/key_manager.rs +++ /dev/null @@ -1,192 +0,0 @@ -use anyhow::{Context, Result}; -use keyring::Entry; -use std::env; -use thiserror::Error; - -#[derive(Error, Debug)] -pub enum KeyManagerError { - #[error("Failed to access keyring: {0}")] - KeyringAccess(String), - - #[error("Failed to save to keyring: {0}")] - KeyringSave(String), - - #[error("Failed to access environment variable: {0}")] - EnvVarAccess(String), -} - -impl From for KeyManagerError { - fn from(err: keyring::Error) -> Self { - KeyManagerError::KeyringAccess(err.to_string()) - } -} - -impl From for KeyManagerError { - fn from(err: env::VarError) -> Self { - KeyManagerError::EnvVarAccess(err.to_string()) - } -} - -#[derive(Debug, Clone, Copy)] -pub enum KeyRetrievalStrategy { - /// Only look in environment variables - EnvironmentOnly, - /// Only look in system keyring - KeyringOnly, - /// Try keyring first, then environment variables (default behavior) - Both, -} - -impl Default for KeyRetrievalStrategy { - fn default() -> Self { - Self::Both - } -} - -pub fn get_keyring_secret(key_name: &str, strategy: KeyRetrievalStrategy) -> Result { - let kr = Entry::new("goose", key_name).context("Failed to create keyring entry")?; - match strategy { - KeyRetrievalStrategy::EnvironmentOnly => { - env::var(key_name).context("Failed to get environment variable") - } - KeyRetrievalStrategy::KeyringOnly => kr - .get_password() - .context("Failed to get password from keyring"), - KeyRetrievalStrategy::Both => { - // Try environment first, then keyring - env::var(key_name).or_else(|_| { - kr.get_password().context(format!( - "Could not find {} key in keyring or environment variables", - key_name - )) - }) - } - } -} - -pub fn key_value_in_environment_variable(key_name: &str) -> Option { - match env::var(key_name) { - Ok(value) => Some(value), - Err(_) => None, - } -} - -pub fn key_value_in_key_chain(key_name: &str) -> Option { - let kr = Entry::new("goose", key_name).ok()?; - - match kr.get_password() { - Ok(password) => Some(password), - Err(_) => None, - } -} - -pub fn save_to_keyring(key_name: &str, key_val: &str) -> std::result::Result<(), KeyManagerError> { - let kr = Entry::new("goose", key_name)?; - kr.set_password(key_val).map_err(KeyManagerError::from) -} - -pub fn delete_from_keyring(key_name: &str) -> std::result::Result<(), KeyManagerError> { - let kr = Entry::new("goose", key_name)?; - kr.delete_credential().map_err(KeyManagerError::from) -} - -#[cfg(test)] -mod tests { - use super::*; - - const TEST_ENV_PREFIX: &str = "GOOSE_TEST_"; - - fn cleanup_env(key: &str) { - std::env::remove_var(key); - } - - fn cleanup_keyring(key: &str) -> Result<(), KeyManagerError> { - let kr = Entry::new("goose", key)?; - kr.delete_credential().map_err(KeyManagerError::from) - } - - #[test] - fn test_delete_from_keyring() { - let key_name = format!("{}{}", TEST_ENV_PREFIX, "DELETE_KEY"); - - // Save a value to the keyring - save_to_keyring(&key_name, "test_value").unwrap(); - - // Verify it exists - let kr = Entry::new("goose", &key_name).unwrap(); - assert_eq!(kr.get_password().unwrap(), "test_value"); - - // Delete the keyring entry - let result = delete_from_keyring(&key_name); - assert!(result.is_ok()); - - // Verify deletion - let kr = Entry::new("goose", &key_name).unwrap(); - let password_result = kr.get_password(); - assert!(password_result.is_err()); - } - - #[test] - fn test_get_key_environment_only() { - let key_name = format!("{}{}", TEST_ENV_PREFIX, "ENV_KEY"); - std::env::set_var(&key_name, "test_value"); - - let result = get_keyring_secret(&key_name, KeyRetrievalStrategy::EnvironmentOnly); - assert_eq!(result.unwrap(), "test_value"); - - cleanup_env(&key_name); - } - - #[test] - fn test_get_key_environment_only_missing() { - let key_name = format!("{}{}", TEST_ENV_PREFIX, "MISSING_KEY"); - - let result = get_keyring_secret(&key_name, KeyRetrievalStrategy::EnvironmentOnly); - assert!(result.is_err()); - } - - #[test] - fn test_get_key_keyring_only() { - let key_name = format!("{}{}", TEST_ENV_PREFIX, "KEYRING_KEY"); - - // First save a value - save_to_keyring(&key_name, "test_value").unwrap(); - - let result = get_keyring_secret(&key_name, KeyRetrievalStrategy::KeyringOnly); - assert_eq!(result.unwrap(), "test_value"); - - cleanup_keyring(&key_name).unwrap(); - } - - #[test] - fn test_get_key_both() { - let key_name = format!("{}{}", TEST_ENV_PREFIX, "BOTH_KEY"); - - // Test environment first - std::env::set_var(&key_name, "env_value"); - let result = get_keyring_secret(&key_name, KeyRetrievalStrategy::Both); - assert_eq!(result.unwrap(), "env_value"); - - // Test keyring takes precedence - save_to_keyring(&key_name, "keyring_value").unwrap(); - let result = get_keyring_secret(&key_name, KeyRetrievalStrategy::Both); - assert_eq!(result.unwrap(), "env_value"); // Environment still takes precedence - - cleanup_env(&key_name); - cleanup_keyring(&key_name).unwrap(); - } - - #[test] - fn test_save_to_keyring() { - let key_name = format!("{}{}", TEST_ENV_PREFIX, "SAVE_KEY"); - - let result = save_to_keyring(&key_name, "test_value"); - assert!(result.is_ok()); - - // Verify the value was saved - let kr = Entry::new("goose", &key_name).unwrap(); - assert_eq!(kr.get_password().unwrap(), "test_value"); - - cleanup_keyring(&key_name).unwrap(); - } -} diff --git a/crates/goose/src/lib.rs b/crates/goose/src/lib.rs index 010e7018f3..00c8f2c3b3 100644 --- a/crates/goose/src/lib.rs +++ b/crates/goose/src/lib.rs @@ -1,6 +1,7 @@ pub mod agents; -pub mod key_manager; +pub mod config; pub mod message; +pub mod model; pub mod prompt_template; pub mod providers; pub mod token_counter; diff --git a/crates/goose/src/providers/configs.rs b/crates/goose/src/model.rs similarity index 100% rename from crates/goose/src/providers/configs.rs rename to crates/goose/src/model.rs diff --git a/crates/goose/src/providers/anthropic.rs b/crates/goose/src/providers/anthropic.rs index e251b95962..332c59e17b 100644 --- a/crates/goose/src/providers/anthropic.rs +++ b/crates/goose/src/providers/anthropic.rs @@ -4,11 +4,11 @@ use reqwest::Client; use serde_json::Value; use std::time::Duration; -use super::base::{Provider, ProviderUsage, Usage}; -use super::configs::ModelConfig; +use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage}; use super::formats::anthropic::{create_request, get_usage, response_to_message}; use super::utils::{emit_debug_trace, get_model, handle_response}; use crate::message::Message; +use crate::model::ModelConfig; use mcp_core::tool::Tool; pub const ANTHROPIC_DEFAULT_MODEL: &str = "claude-3-5-sonnet-latest"; @@ -22,14 +22,20 @@ pub struct AnthropicProvider { model: ModelConfig, } +impl Default for AnthropicProvider { + fn default() -> Self { + let model = ModelConfig::new(AnthropicProvider::metadata().default_model); + AnthropicProvider::from_env(model).expect("Failed to initialize Anthropic provider") + } +} + impl AnthropicProvider { - pub fn from_env() -> Result { - let api_key = - crate::key_manager::get_keyring_secret("ANTHROPIC_API_KEY", Default::default())?; - let host = std::env::var("ANTHROPIC_HOST") + pub fn from_env(model: ModelConfig) -> Result { + let config = crate::config::Config::global(); + let api_key: String = config.get_secret("ANTHROPIC_API_KEY")?; + let host: String = config + .get("ANTHROPIC_HOST") .unwrap_or_else(|_| "https://api.anthropic.com".to_string()); - let model_name = std::env::var("ANTHROPIC_MODEL") - .unwrap_or_else(|_| ANTHROPIC_DEFAULT_MODEL.to_string()); let client = Client::builder() .timeout(Duration::from_secs(600)) @@ -39,7 +45,7 @@ impl AnthropicProvider { client, host, api_key, - model: ModelConfig::new(model_name), + model, }) } @@ -61,8 +67,26 @@ impl AnthropicProvider { #[async_trait] impl Provider for AnthropicProvider { - fn get_model_config(&self) -> &ModelConfig { - &self.model + fn metadata() -> ProviderMetadata { + ProviderMetadata::new( + "anthropic", + "Anthropic", + "Claude and other models from Anthropic", + ANTHROPIC_DEFAULT_MODEL, + vec![ + ConfigKey::new("ANTHROPIC_API_KEY", true, true, None), + ConfigKey::new( + "ANTHROPIC_HOST", + false, + false, + Some("https://api.anthropic.com"), + ), + ], + ) + } + + fn get_model_config(&self) -> ModelConfig { + self.model.clone() } #[tracing::instrument( @@ -82,13 +106,9 @@ impl Provider for AnthropicProvider { // Parse response let message = response_to_message(response.clone())?; - let usage = self.get_usage(&response)?; + let usage = get_usage(&response)?; let model = get_model(&response); emit_debug_trace(self, &payload, &response, &usage); Ok((message, ProviderUsage::new(model, usage))) } - - fn get_usage(&self, data: &Value) -> Result { - get_usage(data) - } } diff --git a/crates/goose/src/providers/base.rs b/crates/goose/src/providers/base.rs index 182ca57a72..a96317e629 100644 --- a/crates/goose/src/providers/base.rs +++ b/crates/goose/src/providers/base.rs @@ -1,11 +1,72 @@ use anyhow::Result; use serde::{Deserialize, Serialize}; -use serde_json::Value; -use super::configs::ModelConfig; use crate::message::Message; +use crate::model::ModelConfig; use mcp_core::tool::Tool; +/// Metadata about a provider's configuration requirements and capabilities +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ProviderMetadata { + /// The unique identifier for this provider + pub name: String, + /// Display name for the provider in UIs + pub display_name: String, + /// Description of the provider's capabilities + pub description: String, + /// The default/recommended model for this provider + pub default_model: String, + /// Required configuration keys + pub config_keys: Vec, +} + +impl ProviderMetadata { + pub fn new( + name: &str, + display_name: &str, + description: &str, + default_model: &str, + config_keys: Vec, + ) -> Self { + Self { + name: name.to_string(), + display_name: display_name.to_string(), + description: description.to_string(), + default_model: default_model.to_string(), + config_keys, + } + } + + pub fn empty() -> Self { + Self { + name: "".to_string(), + display_name: "".to_string(), + description: "".to_string(), + default_model: "".to_string(), + config_keys: vec![], + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ConfigKey { + pub name: String, + pub required: bool, + pub secret: bool, + pub default: Option, +} + +impl ConfigKey { + pub fn new(name: &str, required: bool, secret: bool, default: Option<&str>) -> Self { + Self { + name: name.to_string(), + required, + secret, + default: default.map(|s| s.to_string()), + } + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ProviderUsage { pub model: String, @@ -44,8 +105,10 @@ use async_trait::async_trait; /// Base trait for AI providers (OpenAI, Anthropic, etc) #[async_trait] pub trait Provider: Send + Sync { - /// Get the model configuration - fn get_model_config(&self) -> &ModelConfig; + /// Get the metadata for this provider type + fn metadata() -> ProviderMetadata + where + Self: Sized; /// Generate the next message using the configured model and other parameters /// @@ -63,7 +126,8 @@ pub trait Provider: Send + Sync { tools: &[Tool], ) -> Result<(Message, ProviderUsage)>; - fn get_usage(&self, data: &Value) -> Result; + /// Get the model config from the provider + fn get_model_config(&self) -> ModelConfig; } #[cfg(test)] diff --git a/crates/goose/src/providers/databricks.rs b/crates/goose/src/providers/databricks.rs index e5ef1b90d0..0cafe84cd3 100644 --- a/crates/goose/src/providers/databricks.rs +++ b/crates/goose/src/providers/databricks.rs @@ -5,20 +5,20 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use std::time::Duration; -use super::base::{Provider, ProviderUsage, Usage}; -use super::configs::ModelConfig; +use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage}; use super::formats::openai::{ create_request, get_usage, is_context_length_error, response_to_message, }; use super::oauth; use super::utils::{get_model, handle_response, ImageFormat}; use crate::message::Message; +use crate::model::ModelConfig; use mcp_core::tool::Tool; const DEFAULT_CLIENT_ID: &str = "databricks-cli"; const DEFAULT_REDIRECT_URL: &str = "http://localhost:8020"; const DEFAULT_SCOPES: &[&str] = &["all-apis"]; -pub const DATABRICKS_DEFAULT_MODEL: &str = "claude-3-5-sonnet-2"; +pub const DATABRICKS_DEFAULT_MODEL: &str = "databricks-meta-llama-3-3-70b-instruct"; #[derive(Debug, Clone, Serialize, Deserialize)] pub enum DatabricksAuth { @@ -56,28 +56,29 @@ pub struct DatabricksProvider { image_format: ImageFormat, } +impl Default for DatabricksProvider { + fn default() -> Self { + let model = ModelConfig::new(DatabricksProvider::metadata().default_model); + DatabricksProvider::from_env(model).expect("Failed to initialize Databricks provider") + } +} + impl DatabricksProvider { - pub fn from_env() -> Result { - // Although we don't need host to be stored secretly, we use the keyring to make - // it easier to coordinate with configuration. We could consider a non secret storage tool - // elsewhere in the future - let host = crate::key_manager::get_keyring_secret("DATABRICKS_HOST", Default::default())?; - let model_name = std::env::var("DATABRICKS_MODEL") - .unwrap_or_else(|_| DATABRICKS_DEFAULT_MODEL.to_string()); + pub fn from_env(model: ModelConfig) -> Result { + let config = crate::config::Config::global(); + let host: String = config.get("DATABRICKS_HOST")?; let client = Client::builder() .timeout(Duration::from_secs(600)) .build()?; // If we find a databricks token we prefer that - if let Ok(api_key) = - crate::key_manager::get_keyring_secret("DATABRICKS_TOKEN", Default::default()) - { + if let Ok(api_key) = config.get_secret("DATABRICKS_TOKEN") { return Ok(Self { client, host: host.clone(), auth: DatabricksAuth::token(api_key), - model: ModelConfig::new(model_name), + model, image_format: ImageFormat::Anthropic, }); } @@ -85,9 +86,9 @@ impl DatabricksProvider { // Otherwise use Oauth flow Ok(Self { client, - host: host.clone(), - auth: DatabricksAuth::oauth(host), - model: ModelConfig::new(model_name), + auth: DatabricksAuth::oauth(host.clone()), + host, + model, image_format: ImageFormat::Anthropic, }) } @@ -130,8 +131,21 @@ impl DatabricksProvider { #[async_trait] impl Provider for DatabricksProvider { - fn get_model_config(&self) -> &ModelConfig { - &self.model + fn metadata() -> ProviderMetadata { + ProviderMetadata::new( + "databricks", + "Databricks", + "Models on Databricks AI Gateway", + DATABRICKS_DEFAULT_MODEL, + vec![ + ConfigKey::new("DATABRICKS_HOST", true, false, None), + ConfigKey::new("DATABRICKS_TOKEN", false, true, None), + ], + ) + } + + fn get_model_config(&self) -> ModelConfig { + self.model.clone() } #[tracing::instrument( @@ -163,14 +177,10 @@ impl Provider for DatabricksProvider { // Parse response let message = response_to_message(response.clone())?; - let usage = self.get_usage(&response)?; + let usage = get_usage(&response)?; let model = get_model(&response); super::utils::emit_debug_trace(self, &payload, &response, &usage); Ok((message, ProviderUsage::new(model, usage))) } - - fn get_usage(&self, data: &Value) -> Result { - get_usage(data) - } } diff --git a/crates/goose/src/providers/factory.rs b/crates/goose/src/providers/factory.rs index 813f79d1bc..99a257ae50 100644 --- a/crates/goose/src/providers/factory.rs +++ b/crates/goose/src/providers/factory.rs @@ -1,19 +1,37 @@ use super::{ - anthropic::AnthropicProvider, base::Provider, databricks::DatabricksProvider, - google::GoogleProvider, groq::GroqProvider, ollama::OllamaProvider, openai::OpenAiProvider, + anthropic::AnthropicProvider, + base::{Provider, ProviderMetadata}, + databricks::DatabricksProvider, + google::GoogleProvider, + groq::GroqProvider, + ollama::OllamaProvider, + openai::OpenAiProvider, openrouter::OpenRouterProvider, }; +use crate::model::ModelConfig; use anyhow::Result; -pub fn get_provider(name: &str) -> Result> { +pub fn providers() -> Vec { + vec![ + AnthropicProvider::metadata(), + DatabricksProvider::metadata(), + GoogleProvider::metadata(), + GroqProvider::metadata(), + OllamaProvider::metadata(), + OpenAiProvider::metadata(), + OpenRouterProvider::metadata(), + ] +} + +pub fn create(name: &str, model: ModelConfig) -> Result> { match name { - "openai" => Ok(Box::new(OpenAiProvider::from_env()?)), - "anthropic" => Ok(Box::new(AnthropicProvider::from_env()?)), - "databricks" => Ok(Box::new(DatabricksProvider::from_env()?)), - "groq" => Ok(Box::new(GroqProvider::from_env()?)), - "ollama" => Ok(Box::new(OllamaProvider::from_env()?)), - "openrouter" => Ok(Box::new(OpenRouterProvider::from_env()?)), - "google" => Ok(Box::new(GoogleProvider::from_env()?)), + "openai" => Ok(Box::new(OpenAiProvider::from_env(model)?)), + "anthropic" => Ok(Box::new(AnthropicProvider::from_env(model)?)), + "databricks" => Ok(Box::new(DatabricksProvider::from_env(model)?)), + "groq" => Ok(Box::new(GroqProvider::from_env(model)?)), + "ollama" => Ok(Box::new(OllamaProvider::from_env(model)?)), + "openrouter" => Ok(Box::new(OpenRouterProvider::from_env(model)?)), + "google" => Ok(Box::new(GoogleProvider::from_env(model)?)), _ => Err(anyhow::anyhow!("Unknown provider: {}", name)), } } diff --git a/crates/goose/src/providers/formats/anthropic.rs b/crates/goose/src/providers/formats/anthropic.rs index bb92e4bd63..070ad63655 100644 --- a/crates/goose/src/providers/formats/anthropic.rs +++ b/crates/goose/src/providers/formats/anthropic.rs @@ -1,6 +1,6 @@ use crate::message::{Message, MessageContent}; +use crate::model::ModelConfig; use crate::providers::base::Usage; -use crate::providers::configs::ModelConfig; use anyhow::{anyhow, Result}; use mcp_core::content::Content; use mcp_core::role::Role; diff --git a/crates/goose/src/providers/formats/google.rs b/crates/goose/src/providers/formats/google.rs index 9500379ac2..01060f70d4 100644 --- a/crates/goose/src/providers/formats/google.rs +++ b/crates/goose/src/providers/formats/google.rs @@ -1,6 +1,6 @@ use crate::message::{Message, MessageContent}; +use crate::model::ModelConfig; use crate::providers::base::Usage; -use crate::providers::configs::ModelConfig; use crate::providers::utils::{is_valid_function_name, sanitize_function_name}; use anyhow::Result; use mcp_core::content::Content; diff --git a/crates/goose/src/providers/formats/openai.rs b/crates/goose/src/providers/formats/openai.rs index 5bc9feba40..3a91b979ea 100644 --- a/crates/goose/src/providers/formats/openai.rs +++ b/crates/goose/src/providers/formats/openai.rs @@ -1,6 +1,6 @@ use crate::message::{Message, MessageContent}; +use crate::model::ModelConfig; use crate::providers::base::Usage; -use crate::providers::configs::ModelConfig; use crate::providers::errors::ProviderError; use crate::providers::utils::{ convert_image, is_valid_function_name, sanitize_function_name, ImageFormat, diff --git a/crates/goose/src/providers/google.rs b/crates/goose/src/providers/google.rs index 5705362d56..80bec4888d 100644 --- a/crates/goose/src/providers/google.rs +++ b/crates/goose/src/providers/google.rs @@ -1,6 +1,6 @@ use crate::message::Message; -use crate::providers::base::{Provider, ProviderUsage, Usage}; -use crate::providers::configs::ModelConfig; +use crate::model::ModelConfig; +use crate::providers::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage}; use crate::providers::formats::google::{create_request, get_usage, response_to_message}; use crate::providers::utils::{emit_debug_trace, handle_response, unescape_json_values}; use anyhow::Result; @@ -22,12 +22,20 @@ pub struct GoogleProvider { model: ModelConfig, } +impl Default for GoogleProvider { + fn default() -> Self { + let model = ModelConfig::new(GoogleProvider::metadata().default_model); + GoogleProvider::from_env(model).expect("Failed to initialize Google provider") + } +} + impl GoogleProvider { - pub fn from_env() -> Result { - let api_key = crate::key_manager::get_keyring_secret("GOOGLE_API_KEY", Default::default())?; - let host = std::env::var("GOOGLE_HOST").unwrap_or_else(|_| GOOGLE_API_HOST.to_string()); - let model_name = - std::env::var("GOOGLE_MODEL").unwrap_or_else(|_| GOOGLE_DEFAULT_MODEL.to_string()); + pub fn from_env(model: ModelConfig) -> Result { + let config = crate::config::Config::global(); + let api_key: String = config.get_secret("GOOGLE_API_KEY")?; + let host: String = config + .get("GOOGLE_HOST") + .unwrap_or_else(|_| GOOGLE_API_HOST.to_string()); let client = Client::builder() .timeout(Duration::from_secs(600)) @@ -37,7 +45,7 @@ impl GoogleProvider { client, host, api_key, - model: ModelConfig::new(model_name), + model, }) } @@ -63,8 +71,21 @@ impl GoogleProvider { #[async_trait] impl Provider for GoogleProvider { - fn get_model_config(&self) -> &ModelConfig { - &self.model + fn metadata() -> ProviderMetadata { + ProviderMetadata::new( + "google", + "Google Gemini", + "Gemini models from Google AI", + GOOGLE_DEFAULT_MODEL, + vec![ + ConfigKey::new("GOOGLE_API_KEY", true, true, None), + ConfigKey::new("GOOGLE_HOST", false, false, Some(GOOGLE_API_HOST)), + ], + ) + } + + fn get_model_config(&self) -> ModelConfig { + self.model.clone() } #[tracing::instrument( @@ -84,7 +105,7 @@ impl Provider for GoogleProvider { // Parse response let message = response_to_message(unescape_json_values(&response))?; - let usage = self.get_usage(&response)?; + let usage = get_usage(&response)?; let model = match response.get("modelVersion") { Some(model_version) => model_version.as_str().unwrap_or_default().to_string(), None => self.model.model_name.clone(), @@ -93,8 +114,4 @@ impl Provider for GoogleProvider { let provider_usage = ProviderUsage::new(model, usage); Ok((message, provider_usage)) } - - fn get_usage(&self, data: &Value) -> Result { - get_usage(data) - } } diff --git a/crates/goose/src/providers/groq.rs b/crates/goose/src/providers/groq.rs index c84c79ddaf..0a68922d47 100644 --- a/crates/goose/src/providers/groq.rs +++ b/crates/goose/src/providers/groq.rs @@ -1,6 +1,6 @@ use crate::message::Message; -use crate::providers::base::{Provider, ProviderUsage, Usage}; -use crate::providers::configs::ModelConfig; +use crate::model::ModelConfig; +use crate::providers::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage}; use crate::providers::formats::openai::{create_request, get_usage, response_to_message}; use crate::providers::utils::{get_model, handle_response}; use anyhow::Result; @@ -22,12 +22,20 @@ pub struct GroqProvider { model: ModelConfig, } +impl Default for GroqProvider { + fn default() -> Self { + let model = ModelConfig::new(GroqProvider::metadata().default_model); + GroqProvider::from_env(model).expect("Failed to initialize Groq provider") + } +} + impl GroqProvider { - pub fn from_env() -> Result { - let api_key = crate::key_manager::get_keyring_secret("GROQ_API_KEY", Default::default())?; - let host = std::env::var("GROQ_HOST").unwrap_or_else(|_| GROQ_API_HOST.to_string()); - let model_name = - std::env::var("GROQ_MODEL").unwrap_or_else(|_| GROQ_DEFAULT_MODEL.to_string()); + pub fn from_env(model: ModelConfig) -> Result { + let config = crate::config::Config::global(); + let api_key: String = config.get_secret("GROQ_API_KEY")?; + let host: String = config + .get("GROQ_HOST") + .unwrap_or_else(|_| GROQ_API_HOST.to_string()); let client = Client::builder() .timeout(Duration::from_secs(600)) @@ -37,7 +45,7 @@ impl GroqProvider { client, host, api_key, - model: ModelConfig::new(model_name), + model, }) } @@ -61,8 +69,21 @@ impl GroqProvider { #[async_trait] impl Provider for GroqProvider { - fn get_model_config(&self) -> &ModelConfig { - &self.model + fn metadata() -> ProviderMetadata { + ProviderMetadata::new( + "groq", + "Groq", + "Fast inference with Groq hardware", + GROQ_DEFAULT_MODEL, + vec![ + ConfigKey::new("GROQ_API_KEY", true, true, None), + ConfigKey::new("GROQ_HOST", false, false, Some(GROQ_API_HOST)), + ], + ) + } + + fn get_model_config(&self) -> ModelConfig { + self.model.clone() } #[tracing::instrument( @@ -86,13 +107,9 @@ impl Provider for GroqProvider { let response = self.post(payload.clone()).await?; let message = response_to_message(response.clone())?; - let usage = self.get_usage(&response)?; + let usage = get_usage(&response)?; let model = get_model(&response); super::utils::emit_debug_trace(self, &payload, &response, &usage); Ok((message, ProviderUsage::new(model, usage))) } - - fn get_usage(&self, data: &Value) -> anyhow::Result { - get_usage(data) - } } diff --git a/crates/goose/src/providers/mod.rs b/crates/goose/src/providers/mod.rs index 53c39520b4..52448a58c2 100644 --- a/crates/goose/src/providers/mod.rs +++ b/crates/goose/src/providers/mod.rs @@ -1,9 +1,8 @@ pub mod anthropic; pub mod base; -pub mod configs; pub mod databricks; pub mod errors; -pub mod factory; +mod factory; pub mod formats; pub mod google; pub mod groq; @@ -13,4 +12,4 @@ pub mod openai; pub mod openrouter; pub mod utils; -pub use factory::get_provider; +pub use factory::{create, providers}; diff --git a/crates/goose/src/providers/ollama.rs b/crates/goose/src/providers/ollama.rs index f49832e58e..20e3750d5d 100644 --- a/crates/goose/src/providers/ollama.rs +++ b/crates/goose/src/providers/ollama.rs @@ -1,7 +1,7 @@ -use super::base::{Provider, ProviderUsage, Usage}; -use super::configs::ModelConfig; +use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage}; use super::utils::{get_model, handle_response}; use crate::message::Message; +use crate::model::ModelConfig; use crate::providers::formats::openai::{create_request, get_usage, response_to_message}; use anyhow::Result; use async_trait::async_trait; @@ -21,14 +21,19 @@ pub struct OllamaProvider { model: ModelConfig, } +impl Default for OllamaProvider { + fn default() -> Self { + let model = ModelConfig::new(OllamaProvider::metadata().default_model); + OllamaProvider::from_env(model).expect("Failed to initialize Ollama provider") + } +} + impl OllamaProvider { - pub fn from_env() -> Result { - // Although we don't need host to be stored secretly, we use the keyring to make - // it easier to coordinate with configuration. We could consider a non secret storage tool - // elsewhere in the future - let host = crate::key_manager::get_keyring_secret("OLLAMA_HOST", Default::default()) + pub fn from_env(model: ModelConfig) -> Result { + let config = crate::config::Config::global(); + let host: String = config + .get("OLLAMA_HOST") .unwrap_or_else(|_| OLLAMA_HOST.to_string()); - let model_name = std::env::var("OLLAMA_MODEL").unwrap_or_else(|_| OLLAMA_MODEL.to_string()); let client = Client::builder() .timeout(Duration::from_secs(600)) @@ -37,7 +42,7 @@ impl OllamaProvider { Ok(Self { client, host, - model: ModelConfig::new(model_name), + model, }) } @@ -52,8 +57,23 @@ impl OllamaProvider { #[async_trait] impl Provider for OllamaProvider { - fn get_model_config(&self) -> &ModelConfig { - &self.model + fn metadata() -> ProviderMetadata { + ProviderMetadata::new( + "ollama", + "Ollama", + "Local open source models", + OLLAMA_MODEL, + vec![ConfigKey::new( + "OLLAMA_HOST", + false, + false, + Some(OLLAMA_HOST), + )], + ) + } + + fn get_model_config(&self) -> ModelConfig { + self.model.clone() } #[tracing::instrument( @@ -78,13 +98,9 @@ impl Provider for OllamaProvider { // Parse response let message = response_to_message(response.clone())?; - let usage = self.get_usage(&response)?; + let usage = get_usage(&response)?; let model = get_model(&response); super::utils::emit_debug_trace(self, &payload, &response, &usage); Ok((message, ProviderUsage::new(model, usage))) } - - fn get_usage(&self, data: &Value) -> Result { - get_usage(data) - } } diff --git a/crates/goose/src/providers/openai.rs b/crates/goose/src/providers/openai.rs index 6df238e746..974ba0a62d 100644 --- a/crates/goose/src/providers/openai.rs +++ b/crates/goose/src/providers/openai.rs @@ -4,13 +4,13 @@ use reqwest::Client; use serde_json::Value; use std::time::Duration; -use super::base::{Provider, ProviderUsage, Usage}; -use super::configs::ModelConfig; +use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage}; use super::formats::openai::{ create_request, get_usage, is_context_length_error, response_to_message, }; use super::utils::{emit_debug_trace, get_model, handle_response, ImageFormat}; use crate::message::Message; +use crate::model::ModelConfig; use mcp_core::tool::Tool; pub const OPEN_AI_DEFAULT_MODEL: &str = "gpt-4o"; @@ -24,14 +24,20 @@ pub struct OpenAiProvider { model: ModelConfig, } -impl OpenAiProvider { - pub fn from_env() -> Result { - let api_key = crate::key_manager::get_keyring_secret("OPENAI_API_KEY", Default::default())?; - let host = - std::env::var("OPENAI_HOST").unwrap_or_else(|_| "https://api.openai.com".to_string()); - let model_name = - std::env::var("OPENAI_MODEL").unwrap_or_else(|_| OPEN_AI_DEFAULT_MODEL.to_string()); +impl Default for OpenAiProvider { + fn default() -> Self { + let model = ModelConfig::new(OpenAiProvider::metadata().default_model); + OpenAiProvider::from_env(model).expect("Failed to initialize OpenAI provider") + } +} +impl OpenAiProvider { + pub fn from_env(model: ModelConfig) -> Result { + let config = crate::config::Config::global(); + let api_key: String = config.get_secret("OPENAI_API_KEY")?; + let host: String = config + .get("OPENAI_HOST") + .unwrap_or_else(|_| "https://api.openai.com".to_string()); let client = Client::builder() .timeout(Duration::from_secs(600)) .build()?; @@ -40,7 +46,7 @@ impl OpenAiProvider { client, host, api_key, - model: ModelConfig::new(model_name), + model, }) } @@ -61,8 +67,21 @@ impl OpenAiProvider { #[async_trait] impl Provider for OpenAiProvider { - fn get_model_config(&self) -> &ModelConfig { - &self.model + fn metadata() -> ProviderMetadata { + ProviderMetadata::new( + "openai", + "OpenAI", + "GPT-4 and other OpenAI models", + OPEN_AI_DEFAULT_MODEL, + vec![ + ConfigKey::new("OPENAI_API_KEY", true, true, None), + ConfigKey::new("OPENAI_HOST", false, false, Some("https://api.openai.com")), + ], + ) + } + + fn get_model_config(&self) -> ModelConfig { + self.model.clone() } #[tracing::instrument( @@ -90,31 +109,9 @@ impl Provider for OpenAiProvider { // Parse response let message = response_to_message(response.clone())?; - let usage = self.get_usage(&response)?; + let usage = get_usage(&response)?; let model = get_model(&response); emit_debug_trace(self, &payload, &response, &usage); Ok((message, ProviderUsage::new(model, usage))) } - - fn get_usage(&self, data: &Value) -> Result { - get_usage(data) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_provider_construction() -> Result<()> { - std::env::set_var("OPENAI_API_KEY", "test-key"); - std::env::set_var("OPENAI_HOST", "https://test.openai.com"); - std::env::set_var("OPENAI_MODEL", "gpt-4o"); - - let provider = OpenAiProvider::from_env()?; - assert_eq!(provider.host, "https://test.openai.com"); - assert_eq!(provider.model.model_name, "gpt-4o"); - - Ok(()) - } } diff --git a/crates/goose/src/providers/openrouter.rs b/crates/goose/src/providers/openrouter.rs index 0690111d2d..3cb409041b 100644 --- a/crates/goose/src/providers/openrouter.rs +++ b/crates/goose/src/providers/openrouter.rs @@ -4,10 +4,10 @@ use reqwest::Client; use serde_json::{json, Value}; use std::time::Duration; -use super::base::{Provider, ProviderUsage, Usage}; -use super::configs::ModelConfig; +use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage}; use super::utils::{emit_debug_trace, get_model, handle_response}; use crate::message::Message; +use crate::model::ModelConfig; use crate::providers::formats::openai::{ create_request, get_usage, is_context_length_error, response_to_message, }; @@ -25,14 +25,20 @@ pub struct OpenRouterProvider { model: ModelConfig, } +impl Default for OpenRouterProvider { + fn default() -> Self { + let model = ModelConfig::new(OpenRouterProvider::metadata().default_model); + OpenRouterProvider::from_env(model).expect("Failed to initialize OpenRouter provider") + } +} + impl OpenRouterProvider { - pub fn from_env() -> Result { - let api_key = - crate::key_manager::get_keyring_secret("OPENROUTER_API_KEY", Default::default())?; - let host = std::env::var("OPENROUTER_HOST") + pub fn from_env(model: ModelConfig) -> Result { + let config = crate::config::Config::global(); + let api_key: String = config.get_secret("OPENROUTER_API_KEY")?; + let host: String = config + .get("OPENROUTER_HOST") .unwrap_or_else(|_| "https://openrouter.ai".to_string()); - let model_name = std::env::var("OPENROUTER_MODEL") - .unwrap_or_else(|_| OPENROUTER_DEFAULT_MODEL.to_string()); let client = Client::builder() .timeout(Duration::from_secs(600)) @@ -42,7 +48,7 @@ impl OpenRouterProvider { client, host, api_key, - model: ModelConfig::new(model_name), + model, }) } @@ -149,8 +155,26 @@ fn create_request_based_on_model( #[async_trait] impl Provider for OpenRouterProvider { - fn get_model_config(&self) -> &ModelConfig { - &self.model + fn metadata() -> ProviderMetadata { + ProviderMetadata::new( + "openrouter", + "OpenRouter", + "Router for many model providers", + OPENROUTER_DEFAULT_MODEL, + vec![ + ConfigKey::new("OPENROUTER_API_KEY", true, true, None), + ConfigKey::new( + "OPENROUTER_HOST", + false, + false, + Some("https://openrouter.ai"), + ), + ], + ) + } + + fn get_model_config(&self) -> ModelConfig { + self.model.clone() } #[tracing::instrument( @@ -179,13 +203,9 @@ impl Provider for OpenRouterProvider { // Parse response let message = response_to_message(response.clone())?; - let usage = self.get_usage(&response)?; + let usage = get_usage(&response)?; let model = get_model(&response); emit_debug_trace(self, &payload, &response, &usage); Ok((message, ProviderUsage::new(model, usage))) } - - fn get_usage(&self, data: &Value) -> Result { - get_usage(data) - } } diff --git a/crates/goose/src/token_counter.rs b/crates/goose/src/token_counter.rs index 23e954c2e5..b719433fba 100644 --- a/crates/goose/src/token_counter.rs +++ b/crates/goose/src/token_counter.rs @@ -235,7 +235,7 @@ impl TokenCounter { mod tests { use super::*; use crate::message::{Message, MessageContent}; // or however your `Message` is imported - use crate::providers::configs::{CLAUDE_TOKENIZER, GPT_4O_TOKENIZER}; + use crate::model::{CLAUDE_TOKENIZER, GPT_4O_TOKENIZER}; use mcp_core::role::Role; use mcp_core::tool::Tool; use serde_json::json; diff --git a/crates/goose/tests/providers.rs b/crates/goose/tests/providers.rs index 7e01375b0f..f5caffe5d7 100644 --- a/crates/goose/tests/providers.rs +++ b/crates/goose/tests/providers.rs @@ -219,7 +219,7 @@ async fn test_provider( provider_fn: F, ) -> Result<()> where - F: FnOnce() -> Result, + F: FnOnce() -> T, T: Provider + Send + Sync + 'static, { // We start off as failed, so that if the process panics it is seen as a failure @@ -257,6 +257,12 @@ where // Setup the provider let missing_vars = required_vars.iter().any(|var| std::env::var(var).is_err()); + if missing_vars { + println!("Skipping {} tests - credentials not configured", name); + TEST_REPORT.record_skip(name); + return Ok(()); + } + let provider = provider_fn(); // Restore original environment @@ -273,19 +279,7 @@ where std::mem::drop(lock); - if missing_vars { - println!("Skipping {} tests - credentials not configured", name); - TEST_REPORT.record_skip(name); - return Ok(()); - } - - if provider.is_err() { - println!("Could not setup {} from env", name); - TEST_REPORT.record_fail(name); - return Err(provider.err().expect("is error")); - } - - let tester = ProviderTester::new(provider.expect("already checked"), name.to_string()); + let tester = ProviderTester::new(provider, name.to_string()); match tester.run_test_suite().await { Ok(_) => { TEST_REPORT.record_pass(name); @@ -305,7 +299,7 @@ async fn test_openai_provider() -> Result<()> { "OpenAI", &["OPENAI_API_KEY"], None, - openai::OpenAiProvider::from_env, + openai::OpenAiProvider::default, ) .await } @@ -316,7 +310,7 @@ async fn test_databricks_provider() -> Result<()> { "Databricks", &["DATABRICKS_HOST", "DATABRICKS_TOKEN"], None, - databricks::DatabricksProvider::from_env, + databricks::DatabricksProvider::default, ) .await } @@ -330,7 +324,7 @@ async fn test_databricks_provider_oauth() -> Result<()> { "Databricks OAuth", &["DATABRICKS_HOST"], Some(env_mods), - databricks::DatabricksProvider::from_env, + databricks::DatabricksProvider::default, ) .await } @@ -341,20 +335,14 @@ async fn test_ollama_provider() -> Result<()> { "Ollama", &["OLLAMA_HOST"], None, - ollama::OllamaProvider::from_env, + ollama::OllamaProvider::default, ) .await } #[tokio::test] async fn test_groq_provider() -> Result<()> { - test_provider( - "Groq", - &["GROQ_API_KEY"], - None, - groq::GroqProvider::from_env, - ) - .await + test_provider("Groq", &["GROQ_API_KEY"], None, groq::GroqProvider::default).await } #[tokio::test] @@ -363,7 +351,7 @@ async fn test_anthropic_provider() -> Result<()> { "Anthropic", &["ANTHROPIC_API_KEY"], None, - anthropic::AnthropicProvider::from_env, + anthropic::AnthropicProvider::default, ) .await } @@ -374,7 +362,7 @@ async fn test_openrouter_provider() -> Result<()> { "OpenRouter", &["OPENROUTER_API_KEY"], None, - openrouter::OpenRouterProvider::from_env, + openrouter::OpenRouterProvider::default, ) .await } @@ -385,7 +373,7 @@ async fn test_google_provider() -> Result<()> { "Google", &["GOOGLE_API_KEY"], None, - google::GoogleProvider::from_env, + google::GoogleProvider::default, ) .await } diff --git a/download_cli.sh b/download_cli.sh index 7e224d8eb8..3c3ce82dc6 100755 --- a/download_cli.sh +++ b/download_cli.sh @@ -15,8 +15,9 @@ set -euo pipefail # # Environment variables: # GOOSE_BIN_DIR - Directory to which Goose will be installed (default: $HOME/.local/bin) -# GOOSE_PROVIDER - Optional: provider for goose (passed to "goose configure") -# GOOSE_MODEL - Optional: model for goose (passed to "goose configure") +# GOOSE_PROVIDER - Optional: provider for goose +# GOOSE_MODEL - Optional: model for goose +# ** other provider specific environment variables (eg. DATABRICKS_HOST) ############################################################################## # --- 1) Check for curl --- @@ -98,24 +99,11 @@ if [[ ":$PATH:" != *":$GOOSE_BIN_DIR:"* ]]; then echo "" fi -# --- 7) Auto-configure Goose (Optional) --- -CONFIG_ARGS="" -if [ -n "${GOOSE_PROVIDER:-}" ]; then - CONFIG_ARGS="$CONFIG_ARGS -p $GOOSE_PROVIDER" -fi -if [ -n "${GOOSE_MODEL:-}" ]; then - CONFIG_ARGS="$CONFIG_ARGS -m $GOOSE_MODEL" -fi - -# Print a different message based on whether CONFIG_ARGS is set +# --- 7) Configure Goose (Optional) --- echo "" -if [ -n "$CONFIG_ARGS" ]; then - echo "Configuring Goose with: '$CONFIG_ARGS'" -else - echo "Configuring Goose" -fi +echo "Configuring Goose" echo "" -"$GOOSE_BIN_DIR/$OUT_FILE" configure $CONFIG_ARGS +"$GOOSE_BIN_DIR/$OUT_FILE" configure echo "" echo "Goose installed successfully! Run '$OUT_FILE session' to get started." diff --git a/ui/desktop/src/components/settings/providers/ProviderGrid.tsx b/ui/desktop/src/components/settings/providers/ProviderGrid.tsx index 61722a7ef2..974e61ef2c 100644 --- a/ui/desktop/src/components/settings/providers/ProviderGrid.tsx +++ b/ui/desktop/src/components/settings/providers/ProviderGrid.tsx @@ -130,11 +130,12 @@ export function ProviderGrid({ onSubmit }: ProviderGridProps) { const handleConfigure = async (provider) => { const providerId = provider.id.toLowerCase(); - await initializeSystem(providerId, null); const modelName = getDefaultModel(providerId); const model = createSelectedModel(providerId, modelName); + await initializeSystem(providerId, model.name); + switchModel(model); addRecentModel(model); localStorage.setItem('GOOSE_PROVIDER', providerId);