feat(ocr): add Cohere Parse support for cohere and azure_ai

litellm/litellm_core_utils/health_check_helpers.py

@@ -6,14 +6,13 @@
 from collections.abc import Awaitable, Callable
 from typing import TYPE_CHECKING, Final, Literal
 
-from litellm.types.utils import LIST_BATCHES_SUPPORTED_PROVIDERS
+from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, DocumentType
+from litellm.types.utils import LIST_BATCHES_SUPPORTED_PROVIDERS, LlmProviders
 
 if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging
     from litellm.types.utils import ImageResponse
 
-# Minimal PDF for health checks - base64 encoded 1-page PDF with just "test"
-TEST_PDF_URL = "data:application/pdf;base64,JVBERi0xLjQKJeLjz9MKMyAwIG9iago8PC9UeXBlIC9QYWdlCi9QYXJlbnQgMSAwIFIKL01lZGlhQm94IFswIDAgNjEyIDc5Ml0KL0NvbnRlbnRzIDQgMCBSCi9SZXNvdXJjZXMgPDwvRm9udCA8PC9GMSAyIDAgUj4+Pj4+PgplbmRvYmoKNCAwIG9iago8PC9MZW5ndGggNDQ+PgpzdHJlYW0KQlQKL0YxIDI0IFRmCjEwMCA3MDAgVGQKKHRlc3QpIFRqCkVUCmVuZHN0cmVhbQplbmRvYmoKMiAwIG9iago8PC9UeXBlIC9Gb250Ci9TdWJ0eXBlIC9UeXBlMQovQmFzZUZvbnQgL0hlbHZldGljYT4+CmVuZG9iagoxIDAgb2JqCjw8L1R5cGUgL1BhZ2VzCi9LaWRzIFszIDAgUl0KL0NvdW50IDE+PgplbmRvYmoKNSAwIG9iago8PC9UeXBlIC9DYXRhbG9nCi9QYWdlcyAxIDAgUj4+CmVuZG9iagp0cmFpbGVyCjw8L1NpemUgNgovUm9vdCA1IDAgUj4+CnN0YXJ0eHJlZgozMjQKJSVFT0Y="
 
 # Minimal image for health checks - base64 encoded 512x512 blue circle on a white background PNG
 TEST_IMAGE_BASE64 = "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAJk0lEQVR42u3VQREAIRADwVWCOmTjBVzwSLorCri6nbkAVBpPACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAAAgAAAIAZdY+HgEBgIRr/meeGgGA8EMvDAgAuPh6gACAi68HCAA4+mKAAICjLwYIALj7SoAAgLuvBAgA7r4pAQKAu29KgADg7psSIAA4/SYDCADuvikBAoDTbzKAAOD0mwwgADj9JgMIAE6/yQACgNNvMoAA4PSbDCAAOP0mAwgATr/JAAKA668BIAA4/TKAAOD0mwwgALj+pgEIAE6/yQACgOtvGoAA4PSbDCAAuP6mAQgATr/JAAKA628agADg9JsMIAC4/qYBCACuv2kAAoDrbxqAAOD0mwwgALj+pgEIAK6/aQACgOtvGoAA4PqbBiAArr+ZBiAATr+ZDCAArr+ZBiAArr+ZBiAArr+ZBiAArr+ZBggArr+ZBggArr+ZBggArr+ZBggArr+ZBggArr+ZBggArr+ZBggAAmAmAAKA62+mAQKA62+mAQKA62/m1xYAXH/TAAQA1980AAHA9TcNQAAEwEwAEADX30wDEADX30wDEADX30wDEADX30wDEAABMBMABMD1N9MABMD1N9MABEAAzAQAAXD9zTQAAXD9zTRAABAAMwEQAFx/Mw0QAFx/Mw0QAATATAAEANffTAMEANffTAMEAAEwEwABwPU30wABQADMBEAAXH8z0wABcP3NTAMEQADMTAAEwPU3Mw0QAAEwMwEQANffTAMQAAEwEwAEwPU30wAEQADMBAABcP3NNAABEAAzAUAAXH8zDUAABMBMABAA199MAwQAATATAAHA9TfTAAFAAMwEQABw/c00QAAQADMBEAAEwEwABMD1NzMNEAABMDMBEADX38w0QAAEwMwEQAAEwMwEQABcfzPTAAEQADMTAAEQADMTAAFw/c1MAwRAAMxMAARAAMxMAATA9TczDRAAATATAARAAMwEAAFw/c00AAEQADMBQAAEwEwAEADX30wDBAABMBMAAUAAzARAABAAMwEQAFx/Mw0QAAEwMwEQAAEwMwEQAAEwMwEQANffzDRAAATAzARAAATAzARAAATAzARAAATAzARAAFx/M9MAARAAMxMAARAAMxMAARAAMxMAARAAMxMAAXD9zUwDBEAAzEwABEAAzAQAARAAMwFAAATATAAQAAEwEwABQADMBEAAEAAzARAAXH8zDRAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAAATAzARAA19/MNEAANMDM9UcABMBMABAAATATAAHwBAJgJgACgACYCYAAIABmAiAA+IvMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAABMDMBEAANMDMXH8BEAAzEwABEAAzEwABEAAzEwABEAAzEwABEAAzEwABEAAzAUAABMBMABAADTBz/REAATATAAFAAMwEQAAQADMBEAAEwEwABAANMHP9BUAAzEwABEAAzEwABEAAzEwABEAAzEwABEADzMz1FwABMDMBEAABMDMBEAABMDMBEAANMDPXXwAEwMwEQAAEwMwEQAAEwMwEQAA0wMxcfwEQADMTAAEQADMBQAA0wMz1RwAEwEwAEAABMBMABEADzFx/AUAAzARAABAAMwEQADTAzPUXAATATAAEAAEwEwABQAPMXH8BEAAzEwABEAAzEwAB0AAzc/0FQADMTAAEQAPMzPUXAAEwMwEQAAEwMwEQAA0wM9dfAATAzARAADTAzFx/ARAAMwFAADTAzPVHAATATAAQAA0wc/0RAAEwEwAEQAPMXH8EQADMBAAB0AAz118AEAAzARAANMDM9RcABMBMAAQADTBz/QUAATATAAFAA8xcfwFAA8xcfwFAAMwEQADQADPXXwAEwMwEQAA0wMxcfwHQADNz/QVAAMxMAARAA8xcfwRAA8xcfwRAAMwEAAHQADPXHwHQADPXHwEQADMBQAA0wMz1RwA0wMz1RwAEwEwAEAANMHP9BQANMHP9BQANMHP9BQANMHP9BQABMBMAAUADzFx/AUADzFx/AUADzFx/AUADzFx/AUADzFx/AUADzPVHABAAEwAEAA0w1x8BQAPM9UcA0ABz/REANMBcfwQADTDXHwHQADPXHwHQADPXHwHQADPXHwHQADPXHwHQADPXHwGQATOnHwHQADPXHwHQADPXXwDQADPXXwDQADPXXwDQADPXXwCQAXP6EQA0wFx/BAANMNcfAUADzPVHAJABc/oRADTAXH8EABkwpx8BQAPM9UcAkAFz+hEANMBcfwQAGTCnHwFAA8z1RwCQAXP6EQBkwJx+BAANMNcfAUAGzOlHAJABc/oRAGTA6QcBQAacfhAAZMDpRwBABpx+BABkwOlHAEAGnH4EAJTA3UcAQAacfgQAlMDdRwBACdx9BACUwN1HAEAJ3H0EAJTA3UcAQAwcfQQAmmLgsyIA0NIDHw4BgJYe+DQIAISHwVMjAJDQDI+AAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACACAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAgAAAIAAACAIAAACAAAAgAAAIAgAAAIAAACAAAAgCAAAAIAABNHpialFcmLajuAAAAAElFTkSuQmCC"
