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:
13 files changed
+194
-43
No files matched your search
@@ -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
|
||||
|
||||
@@ -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>,
|
||||
}
|
||||
|
||||
@@ -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
@@ -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");
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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>,
|
||||
}
|
||||
|
||||
@@ -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>,
|
||||
}
|
||||
|
||||
+14
-5
@@ -51,15 +51,24 @@ macro_rules! register_client {
|
||||
|
||||
pub fn list_models(local_config: &$config) -> Vec<Model> {
|
||||
let client_name = Self::name(local_config);
|
||||
if local_config.models.is_empty() {
|
||||
if let Some(v) = $crate::client::ALL_PROVIDER_MODELS.iter().find(|v| {
|
||||
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())
|
||||
}) {
|
||||
return Model::from_config(client_name, &v.models);
|
||||
});
|
||||
if local_config.models.is_empty() {
|
||||
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)
|
||||
}
|
||||
|
||||
@@ -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
@@ -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()
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -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>,
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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>,
|
||||
|
||||
Reference in new issue
Block a user