diff --git a/assets/config-template.yaml b/assets/config-template.yaml index 62a83a5..d72c31b 100644 --- a/assets/config-template.yaml +++ b/assets/config-template.yaml @@ -277,6 +277,7 @@ sync_models_ttl_hours: 24 # Hours before the synced model catalog is cons # All clients have the following configuration: # - type: xxxx # name: xxxx # Only use it to distinguish clients with the same client type. Optional +# extend_models: true # Merge the models below on top of the catalog models for this client instead of replacing them. Optional, default false # models: # - name: xxxx # Chat model # max_input_tokens: 100000 diff --git a/src/client/azure_openai.rs b/src/client/azure_openai.rs index 3431d3f..9e2cc23 100644 --- a/src/client/azure_openai.rs +++ b/src/client/azure_openai.rs @@ -4,13 +4,15 @@ use super::*; use anyhow::Result; use serde::Deserialize; -#[derive(Debug, Clone, Deserialize)] +#[derive(Debug, Clone, Deserialize, Default)] pub struct AzureOpenAIConfig { pub name: Option, pub api_base: Option, pub api_key: Option, #[serde(default)] pub models: Vec, + #[serde(default)] + pub extend_models: bool, pub patch: Option, pub extra: Option, } diff --git a/src/client/bedrock.rs b/src/client/bedrock.rs index 1a1e63d..c359612 100644 --- a/src/client/bedrock.rs +++ b/src/client/bedrock.rs @@ -16,7 +16,7 @@ use reqwest::{Client as ReqwestClient, Method, RequestBuilder}; use serde::Deserialize; use serde_json::{Value, json}; -#[derive(Debug, Clone, Deserialize)] +#[derive(Debug, Clone, Deserialize, Default)] pub struct BedrockConfig { pub name: Option, pub access_key_id: Option, @@ -25,6 +25,8 @@ pub struct BedrockConfig { pub session_token: Option, #[serde(default)] pub models: Vec, + #[serde(default)] + pub extend_models: bool, pub prompt_cache: Option, pub patch: Option, pub extra: Option, diff --git a/src/client/catalog.rs b/src/client/catalog.rs index 48b5eb1..f36cd98 100644 --- a/src/client/catalog.rs +++ b/src/client/catalog.rs @@ -572,7 +572,7 @@ pub(crate) fn name_overlap(old: &[String], new: &[String]) -> f64 { kept as f64 / old.len() as f64 } -fn overlay_hand_owned(hand: ModelData, api: ModelData) -> ModelData { +pub(crate) fn overlay_hand_owned(hand: ModelData, api: ModelData) -> ModelData { ModelData { name: hand.name, model_type: hand.model_type, @@ -604,6 +604,23 @@ fn overlay_hand_owned(hand: ModelData, api: ModelData) -> ModelData { } } +pub(crate) fn extend_catalog_models( + catalog: &[ModelData], + overrides: &[ModelData], +) -> Vec { + let mut models = catalog.to_vec(); + for override_model in overrides { + match models.iter_mut().find(|model| { + model.name == override_model.name && model.model_type == override_model.model_type + }) { + Some(model) => *model = overlay_hand_owned(override_model.clone(), model.clone()), + None => models.push(override_model.clone()), + } + } + + models +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct ProviderQuirks { @@ -1724,4 +1741,47 @@ rules: ); } } + + fn typed_model(name: &str, model_type: &str) -> ModelData { + ModelData { + name: name.to_string(), + model_type: model_type.to_string(), + ..Default::default() + } + } + + #[test] + fn extend_catalog_models_merges_collisions_field_wise() { + let mut catalog_model = ModelData::new("m"); + catalog_model.max_input_tokens = Some(100_000); + catalog_model.input_price = Some(1.0); + catalog_model.supports_vision = true; + let mut override_model = ModelData::new("m"); + override_model.input_price = Some(9.0); + + let merged = extend_catalog_models(&[catalog_model], &[override_model]); + assert_eq!(merged.len(), 1); + assert_eq!(merged[0].input_price, Some(9.0)); + assert_eq!(merged[0].max_input_tokens, Some(100_000)); + assert!(merged[0].supports_vision); + } + + #[test] + fn extend_catalog_models_appends_net_new_after_catalog() { + let catalog = [ModelData::new("a"), ModelData::new("b")]; + let merged = extend_catalog_models(&catalog, &[ModelData::new("c")]); + let names: Vec<&str> = merged.iter().map(|model| model.name.as_str()).collect(); + assert_eq!(names, ["a", "b", "c"]); + } + + #[test] + fn extend_catalog_models_collides_on_name_and_type() { + let mut catalog_model = typed_model("m", "embedding"); + catalog_model.default_chunk_size = Some(1500); + let merged = extend_catalog_models(&[catalog_model], &[typed_model("m", "reranker")]); + assert_eq!(merged.len(), 2); + assert_eq!(merged[0].model_type, "embedding"); + assert_eq!(merged[0].default_chunk_size, Some(1500)); + assert_eq!(merged[1].model_type, "reranker"); + } } diff --git a/src/client/claude.rs b/src/client/claude.rs index 78f12a1..83e3c84 100644 --- a/src/client/claude.rs +++ b/src/client/claude.rs @@ -15,7 +15,7 @@ use serde_json::{Value, json}; const API_BASE: &str = "https://api.anthropic.com/v1"; const CLAUDE_CODE_PREFIX: &str = "You are Claude Code, Anthropic's official CLI for Claude."; -#[derive(Debug, Clone, Deserialize)] +#[derive(Debug, Clone, Deserialize, Default)] pub struct ClaudeConfig { pub name: Option, pub api_key: Option, @@ -24,6 +24,8 @@ pub struct ClaudeConfig { pub oauth: Option>, #[serde(default)] pub models: Vec, + #[serde(default)] + pub extend_models: bool, pub prompt_cache: Option, pub patch: Option, pub extra: Option, @@ -1122,16 +1124,12 @@ mod tests { fn claude_config(name: &str, auth: Option<&str>, oauth: Option) -> ClaudeConfig { ClaudeConfig { name: Some(name.into()), - api_key: None, api_base: oauth .as_ref() .map(|_| "https://gateway.example/v1".to_string()), auth: auth.map(str::to_string), oauth: oauth.map(Box::new), - models: vec![], - prompt_cache: None, - patch: None, - extra: None, + ..Default::default() } } diff --git a/src/client/cohere.rs b/src/client/cohere.rs index 0ae2f6b..8fe8f21 100644 --- a/src/client/cohere.rs +++ b/src/client/cohere.rs @@ -16,6 +16,8 @@ pub struct CohereConfig { pub api_base: Option, #[serde(default)] pub models: Vec, + #[serde(default)] + pub extend_models: bool, pub patch: Option, pub extra: Option, } diff --git a/src/client/gemini.rs b/src/client/gemini.rs index daab901..57ba9b6 100644 --- a/src/client/gemini.rs +++ b/src/client/gemini.rs @@ -20,6 +20,8 @@ pub struct GeminiConfig { pub oauth: Option>, #[serde(default)] pub models: Vec, + #[serde(default)] + pub extend_models: bool, pub patch: Option, pub extra: Option, } diff --git a/src/client/macros.rs b/src/client/macros.rs index 8e8bbe4..dbfd138 100644 --- a/src/client/macros.rs +++ b/src/client/macros.rs @@ -51,15 +51,24 @@ macro_rules! register_client { pub fn list_models(local_config: &$config) -> Vec { let client_name = Self::name(local_config); + let find_catalog = || $crate::client::ALL_PROVIDER_MODELS.iter().find(|v| { + v.provider == $name || + ($name == OpenAICompatibleClient::NAME + && local_config.name.as_ref().map(|name| name.starts_with(&v.provider)).unwrap_or_default()) + }); if local_config.models.is_empty() { - if let Some(v) = $crate::client::ALL_PROVIDER_MODELS.iter().find(|v| { - v.provider == $name || - ($name == OpenAICompatibleClient::NAME - && local_config.name.as_ref().map(|name| name.starts_with(&v.provider)).unwrap_or_default()) - }) { - return Model::from_config(client_name, &v.models); + match find_catalog() { + Some(v) => Model::from_config(client_name, &v.models), + None => vec![], + } + } else if local_config.extend_models { + match find_catalog() { + Some(v) => Model::from_config( + client_name, + &$crate::client::catalog::extend_catalog_models(&v.models, &local_config.models), + ), + None => Model::from_config(client_name, &local_config.models), } - vec![] } else { Model::from_config(client_name, &local_config.models) } diff --git a/src/client/model.rs b/src/client/model.rs index 0558cb8..ab581cf 100644 --- a/src/client/model.rs +++ b/src/client/model.rs @@ -557,4 +557,96 @@ reasoning_levels: [none, low, medium, high] "oauth blocks in the shipped models.yaml are limited to providers with no built-in stock oauth path (currently only xai's bundled defaults); adding one for claude/openai/gemini would silently swap stock oauth users onto a non-stock provider" ); } + + #[test] + fn config_models_replace_catalog_by_default() { + let ClientConfig::ClaudeConfig(config) = serde_yaml::from_str::( + "type: claude\nmodels:\n - name: test-replace-model\n", + ) + .unwrap() else { + panic!("expected a claude config"); + }; + assert!(!config.extend_models); + + let models = crate::client::ClaudeClient::list_models(&config); + assert_eq!(models.len(), 1); + assert_eq!(models[0].name(), "test-replace-model"); + } + + #[test] + fn config_models_extend_catalog_when_flag_set() { + let ClientConfig::ClaudeConfig(config) = serde_yaml::from_str::( + "type: claude\nextend_models: true\nmodels:\n - name: test-extra-model\n", + ) + .unwrap() else { + panic!("expected a claude config"); + }; + + let catalog_len = crate::client::ALL_PROVIDER_MODELS + .iter() + .find(|v| v.provider == "claude") + .map(|v| v.models.len()) + .unwrap(); + let models = crate::client::ClaudeClient::list_models(&config); + assert_eq!(models.len(), catalog_len + 1); + assert_eq!(models.last().unwrap().name(), "test-extra-model"); + } + + #[test] + fn extend_models_with_empty_models_returns_catalog_exactly() { + let ClientConfig::ClaudeConfig(config) = + serde_yaml::from_str::("type: claude\nextend_models: true\n").unwrap() + else { + panic!("expected a claude config"); + }; + + let catalog = crate::client::ALL_PROVIDER_MODELS + .iter() + .find(|v| v.provider == "claude") + .unwrap(); + let models = crate::client::ClaudeClient::list_models(&config); + assert_eq!(models.len(), catalog.models.len()); + for (model, catalog_model) in models.iter().zip(catalog.models.iter()) { + assert_eq!(model.name(), catalog_model.name); + } + } + + #[test] + fn extend_models_without_catalog_match_returns_config_models_only() { + let ClientConfig::OpenAICompatibleConfig(config) = serde_yaml::from_str::( + "type: openai-compatible\nname: no-such-provider-xyz\nextend_models: true\nmodels:\n - name: test-only-model\n", + ) + .unwrap() else { + panic!("expected an openai-compatible config"); + }; + + assert!( + !crate::client::ALL_PROVIDER_MODELS + .iter() + .any(|v| "no-such-provider-xyz".starts_with(v.provider.as_str())), + "test premise: the client name must not prefix-match any catalog provider" + ); + let models = crate::client::OpenAICompatibleClient::list_models(&config); + assert_eq!(models.len(), 1); + assert_eq!(models[0].name(), "test-only-model"); + } + + #[test] + fn extend_models_extends_catalog_via_provider_name_prefix_match() { + let ClientConfig::OpenAICompatibleConfig(config) = serde_yaml::from_str::( + "type: openai-compatible\nname: openrouter-custom\nextend_models: true\nmodels:\n - name: test-net-new-model\n", + ) + .unwrap() else { + panic!("expected an openai-compatible config"); + }; + + let catalog_len = crate::client::ALL_PROVIDER_MODELS + .iter() + .find(|v| v.provider == "openrouter") + .map(|v| v.models.len()) + .unwrap(); + let models = crate::client::OpenAICompatibleClient::list_models(&config); + assert_eq!(models.len(), catalog_len + 1); + assert_eq!(models.last().unwrap().name(), "test-net-new-model"); + } } diff --git a/src/client/oauth.rs b/src/client/oauth.rs index a03c05d..2b3d02c 100644 --- a/src/client/oauth.rs +++ b/src/client/oauth.rs @@ -1457,13 +1457,9 @@ echo_pkce_in_token_exchange: true ClientConfig::OpenAICompatibleConfig(OpenAICompatibleConfig { name: Some(name.into()), api_base: Some("https://api.example/v1".into()), - api_key: None, auth: auth.map(str::to_string), oauth: oauth.map(Box::new), - wire_api: None, - models: vec![], - patch: None, - extra: None, + ..Default::default() }) } @@ -1474,14 +1470,9 @@ echo_pkce_in_token_exchange: true ) -> ClientConfig { ClientConfig::ClaudeConfig(ClaudeConfig { name: Some(name.into()), - api_key: None, - api_base: None, auth: auth.map(str::to_string), oauth: oauth.map(Box::new), - models: vec![], - prompt_cache: None, - patch: None, - extra: None, + ..Default::default() }) } @@ -1492,15 +1483,9 @@ echo_pkce_in_token_exchange: true ) -> ClientConfig { ClientConfig::OpenAIConfig(OpenAIConfig { name: Some(name.into()), - api_key: None, - api_base: None, - organization_id: None, auth: auth.map(str::to_string), oauth: oauth.map(Box::new), - wire_api: None, - models: vec![], - patch: None, - extra: None, + ..Default::default() }) } @@ -1511,13 +1496,9 @@ echo_pkce_in_token_exchange: true ) -> ClientConfig { ClientConfig::GeminiConfig(GeminiConfig { name: Some(name.into()), - api_key: None, - api_base: None, auth: auth.map(str::to_string), oauth: oauth.map(Box::new), - models: vec![], - patch: None, - extra: None, + ..Default::default() }) } diff --git a/src/client/openai.rs b/src/client/openai.rs index 7439be6..6ca8f62 100644 --- a/src/client/openai.rs +++ b/src/client/openai.rs @@ -24,6 +24,8 @@ pub struct OpenAIConfig { pub wire_api: Option, #[serde(default)] pub models: Vec, + #[serde(default)] + pub extend_models: bool, pub patch: Option, pub extra: Option, } diff --git a/src/client/openai_compatible.rs b/src/client/openai_compatible.rs index 80c7c36..39c2da3 100644 --- a/src/client/openai_compatible.rs +++ b/src/client/openai_compatible.rs @@ -9,7 +9,7 @@ use serde::Deserialize; use serde_json::{Value, json}; use oauth::OAuthConfig; -#[derive(Debug, Clone, Deserialize)] +#[derive(Debug, Clone, Deserialize, Default)] pub struct OpenAICompatibleConfig { pub name: Option, pub api_base: Option, @@ -19,6 +19,8 @@ pub struct OpenAICompatibleConfig { pub wire_api: Option, #[serde(default)] pub models: Vec, + #[serde(default)] + pub extend_models: bool, pub patch: Option, pub extra: Option, } @@ -274,12 +276,8 @@ mod tests { name: Some("compat-test".into()), api_base: Some("https://compat.example/v1".into()), api_key: Some("sk-test".into()), - auth: None, - oauth: None, wire_api, - models: vec![], - patch: None, - extra: None, + ..Default::default() } } diff --git a/src/client/vertexai.rs b/src/client/vertexai.rs index 39eb39d..b36c5a1 100644 --- a/src/client/vertexai.rs +++ b/src/client/vertexai.rs @@ -18,6 +18,8 @@ pub struct VertexAIConfig { pub adc_file: Option, #[serde(default)] pub models: Vec, + #[serde(default)] + pub extend_models: bool, pub prompt_cache: Option, pub patch: Option, pub extra: Option,