mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
fix(rust): isolate callback tests and preserve messages logging lifecycle
This commit is contained in:
Yujong Lee
parent
aaace234bb
commit
02572d815f
8 changed files with 662 additions and 58 deletions
|
|
@ -199,11 +199,7 @@ fn prepare(py: Python<'_>, arguments: Py<PyDict>) -> PyResult<Py<MessagesState>>
|
|||
let message = PyDict::new(py);
|
||||
message.set_item("role", "user")?;
|
||||
message.set_item("content", serialized)?;
|
||||
let messages = vec![message];
|
||||
logging
|
||||
.bind(py)
|
||||
.call_method1("update_messages", (&messages,))?;
|
||||
kwargs.set_item("input", messages)?;
|
||||
kwargs.set_item("input", vec![message])?;
|
||||
kwargs.set_item("api_key", "")?;
|
||||
kwargs.set_item("additional_args", additional)?;
|
||||
logging
|
||||
|
|
@ -340,7 +336,7 @@ class Host:
|
|||
self.current = arguments
|
||||
self.asynchronous = asynchronous
|
||||
self.logger = arguments.get('litellm_logging_obj')
|
||||
self.deployment_hooks_owned = self.logger is None
|
||||
self.lifecycle_owned = self.logger is None
|
||||
self.state = None
|
||||
self.response = None
|
||||
self.error = None
|
||||
|
|
@ -354,7 +350,7 @@ class Host:
|
|||
self.streaming = self.logger.stream is True
|
||||
|
||||
async def deployment_pre(self):
|
||||
if not self.deployment_hooks_owned:
|
||||
if not self.lifecycle_owned:
|
||||
return
|
||||
modified = await utils.async_pre_call_deployment_hook(self.current, 'anthropic_messages')
|
||||
if modified is not None:
|
||||
|
|
@ -373,7 +369,7 @@ class Host:
|
|||
self.end = datetime.now()
|
||||
|
||||
async def deployment_success(self):
|
||||
if self.deployment_hooks_owned:
|
||||
if self.lifecycle_owned:
|
||||
self.response = await utils.async_post_call_success_deployment_hook(self.current, self.response, CallTypes.aanthropic_messages)
|
||||
if self.streaming:
|
||||
self.response = retain_stream_response(
|
||||
|
|
@ -384,11 +380,11 @@ class Host:
|
|||
)
|
||||
|
||||
async def deployment_failure(self):
|
||||
if self.deployment_hooks_owned:
|
||||
if self.lifecycle_owned:
|
||||
await utils.async_post_call_failure_deployment_hook(self.current, self.error, 'anthropic_messages')
|
||||
|
||||
def terminal(self, action, value):
|
||||
if self.streaming:
|
||||
if self.streaming or not self.lifecycle_owned:
|
||||
return None
|
||||
return invoke_terminal(action, (self.arguments, self.current, self.state), self.logger, None, value, self.start, self.end)
|
||||
|
||||
|
|
@ -398,7 +394,7 @@ class Host:
|
|||
def sync_failure(self): return self.terminal('sync_failure', self.error)
|
||||
def async_failure(self): return self.terminal('async_failure', self.error)
|
||||
def restore(self):
|
||||
if not self.streaming:
|
||||
if not self.streaming and self.lifecycle_owned:
|
||||
utils._restore_correlation_context_if_supported(self.logger)
|
||||
|
||||
def advance(self, outcome, error=None):
|
||||
|
|
|
|||
24
tests/_prometheus_helpers.py
Normal file
24
tests/_prometheus_helpers.py
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
from typing import Final
|
||||
|
||||
from prometheus_client import REGISTRY, CollectorRegistry
|
||||
|
||||
|
||||
def clear_prometheus_registry() -> None:
|
||||
for collector in tuple(REGISTRY._collector_to_names): # pyright: ignore[reportPrivateUsage] # prometheus_client has no public collector enumeration
|
||||
REGISTRY.unregister(collector)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def isolated_prometheus_registry(registry: CollectorRegistry = REGISTRY) -> Iterator[None]:
|
||||
original: Final = tuple(registry._collector_to_names) # pyright: ignore[reportPrivateUsage] # prometheus_client has no public collector enumeration
|
||||
for collector in original:
|
||||
registry.unregister(collector)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
for collector in tuple(registry._collector_to_names): # pyright: ignore[reportPrivateUsage] # remove collectors created inside the isolation scope
|
||||
registry.unregister(collector)
|
||||
for collector in original:
|
||||
registry.register(collector)
|
||||
|
|
@ -6,6 +6,14 @@ from dataclasses import dataclass
|
|||
from typing import Final
|
||||
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker
|
||||
|
||||
|
||||
async def drain_logging(worker: LoggingWorker = GLOBAL_LOGGING_WORKER) -> None:
|
||||
# The wrapper's scheduled helper enqueues without awaiting; let it run before joining the queue.
|
||||
await asyncio.sleep(0)
|
||||
worker.start()
|
||||
await asyncio.wait_for(worker.flush(), timeout=10)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
|
|||
|
|
@ -1,49 +1,106 @@
|
|||
import asyncio
|
||||
import os
|
||||
from collections.abc import Iterator
|
||||
from typing import Final
|
||||
from collections.abc import AsyncIterator, Generator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from types import ModuleType
|
||||
from typing import Final, cast
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
import litellm
|
||||
from litellm.rust_bridge.configuration import reset_rust_configuration
|
||||
from litellm import utils
|
||||
from litellm.litellm_core_utils import litellm_logging
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.rust_bridge.configuration import ( # pyright: ignore[reportPrivateUsage] # share the canonical env parsing with the module under test
|
||||
_parse_env_bool,
|
||||
reset_rust_configuration,
|
||||
)
|
||||
from tests._prometheus_helpers import isolated_prometheus_registry
|
||||
from tests.test_litellm_rust.callback_recorder import drain_logging
|
||||
from tests.test_litellm_rust.integrations import (
|
||||
otel, # noqa: F401 # pytest fixture export
|
||||
prometheus, # noqa: F401 # pytest fixture export
|
||||
provider, # noqa: F401 # pytest fixture export
|
||||
)
|
||||
from tests.test_litellm_rust.recording_server import recording_server # noqa: F401 # pytest fixture export
|
||||
|
||||
CALLBACK_ATTRIBUTES: Final = (
|
||||
"callbacks",
|
||||
"input_callback",
|
||||
"success_callback",
|
||||
"failure_callback",
|
||||
"_async_input_callback",
|
||||
"_async_success_callback",
|
||||
"_async_failure_callback",
|
||||
)
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def isolate_rust_state() -> Iterator[None]:
|
||||
callback_attributes: Final = (
|
||||
"callbacks",
|
||||
"input_callback",
|
||||
"success_callback",
|
||||
"failure_callback",
|
||||
"_async_input_callback",
|
||||
"_async_success_callback",
|
||||
"_async_failure_callback",
|
||||
)
|
||||
original_callbacks: Final = {attribute: list(getattr(litellm, attribute)) for attribute in callback_attributes}
|
||||
original_cache: Final = litellm.cache
|
||||
for attribute in callback_attributes:
|
||||
getattr(litellm, attribute).clear()
|
||||
litellm.cache = None # test-quality-ok: isolate the process-global cache from native extension tests
|
||||
|
||||
def _list_attribute(container: ModuleType, attribute: str) -> list[object]:
|
||||
value: Final = getattr(container, attribute)
|
||||
if not isinstance(value, list):
|
||||
raise AssertionError(f"{container.__name__}.{attribute} is not a list")
|
||||
return cast(list[object], value)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _isolated_list(container: ModuleType, attribute: str) -> Generator[None]:
|
||||
source: Final = _list_attribute(container, attribute)
|
||||
original: Final = list(source)
|
||||
source.clear() # mutable-ok: test isolation mutates global registries by design
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
source.clear()
|
||||
source.extend(original)
|
||||
setattr(container, attribute, source)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _rebound(container: ModuleType, attribute: str, value: object) -> Generator[None]:
|
||||
original: Final[object] = getattr(container, attribute)
|
||||
setattr(container, attribute, value)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
setattr(container, attribute, original)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _rust_mode() -> Generator[None]:
|
||||
reset_rust_configuration()
|
||||
litellm.rust(True)
|
||||
yield
|
||||
for attribute, callbacks in original_callbacks.items():
|
||||
target = getattr(litellm, attribute)
|
||||
target.clear()
|
||||
target.extend(callbacks)
|
||||
litellm.cache = original_cache # test-quality-ok: restore the process-global cache after native extension tests
|
||||
reset_rust_configuration()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
reset_rust_configuration()
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(autouse=True, loop_scope="function")
|
||||
async def isolate_rust_state() -> AsyncIterator[ExitStack]:
|
||||
with ExitStack() as stack:
|
||||
for attribute in CALLBACK_ATTRIBUTES:
|
||||
stack.enter_context(_isolated_list(litellm, attribute))
|
||||
stack.enter_context(_isolated_list(litellm_logging, "_in_memory_loggers")) # pyright: ignore[reportPrivateUsage] # string-name callback cache has no public accessor
|
||||
stack.enter_context(_rebound(utils, "callback_list", [])) # rebind-ok: legacy global registry mutated by set_callbacks
|
||||
stack.enter_context(_rebound(litellm, "cache", None)) # test-quality-ok: isolate the process-global cache from native extension tests
|
||||
stack.enter_context(isolated_prometheus_registry())
|
||||
stack.enter_context(_rust_mode())
|
||||
executor: Final = ThreadPoolExecutor(thread_name_prefix="rust-test-logging")
|
||||
stack.enter_context(_rebound(utils, "executor", executor))
|
||||
try:
|
||||
yield stack
|
||||
finally:
|
||||
try:
|
||||
await drain_logging()
|
||||
finally:
|
||||
await asyncio.to_thread(executor.shutdown, wait=True)
|
||||
await GLOBAL_LOGGING_WORKER.stop()
|
||||
|
||||
|
||||
def pytest_collection_modifyitems(items):
|
||||
rust_enabled = os.environ.get("LITELLM_RUST", "").strip().lower() in {
|
||||
"1",
|
||||
"true",
|
||||
"yes",
|
||||
"on",
|
||||
}
|
||||
if not rust_enabled:
|
||||
if not _parse_env_bool(os.environ.get("LITELLM_RUST")):
|
||||
skip = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension")
|
||||
for item in items:
|
||||
item.add_marker(skip)
|
||||
|
|
|
|||
201
tests/test_litellm_rust/integrations.py
Normal file
201
tests/test_litellm_rust/integrations.py
Normal file
|
|
@ -0,0 +1,201 @@
|
|||
import asyncio
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
from contextlib import ExitStack
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
|
||||
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
|
||||
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
|
||||
from prometheus_client import REGISTRY
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.opentelemetry import LITELLM_REQUEST_SPAN_NAME, OpenTelemetry, OpenTelemetryConfig
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
from tests.test_litellm_rust.contracts import (
|
||||
MESSAGES,
|
||||
MESSAGES_MODEL,
|
||||
MESSAGES_RESPONSE,
|
||||
OCR_RESPONSE,
|
||||
call_native_aocr,
|
||||
call_native_ocr,
|
||||
)
|
||||
from tests.test_litellm_rust.recording_server import RecordingServer, ResponseSpec
|
||||
from tests.test_litellm_rust.callback_recorder import drain_logging
|
||||
|
||||
RouteName = Literal["ocr-sync", "ocr-async", "messages", "messages-stream"]
|
||||
GuardrailObservation = tuple[Literal["request", "response"], tuple[str, ...]]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Route:
|
||||
name: RouteName
|
||||
call_type: str
|
||||
provider_model: str
|
||||
provider_response: dict[str, object]
|
||||
response_text: str
|
||||
provider: str
|
||||
expected_cost: float
|
||||
logging_only_scan: GuardrailObservation
|
||||
fires_async_hooks: bool
|
||||
|
||||
async def invoke(self, server: RecordingServer, **kwargs: object) -> object:
|
||||
match self.name:
|
||||
case "ocr-sync":
|
||||
return await asyncio.to_thread(call_native_ocr, server, **kwargs)
|
||||
case "ocr-async":
|
||||
return await call_native_aocr(server, **kwargs)
|
||||
case "messages":
|
||||
return await _call_messages(server, **kwargs)
|
||||
case "messages-stream":
|
||||
stream: Final = await _call_messages(server, stream=True, **kwargs)
|
||||
return [chunk async for chunk in stream]
|
||||
|
||||
|
||||
async def _call_messages(server: RecordingServer, **kwargs: object) -> object:
|
||||
return await litellm.anthropic.messages.acreate(
|
||||
model=MESSAGES_MODEL,
|
||||
messages=MESSAGES,
|
||||
max_tokens=64,
|
||||
api_key="test-key",
|
||||
api_base=server.base_url,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
MESSAGES_COST: Final = 5 * 3e-06 + 4 * 1.5e-05
|
||||
OCR_COST: Final = 0.004
|
||||
|
||||
OCR_SYNC: Final = Route(
|
||||
name="ocr-sync",
|
||||
call_type="ocr",
|
||||
provider_model="mistral-ocr-latest",
|
||||
provider_response=OCR_RESPONSE,
|
||||
response_text="native OCR response",
|
||||
provider="mistral",
|
||||
expected_cost=OCR_COST,
|
||||
logging_only_scan=("response", ("native OCR response",)),
|
||||
fires_async_hooks=False,
|
||||
)
|
||||
OCR_ASYNC: Final = Route(
|
||||
name="ocr-async",
|
||||
call_type="aocr",
|
||||
provider_model="mistral-ocr-latest",
|
||||
provider_response=OCR_RESPONSE,
|
||||
response_text="native OCR response",
|
||||
provider="mistral",
|
||||
expected_cost=OCR_COST,
|
||||
logging_only_scan=("response", ("native OCR response",)),
|
||||
fires_async_hooks=True,
|
||||
)
|
||||
MESSAGES_ROUTE: Final = Route(
|
||||
name="messages",
|
||||
call_type="anthropic_messages",
|
||||
provider_model="claude-sonnet-4-5-20250929",
|
||||
provider_response=MESSAGES_RESPONSE,
|
||||
response_text="Hello from native Messages",
|
||||
provider="anthropic",
|
||||
expected_cost=MESSAGES_COST,
|
||||
logging_only_scan=("request", ("Hello",)),
|
||||
fires_async_hooks=True,
|
||||
)
|
||||
MESSAGES_STREAM: Final = Route(
|
||||
name="messages-stream",
|
||||
call_type="anthropic_messages",
|
||||
provider_model="claude-sonnet-4-5-20250929",
|
||||
provider_response=MESSAGES_RESPONSE,
|
||||
response_text="Hello from native Messages",
|
||||
provider="anthropic",
|
||||
expected_cost=MESSAGES_COST,
|
||||
logging_only_scan=("request", ("Hello",)),
|
||||
fires_async_hooks=True,
|
||||
)
|
||||
ALL_ROUTES: Final = (OCR_SYNC, OCR_ASYNC, MESSAGES_ROUTE, MESSAGES_STREAM)
|
||||
ASYNC_ROUTES: Final = tuple(route for route in ALL_ROUTES if route.fires_async_hooks)
|
||||
NON_STREAM_ASYNC_ROUTES: Final = (OCR_ASYNC, MESSAGES_ROUTE)
|
||||
|
||||
|
||||
def route_id(route: Route) -> str:
|
||||
return route.name
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def provider(recording_server: RecordingServer, route: Route) -> RecordingServer:
|
||||
recording_server.default_response = ResponseSpec(body=route.provider_response)
|
||||
return recording_server
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class OtelHarness:
|
||||
logger: OpenTelemetry
|
||||
exporter: InMemorySpanExporter
|
||||
|
||||
def spans(self, name: str = LITELLM_REQUEST_SPAN_NAME) -> tuple[ReadableSpan, ...]:
|
||||
return tuple(span for span in self.exporter.get_finished_spans() if span.name == name)
|
||||
|
||||
async def wait_for_spans(
|
||||
self, name: str = LITELLM_REQUEST_SPAN_NAME, count: int = 1, timeout: float = 10
|
||||
) -> tuple[ReadableSpan, ...]:
|
||||
deadline: Final = time.monotonic() + timeout
|
||||
while len(self.spans(name)) < count:
|
||||
if time.monotonic() >= deadline:
|
||||
raise TimeoutError(f"Timed out waiting for {count} {name} spans; saw {self.spans(name)}")
|
||||
await asyncio.sleep(0.01)
|
||||
await drain_logging()
|
||||
return self.spans(name)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def otel(isolate_rust_state: ExitStack) -> OtelHarness:
|
||||
exporter: Final = InMemorySpanExporter()
|
||||
tracer_provider: Final = TracerProvider()
|
||||
isolate_rust_state.callback(tracer_provider.shutdown)
|
||||
tracer_provider.add_span_processor(SimpleSpanProcessor(exporter))
|
||||
logger: Final = OpenTelemetry(config=OpenTelemetryConfig(exporter=exporter), tracer_provider=tracer_provider)
|
||||
return OtelHarness(logger=logger, exporter=exporter)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def prometheus() -> PrometheusLogger:
|
||||
return PrometheusLogger()
|
||||
|
||||
|
||||
def metric_value(name: str, **labels: str) -> float:
|
||||
for metric in REGISTRY.collect():
|
||||
for sample in metric.samples:
|
||||
if sample.name == name and all(sample.labels.get(key) == value for key, value in labels.items()):
|
||||
return sample.value
|
||||
return 0.0
|
||||
|
||||
|
||||
class RecordingGuardrail(CustomGuardrail):
|
||||
def __init__(self, guardrail_name: str = "rust-review", fail_with: Exception | None = None) -> None:
|
||||
super().__init__(
|
||||
guardrail_name=guardrail_name,
|
||||
event_hook=GuardrailEventHooks.logging_only,
|
||||
default_on=True,
|
||||
)
|
||||
self.observations: list[GuardrailObservation] = []
|
||||
self._fail_with = fail_with
|
||||
|
||||
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None) -> GenericGuardrailAPIInputs:
|
||||
self.observations.append((input_type, tuple(inputs.get("texts") or ())))
|
||||
if self._fail_with is not None:
|
||||
raise self._fail_with
|
||||
return inputs
|
||||
|
||||
|
||||
class ReviewGuardrail(CustomGuardrail):
|
||||
def __init__(self, review: Callable[[object], Awaitable[object]]) -> None:
|
||||
super().__init__(guardrail_name="rust-review", event_hook=GuardrailEventHooks.post_call, default_on=True)
|
||||
self._review = review
|
||||
self.call_types: list[object] = []
|
||||
|
||||
async def async_post_call_success_deployment_hook(self, request_data, response, call_type):
|
||||
self.call_types.append(call_type)
|
||||
return await self._review(response)
|
||||
29
tests/test_litellm_rust/test_conftest.py
Normal file
29
tests/test_litellm_rust/test_conftest.py
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
from types import ModuleType
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.test_litellm_rust.conftest import _isolated_list # pyright: ignore[reportPrivateUsage] # exercise fixture restoration without mutating SDK state
|
||||
|
||||
pytestmark = pytest.mark.requires_rust_extension
|
||||
|
||||
|
||||
def test_callback_list_restores_identity_after_rebinding_and_failure() -> None:
|
||||
container: Final = ModuleType("callback_registry")
|
||||
callback: Final = object()
|
||||
original: Final = [callback]
|
||||
setattr(container, "callbacks", original)
|
||||
|
||||
def failing_test() -> None:
|
||||
with _isolated_list(container, "callbacks"):
|
||||
assert getattr(container, "callbacks") is original
|
||||
assert original == []
|
||||
original.append(object())
|
||||
setattr(container, "callbacks", [object()])
|
||||
raise RuntimeError("test failed")
|
||||
|
||||
with pytest.raises(RuntimeError, match="test failed"):
|
||||
failing_test()
|
||||
|
||||
assert getattr(container, "callbacks") is original
|
||||
assert original == [callback]
|
||||
272
tests/test_litellm_rust/test_integrations.py
Normal file
272
tests/test_litellm_rust/test_integrations.py
Normal file
|
|
@ -0,0 +1,272 @@
|
|||
import json
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from opentelemetry.trace import StatusCode
|
||||
from prometheus_client import CollectorRegistry, Counter
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.litellm_core_utils import litellm_logging
|
||||
from litellm.types.utils import CallTypes
|
||||
from tests._prometheus_helpers import isolated_prometheus_registry
|
||||
from tests.test_litellm_rust.callback_recorder import RecordingLogger, drain_logging
|
||||
from tests.test_litellm_rust.integrations import (
|
||||
ALL_ROUTES,
|
||||
ASYNC_ROUTES,
|
||||
MESSAGES_ROUTE,
|
||||
NON_STREAM_ASYNC_ROUTES,
|
||||
OCR_ASYNC,
|
||||
OCR_SYNC,
|
||||
OtelHarness,
|
||||
RecordingGuardrail,
|
||||
ReviewGuardrail,
|
||||
Route,
|
||||
metric_value,
|
||||
route_id,
|
||||
)
|
||||
from tests.test_litellm_rust.recording_server import RecordingServer, ResponseSpec
|
||||
|
||||
pytestmark = pytest.mark.requires_rust_extension
|
||||
|
||||
FAILURE_RESPONSE: Final = ResponseSpec(body={"message": "provider unavailable"}, status=500)
|
||||
|
||||
|
||||
def test_prometheus_registry_restores_collectors_after_failure() -> None:
|
||||
registry: Final = CollectorRegistry()
|
||||
original: Final = Counter("original", "Original collector", registry=registry)
|
||||
original.inc(2)
|
||||
|
||||
def failing_test() -> None:
|
||||
with isolated_prometheus_registry(registry):
|
||||
assert registry.get_sample_value("original_total") is None
|
||||
Counter("original", "Temporary replacement", registry=registry).inc(7)
|
||||
Counter("temporary", "Temporary collector", registry=registry).inc()
|
||||
raise RuntimeError("test failed")
|
||||
|
||||
with pytest.raises(RuntimeError, match="test failed"):
|
||||
failing_test()
|
||||
|
||||
assert registry.get_sample_value("original_total") == 2
|
||||
assert registry.get_sample_value("temporary_total") is None
|
||||
registry.unregister(original)
|
||||
assert registry.get_sample_value("original_total") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("route", ALL_ROUTES, ids=route_id)
|
||||
async def test_otel_emits_one_request_span_on_success(
|
||||
route: Route, provider: RecordingServer, otel: OtelHarness
|
||||
) -> None:
|
||||
await route.invoke(provider, callbacks=[otel.logger])
|
||||
|
||||
spans: Final = await otel.wait_for_spans()
|
||||
assert len(spans) == 1
|
||||
span: Final = spans[0]
|
||||
assert span.status.status_code is StatusCode.OK
|
||||
assert span.attributes["llm.request.type"] == route.call_type
|
||||
assert span.attributes["gen_ai.request.model"] == route.provider_model
|
||||
assert json.loads(span.attributes["hidden_params"])["response_cost"] == pytest.approx(route.expected_cost)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_otel_stream_span_appears_only_after_exhaustion(
|
||||
otel: OtelHarness, recording_server: RecordingServer
|
||||
) -> None:
|
||||
recording_server.default_response = ResponseSpec(body=MESSAGES_ROUTE.provider_response)
|
||||
|
||||
stream: Final = await MESSAGES_ROUTE.invoke(recording_server, callbacks=[otel.logger], stream=True)
|
||||
await drain_logging()
|
||||
assert otel.spans() == ()
|
||||
|
||||
chunks: Final = [chunk async for chunk in stream]
|
||||
|
||||
assert chunks
|
||||
assert len(await otel.wait_for_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(
|
||||
route: Route, provider: RecordingServer, otel: OtelHarness
|
||||
) -> None:
|
||||
provider.enqueue(FAILURE_RESPONSE)
|
||||
|
||||
with pytest.raises(litellm.InternalServerError):
|
||||
await route.invoke(provider, callbacks=[otel.logger])
|
||||
|
||||
spans: Final = await otel.wait_for_spans()
|
||||
assert len(spans) == 1
|
||||
assert spans[0].status.status_code is StatusCode.ERROR
|
||||
exception_events: Final = [event for event in spans[0].events if event.name == "exception"]
|
||||
assert len(exception_events) == 1
|
||||
assert "InternalServerError" in exception_events[0].attributes["exception.type"]
|
||||
assert spans[0].attributes["llm.request.type"] == route.call_type
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("route", ALL_ROUTES, ids=route_id)
|
||||
async def test_otel_and_custom_logger_observe_same_standard_logging_object(
|
||||
route: Route, provider: RecordingServer, otel: OtelHarness
|
||||
) -> None:
|
||||
recorder: Final = RecordingLogger()
|
||||
|
||||
await route.invoke(provider, callbacks=[otel.logger, recorder])
|
||||
|
||||
events: Final = (
|
||||
await recorder.wait_for_async("async_log_success_event")
|
||||
if route.fires_async_hooks
|
||||
else recorder.wait_for("log_success_event")
|
||||
)
|
||||
payload: Final = events[0].kwargs["standard_logging_object"]
|
||||
span: Final = (await otel.wait_for_spans())[0]
|
||||
assert payload["call_type"] == route.call_type
|
||||
assert payload["model"] == route.provider_model
|
||||
assert json.loads(span.attributes["hidden_params"]) == payload["hidden_params"]
|
||||
assert span.attributes["llm.request.type"] == payload["call_type"]
|
||||
assert span.attributes["litellm.provider.model"] == payload["model"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("route", ASYNC_ROUTES, ids=route_id)
|
||||
async def test_prometheus_counts_one_successful_request(
|
||||
route: Route, provider: RecordingServer, prometheus: PrometheusLogger
|
||||
) -> None:
|
||||
before: Final = metric_value("litellm_requests_metric_total", model=route.provider_model)
|
||||
recorder: Final = RecordingLogger()
|
||||
|
||||
await route.invoke(provider, callbacks=[prometheus, recorder])
|
||||
await recorder.wait_for_async("async_log_success_event")
|
||||
|
||||
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
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prometheus_counts_tokens_from_messages_usage(
|
||||
recording_server: RecordingServer, prometheus: PrometheusLogger
|
||||
) -> None:
|
||||
recording_server.default_response = ResponseSpec(body=MESSAGES_ROUTE.provider_response)
|
||||
recorder: Final = RecordingLogger()
|
||||
|
||||
await MESSAGES_ROUTE.invoke(recording_server, callbacks=[prometheus, recorder])
|
||||
await recorder.wait_for_async("async_log_success_event")
|
||||
|
||||
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
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prometheus_counts_one_failed_request(
|
||||
otel: OtelHarness, prometheus: PrometheusLogger, recording_server: RecordingServer
|
||||
) -> None:
|
||||
recording_server.enqueue(FAILURE_RESPONSE)
|
||||
|
||||
with pytest.raises(litellm.InternalServerError):
|
||||
await OCR_ASYNC.invoke(recording_server, callbacks=[prometheus, otel.logger])
|
||||
|
||||
assert len(await otel.wait_for_spans()) == 1
|
||||
assert metric_value("litellm_llm_api_failed_requests_metric_total", model=OCR_ASYNC.provider_model) == 1
|
||||
assert metric_value("litellm_requests_metric_total", model=OCR_ASYNC.provider_model) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("route", NON_STREAM_ASYNC_ROUTES, ids=route_id)
|
||||
async def test_prometheus_by_string_name_is_initialized_once(route: Route, provider: RecordingServer) -> None:
|
||||
provider.expected_requests = 2
|
||||
litellm.success_callback = ["prometheus"] # test-quality-ok: tests public string-name registration; isolate_rust_state restores this registry
|
||||
recorder: Final = RecordingLogger()
|
||||
|
||||
await route.invoke(provider, callbacks=[recorder])
|
||||
await route.invoke(provider, callbacks=[recorder])
|
||||
await recorder.wait_for_async("async_log_success_event", 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
|
||||
assert metric_value("litellm_requests_metric_total", model=route.provider_model) == 2
|
||||
assert "prometheus" not in litellm.success_callback
|
||||
assert instances[0] in litellm._async_success_callback # pyright: ignore[reportPrivateUsage] # callback registry has no public accessor
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
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])
|
||||
|
||||
assert len(await otel.wait_for_spans()) == 1
|
||||
assert metric_value("litellm_requests_metric_total", model=OCR_SYNC.provider_model) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("route", NON_STREAM_ASYNC_ROUTES, ids=route_id)
|
||||
async def test_logging_only_guardrail_verdict_reaches_otel_and_custom_logger(
|
||||
route: Route, provider: RecordingServer, otel: OtelHarness
|
||||
) -> None:
|
||||
guardrail: Final = RecordingGuardrail()
|
||||
recorder: Final = RecordingLogger()
|
||||
|
||||
await route.invoke(provider, callbacks=[guardrail, otel.logger, recorder])
|
||||
|
||||
payload: Final = (await recorder.wait_for_async("async_log_success_event"))[0].kwargs["standard_logging_object"]
|
||||
assert guardrail.observations == [route.logging_only_scan]
|
||||
verdicts: Final = payload["guardrail_information"]
|
||||
assert [verdict["guardrail_name"] for verdict in verdicts] == ["rust-review"]
|
||||
assert verdicts[0]["guardrail_status"] == "success"
|
||||
guardrail_spans: Final = await otel.wait_for_spans("guardrail")
|
||||
assert len(guardrail_spans) == 1
|
||||
assert guardrail_spans[0].attributes["guardrail_name"] == "rust-review"
|
||||
assert guardrail_spans[0].attributes["guardrail_status"] == "success"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("route", NON_STREAM_ASYNC_ROUTES, ids=route_id)
|
||||
async def test_logging_only_guardrail_failure_does_not_block_loggers(
|
||||
route: Route, provider: RecordingServer, otel: OtelHarness, prometheus: PrometheusLogger
|
||||
) -> None:
|
||||
guardrail: Final = RecordingGuardrail(fail_with=RuntimeError("review service unavailable"))
|
||||
recorder: Final = RecordingLogger()
|
||||
|
||||
response: Final = await route.invoke(provider, callbacks=[guardrail, otel.logger, prometheus, recorder])
|
||||
|
||||
assert response is not None
|
||||
assert len(await otel.wait_for_spans()) == 1
|
||||
assert metric_value("litellm_requests_metric_total", model=route.provider_model) == 1
|
||||
payload: Final = (await recorder.wait_for_async("async_log_success_event"))[0].kwargs["standard_logging_object"]
|
||||
assert payload["guardrail_information"][0]["guardrail_status"] == "guardrail_failed_to_respond"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("route", NON_STREAM_ASYNC_ROUTES, ids=route_id)
|
||||
async def test_post_call_guardrail_replacement_is_what_loggers_see(
|
||||
route: Route, provider: RecordingServer, otel: OtelHarness
|
||||
) -> None:
|
||||
async def review(response: object) -> object:
|
||||
match route.name:
|
||||
case "ocr-async":
|
||||
return response.model_copy(
|
||||
update={"pages": [response.pages[0].model_copy(update={"markdown": "Reviewed OCR"})]}
|
||||
)
|
||||
case _:
|
||||
return {**response, "content": [{"type": "text", "text": "Reviewed Messages"}]}
|
||||
|
||||
guardrail: Final = ReviewGuardrail(review)
|
||||
recorder: Final = RecordingLogger()
|
||||
litellm.callbacks.append(guardrail)
|
||||
|
||||
response: Final = await route.invoke(provider, callbacks=[otel.logger, recorder], guardrails=["rust-review"])
|
||||
|
||||
assert guardrail.call_types == [CallTypes(route.call_type)]
|
||||
event: Final = (await recorder.wait_for_async("async_log_success_event"))[0]
|
||||
span: Final = (await otel.wait_for_spans())[0]
|
||||
match route.name:
|
||||
case "ocr-async":
|
||||
assert response.pages[0].markdown == "Reviewed OCR"
|
||||
assert event.response.pages[0].markdown == "Reviewed OCR"
|
||||
case _:
|
||||
assert response["content"][0]["text"] == "Reviewed Messages"
|
||||
assert event.response.choices[0].message.content == "Reviewed Messages"
|
||||
assert "Reviewed Messages" in json.dumps(dict(span.attributes))
|
||||
assert "guardrails" not in provider.requests[0].body
|
||||
|
|
@ -10,7 +10,7 @@ import litellm
|
|||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.types.utils import CallTypes
|
||||
from tests.test_litellm_rust.callback_recorder import RecordingLogger
|
||||
from tests.test_litellm_rust.callback_recorder import RecordingLogger, drain_logging
|
||||
from tests.test_litellm_rust.contracts import (
|
||||
MESSAGES,
|
||||
MESSAGES_MODEL,
|
||||
|
|
@ -54,18 +54,7 @@ async def test_messages_pre_call_receives_expected_provider_request(messages_ser
|
|||
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 messages == MESSAGES
|
||||
assert additional_args["api_base"] == f"{messages_server.base_url}/v1/messages"
|
||||
assert additional_args["complete_input_dict"] == {
|
||||
"model": "claude-sonnet-4-5-20250929",
|
||||
|
|
@ -155,6 +144,32 @@ async def test_messages_callbacks_run_once(messages_server: RecordingServer) ->
|
|||
assert "async_log_failure_event" not in recorder.names
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_messages_logging_drain_waits_for_suspended_callback(messages_server: RecordingServer) -> None:
|
||||
started: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
finished: Final = asyncio.Event()
|
||||
|
||||
class SuspendedLogger(CustomLogger):
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
started.set()
|
||||
await release.wait()
|
||||
finished.set()
|
||||
|
||||
await call_messages(messages_server, [SuspendedLogger()])
|
||||
draining: Final = asyncio.create_task(drain_logging())
|
||||
try:
|
||||
await asyncio.wait_for(started.wait(), timeout=10)
|
||||
await asyncio.sleep(0)
|
||||
assert not draining.done()
|
||||
assert not finished.is_set()
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.wait_for(draining, timeout=10)
|
||||
|
||||
assert finished.is_set()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.xfail(strict=True, reason="accepted native Messages errors are replayed through the Python transport")
|
||||
async def test_messages_failure_callbacks_receive_original_provider_error(messages_server: RecordingServer) -> None:
|
||||
|
|
@ -268,6 +283,7 @@ async def test_messages_post_call_guardrail_replacement_reaches_caller_and_loggi
|
|||
@pytest.mark.asyncio
|
||||
async def test_messages_logging_hook_replacement_reaches_later_loggers_only(messages_server: RecordingServer) -> None:
|
||||
observations: Final = []
|
||||
exported: Final = asyncio.Event()
|
||||
|
||||
class RecordGuardrailVerdict(CustomLogger):
|
||||
async def async_logging_hook(self, kwargs, result, call_type):
|
||||
|
|
@ -280,9 +296,10 @@ async def test_messages_logging_hook_replacement_reaches_later_loggers_only(mess
|
|||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
observations.append(("export", kwargs["guardrail-verdict"]))
|
||||
exported.set()
|
||||
|
||||
response: Final = await call_messages(messages_server, [RecordGuardrailVerdict(), ExportLog()])
|
||||
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10)
|
||||
await asyncio.wait_for(exported.wait(), timeout=10)
|
||||
|
||||
assert observations == ["guardrail", ("export", "allowed")]
|
||||
assert response["content"][0]["text"] == "Hello from native Messages"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue