mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
Add OCR guardrail_translation handler and support
Enables guardrail processing for OCR requests and responses. Adds OCR handler under litellm/llms/mistral/ocr/guardrail_translation/ to process document URLs on input and extracted page markdown on output. Includes route-to-call-type mappings for /ocr and /v1/ocr endpoints. Adds 14 unit tests and 4 e2e tests verifying handler discovery, input/output processing, and integration with UnifiedLLMGuardrails. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Krrish Dholakia
parent
9806e21871
commit
f52412b4bf
7 changed files with 637 additions and 1 deletions
11
litellm/llms/mistral/ocr/guardrail_translation/__init__.py
Normal file
11
litellm/llms/mistral/ocr/guardrail_translation/__init__.py
Normal file
|
|
@ -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"]
|
||||
155
litellm/llms/mistral/ocr/guardrail_translation/handler.py
Normal file
155
litellm/llms/mistral/ocr/guardrail_translation/handler.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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],
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
0
tests/test_litellm/llms/ocr/__init__.py
Normal file
0
tests/test_litellm/llms/ocr/__init__.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue