diff --git a/litellm/llms/mistral/ocr/guardrail_translation/__init__.py b/litellm/llms/mistral/ocr/guardrail_translation/__init__.py new file mode 100644 index 00000000000..da7b6ee6bf0 --- /dev/null +++ b/litellm/llms/mistral/ocr/guardrail_translation/__init__.py @@ -0,0 +1,11 @@ +"""Mistral OCR handler for Unified Guardrails.""" + +from litellm.llms.mistral.ocr.guardrail_translation.handler import OCRHandler +from litellm.types.utils import CallTypes + +guardrail_translation_mappings = { + CallTypes.ocr: OCRHandler, + CallTypes.aocr: OCRHandler, +} + +__all__ = ["guardrail_translation_mappings", "OCRHandler"] diff --git a/litellm/llms/mistral/ocr/guardrail_translation/handler.py b/litellm/llms/mistral/ocr/guardrail_translation/handler.py new file mode 100644 index 00000000000..87d79a3ce60 --- /dev/null +++ b/litellm/llms/mistral/ocr/guardrail_translation/handler.py @@ -0,0 +1,155 @@ +""" +OCR Handler for Unified Guardrails + +Provides guardrail translation support for the OCR endpoint. +Processes the extracted markdown text from OCR pages. +""" + +from typing import TYPE_CHECKING, Any, List, Optional + +from litellm._logging import verbose_proxy_logger +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.llms.base_llm.ocr.transformation import OCRResponse + + +class OCRHandler(BaseTranslation): + """ + Handler for processing OCR requests/responses with guardrails. + + Input: The OCR input is a document URL/reference - not text content. + We pass the document URL as text for guardrails that may want to + validate or filter document sources. + + Output: OCR responses contain extracted markdown text per page. + The handler extracts all page markdown, applies guardrails, + and maps the guardrailed text back to the pages. + """ + + async def process_input_messages( + self, + data: dict, + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional[Any] = None, + ) -> Any: + """ + Process OCR input by applying guardrails to the document reference. + + The OCR input contains a document dict with a URL. We extract + the URL and pass it to the guardrail for validation. + + Args: + data: Request data containing 'document' parameter + guardrail_to_apply: The guardrail instance to apply + litellm_logging_obj: Optional logging object + + Returns: + Modified data with guardrails applied + """ + document = data.get("document") + if document is None or not isinstance(document, dict): + verbose_proxy_logger.debug( + "OCR guardrail: No valid document found in request data" + ) + return data + + # Extract the document URL for guardrail checking + texts_to_check: List[str] = [] + doc_type = document.get("type") + if doc_type == "document_url": + url = document.get("document_url") + if url and isinstance(url, str): + texts_to_check.append(url) + elif doc_type == "image_url": + url = document.get("image_url") + if url and isinstance(url, str): + texts_to_check.append(url) + + if not texts_to_check: + return data + + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + model = data.get("model") + if model: + inputs["model"] = model + + await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=data, + input_type="request", + logging_obj=litellm_logging_obj, + ) + + return data + + async def process_output_response( + self, + response: "OCRResponse", + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional[Any] = None, + user_api_key_dict: Optional[Any] = None, + ) -> Any: + """ + Process OCR output by applying guardrails to extracted page text. + + Extracts markdown text from each OCR page, applies guardrails, + and maps the guardrailed text back to the pages. + + Args: + response: OCRResponse with pages containing markdown text + guardrail_to_apply: The guardrail instance to apply + litellm_logging_obj: Optional logging object + user_api_key_dict: User API key metadata + + Returns: + Modified OCRResponse with guardrailed page text + """ + if not hasattr(response, "pages") or not response.pages: + verbose_proxy_logger.debug( + "OCR guardrail: No pages found in OCR response" + ) + return response + + # Extract markdown text from all pages + texts_to_check: List[str] = [] + page_indices: List[int] = [] + for i, page in enumerate(response.pages): + if hasattr(page, "markdown") and page.markdown: + texts_to_check.append(page.markdown) + page_indices.append(i) + + if not texts_to_check: + return response + + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + model = getattr(response, "model", None) + if model: + inputs["model"] = model + + # Add user metadata if available + if user_api_key_dict is not None: + metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) + inputs.update(metadata) # type: ignore + + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + + # Map guardrailed text back to pages + guardrailed_texts = guardrailed_inputs.get("texts", []) + for idx, page_idx in enumerate(page_indices): + if idx < len(guardrailed_texts): + response.pages[page_idx].markdown = guardrailed_texts[idx] + + verbose_proxy_logger.debug( + "OCR guardrail: Applied guardrail to %d pages", + len(guardrailed_texts), + ) + + return response diff --git a/litellm/types/utils.py b/litellm/types/utils.py index dda32d98383..c586c415db8 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -800,6 +800,9 @@ API_ROUTE_TO_CALL_TYPES = { CallTypes.allm_passthrough_route, ], "/v1/messages": [CallTypes.anthropic_messages], + # OCR + "/ocr": [CallTypes.aocr, CallTypes.ocr], + "/v1/ocr": [CallTypes.aocr, CallTypes.ocr], } diff --git a/tests/test_litellm/llms/ocr/__init__.py b/tests/test_litellm/llms/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/ocr/guardrail_translation/__init__.py b/tests/test_litellm/llms/ocr/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py b/tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py new file mode 100644 index 00000000000..f6151497e1c --- /dev/null +++ b/tests/test_litellm/llms/ocr/guardrail_translation/test_ocr_guardrail_handler.py @@ -0,0 +1,303 @@ +""" +Unit tests for OCR Guardrail Translation Handler +""" + +import os +import re +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms import get_guardrail_translation_mapping +from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo +from litellm.llms.mistral.ocr.guardrail_translation.handler import OCRHandler +from litellm.types.utils import CallTypes + + +class MockGuardrail(CustomGuardrail): + """Mock guardrail for testing""" + + async def apply_guardrail( + self, inputs: dict, request_data: dict, input_type: str, **kwargs + ) -> dict: + texts = inputs.get("texts", []) + return {"texts": [f"{text} [GUARDRAILED]" for text in texts]} + + +class BlockingGuardrail(CustomGuardrail): + """Mock guardrail that raises on forbidden content""" + + async def apply_guardrail( + self, inputs: dict, request_data: dict, input_type: str, **kwargs + ) -> dict: + texts = inputs.get("texts", []) + for text in texts: + if "FORBIDDEN" in text: + raise ValueError("Content blocked by guardrail") + return {"texts": texts} + + +class TestHandlerDiscovery: + """Test that the handler is properly discovered""" + + def test_handler_discovered_for_ocr(self): + """Test that ocr CallType is mapped to handler""" + handler_class = get_guardrail_translation_mapping(CallTypes.ocr) + assert handler_class == OCRHandler + + def test_handler_discovered_for_aocr(self): + """Test that aocr CallType is mapped to handler""" + handler_class = get_guardrail_translation_mapping(CallTypes.aocr) + assert handler_class == OCRHandler + + +class TestInputProcessing: + """Test input processing functionality""" + + @pytest.mark.asyncio + async def test_process_document_url(self): + """Test processing a document_url input""" + handler = OCRHandler() + guardrail = MockGuardrail(guardrail_name="test") + + data = { + "model": "mistral/mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://arxiv.org/pdf/2201.04234", + }, + } + + result = await handler.process_input_messages(data, guardrail) + + # Document should be unchanged (guardrail can reject but not modify URL) + assert result["document"]["document_url"] == "https://arxiv.org/pdf/2201.04234" + assert result["model"] == "mistral/mistral-ocr-latest" + + @pytest.mark.asyncio + async def test_process_image_url(self): + """Test processing an image_url input""" + handler = OCRHandler() + guardrail = MockGuardrail(guardrail_name="test") + + data = { + "model": "mistral/mistral-ocr-latest", + "document": { + "type": "image_url", + "image_url": "https://example.com/image.png", + }, + } + + result = await handler.process_input_messages(data, guardrail) + + assert result["document"]["image_url"] == "https://example.com/image.png" + + @pytest.mark.asyncio + async def test_process_no_document(self): + """Test processing when no document is provided""" + handler = OCRHandler() + guardrail = MockGuardrail(guardrail_name="test") + + data = {"model": "mistral/mistral-ocr-latest"} + + result = await handler.process_input_messages(data, guardrail) + + assert result == data + assert "document" not in result + + @pytest.mark.asyncio + async def test_process_invalid_document(self): + """Test processing when document is not a dict""" + handler = OCRHandler() + guardrail = MockGuardrail(guardrail_name="test") + + data = {"model": "mistral/mistral-ocr-latest", "document": "not_a_dict"} + + result = await handler.process_input_messages(data, guardrail) + + assert result == data + + @pytest.mark.asyncio + async def test_input_blocking_guardrail(self): + """Test that a blocking guardrail can reject OCR input""" + handler = OCRHandler() + guardrail = BlockingGuardrail(guardrail_name="blocker") + + data = { + "model": "mistral/mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://example.com/FORBIDDEN_document.pdf", + }, + } + + with pytest.raises(ValueError, match="Content blocked by guardrail"): + await handler.process_input_messages(data, guardrail) + + +class TestOutputProcessing: + """Test output processing functionality""" + + @pytest.mark.asyncio + async def test_process_single_page(self): + """Test processing OCR response with a single page""" + handler = OCRHandler() + guardrail = MockGuardrail(guardrail_name="test") + + response = OCRResponse( + pages=[OCRPage(index=0, markdown="Hello world from OCR")], + model="mistral/mistral-ocr-latest", + ) + + result = await handler.process_output_response(response, guardrail) + + assert result.pages[0].markdown == "Hello world from OCR [GUARDRAILED]" + + @pytest.mark.asyncio + async def test_process_multiple_pages(self): + """Test processing OCR response with multiple pages""" + handler = OCRHandler() + guardrail = MockGuardrail(guardrail_name="test") + + response = OCRResponse( + pages=[ + OCRPage(index=0, markdown="Page one content"), + OCRPage(index=1, markdown="Page two content"), + OCRPage(index=2, markdown="Page three content"), + ], + model="mistral/mistral-ocr-latest", + ) + + result = await handler.process_output_response(response, guardrail) + + assert result.pages[0].markdown == "Page one content [GUARDRAILED]" + assert result.pages[1].markdown == "Page two content [GUARDRAILED]" + assert result.pages[2].markdown == "Page three content [GUARDRAILED]" + + @pytest.mark.asyncio + async def test_process_empty_pages(self): + """Test processing OCR response with no pages""" + handler = OCRHandler() + guardrail = MockGuardrail(guardrail_name="test") + + response = OCRResponse( + pages=[], + model="mistral/mistral-ocr-latest", + ) + + result = await handler.process_output_response(response, guardrail) + + assert result.pages == [] + + @pytest.mark.asyncio + async def test_process_page_with_empty_markdown(self): + """Test processing page where markdown is empty""" + handler = OCRHandler() + guardrail = MockGuardrail(guardrail_name="test") + + response = OCRResponse( + pages=[ + OCRPage(index=0, markdown=""), + OCRPage(index=1, markdown="Non-empty content"), + ], + model="mistral/mistral-ocr-latest", + ) + + result = await handler.process_output_response(response, guardrail) + + # Empty markdown page should be skipped + assert result.pages[0].markdown == "" + # Non-empty page should be guardrailed + assert result.pages[1].markdown == "Non-empty content [GUARDRAILED]" + + @pytest.mark.asyncio + async def test_process_preserves_page_metadata(self): + """Test that guardrail processing preserves page metadata""" + handler = OCRHandler() + guardrail = MockGuardrail(guardrail_name="test") + + response = OCRResponse( + pages=[ + OCRPage(index=0, markdown="Page content"), + ], + model="mistral/mistral-ocr-latest", + usage_info=OCRUsageInfo(pages_processed=1, doc_size_bytes=1024), + ) + + result = await handler.process_output_response(response, guardrail) + + assert result.pages[0].index == 0 + assert result.pages[0].markdown == "Page content [GUARDRAILED]" + assert result.model == "mistral/mistral-ocr-latest" + assert result.usage_info.pages_processed == 1 + + @pytest.mark.asyncio + async def test_output_blocking_guardrail(self): + """Test that a blocking guardrail can reject OCR output""" + handler = OCRHandler() + guardrail = BlockingGuardrail(guardrail_name="blocker") + + response = OCRResponse( + pages=[OCRPage(index=0, markdown="This contains FORBIDDEN text")], + model="mistral/mistral-ocr-latest", + ) + + with pytest.raises(ValueError, match="Content blocked by guardrail"): + await handler.process_output_response(response, guardrail) + + +class TestPIIMaskingScenario: + """Test real-world scenario: PII masking in OCR output""" + + @pytest.mark.asyncio + async def test_pii_masking_in_ocr_pages(self): + """Test that PII can be masked from OCR extracted text""" + + class PIIMaskingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, inputs: dict, request_data: dict, input_type: str, **kwargs + ) -> dict: + texts = inputs.get("texts", []) + masked_texts = [] + for text in texts: + masked = re.sub( + r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b", + "[EMAIL_REDACTED]", + text, + ) + masked = re.sub( + r"\b\d{3}-\d{2}-\d{4}\b", + "[SSN_REDACTED]", + masked, + ) + masked_texts.append(masked) + return {"texts": masked_texts} + + handler = OCRHandler() + guardrail = PIIMaskingGuardrail(guardrail_name="mask_pii") + + response = OCRResponse( + pages=[ + OCRPage( + index=0, + markdown="Name: John Doe\nEmail: john@example.com\nSSN: 123-45-6789", + ), + OCRPage( + index=1, + markdown="Contact: jane@corp.com for details", + ), + ], + model="mistral/mistral-ocr-latest", + ) + + result = await handler.process_output_response(response, guardrail) + + assert "john@example.com" not in result.pages[0].markdown + assert "123-45-6789" not in result.pages[0].markdown + assert "[EMAIL_REDACTED]" in result.pages[0].markdown + assert "[SSN_REDACTED]" in result.pages[0].markdown + assert "jane@corp.com" not in result.pages[1].markdown + assert "[EMAIL_REDACTED]" in result.pages[1].markdown diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index b41cded1d0a..7c29c8161db 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -5,11 +5,15 @@ import pytest from litellm.caching import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse +from litellm.llms.mistral.ocr.guardrail_translation.handler import OCRHandler from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import ( MCPGuardrailTranslationHandler, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import unified_guardrail as unified_module +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import ( + unified_guardrail as unified_module, +) from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) @@ -23,12 +27,14 @@ class RecordingGuardrail(CustomGuardrail): def __init__(self): super().__init__(guardrail_name="recording-guardrail") self.event_history = [] + self.apply_calls = [] def should_run_guardrail(self, data, event_type): # type: ignore[override] self.event_history.append(event_type) return True async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + self.apply_calls.append({"inputs": inputs, "input_type": input_type}) return {"texts": inputs.get("texts", [])} @@ -54,6 +60,8 @@ def _inject_mcp_handler_mapping(): unified_module.endpoint_guardrail_translation_mappings = { CallTypes.call_mcp_tool: MCPGuardrailTranslationHandler, CallTypes.anthropic_messages: _NoopTranslation, + CallTypes.ocr: OCRHandler, + CallTypes.aocr: OCRHandler, } yield unified_module.endpoint_guardrail_translation_mappings = None @@ -229,3 +237,159 @@ class TestUnifiedLLMGuardrails: f"Chunk {i} lost its content (got {content!r}). " f"Expected non-empty content for every streamed chunk." ) + + class TestOCRGuardrailE2E: + """End-to-end tests: UnifiedLLMGuardrails -> OCRHandler.""" + + @pytest.mark.asyncio + async def test_pre_call_hook_invokes_ocr_handler_for_input(self): + """ + Verify that async_pre_call_hook with call_type=aocr routes through + the OCR handler and calls apply_guardrail with the document URL. + """ + handler = UnifiedLLMGuardrails() + guardrail = RecordingGuardrail() + cache = DualCache() + + data = { + "guardrail_to_apply": guardrail, + "model": "mistral/mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://arxiv.org/pdf/2201.04234", + }, + } + + result = await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + cache=cache, + data=data, + call_type=CallTypes.aocr.value, + ) + + # Guardrail should have been checked and invoked + assert guardrail.event_history == [GuardrailEventHooks.pre_call] + assert len(guardrail.apply_calls) == 1 + assert guardrail.apply_calls[0]["input_type"] == "request" + assert "https://arxiv.org/pdf/2201.04234" in guardrail.apply_calls[0]["inputs"]["texts"] + + # Data should be returned with document intact + assert result["document"]["document_url"] == "https://arxiv.org/pdf/2201.04234" + + @pytest.mark.asyncio + async def test_moderation_hook_invokes_ocr_handler(self): + """ + Verify that async_moderation_hook with call_type=aocr routes through + the OCR handler correctly. + """ + handler = UnifiedLLMGuardrails() + guardrail = RecordingGuardrail() + + data = { + "guardrail_to_apply": guardrail, + "model": "mistral/mistral-ocr-latest", + "document": { + "type": "image_url", + "image_url": "https://example.com/scan.png", + }, + } + + await handler.async_moderation_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + call_type=CallTypes.aocr.value, + ) + + assert guardrail.event_history == [GuardrailEventHooks.during_call] + assert len(guardrail.apply_calls) == 1 + assert "https://example.com/scan.png" in guardrail.apply_calls[0]["inputs"]["texts"] + + @pytest.mark.asyncio + async def test_post_call_success_hook_guardrails_ocr_output(self): + """ + Verify that async_post_call_success_hook resolves the OCR route + to the OCR handler and applies guardrails to page markdown. + """ + + class TextModifyingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__(guardrail_name="text-modifier") + + def should_run_guardrail(self, data, event_type): # type: ignore[override] + return True + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + texts = inputs.get("texts", []) + return {"texts": [t.replace("SECRET", "[REDACTED]") for t in texts]} + + handler = UnifiedLLMGuardrails() + guardrail = TextModifyingGuardrail() + + ocr_response = OCRResponse( + pages=[ + OCRPage(index=0, markdown="Page 1 has a SECRET value"), + OCRPage(index=1, markdown="Page 2 is clean"), + OCRPage(index=2, markdown="Page 3 also has SECRET data"), + ], + model="mistral/mistral-ocr-latest", + ) + + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + request_route="/v1/ocr", + ) + + data = { + "guardrail_to_apply": guardrail, + "model": "mistral/mistral-ocr-latest", + } + + result = await handler.async_post_call_success_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=ocr_response, + ) + + # Verify the SECRET text was redacted across pages + assert result.pages[0].markdown == "Page 1 has a [REDACTED] value" + assert result.pages[1].markdown == "Page 2 is clean" + assert result.pages[2].markdown == "Page 3 also has [REDACTED] data" + + @pytest.mark.asyncio + async def test_post_call_success_hook_ocr_route_resolves_call_type(self): + """ + Verify that request_route=/v1/ocr correctly resolves to the OCR + call type and the handler is invoked (not skipped). + """ + handler = UnifiedLLMGuardrails() + guardrail = RecordingGuardrail() + + ocr_response = OCRResponse( + pages=[OCRPage(index=0, markdown="Some text")], + model="mistral/mistral-ocr-latest", + ) + + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + request_route="/v1/ocr", + ) + + data = { + "guardrail_to_apply": guardrail, + "model": "mistral/mistral-ocr-latest", + } + + result = await handler.async_post_call_success_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=ocr_response, + ) + + # Guardrail was invoked + assert guardrail.event_history == [GuardrailEventHooks.post_call] + assert len(guardrail.apply_calls) == 1 + assert guardrail.apply_calls[0]["input_type"] == "response" + assert guardrail.apply_calls[0]["inputs"]["texts"] == ["Some text"] + + # Response returned with pages intact + assert result.pages[0].markdown == "Some text"