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

+3 -1
View File
@@ -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<String>,
pub api_base: Option<String>,
pub api_key: Option<String>,
#[serde(default)]
pub models: Vec<ModelData>,
#[serde(default)]
pub extend_models: bool,
pub patch: Option<RequestPatch>,
pub extra: Option<ExtraConfig>,
}
+3 -1
View File
@@ -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<String>,
pub access_key_id: Option<String>,
@@ -25,6 +25,8 @@ pub struct BedrockConfig {
pub session_token: Option<String>,
#[serde(default)]
pub models: Vec<ModelData>,
#[serde(default)]
pub extend_models: bool,
pub prompt_cache: Option<bool>,
pub patch: Option<RequestPatch>,
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
}
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<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)]
#[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");
}
}
+4 -6
View File
@@ -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<String>,
pub api_key: Option<String>,
@@ -24,6 +24,8 @@ pub struct ClaudeConfig {
pub oauth: Option<Box<OAuthConfig>>,
#[serde(default)]
pub models: Vec<ModelData>,
#[serde(default)]
pub extend_models: bool,
pub prompt_cache: Option<bool>,
pub patch: Option<RequestPatch>,
pub extra: Option<ExtraConfig>,
@@ -1122,16 +1124,12 @@ mod tests {
fn claude_config(name: &str, auth: Option<&str>, oauth: Option<OAuthConfig>) -> 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()
}
}
+2
View File
@@ -16,6 +16,8 @@ pub struct CohereConfig {
pub api_base: Option<String>,
#[serde(default)]
pub models: Vec<ModelData>,
#[serde(default)]
pub extend_models: bool,
pub patch: Option<RequestPatch>,
pub extra: Option<ExtraConfig>,
}
+2
View File
@@ -20,6 +20,8 @@ pub struct GeminiConfig {
pub oauth: Option<Box<OAuthConfig>>,
#[serde(default)]
pub models: Vec<ModelData>,
#[serde(default)]
pub extend_models: bool,
pub patch: Option<RequestPatch>,
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> {
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)
}
+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"
);
}
#[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 {
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()
})
}
+2
View File
@@ -24,6 +24,8 @@ pub struct OpenAIConfig {
pub wire_api: Option<WireApi>,
#[serde(default)]
pub models: Vec<ModelData>,
#[serde(default)]
pub extend_models: bool,
pub patch: Option<RequestPatch>,
pub extra: Option<ExtraConfig>,
}
+4 -6
View File
@@ -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<String>,
pub api_base: Option<String>,
@@ -19,6 +19,8 @@ pub struct OpenAICompatibleConfig {
pub wire_api: Option<WireApi>,
#[serde(default)]
pub models: Vec<ModelData>,
#[serde(default)]
pub extend_models: bool,
pub patch: Option<RequestPatch>,
pub extra: Option<ExtraConfig>,
}
@@ -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()
}
}
+2
View File
@@ -18,6 +18,8 @@ pub struct VertexAIConfig {
pub adc_file: Option<String>,
#[serde(default)]
pub models: Vec<ModelData>,
#[serde(default)]
pub extend_models: bool,
pub prompt_cache: Option<bool>,
pub patch: Option<RequestPatch>,
pub extra: Option<ExtraConfig>,