From 42e31fcc085d015de62bd224af56f3ba8427d212 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 7 Nov 2026 06:35:36 +0000 Subject: [PATCH] feat(cost-map): read alias_of pointer rows in the catalog and the loader Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/model-catalog/src/catalog.rs | 80 ++++++++++++++++--- .../crates/model-catalog/src/error.rs | 10 +++ .../model-catalog/src/wire/attributes.rs | 2 + .../crates/model-catalog/tests/catalog.rs | 51 ++++++++++++ .../litellm_core_utils/get_model_cost_map.py | 26 ++++-- .../test_get_model_cost_map.py | 31 +++++++ tests/unit/test_utils.py | 2 + 7 files changed, 184 insertions(+), 18 deletions(-) diff --git a/litellm-rust/crates/model-catalog/src/catalog.rs b/litellm-rust/crates/model-catalog/src/catalog.rs index 98af785b17e..7cffd26d61f 100644 --- a/litellm-rust/crates/model-catalog/src/catalog.rs +++ b/litellm-rust/crates/model-catalog/src/catalog.rs @@ -142,6 +142,7 @@ impl Catalog { let mut sections = Vec::new(); let mut entries = IndexMap::with_capacity(document.len()); let mut alias_lists = Vec::new(); + let mut pointer_rows = Vec::new(); for (index, (name, value)) in document.into_iter().enumerate() { let fields = match value { Value::Object(fields) if !is_section(&name) => fields, @@ -157,22 +158,44 @@ impl Catalog { continue; } }; - let (model, aliases, row_problems) = read_row(&fields); + let row = read_row(&fields); problems.extend( - row_problems + row.problems .into_iter() .map(|problem| problem.prepend(&name)), ); - if let Some(aliases) = aliases { + if let Some(target) = row.alias_of { + pointer_rows.push((name.clone(), target)); + sections.push(Section { + index, + key: name, + value: Value::Object(fields), + }); + continue; + } + if let Some(aliases) = row.aliases { alias_lists.push((name.clone(), aliases)); } - entries.insert(name, ModelEntry { fields, model }); + entries.insert( + name, + ModelEntry { + fields, + model: row.model, + }, + ); } if entries.is_empty() { return Err(Error::new("", Repr::Empty)); } let mut aliases: IndexMap = IndexMap::new(); + for (alias, target) in pointer_rows { + if entries.contains_key(&target) { + aliases.insert(alias, target); + } else { + problems.push(Error::from_key(&alias, Repr::AliasTargetMissing { target })); + } + } for (model, names) in alias_lists { for alias in names { if entries.contains_key(&alias) { @@ -242,8 +265,8 @@ impl Catalog { } pub fn row_errors(model: &str, fields: &Fields) -> Vec { - let (_, _, errors) = read_row(fields); - errors + read_row(fields) + .problems .into_iter() .map(|error| error.prepend(model)) .collect() @@ -255,7 +278,13 @@ impl Catalog { (0..self.rows().entries.len() + self.sections.len()).map_while(move |index| match sections .next_if(|section| section.index == index) { - Some(section) => Some((section.key.as_str(), Written::Section(§ion.value))), + Some(section) => Some(( + section.key.as_str(), + match §ion.value { + Value::Object(fields) if !is_section(§ion.key) => Written::Entry(fields), + value => Written::Section(value), + }, + )), None => entries .next() .map(|(name, entry)| (name.as_str(), Written::Entry(&entry.fields))), @@ -333,7 +362,8 @@ impl Catalog { let mut catalog = self; catalog.store.edit(|rows| { rows.entries.retain(|name, _| !rejected.contains(name)); - rows.aliases.retain(|_, owner| !rejected.contains(owner)); + rows.aliases + .retain(|alias, owner| !rejected.contains(owner) && !rejected.contains(alias)); }); catalog .sections @@ -406,16 +436,42 @@ impl Catalog { } } -fn read_row(fields: &Fields) -> (Option, Option>, Vec) { +/// One row as read: a model, or an `alias_of` pointer at one +struct RowRead { + model: Option, + aliases: Option>, + alias_of: Option, + problems: Vec, +} + +fn read_row(fields: &Fields) -> RowRead { let wire = EntryWire::read(fields); + if let Some(target) = wire.attributes.alias_of.clone() { + let other_fields = fields + .keys() + .filter(|key| *key != "alias_of") + .map(|key| Error::from_key(key, Repr::AliasRowField)); + return RowRead { + model: None, + aliases: None, + alias_of: Some(target), + problems: wire.problems.into_iter().chain(other_fields).collect(), + }; + } let aliases = wire.attributes.aliases.clone(); let timezone_error = unknown_timezone(&wire); - match Model::try_from(wire) { - Ok(model) => (Some(model), aliases, timezone_error.into_iter().collect()), + let (model, problems) = match Model::try_from(wire) { + Ok(model) => (Some(model), timezone_error.into_iter().collect()), Err(mut errors) => { errors.extend(timezone_error); - (None, aliases, errors) + (None, errors) } + }; + RowRead { + model, + aliases, + alias_of: None, + problems, } } diff --git a/litellm-rust/crates/model-catalog/src/error.rs b/litellm-rust/crates/model-catalog/src/error.rs index 89e27c05f46..846398f7155 100644 --- a/litellm-rust/crates/model-catalog/src/error.rs +++ b/litellm-rust/crates/model-catalog/src/error.rs @@ -67,6 +67,12 @@ pub(crate) enum Repr { #[error("alias {alias:?} already belongs to {taken_by}")] #[strum_discriminants(strum(serialize = "alias_taken"))] AliasTaken { alias: String, taken_by: String }, + #[error("alias_of {target:?} is not a model name")] + #[strum_discriminants(strum(serialize = "alias_target_missing"))] + AliasTargetMissing { target: String }, + #[error("an alias_of row carries no other field")] + #[strum_discriminants(strum(serialize = "alias_row_field"))] + AliasRowField, #[error("{first} and {second} price the same meter under the same conditions")] #[strum_discriminants(strum(serialize = "duplicate_price"))] DuplicatePrice { first: String, second: String }, @@ -141,6 +147,8 @@ impl Error { | Repr::UnknownTimezone { .. } | Repr::AliasShadowsModel { .. } | Repr::AliasTaken { .. } + | Repr::AliasTargetMissing { .. } + | Repr::AliasRowField | Repr::DuplicatePrice { .. } | Repr::SuffixInsideBlock { .. } | Repr::TierRange @@ -327,6 +335,8 @@ mod tests { #[case::unknown_timezone(Repr::UnknownTimezone { zone: "invalid".into() }, Rejects::Row)] #[case::alias_shadows_model(Repr::AliasShadowsModel { alias: "model".into() }, Rejects::Row)] #[case::alias_taken(Repr::AliasTaken { alias: "alias".into(), taken_by: "owner".into() }, Rejects::Row)] + #[case::alias_target_missing(Repr::AliasTargetMissing { target: "owner".into() }, Rejects::Row)] + #[case::alias_row_field(Repr::AliasRowField, Rejects::Row)] #[case::duplicate_price(Repr::DuplicatePrice { first: "a".into(), second: "b".into() }, Rejects::Row)] #[case::suffix_inside_block(Repr::SuffixInsideBlock { block: "tiered_pricing" }, Rejects::Row)] #[case::tier_range(Repr::TierRange, Rejects::Row)] diff --git a/litellm-rust/crates/model-catalog/src/wire/attributes.rs b/litellm-rust/crates/model-catalog/src/wire/attributes.rs index 401634f8c62..52cd2e8525f 100644 --- a/litellm-rust/crates/model-catalog/src/wire/attributes.rs +++ b/litellm-rust/crates/model-catalog/src/wire/attributes.rs @@ -47,6 +47,8 @@ macro_rules! attributes { } attributes! { + /// Model this row is another name for; such a row carries no other key. + alias_of: String, /// Other names this entry is served under. aliases: Vec, audio_transcription_config: String, diff --git a/litellm-rust/crates/model-catalog/tests/catalog.rs b/litellm-rust/crates/model-catalog/tests/catalog.rs index 31f907689d3..3454b1f9a52 100644 --- a/litellm-rust/crates/model-catalog/tests/catalog.rs +++ b/litellm-rust/crates/model-catalog/tests/catalog.rs @@ -139,6 +139,18 @@ fn a_bad_entry_is_reported_by_model_and_key_and_kept_raw() { json!({"First": {"litellm_provider": "t", "aliases": ["shared"]}, "Second": {"litellm_provider": "t", "aliases": ["shared"]}}), r#"Second: alias "shared" already belongs to First"# )] +#[case::alias_of_target_missing( + json!({"Real": {"litellm_provider": "t"}, "ghost-latest": {"alias_of": "Ghost"}}), + r#"ghost-latest: alias_of "Ghost" is not a model name"# +)] +#[case::alias_of_row_with_other_fields( + json!({"Real": {"litellm_provider": "t"}, "real-latest": {"alias_of": "Real", "mode": "chat"}}), + "real-latest.mode: an alias_of row carries no other field" +)] +#[case::alias_of_claimed_by_a_list( + json!({"Real": {"litellm_provider": "t"}, "real-latest": {"alias_of": "Real"}, "Other": {"litellm_provider": "t", "aliases": ["real-latest"]}}), + r#"Other: alias "real-latest" already belongs to Real"# +)] #[case::aliases_not_a_list( json!({"a": {"litellm_provider": "t", "aliases": "bad"}}), "a.aliases: invalid value: invalid type: string \"bad\", expected a sequence" @@ -399,3 +411,42 @@ fn a_bad_value_keeps_its_json_error_as_the_source_and_names_its_key_and_field() assert_eq!(error.key().as_deref(), Some("open/router")); assert_eq!(error.field().as_deref(), Some("input_cost_per_token")); } + +#[rstest] +fn an_alias_of_row_is_served_as_its_target_and_written_as_itself() { + let catalog = catalog(json!({ + "Real": {"litellm_provider": "t", "input_cost_per_token": 1e-6}, + "real-latest": {"alias_of": "Real"}, + "After": {"litellm_provider": "t"} + })); + let found = catalog.lookup("REAL-LATEST").unwrap(); + assert_eq!( + (found.matched_key, found.canonical_key), + ("real-latest", "Real") + ); + assert_eq!( + found.entry.field("input_cost_per_token"), + Some(&json!(1e-6)) + ); + assert_eq!(catalog.model_count(), 2); + assert_eq!( + catalog.served().map(|(key, _)| key).collect::>(), + ["Real", "After", "real-latest"] + ); + assert_eq!( + written(&catalog), + r#"{"Real":{"litellm_provider":"t","input_cost_per_token":1e-6},"real-latest":{"alias_of":"Real"},"After":{"litellm_provider":"t"}}"# + ); +} + +#[rstest] +fn a_dangling_alias_of_row_is_rejected_and_dropped_from_the_written_document() { + let parsed = loaded(json!({ + "Real": {"litellm_provider": "t"}, + "ghost-latest": {"alias_of": "Ghost"} + })); + let (catalog, removed) = parsed.value.without_rejected(&parsed.errors); + assert_eq!(removed.len(), 1); + assert!(catalog.lookup("ghost-latest").is_none()); + assert_eq!(written(&catalog), r#"{"Real":{"litellm_provider":"t"}}"#); +} diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index e452befe295..b0df3c9bf38 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -580,18 +580,30 @@ def _alias_groups(alias_owners: Mapping[str, str]) -> Mapping[str, tuple[str, .. def _expand_model_aliases(model_cost: dict[str, dict[str, object]]) -> dict[str, dict[str, object]]: """ - Expand ``aliases`` lists in model cost entries into top-level entries. + Expand alias rows into top-level entries that share their owner's ``AliasedCostEntry``. - Each alias gets a reference to the **same** ``AliasedCostEntry`` as the - canonical entry, which carries every name it is served under. The - ``aliases`` key is removed from the entry so downstream code never sees it. + A row whose ``alias_of`` names another row is served as that row, and an + ``aliases`` list on a row adds further names for it. Every name in a group + references the **same** ``AliasedCostEntry``, which carries every name it is + served under, and neither key survives into the served entries. - If an alias collides with an existing canonical entry the alias is - skipped and a warning is logged. + An alias that collides with a canonical entry or another alias, or a row whose + ``alias_of`` names nothing, is skipped and a warning is logged. """ + pointer_rows: Final = { + name: target for name, info in model_cost.items() if isinstance(target := info.get("alias_of"), str) + } alias_owners: Final[dict[str, str]] = {} keys_with_aliases: Final[list[str]] = [] + for alias, target in pointer_rows.items(): + if target not in model_cost or target in pointer_rows: + verbose_logger.warning( + "LiteLLM model alias '%s' points at '%s', which is not a model entry, dropping it.", alias, target + ) + continue + alias_owners[alias] = target + for model_name, model_info in model_cost.items(): aliases: object = model_info.get("aliases") if aliases is None: @@ -628,6 +640,8 @@ def _expand_model_aliases(model_cost: dict[str, dict[str, object]]) -> dict[str, # Remove the ``aliases`` key from entries so it doesn't pollute model info for key in keys_with_aliases: model_cost[key].pop("aliases", None) + for alias in pointer_rows.keys() - alias_owners.keys(): + model_cost.pop(alias) for owner, names in _alias_groups(alias_owners).items(): model_cost.update(dict.fromkeys(names, AliasedCostEntry(model_cost[owner], names))) diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index d44169ae28a..509a777d137 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -130,6 +130,22 @@ def test_finalize_pops_key_and_installs_rules(): set_fallback_generalizations(previous) +def test_alias_of_rows_are_served_as_their_target_and_dangling_ones_are_dropped(caplog) -> None: + raw: Final = { + "real": {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 1e-6}, + "real-latest": {"alias_of": "real"}, + "ghost-latest": {"alias_of": "ghost"}, + "chain-latest": {"alias_of": "real-latest"}, + } + with caplog.at_level("WARNING"): + finalized: Final = _finalize_model_cost_map(raw) + assert finalized["real-latest"] == finalized["real"] + assert "alias_of" not in finalized["real"] + assert set(finalized) == {"real", "real-latest"} + assert finalized["real-latest"]["litellm_provider"] == "openai" + assert {"ghost-latest", "chain-latest"} <= {record.getMessage().split("'")[1] for record in caplog.records} + + def test_finalize_with_no_block_clears_rules(): previous = list(get_fallback_generalization_rules()) try: @@ -806,6 +822,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="session") def event_loop(): try: @@ -815,6 +832,7 @@ def event_loop(): yield loop loop.close() + @pytest.fixture(scope="function") def setup_and_teardown(event_loop): import litellm @@ -848,6 +866,7 @@ def setup_and_teardown(event_loop): if pending: event_loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "set_verbose": getattr(litellm, "set_verbose", False), @@ -862,6 +881,7 @@ _SCALAR_DEFAULTS = { "cohere_key": getattr(litellm, "cohere_key", None), } + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestCheckIsValidDict: """Unit tests for _check_is_valid_dict.""" @@ -886,6 +906,7 @@ class TestCheckIsValidDict: """Non-empty dict should pass.""" assert GetModelCostMap._check_is_valid_dict({"model": {}}) is True + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestCheckModelCountNotReduced: """Unit tests for _check_model_count_not_reduced.""" @@ -940,6 +961,7 @@ class TestCheckModelCountNotReduced: is True ) + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestValidateModelCostMap: """Unit tests for validate_model_cost_map (combines both checks).""" @@ -976,6 +998,7 @@ class TestValidateModelCostMap: is True ) + def _fetch_with_single_outcome(monkeypatch, outcome): from litellm.litellm_core_utils import get_model_cost_map as module @@ -1051,6 +1074,7 @@ class TestGetModelCostMapFallback: assert isinstance(result, dict) assert len(result) > 0 + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestBackupModelCostMapExists: """Validates the local backup file is always present and valid.""" @@ -1066,6 +1090,7 @@ class TestBackupModelCostMapExists: backup = GetModelCostMap.load_local_model_cost_map() assert len(backup) > 100, f"Backup has only {len(backup)} models, expected > 100" + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestBadHostedModelCostMap: """ @@ -1130,12 +1155,14 @@ class TestBadHostedModelCostMap: finally: litellm.model_cost = original + @pytest.fixture() def _vcr_outcome_gate_local_testing(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -1188,6 +1215,7 @@ def isolate_litellm_state(): setattr(litellm, attr, original_value) _invalidate_model_cost_lowercase_map() + _SCALAR_DEFAULTS_local_testing = { "num_retries": getattr(litellm, "num_retries", None), "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), @@ -1208,6 +1236,7 @@ _SCALAR_DEFAULTS_local_testing = { "api_key": getattr(litellm, "api_key", None), } + @pytest.fixture(scope="module") def setup_and_teardown_local_testing(): """ @@ -1230,6 +1259,7 @@ def setup_and_teardown_local_testing(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -1241,6 +1271,7 @@ def test_get_model_cost_map(): except Exception as e: pytest.fail(f"An exception occurred: {e}") + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 700307f48f5..e4d760ca7bc 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1028,6 +1028,8 @@ _COST_MAP_ROW_SCHEMA: Final = { "audio_transcription_config": {"type": "string"}, "deprecation_date": {"type": "string"}, "successor_model": {"type": "string"}, + "aliases": {"type": "array", "items": {"type": "string"}}, + "alias_of": {"type": "string"}, "deprecation_source": {"type": "string"}, "input_cost_per_audio_per_second": {"type": "number"}, "input_cost_per_audio_per_second_above_128k_tokens": {"type": "number"},