From 004a8201167eb0bd85d52e86473adb9b9b8d1ee7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 22:52:12 -0700 Subject: [PATCH] fix(ocr): send each provider a health-check document it accepts Health checks probed every OCR deployment with a PDF, which Cohere Parse rejects, so /health, background health checks, and the UI Test Connection button marked Cohere Parse deployments unhealthy. BaseOCRConfig gains a get_health_check_document hook (PDF by default) that CohereParseConfig overrides with a 1x1 PNG data URI. cohere also gains ocr in the provider endpoint matrix --- .../health_check_helpers.py | 18 +++++++----- litellm/llms/base_llm/ocr/transformation.py | 8 +++++ litellm/llms/cohere/ocr/transformation.py | 9 ++++++ .../provider_endpoints_support_backup.json | 1 + provider_endpoints_support.json | 1 + .../test_health_check_helpers.py | 29 +++++++++++++++++++ ...st_azure_ai_cohere_parse_transformation.py | 16 ++++++++++ .../ocr/test_cohere_parse_transformation.py | 12 ++++++++ 8 files changed, 87 insertions(+), 7 deletions(-) diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index c745bbea5c4..9f8878d36f0 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -6,14 +6,13 @@ import base64 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 @@ class HealthCheckHelpers: ), "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), ), } diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index 08ae077cb2f..8111f9a194a 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -33,6 +33,8 @@ OCR_REQUEST_FORMAT_HEADER: Final = "x-req-format" 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": @@ -146,6 +148,12 @@ class BaseOCRConfig: """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, diff --git a/litellm/llms/cohere/ocr/transformation.py b/litellm/llms/cohere/ocr/transformation.py index 87454980aa4..dd15d5360a6 100644 --- a/litellm/llms/cohere/ocr/transformation.py +++ b/litellm/llms/cohere/ocr/transformation.py @@ -34,6 +34,9 @@ COHERE_PARSE_OUTPUT_FORMAT_PARAM: Final = "output_format" COHERE_PARSE_OUTPUT_FORMATS: Final = ("markdown", "blocks") COHERE_PARSE_DEFAULT_OUTPUT_FORMAT: Final = "markdown" COHERE_PARSE_SUPPORTED_PARAMS: Final = (COHERE_PARSE_OUTPUT_FORMAT_PARAM, OCR_REQUEST_FORMAT_PARAM) +COHERE_PARSE_HEALTH_CHECK_IMAGE_DATA_URI: Final = ( + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC" +) COHERE_PARSE_IMAGE_ONLY_MESSAGE: Final = ( "Cohere Parse only accepts `image_url` documents (an image URL or a base64 image data URI); " "`document_url` and PDF inputs are not supported." @@ -144,6 +147,12 @@ class CohereParseConfig(BaseOCRConfig): def supports_rust_bridge(self) -> bool: return False + def get_health_check_document(self) -> DocumentType: + return { # mutable-ok: litellm.aocr rejects any document that is not a dict + "type": "image_url", + "image_url": COHERE_PARSE_HEALTH_CHECK_IMAGE_DATA_URI, + } + def _llm_provider(self) -> str: return "cohere" diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index 9d6b1e18f59..dbeaccdda2d 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -559,6 +559,7 @@ "moderations": false, "batches": false, "rerank": true, + "ocr": true, "a2a": true, "interactions": true } diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 41ed8e1d975..c71f4a82a4a 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -594,6 +594,7 @@ "moderations": false, "batches": false, "rerank": true, + "ocr": true, "a2a": true, "interactions": true } diff --git a/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py b/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py index ee2a31beff7..89b377af3a0 100644 --- a/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py @@ -453,3 +453,32 @@ async def test_realtime_health_check_uses_model_level_vertex_params(): "Authorization": "Bearer model-level-token", "x-goog-user-project": "model-level-project", } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model, custom_llm_provider, expected_document_type, expected_uri_prefix", + [ + ("mistral/mistral-ocr-latest", "mistral", "document_url", "data:application/pdf;base64,"), + ("azure_ai/mistral-document-ai-2512", "azure_ai", "document_url", "data:application/pdf;base64,"), + ("cohere/parse-v5.0", "cohere", "image_url", "data:image/png;base64,"), + ("azure_ai/Cohere-parse-v5", "azure_ai", "image_url", "data:image/png;base64,"), + ], +) +async def test_ocr_health_check_sends_the_document_kind_the_provider_config_accepts( + model, custom_llm_provider, expected_document_type, expected_uri_prefix +): + handlers = HealthCheckHelpers.get_mode_handlers( + model=model, + custom_llm_provider=custom_llm_provider, + model_params={"model": model, "api_key": "sk-test"}, + ) + + with patch( # test-quality-ok: the public health-check path has no dependency injection seam + "litellm.aocr", new_callable=AsyncMock, return_value={} + ) as mock_aocr: + await handlers["ocr"]() + + document = mock_aocr.call_args.kwargs["document"] + assert document["type"] == expected_document_type + assert document[expected_document_type].startswith(expected_uri_prefix) diff --git a/tests/test_litellm/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py b/tests/test_litellm/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py index 8c4dd25aa77..3f98e9b6a2d 100644 --- a/tests/test_litellm/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py +++ b/tests/test_litellm/llms/azure_ai/ocr/test_azure_ai_cohere_parse_transformation.py @@ -166,3 +166,19 @@ async def test_aocr_rejects_pdf_before_calling_foundry(disable_aiohttp_transport assert exc_info.value.llm_provider == "azure_ai" assert not route.called + + +@pytest.mark.asyncio +async def test_ahealth_check_ocr_sends_an_image_to_the_foundry_cohere_parse_deployment( + disable_aiohttp_transport, respx_mock +): + route = respx_mock.post(PARSE_URL).respond(json=_parse_response()) + + result = await litellm.ahealth_check( + model_params={"model": MODEL, "api_base": API_BASE, "api_key": "test-key"}, mode="ocr" + ) + + document = json.loads(route.calls.last.request.content)["document"] + assert document["type"] == "image_url" + assert document["image_url"].startswith("data:image/png;base64,") + assert "error" not in result diff --git a/tests/test_litellm/llms/cohere/ocr/test_cohere_parse_transformation.py b/tests/test_litellm/llms/cohere/ocr/test_cohere_parse_transformation.py index 64f2bb383c2..cb9af56f5e0 100644 --- a/tests/test_litellm/llms/cohere/ocr/test_cohere_parse_transformation.py +++ b/tests/test_litellm/llms/cohere/ocr/test_cohere_parse_transformation.py @@ -215,3 +215,15 @@ async def test_aocr_without_api_key_names_the_env_var(disable_aiohttp_transport, await litellm.aocr(model=MODEL, document=IMAGE_DOCUMENT) assert not route.called + + +@pytest.mark.asyncio +async def test_ahealth_check_ocr_sends_an_image_cohere_parse_accepts(disable_aiohttp_transport, respx_mock): + route = respx_mock.post(PARSE_URL).respond(json=_markdown_response()) + + result = await litellm.ahealth_check(model_params={"model": MODEL, "api_key": "test-key"}, mode="ocr") + + document = json.loads(route.calls.last.request.content)["document"] + assert document["type"] == "image_url" + assert document["image_url"].startswith("data:image/png;base64,") + assert "error" not in result