scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-22 00:31:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 96

refactor(batches): drop the alias swap now that credential resolution returns the deployment model

Browse source 
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.
This commit is contained in:
Kent 2026-09-14 12:06:17 +08:00
parent 3344be85db
commit 6c4fd246d3
6 changed files with 7 additions and 225 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
litellm/proxy/batches_endpoints/endpoints.py
View file

@ -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(

18
litellm/router.py
View file

@ -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:
"""

8
tests/test_litellm/proxy/batches_endpoints/test_endpoints.py
View file

@ -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()))

105
tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py
View file

@ -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

56
tests/test_litellm/test_router.py
View file

@ -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():

2
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -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)