Merge 10cbe497fd into dfabf32cd5

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

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

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"]