mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
feat(proxy): let teams use their own allowed S3 buckets for batch files (#50571)
* feat(proxy): let teams send allowed S3 buckets for batch files and read them back A team's metadata can now list allowed_s3_buckets. Keys on that team may send one of those buckets as s3_bucket_name on POST /v1/files and as s3_output_bucket_name on POST /v1/batches, and the upload and batch-create handlers check the bucket against the team list plus general_settings.allowed_client_side_s3_buckets. The early auth body check still rejects bucket fields on every other route and method Reads of managed files now trust the bucket LiteLLM recorded in the stored s3:// uri through an internal stored_s3_bucket_name param, so content, retrieve, delete and the batch cost job can read files written to a request-chosen bucket. Callers can never send that param, and files in the model bucket resolve exactly as before * test(proxy): assert outcomes in team bucket tests and type the managed content read Gives the admit-only tests a real assertion, has the refresh test check what the caller gets back, and narrows the provider file id before reading its bucket so the type gate stays at its base count * fix(proxy): keep team S3 bucket lists proxy-admin only and read prefixed request buckets back A team admin allowed to edit team metadata could add any bucket to the team's allowed_s3_buckets, so the key now sits with the other proxy-controlled team metadata keys. The stored bucket for a managed S3 file now keeps any prefix that comes before LiteLLM's own path, so a file uploaded to a request bucket like team-bucket/tenant reads, refreshes and deletes from that prefix * fix(proxy): keep the request's approved S3 bucket when merging deployment credentials Model-routed uploads and batch creates merge the deployment's credentials into the request, and that merge let a deployment's s3_bucket_name or s3_output_bucket_name replace the bucket the request had already been approved for, so Bedrock wrote to the deployment bucket. The merge now keeps any bucket the request carries and fills it from the deployment only when the request has none * fix(proxy): keep the request bucket only on upload and batch create, not on reads The previous commit let a request's S3 bucket win over the deployment's on every credential merge, which undid #47496 on reads, where the deployment's bucket bounds which buckets a model-routed file id may point at. Deployment credentials win again by default, and only model-routed and provider-only upload and batch create opt in to keep the bucket the request was approved for * fix(proxy): only proxy admins can set a team's allowed_s3_buckets Org admins and users allowed to create teams could put any bucket on a team's allowed_s3_buckets through /team/new or a team update, which grants that team the proxy's AWS access to the bucket. The key now goes through the same proxy-admin-only metadata check as the batch limits, so only a proxy admin can add, change or remove it, and resending the stored value still passes
This commit is contained in:
mrinal-berri
GitHub
parent
59a3e75368
commit
9ba551b871
22 changed files with 850 additions and 44 deletions
|
|
@ -20,6 +20,7 @@ from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
|
|||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.litellm_core_utils.cloud_storage_security import stored_s3_bucket_kwargs
|
||||
from litellm.constants import (
|
||||
BATCH_LINE_ITEMS_MAX_PART_BYTES,
|
||||
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
|
||||
|
|
@ -765,7 +766,7 @@ class CheckBatchCost:
|
|||
credentials: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {}
|
||||
_file_content = await afile_content(
|
||||
file_id=cast(str, raw_output_file_id),
|
||||
**cast(FileContentCallOptions, credentials),
|
||||
**cast(FileContentCallOptions, {**credentials, **stored_s3_bucket_kwargs(cast(str, raw_output_file_id))}),
|
||||
)
|
||||
|
||||
content_bytes: Final = await _file_content_bytes(_file_content)
|
||||
|
|
|
|||
|
|
@ -49,6 +49,7 @@ from litellm.constants import (
|
|||
)
|
||||
from litellm.files.types import FileRetrieveCallOptions, FileRetrieveProvider
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.cloud_storage_security import stored_s3_bucket_kwargs
|
||||
from litellm.litellm_core_utils.core_helpers import require_attributes
|
||||
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR
|
||||
from litellm.litellm_core_utils.internal_call_metadata import requests_background
|
||||
|
|
@ -396,6 +397,7 @@ class _CursorPageArgs(TypedDict, total=False):
|
|||
class _RouterFileCallKwargs(TypedDict, total=False):
|
||||
client: ReadOnly[AsyncOpenAI | None]
|
||||
custom_llm_provider: ReadOnly[str | None]
|
||||
stored_s3_bucket_name: ReadOnly[str]
|
||||
|
||||
|
||||
_MODEL_FILE_ID_MAPPING: Final = TypeAdapter(dict[str, str])
|
||||
|
|
@ -463,7 +465,7 @@ async def _delete_provider_file(
|
|||
await llm_router.afile_delete(
|
||||
model=model_id,
|
||||
file_id=provider_file_id,
|
||||
**cast(_RouterFileCallKwargs, delete_data),
|
||||
**cast(_RouterFileCallKwargs, {**delete_data, **stored_s3_bucket_kwargs(provider_file_id)}),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1157,7 +1159,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
retrieve_file: Final = cast(_ManagedFileRetrieve, litellm.afile_retrieve)
|
||||
retrieve_options: Final = cast(
|
||||
FileRetrieveCallOptions,
|
||||
{**call_options, "max_retries": 0},
|
||||
{**call_options, **stored_s3_bucket_kwargs(provider_file_id), "max_retries": 0},
|
||||
)
|
||||
for attempt in range(len(_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS) + 1):
|
||||
if attempt > 0:
|
||||
|
|
@ -1795,8 +1797,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
if potential_file_id and "llm_output_file_id," in potential_file_id:
|
||||
model_id = self.get_model_id_from_unified_file_id(potential_file_id)
|
||||
if model_id:
|
||||
provider_file_id: Final = self.get_output_file_id_from_unified_file_id(potential_file_id)
|
||||
data["model"] = model_id
|
||||
data["file_id"] = self.get_output_file_id_from_unified_file_id(potential_file_id)
|
||||
data["file_id"] = provider_file_id
|
||||
data.update(stored_s3_bucket_kwargs(provider_file_id))
|
||||
elif call_type == CallTypes.acreate_batch.value:
|
||||
input_file_id = cast(Optional[str], data.get("input_file_id"))
|
||||
if input_file_id:
|
||||
|
|
@ -2894,7 +2898,17 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
await llm_router.afile_content(
|
||||
model=model_id,
|
||||
file_id=provider_file_id,
|
||||
**cast(_RouterFileCallKwargs, data),
|
||||
**cast(
|
||||
_RouterFileCallKwargs,
|
||||
{
|
||||
**data,
|
||||
**(
|
||||
stored_s3_bucket_kwargs(provider_file_id)
|
||||
if isinstance(provider_file_id, str)
|
||||
else {}
|
||||
),
|
||||
},
|
||||
),
|
||||
),
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from pydantic import JsonValue, TypeAdapter
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils import batch_utils as core_batch_utils
|
||||
from litellm.litellm_core_utils.cloud_storage_security import stored_s3_bucket_kwargs
|
||||
from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo
|
||||
|
|
@ -551,7 +552,7 @@ async def fetch_batch_managed_file_content(
|
|||
"file_id": provider_output_file_id,
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
**cast( # cast-ok: preserve dynamic provider credentials without validation
|
||||
FileContentCallOptions, credentials
|
||||
FileContentCallOptions, {**credentials, **stored_s3_bucket_kwargs(provider_output_file_id)}
|
||||
),
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -17,6 +17,28 @@ BEDROCK_MANAGED_S3_PREFIXES: Final = (
|
|||
BEDROCK_MANAGED_S3_OUTPUT_PREFIX,
|
||||
)
|
||||
MANAGED_CLOUD_STORAGE_SCHEMES: Final = ("s3://", "gs://")
|
||||
STORED_S3_BUCKET_PARAM: Final = "stored_s3_bucket_name"
|
||||
REQUEST_S3_BUCKET_PARAMS: Final = frozenset({"s3_bucket_name", "s3_output_bucket_name"})
|
||||
|
||||
|
||||
def stored_s3_bucket_kwargs(stored_file_id: str) -> Mapping[str, str]:
|
||||
"""The bucket (and any prefix before LiteLLM's own path) of an ``s3://`` file id LiteLLM recorded itself,
|
||||
as the param that lets Bedrock read it back.
|
||||
|
||||
Only pass ids from LiteLLM's own records (managed file rows, provider batch responses), never one a caller sent.
|
||||
"""
|
||||
if not stored_file_id.startswith("s3://"):
|
||||
return {}
|
||||
bucket, _, object_key = stored_file_id.removeprefix("s3://").partition("/")
|
||||
if not bucket:
|
||||
return {}
|
||||
segments: Final = object_key.split("/")
|
||||
prefix_depth: Final = next(
|
||||
(depth for depth in range(len(segments)) if "/".join(segments[depth:]).startswith(BEDROCK_MANAGED_S3_PREFIXES)),
|
||||
0,
|
||||
)
|
||||
prefix: Final = "/".join(segments[:prefix_depth])
|
||||
return {STORED_S3_BUCKET_PARAM: f"{bucket}/{prefix}" if prefix else bucket}
|
||||
|
||||
|
||||
MAX_FILE_ID_DECODE_PASSES: Final = 8
|
||||
|
|
|
|||
|
|
@ -62,6 +62,7 @@ OPTIONAL_KWARGS_KEYS: Final = (
|
|||
"bucket_name",
|
||||
"s3_bucket_name",
|
||||
"s3_output_bucket_name",
|
||||
"stored_s3_bucket_name",
|
||||
"allow_legacy_cloud_file_ids",
|
||||
"s3_endpoint_url",
|
||||
"s3_region_name",
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ from litellm.litellm_core_utils.cloud_storage_security import (
|
|||
BEDROCK_MANAGED_S3_OUTPUT_PREFIX,
|
||||
BEDROCK_MANAGED_S3_PREFIXES,
|
||||
BEDROCK_MANAGED_S3_UPLOAD_PREFIX,
|
||||
STORED_S3_BUCKET_PARAM,
|
||||
build_managed_cloud_object_name,
|
||||
encode_s3_object_key_for_url,
|
||||
sanitize_cloud_object_component,
|
||||
|
|
@ -313,11 +314,19 @@ def get_configured_s3_bucket_names(litellm_params: Mapping[str, object]) -> tupl
|
|||
Resolve the S3 buckets a Bedrock file id may live in.
|
||||
|
||||
Bedrock batch outputs land in ``s3_output_bucket_name`` when it differs from
|
||||
the input bucket, so retrieval validates against both.
|
||||
the input bucket, so retrieval validates against both, plus the bucket LiteLLM
|
||||
recorded for a managed file uploaded or written to a request-chosen bucket.
|
||||
"""
|
||||
input_bucket: Final = _input_bucket_name(litellm_params)
|
||||
output_bucket: Final = _output_bucket_name(litellm_params)
|
||||
buckets: Final = tuple(dict.fromkeys(bucket for bucket in (input_bucket, output_bucket) if bucket))
|
||||
stored_bucket: Final = litellm_params.get(STORED_S3_BUCKET_PARAM)
|
||||
buckets: Final = tuple(
|
||||
dict.fromkeys(
|
||||
bucket
|
||||
for bucket in (input_bucket, output_bucket, stored_bucket if isinstance(stored_bucket, str) else None)
|
||||
if bucket
|
||||
)
|
||||
)
|
||||
if not buckets:
|
||||
raise ValueError(_S3_BUCKET_REQUIRED_ERROR)
|
||||
return buckets
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import os
|
|||
import re
|
||||
import sys
|
||||
from collections.abc import Collection, Iterator, Mapping, Sequence
|
||||
from functools import lru_cache
|
||||
from functools import lru_cache, partial
|
||||
from itertools import accumulate, chain, repeat
|
||||
from logging import Logger
|
||||
from types import MappingProxyType
|
||||
|
|
@ -22,6 +22,7 @@ from litellm.constants import (
|
|||
MINIMUM_CUSTOM_KEY_LENGTH,
|
||||
STANDARD_CUSTOMER_ID_HEADERS,
|
||||
)
|
||||
from litellm.litellm_core_utils.cloud_storage_security import REQUEST_S3_BUCKET_PARAMS, STORED_S3_BUCKET_PARAM
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor, reject_api_key_secret_reference
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.litellm_core_utils.url_utils import (
|
||||
|
|
@ -444,6 +445,7 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = (
|
|||
"s3_endpoint_url",
|
||||
"s3_bucket_name",
|
||||
"s3_output_bucket_name",
|
||||
STORED_S3_BUCKET_PARAM,
|
||||
"allow_legacy_cloud_file_ids",
|
||||
"sagemaker_base_url",
|
||||
"deployment_url",
|
||||
|
|
@ -494,6 +496,9 @@ _PROXY_ALLOWLIST_SETTING_BY_PARAM: Final[Mapping[str, _ClientSideAllowlistSettin
|
|||
}
|
||||
)
|
||||
|
||||
TEAM_S3_BUCKET_ALLOWLIST_METADATA_KEY: Final = "allowed_s3_buckets"
|
||||
_BUCKET_NAMES: Final = TypeAdapter(tuple[str, ...])
|
||||
|
||||
|
||||
class _ClientSideStorageAllowlist(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
|
@ -511,8 +516,15 @@ class _ClientSideStorageAllowlist(BaseModel):
|
|||
|
||||
def _banned_param_error(param: str) -> ValueError:
|
||||
setting: Final = _PROXY_ALLOWLIST_SETTING_BY_PARAM.get(param)
|
||||
team_hint: Final = (
|
||||
f" or in the team's `metadata.{TEAM_S3_BUCKET_ALLOWLIST_METADATA_KEY}`"
|
||||
if param in REQUEST_S3_BUCKET_PARAMS
|
||||
else ""
|
||||
)
|
||||
allowlist_hint: Final = (
|
||||
"" if setting is None else f" To allow specific values only, list them in `general_settings.{setting}`."
|
||||
""
|
||||
if setting is None
|
||||
else f" To allow specific values only, list them in `general_settings.{setting}`{team_hint}."
|
||||
)
|
||||
return ValueError(
|
||||
f"Rejected Request: {param} is not allowed in request body. "
|
||||
|
|
@ -532,11 +544,13 @@ def _check_banned_params(
|
|||
*,
|
||||
manages_deployments: bool = False,
|
||||
mcp_server_fields_allowed: bool = False,
|
||||
defer_s3_buckets: bool = False,
|
||||
) -> None:
|
||||
"""Raise ``ValueError`` if ``body`` carries a banned param without admin opt-in.
|
||||
|
||||
Shared between the root-level check and the nested-config check so a
|
||||
new banned param only needs to be added in one place.
|
||||
new banned param only needs to be added in one place. ``defer_s3_buckets`` lets
|
||||
bucket params through for handlers that check them against the caller team's allowlist.
|
||||
"""
|
||||
wif_body: Final[Mapping[str, object]] = (
|
||||
MappingProxyType({key: value for key, value in body.items() if key not in oauth_token_exchange_litellm_params})
|
||||
|
|
@ -557,6 +571,8 @@ def _check_banned_params(
|
|||
setting := _PROXY_ALLOWLIST_SETTING_BY_PARAM.get(param)
|
||||
) is not None and _ClientSideStorageAllowlist.model_validate(general_settings).admits(setting, body[param]):
|
||||
continue
|
||||
if defer_s3_buckets and param in REQUEST_S3_BUCKET_PARAMS:
|
||||
continue
|
||||
if (
|
||||
_allow_model_level_clientside_configurable_parameters(
|
||||
model=model,
|
||||
|
|
@ -704,6 +720,7 @@ def is_request_body_safe(
|
|||
model: str,
|
||||
*,
|
||||
route: str | None = None,
|
||||
defer_s3_buckets: bool = False,
|
||||
) -> bool:
|
||||
"""
|
||||
Check if the request body is safe.
|
||||
|
|
@ -733,33 +750,28 @@ def is_request_body_safe(
|
|||
if "model_list" in request_body:
|
||||
raise ValueError("Rejected Request: model_list is not allowed in the request body.")
|
||||
manages_deployments: Final = route_manages_deployments(route)
|
||||
_check_banned_params(
|
||||
request_body,
|
||||
general_settings,
|
||||
llm_router,
|
||||
model,
|
||||
check_banned_params: Final = partial(
|
||||
_check_banned_params,
|
||||
general_settings=general_settings,
|
||||
llm_router=llm_router,
|
||||
model=model,
|
||||
manages_deployments=manages_deployments,
|
||||
mcp_server_fields_allowed=route_manages_mcp_servers(route),
|
||||
defer_s3_buckets=defer_s3_buckets,
|
||||
)
|
||||
check_banned_params(request_body, mcp_server_fields_allowed=route_manages_mcp_servers(route))
|
||||
for nested_key in _NESTED_CONFIG_KEYS:
|
||||
nested = _coerce_metadata_to_dict(request_body.get(nested_key))
|
||||
if nested is not None:
|
||||
_check_banned_params(nested, general_settings, llm_router, model, manages_deployments=manages_deployments)
|
||||
check_banned_params(nested)
|
||||
for metadata_key in _NESTED_METADATA_KEYS:
|
||||
metadata = _coerce_metadata_to_dict(request_body.get(metadata_key))
|
||||
if metadata is not None:
|
||||
_check_banned_params(metadata, general_settings, llm_router, model, manages_deployments=manages_deployments)
|
||||
check_banned_params(metadata)
|
||||
if any(isinstance(key, str) and key.startswith(f"{metadata_key}[") for key in request_body):
|
||||
_check_banned_params(
|
||||
extract_nested_form_metadata(form_data=request_body, prefix=f"{metadata_key}["),
|
||||
general_settings,
|
||||
llm_router,
|
||||
model,
|
||||
manages_deployments=manages_deployments,
|
||||
)
|
||||
check_banned_params(extract_nested_form_metadata(form_data=request_body, prefix=f"{metadata_key}["))
|
||||
for target in iter_request_fallback_targets(request_body):
|
||||
if isinstance(target, dict):
|
||||
_check_banned_params(target, general_settings, llm_router, model, manages_deployments=manages_deployments)
|
||||
check_banned_params(target)
|
||||
target_model = target.get("model")
|
||||
if isinstance(target_model, str):
|
||||
_reject_url_valued_fallback_target(target_model)
|
||||
|
|
@ -772,9 +784,7 @@ def is_request_body_safe(
|
|||
_reject_request_credential_params(litellm_params, manages_deployments=manages_deployments)
|
||||
litellm_params_metadata: Final = _coerce_metadata_to_dict(litellm_params.get("metadata"))
|
||||
if litellm_params_metadata is not None:
|
||||
_check_banned_params(
|
||||
litellm_params_metadata, general_settings, llm_router, model, manages_deployments=manages_deployments
|
||||
)
|
||||
check_banned_params(litellm_params_metadata)
|
||||
for user_config_params in _user_config_litellm_params(request_body): # pyright: ignore[reportUnknownArgumentType] # request_body is an untyped dict
|
||||
reject_api_key_secret_reference(user_config_params)
|
||||
return True
|
||||
|
|
@ -846,6 +856,48 @@ def _coerce_metadata_to_dict(value: object) -> dict[str, object] | None:
|
|||
return None
|
||||
|
||||
|
||||
def _team_s3_buckets(team_metadata: Mapping[str, object] | None) -> tuple[str, ...]:
|
||||
if team_metadata is None:
|
||||
return ()
|
||||
try:
|
||||
return _BUCKET_NAMES.validate_python(team_metadata.get(TEAM_S3_BUCKET_ALLOWLIST_METADATA_KEY) or ())
|
||||
except ValidationError:
|
||||
return ()
|
||||
|
||||
|
||||
_TEAM_S3_BUCKET_ROUTE_TEMPLATES: Final = frozenset(
|
||||
{"/v1/files", "/files", "/{provider}/v1/files", "/v1/batches", "/batches", "/{provider}/v1/batches"}
|
||||
)
|
||||
|
||||
|
||||
def checks_team_s3_buckets_in_handler(request: Request) -> bool:
|
||||
"""File upload and batch create check request buckets against the caller team's allowlist themselves."""
|
||||
return get_request_route_template(request) in _TEAM_S3_BUCKET_ROUTE_TEMPLATES and request.method == "POST"
|
||||
|
||||
|
||||
def enforce_team_s3_bucket_allowlist(
|
||||
request_body: Mapping[str, object],
|
||||
general_settings: Mapping[str, object],
|
||||
llm_router: Router | None,
|
||||
team_metadata: Mapping[str, object] | None,
|
||||
) -> None:
|
||||
"""Re-run the request body check that auth relaxed for bucket params, now admitting the caller team's buckets."""
|
||||
proxy_buckets: Final = _ClientSideStorageAllowlist.model_validate(general_settings).allowed_client_side_s3_buckets
|
||||
model: Final = request_body.get("model")
|
||||
try:
|
||||
is_request_body_safe(
|
||||
request_body=dict(request_body),
|
||||
general_settings={
|
||||
**general_settings,
|
||||
"allowed_client_side_s3_buckets": (*(proxy_buckets or ()), *_team_s3_buckets(team_metadata)),
|
||||
},
|
||||
llm_router=llm_router,
|
||||
model=model if isinstance(model, str) else "",
|
||||
)
|
||||
except ValueError as rejection:
|
||||
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail=f"Authentication Error, {rejection}")
|
||||
|
||||
|
||||
async def pre_db_read_auth_checks(
|
||||
request: Request,
|
||||
request_data: dict,
|
||||
|
|
@ -875,6 +927,7 @@ async def pre_db_read_auth_checks(
|
|||
llm_router=llm_router,
|
||||
model=request_data.get("model", ""), # [TODO] use model passed in url as well (azure openai routes)
|
||||
route=route,
|
||||
defer_s3_buckets=checks_team_s3_buckets_in_handler(request),
|
||||
)
|
||||
|
||||
# Check 3. Check if IP address is allowed
|
||||
|
|
@ -1522,6 +1575,9 @@ class BatchLimitRequest(Protocol):
|
|||
def model_fields_set(self) -> Collection[str]: ...
|
||||
|
||||
|
||||
_ADMIN_ONLY_METADATA_KEYS: Final = (*ADMIN_ONLY_BATCH_LIMIT_METADATA_KEYS, TEAM_S3_BUCKET_ALLOWLIST_METADATA_KEY)
|
||||
|
||||
|
||||
def enforce_batch_limits_are_admin_only(
|
||||
data: BatchLimitRequest,
|
||||
existing_metadata: Mapping[str, object] | None,
|
||||
|
|
@ -1532,7 +1588,8 @@ def enforce_batch_limits_are_admin_only(
|
|||
|
||||
Every key in ``ADMIN_ONLY_BATCH_LIMIT_METADATA_KEYS`` caps what the holder
|
||||
may do with batches, so a holder-writable copy would let a caller lift their
|
||||
own quota. Gated on the resulting value rather than on presence, so a form
|
||||
own quota, and the S3 bucket list grants access to buckets with the proxy's
|
||||
own credentials. Gated on the resulting value rather than on presence, so a form
|
||||
resending the stored value stays a no-op.
|
||||
"""
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||||
|
|
@ -1542,17 +1599,19 @@ def enforce_batch_limits_are_admin_only(
|
|||
(data.metadata or EMPTY_MAPPING) if "metadata" in data.model_fields_set else stored
|
||||
)
|
||||
changed: Final = next(
|
||||
(key for key in ADMIN_ONLY_BATCH_LIMIT_METADATA_KEYS if requested.get(key) != stored.get(key)),
|
||||
(key for key in _ADMIN_ONLY_METADATA_KEYS if requested.get(key) != stored.get(key)),
|
||||
None,
|
||||
)
|
||||
if changed is None:
|
||||
return
|
||||
reason: Final = (
|
||||
"It lets the holder read and write those buckets with the proxy's AWS credentials."
|
||||
if changed == TEAM_S3_BUCKET_ALLOWLIST_METADATA_KEY
|
||||
else "It limits what the holder can do with batches, so the holder cannot raise it."
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": f"Only proxy admins can set {changed} on a {entity}. "
|
||||
"It limits what the holder can do with batches, so the holder cannot raise it."
|
||||
},
|
||||
detail={"error": f"Only proxy admins can set {changed} on a {entity}. {reason}"},
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from types import MappingProxyType
|
|||
from typing import TYPE_CHECKING, Final, Literal, Protocol, cast
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -20,6 +21,7 @@ from litellm.litellm_core_utils.core_helpers import attribute_or_none, has_attri
|
|||
from litellm.litellm_core_utils.hidden_params import get_hidden_params, set_hidden_param
|
||||
from litellm.llms.openai.workload_identity import resolve_openai_workload_identity_config
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.auth.auth_utils import enforce_team_s3_bucket_allowlist
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit
|
||||
from litellm.proxy.batches_endpoints.litellm_executed_batches import (
|
||||
|
|
@ -81,6 +83,7 @@ if TYPE_CHECKING:
|
|||
from prisma.models import LiteLLM_ManagedObjectTable
|
||||
|
||||
router: Final = APIRouter()
|
||||
_REQUEST_BODY_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
class _UserBatchLister(Protocol):
|
||||
|
|
@ -285,6 +288,7 @@ async def create_batch(
|
|||
"""
|
||||
from litellm.proxy.proxy_server import (
|
||||
general_settings,
|
||||
general_settings_view,
|
||||
llm_router,
|
||||
proxy_config,
|
||||
proxy_logging_obj,
|
||||
|
|
@ -294,6 +298,12 @@ async def create_batch(
|
|||
data: dict = {}
|
||||
try:
|
||||
data = await read_request_body(request=request) # rebind-ok: pre-existing rebinding on a rename-only line
|
||||
enforce_team_s3_bucket_allowlist(
|
||||
request_body=_REQUEST_BODY_ADAPTER.validate_python(data),
|
||||
general_settings=general_settings_view(),
|
||||
llm_router=llm_router,
|
||||
team_metadata=user_api_key_dict.team_metadata,
|
||||
)
|
||||
_raise_when_routing_headers_disagree(request, data.get("model"))
|
||||
verbose_proxy_logger.debug(
|
||||
"Request received by LiteLLM:\n%s",
|
||||
|
|
@ -384,6 +394,7 @@ async def create_batch(
|
|||
prepare_data_with_credentials(
|
||||
data=_create_batch_data,
|
||||
credentials=credentials,
|
||||
keep_request_buckets=True,
|
||||
)
|
||||
|
||||
# Create batch using model credentials
|
||||
|
|
@ -488,6 +499,7 @@ async def create_batch(
|
|||
prepare_data_with_credentials(
|
||||
data=_create_batch_data,
|
||||
credentials=credentials,
|
||||
keep_request_buckets=True,
|
||||
)
|
||||
|
||||
# Create batch using model credentials
|
||||
|
|
@ -506,6 +518,7 @@ async def create_batch(
|
|||
llm_router=llm_router,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
keep_request_buckets=True,
|
||||
)
|
||||
_raise_not_found_when_openai_fallback_unservable(
|
||||
requested_provider=requested_provider,
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ from litellm.proxy._types import (
|
|||
UpdateKeyRequest,
|
||||
UpdateTeamRequest,
|
||||
)
|
||||
from litellm.proxy.auth.auth_utils import TEAM_S3_BUCKET_ALLOWLIST_METADATA_KEY
|
||||
from litellm.types.proxy.management_endpoints.scim_v2 import (
|
||||
SCIM_MANAGED_TEAM_METADATA_KEY,
|
||||
SCIM_TEAM_DATA_METADATA_KEY,
|
||||
|
|
@ -97,6 +98,7 @@ PROXY_CONTROLLED_TEAM_METADATA_KEYS: Final[frozenset[str]] = frozenset(
|
|||
"priority",
|
||||
"zguard_policy_id",
|
||||
"allowed_tools",
|
||||
TEAM_S3_BUCKET_ALLOWLIST_METADATA_KEY,
|
||||
"allow_client_mock_response",
|
||||
"allow_client_message_redaction_opt_out",
|
||||
"allow_client_pricing_override",
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import MAX_FILE_LIST_LIMIT
|
||||
from litellm.litellm_core_utils.batch_utils import batch_cost_is_final
|
||||
from litellm.litellm_core_utils.cloud_storage_security import REQUEST_S3_BUCKET_PARAMS
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.llms.base_llm.managed_resources.utils import is_base64_encoded_unified_id
|
||||
from litellm.proxy._types import ProxyException
|
||||
|
|
@ -884,6 +885,8 @@ def apply_team_provider_credentials(
|
|||
llm_router: Optional["Router"],
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
custom_llm_provider: str,
|
||||
*,
|
||||
keep_request_buckets: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Resolve credentials for a provider-only request (no model pinned) via
|
||||
|
|
@ -901,7 +904,11 @@ def apply_team_provider_credentials(
|
|||
)
|
||||
if credentials is None:
|
||||
return
|
||||
prepare_data_with_credentials(data=data, credentials={**credentials, **explicit_request_credentials(data)})
|
||||
prepare_data_with_credentials(
|
||||
data=data,
|
||||
credentials={**credentials, **explicit_request_credentials(data)},
|
||||
keep_request_buckets=keep_request_buckets,
|
||||
)
|
||||
|
||||
|
||||
def add_deployment_model_info(
|
||||
|
|
@ -925,10 +932,16 @@ def add_deployment_model_info(
|
|||
}
|
||||
|
||||
|
||||
_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object])
|
||||
_NO_REQUEST_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
def prepare_data_with_credentials(
|
||||
data: dict,
|
||||
credentials: dict,
|
||||
file_id: str | None = None,
|
||||
*,
|
||||
keep_request_buckets: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Update data dictionary with model credentials (in-place).
|
||||
|
|
@ -937,8 +950,15 @@ def prepare_data_with_credentials(
|
|||
data: Data dictionary to update
|
||||
credentials: Credentials from router
|
||||
file_id: Optional original file_id to set (for decoded file IDs)
|
||||
keep_request_buckets: Keep an S3 bucket the request already carries over the deployment's. Only
|
||||
file upload and batch create set it, after checking the bucket against the caller's allowlist.
|
||||
Reads keep the deployment's bucket so a request cannot widen which buckets a file id may point at.
|
||||
"""
|
||||
data.update(credentials)
|
||||
request_fields: Final = _REQUEST_FIELDS.validate_python(data) if keep_request_buckets else _NO_REQUEST_FIELDS
|
||||
request_buckets: Final = MappingProxyType(
|
||||
{key: request_fields[key] for key in REQUEST_S3_BUCKET_PARAMS if key in request_fields}
|
||||
)
|
||||
data.update({**credentials, **request_buckets})
|
||||
data.pop("custom_llm_provider", None)
|
||||
|
||||
if file_id is not None:
|
||||
|
|
|
|||
|
|
@ -40,6 +40,7 @@ from litellm.llms.base_llm.files.litellm_db_storage_backend import LITELLM_DB_ST
|
|||
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints, FileListModelScope
|
||||
from litellm.llms.base_llm.managed_resources.isolation import build_list_page
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.auth.auth_utils import enforce_team_s3_bucket_allowlist
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.batches_endpoints.litellm_executed_batches import (
|
||||
litellm_executed_provider_of,
|
||||
|
|
@ -181,6 +182,7 @@ async def _litellm_executed_batch_input_model(
|
|||
|
||||
_MAX_BATCH_FILE_SIZE_MB_ADAPTER: Final = TypeAdapter(int | None)
|
||||
_LISTED_FILES_ADAPTER: Final = TypeAdapter(list[OpenAIFileObject])
|
||||
_REQUEST_BODY_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
class UploadedFileInfo(TypedDict):
|
||||
|
|
@ -465,6 +467,7 @@ async def route_create_file(
|
|||
prepare_data_with_credentials(
|
||||
data=_create_file_request,
|
||||
credentials={**credentials, **deployment_drop_params(llm_router, model, _create_file_request)},
|
||||
keep_request_buckets=True,
|
||||
)
|
||||
|
||||
# Create the file with model credentials
|
||||
|
|
@ -528,6 +531,7 @@ async def route_create_file(
|
|||
llm_router=llm_router,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
keep_request_buckets=True,
|
||||
)
|
||||
# get configs for custom_llm_provider
|
||||
llm_provider_config: Final = get_files_provider_config(custom_llm_provider=custom_llm_provider)
|
||||
|
|
@ -630,6 +634,12 @@ async def create_file(
|
|||
|
||||
# Extract file creation parameters using utility function
|
||||
request_body: Final = await read_request_body(request=request) or {}
|
||||
enforce_team_s3_bucket_allowlist(
|
||||
request_body=_REQUEST_BODY_ADAPTER.validate_python(request_body),
|
||||
general_settings=general_settings_view(),
|
||||
llm_router=llm_router,
|
||||
team_metadata=user_api_key_dict.team_metadata,
|
||||
)
|
||||
file_params: Final = await extract_file_creation_params(
|
||||
request=request,
|
||||
request_body=request_body,
|
||||
|
|
@ -707,6 +717,8 @@ async def create_file(
|
|||
# Add litellm_metadata to data if provided (from form field)
|
||||
if extracted_litellm_metadata is not None:
|
||||
data["litellm_metadata"] = extracted_litellm_metadata
|
||||
if isinstance(s3_bucket_name := form_data_raw.get("s3_bucket_name"), str) and s3_bucket_name:
|
||||
data["s3_bucket_name"] = s3_bucket_name
|
||||
|
||||
if expires_after_anchor is not None or expires_after_seconds_str is not None:
|
||||
if expires_after_anchor is None or expires_after_seconds_str is None:
|
||||
|
|
|
|||
|
|
@ -1072,6 +1072,28 @@ def respx_interceptable_httpx_client(monkeypatch):
|
|||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_bedrock_output_written_to_a_request_bucket_is_read_from_that_bucket(respx_interceptable_httpx_client):
|
||||
route: Final = respx.get(
|
||||
"https://s3.us-west-2.amazonaws.com/team-a-bucket/litellm-batch-outputs/job-1/input.jsonl.out"
|
||||
).mock(return_value=httpx.Response(200, content=b'{"recordId": "1"}\n'))
|
||||
|
||||
file_content: Final = await bu.fetch_batch_managed_file_content(
|
||||
"s3://team-a-bucket/litellm-batch-outputs/job-1/input.jsonl.out",
|
||||
custom_llm_provider="bedrock", # pyright: ignore[reportArgumentType] # the cost job passes bedrock at runtime
|
||||
litellm_params={
|
||||
"s3_bucket_name": "model-bucket",
|
||||
"aws_access_key_id": "AKIAEXAMPLE",
|
||||
"aws_secret_access_key": "secret",
|
||||
"aws_region_name": "us-west-2",
|
||||
},
|
||||
)
|
||||
|
||||
assert route.call_count == 1
|
||||
assert file_content == b'{"recordId": "1"}\n'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_output_file_content_vertex_managed_uri_accepted_by_real_validation(respx_interceptable_httpx_client):
|
||||
|
|
|
|||
|
|
@ -3520,6 +3520,87 @@ async def test_afile_retrieve_fetches_from_provider_when_file_object_none():
|
|||
assert managed_file_table.upsert_calls == []
|
||||
|
||||
|
||||
TEAM_BUCKET_FILE_ID: Final = "s3://team-a-bucket/litellm-batch-outputs/job-1/input.jsonl.out"
|
||||
|
||||
|
||||
def _managed_files_holding_a_team_bucket_file() -> PROXY_LiteLLMManagedFiles:
|
||||
internal_usage_cache: Final = MagicMock()
|
||||
internal_usage_cache.async_get_cache = AsyncMock(
|
||||
return_value={
|
||||
"unified_file_id": _managed_output_file_id(TEAM_BUCKET_FILE_ID),
|
||||
"model_mappings": {"model-1": TEAM_BUCKET_FILE_ID},
|
||||
"flat_model_file_ids": [TEAM_BUCKET_FILE_ID],
|
||||
"file_object": None,
|
||||
"created_by": "user-a",
|
||||
"updated_by": "user-a",
|
||||
}
|
||||
)
|
||||
internal_usage_cache.async_set_cache = AsyncMock()
|
||||
prisma_client: Final = AsyncMock()
|
||||
prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None)
|
||||
managed_files: Final = PROXY_LiteLLMManagedFiles(
|
||||
internal_usage_cache=internal_usage_cache, prisma_client=prisma_client
|
||||
)
|
||||
managed_files.delete_unified_file_id = AsyncMock(return_value=None)
|
||||
return managed_files
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deleting_a_managed_file_in_a_team_bucket_names_that_bucket_to_the_provider():
|
||||
router: Final = MagicMock()
|
||||
router.afile_delete = AsyncMock(return_value=MagicMock())
|
||||
|
||||
await _managed_files_holding_a_team_bucket_file().afile_delete(
|
||||
file_id=_managed_output_file_id(TEAM_BUCKET_FILE_ID), litellm_parent_otel_span=None, llm_router=router
|
||||
)
|
||||
|
||||
assert router.afile_delete.call_args.kwargs["file_id"] == TEAM_BUCKET_FILE_ID
|
||||
assert router.afile_delete.call_args.kwargs["stored_s3_bucket_name"] == "team-a-bucket"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reading_a_managed_file_in_a_team_bucket_names_that_bucket_to_the_provider():
|
||||
router: Final = MagicMock()
|
||||
router.afile_content = AsyncMock(return_value=MagicMock())
|
||||
|
||||
await _managed_files_holding_a_team_bucket_file().afile_content(
|
||||
file_id=_managed_output_file_id(TEAM_BUCKET_FILE_ID), litellm_parent_otel_span=None, llm_router=router
|
||||
)
|
||||
|
||||
assert router.afile_content.call_args.kwargs["file_id"] == TEAM_BUCKET_FILE_ID
|
||||
assert router.afile_content.call_args.kwargs["stored_s3_bucket_name"] == "team-a-bucket"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refreshing_a_managed_file_in_a_team_bucket_names_that_bucket_to_the_provider():
|
||||
from litellm.types.llms.openai import OpenAIFileObject
|
||||
|
||||
stored_file: Final = LiteLLM_ManagedFileTable(
|
||||
unified_file_id="unified-team-file",
|
||||
file_object=None,
|
||||
model_mappings={"model-1": TEAM_BUCKET_FILE_ID},
|
||||
flat_model_file_ids=[TEAM_BUCKET_FILE_ID],
|
||||
created_by="user-a",
|
||||
)
|
||||
proxy_managed_files, _ = _managed_files_with_fake_prisma(stored_file)
|
||||
router: Final = MagicMock()
|
||||
router.registry.return_value.credentials = MagicMock(
|
||||
return_value={"custom_llm_provider": "bedrock", "s3_bucket_name": "model-bucket"}
|
||||
)
|
||||
provider_file: Final = OpenAIFileObject(
|
||||
id=TEAM_BUCKET_FILE_ID, object="file", bytes=1, created_at=1, filename="input.jsonl.out", purpose="batch_output"
|
||||
)
|
||||
|
||||
with patch("litellm.afile_retrieve", new_callable=AsyncMock, return_value=provider_file) as provider_retrieve:
|
||||
refreshed: Final = await proxy_managed_files.afile_retrieve(
|
||||
file_id="unified-team-file", litellm_parent_otel_span=None, llm_router=router
|
||||
)
|
||||
|
||||
assert (refreshed.id, refreshed.filename) == ("unified-team-file", "input.jsonl.out")
|
||||
assert provider_retrieve.call_args.kwargs["file_id"] == TEAM_BUCKET_FILE_ID
|
||||
assert provider_retrieve.call_args.kwargs["stored_s3_bucket_name"] == "team-a-bucket"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_retrieve_raises_error_when_no_router_and_file_object_none():
|
||||
"""
|
||||
|
|
@ -5665,6 +5746,44 @@ async def test_provider_format_file_without_ownership_row_stays_accessible():
|
|||
assert result["file_id"] == RAW_PROVIDER_FILE_ID
|
||||
|
||||
|
||||
def _managed_output_file_id(provider_file_id: str) -> str:
|
||||
unified: Final = (
|
||||
"litellm_proxy:application/json;unified_id,u-1;target_model_names,claude-batch;"
|
||||
f"llm_output_file_id,{provider_file_id};llm_output_file_model_id,model-1"
|
||||
)
|
||||
return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"provider_file_id,stored_bucket",
|
||||
[("s3://team-a-bucket/litellm-batch-outputs/job-1/input.jsonl.out", "team-a-bucket"), ("file-openai-1", None)],
|
||||
)
|
||||
async def test_owned_managed_file_content_is_read_from_the_bucket_litellm_recorded(
|
||||
provider_file_id: str, stored_bucket: str | None
|
||||
):
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
prisma_client = AsyncMock()
|
||||
prisma_client.db.litellm_managedfiletable.find_first.return_value = _owned_record(
|
||||
created_by="user_a", team_id="team_a"
|
||||
)
|
||||
proxy_managed_files = PROXY_LiteLLMManagedFiles(MagicMock(), prisma_client=prisma_client)
|
||||
|
||||
result = await proxy_managed_files.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="user_a", team_id="team_a", parent_otel_span=MagicMock()),
|
||||
cache=MagicMock(),
|
||||
data={"file_id": _managed_output_file_id(provider_file_id)},
|
||||
call_type="afile_content",
|
||||
)
|
||||
|
||||
assert (result["model"], result["file_id"], result.get("stored_s3_bucket_name")) == (
|
||||
"model-1",
|
||||
provider_file_id,
|
||||
stored_bucket,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("batch_id", [MODEL_ENCODED_BATCH_ID, RAW_PROVIDER_BATCH_ID])
|
||||
async def test_post_call_batch_create_stores_ownership_row(batch_id):
|
||||
|
|
|
|||
|
|
@ -2928,6 +2928,88 @@ class TestBedrockFileContentTransformation:
|
|||
assert response.content == b'{"recordId": "x"}'
|
||||
|
||||
|
||||
class TestBedrockFileContentFromTheStoredBucket:
|
||||
TEAM_OUTPUT_URI: Final = "s3://team-a-bucket/litellm-batch-outputs/job-1/input.jsonl.out"
|
||||
TEAM_OUTPUT_URL: Final = "https://s3.us-west-2.amazonaws.com/team-a-bucket/litellm-batch-outputs/job-1/input.jsonl.out"
|
||||
CREDENTIALS: Final = {
|
||||
"aws_access_key_id": "AKIAEXAMPLE",
|
||||
"aws_secret_access_key": "secret",
|
||||
"aws_region_name": "us-west-2",
|
||||
}
|
||||
|
||||
def test_file_in_another_bucket_is_refused_without_the_stored_bucket(self) -> None:
|
||||
with pytest.raises(Exception, match="does not match the configured storage bucket"):
|
||||
litellm.file_content(
|
||||
file_id=self.TEAM_OUTPUT_URI,
|
||||
custom_llm_provider="bedrock",
|
||||
s3_bucket_name="model-bucket",
|
||||
**self.CREDENTIALS,
|
||||
)
|
||||
|
||||
def test_file_in_the_stored_bucket_is_read_from_that_bucket(self) -> None:
|
||||
import httpx
|
||||
import respx
|
||||
|
||||
with respx.mock:
|
||||
route: Final = respx.get(self.TEAM_OUTPUT_URL).mock(return_value=httpx.Response(200, content=b"{}"))
|
||||
response: Final = litellm.file_content(
|
||||
file_id=self.TEAM_OUTPUT_URI,
|
||||
custom_llm_provider="bedrock",
|
||||
s3_bucket_name="model-bucket",
|
||||
stored_s3_bucket_name="team-a-bucket",
|
||||
**self.CREDENTIALS,
|
||||
)
|
||||
|
||||
assert route.call_count == 1
|
||||
assert response.content == b"{}"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"stored_file_id",
|
||||
[
|
||||
"s3://team-a-bucket/tenant/litellm-batch-outputs/job-1/input.jsonl.out",
|
||||
"s3://team-a-bucket/tenant/litellm-bedrock-files-nova-1.jsonl",
|
||||
"s3://team-a-bucket/tenant/deeper/litellm-bedrock-files/nova-1.jsonl",
|
||||
],
|
||||
)
|
||||
def test_file_written_under_a_prefixed_request_bucket_is_read_back_from_it(self, stored_file_id: str) -> None:
|
||||
import httpx
|
||||
import respx
|
||||
|
||||
from litellm.litellm_core_utils.cloud_storage_security import stored_s3_bucket_kwargs
|
||||
|
||||
url: Final = f"https://s3.us-west-2.amazonaws.com/{stored_file_id.removeprefix('s3://')}"
|
||||
with respx.mock:
|
||||
route: Final = respx.get(url).mock(return_value=httpx.Response(200, content=b"{}"))
|
||||
litellm.file_content(
|
||||
file_id=stored_file_id,
|
||||
custom_llm_provider="bedrock",
|
||||
s3_bucket_name="model-bucket",
|
||||
**stored_s3_bucket_kwargs(stored_file_id),
|
||||
**self.CREDENTIALS,
|
||||
)
|
||||
|
||||
assert route.call_count == 1
|
||||
|
||||
def test_file_under_the_model_bucket_prefix_keeps_its_url_when_a_stored_bucket_is_passed(self) -> None:
|
||||
import httpx
|
||||
import respx
|
||||
|
||||
prefixed_url: Final = (
|
||||
"https://s3.us-west-2.amazonaws.com/model-bucket/tenant/litellm-batch-outputs/job-1/input.jsonl.out"
|
||||
)
|
||||
with respx.mock:
|
||||
route: Final = respx.get(prefixed_url).mock(return_value=httpx.Response(200, content=b"{}"))
|
||||
litellm.file_content(
|
||||
file_id="s3://model-bucket/tenant/litellm-batch-outputs/job-1/input.jsonl.out",
|
||||
custom_llm_provider="bedrock",
|
||||
s3_bucket_name="model-bucket/tenant",
|
||||
stored_s3_bucket_name="model-bucket",
|
||||
**self.CREDENTIALS,
|
||||
)
|
||||
|
||||
assert route.call_count == 1
|
||||
|
||||
|
||||
class TestBedrockFilesS3SignatureEncoding:
|
||||
"""
|
||||
S3 rebuilds the canonical request from the wire path with single percent-encoding,
|
||||
|
|
|
|||
|
|
@ -4,7 +4,8 @@ Unit tests for auth_utils functions related to rate limiting and customer ID ext
|
|||
|
||||
import base64
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Callable, Mapping
|
||||
from types import SimpleNamespace
|
||||
from typing import Final, Optional
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
|
@ -12,7 +13,7 @@ import pytest
|
|||
from fastapi import HTTPException, Request
|
||||
|
||||
import litellm
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_utils import (
|
||||
_allow_model_level_clientside_configurable_parameters,
|
||||
_get_customer_id_from_standard_headers,
|
||||
|
|
@ -33,7 +34,9 @@ from litellm.proxy.auth.auth_utils import (
|
|||
get_project_model_rpm_limit,
|
||||
get_project_model_tpm_limit,
|
||||
get_request_route_template,
|
||||
enforce_team_s3_bucket_allowlist,
|
||||
is_request_body_safe,
|
||||
pre_db_read_auth_checks,
|
||||
)
|
||||
from litellm.router import Router
|
||||
from litellm.types.utils import oauth_token_exchange_litellm_params, server_owned_wif_litellm_params
|
||||
|
|
@ -5429,6 +5432,189 @@ def test_other_blocked_fields_do_not_point_at_the_storage_allowlist() -> None:
|
|||
assert "allowed_client_side" not in str(error.value)
|
||||
|
||||
|
||||
_TEAM_A_METADATA: Final = {"allowed_s3_buckets": ["team-a-bucket"]}
|
||||
|
||||
|
||||
def _enforce_for_team(
|
||||
body: Mapping[str, object], team_metadata: Mapping[str, object] | None, **general_settings: object
|
||||
) -> None:
|
||||
enforce_team_s3_bucket_allowlist(
|
||||
request_body=body,
|
||||
general_settings=general_settings,
|
||||
llm_router=None,
|
||||
team_metadata=team_metadata,
|
||||
)
|
||||
|
||||
|
||||
def _rejection_for_team(body: Mapping[str, object], team_metadata: Mapping[str, object] | None) -> HTTPException:
|
||||
with pytest.raises(HTTPException) as rejection:
|
||||
_enforce_for_team(body, team_metadata)
|
||||
assert rejection.value.status_code == 401, rejection.value.detail
|
||||
return rejection.value
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["s3_bucket_name", "s3_output_bucket_name"])
|
||||
def test_team_allowlist_admits_its_own_bucket_and_rejects_another_teams(field: str) -> None:
|
||||
_enforce_for_team({field: "team-a-bucket"}, _TEAM_A_METADATA)
|
||||
rejection: Final = _rejection_for_team({field: "team-b-bucket"}, _TEAM_A_METADATA)
|
||||
assert f"{field} is not allowed in request body" in str(rejection.detail)
|
||||
assert "metadata.allowed_s3_buckets" in str(rejection.detail)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"team_metadata",
|
||||
[None, {}, {"allowed_s3_buckets": None}, {"allowed_s3_buckets": "team-a-bucket"}, {"allowed_s3_buckets": [1]}],
|
||||
)
|
||||
def test_team_without_a_usable_allowlist_cannot_pick_a_bucket(team_metadata: Mapping[str, object] | None) -> None:
|
||||
_rejection_for_team({"s3_bucket_name": "team-a-bucket"}, team_metadata)
|
||||
|
||||
|
||||
def _team_bucket_admitted(
|
||||
body: Mapping[str, object], team_metadata: Mapping[str, object] | None, **general_settings: object
|
||||
) -> bool:
|
||||
try:
|
||||
_enforce_for_team(body, team_metadata, **general_settings)
|
||||
except HTTPException:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def test_proxy_wide_allowlist_still_applies_alongside_the_team_allowlist() -> None:
|
||||
assert _team_bucket_admitted(
|
||||
{"s3_bucket_name": "team-a-bucket", "s3_output_bucket_name": "shared-bucket"},
|
||||
_TEAM_A_METADATA,
|
||||
allowed_client_side_s3_buckets=["shared-bucket"],
|
||||
)
|
||||
assert not _team_bucket_admitted({"s3_output_bucket_name": "shared-bucket"}, _TEAM_A_METADATA)
|
||||
|
||||
|
||||
def test_team_allowlist_is_checked_for_buckets_nested_in_the_body() -> None:
|
||||
_enforce_for_team({"extra_body": {"s3_bucket_name": "team-a-bucket"}}, _TEAM_A_METADATA)
|
||||
_rejection_for_team({"litellm_metadata": {"s3_bucket_name": "team-b-bucket"}}, _TEAM_A_METADATA)
|
||||
|
||||
|
||||
def test_team_allowlist_admits_no_other_blocked_field() -> None:
|
||||
rejection: Final = _rejection_for_team({"s3_endpoint_url": "team-a-bucket"}, _TEAM_A_METADATA)
|
||||
assert "s3_endpoint_url" in str(rejection.detail)
|
||||
|
||||
|
||||
def _routed_request(method: str, route_template: str | None, path: str) -> Request:
|
||||
scope: Final[dict[str, object]] = {"type": "http", "method": method, "path": path, "headers": [], "query_string": b""}
|
||||
return Request(scope if route_template is None else {**scope, "route": SimpleNamespace(path=route_template)})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"route_template,path",
|
||||
[
|
||||
("/v1/files", "/v1/files"),
|
||||
("/files", "/files"),
|
||||
("/{provider}/v1/files", "/bedrock/v1/files"),
|
||||
("/v1/batches", "/v1/batches"),
|
||||
("/batches", "/batches"),
|
||||
("/{provider}/v1/batches", "/bedrock/v1/batches"),
|
||||
],
|
||||
)
|
||||
async def test_auth_leaves_request_buckets_to_the_upload_and_batch_create_handlers(
|
||||
route_template: str, path: str
|
||||
) -> None:
|
||||
assert await _auth_admits(
|
||||
_routed_request("POST", route_template, path),
|
||||
{"s3_bucket_name": "team-b-bucket", "s3_output_bucket_name": "team-b-bucket"},
|
||||
path,
|
||||
)
|
||||
|
||||
|
||||
async def _auth_admits(request: Request, request_data: dict[str, object], route: str) -> bool:
|
||||
try:
|
||||
await pre_db_read_auth_checks(request=request, request_data=request_data, route=route)
|
||||
except ValueError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"method,route_template,path",
|
||||
[
|
||||
("GET", "/v1/files", "/v1/files"),
|
||||
("POST", "/v1/files/{file_id}/content", "/v1/files/file-1/content"),
|
||||
("POST", "/v1/chat/completions", "/v1/chat/completions"),
|
||||
("POST", None, "/v1/batches"),
|
||||
],
|
||||
)
|
||||
async def test_auth_still_rejects_request_buckets_on_every_other_route(
|
||||
method: str, route_template: str | None, path: str
|
||||
) -> None:
|
||||
with pytest.raises(ValueError, match="s3_bucket_name is not allowed"):
|
||||
await pre_db_read_auth_checks(
|
||||
request=_routed_request(method, route_template, path),
|
||||
request_data={"s3_bucket_name": "team-b-bucket"},
|
||||
route=path,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_caller_can_never_send_the_stored_bucket_param() -> None:
|
||||
with pytest.raises(ValueError, match="stored_s3_bucket_name is not allowed"):
|
||||
await pre_db_read_auth_checks(
|
||||
request=_routed_request("POST", "/v1/batches", "/v1/batches"),
|
||||
request_data={"stored_s3_bucket_name": "team-b-bucket"},
|
||||
route="/v1/batches",
|
||||
)
|
||||
rejection: Final = _rejection_for_team(
|
||||
{"stored_s3_bucket_name": "team-a-bucket"}, {"allowed_s3_buckets": ["team-a-bucket"]}
|
||||
)
|
||||
assert "stored_s3_bucket_name" in str(rejection.detail)
|
||||
with pytest.raises(ValueError, match="stored_s3_bucket_name"):
|
||||
_batches_body_is_safe({"stored_s3_bucket_name": "team-a-bucket"}, _BUCKET_ALLOWLIST)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_still_rejects_other_blocked_fields_on_the_upload_and_batch_create_routes() -> None:
|
||||
with pytest.raises(ValueError, match="api_base"):
|
||||
await pre_db_read_auth_checks(
|
||||
request=_routed_request("POST", "/v1/batches", "/v1/batches"),
|
||||
request_data={"s3_bucket_name": "team-b-bucket", "api_base": "https://attacker.example"},
|
||||
route="/v1/batches",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"role,stored,sent,refused",
|
||||
[
|
||||
(LitellmUserRoles.PROXY_ADMIN, None, ["any-bucket"], False),
|
||||
(LitellmUserRoles.INTERNAL_USER, None, ["any-bucket"], True),
|
||||
(LitellmUserRoles.INTERNAL_USER, ["team-a-bucket"], ["team-a-bucket", "other"], True),
|
||||
(LitellmUserRoles.INTERNAL_USER, ["team-a-bucket"], None, True),
|
||||
(LitellmUserRoles.INTERNAL_USER, ["team-a-bucket"], ["team-a-bucket"], False),
|
||||
],
|
||||
)
|
||||
def test_only_proxy_admins_change_a_teams_s3_bucket_list(
|
||||
role: LitellmUserRoles, stored: list[str] | None, sent: list[str] | None, refused: bool
|
||||
) -> None:
|
||||
from litellm.proxy._types import UpdateTeamRequest
|
||||
from litellm.proxy.auth.auth_utils import enforce_batch_limits_are_admin_only
|
||||
|
||||
existing: Final = None if stored is None else {"allowed_s3_buckets": stored, "project": "x"}
|
||||
request: Final = UpdateTeamRequest(
|
||||
team_id="t1", metadata={"project": "x"} if sent is None else {"allowed_s3_buckets": sent, "project": "x"}
|
||||
)
|
||||
caller: Final = UserAPIKeyAuth(user_role=role, user_id="u1")
|
||||
|
||||
if not refused:
|
||||
enforce_batch_limits_are_admin_only(
|
||||
data=request, existing_metadata=existing, user_api_key_dict=caller, entity="team"
|
||||
)
|
||||
return
|
||||
with pytest.raises(HTTPException) as rejection:
|
||||
enforce_batch_limits_are_admin_only(
|
||||
data=request, existing_metadata=existing, user_api_key_dict=caller, entity="team"
|
||||
)
|
||||
assert rejection.value.status_code == 403
|
||||
assert "Only proxy admins can set allowed_s3_buckets on a team." in str(rejection.value.detail)
|
||||
|
||||
|
||||
def _request_with_content_length(content_length: int) -> Request:
|
||||
return Request(
|
||||
{
|
||||
|
|
|
|||
|
|
@ -85,6 +85,12 @@ CREDS: Dict[str, Dict[str, str]] = {
|
|||
"api_base": "http://vllm.test/v1",
|
||||
"model": "hosted_vllm/qwen",
|
||||
},
|
||||
"bedrock-batch": {
|
||||
"custom_llm_provider": "bedrock",
|
||||
"model": "bedrock/us.amazon.nova-lite-v1:0",
|
||||
"s3_bucket_name": "model-bucket",
|
||||
"s3_output_bucket_name": "model-bucket",
|
||||
},
|
||||
}
|
||||
|
||||
# A real model-encoded file id: decodes to "azure/gpt-4o", strips to "file-original123".
|
||||
|
|
@ -426,6 +432,83 @@ async def test_create__model_from_body(harness):
|
|||
assert resp.id == encode_file_id_with_model("batch-provider-id", "vertex-model", id_type="batch")
|
||||
|
||||
|
||||
_TEAM_A_KEY: Final = UserAPIKeyAuth(
|
||||
api_key="sk-team-a", team_id="team-a", team_metadata={"allowed_s3_buckets": ["team-a-bucket"]}
|
||||
)
|
||||
|
||||
|
||||
def _batch_body_with_output_bucket(bucket: str) -> dict[str, object]:
|
||||
return {
|
||||
"input_file_id": "file-plain",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
"model": "vertex-model",
|
||||
"s3_output_bucket_name": bucket,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create__output_bucket_on_the_team_allowlist_reaches_the_provider(harness):
|
||||
set_body(harness, _batch_body_with_output_bucket("team-a-bucket"))
|
||||
|
||||
await call_create(harness, user=_TEAM_A_KEY)
|
||||
|
||||
assert harness.acreate_kwargs()["s3_output_bucket_name"] == "team-a-bucket"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create__team_with_an_allowlist_and_no_request_bucket_leaves_the_bucket_to_the_model(harness):
|
||||
set_body(
|
||||
harness,
|
||||
{
|
||||
"input_file_id": "file-plain",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
"model": "vertex-model",
|
||||
},
|
||||
)
|
||||
|
||||
await call_create(harness, user=_TEAM_A_KEY)
|
||||
|
||||
assert "s3_output_bucket_name" not in harness.acreate_kwargs()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"request_bucket,expected_bucket", [("team-a-bucket", "team-a-bucket"), (None, "model-bucket")]
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"routing",
|
||||
[
|
||||
{"input_file_id": "file-plain", "model": "bedrock-batch"},
|
||||
{"input_file_id": encode_file_id_with_model("file-plain", "bedrock-batch")},
|
||||
],
|
||||
ids=["model_in_body", "model_encoded_file_id"],
|
||||
)
|
||||
async def test_create__model_routed_batch_keeps_the_team_bucket_over_the_deployment_bucket(
|
||||
harness, routing: dict[str, str], request_bucket: str | None, expected_bucket: str
|
||||
):
|
||||
body: Final = {"endpoint": "/v1/chat/completions", "completion_window": "24h", **routing}
|
||||
set_body(harness, body if request_bucket is None else {**body, "s3_output_bucket_name": request_bucket})
|
||||
|
||||
await call_create(harness, user=_TEAM_A_KEY)
|
||||
|
||||
assert harness.acreate_kwargs()["s3_output_bucket_name"] == expected_bucket
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create__output_bucket_off_the_team_allowlist_is_rejected_before_the_provider(harness):
|
||||
set_body(harness, _batch_body_with_output_bucket("team-b-bucket"))
|
||||
|
||||
with pytest.raises(ProxyException) as rejection:
|
||||
await call_create(harness, user=_TEAM_A_KEY)
|
||||
|
||||
assert str(rejection.value.code) == "401"
|
||||
assert "s3_output_bucket_name is not allowed in request body" in rejection.value.message
|
||||
harness.litellm_acreate.assert_not_called()
|
||||
harness.router_acreate.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create__model_from_header(harness):
|
||||
set_body(
|
||||
|
|
|
|||
|
|
@ -2297,6 +2297,10 @@ class TestManagedOutputFileIdEncodesPublicModelGroup:
|
|||
return job
|
||||
|
||||
async def _run(self, job: MagicMock) -> str:
|
||||
output_file_id, _ = await self._run_capturing(job, self._RAW_OUTPUT_FILE_ID)
|
||||
return output_file_id
|
||||
|
||||
async def _run_capturing(self, job: MagicMock, raw_output_file_id: str) -> tuple[str, Mapping[str, object]]:
|
||||
from litellm_enterprise.proxy.common_utils.check_batch_cost import (
|
||||
CheckBatchCost,
|
||||
)
|
||||
|
|
@ -2342,7 +2346,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup:
|
|||
object="batch",
|
||||
status="completed",
|
||||
)
|
||||
response.output_file_id = self._RAW_OUTPUT_FILE_ID
|
||||
response.output_file_id = raw_output_file_id
|
||||
|
||||
file_content = MagicMock()
|
||||
file_content.content = b'{"id":"req-1"}'
|
||||
|
|
@ -2352,7 +2356,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup:
|
|||
"litellm.files.main.afile_content",
|
||||
new_callable=AsyncMock,
|
||||
return_value=file_content,
|
||||
),
|
||||
) as provider_file_content,
|
||||
patch(
|
||||
"litellm.batches.batch_utils.get_file_content_as_dictionary",
|
||||
return_value=[{"id": "req-1"}],
|
||||
|
|
@ -2376,7 +2380,17 @@ class TestManagedOutputFileIdEncodesPublicModelGroup:
|
|||
prom_logger=None,
|
||||
)
|
||||
|
||||
return response.output_file_id
|
||||
return response.output_file_id, provider_file_content.call_args.kwargs
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_output_written_to_a_team_bucket_is_billed_from_that_bucket(self):
|
||||
team_output: Final = "s3://team-a-bucket/litellm-batch-outputs/job-1/input.jsonl.out"
|
||||
|
||||
_, download = await self._run_capturing(
|
||||
self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)), team_output
|
||||
)
|
||||
|
||||
assert (download["file_id"], download["stored_s3_bucket_name"]) == (team_output, "team-a-bucket")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_target_model_names_comes_from_input_file_not_provider_model(self):
|
||||
|
|
|
|||
|
|
@ -322,6 +322,7 @@ _PROXY_READ_TEAM_METADATA_KEYS: Final = (
|
|||
"allowed_vector_store_indexes",
|
||||
"enforced_file_expires_after",
|
||||
"enforced_batch_output_expires_after",
|
||||
"allowed_s3_buckets",
|
||||
"require_trace_id",
|
||||
"allowed_passthrough_routes",
|
||||
"denied_passthrough_routes",
|
||||
|
|
@ -384,7 +385,9 @@ class TestTeamMetadataEdits:
|
|||
verdict: Final = _metadata_verdict(_STORED_METADATA, {**_STORED_METADATA, key: "set-by-team-admin"})
|
||||
assert verdict == TeamAdminMetadataKeyProxyControlled(key=key)
|
||||
|
||||
@pytest.mark.parametrize("key", ["opted_out_global_guardrails", "callback_settings", "priority", "guardrails"])
|
||||
@pytest.mark.parametrize(
|
||||
"key", ["opted_out_global_guardrails", "callback_settings", "priority", "guardrails", "allowed_s3_buckets"]
|
||||
)
|
||||
def test_changing_or_removing_a_stored_proxy_controlled_key_names_it(self, key: str):
|
||||
stored: Final = {**_STORED_METADATA, key: ["stored"]}
|
||||
changed: Final = _metadata_verdict(stored, {**stored, key: ["changed"]})
|
||||
|
|
|
|||
|
|
@ -13062,6 +13062,67 @@ async def test_update_team_blocks_org_admin_require_trace_id_change(mock_db_clie
|
|||
assert "Only proxy admins can set `require_trace_id` on a team." in str(exc.value.message)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_team_blocks_org_admin_granting_s3_buckets(mock_db_client: MagicMock) -> None:
|
||||
mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0)
|
||||
|
||||
with (
|
||||
_org_admins(("u-team-admin", "org-1")),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.team_endpoints._should_auto_add_team_creator",
|
||||
return_value=False,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_org_object",
|
||||
new=AsyncMock(return_value=MagicMock()),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.team_endpoints._check_org_team_limits",
|
||||
new=AsyncMock(return_value=None),
|
||||
),
|
||||
):
|
||||
exc: Final[pytest.ExceptionInfo[ProxyException]]
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await new_team(
|
||||
data=NewTeamRequest(
|
||||
team_alias="t",
|
||||
organization_id="org-1",
|
||||
metadata={"allowed_s3_buckets": ["another-orgs-bucket"]},
|
||||
),
|
||||
http_request=MagicMock(spec=Request),
|
||||
user_api_key_dict=_non_admin_auth(),
|
||||
)
|
||||
|
||||
assert str(exc.value.code) == "403"
|
||||
assert "Only proxy admins can set allowed_s3_buckets on a team." in str(exc.value.message)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_team_blocks_org_admin_widening_s3_buckets(mock_db_client: MagicMock) -> None:
|
||||
existing: Final[MagicMock] = MagicMock()
|
||||
existing.metadata = {"allowed_s3_buckets": ["team-a-bucket"]}
|
||||
existing.model_dump.return_value = {
|
||||
"team_id": "t1",
|
||||
"organization_id": "org-1",
|
||||
"metadata": {"allowed_s3_buckets": ["team-a-bucket"]},
|
||||
}
|
||||
mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing)
|
||||
|
||||
with _org_admins(("u-team-admin", "org-1")):
|
||||
exc: Final[pytest.ExceptionInfo[ProxyException]]
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await update_team(
|
||||
data=UpdateTeamRequest(
|
||||
team_id="t1", metadata={"allowed_s3_buckets": ["team-a-bucket", "another-orgs-bucket"]}
|
||||
),
|
||||
http_request=MagicMock(spec=Request),
|
||||
user_api_key_dict=_non_admin_auth(),
|
||||
)
|
||||
|
||||
assert str(exc.value.code) == "403"
|
||||
assert "Only proxy admins can set allowed_s3_buckets on a team." in str(exc.value.message)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_team_blocks_org_admin_null_metadata_clear_require_trace_id(
|
||||
mock_db_client: MagicMock,
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
|
|||
get_team_provider_deployment,
|
||||
is_litellm_executed_batch,
|
||||
map_raw_file_ids_to_unified,
|
||||
prepare_data_with_credentials,
|
||||
)
|
||||
from litellm.proxy.route_llm_request import ProxyModelNotFoundError
|
||||
from litellm.proxy.utils import InternalUsageCache, handle_exception_on_proxy
|
||||
|
|
@ -952,3 +953,19 @@ def test_a_provider_wildcard_grant_finds_only_deployments_named_under_it(deploye
|
|||
)
|
||||
|
||||
assert (None if credentials is None else credentials["api_key"]) == api_key
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field", ["s3_bucket_name", "s3_output_bucket_name"])
|
||||
def test_create_paths_keep_a_request_bucket_and_reads_keep_the_deployment_bucket(field: str) -> None:
|
||||
deployment: Final = {"custom_llm_provider": "bedrock", "api_key": "sk-deployment", field: "model-bucket"}
|
||||
create_with_bucket: Final[dict[str, object]] = {field: "team-a-bucket"}
|
||||
create_without_bucket: Final[dict[str, object]] = {}
|
||||
read_with_bucket: Final[dict[str, object]] = {field: "team-a-bucket"}
|
||||
|
||||
prepare_data_with_credentials(data=create_with_bucket, credentials=dict(deployment), keep_request_buckets=True)
|
||||
prepare_data_with_credentials(data=create_without_bucket, credentials=dict(deployment), keep_request_buckets=True)
|
||||
prepare_data_with_credentials(data=read_with_bucket, credentials=dict(deployment))
|
||||
|
||||
assert create_with_bucket == {field: "team-a-bucket", "api_key": "sk-deployment"}
|
||||
assert create_without_bucket == {field: "model-bucket", "api_key": "sk-deployment"}
|
||||
assert read_with_bucket == {field: "model-bucket", "api_key": "sk-deployment"}
|
||||
|
|
|
|||
|
|
@ -951,6 +951,45 @@ def test_batch_upload_for_a_provider_model_still_goes_to_the_provider(batch_uplo
|
|||
assert provider_upload.call_args.kwargs["custom_llm_provider"] == "gemini"
|
||||
|
||||
|
||||
def _upload_as_team_a(bucket: str | None):
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
api_key="sk-team-a", team_id="team-a", team_metadata={"allowed_s3_buckets": ["team-a-bucket"]}
|
||||
)
|
||||
return _upload_batch_file(
|
||||
{"x-litellm-model": "gemini-2.0-flash"}, {} if bucket is None else {"s3_bucket_name": bucket}
|
||||
)
|
||||
|
||||
|
||||
def test_batch_upload_from_a_team_with_an_allowlist_and_no_bucket_leaves_the_bucket_to_the_model(batch_upload_seams):
|
||||
_, provider_upload, _ = batch_upload_seams
|
||||
|
||||
response = _upload_as_team_a(None)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert "s3_bucket_name" not in provider_upload.call_args.kwargs
|
||||
|
||||
|
||||
def test_batch_upload_sends_a_bucket_on_the_team_allowlist_to_the_provider_upload(batch_upload_seams):
|
||||
_, provider_upload, _ = batch_upload_seams
|
||||
|
||||
response = _upload_as_team_a("team-a-bucket")
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert provider_upload.call_args.kwargs["s3_bucket_name"] == "team-a-bucket"
|
||||
|
||||
|
||||
def test_batch_upload_rejects_a_bucket_off_the_team_allowlist_before_the_provider_upload(batch_upload_seams):
|
||||
_, provider_upload, _ = batch_upload_seams
|
||||
|
||||
response = _upload_as_team_a("team-b-bucket")
|
||||
|
||||
assert response.status_code == 401, response.text
|
||||
assert "s3_bucket_name is not allowed in request body" in response.text
|
||||
provider_upload.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model, expected_drop_params",
|
||||
[("vertex-batch-drop", True), ("vertex-batch-keep", False)],
|
||||
|
|
|
|||
|
|
@ -125,6 +125,32 @@ async def test_acreate_file_sends_target_model_names_to_a_litellm_proxy_deployme
|
|||
assert dict(host.calls[0].kwargs["extra_body"]) == {"keep": 1, "target_model_names": "remote-model"} # pyright: ignore[reportArgumentType] # a mapping
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"request_bucket,expected_bucket", [({}, "model-bucket"), ({"s3_bucket_name": "team-a-bucket"}, "team-a-bucket")]
|
||||
)
|
||||
async def test_acreate_file_prefers_the_request_s3_bucket_over_the_deployment_bucket(
|
||||
make_core: MakeCore, host: EndpointHost, request_bucket: dict[str, str], expected_bucket: str
|
||||
):
|
||||
core: Final = make_core(model_list=[deployment("g", "d1", s3_bucket_name="model-bucket")])
|
||||
await core.acreate_file(model="g", file=_jsonl(), purpose="batch", **request_bucket)
|
||||
assert host.calls[0].kwargs["s3_bucket_name"] == expected_bucket
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"request_bucket,expected_bucket",
|
||||
[({}, "model-bucket"), ({"s3_output_bucket_name": "team-a-bucket"}, "team-a-bucket")],
|
||||
)
|
||||
async def test_acreate_batch_prefers_the_request_output_bucket_over_the_deployment_bucket(
|
||||
make_core: MakeCore, host: EndpointHost, request_bucket: dict[str, str], expected_bucket: str
|
||||
):
|
||||
host.answer("openai/d1", _batch())
|
||||
core: Final = make_core(model_list=[deployment("g", "d1", s3_output_bucket_name="model-bucket")])
|
||||
await core.acreate_batch(
|
||||
model="g", input_file_id="f", endpoint="/v1/chat/completions", completion_window="24h", **request_bucket
|
||||
)
|
||||
assert host.calls[0].kwargs["s3_output_bucket_name"] == expected_bucket
|
||||
|
||||
|
||||
async def test_acreate_file_sends_the_deployment_gcs_bucket_in_litellm_metadata(
|
||||
make_core: MakeCore, host: EndpointHost
|
||||
):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue