diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index 8ca2afb9a99..e6e63666e52 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -15491,6 +15491,8 @@ class PythonCore(DeprecatedRouterMembers): deployments=_returned_deployments, ) request_settings: Final[Mapping[str, object]] = request_kwargs if request_kwargs is not None else {} + original_function_name: Final = getattr(request_settings.get("original_function"), "__name__", None) + requires_ocr_capability: Final = original_function_name in {"ocr", "aocr"} requested_params: Final[frozenset[str]] = ( self._pre_call_check_requested_chat_params(request_settings) if request_kwargs is not None and self._host.dynamic().drop_params is False @@ -15503,6 +15505,16 @@ class PythonCore(DeprecatedRouterMembers): _model_info = deployment.get("model_info", {}) max_input_tokens = deployment_max_input_tokens[idx] + if requires_ocr_capability and ( + not isinstance(_model_info, Mapping) or _model_info.get("mode") != "ocr" + ): + invalid_model_indices.add(idx) + verbose_router_logger.debug( + "Skipping deployment %s for OCR request: declared mode=%s", + _litellm_params.get("model"), + _model_info.get("mode") if isinstance(_model_info, Mapping) else None, + ) + continue if ( max_input_tokens is not None and has_countable_input diff --git a/litellm/types/router.py b/litellm/types/router.py index 3363d113fc7..454be6662fc 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -1050,7 +1050,7 @@ class ModelGroupInfo(LiteLLMBaseModel): input_cost_per_pixel: float | None = None mode: ( str - | Literal["chat", "embedding", "completion", "image_generation", "audio_transcription", "rerank", "moderations"] + | Literal["chat", "embedding", "completion", "image_generation", "audio_transcription", "ocr", "rerank", "moderations"] | None ) = Field(default="chat") tpm: int | None = None diff --git a/tests/unit/test_router/test_router_ocr_capability.py b/tests/unit/test_router/test_router_ocr_capability.py new file mode 100644 index 00000000000..eec4cfa545e --- /dev/null +++ b/tests/unit/test_router/test_router_ocr_capability.py @@ -0,0 +1,48 @@ +from litellm import Router + + +def test_router_ocr_capability_gate_skips_chat_fallbacks() -> None: + router = Router( + model_list=[ + { + "model_name": "document-ocr", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "ocr", "mode": "ocr"}, + }, + { + "model_name": "document-ocr", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + "model_info": {"id": "chat", "mode": "chat"}, + }, + ], + enable_pre_call_checks=True, + ) + + async def ocr(*args, **kwargs): + return None + + async def aocr(*args, **kwargs): + return None + + async def chat(*args, **kwargs): + return None + + deployments = router.get_model_list(model_name="document-ocr") + chat_deployments = router._pre_call_checks( + model="document-ocr", + healthy_deployments=deployments, + messages=None, + input=None, + request_kwargs={"original_function": chat}, + ) + assert [deployment["model_info"]["id"] for deployment in chat_deployments] == ["ocr", "chat"] + + for ocr_function in (ocr, aocr): + ocr_deployments = router._pre_call_checks( + model="document-ocr", + healthy_deployments=deployments, + messages=None, + input=None, + request_kwargs={"original_function": ocr_function}, + ) + assert [deployment["model_info"]["id"] for deployment in ocr_deployments] == ["ocr"]