diff --git a/litellm/proxy/video_endpoints/endpoints.py b/litellm/proxy/video_endpoints/endpoints.py index 5809967cff5..3bff5f8a157 100644 --- a/litellm/proxy/video_endpoints/endpoints.py +++ b/litellm/proxy/video_endpoints/endpoints.py @@ -20,7 +20,9 @@ from litellm.proxy.video_endpoints.utils import ( encode_character_id_in_response, extract_model_from_target_model_names, get_custom_provider_from_data, + reencode_video_id_with_model_id, ) +from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.types.videos.utils import ( decode_character_id_with_provider, decode_video_id_with_provider, @@ -29,6 +31,17 @@ from litellm.types.videos.utils import ( router = APIRouter() +def _pin_video_id_to_deployment(response: Any, data: dict[str, Any]) -> Any: + hidden_params = get_hidden_params_dict(response) + model_id = ProxyBaseLLMRequestProcessing._get_model_id_from_response(hidden_params, data) + provider = hidden_params.get("custom_llm_provider") or data.get("custom_llm_provider") + return reencode_video_id_with_model_id( + response=response, + custom_llm_provider=provider if isinstance(provider, str) else None, + model_id=model_id, + ) + + @router.post( "/v1/videos", dependencies=[Depends(user_api_key_auth)], @@ -88,7 +101,7 @@ async def video_generation( # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) try: - return await processor.base_process_llm_request( + response = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -106,6 +119,7 @@ async def video_generation( user_api_base=user_api_base, version=version, ) + return _pin_video_id_to_deployment(response, data) except Exception as e: raise await processor._handle_llm_api_exception( e=e, @@ -479,7 +493,7 @@ async def video_remix( # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) try: - return await processor.base_process_llm_request( + response = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -497,6 +511,7 @@ async def video_remix( user_api_base=user_api_base, version=version, ) + return _pin_video_id_to_deployment(response, data) except Exception as e: raise await processor._handle_llm_api_exception( e=e, @@ -788,7 +803,7 @@ async def video_edit( processor = ProxyBaseLLMRequestProcessing(data=data) try: - return await processor.base_process_llm_request( + response = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -806,6 +821,7 @@ async def video_edit( user_api_base=user_api_base, version=version, ) + return _pin_video_id_to_deployment(response, data) except Exception as e: raise await processor._handle_llm_api_exception( e=e, @@ -888,7 +904,7 @@ async def video_extension( processor = ProxyBaseLLMRequestProcessing(data=data) try: - return await processor.base_process_llm_request( + response = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -906,6 +922,7 @@ async def video_extension( user_api_base=user_api_base, version=version, ) + return _pin_video_id_to_deployment(response, data) except Exception as e: raise await processor._handle_llm_api_exception( e=e, diff --git a/litellm/proxy/video_endpoints/utils.py b/litellm/proxy/video_endpoints/utils.py index 412a0e87d88..2deee755b1a 100644 --- a/litellm/proxy/video_endpoints/utils.py +++ b/litellm/proxy/video_endpoints/utils.py @@ -2,7 +2,11 @@ from typing import Any, Dict, Optional import orjson -from litellm.types.videos.utils import encode_character_id_with_provider +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) -> Optional[str]: @@ -52,3 +56,48 @@ def encode_character_id_in_response(response: Any, custom_llm_provider: str, mod model_id=model_id, ) return response + + +def reencode_video_id_with_model_id(response: Any, custom_llm_provider: str | None, model_id: str | None) -> Any: + """ + Re-encode a returned video id so it carries the router-selected deployment id. + + The provider transformation layer encodes the client-facing model/group name + into the video id because it runs before the router attaches the deployment id + to ``_hidden_params``. That is enough for single-deployment groups, but for a + group backed by several deployments the status/content round-trip decodes the + group name and re-routes through load balancing instead of pinning to the + deployment that created the job. Preferring ``model_id`` (the deployment id) + here keeps the follow-up calls on the deployment that owns the video + """ + if not model_id: + return response + + if isinstance(response, dict): + current_id = response.get("id") + else: + current_id = getattr(response, "id", None) + + if not isinstance(current_id, str) or not current_id: + return response + + decoded = decode_video_id_with_provider(current_id) + if decoded.get("model_id") == model_id: + return response + + provider = decoded.get("custom_llm_provider") or custom_llm_provider + if not provider: + return response + + raw_video_id = decoded.get("video_id") or current_id + new_id = encode_video_id_with_provider( + video_id=raw_video_id, + provider=provider, + model_id=model_id, + ) + + if isinstance(response, dict): + response["id"] = new_id + else: + response.id = new_id + return response diff --git a/tests/test_litellm/proxy/video_endpoints/test_endpoints.py b/tests/test_litellm/proxy/video_endpoints/test_endpoints.py index 40a26fad3c3..8c81d5d5ab0 100644 --- a/tests/test_litellm/proxy/video_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/video_endpoints/test_endpoints.py @@ -654,3 +654,121 @@ async def test_extension__extracts_nested_video_id_full_contract(harness): "custom_llm_provider": "azure", "model": "azure-sora", } + + +# =========================================================================== # +# Deployment pinning: the returned video id must carry the router-selected # +# deployment id, not the client-facing model group name, so status/content # +# round-trips pin to the deployment that owns the job (issue #33740). # +# =========================================================================== # + +# What a provider transformation produces before the router attaches the +# deployment id: the client-facing group name ("sora-2") sits in the model_id +# slot of the encoded video id. +GROUP_ENCODED_ID = encode_video_id_with_provider("video_raw", "azure", "sora-2") +# What the round-trip needs instead: the same raw id pinned to the deployment. +DEPLOYMENT_ENCODED_ID = encode_video_id_with_provider( + "video_raw", "azure", VIDEO_MODEL_ID +) + + +@pytest.mark.asyncio +async def test_generation__repins_video_id_to_deployment_from_hidden_params(harness): + harness.base_process.return_value = { + "id": GROUP_ENCODED_ID, + "_hidden_params": {"model_id": VIDEO_MODEL_ID, "custom_llm_provider": "azure"}, + } + + resp = await call_generation(harness, body={"model": "sora-2", "prompt": "x"}) + + # the group name in the id is replaced by the deployment id the router picked. + assert resp["id"] == DEPLOYMENT_ENCODED_ID + + +@pytest.mark.asyncio +async def test_generation__repins_from_litellm_metadata_model_info_id(harness): + """The reporter's core ask: fall back to litellm_metadata.model_info.id (via + _get_model_id_from_response), never to the client-facing data['model']. Here + _hidden_params carries no model_id, so a data['model'] fallback would re-stamp + the group name; only the metadata lookup yields the deployment id.""" + harness.base_process.return_value = { + "id": GROUP_ENCODED_ID, + "_hidden_params": {}, + } + + resp = await call_generation( + harness, + body={ + "model": "sora-2", + "prompt": "x", + "litellm_metadata": {"model_info": {"id": VIDEO_MODEL_ID}}, + }, + ) + + assert resp["id"] == DEPLOYMENT_ENCODED_ID + + +@pytest.mark.asyncio +async def test_generation__no_deployment_id_leaves_video_id_untouched(harness): + harness.base_process.return_value = { + "id": GROUP_ENCODED_ID, + "_hidden_params": {}, + } + + resp = await call_generation(harness, body={"model": "sora-2", "prompt": "x"}) + + # nothing to pin to -> the id is left exactly as the transformation made it. + assert resp["id"] == GROUP_ENCODED_ID + + +@pytest.mark.asyncio +async def test_generation__repins_video_object_response(harness): + from litellm.types.videos.main import VideoObject + + video = VideoObject(id=GROUP_ENCODED_ID, object="video", status="processing") + video._hidden_params = {"model_id": VIDEO_MODEL_ID, "custom_llm_provider": "azure"} + harness.base_process.return_value = video + + resp = await call_generation(harness, body={"model": "sora-2", "prompt": "x"}) + + assert resp.id == DEPLOYMENT_ENCODED_ID + + +@pytest.mark.asyncio +async def test_remix__repins_video_id_to_deployment(harness): + harness.base_process.return_value = { + "id": GROUP_ENCODED_ID, + "_hidden_params": {"model_id": VIDEO_MODEL_ID, "custom_llm_provider": "azure"}, + } + + resp = await call_remix(harness, AZURE_VIDEO_ID, body={"prompt": "new colors"}) + + assert resp["id"] == DEPLOYMENT_ENCODED_ID + + +@pytest.mark.asyncio +async def test_edit__repins_video_id_to_deployment(harness): + harness.base_process.return_value = { + "id": GROUP_ENCODED_ID, + "_hidden_params": {"model_id": VIDEO_MODEL_ID, "custom_llm_provider": "azure"}, + } + + resp = await call_edit( + harness, body={"prompt": "brighter", "video": {"id": AZURE_VIDEO_ID}} + ) + + assert resp["id"] == DEPLOYMENT_ENCODED_ID + + +@pytest.mark.asyncio +async def test_extension__repins_video_id_to_deployment(harness): + harness.base_process.return_value = { + "id": GROUP_ENCODED_ID, + "_hidden_params": {"model_id": VIDEO_MODEL_ID, "custom_llm_provider": "azure"}, + } + + resp = await call_extension( + harness, body={"prompt": "continue", "video": {"id": AZURE_VIDEO_ID}} + ) + + assert resp["id"] == DEPLOYMENT_ENCODED_ID