From 12f80edee6e8d596db1609fb833bdc0e649851a4 Mon Sep 17 00:00:00 2001 From: Rob Gray <1772089+robjgray@users.noreply.github.com> Date: Tue, 7 Jul 2026 20:48:55 -0700 Subject: [PATCH] feat(providers): add OllamaCloudProvider with dynamic model discovery and context limits (#10264) Co-authored-by: Angie Jones --- crates/goose-providers/src/ollama.rs | 82 +-- .../goose/src/config/declarative_providers.rs | 16 + crates/goose/src/providers/mod.rs | 1 + crates/goose/src/providers/ollama_cloud.rs | 485 ++++++++++++++++++ 4 files changed, 545 insertions(+), 39 deletions(-) create mode 100644 crates/goose/src/providers/ollama_cloud.rs diff --git a/crates/goose-providers/src/ollama.rs b/crates/goose-providers/src/ollama.rs index 9d7fb35cf1..f6ac65fcc0 100644 --- a/crates/goose-providers/src/ollama.rs +++ b/crates/goose-providers/src/ollama.rs @@ -184,48 +184,52 @@ impl OllamaProvider { } async fn fetch_models_from_api(&self) -> Result, ProviderError> { - let response = self - .api_client - .request("api/tags") - .response_get() - .await - .map_err(|e| ProviderError::RequestFailed(format!("Failed to fetch models: {}", e)))?; - - if response.status() == StatusCode::NOT_FOUND { - return Err(ProviderError::EndpointNotFound( - "Ollama models endpoint not found".to_string(), - )); - } - - if !response.status().is_success() { - return Err(ProviderError::RequestFailed(format!( - "Failed to fetch models: HTTP {}", - response.status() - ))); - } - - let json_response = response.json::().await.map_err(|e| { - ProviderError::RequestFailed(format!("Failed to parse response: {}", e)) - })?; - - let models = json_response - .get("models") - .and_then(|m| m.as_array()) - .ok_or_else(|| { - ProviderError::RequestFailed("No models array in response".to_string()) - })?; - - let mut model_names: Vec = models - .iter() - .filter_map(|model| model.get("name").and_then(|n| n.as_str()).map(String::from)) - .collect(); - - model_names.sort(); - - Ok(model_names) + fetch_ollama_model_names(&self.api_client) + .await? + .ok_or_else(|| ProviderError::RequestFailed("No models array in response".to_string())) } } +pub async fn fetch_ollama_model_names( + client: &ApiClient, +) -> Result>, ProviderError> { + let response = client + .request("api/tags") + .response_get() + .await + .map_err(|e| ProviderError::RequestFailed(format!("Failed to fetch models: {}", e)))?; + + if response.status() == StatusCode::NOT_FOUND { + return Err(ProviderError::EndpointNotFound( + "Ollama models endpoint not found".to_string(), + )); + } + + if !response.status().is_success() { + return Err(ProviderError::RequestFailed(format!( + "Failed to fetch models: HTTP {}", + response.status() + ))); + } + + let json: Value = response + .json() + .await + .map_err(|e| ProviderError::RequestFailed(format!("Failed to parse response: {}", e)))?; + + let Some(models) = json.get("models").and_then(|m| m.as_array()) else { + return Ok(None); + }; + + let mut names: Vec = models + .iter() + .filter_map(|m| m.get("name").and_then(|n| n.as_str()).map(String::from)) + .collect(); + + names.sort(); + Ok(Some(names)) +} + fn resolve_ollama_num_ctx(options: &OllamaOptions, model_config: &ModelConfig) -> Option { options.input_limit.or(model_config.context_limit) } diff --git a/crates/goose/src/config/declarative_providers.rs b/crates/goose/src/config/declarative_providers.rs index 88981fea50..b5b663f6dd 100644 --- a/crates/goose/src/config/declarative_providers.rs +++ b/crates/goose/src/config/declarative_providers.rs @@ -430,6 +430,22 @@ pub fn register_declarative_provider( huggingface_declarative_inventory_configured(&cfg) }, ); + } else if crate::providers::ollama_cloud::OllamaCloudProvider::matches_declarative_config(&config) { + registry.register_with_name::( + &config, + provider_type, + config.dynamic_models.unwrap_or(false), + move |tls_config| { + let mut cfg = captured.clone(); + resolve_config(&mut cfg)?; + crate::providers::ollama_cloud::OllamaCloudProvider::from_custom_config(cfg, tls_config) + }, + move || { + let mut cfg = identity_config.clone(); + resolve_config(&mut cfg)?; + declarative_inventory_identity(&cfg) + }, + ); } else { registry.register_with_name::( &config, diff --git a/crates/goose/src/providers/mod.rs b/crates/goose/src/providers/mod.rs index c0a961b01b..e0f0a8e64b 100644 --- a/crates/goose/src/providers/mod.rs +++ b/crates/goose/src/providers/mod.rs @@ -58,6 +58,7 @@ pub mod oauth_device_flow; pub mod ollama { pub use goose_providers::ollama::*; } +pub mod ollama_cloud; pub mod ollama_def; pub mod openai { pub use goose_providers::openai::*; diff --git a/crates/goose/src/providers/ollama_cloud.rs b/crates/goose/src/providers/ollama_cloud.rs new file mode 100644 index 0000000000..8a36e35e90 --- /dev/null +++ b/crates/goose/src/providers/ollama_cloud.rs @@ -0,0 +1,485 @@ +use super::base::{ConfigKey, MessageStream, Provider, ProviderDef, ProviderMetadata}; +use crate::config::declarative_providers::DeclarativeProviderConfig; +use crate::config::Config; +use crate::conversation::message::Message; +use anyhow::Result; +use futures::future::BoxFuture; +use goose_providers::api_client::{ApiClient, AuthMethod, TlsConfig}; +use goose_providers::base::ProviderDescriptor; +use goose_providers::errors::ProviderError; +use goose_providers::model::ModelConfig; +use goose_providers::ollama::fetch_ollama_model_names; +use goose_providers::openai::OpenAiProvider; +use rmcp::model::Tool; +use serde_json::Value; +use std::collections::HashMap; +use std::sync::{LazyLock, Mutex}; +use tokio::sync::OnceCell; + +const OLLAMA_CLOUD_PROVIDER_NAME: &str = "ollama_cloud"; + +static SHOW_INFO_CACHE: LazyLock>> = + LazyLock::new(|| Mutex::new(HashMap::new())); + +pub struct OllamaCloudProvider { + inner: OpenAiProvider, + ollama_api_client: ApiClient, + model_names: OnceCell>, + custom_models: Option>, + dynamic_models: Option, +} + +impl OllamaCloudProvider { + pub fn matches_declarative_config(config: &DeclarativeProviderConfig) -> bool { + config.name == OLLAMA_CLOUD_PROVIDER_NAME + || config.catalog_provider_id.as_deref() == Some(OLLAMA_CLOUD_PROVIDER_NAME) + } + + pub fn from_custom_config( + config: DeclarativeProviderConfig, + tls_config: Option, + ) -> Result { + let inner = + crate::providers::openai_def::from_custom_config(config.clone(), tls_config.clone())?; + + let custom_models = if !config.models.is_empty() { + Some( + config + .models + .iter() + .map(|m| m.name.clone()) + .collect::>(), + ) + } else { + None + }; + + if config.dynamic_models == Some(false) && custom_models.is_none() { + return Err(anyhow::anyhow!( + "Provider '{}' has dynamic_models: false but no static models listed; \ + at least one entry in `models` is required.", + config.name + )); + } + + let ollama_api_client = build_ollama_api_client(&config, tls_config)?; + + Ok(Self { + inner, + ollama_api_client, + model_names: OnceCell::new(), + custom_models, + dynamic_models: config.dynamic_models, + }) + } + + async fn get_or_fetch_model_names(&self) -> Result, ProviderError> { + self.model_names + .get_or_try_init(|| { + Box::pin(async { + Ok(fetch_ollama_model_names(&self.ollama_api_client) + .await? + .unwrap_or_default()) + }) + }) + .await + .map(|v| v.to_vec()) + } + + async fn fetch_context_limit_from_show(&self, model_name: &str) -> Option { + let payload = serde_json::json!({ "model": model_name }); + let response = self + .ollama_api_client + .request("api/show") + .response_post(&payload) + .await + .ok()?; + + if !response.status().is_success() { + return None; + } + + let json: Value = response.json().await.ok()?; + json.get("model_info") + .and_then(|info| info.as_object()) + .and_then(|obj| { + obj.iter().find_map(|(key, value)| { + key.ends_with(".context_length") + .then(|| value.as_u64().map(|n| n as usize)) + .flatten() + }) + }) + } +} + +fn build_ollama_api_client( + config: &DeclarativeProviderConfig, + tls_config: Option, +) -> Result { + let normalized_base_url = goose_providers::openai::ensure_url_scheme(&config.base_url); + let url = url::Url::parse(&normalized_base_url) + .map_err(|e| anyhow::anyhow!("Invalid base URL '{}': {}", config.base_url, e))?; + let host = url[..url::Position::BeforePath].to_string(); + + let api_key = crate::providers::openai_def::resolve_api_key(config, &|key| { + Config::global().get_secret(key) + })?; + + let timeout_secs = config + .timeout_seconds + .unwrap_or(crate::providers::base::DEFAULT_PROVIDER_TIMEOUT_SECS); + + let auth = match api_key { + Some(key) if !key.is_empty() => AuthMethod::BearerToken(key), + _ => AuthMethod::NoAuth, + }; + + let mut api_client = ApiClient::with_timeout_and_tls( + host, + auth, + std::time::Duration::from_secs(timeout_secs), + tls_config, + )?; + + if let Some(query) = url.query() { + let query_params = url::form_urlencoded::parse(query.as_bytes()) + .map(|(key, value)| (key.into_owned(), value.into_owned())) + .collect(); + api_client = api_client.with_query(query_params); + } + + 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(api_client.with_request_builder(crate::session_context::session_id_request_builder())) +} + +#[async_trait::async_trait] +impl Provider for OllamaCloudProvider { + fn get_name(&self) -> &str { + self.inner.get_name() + } + + async fn stream( + &self, + model_config: &ModelConfig, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result { + self.inner + .stream(model_config, system, messages, tools) + .await + } + + fn skip_canonical_filtering(&self) -> bool { + self.inner.skip_canonical_filtering() + } + + fn retry_config(&self) -> goose_providers::retry::RetryConfig { + self.inner.retry_config() + } + + async fn fetch_supported_models(&self) -> Result, ProviderError> { + if let Some(custom_models) = &self.custom_models { + if self.dynamic_models == Some(false) { + return Ok(custom_models.clone()); + } + + match self.get_or_fetch_model_names().await { + Ok(models) => return Ok(models), + Err(e) if e.is_endpoint_not_found() => { + tracing::debug!( + "Ollama api/tags not available for provider '{}', using static model list", + self.inner.get_name(), + ); + return Ok(custom_models.clone()); + } + Err(e) => return Err(e), + } + } + + self.get_or_fetch_model_names().await + } + + async fn get_context_limit(&self, model_config: &ModelConfig) -> Result { + if let Some(limit) = model_config.context_limit { + return Ok(limit); + } + + if let Some(cached) = SHOW_INFO_CACHE + .lock() + .ok() + .and_then(|cache| cache.get(&model_config.model_name).copied()) + { + return Ok(cached); + } + + if let Some(limit) = self + .fetch_context_limit_from_show(&model_config.model_name) + .await + { + if let Ok(mut cache) = SHOW_INFO_CACHE.lock() { + cache.insert(model_config.model_name.clone(), limit); + } + + return Ok(limit); + } + + Ok(model_config.context_limit()) + } +} + +impl ProviderDescriptor for OllamaCloudProvider { + fn metadata() -> ProviderMetadata { + ProviderMetadata::new( + OLLAMA_CLOUD_PROVIDER_NAME, + "Ollama Cloud", + "Access hosted models on ollama.com via OpenAI-compatible API", + "qwen3-coder:480b-cloud", + vec![], + "https://ollama.com/library", + vec![ConfigKey::new( + "ollama_cloud_api_key", + false, + true, + None, + true, + )], + ) + } +} + +impl ProviderDef for OllamaCloudProvider { + type Provider = Self; + + fn from_env( + _extensions: Vec, + _tls_config: Option, + ) -> BoxFuture<'static, Result> { + Box::pin(async { + anyhow::bail!( + "Ollama Cloud must be configured as a declarative provider. \ + Run `goose configure` to set it up." + ) + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::declarative_providers::ProviderEngine; + use crate::providers::base::ModelInfo; + + #[test] + fn declarative_matching_accepts_name_or_catalog_provider_id() { + let mut config = test_config(); + config.name = "custom_ollama".to_string(); + assert!(!OllamaCloudProvider::matches_declarative_config(&config)); + + config.name = OLLAMA_CLOUD_PROVIDER_NAME.to_string(); + assert!(OllamaCloudProvider::matches_declarative_config(&config)); + + config.name = "custom_ollama".to_string(); + config.catalog_provider_id = Some(OLLAMA_CLOUD_PROVIDER_NAME.to_string()); + assert!(OllamaCloudProvider::matches_declarative_config(&config)); + } + + #[tokio::test] + async fn fetch_supported_models_uses_static_models_when_dynamic_models_false() { + let server = mock_api_server(vec![], None).await; + let provider = build_provider( + server.uri(), + Some(false), + vec![ModelInfo::new("static-model", 4096)], + ); + + assert_eq!( + provider.fetch_supported_models().await.unwrap(), + vec!["static-model".to_string()] + ); + } + + #[tokio::test] + async fn fetch_supported_models_falls_back_to_static_on_404() { + let server = mock_api_server(vec![], Some(404)).await; + let provider = build_provider( + server.uri(), + None, + vec![ModelInfo::new("static-model", 4096)], + ); + + assert_eq!( + provider.fetch_supported_models().await.unwrap(), + vec!["static-model".to_string()] + ); + } + + #[tokio::test] + async fn fetch_supported_models_uses_api_when_dynamic_models_true() { + let server = mock_api_server(vec!["api-model-1", "api-model-2"], None).await; + let provider = build_provider( + server.uri(), + Some(true), + vec![ModelInfo::new("static-model", 4096)], + ); + + let models = provider.fetch_supported_models().await.unwrap(); + assert_eq!(models, vec!["api-model-1", "api-model-2"]); + } + + #[tokio::test] + async fn fetch_supported_models_uses_api_when_no_static_models() { + let server = mock_api_server(vec!["api-model"], None).await; + let provider = build_provider(server.uri(), None, vec![]); + + let models = provider.fetch_supported_models().await.unwrap(); + assert_eq!(models, vec!["api-model"]); + } + + #[tokio::test] + async fn get_context_limit_extracts_from_flat_model_info() { + let server = mock_show_server("gemma3", 131072).await; + let provider = build_provider(server.uri(), Some(true), vec![]); + + let model_config = ModelConfig::new("gemma3:4b"); + let limit = provider.get_context_limit(&model_config).await.unwrap(); + assert_eq!(limit, 131072); + } + + #[tokio::test] + async fn get_context_limit_extracts_from_arch_prefixed_key() { + let server = mock_show_server("qwen3moe", 262144).await; + let provider = build_provider(server.uri(), Some(true), vec![]); + + let model_config = ModelConfig::new("qwen3-coder:480b"); + let limit = provider.get_context_limit(&model_config).await.unwrap(); + assert_eq!(limit, 262144); + } + + #[tokio::test] + async fn get_context_limit_falls_back_on_missing_model_info() { + let server = mock_show_server_no_model_info().await; + let provider = build_provider(server.uri(), Some(true), vec![]); + + let model_config = ModelConfig::new("unknown-model").with_context_limit(Some(8000)); + let limit = provider.get_context_limit(&model_config).await.unwrap(); + assert_eq!(limit, 8000); + } + + fn build_provider( + base_url: String, + dynamic_models: Option, + models: Vec, + ) -> OllamaCloudProvider { + let config = test_config_with(base_url, dynamic_models, models); + OllamaCloudProvider::from_custom_config(config, None).unwrap() + } + + async fn mock_api_server( + model_names: Vec<&str>, + status_override: Option, + ) -> wiremock::MockServer { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let server = MockServer::start().await; + let response = match status_override { + Some(404) => ResponseTemplate::new(404), + Some(status) => ResponseTemplate::new(status), + None => { + let models_json: Vec = model_names + .iter() + .map(|n| serde_json::json!({"name": n, "model": n})) + .collect(); + ResponseTemplate::new(200).set_body_json(serde_json::json!({"models": models_json})) + } + }; + Mock::given(method("GET")) + .and(path("/api/tags")) + .respond_with(response) + .mount(&server) + .await; + server + } + + async fn mock_show_server(architecture: &str, context_length: u64) -> wiremock::MockServer { + use wiremock::matchers::{body_partial_json, method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let server = MockServer::start().await; + let key = format!("{}.context_length", architecture); + let model_info = serde_json::json!({ + "general.architecture": architecture, + key: context_length, + }); + Mock::given(method("POST")) + .and(path("/api/show")) + .and(body_partial_json(serde_json::json!({}))) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"model_info": model_info})), + ) + .mount(&server) + .await; + server + } + + async fn mock_show_server_no_model_info() -> wiremock::MockServer { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/show")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({}))) + .mount(&server) + .await; + server + } + + fn test_config_with( + base_url: String, + dynamic_models: Option, + models: Vec, + ) -> DeclarativeProviderConfig { + DeclarativeProviderConfig { + name: OLLAMA_CLOUD_PROVIDER_NAME.to_string(), + engine: ProviderEngine::OpenAI, + display_name: "Ollama Cloud".to_string(), + description: None, + api_key_env: String::new(), + base_url, + models, + headers: None, + timeout_seconds: None, + supports_streaming: Some(true), + requires_auth: false, + catalog_provider_id: None, + base_path: None, + env_vars: None, + dynamic_models, + skip_canonical_filtering: false, + model_doc_link: None, + setup_steps: vec![], + fast_model: None, + preserves_thinking: true, + } + } + + fn test_config() -> DeclarativeProviderConfig { + test_config_with( + "https://ollama.com/v1/chat/completions".to_string(), + Some(true), + vec![], + ) + } +}