Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 72 additions & 5 deletions src/providers/registry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -419,9 +419,13 @@ impl ProviderRegistry {

for model in models {
if let Some(first_mapping) = model.mappings.first() {
registry
.model_to_provider
.insert(model.name.clone(), first_mapping.provider.clone());
// Canonical key so `provider_for_model` matches spelling-agnostically
// (`gpt-5.5` config vs a `gpt-5-5` query), mirroring `find_model`.
registry.model_to_provider.insert(
crate::routing::classify::model_name::canonicalize_model_name(&model.name)
.into_owned(),
first_mapping.provider.clone(),
);
}
}

Expand All @@ -440,8 +444,10 @@ impl ProviderRegistry {
/// Returns [`ProviderError::ModelNotSupported`] if no registered
/// provider handles the given model name.
pub fn provider_for_model(&self, model: &str) -> Result<Arc<dyn LlmProvider>, ProviderError> {
// First, check if we have a direct model → provider mapping
if let Some(provider_name) = self.model_to_provider.get(model) {
// First, check if we have a direct model → provider mapping. The index is
// keyed by canonical name, so canonicalize the query to match it.
let canonical = crate::routing::classify::model_name::canonicalize_model_name(model);
if let Some(provider_name) = self.model_to_provider.get(canonical.as_ref()) {
if let Some(provider) = self.providers.get(provider_name) {
return Ok(provider.clone());
}
Expand Down Expand Up @@ -607,6 +613,67 @@ mod tests {
assert!(result.is_err());
}

#[test]
fn provider_for_model_is_spelling_agnostic() {
use crate::providers::{AuthType, ProviderConfig};

// A `[[models]]` entry spelled with a dot, as an operator would write it.
let providers = vec![ProviderConfig {
name: "openai".to_string(),
provider_type: "openai".to_string(),
auth_type: AuthType::ApiKey,
api_key: Some(SecretString::new("test-key".to_string())),
base_url: None,
models: vec![],
enabled: Some(true),
oauth_provider: None,
project_id: None,
location: None,
headers: None,
budget_usd: None,
region: None,
pass_through: None,
tls_cert: None,
tls_key: None,
tls_ca: None,
pool: None,
reasoning_effort: None,
service_tier: None,
codex: Default::default(),
circuit_breaker: None,
health_check: None,
max_retries: None,
}];
let models = vec![crate::cli::ModelConfig {
name: "gpt-5.5".to_string(),
mappings: vec![crate::cli::ModelMapping {
priority: 1,
provider: "openai".to_string(),
actual_model: "gpt-5.5".to_string(),
inject_continuation_prompt: false,
}],
budget_usd: None,
context_window_tokens: None,
strategy: Default::default(),
fan_out: None,
deprecated: None,
}];

let registry = ProviderRegistry::from_configs_with_models(
&providers,
&crate::storage::secrets::EnvBackend,
None,
&models,
&TimeoutConfig::default(),
)
.unwrap();

// The dashed form the router canonicalizes to must resolve to the same
// provider as the dotted config name.
assert!(registry.provider_for_model("gpt-5-5").is_ok());
assert!(registry.provider_for_model("gpt-5.5").is_ok());
}

#[test]
fn test_model_counting_with_configs() {
use crate::providers::{AuthType, ProviderConfig};
Expand Down
10 changes: 8 additions & 2 deletions src/routing/classify/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -346,8 +346,14 @@ impl Router {
// name "default-model" to the upstream provider.
if let Some(ref mapper) = self.auto_mapper {
if mapper.is_match(&request.model) {
let has_explicit_virtual =
self.config.models.iter().any(|m| m.name == request.model);
// `request.model` was canonicalized above; compare the config
// names in canonical form too, so a dotted `[[models]]` entry
// still counts as an explicit match and is not auto-mapped away.
let has_explicit_virtual = self
.config
.models
.iter()
.any(|m| model_name::canonicalize_model_name(&m.name) == request.model);
if has_explicit_virtual {
info!(
"Auto-map skipped for '{}' — explicit [[models]] entry takes precedence",
Expand Down
17 changes: 14 additions & 3 deletions src/server/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -102,11 +102,17 @@ impl ReloadableState {
router: Router,
provider_registry: Arc<ProviderRegistry>,
) -> Self {
// Key by the canonical model name so a dotted config entry (`gpt-5.5`)
// matches a request whose name the router already canonicalized to the
// dashed form (`gpt-5-5`) — routing stays agnostic to the spelling, the
// same way it is for text. Canonicalization borrows unchanged for names
// outside a known family, so existing configs are untouched.
use crate::routing::classify::model_name::canonicalize_model_name;
let model_index = config
.models
.iter()
.enumerate()
.map(|(i, m)| (m.name.to_lowercase(), i))
.map(|(i, m)| (canonicalize_model_name(&m.name).to_lowercase(), i))
.collect();
#[cfg(feature = "policies")]
let policy_matcher = init::init_policies(&config);
Expand All @@ -120,10 +126,15 @@ impl ReloadableState {
}
}

/// O(1) model config lookup by name (case-insensitive)
/// O(1) model config lookup by name (case-insensitive, spelling-agnostic).
///
/// The query is canonicalized to match the index keys, so `gpt-5.5` and
/// `gpt-5-5` resolve to the same `[[models]]` entry.
pub fn find_model(&self, name: &str) -> Option<&crate::cli::ModelConfig> {
let key =
crate::routing::classify::model_name::canonicalize_model_name(name).to_lowercase();
self.model_index
.get(&name.to_lowercase())
.get(&key)
.map(|&idx| &self.config.models[idx])
}
}
Expand Down
Loading