From ceb7819a9d3e532238c6edd3e2f6f04d2730c303 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 31 Oct 2026 17:21:20 -0700 Subject: [PATCH] feat(proxy): honor model_info.max_input_tokens in the Codex catalog for stock models A Codex model catalog entry copied from Codex's stock entry (a stock slug, or an alias whose upstream is one) now takes the max_input_tokens every routable deployment of the group sets as its context_window, capped at the stock entry's max_context_window like Codex's own model_context_window. With nothing set, the stock window stays. --- .../proxy/common_utils/codex_model_catalog.py | 59 ++++++++--- litellm/router_api.py | 18 +++- litellm/router_python_core.py | 17 ++++ .../compatibility/test_codex_model_catalog.py | 28 ++++++ .../common_utils/test_codex_model_catalog.py | 99 ++++++++++++++++++- 5 files changed, 206 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/common_utils/codex_model_catalog.py b/litellm/proxy/common_utils/codex_model_catalog.py index 63fb27d62cf..27872cbba7b 100644 --- a/litellm/proxy/common_utils/codex_model_catalog.py +++ b/litellm/proxy/common_utils/codex_model_catalog.py @@ -7,8 +7,10 @@ never the OpenAI list shape. This module builds that body from the rows ``/v1/mo lists: a model Codex knows keeps Codex's own entry for it, any other model gets the fallback entry Codex itself uses for an unknown slug, and ``model_info.service_tiers`` becomes the entry's ``service_tiers``, which Codex turns into slash commands (``/ultrafast``) that send ``service_tier`` -upstream. A fallback entry offers the model group's ``supported_reasoning_efforts`` as its reasoning -levels, trimmed to the ones the requesting Codex release can parse. +upstream. A ``model_info.max_input_tokens`` the operator sets replaces a known model's +``context_window``, up to the most Codex lets its own ``model_context_window`` raise that model to. +A fallback entry offers the model group's ``supported_reasoning_efforts`` as its reasoning levels, +trimmed to the ones the requesting Codex release can parse. The stock entries are Codex 0.159.3's bundled catalog, vendored unchanged from https://github.com/openai/codex/blob/rust-v0.159.3/codex-rs/models-manager/models.json (Apache-2.0) @@ -158,6 +160,7 @@ class CodexStockModel(LiteLLMBaseModel): supports_search_tool: bool = False supported_reasoning_levels: tuple[CodexReasoningLevel, ...] = () context_window: int | None = None + max_context_window: int | None = None class CodexStockCatalog(LiteLLMBaseModel): @@ -173,8 +176,9 @@ class CodexCatalogRow: """What the catalog needs to know about one listed model: its public id, the listing's mode and input limit, the upstream model behind it, the reasoning effort its deployments pin, and the `model_info` values the operator set; `service_tiers` holds one raw `model_info.service_tiers` - value per deployment behind the id, None where a deployment sets none. `supported_reasoning_efforts` - is the model group's, None when unknown.""" + value per deployment behind the id, None where a deployment sets none, and + `configured_max_input_tokens` one `model_info.max_input_tokens` per deployment the same way. + `supported_reasoning_efforts` is the model group's, None when unknown.""" id: str mode: str | None = None @@ -185,6 +189,7 @@ class CodexCatalogRow: service_tiers: tuple[object, ...] = () supported_reasoning_efforts: tuple[str, ...] | None = None search_tool_support: tuple[bool | None, ...] = () + configured_max_input_tokens: tuple[int | None, ...] = () @cache @@ -333,6 +338,32 @@ def _alias_context_window(row: CodexCatalogRow, stock_model: CodexStockModel) -> return MappingProxyType({"context_window": min(stock_model.context_window, row.max_input_tokens)}) +def _configured_window(row: CodexCatalogRow, stock_model: CodexStockModel) -> Mapping[str, int]: + """The `context_window` the operator's `model_info.max_input_tokens` gives a known model's entry, + or nothing, which keeps `stock_model`'s own. + + A request to the row can route to any of its deployments, so the window is the smallest value they + set, and only when every one sets a value: a deployment that sets none leaves Codex's own window in + place, which is logged. The window never goes past `stock_model`'s `max_context_window`, the most + Codex lets its own `model_context_window` setting raise that model to.""" + per_deployment: Final = row.configured_max_input_tokens + declared: Final = tuple(limit for limit in per_deployment if limit is not None) + if not declared: + return MappingProxyType({}) + if len(declared) < len(per_deployment): + verbose_proxy_logger.warning( + "model_info.max_input_tokens for %s is set on %d of its %d deployments, so Codex keeps its own " + "context window: the window follows max_input_tokens only when every deployment sets it", + row.id, + len(declared), + len(per_deployment), + ) + return MappingProxyType({}) + ceiling: Final = stock_model.max_context_window + smallest: Final = min(declared) + return MappingProxyType({"context_window": smallest if ceiling is None else min(smallest, ceiling)}) + + def _alias_reasoning_level(row: CodexCatalogRow, stock_model: CodexStockModel) -> str | None: efforts: Final = frozenset(level.effort for level in stock_model.supported_reasoning_levels) return row.reasoning_effort if row.reasoning_effort in efforts else None @@ -398,6 +429,7 @@ def _entry( "upgrade": upgrade, **({} if search_tool is None else {"supports_search_tool": search_tool}), **(_deployment_defaults(row, stock_model) if is_alias and alias_defaults == "deployment" else {}), + **_configured_window(row, stock_model), } if configured is None: return stock_model.model_copy(update=listing) @@ -473,10 +505,12 @@ def codex_models_response_json( byte limit silently and keeps its bundled one, so the body holds only the entries that fit, the ones offering a service tier taken first and the rest in listing order, an entry too large for the bytes left passed over, every kept entry at its listing position; `byte_limit=None` - keeps every entry. A row listed under another name than the Codex model its `upstream_model` - names (an alias) takes that model's entry; `alias_defaults="deployment"` drops the defaults - Codex picks for that model in favor of what the alias's deployments pin, while `"codex"` keeps - them, as Codex would use them for the model's own name. `client_version` is the requesting Codex's version, which decides the reasoning + keeps every entry. A row whose deployments all set `model_info.max_input_tokens` gets the smallest + as its known model's `context_window`, capped at that model's `max_context_window`. A row listed + under another name than the Codex model its `upstream_model` names (an alias) takes that model's + entry; `alias_defaults="deployment"` drops the defaults Codex picks for that model in favor of + what the alias's deployments pin, while `"codex"` keeps them, as Codex would use them for the + model's own name. `client_version` is the requesting Codex's version, which decides the reasoning levels a fallback entry may offer; None counts as a version too old for custom levels. """ known: Final = bundled_codex_models() if stock is None else stock @@ -519,6 +553,9 @@ def _catalog_row( search_tool_support=( registry.resolve(lookup_id, kind="search_tool_support", team_id=team_id) if registry is not None else () ), + configured_max_input_tokens=( + registry.group(lookup_id, field="max_input_tokens", team_id=team_id) if registry is not None else () + ), ) @@ -532,9 +569,9 @@ def codex_catalog_rows( """`rows` joined with the router's configured metadata, looked up by each entry's internal id so team-scoped rows resolve the way the Anthropic listing's display names do, a `model_group_alias` reading its target's deployments under its own name; `team_id` is the requesting key's team, so the - upstream model that picks Codex's stock entry and the tiers are read only off the deployments its - requests can route to, and the reasoning efforts are the ones `/model_group/info` reports for the - entry. Display names are read under `label_entries`' ids when given.""" + upstream model that picks Codex's stock entry, the tiers, and the configured input limits are read + only off the deployments its requests can route to, and the reasoning efforts are the ones + `/model_group/info` reports for the entry. Display names are read under `label_entries`' ids when given.""" lookup_ids: Final = MappingProxyType(dict(entries)) label_ids: Final = lookup_ids if label_entries is None else MappingProxyType(dict(label_entries)) return tuple( diff --git a/litellm/router_api.py b/litellm/router_api.py index c013a99573f..3007b5345c7 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -155,6 +155,7 @@ GroupField: TypeAlias = Literal[ "display_name", "mode", "service_tiers", + "max_input_tokens", "token_limits", "listing_info", ] @@ -855,6 +856,16 @@ class RouterRegistry: team_id: str | None = None, ) -> tuple[object, ...]: ... + @overload + def group( + self, + name: str, + *, + field: Literal["max_input_tokens"], + include_hidden: bool = False, + team_id: str | None = None, + ) -> tuple[int | None, ...]: ... + @overload def group( self, @@ -882,9 +893,7 @@ class RouterRegistry: field: GroupField = "info", include_hidden: bool = False, team_id: str | None = None, - ) -> ( - "ModelGroupInfo | DeploymentModelListingInfo | tuple[int | None, int | None] | tuple[object, ...] | str | None" - ): + ) -> "ModelGroupInfo | DeploymentModelListingInfo | tuple[int | None, ...] | tuple[object, ...] | str | None": """ What the router knows about a model group. ``field``: @@ -894,6 +903,7 @@ class RouterRegistry: (were ``get_configured_display_name``, ``get_configured_mode``, ``get_configured_token_limits``, ``get_model_listing_info``) - ``service_tiers``: (was ``get_configured_service_tiers``) + - ``max_input_tokens``: each routable deployment's own ``model_info.max_input_tokens`` """ r: Final = self._router if field == "info": @@ -906,6 +916,8 @@ class RouterRegistry: return r._core_get_configured_mode(name) if field == "service_tiers": return r._core_get_configured_service_tiers(name, team_id=team_id) + if field == "max_input_tokens": + return r._core_get_configured_max_input_tokens(name, team_id=team_id) if field == "token_limits": return r._core_get_configured_token_limits(name) if field == "listing_info": diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index 30b05456987..bdb2738e7aa 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -12680,6 +12680,23 @@ class PythonCore(DeprecatedRouterMembers): for deployment in self._routable_deployments(model_name, team_id) ) + def _core_get_configured_max_input_tokens( + self, model_name: str, team_id: str | None = None + ) -> tuple[int | None, ...]: + """ + Return the max_input_tokens each deployment a request from team_id to + model_name can route to (see _routable_deployments) sets in its own + model_info, in model_list order, None for a deployment that sets none + or a malformed one. Unlike get_model_listing_info, neither the cost map + nor a discovered limit is read, so only an operator's own value counts. + + Returns an empty tuple for wildcard-expanded or unknown names. + """ + return tuple( + coerce_token_limit(_configured_model_info(deployment).get("max_input_tokens")) + for deployment in self._routable_deployments(model_name, team_id) + ) + def _core_get_search_tool_support(self, model_name: str, team_id: str | None = None) -> tuple[bool | None, ...]: """ Whether each deployment a request from team_id to model_name can route diff --git a/tests/integration/compatibility/test_codex_model_catalog.py b/tests/integration/compatibility/test_codex_model_catalog.py index a3ba032e612..0cc10a2d145 100644 --- a/tests/integration/compatibility/test_codex_model_catalog.py +++ b/tests/integration/compatibility/test_codex_model_catalog.py @@ -351,6 +351,34 @@ def test_stock_model_keeps_codex_entry(gateway: Gateway) -> None: assert catalog[name] == expected, catalog[name] +def test_stock_model_alias_takes_configured_max_input_tokens(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + stock: Final = _bundled_stock()["gpt-6-sol"] + stock_window: Final = stock["context_window"] + ceiling: Final = stock["max_context_window"] + assert isinstance(stock_window, int) and isinstance(ceiling, int) and stock_window < ceiling, stock + raised: Final = _model_name("codex-raised-window") + lowered: Final = _model_name("codex-lowered-window") + half_set: Final = _model_name("codex-half-set-window") + unset: Final = _model_name("codex-unset-window") + _new_model(scenario, raised, model="openai/gpt-6-sol", model_info={"max_input_tokens": ceiling + 50000}) + _new_model(scenario, lowered, model="openai/gpt-6-sol", model_info={"max_input_tokens": stock_window // 2}) + _new_model(scenario, half_set, model="openai/gpt-6-sol", model_info={"max_input_tokens": ceiling}) + _new_model(scenario, half_set, model="openai/gpt-6-sol") + _new_model(scenario, unset, model="openai/gpt-6-sol") + windows: Final = _settled( + lambda: {slug: entry["context_window"] for slug, entry in _by_slug(_catalog(gateway)).items()}, + lambda value: all(name in value for name in (raised, lowered, half_set, unset)) + and value[raised] == ceiling, + ) + assert {name: windows[name] for name in (raised, lowered, half_set, unset)} == { + raised: ceiling, + lowered: stock_window // 2, + half_set: stock_window, + unset: stock_window, + }, windows + + def test_stock_model_configured_tiers_override(gateway: Gateway) -> None: with gateway.scenario() as scenario: stock: Final = _bundled_stock()["gpt-6-sol"] diff --git a/tests/unit/proxy/common_utils/test_codex_model_catalog.py b/tests/unit/proxy/common_utils/test_codex_model_catalog.py index 509303d9635..01703029d3d 100644 --- a/tests/unit/proxy/common_utils/test_codex_model_catalog.py +++ b/tests/unit/proxy/common_utils/test_codex_model_catalog.py @@ -233,6 +233,97 @@ def test_alias_with_deployment_defaults_caps_codexs_window_at_its_listed_max_inp assert json.loads(body.json)["models"][0]["context_window"] == expected +_WIDE_STOCK = MappingProxyType( + { + "gpt-6-sol": _stock("gpt-6-sol").model_copy(update={"context_window": 272000, "max_context_window": 872000}), + "gpt-5.5": _stock("gpt-5.5").model_copy(update={"context_window": 272000}), + } +) + + +@pytest.mark.parametrize( + ("configured", "expected"), + [ + ((500000,), 500000), + ((100000,), 100000), + ((1050000,), 872000), + ((600000, 400000), 400000), + ((600000, None), 272000), + ((None,), 272000), + ((), 272000), + ], +) +@pytest.mark.parametrize("alias_defaults", ["codex", "deployment"]) +def test_a_known_model_and_its_alias_take_the_max_input_tokens_every_deployment_sets( + configured, expected, alias_defaults +): + rows = ( + _row("gpt-6-sol", max_input_tokens=1050000, configured_max_input_tokens=configured), + _row( + "sol-1m", + upstream_model="openai/gpt-6-sol", + max_input_tokens=1050000, + configured_max_input_tokens=configured, + ), + ) + + body = codex_models_response_json(rows, stock=_WIDE_STOCK, instructions=_PROMPT, alias_defaults=alias_defaults) + + assert {entry["slug"]: entry["context_window"] for entry in json.loads(body.json)["models"]} == { + "gpt-6-sol": expected, + "sol-1m": expected, + } + + +def test_a_configured_window_goes_past_codexs_own_when_the_stock_entry_sets_no_ceiling(): + (entry,) = json.loads( + codex_models_response_json( + (_row("gpt-5.5", configured_max_input_tokens=(1050000,)),), stock=_WIDE_STOCK, instructions=_PROMPT + ).json + )["models"] + + assert entry["context_window"] == 1050000 + + +def test_model_list_body_sizes_a_known_models_window_off_the_deployments_the_key_team_routes_to(): + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra"}, + "model_info": {"max_input_tokens": 922000}, + }, + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://team-2.example"}, + "model_info": {"team_id": "team-2", "max_input_tokens": 400000}, + }, + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://team-3.example"}, + "model_info": {"team_id": "team-3"}, + }, + ], + model_group_alias={"astra-1m": "gpt-6-astra"}, + ) + listing = tuple( + {"id": name, "object": "model", "created": 0, "owned_by": "openai", "mode": "chat"} + for name in ("gpt-6-astra", "astra-1m") + ) + entries = (("gpt-6-astra", "gpt-6-astra"), ("astra-1m", "astra-1m")) + stock = bundled_codex_models()["gpt-6-astra"] + assert stock.context_window is not None and stock.max_context_window is not None + assert stock.context_window < 400000 < stock.max_context_window < 922000 + + def windows(team_id): + models = json.loads(codex_model_list_body(listing, entries, router, team_id))["models"] + return {entry["slug"]: entry["context_window"] for entry in models} + + assert windows("team-1") == {"gpt-6-astra": stock.max_context_window, "astra-1m": stock.max_context_window} + assert windows("team-2") == {"gpt-6-astra": 400000, "astra-1m": 400000} + assert windows("team-3") == {"gpt-6-astra": stock.context_window, "astra-1m": stock.context_window} + + def test_alias_with_deployment_defaults_drops_codexs_tier_and_upgrade_while_the_native_name_keeps_them(): rows = (_row("gpt-6-sol"), _row("team-sol", upstream_model="gpt-6-sol"), _row("team-55", upstream_model="gpt-5.5")) codex = {entry["slug"]: entry for entry in _models(*rows)} @@ -599,9 +690,14 @@ def test_catalog_rows_read_the_router_by_each_entry_lookup_id(): "model_name_team-1_c0ffee" ).supported_reasoning_efforts, search_tool_support=(None, None), + configured_max_input_tokens=(None, None), ), CodexCatalogRow( - id="plain", upstream_model="openai/some-unmapped-model", service_tiers=(None,), search_tool_support=(None,) + id="plain", + upstream_model="openai/some-unmapped-model", + service_tiers=(None,), + search_tool_support=(None,), + configured_max_input_tokens=(None,), ), ) @@ -631,6 +727,7 @@ def test_catalog_rows_read_an_alias_off_its_target_under_the_alias_name(): service_tiers=(["ultrafast"],), supported_reasoning_efforts=router.get_model_group_info("gpt-6-astra").supported_reasoning_efforts, search_tool_support=(None,), + configured_max_input_tokens=(None,), ) (entry,) = json.loads(codex_model_list_body(listing, (("gpt-6", "gpt-6"),), router))["models"] assert (entry["slug"], entry["display_name"], [tier["id"] for tier in entry["service_tiers"]]) == (