mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-23 00:41:40 +00:00
63d994ade4
Extracted from #41733 without the router loop, the cache machine layer, streaming, or the error, timeout and route-pruning work that moved to #41745 litellm-callbacks holds the contract a native call and its host share: Machine, HostOp, CallEvent, the in-process run loop, and Passthrough, which is built only by comparing the caller's inputs with the body the route sends, so a route can never mark a key it rewrote. litellm-host-python (formerly python-interop) owns the CPython driver and the Execution handle, and litellm-callbacks-legacy is the @client wrapper as the native call sees it: function_setup, the deployment hooks, pre_call and post_call, the success and failure fan-out and the deferred proxy release. OCR is the one route on it, and the old core and bridge lifecycles are gone The passthrough rule is the structural fix for the bug #41719 patched in core and #41716 reworks: an inlined remote document no longer counts as the caller's value, so the legacy adapter never hands the caller's URL back into the body. core/tests/ocr/passthrough.rs pins it for every route and document source, including that unchanged values stay passthrough, and callbacks-legacy/tests/payload.rs pins the adapter side with a real pre_call callback Python OCR integration tests that only exercised core behavior now live as Rust tests, so tests/test_litellm_rust keeps the cases that need the full Python stack
74 lines
2.9 KiB
Python
74 lines
2.9 KiB
Python
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_native, call_native_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_native_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_native_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_native_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_native_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
|