feat: Support the merging of custom user-defined models on clients configured in the global configuration file with the built in models

This commit is contained in:
Dark-Alex-17 committed 2026-09-21 12:08:53 -06:00
1 parent e237845827
commit a65280754e
13 files changed
+196 -45

No files matched your search

+1
View File
@@ -277,6 +277,7 @@ sync_models_ttl_hours: 24 # Hours before the synced model catalog is cons
# All clients have the following configuration: # All clients have the following configuration:
# - type: xxxx # - type: xxxx
# name: xxxx # Only use it to distinguish clients with the same client type. Optional # 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: # models:
# - name: xxxx # Chat model # - name: xxxx # Chat model
# max_input_tokens: 100000 # max_input_tokens: 100000
+3 -1
View File
@@ -4,13 +4,15 @@ use super::*;
use anyhow::Result; use anyhow::Result;
use serde::Deserialize; use serde::Deserialize;
#[derive(Debug, Clone, Deserialize)] #[derive(Debug, Clone, Deserialize, Default)]
pub struct AzureOpenAIConfig { pub struct AzureOpenAIConfig {
pub name: Option<String>, pub name: Option<String>,
pub api_base: Option<String>, pub api_base: Option<String>,
pub api_key: Option<String>, pub api_key: Option<String>,
#[serde(default)] #[serde(default)]
pub models: Vec<ModelData>, pub models: Vec<ModelData>,
#[serde(default)]
pub extend_models: bool,
pub patch: Option<RequestPatch>, pub patch: Option<RequestPatch>,
pub extra: Option<ExtraConfig>, pub extra: Option<ExtraConfig>,
} }
+3 -1
View File
@@ -16,7 +16,7 @@ use reqwest::{Client as ReqwestClient, Method, RequestBuilder};
use serde::Deserialize; use serde::Deserialize;
use serde_json::{Value, json}; use serde_json::{Value, json};
#[derive(Debug, Clone, Deserialize)] #[derive(Debug, Clone, Deserialize, Default)]
pub struct BedrockConfig { pub struct BedrockConfig {
pub name: Option<String>, pub name: Option<String>,
pub access_key_id: Option<String>, pub access_key_id: Option<String>,
@@ -25,6 +25,8 @@ pub struct BedrockConfig {
pub session_token: Option<String>, pub session_token: Option<String>,
#[serde(default)] #[serde(default)]
pub models: Vec<ModelData>, pub models: Vec<ModelData>,
#[serde(default)]
pub extend_models: bool,
pub prompt_cache: Option<bool>, pub prompt_cache: Option<bool>,
pub patch: Option<RequestPatch>, pub patch: Option<RequestPatch>,
pub extra: Option<ExtraConfig>, pub extra: Option<ExtraConfig>,
+61 -1
View File
@@ -572,7 +572,7 @@ pub(crate) fn name_overlap(old: &[String], new: &[String]) -> f64 {
kept as f64 / old.len() as 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 { ModelData {
name: hand.name, name: hand.name,
model_type: hand.model_type, 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<ModelData> {
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)] #[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)] #[serde(deny_unknown_fields)]
pub struct ProviderQuirks { 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");
}
} }
+4 -6
View File
@@ -15,7 +15,7 @@ use serde_json::{Value, json};
const API_BASE: &str = "https://api.anthropic.com/v1"; const API_BASE: &str = "https://api.anthropic.com/v1";
const CLAUDE_CODE_PREFIX: &str = "You are Claude Code, Anthropic's official CLI for Claude."; 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 struct ClaudeConfig {
pub name: Option<String>, pub name: Option<String>,
pub api_key: Option<String>, pub api_key: Option<String>,
@@ -24,6 +24,8 @@ pub struct ClaudeConfig {
pub oauth: Option<Box<OAuthConfig>>, pub oauth: Option<Box<OAuthConfig>>,
#[serde(default)] #[serde(default)]
pub models: Vec<ModelData>, pub models: Vec<ModelData>,
#[serde(default)]
pub extend_models: bool,
pub prompt_cache: Option<bool>, pub prompt_cache: Option<bool>,
pub patch: Option<RequestPatch>, pub patch: Option<RequestPatch>,
pub extra: Option<ExtraConfig>, pub extra: Option<ExtraConfig>,
@@ -1122,16 +1124,12 @@ mod tests {
fn claude_config(name: &str, auth: Option<&str>, oauth: Option<OAuthConfig>) -> ClaudeConfig { fn claude_config(name: &str, auth: Option<&str>, oauth: Option<OAuthConfig>) -> ClaudeConfig {
ClaudeConfig { ClaudeConfig {
name: Some(name.into()), name: Some(name.into()),
api_key: None,
api_base: oauth api_base: oauth
.as_ref() .as_ref()
.map(|_| "https://gateway.example/v1".to_string()), .map(|_| "https://gateway.example/v1".to_string()),
auth: auth.map(str::to_string), auth: auth.map(str::to_string),
oauth: oauth.map(Box::new), oauth: oauth.map(Box::new),
models: vec![], ..Default::default()
prompt_cache: None,
patch: None,
extra: None,
} }
} }
+2
View File
@@ -16,6 +16,8 @@ pub struct CohereConfig {
pub api_base: Option<String>, pub api_base: Option<String>,
#[serde(default)] #[serde(default)]
pub models: Vec<ModelData>, pub models: Vec<ModelData>,
#[serde(default)]
pub extend_models: bool,
pub patch: Option<RequestPatch>, pub patch: Option<RequestPatch>,
pub extra: Option<ExtraConfig>, pub extra: Option<ExtraConfig>,
} }
+2
View File
@@ -20,6 +20,8 @@ pub struct GeminiConfig {
pub oauth: Option<Box<OAuthConfig>>, pub oauth: Option<Box<OAuthConfig>>,
#[serde(default)] #[serde(default)]
pub models: Vec<ModelData>, pub models: Vec<ModelData>,
#[serde(default)]
pub extend_models: bool,
pub patch: Option<RequestPatch>, pub patch: Option<RequestPatch>,
pub extra: Option<ExtraConfig>, pub extra: Option<ExtraConfig>,
} }
+16 -7
View File
@@ -51,15 +51,24 @@ macro_rules! register_client {
pub fn list_models(local_config: &$config) -> Vec<Model> { pub fn list_models(local_config: &$config) -> Vec<Model> {
let client_name = Self::name(local_config); 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 local_config.models.is_empty() {
if let Some(v) = $crate::client::ALL_PROVIDER_MODELS.iter().find(|v| { match find_catalog() {
v.provider == $name || Some(v) => Model::from_config(client_name, &v.models),
($name == OpenAICompatibleClient::NAME None => vec![],
&& local_config.name.as_ref().map(|name| name.starts_with(&v.provider)).unwrap_or_default()) }
}) { } else if local_config.extend_models {
return Model::from_config(client_name, &v.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 { } else {
Model::from_config(client_name, &local_config.models) Model::from_config(client_name, &local_config.models)
} }
+92
View File
@@ -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" "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::<ClientConfig>(
"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::<ClientConfig>(
"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::<ClientConfig>("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::<ClientConfig>(
"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::<ClientConfig>(
"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");
}
} }
+4 -23
View File
@@ -1457,13 +1457,9 @@ echo_pkce_in_token_exchange: true
ClientConfig::OpenAICompatibleConfig(OpenAICompatibleConfig { ClientConfig::OpenAICompatibleConfig(OpenAICompatibleConfig {
name: Some(name.into()), name: Some(name.into()),
api_base: Some("https://api.example/v1".into()), api_base: Some("https://api.example/v1".into()),
api_key: None,
auth: auth.map(str::to_string), auth: auth.map(str::to_string),
oauth: oauth.map(Box::new), oauth: oauth.map(Box::new),
wire_api: None, ..Default::default()
models: vec![],
patch: None,
extra: None,
}) })
} }
@@ -1474,14 +1470,9 @@ echo_pkce_in_token_exchange: true
) -> ClientConfig { ) -> ClientConfig {
ClientConfig::ClaudeConfig(ClaudeConfig { ClientConfig::ClaudeConfig(ClaudeConfig {
name: Some(name.into()), name: Some(name.into()),
api_key: None,
api_base: None,
auth: auth.map(str::to_string), auth: auth.map(str::to_string),
oauth: oauth.map(Box::new), oauth: oauth.map(Box::new),
models: vec![], ..Default::default()
prompt_cache: None,
patch: None,
extra: None,
}) })
} }
@@ -1492,15 +1483,9 @@ echo_pkce_in_token_exchange: true
) -> ClientConfig { ) -> ClientConfig {
ClientConfig::OpenAIConfig(OpenAIConfig { ClientConfig::OpenAIConfig(OpenAIConfig {
name: Some(name.into()), name: Some(name.into()),
api_key: None,
api_base: None,
organization_id: None,
auth: auth.map(str::to_string), auth: auth.map(str::to_string),
oauth: oauth.map(Box::new), oauth: oauth.map(Box::new),
wire_api: None, ..Default::default()
models: vec![],
patch: None,
extra: None,
}) })
} }
@@ -1511,13 +1496,9 @@ echo_pkce_in_token_exchange: true
) -> ClientConfig { ) -> ClientConfig {
ClientConfig::GeminiConfig(GeminiConfig { ClientConfig::GeminiConfig(GeminiConfig {
name: Some(name.into()), name: Some(name.into()),
api_key: None,
api_base: None,
auth: auth.map(str::to_string), auth: auth.map(str::to_string),
oauth: oauth.map(Box::new), oauth: oauth.map(Box::new),
models: vec![], ..Default::default()
patch: None,
extra: None,
}) })
} }
+2
View File
@@ -24,6 +24,8 @@ pub struct OpenAIConfig {
pub wire_api: Option<WireApi>, pub wire_api: Option<WireApi>,
#[serde(default)] #[serde(default)]
pub models: Vec<ModelData>, pub models: Vec<ModelData>,
#[serde(default)]
pub extend_models: bool,
pub patch: Option<RequestPatch>, pub patch: Option<RequestPatch>,
pub extra: Option<ExtraConfig>, pub extra: Option<ExtraConfig>,
} }
+4 -6
View File
@@ -9,7 +9,7 @@ use serde::Deserialize;
use serde_json::{Value, json}; use serde_json::{Value, json};
use oauth::OAuthConfig; use oauth::OAuthConfig;
#[derive(Debug, Clone, Deserialize)] #[derive(Debug, Clone, Deserialize, Default)]
pub struct OpenAICompatibleConfig { pub struct OpenAICompatibleConfig {
pub name: Option<String>, pub name: Option<String>,
pub api_base: Option<String>, pub api_base: Option<String>,
@@ -19,6 +19,8 @@ pub struct OpenAICompatibleConfig {
pub wire_api: Option<WireApi>, pub wire_api: Option<WireApi>,
#[serde(default)] #[serde(default)]
pub models: Vec<ModelData>, pub models: Vec<ModelData>,
#[serde(default)]
pub extend_models: bool,
pub patch: Option<RequestPatch>, pub patch: Option<RequestPatch>,
pub extra: Option<ExtraConfig>, pub extra: Option<ExtraConfig>,
} }
@@ -274,12 +276,8 @@ mod tests {
name: Some("compat-test".into()), name: Some("compat-test".into()),
api_base: Some("https://compat.example/v1".into()), api_base: Some("https://compat.example/v1".into()),
api_key: Some("sk-test".into()), api_key: Some("sk-test".into()),
auth: None,
oauth: None,
wire_api, wire_api,
models: vec![], ..Default::default()
patch: None,
extra: None,
} }
} }
+2
View File
@@ -18,6 +18,8 @@ pub struct VertexAIConfig {
pub adc_file: Option<String>, pub adc_file: Option<String>,
#[serde(default)] #[serde(default)]
pub models: Vec<ModelData>, pub models: Vec<ModelData>,
#[serde(default)]
pub extend_models: bool,
pub prompt_cache: Option<bool>, pub prompt_cache: Option<bool>,
pub patch: Option<RequestPatch>, pub patch: Option<RequestPatch>,
pub extra: Option<ExtraConfig>, pub extra: Option<ExtraConfig>,