From 2f7ba0b33366b9e8a626848f4fd10e0d3dda0c9a Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 9 Sep 2026 12:02:22 -0700 Subject: [PATCH] test(ocr): restore guardrail contracts --- .github/workflows/test-rust.yml | 2 + tests/test_litellm_rust/README.md | 2 +- tests/test_litellm_rust/conftest.py | 1 + .../test_litellm_rust/ocr/test_guardrails.py | 74 +++++++++++++++++++ 4 files changed, 78 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm_rust/ocr/test_guardrails.py diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 461722d605c..17b6481a2bf 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -9,6 +9,7 @@ on: - "litellm/integrations/custom_logger.py" - "litellm/litellm_core_utils/litellm_logging.py" - "litellm/litellm_core_utils/logging_worker.py" + - "litellm/proxy/guardrails/**" - "litellm/utils.py" - "litellm/ocr/**" - "litellm/llms/base_llm/ocr/**" @@ -38,6 +39,7 @@ on: - "litellm/integrations/custom_logger.py" - "litellm/litellm_core_utils/litellm_logging.py" - "litellm/litellm_core_utils/logging_worker.py" + - "litellm/proxy/guardrails/**" - "litellm/utils.py" - "litellm/ocr/**" - "litellm/llms/base_llm/ocr/**" diff --git a/tests/test_litellm_rust/README.md b/tests/test_litellm_rust/README.md index 157235f88a8..e1eadbbe5d6 100644 --- a/tests/test_litellm_rust/README.md +++ b/tests/test_litellm_rust/README.md @@ -4,7 +4,7 @@ This suite covers OCR requests through LiteLLM's compiled Rust extension. OCR be A test name identifies the OCR entrypoint or callback under test and its expected observable result. Parameter IDs state the execution mode or credential case. Keep multiple assertions together only when they prove one request, mutation, failure, or callback lifecycle behavior. Record callback observations and assert them after the callback returns because production logging can swallow callback exceptions -`ocr/test_requests.py` covers provider payloads, endpoint and credential resolution, normalized responses, errors, timeouts, and Azure token-provider behavior. `ocr/test_callbacks.py` covers OCR callback inputs, mutations, ordering, context, failure handling, concurrency, and cleanup. `ocr/test_dispatch.py` covers public sync and async native dispatch, explicit Python dispatch, fallback, and the native compression header. `test_ocr.py` is the strict wire-level smoke test +`ocr/test_requests.py` covers provider payloads, endpoint and credential resolution, normalized responses, errors, timeouts, and Azure token-provider behavior. `ocr/test_callbacks.py` covers OCR callback inputs, mutations, ordering, context, failure handling, concurrency, and cleanup. `ocr/test_guardrails.py` covers OCR post-call blocking and response replacement. `ocr/test_dispatch.py` covers public sync and async native dispatch, explicit Python dispatch, fallback, and the native compression header. `test_ocr.py` is the strict wire-level smoke test Run `make test-rust-extension` as the acceptance command. It builds a fresh wheel, installs that wheel into a temporary environment, requires `LITELLM_RUST=1`, and runs this suite with isolated Python imports diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index 4e41ecf4121..c7ee9895aa0 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -32,6 +32,7 @@ CALLBACK_ATTRIBUTES: Final = ( EXPECTED_FAILURE_REASONS: Final = { "ocr/test_callbacks.py": "requires the OCR callback lifecycle implementation from #40070", "ocr/test_dispatch.py": "requires the OCR native dispatch implementation from #40070", + "ocr/test_guardrails.py": "requires the OCR guardrail lifecycle implementation from #40070", "ocr/test_requests.py": "requires the OCR request and Azure authentication implementation from #40070", } diff --git a/tests/test_litellm_rust/ocr/test_guardrails.py b/tests/test_litellm_rust/ocr/test_guardrails.py new file mode 100644 index 00000000000..0a367d33854 --- /dev/null +++ b/tests/test_litellm_rust/ocr/test_guardrails.py @@ -0,0 +1,74 @@ +from typing import Final + +import pytest +from fastapi import HTTPException + +import litellm +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail +from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks +from litellm.types.utils import CallTypes +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.requests import OCR_RESPONSE, call_aocr + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.fixture +def ocr_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=OCR_RESPONSE) + return recording_server + + +class ReplaceOCRMarkdown(CustomGuardrail): + def __init__(self) -> None: + super().__init__( + guardrail_name="replace-ocr-markdown", event_hook=GuardrailEventHooks.post_call, default_on=True + ) + self.call_types: list[CallTypes] = [] + + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + self.call_types.append(call_type) + reviewed_page: Final = response.pages[0].model_copy(update={"markdown": "Reviewed OCR"}) + return response.model_copy(update={"pages": [reviewed_page]}) + + +@pytest.mark.asyncio +async def test_public_aocr_post_call_content_filter_blocks_matching_markdown( + ocr_server: RecordingServer, +) -> None: + guardrail: Final = ContentFilterGuardrail( + guardrail_name="block-native-ocr-markdown", + event_hook=GuardrailEventHooks.post_call, + blocked_words=[BlockedWord(keyword="native OCR response", action=ContentFilterAction.BLOCK)], + ) + litellm.callbacks.append(guardrail) + + with pytest.raises(HTTPException, match="Content blocked") as blocked: + await call_aocr(ocr_server, guardrails=[guardrail.guardrail_name]) + + assert blocked.value.status_code == 400 + assert len(ocr_server.requests) == 1 + + +@pytest.mark.asyncio +async def test_public_aocr_post_call_replacement_reaches_caller_and_success_callback( + ocr_server: RecordingServer, +) -> None: + guardrail: Final = ReplaceOCRMarkdown() + recorder: Final = RecordingLogger() + litellm.callbacks.append(guardrail) + + response: Final = await call_aocr( + ocr_server, + callbacks=[recorder], + guardrails=[guardrail.guardrail_name], + ) + success_events: Final = await recorder.wait_for_async("async_log_success_event") + + assert guardrail.call_types == [CallTypes.aocr] + assert response.pages[0].markdown == "Reviewed OCR" + assert len(success_events) == 1 + assert success_events[0].response.pages[0].markdown == "Reviewed OCR" + assert "guardrails" not in ocr_server.requests[0].body