From c720b23238ab12a574617e6d4d5d8c71cb6c697c Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Mon, 2 Nov 2026 13:27:58 -0800 Subject: [PATCH] feat(proxy)!: honor model_info.max_input_tokens in the Codex catalog for stock models (#49701) * 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. * feat(proxy): honor the smallest max_input_tokens a group's deployments set in the Codex catalog A deployment without a limit counts as Codex's own window, so a limit on only some deployments can lower the window but never raise it * fix(proxy): skip a zero or negative max_input_tokens when sizing a Codex catalog window * test(proxy): stub the max_input_tokens group field and type the Codex window tests * fix(proxy): skip a max_input_tokens a saved /model/info response copied from the cost map --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/client/cli/commands/agents.py | 7 +- .../proxy/common_utils/codex_model_catalog.py | 52 +++++-- litellm/router_api.py | 18 ++- litellm/router_python_core.py | 32 ++++ .../compatibility/test_codex_model_catalog.py | 35 +++++ .../common_utils/test_codex_model_catalog.py | 145 +++++++++++++++++- .../proxy/proxy_server/test_routes_models.py | 1 + tests/unit/proxy/router_entry_point_stubs.py | 2 + 8 files changed, 274 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 12d97ef6cce..b70432d1e72 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -770,9 +770,10 @@ def codex_model_catalog( A model Codex does not know gets the fallback entry, with the same base instructions Codex itself uses so the agent never runs without a system prompt, and offers the reasoning efforts `groups` gives its model group - that the installed Codex, at `client_version`, can parse. Apart from aliases, the catalog is - the one `/v1/models?client_version=...` serves, built from the installed Codex's own entries - instead of the proxy's vendored ones. + that the installed Codex, at `client_version`, can parse. Apart from aliases, and the window the + deployments' `model_info.max_input_tokens` gives a model Codex knows, the catalog is the one + `/v1/models?client_version=...` serves, built from the installed Codex's own entries instead of + the proxy's vendored ones. """ rows: Final = tuple(_codex_row(m, upstreams.get(m.id), groups) for m in models) body: Final = codex_models_response_json( diff --git a/litellm/proxy/common_utils/codex_model_catalog.py b/litellm/proxy/common_utils/codex_model_catalog.py index 63fb27d62cf..7b48aa07365 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,25 @@ 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 one of them + takes: its `max_input_tokens` where the operator set a positive one, Codex's own window where not. A + limit set on only some deployments can therefore lower the window but never raise it. 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 and limit > 0) + unset_window: Final = () if len(declared) == len(per_deployment) else (stock_model.context_window,) + if not declared or None in unset_window: + return MappingProxyType({}) + smallest: Final = min(limit for limit in (*declared, *unset_window) if limit is not None) + ceiling: Final = stock_model.max_context_window + 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 +422,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 +498,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 set `model_info.max_input_tokens` gets the smallest + window any of them takes 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 +546,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 +562,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 c63e538c998..98c9e641278 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -141,6 +141,7 @@ from litellm.litellm_core_utils.get_llm_provider_logic import ( is_registered_custom_provider, resolve_model_and_provider_for_metadata, ) +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.get_retry_from_policy import ( get_num_retries_from_retry_policy as _get_num_retries_from_retry_policy, ) @@ -475,6 +476,7 @@ from litellm.types.router import ( from litellm.types.services import ServiceTypes from litellm.types.utils import ( AUTOROUTER_CLASSIFIER_CALL_ORIGIN, + COST_MAP_LOOKUP_KEY, PROMPT_QUOTING_ROUTING_DECISION_FIELDS, CustomPricingLiteLLMParams, GenericBudgetConfigType, @@ -486,6 +488,7 @@ from litellm.types.utils import ( StandardLoggingRoutingDecision, Usage, all_litellm_params, + echoed_cost_map_fields, is_server_derived_pricing_key, shared_backend_model_info, ) @@ -772,6 +775,17 @@ def _configured_model_info(deployment: Mapping[str, object]) -> Mapping[str, obj return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {}) +def _operator_max_input_tokens(deployment: Mapping[str, object]) -> int | None: + model_info: Final = _configured_model_info(deployment) + catalog_key: Final = model_info.get(COST_MAP_LOOKUP_KEY) + catalog_entry: Final = ( + GetModelCostMap.loaded_model_cost_map().get(catalog_key) if isinstance(catalog_key, str) else None + ) + if catalog_entry is not None and "max_input_tokens" in echoed_cost_map_fields(model_info, catalog_entry): + return None + return coerce_token_limit(model_info.get("max_input_tokens")) + + def _search_tool_support(deployment: DeploymentTypedDict) -> bool | None: from litellm.responses.litellm_completion_transformation.tool_search import chat_config_builds_anthropic_request @@ -12693,6 +12707,24 @@ 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, + and a value a saved /model/info response copied from the cost map is + skipped too. + + Returns an empty tuple for wildcard-expanded or unknown names. + """ + return tuple( + _operator_max_input_tokens(deployment) 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..fbf05b0f526 100644 --- a/tests/integration/compatibility/test_codex_model_catalog.py +++ b/tests/integration/compatibility/test_codex_model_catalog.py @@ -351,6 +351,41 @@ 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") + half_lowered: Final = _model_name("codex-half-lowered-window") + unset: Final = _model_name("codex-unset-window") + names: Final = (raised, lowered, half_set, half_lowered, unset) + _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, half_lowered, model="openai/gpt-6-sol", model_info={"max_input_tokens": stock_window // 3}) + _new_model(scenario, half_lowered, model="openai/gpt-6-sol") + _new_model(scenario, unset, model="openai/gpt-6-sol") + expected: Final = { + raised: ceiling, + lowered: stock_window // 2, + half_set: stock_window, + half_lowered: stock_window // 3, + unset: stock_window, + } + windows: Final = _settled( + lambda: { + slug: entry["context_window"] for slug, entry in _by_slug(_catalog(gateway)).items() if slug in names + }, + lambda value: value == expected, + ) + assert windows == expected, 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..e3d941a6aa4 100644 --- a/tests/unit/proxy/common_utils/test_codex_model_catalog.py +++ b/tests/unit/proxy/common_utils/test_codex_model_catalog.py @@ -1,9 +1,12 @@ import json +from collections.abc import Mapping from types import MappingProxyType +from typing import Final, Literal import pytest import litellm +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.proxy.common_utils.codex_model_catalog import ( CODEX_CATALOG_BYTE_LIMIT, CodexCatalogRow, @@ -233,6 +236,140 @@ 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), + ((100000, None), 100000), + ((None,), 272000), + ((), 272000), + ((0,), 272000), + ((-5, 400000), 272000), + ((0, 100000), 100000), + ], +) +@pytest.mark.parametrize("alias_defaults", ["codex", "deployment"]) +def test_a_known_model_and_its_alias_take_the_smallest_window_any_deployment_takes( + configured: tuple[int | None, ...], expected: int, alias_defaults: Literal["codex", "deployment"] +) -> None: + rows: Final = ( + _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: Final = 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() -> None: + (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_a_deployment_without_a_limit_keeps_a_windowless_stock_entry_as_codex_ships_it() -> None: + windowless: Final = MappingProxyType({"gpt-5.5": _stock("gpt-5.5").model_copy(update={"context_window": None})}) + + (entry,) = json.loads( + codex_models_response_json( + (_row("gpt-5.5", configured_max_input_tokens=(1050000, None)),), stock=windowless, instructions=_PROMPT + ).json + )["models"] + + assert entry["context_window"] is None + + +def test_model_list_body_sizes_a_known_models_window_off_the_deployments_the_key_team_routes_to() -> None: + router: Final = 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: Final = tuple( + {"id": name, "object": "model", "created": 0, "owned_by": "openai", "mode": "chat"} + for name in ("gpt-6-astra", "astra-1m") + ) + entries: Final = (("gpt-6-astra", "gpt-6-astra"), ("astra-1m", "astra-1m")) + stock: Final = 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: str) -> dict[str, int | None]: + models: Final = 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_model_list_body_skips_a_limit_a_saved_model_info_response_copied_from_the_cost_map() -> None: + catalog_limit: Final = GetModelCostMap.loaded_model_cost_map()["gpt-6-astra"]["max_input_tokens"] + stock: Final = bundled_codex_models()["gpt-6-astra"] + assert isinstance(catalog_limit, int) and stock.max_context_window is not None + assert stock.context_window is not None and stock.context_window < 600000 < stock.max_context_window + + def window(model_info: Mapping[str, object]) -> int | None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra"}, + "model_info": dict(model_info), + } + ] + ) + listing: Final = ({"id": "gpt-6-astra", "object": "model", "created": 0, "owned_by": "openai", "mode": "chat"},) + (entry,) = json.loads(codex_model_list_body(listing, (("gpt-6-astra", "gpt-6-astra"),), router))["models"] + return entry["context_window"] + + assert window({"key": "gpt-6-astra", "max_input_tokens": catalog_limit}) == stock.context_window + assert window({"key": "gpt-6-astra", "max_input_tokens": 600000}) == 600000 + assert window({"max_input_tokens": catalog_limit}) == min(catalog_limit, stock.max_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 +736,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 +773,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"]]) == ( diff --git a/tests/unit/proxy/proxy_server/test_routes_models.py b/tests/unit/proxy/proxy_server/test_routes_models.py index 874f471f3e7..ff74cd3fa1d 100644 --- a/tests/unit/proxy/proxy_server/test_routes_models.py +++ b/tests/unit/proxy/proxy_server/test_routes_models.py @@ -58,6 +58,7 @@ def patched_models(monkeypatch): router.get_routable_upstream_model = MagicMock(return_value="gpt-4") router.get_configured_display_name = MagicMock(return_value=None) router.get_configured_service_tiers = MagicMock(return_value=()) + router.get_configured_max_input_tokens = MagicMock(return_value=()) router.get_model_group_info = MagicMock(return_value=None) router.get_search_tool_support = MagicMock(return_value=()) route_entry_points_to_stubs(router) diff --git a/tests/unit/proxy/router_entry_point_stubs.py b/tests/unit/proxy/router_entry_point_stubs.py index 6912049d225..51b5bf13f16 100644 --- a/tests/unit/proxy/router_entry_point_stubs.py +++ b/tests/unit/proxy/router_entry_point_stubs.py @@ -123,6 +123,8 @@ def route_entry_points_to_stubs(router: _RouterDouble) -> _RouterDouble: return router.get_configured_display_name(name) if field == "service_tiers": return router.get_configured_service_tiers(name, **_given(team_id=team_id)) + if field == "max_input_tokens": + return router.get_configured_max_input_tokens(name, **_given(team_id=team_id)) if field == "listing_info": return router.get_model_listing_info(name) if field == "mode":