fix(proxy): price RAG routes by the model the router's global alias serves

This commit is contained in:
mateo-berri 2026-09-28 19:42:53 -07:00
parent 45d880c4f5
commit b67411f495
3 changed files with 13 additions and 2 deletions

View file

@ -472,6 +472,10 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None
)
def is_router_resolved_model_cost_zero(model: str | list[str] | None, llm_router: Router | None) -> bool:
return _is_model_cost_zero(model=model, llm_router=llm_router)
def is_requested_model_cost_zero(model: str | list[str] | None, llm_router: Router | None) -> bool:
if llm_router is None:
return False

View file

@ -66,6 +66,7 @@ from litellm.proxy.auth.auth_checks import (
get_user_object,
is_dispatched_model_cost_zero,
is_requested_model_cost_zero,
is_router_resolved_model_cost_zero,
is_valid_fallback_model,
jwt_key_mapping_cache_key,
key_model_aliases_for_auth_check,
@ -3160,6 +3161,8 @@ async def _is_dispatched_model_cost_zero(
request: Request | None,
route: str,
) -> bool:
if route in LiteLLMRoutes.litellm_native_routes.value:
return is_router_resolved_model_cost_zero(model=model, llm_router=llm_router)
if not RouteChecks.is_unified_llm_api_route(route):
return is_requested_model_cost_zero(model=model, llm_router=llm_router)

View file

@ -1881,10 +1881,14 @@ async def test_budget_skip_judges_the_model_the_key_aliases_dispatch_to(
("/anthropic/v1/messages", "my-alias", {"my-alias": "free-model"}, None, False),
("/anthropic/v1/messages", "rs-alias", {}, {"model_group_alias": {"rs-alias": "free-model"}}, False),
("/anthropic/v1/messages", "free-model", {"free-model": "paid-model"}, None, True),
("/rag/query", "free-alias", {}, None, False),
("/rag/query", "free-alias", {}, None, True),
("/v1/rag/query", "free-alias", {}, None, True),
("/rag/query", "my-alias", {"my-alias": "free-model"}, None, False),
("/rag/query", "rs-alias", {}, {"model_group_alias": {"rs-alias": "free-model"}}, False),
("/rag/query", "free-model", {"free-model": "paid-model"}, None, True),
],
)
async def test_budget_skip_prices_the_requested_name_on_routes_that_forward_it_unaliased(
async def test_budget_skip_off_the_unified_routes_prices_the_model_each_route_dispatches(
route: str,
model: str,
key_aliases: dict[str, str],