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

Custom providers update (#4099)

Browse Source 
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>
This commit is contained in:
Douwe Osinga
2025-08-19 18:03:10 -04:00
committed by GitHub
parent 1d93a59f31
commit 942ef5b0a3
31 changed files with 1704 additions and 567 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/Cargo.toml
+1
View File
@@ -47,6 +47,7 @@ shlex = "1.3.0"
async-trait = "0.1.86"
base64 = "0.22.1"
regex = "1.11.1"
uuid = { version = "1.11", features = ["v4"] }
nix = { version = "0.30.1", features = ["process", "signal"] }
tar = "0.4"
# Web server dependencies
crates/goose-cli/src/commands/configure.rs
+139 -12
View File
@@ -1,3 +1,4 @@
use crate::recipes::github_recipe::GOOSE_RECIPE_GITHUB_REPO_CONFIG_KEY;
use cliclack::spinner;
use console::style;
use goose::agents::extension::ToolInfo;
@@ -7,6 +8,7 @@ use goose::agents::platform_tools::{
};
use goose::agents::Agent;
use goose::agents::{extension::Envs, ExtensionConfig};
use goose::config::custom_providers::CustomProviderConfig;
use goose::config::extensions::name_to_key;
use goose::config::permission::PermissionLevel;
use goose::config::{
@@ -14,6 +16,7 @@ use goose::config::{
PermissionManager,
};
use goose::conversation::message::Message;
use goose::model::ModelConfig;
use goose::providers::{create, providers};
use rmcp::model::{Tool, ToolAnnotations};
use rmcp::object;
@@ -21,8 +24,6 @@ use serde_json::Value;
use std::collections::HashMap;
use std::error::Error;
use crate::recipes::github_recipe::GOOSE_RECIPE_GITHUB_REPO_CONFIG_KEY;
// useful for light themes where there is no dicernible colour contrast between
// cursor-selected and cursor-unselected items.
const MULTISELECT_VISIBILITY_HINT: &str = "<";
@@ -221,6 +222,11 @@ pub async fn handle_configure() -> Result<(), Box<dyn Error>> {
"Configure Providers",
"Change provider or update credentials",
)
.item(
"custom_providers",
"Custom Providers",
"Add custom provider with compatible API",
)
.item("add", "Add Extension", "Connect to a new extension")
.item(
"toggle",
@@ -241,6 +247,7 @@ pub async fn handle_configure() -> Result<(), Box<dyn Error>> {
"remove" => remove_extension_dialog(),
"settings" => configure_settings_dialog().await.and(Ok(())),
"providers" => configure_provider_dialog().await.and(Ok(())),
"custom_providers" => configure_custom_provider_dialog(),
_ => unreachable!(),
}
}
@@ -250,10 +257,7 @@ pub async fn handle_configure() -> Result<(), Box<dyn Error>> {
async fn handle_oauth_configuration(
provider_name: &str,
key_name: &str,
) -> Result<(), Box<dyn std::error::Error>> {
use goose::model::ModelConfig;
use goose::providers::create;
) -> Result<(), Box<dyn Error>> {
let _ = cliclack::log::info(format!(
"Configuring {} using OAuth device code flow...",
key_name
@@ -279,8 +283,7 @@ async fn handle_oauth_configuration(
}
}
/// Interactive model search that truncates the list to improve UX
fn interactive_model_search(models: &[String]) -> Result<String, Box<dyn std::error::Error>> {
fn interactive_model_search(models: &[String]) -> Result<String, Box<dyn Error>> {
const MAX_VISIBLE: usize = 30;
let mut query = String::new();
@@ -553,7 +556,7 @@ pub async fn configure_provider_dialog() -> Result<bool, Box<dyn Error>> {
let spin = spinner();
spin.start("Attempting to fetch supported models...");
let models_res = {
let temp_model_config = goose::model::ModelConfig::new(&provider_meta.default_model)?;
let temp_model_config = ModelConfig::new(&provider_meta.default_model)?;
let temp_provider = create(provider_name, temp_model_config)?;
temp_provider.fetch_supported_models().await
};
@@ -585,7 +588,7 @@ pub async fn configure_provider_dialog() -> Result<bool, Box<dyn Error>> {
.map(|val| val == "1" || val.to_lowercase() == "true")
.unwrap_or(false);
let model_config = goose::model::ModelConfig::new(&model)?
let model_config = ModelConfig::new(&model)?
.with_max_tokens(Some(50))
.with_toolshim(toolshim_enabled)
.with_toolshim_model(std::env::var("GOOSE_TOOLSHIM_OLLAMA_MODEL").ok());
@@ -1429,7 +1432,7 @@ pub async fn configure_tool_permissions_dialog() -> Result<(), Box<dyn Error>> {
let model: String = config
.get_param("GOOSE_MODEL")
.expect("No model configured. Please set model first");
let model_config = goose::model::ModelConfig::new(&model)?;
let model_config = ModelConfig::new(&model)?;
// Create the agent
let agent = Agent::new();
@@ -1569,7 +1572,6 @@ fn configure_recipe_dialog() -> Result<(), Box<dyn Error>> {
recipe_repo_input = recipe_repo_input.default_input(&recipe_repo);
}
let input_value: String = recipe_repo_input.interact()?;
// if input is blank, it clears the recipe github repo settings in the config file
if input_value.clone().trim().is_empty() {
config.delete(key_name)?;
} else {
@@ -1767,3 +1769,128 @@ pub async fn handle_openrouter_auth() -> Result<(), Box<dyn Error>> {
Ok(())
}
fn add_provider() -> Result<(), Box<dyn Error>> {
let provider_type = cliclack::select("What type of API is this?")
.item(
"openai_compatible",
"OpenAI Compatible",
"Uses OpenAI API format",
)
.item(
"anthropic_compatible",
"Anthropic Compatible",
"Uses Anthropic API format",
)
.item(
"ollama_compatible",
"Ollama Compatible",
"Uses Ollama API format",
)
.interact()?;
let display_name: String = cliclack::input("What should we call this provider?")
.placeholder("Your Provider Name")
.validate(|input: &String| {
if input.is_empty() {
Err("Please enter a name")
} else {
Ok(())
}
})
.interact()?;
let api_url: String = cliclack::input("Provider API URL:")
.placeholder("https://api.example.com/v1/messages")
.validate(|input: &String| {
if !input.starts_with("http://") && !input.starts_with("https://") {
Err("URL must start with either http:// or https://")
} else {
Ok(())
}
})
.interact()?;
let api_key: String = cliclack::password("API key:").mask('▪').interact()?;
let models_input: String = cliclack::input("Available models (seperate with commas):")
.placeholder("model-a, model-b, model-c")
.validate(|input: &String| {
if input.trim().is_empty() {
Err("Please enter at least one model name")
} else {
Ok(())
}
})
.interact()?;
let models: Vec<String> = models_input
.split(',')
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.collect();
let supports_streaming = cliclack::confirm("Does this provider support streaming responses?")
.initial_value(true)
.interact()?;
CustomProviderConfig::create_and_save(
provider_type,
display_name.clone(),
api_url,
api_key,
models,
Some(supports_streaming),
)?;
cliclack::outro(format!("Custom provider added: {}", display_name))?;
Ok(())
}
fn remove_provider() -> Result<(), Box<dyn Error>> {
let custom_providers_dir = goose::config::custom_providers::custom_providers_dir();
let custom_providers = if custom_providers_dir.exists() {
goose::config::custom_providers::load_custom_providers(&custom_providers_dir)?
} else {
Vec::new()
};
if custom_providers.is_empty() {
cliclack::outro("No custom providers added just yet.")?;
return Ok(());
}
let provider_items: Vec<_> = custom_providers
.iter()
.map(|p| (p.name.as_str(), p.display_name.as_str(), "Custom provider"))
.collect();
let selected_id = cliclack::select("Which custom provider would you like to remove?")
.items(&provider_items)
.interact()?;
CustomProviderConfig::remove(selected_id)?;
cliclack::outro(format!("Removed custom provider: {}", selected_id))?;
Ok(())
}
pub fn configure_custom_provider_dialog() -> Result<(), Box<dyn Error>> {
let action = cliclack::select("What would you like to do?")
.item(
"add",
"Add A Custom Provider",
"Add a new OpenAI/Anthropic/Ollama compatible Provider",
)
.item(
"remove",
"Remove Custom Provider",
"Remove an existing custom provider",
)
.interact()?;
match action {
"add" => add_provider(),
"remove" => remove_provider(),
_ => unreachable!(),
}
}
crates/goose-server/Cargo.toml
+1
View File
@@ -40,6 +40,7 @@ serde_yaml = "0.9.34"
utoipa = { version = "4.1", features = ["axum_extras", "chrono"] }
reqwest = { version = "0.12.9", features = ["json", "rustls-tls", "blocking", "multipart"], default-features = false }
tokio-util = "0.7.15"
uuid = { version = "1.11", features = ["v4"] }
[[bin]]
name = "goosed"
crates/goose-server/src/openapi.rs
+3
View File
@@ -367,6 +367,8 @@ impl<'__s> ToSchema<'__s> for AnnotatedSchema {
super::routes::config_management::read_all_config,
super::routes::config_management::providers,
super::routes::config_management::upsert_permissions,
super::routes::config_management::create_custom_provider,
super::routes::config_management::remove_custom_provider,
super::routes::agent::get_tools,
super::routes::agent::add_sub_recipes,
super::routes::agent::extend_prompt,
@@ -402,6 +404,7 @@ impl<'__s> ToSchema<'__s> for AnnotatedSchema {
super::routes::config_management::ExtensionQuery,
super::routes::config_management::ToolPermission,
super::routes::config_management::UpsertPermissionsQuery,
super::routes::config_management::CreateCustomProviderRequest,
super::routes::reply::PermissionConfirmationRequest,
super::routes::context::ContextManageRequest,
super::routes::context::ContextManageResponse,
crates/goose-server/src/routes/agent.rs
+2 -2
View File
@@ -192,7 +192,7 @@ async fn update_agent_provider(
let agent = state
.get_agent()
.await
.map_err(|_| StatusCode::PRECONDITION_FAILED)?;
.map_err(|_e| StatusCode::PRECONDITION_FAILED)?;
let config = Config::global();
let model = match payload
@@ -210,7 +210,7 @@ async fn update_agent_provider(
agent
.update_provider(new_provider)
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
.map_err(|_e| StatusCode::INTERNAL_SERVER_ERROR)?;
Ok(StatusCode::OK)
}
crates/goose-server/src/routes/config_management.rs
+134 -5
View File
@@ -8,7 +8,6 @@ use axum::{
};
use etcetera::{choose_app_strategy, AppStrategy};
use goose::config::APP_STRATEGY;
use goose::config::{extensions::name_to_key, PermissionManager};
use goose::config::{Config, ConfigError};
use goose::config::{ExtensionConfigManager, ExtensionEntry};
use goose::model::ModelConfig;
@@ -78,6 +77,16 @@ pub struct UpsertPermissionsQuery {
pub tool_permissions: Vec<ToolPermission>,
}
#[derive(Deserialize, ToSchema)]
pub struct CreateCustomProviderRequest {
pub provider_type: String,
pub display_name: String,
pub api_url: String,
pub api_key: String,
pub models: Vec<String>,
pub supports_streaming: Option<bool>,
}
#[utoipa::path(
post,
path = "/config/upsert",
@@ -227,7 +236,7 @@ pub async fn add_extension(
let extensions =
ExtensionConfigManager::get_all().map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
let key = name_to_key(&extension_query.name);
let key = goose::config::extensions::name_to_key(&extension_query.name);
let is_update = extensions.iter().any(|e| e.config.key() == key);
@@ -262,7 +271,7 @@ pub async fn remove_extension(
) -> Result<Json<String>, StatusCode> {
verify_secret_key(&headers, &state)?;
let key = name_to_key(&name);
let key = goose::config::extensions::name_to_key(&name);
match ExtensionConfigManager::remove(&key) {
Ok(_) => Ok(Json(format!("Removed extension {}", name))),
Err(_) => Err(StatusCode::NOT_FOUND),
@@ -304,7 +313,62 @@ pub async fn providers(
) -> Result<Json<Vec<ProviderDetails>>, StatusCode> {
verify_secret_key(&headers, &state)?;
let providers_metadata = get_providers();
let mut providers_metadata = get_providers();
let custom_providers_dir = goose::config::custom_providers::custom_providers_dir();
if custom_providers_dir.exists() {
if let Ok(entries) = std::fs::read_dir(&custom_providers_dir) {
for entry in entries.flatten() {
if let Some(extension) = entry.path().extension() {
if extension == "json" {
if let Ok(content) = std::fs::read_to_string(entry.path()) {
if let Ok(custom_provider) = serde_json::from_str::<
goose::config::custom_providers::CustomProviderConfig,
>(&content)
{
// CustomProviderConfig => ProviderMetadata
let default_model = custom_provider
.models
.first()
.map(|m| m.name.clone())
.unwrap_or_default();
let metadata = goose::providers::base::ProviderMetadata {
name: custom_provider.name.clone(),
display_name: custom_provider.display_name.clone(),
description: custom_provider
.description
.clone()
.unwrap_or_else(|| {
format!("{} (custom)", custom_provider.display_name)
}),
default_model,
known_models: custom_provider.models.clone(),
model_doc_link: "Custom provider".to_string(),
config_keys: vec![
goose::providers::base::ConfigKey::new(
&custom_provider.api_key_env,
true,
true,
None,
),
goose::providers::base::ConfigKey::new(
"CUSTOM_PROVIDER_BASE_URL",
true,
false,
Some(&custom_provider.base_url),
),
],
};
providers_metadata.push(metadata);
}
}
}
}
}
}
}
let providers_response: Vec<ProviderDetails> = providers_metadata
.into_iter()
@@ -491,7 +555,7 @@ pub async fn upsert_permissions(
) -> Result<Json<String>, StatusCode> {
verify_secret_key(&headers, &state)?;
let mut permission_manager = PermissionManager::default();
let mut permission_manager = goose::config::PermissionManager::default();
for tool_permission in &query.tool_permissions {
permission_manager.update_user_permission(
@@ -637,6 +701,66 @@ pub async fn get_current_model(
})))
}
#[utoipa::path(
post,
path = "/config/custom-providers",
request_body = CreateCustomProviderRequest,
responses(
(status = 200, description = "Custom provider created successfully", body = String),
(status = 400, description = "Invalid request"),
(status = 500, description = "Internal server error")
)
)]
pub async fn create_custom_provider(
State(state): State<Arc<AppState>>,
headers: HeaderMap,
Json(request): Json<CreateCustomProviderRequest>,
) -> Result<Json<String>, StatusCode> {
verify_secret_key(&headers, &state)?;
let config = goose::config::custom_providers::CustomProviderConfig::create_and_save(
&request.provider_type,
request.display_name,
request.api_url,
request.api_key,
request.models,
request.supports_streaming,
)
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
if let Err(e) = goose::providers::refresh_custom_providers() {
tracing::warn!("Failed to refresh custom providers after creation: {}", e);
}
Ok(Json(format!("Custom provider added - ID: {}", config.id())))
}
#[utoipa::path(
delete,
path = "/config/custom-providers/{id}",
responses(
(status = 200, description = "Custom provider removed successfully", body = String),
(status = 404, description = "Provider not found"),
(status = 500, description = "Internal server error")
)
)]
pub async fn remove_custom_provider(
State(state): State<Arc<AppState>>,
headers: HeaderMap,
axum::extract::Path(id): axum::extract::Path<String>,
) -> Result<Json<String>, StatusCode> {
verify_secret_key(&headers, &state)?;
goose::config::custom_providers::CustomProviderConfig::remove(&id)
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
if let Err(e) = goose::providers::refresh_custom_providers() {
tracing::warn!("Failed to refresh custom providers after deletion: {}", e);
}
Ok(Json(format!("Removed custom provider: {}", id)))
}
pub fn routes(state: Arc<AppState>) -> Router {
Router::new()
.route("/config", get(read_all_config))
@@ -654,6 +778,11 @@ pub fn routes(state: Arc<AppState>) -> Router {
.route("/config/validate", get(validate_config))
.route("/config/permissions", post(upsert_permissions))
.route("/config/current-model", get(get_current_model))
.route("/config/custom-providers", post(create_custom_provider))
.route(
"/config/custom-providers/{id}",
delete(remove_custom_provider),
)
.with_state(state)
}
crates/goose/src/config/custom_providers.rs
+215
View File
@@ -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(())
}
crates/goose/src/config/mod.rs
+2
View File
@@ -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;
crates/goose/src/providers/anthropic.rs
+29 -2
View File
@@ -15,6 +15,7 @@ use super::formats::anthropic::{
create_request, get_usage, response_to_message, response_to_streaming_message,
};
use super::utils::{emit_debug_trace, get_model, map_http_error_to_provider_error};
use crate::config::custom_providers::CustomProviderConfig;
use crate::conversation::message::Message;
use crate::impl_provider_default;
use crate::model::ModelConfig;
@@ -42,6 +43,7 @@ pub struct AnthropicProvider {
#[serde(skip)]
api_client: ApiClient,
model: ModelConfig,
supports_streaming: bool,
}
impl_provider_default!(AnthropicProvider);
@@ -62,7 +64,32 @@ impl AnthropicProvider {
let api_client =
ApiClient::new(host, auth)?.with_header("anthropic-version", ANTHROPIC_API_VERSION)?;
Ok(Self { api_client, model })
Ok(Self {
api_client,
model,
supports_streaming: true,
})
}
pub fn from_custom_config(model: ModelConfig, config: CustomProviderConfig) -> Result<Self> {
let global_config = crate::config::Config::global();
let api_key: String = global_config
.get_secret(&config.api_key_env)
.map_err(|_| anyhow::anyhow!("Missing API key: {}", config.api_key_env))?;
let auth = AuthMethod::ApiKey {
header_name: "x-api-key".to_string(),
key: api_key,
};
let api_client = ApiClient::new(config.base_url, auth)?
.with_header("anthropic-version", ANTHROPIC_API_VERSION)?;
Ok(Self {
api_client,
model,
supports_streaming: config.supports_streaming.unwrap_or(true),
})
}
fn get_conditional_headers(&self) -> Vec<(&str, &str)> {
@@ -260,6 +287,6 @@ impl Provider for AnthropicProvider {
}
fn supports_streaming(&self) -> bool {
true
self.supports_streaming
}
}
crates/goose/src/providers/factory.rs
+160 -249
View File
@@ -1,4 +1,4 @@
use std::sync::Arc;
use std::sync::{Arc, RwLock};
use super::{
anthropic::AnthropicProvider,
@@ -17,66 +17,87 @@ use super::{
ollama::OllamaProvider,
openai::OpenAiProvider,
openrouter::OpenRouterProvider,
provider_registry::ProviderRegistry,
sagemaker_tgi::SageMakerTgiProvider,
snowflake::SnowflakeProvider,
venice::VeniceProvider,
xai::XaiProvider,
};
use crate::config::custom_providers::{custom_providers_dir, register_custom_providers};
use crate::model::ModelConfig;
use anyhow::Result;
use once_cell::sync::Lazy;
#[cfg(test)]
use super::errors::ProviderError;
#[cfg(test)]
use rmcp::model::Tool;
fn default_lead_turns() -> usize {
3
}
fn default_failure_threshold() -> usize {
2
}
fn default_fallback_turns() -> usize {
2
const DEFAULT_LEAD_TURNS: usize = 3;
const DEFAULT_FAILURE_THRESHOLD: usize = 2;
const DEFAULT_FALLBACK_TURNS: usize = 2;
static REGISTRY: Lazy<RwLock<ProviderRegistry>> = Lazy::new(|| {
let registry = ProviderRegistry::new().with_providers(|registry| {
registry.register::<AnthropicProvider, _>(AnthropicProvider::from_env);
registry.register::<AzureProvider, _>(AzureProvider::from_env);
registry.register::<BedrockProvider, _>(BedrockProvider::from_env);
registry.register::<ClaudeCodeProvider, _>(ClaudeCodeProvider::from_env);
registry.register::<CursorAgentProvider, _>(CursorAgentProvider::from_env);
registry.register::<DatabricksProvider, _>(DatabricksProvider::from_env);
registry.register::<GcpVertexAIProvider, _>(GcpVertexAIProvider::from_env);
registry.register::<GeminiCliProvider, _>(GeminiCliProvider::from_env);
registry.register::<GoogleProvider, _>(GoogleProvider::from_env);
registry.register::<GroqProvider, _>(GroqProvider::from_env);
registry.register::<LiteLLMProvider, _>(LiteLLMProvider::from_env);
registry.register::<OllamaProvider, _>(OllamaProvider::from_env);
registry.register::<OpenAiProvider, _>(OpenAiProvider::from_env);
registry.register::<OpenRouterProvider, _>(OpenRouterProvider::from_env);
registry.register::<SageMakerTgiProvider, _>(SageMakerTgiProvider::from_env);
registry.register::<SnowflakeProvider, _>(SnowflakeProvider::from_env);
registry.register::<VeniceProvider, _>(VeniceProvider::from_env);
registry.register::<XaiProvider, _>(XaiProvider::from_env);
if let Err(e) = load_custom_providers_into_registry(registry) {
tracing::warn!("Failed to load custom providers: {}", e);
}
});
RwLock::new(registry)
});
fn load_custom_providers_into_registry(registry: &mut ProviderRegistry) -> Result<()> {
let config_dir = custom_providers_dir();
register_custom_providers(registry, &config_dir)
}
pub fn providers() -> Vec<ProviderMetadata> {
vec![
AnthropicProvider::metadata(),
AzureProvider::metadata(),
BedrockProvider::metadata(),
ClaudeCodeProvider::metadata(),
CursorAgentProvider::metadata(),
DatabricksProvider::metadata(),
GcpVertexAIProvider::metadata(),
GeminiCliProvider::metadata(),
// GithubCopilotProvider::metadata(),
GoogleProvider::metadata(),
GroqProvider::metadata(),
LiteLLMProvider::metadata(),
OllamaProvider::metadata(),
OpenAiProvider::metadata(),
OpenRouterProvider::metadata(),
SageMakerTgiProvider::metadata(),
VeniceProvider::metadata(),
SnowflakeProvider::metadata(),
XaiProvider::metadata(),
]
REGISTRY.read().unwrap().all_metadata()
}
pub fn refresh_custom_providers() -> Result<()> {
let mut registry = REGISTRY.write().unwrap();
registry.remove_custom_providers();
if let Err(e) = load_custom_providers_into_registry(&mut registry) {
tracing::warn!("Failed to refresh custom providers: {}", e);
return Err(e);
}
tracing::info!("Custom providers refreshed");
Ok(())
}
pub fn create(name: &str, model: ModelConfig) -> Result<Arc<dyn Provider>> {
let config = crate::config::Config::global();
// Check for lead model environment variables
if let Ok(lead_model_name) = config.get_param::<String>("GOOSE_LEAD_MODEL") {
tracing::info!("Creating lead/worker provider from environment variables");
return create_lead_worker_from_env(name, &model, &lead_model_name);
}
create_provider(name, model)
REGISTRY.read().unwrap().create(name, model)
}
/// Create a lead/worker provider from environment variables
fn create_lead_worker_from_env(
default_provider_name: &str,
default_model: &ModelConfig,
@@ -84,61 +105,36 @@ fn create_lead_worker_from_env(
) -> Result<Arc<dyn Provider>> {
let config = crate::config::Config::global();
// Get lead provider (optional, defaults to main provider)
let lead_provider_name = config
.get_param::<String>("GOOSE_LEAD_PROVIDER")
.unwrap_or_else(|_| default_provider_name.to_string());
// Get configuration parameters with defaults
let lead_turns = config
.get_param::<usize>("GOOSE_LEAD_TURNS")
.unwrap_or(default_lead_turns());
.unwrap_or(DEFAULT_LEAD_TURNS);
let failure_threshold = config
.get_param::<usize>("GOOSE_LEAD_FAILURE_THRESHOLD")
.unwrap_or(default_failure_threshold());
.unwrap_or(DEFAULT_FAILURE_THRESHOLD);
let fallback_turns = config
.get_param::<usize>("GOOSE_LEAD_FALLBACK_TURNS")
.unwrap_or(default_fallback_turns());
.unwrap_or(DEFAULT_FALLBACK_TURNS);
let lead_model_config = ModelConfig::new_with_context_env(
lead_model_name.to_string(),
Some("GOOSE_LEAD_CONTEXT_LIMIT"),
)?;
// For worker model, preserve the original context_limit from config (highest precedence)
// while still allowing environment variable overrides
let worker_model_config = {
// Start with a clone of the original model to preserve user-specified settings
let mut worker_config = ModelConfig::new_or_fail(default_model.model_name.as_str())
.with_context_limit(default_model.context_limit)
.with_temperature(default_model.temperature)
.with_max_tokens(default_model.max_tokens)
.with_toolshim(default_model.toolshim)
.with_toolshim_model(default_model.toolshim_model.clone());
let worker_model_config = create_worker_model_config(default_model)?;
// Apply environment variable overrides with proper precedence
let global_config = crate::config::Config::global();
let lead_provider = REGISTRY
.read()
.unwrap()
.create(&lead_provider_name, lead_model_config)?;
let worker_provider = REGISTRY
.read()
.unwrap()
.create(default_provider_name, worker_model_config)?;
// Check for worker-specific context limit
if let Ok(limit_str) = global_config.get_param::<String>("GOOSE_WORKER_CONTEXT_LIMIT") {
if let Ok(limit) = limit_str.parse::<usize>() {
worker_config = worker_config.with_context_limit(Some(limit));
}
} else if let Ok(limit_str) = global_config.get_param::<String>("GOOSE_CONTEXT_LIMIT") {
// Check for general context limit if worker-specific is not set
if let Ok(limit) = limit_str.parse::<usize>() {
worker_config = worker_config.with_context_limit(Some(limit));
}
}
worker_config
};
// Create the providers
let lead_provider = create_provider(&lead_provider_name, lead_model_config)?;
let worker_provider = create_provider(default_provider_name, worker_model_config)?;
// Create the lead/worker provider with configured settings
Ok(Arc::new(LeadWorkerProvider::new_with_settings(
lead_provider,
worker_provider,
@@ -148,30 +144,27 @@ fn create_lead_worker_from_env(
)))
}
fn create_provider(name: &str, model: ModelConfig) -> Result<Arc<dyn Provider>> {
// We use Arc instead of Box to be able to clone for multiple async tasks
match name {
"anthropic" => Ok(Arc::new(AnthropicProvider::from_env(model)?)),
"aws_bedrock" => Ok(Arc::new(BedrockProvider::from_env(model)?)),
"azure_openai" => Ok(Arc::new(AzureProvider::from_env(model)?)),
"claude-code" => Ok(Arc::new(ClaudeCodeProvider::from_env(model)?)),
"cursor-agent" => Ok(Arc::new(CursorAgentProvider::from_env(model)?)),
"databricks" => Ok(Arc::new(DatabricksProvider::from_env(model)?)),
"gcp_vertex_ai" => Ok(Arc::new(GcpVertexAIProvider::from_env(model)?)),
"gemini-cli" => Ok(Arc::new(GeminiCliProvider::from_env(model)?)),
// "github_copilot" => Ok(Arc::new(GithubCopilotProvider::from_env(model)?)),
"google" => Ok(Arc::new(GoogleProvider::from_env(model)?)),
"groq" => Ok(Arc::new(GroqProvider::from_env(model)?)),
"litellm" => Ok(Arc::new(LiteLLMProvider::from_env(model)?)),
"ollama" => Ok(Arc::new(OllamaProvider::from_env(model)?)),
"openai" => Ok(Arc::new(OpenAiProvider::from_env(model)?)),
"openrouter" => Ok(Arc::new(OpenRouterProvider::from_env(model)?)),
"sagemaker_tgi" => Ok(Arc::new(SageMakerTgiProvider::from_env(model)?)),
"snowflake" => Ok(Arc::new(SnowflakeProvider::from_env(model)?)),
"venice" => Ok(Arc::new(VeniceProvider::from_env(model)?)),
"xai" => Ok(Arc::new(XaiProvider::from_env(model)?)),
_ => Err(anyhow::anyhow!("Unknown provider: {}", name)),
fn create_worker_model_config(default_model: &ModelConfig) -> Result<ModelConfig> {
let mut worker_config = ModelConfig::new_or_fail(&default_model.model_name)
.with_context_limit(default_model.context_limit)
.with_temperature(default_model.temperature)
.with_max_tokens(default_model.max_tokens)
.with_toolshim(default_model.toolshim)
.with_toolshim_model(default_model.toolshim_model.clone());
let global_config = crate::config::Config::global();
if let Ok(limit_str) = global_config.get_param::<String>("GOOSE_WORKER_CONTEXT_LIMIT") {
if let Ok(limit) = limit_str.parse::<usize>() {
worker_config = worker_config.with_context_limit(Some(limit));
}
} else if let Ok(limit_str) = global_config.get_param::<String>("GOOSE_CONTEXT_LIMIT") {
if let Ok(limit) = limit_str.parse::<usize>() {
worker_config = worker_config.with_context_limit(Some(limit));
}
}
Ok(worker_config)
}
#[cfg(test)]
@@ -183,7 +176,6 @@ mod tests {
use rmcp::model::{AnnotateAble, RawTextContent, Role};
use std::env;
#[allow(dead_code)]
#[derive(Clone)]
struct MockTestProvider {
name: String,
@@ -233,222 +225,141 @@ mod tests {
}
}
struct EnvVarGuard {
vars: Vec<(String, Option<String>)>,
}
impl EnvVarGuard {
fn new(vars: &[&str]) -> Self {
let saved_vars = vars
.iter()
.map(|&var| (var.to_string(), env::var(var).ok()))
.collect();
for &var in vars {
env::remove_var(var);
}
Self { vars: saved_vars }
}
fn set(&self, key: &str, value: &str) {
env::set_var(key, value);
}
}
impl Drop for EnvVarGuard {
fn drop(&mut self) {
for (key, value) in &self.vars {
match value {
Some(val) => env::set_var(key, val),
None => env::remove_var(key),
}
}
}
}
#[test]
fn test_create_lead_worker_provider() {
// Save current env vars
let saved_lead = env::var("GOOSE_LEAD_MODEL").ok();
let saved_provider = env::var("GOOSE_LEAD_PROVIDER").ok();
let saved_turns = env::var("GOOSE_LEAD_TURNS").ok();
let _guard = EnvVarGuard::new(&[
"GOOSE_LEAD_MODEL",
"GOOSE_LEAD_PROVIDER",
"GOOSE_LEAD_TURNS",
]);
// Test with basic lead model configuration
env::set_var("GOOSE_LEAD_MODEL", "gpt-4o");
_guard.set("GOOSE_LEAD_MODEL", "gpt-4o");
// This will try to create a lead/worker provider
let gpt4mini_config = ModelConfig::new_or_fail("gpt-4o-mini");
let result = create("openai", gpt4mini_config.clone());
// The creation might succeed or fail depending on API keys, but we can verify the logic path
match result {
Ok(_) => {
// If it succeeds, it means we created a lead/worker provider successfully
// This would happen if API keys are available in the test environment
}
Ok(_) => {}
Err(error) => {
// If it fails, it should be due to missing API keys, confirming we tried to create providers
let error_msg = error.to_string();
assert!(error_msg.contains("OPENAI_API_KEY") || error_msg.contains("secret"));
}
}
// Test with different lead provider
env::set_var("GOOSE_LEAD_PROVIDER", "anthropic");
env::set_var("GOOSE_LEAD_TURNS", "5");
_guard.set("GOOSE_LEAD_PROVIDER", "anthropic");
_guard.set("GOOSE_LEAD_TURNS", "5");
let _result = create("openai", gpt4mini_config);
// Similar validation as above - will fail due to missing API keys but confirms the logic
// Restore env vars
match saved_lead {
Some(val) => env::set_var("GOOSE_LEAD_MODEL", val),
None => env::remove_var("GOOSE_LEAD_MODEL"),
}
match saved_provider {
Some(val) => env::set_var("GOOSE_LEAD_PROVIDER", val),
None => env::remove_var("GOOSE_LEAD_PROVIDER"),
}
match saved_turns {
Some(val) => env::set_var("GOOSE_LEAD_TURNS", val),
None => env::remove_var("GOOSE_LEAD_TURNS"),
}
}
#[test]
fn test_lead_model_env_vars_with_defaults() {
// Save current env vars
let saved_vars = [
("GOOSE_LEAD_MODEL", env::var("GOOSE_LEAD_MODEL").ok()),
("GOOSE_LEAD_PROVIDER", env::var("GOOSE_LEAD_PROVIDER").ok()),
("GOOSE_LEAD_TURNS", env::var("GOOSE_LEAD_TURNS").ok()),
(
"GOOSE_LEAD_FAILURE_THRESHOLD",
env::var("GOOSE_LEAD_FAILURE_THRESHOLD").ok(),
),
(
"GOOSE_LEAD_FALLBACK_TURNS",
env::var("GOOSE_LEAD_FALLBACK_TURNS").ok(),
),
];
let _guard = EnvVarGuard::new(&[
"GOOSE_LEAD_MODEL",
"GOOSE_LEAD_PROVIDER",
"GOOSE_LEAD_TURNS",
"GOOSE_LEAD_FAILURE_THRESHOLD",
"GOOSE_LEAD_FALLBACK_TURNS",
]);
// Clear all lead env vars
for (key, _) in &saved_vars {
env::remove_var(key);
}
_guard.set("GOOSE_LEAD_MODEL", "grok-3");
// Set only the required lead model
env::set_var("GOOSE_LEAD_MODEL", "grok-3");
// This should use defaults for all other values
let result = create("openai", ModelConfig::new_or_fail("gpt-4o-mini"));
// Should attempt to create lead/worker provider (will fail due to missing API keys but confirms logic)
match result {
Ok(_) => {
// Success means we have API keys and created the provider
}
Ok(_) => {}
Err(error) => {
// Should fail due to missing API keys, confirming we tried to create providers
let error_msg = error.to_string();
assert!(error_msg.contains("OPENAI_API_KEY") || error_msg.contains("secret"));
}
}
// Test with custom values
env::set_var("GOOSE_LEAD_TURNS", "7");
env::set_var("GOOSE_LEAD_FAILURE_THRESHOLD", "4");
env::set_var("GOOSE_LEAD_FALLBACK_TURNS", "3");
_guard.set("GOOSE_LEAD_TURNS", "7");
_guard.set("GOOSE_LEAD_FAILURE_THRESHOLD", "4");
_guard.set("GOOSE_LEAD_FALLBACK_TURNS", "3");
let _result = create("openai", ModelConfig::new_or_fail("gpt-4o-mini"));
// Should still attempt to create lead/worker provider with custom settings
// Restore all env vars
for (key, value) in saved_vars {
match value {
Some(val) => env::set_var(key, val),
None => env::remove_var(key),
}
}
}
#[test]
fn test_create_regular_provider_without_lead_config() {
// Save current env vars
let saved_lead = env::var("GOOSE_LEAD_MODEL").ok();
let saved_provider = env::var("GOOSE_LEAD_PROVIDER").ok();
let saved_turns = env::var("GOOSE_LEAD_TURNS").ok();
let saved_threshold = env::var("GOOSE_LEAD_FAILURE_THRESHOLD").ok();
let saved_fallback = env::var("GOOSE_LEAD_FALLBACK_TURNS").ok();
let _guard = EnvVarGuard::new(&[
"GOOSE_LEAD_MODEL",
"GOOSE_LEAD_PROVIDER",
"GOOSE_LEAD_TURNS",
"GOOSE_LEAD_FAILURE_THRESHOLD",
"GOOSE_LEAD_FALLBACK_TURNS",
]);
// Ensure all GOOSE_LEAD_* variables are not set
env::remove_var("GOOSE_LEAD_MODEL");
env::remove_var("GOOSE_LEAD_PROVIDER");
env::remove_var("GOOSE_LEAD_TURNS");
env::remove_var("GOOSE_LEAD_FAILURE_THRESHOLD");
env::remove_var("GOOSE_LEAD_FALLBACK_TURNS");
// This should try to create a regular provider
let result = create("openai", ModelConfig::new_or_fail("gpt-4o-mini"));
// The creation might succeed or fail depending on API keys
match result {
Ok(_) => {
// If it succeeds, it means we created a regular provider successfully
// This would happen if API keys are available in the test environment
}
Ok(_) => {}
Err(error) => {
// If it fails, it should be due to missing API keys
let error_msg = error.to_string();
assert!(error_msg.contains("OPENAI_API_KEY") || error_msg.contains("secret"));
}
}
if let Some(val) = saved_lead {
env::set_var("GOOSE_LEAD_MODEL", val);
}
if let Some(val) = saved_provider {
env::set_var("GOOSE_LEAD_PROVIDER", val);
}
if let Some(val) = saved_turns {
env::set_var("GOOSE_LEAD_TURNS", val);
}
if let Some(val) = saved_threshold {
env::set_var("GOOSE_LEAD_FAILURE_THRESHOLD", val);
}
if let Some(val) = saved_fallback {
env::set_var("GOOSE_LEAD_FALLBACK_TURNS", val);
}
}
#[test]
fn test_worker_model_preserves_original_context_limit() {
use std::env;
let _guard = EnvVarGuard::new(&[
"GOOSE_LEAD_MODEL",
"GOOSE_WORKER_CONTEXT_LIMIT",
"GOOSE_CONTEXT_LIMIT",
]);
// Save current env vars
let saved_vars = [
("GOOSE_LEAD_MODEL", env::var("GOOSE_LEAD_MODEL").ok()),
(
"GOOSE_WORKER_CONTEXT_LIMIT",
env::var("GOOSE_WORKER_CONTEXT_LIMIT").ok(),
),
("GOOSE_CONTEXT_LIMIT", env::var("GOOSE_CONTEXT_LIMIT").ok()),
];
_guard.set("GOOSE_LEAD_MODEL", "gpt-4o");
// Clear env vars to ensure clean test
for (key, _) in &saved_vars {
env::remove_var(key);
}
// Set up lead model to trigger lead/worker mode
env::set_var("GOOSE_LEAD_MODEL", "gpt-4o");
// Create a default model with explicit context_limit
let default_model =
ModelConfig::new_or_fail("gpt-3.5-turbo").with_context_limit(Some(16_000));
// Test case 1: No environment variables - should preserve original context_limit
let result = create_lead_worker_from_env("openai", &default_model, "gpt-4o");
// Test case 2: With GOOSE_WORKER_CONTEXT_LIMIT - should override original
env::set_var("GOOSE_WORKER_CONTEXT_LIMIT", "32000");
_guard.set("GOOSE_WORKER_CONTEXT_LIMIT", "32000");
let _result = create_lead_worker_from_env("openai", &default_model, "gpt-4o");
env::remove_var("GOOSE_WORKER_CONTEXT_LIMIT");
// Test case 3: With GOOSE_CONTEXT_LIMIT - should override original
env::set_var("GOOSE_CONTEXT_LIMIT", "64000");
_guard.set("GOOSE_CONTEXT_LIMIT", "64000");
let _result = create_lead_worker_from_env("openai", &default_model, "gpt-4o");
env::remove_var("GOOSE_CONTEXT_LIMIT");
// Restore env vars
for (key, value) in saved_vars {
match value {
Some(val) => env::set_var(key, val),
None => env::remove_var(key),
}
}
// The main verification is that the function doesn't panic and handles
// the context limit preservation logic correctly. More detailed testing
// would require mocking the provider creation.
// The result could be Ok or Err depending on whether API keys are available
// in the test environment - both are acceptable for this test
match result {
Ok(_) => {
// Success means API keys are available and lead/worker provider was created
// This confirms our logic path is working
}
Err(_) => {
// Error is expected if API keys are not available
// This also confirms our logic path is working
}
Ok(_) => {}
Err(_) => {}
}
}
}
crates/goose/src/providers/mod.rs
+2 -1
View File
@@ -24,6 +24,7 @@ pub mod ollama;
pub mod openai;
pub mod openrouter;
pub mod pricing;
pub mod provider_registry;
mod retry;
pub mod sagemaker_tgi;
pub mod snowflake;
@@ -35,4 +36,4 @@ pub mod utils_universal_openai_stream;
pub mod venice;
pub mod xai;
pub use factory::{create, providers};
pub use factory::{create, providers, refresh_custom_providers};
crates/goose/src/providers/ollama.rs
+47 -1
View File
@@ -3,6 +3,7 @@ use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage};
use super::errors::ProviderError;
use super::retry::ProviderRetry;
use super::utils::{get_model, handle_response_openai_compat};
use crate::config::custom_providers::CustomProviderConfig;
use crate::conversation::message::Message;
use crate::conversation::Conversation;
use crate::impl_provider_default;
@@ -30,6 +31,7 @@ pub struct OllamaProvider {
#[serde(skip)]
api_client: ApiClient,
model: ModelConfig,
supports_streaming: bool,
}
impl_provider_default!(OllamaProvider);
@@ -73,7 +75,47 @@ impl OllamaProvider {
let auth = AuthMethod::Custom(Box::new(NoAuth));
let api_client = ApiClient::with_timeout(base_url.to_string(), auth, timeout)?;
Ok(Self { api_client, model })
Ok(Self {
api_client,
model,
supports_streaming: true,
})
}
pub fn from_custom_config(model: ModelConfig, config: CustomProviderConfig) -> Result<Self> {
let timeout = Duration::from_secs(config.timeout_seconds.unwrap_or(OLLAMA_TIMEOUT));
// Parse and normalize the custom URL
let base =
if config.base_url.starts_with("http://") || config.base_url.starts_with("https://") {
config.base_url.clone()
} else {
format!("http://{}", config.base_url)
};
let mut base_url = Url::parse(&base)
.map_err(|e| anyhow::anyhow!("Invalid base URL '{}': {}", config.base_url, e))?;
// Set default port if missing and not using standard ports
let explicit_default_port =
config.base_url.ends_with(":80") || config.base_url.ends_with(":443");
let is_https = base_url.scheme() == "https";
if base_url.port().is_none() && !explicit_default_port && !is_https {
base_url
.set_port(Some(OLLAMA_DEFAULT_PORT))
.map_err(|_| anyhow::anyhow!("Failed to set default port"))?;
}
// No authentication for Ollama
let auth = AuthMethod::Custom(Box::new(NoAuth));
let api_client = ApiClient::with_timeout(base_url.to_string(), auth, timeout)?;
Ok(Self {
api_client,
model,
supports_streaming: config.supports_streaming.unwrap_or(true),
})
}
async fn post(&self, payload: &Value) -> Result<Value, ProviderError> {
@@ -181,6 +223,10 @@ impl Provider for OllamaProvider {
Ok(safe_truncate(&description, 100))
}
fn supports_streaming(&self) -> bool {
self.supports_streaming
}
}
impl OllamaProvider {
crates/goose/src/providers/openai.rs
+48 -1
View File
@@ -20,6 +20,7 @@ use super::utils::{
emit_debug_trace, get_model, handle_response_openai_compat, handle_status_openai_compat,
ImageFormat,
};
use crate::config::custom_providers::CustomProviderConfig;
use crate::conversation::message::Message;
use crate::impl_provider_default;
use crate::model::ModelConfig;
@@ -51,6 +52,7 @@ pub struct OpenAiProvider {
project: Option<String>,
model: ModelConfig,
custom_headers: Option<HashMap<String, String>>,
supports_streaming: bool,
}
impl_provider_default!(OpenAiProvider);
@@ -103,6 +105,51 @@ impl OpenAiProvider {
project,
model,
custom_headers,
supports_streaming: true,
})
}
pub fn from_custom_config(model: ModelConfig, config: CustomProviderConfig) -> Result<Self> {
let global_config = crate::config::Config::global();
let api_key: String = global_config
.get_secret(&config.api_key_env)
.map_err(|_e| anyhow::anyhow!("Missing API key: {}", config.api_key_env))?;
let url = url::Url::parse(&config.base_url)
.map_err(|e| anyhow::anyhow!("Invalid base URL '{}': {}", config.base_url, e))?;
let host = format!("{}://{}", url.scheme(), url.host_str().unwrap_or(""));
let base_path = url.path().trim_start_matches('/').to_string();
let base_path = if base_path.is_empty() {
"v1/chat/completions".to_string()
} else {
base_path
};
let timeout_secs = config.timeout_seconds.unwrap_or(600);
let auth = AuthMethod::BearerToken(api_key);
let mut api_client =
ApiClient::with_timeout(host, auth, std::time::Duration::from_secs(timeout_secs))?;
// Add custom headers if present
if let Some(headers) = &config.headers {
let mut header_map = reqwest::header::HeaderMap::new();
for (key, value) in headers {
let header_name = reqwest::header::HeaderName::from_bytes(key.as_bytes())?;
let header_value = reqwest::header::HeaderValue::from_str(value)?;
header_map.insert(header_name, header_value);
}
api_client = api_client.with_headers(header_map)?;
}
Ok(Self {
api_client,
base_path,
organization: None,
project: None,
model,
custom_headers: config.headers,
supports_streaming: config.supports_streaming.unwrap_or(true),
})
}
@@ -206,7 +253,7 @@ impl Provider for OpenAiProvider {
}
fn supports_streaming(&self) -> bool {
true
self.supports_streaming
}
async fn stream(
crates/goose/src/providers/provider_registry.rs
+102
View File
@@ -0,0 +1,102 @@
use super::base::{Provider, ProviderMetadata};
use crate::model::ModelConfig;
use anyhow::Result;
use std::collections::HashMap;
use std::sync::Arc;
type ProviderConstructor = Box<dyn Fn(ModelConfig) -> Result<Arc<dyn Provider>> + Send + Sync>;
struct ProviderEntry {
metadata: ProviderMetadata,
constructor: ProviderConstructor,
}
#[derive(Default)]
pub struct ProviderRegistry {
entries: HashMap<String, ProviderEntry>,
}
impl ProviderRegistry {
pub fn new() -> Self {
Self {
entries: HashMap::new(),
}
}
pub fn register<P, F>(&mut self, constructor: F)
where
P: Provider + 'static,
F: Fn(ModelConfig) -> Result<P> + Send + Sync + 'static,
{
let metadata = P::metadata();
let name = metadata.name.clone();
self.entries.insert(
name,
ProviderEntry {
metadata,
constructor: Box::new(move |model| Ok(Arc::new(constructor(model)?))),
},
);
}
/// create provider with custom name
pub fn register_with_name<P, F>(
&mut self,
custom_name: String,
display_name: String,
description: String,
default_model: String,
known_models: Vec<super::base::ModelInfo>,
constructor: F,
) where
P: Provider + 'static,
F: Fn(ModelConfig) -> Result<P> + Send + Sync + 'static,
{
let base_metadata = P::metadata();
let custom_metadata = ProviderMetadata {
name: custom_name.clone(),
display_name,
description,
default_model,
known_models,
model_doc_link: base_metadata.model_doc_link,
config_keys: base_metadata.config_keys,
};
self.entries.insert(
custom_name,
ProviderEntry {
metadata: custom_metadata,
constructor: Box::new(move |model| Ok(Arc::new(constructor(model)?))),
},
);
}
pub fn with_providers<F>(mut self, setup: F) -> Self
where
F: FnOnce(&mut Self),
{
setup(&mut self);
self
}
pub fn create(&self, name: &str, model: ModelConfig) -> Result<Arc<dyn Provider>> {
let _available_providers: Vec<_> = self.entries.keys().collect();
let entry = self
.entries
.get(name)
.ok_or_else(|| anyhow::anyhow!("Unknown provider: {}", name))?;
(entry.constructor)(model)
}
pub fn all_metadata(&self) -> Vec<ProviderMetadata> {
self.entries.values().map(|e| e.metadata.clone()).collect()
}
pub fn remove_custom_providers(&mut self) {
self.entries.retain(|name, _| !name.starts_with("custom_"));
}
}