scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

feat(cost-map): read alias_of pointer rows in the catalog and the loader

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-11-07 06:35:36 +00:00
parent 5103d8a0ee
commit 42e31fcc08
7 changed files with 184 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

80
litellm-rust/crates/model-catalog/src/catalog.rs
View file

@ -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<String, String> = 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<Error> {
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(&section.value))),
Some(section) => Some((
section.key.as_str(),
match &section.value {
Value::Object(fields) if !is_section(&section.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<Model>, Option<Vec<String>>, Vec<Error>) {
/// One row as read: a model, or an `alias_of` pointer at one
struct RowRead {
model: Option<Model>,
aliases: Option<Vec<String>>,
alias_of: Option<String>,
problems: Vec<Error>,
}
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,
}
}

10
litellm-rust/crates/model-catalog/src/error.rs
View file

@ -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)]

2
litellm-rust/crates/model-catalog/src/wire/attributes.rs
View file

@ -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<String>,
audio_transcription_config: String,

51
litellm-rust/crates/model-catalog/tests/catalog.rs
View file

@ -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::<Vec<_>>(),
["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"}}"#);
}

26
litellm/litellm_core_utils/get_model_cost_map.py
View file

@ -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)))

31
tests/unit/litellm_core_utils/test_get_model_cost_map.py
View file

@ -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",

2
tests/unit/test_utils.py
View file

@ -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"},