tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

Declarative providers (#5084)

Browse Source 
Co-authored-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: Michael Neale <michael.neale@gmail.com>
This commit is contained in:
Douwe Osinga
2025-10-15 09:48:14 -04:00
committed by GitHub
parent 925a042cb1
commit 9251da4314
28 changed files with 1219 additions and 802 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/config/custom_providers.rs
-212
View File
@@ -1,212 +0,0 @@
use crate::config::paths::Paths;
use crate::config::Config;
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 serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;
pub fn custom_providers_dir() -> std::path::PathBuf {
Paths::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(())
}
crates/goose/src/config/declarative_providers.rs
+317
View File
@@ -0,0 +1,317 @@
use crate::config::paths::Paths;
use crate::config::Config;
use crate::providers::anthropic::AnthropicProvider;
use crate::providers::base::{ModelInfo, ProviderType};
use crate::providers::ollama::OllamaProvider;
use crate::providers::openai::OpenAiProvider;
use anyhow::Result;
use include_dir::{include_dir, Dir};
use once_cell::sync::Lazy;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;
use std::sync::Mutex;
use utoipa::ToSchema;
static FIXED_PROVIDERS: Dir = include_dir!("$CARGO_MANIFEST_DIR/src/providers/declarative");
pub fn custom_providers_dir() -> std::path::PathBuf {
Paths::config_dir().join("custom_providers")
}
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
#[serde(rename_all = "lowercase")]
pub enum ProviderEngine {
OpenAI,
Ollama,
Anthropic,
}
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub struct DeclarativeProviderConfig {
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 DeclarativeProviderConfig {
pub fn id(&self) -> &str {
&self.name
}
pub fn display_name(&self) -> &str {
&self.display_name
}
pub fn models(&self) -> &[ModelInfo] {
&self.models
}
}
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub struct LoadedProvider {
pub config: DeclarativeProviderConfig,
pub is_editable: bool,
}
static ID_GENERATION_LOCK: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));
pub fn generate_id(display_name: &str) -> String {
let _guard = ID_GENERATION_LOCK.lock().unwrap();
let normalized = display_name.to_lowercase().replace(' ', "_");
let base_id = format!("custom_{}", normalized);
let custom_dir = custom_providers_dir();
let mut candidate_id = base_id.clone();
let mut counter = 1;
while custom_dir.join(format!("{}.json", candidate_id)).exists() {
candidate_id = format!("{}_{}", base_id, counter);
counter += 1;
}
candidate_id
}
pub fn generate_api_key_name(id: &str) -> String {
format!("{}_API_KEY", id.to_uppercase())
}
pub fn create_custom_provider(
engine: &str,
display_name: String,
api_url: String,
api_key: String,
models: Vec<String>,
supports_streaming: Option<bool>,
) -> Result<DeclarativeProviderConfig> {
let id = generate_id(&display_name);
let api_key_name = 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 = DeclarativeProviderConfig {
name: id.clone(),
engine: match engine {
"openai_compatible" => ProviderEngine::OpenAI,
"anthropic_compatible" => ProviderEngine::Anthropic,
"ollama_compatible" => ProviderEngine::Ollama,
_ => return Err(anyhow::anyhow!("Invalid provider type: {}", engine)),
},
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,
};
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 update_custom_provider(
id: &str,
provider_type: &str,
display_name: String,
api_url: String,
api_key: String,
models: Vec<String>,
supports_streaming: Option<bool>,
) -> Result<()> {
let loaded_provider = load_provider(id)?;
let existing_config = loaded_provider.config;
let editable = loaded_provider.is_editable;
let config = Config::global();
if !api_key.is_empty() {
config.set_secret(
&existing_config.api_key_env,
serde_json::Value::String(api_key),
)?;
}
if editable {
let model_infos: Vec<ModelInfo> = models
.into_iter()
.map(|name| ModelInfo::new(name, 128000))
.collect();
let updated_config = DeclarativeProviderConfig {
name: id.to_string(),
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,
description: existing_config.description,
api_key_env: existing_config.api_key_env,
base_url: api_url,
models: model_infos,
headers: existing_config.headers,
timeout_seconds: existing_config.timeout_seconds,
supports_streaming,
};
let file_path = custom_providers_dir().join(format!("{}.json", id));
let json_content = serde_json::to_string_pretty(&updated_config)?;
std::fs::write(file_path, json_content)?;
}
Ok(())
}
pub fn remove_custom_provider(id: &str) -> Result<()> {
let config = Config::global();
let api_key_name = 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_provider(id: &str) -> Result<LoadedProvider> {
let custom_file_path = custom_providers_dir().join(format!("{}.json", id));
if custom_file_path.exists() {
let content = std::fs::read_to_string(&custom_file_path)?;
let config: DeclarativeProviderConfig = serde_json::from_str(&content)?;
return Ok(LoadedProvider {
config,
is_editable: true,
});
}
for file in FIXED_PROVIDERS.files() {
if file.path().extension().and_then(|s| s.to_str()) != Some("json") {
continue;
}
let content = file
.contents_utf8()
.ok_or_else(|| anyhow::anyhow!("Failed to read file as UTF-8: {:?}", file.path()))?;
let config: DeclarativeProviderConfig = serde_json::from_str(content)?;
if config.name == id {
return Ok(LoadedProvider {
config,
is_editable: false,
});
}
}
Err(anyhow::anyhow!("Provider not found: {}", id))
}
pub fn load_custom_providers(dir: &Path) -> Result<Vec<DeclarativeProviderConfig>> {
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()
}
fn load_fixed_providers() -> Result<Vec<DeclarativeProviderConfig>> {
let mut res = Vec::new();
for file in FIXED_PROVIDERS.files() {
if file.path().extension().and_then(|s| s.to_str()) != Some("json") {
continue;
}
let content = file
.contents_utf8()
.ok_or_else(|| anyhow::anyhow!("Failed to read file as UTF-8: {:?}", file.path()))?;
let config: DeclarativeProviderConfig = serde_json::from_str(content)?;
res.push(config)
}
Ok(res)
}
pub fn register_declarative_providers(
registry: &mut crate::providers::provider_registry::ProviderRegistry,
) -> Result<()> {
let dir = custom_providers_dir();
let custom_providers = load_custom_providers(&dir)?;
let fixed_providers = load_fixed_providers()?;
for config in fixed_providers {
register_declarative_provider(registry, config, ProviderType::Declarative);
}
for config in custom_providers {
register_declarative_provider(registry, config, ProviderType::Custom);
}
Ok(())
}
pub fn register_declarative_provider(
registry: &mut crate::providers::provider_registry::ProviderRegistry,
config: DeclarativeProviderConfig,
provider_type: ProviderType,
) {
let config_clone = config.clone();
match config.engine {
ProviderEngine::OpenAI => {
registry.register_with_name::<OpenAiProvider, _>(
&config,
provider_type,
move |model| OpenAiProvider::from_custom_config(model, config_clone.clone()),
);
}
ProviderEngine::Ollama => {
registry.register_with_name::<OllamaProvider, _>(
&config,
provider_type,
move |model| OllamaProvider::from_custom_config(model, config_clone.clone()),
);
}
ProviderEngine::Anthropic => {
registry.register_with_name::<AnthropicProvider, _>(
&config,
provider_type,
move |model| AnthropicProvider::from_custom_config(model, config_clone.clone()),
);
}
}
}
crates/goose/src/config/mod.rs
+2 -2
View File
@@ -1,5 +1,5 @@
pub mod base;
pub mod custom_providers;
pub mod declarative_providers;
mod experiments;
pub mod extensions;
pub mod paths;
@@ -9,7 +9,7 @@ pub mod signup_tetrate;
pub use crate::agents::ExtensionConfig;
pub use base::{Config, ConfigError};
pub use custom_providers::CustomProviderConfig;
pub use declarative_providers::DeclarativeProviderConfig;
pub use experiments::ExperimentManager;
pub use extensions::{
get_all_extension_names, get_all_extensions, get_enabled_extensions, get_extension_by_name,