mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
fix(router): scope alias-to-model resolution to the caller's team
The alias resolver picked the first deployment under a model group while the credential resolver skipped deployments owned by other teams, so a team-owned deployment listed before a shared one under the same alias leaked its private model id into an outside caller's batch while the request ran on the shared deployment's credentials _resolve_unblocked_deployment now carries the full team-aware lookup (team-usable name match, exact team public model name, team wildcard before global wildcard) and both resolvers delegate to it, so model and credentials always come from the same deployment by construction. The batch endpoints pass user_api_key_dict.team_id through both paths
This commit is contained in:
parent
b1c409b94b
commit
adbb3e1bce
6 changed files with 126 additions and 60 deletions
|
|
@ -52,6 +52,7 @@ def _swap_alias_for_deployment_model(
|
|||
create_batch_data: LiteLLMBatchCreateRequest,
|
||||
alias: str,
|
||||
llm_router: Optional["Router"],
|
||||
team_id: "str | None",
|
||||
) -> None:
|
||||
"""
|
||||
Replace a proxy model-group alias on the batch request with the
|
||||
|
|
@ -61,11 +62,14 @@ def _swap_alias_for_deployment_model(
|
|||
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.
|
||||
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 = llm_router.get_deployment_model_for_alias(model_id=alias)
|
||||
resolved_model = 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
|
||||
|
||||
|
|
@ -214,6 +218,7 @@ async def create_batch(
|
|||
llm_router=llm_router,
|
||||
model_id=model_from_file_id,
|
||||
operation_context="batch creation (file created with model)",
|
||||
team_id=user_api_key_dict.team_id,
|
||||
)
|
||||
|
||||
original_file_id = get_original_file_id(input_file_id)
|
||||
|
|
@ -226,6 +231,7 @@ async def create_batch(
|
|||
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
|
||||
|
|
@ -308,6 +314,7 @@ async def create_batch(
|
|||
llm_router=llm_router,
|
||||
model_id=model_param,
|
||||
operation_context="batch creation",
|
||||
team_id=user_api_key_dict.team_id,
|
||||
)
|
||||
|
||||
prepare_data_with_credentials(
|
||||
|
|
@ -318,6 +325,7 @@ async def create_batch(
|
|||
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
|
||||
|
|
@ -518,6 +526,7 @@ async def retrieve_batch(
|
|||
llm_router=llm_router,
|
||||
model_id=model_from_id,
|
||||
operation_context="batch retrieval (batch created with model)",
|
||||
team_id=user_api_key_dict.team_id,
|
||||
)
|
||||
|
||||
original_batch_id = get_original_file_id(batch_id)
|
||||
|
|
@ -724,6 +733,7 @@ async def list_batches(
|
|||
llm_router=llm_router,
|
||||
model_id=model_param,
|
||||
operation_context="batch listing",
|
||||
team_id=user_api_key_dict.team_id,
|
||||
)
|
||||
|
||||
data.update(credentials)
|
||||
|
|
@ -905,6 +915,7 @@ async def cancel_batch(
|
|||
llm_router=llm_router,
|
||||
model_id=model_from_id,
|
||||
operation_context="batch cancellation (batch created with model)",
|
||||
team_id=user_api_key_dict.team_id,
|
||||
)
|
||||
|
||||
original_batch_id = get_original_file_id(batch_id)
|
||||
|
|
|
|||
|
|
@ -259,6 +259,7 @@ def get_credentials_for_model(
|
|||
llm_router, # Router instance
|
||||
model_id: str,
|
||||
operation_context: str = "file operation",
|
||||
team_id: "str | None" = None,
|
||||
):
|
||||
"""
|
||||
Retrieve API credentials for a model from the LLM Router.
|
||||
|
|
@ -267,6 +268,8 @@ def get_credentials_for_model(
|
|||
llm_router: LiteLLM Router instance
|
||||
model_id: Model name or deployment ID
|
||||
operation_context: Description for error messages (e.g., "file upload", "batch creation")
|
||||
team_id: Caller's team id; unlocks that team's own deployments and keeps
|
||||
shared model names from resolving another team's credentials
|
||||
|
||||
Returns:
|
||||
Dictionary with credentials (api_key, api_base, custom_llm_provider, etc.)
|
||||
|
|
@ -282,7 +285,7 @@ def get_credentials_for_model(
|
|||
detail={"error": "Router not initialized. Cannot use model-based routing."},
|
||||
)
|
||||
|
||||
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
|
||||
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id, team_id=team_id)
|
||||
|
||||
if credentials is None:
|
||||
raise HTTPException(
|
||||
|
|
|
|||
|
|
@ -8639,19 +8639,54 @@ class Router:
|
|||
raise Exception("Model Name invalid - {}".format(type(model)))
|
||||
return None
|
||||
|
||||
def _resolve_unblocked_deployment(self, model_id: str) -> "Deployment | None":
|
||||
def _resolve_unblocked_deployment(self, model_id: str, team_id: "str | None" = None) -> "Deployment | None":
|
||||
"""
|
||||
Resolve a model id, model-group alias, or wildcard pattern to a single
|
||||
deployment, returning None when nothing matches or the match is paused
|
||||
via ``LiteLLM_ProxyModelTable.blocked``.
|
||||
|
||||
Both the credential resolver and the alias-to-model resolver delegate
|
||||
here so a mixed model group can never hand one caller the credentials
|
||||
of one deployment and the model of another. Name and wildcard lookups
|
||||
skip deployments owned by a team other than ``team_id``; passing
|
||||
``team_id`` also unlocks that team's own deployments (exact team public
|
||||
model name and team wildcard patterns).
|
||||
"""
|
||||
deployment = self.get_deployment(model_id=model_id)
|
||||
|
||||
if deployment is None:
|
||||
deployment = self.get_deployment_by_model_group_name(model_group_name=model_id)
|
||||
deployment = self._get_model_group_deployment_usable_by_team(model_group_name=model_id, team_id=team_id)
|
||||
|
||||
# Team-scoped deployments whose team public model name exactly matches
|
||||
# model_id (wildcard team names are matched via team_pattern_routers
|
||||
# below).
|
||||
if deployment is None and team_id is not None:
|
||||
team_indices = self.team_model_to_deployment_indices.get((team_id, model_id)) or ()
|
||||
team_match = next((self.model_list[idx] for idx in team_indices), None)
|
||||
if isinstance(team_match, dict):
|
||||
deployment = Deployment(**team_match)
|
||||
elif isinstance(team_match, Deployment):
|
||||
deployment = team_match
|
||||
|
||||
# Wildcard pattern matches. Team wildcard matches take priority so a
|
||||
# global pattern (e.g. "openai/*") doesn't shadow the team's own entry.
|
||||
if deployment is None:
|
||||
wildcard_match = next(iter(self.pattern_router.route(model_id) or ()), None)
|
||||
team_pattern_router = self.team_pattern_routers.get(team_id) if team_id is not None else None
|
||||
team_wildcard_match = (
|
||||
next(iter(team_pattern_router.route(model_id) or ()), None) if team_pattern_router else None
|
||||
)
|
||||
wildcard_match = (
|
||||
team_wildcard_match
|
||||
if team_wildcard_match is not None
|
||||
else next(
|
||||
(
|
||||
wildcard_model
|
||||
for wildcard_model in (self.pattern_router.route(model_id) or ())
|
||||
if self._deployment_usable_by_team(wildcard_model, team_id)
|
||||
),
|
||||
None,
|
||||
)
|
||||
)
|
||||
if isinstance(wildcard_match, dict):
|
||||
deployment = Deployment(**wildcard_match)
|
||||
elif isinstance(wildcard_match, Deployment):
|
||||
|
|
@ -8661,7 +8696,7 @@ class Router:
|
|||
return None
|
||||
return deployment
|
||||
|
||||
def get_deployment_model_for_alias(self, model_id: str) -> "str | None":
|
||||
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``.
|
||||
|
|
@ -8671,8 +8706,10 @@ class Router:
|
|||
``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 = self._resolve_unblocked_deployment(model_id=model_id)
|
||||
deployment = self._resolve_unblocked_deployment(model_id=model_id, team_id=team_id)
|
||||
if deployment is None:
|
||||
return None
|
||||
return deployment.litellm_params.model
|
||||
|
|
@ -8753,43 +8790,8 @@ class Router:
|
|||
credentials = router.get_deployment_credentials_with_provider("gpt-4o-litellm")
|
||||
# Returns: {"api_key": "sk-...", "custom_llm_provider": "openai", ...}
|
||||
"""
|
||||
# Try to get deployment by model_id first
|
||||
deployment = self.get_deployment(model_id=model_id)
|
||||
|
||||
# If not found, try by model_group_name
|
||||
deployment = self._resolve_unblocked_deployment(model_id=model_id, team_id=team_id)
|
||||
if deployment is None:
|
||||
deployment = self._get_model_group_deployment_usable_by_team(model_group_name=model_id, team_id=team_id)
|
||||
|
||||
# If not found, check team-scoped deployments whose team public model
|
||||
# name exactly matches model_id (wildcard team names are matched via
|
||||
# team_pattern_routers below).
|
||||
if deployment is None and team_id is not None:
|
||||
team_indices = self.team_model_to_deployment_indices.get((team_id, model_id), [])
|
||||
if team_indices:
|
||||
team_model = self.model_list[team_indices[0]]
|
||||
deployment = Deployment(**team_model) if isinstance(team_model, dict) else team_model
|
||||
|
||||
# If still not found, check for wildcard pattern matches. Team wildcard
|
||||
# matches take priority so a global pattern (e.g. "openai/*") doesn't
|
||||
# shadow the team's own entry.
|
||||
if deployment is None:
|
||||
team_pattern_router = self.team_pattern_routers.get(team_id) if team_id is not None else None
|
||||
team_wildcard_models = (team_pattern_router.route(model_id) or []) if team_pattern_router else []
|
||||
global_wildcard_models = [
|
||||
wildcard_model
|
||||
for wildcard_model in (self.pattern_router.route(model_id) or [])
|
||||
if self._deployment_usable_by_team(wildcard_model, team_id)
|
||||
]
|
||||
potential_wildcard_models = team_wildcard_models or global_wildcard_models
|
||||
if potential_wildcard_models:
|
||||
# Use the first matching wildcard deployment
|
||||
deployment_dict = potential_wildcard_models[0]
|
||||
if isinstance(deployment_dict, dict):
|
||||
deployment = Deployment(**deployment_dict)
|
||||
elif isinstance(deployment_dict, Deployment):
|
||||
deployment = deployment_dict
|
||||
|
||||
if deployment is None or self._is_deployment_blocked(deployment):
|
||||
return None
|
||||
|
||||
# Get basic credentials
|
||||
|
|
|
|||
|
|
@ -146,12 +146,12 @@ class Harness:
|
|||
return dict(self.router_acreate.call_args.kwargs)
|
||||
|
||||
|
||||
def _creds_lookup(*, model_id: str) -> Dict[str, str]:
|
||||
def _creds_lookup(*, model_id: str, team_id: Optional[str] = None) -> Dict[str, str]:
|
||||
# KeyError on an unknown/hardcoded model_id - the bug cannot hide.
|
||||
return dict(CREDS[model_id])
|
||||
|
||||
|
||||
def _alias_lookup(*, model_id: str) -> str:
|
||||
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.
|
||||
|
|
@ -267,7 +267,7 @@ async def test_create__model_encoded_file_id(harness):
|
|||
harness.router_acreate.assert_not_called()
|
||||
|
||||
# 2. CREDENTIALS - resolved for the model decoded FROM the file id.
|
||||
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
||||
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
|
||||
|
||||
# 3. SEAM PAYLOAD - exact, whole dict. A new forwarded key breaks this.
|
||||
assert harness.acreate_kwargs() == {
|
||||
|
|
@ -323,7 +323,7 @@ async def test_create__model_encoded_file_id__resolver_gets_decoded_model(harnes
|
|||
|
||||
await call_create(harness)
|
||||
|
||||
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
||||
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
|
|
@ -347,7 +347,7 @@ async def test_create__model_from_body(harness):
|
|||
|
||||
assert harness.litellm_acreate.call_count == 1
|
||||
harness.router_acreate.assert_not_called()
|
||||
harness.creds_resolver.assert_called_once_with(model_id="vertex-model")
|
||||
harness.creds_resolver.assert_called_once_with(model_id="vertex-model", team_id=None)
|
||||
payload = harness.acreate_kwargs()
|
||||
assert payload["custom_llm_provider"] == "vertex_ai"
|
||||
assert payload["input_file_id"] == "file-plain"
|
||||
|
|
@ -367,7 +367,7 @@ async def test_create__model_from_header(harness):
|
|||
|
||||
await call_create(harness, headers={"x-litellm-model": "vertex-model"})
|
||||
|
||||
harness.creds_resolver.assert_called_once_with(model_id="vertex-model")
|
||||
harness.creds_resolver.assert_called_once_with(model_id="vertex-model", team_id=None)
|
||||
harness.router_acreate.assert_not_called()
|
||||
|
||||
|
||||
|
|
@ -384,7 +384,7 @@ async def test_create__model_from_query(harness):
|
|||
|
||||
await call_create(harness, query={"model": "vertex-model"})
|
||||
|
||||
harness.creds_resolver.assert_called_once_with(model_id="vertex-model")
|
||||
harness.creds_resolver.assert_called_once_with(model_id="vertex-model", team_id=None)
|
||||
harness.router_acreate.assert_not_called()
|
||||
|
||||
|
||||
|
|
@ -407,7 +407,7 @@ async def test_create__body_model_beats_header_and_query(harness):
|
|||
query={"model": "vertex-model"},
|
||||
)
|
||||
|
||||
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
||||
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
|
|
@ -715,7 +715,7 @@ async def test_create__model_encoded_beats_unified(harness):
|
|||
|
||||
assert harness.litellm_acreate.call_count == 1
|
||||
harness.router_acreate.assert_not_called()
|
||||
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
||||
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
|
|
@ -761,7 +761,7 @@ async def test_create__model_encoded_beats_loadbalancing(harness):
|
|||
|
||||
assert harness.litellm_acreate.call_count == 1
|
||||
harness.router_acreate.assert_not_called()
|
||||
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
||||
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
|
|
@ -1072,7 +1072,7 @@ async def test_retrieve__model_encoded_id(retrieve_harness):
|
|||
retrieve_harness.router_aretrieve.assert_not_called()
|
||||
|
||||
# 2. CREDENTIALS - resolved for the model decoded FROM the batch id.
|
||||
retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
||||
retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
|
||||
|
||||
# 3. SEAM PAYLOAD - exact, whole dict forwarded to the provider call.
|
||||
# Note `model` is the DECODED model, not the deployment from creds: the
|
||||
|
|
@ -1130,7 +1130,7 @@ async def test_retrieve__model_encoded_beats_loadbalancing(retrieve_harness):
|
|||
|
||||
assert retrieve_harness.litellm_aretrieve.call_count == 1
|
||||
retrieve_harness.router_aretrieve.assert_not_called()
|
||||
retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
||||
retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
|
@ -1546,7 +1546,7 @@ async def test_list__model_from_body_routes_and_encodes(list_harness):
|
|||
|
||||
assert list_harness.litellm_alist.call_count == 1
|
||||
list_harness.router_alist.assert_not_called()
|
||||
list_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
||||
list_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
|
||||
assert resp.data[0].id == encode_file_id_with_model("batch-1", "azure/gpt-4o", id_type="batch")
|
||||
assert resp.data[1].id == encode_file_id_with_model("batch-2", "azure/gpt-4o", id_type="batch")
|
||||
|
||||
|
|
@ -1836,7 +1836,7 @@ async def test_cancel__model_encoded_id(cancel_harness):
|
|||
cancel_harness.router_acancel.assert_not_called()
|
||||
|
||||
# CREDENTIALS - resolved for the model decoded from the batch id.
|
||||
cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
||||
cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
|
||||
|
||||
# SEAM PAYLOAD - exact dict. NOTE current behavior: `model` is the
|
||||
# DEPLOYMENT name from creds, NOT the decoded model (cancel, unlike
|
||||
|
|
@ -1874,7 +1874,7 @@ async def test_cancel__model_encoded_beats_unified(cancel_harness):
|
|||
|
||||
assert cancel_harness.litellm_acancel.call_count == 1
|
||||
cancel_harness.router_acancel.assert_not_called()
|
||||
cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
||||
cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
|
|
|||
|
|
@ -444,6 +444,7 @@ async def test_create_batch_swaps_alias_for_deployment_model_before_provider_cal
|
|||
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()
|
||||
|
|
@ -510,7 +511,7 @@ async def test_create_batch_swaps_alias_for_deployment_model_before_provider_cal
|
|||
user_api_key_dict=mock_user_api_key_dict,
|
||||
)
|
||||
|
||||
mock_router.get_deployment_model_for_alias.assert_called_once_with(model_id=alias)
|
||||
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 "
|
||||
|
|
|
|||
|
|
@ -5844,6 +5844,55 @@ def test_get_deployment_model_for_alias_returns_none_for_blocked_deployment():
|
|||
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():
|
||||
"""
|
||||
Model and credential resolution must pick the SAME deployment for a caller.
|
||||
|
||||
Regression: with a team-owned deployment listed before a shared one under
|
||||
the same alias, an unscoped alias lookup returned the team deployment's
|
||||
model while the team-aware credential resolver returned the shared
|
||||
deployment's credentials, so an outside caller's batch reached the provider
|
||||
with team A's private model id on the shared account.
|
||||
"""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "bedrock-batch",
|
||||
"litellm_params": {
|
||||
"model": "bedrock/team-a-private-model",
|
||||
"aws_region_name": "team-a-region",
|
||||
},
|
||||
"model_info": {
|
||||
"id": "team-a-dep",
|
||||
"team_id": "team-a",
|
||||
"team_public_model_name": "bedrock-batch",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "bedrock-batch",
|
||||
"litellm_params": {
|
||||
"model": "bedrock/shared-model",
|
||||
"aws_region_name": "shared-region",
|
||||
},
|
||||
"model_info": {"id": "shared-dep"},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
for team_id, expected_model, expected_region in [
|
||||
(None, "bedrock/shared-model", "shared-region"),
|
||||
("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["aws_region_name"] == expected_region, (
|
||||
f"team_id={team_id}: credentials came from a different deployment than the model"
|
||||
)
|
||||
|
||||
|
||||
def test_resolve_unblocked_deployment_resolves_alias_id_and_wildcard():
|
||||
"""
|
||||
_resolve_unblocked_deployment underpins both the credential resolver and the
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue