diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 069dcffd21a..93b45be4922 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2466,7 +2466,7 @@ class Logging(LiteLLMLoggingBaseClass): await self.async_failure_handler(exception, traceback_exception) finally: if self._should_run_sync_failure_callbacks_for_async_calls(): - executor.submit(self.failure_handler, exception, traceback_exception) + executor.submit(self.failure_handler, exception, traceback_exception, for_async_request=True) def should_run_logging( self, @@ -4090,15 +4090,21 @@ class Logging(LiteLLMLoggingBaseClass): traceback_exception: str, start_time: datetime.datetime | None = None, end_time: datetime.datetime | None = None, + *, + for_async_request: bool = False, ) -> None: """Restores trace_id/session_id contextvars once this attempt's own failure - logging (including any nested calls its callbacks trigger) is fully done.""" + logging (including any nested calls its callbacks trigger) is fully done. + + ``for_async_request`` marks the sync run of a failure whose async run also logs it, so a + CustomLogger that gets ``async_log_failure_event`` from that async run skips its sync hook.""" try: return self._failure_handler_body( exception=exception, traceback_exception=traceback_exception, start_time=start_time, end_time=end_time, + for_async_request=for_async_request, ) finally: self._restore_correlation_context() @@ -4109,6 +4115,7 @@ class Logging(LiteLLMLoggingBaseClass): traceback_exception: str, start_time: datetime.datetime | None = None, end_time: datetime.datetime | None = None, + for_async_request: bool = False, ) -> None: verbose_logger.debug("Logging Details LiteLLM-Failure Call: %s", litellm.failure_callback) if not self.should_run_logging(event_type="sync_failure"): # prevent double logging @@ -4201,6 +4208,7 @@ class Logging(LiteLLMLoggingBaseClass): isinstance(callback, CustomLogger) and is_sync_request and self.call_type != CallTypes.pass_through.value + and not (for_async_request and self._async_handler_logs_to(callback)) ): # custom logger class callback.log_failure_event( start_time=start_time, @@ -4563,6 +4571,12 @@ class Logging(LiteLLMLoggingBaseClass): cb_name: Final = self._get_callback_name(cb) return any(prefix in cb_name for prefix in INTERNAL_PREFIXES) + def _async_handler_logs_to(self, logger: CustomLogger) -> bool: + """True when the async failure run of this call hands ``logger`` its own ``async_log_failure_event``.""" + if type(logger).async_log_failure_event is CustomLogger.async_log_failure_event: # pyright: ignore[reportUnknownMemberType] # CustomLogger's hook signatures are untyped + return False + return logger in (self.dynamic_async_failure_callbacks or []) or logger in litellm._async_failure_callback # pyright: ignore[reportUnknownMemberType, reportPrivateUsage] # callback lists are untyped and have no public accessor + def _remove_internal_custom_logger_callbacks(self, callbacks: list) -> list: """ Removes internal custom logger callbacks from the list. diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index aca222c31ef..c5983077d4c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3865,6 +3865,7 @@ class ProxyLogging: original_exception, traceback_str, ), + kwargs={"for_async_request": True}, daemon=True, ).start() diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index a82e734b47d..ae913882968 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -11,6 +11,7 @@ import contextvars import datetime import traceback import uuid +import weakref from collections.abc import Awaitable, Coroutine, Mapping from dataclasses import dataclass from typing import ( @@ -61,9 +62,21 @@ def setup( } supplied: Final = arguments.get("litellm_logging_obj") if isinstance(supplied, Logging): - return _claim_budget_reservation(CallSetup(supplied, arguments), asynchronous) + return _begin_call(CallSetup(supplied, arguments), asynchronous) logger, prepared = function_setup(call_type, Rules(), start_time, *args, is_async_call=asynchronous, **arguments) - return _claim_budget_reservation(CallSetup(logger, prepared), asynchronous) + return _begin_call(CallSetup(logger, prepared), asynchronous) + + +# Loggers whose current native call is async: their sync failure step is followed by the async one +_ASYNC_CALLS: Final[weakref.WeakSet[Logging]] = weakref.WeakSet() + + +def _begin_call(call_setup: CallSetup, asynchronous: bool) -> CallSetup: + if asynchronous: + _ASYNC_CALLS.add(call_setup.logger) + else: + _ASYNC_CALLS.discard(call_setup.logger) + return _claim_budget_reservation(call_setup, asynchronous) def _claim_budget_reservation(call_setup: CallSetup, asynchronous: bool) -> CallSetup: @@ -135,6 +148,8 @@ class LoggingSurface(Protocol): traceback_exception: str, start_time: datetime.datetime | None = None, end_time: datetime.datetime | None = None, + *, + for_async_request: bool = False, ) -> None: ... def async_failure_handler( @@ -261,7 +276,7 @@ def failure_handler( if not is_internal_call(): unbind_budget_reservation_from_callbacks(logger.litellm_params) return logger.async_failure_handler(error, trace, start, end) - logger.failure_handler(error, trace, start, end) + logger.failure_handler(error, trace, start, end, for_async_request=logger in _ASYNC_CALLS) return None diff --git a/litellm/utils.py b/litellm/utils.py index 987139b7530..a9079f470c7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2371,7 +2371,7 @@ def client(original_function): unbind_budget_reservation_from_callbacks(logging_obj.litellm_params) try: logging_obj.failure_handler( - e, traceback_exception, start_time, end_time + e, traceback_exception, start_time, end_time, for_async_request=True ) # DO NOT MAKE THREADED - router retry fallback relies on this! except Exception as e: raise e diff --git a/tests/integration/_support/hook_count_probe.py b/tests/integration/_support/hook_count_probe.py new file mode 100644 index 00000000000..0d4056e8fb9 --- /dev/null +++ b/tests/integration/_support/hook_count_probe.py @@ -0,0 +1,91 @@ +import json +import os +from collections.abc import Mapping +from datetime import datetime +from typing import Final + +import httpx +from fastapi import HTTPException + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger + + +def _record(hook: str, kwargs: Mapping[str, object]) -> bytes: + return json.dumps( + {"hook": hook, "litellm_call_id": kwargs.get("litellm_call_id"), "call_type": kwargs.get("call_type")} + ).encode() + + +def _send(hook: str, kwargs: Mapping[str, object]) -> None: + with httpx.Client(trust_env=False) as client: + client.post(os.environ["HOOK_COUNT_PROBE_URL"], content=_record(hook, kwargs)).raise_for_status() + + +async def _send_async(hook: str, kwargs: Mapping[str, object]) -> None: + async with httpx.AsyncClient(trust_env=False) as client: + response: Final = await client.post(os.environ["HOOK_COUNT_PROBE_URL"], content=_record(hook, kwargs)) + response.raise_for_status() + + +class HookCountProbe(CustomLogger): + """Reports every sync and async success and failure hook it receives, one record per call.""" + + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + _send("sync_success", kwargs) + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + await _send_async("async_success", kwargs) + + def log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + _send("sync_failure", kwargs) + + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + await _send_async("async_failure", kwargs) + + +class SyncOnlyHookCountProbe(CustomLogger): + """Overrides only the sync success and failure hooks.""" + + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + _send("sync_success", kwargs) + + def log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + _send("sync_failure", kwargs) + + +class BlockingPostCallGuardrail(CustomGuardrail): + """Rejects every response with a 400 after the provider answered.""" + + async def async_post_call_success_hook( + self, data: Mapping[str, object], user_api_key_dict: object, response: object + ) -> None: + raise HTTPException(status_code=400, detail="blocked by the hook-count guardrail") + + +probe: Final = HookCountProbe() +sync_only_probe: Final = SyncOnlyHookCountProbe() + + +def sync_success_function( + kwargs: Mapping[str, object], completion_response: object, start_time: datetime, end_time: datetime +) -> None: + _send("function_success", kwargs) + + +def sync_failure_function( + kwargs: Mapping[str, object], completion_response: object, start_time: datetime, end_time: datetime +) -> None: + _send("function_failure", kwargs) diff --git a/tests/integration/_support/hook_counts.py b/tests/integration/_support/hook_counts.py new file mode 100644 index 00000000000..6b7f482bd85 --- /dev/null +++ b/tests/integration/_support/hook_counts.py @@ -0,0 +1,329 @@ +"""Scripted provider, proxy config and request helpers shared by the CustomLogger hook count tests.""" + +import base64 +import json +import threading +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import Final + +import httpx +import yaml +from integration._support.client import Gateway, held +from integration._support.wire import Reply, Request, Wire + +WAV_BYTES: Final = ( + b"RIFF\x30\x00\x00\x00WAVEfmt \x10\x00\x00\x00\x01\x00\x01\x00" + b"\x40\x1f\x00\x00\x40\x1f\x00\x00\x01\x00\x08\x00data\x08\x00\x00\x00" + b"\x00\x01\x02\x03\x04\x05\x06\x07" +) +PNG_BYTES: Final = ( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00" + b"\x1f\x15\xc4\x89\x00\x00\x00\rIDAT\x08\xd7c\xf8\xcf\xc0\xf0\x1f\x00\x05\x00\x01\xff" + b"\x89\x99=\x1d\x00\x00\x00\x00IEND\xaeB`\x82" +) +_IMAGE_REPLY: Final = json.dumps({"created": 1700000000, "data": [{"b64_json": base64.b64encode(PNG_BYTES).decode()}]}) +_CHAT_REPLY: Final = json.dumps( + { + "id": "chatcmpl-hook-count", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, + } +) +_EMBEDDING_REPLY: Final = json.dumps( + { + "object": "list", + "model": "text-embedding-3-small", + "data": [{"object": "embedding", "index": 0, "embedding": [0.25, 0.5]}], + "usage": {"prompt_tokens": 3, "total_tokens": 3}, + } +) +_MESSAGE_REPLY: Final = json.dumps( + { + "id": "msg_hook_count", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 1}, + } +) +_RESPONSE_REPLY: Final = json.dumps( + { + "id": "resp_hook_count", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": "msg_hook_count", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4}, + } +) +_OCR_REPLY: Final = json.dumps( + { + "pages": [ + {"index": 0, "markdown": "page", "images": [], "dimensions": {"dpi": 200, "height": 1100, "width": 850}} + ], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1, "doc_size_bytes": 64}, + "document_annotation": None, + } +) +_SEARCH_REPLY: Final = json.dumps( + {"results": [{"title": "hook count", "url": "https://result.invalid/a", "snippet": "hit"}]} +) +_CHAT_FRAMES: Final = ( + 'data: {"id":"chatcmpl-hook-count","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",' + '"choices":[{"index":0,"delta":{"role":"assistant","content":"ok"},"finish_reason":null}]}', + 'data: {"id":"chatcmpl-hook-count","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",' + '"choices":[{"index":0,"delta":{},"finish_reason":"stop"}],' + '"usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4}}', + "data: [DONE]", +) +_MESSAGE_FRAMES: Final = ( + "event: message_start\n" + 'data: {"type":"message_start","message":{"id":"msg_hook_count","type":"message","role":"assistant",' + '"model":"claude-haiku-4-5","content":[],"stop_reason":null,"stop_sequence":null,' + '"usage":{"input_tokens":3,"output_tokens":0}}}', + 'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}', + "event: content_block_delta\n" + 'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"ok"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}', + "event: message_delta\n" + 'data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":1}}', + 'event: message_stop\ndata: {"type":"message_stop"}', +) +_RESPONSE_FRAMES: Final = ( + "event: response.created\n" + 'data: {"type":"response.created","sequence_number":0,"response":{"id":"resp_hook_count","object":"response",' + '"created_at":1,"status":"in_progress","model":"gpt-4o-mini","output":[],"usage":null}}', + "event: response.output_text.delta\n" + 'data: {"type":"response.output_text.delta","sequence_number":1,"item_id":"msg_hook_count","output_index":0,' + '"content_index":0,"delta":"ok"}', + "event: response.completed\n" + 'data: {"type":"response.completed","sequence_number":2,"response":{"id":"resp_hook_count","object":"response",' + '"created_at":1,"status":"completed","model":"gpt-4o-mini","output":[{"type":"message","id":"msg_hook_count",' + '"status":"completed","role":"assistant","content":[{"type":"output_text","text":"ok","annotations":[]}]}],' + '"usage":{"input_tokens":3,"output_tokens":1,"total_tokens":4}}}', +) +_STREAM_FRAMES: Final = { + "/v1/chat/completions": _CHAT_FRAMES, + "/v1/messages": _MESSAGE_FRAMES, + "/v1/responses": _RESPONSE_FRAMES, +} +SURFACE_CALL_TYPES: Final = { + "chat": "acompletion", + "embedding": "aembedding", + "messages": "anthropic_messages", + "responses": "aresponses", + "speech": "aspeech", + "transcription": "atranscription", + "image_generation": "aimage_generation", + "image_edit": "aimage_edit", + "ocr": "aocr", + "search": "asearch", +} +IN_THE_ASYNC_FLAG_LIST: Final = frozenset({"chat", "embedding", "messages", "responses"}) +STREAMING_SURFACES: Final = ("chat", "messages", "responses") + + +def _streams(request: Request) -> bool: + if request.headers.get("content-type", "").startswith("application/json"): + return json.loads(request.body or b"{}").get("stream") is True + return False + + +def provider(request: Request) -> Reply: + """Answers every surface the hook count tests drive; any target under /fail/ is a provider 400.""" + if request.method == "GET": + return Reply(body=b'{"object":"list","data":[]}') + if request.target.startswith("/fail/"): + return Reply( + status=400, + body=b'{"error":{"message":"scripted provider rejection","type":"invalid_request_error","code":null}}', + ) + if request.target in _STREAM_FRAMES and _streams(request): + return Reply( + chunks=tuple(f"{frame}\n\n".encode() for frame in _STREAM_FRAMES[request.target]), + content_type="text/event-stream", + ) + replies: Final[dict[str, Callable[[], Reply]]] = { + "/v1/chat/completions": lambda: Reply(body=_CHAT_REPLY.encode()), + "/v1/embeddings": lambda: Reply(body=_EMBEDDING_REPLY.encode()), + "/v1/messages": lambda: Reply(body=_MESSAGE_REPLY.encode()), + "/v1/responses": lambda: Reply(body=_RESPONSE_REPLY.encode()), + "/v1/audio/speech": lambda: Reply(body=WAV_BYTES, content_type="audio/wav"), + "/v1/audio/transcriptions": lambda: Reply(body=b'{"text":"hook count transcript"}'), + "/v1/images/generations": lambda: Reply(body=_IMAGE_REPLY.encode()), + "/v1/images/edits": lambda: Reply(body=_IMAGE_REPLY.encode()), + "/v1/ocr": lambda: Reply(body=_OCR_REPLY.encode()), + "/perplexity/search": lambda: Reply(body=_SEARCH_REPLY.encode()), + } + return replies[request.target]() + + +@dataclass +class HookRecords: + """Every hook record the probe reported, kept across drains so cells can read their own call id.""" + + wire: Wire + function_callbacks: bool + seen: list[dict[str, str]] = field(default_factory=list) + lock: threading.Lock = field(default_factory=threading.Lock) + + def drained(self) -> tuple[dict[str, str], ...]: + with self.lock: + self.seen.extend(json.loads(request.body) for request in self.wire.drain()) + return tuple(self.seen) + + def for_call(self, call_id: str) -> tuple[dict[str, str], ...]: + return tuple(record for record in self.drained() if record["litellm_call_id"] == call_id) + + +def _model(name: str, model: str, api_base: str, **extra: object) -> dict[str, object]: + return { + "model_name": name, + "litellm_params": {"model": model, "api_base": api_base, "api_key": "synthetic-provider-key"}, + **extra, + } + + +def write_hook_count_config( + path: Path, + provider_url: str, + litellm_settings: Mapping[str, object], + guardrails: tuple[Mapping[str, object], ...] = (), +) -> Path: + """One success and one failure deployment per surface, all answered by `provider`.""" + models: Final = [] + tools: Final = [] + for outcome, prefix in (("success", ""), ("failure", "/fail")): + base: str = provider_url + prefix + models.extend( + ( + _model(f"hooks-chat-{outcome}", "openai/gpt-4o-mini", base + "/v1"), + _model(f"hooks-embedding-{outcome}", "openai/text-embedding-3-small", base + "/v1"), + _model(f"hooks-messages-{outcome}", "anthropic/claude-haiku-4-5", base), + _model(f"hooks-responses-{outcome}", "openai/gpt-4o-mini", base + "/v1"), + _model(f"hooks-speech-{outcome}", "openai/gpt-4o-mini-tts", base + "/v1"), + _model(f"hooks-transcription-{outcome}", "openai/whisper-1", base + "/v1"), + _model(f"hooks-image_generation-{outcome}", "openai/gpt-image-1", base + "/v1"), + _model(f"hooks-image_edit-{outcome}", "openai/gpt-image-1", base + "/v1"), + _model(f"hooks-ocr-{outcome}", "mistral/mistral-ocr-latest", base + "/v1", model_info={"mode": "ocr"}), + ) + ) + tools.append( + { + "search_tool_name": f"hooks-search-{outcome}", + "litellm_params": { + "search_provider": "perplexity", + "api_key": "synthetic-provider-key", + "api_base": base + "/perplexity", + }, + } + ) + path.write_text( + yaml.safe_dump( + { + "model_list": models, + "search_tools": tools, + "litellm_settings": dict(litellm_settings), + **({"guardrails": [dict(guardrail) for guardrail in guardrails]} if guardrails else {}), + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "disable_model_info_refresh": True, + }, + "router_settings": {"num_retries": 0}, + } + ) + ) + return path + + +def send( + proxy: Gateway, surface: str, outcome: str, call_id: str, *, key: str | None = None, stream: bool = False +) -> httpx.Response: + model: Final = f"hooks-{surface}-{outcome}" + nonce: Final = f"hook-count-{call_id}" + headers: Final = {"Authorization": f"Bearer {key or proxy.key}", "x-litellm-call-id": call_id} + if surface == "transcription": + return proxy.client.post( + "/v1/audio/transcriptions", + data={"model": model, "prompt": nonce}, + files={"file": ("a.wav", WAV_BYTES + call_id.encode(), "audio/wav")}, + headers=headers, + ) + if surface == "image_edit": + return proxy.client.post( + "/v1/images/edits", + data={"model": model, "prompt": nonce}, + files={"image": ("a.png", PNG_BYTES, "image/png")}, + headers=headers, + ) + streaming: Final[dict[str, object]] = {"stream": True} if stream else {} + routes: Final[dict[str, tuple[str, dict[str, object]]]] = { + "chat": ( + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": nonce}], **streaming}, + ), + "embedding": ("/v1/embeddings", {"model": model, "input": nonce}), + "messages": ( + "/v1/messages", + {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": nonce}], **streaming}, + ), + "responses": ("/v1/responses", {"model": model, "input": nonce, **streaming}), + "speech": ("/v1/audio/speech", {"model": model, "input": nonce, "voice": "alloy", "response_format": "wav"}), + "image_generation": ("/v1/images/generations", {"model": model, "prompt": nonce, "n": 1}), + "ocr": ( + "/v1/ocr", + { + "model": model, + "document": { + "type": "document_url", + "document_url": "data:application/pdf;base64," + base64.b64encode(nonce.encode()).decode(), + }, + }, + ), + "search": ("/v1/search", {"search_tool_name": f"hooks-search-{outcome}", "query": nonce}), + } + path, body = routes[surface] + return proxy.client.post(path, json=body, headers=headers) + + +def counts_function_hooks(records: HookRecords, surface: str) -> bool: + return records.function_callbacks and surface not in IN_THE_ASYNC_FLAG_LIST + + +def probe_hooks(records: HookRecords, surface: str, found: tuple[dict[str, str], ...]) -> list[str]: + counted: Final = counts_function_hooks(records, surface) + return sorted(record["hook"] for record in found if counted or not record["hook"].startswith("function_")) + + +def one_async_hook(records: HookRecords, surface: str, outcome: str) -> list[str]: + function_hook: Final = (f"function_{outcome}",) if counts_function_hooks(records, surface) else () + return sorted((f"async_{outcome}", *function_hook)) + + +def assert_hooks(records: HookRecords, surface: str, call_id: str, expected: list[str]) -> None: + settled: Final = held( + lambda: records.for_call(call_id), + lambda found: probe_hooks(records, surface, found) == expected, + holding=3, + seconds=20, + ) + assert {record["call_type"] for record in settled} == {SURFACE_CALL_TYPES[surface]}, settled diff --git a/tests/integration/observability/test_custom_logger_hook_counts.py b/tests/integration/observability/test_custom_logger_hook_counts.py new file mode 100644 index 00000000000..420ac2a5fcb --- /dev/null +++ b/tests/integration/observability/test_custom_logger_hook_counts.py @@ -0,0 +1,96 @@ +import uuid +from collections.abc import Iterator +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.hook_counts import ( + STREAMING_SURFACES, + SURFACE_CALL_TYPES, + HookRecords, + assert_hooks, + one_async_hook, + provider, + send, + write_hook_count_config, +) +from integration._support.process import owned_proxy +from integration._support.wire import Reply, wire_server + +pytestmark = pytest.mark.timeout(240) + + +@pytest.fixture(scope="module", params=["cache", "cache_and_sync_callback"]) +def hook_proxy( + request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory +) -> Iterator[tuple[Gateway, HookRecords]]: + tmp_path: Final = tmp_path_factory.mktemp("custom-logger-hook-counts") + legacy_callbacks: Final = ( + { + "success_callback": ["integration._support.hook_count_probe.sync_success_function"], + "failure_callback": ["integration._support.hook_count_probe.sync_failure_function"], + } + if request.param == "cache_and_sync_callback" + else {} + ) + with ( + wire_server(provider) as upstream, + wire_server(lambda request: Reply()) as sink, + gateway_from_environment() as rig, + ): + config: Final = write_hook_count_config( + tmp_path / "hook-counts.yaml", + upstream.url, + { + "callbacks": ["integration._support.hook_count_probe.probe"], + **legacy_callbacks, + "cache": True, + "cache_params": {"type": "redis", "host": "os.environ/REDIS_HOST", "port": "os.environ/REDIS_PORT"}, + }, + ) + with owned_proxy( + rig, tmp_path, {"HOOK_COUNT_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"}, config=config, workers=2 + ) as proxy: + yield proxy, HookRecords(sink, function_callbacks=request.param == "cache_and_sync_callback") + + +@pytest.mark.parametrize("outcome", ["success", "failure"]) +@pytest.mark.parametrize("surface", sorted(SURFACE_CALL_TYPES)) +def test_custom_logger_gets_one_async_hook_per_request_with_the_response_cache_on( + hook_proxy: tuple[Gateway, HookRecords], surface: str, outcome: str +) -> None: + proxy, records = hook_proxy + call_id: Final = str(uuid.uuid4()) + response: Final = send(proxy, surface, outcome, call_id) + assert response.status_code == (200 if outcome == "success" else 400), response.text + assert_hooks(records, surface, call_id, one_async_hook(records, surface, outcome)) + + +@pytest.mark.parametrize("outcome", ["success", "failure"]) +@pytest.mark.parametrize("surface", STREAMING_SURFACES) +def test_custom_logger_gets_one_async_hook_per_streamed_request( + hook_proxy: tuple[Gateway, HookRecords], surface: str, outcome: str +) -> None: + proxy, records = hook_proxy + call_id: Final = str(uuid.uuid4()) + response: Final = send(proxy, surface, outcome, call_id, stream=True) + assert response.status_code == (200 if outcome == "success" else 400), response.text + if outcome == "success": + assert "ok" in response.text, response.text + assert_hooks(records, surface, call_id, one_async_hook(records, surface, outcome)) + + +@pytest.mark.parametrize("surface", sorted(SURFACE_CALL_TYPES)) +def test_custom_logger_gets_one_async_hook_per_request_rejected_by_key_auth( + hook_proxy: tuple[Gateway, HookRecords], surface: str +) -> None: + proxy, records = hook_proxy + known_calls: Final = {record["litellm_call_id"] for record in records.drained()} + response: Final = send(proxy, surface, "success", str(uuid.uuid4()), key="sk-not-a-virtual-key") + assert response.status_code == 401, response.text + + def new_calls(found: tuple[dict[str, str], ...]) -> set[str]: + return {record["litellm_call_id"] for record in found if record["litellm_call_id"] not in known_calls} + + (rejected_call,) = eventually(lambda: new_calls(records.drained()), bool, seconds=20) + assert_hooks(records, surface, rejected_call, one_async_hook(records, surface, "failure")) diff --git a/tests/integration/observability/test_custom_logger_hook_counts_chaos.py b/tests/integration/observability/test_custom_logger_hook_counts_chaos.py new file mode 100644 index 00000000000..7b19ef3b1ff --- /dev/null +++ b/tests/integration/observability/test_custom_logger_hook_counts_chaos.py @@ -0,0 +1,181 @@ +"""One hook per request under a concurrent burst while the logger's destination stalls and fails, and while a +worker is killed.""" + +import threading +import time +import uuid +from collections.abc import Callable +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, eventually, gateway_from_environment, held +from integration._support.hook_counts import ( + STREAMING_SURFACES, + SURFACE_CALL_TYPES, + HookRecords, + assert_hooks, + one_async_hook, + probe_hooks, + provider, + send, + write_hook_count_config, +) +from integration._support.process import graceful_stop_seconds, owned_proxy, owned_proxy_process +from integration._support.wire import Reply, Request, wire_server + +pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 300) + +RecordProperty = Callable[[str, object], None] + +_SETTINGS: Final = {"callbacks": ["integration._support.hook_count_probe.probe"]} +_HOLD_SECONDS: Final = 2.0 + + +@dataclass(frozen=True) +class _Call: + surface: str + outcome: str + stream: bool + call_id: str + + +@dataclass(frozen=True) +class _Outcome: + call: _Call + status: int | None + error: str | None + + +def _burst(rounds: int) -> tuple[_Call, ...]: + shapes: Final = ( + *((surface, outcome, False) for surface in sorted(SURFACE_CALL_TYPES) for outcome in ("success", "failure")), + *((surface, "success", True) for surface in STREAMING_SURFACES), + ) + return tuple( + _Call(surface, outcome, stream, str(uuid.uuid4())) for _ in range(rounds) for surface, outcome, stream in shapes + ) + + +def _attempt(proxy: Gateway, call: _Call) -> _Outcome: + try: + response: Final = send(proxy, call.surface, call.outcome, call.call_id, stream=call.stream) + except httpx.HTTPError as error: + return _Outcome(call, None, repr(error)) + return _Outcome(call, response.status_code, None) + + +def _expected_status(call: _Call) -> int: + return 200 if call.outcome == "success" else 400 + + +def _done(futures: tuple[Future[_Outcome], ...]) -> int: + return sum(future.done() for future in futures) + + +class _Destination: + """The probe's sink: answers at once, holds each report until released, or refuses it with a 503.""" + + def __init__(self) -> None: + self.mode = "ok" + self.release = threading.Event() + + def __call__(self, request: Request) -> Reply: + if self.mode == "hold": + self.release.wait(timeout=_HOLD_SECONDS + 1) + if self.mode == "fail": + return Reply(status=503, body=b'{"error":"destination unavailable"}') + return Reply() + + +def test_logger_destination_stalling_then_failing_mid_burst_gets_one_hook_per_request( + tmp_path: Path, record_property: RecordProperty +) -> None: + destination: Final = _Destination() + with ( + wire_server(provider) as upstream, + wire_server(destination) as sink, + gateway_from_environment() as rig, + ): + config: Final = write_hook_count_config(tmp_path / "hook-count-chaos.yaml", upstream.url, _SETTINGS) + with owned_proxy( + rig, tmp_path, {"HOOK_COUNT_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"}, config=config, workers=2 + ) as proxy: + records: Final = HookRecords(sink, function_callbacks=False) + calls: Final = _burst(2) + destination.mode = "hold" + with ThreadPoolExecutor(max_workers=len(calls)) as pool: + futures: Final = tuple(pool.submit(_attempt, proxy, call) for call in calls) + eventually(lambda: _done(futures), lambda done: done >= 1, seconds=60) + held_liveliness: Final = proxy.client.get("/health/liveliness").status_code + time.sleep(_HOLD_SECONDS) + destination.mode = "fail" + destination.release.set() + time.sleep(_HOLD_SECONDS) + destination.mode = "ok" + outcomes: Final = tuple(future.result(timeout=180) for future in futures) + record_property("burst_size", len(calls)) + assert held_liveliness == 200, held_liveliness + assert [outcome for outcome in outcomes if outcome.status != _expected_status(outcome.call)] == [] + for call in calls: + assert_hooks(records, call.surface, call.call_id, one_async_hook(records, call.surface, call.outcome)) + + +def test_worker_killed_mid_burst_leaves_at_most_one_hook_per_request( + tmp_path: Path, record_property: RecordProperty +) -> None: + with ( + wire_server(provider) as upstream, + wire_server(lambda request: Reply()) as sink, + gateway_from_environment() as rig, + ): + config: Final = write_hook_count_config(tmp_path / "hook-count-kill.yaml", upstream.url, _SETTINGS) + with owned_proxy_process( + rig, tmp_path, {"HOOK_COUNT_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"}, config=config, workers=2 + ) as owned: + proxy: Final = owned.gateway + records: Final = HookRecords(sink, function_callbacks=False) + victim: Final = min(eventually(owned.worker_pids, lambda found: len(found) == 2, seconds=60)) + calls: Final = _burst(2) + with ThreadPoolExecutor(max_workers=len(calls)) as pool: + futures: Final = tuple(pool.submit(_attempt, proxy, call) for call in calls) + eventually(lambda: _done(futures), lambda done: done >= 1, seconds=60) + with owned.replacing_worker(victim) as kill: + kill() + done_at_kill: Final = _done(futures) + survived: Final = eventually( + lambda: _attempt(proxy, _Call("speech", "success", False, str(uuid.uuid4()))), + lambda outcome: outcome.status == 200, + seconds=30, + ) + outcomes: Final = tuple(future.result(timeout=180) for future in futures) + respawned: Final = _Call("speech", "success", False, str(uuid.uuid4())) + after_respawn: Final = _attempt(proxy, respawned) + completed: Final = tuple(outcome for outcome in outcomes if outcome.status is not None) + interrupted: Final = tuple(outcome for outcome in outcomes if outcome.status is None) + assert after_respawn.status == 200, after_respawn + assert [outcome for outcome in completed if outcome.status != _expected_status(outcome.call)] == [] + for call in (survived.call, respawned): + assert_hooks(records, call.surface, call.call_id, one_async_hook(records, call.surface, call.outcome)) + burst_hooks: Final = held( + lambda: _hooks_of(records, calls), + lambda found: all( + hooks in ([], [f"async_{call.outcome}"]) for call, hooks in zip(calls, found, strict=True) + ), + holding=3, + seconds=20, + ) + record_property("burst_size", len(calls)) + record_property("killed_worker_pid", victim) + record_property("done_at_kill", done_at_kill) + record_property("interrupted_calls", str(tuple(outcome.error for outcome in interrupted))) + record_property("burst_calls_without_a_hook", sum(not hooks for hooks in burst_hooks)) + + +def _hooks_of(records: HookRecords, calls: tuple[_Call, ...]) -> tuple[list[str], ...]: + """A burst call the killed worker served may lose its logging with the worker, whether or not its response + arrived first, so a burst call gets its one async hook or none, never two.""" + return tuple(probe_hooks(records, call.surface, records.for_call(call.call_id)) for call in calls) diff --git a/tests/integration/observability/test_custom_logger_hook_counts_modes.py b/tests/integration/observability/test_custom_logger_hook_counts_modes.py new file mode 100644 index 00000000000..e143559040e --- /dev/null +++ b/tests/integration/observability/test_custom_logger_hook_counts_modes.py @@ -0,0 +1,95 @@ +"""A logger the async handler never reaches keeps its sync hook on async calls outside the async flag list. + +A logger with both hooks registered only under `success_callback` and `failure_callback` is never reached by the async +failure handler, so it keeps its sync failure hook. The async success handler does reach it, and gives it its async +success hook instead of the sync one. + +The sync success handler only runs for an async call when a legacy sync callback is configured, so every +registration here carries the legacy function callbacks too.""" + +import uuid +from collections.abc import Iterator +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.hook_counts import ( + IN_THE_ASYNC_FLAG_LIST, + SURFACE_CALL_TYPES, + HookRecords, + assert_hooks, + provider, + send, + write_hook_count_config, +) +from integration._support.process import owned_proxy +from integration._support.wire import Reply, wire_server + +pytestmark = pytest.mark.timeout(240) + +_PROBES: Final = "integration._support.hook_count_probe" +_REGISTRATIONS: Final[dict[str, dict[str, object]]] = { + "sync_only_logger": { + "callbacks": [f"{_PROBES}.sync_only_probe"], + "success_callback": [f"{_PROBES}.sync_success_function"], + "failure_callback": [f"{_PROBES}.sync_failure_function"], + }, + "sync_event_lists": { + "success_callback": [f"{_PROBES}.probe", f"{_PROBES}.sync_success_function"], + "failure_callback": [f"{_PROBES}.probe", f"{_PROBES}.sync_failure_function"], + }, +} +_OUTSIDE_THE_ASYNC_FLAG_LIST: Final = sorted(SURFACE_CALL_TYPES.keys() - IN_THE_ASYNC_FLAG_LIST) + + +@pytest.fixture(scope="module", params=sorted(_REGISTRATIONS)) +def sync_hook_proxy( + request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory +) -> Iterator[tuple[Gateway, HookRecords, str]]: + tmp_path: Final = tmp_path_factory.mktemp("custom-logger-hook-count-modes") + with ( + wire_server(provider) as upstream, + wire_server(lambda request: Reply()) as sink, + gateway_from_environment() as rig, + ): + config: Final = write_hook_count_config( + tmp_path / "hook-count-modes.yaml", upstream.url, _REGISTRATIONS[request.param] + ) + with owned_proxy( + rig, tmp_path, {"HOOK_COUNT_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"}, config=config, workers=2 + ) as proxy: + yield proxy, HookRecords(sink, function_callbacks=True), request.param + + +@pytest.mark.parametrize("outcome", ["success", "failure"]) +@pytest.mark.parametrize("surface", _OUTSIDE_THE_ASYNC_FLAG_LIST) +def test_logger_the_async_handler_skips_gets_one_sync_hook_per_request( + sync_hook_proxy: tuple[Gateway, HookRecords, str], surface: str, outcome: str +) -> None: + proxy, records, registration = sync_hook_proxy + reached_by_the_async_success: Final = registration == "sync_event_lists" and outcome == "success" + call_id: Final = str(uuid.uuid4()) + response: Final = send(proxy, surface, outcome, call_id) + assert response.status_code == (200 if outcome == "success" else 400), response.text + assert_hooks( + records, + surface, + call_id, + sorted([f"function_{outcome}", f"async_{outcome}" if reached_by_the_async_success else f"sync_{outcome}"]), + ) + + +@pytest.mark.parametrize("surface", _OUTSIDE_THE_ASYNC_FLAG_LIST) +def test_logger_the_async_handler_skips_gets_one_sync_hook_per_request_rejected_by_key_auth( + sync_hook_proxy: tuple[Gateway, HookRecords, str], surface: str +) -> None: + proxy, records, _ = sync_hook_proxy + known_calls: Final = {record["litellm_call_id"] for record in records.drained()} + response: Final = send(proxy, surface, "success", str(uuid.uuid4()), key="sk-not-a-virtual-key") + assert response.status_code == 401, response.text + + def new_calls(found: tuple[dict[str, str], ...]) -> set[str]: + return {record["litellm_call_id"] for record in found if record["litellm_call_id"] not in known_calls} + + (rejected_call,) = eventually(lambda: new_calls(records.drained()), bool, seconds=20) + assert_hooks(records, surface, rejected_call, ["function_failure", "sync_failure"]) diff --git a/tests/integration/observability/test_custom_logger_hook_counts_post_call_block.py b/tests/integration/observability/test_custom_logger_hook_counts_post_call_block.py new file mode 100644 index 00000000000..1b708db64d8 --- /dev/null +++ b/tests/integration/observability/test_custom_logger_hook_counts_post_call_block.py @@ -0,0 +1,75 @@ +"""A CustomLogger with both hooks gets one failure hook when a post-call guardrail blocks the response. + +The proxy drops the async success logging it parked for the guardrails and logs the block as a failure through both +the async and the sync failure handlers. The logger gets only its async failure hook from them. Speech and +transcription responses never reach post-call guardrails, and image generation logs its success before they run, so +neither parks. + +A legacy sync callback is configured, because that is what makes the sync handler run for an async call.""" + +import uuid +from collections.abc import Iterator +from typing import Final + +import pytest +from integration._support.client import Gateway, gateway_from_environment, held +from integration._support.hook_counts import ( + SURFACE_CALL_TYPES, + HookRecords, + probe_hooks, + provider, + send, + write_hook_count_config, +) +from integration._support.process import owned_proxy +from integration._support.wire import Reply, wire_server + +pytestmark = pytest.mark.timeout(240) + +_PROBES: Final = "integration._support.hook_count_probe" +_PARKED_FOR_POST_CALL_GUARDRAILS: Final = sorted( + SURFACE_CALL_TYPES.keys() - {"speech", "transcription", "image_generation"} +) + + +@pytest.fixture(scope="module") +def post_call_block(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, HookRecords]]: + tmp_path: Final = tmp_path_factory.mktemp("hook-count-post-call-block") + guardrail: Final = { + "guardrail_name": "hook-count-post-call-block", + "litellm_params": { + "guardrail": f"{_PROBES}.BlockingPostCallGuardrail", + "mode": "post_call", + "default_on": True, + }, + } + with ( + wire_server(provider) as upstream, + wire_server(lambda request: Reply()) as sink, + gateway_from_environment() as rig, + ): + config: Final = write_hook_count_config( + tmp_path / "hook-count-post-call-block.yaml", + upstream.url, + {"callbacks": [f"{_PROBES}.probe"], "success_callback": [f"{_PROBES}.sync_success_function"]}, + (guardrail,), + ) + overrides: Final = {"HOOK_COUNT_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"} + with owned_proxy(rig, tmp_path, overrides, config=config, workers=2) as proxy: + yield proxy, HookRecords(sink, function_callbacks=False) + + +@pytest.mark.parametrize("surface", _PARKED_FOR_POST_CALL_GUARDRAILS) +def test_response_a_post_call_guardrail_blocks_gives_the_logger_one_failure_hook( + post_call_block: tuple[Gateway, HookRecords], surface: str +) -> None: + proxy, records = post_call_block + call_id: Final = str(uuid.uuid4()) + response: Final = send(proxy, surface, "success", call_id) + assert response.status_code == 400, response.text + held( + lambda: probe_hooks(records, surface, records.for_call(call_id)), + lambda hooks: hooks == ["async_failure"], + holding=3, + seconds=20, + ) diff --git a/tests/integration/observability/test_custom_logger_hook_counts_vendors.py b/tests/integration/observability/test_custom_logger_hook_counts_vendors.py new file mode 100644 index 00000000000..3c9be9e2a3f --- /dev/null +++ b/tests/integration/observability/test_custom_logger_hook_counts_vendors.py @@ -0,0 +1,128 @@ +"""Built-in loggers that override both hooks deliver one record per async request outside the async flag list: +Datadog one log entry and LangSmith one run. OpenTelemetry keeps its own emit-once guard, so it never exports two +spans for a request either way; its speech, OCR and search spans carry no `litellm.call_id`, so only "never two" +is checked for it. + +A legacy sync callback is configured, because that is what makes the sync handler run for an async call.""" + +import gzip +import json +import uuid +import zlib +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Final + +import pytest +from integration._support.client import Gateway, gateway_from_environment, held +from integration._support.hook_counts import IN_THE_ASYNC_FLAG_LIST, SURFACE_CALL_TYPES, provider, send +from integration._support.hook_counts import write_hook_count_config +from integration._support.otlp_sink import decode_spans +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server + +pytestmark = pytest.mark.timeout(240) + +_PROBES: Final = "integration._support.hook_count_probe" +_FLUSH_SECONDS: Final = 1 +_OUTSIDE_THE_ASYNC_FLAG_LIST: Final = sorted(SURFACE_CALL_TYPES.keys() - IN_THE_ASYNC_FLAG_LIST) + + +@dataclass(frozen=True) +class _Vendors: + proxy: Gateway + sink: Wire + + +@pytest.fixture(scope="module") +def vendors(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Vendors]: + tmp_path: Final = tmp_path_factory.mktemp("custom-logger-hook-count-vendors") + with ( + wire_server(provider) as upstream, + wire_server(lambda request: Reply(status=202)) as sink, + gateway_from_environment() as rig, + ): + config: Final = write_hook_count_config( + tmp_path / "hook-count-vendors.yaml", + upstream.url, + { + "callbacks": ["datadog", "langsmith", "otel"], + "success_callback": [f"{_PROBES}.sync_success_function"], + "failure_callback": [f"{_PROBES}.sync_failure_function"], + }, + ) + overrides: Final = { + "HOOK_COUNT_PROBE_URL": sink.url + "/probe", + "STORE_MODEL_IN_DB": "False", + "DEFAULT_FLUSH_INTERVAL_SECONDS": str(_FLUSH_SECONDS), + "DD_API_KEY": "synthetic-datadog-key", + "DD_SITE": "us5.datadoghq.com", + "DD_BASE_URL": sink.url + "/datadog", + "LANGSMITH_API_KEY": "synthetic-langsmith-key", + "LANGSMITH_PROJECT": "hook-counts", + "LANGSMITH_BASE_URL": sink.url + "/langsmith", + "OTEL_EXPORTER": "otlp_http", + "OTEL_ENDPOINT": sink.url + "/otel/v1/traces", + } + with owned_proxy(rig, tmp_path, overrides, config=config, workers=2) as proxy: + yield _Vendors(proxy, sink) + + +def _unzipped(body: bytes) -> bytes: + if body[:2] == b"\x1f\x8b": + return gzip.decompress(body) + try: + return zlib.decompress(body, zlib.MAX_WBITS | 16) + except zlib.error: + return body + + +def _entries(body: bytes) -> tuple[object, ...]: + decoded: Final = json.loads(_unzipped(body)) + return tuple(decoded) if isinstance(decoded, list) else (decoded,) + + +def _records(request: Request) -> tuple[tuple[str, str], ...]: + """Every record one delivery carries, as (destination, serialized record). A LangSmith batch carries its new + runs under `post`, and an OpenTelemetry export carries every span of the request's trace.""" + if request.target.startswith("/datadog"): + return tuple(("datadog", json.dumps(entry)) for entry in _entries(request.body)) + if request.target.startswith("/langsmith") and request.method == "POST": + runs: Final = tuple( + run + for entry in _entries(request.body) + for run in (entry.get("post", []) if isinstance(entry, dict) and "post" in entry else [entry]) + ) + return tuple(("langsmith", json.dumps(run)) for run in runs) + if request.target.startswith("/otel"): + spans: Final = decode_spans(request.body, request.headers.get("content-type", "")) + return tuple(("otel", json.dumps(span)) for span in spans if "litellm.call_id" in span["attributes"]) + return () + + +@dataclass +class _Deliveries: + sink: Wire + seen: tuple[tuple[str, str], ...] = () + + def per_destination(self, call_id: str) -> dict[str, int]: + self.seen = (*self.seen, *(record for request in self.sink.drain() for record in _records(request))) + destinations: Final = ("datadog", "langsmith", "otel") + return { + destination: sum(1 for name, record in self.seen if name == destination and call_id in record) + for destination in destinations + } + + +@pytest.mark.parametrize("surface", _OUTSIDE_THE_ASYNC_FLAG_LIST) +def test_built_in_loggers_deliver_one_record_per_async_request(vendors: _Vendors, surface: str) -> None: + deliveries: Final = _Deliveries(vendors.sink) + call_id: Final = str(uuid.uuid4()) + response: Final = send(vendors.proxy, surface, "success", call_id) + assert response.status_code == 200, response.text + held( + lambda: deliveries.per_destination(call_id), + lambda counts: counts["datadog"] == 1 and counts["langsmith"] == 1 and counts["otel"] <= 1, + holding=4 * _FLUSH_SECONDS, + seconds=60, + ) diff --git a/tests/integration/routing/test_router_strategies_count_async_requests_once.py b/tests/integration/routing/test_router_strategies_count_async_requests_once.py new file mode 100644 index 00000000000..f7403ad0437 --- /dev/null +++ b/tests/integration/routing/test_router_strategies_count_async_requests_once.py @@ -0,0 +1,214 @@ +"""Router strategies that keep both a sync and an async hook count a finished async request once: least-busy +keeps counting a request still in flight while others finish around it, and usage-based routing v2 adds a +request's tokens to its deployment once. + +A legacy sync callback is configured, because that is what makes the sync handler run for an async call.""" + +import base64 +import json +import os +import threading +import uuid +from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment, held +from integration._support.hook_counts import PNG_BYTES, WAV_BYTES, HookRecords +from integration._support.process import OwnedState, owned_proxy +from integration._support.wire import Reply, Request, wire_server +from redis import Redis + +pytestmark = pytest.mark.timeout(240) + +_MODEL: Final = "router-strategy-model" +_PROBES: Final = "integration._support.hook_count_probe" +_REJECTION: Final = b'{"error":{"message":"scripted provider rejection","type":"invalid_request_error"}}' +_IMAGE_TOKENS: Final = 12 +_IMAGE_REPLY: Final = json.dumps( + { + "created": 1700000000, + "data": [{"b64_json": base64.b64encode(PNG_BYTES).decode()}], + "usage": { + "input_tokens": 5, + "output_tokens": 7, + "total_tokens": _IMAGE_TOKENS, + "input_tokens_details": {"text_tokens": 5, "image_tokens": 0}, + }, + } +).encode() + + +@dataclass(frozen=True) +class _Rig: + proxy: Gateway + records: HookRecords + release: threading.Event + deployment_id: str + + +def _provider(release: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == "/v1/images/generations": + return Reply(body=_IMAGE_REPLY) + spoken: Final = str(json.loads(request.body or b"{}").get("input", "")) + if spoken.startswith("fail"): + return Reply(status=400, body=_REJECTION) + if spoken.startswith("hold"): + return Reply(chunks=(WAV_BYTES,), content_type="audio/wav", gate_before_first=release) + return Reply(body=WAV_BYTES, content_type="audio/wav") + + return respond + + +def _write_config(path: Path, provider_url: str, strategy: str, model: str, deployment_id: str) -> Path: + path.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": _MODEL, + "litellm_params": { + "model": model, + "api_base": provider_url + "/v1", + "api_key": "synthetic-provider-key", + }, + "model_info": {"id": deployment_id}, + } + ], + "litellm_settings": { + "success_callback": [f"{_PROBES}.sync_success_function"], + "failure_callback": [f"{_PROBES}.sync_failure_function"], + }, + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "disable_model_info_refresh": True, + }, + "router_settings": { + "routing_strategy": strategy, + "routing_strategy_args": {"ttl": 600}, + "num_retries": 0, + "disable_cooldowns": True, + "redis_host": "os.environ/REDIS_HOST", + "redis_port": "os.environ/REDIS_PORT", + }, + } + ) + ) + return path + + +@contextmanager +def _router_proxy(tmp_path: Path, strategy: str, model: str) -> Iterator[_Rig]: + release: Final = threading.Event() + deployment_id: Final = f"{strategy}-{uuid.uuid4().hex}" + with ( + wire_server(_provider(release)) as upstream, + wire_server(lambda request: Reply()) as sink, + gateway_from_environment() as rig, + ): + config: Final = _write_config(tmp_path / f"{strategy}.yaml", upstream.url, strategy, model, deployment_id) + try: + with owned_proxy( + rig, + tmp_path, + {"HOOK_COUNT_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"}, + config=config, + workers=2, + state=OwnedState.SHARED_WITH_WORKER, + ) as proxy: + yield _Rig(proxy, HookRecords(sink, function_callbacks=True), release, deployment_id) + finally: + release.set() + + +@pytest.fixture(scope="module") +def least_busy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + with _router_proxy(tmp_path_factory.mktemp("least-busy"), "least-busy", "openai/gpt-4o-mini-tts") as rig: + yield rig + + +@pytest.fixture(scope="module") +def usage_based(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + with _router_proxy(tmp_path_factory.mktemp("usage-based"), "usage-based-routing-v2", "openai/gpt-image-1") as rig: + yield rig + + +def _redis() -> Redis: + return Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) + + +def _in_flight(rig: _Rig) -> int: + with _redis() as cache: + stored: Final = cache.get(f"{_MODEL}_request_count:{rig.deployment_id}") + return 0 if stored is None else int(stored) + + +def _tokens_counted(rig: _Rig) -> int: + with _redis() as cache: + keys: Final = tuple(cache.scan_iter(match=f"{rig.deployment_id}:*:tpm:*")) + return sum(int(cache.get(key) or 0) for key in keys) + + +def _post(rig: _Rig, path: str, body: dict[str, object], call_id: str) -> httpx.Response: + return rig.proxy.client.post( + path, + json={"model": _MODEL, **body}, + headers={"Authorization": f"Bearer {rig.proxy.key}", "x-litellm-call-id": call_id}, + timeout=120, + ) + + +def _speak(rig: _Rig, text: str, call_id: str) -> httpx.Response: + return _post(rig, "/v1/audio/speech", {"input": text, "voice": "alloy", "response_format": "wav"}, call_id) + + +def _logged(rig: _Rig, call_id: str) -> None: + """The request's success or failure logging, the legacy sync callback included, has run.""" + eventually(lambda: rig.records.for_call(call_id), bool, seconds=20) + + +def _finish_a_speech_request(rig: _Rig, outcome: str) -> None: + call_id: Final = str(uuid.uuid4()) + finished: Final = _speak(rig, f"{'ok' if outcome == 'success' else 'fail'} {call_id}", call_id) + assert finished.status_code == (200 if outcome == "success" else 400), finished.text + _logged(rig, call_id) + held(lambda: _in_flight(rig), lambda count: count == 1, holding=3, seconds=20) + + +def _generate_an_image(rig: _Rig) -> None: + call_id: Final = str(uuid.uuid4()) + response: Final = _post(rig, "/v1/images/generations", {"prompt": f"tokens {call_id}", "n": 1}, call_id) + assert response.status_code == 200, response.text + _logged(rig, call_id) + + +@pytest.mark.parametrize("outcome", ["success", "failure"]) +def test_least_busy_keeps_counting_a_request_in_flight_while_async_speech_requests_finish( + least_busy: _Rig, outcome: str +) -> None: + held(lambda: _in_flight(least_busy), lambda count: count == 0, holding=1, seconds=30) + least_busy.release.clear() + with ThreadPoolExecutor(max_workers=1) as pool: + in_flight: Final = pool.submit(_speak, least_busy, f"hold {uuid.uuid4().hex}", str(uuid.uuid4())) + eventually(lambda: _in_flight(least_busy), lambda count: count == 1, seconds=30) + try: + for _ in range(3): + _finish_a_speech_request(least_busy, outcome) + finally: + least_busy.release.set() + assert in_flight.result(timeout=60).status_code == 200 + held(lambda: _in_flight(least_busy), lambda count: count == 0, holding=3, seconds=30) + + +def test_usage_based_routing_counts_an_async_image_request_s_tokens_once(usage_based: _Rig) -> None: + for _ in range(3): + _generate_an_image(usage_based) + held(lambda: _tokens_counted(usage_based), lambda tokens: tokens == 3 * _IMAGE_TOKENS, holding=3, seconds=30) diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py index e8d49453065..cd5478dfc3a 100644 --- a/tests/test_litellm_rust/messages/test_callbacks.py +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -67,7 +67,7 @@ async def test_success_hook_fires_the_expected_number_of_times( pytest.param("async_log_failure_event", 1), pytest.param("async_post_call_failure_deployment_hook", 1), pytest.param("log_model_group_rate_limit_error", 1), - pytest.param("log_failure_event", 0, marks=SYNC_HANDLERS_ON_ASYNC_CALLS), + pytest.param("log_failure_event", 0), ), ) async def test_failure_hook_fires_the_expected_number_of_times( diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py index 5f961288800..18925599acf 100644 --- a/tests/test_litellm_rust/ocr/test_callbacks.py +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -237,7 +237,7 @@ async def test_native_aocr_failure_callbacks_receive_call_type_error_and_no_resp with pytest.raises(litellm.InternalServerError): await call_aocr(server, callbacks=[Observe()]) - assert [observation[0] for observation in observations] == ["sync", "async"] + assert [observation[0] for observation in observations] == ["async"] assert all(observation[1] == "aocr" for observation in observations) assert all(isinstance(observation[2], litellm.InternalServerError) for observation in observations) assert all(observation[3] is None for observation in observations) @@ -278,7 +278,7 @@ async def test_native_aocr_failure_callbacks_receive_state_added_by_pre_call_cal with pytest.raises(litellm.InternalServerError): await call_aocr(server, callbacks=[TrackInFlightRequest()]) - assert [event for event, _ in observed] == ["sync", "async"] + assert [event for event, _ in observed] == ["async"] assert all(observed_token is token for _, observed_token in observed) @@ -445,7 +445,7 @@ async def test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_f await call_aocr(server) await drain_logging() - assert observed == [("sync", token, raised.value), ("async", token, raised.value)] + assert observed == [("async", token, raised.value)] @pytest.mark.asyncio diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index c35f6836144..611316fb52c 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -144,7 +144,7 @@ async def test_mapped_failure_identity_and_deployment_snapshot(server: Recording with pytest.raises(litellm.InternalServerError) as caught: await call_aocr(server, callbacks=[recorder]) failures: Final = tuple(event for event in recorder.events if "failure" in event.name) - assert [event.name for event in failures] == ["log_failure_event", "async_log_failure_event"] + assert [event.name for event in failures] == ["async_log_failure_event"] assert all(event.kwargs["exception"] is caught.value for event in failures) assert caught.value.status_code == 500 assert snapshots[0] is not caught.value @@ -324,24 +324,26 @@ async def test_failure_callbacks_continue_within_both_families(server: Recording server.enqueue(ResponseSpec(body={"message": "failed"}, status=500)) observed: Final = [] - class Broken(CustomLogger): + class BrokenSync(CustomLogger): def log_failure_event(self, kwargs, response_obj, start_time, end_time): observed.append(("broken-sync", kwargs["exception"])) raise RuntimeError("sync observer") + class FollowingSync(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("following-sync", kwargs["exception"])) + + class BrokenAsync(CustomLogger): async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): observed.append(("broken-async", kwargs["exception"])) raise RuntimeError("async observer") - class Following(CustomLogger): - def log_failure_event(self, kwargs, response_obj, start_time, end_time): - observed.append(("following-sync", kwargs["exception"])) - + class FollowingAsync(CustomLogger): async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): observed.append(("following-async", kwargs["exception"])) with pytest.raises(litellm.InternalServerError) as caught: - await call_aocr(server, callbacks=[Broken(), Following()]) + await call_aocr(server, callbacks=[BrokenSync(), FollowingSync(), BrokenAsync(), FollowingAsync()]) assert [name for name, _ in observed] == ["broken-sync", "following-sync", "broken-async", "following-async"] assert all(error is caught.value for _, error in observed) diff --git a/tests/test_litellm_rust/responses/test_callbacks.py b/tests/test_litellm_rust/responses/test_callbacks.py index a4abdd43529..0e921e9d5db 100644 --- a/tests/test_litellm_rust/responses/test_callbacks.py +++ b/tests/test_litellm_rust/responses/test_callbacks.py @@ -95,7 +95,7 @@ async def test_success_hook_fires_the_expected_number_of_times( "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", marks=SYNC_HANDLERS_ON_ASYNC_CALLS), + pytest.param("log_failure_event", True, 0, id="async-log_failure_event"), pytest.param("log_failure_event", False, 1, id="sync-log_failure_event"), ), ) diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 981c43a9792..b3414ebe9d7 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -3088,7 +3088,7 @@ async def test_dispatch_failure_handlers_submits_sync_handler_for_failure_only_c prefer_async_handlers=True, ) - mock_submit.assert_called_once_with(mock_sync, exception, traceback_exception) + mock_submit.assert_called_once_with(mock_sync, exception, traceback_exception, for_async_request=True) @pytest.mark.asyncio @@ -8188,6 +8188,155 @@ def test_cache_callback_alone_does_not_start_sync_callbacks_for_async_requests( assert logging_obj._should_run_sync_callbacks_for_async_calls() is expected +class _FailureHookRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.hooks: list[str] = [] + + def log_failure_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.hooks.append("sync_failure") + + +class _BothFailureHooksRecorder(_FailureHookRecorder): + async def async_log_failure_event( + self, kwargs: object, response_obj: object, start_time: object, end_time: object + ) -> None: + self.hooks.append("async_failure") + + +def _failure_hook_count_logging_obj( + call_type: str, recorder: _FailureHookRecorder, *, async_failure_listed: bool +) -> LitellmLogging: + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini-tts", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type=call_type, + start_time=datetime.datetime.now(), + litellm_call_id=f"failure-hook-count-{call_type}", + function_id=f"failure-hook-count-{call_type}", + dynamic_failure_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder] if async_failure_listed else [], + ) + logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={}) + return logging_obj + + +@pytest.fixture +def no_global_callbacks(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ( + "callbacks", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", + ): + monkeypatch.setattr(litellm, name, []) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("recorder_type", "async_listed", "expected"), + [ + (_BothFailureHooksRecorder, True, ["async_failure"]), + (_FailureHookRecorder, True, ["sync_failure"]), + (_BothFailureHooksRecorder, False, ["sync_failure"]), + ], +) +@pytest.mark.parametrize("call_type", ["aspeech", "atranscription", "aimage_edit", "asearch"]) +@pytest.mark.usefixtures("no_global_callbacks") +async def test_async_call_outside_the_async_flag_list_gets_one_custom_logger_failure_hook( + call_type: str, recorder_type: type[_FailureHookRecorder], async_listed: bool, expected: list[str] +) -> None: + recorder: Final = recorder_type() + logging_obj: Final = _failure_hook_count_logging_obj(call_type, recorder, async_failure_listed=async_listed) + + logging_obj.failure_handler(ValueError("upstream rejected"), "traceback", for_async_request=True) + await logging_obj.async_failure_handler(ValueError("upstream rejected"), "traceback") + + assert recorder.hooks == expected + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("no_global_callbacks") +async def test_logger_in_the_global_failure_callback_lists_gets_one_failure_hook( + monkeypatch: pytest.MonkeyPatch, +) -> None: + recorder: Final = _BothFailureHooksRecorder() + monkeypatch.setattr(litellm, "failure_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_failure_callback", [recorder]) + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini-tts", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="aspeech", + start_time=datetime.datetime.now(), + litellm_call_id="global-failure", + function_id="global-failure", + ) + logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={}) + + logging_obj.failure_handler(ValueError("upstream rejected"), "traceback", for_async_request=True) + await logging_obj.async_failure_handler(ValueError("upstream rejected"), "traceback") + + assert recorder.hooks == ["async_failure"] + + +@pytest.mark.usefixtures("no_global_callbacks") +def test_failure_handler_run_without_an_async_failure_run_keeps_the_sync_hook() -> None: + recorder: Final = _BothFailureHooksRecorder() + + _failure_hook_count_logging_obj("aspeech", recorder, async_failure_listed=True).failure_handler( + ValueError("upstream rejected"), "traceback" + ) + + assert recorder.hooks == ["sync_failure"] + + +@pytest.mark.usefixtures("no_global_callbacks") +def test_async_call_outside_the_async_flag_list_still_runs_sync_failure_function_callbacks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm_logging, "customLogger", CustomLogger()) + calls: Final[list[str]] = [] + + def failure_function(kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + calls.append("function_failure") + + recorder: Final = _BothFailureHooksRecorder() + logging_obj: Final = _failure_hook_count_logging_obj("aspeech", recorder, async_failure_listed=True) + logging_obj.dynamic_failure_callbacks = [recorder, failure_function] + + logging_obj.failure_handler(ValueError("upstream rejected"), "traceback", for_async_request=True) + + assert calls == ["function_failure"] + assert recorder.hooks == [] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("no_global_callbacks") +async def test_client_wrapped_failure_gives_a_logger_in_litellm_callbacks_one_hook( + monkeypatch: pytest.MonkeyPatch, +) -> None: + recorder: Final = _BothFailureHooksRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + + @litellm.utils.client + async def aspeech(model: str, input: str, voice: str, **kwargs: object) -> None: + raise litellm.BadRequestError(message="upstream rejected", model=model, llm_provider="openai") + + @litellm.utils.client + def speech(model: str, input: str, voice: str, **kwargs: object) -> None: + raise litellm.BadRequestError(message="upstream rejected", model=model, llm_provider="openai") + + with pytest.raises(litellm.BadRequestError): + await aspeech(model="gpt-4o-mini-tts", input="hi", voice="alloy") + with pytest.raises(litellm.BadRequestError): + speech(model="gpt-4o-mini-tts", input="hi", voice="alloy") + + assert recorder.hooks == ["async_failure", "sync_failure"] + + @pytest.mark.asyncio @pytest.mark.parametrize("turn_off_message_logging", [False, True]) @pytest.mark.parametrize("make_response", [ModelResponse, _responses_api_response_with_text]) diff --git a/tests/unit/proxy/test_proxy_reject_logging.py b/tests/unit/proxy/test_proxy_reject_logging.py index 6e8545a58b3..b300d48808f 100644 --- a/tests/unit/proxy/test_proxy_reject_logging.py +++ b/tests/unit/proxy/test_proxy_reject_logging.py @@ -209,4 +209,4 @@ async def test_chat_completion_request_with_redaction(route, body, router, monke assert _test_logger.reaches_async_failure_event is True - assert _test_logger.reaches_sync_failure_event is True + assert _test_logger.reaches_sync_failure_event is False diff --git a/tests/unit/proxy/test_proxy_utils.py b/tests/unit/proxy/test_proxy_utils.py index 2fc1b72c932..02f31bf07c6 100644 --- a/tests/unit/proxy/test_proxy_utils.py +++ b/tests/unit/proxy/test_proxy_utils.py @@ -2652,6 +2652,58 @@ async def test_post_call_failure_hook_auth_error_llm_api_route(): mock_handle_logging.assert_called_once() +class _InlineThread: + def __init__(self, target, args=(), kwargs=None, daemon=None) -> None: + self._run = lambda: target(*args, **(kwargs or {})) + + def start(self) -> None: + self._run() + + +@pytest.mark.asyncio +async def test_proxy_only_failure_gives_a_logger_with_both_failure_hooks_only_the_async_one( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import types + + import litellm.proxy.utils as proxy_utils + from litellm.integrations.custom_logger import CustomLogger + from litellm.litellm_core_utils.litellm_logging import Logging + + hooks: Final[list[str]] = [] + + class Recorder(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + hooks.append("sync_failure") + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + hooks.append("async_failure") + + for name in ("failure_callback", "_async_failure_callback"): + monkeypatch.setattr(litellm, name, []) + monkeypatch.setattr(proxy_utils, "threading", types.SimpleNamespace(Thread=_InlineThread)) + recorder: Final = Recorder() + logging_obj: Final = Logging( + model="gpt-4o-mini-tts", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="aspeech", + start_time=datetime.now(), + litellm_call_id="proxy-only-failure", + function_id="proxy-only-failure", + dynamic_failure_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={}) + + await proxy_utils.ProxyLogging._dispatch_proxy_only_failure_handlers( + litellm_logging_obj=logging_obj, + original_exception=HTTPException(status_code=401, detail="Authentication Error: invalid user key"), + ) + + assert hooks == ["async_failure"] + + @pytest.mark.asyncio @pytest.mark.parametrize( "request_data, route, expected_call_type", diff --git a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py index eb1bfd82715..a40f8a77105 100644 --- a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py +++ b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py @@ -9,6 +9,7 @@ import pytest from pydantic import TypeAdapter from litellm._internal_context import is_internal_call +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.rust_bridge.host import callbacks_legacy_python as legacy from litellm.rust_bridge.host.callbacks_legacy_python import failure_handler, setup @@ -167,6 +168,75 @@ def test_failure_handler_of_an_internal_call_leaves_the_outer_budget_reservation pending.close() +class _Hooks(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.hooks: list[str] = [] + + def log_success_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.hooks.append("sync_success") + + async def async_log_success_event( + self, kwargs: object, response_obj: object, start_time: object, end_time: object + ) -> None: + self.hooks.append("async_success") + + def log_failure_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.hooks.append("sync_failure") + + async def async_log_failure_event( + self, kwargs: object, response_obj: object, start_time: object, end_time: object + ) -> None: + self.hooks.append("async_failure") + + +def _logger_with_hooks(hooks: _Hooks) -> Logging: + logger: Final = Logging( + model="mistral/mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="hooks", + function_id="hooks", + dynamic_success_callbacks=[hooks], + dynamic_async_success_callbacks=[hooks], + dynamic_failure_callbacks=[hooks], + dynamic_async_failure_callbacks=[hooks], + ) + logger.update_environment_variables(litellm_params={"metadata": {}}, optional_params={}) + return logger + + +@pytest.mark.asyncio +async def test_failure_of_an_async_call_gives_a_logger_with_both_failure_hooks_only_the_async_one() -> None: + hooks: Final = _Hooks() + now: Final = datetime.datetime.now() + kwargs: Final = {**_OCR_KWARGS, "litellm_logging_obj": _logger_with_hooks(hooks)} + result: Final = setup("aocr", (), kwargs, now, asynchronous=True) + + assert failure_handler(result.logger, RuntimeError("upstream refused"), now, now, asynchronous=False) is None + pending: Final = failure_handler(result.logger, RuntimeError("upstream refused"), now, now, asynchronous=True) + assert pending is not None + await pending + + assert hooks.hooks == ["async_failure"] + + +@pytest.mark.parametrize("earlier_async_call", [False, True]) +def test_failure_of_a_sync_call_keeps_the_sync_failure_hook(earlier_async_call: bool) -> None: + hooks: Final = _Hooks() + now: Final = datetime.datetime.now() + kwargs: Final = {**_OCR_KWARGS, "litellm_logging_obj": _logger_with_hooks(hooks)} + if earlier_async_call: + setup("aocr", (), kwargs, now, asynchronous=True) + result: Final = setup("ocr", (), kwargs, now, asynchronous=False) + + assert failure_handler(result.logger, RuntimeError("upstream refused"), now, now, asynchronous=False) is None + + assert hooks.hooks == ["sync_failure"] + + CONTRACT_PATH: Final = Path(__file__).parents[4] / "litellm-rust/crates/callbacks-legacy-python/python_contract.json"