diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 763737191ef..6d1f79e1278 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -52,9 +52,10 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( # noqa: F401 # apply_team_provider_credentials, authorize_model_for_key, batch_cost_poller_is_active, + batch_deployment_id, decode_model_from_file_id, encode_batch_response_ids, - encode_file_id_with_model, + encode_file_id_with_model, # pyright: ignore[reportUnusedImport] # backwards-compatible package export ensure_batch_response_managed_file_ids, get_authorized_credentials_for_model, get_batch_from_database, @@ -62,6 +63,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( # noqa: F401 # get_model_id_from_unified_batch_id, get_models_from_unified_file_id, get_original_file_id, + hinted_batch_deployment_id, is_base64_encoded_unified_file_id, is_litellm_executed_batch, prepare_data_with_credentials, @@ -360,11 +362,13 @@ async def create_batch( # SCENARIO 1: File ID is encoded with model info if model_from_file_id is not None and input_file_id: + input_file_deployment_id: Final = hinted_batch_deployment_id(llm_router, model_from_file_id, input_file_id) credentials = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_from_file_id, user_api_key_dict=user_api_key_dict, operation_context="batch creation (file created with model)", + deployment_id=input_file_deployment_id, ) original_file_id: Final = get_original_file_id(input_file_id) @@ -380,35 +384,11 @@ async def create_batch( **_create_batch_data, ) - # Encode the batch ID and related file IDs with model information - if response and hasattr(response, "id") and response.id: - original_batch_id: Final = response.id - encoded_batch_id: Final = encode_file_id_with_model( - file_id=original_batch_id, - model=model_from_file_id, - id_type="batch", - ) - response.id = encoded_batch_id - - if hasattr(response, "output_file_id") and response.output_file_id: - response.output_file_id = encode_file_id_with_model( - file_id=response.output_file_id, model=model_from_file_id - ) - - if hasattr(response, "error_file_id") and response.error_file_id: - response.error_file_id = encode_file_id_with_model( - file_id=response.error_file_id, model=model_from_file_id - ) - - verbose_proxy_logger.debug( - "Created batch using model: %s, original_batch_id: %s, encoded: %s", - model_from_file_id, - original_batch_id, - encoded_batch_id, - ) - + encode_batch_response_ids(response, model=model_from_file_id, deployment_id=input_file_deployment_id) response.input_file_id = input_file_id + verbose_proxy_logger.debug("Created batch using model: %s", model_from_file_id) + elif litellm.enable_loadbalancing_on_batch_endpoints is True and is_router_model and router_model is not None: if llm_router is None: raise HTTPException( @@ -465,11 +445,13 @@ async def create_batch( # SCENARIO 2 & 3: Model from header/query OR custom_llm_provider fallback if model_param: # SCENARIO 2: Use model-based routing from header/query/body + model_param_deployment_id: Final = batch_deployment_id(llm_router, model_param) credentials = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_param, user_api_key_dict=user_api_key_dict, operation_context="batch creation", + deployment_id=model_param_deployment_id, ) await _raise_when_input_file_must_be_managed(model_param, credentials) @@ -484,7 +466,7 @@ async def create_batch( **_create_batch_data, ) - encode_batch_response_ids(response, model=model_param) + encode_batch_response_ids(response, model=model_param, deployment_id=model_param_deployment_id) verbose_proxy_logger.debug("Created batch using model: %s", model_param) else: @@ -709,11 +691,13 @@ async def retrieve_batch( # Retrieve from provider (for non-terminal states or if DB lookup failed) # SCENARIO 1: Batch ID is encoded with model info if model_from_id is not None: + retrieve_deployment_id: Final = hinted_batch_deployment_id(llm_router, model_from_id, batch_id) credentials: Final = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_from_id, user_api_key_dict=user_api_key_dict, operation_context="batch retrieval (batch created with model)", + deployment_id=retrieve_deployment_id, ) original_batch_id: Final = get_original_file_id(batch_id) @@ -728,7 +712,9 @@ async def retrieve_batch( # so litellm.aretrieve_batch can load BedrockBatchesConfig. Without # it the call falls into the legacy provider switch and 400s. data["model"] = model_from_id - add_deployment_model_info(data=data, llm_router=llm_router, model_id=model_from_id) + add_deployment_model_info( + data=data, llm_router=llm_router, model_id=retrieve_deployment_id or model_from_id + ) _litellm_metadata_of(data).setdefault("model_group", model_from_id) # Retrieve batch using model credentials @@ -737,7 +723,7 @@ async def retrieve_batch( **data, ) - encode_batch_response_ids(response, model=model_from_id) + encode_batch_response_ids(response, model=model_from_id, deployment_id=retrieve_deployment_id) verbose_proxy_logger.debug( "Retrieved batch using model: %s, original_id: %s", model_from_id, original_batch_id @@ -939,11 +925,15 @@ async def list_batches( data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") ): # SCENARIO 2: Use model-based routing from header/query/body + list_deployment_id: Final = ( + batch_deployment_id(llm_router, model_param) if isinstance(model_param, str) else None + ) credentials: Final = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_param, user_api_key_dict=user_api_key_dict, operation_context="batch listing", + deployment_id=list_deployment_id, ) prepare_data_with_credentials(data=data, credentials=credentials) @@ -960,7 +950,7 @@ async def list_batches( response_data: Final = getattr(response, "data", None) if response_data: for batch in response_data: - encode_batch_response_ids(batch, model=model_param) + encode_batch_response_ids(batch, model=model_param, deployment_id=list_deployment_id) verbose_proxy_logger.debug("Listed batches using model: %s", model_param) @@ -1140,11 +1130,13 @@ async def cancel_batch( # SCENARIO 1: Batch ID is encoded with model info if model_from_id is not None: + cancel_deployment_id: Final = hinted_batch_deployment_id(llm_router, model_from_id, batch_id) credentials: Final = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_from_id, user_api_key_dict=user_api_key_dict, operation_context="batch cancellation (batch created with model)", + deployment_id=cancel_deployment_id, ) original_batch_id: Final = get_original_file_id(batch_id) @@ -1162,7 +1154,7 @@ async def cancel_batch( **data, ) - encode_batch_response_ids(response, model=model_from_id) + encode_batch_response_ids(response, model=model_from_id, deployment_id=cancel_deployment_id) verbose_proxy_logger.debug( "Cancelled batch using model: %s, original_id: %s", model_from_id, original_batch_id diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index f7289445d3a..5acd2e181ea 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -229,11 +229,13 @@ def is_litellm_executed_batch(decoded_unified_batch_id: str) -> bool: return bool(marker) and batch_id.startswith(LITELLM_EXECUTED_BATCH_ID_PREFIX) -def encode_file_id_with_model(file_id: str, model: str, id_type: Literal["file", "batch"] = "file") -> str: +def encode_file_id_with_model( + file_id: str, model: str, id_type: Literal["file", "batch"] = "file", deployment_id: str | None = None +) -> str: """ Encode a file/batch ID with model routing information. - Format: ;model,)> + Format: ;model,[;deployment,])> The result preserves the original prefix (file-, batch_, etc.) for OpenAI compliance. Args: @@ -242,6 +244,8 @@ def encode_file_id_with_model(file_id: str, model: str, id_type: Literal["file", id_type: Type of ID being encoded. Used to determine the correct prefix when the raw ID lacks a recognizable prefix (e.g., Vertex AI numeric IDs). Defaults to "file" for backward compatibility. + deployment_id: The group deployment that holds the resource, set only when the group's + first deployment cannot serve it (a mixed Azure chat and batch group). Returns: Encoded ID starting with appropriate prefix and containing routing information @@ -256,7 +260,8 @@ def encode_file_id_with_model(file_id: str, model: str, id_type: Literal["file", encode_file_id_with_model("3814889423749775360", "gemini-2.5-pro", id_type="batch") -> "batch_bGl0ZWxsbTozODE0ODg5NDIzNzQ5Nzc1MzYwO21vZGVsLGdlbWluaS0yLjUtcHJv" """ - encoded_str: Final = f"litellm:{file_id};model,{model}" + deployment_segment: Final = f";deployment,{deployment_id}" if deployment_id else "" + encoded_str: Final = f"litellm:{file_id};model,{model}{deployment_segment}" encoded_bytes: Final = base64.urlsafe_b64encode(encoded_str.encode()) encoded_b64: Final = encoded_bytes.decode().rstrip("=") @@ -273,47 +278,46 @@ def encode_file_id_with_model(file_id: str, model: str, id_type: Literal["file", return f"{prefix}{encoded_b64}" -def encode_batch_response_ids(response, model: str) -> None: +def encode_batch_response_ids(response, model: str, deployment_id: str | None = None) -> None: """Encode all IDs in a batch response with model routing info (in-place).""" if not response or not hasattr(response, "id") or not response.id: return - response.id = encode_file_id_with_model(file_id=response.id, model=model, id_type="batch") + response.id = encode_file_id_with_model( + file_id=response.id, model=model, id_type="batch", deployment_id=deployment_id + ) for attr in ("output_file_id", "error_file_id", "input_file_id"): if hasattr(response, attr) and getattr(response, attr): setattr( response, attr, - encode_file_id_with_model(file_id=getattr(response, attr), model=model), + encode_file_id_with_model(file_id=getattr(response, attr), model=model, deployment_id=deployment_id), ) -def decode_model_from_file_id(encoded_id: str) -> str | None: +def _model_embedded_payload(encoded_id: object) -> str | None: + if not isinstance(encoded_id, str): + return None + prefix: Final = next((prefix for prefix in ("file-", "batch_") if encoded_id.startswith(prefix)), "") + b64_part: Final = encoded_id[len(prefix) :] + try: + decoded: Final = base64.urlsafe_b64decode(b64_part + "=" * (-len(b64_part) % 4)).decode() + except ValueError: + return None + return decoded if decoded.startswith("litellm:") and ";model," in decoded else None + + +def decode_model_from_file_id(encoded_id: object) -> str | None: """ Extract model name from an encoded file/batch ID. Handles IDs that start with "file-" or "batch_" prefix. """ - try: - if not isinstance(encoded_id, str): - return None + match: Final = re.search(r";model,([^;]+)", _model_embedded_payload(encoded_id) or "") + return match.group(1).strip() if match else None - # Remove prefix if present (file-, batch_, etc.) - if encoded_id.startswith("file-"): - b64_part = encoded_id[5:] # Remove "file-" - elif encoded_id.startswith("batch_"): - b64_part = encoded_id[6:] # Remove "batch_" - else: - b64_part = encoded_id - padded: Final = b64_part + "=" * (-len(b64_part) % 4) - decoded: Final = base64.urlsafe_b64decode(padded).decode() - if decoded.startswith("litellm:") and ";model," in decoded: - match: Final = re.search(r";model,([^;]+)", decoded) - if match: - return match.group(1).strip() - - return None - except Exception: - return None +def decode_deployment_from_file_id(encoded_id: object) -> str | None: + match: Final = re.search(r";deployment,([^;]+)", _model_embedded_payload(encoded_id) or "") + return match.group(1) if match else None def get_original_file_id(encoded_id: str) -> str: @@ -321,29 +325,8 @@ def get_original_file_id(encoded_id: str) -> str: Extract the original provider file/batch ID from an encoded ID. Handles IDs that start with "file-" or "batch_" prefix. """ - try: - if not isinstance(encoded_id, str): - return encoded_id - - # Remove prefix if present (file-, batch_, etc.) - if encoded_id.startswith("file-"): - b64_part = encoded_id[5:] # Remove "file-" - elif encoded_id.startswith("batch_"): - b64_part = encoded_id[6:] # Remove "batch_" - else: - b64_part = encoded_id - - padded: Final = b64_part + "=" * (-len(b64_part) % 4) - decoded: Final = base64.urlsafe_b64decode(padded).decode() - - if decoded.startswith("litellm:") and ";model," in decoded: - match: Final = re.search(r"litellm:([^;]+);model,", decoded) - if match: - return match.group(1) - - return encoded_id - except Exception: - return encoded_id + match: Final = re.search(r"litellm:([^;]+);model,", _model_embedded_payload(encoded_id) or "") + return match.group(1) if match else encoded_id def is_model_embedded_id(file_id: str) -> bool: @@ -467,16 +450,32 @@ async def get_authorized_credentials_for_model( model_id: str, user_api_key_dict: "UserAPIKeyAuth", operation_context: str = "file operation", + deployment_id: str | None = None, ) -> dict: # mutable-ok: same contract as get_credentials_for_model, callers merge it into request data """``get_credentials_for_model`` gated by ``authorize_model_for_key``.""" await authorize_model_for_key(model_id=model_id, llm_router=llm_router, user_api_key_dict=user_api_key_dict) return get_credentials_for_model( llm_router=llm_router, - model_id=model_id, + model_id=deployment_id or model_id, operation_context=operation_context, ) +def batch_deployment_id( + llm_router: Optional["Router"], model: str, preferred_deployment_id: str | None = None +) -> str | None: + return ( + llm_router.get_batch_deployment_id(model_group_name=model, preferred_deployment_id=preferred_deployment_id) + if llm_router is not None + else None + ) + + +def hinted_batch_deployment_id(llm_router: Optional["Router"], model: str, encoded_id: str) -> str | None: + hint: Final = decode_deployment_from_file_id(encoded_id) + return batch_deployment_id(llm_router, model, preferred_deployment_id=hint) if hint else None + + def deployment_drop_params( llm_router: Optional["Router"], model_id: str, request: Mapping[str, object] ) -> Mapping[str, bool | str]: @@ -801,7 +800,7 @@ def prepare_data_with_credentials( async def handle_model_based_routing( file_id: str, request, # FastAPI Request object - llm_router, # Router instance + llm_router: Optional["Router"], data: dict, user_api_key_dict: "UserAPIKeyAuth", check_file_id_encoding: bool = True, @@ -845,6 +844,7 @@ async def handle_model_based_routing( model_id=model_from_id, user_api_key_dict=user_api_key_dict, operation_context="file operation (file created with model)", + deployment_id=hinted_batch_deployment_id(llm_router, model_from_id, file_id), ) original_file_id: Final = get_original_file_id(file_id) return True, model_from_id, original_file_id, credentials diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 5d63c3db222..0a16af6efce 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -80,6 +80,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( # noqa: F401 # add_internal_model_credentials, apply_team_provider_credentials, authorize_model_for_key, + batch_deployment_id, deployment_drop_params, encode_file_id_with_model, extract_file_creation_params, @@ -454,18 +455,23 @@ async def route_create_file( # NEW: Handle model-based routing (no DB required) if model is not None: + deployment_id: Final = batch_deployment_id(llm_router, model) if purpose == "batch" else None # Get credentials from model_list via router credentials: Final = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model, user_api_key_dict=user_api_key_dict, operation_context="file upload", + deployment_id=deployment_id, ) # Merge credentials into the request prepare_data_with_credentials( data=_create_file_request, - credentials={**credentials, **deployment_drop_params(llm_router, model, _create_file_request)}, + credentials={ + **credentials, + **deployment_drop_params(llm_router, deployment_id or model, _create_file_request), + }, ) # Create the file with model credentials @@ -477,7 +483,7 @@ async def route_create_file( # Encode the file ID with model information if response and hasattr(response, "id") and response.id: original_id: Final = response.id - encoded_id: Final = encode_file_id_with_model(file_id=original_id, model=model) + encoded_id: Final = encode_file_id_with_model(file_id=original_id, model=model, deployment_id=deployment_id) response.id = encoded_id verbose_proxy_logger.debug("Encoded file ID: %s -> %s (model: %s)", original_id, encoded_id, model) diff --git a/litellm/router.py b/litellm/router.py index 3809a5d6555..8eb31775efb 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -724,6 +724,14 @@ def _configured_model_info(deployment: Mapping[str, object]) -> Mapping[str, obj return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {}) +def _unpaused_deployment_id(deployment: Mapping[str, object]) -> str | None: + model_info: Final = _configured_model_info(deployment) + deployment_id: Final = model_info.get("id") + if model_info.get("blocked") is True or not isinstance(deployment_id, str): + return None + return deployment_id + + def _search_tool_support(deployment: DeploymentTypedDict) -> bool | None: from litellm.responses.litellm_completion_transformation.tool_search import chat_config_builds_anthropic_request @@ -12021,6 +12029,26 @@ class Router: return None return Deployment(**first_usable) if isinstance(first_usable, dict) else first_usable + def get_batch_deployment_id(self, model_group_name: str, preferred_deployment_id: str | None = None) -> str | None: + """ + The shared deployment of ``model_group_name`` that file and batch calls + should reach when the group mixes batch and non-batch deployments, or + None when the batch rule drops nothing and the group name resolves as before. + """ + group: Final = tuple( + deployment + for deployment in self.deployments_named(model_group_name) or () + if self.deployment_usable_by_team(deployment, None) + ) + with batch_create_routing(): + batch_ready: Final = without_deployments_that_cannot_batch(group, caller_deployments=lambda: group) + if len(batch_ready) == len(group): + return None + ready_ids: Final = tuple( + deployment_id for deployment_id in map(_unpaused_deployment_id, batch_ready) if deployment_id is not None + ) + return preferred_deployment_id if preferred_deployment_id in ready_ids else next(iter(ready_ids), None) + async def arefresh_model_info(self, *, client: AsyncHTTPHandler | None = None) -> None: """Refresh token limits advertised by configured OpenAI-compatible deployments.""" deployments: Final = iter(tuple(self.model_list)) diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index c7781819ae4..c389f0ed27d 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -95,6 +95,7 @@ ignored_function_names = [ "arm_routing_read_prefetch", # Tested in tests/unit/caching/test_request_redis_batch_pre_call.py (file lacks "router" in name) "_deployment_params_base_model", # Tested through log_retry in test_router_helper_utils.py "_configured_model_info", # Tested through get_configured_service_tiers in test_router.py + "_unpaused_deployment_id", # Tested through the paused batch deployment test of get_batch_deployment_id in test_router.py "_every_deployment_refused", # Tested through the unsupported transcription retry tests in test_router.py "_routable_deployments", # Tested through get_configured_service_tiers and get_routable_upstream_model in test_router.py "_async_get_available_deployment", # Body of the `route {model}` phase wrapper, exercised through async_get_available_deployment in test_router.py diff --git a/tests/unit/proxy/batches_endpoints/test_endpoints.py b/tests/unit/proxy/batches_endpoints/test_endpoints.py index 37d64a1fc79..00af34d9d53 100644 --- a/tests/unit/proxy/batches_endpoints/test_endpoints.py +++ b/tests/unit/proxy/batches_endpoints/test_endpoints.py @@ -32,6 +32,7 @@ cannot drift without a test failure. import base64 import json import logging +from collections.abc import Mapping from contextlib import ExitStack from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -50,7 +51,10 @@ import litellm.proxy.proxy_server as proxy_server from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.openai_files_endpoints.common_utils import ( + decode_deployment_from_file_id, + decode_model_from_file_id, encode_file_id_with_model, + get_original_file_id, ) from litellm.proxy.utils import ProxyLogging from litellm.router import Router @@ -198,6 +202,7 @@ def harness(monkeypatch: pytest.MonkeyPatch): router.model_list = [] router.acreate_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) + router.get_batch_deployment_id = MagicMock(return_value=None) read_body = AsyncMock(side_effect=lambda request: body_holder["body"]) pre_call = AsyncMock(side_effect=lambda **kw: (body_holder["body"], MagicMock())) @@ -1457,6 +1462,7 @@ def retrieve_harness(): router.model_list = [] router.aretrieve_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) + router.get_batch_deployment_id = MagicMock(return_value=None) router.get_credential_deployment = MagicMock(return_value=None) pre_call = AsyncMock(side_effect=lambda **kw: (data_holder["data"], MagicMock())) @@ -2040,6 +2046,7 @@ def list_harness(): router.model_list = [] router.alist_batches = AsyncMock(return_value=FakeListPage([])) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) + router.get_batch_deployment_id = MagicMock(return_value=None) read_body = AsyncMock(side_effect=lambda request: body_holder["body"]) pre_call = AsyncMock(side_effect=lambda **kw: (body_holder["body"], MagicMock())) @@ -2458,6 +2465,7 @@ def cancel_harness(): router.model_list = [] router.acancel_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) + router.get_batch_deployment_id = MagicMock(return_value=None) pre_call = AsyncMock(side_effect=lambda **kw: (data_holder["data"], MagicMock())) # add_litellm_data_to_request is a passthrough that returns the data it got. @@ -3482,3 +3490,166 @@ async def test_cancel__executed_batch_rejects_key_without_model_grant(cancel_har factory.assert_not_called() runner.cancel.assert_not_called() cancel_harness.router_acancel.assert_not_called() + + +def _mixed_azure_group_router() -> Router: + return Router( + model_list=[ + { + "model_name": model_group, + "litellm_params": { + "model": f"azure/{sku}", + "api_base": f"https://{sku}.openai.azure.com", + "api_key": f"{sku}-key", + "api_version": "2025-04-01-preview", + }, + "model_info": {"id": deployment_id, "mode": mode}, + } + for model_group, sku, mode, deployment_id in ( + ("gpt-mixed", "chat-sku", "chat", "chat-id"), + ("gpt-mixed", "batch-sku", "batch", "batch-id"), + ("gpt-other", "other-sku", "batch", "other-batch-id"), + ) + ] + ) + + +def _decoded(encoded_id: str) -> tuple[str | None, str, str | None]: + return ( + decode_model_from_file_id(encoded_id), + get_original_file_id(encoded_id), + decode_deployment_from_file_id(encoded_id), + ) + + +def _reached(provider_kwargs: Mapping[str, object]) -> tuple[object, object]: + return provider_kwargs["api_base"], provider_kwargs["api_key"] + + +_BATCH_SKU: Final = ("https://batch-sku.openai.azure.com", "batch-sku-key") +_CHAT_SKU: Final = ("https://chat-sku.openai.azure.com", "chat-sku-key") +_MIXED_GROUP_HINTS: Final = pytest.mark.parametrize( + ("deployment_hint", "reached", "kept_hint"), + [ + ("batch-id", _BATCH_SKU, "batch-id"), + ("other-batch-id", _BATCH_SKU, "batch-id"), + (None, _CHAT_SKU, None), + ], + ids=["hinted", "hint-naming-another-group", "group-only-id-from-before-the-upgrade"], +) + + +@pytest.mark.asyncio +@_MIXED_GROUP_HINTS +async def test_create__mixed_group_input_file_reaches_the_deployment_its_hint_names( + harness, deployment_hint: str | None, reached: tuple[str, str], kept_hint: str | None +): + input_file_id: Final = encode_file_id_with_model("file-original123", "gpt-mixed", deployment_id=deployment_hint) + set_body(harness, {"input_file_id": input_file_id, "endpoint": "/v1/chat/completions", "completion_window": "24h"}) + provider_create: Final = AsyncMock(return_value=make_batch(output_file_id="file-out", error_file_id="file-err")) + + with ( + patch.object(proxy_server, "llm_router", _mixed_azure_group_router()), + patch.object(litellm, "acreate_batch", provider_create), + ): + resp: Final = await call_create(harness) + + assert _reached(provider_create.call_args.kwargs) == reached + assert provider_create.call_args.kwargs["input_file_id"] == "file-original123" + assert tuple(map(_decoded, (resp.id, resp.output_file_id, resp.error_file_id))) == ( + ("gpt-mixed", "batch-provider-id", kept_hint), + ("gpt-mixed", "file-out", kept_hint), + ("gpt-mixed", "file-err", kept_hint), + ) + assert resp.input_file_id == input_file_id + + +@pytest.mark.asyncio +async def test_create__model_header_on_a_mixed_group_reaches_its_batch_deployment(harness): + set_body(harness, {"input_file_id": "file-raw", "endpoint": "/v1/chat/completions", "completion_window": "24h"}) + + with patch.object(proxy_server, "llm_router", _mixed_azure_group_router()): + resp: Final = await call_create(harness, headers={"x-litellm-model": "gpt-mixed"}) + + assert _reached(harness.acreate_kwargs()) == _BATCH_SKU + assert _decoded(resp.id) == ("gpt-mixed", "batch-provider-id", "batch-id") + assert _decoded(resp.input_file_id) == ("gpt-mixed", "file-provider-input", "batch-id") + + +@pytest.mark.asyncio +async def test_list__model_header_on_a_mixed_group_reaches_its_batch_deployment(list_harness): + provider_list: Final = AsyncMock(return_value=FakeListPage([make_batch(id="batch-1")])) + + with ( + patch.object(proxy_server, "llm_router", _mixed_azure_group_router()), + patch.object(litellm, "alist_batches", provider_list), + ): + resp: Final = await call_list(list_harness, headers={"x-litellm-model": "gpt-mixed"}) + + assert _reached(provider_list.call_args.kwargs) == _BATCH_SKU + assert _decoded(resp.data[0].id) == ("gpt-mixed", "batch-1", "batch-id") + + +@pytest.mark.asyncio +@_MIXED_GROUP_HINTS +async def test_retrieve__mixed_group_batch_reaches_the_deployment_its_hint_names( + retrieve_harness, deployment_hint: str | None, reached: tuple[str, str], kept_hint: str | None +): + batch_id: Final = encode_file_id_with_model( + "batch-provider-id", "gpt-mixed", id_type="batch", deployment_id=deployment_hint + ) + provider_retrieve: Final = AsyncMock(return_value=make_batch(output_file_id="file-out")) + + with ( + patch.object(proxy_server, "llm_router", _mixed_azure_group_router()), + patch.object(litellm, "aretrieve_batch", provider_retrieve), + ): + resp: Final = await call_retrieve(retrieve_harness, batch_id) + + assert _reached(provider_retrieve.call_args.kwargs) == reached + assert provider_retrieve.call_args.kwargs["litellm_metadata"]["model_info"]["id"] == (kept_hint or "chat-id") + assert _decoded(resp.id) == ("gpt-mixed", "batch-provider-id", kept_hint) + assert _decoded(resp.output_file_id) == ("gpt-mixed", "file-out", kept_hint) + + +@pytest.mark.asyncio +@_MIXED_GROUP_HINTS +async def test_cancel__mixed_group_batch_reaches_the_deployment_its_hint_names( + cancel_harness, deployment_hint: str | None, reached: tuple[str, str], kept_hint: str | None +): + batch_id: Final = encode_file_id_with_model( + "batch-provider-id", "gpt-mixed", id_type="batch", deployment_id=deployment_hint + ) + + with patch.object(proxy_server, "llm_router", _mixed_azure_group_router()): + resp: Final = await call_cancel(cancel_harness, batch_id) + + assert _reached(cancel_harness.acancel_kwargs()) == reached + assert _decoded(resp.id) == ("gpt-mixed", "batch-provider-id", kept_hint) + + +@pytest.mark.asyncio +async def test_retrieve__a_key_granted_only_the_group_may_use_its_hinted_batch_id(retrieve_harness): + batch_id: Final = encode_file_id_with_model( + "batch-provider-id", "gpt-mixed", id_type="batch", deployment_id="batch-id" + ) + + with patch.object(proxy_server, "llm_router", _mixed_azure_group_router()): + resp: Final = await call_retrieve(retrieve_harness, batch_id, user=_key_restricted_to("gpt-mixed")) + + assert _reached(retrieve_harness.aretrieve_kwargs()) == _BATCH_SKU + assert _decoded(resp.id) == ("gpt-mixed", "batch-provider-id", "batch-id") + + +@pytest.mark.asyncio +async def test_retrieve__a_hint_naming_a_granted_deployment_never_authorizes_a_group_the_key_lacks(retrieve_harness): + batch_id: Final = encode_file_id_with_model( + "batch-provider-id", "gpt-mixed", id_type="batch", deployment_id="other-batch-id" + ) + + with patch.object(proxy_server, "llm_router", _mixed_azure_group_router()): + with pytest.raises(ProxyException) as exc_info: + await call_retrieve(retrieve_harness, batch_id, user=_key_restricted_to("gpt-other")) + + assert exc_info.value.code == "403" + retrieve_harness.litellm_aretrieve.assert_not_called() diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py index f6c6b51fc93..2ec62e36609 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py @@ -1,6 +1,7 @@ +import base64 from datetime import datetime, timezone from types import MappingProxyType -from typing import Final +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock import pytest @@ -13,7 +14,11 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.openai_files_endpoints.common_utils import ( apply_team_provider_credentials, apply_unified_file_ids, + decode_deployment_from_file_id, + decode_model_from_file_id, + encode_file_id_with_model, get_credentials_for_model, + get_original_file_id, get_team_provider_credentials, is_litellm_executed_batch, map_raw_file_ids_to_unified, @@ -851,3 +856,37 @@ def test_a_team_keeps_its_first_config_file_openai_deployment_over_older_db_ones assert credentials is not None assert credentials["api_key"] == "deployment-z-key" + + +@pytest.mark.parametrize(("raw_id", "id_type"), [("file-abc123", "file"), ("batch_abc123", "batch")]) +def test_a_deployment_hinted_id_decodes_to_its_group_original_id_and_deployment( + raw_id: str, id_type: Literal["file", "batch"] +): + encoded_id: Final = encode_file_id_with_model(raw_id, "gpt-mixed", id_type=id_type, deployment_id="batch-id") + + assert ( + decode_model_from_file_id(encoded_id), + get_original_file_id(encoded_id), + decode_deployment_from_file_id(encoded_id), + ) == ("gpt-mixed", raw_id, "batch-id") + + +def test_an_unhinted_id_keeps_the_encoding_ids_created_before_the_hint_existed(): + encoded_before_the_hint: Final = "file-" + base64.urlsafe_b64encode( + b"litellm:file-abc123;model,gpt-mixed" + ).decode().rstrip("=") + + assert encode_file_id_with_model("file-abc123", "gpt-mixed") == encoded_before_the_hint + assert decode_deployment_from_file_id(encoded_before_the_hint) is None + assert get_original_file_id(encoded_before_the_hint) == "file-abc123" + + +@pytest.mark.parametrize("raw_id", ["file-abc123", "batch_abc123", "file-é", ""]) +def test_a_provider_id_carries_no_deployment_hint(raw_id: str): + assert decode_deployment_from_file_id(raw_id) is None + assert get_original_file_id(raw_id) == raw_id + + +@pytest.mark.parametrize("unvalidated_id", [None, 123, ["file-abc123"]]) +def test_a_non_string_id_from_a_request_body_decodes_to_no_model_or_deployment(unvalidated_id: object): + assert (decode_model_from_file_id(unvalidated_id), decode_deployment_from_file_id(unvalidated_id)) == (None, None) diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py index ff63e37bcd8..94ccfa6fcf3 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py @@ -1,5 +1,6 @@ import itertools import json +from collections.abc import Mapping from typing import Final, List, Optional from urllib.parse import quote from unittest.mock import ANY, AsyncMock, MagicMock @@ -18,6 +19,12 @@ from litellm.files.types import FileContentStreamingResult from litellm.proxy._types import LiteLLM_UserTableFiltered, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.hooks import get_proxy_hook from litellm.proxy.management_endpoints.internal_user_endpoints import ui_view_users +from litellm.proxy.openai_files_endpoints.common_utils import ( + decode_deployment_from_file_id, + decode_model_from_file_id, + encode_file_id_with_model, + get_original_file_id, +) from litellm.proxy.openai_files_endpoints.file_content_streaming_handler import ( FileContentStreamingHandler, ) @@ -984,6 +991,108 @@ def test_batch_upload_keeps_the_callers_drop_params_over_the_deployment(batch_up assert provider_upload.call_args.kwargs["drop_params"] is True +_MIXED_GROUP_SKUS: Final = ("chat-sku", "batch-sku", "other-sku") + + +@pytest.fixture +def mixed_azure_group(monkeypatch): + import litellm.proxy.proxy_server as ps + + router: Final = Router( + model_list=[ + { + "model_name": model_group, + "litellm_params": { + "model": f"azure/{sku}", + "api_base": f"https://{sku}.openai.azure.com", + "api_key": f"{sku}-key", + "api_version": "2025-04-01-preview", + }, + "model_info": {"id": deployment_id, "mode": mode}, + } + for model_group, sku, mode, deployment_id in ( + ("gpt-mixed", "chat-sku", "chat", "chat-id"), + ("gpt-mixed", "batch-sku", "batch", "batch-id"), + ("gpt-other", "other-sku", "batch", "other-batch-id"), + ) + ] + ) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + setup_proxy_logging_object(monkeypatch, router) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user" + ) + try: + with respx.mock(assert_all_called=False) as azure: + yield azure + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +def _azure_upload_routes(azure: respx.MockRouter) -> Mapping[str, respx.Route]: + uploaded: Final = { + "id": "file-up", + "object": "file", + "bytes": len(BATCH_JSONL_LINE), + "created_at": 0, + "filename": "batch.jsonl", + "purpose": "batch", + "status": "processed", + } + return { + sku: azure.post(host=f"{sku}.openai.azure.com").mock(return_value=httpx.Response(200, json=uploaded)) + for sku in _MIXED_GROUP_SKUS + } + + +@pytest.mark.parametrize( + ("purpose", "reached_sku", "deployment_hint"), + [("batch", "batch-sku", "batch-id"), ("assistants", "chat-sku", None)], + ids=["batch-file", "other-purpose"], +) +def test_upload_naming_a_mixed_azure_group_reaches_its_batch_deployment_only_for_batch_files( + mixed_azure_group: respx.MockRouter, purpose: str, reached_sku: str, deployment_hint: str | None +): + upload_routes: Final = _azure_upload_routes(mixed_azure_group) + + response: Final = _upload_batch_file({"x-litellm-model": "gpt-mixed"}, {"purpose": purpose}) + + assert response.status_code == 200, response.text + assert {sku: route.call_count for sku, route in upload_routes.items()} == { + sku: int(sku == reached_sku) for sku in _MIXED_GROUP_SKUS + } + assert upload_routes[reached_sku].calls.last.request.headers["api-key"] == f"{reached_sku}-key" + encoded_id: Final = response.json()["id"] + assert ( + decode_model_from_file_id(encoded_id), + get_original_file_id(encoded_id), + decode_deployment_from_file_id(encoded_id), + ) == ("gpt-mixed", "file-up", deployment_hint) + + +@pytest.mark.parametrize( + ("deployment_hint", "reached_sku"), + [("batch-id", "batch-sku"), ("other-batch-id", "batch-sku"), (None, "chat-sku")], + ids=["hinted", "hint-naming-another-group", "group-only-id-from-before-the-upgrade"], +) +def test_file_content_of_a_mixed_azure_group_id_comes_from_the_deployment_its_hint_names( + mixed_azure_group: respx.MockRouter, deployment_hint: str | None, reached_sku: str +): + for sku in _MIXED_GROUP_SKUS: + mixed_azure_group.get(host=f"{sku}.openai.azure.com", path__regex=r"/files/file-out/content$").mock( + return_value=httpx.Response(200, content=f"rows from {sku}".encode()) + ) + output_file_id: Final = encode_file_id_with_model("file-out", "gpt-mixed", deployment_id=deployment_hint) + + response: Final = client.get(f"/v1/files/{output_file_id}/content", headers={"Authorization": "Bearer test-key"}) + + assert response.status_code == 200, response.text + assert response.content == f"rows from {reached_sku}".encode() + + @pytest.mark.skip(reason="mock respx fails on ci/cd - unclear why") def test_create_file_and_call_chat_completion_e2e( mocker: MockerFixture, monkeypatch, llm_router: Router diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index b55b3c5efc2..ccd421ce413 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -14168,6 +14168,71 @@ def test_credentials_for_a_mixed_group_name_keep_resolving_to_its_first_deployme assert credentials is not None and credentials["api_base"] == "https://chat-sku.openai.azure.com" +def test_batch_deployment_of_a_mixed_group_resolves_to_its_batch_deployment_credentials(): + router: Final = Router( + model_list=[ + _azure_sku_deployment("gpt-mixed", "chat-sku", "chat", id="chat-id"), + _azure_sku_deployment("gpt-mixed", "batch-sku", "batch", id="batch-id"), + ] + ) + + deployment_id: Final = router.get_batch_deployment_id("gpt-mixed") + credentials: Final = router.get_deployment_credentials_with_provider(model_id=deployment_id or "gpt-mixed") + + assert deployment_id == "batch-id" + assert credentials is not None and credentials["api_base"] == "https://batch-sku.openai.azure.com" + + +def test_batch_deployment_is_none_for_a_group_without_a_batch_deployment(): + router: Final = Router( + model_list=[ + _azure_sku_deployment("gpt-mixed", "chat-sku", "chat", id="chat-id"), + _azure_sku_deployment("gpt-mixed", "other-chat-sku", "chat", id="other-chat-id"), + ] + ) + + assert router.get_batch_deployment_id("gpt-mixed") is None + + +def test_batch_deployment_is_none_when_the_batch_rule_drops_no_deployment_of_the_group(): + router: Final = Router( + model_list=[ + {"model_name": "gpt-mixed", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "dummy-key"}}, + _azure_sku_deployment("gpt-mixed", "batch-sku", "batch", id="batch-id"), + ] + ) + + assert router.get_batch_deployment_id("gpt-mixed") is None + + +@pytest.mark.parametrize( + ("preferred_deployment_id", "expected"), + [ + ("second-batch-id", "second-batch-id"), + (None, "first-batch-id"), + ("removed-id", "first-batch-id"), + ("other-group-batch-id", "first-batch-id"), + ("chat-id", "first-batch-id"), + ("paused-batch-id", "first-batch-id"), + ], + ids=["still-batch-ready", "no-hint", "removed", "other-group", "chat-deployment", "paused"], +) +def test_batch_deployment_keeps_a_hinted_deployment_only_while_it_is_batch_ready_in_the_group( + preferred_deployment_id: str | None, expected: str +): + router: Final = Router( + model_list=[ + _azure_sku_deployment("gpt-mixed", "chat-sku", "chat", id="chat-id"), + _azure_sku_deployment("gpt-mixed", "paused-sku", "batch", id="paused-batch-id", blocked=True), + _azure_sku_deployment("gpt-mixed", "first-batch-sku", "batch", id="first-batch-id"), + _azure_sku_deployment("gpt-mixed", "second-batch-sku", "batch", id="second-batch-id"), + _azure_sku_deployment("gpt-other", "other-batch-sku", "batch", id="other-group-batch-id"), + ] + ) + + assert router.get_batch_deployment_id("gpt-mixed", preferred_deployment_id=preferred_deployment_id) == expected + + _AZURE_CHAT_SKU_REJECTS_BATCH: Final = httpx.Response( 400, json={"error": {"code": "invalid_deployment_type", "message": "The current deployment is not supported for batch jobs"}},