feat: general purpose config (#701)

Co-authored-by: Salman Mohammed <smohammed@squareup.com>
This commit is contained in:
Bradley Axen
2025-01-23 10:49:18 -08:00
committed by GitHub
parent c53412bd28
commit d230c6eb0f
37 changed files with 1231 additions and 1105 deletions
+160 -175
View File
@@ -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<dyn Error>> {
let config = Config::global();
pub async fn handle_configure(
provided_provider: Option<String>,
provided_model: Option<String>,
) -> Result<(), Box<dyn Error>> {
// 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<String>,
provided_model: Option<String>,
) -> Result<(), Box<dyn Error>> {
// Load existing config or create new one
let mut config = Config::load().unwrap_or_default();
pub async fn configure_provider_dialog() -> Result<(), Box<dyn Error>> {
// 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<String> = 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<String, _> = 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<dyn Error>> {
// 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<dyn Error>> {
}
// 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<dyn Error>> {
.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<dyn Error>> {
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<dyn Error>> {
.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<dyn Error>> {
"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<dyn Error>> {
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<dyn Error>> {
}
}
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<dyn Error>> {
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<dyn Error>> {
})
.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<dyn Error>> {
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<dyn Error>> {
}
}
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<dyn Error>> {
envs: Envs::new(envs),
},
},
);
)?;
cliclack::outro(format!("Added {} extension", style(name).green()))?;
}
_ => unreachable!(),
};
config.save()?;
Ok(())
}
+62 -155
View File
@@ -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<String>,
cli_model: Option<String>,
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<String>, 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<String>,
provider: Option<String>,
model: Option<String>,
agent_version: Option<String>,
resume: bool,
) -> Box<Session<'a>> {
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<String> = 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<dyn Prompt>,
_ => Box::new(RustylinePrompt::new()) as Box<dyn Prompt>,
},
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<String>,
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::<String>()
.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)
}
-203
View File
@@ -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<String, ExtensionEntry>,
}
/// 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<PathBuf> {
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<Self> {
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<ExtensionConfig> {
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());
}
}
+8 -121
View File
@@ -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<String>,
/// Model to use
#[arg(
short,
long,
help = "Model to use (e.g., 'gpt-4', 'llama2')",
long_help = "Specify which model to use."
)]
model: Option<String>,
},
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<String>,
/// 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<String>,
/// 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<String>,
/// 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<String>,
/// Resume a previous session
#[arg(
short,
@@ -134,24 +87,6 @@ enum Command {
)]
input_text: Option<String>,
/// 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<String>,
/// 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<String>,
/// Name for this run session
#[arg(
short,
@@ -162,15 +97,6 @@ enum Command {
)]
name: Option<String>,
/// 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<String>,
/// 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(),
+11 -2
View File
@@ -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
+6 -3
View File
@@ -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);
}
+7 -7
View File
@@ -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<Usage> {
Ok(Usage::new(None, None, None))
}
}
#[test]
+5 -6
View File
@@ -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<HashMap<String, ProviderConfig>> = Lazy::
fn check_key_status(key: &str) -> (bool, Option<String>) {
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::<String>(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),
}
+7 -6
View File
@@ -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
harness = false
+1 -1
View File
@@ -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");
+1 -1
View File
@@ -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.");
+2 -2
View File
@@ -19,8 +19,8 @@ async fn main() -> Result<()> {
// Create providers
let providers: Vec<Box<dyn goose::providers::base::Provider + Send + Sync>> = vec![
Box::new(DatabricksProvider::from_env()?),
Box::new(OpenAiProvider::from_env()?),
Box::new(DatabricksProvider::default()),
Box::new(OpenAiProvider::default()),
];
for provider in providers {
+8 -8
View File
@@ -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<Usage> {
Ok(Usage::new(None, None, None))
}
}
struct MockClient {}
+16 -16
View File
@@ -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<Usage> {
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(())
}
}
+462
View File
@@ -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<serde_json::Error> for ConfigError {
fn from(err: serde_json::Error) -> Self {
ConfigError::DeserializeError(err.to_string())
}
}
impl From<serde_yaml::Error> for ConfigError {
fn from(err: serde_yaml::Error) -> Self {
ConfigError::DeserializeError(err.to_string())
}
}
impl From<keyring::Error> 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<Config> = 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<P: AsRef<Path>>(config_path: P, service: &str) -> Result<Self, ConfigError> {
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<HashMap<String, Value>, 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<HashMap<String, Value>, ConfigError> {
let entry = Entry::new(&self.keyring_service, KEYRING_USERNAME)?;
match entry.get_password() {
Ok(content) => {
let values: HashMap<String, Value> = 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<T: for<'de> Deserialize<'de>>(&self, key: &str) -> Result<T, ConfigError> {
// 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<T: for<'de> Deserialize<'de>>(&self, key: &str) -> Result<T, ConfigError> {
// 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<String, ConfigError> = 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<String, ConfigError> = 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<String, ConfigError> = config.get_secret("key1");
let value2: String = config.get_secret("key2")?;
assert!(matches!(result1, Err(ConfigError::NotFound(_))));
assert_eq!(value2, "secret2");
cleanup_keyring()?;
Ok(())
}
}
+106
View File
@@ -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<Option<ExtensionConfig>> {
let config = Config::global();
// Try to get the extension entry
let extensions: HashMap<String, ExtensionEntry> = 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<String, ExtensionEntry> =
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<String, ExtensionEntry> =
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<String, ExtensionEntry> =
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<HashMap<String, ExtensionEntry>> {
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<bool> {
let config = Config::global();
let extensions: HashMap<String, ExtensionEntry> =
config.get("extensions").unwrap_or_else(|_| HashMap::new());
Ok(extensions.get(name).map(|e| e.enabled).unwrap_or(false))
}
}
+6
View File
@@ -0,0 +1,6 @@
mod base;
mod extensions;
pub use crate::agents::ExtensionConfig;
pub use base::{Config, ConfigError};
pub use extensions::{ExtensionEntry, ExtensionManager};
-192
View File
@@ -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<keyring::Error> for KeyManagerError {
fn from(err: keyring::Error) -> Self {
KeyManagerError::KeyringAccess(err.to_string())
}
}
impl From<env::VarError> 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<String> {
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<String> {
match env::var(key_name) {
Ok(value) => Some(value),
Err(_) => None,
}
}
pub fn key_value_in_key_chain(key_name: &str) -> Option<String> {
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();
}
}
+2 -1
View File
@@ -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;
+36 -16
View File
@@ -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<Self> {
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<Self> {
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<Usage> {
get_usage(data)
}
}
+69 -5
View File
@@ -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<ConfigKey>,
}
impl ProviderMetadata {
pub fn new(
name: &str,
display_name: &str,
description: &str,
default_model: &str,
config_keys: Vec<ConfigKey>,
) -> 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<String>,
}
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<Usage>;
/// Get the model config from the provider
fn get_model_config(&self) -> ModelConfig;
}
#[cfg(test)]
+34 -24
View File
@@ -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<Self> {
// 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<Self> {
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<Usage> {
get_usage(data)
}
}
+28 -10
View File
@@ -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<Box<dyn Provider + Send + Sync>> {
pub fn providers() -> Vec<ProviderMetadata> {
vec![
AnthropicProvider::metadata(),
DatabricksProvider::metadata(),
GoogleProvider::metadata(),
GroqProvider::metadata(),
OllamaProvider::metadata(),
OpenAiProvider::metadata(),
OpenRouterProvider::metadata(),
]
}
pub fn create(name: &str, model: ModelConfig) -> Result<Box<dyn Provider + Send + Sync>> {
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)),
}
}
@@ -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;
+1 -1
View File
@@ -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;
+1 -1
View File
@@ -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,
+32 -15
View File
@@ -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<Self> {
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<Self> {
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<Usage> {
get_usage(data)
}
}
+32 -15
View File
@@ -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<Self> {
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<Self> {
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<Usage> {
get_usage(data)
}
}
+2 -3
View File
@@ -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};
+32 -16
View File
@@ -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<Self> {
// 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<Self> {
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<Usage> {
get_usage(data)
}
}
+32 -35
View File
@@ -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<Self> {
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<Self> {
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<Usage> {
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(())
}
}
+36 -16
View File
@@ -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<Self> {
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<Self> {
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<Usage> {
get_usage(data)
}
}
+1 -1
View File
@@ -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;
+16 -28
View File
@@ -219,7 +219,7 @@ async fn test_provider<F, T>(
provider_fn: F,
) -> Result<()>
where
F: FnOnce() -> Result<T>,
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
}
+6 -18
View File
@@ -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."
@@ -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);