From deb7697a8e515671774b3ccda6dd6d67123e41cc Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Wed, 11 Nov 2026 12:22:12 -0800 Subject: [PATCH] fix(auth): enforce model access on Google and image URL routes (#51226) * fix(auth): authorize URL-selected Google and image models * fix(auth): preserve server model selection and budget fallbacks * test(auth): isolate URL model authorization settings * fix(auth): type handler model selection boundaries * fix(auth): validate request model selection inputs * fix(auth): tolerate absent request path metadata * fix(auth): support requests without path metadata --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- litellm/proxy/auth/auth_utils.py | 46 +++++++- litellm/proxy/auth/user_api_key_auth.py | 3 + .../proxy/common_utils/http_parsing_utils.py | 6 ++ litellm/proxy/proxy_server.py | 4 + .../spend_tracking/budget_reservation.py | 24 ++++- tests/unit/proxy/auth/test_auth_utils.py | 101 ++++++++++++++++++ .../unit/proxy/auth/test_user_api_key_auth.py | 48 +++++++++ .../test_user_api_key_auth_request_flow.py | 40 +++++++ .../common_utils/test_http_parsing_utils.py | 20 ++++ .../spend_tracking/test_budget_reservation.py | 34 ++++++ 10 files changed, 320 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 7413490e44f..900e993cadb 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -37,7 +37,11 @@ from litellm.llms.azure.passthrough.transformation import azure_router_model_in_ from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path from litellm.proxy._types import * from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.common_utils.http_parsing_utils import extract_nested_form_metadata +from litellm.proxy.common_utils.http_parsing_utils import ( + extract_nested_form_metadata, + read_request_path_params, + resolve_inference_model, +) from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_ENDPOINT_MARKER, ) @@ -2672,6 +2676,7 @@ def get_model_from_request( llm_router: Router | None = None, request: Request | None = None, team_id: str | None = None, + path_params: Mapping[str, object] | None = None, ) -> str | list[str] | None: """Resolve the model(s) a request targets, for model-access and budget checks. @@ -2690,6 +2695,12 @@ def get_model_from_request( if request_dispatched_to_pass_through_endpoint(request): return None + handler_request_data: Final = _MODEL_REQUEST_DATA.validate_python(request_data) + handler_path_params: Final = read_request_path_params(request) if request is not None else path_params + handler_model: Final = _handler_selected_model(route, handler_request_data, handler_path_params) + if handler_model is not None: + return handler_model + candidates: Final = _extract_model_candidates_from_request( request_data=request_data, route=route, @@ -2719,6 +2730,39 @@ def get_model_from_request( return model +_HandlerModelKind: TypeAlias = Literal["completion", "image_generation", "image_edit"] +_MODEL_REQUEST_DATA: Final = TypeAdapter(Mapping[str, object]) +_HANDLER_MODEL_ROUTES: Final[tuple[tuple[re.Pattern[str], str, _HandlerModelKind], ...]] = ( + (re.compile(r"^(?:/v1beta)?/models/(.+):(?:generateContent|streamGenerateContent)$"), "model_name", "completion"), + (re.compile(r"^/openai/deployments/(.+)/images/generations$"), "model", "image_generation"), + (re.compile(r"^/openai/deployments/(.+)/images/edits$"), "model", "image_edit"), +) + + +def _handler_selected_model( + route: str, request_data: Mapping[str, object], path_params: Mapping[str, object] | None +) -> str | None: + selection: Final[tuple[re.Match[str], str, _HandlerModelKind] | None] = next( + ((match, name, kind) for pattern, name, kind in _HANDLER_MODEL_ROUTES if (match := pattern.match(route))), + None, + ) + if selection is None: + return None + from litellm.proxy.proxy_server import general_settings_view, user_model_view + + route_match, path_param_name, kind = selection + route_model: Final = route_match.group(1) + path_model: Final = path_params.get(path_param_name) if path_params is not None else None + model: Final = resolve_inference_model( + request_data.get("model"), + general_settings_view(), + user_model_view(), + path_model if isinstance(path_model, str) and path_model else route_model, + kind=kind, + ) + return model if isinstance(model, str) and model else route_model + + def _model_from_route_path(route: str) -> str | None: """ Google generateContent-style routes allow "/" inside the model id, e.g. diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index d7a8d4a0700..1b2e4e791e3 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -153,6 +153,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg populate_request_with_path_params, read_raw_json_body, read_request_body, + read_request_path_params, resolve_realtime_route_model, rewrite_request_model, safe_get_request_headers, @@ -4252,6 +4253,7 @@ async def _reserve_budget_after_common_checks( reserve_budget_for_request, ) + reservation_path_params: Final = read_request_path_params(request) user_api_key_auth_obj.budget_reservation = await reserve_budget_for_request( request_body=request_data, route=route, @@ -4270,6 +4272,7 @@ async def _reserve_budget_after_common_checks( or user_api_key_auth_obj.billing_agent_policy is not None ), raw_body=await read_raw_json_body(request=request), + path_params=reservation_path_params, ) if request is not None: reservation: Final = user_api_key_auth_obj.budget_reservation diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 80ba2f49745..2786b4271e4 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -37,6 +37,12 @@ _BINARY_CONTENT_TYPES: Final[frozenset[str]] = frozenset({"application/x-protobu _ANNOTATION_QUALIFIERS: Final[frozenset[object]] = frozenset({Annotated, NotRequired, ReadOnly, Required}) _DECODED_JSON: Final = TypeAdapter(object) _JSON_REQUEST_BODY: Final = TypeAdapter(Mapping[str, object]) +_REQUEST_PATH_PARAMS: Final = TypeAdapter(Mapping[str, object]) + + +def read_request_path_params(request: Request | None) -> Mapping[str, object] | None: + path_params: Final = _DECODED_JSON.validate_python(getattr(request, "path_params", None)) + return _REQUEST_PATH_PARAMS.validate_python(path_params) if isinstance(path_params, Mapping) else None class _ReadableUpload(Protocol): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 23d45e07dfa..6fbe562b5d1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3183,6 +3183,10 @@ def general_settings_view() -> Mapping[str, object]: return _GENERAL_SETTINGS_VIEW.validate_python(general_settings) +def user_model_view() -> str | None: + return user_model if isinstance(user_model, str) else None + + config_passthrough_endpoints: list[dict[str, object]] | None = None log_file: Final = "api_log.json" worker_config: Final = None diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 8f0584bd0d0..096578264e3 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -262,12 +262,13 @@ async def reserve_budget_for_request( apply_user_budget_to_team_keys: bool = False, fail_closed_budget_enforcement: bool = False, raw_body: bytes | None = None, + path_params: Mapping[str, object] | None = None, ) -> dict | None: if valid_token is None or not RouteChecks.is_llm_api_route(route=route): return None if RouteChecks.is_unbilled_route(route): return None - if get_model_from_request(request_body, route, llm_router=llm_router) is None: + if get_model_from_request(request_body, route, llm_router=llm_router, path_params=path_params) is None: return None counters: Final = await _get_budget_counters( @@ -291,6 +292,7 @@ async def reserve_budget_for_request( llm_router=llm_router, team_id=valid_token.team_id, raw_body=raw_body, + path_params=path_params, ) reservation_cost = estimate_request_max_cost( @@ -298,6 +300,7 @@ async def reserve_budget_for_request( route=route, llm_router=llm_router, input_token_counts=input_token_counts, + path_params=path_params, ) # estimate_request_max_cost still returns None when the model is unknown # to the cost map (no token-priced cost fields, e.g. image/audio routes). @@ -334,6 +337,7 @@ async def reserve_budget_for_request( route=route, llm_router=llm_router, input_token_counts=input_token_counts, + path_params=path_params, ) budget_reservation: Final = { "reserved_cost": reservation_cost, @@ -1319,6 +1323,7 @@ def estimate_request_max_cost( route: str, llm_router: Router | None, input_token_counts: Mapping[str, int] | None = None, + path_params: Mapping[str, object] | None = None, ) -> float | None: estimates = [ _estimate_request_max_cost_for_model( @@ -1328,7 +1333,9 @@ def estimate_request_max_cost( llm_router=llm_router, input_tokens=(input_token_counts or {}).get(model_name), ) - for model_name in _get_request_models(request_body=request_body, route=route, llm_router=llm_router) + for model_name in _get_request_models( + request_body=request_body, route=route, llm_router=llm_router, path_params=path_params + ) ] estimates = [estimate for estimate in estimates if estimate is not None] if not estimates: @@ -1341,6 +1348,7 @@ def estimate_request_input_cost( route: str, llm_router: Router | None, input_token_counts: Mapping[str, int] | None = None, + path_params: Mapping[str, object] | None = None, ) -> float | None: """Cost of the request's input tokens alone. @@ -1357,7 +1365,9 @@ def estimate_request_input_cost( llm_router=llm_router, input_tokens=(input_token_counts or {}).get(model_name), ) - for model_name in _get_request_models(request_body=request_body, route=route, llm_router=llm_router) + for model_name in _get_request_models( + request_body=request_body, route=route, llm_router=llm_router, path_params=path_params + ) ] estimates = [estimate for estimate in estimates if estimate is not None] if not estimates: @@ -1609,8 +1619,9 @@ def _get_request_models( request_body: dict, route: str, llm_router: Router | None, + path_params: Mapping[str, object] | None = None, ) -> Sequence[str]: - model: Final = get_model_from_request(request_body, route, llm_router=llm_router) + model: Final = get_model_from_request(request_body, route, llm_router=llm_router, path_params=path_params) if model is None: return () return (model,) if isinstance(model, str) else tuple(model) @@ -1622,11 +1633,14 @@ async def count_request_input_tokens( llm_router: Router | None, team_id: str | None, raw_body: bytes | None = None, + path_params: Mapping[str, object] | None = None, ) -> Mapping[str, int]: """Input-token count per candidate model, counted once per request. The counts are reused by both the max-cost and the input-cost estimate.""" - models: Final = _get_request_models(request_body=request_body, route=route, llm_router=llm_router) + models: Final = _get_request_models( + request_body=request_body, route=route, llm_router=llm_router, path_params=path_params + ) if not models: return MappingProxyType({}) counted_models: Final = MappingProxyType( diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 73a2c534628..58a78aede3c 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -1379,6 +1379,107 @@ def test_get_model_from_request_supports_google_model_names_with_slashes(): ) +_URL_SELECTED_MODEL_ROUTES: Final = ( + ("/v1beta/models/premium-model:generateContent", "premium-model"), + ("/models/premium-model:generateContent", "premium-model"), + ("/v1beta/models/premium-model:streamGenerateContent", "premium-model"), + ("/models/provider/premium-model:7b:streamGenerateContent", "provider/premium-model:7b"), + ("/openai/deployments/premium-model/images/generations", "premium-model"), + ("/openai/deployments/provider/premium-model/images/edits", "provider/premium-model"), +) + + +@pytest.fixture +def no_inference_model_override(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "user_model", None) + + +@pytest.mark.usefixtures("no_inference_model_override") +@pytest.mark.parametrize("route, url_model", _URL_SELECTED_MODEL_ROUTES) +@pytest.mark.parametrize( + "request_data", + [ + {"model": "cheap-model"}, + {"target_model_names": ["cheap-model"]}, + {"model": "cheap-model", "target_model_names": "cheap-model"}, + ], +) +def test_get_model_from_request_checks_the_url_model_on_routes_served_from_the_url( + route: str, url_model: str, request_data: dict[str, object] +) -> None: + assert get_model_from_request(request_data=request_data, route=route) == url_model + + +@pytest.mark.usefixtures("no_inference_model_override") +@pytest.mark.parametrize( + "route, path_params", + [ + ("/v1beta/models/premium-model:generateContent", {"model_name": "cheap-model"}), + ("/models/premium-model:streamGenerateContent", {"model_name": "cheap-model"}), + ("/openai/deployments/premium-model/images/generations", {"model": "cheap-model"}), + ("/openai/deployments/premium-model/images/edits", {"model": "cheap-model"}), + ], +) +def test_get_model_from_request_follows_the_handlers_path_param_over_route_and_body( + route: str, path_params: dict[str, str] +) -> None: + request: Final = Request(scope={"type": "http", "path": route, "path_params": path_params, "headers": []}) + body: Final = {"model": "body-model"} + + assert get_model_from_request(request_data=body, route=route, request=request) == "cheap-model" + assert get_model_from_request(request_data=body, route=route, path_params=path_params) == "cheap-model" + assert get_model_from_request(request_data=body, route=route, path_params={"other": "x"}) == "premium-model" + + +@pytest.mark.parametrize( + "route, settings, cli_model, expected", + [ + ("/v1beta/models/url-model:generateContent", {"completion_model": "configured"}, None, "configured"), + ("/models/url-model:streamGenerateContent", {}, "cli", "cli"), + ("/v1beta/models/url-model:generateContent", {"image_generation_model": "configured"}, None, "url-model"), + ("/openai/deployments/url-model/images/generations", {}, "cli", "cli"), + ("/openai/deployments/url-model/images/generations", {"completion_model": "configured"}, None, "url-model"), + ( + "/openai/deployments/url-model/images/generations", + {"image_generation_model": "configured"}, + None, + "url-model", + ), + ("/openai/deployments/url-model/images/edits", {"completion_model": "configured"}, "cli", "configured"), + ("/openai/deployments/url-model/images/edits", {}, "cli", "cli"), + ("/openai/deployments/url-model/images/edits", {"image_generation_model": "configured"}, None, "url-model"), + ], +) +def test_get_model_from_request_applies_the_handlers_server_model_precedence( + monkeypatch: pytest.MonkeyPatch, + route: str, + settings: dict[str, str], + cli_model: str | None, + expected: str, +) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", settings) + monkeypatch.setattr(proxy_server, "user_model", cli_model) + + assert get_model_from_request(request_data={"model": "body-model"}, route=route) == expected + + +@pytest.mark.usefixtures("no_inference_model_override") +@pytest.mark.parametrize("route", ["/v1/images/generations", "/images/edits", "/v1/chat/completions"]) +def test_get_model_from_request_keeps_the_body_model_where_the_url_names_none(route: str) -> None: + assert get_model_from_request(request_data={"model": "cheap-model"}, route=route) == "cheap-model" + + +@pytest.mark.parametrize("route, url_model", _URL_SELECTED_MODEL_ROUTES) +def test_get_model_from_request_skips_url_model_on_configured_pass_through(route: str, url_model: str) -> None: + request: Final = _request_dispatched_to(_pass_through_endpoint()) + assert get_model_from_request(request_data={"model": url_model}, route=route, request=request) is None + + def test_get_model_from_request_vertex_passthrough_still_works(): route = "/vertex_ai/v1/projects/p/locations/l/publishers/google/models/gemini-1.5-pro:generateContent" assert get_model_from_request(request_data={}, route=route) == "gemini-1.5-pro" diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index ce4a844550b..7746fc96673 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -1293,6 +1293,54 @@ def test_get_model_from_request(route, request_data, expected_model): assert get_model_from_request(request_data, route) == expected_model +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route_template", + [ + "/v1beta/models/{model}:generateContent", + "/models/{model}:streamGenerateContent", + "/openai/deployments/{model}/images/generations", + "/openai/deployments/{model}/images/edits", + ], +) +@pytest.mark.parametrize("body", [{"model": "cheap-model"}, {"target_model_names": ["cheap-model"]}]) +async def test_key_model_allowlist_applies_to_the_model_selected_by_the_url( + monkeypatch: pytest.MonkeyPatch, route_template: str, body: dict[str, object] +) -> None: + from litellm.proxy._types import ProxyErrorTypes + from litellm.proxy.auth.user_api_key_auth import enforce_key_and_fallback_model_access + + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", None) + monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {}) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_model", None) + router: Final = litellm.Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": f"openai/{name}", "api_key": "sk-fake"}} + for name in ("cheap-model", "premium-model") + ] + ) + token: Final = UserAPIKeyAuth(models=["cheap-model"]) + + async def authorize(url_model: str) -> None: + route: Final = route_template.format(model=url_model) + request: Final = StarletteRequest( + {"type": "http", "method": "POST", "path": route, "headers": [], "query_string": b""} + ) + await enforce_key_and_fallback_model_access( + valid_token=token, + request_data=dict(body), + route=route, + request=request, + llm_model_list=list(router.model_list), + llm_router=router, + ) + + await authorize("cheap-model") + with pytest.raises(ProxyException) as exc_info: + await authorize("premium-model") + assert exc_info.value.type == ProxyErrorTypes.key_model_access_denied + + @pytest.mark.asyncio async def test_jwt_non_admin_team_route_access(monkeypatch): """ diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 0d2433d9707..5824b5fccab 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -10952,6 +10952,42 @@ class TestCheckTeamMemberModelBudgetWithFallback: assert request_data["model"] == "gpt-4o-mini" + @pytest.mark.asyncio + @pytest.mark.parametrize("operation", ["generations", "edits"]) + async def test_image_deployment_key_fallback_is_the_model_later_checks_and_reservation_see( + self, monkeypatch: pytest.MonkeyPatch, operation: str + ) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.auth.auth_utils import get_model_from_request + + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "user_model", None) + route: Final = f"/openai/deployments/gpt-4o/images/{operation}" + valid_token: Final = UserAPIKeyAuth( + token="test-key", + models=["gpt-4o", "gpt-4o-mini"], + budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}, + model_max_budget={ + "gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}, + "gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ) + limiter, _ = await _real_model_budget_limiter(key_spend={"gpt-4o": 10.0}) + request_data: Final[dict[str, object]] = {"model": "gpt-4o"} + request: Final = self._request("gpt-4o") + assert get_model_from_request(request_data, route, request=request) == "gpt-4o" + + await _check_key_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + model_name="gpt-4o", + request_data=request_data, + request=request, + ) + + assert get_model_from_request(request_data, route, request=request) == "gpt-4o-mini" + assert get_model_from_request(request_data, route, path_params=request.path_params) == "gpt-4o-mini" + @pytest.mark.asyncio async def test_fallback_outside_member_allowed_models_is_rejected(self): valid_token: Final = UserAPIKeyAuth( @@ -13440,6 +13476,10 @@ async def test_managed_actor_cannot_access_provider_resource_routes(monkeypatch, ("/v1/chat/completions", {}, "forbidden-model"), ("/v1/images/generations", {"image_generation_model": "forbidden-model"}, None), ("/v1/images/edits", {"image_generation_model": "forbidden-model"}, None), + ("/v1beta/models/test-model:generateContent", {"completion_model": "forbidden-model"}, None), + ("/models/test-model:streamGenerateContent", {}, "forbidden-model"), + ("/openai/deployments/test-model/images/generations", {}, "forbidden-model"), + ("/openai/deployments/test-model/images/edits", {"completion_model": "forbidden-model"}, None), ], ) async def test_managed_agent_cannot_bypass_grants_with_server_default( diff --git a/tests/unit/proxy/common_utils/test_http_parsing_utils.py b/tests/unit/proxy/common_utils/test_http_parsing_utils.py index 8423dee9886..b8a686d3f67 100644 --- a/tests/unit/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/unit/proxy/common_utils/test_http_parsing_utils.py @@ -31,6 +31,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( numeric_form_fields, populate_request_with_path_params, read_raw_json_body, + read_request_path_params, ) from fastapi import Request as Request_http_parsing from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -1786,3 +1787,22 @@ async def test_convert_upload_files_to_file_data_reads_every_upload_into_file_tu "tags": ["red", "blue"], "display_title": "My skill", } + + +@pytest.mark.parametrize( + "scope_path_params, expected", + [ + ({"model": "deployment/name"}, {"model": "deployment/name"}), + ({}, {}), + (None, None), + ], +) +def test_read_request_path_params_returns_mappings_and_drops_non_mapping_metadata( + scope_path_params: Mapping[str, object] | None, expected: Mapping[str, object] | None +) -> None: + request: Final = Request({"type": "http", "path_params": scope_path_params}) + assert read_request_path_params(request) == expected + + +def test_read_request_path_params_without_request_is_none() -> None: + assert read_request_path_params(None) is None diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 72927fa0197..087434c981f 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -223,6 +223,40 @@ def test_deployment_pricing_update_invalidates_cached_estimate() -> None: assert math.isclose(after, before * 1000) +def _priced_group(model_name: str, cost_per_token: float) -> dict[str, object]: + return Deployment( + model_name=model_name, + litellm_params=LiteLLM_Params(model="dashscope/qwen3-max", api_key="sk-fake"), + model_info=ModelInfo( + id=model_name, + max_output_tokens=1000, + tiered_pricing=[ + {"input_cost_per_token": cost_per_token, "output_cost_per_token": cost_per_token, "range": [0, 128000]} + ], + ), + ).model_dump() + + +def test_google_route_reserves_on_the_model_its_path_param_selects(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "user_model", None) + router: Final = Router(model_list=[_priced_group("cheap-group", 1e-06), _priced_group("premium-group", 1e-03)]) + route: Final = "/v1beta/models/premium-group:generateContent" + body: Final = {"model": "cheap-group", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 10} + + url_model_cost: Final = estimate_request_max_cost(request_body=body, route=route, llm_router=router) + path_param_cost: Final = estimate_request_max_cost( + request_body=body, route=route, llm_router=router, path_params={"model_name": "cheap-group"} + ) + cheap_cost: Final = estimate_request_max_cost( + request_body=body, route="/v1beta/models/cheap-group:generateContent", llm_router=router + ) + + assert url_model_cost is not None and cheap_cost is not None + assert path_param_cost == cheap_cost + assert math.isclose(url_model_cost, cheap_cost * 1000) + + @pytest.mark.asyncio @pytest.mark.parametrize( "expiry_offset, expected_max_budget",