mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-20 00:11:50 +00:00
test(rust): cover messages callbacks and streaming
This commit is contained in:
Yujong Lee
parent
969fa34968
commit
f1a5a6f47a
7 changed files with 562 additions and 75 deletions
110
tests/test_litellm_rust/callback_recorder.py
Normal file
110
tests/test_litellm_rust/callback_recorder.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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()
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
131
tests/test_litellm_rust/test_messages.py
Normal file
131
tests/test_litellm_rust/test_messages.py
Normal file
|
|
@ -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}
|
||||
224
tests/test_litellm_rust/test_messages_callbacks.py
Normal file
224
tests/test_litellm_rust/test_messages_callbacks.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue