mirror of
https://github.com/block/goose.git
synced 2026-07-17 12:56:20 +02:00
feat: general purpose config (#701)
Co-authored-by: Salman Mohammed <smohammed@squareup.com>
This commit is contained in:
@@ -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(())
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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(),
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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),
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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");
|
||||
|
||||
@@ -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.");
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {}
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
mod base;
|
||||
mod extensions;
|
||||
|
||||
pub use crate::agents::ExtensionConfig;
|
||||
pub use base::{Config, ConfigError};
|
||||
pub use extensions::{ExtensionEntry, ExtensionManager};
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,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,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,
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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};
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user