From f1a5a6f47a61771fb43be23fd6fc1e048c0e7be5 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 7 Sep 2026 21:32:30 -0700 Subject: [PATCH] test(rust): cover messages callbacks and streaming --- tests/test_litellm_rust/callback_recorder.py | 110 +++++++++ tests/test_litellm_rust/conftest.py | 8 +- ...ocr_test_server.py => recording_server.py} | 33 ++- tests/test_litellm_rust/test_messages.py | 131 ++++++++++ .../test_messages_callbacks.py | 224 ++++++++++++++++++ tests/test_litellm_rust/test_ocr.py | 47 ++-- tests/test_litellm_rust/test_ocr_callbacks.py | 84 +++---- 7 files changed, 562 insertions(+), 75 deletions(-) create mode 100644 tests/test_litellm_rust/callback_recorder.py rename tests/test_litellm_rust/{ocr_test_server.py => recording_server.py} (74%) create mode 100644 tests/test_litellm_rust/test_messages.py create mode 100644 tests/test_litellm_rust/test_messages_callbacks.py diff --git a/tests/test_litellm_rust/callback_recorder.py b/tests/test_litellm_rust/callback_recorder.py new file mode 100644 index 00000000000..6924a765ea6 --- /dev/null +++ b/tests/test_litellm_rust/callback_recorder.py @@ -0,0 +1,110 @@ +import asyncio +import copy +import threading +import time +from dataclasses import dataclass +from typing import Final + +from litellm.integrations.custom_logger import CustomLogger + + +@dataclass(frozen=True, slots=True) +class HookEvent: + name: str + call_type: str | None + stream: bool | None + thread: threading.Thread + loop: asyncio.AbstractEventLoop | None + has_running_loop: bool + kwargs: object + response: object + + +class RecordingLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self._events: list[HookEvent] = [] + self._condition = threading.Condition() + + @property + def events(self) -> tuple[HookEvent, ...]: + with self._condition: + return tuple(self._events) + + @property + def names(self) -> tuple[str, ...]: + return tuple(event.name for event in self.events) + + def _record(self, name: str, kwargs: object = None, response: object = None) -> None: + details: Final = kwargs if isinstance(kwargs, dict) else {} + try: + loop: Final = asyncio.get_running_loop() + has_running_loop: Final = True + except RuntimeError: + loop = None + has_running_loop = False + event: Final = HookEvent( + name=name, + call_type=details.get("call_type"), + stream=details.get("stream"), + thread=threading.current_thread(), + loop=loop, + has_running_loop=has_running_loop, + kwargs=copy.deepcopy(details), + response=response, + ) + with self._condition: + self._events.append(event) + self._condition.notify_all() + + def wait_for(self, name: str, count: int = 1, timeout: float = 10) -> tuple[HookEvent, ...]: + deadline: Final = time.monotonic() + timeout + with self._condition: + while sum(event.name == name for event in self._events) < count: + remaining: Final = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"Timed out waiting for {count} {name} events; saw {self.names}") + self._condition.wait(remaining) + return tuple(event for event in self._events if event.name == name) + + async def wait_for_async(self, name: str, count: int = 1, timeout: float = 10) -> tuple[HookEvent, ...]: + return await asyncio.wait_for(asyncio.to_thread(self.wait_for, name, count, timeout), timeout=timeout + 1) + + def log_pre_api_call(self, model, messages, kwargs): + self._record("log_pre_api_call", kwargs) + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + self._record("log_success_event", kwargs, response_obj) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self._record("async_log_success_event", kwargs, response_obj) + + def log_stream_event(self, kwargs, response_obj, start_time, end_time): + self._record("log_stream_event", kwargs, response_obj) + + async def async_log_stream_event(self, kwargs, response_obj, start_time, end_time): + self._record("async_log_stream_event", kwargs, response_obj) + + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + self._record("log_failure_event", kwargs, response_obj) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + self._record("async_log_failure_event", kwargs, response_obj) + + def logging_hook(self, kwargs, result, call_type): + self._record("logging_hook", kwargs, result) + return kwargs, result + + async def async_logging_hook(self, kwargs, result, call_type): + self._record("async_logging_hook", kwargs, result) + return kwargs, result + + async def async_pre_call_deployment_hook(self, kwargs, call_type): + self._record("async_pre_call_deployment_hook", kwargs) + + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + self._record("async_post_call_success_deployment_hook", request_data, response) + return response + + async def async_post_call_failure_deployment_hook(self, request_data, exception, call_type, fallback_depth=None): + self._record("async_post_call_failure_deployment_hook", request_data, exception) diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index 353f3cc9a5b..e1fab29e1b9 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -9,11 +9,11 @@ import litellm from litellm.rust_bridge import ocr as native_ocr from litellm.rust_bridge.configuration import reset_rust_configuration from litellm.rust_bridge.configuration import rust_enabled -from tests.test_litellm_rust.ocr_test_server import ocr_server # noqa: F401 # pytest fixture export +from tests.test_litellm_rust.recording_server import recording_server # noqa: F401 # pytest fixture export @pytest.fixture(autouse=True) -def isolate_rust_ocr_state(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: +def isolate_rust_state(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: callback_attributes: Final = ( "callbacks", "input_callback", @@ -27,7 +27,7 @@ def isolate_rust_ocr_state(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: original_cache: Final = litellm.cache for attribute in callback_attributes: getattr(litellm, attribute).clear() - litellm.cache = None + litellm.cache = None # test-quality-ok: isolate the process-global cache from native extension tests reset_rust_configuration() litellm.rust(True) python_ocr: Final = litellm.ocr @@ -58,7 +58,7 @@ def isolate_rust_ocr_state(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: target = getattr(litellm, attribute) target.clear() target.extend(callbacks) - litellm.cache = original_cache + litellm.cache = original_cache # test-quality-ok: restore the process-global cache after native extension tests reset_rust_configuration() diff --git a/tests/test_litellm_rust/ocr_test_server.py b/tests/test_litellm_rust/recording_server.py similarity index 74% rename from tests/test_litellm_rust/ocr_test_server.py rename to tests/test_litellm_rust/recording_server.py index 49f7f07098a..6ecaf25f3b7 100644 --- a/tests/test_litellm_rust/ocr_test_server.py +++ b/tests/test_litellm_rust/recording_server.py @@ -1,4 +1,5 @@ import json +import copy import threading import time from collections.abc import Iterator @@ -8,12 +9,6 @@ from typing import Final import pytest -OCR_RESPONSE: Final = { - "pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}], - "model": "mistral-ocr-latest", - "usage_info": {"pages_processed": 1, "doc_size_bytes": 3}, -} - @dataclass class RecordedRequest: @@ -25,17 +20,19 @@ class RecordedRequest: @dataclass class ResponseSpec: + body: object status: int = 200 - body: object = field(default_factory=lambda: dict(OCR_RESPONSE)) headers: dict[str, str] = field(default_factory=dict) delay: float = 0 + events: tuple[tuple[str, object], ...] = () @dataclass -class OCRTestServer: +class RecordingServer: server: ThreadingHTTPServer requests: list[RecordedRequest] responses: list[ResponseSpec] + default_response: ResponseSpec @property def base_url(self) -> str: @@ -47,7 +44,7 @@ class OCRTestServer: @pytest.fixture -def ocr_server() -> Iterator[OCRTestServer]: +def recording_server() -> Iterator[RecordingServer]: requests: list[RecordedRequest] = [] responses: list[ResponseSpec] = [] @@ -64,12 +61,16 @@ def ocr_server() -> Iterator[OCRTestServer]: body=body, ) ) - response: Final = responses.pop(0) if responses else ResponseSpec() + response: Final = responses.pop(0) if responses else copy.deepcopy(recording_server.default_response) if response.delay: time.sleep(response.delay) - payload: Final = json.dumps(response.body).encode() + payload: Final = ( + b"".join(f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in response.events) + if response.events + else json.dumps(response.body).encode() + ) self.send_response(response.status) - self.send_header("Content-Type", "application/json") + self.send_header("Content-Type", "text/event-stream" if response.events else "application/json") self.send_header("Content-Length", str(len(payload))) for name, value in response.headers.items(): self.send_header(name, value) @@ -89,7 +90,13 @@ def ocr_server() -> Iterator[OCRTestServer]: thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True) thread.start() try: - yield OCRTestServer(server=server, requests=requests, responses=responses) + recording_server = RecordingServer( + server=server, + requests=requests, + responses=responses, + default_response=ResponseSpec(body={}), + ) + yield recording_server finally: server.shutdown() server.server_close() diff --git a/tests/test_litellm_rust/test_messages.py b/tests/test_litellm_rust/test_messages.py new file mode 100644 index 00000000000..f2ce7e135f3 --- /dev/null +++ b/tests/test_litellm_rust/test_messages.py @@ -0,0 +1,131 @@ +import json +from collections.abc import AsyncIterator +from typing import Final, cast + +import pytest + +import litellm +from tests.test_litellm_rust.recording_server import RecordingServer, ResponseSpec + +pytestmark = pytest.mark.requires_rust_extension + +MODEL: Final = "anthropic/claude-sonnet-4-5-20250929" +MESSAGES: Final = [{"role": "user", "content": "Hello"}] +MESSAGES_RESPONSE: Final = { + "id": "msg_native", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [{"type": "text", "text": "Hello from native Messages"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 4}, +} + + +@pytest.fixture +def messages_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) + return recording_server + + +async def call_messages(server: RecordingServer, **kwargs: object): + return await litellm.anthropic.messages.acreate( + model=MODEL, + messages=MESSAGES, + max_tokens=64, + api_key="test-key", + api_base=server.base_url, + **kwargs, + ) + + +def assert_native_request(server: RecordingServer) -> None: + assert len(server.requests) == 1 + assert "accept-encoding" not in server.requests[0].headers + + +@pytest.mark.asyncio +async def test_messages_sends_expected_provider_request(messages_server: RecordingServer) -> None: + response: Final = await call_messages(messages_server) + + assert response["content"] == [{"type": "text", "text": "Hello from native Messages"}] + assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"} + assert_native_request(messages_server) + request: Final = messages_server.requests[0] + assert request.path == "/v1/messages" + assert request.headers["x-api-key"] == "test-key" + assert request.headers["anthropic-version"] == "2023-06-01" + assert request.body == { + "model": "claude-sonnet-4-5-20250929", + "messages": MESSAGES, + "max_tokens": 64, + } + + +@pytest.mark.asyncio +async def test_messages_sends_custom_headers(messages_server: RecordingServer) -> None: + await call_messages(messages_server, extra_headers={"x-trace-id": "trace-1"}) + + assert messages_server.requests[0].headers["x-trace-id"] == "trace-1" + + +@pytest.mark.asyncio +async def test_messages_resolves_provider_credentials( + messages_server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("ANTHROPIC_API_KEY", "environment-key") + + await litellm.anthropic.messages.acreate( + model=MODEL, + messages=MESSAGES, + max_tokens=64, + api_base=messages_server.base_url, + ) + + assert messages_server.requests[0].headers["x-api-key"] == "environment-key" + + +@pytest.mark.asyncio +async def test_messages_explicit_credentials_override_defaults( + messages_server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("ANTHROPIC_API_KEY", "environment-key") + + await call_messages(messages_server) + + assert messages_server.requests[0].headers["x-api-key"] == "test-key" + + +@pytest.mark.asyncio +async def test_azure_messages_uses_foundry_endpoint_and_credentials(messages_server: RecordingServer) -> None: + await litellm.anthropic.messages.acreate( + model="azure_ai/claude-opus-4.5", + messages=MESSAGES, + max_tokens=64, + api_key="azure-key", + api_base=f"{messages_server.base_url}/anthropic", + ) + + assert_native_request(messages_server) + assert messages_server.requests[0].path == "/anthropic/v1/messages" + assert messages_server.requests[0].headers["x-api-key"] == "azure-key" + + +@pytest.mark.asyncio +async def test_messages_stream_yields_anthropic_events(messages_server: RecordingServer) -> None: + stream: Final = cast(AsyncIterator[bytes], await call_messages(messages_server, stream=True)) + payload: Final = b"".join([chunk async for chunk in stream]) + + assert_native_request(messages_server) + assert "stream" not in messages_server.requests[0].body + assert b"event: message_start" in payload + assert b"event: content_block_delta" in payload + assert b"Hello from native Messages" in payload + assert b"event: message_stop" in payload + message_delta: Final = next( + json.loads(block.split(b"data: ", 1)[1]) + for block in payload.split(b"\n\n") + if block.startswith(b"event: message_delta") + ) + assert message_delta["usage"] == {"input_tokens": 5, "output_tokens": 4} diff --git a/tests/test_litellm_rust/test_messages_callbacks.py b/tests/test_litellm_rust/test_messages_callbacks.py new file mode 100644 index 00000000000..7e6c910d500 --- /dev/null +++ b/tests/test_litellm_rust/test_messages_callbacks.py @@ -0,0 +1,224 @@ +import asyncio +import copy +import json +import threading +from typing import Final + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from tests.test_litellm_rust.callback_recorder import RecordingLogger +from tests.test_litellm_rust.recording_server import RecordingServer, ResponseSpec + +pytestmark = pytest.mark.requires_rust_extension + +MODEL: Final = "anthropic/claude-sonnet-4-5-20250929" +MESSAGES: Final = [{"role": "user", "content": "Hello"}] +MESSAGES_RESPONSE: Final = { + "id": "msg_native", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [{"type": "text", "text": "Hello from native Messages"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 4}, +} + + +@pytest.fixture +def messages_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) + return recording_server + + +async def call_messages(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object): + return await litellm.anthropic.messages.acreate( + model=MODEL, + messages=MESSAGES, + max_tokens=64, + api_key="test-key", + api_base=server.base_url, + callbacks=callbacks, + **kwargs, + ) + + +def request_body(kwargs: dict) -> dict: + return kwargs["additional_args"]["complete_input_dict"] + + +def request_headers(kwargs: dict) -> dict: + return kwargs["additional_args"]["headers"] + + +@pytest.mark.asyncio +@pytest.mark.xfail(strict=True, reason="UC-MSG-PRECALL-VIEW: native pre-call arguments differ from legacy") +async def test_messages_pre_call_receives_expected_provider_request(messages_server: RecordingServer) -> None: + observations: Final = [] + + class Observe(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observations.append((model, messages, copy.deepcopy(kwargs["additional_args"]))) + + await call_messages(messages_server, [Observe()]) + + assert len(observations) == 1 + model, messages, additional_args = observations[0] + assert model == "claude-sonnet-4-5-20250929" + assert messages == [ + { + "role": "user", + "content": json.dumps( + { + "model": "claude-sonnet-4-5-20250929", + "messages": MESSAGES, + "max_tokens": 64, + } + ), + } + ] + assert additional_args["api_base"] == f"{messages_server.base_url}/v1/messages" + assert additional_args["complete_input_dict"] == { + "model": "claude-sonnet-4-5-20250929", + "messages": MESSAGES, + "max_tokens": 64, + } + assert additional_args["headers"]["x-api-key"] == "test-key" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("raise_after_edit", [False, True]) +@pytest.mark.xfail( + strict=True, + reason="UC-MSG-PRECALL-MUTATION: native transport snapshots body and headers before pre-call", +) +async def test_messages_pre_call_edits_reach_later_callbacks_and_provider( + messages_server: RecordingServer, raise_after_edit: bool +) -> None: + observed: Final = [] + + class Edit(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + request_body(kwargs)["temperature"] = 0.25 + request_headers(kwargs)["x-audit-tag"] = "reviewed" + if raise_after_edit: + raise RuntimeError("audit exporter unavailable") + + class Observe(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observed.append((copy.deepcopy(request_body(kwargs)), dict(request_headers(kwargs)))) + + await call_messages(messages_server, [Edit(), Observe()]) + + assert observed[0][0]["temperature"] == 0.25 + assert observed[0][1]["x-audit-tag"] == "reviewed" + assert messages_server.requests[0].body["temperature"] == 0.25 + assert messages_server.requests[0].headers["x-audit-tag"] == "reviewed" + + +@pytest.mark.asyncio +async def test_messages_pre_call_rebinding_does_not_replace_inflight_request( + messages_server: RecordingServer, +) -> None: + observed: Final = [] + + class Rebind(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + kwargs["additional_args"]["complete_input_dict"] = {"replacement": True} + + class Observe(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observed.append(request_body(kwargs)) + + await call_messages(messages_server, [Rebind(), Observe()]) + + assert observed == [{"replacement": True}] + assert messages_server.requests[0].body["messages"] == MESSAGES + + +@pytest.mark.asyncio +async def test_messages_pre_call_state_reaches_terminal_callbacks(messages_server: RecordingServer) -> None: + token: Final = object() + observed: Final = [] + finished: Final = asyncio.Event() + + class Stash(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + kwargs["test-token"] = token + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + observed.append(kwargs["test-token"]) + finished.set() + + await call_messages(messages_server, [Stash()]) + await asyncio.wait_for(finished.wait(), timeout=10) + + assert observed == [token] + + +@pytest.mark.asyncio +async def test_messages_callbacks_run_once(messages_server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + await call_messages(messages_server, [recorder]) + await recorder.wait_for_async("async_log_success_event") + + assert recorder.names.count("log_pre_api_call") == 1 + assert recorder.names.count("async_logging_hook") == 1 + assert recorder.names.count("async_log_success_event") == 1 + + +@pytest.mark.asyncio +async def test_messages_success_callbacks_receive_expected_context(messages_server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + await call_messages( + messages_server, + [recorder], + litellm_call_id="messages-success", + metadata={"source": "callback-test"}, + ) + async_events: Final = await recorder.wait_for_async("async_log_success_event") + + assert len(async_events) == 1 + event: Final = async_events[0] + assert event.call_type == "anthropic_messages" + assert event.kwargs["litellm_call_id"] == "messages-success" + assert event.kwargs["litellm_params"]["metadata"]["source"] == "callback-test" + assert event.response.choices[0].message.content == "Hello from native Messages" + + +@pytest.mark.asyncio +async def test_messages_pre_call_runs_in_callers_execution_context(messages_server: RecordingServer) -> None: + caller_thread: Final = threading.current_thread() + observations: Final = [] + + class Observe(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + observations.append((asyncio.get_running_loop(), threading.current_thread())) + + caller_loop: Final = asyncio.get_running_loop() + await call_messages(messages_server, [Observe()]) + + assert observations == [(caller_loop, caller_thread)] + + +@pytest.mark.asyncio +@pytest.mark.xfail( + strict=True, + reason="UC-MSG-STREAM-COMPLETION: native fake stream logs before assembled stream finalization", +) +async def test_messages_stream_logs_success_after_exhaustion(messages_server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + stream: Final = await call_messages(messages_server, [recorder], stream=True) + + assert "async_log_success_event" not in recorder.names + chunks: Final = [chunk async for chunk in stream] + events: Final = await recorder.wait_for_async("async_log_success_event") + + assert chunks + assert len(events) == 1 + assert "async_log_stream_event" not in recorder.names + assert events[0].kwargs["complete_streaming_response"] is not None diff --git a/tests/test_litellm_rust/test_ocr.py b/tests/test_litellm_rust/test_ocr.py index 647e7165650..88cddb38a80 100644 --- a/tests/test_litellm_rust/test_ocr.py +++ b/tests/test_litellm_rust/test_ocr.py @@ -5,15 +5,26 @@ import pytest import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse -from tests.test_litellm_rust.ocr_test_server import OCRTestServer, ResponseSpec +from tests.test_litellm_rust.recording_server import RecordingServer, ResponseSpec pytestmark = pytest.mark.requires_rust_extension DOCUMENT: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"} MODEL: Final = "mistral/mistral-ocr-latest" +OCR_RESPONSE: Final = { + "pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1, "doc_size_bytes": 3}, +} -def call_ocr(server: OCRTestServer, **kwargs: object) -> OCRResponse: +@pytest.fixture +def ocr_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=OCR_RESPONSE) + return recording_server + + +def call_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse: return litellm.ocr( model=MODEL, document=dict(DOCUMENT), @@ -23,12 +34,12 @@ def call_ocr(server: OCRTestServer, **kwargs: object) -> OCRResponse: ) -def assert_native_request(server: OCRTestServer) -> None: +def assert_native_request(server: RecordingServer) -> None: assert len(server.requests) == 1 assert not server.requests[0].headers.get("user-agent", "").startswith("python-httpx") -def test_ocr_sends_expected_provider_request(ocr_server: OCRTestServer) -> None: +def test_ocr_sends_expected_provider_request(ocr_server: RecordingServer) -> None: response: Final = call_ocr(ocr_server) assert response.pages[0].markdown == "native OCR response" @@ -37,7 +48,7 @@ def test_ocr_sends_expected_provider_request(ocr_server: OCRTestServer) -> None: assert ocr_server.requests[0].body == {"model": "mistral-ocr-latest", "document": DOCUMENT} -def test_ocr_rejects_unsupported_file_document_before_callbacks(ocr_server: OCRTestServer) -> None: +def test_ocr_rejects_unsupported_file_document_before_callbacks(ocr_server: RecordingServer) -> None: with pytest.raises(NotImplementedError, match="OCR file document preparation"): litellm.ocr( model=MODEL, @@ -49,21 +60,21 @@ def test_ocr_rejects_unsupported_file_document_before_callbacks(ocr_server: OCRT assert ocr_server.requests == [] -def test_ocr_sends_optional_parameters(ocr_server: OCRTestServer) -> None: +def test_ocr_sends_optional_parameters(ocr_server: RecordingServer) -> None: call_ocr(ocr_server, pages=[0, 2], include_image_base64=True) assert ocr_server.requests[0].body["pages"] == [0, 2] assert ocr_server.requests[0].body["include_image_base64"] is True -def test_ocr_sends_custom_headers(ocr_server: OCRTestServer) -> None: +def test_ocr_sends_custom_headers(ocr_server: RecordingServer) -> None: call_ocr(ocr_server, extra_headers={"x-trace-id": "trace-1"}) assert ocr_server.requests[0].headers["authorization"] == "Bearer test-key" assert ocr_server.requests[0].headers["x-trace-id"] == "trace-1" -def test_ocr_resolves_provider_credentials(ocr_server: OCRTestServer, monkeypatch: pytest.MonkeyPatch) -> None: +def test_ocr_resolves_provider_credentials(ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("MISTRAL_API_KEY", "environment-key") litellm.ocr(model=MODEL, document=DOCUMENT, api_base=ocr_server.base_url) @@ -71,7 +82,9 @@ def test_ocr_resolves_provider_credentials(ocr_server: OCRTestServer, monkeypatc assert ocr_server.requests[0].headers["authorization"] == "Bearer environment-key" -def test_ocr_explicit_credentials_override_defaults(ocr_server: OCRTestServer, monkeypatch: pytest.MonkeyPatch) -> None: +def test_ocr_explicit_credentials_override_defaults( + ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: monkeypatch.setenv("MISTRAL_API_KEY", "environment-key") call_ocr(ocr_server) @@ -79,7 +92,7 @@ def test_ocr_explicit_credentials_override_defaults(ocr_server: OCRTestServer, m assert ocr_server.requests[0].headers["authorization"] == "Bearer test-key" -def test_ocr_resolves_provider_endpoint(ocr_server: OCRTestServer, monkeypatch: pytest.MonkeyPatch) -> None: +def test_ocr_resolves_provider_endpoint(ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("AZURE_AI_API_KEY", "azure-key") monkeypatch.setenv("AZURE_AI_API_BASE", ocr_server.base_url) @@ -90,7 +103,7 @@ def test_ocr_resolves_provider_endpoint(ocr_server: OCRTestServer, monkeypatch: assert ocr_server.requests[0].headers["api-key"] == "azure-key" -def test_ocr_resolves_vertex_project_and_location(ocr_server: OCRTestServer) -> None: +def test_ocr_resolves_vertex_project_and_location(ocr_server: RecordingServer) -> None: litellm.ocr( model="vertex_ai/mistral-ocr-2505", document=DOCUMENT, @@ -106,7 +119,7 @@ def test_ocr_resolves_vertex_project_and_location(ocr_server: OCRTestServer) -> ) -def test_ocr_returns_normalized_response(ocr_server: OCRTestServer) -> None: +def test_ocr_returns_normalized_response(ocr_server: RecordingServer) -> None: response: Final = call_ocr(ocr_server) assert isinstance(response, OCRResponse) @@ -114,8 +127,8 @@ def test_ocr_returns_normalized_response(ocr_server: OCRTestServer) -> None: assert response.usage_info.pages_processed == 1 -def test_ocr_provider_error_preserves_status_and_context(ocr_server: OCRTestServer) -> None: - ocr_server.enqueue(ResponseSpec(status=400, body={"message": "invalid OCR request"})) +def test_ocr_provider_error_preserves_status_and_context(ocr_server: RecordingServer) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "invalid OCR request"}, status=400)) with pytest.raises(litellm.BadRequestError) as caught: call_ocr(ocr_server) @@ -126,8 +139,8 @@ def test_ocr_provider_error_preserves_status_and_context(ocr_server: OCRTestServ assert "invalid OCR request" not in str(caught.value) -def test_ocr_honors_request_timeout(ocr_server: OCRTestServer) -> None: - ocr_server.enqueue(ResponseSpec(delay=0.2)) +def test_ocr_honors_request_timeout(ocr_server: RecordingServer) -> None: + ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.2)) started_at: Final = time.monotonic() with pytest.raises(RuntimeError, match="OCR transport failed"): @@ -137,7 +150,7 @@ def test_ocr_honors_request_timeout(ocr_server: OCRTestServer) -> None: assert len(ocr_server.requests) == 1 -def test_ocr_respects_runtime_toggle(ocr_server: OCRTestServer) -> None: +def test_ocr_respects_runtime_toggle(ocr_server: RecordingServer) -> None: litellm.rust(False) call_ocr(ocr_server) litellm.rust(True) diff --git a/tests/test_litellm_rust/test_ocr_callbacks.py b/tests/test_litellm_rust/test_ocr_callbacks.py index 55d702a18f4..c5ddc618794 100644 --- a/tests/test_litellm_rust/test_ocr_callbacks.py +++ b/tests/test_litellm_rust/test_ocr_callbacks.py @@ -9,15 +9,27 @@ import pytest import litellm from litellm.integrations.custom_logger import CustomLogger -from tests.test_litellm_rust.ocr_test_server import OCRTestServer, ResponseSpec +from tests.test_litellm_rust.callback_recorder import RecordingLogger +from tests.test_litellm_rust.recording_server import RecordingServer, ResponseSpec pytestmark = pytest.mark.requires_rust_extension DOCUMENT: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"} MODEL: Final = "mistral/mistral-ocr-latest" +OCR_RESPONSE: Final = { + "pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1, "doc_size_bytes": 3}, +} -def call_ocr(server: OCRTestServer, callbacks: list[CustomLogger], **kwargs: object): +@pytest.fixture +def ocr_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=OCR_RESPONSE) + return recording_server + + +def call_ocr(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object): return litellm.ocr( model=MODEL, document=dict(DOCUMENT), @@ -28,7 +40,7 @@ def call_ocr(server: OCRTestServer, callbacks: list[CustomLogger], **kwargs: obj ) -async def call_aocr(server: OCRTestServer, callbacks: list[CustomLogger], **kwargs: object): +async def call_aocr(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object): return await litellm.aocr( model=MODEL, document=dict(DOCUMENT), @@ -47,7 +59,7 @@ def request_headers(kwargs: dict) -> dict: return kwargs["additional_args"]["headers"] -def test_pre_call_receives_expected_provider_request(ocr_server: OCRTestServer) -> None: +def test_pre_call_receives_expected_provider_request(ocr_server: RecordingServer) -> None: observations: Final = [] class Observe(CustomLogger): @@ -70,7 +82,7 @@ def test_pre_call_receives_expected_provider_request(ocr_server: OCRTestServer) @pytest.mark.parametrize("raise_after_edit", [False, True]) def test_pre_call_body_edits_reach_later_callbacks_and_provider( - ocr_server: OCRTestServer, raise_after_edit: bool + ocr_server: RecordingServer, raise_after_edit: bool ) -> None: observed: Final = [] @@ -90,7 +102,7 @@ def test_pre_call_body_edits_reach_later_callbacks_and_provider( assert ocr_server.requests[0].body["include_image_base64"] is True -def test_pre_call_header_edits_reach_later_callbacks_and_provider(ocr_server: OCRTestServer) -> None: +def test_pre_call_header_edits_reach_later_callbacks_and_provider(ocr_server: RecordingServer) -> None: observed: Final = [] class Edit(CustomLogger): @@ -107,7 +119,7 @@ def test_pre_call_header_edits_reach_later_callbacks_and_provider(ocr_server: OC assert ocr_server.requests[0].headers["x-audit-tag"] == "reviewed" -def test_pre_call_nested_mutation_updates_retained_references(ocr_server: OCRTestServer) -> None: +def test_pre_call_nested_mutation_updates_retained_references(ocr_server: RecordingServer) -> None: original: Final = dict(DOCUMENT) replacement_url: Final = "data:application/pdf;base64,ZGVm" retained: Final = [] @@ -133,7 +145,7 @@ def test_pre_call_nested_mutation_updates_retained_references(ocr_server: OCRTes assert ocr_server.requests[0].body["document"]["document_url"] == replacement_url -def test_pre_call_field_replacement_preserves_original_references(ocr_server: OCRTestServer) -> None: +def test_pre_call_field_replacement_preserves_original_references(ocr_server: RecordingServer) -> None: original: Final = dict(DOCUMENT) replacement: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,ZGVm"} retained: Final = [] @@ -157,7 +169,7 @@ def test_pre_call_field_replacement_preserves_original_references(ocr_server: OC assert ocr_server.requests[0].body["document"] == replacement -def test_pre_call_body_rebinding_does_not_replace_inflight_request(ocr_server: OCRTestServer) -> None: +def test_pre_call_body_rebinding_does_not_replace_inflight_request(ocr_server: RecordingServer) -> None: observed: Final = [] class Rebind(CustomLogger): @@ -174,7 +186,7 @@ def test_pre_call_body_rebinding_does_not_replace_inflight_request(ocr_server: O assert ocr_server.requests[0].body == {"model": "mistral-ocr-latest", "document": DOCUMENT} -def test_queued_payload_observes_later_callback_mutations(ocr_server: OCRTestServer) -> None: +def test_queued_payload_observes_later_callback_mutations(ocr_server: RecordingServer) -> None: queued: Final = [] class QueuePayload(CustomLogger): @@ -190,7 +202,7 @@ def test_queued_payload_observes_later_callback_mutations(ocr_server: OCRTestSer assert queued[0]["queued-edit"] is True -def test_callback_copies_preserve_expected_sharing(ocr_server: OCRTestServer) -> None: +def test_callback_copies_preserve_expected_sharing(ocr_server: RecordingServer) -> None: copies: Final = {} replacement_url: Final = "data:application/pdf;base64,ZGVm" @@ -212,7 +224,7 @@ def test_callback_copies_preserve_expected_sharing(ocr_server: OCRTestServer) -> assert json.loads(copies["serialized"])["document"]["document_url"] == "data:application/pdf;base64,YWJj" -def test_pre_call_state_reaches_terminal_callbacks(ocr_server: OCRTestServer) -> None: +def test_pre_call_state_reaches_terminal_callbacks(ocr_server: RecordingServer) -> None: token: Final = object() terminal_tokens: queue.SimpleQueue[object] = queue.SimpleQueue() finished: Final = threading.Event() @@ -232,36 +244,27 @@ def test_pre_call_state_reaches_terminal_callbacks(ocr_server: OCRTestServer) -> @pytest.mark.asyncio -async def test_success_callbacks_receive_expected_context_and_response(ocr_server: OCRTestServer) -> None: - observations: Final = [] - finished: Final = asyncio.Event() - - class Observe(CustomLogger): - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - observations.append( - ( - kwargs["call_type"], - kwargs["litellm_call_id"], - kwargs["litellm_params"]["metadata"]["source"], - response_obj.pages[0].markdown, - ) - ) - finished.set() +async def test_success_callbacks_receive_expected_context_and_response(ocr_server: RecordingServer) -> None: + recorder: Final = RecordingLogger() await call_aocr( ocr_server, - [Observe()], + [recorder], litellm_call_id="ocr-success", metadata={"source": "callback-test"}, ) - await asyncio.wait_for(finished.wait(), timeout=10) + events: Final = await recorder.wait_for_async("async_log_success_event") - assert observations == [("aocr", "ocr-success", "callback-test", "native OCR response")] + assert len(events) == 1 + assert events[0].call_type == "aocr" + assert events[0].kwargs["litellm_call_id"] == "ocr-success" + assert events[0].kwargs["litellm_params"]["metadata"]["source"] == "callback-test" + assert events[0].response.pages[0].markdown == "native OCR response" @pytest.mark.asyncio -async def test_failure_callbacks_receive_expected_context_and_error(ocr_server: OCRTestServer) -> None: - ocr_server.enqueue(ResponseSpec(status=500, body={"message": "provider unavailable"})) +async def test_failure_callbacks_receive_expected_context_and_error(ocr_server: RecordingServer) -> None: + ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500)) observations: Final = [] class Observe(CustomLogger): @@ -280,7 +283,7 @@ async def test_failure_callbacks_receive_expected_context_and_error(ocr_server: assert all(observation[3] is None for observation in observations) -def test_background_callback_can_mutate_retained_state_after_return(ocr_server: OCRTestServer) -> None: +def test_background_callback_can_mutate_retained_state_after_return(ocr_server: RecordingServer) -> None: release: Final = threading.Event() finished: Final = threading.Event() retained: Final = [] @@ -306,15 +309,14 @@ def test_background_callback_can_mutate_retained_state_after_return(ocr_server: @pytest.mark.asyncio -async def test_pre_call_runs_in_callers_execution_context(ocr_server: OCRTestServer) -> None: +async def test_pre_call_runs_in_callers_execution_context(ocr_server: RecordingServer) -> None: caller_loop: Final = asyncio.get_running_loop() caller_thread: Final = threading.current_thread() - observations: Final = [] + recorder: Final = RecordingLogger() - class Observe(CustomLogger): - def log_pre_api_call(self, model, messages, kwargs): - observations.append((asyncio.get_running_loop(), threading.current_thread())) + await call_aocr(ocr_server, [recorder]) - await call_aocr(ocr_server, [Observe()]) - - assert observations == [(caller_loop, caller_thread)] + events: Final = recorder.wait_for("log_pre_api_call") + assert len(events) == 1 + assert events[0].loop is caller_loop + assert events[0].thread is caller_thread