diff --git a/crates/protocols/src/models.rs b/crates/protocols/src/models.rs index 664dbc1dde..dff834abf5 100644 --- a/crates/protocols/src/models.rs +++ b/crates/protocols/src/models.rs @@ -10,6 +10,12 @@ use serde_json::Value; use super::{model_card::ModelCard, worker::ProviderType}; +#[derive(Debug, Clone)] +struct UpstreamModelInfo { + id: String, + created: Option, +} + /// A single model entry in the `/v1/models` response. #[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)] pub struct ModelObject { @@ -58,13 +64,209 @@ impl ListModelsResponse { let Some(data) = json.get("data").and_then(|d| d.as_array()) else { return Vec::new(); }; - data.iter() - .filter_map(|m| m.get("id").and_then(|id| id.as_str())) - .map(|id| { - let mut card = ModelCard::new(id); - card.provider.clone_from(&provider); - card + let models: Vec<_> = data + .iter() + .filter_map(|model| { + let id = model.get("id").and_then(|id| id.as_str())?; + let created = model.get("created").and_then(Value::as_u64); + Some(UpstreamModelInfo { + id: id.to_string(), + created, + }) + }) + .collect(); + + group_upstream_models_into_cards(models, provider) + } +} + +fn is_ascii_digits(value: &str, len: usize) -> bool { + value.len() == len && value.bytes().all(|b| b.is_ascii_digit()) +} + +fn strip_date_suffix(id: &str) -> Option { + let parts: Vec<_> = id.split('-').collect(); + + if parts.len() >= 4 + && is_ascii_digits(parts[parts.len() - 3], 4) + && is_ascii_digits(parts[parts.len() - 2], 2) + && is_ascii_digits(parts[parts.len() - 1], 2) + { + return Some(parts[..parts.len() - 3].join("-")); + } + + if parts.len() >= 3 + && is_ascii_digits(parts[parts.len() - 2], 4) + && is_ascii_digits(parts[parts.len() - 1], 2) + { + return Some(parts[..parts.len() - 2].join("-")); + } + + None +} + +fn xai_group_key_and_rank(id: &str) -> Option<(String, u16)> { + let parts: Vec<_> = id.split('-').collect(); + if parts.len() == 3 + && parts[0] == "grok" + && !parts[1].is_empty() + && parts[1].bytes().all(|b| b.is_ascii_digit()) + && is_ascii_digits(parts[2], 4) + { + return parts[2] + .parse::() + .ok() + .map(|rank| (format!("{}-{}", parts[0], parts[1]), rank)); + } + + None +} + +fn alias_group_key(id: &str) -> String { + if let Some((group_key, _)) = xai_group_key_and_rank(id) { + return group_key; + } + + strip_date_suffix(id).unwrap_or_else(|| id.to_string()) +} + +fn xai_revision_rank(id: &str) -> Option { + xai_group_key_and_rank(id).map(|(_, rank)| rank) +} + +fn select_primary_and_aliases( + group_key: &str, + variants: &[UpstreamModelInfo], +) -> (String, Vec) { + let primary_id = if variants.iter().any(|variant| variant.id == group_key) { + group_key.to_string() + } else if variants + .iter() + .all(|variant| xai_revision_rank(&variant.id).is_some()) + { + variants + .iter() + .max_by_key(|variant| { + ( + variant.created.unwrap_or(0), + xai_revision_rank(&variant.id).unwrap_or(0), + variant.id.as_str(), + ) }) - .collect() + .map(|variant| variant.id.clone()) + .unwrap_or_else(|| group_key.to_string()) + } else { + variants + .iter() + .map(|variant| variant.id.as_str()) + .min_by(|a, b| a.len().cmp(&b.len()).then_with(|| a.cmp(b))) + .map(ToOwned::to_owned) + .unwrap_or_else(|| group_key.to_string()) + }; + + let mut aliases: Vec = variants + .iter() + .filter_map(|variant| (variant.id != primary_id).then(|| variant.id.clone())) + .collect(); + + if group_key != primary_id && !aliases.iter().any(|alias| alias == group_key) { + aliases.push(group_key.to_string()); + } + + aliases.sort(); + aliases.dedup(); + + (primary_id, aliases) +} + +fn group_upstream_models_into_cards( + models: Vec, + provider: Option, +) -> Vec { + let mut groups = std::collections::BTreeMap::>::new(); + for model in models { + groups + .entry(alias_group_key(&model.id)) + .or_default() + .push(model); + } + + groups + .into_iter() + .map(|(group_key, variants)| { + let (primary_id, aliases) = select_primary_and_aliases(&group_key, &variants); + let mut card = ModelCard::new(primary_id).with_aliases(aliases); + card.provider.clone_from(&provider); + card + }) + .collect() +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::ListModelsResponse; + use crate::worker::ProviderType; + + #[test] + fn parse_upstream_groups_openai_date_variants_under_stable_name() { + let json = json!({ + "object": "list", + "data": [ + {"id": "gpt-4o", "object": "model"}, + {"id": "gpt-4o-2024-08-06", "object": "model"}, + {"id": "gpt-4o-2024-11-20", "object": "model"} + ] + }); + + let cards = ListModelsResponse::parse_upstream(&json, Some(ProviderType::OpenAI)); + + assert_eq!(cards.len(), 1); + assert_eq!(cards[0].id, "gpt-4o"); + assert!(cards[0] + .aliases + .iter() + .any(|alias| alias == "gpt-4o-2024-08-06")); + assert!(cards[0] + .aliases + .iter() + .any(|alias| alias == "gpt-4o-2024-11-20")); + assert_eq!(cards[0].provider, Some(ProviderType::OpenAI)); + } + + #[test] + fn parse_upstream_adds_xai_family_alias_for_revisioned_model() { + let json = json!({ + "object": "list", + "data": [ + {"id": "grok-4-0709", "object": "model", "created": 100} + ] + }); + + let cards = ListModelsResponse::parse_upstream(&json, Some(ProviderType::XAI)); + + assert_eq!(cards.len(), 1); + assert_eq!(cards[0].id, "grok-4-0709"); + assert!(cards[0].aliases.iter().any(|alias| alias == "grok-4")); + assert_eq!(cards[0].provider, Some(ProviderType::XAI)); + } + + #[test] + fn parse_upstream_picks_latest_xai_revision_as_primary_model() { + let json = json!({ + "object": "list", + "data": [ + {"id": "grok-4-0709", "object": "model", "created": 100}, + {"id": "grok-4-0812", "object": "model", "created": 200} + ] + }); + + let cards = ListModelsResponse::parse_upstream(&json, Some(ProviderType::XAI)); + + assert_eq!(cards.len(), 1); + assert_eq!(cards[0].id, "grok-4-0812"); + assert!(cards[0].aliases.iter().any(|alias| alias == "grok-4")); + assert!(cards[0].aliases.iter().any(|alias| alias == "grok-4-0709")); } } diff --git a/model_gateway/src/routers/openai/chat.rs b/model_gateway/src/routers/openai/chat.rs index 391875d74b..1a3238c21c 100644 --- a/model_gateway/src/routers/openai/chat.rs +++ b/model_gateway/src/routers/openai/chat.rs @@ -16,7 +16,7 @@ use tokio_stream::wrappers::UnboundedReceiverStream; use super::{ context::{ComponentRefs, PayloadState, RequestContext, SharedComponents, WorkerSelection}, provider::ProviderRegistry, - router::resolve_provider, + router::{request_provider, resolve_effective_model, resolve_provider}, }; use crate::{ config::types::RetryConfig, @@ -26,7 +26,7 @@ use crate::{ header_utils::{apply_provider_headers, extract_auth_header}, worker_selection::{SelectWorkerRequest, WorkerSelector}, }, - worker::{is_retryable_status, Endpoint, ProviderType, RetryExecutor, WorkerRegistry}, + worker::{is_retryable_status, Endpoint, RetryExecutor, WorkerRegistry}, }; /// Shared context passed to chat routing functions. @@ -62,7 +62,7 @@ pub(super) async fn route_chat( .select_worker(&SelectWorkerRequest { model_id: model, headers, - provider: Some(ProviderType::OpenAI), + provider: Some(request_provider(model)), ..Default::default() }) .await @@ -80,6 +80,7 @@ pub(super) async fn route_chat( return response; } }; + let effective_model = resolve_effective_model(worker.as_ref(), model); let mut payload = match to_value(body) { Ok(v) => v, @@ -100,9 +101,9 @@ pub(super) async fn route_chat( }; // Patch the serialized payload to use the effective model consistently. - payload["model"] = serde_json::Value::String(model.to_owned()); + payload["model"] = serde_json::Value::String(effective_model.clone()); - let provider = resolve_provider(deps.provider_registry, worker.as_ref(), model); + let provider = resolve_provider(deps.provider_registry, worker.as_ref(), &effective_model); if let Err(e) = provider.transform_request(&mut payload, Endpoint::Chat) { Metrics::record_router_error( metrics_labels::ROUTER_OPENAI, diff --git a/model_gateway/src/routers/openai/responses/route.rs b/model_gateway/src/routers/openai/responses/route.rs index 3d3100fce6..2063565a40 100644 --- a/model_gateway/src/routers/openai/responses/route.rs +++ b/model_gateway/src/routers/openai/responses/route.rs @@ -16,7 +16,7 @@ use super::{ ComponentRefs, PayloadState, RequestContext, ResponsesComponents, WorkerSelection, }, provider::ProviderRegistry, - router::resolve_provider, + router::{request_provider, resolve_effective_model, resolve_provider}, }, handle_non_streaming_response, handle_streaming_response, }; @@ -26,7 +26,7 @@ use crate::{ error, worker_selection::{SelectWorkerRequest, WorkerSelector}, }, - worker::{Endpoint, ProviderType, WorkerRegistry}, + worker::{Endpoint, WorkerRegistry}, }; /// Shared context passed to responses routing functions. @@ -63,7 +63,7 @@ pub(in crate::routers::openai) async fn route_responses( .select_worker(&SelectWorkerRequest { model_id: model, headers, - provider: Some(ProviderType::OpenAI), + provider: Some(request_provider(model)), ..Default::default() }) .await @@ -81,6 +81,7 @@ pub(in crate::routers::openai) async fn route_responses( return response; } }; + let effective_model = resolve_effective_model(worker.as_ref(), model); // Validate mutual exclusivity of conversation and previous_response_id // Treat empty strings as unset to match other metadata paths @@ -105,7 +106,7 @@ pub(in crate::routers::openai) async fn route_responses( } let mut request_body = body.clone(); - request_body.model = model_id.to_string(); + request_body.model = effective_model.clone(); request_body.conversation = None; let original_previous_response_id = match super::history::load_input_history( @@ -143,7 +144,7 @@ pub(in crate::routers::openai) async fn route_responses( } }; - let provider = resolve_provider(deps.provider_registry, worker.as_ref(), model); + let provider = resolve_provider(deps.provider_registry, worker.as_ref(), &effective_model); if let Err(e) = provider.transform_request(&mut payload, Endpoint::Responses) { Metrics::record_router_error( metrics_labels::ROUTER_OPENAI, diff --git a/model_gateway/src/routers/openai/router.rs b/model_gateway/src/routers/openai/router.rs index a344c5d4cc..cd6eab8e62 100644 --- a/model_gateway/src/routers/openai/router.rs +++ b/model_gateway/src/routers/openai/router.rs @@ -36,11 +36,33 @@ use crate::{ /// /// Checks (in order): worker's per-model provider, model name heuristic, /// then falls back to the default provider. +pub(super) fn request_provider(model: &str) -> ProviderType { + ProviderType::from_model_name(model).unwrap_or(ProviderType::OpenAI) +} + +/// Resolve the canonical upstream model ID from the worker's live model view. +pub(super) fn resolve_effective_model(worker: &dyn Worker, model: &str) -> String { + worker + .models() + .into_iter() + .find(|candidate| candidate.matches(model)) + .map(|candidate| candidate.id) + .unwrap_or_else(|| worker.canonical_model_id(model).to_string()) +} + pub(super) fn resolve_provider( registry: &ProviderRegistry, worker: &dyn Worker, model: &str, ) -> Arc { + if let Some(pt) = worker + .models() + .into_iter() + .find(|candidate| candidate.matches(model)) + .and_then(|candidate| candidate.provider) + { + return registry.get_arc(&pt); + } if let Some(pt) = worker.provider_for_model(model) { return registry.get_arc(pt); } @@ -117,7 +139,7 @@ impl OpenAIRouter { .select_worker(&SelectWorkerRequest { model_id, headers, - provider: Some(ProviderType::OpenAI), + provider: Some(request_provider(model_id)), ..Default::default() }) .await diff --git a/model_gateway/src/routers/worker_selection.rs b/model_gateway/src/routers/worker_selection.rs index bd5c7e46f7..6706f57fe1 100644 --- a/model_gateway/src/routers/worker_selection.rs +++ b/model_gateway/src/routers/worker_selection.rs @@ -2,7 +2,7 @@ //! //! Single public API: [`WorkerSelector::select_worker`]. -use std::{sync::Arc, time::Duration}; +use std::{collections::HashSet, sync::Arc, time::Duration}; use axum::{ http::{HeaderMap, HeaderValue}, @@ -114,7 +114,7 @@ impl<'a> WorkerSelector<'a> { let candidates: Vec<_> = workers.into_iter().filter(|w| w.is_available()).collect(); match &req.provider { - Some(provider) => filter_by_provider(candidates, provider), + Some(provider) => filter_by_provider(candidates, provider, req.model_id), None => candidates, } } @@ -137,7 +137,7 @@ impl<'a> WorkerSelector<'a> { true, // healthy only — model exists even if circuit-broken ); let candidates = match &req.provider { - Some(p) => filter_by_provider(workers, p), + Some(p) => filter_by_provider(workers, p, req.model_id), None => workers, }; candidates.iter().any(|w| w.supports_model(req.model_id)) @@ -160,7 +160,7 @@ impl<'a> WorkerSelector<'a> { // Only refresh workers matching the request's provider to avoid sending // e.g. an OpenAI key to Anthropic workers during model discovery. if let Some(p) = provider { - external_workers.retain(|w| matches!(w.default_provider(), Some(wp) if wp == p)); + external_workers.retain(|w| worker_matches_provider(w, p, None)); } if external_workers.is_empty() { @@ -189,29 +189,54 @@ impl<'a> WorkerSelector<'a> { fn filter_by_provider( workers: Vec>, target: &ProviderType, + model_id: &str, ) -> Vec> { - let mut first_provider: Option> = None; - let has_multiple_providers = workers.iter().any(|w| { - let provider = w.default_provider().cloned(); - match first_provider { - None => { - first_provider = Some(provider); - false - } - Some(ref first) => *first != provider, - } - }); + let providers: HashSet<_> = workers + .iter() + .filter_map(|worker| worker_provider_hint(worker)) + .collect(); + let has_multiple_providers = providers.len() > 1; if has_multiple_providers { workers .into_iter() - .filter(|w| matches!(w.default_provider(), Some(p) if p == target)) + .filter(|worker| worker_matches_provider(worker, target, Some(model_id))) .collect() } else { workers } } +fn worker_provider_hint(worker: &Arc) -> Option { + worker + .default_provider() + .cloned() + .or_else(|| worker.models().into_iter().find_map(|model| model.provider)) + .or_else(|| ProviderType::from_url(worker.url())) +} + +fn worker_provider_for_request(worker: &Arc, model_id: &str) -> Option { + worker + .models() + .into_iter() + .find(|model| model.matches(model_id)) + .and_then(|model| model.provider) + .or_else(|| worker.provider_for_model(model_id).cloned()) + .or_else(|| ProviderType::from_model_name(model_id)) + .or_else(|| worker_provider_hint(worker)) +} + +fn worker_matches_provider( + worker: &Arc, + target: &ProviderType, + model_id: Option<&str>, +) -> bool { + model_id + .and_then(|requested_model| worker_provider_for_request(worker, requested_model)) + .or_else(|| worker_provider_hint(worker)) + .is_some_and(|provider| provider == *target) +} + /// Refresh a single worker's model list by calling its `/v1/models` endpoint. /// /// Auth headers are adapted per-vendor via [`apply_provider_headers`] (e.g. diff --git a/model_gateway/src/worker/registry.rs b/model_gateway/src/worker/registry.rs index 914b330814..cab8319c07 100644 --- a/model_gateway/src/worker/registry.rs +++ b/model_gateway/src/worker/registry.rs @@ -298,12 +298,19 @@ impl WorkerRegistry { pub fn worker_model_ids(worker: &Arc) -> Vec { let mut seen = HashSet::new(); - let mut model_ids: Vec = worker - .models() - .into_iter() - .map(|model| model.id) - .filter(|model_id| seen.insert(model_id.clone())) - .collect(); + let mut model_ids = Vec::new(); + + for model in worker.models() { + if seen.insert(model.id.clone()) { + model_ids.push(model.id.clone()); + } + + for alias in model.aliases { + if seen.insert(alias.clone()) { + model_ids.push(alias); + } + } + } if model_ids.is_empty() { model_ids.push(worker.model_id().to_string()); @@ -1419,6 +1426,33 @@ mod tests { assert_eq!(stats.total_models, 2); } + #[test] + fn test_worker_registry_indexes_aliases() { + let registry = WorkerRegistry::new(); + + let worker: Arc = Arc::new( + BasicWorkerBuilder::new("https://cloud9.api.x.ai") + .worker_type(WorkerType::Regular) + .models(vec![ModelCard::new("grok-4-0709") + .with_alias("grok-4") + .with_provider(openai_protocol::worker::ProviderType::XAI)]) + .circuit_breaker_config(CircuitBreakerConfig::default()) + .build(), + ); + + registry.register(worker).unwrap(); + + assert_eq!(registry.get_by_model("grok-4-0709").len(), 1); + assert_eq!(registry.get_by_model("grok-4").len(), 1); + + let mut models = registry.get_models(); + models.sort(); + assert_eq!( + models, + vec!["grok-4".to_string(), "grok-4-0709".to_string()] + ); + } + #[test] fn test_replace_same_url_refreshes_all_model_indexes() { let registry = WorkerRegistry::new(); diff --git a/model_gateway/src/worker/worker.rs b/model_gateway/src/worker/worker.rs index e072cbf569..5db9d60308 100644 --- a/model_gateway/src/worker/worker.rs +++ b/model_gateway/src/worker/worker.rs @@ -1,5 +1,6 @@ use std::{ any::Any, + collections::HashSet, fmt, sync::{ atomic::{AtomicU64, AtomicU8, AtomicUsize, Ordering}, @@ -397,6 +398,14 @@ pub trait Worker: Send + Sync + fmt::Debug + 'static { .and_then(|m| m.chat_template.as_deref()) } + /// Resolve the canonical model ID for a request. + /// + /// If the requested model matches a `ModelCard` alias, this returns the + /// card's primary `id`; otherwise it returns the requested model unchanged. + fn canonical_model_id<'a>(&'a self, model_id: &'a str) -> &'a str { + self.metadata().canonical_model_id(model_id) + } + /// Get the default provider type for this worker. /// `None` means native/passthrough. fn default_provider(&self) -> Option<&ProviderType> { @@ -537,6 +546,13 @@ impl WorkerMetadata { self.spec.models.find(model_id) } + /// Resolve the canonical model ID for a request. + pub fn canonical_model_id<'a>(&'a self, model_id: &'a str) -> &'a str { + self.find_model(model_id) + .map(|model| model.id.as_str()) + .unwrap_or(model_id) + } + /// Check if this worker can serve a given model. /// Wildcard workers accept any model. pub fn supports_model(&self, model_id: &str) -> bool { @@ -686,6 +702,46 @@ impl BasicWorker { self.circuit_breaker.store(other_cb); } } + + fn merge_model_card(primary: &mut ModelCard, fallback: &ModelCard) { + let mut seen_aliases: HashSet = primary.aliases.iter().cloned().collect(); + for alias in &fallback.aliases { + if seen_aliases.insert(alias.clone()) { + primary.aliases.push(alias.clone()); + } + } + + if primary.provider.is_none() { + primary.provider = fallback.provider.clone(); + } + } + + fn effective_models(&self) -> Vec { + let overridden = self.models_override.load(); + let metadata_models = self.metadata.spec.models.all(); + + if overridden.is_wildcard() { + return metadata_models.to_vec(); + } + + if self.metadata.spec.models.is_wildcard() { + return overridden.all().to_vec(); + } + + let mut merged_models = overridden.all().to_vec(); + for metadata_model in metadata_models { + if let Some(existing) = merged_models + .iter_mut() + .find(|model| model.id == metadata_model.id) + { + Self::merge_model_card(existing, metadata_model); + } else { + merged_models.push(metadata_model.clone()); + } + } + + merged_models + } } #[async_trait] @@ -903,20 +959,17 @@ impl Worker for BasicWorker { fn supports_model(&self, model_id: &str) -> bool { let overridden = self.models_override.load(); - if !overridden.is_wildcard() { - return overridden.supports(model_id); + if overridden.is_wildcard() && self.metadata.spec.models.is_wildcard() { + return true; } - self.metadata.supports_model(model_id) + + self.effective_models() + .into_iter() + .any(|model| model.matches(model_id)) } fn models(&self) -> Vec { - let overridden = self.models_override.load(); - let source = if overridden.is_wildcard() { - self.metadata.spec.models.all() - } else { - overridden.all() - }; - source.to_vec() + self.effective_models() } fn set_models(&self, models: Vec) { @@ -1150,7 +1203,8 @@ impl http_body::Body for AttachedBody { /// everything to a routability bool for backwards compatibility. pub fn worker_to_info(worker: &Arc) -> WorkerInfo { let metadata = worker.metadata(); - let spec = metadata.spec.clone(); + let mut spec = metadata.spec.clone(); + spec.models = WorkerModels::from(worker.models()); let status = worker.status(); WorkerInfo { @@ -1884,6 +1938,28 @@ mod tests { assert!(metadata.find_model("unknown-model").is_none()); } + #[test] + fn test_worker_metadata_canonical_model_id() { + use super::ModelCard; + + let model = ModelCard::new("grok-4-0709").with_alias("grok-4"); + + let mut spec = WorkerSpec::new("http://test:8080"); + spec.models = WorkerModels::from(vec![model]); + let metadata = WorkerMetadata { + spec, + health_config: HealthCheckConfig::default(), + health_endpoint: "/health".to_string(), + }; + + assert_eq!(metadata.canonical_model_id("grok-4"), "grok-4-0709"); + assert_eq!(metadata.canonical_model_id("grok-4-0709"), "grok-4-0709"); + assert_eq!( + metadata.canonical_model_id("unknown-model"), + "unknown-model" + ); + } + #[test] fn test_worker_routing_key_load_increment_decrement() { let load = WorkerRoutingKeyLoad::new("http://test:8000"); @@ -2056,4 +2132,22 @@ mod tests { assert!(!worker.supports_model("non-existent-model")); assert!(worker.has_models_discovered()); } + + #[test] + fn test_lazy_discovered_models_preserve_static_aliases() { + let worker = BasicWorkerBuilder::new("http://test:8080") + .model(ModelCard::new("grok-4-0709").with_alias("grok-4")) + .build(); + + worker.set_models(vec![ModelCard::new("grok-4-0709")]); + + assert!(worker.supports_model("grok-4")); + + let grok_model = worker + .models() + .into_iter() + .find(|model| model.id == "grok-4-0709") + .expect("grok-4-0709 should still be present"); + assert_eq!(grok_model.aliases, vec!["grok-4".to_string()]); + } } diff --git a/model_gateway/src/workflow/steps/external/discover_models.rs b/model_gateway/src/workflow/steps/external/discover_models.rs index b57ed2c409..83cb777353 100644 --- a/model_gateway/src/workflow/steps/external/discover_models.rs +++ b/model_gateway/src/workflow/steps/external/discover_models.rs @@ -33,6 +33,14 @@ static HTTP_CLIENT: Lazy = Lazy::new(|| { static DATE_SUFFIX_PATTERN: Lazy = Lazy::new(|| Regex::new(r"-\d{4}-\d{2}(-\d{2})?$").expect("Invalid date regex")); +// xAI Grok revision suffixes commonly look like -MMDD (for example grok-4-0709). +#[expect( + clippy::expect_used, + reason = "Lazy static initialization — compile-time constant regex pattern cannot fail" +)] +static XAI_REVISION_SUFFIX_PATTERN: Lazy = + Lazy::new(|| Regex::new(r"^(grok-\d+)-(\d{4})$").expect("Invalid xAI revision regex")); + /// OpenAI /v1/models response format. #[derive(Debug, Clone, Deserialize)] pub struct ModelsResponse { @@ -53,29 +61,83 @@ pub struct ModelInfo { pub owned_by: Option, } -/// Group models by base name (stripping date suffixes) and create ModelCards with aliases. +fn alias_group_key(id: &str) -> String { + if let Some(captures) = XAI_REVISION_SUFFIX_PATTERN.captures(id) { + if let Some(base) = captures.get(1) { + return base.as_str().to_string(); + } + } + + DATE_SUFFIX_PATTERN.replace(id, "").to_string() +} + +fn xai_revision_rank(id: &str) -> Option { + XAI_REVISION_SUFFIX_PATTERN + .captures(id) + .and_then(|captures| captures.get(2)) + .and_then(|suffix| suffix.as_str().parse::().ok()) +} + +fn select_primary_and_aliases(group_key: &str, variants: Vec) -> (String, Vec) { + let primary_id = if variants.iter().any(|variant| variant.id == group_key) { + group_key.to_string() + } else if variants + .iter() + .all(|variant| xai_revision_rank(&variant.id).is_some()) + { + variants + .iter() + .max_by_key(|variant| { + ( + variant.created.unwrap_or(0), + xai_revision_rank(&variant.id).unwrap_or(0), + variant.id.as_str(), + ) + }) + .map(|variant| variant.id.clone()) + .unwrap_or_else(|| group_key.to_string()) + } else { + variants + .iter() + .map(|variant| variant.id.as_str()) + .min_by(|a, b| a.len().cmp(&b.len()).then_with(|| a.cmp(b))) + .map(ToOwned::to_owned) + .unwrap_or_else(|| group_key.to_string()) + }; + + let mut aliases: Vec = variants + .iter() + .filter_map(|variant| (variant.id != primary_id).then(|| variant.id.clone())) + .collect(); + + if group_key != primary_id && !aliases.iter().any(|alias| alias == group_key) { + aliases.push(group_key.to_string()); + } + + aliases.sort(); + aliases.dedup(); + + (primary_id, aliases) +} + +/// Group models by base name and create ModelCards with aliases. /// /// # Example /// Input: `["gpt-4o", "gpt-4o-2024-05-13", "gpt-4o-2024-08-06", "gpt-4o-2024-11-20"]` /// Output: `ModelCard { id: "gpt-4o", aliases: ["gpt-4o-2024-05-13", "gpt-4o-2024-08-06", "gpt-4o-2024-11-20"] }` pub fn group_models_into_cards(models: Vec) -> Vec { - // Group model IDs by base name (with date stripped) - let mut groups: HashMap> = HashMap::new(); - for model in &models { - let base = DATE_SUFFIX_PATTERN.replace(&model.id, "").to_string(); - groups.entry(base).or_default().push(model.id.clone()); + // Group model IDs by alias family key. + let mut groups: HashMap> = HashMap::new(); + for model in models { + let base = alias_group_key(&model.id); + groups.entry(base).or_default().push(model); } // Create ModelCard for each group groups - .into_values() - .map(|mut variants| { - // Sort: shortest first (base name), then alphabetically - variants.sort_by(|a, b| a.len().cmp(&b.len()).then_with(|| a.cmp(b))); - - let primary_id = variants.remove(0); // shortest = primary ID - let aliases = variants; // rest = aliases - + .into_iter() + .map(|(group_key, variants)| { + let (primary_id, aliases) = select_primary_and_aliases(&group_key, variants); let model_type = infer_model_type_from_id(&primary_id); let provider = infer_provider_from_id(&primary_id); @@ -92,6 +154,64 @@ pub fn group_models_into_cards(models: Vec) -> Vec { .collect() } +#[cfg(test)] +mod tests { + use super::{group_models_into_cards, ModelInfo}; + + fn model(id: &str) -> ModelInfo { + ModelInfo { + id: id.to_string(), + object: "model".to_string(), + created: None, + owned_by: None, + } + } + + #[test] + fn groups_openai_date_variants_under_stable_name() { + let cards = group_models_into_cards(vec![ + model("gpt-4o"), + model("gpt-4o-2024-08-06"), + model("gpt-4o-2024-11-20"), + ]); + + assert_eq!(cards.len(), 1); + assert_eq!(cards[0].id, "gpt-4o"); + assert!(cards[0] + .aliases + .iter() + .any(|alias| alias == "gpt-4o-2024-08-06")); + assert!(cards[0] + .aliases + .iter() + .any(|alias| alias == "gpt-4o-2024-11-20")); + } + + #[test] + fn groups_xai_revisioned_model_under_family_alias() { + let cards = group_models_into_cards(vec![model("grok-4-0709")]); + + assert_eq!(cards.len(), 1); + assert_eq!(cards[0].id, "grok-4-0709"); + assert!(cards[0].aliases.iter().any(|alias| alias == "grok-4")); + } + + #[test] + fn picks_latest_xai_revision_as_primary_model() { + let mut older = model("grok-4-0709"); + older.created = Some(100); + let mut newer = model("grok-4-0812"); + newer.created = Some(200); + + let cards = group_models_into_cards(vec![older, newer]); + + assert_eq!(cards.len(), 1); + assert_eq!(cards[0].id, "grok-4-0812"); + assert!(cards[0].aliases.iter().any(|alias| alias == "grok-4")); + assert!(cards[0].aliases.iter().any(|alias| alias == "grok-4-0709")); + } +} + /// Infer ModelType from model ID string. pub fn infer_model_type_from_id(id: &str) -> ModelType { let id_lower = id.to_lowercase(); diff --git a/model_gateway/tests/routing/test_openai_routing.rs b/model_gateway/tests/routing/test_openai_routing.rs index c3ed94cf92..746fd9a348 100644 --- a/model_gateway/tests/routing/test_openai_routing.rs +++ b/model_gateway/tests/routing/test_openai_routing.rs @@ -21,7 +21,9 @@ use openai_protocol::{ common::StringOrArray, completion::CompletionRequest, generate::GenerateRequest, + model_card::ModelCard, responses::{ResponseInput, ResponsesRequest}, + worker::ProviderType, }; use serde_json::json; use smg::{ @@ -41,7 +43,9 @@ use tower::ServiceExt; use crate::common::{ mock_openai_server::MockOpenAIServer, - test_app::{create_test_app_context, register_external_worker}, + test_app::{ + create_test_app_context, register_external_worker, register_external_worker_with_card, + }, }; /// Helper function to create a minimal chat completion request for testing @@ -689,6 +693,45 @@ async fn test_openai_router_chat_completion_with_mock() { assert!(!chat_response["choices"].as_array().unwrap().is_empty()); } +/// Test that alias requests are rewritten to the canonical discovered model ID +/// before being forwarded upstream for chat completions. +#[tokio::test] +async fn test_openai_router_chat_rewrites_alias_to_canonical_model() { + let mock_server = MockOpenAIServer::new().await; + let base_url = mock_server.base_url(); + + let ctx = create_test_app_context().await; + register_external_worker_with_card( + &ctx, + &base_url, + ModelCard::new("grok-4-0709") + .with_alias("grok-4") + .with_provider(ProviderType::XAI), + ); + let router = OpenAIRouter::new(&ctx).await.unwrap(); + + let val = json!({ + "model": "grok-4", + "messages": [ + {"role": "user", "content": "Hello"} + ], + "max_tokens": 20 + }); + let chat_request: ChatCompletionRequest = serde_json::from_value(val).unwrap(); + + let response = router + .route_chat(None, &chat_request, &chat_request.model) + .await; + assert_eq!(response.status(), StatusCode::OK); + + let (_, body) = response.into_parts(); + let body_bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap(); + let body_str = String::from_utf8(body_bytes.to_vec()).unwrap(); + let chat_response: serde_json::Value = serde_json::from_str(&body_str).unwrap(); + + assert_eq!(chat_response["model"], "grok-4-0709"); +} + /// Test full E2E flow with Axum server #[tokio::test] async fn test_openai_e2e_with_server() { @@ -801,6 +844,77 @@ async fn test_openai_router_chat_streaming_with_mock() { assert!(text.contains("[DONE]")); } +/// Test that alias requests are rewritten to the canonical discovered model ID +/// before being forwarded upstream for responses requests. +#[expect(clippy::disallowed_methods, reason = "test infrastructure")] +#[tokio::test] +async fn test_openai_router_responses_rewrite_alias_to_canonical_model() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + let app = Router::new().route( + "/v1/responses", + post(|Json(request): Json| async move { + let model = request + .get("model") + .and_then(|v| v.as_str()) + .unwrap_or("unknown"); + Json(json!({ + "id": "resp_alias_test", + "object": "response", + "created_at": 1_700_000_000, + "status": "completed", + "model": model, + "output": [{ + "type": "message", + "id": "msg_alias_test", + "role": "assistant", + "status": "completed", + "content": [{ + "type": "output_text", + "text": "ok", + "annotations": [] + }] + }], + "metadata": {} + })) + }), + ); + + let server = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + + let base_url = format!("http://{addr}"); + let ctx = create_test_app_context().await; + register_external_worker_with_card( + &ctx, + &base_url, + ModelCard::new("grok-4-0709") + .with_alias("grok-4") + .with_provider(ProviderType::XAI), + ); + let router = OpenAIRouter::new(&ctx).await.unwrap(); + + let request = ResponsesRequest { + model: "grok-4".to_string(), + input: ResponseInput::Text("Say hi".to_string()), + store: Some(false), + ..Default::default() + }; + + let response = router.route_responses(None, &request, &request.model).await; + assert_eq!(response.status(), StatusCode::OK); + + let body_bytes = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + let body: serde_json::Value = serde_json::from_slice(&body_bytes).unwrap(); + assert_eq!(body["model"], "grok-4-0709"); + + server.abort(); +} + /// Test circuit breaker functionality #[tokio::test] async fn test_openai_router_circuit_breaker() { @@ -836,9 +950,7 @@ async fn test_openai_router_models_from_registry() { BasicWorkerBuilder::new("http://localhost:8000") .worker_type(WorkerType::Regular) .runtime_type(RuntimeType::Sglang) - .models(vec![openai_protocol::model_card::ModelCard::new( - "gpt-3.5-turbo", - )]) + .models(vec![ModelCard::new("gpt-3.5-turbo")]) .health_config(openai_protocol::worker::HealthCheckConfig { disable_health_check: true, ..Default::default()