scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 267
litellm/tests/test_litellm_rust/responses/test_callbacks.py
yucheng-berri dd511a2b22
fix(logging)!: fire one CustomLogger failure hook per async call outside the async flag list (#47124) 
* fix(logging): fire one CustomLogger hook per async call outside the async flag list

aspeech, atranscription, aimage_generation, aimage_edit, aocr and asearch are
missing from is_sync_litellm_request's flag list, so the sync success and
failure handlers called a CustomLogger's sync hook on top of the async hook the
async handler had already called for the same request.

The sync handlers now skip a CustomLogger's log_success_event and
log_failure_event when that run is paired with an async handler run
(for_async_request) and the async handler gives the logger its own async hook.
The async client wrapper, dispatch_failure_handlers, the proxy-only failure
dispatch and the native bridge's sync failure step for an async call mark their
sync failure run as paired.

Resolves LIT-9486

* fix(logging): keep the sync CustomLogger hook when no async success runs

The async client wrapper skips async success logging for calls carrying a
fallbacks kwarg, and the native bridge skips it for internal and fallbacks
calls, yet the paired sync success run still dropped a both-hooks logger's
sync hook, so that logger got no success event at all.

The sync success run now suppresses the sync hook only when the async success
handler was actually scheduled for the call: the client wrapper passes
async_handler_runs=not is_completion_with_fallbacks, and the native host
records per call whether it enqueued or deferred the async success.

Tests that pinned the doubled failure hook on native aocr, messages and
responses calls and on proxy-rejected requests now expect the async hook
alone. The integration test adds /v1/messages and /v1/responses cells and
runs the key-auth rejection on every route.

* fix(logging): type dispatch_success_handlers kwargs

The new async_handler_runs keyword on success_handler made the untyped
**kwargs that dispatch_success_handlers forwards an unknown argument for it,
which grew the reportUnknownArgumentType count past its merge-base ceiling.
Typing the forwarded kwargs as Any, as _success_handler_body already does,
removes those diagnostics without changing what is forwarded.

* test(integration): cover CustomLogger hook counts across modes, chaos and router strategies

Moves the hook count rig into tests/integration/_support/hook_counts.py and
extends the cells it drives:

- streamed chat, Messages and Responses requests, success and failure
- loggers the async handler never reaches (sync-only overrides, sync event
  lists only) keep one sync hook per request on every surface outside the
  async flag list, key-auth rejections included
- a concurrent burst while the logger's destination stalls and then fails,
  and a burst while a worker is killed, never give a request two hooks
- least-busy keeps counting a request in flight while async speech requests
  finish around it, and usage-based routing v2 counts an async image
  request's tokens once

* test(integration): Datadog and LangSmith deliver one record per async request

Points the built-in Datadog, LangSmith and OpenTelemetry loggers at a local
sink and checks every async surface outside the flag list: one Datadog log
entry and one LangSmith run per request (two on the merge base), and never two
OpenTelemetry spans, which its own emit-once guard already ensured.

* fix(logging): give a both-hooks logger its sync hook when the async run never reaches it

The sync hook was skipped whenever the logger sat in the async callback list, so a
logger the async run never reached got no hook at all: a sibling stalled past the
logging worker timeout, a batch poll that skips default logging, or an async
failure run cut off. The sync run now holds the hook and the async run hands it
back unless it reached the logger

* fix(logging): run the held sync success hooks when the proxy drops a blocked response's success

A post-call guardrail block drops the parked async success, which left the sync
success hook held for it unrun. Base gave a both-hooks logger that hook on these
call types, so the drop now ends the pairing and the held hook runs as before

* fix(logging): run held sync hooks only after every sync run of the request has closed

A held hook released by the async run could run on another thread while the sync
run was still walking its other callbacks. The pairing now waits until the async
run has ended and every sync run has closed, and the last sync run to close runs
the held hooks itself

* fix(logging): end the async run when its coroutine is dropped before it starts

A logging queue that drops the async success coroutine, or a deferred success
that fails to schedule, left the sync hooks held for it unrun. The enqueued async
run now ends when its coroutine is discarded, and the proxy ends it when firing
the deferred success raises

* fix(logging): leave a still-queued async run to the exit flush

The discard finalizer ran at interpreter exit before the logging worker's exit
flush, so a queued request could get its held sync hook and then its async hook

* test: keep the router and file-stream tests in step with main

Main gives each owned proxy a private Redis by default, so the router cells
share the worker's Redis to read its counters, and the file-stream stub takes
the hook pairing the stream now passes

* fix(logging): run held sync hooks when a during_call guardrail raises or the exit flush gives up

A during_call guardrail that raised after the LLM call parked its deferred
success skipped the proxy's flush, and async runs the exit flush never started
were never ended, so a both-hooks logger lost its sync hook in both cases

* fix(logging): skip the sync failure hook when the async failure run logs to the logger

An async call outside is_sync_litellm_request's flag list ran both failure
handlers, so a CustomLogger with both failure hooks got both. The sync run
of an async request now skips log_failure_event for a logger that overrides
async_log_failure_event and is registered for async failures. The four sync
runs of async failures (dispatch_failure_handlers, the @client async wrapper,
the proxy-only failure thread and the Rust host) pass for_async_request=True.

Drops the HookPairing fallback. Successes are left to #47192, which already
gives such a logger only async_log_success_event.

---------

Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>
2026-10-20 15:19:14 -07:00

408 lines
17 KiB
Python
Raw Blame History

import asyncio
from collections.abc import Iterator
from contextvars import ContextVar
from datetime import datetime
from typing import Final
import pytest
from pydantic import TypeAdapter
import litellm
from litellm.integrations.argilla import ArgillaLogger
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.integrations.argilla import ArgillaItem
from litellm.types.utils import StandardLoggingPayload
from tests.test_litellm_rust.responses.data import (
BAD_REQUEST,
RATE_LIMITED,
RESPONSES_RESPONSE,
SYNC_DEPLOYMENT_HOOK,
call,
)
from tests.test_litellm_rust.support.callback_recorder import HookRecorder, RecordingLogger, request_body
from tests.test_litellm_rust.support.gaps import gap
from tests.test_litellm_rust.support.isolation import rebound
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed
from tests.test_litellm_rust.support.results import response_id, with_id
from tests.test_litellm_rust.support.routing import assert_served_natively
pytestmark = pytest.mark.requires_rust_extension
SYNC_HANDLERS_ON_ASYNC_CALLS: Final = gap("native async calls also run the sync success and failure handlers")
_OBJECT: Final = TypeAdapter(dict[str, object])
_HEADERS: Final = TypeAdapter(dict[str, str])
def success_signal(asynchronous: bool) -> str:
return "async_log_success_event" if asynchronous else "log_success_event"
@pytest.fixture
def server(server: RecordingServer) -> Iterator[RecordingServer]:
yield from relaxed(server)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"hook,asynchronous,expected",
(
pytest.param("log_pre_api_call", True, 1, id="async-log_pre_api_call"),
pytest.param("log_post_api_call", True, 1, id="async-log_post_api_call"),
pytest.param("async_log_success_event", True, 1, id="async-async_log_success_event"),
pytest.param("async_logging_hook", True, 1, id="async-async_logging_hook"),
pytest.param("redact_standard_logging_payload_from_model_call_details", True, 1, id="async-redact"),
pytest.param("async_pre_call_deployment_hook", True, 1, id="async-async_pre_call_deployment_hook"),
pytest.param(
"async_post_call_success_deployment_hook", True, 1, id="async-async_post_call_success_deployment_hook"
),
pytest.param("log_success_event", True, 0, id="async-log_success_event"),
pytest.param("logging_hook", True, 0, id="async-logging_hook"),
pytest.param("log_pre_api_call", False, 1, id="sync-log_pre_api_call"),
pytest.param("log_post_api_call", False, 1, id="sync-log_post_api_call"),
pytest.param("log_success_event", False, 1, id="sync-log_success_event"),
pytest.param("logging_hook", False, 1, id="sync-logging_hook"),
pytest.param(
"async_pre_call_deployment_hook",
False,
1,
id="sync-async_pre_call_deployment_hook",
marks=SYNC_DEPLOYMENT_HOOK,
),
pytest.param(
"async_post_call_success_deployment_hook", False, 0, id="sync-async_post_call_success_deployment_hook"
),
),
)
async def test_success_hook_fires_the_expected_number_of_times(
server: RecordingServer, hook: str, asynchronous: bool, expected: int
) -> None:
recorder: Final = HookRecorder()
with rebound(litellm, "callbacks", [recorder]):
await call(server, asynchronous=asynchronous)
await recorder.wait_for_async(success_signal(asynchronous))
assert_served_natively(server)
assert recorder.names.count(hook) == expected
@pytest.mark.asyncio
@pytest.mark.parametrize(
"hook,asynchronous,expected",
(
pytest.param("async_log_failure_event", True, 1, id="async-async_log_failure_event"),
pytest.param(
"async_post_call_failure_deployment_hook", True, 1, id="async-async_post_call_failure_deployment_hook"
),
pytest.param("log_model_group_rate_limit_error", True, 1, id="async-log_model_group_rate_limit_error"),
pytest.param("log_failure_event", True, 0, id="async-log_failure_event"),
pytest.param("log_failure_event", False, 1, id="sync-log_failure_event"),
),
)
async def test_failure_hook_fires_the_expected_number_of_times(
server: RecordingServer, hook: str, asynchronous: bool, expected: int
) -> None:
server.default_response = ResponseSpec(body=RATE_LIMITED, status=429)
recorder: Final = HookRecorder()
with rebound(litellm, "callbacks", [recorder]):
with pytest.raises(litellm.RateLimitError):
await call(server, asynchronous=asynchronous, metadata={"model_group_size": 1})
await recorder.wait_for_async("async_log_failure_event" if asynchronous else "log_failure_event")
assert_served_natively(server)
assert recorder.names.count(hook) == expected
@pytest.mark.asyncio
async def test_failure_events_carry_the_error_raised_to_the_caller(server: RecordingServer) -> None:
server.default_response = ResponseSpec(body=BAD_REQUEST, status=400)
recorder: Final = RecordingLogger()
with rebound(litellm, "callbacks", [recorder]):
with pytest.raises(litellm.BadRequestError) as raised:
await call(server)
events: Final = await recorder.wait_for_async("async_log_failure_event")
assert_served_natively(server)
assert isinstance(events[0].kwargs, dict)
assert events[0].kwargs["exception"] is raised.value
@pytest.mark.asyncio
async def test_pre_api_call_sees_the_wire_body(server: RecordingServer) -> None:
recorder: Final = RecordingLogger()
await call(server, callbacks=[recorder])
assert_served_natively(server)
pre_call: Final = recorder.wait_for("log_pre_api_call")[0].kwargs
assert isinstance(pre_call, dict)
assert request_body(pre_call) == server.requests[0].body
@pytest.mark.asyncio
async def test_pre_api_call_edits_reach_the_provider(server: RecordingServer) -> None:
class Edit(CustomLogger):
def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None:
request_body(kwargs)["temperature"] = 0.25
await call(server, callbacks=[Edit()])
assert_served_natively(server)
assert isinstance(server.requests[0].body, dict)
assert server.requests[0].body["temperature"] == 0.25
@pytest.mark.asyncio
async def test_logging_hook_and_redaction_edits_reach_the_async_success_event(server: RecordingServer) -> None:
class Rewrite(HookRecorder):
async def async_logging_hook(
self, kwargs: dict[str, object], result: object, call_type: str
) -> tuple[dict[str, object], object]:
return {**kwargs, "logging_hook_ran": True}, result
def redact_standard_logging_payload_from_model_call_details(
self, model_call_details: dict[str, object]
) -> dict[str, object]:
return {**model_call_details, "redacted": True}
recorder: Final = Rewrite()
with rebound(litellm, "callbacks", [recorder]):
await call(server)
success: Final = await recorder.wait_for_async("async_log_success_event")
assert isinstance(success[0].kwargs, dict)
assert success[0].kwargs["logging_hook_ran"] is True
assert success[0].kwargs["redacted"] is True
@pytest.mark.asyncio
async def test_pre_call_deployment_hooks_compose_the_wire_request_in_registration_order(
server: RecordingServer,
) -> None:
class Rewrite(CustomLogger):
def __init__(self, marker: str) -> None:
super().__init__()
self.marker = marker
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, object], call_type: object
) -> dict[str, object]:
headers: Final = _HEADERS.validate_python(kwargs["extra_headers"])
return {
**kwargs,
"extra_headers": {**headers, "x-deployment-hook": headers["x-deployment-hook"] + self.marker},
}
with rebound(litellm, "callbacks", [Rewrite("first"), Rewrite("-second")]):
await call(server, extra_headers={"x-deployment-hook": "caller-", "x-kept": "unchanged"})
assert_served_natively(server)
assert server.requests[0].headers["x-deployment-hook"] == "caller-first-second"
assert server.requests[0].headers["x-kept"] == "unchanged"
@pytest.mark.asyncio
async def test_pre_call_deployment_rejection_logs_failure_without_provider_io(server: RecordingServer) -> None:
rejection: Final = litellm.BadRequestError(message="deployment rejected", model="test-model", llm_provider="openai")
recorder: Final = RecordingLogger()
class Reject(CustomLogger):
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, object], call_type: object
) -> dict[str, object]:
raise rejection
with rebound(litellm, "callbacks", [Reject(), recorder]):
with pytest.raises(litellm.BadRequestError, match="deployment rejected") as raised:
await call(server)
failures: Final = await recorder.wait_for_async("async_log_failure_event")
assert raised.value is rejection
assert server.requests == []
assert len(failures) == 1
assert _OBJECT.validate_python(failures[0].kwargs)["exception"] is raised.value
assert not {"log_pre_api_call", "log_success_event", "async_log_success_event"} & set(recorder.names)
@pytest.mark.asyncio
async def test_pre_call_deployment_hook_can_delete_keywords(server: RecordingServer) -> None:
class DropKeywords(CustomLogger):
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, object], call_type: object
) -> dict[str, object]:
return {name: value for name, value in kwargs.items() if name not in ("instructions", "extra_headers")}
with rebound(litellm, "callbacks", [DropKeywords()]):
await call(
server,
instructions="remove this instruction",
extra_headers={"x-deleted": "remove this header"},
)
assert_served_natively(server)
sent: Final = server.requests[0]
assert isinstance(sent.body, dict)
assert "instructions" not in sent.body
assert "x-deleted" not in sent.headers
@pytest.mark.asyncio
async def test_post_call_success_deployment_hook_replaces_the_response(server: RecordingServer) -> None:
recorder: Final = RecordingLogger()
class Replace(CustomLogger):
async def async_post_call_success_deployment_hook(
self, request_data: dict[str, object], response: object, call_type: object
) -> object:
return with_id(response, "replaced")
with rebound(litellm, "callbacks", [Replace(), recorder]):
response: Final = await call(server)
success: Final = await recorder.wait_for_async("async_log_success_event")
assert_served_natively(server)
assert response_id(response) == "replaced"
assert len(success) == 1
assert response_id(success[0].response) == response_id(response)
@pytest.mark.asyncio
async def test_concurrent_calls_keep_callback_context_metadata_and_responses_isolated(server: RecordingServer) -> None:
context: Final = ContextVar("responses-callback-context", default="outside")
first_entered: Final = asyncio.Event()
second_entered: Final = asyncio.Event()
server.enqueue(ResponseSpec(body={**RESPONSES_RESPONSE, "id": "resp_first"}))
server.enqueue(ResponseSpec(body={**RESPONSES_RESPONSE, "id": "resp_second"}))
class Observe(RecordingLogger):
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, object], call_type: object
) -> dict[str, object]:
arrived, other = (
(first_entered, second_entered)
if kwargs["litellm_call_id"] == "first"
else (second_entered, first_entered)
)
arrived.set()
await asyncio.wait_for(other.wait(), timeout=5)
return kwargs
def log_pre_api_call(self, model: object, messages: object, kwargs: dict[str, object]) -> None:
self._record("log_pre_api_call", {**kwargs, "caller_context": context.get()})
async def async_log_success_event(
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
self._record("async_log_success_event", {**kwargs, "caller_context": context.get()}, response_obj)
recorder: Final = Observe()
async def invoke(marker: str) -> object:
token: Final = context.set(marker)
try:
return await call(server, litellm_call_id=marker, metadata={"marker": marker})
finally:
context.reset(token)
with rebound(litellm, "callbacks", [recorder]):
responses: Final = await asyncio.gather(invoke("first"), invoke("second"))
success: Final = await recorder.wait_for_async("async_log_success_event", count=2)
assert_served_natively(server, count=2)
expected: Final = {"first": responses[0], "second": responses[1]}
pre_call: Final = recorder.wait_for("log_pre_api_call", count=2)
assert len(pre_call) == len(success) == 2
assert {_OBJECT.validate_python(event.kwargs)["litellm_call_id"] for event in success} == expected.keys()
for event in (*pre_call, *success):
details: Final = _OBJECT.validate_python(event.kwargs)
call_id: Final = str(details["litellm_call_id"])
params: Final = _OBJECT.validate_python(details["litellm_params"])
assert details["caller_context"] == call_id
assert _OBJECT.validate_python(params["metadata"])["marker"] == call_id
for event in success:
logged_call_id: Final = str(_OBJECT.validate_python(event.kwargs)["litellm_call_id"])
assert response_id(event.response) == response_id(expected[logged_call_id])
assert context.get() == "outside"
@pytest.mark.asyncio
async def test_a_failing_logger_does_not_fail_the_call(server: RecordingServer) -> None:
class Broken(CustomLogger):
def log_pre_api_call(self, model: str, messages: object, kwargs: dict[str, object]) -> None:
raise RuntimeError("logger exploded")
response: Final = await call(server, callbacks=[Broken()])
assert_served_natively(server)
assert response_id(response) == RESPONSES_RESPONSE["id"]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"event,asynchronous,dynamic,expected",
(
pytest.param("input", False, False, 1, id="global-input"),
pytest.param("success", True, False, 1, id="global-async-success"),
pytest.param("success", True, True, 1, id="request-async-success"),
pytest.param("success", False, False, 0, id="global-sync-success", marks=SYNC_HANDLERS_ON_ASYNC_CALLS),
pytest.param("success", False, True, 0, id="request-sync-success", marks=SYNC_HANDLERS_ON_ASYNC_CALLS),
pytest.param("failure", True, False, 1, id="global-async-failure"),
pytest.param("failure", False, False, 1, id="global-sync-failure"),
pytest.param("failure", False, True, 1, id="request-sync-failure"),
),
)
async def test_function_callback_fires_the_expected_number_of_times(
server: RecordingServer, event: str, asynchronous: bool, dynamic: bool, expected: int
) -> None:
recorder: Final = RecordingLogger()
def callback(kwargs: dict[str, object], *_: object) -> None:
recorder._record(event, kwargs)
async def async_callback(kwargs: dict[str, object], *_: object) -> None:
recorder._record(event, kwargs)
registered: Final = async_callback if asynchronous else callback
global_list: Final = f"_async_{event}_callback" if asynchronous else f"{event}_callback"
request_callbacks: Final = {f"{event}_callback": [registered]} if dynamic else {}
if event == "failure":
server.default_response = ResponseSpec(body=BAD_REQUEST, status=400)
with rebound(litellm, "callbacks", [recorder]), rebound(litellm, global_list, [] if dynamic else [registered]):
try:
await call(server, **request_callbacks)
except litellm.BadRequestError:
assert event == "failure"
await recorder.wait_for_async(f"async_log_{'failure' if event == 'failure' else 'success'}_event")
assert_served_natively(server)
assert recorder.names.count(event) == expected
@pytest.mark.asyncio
async def test_dataset_hook_returning_none_drops_the_record_from_the_success_sink(server: RecordingServer) -> None:
datasets: Final = RecordingLogger()
class Drop(CustomLogger):
async def async_dataset_hook(
self, logged_data: ArgillaItem, standard_logging_payload: StandardLoggingPayload | None
) -> None:
datasets._record("async_dataset_hook", response=logged_data)
server.enqueue(ResponseSpec(body={"items": [{"id": "callback-dataset"}]}))
with rebound(litellm, "argilla_transformation_object", {"prompt": "messages", "answer": "response"}):
sink: Final = ArgillaLogger(
argilla_api_key="test-key",
argilla_base_url=server.base_url,
argilla_dataset_name="callback-dataset",
flush_interval=3600,
)
with rebound(litellm, "callbacks", [sink, Drop()]):
await call(server)
await datasets.wait_for_async("async_dataset_hook")
await sink.flush_queue()
assert [request.method for request in server.requests] == ["GET", "POST"]
Reference in a new issue View git blame Copy permalink