Custom providers update (#4099)

Co-authored-by: developerayo <shodipovi@gmail.com> Co-authored-by: Douwe Osinga <douwe@squareup.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Zane Staggs <zane@squareup.com>
2025-08-19 18:03:10 -04:00
parent 1d93a59f31
commit 942ef5b0a3
31 changed files with 1704 additions and 567 deletions
@@ -0,0 +1,215 @@
+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<String>,
+    pub api_key_env: String,
+    pub base_url: String,
+    pub models: Vec<ModelInfo>,
+    pub headers: Option<HashMap<String, String>>,
+    pub timeout_seconds: Option<u64>,
+    pub supports_streaming: Option<bool>,
+}
+
+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<String>,
+        supports_streaming: Option<bool>,
+    ) -> Result<Self> {
+        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<ModelInfo> = 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<Vec<CustomProviderConfig>> {
+    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<ModelInfo> = 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::<OpenAiProvider, _>(
+                    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::<OllamaProvider, _>(
+                    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::<AnthropicProvider, _>(
+                    config.name.clone(),
+                    config.display_name.clone(),
+                    description,
+                    default_model,
+                    known_models,
+                    move |model: ModelConfig| {
+                        AnthropicProvider::from_custom_config(model, config_clone.clone())
+                    },
+                );
+            }
+        }
+    }
+    Ok(())
+}
@@ -1,4 +1,5 @@
 pub mod base;
+pub mod custom_providers;
 mod experiments;
 pub mod extensions;
 pub mod permission;
@@ -6,6 +7,7 @@ pub mod signup_openrouter;

 pub use crate::agents::ExtensionConfig;
 pub use base::{Config, ConfigError, APP_STRATEGY};
+pub use custom_providers::CustomProviderConfig;
 pub use experiments::ExperimentManager;
 pub use extensions::{ExtensionConfigManager, ExtensionEntry};
 pub use permission::PermissionManager;