mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
fix(router): let an exact base_model widen a deployment's context window (#46197)
* fix(router): honor explicit bedrock base model capacity * fix(router): honor an exact base_model's capacity for every non-Azure provider Extends the Bedrock-only check to every non-Azure provider and only uses the base_model when it has its own exact cost-map entry, so an unmapped or wildcard base_model falls back to the route id instead of raising in the router. * test(router): call _maps_to_exact_model_info directly in its cost-map tests * fix(router): let an exact base_model only widen a deployment's context window A base_model used as a pricing stand-in with a smaller window no longer narrows the route's limit, and a route with no window never takes one from base_model * test(router): count the base_model window helpers as covered through get_router_model_info * fix(router): widen only from a model_info base_model naming the same model A stand-in base_model for a different, larger model widened a smaller deployment's window, so pre-call checks sent it prompts the provider rejects with 400. The window now widens only when model_info.base_model names the same provider model id as the route, which is also the field /model/info reads. * fix(router): match a base_model's model id across Bedrock region prefixes and @ qualifiers * fix(router): read litellm_params.base_model too and look the base row up as written * test(router): cover a same-model base_model window on chat, messages, responses, and complexity escalation --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Michelangelo Setaro
GitHub
parent
c8b73cde84
commit
c288dc7da0
4 changed files with 592 additions and 130 deletions
|
|
@ -128,6 +128,7 @@ from litellm.llms.base_llm.vector_store.transformation import (
|
|||
RouterVectorStoreEmbeddingExecutor,
|
||||
vector_store_request_metadata,
|
||||
)
|
||||
from litellm.llms.bedrock.common_utils import get_bedrock_base_model
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
|
||||
from litellm.llms.openai_like.json_loader import JSONProviderRegistry
|
||||
from litellm.llms.openai_like.model_info import (
|
||||
|
|
@ -374,6 +375,7 @@ from litellm.utils import (
|
|||
get_non_default_completion_params,
|
||||
get_secret,
|
||||
get_utc_datetime,
|
||||
is_generalized_model_info,
|
||||
is_region_allowed,
|
||||
provider_rejectable_params,
|
||||
set_live_deployment_replay,
|
||||
|
|
@ -991,6 +993,49 @@ def as_output_cap(value: object) -> int | None:
|
|||
return cap if cap >= 0 else None
|
||||
|
||||
|
||||
def _exact_base_model_max_input_tokens(base_model: str) -> int | None:
|
||||
try:
|
||||
model_info: Final = litellm.get_model_info(model=base_model)
|
||||
except Exception: # noqa: BLE001 # get_model_info raises ModelNotMappedError or a bare Exception for an unmapped model
|
||||
return None
|
||||
if is_generalized_model_info(model_info):
|
||||
return None
|
||||
return model_info.get("max_input_tokens")
|
||||
|
||||
|
||||
def _provider_model_id(model_name: str, custom_llm_provider: str) -> str:
|
||||
model_id: Final = model_name.rsplit("/", 1)[-1]
|
||||
return get_bedrock_base_model(model_id) if custom_llm_provider.startswith("bedrock") else model_id
|
||||
|
||||
|
||||
def _names_route_model(base_model: str, route_key: str, custom_llm_provider: str) -> bool:
|
||||
route_model_id: Final = _provider_model_id(route_key, custom_llm_provider)
|
||||
base_model_id: Final = _provider_model_id(base_model, custom_llm_provider)
|
||||
return route_model_id in (base_model_id, base_model_id.split("@", 1)[0])
|
||||
|
||||
|
||||
def _declared_base_model(user_model_info: Mapping[str, object]) -> str | None:
|
||||
match user_model_info:
|
||||
case {"base_model": str() as base_model}:
|
||||
return base_model
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def _base_model_widened_window(
|
||||
route_model_info: ModelMapInfo, base_model: str | None, custom_llm_provider: str
|
||||
) -> Mapping[str, int]:
|
||||
route_max_input_tokens: Final = route_model_info.get("max_input_tokens")
|
||||
if base_model is None or route_max_input_tokens is None:
|
||||
return MappingProxyType({})
|
||||
if not _names_route_model(base_model, route_model_info["key"], custom_llm_provider):
|
||||
return MappingProxyType({})
|
||||
base_max_input_tokens: Final = _exact_base_model_max_input_tokens(base_model)
|
||||
if base_max_input_tokens is None or base_max_input_tokens <= route_max_input_tokens:
|
||||
return MappingProxyType({})
|
||||
return MappingProxyType({"max_input_tokens": base_max_input_tokens})
|
||||
|
||||
|
||||
class Router:
|
||||
@property
|
||||
def _routing_groups(self) -> dict[str, RoutingGroup]:
|
||||
|
|
@ -11343,6 +11388,7 @@ class Router:
|
|||
Augment litellm info with additional params set in `model_info`.
|
||||
|
||||
For azure models, ignore the `model:`. Only set max tokens, cost values if base_model is set.
|
||||
For every other provider, a base_model naming the same model as `model:` can only widen its context window.
|
||||
|
||||
Returns
|
||||
- ModelInfo - If found -> typed dict with max tokens, input cost, etc.
|
||||
|
|
@ -11439,20 +11485,23 @@ class Router:
|
|||
|
||||
## CHECK USER SET MODEL INFO
|
||||
raw_user_model_info: Final = deployment.get("model_info")
|
||||
user_model_info: Final = (
|
||||
user_model_info: Final = _MODEL_INFO_ADAPTER.validate_python(
|
||||
raw_user_model_info.model_dump(exclude_none=True)
|
||||
if isinstance(raw_user_model_info, BaseModel)
|
||||
else raw_user_model_info
|
||||
else raw_user_model_info or {}
|
||||
)
|
||||
|
||||
# get_model_info() hands back an lru_cache'd dict, so merge into a copy; unset
|
||||
# values are skipped or Deployment's None pricing defaults would erase the map's
|
||||
merged_model_info: Final[ModelMapInfo] = {
|
||||
**copy.deepcopy(model_info),
|
||||
**self.get_discovered_model_info((deployment.get("model_info") or {}).get("id")),
|
||||
**MappingProxyType(
|
||||
{key: value for key, value in (user_model_info or MappingProxyType({})).items() if value is not None}
|
||||
**_base_model_widened_window(
|
||||
model_info,
|
||||
_declared_base_model(user_model_info) or _declared_base_model(litellm_params.model_extra or {}),
|
||||
custom_llm_provider,
|
||||
),
|
||||
**self.get_discovered_model_info((deployment.get("model_info") or {}).get("id")),
|
||||
**MappingProxyType({key: value for key, value in user_model_info.items() if value is not None}),
|
||||
}
|
||||
|
||||
return merged_model_info
|
||||
|
|
|
|||
|
|
@ -114,6 +114,11 @@ ignored_function_names = [
|
|||
"_deployment_that_just_failed", # Tested through the same-boundary hop test in test_router_order_fallback.py
|
||||
"_collect_ptu_share_teams", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py
|
||||
"_resolve_team_ptu_ceiling", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py
|
||||
"_exact_base_model_max_input_tokens", # Tested through get_router_model_info in test_router_model_cost_isolation.py
|
||||
"_base_model_widened_window", # Tested through get_router_model_info in test_router_model_cost_isolation.py
|
||||
"_provider_model_id", # Tested through get_router_model_info in test_router_model_cost_isolation.py
|
||||
"_names_route_model", # Tested through get_router_model_info in test_router_model_cost_isolation.py
|
||||
"_declared_base_model", # Tested through get_router_model_info in test_router_model_cost_isolation.py
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,171 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import gateway_from_environment
|
||||
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from integration.providers._cache_control_marks_support import marker_of, owned_config
|
||||
from pydantic import JsonValue
|
||||
|
||||
_TOKEN: Final = "synthetic-bedrock-bearer"
|
||||
_KIMI: Final = "moonshotai.kimi-k2.5"
|
||||
_SONNET: Final = "global.anthropic.claude-sonnet-4-6"
|
||||
_KIMI_BASE: Final = f"bedrock/us-east-1/{_KIMI}"
|
||||
_KIMI_TARGET: Final = f"/model/{_KIMI}/converse"
|
||||
_SONNET_TARGET: Final = f"/model/{_SONNET}/converse"
|
||||
_PAST_ROUTE_ROW_LINES: Final = 16250
|
||||
_PAST_ESCALATION_BUFFER_LINES: Final = 15480
|
||||
_SMALL_LINES: Final = 1300
|
||||
_REPLY: Final = json.dumps(
|
||||
{
|
||||
"output": {"message": {"role": "assistant", "content": [{"text": "8115"}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": 9, "outputTokens": 2, "totalTokens": 11},
|
||||
"metrics": {"latencyMs": 1},
|
||||
}
|
||||
).encode()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Rig:
|
||||
owned: OwnedProxy
|
||||
wire: Wire
|
||||
|
||||
|
||||
def _converse_peer(request: Request) -> Reply:
|
||||
assert request.headers["authorization"] == f"Bearer {_TOKEN}", request.headers
|
||||
return Reply(body=_REPLY)
|
||||
|
||||
|
||||
def _deployment(
|
||||
name: str,
|
||||
model: str,
|
||||
api_base: str,
|
||||
*,
|
||||
model_info: Mapping[str, JsonValue] | None = None,
|
||||
**params: JsonValue,
|
||||
) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model_name": name,
|
||||
"litellm_params": {
|
||||
"model": f"bedrock/converse/{model}",
|
||||
"api_key": _TOKEN,
|
||||
"aws_region_name": "us-east-1",
|
||||
"aws_bedrock_runtime_endpoint": api_base,
|
||||
**params,
|
||||
},
|
||||
"model_info": {"id": f"{name}-{uuid.uuid4().hex}", **(model_info or {})},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp("base-model-window")
|
||||
with gateway_from_environment() as gateway, wire_server(_converse_peer) as wire:
|
||||
config: Final = owned_config(
|
||||
directory,
|
||||
[
|
||||
_deployment("base-in-model-info", _KIMI, wire.url, model_info={"base_model": _KIMI_BASE}),
|
||||
_deployment("base-in-litellm-params", _KIMI, wire.url, base_model=_KIMI_BASE),
|
||||
_deployment("no-base", _KIMI, wire.url),
|
||||
_deployment("other-model-base", _KIMI, wire.url, model_info={"base_model": _SONNET}),
|
||||
_deployment(
|
||||
"declared-window",
|
||||
_KIMI,
|
||||
wire.url,
|
||||
model_info={"base_model": _KIMI_BASE, "max_input_tokens": 257000},
|
||||
),
|
||||
_deployment(
|
||||
"unmapped-base", _KIMI, wire.url, model_info={"base_model": f"{_KIMI_BASE}-not-in-cost-map"}
|
||||
),
|
||||
_deployment(
|
||||
"smaller-other-model-base",
|
||||
_KIMI,
|
||||
wire.url,
|
||||
model_info={"base_model": "bedrock/us-east-1/meta.llama3-8b-instruct-v1:0"},
|
||||
),
|
||||
_deployment("one-million", _SONNET, wire.url),
|
||||
{
|
||||
"model_name": "complexity",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": {
|
||||
"tiers": {
|
||||
"SIMPLE": "base-in-model-info",
|
||||
"MEDIUM": "base-in-model-info",
|
||||
"COMPLEX": "one-million",
|
||||
"REASONING": "one-million",
|
||||
},
|
||||
"enable_context_window_escalation": True,
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
litellm_settings={"cache": False},
|
||||
router_settings={"enable_pre_call_checks": True},
|
||||
)
|
||||
with owned_proxy_process(gateway, directory, {}, config=config, workers=2, client_timeout=120) as owned:
|
||||
yield _Rig(owned, wire)
|
||||
|
||||
|
||||
def _prompt(lines: int, marker: str) -> str:
|
||||
listing: Final = "\n".join(f"Line {i}: the value of item {i} is {(i * 7) % 1000}." for i in range(1, lines + 1))
|
||||
return f"{listing}\n\nWhat is the value of item 12345? Answer with the number only. marker-{marker}"
|
||||
|
||||
|
||||
def _request(surface: str, model: str, text: str) -> tuple[str, dict[str, JsonValue]]:
|
||||
if surface == "messages":
|
||||
return "/v1/messages", {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": text}]}
|
||||
if surface == "responses":
|
||||
return "/v1/responses", {"model": model, "max_output_tokens": 16, "input": text}
|
||||
return "/v1/chat/completions", {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": text}]}
|
||||
|
||||
|
||||
def _send(rig: _Rig, surface: str, model: str, lines: int) -> tuple[int, str, tuple[str, ...]]:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
path, body = _request(surface, model, _prompt(lines, marker))
|
||||
response: Final = rig.owned.gateway.request("POST", path, body)
|
||||
targets: Final = tuple(request.target for request in rig.wire.drain() if marker_of(request) == marker)
|
||||
return response.status_code, response.text, targets
|
||||
|
||||
|
||||
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
@pytest.mark.parametrize("surface", ["chat", "messages", "responses"])
|
||||
def test_a_same_model_base_model_row_widens_the_window_pre_call_checks_enforce(rig: _Rig, surface: str) -> None:
|
||||
status, text, targets = _send(rig, surface, "base-in-model-info", _PAST_ROUTE_ROW_LINES)
|
||||
assert status == 200, text
|
||||
assert targets == (_KIMI_TARGET,)
|
||||
|
||||
|
||||
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
def test_a_same_model_base_model_in_litellm_params_widens_the_window_too(rig: _Rig) -> None:
|
||||
status, text, targets = _send(rig, "chat", "base-in-litellm-params", _PAST_ROUTE_ROW_LINES)
|
||||
assert status == 200, text
|
||||
assert targets == (_KIMI_TARGET,)
|
||||
|
||||
|
||||
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
@pytest.mark.parametrize("model", ["no-base", "other-model-base", "declared-window", "unmapped-base"])
|
||||
def test_a_base_model_that_cannot_widen_keeps_the_route_window_and_calls_no_upstream(rig: _Rig, model: str) -> None:
|
||||
status, text, targets = _send(rig, "chat", model, _PAST_ROUTE_ROW_LINES)
|
||||
assert status == 400, text
|
||||
assert "ContextWindowExceededError" in text, text
|
||||
assert targets == ()
|
||||
|
||||
|
||||
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
def test_a_smaller_other_model_base_model_never_narrows_the_route_window(rig: _Rig) -> None:
|
||||
status, text, targets = _send(rig, "chat", "smaller-other-model-base", _SMALL_LINES)
|
||||
assert status == 200, text
|
||||
assert targets == (_KIMI_TARGET,)
|
||||
|
||||
|
||||
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
def test_complexity_router_keeps_a_prompt_on_the_tier_whose_base_model_row_holds_it(rig: _Rig) -> None:
|
||||
status, text, targets = _send(rig, "chat", "complexity", _PAST_ESCALATION_BUFFER_LINES)
|
||||
assert status == 200, text
|
||||
assert targets == (_KIMI_TARGET,)
|
||||
|
|
@ -119,27 +119,33 @@ async def test_discovered_limits_survive_deployment_growth_and_removal(
|
|||
|
||||
async def test_discovery_discards_metadata_for_a_replaced_deployment(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost))
|
||||
router: Final = Router(model_list=[{
|
||||
"model_name": "local",
|
||||
"litellm_params": {
|
||||
"model": "hosted_vllm/local-model",
|
||||
"api_base": "https://original.test/v1",
|
||||
"api_key": "local-key",
|
||||
},
|
||||
"model_info": {"id": "replaced-deployment"},
|
||||
}])
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "local",
|
||||
"litellm_params": {
|
||||
"model": "hosted_vllm/local-model",
|
||||
"api_base": "https://original.test/v1",
|
||||
"api_key": "local-key",
|
||||
},
|
||||
"model_info": {"id": "replaced-deployment"},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
if request.url.host == "original.test":
|
||||
router.upsert_deployment(Deployment(
|
||||
model_name="local",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="hosted_vllm/local-model",
|
||||
api_base="https://replacement.test/v1",
|
||||
api_key="local-key",
|
||||
),
|
||||
model_info=ModelInfo(id="replaced-deployment"),
|
||||
))
|
||||
router.upsert_deployment(
|
||||
Deployment(
|
||||
model_name="local",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="hosted_vllm/local-model",
|
||||
api_base="https://replacement.test/v1",
|
||||
api_key="local-key",
|
||||
),
|
||||
model_info=ModelInfo(id="replaced-deployment"),
|
||||
)
|
||||
)
|
||||
return httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]})
|
||||
assert request.url.host == "replacement.test"
|
||||
return httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 2048}]})
|
||||
|
|
@ -158,15 +164,19 @@ async def test_discovery_discards_metadata_for_a_replaced_deployment(monkeypatch
|
|||
async def test_discovery_is_isolated_across_routers_and_reused_ids(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost))
|
||||
first, second = tuple(
|
||||
Router(model_list=[{
|
||||
"model_name": "local",
|
||||
"litellm_params": {
|
||||
"model": "hosted_vllm/local-model",
|
||||
"api_base": f"https://{host}.test/v1",
|
||||
"api_key": "local-key",
|
||||
},
|
||||
"model_info": {"id": "shared-discovery-id"},
|
||||
}])
|
||||
Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "local",
|
||||
"litellm_params": {
|
||||
"model": "hosted_vllm/local-model",
|
||||
"api_base": f"https://{host}.test/v1",
|
||||
"api_key": "local-key",
|
||||
},
|
||||
"model_info": {"id": "shared-discovery-id"},
|
||||
}
|
||||
]
|
||||
)
|
||||
for host in ("first", "second")
|
||||
)
|
||||
|
||||
|
|
@ -187,15 +197,17 @@ async def test_discovery_is_isolated_across_routers_and_reused_ids(monkeypatch:
|
|||
assert first.get_configured_token_limits("local") == (8192, 8192)
|
||||
assert second.get_configured_token_limits("local") == (2048, 2048)
|
||||
assert litellm.model_cost["shared-discovery-id"].get("max_input_tokens") is None
|
||||
first.upsert_deployment(Deployment(
|
||||
model_name="local",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="hosted_vllm/local-model",
|
||||
api_base="https://unavailable.test/v1",
|
||||
api_key="local-key",
|
||||
),
|
||||
model_info=ModelInfo(id="shared-discovery-id"),
|
||||
))
|
||||
first.upsert_deployment(
|
||||
Deployment(
|
||||
model_name="local",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="hosted_vllm/local-model",
|
||||
api_base="https://unavailable.test/v1",
|
||||
api_key="local-key",
|
||||
),
|
||||
model_info=ModelInfo(id="shared-discovery-id"),
|
||||
)
|
||||
)
|
||||
assert first.get_configured_token_limits("local") == (None, None)
|
||||
await first.arefresh_model_info(client=handler)
|
||||
assert first.get_configured_token_limits("local") == (None, None)
|
||||
|
|
@ -206,17 +218,19 @@ async def test_discovery_is_isolated_across_routers_and_reused_ids(monkeypatch:
|
|||
async def test_discovery_refreshes_other_endpoints_while_one_is_pending(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost))
|
||||
second_started: Final = asyncio.Event()
|
||||
router: Final = Router(model_list=[
|
||||
{
|
||||
"model_name": host,
|
||||
"litellm_params": {
|
||||
"model": "hosted_vllm/local-model",
|
||||
"api_base": f"https://{host}.test/v1",
|
||||
"api_key": "local-key",
|
||||
},
|
||||
}
|
||||
for host in ("first", "second", "third")
|
||||
])
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": host,
|
||||
"litellm_params": {
|
||||
"model": "hosted_vllm/local-model",
|
||||
"api_base": f"https://{host}.test/v1",
|
||||
"api_key": "local-key",
|
||||
},
|
||||
}
|
||||
for host in ("first", "second", "third")
|
||||
]
|
||||
)
|
||||
|
||||
async def respond(request: httpx.Request) -> httpx.Response:
|
||||
if request.url.host == "first.test":
|
||||
|
|
@ -240,21 +254,27 @@ async def test_discovery_refreshes_other_endpoints_while_one_is_pending(monkeypa
|
|||
async def test_discovered_limits_expire_after_the_last_successful_refresh(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost))
|
||||
clock: Final = Mock(return_value=0.0)
|
||||
router: Final = Router(model_list=[{
|
||||
"model_name": "local",
|
||||
"litellm_params": {
|
||||
"model": "hosted_vllm/local-model",
|
||||
"api_base": "https://expiry.test/v1",
|
||||
"api_key": "local-key",
|
||||
},
|
||||
"model_info": {"id": "expiring-discovery"},
|
||||
}])
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "local",
|
||||
"litellm_params": {
|
||||
"model": "hosted_vllm/local-model",
|
||||
"api_base": "https://expiry.test/v1",
|
||||
"api_key": "local-key",
|
||||
},
|
||||
"model_info": {"id": "expiring-discovery"},
|
||||
}
|
||||
]
|
||||
)
|
||||
router._discovered_model_info_cache = InMemoryCache(clock=clock, default_ttl=2 * MODEL_INFO_REFRESH_SECONDS)
|
||||
responses: Final = iter((
|
||||
httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}),
|
||||
httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]}),
|
||||
httpx.Response(503),
|
||||
))
|
||||
responses: Final = iter(
|
||||
(
|
||||
httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}),
|
||||
httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]}),
|
||||
httpx.Response(503),
|
||||
)
|
||||
)
|
||||
handler: Final = AsyncHTTPHandler()
|
||||
await handler.client.aclose()
|
||||
async with httpx.AsyncClient(transport=httpx.MockTransport(lambda request: next(responses))) as client:
|
||||
|
|
@ -343,10 +363,12 @@ async def test_discovered_limits_are_isolated_overridable_and_refreshable(
|
|||
|
||||
async def test_discovery_preserves_input_overrides_and_survives_outages(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost))
|
||||
responses: Final = iter((
|
||||
httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}),
|
||||
httpx.Response(503),
|
||||
))
|
||||
responses: Final = iter(
|
||||
(
|
||||
httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}),
|
||||
httpx.Response(503),
|
||||
)
|
||||
)
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
assert request.url.host == "backend.test"
|
||||
|
|
@ -354,27 +376,29 @@ async def test_discovery_preserves_input_overrides_and_survives_outages(monkeypa
|
|||
assert request.headers["x-tenant"] == "tenant"
|
||||
return next(responses)
|
||||
|
||||
router: Final = Router(model_list=[
|
||||
{
|
||||
"model_name": "configured",
|
||||
"litellm_params": {
|
||||
"model": "hosted_vllm/local-model",
|
||||
"api_base": "https://backend.test/v1",
|
||||
"api_key": "unused-key",
|
||||
"extra_headers": {"authorization": "Bearer local-key", "X-Tenant": "tenant"},
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "configured",
|
||||
"litellm_params": {
|
||||
"model": "hosted_vllm/local-model",
|
||||
"api_base": "https://backend.test/v1",
|
||||
"api_key": "unused-key",
|
||||
"extra_headers": {"authorization": "Bearer local-key", "X-Tenant": "tenant"},
|
||||
},
|
||||
"model_info": {"id": "configured", "max_input_tokens": 1024},
|
||||
},
|
||||
"model_info": {"id": "configured", "max_input_tokens": 1024},
|
||||
},
|
||||
{
|
||||
"model_name": "byok",
|
||||
"litellm_params": {
|
||||
"model": "openai/local-model",
|
||||
"api_base": "https://caller.test/v1",
|
||||
"use_clientside_credentials": True,
|
||||
{
|
||||
"model_name": "byok",
|
||||
"litellm_params": {
|
||||
"model": "openai/local-model",
|
||||
"api_base": "https://caller.test/v1",
|
||||
"use_clientside_credentials": True,
|
||||
},
|
||||
},
|
||||
},
|
||||
{"model_name": "default-openai", "litellm_params": {"model": "openai/local-model", "api_key": "unused"}},
|
||||
])
|
||||
{"model_name": "default-openai", "litellm_params": {"model": "openai/local-model", "api_key": "unused"}},
|
||||
]
|
||||
)
|
||||
handler: Final = AsyncHTTPHandler()
|
||||
await handler.client.aclose()
|
||||
responder: Final = Mock(side_effect=respond)
|
||||
|
|
@ -944,8 +968,7 @@ def _register_azure_tier_backend() -> None:
|
|||
|
||||
def test_inherit_builtin_service_tier_pricing_fills_only_missing_fields() -> None:
|
||||
model_cost_entries: Final = {
|
||||
key: copy.deepcopy(litellm.model_cost.get(key))
|
||||
for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL)
|
||||
key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL)
|
||||
}
|
||||
try:
|
||||
_register_tier_backend()
|
||||
|
|
@ -974,9 +997,7 @@ def test_inherit_builtin_service_tier_pricing_fills_only_missing_fields() -> Non
|
|||
"output_cost_per_token_flex": _TIER_BACKEND_ENTRY["output_cost_per_token_flex"],
|
||||
"input_cost_per_token_balanced": _TIER_BACKEND_ENTRY["input_cost_per_token_balanced"],
|
||||
"output_cost_per_token_balanced": _TIER_BACKEND_ENTRY["output_cost_per_token_balanced"],
|
||||
"cache_read_input_token_cost_ultrafast": _TIER_BACKEND_ENTRY[
|
||||
"cache_read_input_token_cost_ultrafast"
|
||||
],
|
||||
"cache_read_input_token_cost_ultrafast": _TIER_BACKEND_ENTRY["cache_read_input_token_cost_ultrafast"],
|
||||
"input_cost_per_token_above_272k_tokens_ultrafast": _TIER_BACKEND_ENTRY[
|
||||
"input_cost_per_token_above_272k_tokens_ultrafast"
|
||||
],
|
||||
|
|
@ -991,8 +1012,7 @@ def test_inherit_builtin_service_tier_pricing_fills_only_missing_fields() -> Non
|
|||
|
||||
def test_inherit_builtin_service_tier_pricing_noop_without_base_rate_or_backend() -> None:
|
||||
model_cost_entries: Final = {
|
||||
key: copy.deepcopy(litellm.model_cost.get(key))
|
||||
for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL)
|
||||
key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL)
|
||||
}
|
||||
try:
|
||||
_register_tier_backend()
|
||||
|
|
@ -1029,8 +1049,7 @@ def test_inherit_builtin_service_tier_pricing_noop_without_base_rate_or_backend(
|
|||
def test_router_completion_uses_custom_standard_and_backend_ultrafast_pricing() -> None:
|
||||
model_id: Final = "tier-priced-deployment"
|
||||
model_cost_entries: Final = {
|
||||
key: copy.deepcopy(litellm.model_cost.get(key))
|
||||
for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id)
|
||||
key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id)
|
||||
}
|
||||
try:
|
||||
_register_tier_backend()
|
||||
|
|
@ -1090,8 +1109,7 @@ def test_router_completion_uses_custom_standard_and_backend_ultrafast_pricing()
|
|||
def test_router_completion_uses_backend_ultrafast_long_context_rates() -> None:
|
||||
model_id: Final = "tier-priced-long-context-deployment"
|
||||
model_cost_entries: Final = {
|
||||
key: copy.deepcopy(litellm.model_cost.get(key))
|
||||
for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id)
|
||||
key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id)
|
||||
}
|
||||
try:
|
||||
_register_tier_backend()
|
||||
|
|
@ -1142,8 +1160,7 @@ def test_ptu_service_tier_pricing_is_disabled_only_when_attribution_is_enabled(
|
|||
) -> None:
|
||||
model_id: Final = f"ptu-tier-deployment-{ptu_enabled}"
|
||||
model_cost_entries: Final = {
|
||||
key: copy.deepcopy(litellm.model_cost.get(key))
|
||||
for key in (_TIER_BACKEND_KEY, model_id)
|
||||
key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, model_id)
|
||||
}
|
||||
try:
|
||||
_register_tier_backend()
|
||||
|
|
@ -1200,8 +1217,7 @@ def test_azure_base_model_inherits_service_tier_pricing_for_registration_and_pay
|
|||
model_id: Final = "azure-tier-priced-alias"
|
||||
payload_id: Final = "azure-tier-priced-payload"
|
||||
model_cost_entries: Final = {
|
||||
key: copy.deepcopy(litellm.model_cost.get(key))
|
||||
for key in (_AZURE_TIER_BACKEND_KEY, model_id, payload_id)
|
||||
key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_AZURE_TIER_BACKEND_KEY, model_id, payload_id)
|
||||
}
|
||||
try:
|
||||
_register_azure_tier_backend()
|
||||
|
|
@ -1257,12 +1273,8 @@ def test_azure_base_model_inherits_service_tier_pricing_for_registration_and_pay
|
|||
)
|
||||
)
|
||||
|
||||
assert payload["input_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY[
|
||||
"input_cost_per_token_priority"
|
||||
]
|
||||
assert payload["output_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY[
|
||||
"output_cost_per_token_priority"
|
||||
]
|
||||
assert payload["input_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY["input_cost_per_token_priority"]
|
||||
assert payload["output_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY["output_cost_per_token_priority"]
|
||||
finally:
|
||||
_restore_model_cost_entries(model_cost_entries)
|
||||
litellm.get_model_info.cache_clear()
|
||||
|
|
@ -1364,12 +1376,10 @@ def test_inherit_builtin_base_rates_for_off_peak_carries_threshold_rates():
|
|||
|
||||
assert model_info["input_cost_per_token"] == builtin_info["input_cost_per_token"]
|
||||
assert (
|
||||
model_info["input_cost_per_token_above_200k_tokens"]
|
||||
== builtin_info["input_cost_per_token_above_200k_tokens"]
|
||||
model_info["input_cost_per_token_above_200k_tokens"] == builtin_info["input_cost_per_token_above_200k_tokens"]
|
||||
)
|
||||
assert (
|
||||
model_info["output_cost_per_token_above_200k_tokens"]
|
||||
== builtin_info["output_cost_per_token_above_200k_tokens"]
|
||||
model_info["output_cost_per_token_above_200k_tokens"] == builtin_info["output_cost_per_token_above_200k_tokens"]
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -2260,8 +2270,7 @@ def test_deployment_model_cost_payload_folds_in_litellm_params_pricing():
|
|||
def test_deployment_model_cost_payload_includes_builtin_service_tier_pricing() -> None:
|
||||
model_id: Final = "tier-priced-payload"
|
||||
model_cost_entries: Final = {
|
||||
key: copy.deepcopy(litellm.model_cost.get(key))
|
||||
for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id)
|
||||
key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id)
|
||||
}
|
||||
try:
|
||||
_register_tier_backend()
|
||||
|
|
@ -2278,12 +2287,8 @@ def test_deployment_model_cost_payload_includes_builtin_service_tier_pricing() -
|
|||
)
|
||||
)
|
||||
|
||||
assert (
|
||||
payload["input_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["input_cost_per_token_ultrafast"]
|
||||
)
|
||||
assert (
|
||||
payload["output_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["output_cost_per_token_ultrafast"]
|
||||
)
|
||||
assert payload["input_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["input_cost_per_token_ultrafast"]
|
||||
assert payload["output_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["output_cost_per_token_ultrafast"]
|
||||
assert payload["input_cost_per_token_balanced"] == _TIER_BACKEND_ENTRY["input_cost_per_token_balanced"]
|
||||
assert payload["input_cost_per_token"] == _CUSTOM_STANDARD_INPUT_RATE
|
||||
assert payload["output_cost_per_token"] == _CUSTOM_STANDARD_OUTPUT_RATE
|
||||
|
|
@ -3282,9 +3287,7 @@ def test_deployment_id_colliding_with_another_providers_catalog_key_keeps_its_ow
|
|||
)
|
||||
|
||||
assert isinstance(response, litellm.ModelResponse)
|
||||
assert response._hidden_params["response_cost"] == pytest.approx(
|
||||
1000 * 0.00000096 + 500 * 0.00000302
|
||||
)
|
||||
assert response._hidden_params["response_cost"] == pytest.approx(1000 * 0.00000096 + 500 * 0.00000302)
|
||||
assert litellm.model_cost[_COLLIDING_BUILTIN_KEY] == builtin_row_before
|
||||
finally:
|
||||
_restore_model_cost_entries(model_cost_entries)
|
||||
|
|
@ -3297,8 +3300,7 @@ def test_deployment_id_matching_its_own_providers_catalog_key_still_merges() ->
|
|||
matches against."""
|
||||
model_id: Final = "baseten/zai-org/glm-5.2"
|
||||
model_cost_entries: Final = {
|
||||
key: copy.deepcopy(litellm.model_cost.get(key))
|
||||
for key in (model_id, _COLLIDING_BUILTIN_KEY)
|
||||
key: copy.deepcopy(litellm.model_cost.get(key)) for key in (model_id, _COLLIDING_BUILTIN_KEY)
|
||||
}
|
||||
try:
|
||||
router: Final = Router(model_list=[_colliding_id_deployment(model_id, "baseten")])
|
||||
|
|
@ -3313,9 +3315,244 @@ def test_deployment_id_matching_its_own_providers_catalog_key_still_merges() ->
|
|||
)
|
||||
|
||||
assert isinstance(response, litellm.ModelResponse)
|
||||
assert response._hidden_params["response_cost"] == pytest.approx(
|
||||
1000 * 0.00000096 + 500 * 0.00000302
|
||||
)
|
||||
assert response._hidden_params["response_cost"] == pytest.approx(1000 * 0.00000096 + 500 * 0.00000302)
|
||||
finally:
|
||||
_restore_model_cost_entries(model_cost_entries)
|
||||
litellm.get_model_info.cache_clear()
|
||||
|
||||
|
||||
_CAPACITY_TEST_DEPLOYMENTS: Final = (
|
||||
(
|
||||
"bedrock",
|
||||
"bedrock/converse/anthropic.claude-router-capacity-test",
|
||||
"bedrock/us-east-1/anthropic.claude-router-capacity-test",
|
||||
),
|
||||
(
|
||||
"bedrock_mantle",
|
||||
"bedrock_mantle/openai.router-capacity-test",
|
||||
"bedrock_mantle/us-east-1/openai.router-capacity-test",
|
||||
),
|
||||
("openai", "openai/router-capacity-test", "openai/us/router-capacity-test"),
|
||||
)
|
||||
|
||||
|
||||
def _capacity_test_metadata(provider: str, max_input_tokens: int) -> dict[str, object]:
|
||||
return {
|
||||
"litellm_provider": provider,
|
||||
"mode": "chat",
|
||||
"max_input_tokens": max_input_tokens,
|
||||
"max_output_tokens": 512,
|
||||
"max_tokens": 512,
|
||||
"input_cost_per_token": 0,
|
||||
"output_cost_per_token": 0,
|
||||
}
|
||||
|
||||
|
||||
def _capacity_test_router(route: str, base: str, base_location: str | None, explicit_limit: int | None) -> Router:
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "capacity-test",
|
||||
"litellm_params": {
|
||||
"model": route,
|
||||
**({"base_model": base} if base_location == "litellm_params" else {}),
|
||||
},
|
||||
"model_info": {
|
||||
"id": "capacity-test",
|
||||
**({"base_model": base} if base_location == "model_info" else {}),
|
||||
**({"max_input_tokens": explicit_limit} if explicit_limit is not None else {}),
|
||||
},
|
||||
}
|
||||
],
|
||||
enable_pre_call_checks=True,
|
||||
)
|
||||
|
||||
|
||||
def _deployments_allowed_for_2048_input_tokens(router: Router) -> int:
|
||||
try:
|
||||
return len(
|
||||
router._pre_call_checks(
|
||||
model="capacity-test",
|
||||
healthy_deployments=router.model_list,
|
||||
input="synthetic input",
|
||||
input_token_count=2048,
|
||||
)
|
||||
)
|
||||
except litellm.ContextWindowExceededError:
|
||||
return 0
|
||||
|
||||
|
||||
def _effective_max_input_tokens(router: Router) -> object:
|
||||
return router.get_router_model_info(id="capacity-test", deployment=None, received_model_name="capacity-test")[
|
||||
"max_input_tokens"
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS)
|
||||
@pytest.mark.parametrize("base_location", ("model_info", "litellm_params", None))
|
||||
@pytest.mark.parametrize("explicit_limit", (None, 4096))
|
||||
def test_router_capacity_widens_to_a_larger_base_model_window(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
provider: str,
|
||||
route: str,
|
||||
base: str,
|
||||
base_location: str | None,
|
||||
explicit_limit: int | None,
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"model_cost",
|
||||
{
|
||||
**litellm.model_cost,
|
||||
route: _capacity_test_metadata(provider, 1024),
|
||||
base: _capacity_test_metadata(provider, 8192),
|
||||
},
|
||||
)
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
router: Final = _capacity_test_router(route, base, base_location, explicit_limit)
|
||||
expected: Final = (
|
||||
explicit_limit if explicit_limit is not None else (1024 if base_location is None else 8192)
|
||||
)
|
||||
assert _effective_max_input_tokens(router) == expected
|
||||
assert _deployments_allowed_for_2048_input_tokens(router) == (1 if expected >= 2048 else 0)
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS)
|
||||
@pytest.mark.parametrize("base_location", ("model_info", "litellm_params"))
|
||||
def test_router_capacity_never_narrows_to_a_smaller_base_model_window(
|
||||
monkeypatch: pytest.MonkeyPatch, provider: str, route: str, base: str, base_location: str
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"model_cost",
|
||||
{
|
||||
**litellm.model_cost,
|
||||
route: _capacity_test_metadata(provider, 8192),
|
||||
base: _capacity_test_metadata(provider, 1024),
|
||||
},
|
||||
)
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
router: Final = _capacity_test_router(route, base, base_location, None)
|
||||
assert _effective_max_input_tokens(router) == 8192
|
||||
assert _deployments_allowed_for_2048_input_tokens(router) == 1
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS)
|
||||
def test_router_capacity_never_takes_a_window_from_base_model_when_the_route_is_unmapped(
|
||||
monkeypatch: pytest.MonkeyPatch, provider: str, route: str, base: str
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, base: _capacity_test_metadata(provider, 1024)})
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
router: Final = _capacity_test_router(f"{route}-not-in-cost-map", base, "model_info", None)
|
||||
assert _deployments_allowed_for_2048_input_tokens(router) == 1
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS)
|
||||
def test_router_capacity_never_takes_a_window_from_base_model_when_the_route_declares_none(
|
||||
monkeypatch: pytest.MonkeyPatch, provider: str, route: str, base: str
|
||||
) -> None:
|
||||
route_without_window: Final = {
|
||||
key: value
|
||||
for key, value in _capacity_test_metadata(provider, 0).items()
|
||||
if key not in ("max_input_tokens", "max_tokens")
|
||||
}
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"model_cost",
|
||||
{**litellm.model_cost, route: route_without_window, base: _capacity_test_metadata(provider, 1024)},
|
||||
)
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
router: Final = _capacity_test_router(route, base, "model_info", None)
|
||||
assert _effective_max_input_tokens(router) is None
|
||||
assert _deployments_allowed_for_2048_input_tokens(router) == 1
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
|
||||
_SAME_MODEL_ROUTE: Final = "bedrock/converse/anthropic.claude-routercapacitytest-9-9"
|
||||
_SAME_MODEL_GLOBAL_ROUTE: Final = "bedrock/converse/global.anthropic.claude-routercapacitytest-9-9"
|
||||
_SAME_MODEL_QUALIFIED_ROUTE: Final = "bedrock/converse/anthropic.claude-routercapacitytest-9-9@first"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route", "base_model", "expected"),
|
||||
(
|
||||
(_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9", 8192),
|
||||
(_SAME_MODEL_ROUTE, "us-east-1/anthropic.claude-routercapacitytest-9-9", 1024),
|
||||
(_SAME_MODEL_ROUTE, "anthropic.claude-routercapacitytest-9-9", 1024),
|
||||
(_SAME_MODEL_GLOBAL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9", 8192),
|
||||
(_SAME_MODEL_ROUTE, "bedrock/anthropic.claude-routercapacitytest-9-9@global:us-east-1", 8192),
|
||||
(_SAME_MODEL_GLOBAL_ROUTE, "bedrock/anthropic.claude-routercapacitytest-9-9@global:us-east-1", 8192),
|
||||
(_SAME_MODEL_QUALIFIED_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@first", 8192),
|
||||
(_SAME_MODEL_QUALIFIED_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@second", 1024),
|
||||
(_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-*", 1024),
|
||||
(_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9", 1024),
|
||||
(_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9-extended", 1024),
|
||||
(_SAME_MODEL_ROUTE, "bedrock/eu-west-1/anthropic.claude-routercapacitytest-9-9", 1024),
|
||||
),
|
||||
)
|
||||
def test_router_capacity_widens_only_for_an_exact_base_model_entry_of_the_same_model(
|
||||
monkeypatch: pytest.MonkeyPatch, route: str, base_model: str, expected: int
|
||||
) -> None:
|
||||
larger_window: Final = _capacity_test_metadata("bedrock", 8192)
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"model_cost",
|
||||
{
|
||||
**litellm.model_cost,
|
||||
_SAME_MODEL_ROUTE: _capacity_test_metadata("bedrock", 1024),
|
||||
_SAME_MODEL_GLOBAL_ROUTE: _capacity_test_metadata("bedrock", 1024),
|
||||
_SAME_MODEL_QUALIFIED_ROUTE: _capacity_test_metadata("bedrock", 1024),
|
||||
"bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@first": larger_window,
|
||||
"bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@second": larger_window,
|
||||
"bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9": larger_window,
|
||||
"anthropic.claude-routercapacitytest-9-9": _capacity_test_metadata("bedrock", 1024),
|
||||
"bedrock/anthropic.claude-routercapacitytest-9-9": larger_window,
|
||||
"bedrock/anthropic.claude-routercapacitytest-9-9@global:us-east-1": larger_window,
|
||||
"bedrock/us-east-1/anthropic.claude-routercapacitytest-*": larger_window,
|
||||
"bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9": larger_window,
|
||||
"bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9-extended": larger_window,
|
||||
},
|
||||
)
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
router: Final = _capacity_test_router(route, base_model, "model_info", None)
|
||||
assert _effective_max_input_tokens(router) == expected
|
||||
assert _deployments_allowed_for_2048_input_tokens(router) == (1 if expected >= 2048 else 0)
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
|
||||
def test_router_capacity_reads_a_model_info_base_model_before_a_litellm_params_one(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"model_cost",
|
||||
{
|
||||
**litellm.model_cost,
|
||||
_SAME_MODEL_ROUTE: _capacity_test_metadata("bedrock", 1024),
|
||||
"bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9": _capacity_test_metadata("bedrock", 8192),
|
||||
"bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9": _capacity_test_metadata("bedrock", 8192),
|
||||
},
|
||||
)
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "capacity-test",
|
||||
"litellm_params": {
|
||||
"model": _SAME_MODEL_ROUTE,
|
||||
"base_model": "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9",
|
||||
},
|
||||
"model_info": {
|
||||
"id": "capacity-test",
|
||||
"base_model": "bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9",
|
||||
},
|
||||
}
|
||||
],
|
||||
enable_pre_call_checks=True,
|
||||
)
|
||||
assert _effective_max_input_tokens(router) == 1024
|
||||
assert _deployments_allowed_for_2048_input_tokens(router) == 0
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue