use crate::config::{Config, APP_STRATEGY}; use crate::model::ModelConfig; use crate::providers::anthropic::AnthropicProvider; use crate::providers::base::ModelInfo; use crate::providers::ollama::OllamaProvider; use crate::providers::openai::OpenAiProvider; use anyhow::Result; use etcetera::{choose_app_strategy, AppStrategy}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::path::Path; pub fn custom_providers_dir() -> std::path::PathBuf { choose_app_strategy(APP_STRATEGY.clone()) .expect("goose requires a home dir") .config_dir() .join("custom_providers") } #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum ProviderEngine { OpenAI, Ollama, Anthropic, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct CustomProviderConfig { pub name: String, pub engine: ProviderEngine, pub display_name: String, pub description: Option, pub api_key_env: String, pub base_url: String, pub models: Vec, pub headers: Option>, pub timeout_seconds: Option, pub supports_streaming: Option, } impl CustomProviderConfig { pub fn id(&self) -> &str { &self.name } pub fn display_name(&self) -> &str { &self.display_name } pub fn models(&self) -> &[ModelInfo] { &self.models } pub fn generate_id(display_name: &str) -> String { format!("custom_{}", display_name.to_lowercase().replace(' ', "_")) } pub fn generate_api_key_name(id: &str) -> String { format!("{}_API_KEY", id.to_uppercase()) } pub fn create_and_save( provider_type: &str, display_name: String, api_url: String, api_key: String, models: Vec, supports_streaming: Option, ) -> Result { let id = Self::generate_id(&display_name); let api_key_name = Self::generate_api_key_name(&id); let config = Config::global(); config.set_secret(&api_key_name, serde_json::Value::String(api_key))?; let model_infos: Vec = models .into_iter() .map(|name| ModelInfo::new(name, 128000)) .collect(); let provider_config = CustomProviderConfig { name: id.clone(), engine: match provider_type { "openai_compatible" => ProviderEngine::OpenAI, "anthropic_compatible" => ProviderEngine::Anthropic, "ollama_compatible" => ProviderEngine::Ollama, _ => return Err(anyhow::anyhow!("Invalid provider type: {}", provider_type)), }, display_name: display_name.clone(), description: Some(format!("Custom {} provider", display_name)), api_key_env: api_key_name, base_url: api_url, models: model_infos, headers: None, timeout_seconds: None, supports_streaming, }; // save to JSON file let custom_providers_dir = custom_providers_dir(); std::fs::create_dir_all(&custom_providers_dir)?; let json_content = serde_json::to_string_pretty(&provider_config)?; let file_path = custom_providers_dir.join(format!("{}.json", id)); std::fs::write(file_path, json_content)?; Ok(provider_config) } pub fn remove(id: &str) -> Result<()> { let config = Config::global(); let api_key_name = Self::generate_api_key_name(id); let _ = config.delete_secret(&api_key_name); let custom_providers_dir = custom_providers_dir(); let file_path = custom_providers_dir.join(format!("{}.json", id)); if file_path.exists() { std::fs::remove_file(file_path)?; } Ok(()) } } pub fn load_custom_providers(dir: &Path) -> Result> { if !dir.exists() { return Ok(Vec::new()); } std::fs::read_dir(dir)? .filter_map(|entry| { let path = entry.ok()?.path(); (path.extension()? == "json").then_some(path) }) .map(|path| { let content = std::fs::read_to_string(&path)?; serde_json::from_str(&content) .map_err(|e| anyhow::anyhow!("Failed to parse {}: {}", path.display(), e)) }) .collect() } pub fn register_custom_providers( registry: &mut crate::providers::provider_registry::ProviderRegistry, dir: &Path, ) -> Result<()> { let configs = load_custom_providers(dir)?; for config in configs { let config_clone = config.clone(); let description = config .description .clone() .unwrap_or_else(|| format!("Custom {} provider", config.display_name)); let default_model = config .models .first() .map(|m| m.name.clone()) .unwrap_or_default(); let known_models: Vec = config .models .iter() .map(|m| ModelInfo { name: m.name.clone(), context_limit: m.context_limit, input_token_cost: m.input_token_cost, output_token_cost: m.output_token_cost, currency: m.currency.clone(), supports_cache_control: Some(m.supports_cache_control.unwrap_or(false)), }) .collect(); match config.engine { ProviderEngine::OpenAI => { registry.register_with_name::( config.name.clone(), config.display_name.clone(), description, default_model, known_models, move |model: ModelConfig| { OpenAiProvider::from_custom_config(model, config_clone.clone()) }, ); } ProviderEngine::Ollama => { registry.register_with_name::( config.name.clone(), config.display_name.clone(), description, default_model, known_models, move |model: ModelConfig| { OllamaProvider::from_custom_config(model, config_clone.clone()) }, ); } ProviderEngine::Anthropic => { registry.register_with_name::( config.name.clone(), config.display_name.clone(), description, default_model, known_models, move |model: ModelConfig| { AnthropicProvider::from_custom_config(model, config_clone.clone()) }, ); } } } Ok(()) }