diff --git a/litellm/constants.py b/litellm/constants.py index 938a3c85c29..78ec67a70b3 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1563,6 +1563,7 @@ LITELLM_METADATA_FIELD: Final = "litellm_metadata" OLD_LITELLM_METADATA_FIELD: Final = "metadata" RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name" SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl" +PINNED_DEPLOYMENT_ID_METADATA_KEY: Final = "pinned_deployment_id" OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_tokens", "max_output_tokens"}) CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling" CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags" diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 687df9b0348..1ec98bc6e60 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -41,6 +41,7 @@ from litellm.constants import ( MAX_LITELLM_CALL_ID_LENGTH, MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG, NON_INFERENCE_CALL_TYPES, + PINNED_DEPLOYMENT_ID_METADATA_KEY, RETURN_RAW_MODEL_NAME_METADATA_KEY, STREAM_SSE_DATA_PREFIX, STREAM_SSE_KEEPALIVE_PING_BYTES, @@ -241,6 +242,7 @@ from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( AnthropicStreamModelRestamper, ) from litellm.proxy.litellm_pre_call_utils import ( + _get_metadata_variable_name, add_litellm_data_to_request, refresh_proxy_server_request_body_snapshot, reject_url_valued_destination, @@ -1701,8 +1703,9 @@ def _timing_values( class ProxyBaseLLMRequestProcessing: - def __init__(self, data: dict): + def __init__(self, data: dict, pinned_deployment_id: str | None = None): self.data = data + self.pinned_deployment_id: Final = pinned_deployment_id self._tags_before_guardrails: frozenset[str] | None = None @property @@ -2052,6 +2055,12 @@ class ProxyBaseLLMRequestProcessing: version=version, proxy_config=proxy_config, ) + if self.pinned_deployment_id is not None: + # Set after add_litellm_data_to_request, which strips any pin the client sent + self.data[_get_metadata_variable_name(request)][PINNED_DEPLOYMENT_ID_METADATA_KEY] = ( + self.pinned_deployment_id + ) + self.data["disable_fallbacks"] = True if not general_settings.get("expose_fallback_errors_to_caller"): self.data.pop("include_fallback_errors", None) if route_type in {"aresponses", "_aresponses_websocket"}: @@ -2074,8 +2083,6 @@ class ProxyBaseLLMRequestProcessing: # Store queue time in metadata after add_litellm_data_to_request to ensure it's preserved if queue_time_seconds is not None: - from litellm.proxy.litellm_pre_call_utils import _get_metadata_variable_name - _metadata_variable_name: Final = _get_metadata_variable_name(request) if _metadata_variable_name not in self.data: self.data[_metadata_variable_name] = {} diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 152438d0573..27553bfaa18 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -24,6 +24,7 @@ from litellm.constants import ( INTERNAL_CALL_ORIGIN_METADATA_KEY, LITELLM_PROXY_MASTER_KEY_ALIAS, OTEL_SERVICE_NAME_METADATA_KEYS, + PINNED_DEPLOYMENT_ID_METADATA_KEY, PRE_CALL_EXECUTED_GUARDRAILS_KEY, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, ROUTING_REQUEST_TAGS_METADATA_KEY, @@ -381,6 +382,7 @@ _ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset( "request_retry_count", CLIENT_OUTPUT_CEILING_METADATA_KEY, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, + PINNED_DEPLOYMENT_ID_METADATA_KEY, } ) _ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override" diff --git a/litellm/proxy/video_endpoints/endpoints.py b/litellm/proxy/video_endpoints/endpoints.py index fe966c2e31a..4e9809c327d 100644 --- a/litellm/proxy/video_endpoints/endpoints.py +++ b/litellm/proxy/video_endpoints/endpoints.py @@ -1,6 +1,6 @@ #### Video Endpoints ##### -from typing import Final +from typing import Any, Final, cast # noqa: TID251 # request-body boundary needs Any and a checked-by-source cast from fastapi import APIRouter, Depends, File, Form, Request, Response, UploadFile from fastapi.responses import ORJSONResponse @@ -17,9 +17,13 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( ) from litellm.proxy.image_endpoints.endpoints import batch_to_bytesio from litellm.proxy.video_endpoints.utils import ( + deployment_id_for_encoding, encode_character_id_in_response, + encode_video_id_in_response, extract_model_from_target_model_names, get_custom_provider_from_data, + route_to_encoded_deployment, + video_id_for_provider, video_reference_to_id, ) from litellm.types.videos.utils import ( @@ -80,7 +84,10 @@ async def video_generation( ) # Read request body - data: Final = await _read_request_body(request=request) + data: Final = cast( # cast-ok: _read_request_body is legacy-typed as bare dict + dict[str, Any], + await _read_request_body(request=request), + ) if input_reference is not None: input_reference_file: Final = await batch_to_bytesio([input_reference]) if input_reference_file: @@ -107,6 +114,7 @@ async def video_generation( user_api_base=user_api_base, version=version, ) + encoded_response: Final = encode_video_id_in_response(generated, deployment_id_for_encoding(generated, data)) except Exception as e: raise await processor._handle_llm_api_exception( e=e, @@ -115,7 +123,7 @@ async def video_generation( version=version, ) else: - return generated + return encoded_response @router.get( @@ -268,13 +276,16 @@ async def video_status( # Resolve model_name from model_id if available # This allows the router to automatically inject litellm_params from the model config + pinned_deployment_id: Final = ( + route_to_encoded_deployment(llm_router, model_id_from_decoded, data) + if model_id_from_decoded and llm_router + else None + ) if model_id_from_decoded and llm_router: - resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded) - if resolved_model: - data["model"] = resolved_model + data["video_id"] = video_id_for_provider(llm_router, video_id) # Process request using ProxyBaseLLMRequestProcessing - processor: Final = ProxyBaseLLMRequestProcessing(data=data) + processor: Final = ProxyBaseLLMRequestProcessing(data=data, pinned_deployment_id=pinned_deployment_id) try: status: Final[object] = await processor.base_process_llm_request( request=request, @@ -294,6 +305,9 @@ async def video_status( user_api_base=user_api_base, version=version, ) + encoded_status: Final = encode_video_id_in_response( + status, deployment_id_for_encoding(status, data, pinned_deployment_id) + ) except Exception as e: raise await processor._handle_llm_api_exception( e=e, @@ -302,7 +316,7 @@ async def video_status( version=version, ) else: - return status + return encoded_status @router.get( @@ -368,12 +382,15 @@ async def video_content( # Resolve model_name from model_id if available # This allows the router to automatically inject litellm_params from the model config + pinned_deployment_id: Final = ( + route_to_encoded_deployment(llm_router, model_id_from_decoded, data) + if model_id_from_decoded and llm_router + else None + ) if model_id_from_decoded and llm_router: - resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded) - if resolved_model: - data["model"] = resolved_model + data["video_id"] = video_id_for_provider(llm_router, video_id) # Process request using ProxyBaseLLMRequestProcessing - processor: Final = ProxyBaseLLMRequestProcessing(data=data) + processor: Final = ProxyBaseLLMRequestProcessing(data=data, pinned_deployment_id=pinned_deployment_id) try: # Call the video content function directly to get raw bytes video_bytes: Final = await processor.base_process_llm_request( @@ -458,7 +475,10 @@ async def video_remix( version, ) - data: Final = await _read_request_body(request=request) + data: Final = cast( # cast-ok: _read_request_body is legacy-typed as bare dict + dict[str, Any], + await _read_request_body(request=request), + ) data["video_id"] = video_id decoded: Final = decode_video_id_with_provider(video_id) @@ -476,13 +496,16 @@ async def video_remix( # Resolve model_name from model_id if available # This allows the router to automatically inject litellm_params from the model config + pinned_deployment_id: Final = ( + route_to_encoded_deployment(llm_router, model_id_from_decoded, data) + if model_id_from_decoded and llm_router + else None + ) if model_id_from_decoded and llm_router: - resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded) - if resolved_model: - data["model"] = resolved_model + data["video_id"] = video_id_for_provider(llm_router, video_id) # Process request using ProxyBaseLLMRequestProcessing - processor: Final = ProxyBaseLLMRequestProcessing(data=data) + processor: Final = ProxyBaseLLMRequestProcessing(data=data, pinned_deployment_id=pinned_deployment_id) try: remixed: Final[object] = await processor.base_process_llm_request( request=request, @@ -502,6 +525,9 @@ async def video_remix( user_api_base=user_api_base, version=version, ) + encoded_response: Final = encode_video_id_in_response( + remixed, deployment_id_for_encoding(remixed, data, pinned_deployment_id) + ) except Exception as e: raise await processor._handle_llm_api_exception( e=e, @@ -510,7 +536,7 @@ async def video_remix( version=version, ) else: - return remixed + return encoded_response @router.post( @@ -560,7 +586,10 @@ async def video_create_character( version, ) - data: Final = await _read_request_body(request=request) + data: Final = cast( # cast-ok: _read_request_body is legacy-typed as bare dict + dict[str, Any], + await _read_request_body(request=request), + ) video_file: Final = await batch_to_bytesio([video]) if video_file: data["video"] = video_file[0] @@ -600,7 +629,7 @@ async def video_create_character( if target_model_name: hidden_params: Final = getattr(response, "_hidden_params", {}) or {} provider_for_encoding: Final = hidden_params.get("custom_llm_provider") or custom_llm_provider or "openai" - model_id_for_encoding: Final = hidden_params.get("model_id") or data.get("model") + model_id_for_encoding: Final = deployment_id_for_encoding(response, data) response = encode_character_id_in_response( response=response, custom_llm_provider=provider_for_encoding, @@ -665,7 +694,10 @@ async def video_get_character( decoded: Final = decode_character_id_with_provider(character_id) provider_from_id: Final = decoded.get("custom_llm_provider") - model_id_from_decoded: Final = decoded.get("model_id") + raw_model_id_from_decoded: Final = decoded.get("model_id") + model_id_from_decoded: Final[str | None] = ( + raw_model_id_from_decoded if isinstance(raw_model_id_from_decoded, str) else None + ) decoded_character_id: Final = decoded.get("character_id") if decoded_character_id: data["character_id"] = decoded_character_id @@ -679,12 +711,13 @@ async def video_get_character( ) data["custom_llm_provider"] = custom_llm_provider - if model_id_from_decoded and llm_router: - resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded) - if resolved_model: - data["model"] = resolved_model + pinned_deployment_id: Final = ( + route_to_encoded_deployment(llm_router, model_id_from_decoded, data) + if model_id_from_decoded and llm_router + else None + ) - processor: Final = ProxyBaseLLMRequestProcessing(data=data) + processor: Final = ProxyBaseLLMRequestProcessing(data=data, pinned_deployment_id=pinned_deployment_id) try: response: object = await processor.base_process_llm_request( request=request, @@ -767,7 +800,10 @@ async def video_edit( version, ) - data: Final = await _read_request_body(request=request) + data: Final = cast( # cast-ok: _read_request_body is legacy-typed as bare dict + dict[str, Any], + await _read_request_body(request=request), + ) uploaded_video: Final = data.pop("video", None) if isinstance(uploaded_video, StarletteUploadFile): video_files: Final = await batch_to_bytesio((uploaded_video,)) @@ -790,12 +826,15 @@ async def video_edit( ) data["custom_llm_provider"] = custom_llm_provider + pinned_deployment_id: Final = ( + route_to_encoded_deployment(llm_router, model_id_from_decoded, data) + if model_id_from_decoded and llm_router + else None + ) if model_id_from_decoded and llm_router: - resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded) - if resolved_model: - data["model"] = resolved_model + data["video_id"] = video_id_for_provider(llm_router, data["video_id"]) - processor: Final = ProxyBaseLLMRequestProcessing(data=data) + processor: Final = ProxyBaseLLMRequestProcessing(data=data, pinned_deployment_id=pinned_deployment_id) try: edited: Final[object] = await processor.base_process_llm_request( request=request, @@ -815,6 +854,9 @@ async def video_edit( user_api_base=user_api_base, version=version, ) + encoded_response: Final = encode_video_id_in_response( + edited, deployment_id_for_encoding(edited, data, pinned_deployment_id) + ) except Exception as e: raise await processor._handle_llm_api_exception( e=e, @@ -823,7 +865,7 @@ async def video_edit( version=version, ) else: - return edited + return encoded_response @router.post( @@ -871,7 +913,10 @@ async def video_extension( version, ) - data: Final = await _read_request_body(request=request) + data: Final = cast( # cast-ok: _read_request_body is legacy-typed as bare dict + dict[str, Any], + await _read_request_body(request=request), + ) data["video_id"] = video_reference_to_id(data.pop("video", None)) decoded: Final = decode_video_id_with_provider(data["video_id"]) @@ -887,12 +932,15 @@ async def video_extension( ) data["custom_llm_provider"] = custom_llm_provider + pinned_deployment_id: Final = ( + route_to_encoded_deployment(llm_router, model_id_from_decoded, data) + if model_id_from_decoded and llm_router + else None + ) if model_id_from_decoded and llm_router: - resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded) - if resolved_model: - data["model"] = resolved_model + data["video_id"] = video_id_for_provider(llm_router, data["video_id"]) - processor: Final = ProxyBaseLLMRequestProcessing(data=data) + processor: Final = ProxyBaseLLMRequestProcessing(data=data, pinned_deployment_id=pinned_deployment_id) try: extended: Final[object] = await processor.base_process_llm_request( request=request, @@ -912,6 +960,9 @@ async def video_extension( user_api_base=user_api_base, version=version, ) + encoded_response: Final = encode_video_id_in_response( + extended, deployment_id_for_encoding(extended, data, pinned_deployment_id) + ) except Exception as e: raise await processor._handle_llm_api_exception( e=e, @@ -920,4 +971,4 @@ async def video_extension( version=version, ) else: - return extended + return encoded_response diff --git a/litellm/proxy/video_endpoints/utils.py b/litellm/proxy/video_endpoints/utils.py index a38226cc253..b2feb245e53 100644 --- a/litellm/proxy/video_endpoints/utils.py +++ b/litellm/proxy/video_endpoints/utils.py @@ -1,8 +1,15 @@ +from collections.abc import Mapping from typing import Any, Final import orjson -from litellm.types.videos.utils import encode_character_id_with_provider +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.router import Router +from litellm.types.videos.utils import ( + decode_video_id_with_provider, + encode_character_id_with_provider, + encode_video_id_with_provider, +) def extract_model_from_target_model_names(target_model_names: Any) -> str | None: @@ -47,6 +54,79 @@ def get_custom_provider_from_data(data: dict[str, Any]) -> str | None: return None +def _hidden_param(response: object, key: str) -> str | None: + params: Final = getattr(response, "_hidden_params", None) + if not isinstance(params, dict): + return None + value: Final = params.get(key) + return value if isinstance(value, str) else None + + +def deployment_id_for_encoding( + response: object, data: Mapping[str, Any], pinned_deployment_id: str | None = None +) -> str | None: + """The deployment the request actually routed to, not the public model group. + + The router leaves ``model_id`` out of ``_hidden_params`` on the generic path and + records the selected deployment under ``litellm_metadata.model_info.id`` instead. + Encoding ``data["model"]`` there would embed the group, so a later status or + content call could load-balance to a different deployment and fail to resolve. + """ + litellm_metadata: Final = data.get("litellm_metadata") or {} + model_info: Final = litellm_metadata.get("model_info") or {} + return _hidden_param(response, "model_id") or model_info.get("id") or pinned_deployment_id or data.get("model") + + +def route_to_encoded_deployment(llm_router: Router, model_id: str, data: dict[str, Any]) -> str | None: + """Route by the model group and return the deployment id to pin, if the id encodes one. + + ``data["model"]`` stays the group so guardrails, limits and budgets keyed on it apply. + The router's own access-group and team filters run first, then the pin narrows the result. + """ + resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id) + if resolved_model: + data["model"] = resolved_model + if model_id in llm_router.model_names or not llm_router.has_model_id(model_id): + return None + return model_id + + +def video_id_for_provider(llm_router: Router, video_id: str) -> str: + decoded: Final = decode_video_id_with_provider(video_id) + provider: Final = decoded.get("custom_llm_provider") + deployment: Final = llm_router.get_deployment(model_id=decoded.get("model_id") or "") + if provider is None or deployment is None: + return video_id + provider_model: Final = get_llm_provider( + model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + )[0] + return encode_video_id_with_provider( + video_id=decoded.get("video_id", ""), provider=provider, model_id=provider_model + ) + + +def encode_video_id_in_response(response: object, fallback_model: str | None) -> object: + hidden_provider: Final = _hidden_param(response, "custom_llm_provider") + hidden_model_id: Final = _hidden_param(response, "model_id") + video_id: Final = response.get("id") if isinstance(response, dict) else getattr(response, "id", None) + if not isinstance(video_id, str) or not video_id: + return response + + decoded: Final = decode_video_id_with_provider(video_id) + encoded: Final = encode_video_id_with_provider( + video_id=decoded.get("video_id", ""), + provider=hidden_provider or decoded.get("custom_llm_provider") or "openai", + model_id=hidden_model_id or fallback_model or decoded.get("model_id"), + ) + if isinstance(response, dict): + response["id"] = encoded # rebind-ok: in-place id rewrite, matching encode_character_id_in_response + return response + + response.id = encoded # pyright: ignore[reportAttributeAccessIssue] # rebind-ok: in-place id rewrite on a non-dict response, matching encode_character_id_in_response + return response + + def encode_character_id_in_response(response: Any, custom_llm_provider: str, model_id: str | None) -> Any: if isinstance(response, dict) and response.get("id"): response["id"] = encode_character_id_with_provider( diff --git a/litellm/router.py b/litellm/router.py index 513a5dc7c25..374b3dab72e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -181,11 +181,13 @@ from litellm.router_utils.clientside_credential_handler import ( ) from litellm.router_utils.common_utils import ( _is_proxy_admin_request, + filter_pinned_deployment, filter_team_based_models, filter_web_search_deployments, format_fallback_outcome_message, format_no_fallback_group_message, get_request_team_id, + pinned_deployment_id, provider_for_generic_call, resolve_model_group_alias, truncate_fallback_error_detail, @@ -7983,6 +7985,9 @@ class Router: @tracer.wrap() async def async_function_with_retries(self, *args, **kwargs): + typed_request_kwargs: Final = cast( # cast-ok: Python kwargs always have str keys and arbitrary object values + Mapping[str, object], kwargs + ) verbose_router_logger.debug("Inside async function with retries.") original_function: Final = kwargs.pop("original_function") fallbacks: Final = kwargs.pop("fallbacks", self.fallbacks) @@ -8040,6 +8045,7 @@ class Router: ) = await self._async_get_healthy_deployments( model=kwargs.get("model") or "", parent_otel_span=parent_otel_span, + request_kwargs=typed_request_kwargs, ) # Check retry policy FIRST, before should_retry_this_error @@ -8132,6 +8138,7 @@ class Router: ) = await self._async_get_healthy_deployments( model=_model, parent_otel_span=parent_otel_span, + request_kwargs=typed_request_kwargs, ) else: _healthy_deployments = [] @@ -8908,7 +8915,7 @@ class Router: return healthy_deployments, _all_deployments async def _async_get_healthy_deployments( - self, model: str, parent_otel_span: Span | None + self, model: str, parent_otel_span: Span | None, request_kwargs: Mapping[str, object] | None = None ) -> tuple[list[dict], list[dict]]: """ Returns Tuple of: @@ -8925,6 +8932,9 @@ class Router: return [], _all_deployments except Exception: pass + pinned_id: Final = pinned_deployment_id(request_kwargs) + if pinned_id is not None: + _all_deployments = [d for d in _all_deployments if d["model_info"]["id"] == pinned_id] unhealthy_deployments: Final = await _async_get_cooldown_deployments( litellm_router_instance=self, parent_otel_span=parent_otel_span @@ -13405,6 +13415,11 @@ class Router: if verbose_router_logger.isEnabledFor(logging.DEBUG): verbose_router_logger.debug("healthy_deployments after web search filter: %s", healthy_deployments) + pinned_request_kwargs: Final = cast( # cast-ok: this API's legacy bare dict is a str-keyed request mapping + Mapping[str, object], request_kwargs + ) + healthy_deployments = filter_pinned_deployment(model, healthy_deployments, pinned_request_kwargs) + if isinstance(healthy_deployments, dict): if (healthy_deployments.get("model_info") or {}).get("blocked") is True: raise litellm.ServiceUnavailableError( @@ -14596,6 +14611,14 @@ class Router: specific_deployment=specific_deployment, request_kwargs=request_kwargs, ) + pinned_request_kwargs: Final = ( + cast( # cast-ok: this API's legacy bare dict is a str-keyed request mapping + Mapping[str, object], request_kwargs + ) + if request_kwargs is not None + else None + ) + healthy_deployments = filter_pinned_deployment(model, healthy_deployments, pinned_request_kwargs) strategy, strategy_selector = self._get_routing_context(model, request_kwargs) pick_attributes: Final = _deployment_pick_attributes( model, _ROUTING_KWARGS_ADAPTER.validate_python(request_kwargs) diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 7d111a80264..4141fb06e32 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -9,11 +9,11 @@ if TYPE_CHECKING: import litellm from litellm._logging import verbose_logger, verbose_router_logger -from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS +from litellm.constants import PINNED_DEPLOYMENT_ID_METADATA_KEY, ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.exceptions import BadRequestError from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure -from litellm.types.router import CredentialLiteLLMParams +from litellm.types.router import CredentialLiteLLMParams, RouterErrors from litellm.types.utils import LlmProviders @@ -207,6 +207,36 @@ def filter_team_based_models( ] +def pinned_deployment_id(request_kwargs: Mapping[str, object] | None) -> str | None: + """The deployment a follow-up call must reach, such as the one that created a video.""" + if request_kwargs is None: + return None + buckets: Final = (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")) + pinned_ids: Final = [ + bucket.get(PINNED_DEPLOYMENT_ID_METADATA_KEY) for bucket in buckets if isinstance(bucket, Mapping) + ] + return next((pinned_id for pinned_id in pinned_ids if isinstance(pinned_id, str) and pinned_id), None) + + +def filter_pinned_deployment( + model: str, + healthy_deployments: list[dict] | dict, + request_kwargs: Mapping[str, object] | None, +) -> list[dict] | dict: + """Keep only the pinned deployment, if the request has one.""" + pinned_id: Final = pinned_deployment_id(request_kwargs) + if pinned_id is None or isinstance(healthy_deployments, dict): + return healthy_deployments + pinned: Final = [d for d in healthy_deployments if (d.get("model_info") or {}).get("id") == pinned_id] + if not pinned: + raise BadRequestError( + message=f"You passed in model={model}. {RouterErrors.no_healthy_deployments.value}", + model=model, + llm_provider="", + ) + return pinned + + def _deployment_supports_web_search(deployment: dict) -> bool: """ Check if a deployment supports web search. diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 8dc6c0477b0..3ec16a9265f 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -2045,6 +2045,46 @@ class TestProxyBaseLLMRequestProcessing: logging_obj.start_time.timestamp(), abs=1e-6 ) + @pytest.mark.asyncio + async def test_pinned_deployment_is_set_after_the_client_metadata_is_stripped(self, monkeypatch): + processing_obj = ProxyBaseLLMRequestProcessing( + data={"model": "sora-2", "metadata": {"pinned_deployment_id": "deployment-a"}}, + pinned_deployment_id="deployment-b", + ) + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + mock_request.url = MagicMock() + mock_request.url.path = "/v1/videos/video_123" + + async def strip_client_metadata(*args, **kwargs): + data = kwargs["data"] + data["metadata"] = {} + data["disable_fallbacks"] = False + return data + + async def mock_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False): + return copy.deepcopy(data) + + mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=mock_pre_call_hook) + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "add_litellm_data_to_request", + strip_client_metadata, + ) + + returned_data, _ = await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=mock_proxy_logging_obj, + proxy_config=MagicMock(spec=ProxyConfig), + route_type="avideo_status", + ) + + assert returned_data["metadata"]["pinned_deployment_id"] == "deployment-b" + assert returned_data["disable_fallbacks"] is True + @pytest.mark.asyncio class TestCommonRequestProcessingHelpers: diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index c97a5d1337f..e4920ef3112 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -7858,6 +7858,7 @@ _PLANTED_STAMPS = { "request_retry_count": -100, "_client_output_ceiling": {"api_base": "https://attacker.example"}, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: 10**9, + "pinned_deployment_id": "deployment-b", "client_key": "client_value", } @@ -7891,6 +7892,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo assert "_client_output_ceiling" not in updated["metadata"] assert "request_retry_count" not in updated["metadata"] assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in updated["metadata"] + assert "pinned_deployment_id" not in updated["metadata"] assert updated["metadata"]["client_key"] == "client_value" diff --git a/tests/unit/proxy/video_endpoints/test_endpoints.py b/tests/unit/proxy/video_endpoints/test_endpoints.py index c5996f95f54..5005282813e 100644 --- a/tests/unit/proxy/video_endpoints/test_endpoints.py +++ b/tests/unit/proxy/video_endpoints/test_endpoints.py @@ -33,15 +33,21 @@ from unittest.mock import AsyncMock, MagicMock, patch import orjson import pytest +import respx +import litellm import litellm.proxy.proxy_server as proxy_server import litellm.proxy.video_endpoints.endpoints as endpoints -from litellm.proxy._types import UserAPIKeyAuth +from litellm.llms.fal_ai.videos.transformation import FalAIVideoConfig +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import ProxyLogging, _check_and_merge_model_level_guardrails from litellm.router import Router +from litellm.types.router import GenericLiteLLMParams +from litellm.types.videos.main import VideoObject from litellm.types.videos.utils import ( + decode_video_id_with_provider, encode_character_id_with_provider, encode_video_id_with_provider, ) @@ -67,7 +73,7 @@ AZURE_CHARACTER_ID = encode_character_id_with_provider( RESOLVED_MODELS: Dict[str, str] = {VIDEO_MODEL_ID: "azure-sora"} # Sentinel propagated by base_process for the passthrough endpoints. -SENTINEL = object() +SENTINEL = VideoObject(id="video_raw", object="video", status="processing") class FakeRequest: @@ -105,12 +111,16 @@ class Harness: assert self.base_process.call_count == 1 return dict(self.base_process.call_args.args[0].data) + def pinned_deployment_id(self) -> str | None: + return self.base_process.call_args.args[0].pinned_deployment_id + def route_type(self) -> str: return self.base_process.call_args.kwargs["route_type"] @pytest.fixture def harness(): + SENTINEL.id = "video_raw" logging = MagicMock(spec=ProxyLogging) router = MagicMock(spec=Router) @@ -118,6 +128,8 @@ def harness(): side_effect=lambda model_id: RESOLVED_MODELS.get(model_id) ) router.resolve_model_name_from_model_id = resolve_model + router.has_model_id.return_value = False + router.get_deployment.return_value = None read_body = AsyncMock(return_value={}) batch_to_bytesio = AsyncMock(return_value=[b"filebytes"]) @@ -229,6 +241,51 @@ async def test_generation__route_type_data_and_no_provider_default(harness): harness.batch_to_bytesio.assert_not_called() +@pytest.mark.asyncio +async def test_generation__reencodes_id_with_model_id(harness): + response = VideoObject( + id=encode_video_id_with_provider("video_raw", "openai", None), + object="video", + status="processing", + ) + response._hidden_params = { + "custom_llm_provider": "openai", + "model_id": VIDEO_MODEL_ID, + } + harness.base_process.return_value = response + + resp = await call_generation(harness, body={"model": "sora-2"}) + + decoded_video_id = decode_video_id_with_provider(resp.id) + assert decoded_video_id["model_id"] == VIDEO_MODEL_ID + assert decoded_video_id["video_id"] == "video_raw" + + +@pytest.mark.asyncio +async def test_generation__reencodes_id_with_routed_deployment_id(harness): + # data["model"] is the public group here; only litellm_metadata carries the + # deployment the router picked. + response = VideoObject( + id=encode_video_id_with_provider("video_raw", "openai", None), + object="video", + status="processing", + ) + response._hidden_params = {"custom_llm_provider": "openai"} + harness.base_process.return_value = response + + resp = await call_generation( + harness, + body={ + "model": "sora-2", + "litellm_metadata": {"model_info": {"id": VIDEO_MODEL_ID}}, + }, + ) + + decoded_video_id = decode_video_id_with_provider(resp.id) + assert decoded_video_id["model_id"] == VIDEO_MODEL_ID + assert decoded_video_id["video_id"] == "video_raw" + + @pytest.mark.asyncio async def test_generation__input_reference_attached(harness): body = {"model": "sora-2", "prompt": "a sunset"} @@ -400,6 +457,28 @@ async def test_edit__extracts_nested_video_id_full_contract(harness): } +@pytest.mark.asyncio +async def test_edit__reencodes_id_with_model_id(harness): + response = VideoObject( + id=encode_video_id_with_provider("video_raw", "openai", None), + object="video", + status="processing", + ) + response._hidden_params = { + "custom_llm_provider": "openai", + "model_id": VIDEO_MODEL_ID, + } + harness.base_process.return_value = response + + resp = await call_edit( + harness, body={"prompt": "brighter", "video": {"id": "video_plain"}} + ) + + decoded_video_id = decode_video_id_with_provider(resp.id) + assert decoded_video_id["model_id"] == VIDEO_MODEL_ID + assert decoded_video_id["video_id"] == "video_raw" + + @pytest.mark.asyncio async def test_edit__provider_from_body_data_for_plain_id(harness): """For a plain id, get_custom_provider_from_data (run for real) pulls the @@ -543,6 +622,26 @@ async def test_remix__model_encoded_id_full_contract(harness): } +@pytest.mark.asyncio +async def test_remix__reencodes_id_with_model_id(harness): + response = VideoObject( + id=encode_video_id_with_provider("video_raw", "openai", None), + object="video", + status="processing", + ) + response._hidden_params = { + "custom_llm_provider": "openai", + "model_id": VIDEO_MODEL_ID, + } + harness.base_process.return_value = response + + resp = await call_remix(harness, "video_plain", body={"prompt": "new colors"}) + + decoded_video_id = decode_video_id_with_provider(resp.id) + assert decoded_video_id["model_id"] == VIDEO_MODEL_ID + assert decoded_video_id["video_id"] == "video_raw" + + @pytest.mark.asyncio async def test_remix__provider_from_body_data_not_request_body_reader(harness): """remix resolves the provider from data.get('custom_llm_provider'), never @@ -701,3 +800,317 @@ async def test_extension__extracts_nested_video_id_full_contract(harness): "custom_llm_provider": "azure", "model": "azure-sora", } + + +@pytest.mark.asyncio +async def test_extension__reencodes_id_with_model_id(harness): + response = VideoObject( + id=encode_video_id_with_provider("video_raw", "openai", None), + object="video", + status="processing", + ) + response._hidden_params = { + "custom_llm_provider": "openai", + "model_id": VIDEO_MODEL_ID, + } + harness.base_process.return_value = response + + resp = await call_extension(harness, body={"prompt": "continue"}) + + decoded_video_id = decode_video_id_with_provider(resp.id) + assert decoded_video_id["model_id"] == VIDEO_MODEL_ID + assert decoded_video_id["video_id"] == "video_raw" + + +# =========================================================================== # +# Follow-up calls against a real Router # +# =========================================================================== # + +FAL_MODEL = "fal-ai/kling-video/v2/master/text-to-video" + + +def _video_router() -> Router: + return Router( + model_list=[ + { + "model_name": "sora-2", + "litellm_params": { + "model": "openai/sora-2", + "api_key": "sk-mock-a", + "api_base": "http://a.localhost/v1", + "weight": 1, + }, + "model_info": {"id": "deployment-a"}, + }, + { + "model_name": "sora-2", + "litellm_params": { + "model": "openai/sora-2", + "api_key": "sk-mock-b", + "api_base": "http://b.localhost/v1", + "weight": 0, + }, + "model_info": {"id": "deployment-b"}, + }, + { + "model_name": "kling", + "litellm_params": {"model": f"fal_ai/{FAL_MODEL}", "api_key": "sk-mock-fal"}, + "model_info": {"id": "deployment-fal"}, + }, + ] + ) + + +def _video_id(model_id: str) -> str: + return encode_video_id_with_provider("video_orig", "openai", model_id) + + +FOLLOW_UP_CALLS = { + "status": lambda h, model_id: call_status(h, _video_id(model_id)), + "content": lambda h, model_id: call_content(h, _video_id(model_id)), + "remix": lambda h, model_id: call_remix(h, _video_id(model_id), body={"prompt": "x"}), + "edit": lambda h, model_id: call_edit(h, body={"prompt": "x", "video": {"id": _video_id(model_id)}}), + "extension": lambda h, model_id: call_extension(h, body={"prompt": "x", "video": {"id": _video_id(model_id)}}), + "get_character": lambda h, model_id: call_get_character( + h, encode_character_id_with_provider("char_orig", "openai", model_id) + ), +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("endpoint", sorted(FOLLOW_UP_CALLS)) +@pytest.mark.parametrize( + ("model_id", "expected_model"), + [ + pytest.param("deployment-b", "sora-2", id="deployment_id_pins"), + pytest.param("sora-2", "sora-2", id="group_id_load_balances"), + pytest.param("", None, id="no_model_in_id"), + pytest.param("deployment-gone", None, id="deleted_deployment_calls_provider_directly"), + ], +) +async def test_follow_up__model_comes_from_the_id(harness, endpoint, model_id, expected_model): + harness.base_process.return_value = b"video-bytes" + + with patch.object(proxy_server, "llm_router", _video_router()): + await FOLLOW_UP_CALLS[endpoint](harness, model_id) + + assert harness.processor_data().get("model") == expected_model + assert harness.pinned_deployment_id() == ("deployment-b" if model_id == "deployment-b" else None) + + +@pytest.mark.asyncio +async def test_status__real_router_reaches_the_creating_deployment(harness, monkeypatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + llm_router = _video_router() + + with patch.object(proxy_server, "llm_router", llm_router): + await call_status(harness, _video_id("deployment-b")) + with respx.mock(assert_all_called=True) as respx_mock: + deployment_b = respx_mock.get("http://b.localhost/v1/videos/video_orig").respond( + json={"id": "video_orig", "object": "video", "status": "completed", "created_at": 0} + ) + data = harness.processor_data() + status = await llm_router.avideo_status( + **{**data, "metadata": {**data.get("metadata", {}), "pinned_deployment_id": harness.pinned_deployment_id()}} + ) + + assert deployment_b.call_count == 1 + assert status.status == "completed" + + +@pytest.mark.asyncio +async def test_status__fal_id_keeps_the_model_path_fal_reads(harness): + with patch.object(proxy_server, "llm_router", _video_router()): + await call_status(harness, encode_video_id_with_provider("req-123", "fal_ai", "deployment-fal")) + data = harness.processor_data() + url, _ = FalAIVideoConfig().transform_video_status_retrieve_request( + video_id=data["video_id"], + api_base="https://queue.fal.run", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert data["model"] == "kling" + assert url == "https://queue.fal.run/fal-ai/kling-video/requests/req-123/status" + + +@pytest.mark.asyncio +async def test_status__response_id_keeps_the_deployment_for_content(harness): + harness.base_process.return_value = VideoObject( + id=encode_video_id_with_provider("video_orig", "openai", None), object="video", status="completed" + ) + + with patch.object(proxy_server, "llm_router", _video_router()): + status = await call_status(harness, _video_id("deployment-b")) + harness.base_process.reset_mock() + harness.base_process.return_value = b"video-bytes" + await call_content(harness, status.id) + + assert decode_video_id_with_provider(status.id)["model_id"] == "deployment-b" + assert harness.processor_data().get("model") == "sora-2" + assert harness.pinned_deployment_id() == "deployment-b" + + +@pytest.mark.asyncio +async def test_status__without_router_keeps_the_provider_model_in_the_id(harness): + fal_status_id = encode_video_id_with_provider("req-123", "fal_ai", "fal-ai/kling-video") + harness.base_process.return_value = VideoObject(id=fal_status_id, object="video", status="completed") + + with patch.object(proxy_server, "llm_router", None): + status = await call_status(harness, fal_status_id) + + assert status.id == fal_status_id + + +@pytest.mark.asyncio +@pytest.mark.parametrize("endpoint", sorted(FOLLOW_UP_CALLS)) +async def test_follow_up__pinned_deployment_keeps_its_model_level_guardrails(harness, endpoint): + harness.base_process.return_value = b"video-bytes" + llm_router = Router( + model_list=[ + { + "model_name": "sora-2", + "litellm_params": {"model": "openai/sora-2", "api_key": "sk-mock-b", "guardrails": ["video-guard"]}, + "model_info": {"id": "deployment-b"}, + } + ] + ) + + with patch.object(proxy_server, "llm_router", llm_router): + await FOLLOW_UP_CALLS[endpoint](harness, "deployment-b") + merged = _check_and_merge_model_level_guardrails( + data=harness.processor_data(), llm_router=llm_router, trust_client_model_info=False + ) + + assert merged["metadata"]["guardrails"] == ["video-guard"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("metadata", [False, [], "", '{"tag": "blue"}']) +async def test_remix__pinned_id_leaves_the_client_metadata_for_the_proxy_to_validate(harness, metadata): + harness.base_process.return_value = b"video-bytes" + + with patch.object(proxy_server, "llm_router", _video_router()): + await call_remix(harness, _video_id("deployment-b"), body={"prompt": "x", "metadata": metadata}) + + assert harness.processor_data()["metadata"] == metadata + assert harness.pinned_deployment_id() == "deployment-b" + + +async def _call_router_as_proxy(llm_router: Router, harness, key: UserAPIKeyAuth): + data = harness.processor_data() + metadata = { + **data.get("metadata", {}), + "pinned_deployment_id": harness.pinned_deployment_id(), + "user_api_key_auth": key, + "user_api_key_team_id": key.team_id, + } + return await getattr(llm_router, harness.route_type())(**{**data, "metadata": metadata}) + + +def _access_group_router() -> Router: + return Router( + model_list=[ + { + "model_name": "sora-2", + "litellm_params": {"model": "openai/sora-2", "api_key": "sk-mock-a"}, + "model_info": {"id": "deployment-a", "access_groups": ["group-a"]}, + }, + { + "model_name": "sora-2", + "litellm_params": {"model": "openai/sora-2", "api_key": "sk-mock-b"}, + "model_info": {"id": "deployment-b", "access_groups": ["group-b"]}, + }, + ] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("endpoint", sorted(FOLLOW_UP_CALLS)) +async def test_follow_up__id_outside_the_key_access_group_is_rejected(harness, monkeypatch, endpoint): + key = UserAPIKeyAuth(api_key="sk-test", models=["group-a"]) + monkeypatch.setitem(globals(), "_user", lambda: key) + harness.base_process.return_value = b"video-bytes" + llm_router = _access_group_router() + + with patch.object(proxy_server, "llm_router", llm_router): + await FOLLOW_UP_CALLS[endpoint](harness, "deployment-b") + with pytest.raises(litellm.BadRequestError) as rejected: + await _call_router_as_proxy(llm_router, harness, key) + + assert rejected.value.status_code == 400 + assert "You passed in model=sora-2. There are no healthy deployments for this model" in rejected.value.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("key", "model_id"), + [ + pytest.param(UserAPIKeyAuth(api_key="sk-test", models=["group-a"]), "deployment-a", id="own_access_group"), + pytest.param(UserAPIKeyAuth(api_key="sk-test", models=["group-a", "group-b"]), "deployment-b", id="both_groups"), + pytest.param(UserAPIKeyAuth(api_key="sk-test", models=["sora-2"]), "deployment-b", id="group_name"), + pytest.param(UserAPIKeyAuth(api_key="sk-test", models=["*"]), "deployment-b", id="wildcard"), + pytest.param(UserAPIKeyAuth(api_key="sk-test", models=["all-proxy-models"]), "deployment-b", id="all_proxy"), + pytest.param( + UserAPIKeyAuth(api_key="sk-test", user_role=LitellmUserRoles.PROXY_ADMIN), "deployment-b", id="admin" + ), + ], +) +async def test_status__id_inside_the_key_access_pins_the_deployment(harness, monkeypatch, key, model_id): + monkeypatch.setitem(globals(), "_user", lambda: key) + harness.base_process.return_value = b"video-bytes" + + with patch.object(proxy_server, "llm_router", _access_group_router()): + await call_status(harness, _video_id(model_id)) + + assert harness.processor_data().get("model") == "sora-2" + assert harness.pinned_deployment_id() == model_id + + +def _team_router() -> Router: + return Router( + model_list=[ + { + "model_name": "model_name_team-b_sora", + "litellm_params": {"model": "openai/sora-2", "api_key": "sk-mock-team-b"}, + "model_info": {"id": "team-b-deployment", "team_id": "team-b", "team_public_model_name": "sora-2"}, + } + ] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("endpoint", sorted(FOLLOW_UP_CALLS)) +async def test_follow_up__id_of_another_team_deployment_is_rejected(harness, monkeypatch, endpoint): + key = UserAPIKeyAuth(api_key="sk-test", team_id="team-a") + monkeypatch.setitem(globals(), "_user", lambda: key) + harness.base_process.return_value = b"video-bytes" + llm_router = _team_router() + + with patch.object(proxy_server, "llm_router", llm_router): + await FOLLOW_UP_CALLS[endpoint](harness, "team-b-deployment") + with pytest.raises(litellm.BadRequestError) as rejected: + await _call_router_as_proxy(llm_router, harness, key) + + assert rejected.value.status_code == 400 + assert "You passed in model=model_name_team-b_sora. There are no healthy deployments" in rejected.value.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key", + [ + pytest.param(UserAPIKeyAuth(api_key="sk-test", team_id="team-b"), id="same_team"), + pytest.param(UserAPIKeyAuth(api_key="sk-test", user_role=LitellmUserRoles.PROXY_ADMIN), id="admin"), + ], +) +async def test_status__id_of_a_team_deployment_pins_it_for_that_team(harness, monkeypatch, key): + monkeypatch.setitem(globals(), "_user", lambda: key) + harness.base_process.return_value = b"video-bytes" + + with patch.object(proxy_server, "llm_router", _team_router()): + await call_status(harness, _video_id("team-b-deployment")) + + assert harness.processor_data().get("model") == "model_name_team-b_sora" + assert harness.pinned_deployment_id() == "team-b-deployment" diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index ac18b4889dd..88b335e27fa 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -4,12 +4,14 @@ from unittest.mock import Mock import pytest +import litellm from litellm import Router -from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS +from litellm.constants import PINNED_DEPLOYMENT_ID_METADATA_KEY, ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.proxy._types import UserAPIKeyAuth from litellm.router_utils.common_utils import ( _deployment_supports_web_search, add_model_file_id_mappings, + filter_pinned_deployment, filter_team_based_models, filter_web_search_deployments, provider_for_generic_call, @@ -774,3 +776,88 @@ class TestWarnOnProviderCredentialMismatch: def test_provider_for_generic_call(litellm_params, expected, monkeypatch): monkeypatch.setenv("AZURE_AI_API_BASE", "https://unrelated.openai.azure.com") assert provider_for_generic_call(litellm_params) == expected + + +def _pinned_router(routing_strategy: str) -> Router: + return Router( + model_list=[ + { + "model_name": "sora-2", + "litellm_params": {"model": "openai/sora-2", "api_key": f"sk-mock-{name}", "weight": weight}, + "model_info": {"id": f"deployment-{name}"}, + } + for name, weight in (("a", 1), ("b", 0)) + ], + routing_strategy=routing_strategy, + ) + + +@pytest.mark.parametrize("bucket", ["metadata", "litellm_metadata"]) +def test_filter_pinned_deployment_keeps_only_the_pinned_id(bucket): + deployments = [{"model_info": {"id": "deployment-a"}}, {"model_info": {"id": "deployment-b"}}] + + kept = filter_pinned_deployment("sora-2", deployments, {bucket: {PINNED_DEPLOYMENT_ID_METADATA_KEY: "deployment-b"}}) + + assert kept == [{"model_info": {"id": "deployment-b"}}] + + +def test_filter_pinned_deployment_without_a_pin_keeps_every_deployment(): + deployments = [{"model_info": {"id": "deployment-a"}}, {"model_info": {"id": "deployment-b"}}] + + assert filter_pinned_deployment("sora-2", deployments, {"metadata": {}}) == deployments + + +def test_filter_pinned_deployment_rejects_a_pin_outside_the_candidates(): + with pytest.raises(litellm.BadRequestError) as rejected: + filter_pinned_deployment( + "sora-2", + [{"model_info": {"id": "deployment-a"}}], + {"metadata": {PINNED_DEPLOYMENT_ID_METADATA_KEY: "deployment-b"}}, + ) + + assert rejected.value.status_code == 400 + assert "You passed in model=sora-2. There are no healthy deployments for this model" in rejected.value.message + + +@pytest.mark.asyncio +async def test_async_routing_reaches_the_pinned_deployment(): + router = _pinned_router("simple-shuffle") + request_kwargs = {"metadata": {PINNED_DEPLOYMENT_ID_METADATA_KEY: "deployment-b"}} + + picked = [ + (await router.async_get_available_deployment(model="sora-2", request_kwargs=request_kwargs))["model_info"]["id"] + for _ in range(20) + ] + + assert picked == ["deployment-b"] * 20 + + +def test_sync_routing_reaches_the_pinned_deployment(): + router = _pinned_router("usage-based-routing") + request_kwargs = {"metadata": {PINNED_DEPLOYMENT_ID_METADATA_KEY: "deployment-b"}} + + picked = [ + router.get_available_deployment(model="sora-2", request_kwargs=request_kwargs)["model_info"]["id"] + for _ in range(20) + ] + + assert picked == ["deployment-b"] * 20 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_kwargs", "expected_ids"), + [ + ({"metadata": {PINNED_DEPLOYMENT_ID_METADATA_KEY: "deployment-b"}}, ["deployment-b"]), + ({"metadata": {}}, ["deployment-a", "deployment-b"]), + ], +) +async def test_retry_candidates_of_a_pinned_request_are_only_the_pinned_deployment(request_kwargs, expected_ids): + router = _pinned_router("simple-shuffle") + + healthy, all_deployments = await router._async_get_healthy_deployments( + model="sora-2", parent_otel_span=None, request_kwargs=request_kwargs + ) + + assert [d["model_info"]["id"] for d in healthy] == expected_ids + assert [d["model_info"]["id"] for d in all_deployments] == expected_ids diff --git a/tests/unit/test_video_generation.py b/tests/unit/test_video_generation.py index a1e5a335fd5..227899ec00e 100644 --- a/tests/unit/test_video_generation.py +++ b/tests/unit/test_video_generation.py @@ -2312,6 +2312,7 @@ def test_video_get_character_accepts_encoded_character_id(video_proxy_test_clien ) mock_router = MagicMock() mock_router.resolve_model_name_from_model_id.return_value = "vertex-ai-sora-2" + mock_router.has_model_id.return_value = False with patch("litellm.proxy.proxy_server.llm_router", mock_router): with patch.object( @@ -2469,6 +2470,8 @@ def test_edit_and_extension_route_with_encoded_video_ids( mock_router = MagicMock() mock_router.resolve_model_name_from_model_id.return_value = "vertex-ai-sora-2" + mock_router.has_model_id.return_value = False + mock_router.get_deployment.return_value = None with patch("litellm.proxy.proxy_server.llm_router", mock_router): with patch.object(