mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
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:
parent
87e32e7c6d
commit
ceb7819a9d
5 changed files with 206 additions and 15 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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":
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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"]]) == (
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue