scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

test(rust): clarify retained callback contracts

Browse source
This commit is contained in:
Yujong Lee 2026-09-09 10:51:52 -07:00
parent d701b818f9
commit 348d0aea66
13 changed files with 755 additions and 509 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
.github/workflows/test-rust.yml vendored
View file

@ -4,9 +4,23 @@ on:
push:
paths:
- "litellm-rust/**"
- "litellm/rust_bridge/**"
- "tests/test_litellm_rust/**"
- "litellm/integrations/**"
- "litellm/litellm_core_utils/litellm_logging.py"
- "litellm/litellm_core_utils/logging_worker.py"
- "litellm/proxy/guardrails/**"
- "litellm/proxy/common_request_processing.py"
- "litellm/utils.py"
- "litellm/ocr/**"
- "litellm/llms/base_llm/ocr/**"
- "litellm/llms/custom_httpx/llm_http_handler.py"
- "tests/test_litellm/ocr/**"
- "tests/test_litellm/conftest.py"
- "Makefile"
- ".cargo/**"
- "pyproject.toml"
- "uv.lock"
- "rust-toolchain.toml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"
@ -21,9 +35,23 @@ on:
- "litellm_**"
paths:
- "litellm-rust/**"
- "litellm/rust_bridge/**"
- "tests/test_litellm_rust/**"
- "litellm/integrations/**"
- "litellm/litellm_core_utils/litellm_logging.py"
- "litellm/litellm_core_utils/logging_worker.py"
- "litellm/proxy/guardrails/**"
- "litellm/proxy/common_request_processing.py"
- "litellm/utils.py"
- "litellm/ocr/**"
- "litellm/llms/base_llm/ocr/**"
- "litellm/llms/custom_httpx/llm_http_handler.py"
- "tests/test_litellm/ocr/**"
- "tests/test_litellm/conftest.py"
- "Makefile"
- ".cargo/**"
- "pyproject.toml"
- "uv.lock"
- "rust-toolchain.toml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"

12
tests/test_litellm_rust/README.md
View file

@ -1,11 +1,15 @@
# Rust bridge tests
Put API-specific tests under `ocr/`, `messages/`, or `chat/`. Keep request construction, recording servers, callback recorders, and state isolation in `support/`. Integration tests and their fixtures belong in `integrations/`.
Put API-specific tests under `ocr/`, `messages/`, or `chat/`. Keep request construction, recording servers, callback recorders, and state isolation in `support/`. Under `integrations/`, keep registry inventory in `test_catalogue.py`, Python/Rust comparisons in `test_backend_parity.py`, callback ownership and completion in `test_callback_lifecycle.py`, individual exporters in `test_exporters.py`, and only multi-integration interactions in `test_composition.py`
Use `backend` with `"python"` and `"rust"` for parity tests. Keep callback mutation and lifecycle assertions next to the API surface that owns them. Shared route parametrization belongs in `integrations/routes.py` only when its observable contract is the same for every route.
Use `backend` with `"python"` and `"rust"` for parity tests. A test name must identify the observer and expected result. Keep multiple assertions together only when they are consequences of one mutation, failure, scheduling, or lifecycle event. Record callback observations and assert them after the callback returns because production logging can swallow callback exceptions. Shared route parametrization belongs in `support/routes.py` only when its observable contract is the same for every route
Run the suite with `LITELLM_RUST=1 uv run pytest tests/test_litellm_rust`.
The catalogue stores labels for selected behavioral cases. Those labels are inventory metadata and do not prove that matching pytest cases exist or execute
The retained callback contract modules are marked as non-strict expected failures until the implementation from #40070 lands. Harness tests remain strict, and passing contract cases appear as XPASS so staging coverage stays visible
Run `make test-rust-extension` as the acceptance command. It builds a fresh wheel, installs that wheel into a temporary environment, requires `LITELLM_RUST=1`, and runs this suite with isolated Python imports
Collection fails when `LITELLM_RUST=1` is set but the compiled `_native` module cannot be imported. The autouse fixture selects Rust unless a parity or fallback case explicitly selects Python. Backend selection alone does not prove native execution because a public route can fall back. Tests that claim native request coverage must also assert either the `x-litellm-rust` response marker or a wire property that distinguishes the native client. `test_public_ocr_executes_the_compiled_extension_without_python_fallback` is strict and its server rejects Python HTTPX requests
The retained callback contract modules are marked as non-strict expected failures until the implementation from #40070 lands. Harness tests and the compiled-extension OCR smoke test remain strict. Passing contract cases appear as XPASS so staging coverage stays visible
`isolated_backend` restores the previous backend override and callback state on exit, including after exceptions or cancellation. Scopes can nest in the same task. Run backend comparisons sequentially: overlapping scopes in different tasks raise before changing process-global state. Use separate worker processes for parallel backend comparisons

216
tests/test_litellm_rust/chat/test_callback_mutation.py
View file

@ -44,62 +44,136 @@ def _verify_sigv4(request: RecordedRequest, secret_key: str) -> None:
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("provider", ["anthropic", "bedrock"])
@pytest.mark.parametrize("rebind_logging_view", [False, True])
@pytest.mark.parametrize("status", [200, 429])
@pytest.mark.parametrize("native", [False, True])
async def test_chat_retains_callback_edits_through_public_dispatch(
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
@pytest.mark.parametrize("rebind_logging_view", [False, True], ids=["in-place-view", "rebound-view"])
@pytest.mark.parametrize("native", [False, True], ids=["python", "rust"])
async def test_anthropic_pre_call_body_and_header_mutations_reach_provider(
recording_server: RecordingServer,
asynchronous: bool,
provider: str,
rebind_logging_view: bool,
status: int,
native: bool,
) -> None:
import threading
litellm.rust(native)
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
from tests.test_litellm_rust.support.requests import MESSAGES_RESPONSE
recording_server.default_response = ResponseSpec(
status=status,
body=(
MESSAGES_RESPONSE
if provider == "anthropic"
else {
"output": {"message": {"role": "assistant", "content": [{"text": "native chat"}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 5, "outputTokens": 4, "totalTokens": 9},
"metrics": {"latencyMs": 1},
}
),
)
caller_thread: Final = threading.current_thread()
observations: Final = []
recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE)
callback_received_dict: Final = []
class EditingLogger(RecordingLogger):
def log_pre_api_call(self, model, messages, kwargs):
super().log_pre_api_call(model, messages, kwargs)
body = kwargs["additional_args"]["complete_input_dict"]
headers = kwargs["additional_args"]["headers"]
if provider == "anthropic":
assert isinstance(body, dict)
body["messages"][0]["content"][0]["text"] = "edited by callback"
body["max_tokens"] = 32
else:
assert isinstance(body, str)
assert json.loads(body)["messages"][0]["content"][0]["text"] == "original"
callback_received_dict.append(isinstance(body, dict))
if not isinstance(body, dict):
return
body["messages"][0]["content"][0]["text"] = "edited by callback"
body["max_tokens"] = 32
headers["x-retained-callback"] = "original"
observations.append((threading.current_thread(), body, headers))
if rebind_logging_view:
kwargs["additional_args"]["complete_input_dict"] = {"replacement": True}
kwargs["additional_args"]["headers"] = {"x-retained-callback": "replacement"}
recorder: Final = EditingLogger()
kwargs: Final = {
"model": "anthropic/claude-opus-5",
"messages": [{"role": "user", "content": "original"}],
"max_tokens": 64,
"api_key": "test-key",
"api_base": recording_server.base_url,
"callbacks": [recorder],
"num_retries": 0,
}
response: Final = await litellm.acompletion(**kwargs) if asynchronous else litellm.completion(**kwargs)
assert response._hidden_params.get("additional_headers", {}).get("x-litellm-rust") == ("true" if native else None)
assert callback_received_dict == [True]
assert recorder.names.count("log_pre_api_call") == 1
assert recording_server.requests[0].body["messages"][0]["content"][0]["text"] == "edited by callback"
assert recording_server.requests[0].headers["x-retained-callback"] == "original"
assert recording_server.requests[0].body["max_tokens"] == 32
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
@pytest.mark.parametrize("rebind_logging_view", [False, True], ids=["in-place-view", "rebound-view"])
@pytest.mark.parametrize("native", [False, True], ids=["python", "rust"])
async def test_bedrock_pre_call_rebinding_does_not_replace_signed_transport(
recording_server: RecordingServer,
asynchronous: bool,
rebind_logging_view: bool,
native: bool,
) -> None:
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
litellm.rust(native)
recording_server.default_response = ResponseSpec(
body={
"output": {"message": {"role": "assistant", "content": [{"text": "native chat"}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 5, "outputTokens": 4, "totalTokens": 9},
"metrics": {"latencyMs": 1},
}
)
callback_body_text: Final = []
class EditingLogger(RecordingLogger):
def log_pre_api_call(self, model, messages, kwargs):
super().log_pre_api_call(model, messages, kwargs)
additional: Final = kwargs["additional_args"]
body: Final = additional["complete_input_dict"]
callback_body_text.append(
json.loads(body)["messages"][0]["content"][0]["text"] if isinstance(body, str) else None
)
additional["headers"]["x-retained-callback"] = "original"
if rebind_logging_view:
additional["complete_input_dict"] = {"replacement": True}
additional["headers"] = {"x-retained-callback": "replacement"}
recorder: Final = EditingLogger()
request: Final = {
"model": "bedrock/anthropic.claude-opus-5",
"messages": [{"role": "user", "content": "original"}],
"max_tokens": 64,
"api_base": recording_server.base_url,
"callbacks": [recorder],
"num_retries": 0,
"aws_access_key_id": "test-access",
"aws_secret_access_key": "test-secret",
"aws_region_name": "us-east-1",
}
response: Final = await litellm.acompletion(**request) if asynchronous else litellm.completion(**request)
recorded: Final = recording_server.requests[0]
assert response._hidden_params.get("additional_headers", {}).get("x-litellm-rust") == ("true" if native else None)
assert callback_body_text == ["original"]
assert recorder.names.count("log_pre_api_call") == 1
assert recorded.body["messages"][0]["content"][0]["text"] == "original"
assert "replacement" not in recorded.body
assert recorded.headers["x-retained-callback"] == "original"
assert recorded.body["inferenceConfig"]["maxTokens"] == 64
assert recorded.headers["authorization"].startswith("AWS4-HMAC-SHA256 ")
_verify_sigv4(recorded, "test-secret")
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
@pytest.mark.parametrize("provider", ["anthropic", "bedrock"])
@pytest.mark.parametrize("native", [False, True], ids=["python", "rust"])
async def test_chat_provider_failure_reaches_one_failure_callback(
recording_server: RecordingServer,
asynchronous: bool,
provider: str,
native: bool,
) -> None:
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
litellm.rust(native)
recording_server.default_response = ResponseSpec(status=429, body={"message": "rate limited"})
recorder: Final = RecordingLogger()
request: Final = {
"model": "anthropic/claude-opus-5" if provider == "anthropic" else "bedrock/anthropic.claude-opus-5",
"messages": [{"role": "user", "content": "original"}],
"max_tokens": 64,
@ -113,45 +187,23 @@ async def test_chat_retains_callback_edits_through_public_dispatch(
else {}
),
}
if status != 200:
with pytest.raises((litellm.APIError, litellm.RateLimitError)) as raised:
await litellm.acompletion(**kwargs) if asynchronous else litellm.completion(**kwargs)
failure_event: Final = "async_log_failure_event" if asynchronous else "log_failure_event"
failures: Final = await recorder.wait_for_async(failure_event)
assert raised.value.status_code == status
assert failures[0].kwargs["exception"] is raised.value
assert recorder.names.count(failure_event) == 1
assert recorder.names.count("log_pre_api_call") == 1
assert len(recording_server.requests) == 1
return
response: Final = await litellm.acompletion(**kwargs) if asynchronous else litellm.completion(**kwargs)
event_name: Final = "async_log_success_event" if asynchronous else "log_success_event"
events: Final = await recorder.wait_for_async(event_name)
with pytest.raises((litellm.APIError, litellm.RateLimitError)) as raised:
await litellm.acompletion(**request) if asynchronous else litellm.completion(**request)
failure_event: Final = "async_log_failure_event" if asynchronous else "log_failure_event"
failures: Final = await recorder.wait_for_async(failure_event)
if native:
assert response._hidden_params["additional_headers"]["x-litellm-rust"] == "true"
assert raised.value.status_code == 429
assert failures[0].kwargs["exception"] is raised.value
assert recorder.names.count(failure_event) == 1
assert recorder.names.count("log_pre_api_call") == 1
assert recorder.names.count(event_name) == 1
if asynchronous and provider == "anthropic" and not native:
assert observations[0][0] is not caller_thread
else:
assert observations[0][0] is caller_thread
assert events[0].response is response
assert recording_server.requests[0].body["messages"][0]["content"][0]["text"] == (
"edited by callback" if provider == "anthropic" else "original"
)
assert recording_server.requests[0].headers["x-retained-callback"] == "original"
body: Final = recording_server.requests[0].body
assert (body["max_tokens"] if provider == "anthropic" else body["inferenceConfig"]["maxTokens"]) == (
32 if provider == "anthropic" else 64
)
assert len(recording_server.requests) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
@pytest.mark.parametrize("status", [200, 429])
@pytest.mark.parametrize("native", [False, True])
@pytest.mark.parametrize("native", [False, True], ids=["python", "rust"])
async def test_bedrock_callbacks_share_state_without_replacing_signed_transport(
recording_server: RecordingServer,
monkeypatch: pytest.MonkeyPatch,
@ -173,6 +225,7 @@ async def test_bedrock_callbacks_share_state_without_replacing_signed_transport(
},
)
observations: Final[list[tuple[int, int]]] = []
shared_state: Final[list[tuple[object, object, object]]] = []
replacement_body: Final = {"replacement": True}
rebound_headers: Final = {"x-callback-header": "rebound"}
@ -182,8 +235,15 @@ async def test_bedrock_callbacks_share_state_without_replacing_signed_transport(
callback_body: Final = additional["complete_input_dict"]
original_headers: Final = additional["headers"]
observations.append((id(kwargs), id(additional)))
assert isinstance(callback_body, str)
assert json.loads(callback_body)["messages"][0]["content"][0]["text"] == "original"
shared_state.append(
(
isinstance(callback_body, str),
json.loads(callback_body)["messages"][0]["content"][0]["text"]
if isinstance(callback_body, str)
else None,
None,
)
)
kwargs["callback_marker"] = "visible"
original_headers["x-callback-header"] = "in-place"
additional["complete_input_dict"] = replacement_body
@ -194,9 +254,13 @@ async def test_bedrock_callbacks_share_state_without_replacing_signed_transport(
def log_pre_api_call(self, model, messages, kwargs):
additional: Final = kwargs["additional_args"]
observations.append((id(kwargs), id(additional)))
assert kwargs["callback_marker"] == "visible"
assert additional["complete_input_dict"] is replacement_body
assert additional["headers"] is rebound_headers
shared_state.append(
(
kwargs.get("callback_marker"),
additional["complete_input_dict"] is replacement_body,
additional["headers"] is rebound_headers,
)
)
request: Final = {
"model": "bedrock/anthropic.claude-opus-5",
@ -216,6 +280,7 @@ async def test_bedrock_callbacks_share_state_without_replacing_signed_transport(
await litellm.acompletion(**request) if asynchronous else litellm.completion(**request)
assert observations[0] == observations[1]
assert shared_state == [(True, "original", None), ("visible", True, True)]
recorded: Final = recording_server.requests[0]
assert recorded.body["messages"][0]["content"][0]["text"] == "original"
assert "replacement" not in recorded.body
@ -261,7 +326,7 @@ async def test_native_chat_pre_call_preserves_the_caller_task_and_context(
@pytest.mark.asyncio
async def test_native_chat_concurrent_calls_keep_callback_roots_isolated(
async def test_chat_concurrent_calls_keep_callback_roots_isolated(
recording_server: RecordingServer,
) -> None:
import asyncio
@ -291,8 +356,8 @@ async def test_native_chat_concurrent_calls_keep_callback_roots_isolated(
logger: Final = Correlate()
async def invoke(call_id: str) -> None:
await litellm.acompletion(
async def invoke(call_id: str) -> object:
return await litellm.acompletion(
model="anthropic/claude-opus-5",
messages=[{"role": "user", "content": f"original-{call_id}"}],
max_tokens=16,
@ -308,13 +373,14 @@ async def test_native_chat_concurrent_calls_keep_callback_roots_isolated(
assert all(await asyncio.gather(*(asyncio.to_thread(signal.wait, 3) for signal in accepted)))
finally:
release.set()
await asyncio.gather(*tasks)
responses: Final = await asyncio.gather(*tasks)
await drain_logging()
sent: Final = {request.body["messages"][0]["content"][0]["text"] for request in recording_server.requests}
assert sent == {f"callback-{call_id}" for call_id in call_ids}
assert len(terminal_state) == len(call_ids)
assert all(token is tokens[call_id] for call_id, token in terminal_state)
assert all(response._hidden_params["additional_headers"]["x-litellm-rust"] == "true" for response in responses)
@pytest.mark.asyncio
@ -364,7 +430,7 @@ async def test_native_chat_callback_can_make_a_nested_native_call(
@pytest.mark.asyncio
async def test_native_chat_cancellation_during_io_does_not_publish_a_terminal(
async def test_chat_cancellation_during_io_does_not_publish_a_terminal(
recording_server: RecordingServer,
) -> None:
import asyncio

2
tests/test_litellm_rust/conftest.py
View file

@ -38,6 +38,8 @@ Backend = Literal["python", "rust"]
EXPECTED_FAILURE_FILES: Final = frozenset(
{
"chat/test_callback_mutation.py",
"integrations/test_backend_parity.py",
"integrations/test_callback_lifecycle.py",
"integrations/test_composition.py",
"integrations/test_exporters.py",
"integrations/test_guardrails.py",

4
tests/test_litellm_rust/integrations/catalogue.py
View file

@ -124,7 +124,7 @@ class CoverageObligation:
stable_name: str
registration_names: tuple[str, ...]
dependency_profile: DependencyProfile
behavioral_cases: tuple[str, ...]
behavioral_case_labels: tuple[str, ...]
REQUIRED_LOGGER_BEHAVIOR: Final = frozenset({"generic_api", "gcs_bucket", "literalai", "prometheus", "opentelemetry"})
@ -172,7 +172,7 @@ def _logger_obligations() -> Mapping[str, CoverageObligation]:
stable_name=names[0],
registration_names=names,
dependency_profile="enterprise" if set(names) & ENTERPRISE_LOGGER_NAMES else "required",
behavioral_cases=next(
behavioral_case_labels=next(
(LOGGER_BEHAVIORAL_CASES[name] for name in names if name in LOGGER_BEHAVIORAL_CASES), ()
),
)

92
tests/test_litellm_rust/integrations/test_backend_parity.py Normal file
View file

@ -0,0 +1,92 @@
from typing import Final
import pytest
import litellm
from tests.test_litellm_rust.conftest import Backend, isolated_backend
from tests.test_litellm_rust.integrations import (
MESSAGES_ROUTE,
OCR_ASYNC,
OCR_SYNC,
Route,
RunObservation,
provider_response,
route_id,
wait_for_callback,
)
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
from tests.test_litellm_rust.support.provenance import has_rust_response_marker
from tests.test_litellm_rust.support.recording_server import ResponseSpec, recording_service
from tests.test_litellm_rust.support.requests import CHAT_MESSAGES, CHAT_MODEL, CHAT_RESPONSE
pytestmark = pytest.mark.requires_rust_extension
async def observe_route(backend: Backend, route: Route) -> RunObservation:
async with isolated_backend(backend):
with recording_service() as provider:
provider.default_response = provider_response(route)
recorder: Final = RecordingLogger()
response: Final = await route.invoke(provider, callbacks=[recorder])
event: Final = (await wait_for_callback(route, recorder))[0]
payload: Final = event.kwargs["standard_logging_object"]
provider_body: Final = provider.requests[0].body
if not isinstance(provider_body, dict):
raise TypeError(f"Expected provider object body, got {type(provider_body).__name__}")
return RunObservation(
call_type=payload["call_type"],
model=payload["model"],
response_cost=payload["response_cost"],
response_text=route.response_text(response),
provider_body=provider_body,
rust_dispatch=has_rust_response_marker(response),
)
@pytest.mark.asyncio
@pytest.mark.parametrize("route", (OCR_SYNC, OCR_ASYNC, MESSAGES_ROUTE), ids=route_id)
async def test_python_and_rust_match_public_route_outputs_and_logging(route: Route) -> None:
python: Final = await observe_route("python", route)
rust: Final = await observe_route("rust", route)
python_body: Final = {**python.provider_body, "stream": python.provider_body.get("stream", False)}
rust_body: Final = {**rust.provider_body, "stream": rust.provider_body.get("stream", False)}
assert python.call_type == rust.call_type == route.call_type
assert python.model == rust.model == route.provider_model
assert python.response_cost == pytest.approx(rust.response_cost)
assert python.response_text == rust.response_text == route.expected_text
assert python_body == rust_body
assert python.rust_dispatch is False
assert rust.rust_dispatch is True
async def observe_retry(backend: Backend) -> tuple[tuple[str, ...], bool]:
async with isolated_backend(backend):
with recording_service() as provider:
provider.expected_requests = 2
provider.enqueue(ResponseSpec(body={"message": "retry this attempt"}, status=500))
provider.default_response = ResponseSpec(body=CHAT_RESPONSE)
recorder: Final = RecordingLogger()
response: Final = await litellm.acompletion(
model=CHAT_MODEL,
messages=CHAT_MESSAGES,
api_key="test-key",
api_base=provider.base_url,
callbacks=[recorder],
num_retries=1,
)
await drain_logging()
return recorder.names, has_rust_response_marker(response)
@pytest.mark.asyncio
async def test_python_and_rust_emit_the_same_retry_callback_sequence() -> None:
python_names, python_used_rust = await observe_retry("python")
rust_names, rust_used_rust = await observe_retry("rust")
assert rust_names == python_names
assert rust_names.count("log_pre_api_call") == 2
assert rust_names.count("async_log_failure_event") == 1
assert rust_names.count("async_log_success_event") == 0
assert python_used_rust is False
assert rust_used_rust is True

281
tests/test_litellm_rust/integrations/test_callback_lifecycle.py Normal file
View file

@ -0,0 +1,281 @@
import asyncio
import gc
import json
import threading
import weakref
from collections.abc import Mapping
from datetime import datetime
from typing import Final
import pytest
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.types.utils import CallTypes
from tests.test_litellm_rust.conftest import Backend, isolated_backend
from tests.test_litellm_rust.integrations import MESSAGES_ROUTE, MESSAGES_STREAM, OCR_ASYNC, Route, provider_response
from tests.test_litellm_rust.support.callback_recorder import (
LiveReferenceLogger,
RecordingLogger,
SecondaryLiveReferenceLogger,
drain_logging,
)
from tests.test_litellm_rust.support.provenance import has_rust_response_marker
from tests.test_litellm_rust.support.recording_server import ResponseSpec, recording_service
from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_EVENTS
pytestmark = pytest.mark.requires_rust_extension
@pytest.mark.asyncio
@pytest.mark.parametrize("backend", ("python", "rust"))
async def test_terminal_callbacks_receive_the_same_live_kwargs_and_response(backend: Backend) -> None:
async with isolated_backend(backend):
with recording_service() as provider:
provider.default_response = provider_response(MESSAGES_ROUTE)
first: Final = LiveReferenceLogger()
second: Final = SecondaryLiveReferenceLogger()
response: Final = await MESSAGES_ROUTE.invoke(provider, callbacks=[first, second])
first_event: Final = (await first.wait_for_async())[0]
second_event: Final = (await second.wait_for_async())[0]
assert first_event.kwargs is second_event.kwargs
assert first_event.response is second_event.response
assert has_rust_response_marker(response) is (backend == "rust")
first.release()
second.release()
@pytest.mark.asyncio
@pytest.mark.parametrize("backend", ("python", "rust"))
async def test_redacted_and_plain_loggers_receive_separate_payloads_with_shared_metadata(backend: Backend) -> None:
class Capture(CustomLogger):
def __init__(self, redact: bool) -> None:
super().__init__(turn_off_message_logging=redact)
self.kwargs: object = None
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.kwargs = kwargs
async with isolated_backend(backend):
with recording_service() as provider:
provider.default_response = provider_response(MESSAGES_ROUTE)
redacted: Final = Capture(True)
plain: Final = Capture(False)
response: Final = await MESSAGES_ROUTE.invoke(provider, callbacks=[redacted, plain])
await drain_logging()
redacted_kwargs: Final = redacted.kwargs
plain_kwargs: Final = plain.kwargs
assert isinstance(redacted_kwargs, dict)
assert isinstance(plain_kwargs, dict)
redacted_payload: Final = redacted_kwargs["standard_logging_object"]
plain_payload: Final = plain_kwargs["standard_logging_object"]
assert redacted_kwargs is not plain_kwargs
assert redacted_payload is not plain_payload
assert redacted_payload["metadata"] is plain_payload["metadata"]
assert "redacted-by-litellm" in json.dumps(redacted_payload["messages"])
assert MESSAGES_ROUTE.expected_text in json.dumps(plain_payload["response"])
assert "redacted-by-litellm" not in json.dumps(plain_payload)
assert has_rust_response_marker(response) is (backend == "rust")
@pytest.mark.asyncio
@pytest.mark.parametrize("backend", ("python", "rust"))
@pytest.mark.parametrize("accepted", (True, False), ids=("accepted", "rejected"))
async def test_deferred_terminal_callback_runs_only_after_proxy_completion_decision(
backend: Backend, accepted: bool
) -> None:
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
async with isolated_backend(backend):
with recording_service() as provider:
provider.default_response = provider_response(MESSAGES_ROUTE)
recorder: Final = RecordingLogger()
logger: Final = Logging(
model=MESSAGES_ROUTE.provider_model,
messages=MESSAGES,
stream=False,
call_type=MESSAGES_ROUTE.call_type,
start_time=datetime.now(),
litellm_call_id="deferred-completion",
function_id="deferred-completion",
dynamic_async_success_callbacks=[recorder],
dynamic_async_failure_callbacks=[recorder],
)
logger._defer_async_logging = True
response: Final = await MESSAGES_ROUTE.invoke(provider, litellm_logging_obj=logger)
await asyncio.sleep(0)
assert "async_log_success_event" not in recorder.names
ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(
logging_obj=logger,
exception_raised=not accepted,
)
if accepted:
accepted_events: Final = await recorder.wait_for_async("async_log_success_event")
assert len(accepted_events) == 1
assert "async_log_failure_event" not in recorder.names
else:
error: Final = RuntimeError("post-call guardrail rejected response")
await logger.async_failure_handler(error, str(error))
rejected_events: Final = await recorder.wait_for_async("async_log_failure_event")
assert len(rejected_events) == 1
assert "async_log_success_event" not in recorder.names
assert has_rust_response_marker(response) is (backend == "rust")
@pytest.mark.asyncio
@pytest.mark.parametrize("backend", ("python", "rust"))
async def test_suspended_terminal_callback_retains_metadata_until_completion(backend: Backend) -> None:
class Root:
pass
class Suspended(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.started = asyncio.Event()
self.release = asyncio.Event()
self.saw_retained_root = False
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.saw_retained_root = kwargs["litellm_params"]["metadata"]["retained"] is not None
self.started.set()
await self.release.wait()
async def invoke() -> weakref.ReferenceType[Root]:
async with isolated_backend(backend):
with recording_service() as provider:
provider.default_response = provider_response(MESSAGES_ROUTE)
callback: Final = Suspended()
root = Root()
reference: Final = weakref.ref(root)
response: Final = await MESSAGES_ROUTE.invoke(
provider,
callbacks=[callback],
metadata={"retained": root},
)
del root
await asyncio.wait_for(callback.started.wait(), timeout=2)
assert callback.saw_retained_root is True
gc.collect()
assert reference() is not None
assert has_rust_response_marker(response) is (backend == "rust")
callback.release.set()
await drain_logging()
del response
return reference
reference: Final = await invoke()
await asyncio.sleep(0)
gc.collect()
assert reference() is None
@pytest.mark.asyncio
@pytest.mark.parametrize("backend", ("python", "rust"))
@pytest.mark.parametrize(
("route", "phase"),
(
(OCR_ASYNC, "pre"),
(MESSAGES_ROUTE, "pre"),
(MESSAGES_STREAM, "pre"),
(OCR_ASYNC, "success"),
(MESSAGES_ROUTE, "success"),
),
ids=("ocr-pre", "messages-pre", "messages-stream-pre", "ocr-success", "messages-success"),
)
async def test_deployment_rejection_logs_terminal_failure_without_calling_deployment_failure_hook(
backend: Backend, route: Route, phase: str
) -> None:
class Reject(CustomLogger):
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, object], call_type: CallTypes | None
) -> dict[str, object]:
if phase == "pre":
raise litellm.BadRequestError("deployment rejected", "test-provider", route.provider_model)
return kwargs
async def async_post_call_success_deployment_hook(
self, request_data: Mapping[str, object], response: object, call_type: CallTypes | None
) -> object:
raise litellm.BadRequestError("deployment rejected", "test-provider", route.provider_model)
async with isolated_backend(backend):
with recording_service() as provider:
provider.expected_requests = 0 if phase == "pre" else 1
provider.default_response = provider_response(route)
recorder: Final = RecordingLogger()
litellm.callbacks.extend((Reject(), recorder))
with pytest.raises(litellm.BadRequestError, match="deployment rejected"):
await route.invoke(provider)
await recorder.wait_for_async("async_log_failure_event")
assert len(provider.requests) == provider.expected_requests
assert "async_log_success_event" not in recorder.names
assert recorder.names.count("async_post_call_failure_deployment_hook") == 0
@pytest.mark.asyncio
@pytest.mark.parametrize("ending", ("provider_error", "truncated", "close"))
async def test_established_rust_stream_finishes_once_without_deployment_failure_hook(ending: str) -> None:
with recording_service() as provider:
provider.default_response = ResponseSpec(
body=None,
events=(
(
*MESSAGES_EVENTS[:1],
(
"error",
{"type": "error", "error": {"type": "overloaded_error", "message": "upstream overloaded"}},
),
)
if ending == "provider_error"
else MESSAGES_EVENTS[:-1]
if ending == "truncated"
else MESSAGES_EVENTS
),
)
release: Final = threading.Event()
if ending == "close":
provider.enqueue(
ResponseSpec(
body=None,
chunks=tuple(
f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in MESSAGES_EVENTS
),
release=release,
)
)
recorder: Final = RecordingLogger()
litellm.callbacks.append(recorder)
try:
stream: Final = await MESSAGES_STREAM.open_stream(provider)
assert has_rust_response_marker(stream)
assert recorder.names.count("async_pre_call_deployment_hook") == 1
assert "async_post_call_failure_deployment_hook" not in recorder.names
if ending == "close":
assert await anext(stream)
await asyncio.wait_for(stream.aclose(), timeout=1)
await stream.aclose()
assert "async_log_success_event" not in recorder.names
assert "async_log_failure_event" not in recorder.names
release.set()
await recorder.wait_for_async("async_log_success_event")
else:
try:
async for _ in stream:
pass
except litellm.APIError:
pass
await drain_logging()
finally:
release.set()
assert len(provider.requests) == 1
assert recorder.names.count("async_log_failure_event") == (0 if ending == "close" else 1), recorder.names
assert recorder.names.count("async_log_success_event") == (1 if ending == "close" else 0), recorder.names
assert recorder.names.count("async_post_call_failure_deployment_hook") == 0

54
tests/test_litellm_rust/integrations/test_catalogue.py Normal file
View file

@ -0,0 +1,54 @@
from typing import Final
import pytest
from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry
from litellm.proxy.guardrails.guardrail_registry import guardrail_initializer_registry
from litellm.types.guardrails import SupportedGuardrailIntegrations
from tests.test_litellm_rust.integrations import (
DISCOVERED_ONLY_GUARDRAIL_NAMES,
ENTERPRISE_LOGGER_NAMES,
GUARDRAIL_NAMES,
GUARDRAIL_OBLIGATIONS,
LOGGER_OBLIGATIONS,
OSS_LOGGER_NAMES,
REQUIRED_GUARDRAIL_BEHAVIOR,
REQUIRED_LOGGER_BEHAVIOR,
)
pytestmark = pytest.mark.requires_rust_extension
def test_logger_catalogue_matches_registered_integration_names() -> None:
registered_names: Final = frozenset(CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE)
catalogued_names: Final = frozenset(
name for obligation in LOGGER_OBLIGATIONS.values() for name in obligation.registration_names
)
assert OSS_LOGGER_NAMES <= registered_names
assert registered_names <= OSS_LOGGER_NAMES | ENTERPRISE_LOGGER_NAMES
assert catalogued_names == OSS_LOGGER_NAMES | ENTERPRISE_LOGGER_NAMES
def test_guardrail_catalogue_matches_registered_integration_names() -> None:
enum_names: Final = frozenset(integration.value for integration in SupportedGuardrailIntegrations)
initializer_names: Final = frozenset(guardrail_initializer_registry)
assert enum_names == GUARDRAIL_NAMES
assert initializer_names == GUARDRAIL_NAMES | DISCOVERED_ONLY_GUARDRAIL_NAMES
assert frozenset(GUARDRAIL_OBLIGATIONS) == initializer_names
def test_required_integrations_have_behavioral_case_labels() -> None:
labelled_loggers: Final = frozenset(
name
for obligation in LOGGER_OBLIGATIONS.values()
if obligation.behavioral_case_labels
for name in obligation.registration_names
)
labelled_guardrails: Final = frozenset(
name for name, obligation in GUARDRAIL_OBLIGATIONS.items() if obligation.behavioral_case_labels
)
assert REQUIRED_LOGGER_BEHAVIOR <= labelled_loggers
assert REQUIRED_GUARDRAIL_BEHAVIOR <= labelled_guardrails

411
tests/test_litellm_rust/integrations/test_composition.py
View file

@ -1,66 +1,30 @@
import asyncio
import gc
import json
import threading
import weakref
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime
from typing import Final
import pytest
from opentelemetry.trace import StatusCode
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger
from litellm.integrations.prometheus import PrometheusLogger
from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy.guardrails.guardrail_registry import guardrail_initializer_registry
from litellm.types.guardrails import SupportedGuardrailIntegrations
from litellm.types.utils import CallTypes
from tests.test_litellm_rust.conftest import Backend, isolated_backend
from tests.test_litellm_rust.integrations import (
ASYNC_ROUTES,
DISCOVERED_ONLY_GUARDRAIL_NAMES,
ENTERPRISE_LOGGER_NAMES,
GUARDRAIL_NAMES,
GUARDRAIL_OBLIGATIONS,
LOGGER_OBLIGATIONS,
MESSAGES_ROUTE,
MESSAGES_STREAM,
OCR_ASYNC,
OCR_SYNC,
OSS_LOGGER_NAMES,
REQUIRED_GUARDRAIL_BEHAVIOR,
REQUIRED_LOGGER_BEHAVIOR,
AsyncBoundaryLogger,
MutatingFailingLogger,
OtelHarness,
Route,
RunObservation,
gcs_literalai_harness,
metric_value,
provider_response,
route_id,
wait_for_callback,
)
from tests.test_litellm_rust.support.callback_recorder import (
LiveReferenceLogger,
RecordingLogger,
SecondaryLiveReferenceLogger,
drain_logging,
)
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger
from tests.test_litellm_rust.support.provenance import has_rust_response_marker
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, recording_service
from tests.test_litellm_rust.support.requests import (
CHAT_MESSAGES,
CHAT_MODEL,
CHAT_RESPONSE,
MESSAGES,
MESSAGES_EVENTS,
)
from tests.test_litellm_rust.support.recording_server import RecordingServer, recording_service
pytestmark = pytest.mark.requires_rust_extension
@ -77,62 +41,9 @@ MESSAGES_TOOLS: Final = [
]
async def observe_backend(route: Route, backend: Backend) -> RunObservation:
async with isolated_backend(backend):
with recording_service() as provider:
provider.default_response = provider_response(route)
recorder: Final = RecordingLogger()
response: Final = await route.invoke(provider, callbacks=[recorder])
event: Final = (await wait_for_callback(route, recorder))[0]
payload: Final = event.kwargs["standard_logging_object"]
provider_body: Final = provider.requests[0].body
if not isinstance(provider_body, dict):
raise TypeError(f"Expected provider object body, got {type(provider_body).__name__}")
return RunObservation(
call_type=payload["call_type"],
model=payload["model"],
response_cost=payload["response_cost"],
response_text=route.response_text(response),
provider_body=provider_body,
rust_dispatch=has_rust_response_marker(response),
)
def test_logger_catalogue_reconciles_with_production_registry() -> None:
actual_names: Final = frozenset(CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE)
catalogued_names: Final = frozenset(
name for obligation in LOGGER_OBLIGATIONS.values() for name in obligation.registration_names
)
assert OSS_LOGGER_NAMES <= actual_names
assert actual_names <= OSS_LOGGER_NAMES | ENTERPRISE_LOGGER_NAMES
assert catalogued_names == OSS_LOGGER_NAMES | ENTERPRISE_LOGGER_NAMES
def test_guardrail_catalogue_reconciles_enum_and_runtime_discovery() -> None:
enum_names: Final = frozenset(integration.value for integration in SupportedGuardrailIntegrations)
initializer_names: Final = frozenset(guardrail_initializer_registry)
assert enum_names == GUARDRAIL_NAMES
assert initializer_names == GUARDRAIL_NAMES | DISCOVERED_ONLY_GUARDRAIL_NAMES
assert frozenset(GUARDRAIL_OBLIGATIONS) == initializer_names
def test_required_integrations_have_behavioral_cases() -> None:
logger_names: Final = frozenset(
name
for obligation in LOGGER_OBLIGATIONS.values()
if obligation.behavioral_cases
for name in obligation.registration_names
)
guardrail_names: Final = frozenset(
name for name, obligation in GUARDRAIL_OBLIGATIONS.items() if obligation.behavioral_cases
)
assert REQUIRED_LOGGER_BEHAVIOR <= logger_names
assert REQUIRED_GUARDRAIL_BEHAVIOR <= guardrail_names
@pytest.mark.asyncio
@pytest.mark.parametrize("route", ASYNC_ROUTES, ids=lambda route: f"otel-prometheus-generic-api-{route.name}-success")
async def test_export_composition(
async def test_otel_prometheus_and_generic_api_each_export_one_success(
route: Route, recording_server: RecordingServer, otel: OtelHarness, prometheus: PrometheusLogger
) -> None:
recording_server.default_response = provider_response(route)
@ -141,7 +52,7 @@ async def test_export_composition(
with recording_service() as sink:
logger: Final = GenericAPILogger(endpoint=f"{sink.base_url}/logs", batch_size=1, log_format="single")
await route.invoke(recording_server, callbacks=[otel.logger, prometheus, logger, recorder])
response: Final = await route.invoke(recording_server, callbacks=[otel.logger, prometheus, logger, recorder])
await wait_for_callback(route, recorder)
exports: Final = tuple(request for request in sink.requests if request.path == "/logs")
@ -152,191 +63,7 @@ async def test_export_composition(
assert len(spans) == 1
assert spans[0].status.status_code is StatusCode.OK
assert metric_value("litellm_requests_metric_total", model=route.provider_model) == before + 1
@pytest.mark.asyncio
@pytest.mark.parametrize("route", (OCR_SYNC, OCR_ASYNC, MESSAGES_ROUTE), ids=route_id)
async def test_public_sdk_python_rust_composition_parity(route: Route) -> None:
python: Final = await observe_backend(route, "python")
rust: Final = await observe_backend(route, "rust")
assert python.call_type == rust.call_type == route.call_type
assert python.model == rust.model == route.provider_model
assert python.response_cost == pytest.approx(rust.response_cost)
assert python.response_text == rust.response_text == route.expected_text
python_body: Final = {**python.provider_body, "stream": python.provider_body.get("stream", False)}
rust_body: Final = {**rust.provider_body, "stream": rust.provider_body.get("stream", False)}
assert python_body == rust_body
assert python.rust_dispatch is False
assert rust.rust_dispatch is True
@pytest.mark.asyncio
@pytest.mark.parametrize("backend", ("python", "rust"))
async def test_terminal_callbacks_share_live_objects_within_one_run(backend: Backend) -> None:
async with isolated_backend(backend):
with recording_service() as provider:
provider.default_response = provider_response(MESSAGES_ROUTE)
first: Final = LiveReferenceLogger()
second: Final = SecondaryLiveReferenceLogger()
response: Final = await MESSAGES_ROUTE.invoke(provider, callbacks=[first, second])
first_event: Final = (await first.wait_for_async())[0]
second_event: Final = (await second.wait_for_async())[0]
assert first_event.kwargs is second_event.kwargs
assert first_event.response is second_event.response
assert has_rust_response_marker(response) is (backend == "rust")
first.release()
second.release()
@pytest.mark.asyncio
@pytest.mark.parametrize("backend", ("python", "rust"))
async def test_per_logger_redaction_preserves_the_shared_unredacted_payload(backend: Backend) -> None:
class Capture(CustomLogger):
def __init__(self, redact: bool) -> None:
super().__init__(turn_off_message_logging=redact)
self.kwargs: object = None
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.kwargs = kwargs
async with isolated_backend(backend):
with recording_service() as provider:
provider.default_response = provider_response(MESSAGES_ROUTE)
redacted: Final = Capture(True)
plain: Final = Capture(False)
await MESSAGES_ROUTE.invoke(provider, callbacks=[redacted, plain])
await drain_logging()
redacted_kwargs: Final = redacted.kwargs
plain_kwargs: Final = plain.kwargs
assert isinstance(redacted_kwargs, dict)
assert isinstance(plain_kwargs, dict)
redacted_payload: Final = redacted_kwargs["standard_logging_object"]
plain_payload: Final = plain_kwargs["standard_logging_object"]
assert redacted_kwargs is not plain_kwargs
assert redacted_payload is not plain_payload
assert redacted_payload["metadata"] is plain_payload["metadata"]
assert "redacted-by-litellm" in json.dumps(redacted_payload["messages"])
assert MESSAGES_ROUTE.expected_text in json.dumps(plain_payload["response"])
assert "redacted-by-litellm" not in json.dumps(plain_payload)
@pytest.mark.asyncio
@pytest.mark.parametrize("backend", ("python", "rust"))
@pytest.mark.parametrize("accepted", (True, False), ids=("accepted", "rejected"))
async def test_terminal_respects_proxy_deferred_completion_gate(backend: Backend, accepted: bool) -> None:
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
async with isolated_backend(backend):
with recording_service() as provider:
provider.default_response = provider_response(MESSAGES_ROUTE)
recorder: Final = RecordingLogger()
logger: Final = Logging(
model=MESSAGES_ROUTE.provider_model,
messages=MESSAGES,
stream=False,
call_type=MESSAGES_ROUTE.call_type,
start_time=datetime.now(),
litellm_call_id="deferred-completion",
function_id="deferred-completion",
dynamic_async_success_callbacks=[recorder],
dynamic_async_failure_callbacks=[recorder],
)
logger._defer_async_logging = True
response: Final = await MESSAGES_ROUTE.invoke(provider, litellm_logging_obj=logger)
await asyncio.sleep(0)
assert "async_log_success_event" not in recorder.names
ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(
logging_obj=logger,
exception_raised=not accepted,
)
if accepted:
accepted_events: Final = await recorder.wait_for_async("async_log_success_event")
assert len(accepted_events) == 1
assert "async_log_failure_event" not in recorder.names
else:
error: Final = RuntimeError("post-call guardrail rejected response")
await logger.async_failure_handler(error, str(error))
rejected_events: Final = await recorder.wait_for_async("async_log_failure_event")
assert len(rejected_events) == 1
assert "async_log_success_event" not in recorder.names
assert has_rust_response_marker(response) is (backend == "rust")
@pytest.mark.asyncio
@pytest.mark.parametrize("backend", ("python", "rust"))
async def test_suspended_terminal_callback_owns_payload_until_it_finishes(backend: Backend) -> None:
class Root:
pass
class Suspended(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.started = asyncio.Event()
self.release = asyncio.Event()
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
assert kwargs["litellm_params"]["metadata"]["retained"] is not None
self.started.set()
await self.release.wait()
async def invoke() -> weakref.ReferenceType[Root]:
async with isolated_backend(backend):
with recording_service() as provider:
provider.default_response = provider_response(MESSAGES_ROUTE)
callback: Final = Suspended()
root = Root()
reference: Final = weakref.ref(root)
response: Final = await MESSAGES_ROUTE.invoke(
provider,
callbacks=[callback],
metadata={"retained": root},
)
del root
await asyncio.wait_for(callback.started.wait(), timeout=2)
gc.collect()
assert reference() is not None
callback.release.set()
await drain_logging()
del response
return reference
reference: Final = await invoke()
await asyncio.sleep(0)
gc.collect()
assert reference() is None
async def observe_retry(backend: Backend) -> tuple[tuple[str, ...], bool]:
async with isolated_backend(backend):
with recording_service() as provider:
provider.expected_requests = 2
provider.enqueue(ResponseSpec(body={"message": "retry this attempt"}, status=500))
provider.default_response = ResponseSpec(body=CHAT_RESPONSE)
recorder: Final = RecordingLogger()
response: Final = await litellm.acompletion(
model=CHAT_MODEL,
messages=CHAT_MESSAGES,
api_key="test-key",
api_base=provider.base_url,
callbacks=[recorder],
num_retries=1,
)
await drain_logging()
return recorder.names, has_rust_response_marker(response)
@pytest.mark.asyncio
async def test_retry_attempt_callback_sequence_matches_python() -> None:
python_names, python_rust_dispatch = await observe_retry("python")
rust_names, rust_rust_dispatch = await observe_retry("rust")
assert rust_names == python_names
assert rust_names.count("log_pre_api_call") == 2
assert rust_names.count("async_log_failure_event") == 1
assert rust_names.count("async_log_success_event") == 0
assert python_rust_dispatch is False
assert rust_rust_dispatch is True
assert has_rust_response_marker(response)
@pytest.mark.asyncio
@ -352,6 +79,7 @@ async def test_failing_callback_preserves_prior_mutation_and_later_exporters() -
assert events[0].kwargs["standard_logging_object"]["metadata"]["composition_marker"] == "visible-before-failure"
assert len(exports) == 1
assert exports[0].body["metadata"]["composition_marker"] == "visible-before-failure"
assert has_rust_response_marker(response)
@dataclass(frozen=True, slots=True)
@ -419,130 +147,3 @@ async def test_gcs_literalai_serialization_schedule(
assert len(harness.storage.posts) == 1
assert len(harness.literal_sink.posts) == 1
assert has_rust_response_marker(response) is (backend == "rust")
@pytest.mark.asyncio
@pytest.mark.parametrize("backend", ("python", "rust"))
async def test_interrupted_stream_emits_terminal_only_when_consumer_closes(backend: Backend, otel: OtelHarness) -> None:
async with isolated_backend(backend):
with recording_service() as provider:
provider.default_response = ResponseSpec(body=None, events=MESSAGES_EVENTS)
stream: Final = await MESSAGES_STREAM.open_stream(provider, callbacks=[otel.logger])
first_chunk: Final = await anext(stream)
await drain_logging()
assert otel.spans() == ()
assert has_rust_response_marker(stream) is (backend == "rust")
await stream.aclose()
await drain_logging()
assert first_chunk is not None
assert len(otel.spans()) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("backend", ("python", "rust"))
@pytest.mark.parametrize(
("route", "phase"),
(
(OCR_ASYNC, "pre"),
(MESSAGES_ROUTE, "pre"),
(MESSAGES_STREAM, "pre"),
(OCR_ASYNC, "success"),
(MESSAGES_ROUTE, "success"),
),
ids=("ocr-pre", "messages-pre", "messages-stream-pre", "ocr-success", "messages-success"),
)
async def test_deployment_rejection_logs_failure_without_deployment_failure(
backend: Backend, route: Route, phase: str
) -> None:
class Reject(CustomLogger):
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, object], call_type: CallTypes | None
) -> dict[str, object]:
if phase == "pre":
raise litellm.BadRequestError("deployment rejected", "test-provider", route.provider_model)
return kwargs
async def async_post_call_success_deployment_hook(
self, request_data: Mapping[str, object], response: object, call_type: CallTypes | None
) -> object:
raise litellm.BadRequestError("deployment rejected", "test-provider", route.provider_model)
async with isolated_backend(backend):
with recording_service() as provider:
provider.expected_requests = 0 if phase == "pre" else 1
provider.default_response = provider_response(route)
recorder: Final = RecordingLogger()
litellm.callbacks.extend((Reject(), recorder))
with pytest.raises(litellm.BadRequestError, match="deployment rejected"):
await route.invoke(provider)
await recorder.wait_for_async("async_log_failure_event")
assert len(provider.requests) == provider.expected_requests
assert "async_log_success_event" not in recorder.names
assert recorder.names.count("async_post_call_failure_deployment_hook") == 0
@pytest.mark.asyncio
@pytest.mark.parametrize("backend", ("rust",))
@pytest.mark.parametrize("ending", ("provider_error", "truncated", "close"))
async def test_established_stream_completion_does_not_notify_deployment_failure(backend: Backend, ending: str) -> None:
async with isolated_backend(backend):
with recording_service() as provider:
provider.default_response = ResponseSpec(
body=None,
events=(
(
*MESSAGES_EVENTS[:1],
(
"error",
{"type": "error", "error": {"type": "overloaded_error", "message": "upstream overloaded"}},
),
)
if ending == "provider_error"
else MESSAGES_EVENTS[:-1]
if ending == "truncated"
else MESSAGES_EVENTS
),
)
release: Final = threading.Event()
if ending == "close":
provider.enqueue(
ResponseSpec(
body=None,
chunks=tuple(
f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in MESSAGES_EVENTS
),
release=release,
)
)
recorder: Final = RecordingLogger()
litellm.callbacks.append(recorder)
try:
stream: Final = await MESSAGES_STREAM.open_stream(provider)
assert has_rust_response_marker(stream) is (backend == "rust")
assert recorder.names.count("async_pre_call_deployment_hook") == 1
assert "async_post_call_failure_deployment_hook" not in recorder.names
if ending == "close":
assert await anext(stream)
await asyncio.wait_for(stream.aclose(), timeout=1)
await stream.aclose()
assert "async_log_success_event" not in recorder.names
assert "async_log_failure_event" not in recorder.names
release.set()
await recorder.wait_for_async("async_log_success_event")
else:
try:
async for _ in stream:
pass
except litellm.APIError:
pass
await drain_logging()
finally:
release.set()
assert len(provider.requests) == 1
assert recorder.names.count("async_log_failure_event") == (0 if ending == "close" else 1), recorder.names
assert recorder.names.count("async_log_success_event") == (1 if ending == "close" else 0), recorder.names
assert recorder.names.count("async_post_call_failure_deployment_hook") == 0

48
tests/test_litellm_rust/integrations/test_exporters.py
View file

@ -10,6 +10,7 @@ from litellm.integrations.prometheus import PrometheusLogger
from litellm.litellm_core_utils import litellm_logging
from tests._prometheus_helpers import isolated_prometheus_registry
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
from tests.test_litellm_rust.support.provenance import has_rust_response_marker
from tests.test_litellm_rust.support.requests import MESSAGES_EVENTS
from tests.test_litellm_rust.integrations import (
ALL_ROUTES,
@ -33,6 +34,16 @@ pytestmark = pytest.mark.requires_rust_extension
FAILURE_RESPONSE: Final = ResponseSpec(body={"message": "provider unavailable"}, status=500)
async def invoke_native(route: Route, provider: RecordingServer, **kwargs: object) -> object:
if route.name == "messages-stream":
stream: Final = await route.open_stream(provider, **kwargs)
assert has_rust_response_marker(stream)
return [chunk async for chunk in stream]
response: Final = await route.invoke(provider, **kwargs)
assert has_rust_response_marker(response)
return response
@pytest.mark.asyncio
@pytest.mark.parametrize("route", ASYNC_ROUTES, ids=route_id)
async def test_generic_api_logger_exports_success_over_http(
@ -41,7 +52,7 @@ async def test_generic_api_logger_exports_success_over_http(
generic_api_export: GenericAPIExportHarness,
) -> None:
recorder: Final = RecordingLogger()
await route.invoke(provider, callbacks=[generic_api_export.logger, recorder])
await invoke_native(route, provider, callbacks=[generic_api_export.logger, recorder])
await wait_for_callback(route, recorder)
assert len(generic_api_export.exports) == 1
@ -104,7 +115,7 @@ def test_prometheus_registry_restores_collectors_after_failure() -> None:
async def test_otel_emits_one_request_span_on_success(
route: Route, provider: RecordingServer, otel: OtelHarness
) -> None:
await route.invoke(provider, callbacks=[otel.logger])
await invoke_native(route, provider, callbacks=[otel.logger])
spans: Final = await otel.wait_for_spans()
assert len(spans) == 1
span: Final = spans[0]
@ -127,6 +138,25 @@ async def test_otel_stream_span_appears_only_after_exhaustion(
assert len(await otel.wait_for_spans()) == 1
@pytest.mark.asyncio
async def test_otel_stream_span_appears_after_an_interrupted_consumer_closes(
otel: OtelHarness, recording_server: RecordingServer
) -> None:
recording_server.default_response = ResponseSpec(body=None, events=MESSAGES_EVENTS)
stream: Final = await MESSAGES_STREAM.open_stream(recording_server, callbacks=[otel.logger])
first_chunk: Final = await anext(stream)
await drain_logging()
assert otel.spans() == ()
assert has_rust_response_marker(stream)
await stream.aclose()
await drain_logging()
assert first_chunk is not None
assert len(otel.spans()) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("route", (OCR_SYNC, OCR_ASYNC), ids=route_id)
async def test_otel_emits_one_error_span_on_provider_failure(
@ -146,11 +176,11 @@ async def test_otel_emits_one_error_span_on_provider_failure(
@pytest.mark.asyncio
@pytest.mark.parametrize("route", ALL_ROUTES, ids=route_id)
async def test_otel_and_custom_logger_observe_same_standard_logging_object(
async def test_otel_and_custom_logger_export_matching_standard_logging_values(
route: Route, provider: RecordingServer, otel: OtelHarness
) -> None:
recorder: Final = RecordingLogger()
await route.invoke(provider, callbacks=[otel.logger, recorder])
await invoke_native(route, provider, callbacks=[otel.logger, recorder])
payload: Final = (await wait_for_callback(route, recorder))[0].kwargs["standard_logging_object"]
span: Final = (await otel.wait_for_spans())[0]
assert payload["call_type"] == route.call_type
@ -167,7 +197,7 @@ async def test_prometheus_counts_one_successful_request(
) -> None:
before: Final = metric_value("litellm_requests_metric_total", model=route.provider_model)
recorder: Final = RecordingLogger()
await route.invoke(provider, callbacks=[prometheus, recorder])
await invoke_native(route, provider, callbacks=[prometheus, recorder])
await wait_for_callback(route, recorder)
assert metric_value("litellm_requests_metric_total", model=route.provider_model) == before + 1
assert metric_value("litellm_llm_api_failed_requests_metric_total", model=route.provider_model) == 0
@ -179,7 +209,7 @@ async def test_prometheus_counts_tokens_from_messages_usage(
) -> None:
recording_server.default_response = ResponseSpec(body=MESSAGES_ROUTE.provider_response)
recorder: Final = RecordingLogger()
await MESSAGES_ROUTE.invoke(recording_server, callbacks=[prometheus, recorder])
await invoke_native(MESSAGES_ROUTE, recording_server, callbacks=[prometheus, recorder])
await wait_for_callback(MESSAGES_ROUTE, recorder)
assert metric_value("litellm_input_tokens_metric_total", model=MESSAGES_ROUTE.provider_model) == 5
assert metric_value("litellm_output_tokens_metric_total", model=MESSAGES_ROUTE.provider_model) == 4
@ -203,8 +233,8 @@ async def test_prometheus_by_string_name_is_initialized_once(route: Route, provi
provider.expected_requests = 2
litellm.success_callback = ["prometheus"] # test-quality-ok: public registration; fixture restores globals
recorder: Final = RecordingLogger()
await route.invoke(provider, callbacks=[recorder])
await route.invoke(provider, callbacks=[recorder])
await invoke_native(route, provider, callbacks=[recorder])
await invoke_native(route, provider, callbacks=[recorder])
await wait_for_callback(route, recorder, count=2)
instances: Final = [cb for cb in litellm_logging._in_memory_loggers if isinstance(cb, PrometheusLogger)] # pyright: ignore[reportPrivateUsage] # string-name cache has no public accessor
assert len(instances) == 1
@ -218,6 +248,6 @@ async def test_sync_ocr_reaches_sync_hooks_only(
recording_server: RecordingServer, otel: OtelHarness, prometheus: PrometheusLogger
) -> None:
recording_server.default_response = ResponseSpec(body=OCR_SYNC.provider_response)
await OCR_SYNC.invoke(recording_server, callbacks=[otel.logger, prometheus])
await invoke_native(OCR_SYNC, recording_server, callbacks=[otel.logger, prometheus])
assert len(await otel.wait_for_spans()) == 1
assert metric_value("litellm_requests_metric_total", model=OCR_SYNC.provider_model) == 0

28
tests/test_litellm_rust/messages/test_callback_mutation.py
View file

@ -69,10 +69,10 @@ async def test_messages_pre_call_receives_expected_provider_request(messages_ser
@pytest.mark.asyncio
@pytest.mark.parametrize("stream", [False, True])
@pytest.mark.parametrize("raise_after_edit", [False, True])
@pytest.mark.parametrize("native", [False, True])
async def test_messages_pre_call_edits_reach_later_callbacks_and_provider(
@pytest.mark.parametrize("stream", [False, True], ids=["response", "stream"])
@pytest.mark.parametrize("raise_after_edit", [False, True], ids=["edit-returns", "edit-raises"])
@pytest.mark.parametrize("native", [False, True], ids=["python", "rust"])
async def test_messages_pre_call_edits_reach_later_callbacks_but_only_headers_reach_provider(
messages_server: RecordingServer, raise_after_edit: bool, native: bool, stream: bool
) -> None:
litellm.rust(native)
@ -100,6 +100,7 @@ async def test_messages_pre_call_edits_reach_later_callbacks_and_provider(
assert observed[0][1]["x-audit-tag"] == "reviewed"
assert "temperature" not in messages_server.requests[0].body
assert messages_server.requests[0].headers["x-audit-tag"] == "reviewed"
assert has_rust_response_marker(response) is native
@pytest.mark.asyncio
@ -143,23 +144,28 @@ async def test_messages_pre_call_state_reaches_terminal_callbacks(messages_serve
@pytest.mark.asyncio
@pytest.mark.parametrize("native", [False, True])
@pytest.mark.parametrize("native", [False, True], ids=["python", "rust"])
@pytest.mark.parametrize("provider", ["anthropic", "azure_ai"])
@pytest.mark.parametrize("raise_after_edit", [False, True])
@pytest.mark.parametrize("raise_after_edit", [False, True], ids=["edit-returns", "edit-raises"])
async def test_messages_retained_aliases_preserve_identity_and_snapshot_timing(
messages_server: RecordingServer, native: bool, provider: str, raise_after_edit: bool
) -> None:
litellm.rust(native)
retained: Final = []
observed: Final = []
aliases: Final = []
class Retain(CustomLogger):
def log_pre_api_call(self, model, messages, kwargs):
body = request_body(kwargs)
message = kwargs["messages"][0]
assert body["messages"][0] is message
assert body["messages"][0]["content"] is message["content"]
assert body["messages"][0]["content"][0] is message["content"][0]
aliases.append(
(
body["messages"][0] is message,
body["messages"][0]["content"] is message["content"],
body["messages"][0]["content"][0] is message["content"][0],
)
)
retained.append(message["content"][0])
class Edit(CustomLogger):
@ -181,11 +187,11 @@ async def test_messages_retained_aliases_preserve_identity_and_snapshot_timing(
api_base=messages_server.base_url,
callbacks=[Retain(), Edit(), Observe()],
)
assert aliases == [(True, True, True)]
assert observed == [(True, "changed through retained reference")]
assert retained[0]["text"] == "changed through retained reference"
assert messages_server.requests[0].body["messages"][0]["content"][0]["text"] == "original"
if native:
assert response["_hidden_params"]["additional_headers"]["x-litellm-rust"] == "true"
assert has_rust_response_marker(response) is native
@pytest.mark.asyncio

8
tests/test_litellm_rust/ocr/test_callbacks.py
View file

@ -98,8 +98,8 @@ def test_pre_call_header_edits_reach_later_callbacks_and_provider(ocr_server: Re
@pytest.mark.asyncio
@pytest.mark.parametrize("rust_enabled", [False, True])
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("rust_enabled", [False, True], ids=["python", "rust"])
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
async def test_pre_call_nested_mutation_updates_retained_references(
ocr_server: RecordingServer, rust_enabled: bool, asynchronous: bool
) -> None:
@ -107,10 +107,11 @@ async def test_pre_call_nested_mutation_updates_retained_references(
original: Final = dict(OCR_DOCUMENT)
replacement_url: Final = "data:application/pdf;base64,ZGVm"
retained: Final = []
aliases: Final = []
class Retain(CustomLogger):
def log_pre_api_call(self, model, messages, kwargs):
assert request_body(kwargs)["document"] is original
aliases.append(request_body(kwargs)["document"] is original)
retained.append(request_body(kwargs)["document"])
class Edit(CustomLogger):
@ -126,6 +127,7 @@ async def test_pre_call_nested_mutation_updates_retained_references(
}
response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments)
assert aliases == [True]
assert retained[0]["document_url"] == replacement_url
assert original["document_url"] == replacement_url
assert ocr_server.requests[0].body["document"]["document_url"] == replacement_url

80
tests/test_litellm_rust/test_ocr.py Normal file
View file

@ -0,0 +1,80 @@
import json
import threading
from collections.abc import Generator
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Final
import pytest
import litellm
pytestmark = pytest.mark.requires_rust_extension
@pytest.fixture
def native_only_ocr_server() -> Generator[tuple[ThreadingHTTPServer, list[dict[str, object]]]]:
requests: Final[list[dict[str, object]]] = []
class Handler(BaseHTTPRequestHandler):
def do_POST(self) -> None:
requests.append(
{
"headers": {name.lower(): value for name, value in self.headers.items()},
"body": json.loads(self.rfile.read(int(self.headers["Content-Length"]))),
}
)
if self.headers.get("User-Agent", "").startswith("python-httpx"):
self.send_response(418)
self.end_headers()
return
response: Final = json.dumps(
{
"pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}],
"model": "mistral-ocr-latest",
"usage_info": {"pages_processed": 1, "doc_size_bytes": 3},
}
).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(response)))
self.end_headers()
self.wfile.write(response)
def log_message(self, format: str, *args: object) -> None:
pass
server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
thread: Final = threading.Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True)
thread.start()
try:
yield server, requests
finally:
server.shutdown()
server.server_close()
thread.join()
def test_public_ocr_executes_the_compiled_extension_without_python_fallback(
native_only_ocr_server: tuple[ThreadingHTTPServer, list[dict[str, object]]],
) -> None:
server, requests = native_only_ocr_server
address: Final = server.server_address
host: Final = str(address[0])
port: Final = int(address[1])
response: Final = litellm.ocr(
model="mistral/mistral-ocr-latest",
document={"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
api_key="test-key",
api_base=f"http://{host}:{port}",
)
assert response.pages[0].markdown == "native OCR response"
assert len(requests) == 1
headers: Final = requests[0]["headers"]
assert isinstance(headers, dict)
assert not str(headers.get("user-agent", "")).startswith("python-httpx")
assert requests[0]["body"] == {
"model": "mistral-ocr-latest",
"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
}