From fcef9cd0a3aa69580a87293f724d8afb11de8858 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Mon, 2 Nov 2026 21:28:46 -0800 Subject: [PATCH 1/6] feat(videos): serve custom_provider_map BaseVideoConfig providers on /v1/videos (#49785) * feat(videos): serve custom_provider_map BaseVideoConfig providers on /v1/videos A provider registered in custom_provider_map with a BaseVideoConfig handler now resolves on every video call type (create, status, content, list, remix, edit, extension, characters), the same way CustomLLM handlers already resolve on /v1/images/generations. The video entry points and the HTTP handler stop casting the provider name to LlmProviders, which raised before any lookup ran, and a provider without a video config gets a 400 instead of a 500. * fix(videos): type the custom provider lookups so the type gate holds Pick the CustomLLM for a custom provider through ProviderConfigManager.get_custom_llm_handler, which narrows the now wider custom_provider_map handler union, read custom_provider_map instead of the private _custom_providers list, and annotate the untyped custom_llm_provider parameters on the video entry points * test(videos): cover a CustomLLM shadowing a built-in video provider name and type the new tests Mark custom_handler ReadOnly on CustomLLMItem, type the new tests, and assert that a CustomLLM registered under openai still leaves openai video on the built-in config --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/images/main.py | 11 +- litellm/llms/custom_httpx/llm_http_handler.py | 20 +- litellm/main.py | 12 +- litellm/types/llms/custom_llm.py | 5 +- litellm/utils.py | 36 +++- litellm/videos/main.py | 79 ++++---- .../_support/custom_video_provider.py | 181 ++++++++++++++++++ .../test_custom_video_provider_wire.py | 113 +++++++++++ tests/unit/test_utils.py | 80 +++++--- tests/unit/videos/test_main.py | 74 +++---- 10 files changed, 488 insertions(+), 123 deletions(-) create mode 100644 tests/integration/_support/custom_video_provider.py create mode 100644 tests/integration/providers/test_custom_video_provider_wire.py diff --git a/litellm/images/main.py b/litellm/images/main.py index 165543aadba..f37b252fa45 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -28,7 +28,6 @@ from litellm.litellm_core_utils.mock_functions import mock_image_generation from litellm.llms.base_llm import BaseImageEditConfig, BaseImageGenerationConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.llms.custom_llm import CustomLLM from litellm.utils import exception_type, filter_out_litellm_params, get_litellm_params #################### Initialize provider clients #################### @@ -608,10 +607,7 @@ def image_generation( ) elif custom_llm_provider in litellm._custom_providers: # Assume custom LLM provider # Get the Custom Handler - custom_handler: CustomLLM | None = None - for item in litellm.custom_provider_map: - if item["provider"] == custom_llm_provider: - custom_handler = item["custom_handler"] + custom_handler: Final = ProviderConfigManager.get_custom_llm_handler(custom_llm_provider) if custom_handler is None: raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) @@ -915,10 +911,7 @@ def image_edit( # Check for custom provider if custom_llm_provider in litellm._custom_providers: - custom_handler: CustomLLM | None = None - for item in litellm.custom_provider_map: - if item["provider"] == custom_llm_provider: - custom_handler = item["custom_handler"] + custom_handler: Final = ProviderConfigManager.get_custom_llm_handler(custom_llm_provider) if custom_handler is None: raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 879e4e02e3e..d247c5169bd 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -8105,7 +8105,7 @@ class BaseLLMHTTPHandler: """ if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -8301,7 +8301,7 @@ class BaseLLMHTTPHandler: """ if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -8482,7 +8482,7 @@ class BaseLLMHTTPHandler: """ if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -8646,7 +8646,7 @@ class BaseLLMHTTPHandler: ): if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -8787,7 +8787,7 @@ class BaseLLMHTTPHandler: ): if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -8960,7 +8960,7 @@ class BaseLLMHTTPHandler: ): if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -9147,7 +9147,7 @@ class BaseLLMHTTPHandler: ): if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -9280,7 +9280,7 @@ class BaseLLMHTTPHandler: """ if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -9362,7 +9362,7 @@ class BaseLLMHTTPHandler: """ if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -9547,7 +9547,7 @@ class BaseLLMHTTPHandler: """ if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: diff --git a/litellm/main.py b/litellm/main.py index 38dd9f90be9..b5d1398ecad 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -266,7 +266,7 @@ from .llms.codestral.completion.handler import CodestralTextCompletion from .llms.cohere.embed import handler as cohere_embed from .llms.custom_httpx.aiohttp_handler import BaseLLMAIOHTTPHandler from .llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from .llms.custom_llm import CustomLLM, custom_chat_llm_router +from .llms.custom_llm import custom_chat_llm_router from .llms.databricks.common_utils import DatabricksBase from .llms.databricks.embed.handler import DatabricksEmbeddingHandler from .llms.deprecated_providers import aleph_alpha @@ -5337,10 +5337,7 @@ def _complete_custom_providers( stream: Final = ctx.stream timeout: Final = ctx.timeout - custom_handler: CustomLLM | None = None - for item in litellm.custom_provider_map: - if item["provider"] == custom_llm_provider: - custom_handler = item["custom_handler"] + custom_handler: Final = ProviderConfigManager.get_custom_llm_handler(custom_llm_provider) if custom_handler is None: raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) @@ -7911,10 +7908,7 @@ def embedding( litellm_params={}, ) elif custom_llm_provider in litellm._custom_providers: - custom_handler: CustomLLM | None = None - for item in litellm.custom_provider_map: - if item["provider"] == custom_llm_provider: - custom_handler = item["custom_handler"] + custom_handler: Final = ProviderConfigManager.get_custom_llm_handler(custom_llm_provider) if custom_handler is None: raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) diff --git a/litellm/types/llms/custom_llm.py b/litellm/types/llms/custom_llm.py index 6b583b0a83d..bb12449617d 100644 --- a/litellm/types/llms/custom_llm.py +++ b/litellm/types/llms/custom_llm.py @@ -1,11 +1,12 @@ from typing import TYPE_CHECKING -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict if TYPE_CHECKING: + from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.custom_llm import CustomLLM class CustomLLMItem(TypedDict): provider: str - custom_handler: "CustomLLM" + custom_handler: ReadOnly["CustomLLM | BaseVideoConfig"] diff --git a/litellm/utils.py b/litellm/utils.py index ddfc38ca677..85e669843fe 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -467,6 +467,7 @@ if TYPE_CHECKING: ) from litellm.llms.cohere.common_utils import CohereModelInfo from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.llms.custom_llm import CustomLLM from litellm.types.llms.anthropic import ( ANTHROPIC_API_ONLY_HEADERS, AnthropicThinkingParam, @@ -10958,8 +10959,13 @@ class ProviderConfigManager: @staticmethod def get_provider_video_config( model: str | None, - provider: LlmProviders, + provider: LlmProviders | str, ) -> BaseVideoConfig | None: + if not isinstance(provider, LlmProviders): + custom_video_config: Final = ProviderConfigManager.get_custom_provider_video_config(provider) + if custom_video_config is not None or provider not in LlmProvidersSet: + return custom_video_config + return ProviderConfigManager.get_provider_video_config(model=model, provider=LlmProviders(provider)) if LlmProviders.OPENAI == provider: from litellm.llms.openai.videos.transformation import OpenAIVideoConfig @@ -11010,6 +11016,34 @@ class ProviderConfigManager: return VolcEngineVideoConfig() return None + @staticmethod + def get_custom_llm_handler(provider: str | None) -> CustomLLM | None: + """The `CustomLLM` registered for `provider` in `litellm.custom_provider_map`, if its handler is one.""" + from litellm.llms.custom_llm import CustomLLM + + return next( + ( + item["custom_handler"] + for item in litellm.custom_provider_map + if item["provider"] == provider and isinstance(item["custom_handler"], CustomLLM) + ), + None, + ) + + @staticmethod + def get_custom_provider_video_config(provider: str) -> BaseVideoConfig | None: + """The `BaseVideoConfig` registered for `provider` in `litellm.custom_provider_map`, if its handler is one.""" + from litellm.llms.base_llm.videos.transformation import BaseVideoConfig + + return next( + ( + item["custom_handler"] + for item in litellm.custom_provider_map + if item["provider"] == provider and isinstance(item["custom_handler"], BaseVideoConfig) + ), + None, + ) + @staticmethod def get_provider_container_config( provider: LlmProviders, diff --git a/litellm/videos/main.py b/litellm/videos/main.py index 85e651b687c..66eb10e30db 100644 --- a/litellm/videos/main.py +++ b/litellm/videos/main.py @@ -26,6 +26,17 @@ from litellm.types.videos.utils import decode_video_id_with_provider from litellm.utils import ProviderConfigManager, client from litellm.videos.utils import VideoGenerationRequestUtils + +def _unsupported_video_call( + call_name: str, custom_llm_provider: str, model: str | None = None +) -> litellm.BadRequestError: + return litellm.BadRequestError( + message=f"{call_name} is not supported for {custom_llm_provider}", + model=model or "", + llm_provider=custom_llm_provider, + ) + + #################### Initialize provider clients #################### llm_http_handler: BaseLLMHTTPHandler = BaseLLMHTTPHandler() @@ -40,7 +51,7 @@ async def avideo_generation( size: str | None = None, user: str | None = None, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -170,7 +181,7 @@ def video_generation( size: str | None = None, user: str | None = None, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -208,11 +219,11 @@ def video_generation( # get provider config video_generation_provider_config: BaseVideoConfig | None = ProviderConfigManager.get_provider_video_config( model=model, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if video_generation_provider_config is None: - raise ValueError(f"video generation is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video generation", custom_llm_provider, model) local_vars.update(kwargs) # Get VideoGenerationOptionalRequestParams with only valid parameters @@ -329,11 +340,11 @@ def video_content( # get provider config video_provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if video_provider_config is None: - raise ValueError(f"video support download is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video support download", custom_llm_provider) local_vars.update(kwargs) # For video content download, we don't need complex optional parameter handling @@ -456,7 +467,7 @@ async def avideo_remix( video_id: str, prompt: str, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -559,7 +570,7 @@ def video_remix( video_id: str, prompt: str, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -598,11 +609,11 @@ def video_remix( # get provider config video_remix_provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if video_remix_provider_config is None: - raise ValueError(f"video remix is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video remix", custom_llm_provider) local_vars.update(kwargs) # For video remix, we need the video_id and prompt @@ -660,7 +671,7 @@ async def avideo_list( order: str | None = None, api_key: str | None = None, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -774,7 +785,7 @@ def video_list( limit: int | None = None, order: str | None = None, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -810,11 +821,11 @@ def video_list( # get provider config video_list_provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if video_list_provider_config is None: - raise ValueError(f"video list is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video list", custom_llm_provider) local_vars.update(kwargs) # For video list, we need the query parameters @@ -871,7 +882,7 @@ def video_list( async def avideo_status( video_id: str, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -970,7 +981,7 @@ def video_status( def video_status( video_id: str, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -1031,11 +1042,11 @@ def video_status( # get provider config video_status_provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if video_status_provider_config is None: - raise ValueError(f"video status is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video status", custom_llm_provider) local_vars.update(kwargs) # For video status, we need the video_id @@ -1088,7 +1099,7 @@ async def avideo_create_character( name: str, video: FileTypes, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1143,7 +1154,7 @@ def video_create_character( name: str, video: FileTypes, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1172,11 +1183,11 @@ def video_create_character( provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if provider_config is None: - raise ValueError(f"video create character is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video create character", custom_llm_provider) local_vars.update(kwargs) request_params: Final[dict] = {"name": name} @@ -1218,7 +1229,7 @@ def video_create_character( async def avideo_get_character( character_id: str, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1268,7 +1279,7 @@ async def avideo_get_character( def video_get_character( character_id: str, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1297,11 +1308,11 @@ def video_get_character( provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if provider_config is None: - raise ValueError(f"video get character is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video get character", custom_llm_provider) local_vars.update(kwargs) request_params: Final[dict] = {"character_id": character_id} @@ -1343,7 +1354,7 @@ async def avideo_edit( video_id: str, prompt: str, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1398,7 +1409,7 @@ def video_edit( video_id: str, prompt: str, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1430,11 +1441,11 @@ def video_edit( provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if provider_config is None: - raise ValueError(f"video edit is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video edit", custom_llm_provider) local_vars.update(kwargs) request_params: Final[dict] = {"video_id": video_id, "prompt": prompt} @@ -1480,7 +1491,7 @@ async def avideo_extension( prompt: str, seconds: str, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1534,7 +1545,7 @@ def video_extension( prompt: str, seconds: str, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1564,11 +1575,11 @@ def video_extension( provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if provider_config is None: - raise ValueError(f"video extension is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video extension", custom_llm_provider) local_vars.update(kwargs) request_params: Final[dict] = { diff --git a/tests/integration/_support/custom_video_provider.py b/tests/integration/_support/custom_video_provider.py new file mode 100644 index 00000000000..3513d1e3186 --- /dev/null +++ b/tests/integration/_support/custom_video_provider.py @@ -0,0 +1,181 @@ +"""A customer-owned video provider for the proxy's `custom_provider_map`: a job queue with its own wire shape.""" + +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +import httpx +from httpx._types import RequestFiles + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.videos.transformation import BaseVideoConfig +from litellm.llms.custom_llm import CustomLLM +from litellm.types.router import GenericLiteLLMParams +from litellm.types.videos.main import VideoCreateOptionalRequestParams, VideoObject +from litellm.types.videos.utils import encode_video_id_with_provider, extract_original_video_id + +_STATUS_BY_STATE: Final = MappingProxyType({"queued": "queued", "running": "in_progress", "done": "completed"}) + + +def _video_object(job: Mapping[str, object], model: str | None, custom_llm_provider: str | None) -> VideoObject: + job_id: Final = str(job["job_id"]) + seconds: Final = job.get("seconds") + video: Final = VideoObject( + id=encode_video_id_with_provider(job_id, custom_llm_provider, model) if custom_llm_provider else job_id, + object="video", + status=_STATUS_BY_STATE[str(job["state"])], + created_at=0, + model=model, + seconds=str(seconds) if seconds is not None else None, + ) + return video + + +class QueueVideoConfig(BaseVideoConfig): + def get_supported_openai_params(self, model: str) -> list: # mutable-ok: BaseVideoConfig contract + return ["seconds"] + + def map_openai_params( + self, + video_create_optional_params: VideoCreateOptionalRequestParams, + model: str, + drop_params: bool, + ) -> dict: # mutable-ok: BaseVideoConfig contract + return {key: value for key, value in video_create_optional_params.items() if key == "seconds"} + + def validate_environment( + self, + headers: dict, # mutable-ok: BaseVideoConfig contract + model: str, + api_key: str | None = None, + litellm_params: GenericLiteLLMParams | None = None, + ) -> dict: # mutable-ok: BaseVideoConfig contract + key: Final = (litellm_params.api_key if litellm_params is not None else None) or api_key or "" + return {**headers, "X-Api-Key": key} + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: dict, # mutable-ok: BaseVideoConfig contract + ) -> str: + if api_base is None: + raise ValueError("api_base is required") + return f"{api_base.rstrip('/')}/jobs" + + def transform_video_create_request( + self, + model: str, + prompt: str, + api_base: str, + video_create_optional_request_params: dict, # mutable-ok: BaseVideoConfig contract + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: BaseVideoConfig contract + ) -> tuple[dict, RequestFiles, str]: # mutable-ok: BaseVideoConfig contract + seconds: Final = video_create_optional_request_params.get("seconds") + return ({"model": model, "prompt": prompt, "seconds": seconds}, [], api_base) + + def transform_video_create_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: str | None = None, + request_data: dict | None = None, + ) -> VideoObject: + return _video_object(raw_response.json(), model, custom_llm_provider) + + def transform_video_status_retrieve_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: BaseVideoConfig contract + ) -> tuple[str, dict]: # mutable-ok: BaseVideoConfig contract + return (f"{api_base}/{extract_original_video_id(video_id)}", {}) + + def transform_video_status_retrieve_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: str | None = None, + client: object = None, + ) -> VideoObject: + return _video_object(raw_response.json(), None, custom_llm_provider) + + def transform_video_content_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: BaseVideoConfig contract + variant: str | None = None, + ) -> tuple[str, dict]: # mutable-ok: BaseVideoConfig contract + return (f"{api_base}/{extract_original_video_id(video_id)}/file", {}) + + def transform_video_content_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + client: object = None, + ) -> bytes: + return raw_response.content + + def transform_video_list_request( + self, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: BaseVideoConfig contract + after: str | None = None, + limit: int | None = None, + order: str | None = None, + extra_query: dict[str, object] | None = None, + ) -> tuple[str, dict]: # mutable-ok: BaseVideoConfig contract + return (api_base, {}) + + def transform_video_list_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: str | None = None, + ) -> dict[str, str]: # mutable-ok: BaseVideoConfig contract + return raw_response.json() + + def transform_video_remix_request( + self, + video_id: str, + prompt: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: BaseVideoConfig contract + extra_body: dict[str, object] | None = None, + ) -> tuple[str, dict]: # mutable-ok: BaseVideoConfig contract + raise NotImplementedError("the queue backend has no remix") + + def transform_video_remix_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: str | None = None, + ) -> VideoObject: + raise NotImplementedError("the queue backend has no remix") + + def transform_video_delete_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: BaseVideoConfig contract + ) -> tuple[str, dict]: # mutable-ok: BaseVideoConfig contract + raise NotImplementedError("the queue backend has no delete") + + def transform_video_delete_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + ) -> VideoObject: + raise NotImplementedError("the queue backend has no delete") + + +VIDEO_CONFIG: Final = QueueVideoConfig() +IMAGE_ONLY_HANDLER: Final = CustomLLM() diff --git a/tests/integration/providers/test_custom_video_provider_wire.py b/tests/integration/providers/test_custom_video_provider_wire.py new file mode 100644 index 00000000000..a40ecae3ed4 --- /dev/null +++ b/tests/integration/providers/test_custom_video_provider_wire.py @@ -0,0 +1,113 @@ +"""A provider registered in `custom_provider_map` with a `BaseVideoConfig` serves /v1/videos create, status and content.""" + +import json +import shutil +from contextlib import ExitStack +from pathlib import Path +from typing import Final + +from integration._support.client import Gateway +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server + +_SUPPORT_MODULE: Final = Path(__file__).resolve().parents[1] / "_support" / "custom_video_provider.py" +_PROMPT: Final = "a cat walking across a bridge" +_MP4: Final = b"\x00\x00\x00\x18ftypmp42" + b"queue-video-bytes" + + +def _queue_backend(request: Request) -> Reply: + if (request.method, request.target) == ("POST", "/jobs"): + return Reply(body=json.dumps({"job_id": "job-7", "state": "queued"}).encode()) + if (request.method, request.target) == ("GET", "/jobs/job-7"): + return Reply(body=json.dumps({"job_id": "job-7", "state": "done", "seconds": 8}).encode()) + if (request.method, request.target) == ("GET", "/jobs/job-7/file"): + return Reply(body=_MP4, content_type="video/mp4") + return Reply(status=404, body=json.dumps({"error": f"no route {request.method} {request.target}"}).encode()) + + +def _write_config(path: Path, wire: Wire) -> Path: + shutil.copy(_SUPPORT_MODULE, path.parent / "custom_video_provider.py") + path.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": "queue-video", + "litellm_params": { + "model": "queue-video/seedance-1", + "api_base": wire.url, + "api_key": "queue-secret", + }, + }, + { + "model_name": "image-only-video", + "litellm_params": { + "model": "image-only/seedance-1", + "api_base": wire.url, + "api_key": "unused", + }, + }, + ], + "litellm_settings": { + "custom_provider_map": [ + {"provider": "queue-video", "custom_handler": "custom_video_provider.VIDEO_CONFIG"}, + {"provider": "image-only", "custom_handler": "custom_video_provider.IMAGE_ONLY_HANDLER"}, + ] + }, + } + ), + encoding="utf-8", + ) + return path + + +def _lower_headers(request: Request) -> dict[str, str]: + return {name.lower(): value for name, value in request.headers.items()} + + +def test_custom_provider_map_base_video_config_serves_create_status_and_content( + gateway: Gateway, tmp_path: Path +) -> None: + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(_queue_backend)) + config: Final = _write_config(tmp_path / "custom-video.yaml", wire) + owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config)) + + created: Final = owned.gateway.request( + "POST", "/v1/videos", {"model": "queue-video", "prompt": _PROMPT, "seconds": "8"} + ) + assert created.status_code == 200, created.text + created_body: Final = created.json() + assert created_body["status"] == "queued", created_body + video_id: Final = created_body["id"] + assert isinstance(video_id, str) and video_id.startswith("video_"), created_body + + submits: Final = wire.drain() + assert [(r.method, r.target) for r in submits] == [("POST", "/jobs")], submits + assert json.loads(submits[0].body) == {"model": "seedance-1", "prompt": _PROMPT, "seconds": "8"} + assert _lower_headers(submits[0])["x-api-key"] == "queue-secret" + + status: Final = owned.gateway.request("GET", f"/v1/videos/{video_id}") + assert status.status_code == 200, status.text + assert status.json()["status"] == "completed", status.text + assert status.json()["id"] == video_id, status.text + + content: Final = owned.gateway.request("GET", f"/v1/videos/{video_id}/content") + assert content.status_code == 200, content.text + assert content.content == _MP4 + + polls: Final = wire.drain() + assert [(r.method, r.target) for r in polls] == [("GET", "/jobs/job-7"), ("GET", "/jobs/job-7/file")], polls + assert all(_lower_headers(r)["x-api-key"] == "queue-secret" for r in polls) + + +def test_custom_provider_without_a_video_config_is_refused_with_400(gateway: Gateway, tmp_path: Path) -> None: + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(_queue_backend)) + config: Final = _write_config(tmp_path / "custom-video.yaml", wire) + owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config)) + + refused: Final = owned.gateway.request("POST", "/v1/videos", {"model": "image-only-video", "prompt": _PROMPT}) + assert refused.status_code == 400, refused.text + assert "video generation is not supported for image-only" in refused.text + assert wire.drain() == () diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index ed2f023dd11..010c2ee6f92 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -967,21 +967,11 @@ _COST_MAP_ROW_SCHEMA: Final = { "cache_creation_input_token_cost": {"type": "number"}, "cache_creation_input_token_cost_above_1hr": {"type": "number"}, "cache_creation_input_token_cost_above_12hr": {"type": "number"}, - "cache_creation_input_token_cost_above_12hr_above_272k_tokens": { - "type": "number" - }, - "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": { - "type": "number" - }, - "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": { - "type": "number" - }, - "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": { - "type": "number" - }, - "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": { - "type": "number" - }, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": {"type": "number"}, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": {"type": "number"}, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": {"type": "number"}, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": {"type": "number"}, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": {"type": "number"}, "cache_creation_input_token_cost_above_12hr_flex": {"type": "number"}, "cache_creation_input_token_cost_above_12hr_priority": {"type": "number"}, "cache_creation_input_token_cost_above_12hr_batches": {"type": "number"}, @@ -1671,6 +1661,38 @@ def test_get_provider_rerank_config(): assert isinstance(config, HostedVLLMRerankConfig) +def test_get_provider_video_config_resolves_a_custom_provider_map_base_video_config( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A `BaseVideoConfig` registered in `custom_provider_map` answers for its provider name, a `CustomLLM` + registered there does not, and a built-in name still resolves to its built-in config even when a `CustomLLM` + shadows it in the map.""" + from litellm.llms.custom_llm import CustomLLM + from litellm.llms.openai.videos.transformation import OpenAIVideoConfig + + class QueueVideoConfig(OpenAIVideoConfig): + pass + + video_config: Final = QueueVideoConfig() + monkeypatch.setattr( + litellm, + "custom_provider_map", + [ + {"provider": "queue-video", "custom_handler": video_config}, + {"provider": "image-only", "custom_handler": CustomLLM()}, + {"provider": "openai", "custom_handler": CustomLLM()}, + ], + ) + monkeypatch.setattr(litellm, "_custom_providers", ["queue-video", "image-only", "openai"]) + + assert ProviderConfigManager.get_provider_video_config(model="seedance-1", provider="queue-video") is video_config + assert ProviderConfigManager.get_provider_video_config(model="seedance-1", provider="image-only") is None + assert isinstance( + ProviderConfigManager.get_provider_video_config(model="sora-2", provider="openai"), OpenAIVideoConfig + ) + assert ProviderConfigManager.get_provider_video_config(model="seedance-1", provider="not-registered") is None + + def test_get_provider_text_to_speech_config_vertex_gemini_skips_cloud_tts(): """Regression for LIT-6501: mapping vertex Gemini TTS params through Google Cloud TTS dropped response_format before the speech_to_completion bridge could honor it.""" @@ -7572,9 +7594,10 @@ async def test_oidc_google_api_key_fails_closed_when_the_metadata_server_fails_a side_effect=[metadata_failure(), httpx.Response(200, text=recovered_token)] ) chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY) - with caplog.at_level(logging.WARNING, logger="LiteLLM"), pytest.raises( - (litellm.AuthenticationError, litellm.ServiceUnavailableError) - ) as excinfo: + with ( + caplog.at_level(logging.WARNING, logger="LiteLLM"), + pytest.raises((litellm.AuthenticationError, litellm.ServiceUnavailableError)) as excinfo, + ): await _ping_vllm_through(call_path, "vllm-sa") sent_while_failing: Final = chat.call_count await _ping_vllm_through(call_path, "vllm-sa") @@ -7823,9 +7846,7 @@ async def test_router_resolves_an_oidc_api_key_only_on_a_proxy_level_deployment( try: with respx.mock(assert_all_called=True) as upstream: upstream.post(_MISTRAL_CHAT_URL).respond(200, json=_MISTRAL_REPLY) - await router.acompletion( - model="devstral", messages=_PING, metadata={"user_api_key_team_id": "team-1"} - ) + await router.acompletion(model="devstral", messages=_PING, metadata={"user_api_key_team_id": "team-1"}) sent: Final = upstream.calls.last.request.headers["Authorization"] finally: router.discard() @@ -7893,7 +7914,9 @@ async def test_an_inline_oidc_api_key_reaches_the_endpoint_its_deployment_config with respx.mock(assert_all_called=False) as upstream: upstream.route().respond(200, json=_BEDROCK_REPLY) await router.acompletion( - model="claude", messages=_PING, aws_bedrock_runtime_endpoint="https://BEDROCK-runtime.us-east-1.amazonaws.com" + model="claude", + messages=_PING, + aws_bedrock_runtime_endpoint="https://BEDROCK-runtime.us-east-1.amazonaws.com", ) sent: Final = tuple( (call.request.url.host, call.request.headers["Authorization"]) for call in upstream.calls @@ -8102,7 +8125,9 @@ async def test_a_team_row_naming_an_oidc_credential_reaches_only_the_credentials ( lambda reference: LiteLLM_Params( **litellm.Router( - model_list=[{"model_name": "openai/*", "litellm_params": {"model": "openai/*", "api_key": reference}}] + model_list=[ + {"model_name": "openai/*", "litellm_params": {"model": "openai/*", "api_key": reference}} + ] ).get_model_list(model_name="openai/*")[0]["litellm_params"] ), lambda reference: "Bearer listing-token", @@ -9108,7 +9133,9 @@ def test_basic_trimming_no_max_tokens_specified(): trimmed_messages = trim_messages(messages, model="gpt-4") print("trimmed messages for gpt-4") print(trimmed_messages) - assert (get_token_count(messages=trimmed_messages, model="gpt-4")) <= litellm.model_cost["gpt-4"]["max_input_tokens"] + assert (get_token_count(messages=trimmed_messages, model="gpt-4")) <= litellm.model_cost["gpt-4"][ + "max_input_tokens" + ] def test_multiple_messages_trimming(): @@ -11427,7 +11454,9 @@ def test_declared_mid_conversation_system_support_prefers_the_deployment_over_th ) -def test_declared_mid_conversation_system_support_is_none_when_the_model_lookup_raises(monkeypatch: pytest.MonkeyPatch) -> None: +def test_declared_mid_conversation_system_support_is_none_when_the_model_lookup_raises( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.setitem(litellm.model_cost, "hosted_vllm/unreadable-probe", "not a model info entry") litellm.get_model_info.cache_clear() assert litellm.utils.declared_mid_conversation_system_support("unreadable-probe", "hosted_vllm", None) is None @@ -11448,7 +11477,6 @@ def test_get_optional_params_forwards_a_tool_config_cache_point_to_a_downstream_ assert optional_params["extra_body"] == {"kept": True, "cache_control_injection_points": [point]} - @pytest.mark.parametrize( ("provider", "expected_warnings"), (("streamlake", 1), ("veniceai", 1), ("my-own-llm", 0)), diff --git a/tests/unit/videos/test_main.py b/tests/unit/videos/test_main.py index 22e1e5c05eb..76e9b062476 100644 --- a/tests/unit/videos/test_main.py +++ b/tests/unit/videos/test_main.py @@ -32,7 +32,7 @@ reflect production. from contextlib import ExitStack from dataclasses import dataclass -from typing import Any, Dict +from typing import Any, Dict, Final from unittest.mock import MagicMock, patch import pytest @@ -205,9 +205,7 @@ def test_video_edit__dispatch_and_provider_from_id(seams): def test_video_extension__dispatch_and_provider_from_id(seams): - result = videos_main.video_extension( - video_id=AZURE_VIDEO_ID, prompt="continue", seconds="5" - ) + result = videos_main.video_extension(video_id=AZURE_VIDEO_ID, prompt="continue", seconds="5") seams.assert_only("video_extension_handler") assert result is seams.handler.video_extension_handler.return_value @@ -285,9 +283,7 @@ def test_generation__mock_response_short_circuits(seams): def test_list__mock_response_short_circuits(seams): - resp = videos_main.video_list( - mock_response=[{"id": "v1", "object": "video", "status": "completed"}] - ) + resp = videos_main.video_list(mock_response=[{"id": "v1", "object": "video", "status": "completed"}]) assert isinstance(resp, list) assert resp[0].id == "v1" @@ -318,7 +314,7 @@ def test_get_character__mock_response_short_circuits(seams): def test_unsupported_provider_raises_without_dispatch(seams): seams.get_config.return_value = None - with pytest.raises(litellm.APIConnectionError): + with pytest.raises(litellm.BadRequestError, match="video status is not supported for azure"): videos_main.video_status(video_id=AZURE_VIDEO_ID) seams.handler.video_status_handler.assert_not_called() @@ -333,9 +329,7 @@ def test_unsupported_provider_raises_without_dispatch(seams): async def test_avideo_generation__delegates_with_async_flag(): sentinel = VideoObject(id="v-async", object="video", status="queued") with ( - patch.object( - videos_main, "video_generation", MagicMock(return_value=sentinel) - ) as sync, + patch.object(videos_main, "video_generation", MagicMock(return_value=sentinel)) as sync, patch.object( litellm, "get_llm_provider", @@ -352,9 +346,7 @@ async def test_avideo_generation__delegates_with_async_flag(): @pytest.mark.asyncio async def test_avideo_status__delegates_untouched(): sentinel = VideoObject(id="v-async", object="video", status="queued") - with patch.object( - videos_main, "video_status", MagicMock(return_value=sentinel) - ) as sync: + with patch.object(videos_main, "video_status", MagicMock(return_value=sentinel)) as sync: result = await videos_main.avideo_status(video_id="video_plain") assert result is sentinel @@ -367,9 +359,7 @@ async def test_avideo_content__pre_decodes_provider_before_delegating(): """avideo_content resolves the provider from the encoded id itself before handing off, so the sync worker receives the decoded provider, not None.""" sentinel = b"mp4-bytes" - with patch.object( - videos_main, "video_content", MagicMock(return_value=sentinel) - ) as sync: + with patch.object(videos_main, "video_content", MagicMock(return_value=sentinel)) as sync: result = await videos_main.avideo_content(video_id=AZURE_VIDEO_ID) assert result is sentinel @@ -393,9 +383,7 @@ DB_YAML_CREDS = { CREDENTIAL_OPERATIONS = [ ( "video_generation_handler", - lambda: videos_main.video_generation( - prompt="p", model="sora-2", **DB_YAML_CREDS - ), + lambda: videos_main.video_generation(prompt="p", model="sora-2", **DB_YAML_CREDS), ), ( "video_status_handler", @@ -407,21 +395,15 @@ CREDENTIAL_OPERATIONS = [ ), ( "video_remix_handler", - lambda: videos_main.video_remix( - video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS - ), + lambda: videos_main.video_remix(video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS), ), ( "video_edit_handler", - lambda: videos_main.video_edit( - video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS - ), + lambda: videos_main.video_edit(video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS), ), ( "video_extension_handler", - lambda: videos_main.video_extension( - video_id=AZURE_VIDEO_ID, prompt="p", seconds="5", **DB_YAML_CREDS - ), + lambda: videos_main.video_extension(video_id=AZURE_VIDEO_ID, prompt="p", seconds="5", **DB_YAML_CREDS), ), ( "video_list_handler", @@ -429,9 +411,7 @@ CREDENTIAL_OPERATIONS = [ ), ( "video_create_character_handler", - lambda: videos_main.video_create_character( - name="hero", video=MagicMock(name="vid"), **DB_YAML_CREDS - ), + lambda: videos_main.video_create_character(name="hero", video=MagicMock(name="vid"), **DB_YAML_CREDS), ), ( "video_get_character_handler", @@ -453,3 +433,33 @@ def test_db_yaml_credentials_reach_every_handler(seams, handler_name, invoke): assert litellm_params.get("api_base") == DB_YAML_CREDS["api_base"] assert litellm_params.get("api_version") == DB_YAML_CREDS["api_version"] assert litellm_params.get("vertex_project") == DB_YAML_CREDS["vertex_project"] + + +def test_video_generation__custom_provider_map_base_video_config_is_dispatched(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.llms.openai.videos.transformation import OpenAIVideoConfig + + class QueueVideoConfig(OpenAIVideoConfig): + pass + + video_config: Final = QueueVideoConfig() + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "queue-video", "custom_handler": video_config}]) + monkeypatch.setattr(litellm, "_custom_providers", ["queue-video"]) + handler: Final = MagicMock(spec=BaseLLMHTTPHandler) + with patch.object(videos_main, "base_llm_http_handler", handler): + videos_main.video_generation(model="queue-video/seedance-1", prompt="a cat", api_base="http://backend.test") + kw: Final = handler.video_generation_handler.call_args.kwargs + assert kw["video_generation_provider_config"] is video_config + assert kw["custom_llm_provider"] == "queue-video" + assert kw["model"] == "seedance-1" + + +def test_video_generation__custom_provider_without_a_video_config_is_a_400(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.llms.custom_llm import CustomLLM + + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "image-only", "custom_handler": CustomLLM()}]) + monkeypatch.setattr(litellm, "_custom_providers", ["image-only"]) + handler: Final = MagicMock(spec=BaseLLMHTTPHandler) + with patch.object(videos_main, "base_llm_http_handler", handler): + with pytest.raises(litellm.BadRequestError, match="video generation is not supported for image-only"): + videos_main.video_generation(model="image-only/seedance-1", prompt="a cat") + handler.video_generation_handler.assert_not_called() From 3da6029977d9e90953a250399911117a603592be Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 3 Nov 2026 06:10:00 +0000 Subject: [PATCH 2/6] test(bedrock): assert native structured output when the model rejects forced tools (#49799) Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_vertex_anthropic_structured_output_wire.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py index 4e31122440d..db81ff70263 100644 --- a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py +++ b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py @@ -1488,7 +1488,7 @@ def test_vertex_chat_json_schema_on_a_model_rejecting_forced_tools_maps_reasonin assert _reasoning_fields(structured_sent) == _reasoning_fields(plain_sent), (structured_sent.body, plain_sent.body) -def test_bedrock_invoke_json_schema_sends_the_json_tool_unforced_when_the_model_rejects_forced_tools( +def test_bedrock_invoke_json_schema_sends_native_structured_output_when_the_model_rejects_forced_tools( gateway: Gateway, ) -> None: marker, marker_hex = _marker("plain") @@ -1499,7 +1499,8 @@ def test_bedrock_invoke_json_schema_sends_the_json_tool_unforced_when_the_model_ assert _chat_content(response) == _payload(marker_hex), response.text sent: Final = _one_delivery(wire, marker_hex) assert sent.target == f"/model/{backend}/invoke", sent.target - assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body + assert object_value(_folded_format(sent)["schema"])["properties"] == _SCHEMA["properties"], sent.body + assert "tools" not in sent.body, sent.body assert "tool_choice" not in sent.body, sent.body From 408773bb69c1fc975ef9c936bfbdc8334e2bfe4e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 3 Nov 2026 06:10:08 +0000 Subject: [PATCH 3/6] test(spend): assert Anthropic envelope for Messages budget refusals (#49800) Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../spend/test_partial_stream_live_budget.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/integration/spend/test_partial_stream_live_budget.py b/tests/integration/spend/test_partial_stream_live_budget.py index 89fb66c9bbf..cba51873e2d 100644 --- a/tests/integration/spend/test_partial_stream_live_budget.py +++ b/tests/integration/spend/test_partial_stream_live_budget.py @@ -346,5 +346,15 @@ def test_a_broken_messages_stream_counts_toward_the_budget(gateway: Gateway) -> _broken_messages_stream(gateway, model, key) after: Final = _messages(gateway, model, key, streamed=False) - assert_budget_refusal(after) + assert after.status_code == 422, after.text + body: Final = JSON_OBJECT.validate_json(after.content) + error: Final = JSON_OBJECT.validate_python(body["error"]) + assert body == { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": error["message"], + "litellm_error_type": "budget_exceeded", + }, + }, after.text assert _rows(gateway, key, 2) == (metered, StatusRow(status="failure", spend=0.0)) From b777dd6f14b84090ee537fe745de778f51dffd7c Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 3 Nov 2026 06:12:42 +0000 Subject: [PATCH 4/6] fix(router)!: scope context fallback precedence to the requested group (#49779) * fix(router): scope context fallback precedence to the requested group * test(router): keep existing fallback test formatting * fix(router): preserve context errors after ordered backup failures --- litellm/router_python_core.py | 52 +++--- tests/unit/test_router/core/test_fallbacks.py | 19 +-- .../unit/test_router/test_router_fallbacks.py | 153 ++++++++++++++++++ 3 files changed, 188 insertions(+), 36 deletions(-) diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index 2235abc9897..ee67d82abfa 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -9152,11 +9152,17 @@ class PythonCore(DeprecatedRouterMembers): if include_fallback_errors: input_kwargs["include_fallback_errors"] = True - # ORDER-BASED FALLBACKS: prepend higher order levels to the fallback list - # Skip for error types that have their own dedicated fallback handlers - _skip_order_fallback: Final = isinstance( - e, - (litellm_exceptions.ContextWindowExceededError, litellm_exceptions.ContentPolicyViolationError), + context_window_fallback_model_group: Final = ( + self._get_fallback_model_group_for_lookup_groups( + fallbacks=context_window_fallbacks, + lookup_groups=lookup_groups, + ) + if isinstance(e, litellm_exceptions.ContextWindowExceededError) and context_window_fallbacks is not None + else None + ) + _skip_order_fallback: Final = ( + isinstance(e, litellm_exceptions.ContentPolicyViolationError) + or context_window_fallback_model_group is not None ) _request_team_id: Final[str | None] = (kwargs.get("metadata", {}) or {}).get("user_api_key_team_id") # Use wildcard-aware lookup so order-based fallback also works for model @@ -9202,14 +9208,23 @@ class PythonCore(DeprecatedRouterMembers): "original_model_group": original_model_group, } ) - response = await run_async_fallback( - *args, - **input_kwargs, - ) - return response + try: + return await run_async_fallback( + *args, + **input_kwargs, + ) + except Exception: + if isinstance(original_exception, litellm_exceptions.ContextWindowExceededError): + raise original_exception + raise # Weighted intra-group failover (simple-shuffle only); see _maybe_run_weighted_failover. - if self.enable_weighted_failover and not _skip_order_fallback and original_model_group is not None: + if ( + self.enable_weighted_failover + and not _skip_order_fallback + and not isinstance(e, litellm_exceptions.ContextWindowExceededError) + and original_model_group is not None + ): response = await self._maybe_run_weighted_failover( exception=e, original_model_group=original_model_group, @@ -9227,7 +9242,7 @@ class PythonCore(DeprecatedRouterMembers): # check if client-side fallbacks are used (e.g. fallbacks = ["gpt-3.5-turbo", "claude-3-haiku"] or fallbacks=[{"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hey, how's it going?"}]}] is_non_standard_fallback_format: Final = check_non_standard_fallback_format(fallbacks=fallbacks) - if is_non_standard_fallback_format: + if is_non_standard_fallback_format and context_window_fallback_model_group is None: input_kwargs.update( { "fallback_model_group": fallbacks, @@ -9243,16 +9258,7 @@ class PythonCore(DeprecatedRouterMembers): return response if isinstance(e, litellm_exceptions.ContextWindowExceededError): - if context_window_fallbacks is not None: - context_window_fallback_model_group: Final[list[str] | None] = ( - self._get_fallback_model_group_for_lookup_groups( - fallbacks=context_window_fallbacks, - lookup_groups=lookup_groups, - ) - ) - if context_window_fallback_model_group is None: - raise original_exception - + if context_window_fallback_model_group is not None: input_kwargs.update( { "fallback_model_group": context_window_fallback_model_group, @@ -9266,7 +9272,7 @@ class PythonCore(DeprecatedRouterMembers): ) return response - else: + elif context_window_fallbacks is None: error_message = f"model={model_group}. context_window_fallbacks={mask_sensitive_structure(context_window_fallbacks)}. fallbacks={mask_sensitive_structure(fallbacks)}.\n\nSet 'context_window_fallback' - https://docs.litellm.ai/docs/routing#fallbacks" verbose_router_logger.info( msg=f"Got 'ContextWindowExceededError'. No context_window_fallback set. Defaulting \ diff --git a/tests/unit/test_router/core/test_fallbacks.py b/tests/unit/test_router/core/test_fallbacks.py index e4e5c3120c8..bc6c70543c8 100644 --- a/tests/unit/test_router/core/test_fallbacks.py +++ b/tests/unit/test_router/core/test_fallbacks.py @@ -1,6 +1,7 @@ """Cross-group fallbacks of a router core: where the chain comes from, how it is walked, and what is raised.""" from collections.abc import Callable, Mapping +from typing import Final import pytest @@ -355,21 +356,13 @@ async def test_an_error_specific_chain_wins_over_the_general_chain( @pytest.mark.asyncio -@pytest.mark.parametrize( - ("error", "field", "raised"), - [ - (context_window, "context_window_fallbacks", litellm.ContextWindowExceededError), - (content_policy, "content_policy_fallbacks", litellm.ContentPolicyViolationError), - ], - ids=["context-window", "content-policy"], -) -async def test_an_error_specific_chain_without_the_group_raises_without_trying_the_general_chain( - make_core: MakeCore, host: ScriptedHost, error: Callable[[], Exception], field: str, raised: type[Exception] +async def test_an_unmatched_content_policy_chain_does_not_try_the_general_chain( + make_core: MakeCore, host: ScriptedHost ) -> None: - core = build(make_core, fallbacks=[{"a": ["b"]}], **{field: [{"other": ["c"]}]}) - host.answer("openai/a", error()) + core: Final = build(make_core, fallbacks=[{"a": ["b"]}], content_policy_fallbacks=[{"other": ["c"]}]) + host.answer("openai/a", content_policy()) - with pytest.raises(raised): + with pytest.raises(litellm.ContentPolicyViolationError): await complete(core) assert host.deployments_called("acompletion") == ["a"] diff --git a/tests/unit/test_router/test_router_fallbacks.py b/tests/unit/test_router/test_router_fallbacks.py index bccd6290f44..b25d5c35892 100644 --- a/tests/unit/test_router/test_router_fallbacks.py +++ b/tests/unit/test_router/test_router_fallbacks.py @@ -1207,3 +1207,156 @@ def test_context_stream_fallback_preserves_other_errors_for_existing_continuatio message="provider unavailable", model="primary", llm_provider="openai", generated_content="already visible" ) assert _context_window_stream_fallback_error(error, "primary", "openai", True) is error + + +def _context_fallback_precedence_router(context_key: str | None, ordered: bool) -> Router: + deployments: Final = ( + ("primary", "primary", 1), + ("primary" if ordered else "unused", "ordered", 2), + ("generic", "generic", 1), + ("context", "context", 1), + ) + return Router( + model_list=[ + { + "model_name": group, + "litellm_params": { + "model": f"openai/{model}", + "api_key": "test-key", + "api_base": "https://fallback-migration.local/v1", + "order": order, + }, + } + for group, model, order in deployments + ], + fallbacks=[{"primary": ["generic"]}], + context_window_fallbacks=None if context_key is None else [{context_key: ["context"]}], + num_retries=0, + disable_cooldowns=True, + ) + + +@pytest.mark.parametrize("sync_mode", (False, True)) +@pytest.mark.parametrize( + ("context_key", "ordered", "flat_fallbacks", "expected_model"), + ( + (None, True, False, "ordered"), + ("unrelated", True, False, "ordered"), + ("unrelated", False, False, "generic"), + ("primary", True, False, "context"), + ("*", True, False, "context"), + ("primary", True, True, "context"), + ("primary", False, True, "context"), + ), +) +@pytest.mark.asyncio +async def test_context_fallback_precedence_is_specific_to_the_requested_group( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + sync_mode: bool, + context_key: str | None, + ordered: bool, + flat_fallbacks: bool, + expected_model: str, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = _context_fallback_precedence_router(context_key, ordered) + route: Final = respx_mock.post("https://fallback-migration.local/v1/chat/completions").mock( + side_effect=( + httpx.Response( + 400, + json={ + "error": { + "message": "This model's maximum context length is 10 tokens", + "type": "invalid_request_error", + "code": "context_length_exceeded", + } + }, + ), + _fallback_chat_response(expected_model, f"served by {expected_model}"), + ) + ) + + response: Final = ( + router.completion( + model="primary", + messages=[{"role": "user", "content": "long prompt"}], + fallbacks=["generic"] if flat_fallbacks else [{"primary": ["generic"]}], + ) + if sync_mode + else await router.acompletion( + model="primary", + messages=[{"role": "user", "content": "long prompt"}], + fallbacks=["generic"] if flat_fallbacks else [{"primary": ["generic"]}], + ) + ) + + assert tuple(json.loads(call.request.content)["model"] for call in route.calls) == ("primary", expected_model) + assert response.choices[0].message.content == f"served by {expected_model}" + + +@pytest.mark.parametrize("matching_context", (False, True)) +@pytest.mark.asyncio +async def test_context_fallback_exhaustion_preserves_the_selected_chain( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + matching_context: bool, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = _context_fallback_precedence_router("primary" if matching_context else "unrelated", True) + refusal: Final = httpx.Response( + 400, + json={"error": {"message": "This model's maximum context length is 10 tokens"}}, + ) + route: Final = respx_mock.post("https://fallback-migration.local/v1/chat/completions").mock( + side_effect=(refusal, refusal, _fallback_chat_response("generic", "generic answer")) + ) + + if matching_context: + with pytest.raises(litellm.ContextWindowExceededError): + await router.acompletion(model="primary", messages=[{"role": "user", "content": "long prompt"}]) + assert tuple(json.loads(call.request.content)["model"] for call in route.calls) == ("primary", "context") + return + + response: Final = await router.acompletion(model="primary", messages=[{"role": "user", "content": "long prompt"}]) + assert tuple(json.loads(call.request.content)["model"] for call in route.calls) == ("primary", "ordered", "generic") + assert response.choices[0].message.content == "generic answer" + + +@pytest.mark.parametrize("sync_mode", (False, True)) +@pytest.mark.parametrize("context_error", (False, True)) +@pytest.mark.parametrize("backup_status", (429, 503)) +@pytest.mark.asyncio +async def test_ordered_fallback_failures_preserve_the_original_context_error( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + sync_mode: bool, + context_error: bool, + backup_status: int, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = _context_fallback_precedence_router("unrelated", True) + original_message: Final = ( + "This model's maximum context length is 10 tokens" if context_error else "primary unavailable" + ) + backup_error: Final = litellm.RateLimitError if backup_status == 429 else litellm.ServiceUnavailableError + expected_error: Final = litellm.ContextWindowExceededError if context_error else backup_error + route: Final = respx_mock.post("https://fallback-migration.local/v1/chat/completions").mock( + side_effect=( + httpx.Response(400 if context_error else 503, json={"error": {"message": original_message}}), + httpx.Response(backup_status, json={"error": {"message": "ordered deployment unavailable"}}), + httpx.Response(backup_status, json={"error": {"message": "generic deployment unavailable"}}), + ) + ) + + request: Final = ( + asyncio.to_thread(router.completion, model="primary", messages=[{"role": "user", "content": "long prompt"}]) + if sync_mode + else router.acompletion(model="primary", messages=[{"role": "user", "content": "long prompt"}]) + ) + with pytest.raises(expected_error) as raised: + await request + + assert tuple(json.loads(call.request.content)["model"] for call in route.calls) == ("primary", "ordered", "generic") + assert raised.value.status_code == (400 if context_error else backup_status) + assert (original_message if context_error else "generic deployment unavailable") in str(raised.value) From 4e007c6e8116ba59368e828abc4dc7565ce5519e Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 3 Nov 2026 06:13:56 +0000 Subject: [PATCH 5/6] docs: preserve author review when updating PR summaries (#49803) --- .github/PULL_REQUEST_TEMPLATE/general.md | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/.github/PULL_REQUEST_TEMPLATE/general.md b/.github/PULL_REQUEST_TEMPLATE/general.md index aafd968ee0f..4f1a18e925e 100644 --- a/.github/PULL_REQUEST_TEMPLATE/general.md +++ b/.github/PULL_REQUEST_TEMPLATE/general.md @@ -6,12 +6,13 @@ ## Handwritten 1-3 sentence description - From 4573fa1f312a06e3daed57fd096205e5b50386f4 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 2 Nov 2026 22:29:37 -0800 Subject: [PATCH 6/6] fix(embeddings): honor the caller's timeout on Cohere sync and HuggingFace embedding requests (#48593) * fix(embeddings): honor the caller's timeout on Cohere sync and HuggingFace embedding requests Bedrock Cohere embeddings ran the sync request with the client default instead of the timeout passed to litellm.embedding, and HuggingFace embeddings dropped it on both the sync and async paths. Pass the timeout through to the HTTP request on each path. Fixes #47301 Fixes #47299 * fix(huggingface): pass the caller's timeout to the embedding model lookup too Without input_type the handler first fetches the model's pipeline tag from huggingface.co, and that GET still used the default timeout. Thread the caller's timeout through both lookup helpers and the transform step so every request the embedding makes honors it. Refs #47299 * fix(huggingface): keep the unreachable async lookup path untouched The async model lookup helper and the call_type="async" transform branch are never reached, since the async embedding path transforms through the sync lookup. Leave them as they were so the change only touches requests that actually run. --- litellm/llms/cohere/embed/handler.py | 4 +- litellm/llms/huggingface/embedding/handler.py | 22 ++++++--- litellm/main.py | 1 + .../embed/test_cohere_embedding_handler.py | 35 ++++++++++++++ .../test_huggingface_embedding_handler.py | 46 +++++++++++++++---- 5 files changed, 92 insertions(+), 16 deletions(-) create mode 100644 tests/unit/llms/cohere/embed/test_cohere_embedding_handler.py diff --git a/litellm/llms/cohere/embed/handler.py b/litellm/llms/cohere/embed/handler.py index 2f60d94c08a..34fde444880 100644 --- a/litellm/llms/cohere/embed/handler.py +++ b/litellm/llms/cohere/embed/handler.py @@ -82,7 +82,7 @@ async def async_embedding( ) try: - response: Final = await client.post(api_base, headers=headers, data=json.dumps(data)) + response: Final = await client.post(api_base, headers=headers, data=json.dumps(data), timeout=timeout) except httpx.HTTPStatusError as e: ## LOGGING logging_obj.post_call( @@ -166,7 +166,7 @@ def embedding( if client is None or not isinstance(client, HTTPHandler): client = HTTPHandler(concurrent_limit=1) - response: Final = client.post(embed_url, headers=headers, data=json.dumps(data)) + response: Final = client.post(embed_url, headers=headers, data=json.dumps(data), timeout=timeout) return CohereEmbeddingConfig().transform_response( response=response, diff --git a/litellm/llms/huggingface/embedding/handler.py b/litellm/llms/huggingface/embedding/handler.py index a1d1a3f6a19..5d997babc6e 100644 --- a/litellm/llms/huggingface/embedding/handler.py +++ b/litellm/llms/huggingface/embedding/handler.py @@ -38,7 +38,9 @@ class _SupportsTokenEncode(Protocol): def encode(self, text: str) -> list[int]: ... -def get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: str) -> object: +def get_hf_task_embedding_for_model( + model: str, task_type: str | None, api_base: str, timeout: float | httpx.Timeout | None = None +) -> object: if task_type is not None: if task_type in get_args(hf_tasks_embeddings): return task_type @@ -46,7 +48,7 @@ def get_hf_task_embedding_for_model(model: str, task_type: str | None, api_base: raise Exception(f"Invalid task_type={task_type}. Expected one of={hf_tasks_embeddings}") http_client: Final = HTTPHandler(concurrent_limit=1) - model_info: Final = http_client.get(url=f"{api_base}/api/models/{model}") + model_info: Final = http_client.get(url=f"{api_base}/api/models/{model}", timeout=timeout) model_info_dict: Final = decoded_json_object(_DECODED_JSON.validate_python(model_info.json())) @@ -148,6 +150,7 @@ class HuggingFaceEmbedding(BaseLLM): call_type: Literal["sync", "async"], optional_params: dict, embed_url: str, + timeout: float | httpx.Timeout | None = None, ) -> dict: data: dict = {} @@ -165,7 +168,9 @@ class HuggingFaceEmbedding(BaseLLM): task_type: Final = optional_params.pop("input_type", None) if call_type == "sync": - hf_task: Final = get_hf_task_embedding_for_model(model=model, task_type=task_type, api_base=HF_HUB_URL) + hf_task: Final = get_hf_task_embedding_for_model( + model=model, task_type=task_type, api_base=HF_HUB_URL, timeout=timeout + ) elif call_type == "async": return self._async_transform_input(model=model, task_type=task_type, embed_url=embed_url, input=input) @@ -253,6 +258,7 @@ class HuggingFaceEmbedding(BaseLLM): call_type="sync", optional_params=optional_params, embed_url=api_base, + timeout=timeout, ) ## LOGGING @@ -269,9 +275,10 @@ class HuggingFaceEmbedding(BaseLLM): if client is None: client = get_async_httpx_client( llm_provider=litellm.LlmProviders.HUGGINGFACE, + params={"timeout": timeout}, ) - response: Final = await client.post(api_base, headers=headers, data=json.dumps(data)) + response: Final = await client.post(api_base, headers=headers, data=json.dumps(data), timeout=timeout) ## LOGGING logging_obj.post_call( @@ -321,7 +328,9 @@ class HuggingFaceEmbedding(BaseLLM): litellm_params=litellm_params, ) task_type: Final = optional_params.get("input_type", None) - task: Final = get_hf_task_embedding_for_model(model=model, task_type=task_type, api_base=HF_HUB_URL) + task: Final = get_hf_task_embedding_for_model( + model=model, task_type=task_type, api_base=HF_HUB_URL, timeout=timeout + ) # print_verbose(f"{model}, {task}") embed_url = "" if model.startswith(("http://", "https://")): @@ -359,6 +368,7 @@ class HuggingFaceEmbedding(BaseLLM): call_type="sync", optional_params=optional_params, embed_url=embed_url, + timeout=timeout, ) ## LOGGING @@ -374,7 +384,7 @@ class HuggingFaceEmbedding(BaseLLM): ## COMPLETION CALL if client is None or not isinstance(client, HTTPHandler): client = HTTPHandler(concurrent_limit=1) - response: Final = client.post(embed_url, headers=headers, data=json.dumps(data)) + response: Final = client.post(embed_url, headers=headers, data=json.dumps(data), timeout=timeout) ## LOGGING logging_obj.post_call( diff --git a/litellm/main.py b/litellm/main.py index b5d1398ecad..b9fd86274b4 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -7400,6 +7400,7 @@ def embedding( aembedding=aembedding, litellm_params=litellm_params_dict, headers=headers, + timeout=timeout, ) elif custom_llm_provider == "bedrock": if isinstance(input, str): diff --git a/tests/unit/llms/cohere/embed/test_cohere_embedding_handler.py b/tests/unit/llms/cohere/embed/test_cohere_embedding_handler.py new file mode 100644 index 00000000000..f5cb874c6ef --- /dev/null +++ b/tests/unit/llms/cohere/embed/test_cohere_embedding_handler.py @@ -0,0 +1,35 @@ +from typing import Final + +import httpx +import pytest +import respx + +import litellm + +BEDROCK_COHERE_INVOKE_URL: Final = ( + "https://bedrock-runtime.us-east-1.amazonaws.com/model/cohere.embed-english-v3/invoke" +) +COHERE_EMBED_RESPONSE: Final = {"embeddings": [[0.1, 0.2]], "meta": {"billed_units": {"input_tokens": 1}}} + + +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_bedrock_cohere_embedding_sends_the_caller_timeout( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, sync_mode: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post(BEDROCK_COHERE_INVOKE_URL).mock( + return_value=httpx.Response(200, json=COHERE_EMBED_RESPONSE) + ) + request: Final = { + "model": "bedrock/cohere.embed-english-v3", + "input": ["hello"], + "timeout": 45, + "aws_region_name": "us-east-1", + "aws_access_key_id": "mockaws_access_key_id", + "aws_secret_access_key": "mockaws_secret_access_key", + } + + response: Final = litellm.embedding(**request) if sync_mode else await litellm.aembedding(**request) + + assert route.calls.last.request.extensions["timeout"]["read"] == 45 + assert response.data[0]["embedding"] == [0.1, 0.2] diff --git a/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py b/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py index ea1e04291ca..7040fbe5fcc 100644 --- a/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py +++ b/tests/unit/llms/huggingface/embedding/test_huggingface_embedding_handler.py @@ -1,12 +1,12 @@ import json from typing import Final -from unittest.mock import patch, MagicMock, AsyncMock +from unittest.mock import AsyncMock, MagicMock, patch - -import litellm +import httpx import pytest import respx +import litellm MOCK_EMBEDDING_RESPONSE = [[0.1, 0.2, 0.3, 0.4, 0.5]] @@ -52,12 +52,10 @@ def mock_embedding_async_http_handler(): class TestHuggingFaceEmbedding: @pytest.fixture(autouse=True) def setup(self, mock_embedding_http_handler, mock_embedding_async_http_handler, mock_hf_config_fetch): - self.mock_get_task_patcher = patch( - "litellm.llms.huggingface.embedding.handler.get_hf_task_embedding_for_model" - ) + self.mock_get_task_patcher = patch("litellm.llms.huggingface.embedding.handler.get_hf_task_embedding_for_model") self.mock_get_task = self.mock_get_task_patcher.start() - def mock_get_task_side_effect(model, task_type, api_base): + def mock_get_task_side_effect(model, task_type, api_base, timeout=None): if task_type is not None: return task_type return "sentence-similarity" @@ -180,7 +178,9 @@ async def test_hf_embedding_sentence_sim_sends_one_similarity_request( monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.setenv("HUGGINGFACE_API_KEY", "hf-test-token") with respx.mock(assert_all_called=False) as hf: - hf.get(url__regex=r"https://huggingface\.co/api/models/.*").respond(json={"pipeline_tag": "sentence-similarity"}) + hf.get(url__regex=r"https://huggingface\.co/api/models/.*").respond( + json={"pipeline_tag": "sentence-similarity"} + ) hf.get(url__regex=r"https://huggingface\.co/.*/config\.json").respond(json={"max_position_embeddings": 512}) inference: Final = hf.post( "https://router.huggingface.co/hf-inference/pipeline/sentence-similarity/sentence-transformers/TaylorAI/bge-micro-v2" @@ -197,3 +197,33 @@ async def test_hf_embedding_sentence_sim_sends_one_similarity_request( } assert inference.calls.last.request.headers["Authorization"] == "Bearer hf-test-token" assert response.data == [{"object": "embedding", "index": 0, "embedding": 0.7708950042724609}] + + +HF_MODEL_LOOKUP_URL: Final = "https://huggingface.co/api/models/BAAI/bge-m3" +HF_EMBED_URL: Final = "https://router.huggingface.co/hf-inference/pipeline/feature-extraction/BAAI/bge-m3" + + +@pytest.mark.respx(assert_all_called=False) +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_huggingface_embedding_sends_the_caller_timeout( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, sync_mode: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + respx_mock.get(url__regex=r"https://huggingface\.co/.*/config\.json").respond(json={"max_position_embeddings": 512}) + lookup: Final = respx_mock.get(HF_MODEL_LOOKUP_URL).mock( + return_value=httpx.Response(200, json={"pipeline_tag": "feature-extraction"}) + ) + route: Final = respx_mock.post(HF_EMBED_URL).mock(return_value=httpx.Response(200, json=MOCK_EMBEDDING_RESPONSE)) + request: Final = { + "model": "huggingface/BAAI/bge-m3", + "input": ["hello"], + "api_key": "hf_test", + "timeout": 45, + } + + response: Final = litellm.embedding(**request) if sync_mode else await litellm.aembedding(**request) + + assert all(call.request.extensions["timeout"]["read"] == 45 for call in lookup.calls) + assert lookup.called + assert route.calls.last.request.extensions["timeout"]["read"] == 45 + assert response.data[0]["embedding"] == MOCK_EMBEDDING_RESPONSE[0]