@@ -29,6 +28,14 @@ def get_image_file_for_health_check() -> bytes:
     return base64.b64decode(TEST_IMAGE_BASE64)
 
 
+def _ocr_health_check_document(model: str, custom_llm_provider: str) -> DocumentType:
+    from litellm.utils import ProviderConfigManager
+
+    provider: Final = next((known for known in LlmProviders if known.value == custom_llm_provider), None)
+    config: Final = ProviderConfigManager.get_provider_ocr_config(model=model, provider=provider) if provider else None
+    return (config or BaseOCRConfig()).get_health_check_document()
+
+
 class HealthCheckHelpers:
     @staticmethod
     async def ahealth_check_wildcard_models(
@@ -247,9 +254,6 @@ def get_mode_handlers(
             ),
             "ocr": lambda: litellm.aocr(
                 **_filter_model_params(model_params=model_params),
-                document={
-                    "type": "document_url",
-                    "document_url": TEST_PDF_URL,
-                },
+                document=_ocr_health_check_document(model=model, custom_llm_provider=custom_llm_provider),
             ),
         }

litellm/llms/azure_ai/ocr/__init__.py

@@ -1,12 +1,14 @@
 """Azure AI OCR module."""
 
+from .cohere_parse_transformation import AzureAICohereParseConfig
 from .common_utils import get_azure_ai_ocr_config
 from .document_intelligence.transformation import (
     AzureDocumentIntelligenceOCRConfig,
 )
 from .transformation import AzureAIOCRConfig
 
 __all__ = [
+    "AzureAICohereParseConfig",
     "AzureAIOCRConfig",
     "AzureDocumentIntelligenceOCRConfig",
     "get_azure_ai_ocr_config",

litellm/llms/azure_ai/ocr/cohere_parse_transformation.py

@@ -0,0 +1,91 @@
+"""Cohere Parse served from Azure AI Foundry (`/providers/cohere/v2/parse`)."""
+
+from collections.abc import Mapping
+from typing import Final
+
+import httpx
+
+from litellm.litellm_core_utils.prompt_templates.image_handling import (
+    async_convert_url_to_base64,
+    convert_url_to_base64,
+)
+from litellm.llms.azure_ai.common_utils import get_azure_ai_auth_headers
+from litellm.llms.cohere.ocr.transformation import COHERE_PARSE_PATH, CohereParseConfig
+from litellm.secret_managers.main import get_secret_str
+
+AZURE_AI_API_KEY_ENV_VAR: Final = "AZURE_AI_API_KEY"
+AZURE_AI_API_BASE_ENV_VAR: Final = "AZURE_AI_API_BASE"
+AZURE_AI_COHERE_PROVIDER_PATH: Final = "/providers/cohere"
+AZURE_AI_MODELS_PATH_SUFFIX: Final = "/models"
+
+
+class AzureAICohereParseConfig(CohereParseConfig):
+    """Same request and response shape as Cohere Parse, behind Azure AI auth and URL layout.
+
+    Foundry cannot fetch external URLs, so remote images are inlined as base64 data URIs.
+    """
+
+    def get_api_key_env_var(self) -> str | None:
+        return AZURE_AI_API_KEY_ENV_VAR
+
+    def _llm_provider(self) -> str:
+        return "azure_ai"
+
+    def validate_environment(
+        self,
+        headers: Mapping[str, str],
+        model: str,
+        api_key: str | None = None,
+        api_base: str | None = None,
+        litellm_params: Mapping[str, object] | None = None,
+        **kwargs: object,  # kwargs-ok: BaseOCRConfig.validate_environment signature
+    ) -> dict[str, str]:  # mutable-ok: BaseOCRConfig signature
+        resolved_base: Final = api_base or get_secret_str(AZURE_AI_API_BASE_ENV_VAR)
+        if resolved_base is None:
+            raise ValueError(
+                f"Missing Azure AI API Base - Set {AZURE_AI_API_BASE_ENV_VAR} environment variable "
+                "or pass api_base parameter"
+            )
+        resolved_key: Final = api_key or get_secret_str(AZURE_AI_API_KEY_ENV_VAR)
+        return {  # mutable-ok: BaseOCRConfig signature
+            **get_azure_ai_auth_headers(api_key=resolved_key, litellm_params=litellm_params),
+            "Content-Type": "application/json",
+            **headers,
+        }
+
+    def get_complete_url(
+        self,
+        api_base: str | None,
+        model: str,
+        optional_params: Mapping[str, object],
+        litellm_params: Mapping[str, object] | None = None,
+        **kwargs: object,  # kwargs-ok: BaseOCRConfig.get_complete_url signature
+    ) -> str:
+        resolved_base: Final = api_base or get_secret_str(AZURE_AI_API_BASE_ENV_VAR)
+        if resolved_base is None:
+            raise ValueError(
+                f"Missing Azure AI API Base - Set {AZURE_AI_API_BASE_ENV_VAR} environment variable "
+                "or pass api_base parameter"
+            )
+        url: Final = httpx.URL(resolved_base)
+        if not url.is_absolute_url:
+            raise ValueError(
+                "Azure AI API Base must be an absolute URL including scheme (e.g. "
+                f"'https://<resource>.services.ai.azure.com'). Got api_base={resolved_base!r}."
+            )
+        path: Final = url.path.rstrip("/")
+        if path.endswith(COHERE_PARSE_PATH):
+            return str(url.copy_with(path=path))
+        if path.endswith(f"{AZURE_AI_COHERE_PROVIDER_PATH}/v2"):
+            return str(url.copy_with(path=f"{path}/parse"))
+        return str(
+            url.copy_with(
+                path=f"{path.removesuffix(AZURE_AI_MODELS_PATH_SUFFIX)}{AZURE_AI_COHERE_PROVIDER_PATH}{COHERE_PARSE_PATH}"
+            )
+        )
+
+    def _resolve_image_url_sync(self, image_url: str) -> str:
+        return convert_url_to_base64(image_url)
+
+    async def _resolve_image_url_async(self, image_url: str) -> str:
+        return await async_convert_url_to_base64(image_url)

litellm/llms/azure_ai/ocr/common_utils.py

@@ -24,6 +24,11 @@ def is_azure_document_intelligence_model(model: str) -> bool:
     return "doc-intelligence" in lowered or "documentintelligence" in lowered
 
 
+def is_azure_cohere_parse_model(model: str) -> bool:
+    lowered: Final = model.lower()
+    return "cohere" in lowered and "parse" in lowered
+
+
 def get_azure_ai_ocr_config(model: str) -> Optional["BaseOCRConfig"]:
     """
     Determine which Azure AI OCR configuration to use based on the model name.
@@ -46,6 +51,7 @@ def get_azure_ai_ocr_config(model: str) -> Optional["BaseOCRConfig"]:
         >>> get_azure_ai_ocr_config("azure_ai/pixtral-12b-2409")
         <AzureAIOCRConfig object>
     """
+    from litellm.llms.azure_ai.ocr.cohere_parse_transformation import AzureAICohereParseConfig
     from litellm.llms.azure_ai.ocr.document_intelligence.transformation import (
         AzureDocumentIntelligenceOCRConfig,
     )
@@ -56,6 +62,10 @@ def get_azure_ai_ocr_config(model: str) -> Optional["BaseOCRConfig"]:
         verbose_logger.debug("Routing %s to Azure Document Intelligence OCR config", model)
         return AzureDocumentIntelligenceOCRConfig()
 
+    if is_azure_cohere_parse_model(model):
+        verbose_logger.debug("Routing %s to Azure AI Cohere Parse config", model)
+        return AzureAICohereParseConfig()
+
     # Default to Mistral-based OCR for other azure_ai models
     verbose_logger.debug("Routing %s to Azure AI (Mistral) OCR config", model)
     return AzureAIOCRConfig()

litellm/llms/base_llm/ocr/transformation.py

@@ -33,6 +33,8 @@
 
 PROVIDER_NATIVE_RESPONSE_KEY: Final = "provider_native_response"
 
+HEALTH_CHECK_PDF_DATA_URI: Final = "data:application/pdf;base64,JVBERi0xLjQKJeLjz9MKMyAwIG9iago8PC9UeXBlIC9QYWdlCi9QYXJlbnQgMSAwIFIKL01lZGlhQm94IFswIDAgNjEyIDc5Ml0KL0NvbnRlbnRzIDQgMCBSCi9SZXNvdXJjZXMgPDwvRm9udCA8PC9GMSAyIDAgUj4+Pj4+PgplbmRvYmoKNCAwIG9iago8PC9MZW5ndGggNDQ+PgpzdHJlYW0KQlQKL0YxIDI0IFRmCjEwMCA3MDAgVGQKKHRlc3QpIFRqCkVUCmVuZHN0cmVhbQplbmRvYmoKMiAwIG9iago8PC9UeXBlIC9Gb250Ci9TdWJ0eXBlIC9UeXBlMQovQmFzZUZvbnQgL0hlbHZldGljYT4+CmVuZG9iagoxIDAgb2JqCjw8L1R5cGUgL1BhZ2VzCi9LaWRzIFszIDAgUl0KL0NvdW50IDE+PgplbmRvYmoKNSAwIG9iago8PC9UeXBlIC9DYXRhbG9nCi9QYWdlcyAxIDAgUj4+CmVuZG9iagp0cmFpbGVyCjw8L1NpemUgNgovUm9vdCA1IDAgUj4+CnN0YXJ0eHJlZgozMjQKJSVFT0Y="
+
 
 def parse_ocr_request_format(value: object) -> OCRRequestFormat:
     if value == "litellm":
@@ -142,6 +144,16 @@ def get_api_key_env_var(self) -> str | None:
         """
         return None
 
+    def supports_rust_bridge(self) -> bool:
+        """Whether the Rust OCR bridge may serve this config when it is enabled for the provider."""
+        return True
+
+    def get_health_check_document(self) -> DocumentType:
+        return {  # mutable-ok: litellm.aocr rejects any document that is not a dict
+            "type": "document_url",
+            "document_url": HEALTH_CHECK_PDF_DATA_URI,
+        }
+
     def map_ocr_params(
         self,
         non_default_params: dict,

litellm/llms/cohere/ocr/__init__.py

@@ -0,0 +1,3 @@
+from litellm.llms.cohere.ocr.transformation import CohereParseConfig
+
+__all__ = ("CohereParseConfig",)