From 6c4fd246d37641bfd6824c6972ab6ad48f2758aa Mon Sep 17 00:00:00 2001 From: Kent Date: Mon, 14 Sep 2026 12:06:17 +0800 Subject: [PATCH] refactor(batches): drop the alias swap now that credential resolution returns the deployment model get_deployment_credentials_with_provider returns credentials["model"] and prepare_data_with_credentials does data.update(credentials), so the proxy batch-create path already hands litellm.acreate_batch the deployment's real provider model. _swap_alias_for_deployment_model and Router.get_deployment_model_for_alias were doing that same swap a second time. What is left is the part that is not upstream: batch endpoints pass the caller's team_id into credential resolution, and _resolve_unblocked_deployment applies the team guard to exact deployment-id lookups too. Also regenerates the dashboard API types, which the base left stale. --- litellm/proxy/batches_endpoints/endpoints.py | 43 +------ litellm/router.py | 18 --- .../proxy/batches_endpoints/test_endpoints.py | 8 -- .../test_batch_x_litellm_model_encoding.py | 105 ------------------ tests/test_litellm/test_router.py | 56 +--------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 - 6 files changed, 7 insertions(+), 225 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 08be760ca77..808a40eb02c 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -7,7 +7,7 @@ import asyncio import os from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, cast +from typing import Any, Final, cast from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response @@ -49,9 +49,6 @@ from litellm.proxy.utils import handle_exception_on_proxy, is_known_model from litellm.repositories.table_repositories import ManagedFileRepository from litellm.types.llms.openai import LiteLLMBatchCreateRequest -if TYPE_CHECKING: - from litellm.router import Router - router: Final = APIRouter() @@ -72,32 +69,6 @@ def _raise_not_found_when_openai_fallback_unservable( ) -def _swap_alias_for_deployment_model( - create_batch_data: LiteLLMBatchCreateRequest, - alias: str, - llm_router: "Router | None", - team_id: "str | None", -) -> None: - """ - Replace a proxy model-group alias on the batch request with the - deployment's real provider model (in place). - - ``litellm.create_batch`` runs the model through ``get_llm_provider``, which - cannot resolve a proxy alias, so a provider transform (e.g. Bedrock's, which - forwards ``model`` as the batch ``modelId``) would otherwise receive the - alias and the provider would reject it. Falls back to the alias when the - router is unavailable or the alias resolves to nothing. ``team_id`` keeps - this lookup on the same team-usable deployment the credential resolver - picked, so a team-owned deployment sharing the alias can't leak its model - to callers outside that team. - """ - if llm_router is None: - return - resolved_model: Final = llm_router.get_deployment_model_for_alias(model_id=alias, team_id=team_id) - if resolved_model is not None: - create_batch_data["model"] = resolved_model - - async def _resolve_managed_input_file_storage_url(input_file_id: str) -> "str | None": """Resolve a managed (unified) input_file_id to its backend storage_url. @@ -260,12 +231,6 @@ async def create_batch( data=_create_batch_data, credentials=credentials, ) - _swap_alias_for_deployment_model( - create_batch_data=_create_batch_data, - alias=model_from_file_id, - llm_router=llm_router, - team_id=user_api_key_dict.team_id, - ) # Create batch using model credentials response = await litellm.acreate_batch( @@ -357,12 +322,6 @@ async def create_batch( data=_create_batch_data, credentials=credentials, ) - _swap_alias_for_deployment_model( - create_batch_data=_create_batch_data, - alias=model_param, - llm_router=llm_router, - team_id=user_api_key_dict.team_id, - ) # Create batch using model credentials response = await litellm.acreate_batch( diff --git a/litellm/router.py b/litellm/router.py index 604197a6f65..269baa937b0 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -10270,24 +10270,6 @@ class Router: return None return deployment - def get_deployment_model_for_alias(self, model_id: str, team_id: "str | None" = None) -> "str | None": - """ - Resolve a model-group alias (or deployment id / wildcard) to the - deployment's underlying ``litellm_params.model``. - - Callers that hand a model to provider SDKs (e.g. the proxy batch-create - path) need the real provider model id, not the proxy alias: - ``get_llm_provider`` cannot resolve an alias, so passing it straight - through reaches the provider as an invalid model identifier. Returns - None when the alias resolves to nothing or to a paused deployment. - Pass the caller's ``team_id`` so the deployment picked here is the same - one the credential resolver picks for that caller. - """ - deployment: Final = self._resolve_unblocked_deployment(model_id=model_id, team_id=team_id) - if deployment is None: - return None - return deployment.litellm_params.model - @staticmethod def _deployment_usable_by_team(model: Mapping | Deployment, team_id: str | None) -> bool: """ diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 3343a2b155c..2e8d259ec90 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -165,13 +165,6 @@ def _creds_lookup(*, model_id: str, team_id: Optional[str] = None) -> Dict[str, return dict(CREDS[model_id]) -def _alias_lookup(*, model_id: str, team_id: Optional[str] = None) -> str: - # The endpoint swaps the request model for the deployment's real provider - # model before calling the provider; mirror that with the CREDS model so a - # wrong/hardcoded model_id KeyErrors instead of hiding. - return CREDS[model_id]["model"] - - @pytest.fixture def harness(): """Seam harness. Patches only true I/O boundaries; pure encode/decode/merge @@ -186,7 +179,6 @@ def harness(): router = MagicMock(spec=Router) router.acreate_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) - router.get_deployment_model_for_alias = MagicMock(side_effect=_alias_lookup) read_body = AsyncMock(side_effect=lambda request: body_holder["body"]) pre_call = AsyncMock(side_effect=lambda **kw: (body_holder["body"], MagicMock())) diff --git a/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py b/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py index 9212b257948..3d1831bb4cd 100644 --- a/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py +++ b/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py @@ -418,111 +418,6 @@ class TestBatchIdRoundTripWithRetrieve: assert get_original_file_id(encoded) == raw_id -@pytest.mark.asyncio -async def test_create_batch_swaps_alias_for_deployment_model_before_provider_call(): - """ - SCENARIO 1 (model-encoded input_file_id): the proxy must hand - litellm.acreate_batch the deployment's real provider model, not the proxy - alias. get_llm_provider cannot resolve an alias, so passing it straight - through reaches the Bedrock batch transform as an invalid modelId. The - response IDs must still be encoded with the ALIAS so retrieve routes back. - """ - from litellm.proxy.batches_endpoints.endpoints import create_batch - from litellm.proxy.openai_files_endpoints.common_utils import ( - encode_file_id_with_model, - ) - - alias = "bedrock-batch-haiku" - real_model = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" - raw_input_file_id = "s3://bucket/litellm-bedrock-files/in.jsonl" - encoded_input_file_id = encode_file_id_with_model( - file_id=raw_input_file_id, model=alias - ) - raw_batch_id = "batch_bedrock_123" - - mock_response = _make_batch_response(batch_id=raw_batch_id) - mock_request = _make_mock_request(headers={}) - mock_fastapi_response = MagicMock() - mock_user_api_key_dict = MagicMock() - mock_user_api_key_dict.parent_otel_span = None - mock_user_api_key_dict.user_id = "test_user" - mock_user_api_key_dict.team_id = "team-caller" - mock_user_api_key_dict.team_metadata = {} - - mock_router = MagicMock() - mock_router.get_deployment_model_for_alias = MagicMock(return_value=real_model) - - mock_credentials = { - "custom_llm_provider": "bedrock", - "aws_region_name": "us-east-1", - } - - request_body = { - "input_file_id": encoded_input_file_id, - "endpoint": "/v1/chat/completions", - "completion_window": "24h", - "model": alias, - } - - with ( - patch( - "litellm.proxy.batches_endpoints.endpoints._read_request_body", - new=AsyncMock(return_value=dict(request_body)), - ), - patch( - "litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing" - ) as mock_processor_cls, - patch( - "litellm.proxy.batches_endpoints.endpoints.get_credentials_for_model", - return_value=mock_credentials, - ), - patch( - "litellm.proxy.batches_endpoints.endpoints.prepare_data_with_credentials", - ), - patch( - "litellm.acreate_batch", - new_callable=AsyncMock, - ) as mock_create_batch, - patch( - "litellm.proxy.batches_endpoints.endpoints.is_known_model", - return_value=False, - ), - patch("litellm.proxy.proxy_server.general_settings", {}), - patch("litellm.proxy.proxy_server.llm_router", mock_router), - patch("litellm.proxy.proxy_server.proxy_config", MagicMock()), - patch("litellm.proxy.proxy_server.version", "1.0.0"), - patch( - "litellm.proxy.proxy_server.proxy_logging_obj", - MagicMock( - post_call_success_hook=AsyncMock(return_value=mock_response), - update_request_status=AsyncMock(), - ), - ), - ): - mock_create_batch.return_value = mock_response - mock_processor = MagicMock() - mock_processor.common_processing_pre_call_logic = AsyncMock( - return_value=(dict(request_body), MagicMock()) - ) - mock_processor_cls.return_value = mock_processor - - response = await create_batch( - request=mock_request, - fastapi_response=mock_fastapi_response, - provider=None, - user_api_key_dict=mock_user_api_key_dict, - ) - - mock_router.get_deployment_model_for_alias.assert_called_once_with(model_id=alias, team_id="team-caller") - create_kwargs = mock_create_batch.call_args.kwargs - assert create_kwargs["model"] == real_model, ( - "Bedrock batch transform receives modelId from this 'model'; it must be the " - f"deployment's real model, got {create_kwargs.get('model')!r}" - ) - # The encoded response id must carry the ALIAS so retrieve routes back. - assert decode_model_from_file_id(response.id) == alias - - @pytest.mark.asyncio async def test_cancel_batch_with_unified_id_routes_with_decoded_model_and_batch_id(): from litellm.proxy.batches_endpoints.endpoints import cancel_batch diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index bb6f4151a3e..56c94fdfd9c 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7859,49 +7859,7 @@ def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): ) -def test_get_deployment_model_for_alias_resolves_underlying_model(): - """ - The proxy batch-create path resolves a model-group alias to its deployment - so it can hand the provider the deployment's real model id, not the alias. - get_llm_provider cannot resolve a proxy alias, so without this the Bedrock - batch transform receives the alias as a modelId and AWS rejects it. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "bedrock-batch-haiku", - "litellm_params": { - "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - "aws_region_name": "us-east-1", - }, - "model_info": {"id": "bedrock-batch-dep-0"}, - } - ] - ) - - assert ( - router.get_deployment_model_for_alias(model_id="bedrock-batch-haiku") - == "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" - ) - # Resolving by deployment id returns the same underlying model. - assert ( - router.get_deployment_model_for_alias(model_id="bedrock-batch-dep-0") - == "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" - ) - - -def test_get_deployment_model_for_alias_returns_none_for_unknown_model(): - router = _router_with_two_deployments([False, False]) - assert router.get_deployment_model_for_alias(model_id="does-not-exist") is None - - -def test_get_deployment_model_for_alias_returns_none_for_blocked_deployment(): - router = _router_with_two_deployments([True, False]) - assert router.get_deployment_model_for_alias(model_id="dep-0") is None - assert router.get_deployment_model_for_alias(model_id="dep-1") == "openai/gpt-4o-1" - - -def test_get_deployment_model_for_alias_matches_credential_deployment_per_team(): +def test_deployment_credentials_are_scoped_to_the_callers_team(): """ Model and credential resolution must pick the SAME deployment for a caller. @@ -7941,22 +7899,20 @@ def test_get_deployment_model_for_alias_matches_credential_deployment_per_team() ("team-b", "bedrock/shared-model", "shared-region"), ("team-a", "bedrock/team-a-private-model", "team-a-region"), ]: - resolved_model = router.get_deployment_model_for_alias(model_id="bedrock-batch", team_id=team_id) credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch", team_id=team_id) - assert resolved_model == expected_model, f"team_id={team_id}" assert credentials is not None + assert credentials["model"] == expected_model, f"team_id={team_id}" assert credentials["aws_region_name"] == expected_region, ( f"team_id={team_id}: credentials came from a different deployment than the model" ) # A caller who knows another team's exact deployment id must not resolve - # its model or credentials through it either. + # its credentials through it either. for outsider_team_id in [None, "team-b"]: - assert router.get_deployment_model_for_alias(model_id="team-a-dep", team_id=outsider_team_id) is None assert router.get_deployment_credentials_with_provider(model_id="team-a-dep", team_id=outsider_team_id) is None - assert router.get_deployment_model_for_alias(model_id="team-a-dep", team_id="team-a") == ( - "bedrock/team-a-private-model" - ) + own_team_credentials = router.get_deployment_credentials_with_provider(model_id="team-a-dep", team_id="team-a") + assert own_team_credentials is not None + assert own_team_credentials["model"] == "bedrock/team-a-private-model" def test_resolve_unblocked_deployment_resolves_alias_id_and_wildcard(): diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..839aa52fa84 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16781,7 +16781,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16886,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)