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.
This commit is contained in:
mateo-berri 2026-10-31 17:21:20 -07:00
parent 87e32e7c6d
commit ceb7819a9d
5 changed files with 206 additions and 15 deletions

View file

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

View file

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

View file

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

View file

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

View file

@ -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"]]) == (