From 20b6e855bde9cb1df1cd689c679b71d0d0c497b6 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Fri, 13 Nov 2026 14:08:05 -0800 Subject: [PATCH] feat(custom-logger-v2): add types, dispatcher and route callback call sites (step 0) (#51214) --- litellm/integrations/custom_logger.py | 8 + .../integrations/custom_logger_v2/__init__.py | 74 + litellm/integrations/custom_logger_v2/call.py | 55 + .../custom_logger_v2/dispatcher.py | 400 ++++ .../custom_logger_v2/enforcement.py | 95 + .../integrations/custom_logger_v2/enums.py | 97 + .../custom_logger_v2/hook_point_fields.py | 98 + .../integrations/custom_logger_v2/hooks.py | 211 ++ litellm/integrations/custom_logger_v2/host.py | 62 + .../integrations/custom_logger_v2/legacy.py | 434 ++++ .../integrations/custom_logger_v2/logger.py | 109 + .../integrations/custom_logger_v2/manifest.py | 73 + .../integrations/custom_logger_v2/offpath.py | 27 + .../integrations/custom_logger_v2/outcomes.py | 106 + .../integrations/custom_logger_v2/patches.py | 70 + .../integrations/custom_logger_v2/planner.py | 77 + .../integrations/custom_logger_v2/registry.py | 51 + .../custom_logger_v2/selection.py | 36 + litellm/integrations/otel/langfuse_logger.py | 15 +- litellm/litellm_core_utils/litellm_logging.py | 1972 ++++++++++------- .../llms/anthropic/messages/native/handler.py | 53 +- litellm/llms/custom_httpx/llm_http_handler.py | 13 +- litellm/proxy/common_request_processing.py | 122 +- .../policy_endpoints/endpoints.py | 3 +- .../proxy/policy_engine/pipeline_executor.py | 7 +- litellm/proxy/utils.py | 1261 ++++++----- litellm/router_rust/steps/callbacks.py | 85 +- .../router_utils/fallback_event_handlers.py | 143 +- .../host/callbacks_legacy_python.py | 21 +- litellm/utils.py | 231 +- .../integrations/custom_logger_v2/__init__.py | 0 .../custom_logger_v2/test_dispatcher.py | 547 +++++ .../custom_logger_v2/test_enums.py | 73 + .../test_hook_point_fields.py | 76 + .../custom_logger_v2/test_hooks.py | 129 ++ .../custom_logger_v2/test_host.py | 145 ++ .../custom_logger_v2/test_legacy.py | 332 +++ .../custom_logger_v2/test_logger.py | 242 ++ .../custom_logger_v2/test_manifest.py | 133 ++ .../custom_logger_v2/test_outcomes.py | 52 + .../custom_logger_v2/test_patches.py | 39 + .../custom_logger_v2/test_planner.py | 142 ++ .../custom_logger_v2/test_selection.py | 55 + .../test_custom_guardrail_post_call_once.py | 1 - tests/unit/integrations/test_custom_logger.py | 47 +- .../test_litellm_logging.py | 498 ++++- .../native/test_anthropic_messages_handler.py | 137 +- .../custom_httpx/test_llm_http_handler.py | 68 +- .../test_proxy_logging_hook_detection.py | 57 + tests/unit/proxy/test_v2_hook_routing.py | 1212 ++++++++++ .../test_fallback_event_v2_dispatch.py | 141 ++ .../core/test_callback_filters_v2_dispatch.py | 232 ++ tests/unit/test_utils.py | 391 +++- 53 files changed, 9155 insertions(+), 1603 deletions(-) create mode 100644 litellm/integrations/custom_logger_v2/__init__.py create mode 100644 litellm/integrations/custom_logger_v2/call.py create mode 100644 litellm/integrations/custom_logger_v2/dispatcher.py create mode 100644 litellm/integrations/custom_logger_v2/enforcement.py create mode 100644 litellm/integrations/custom_logger_v2/enums.py create mode 100644 litellm/integrations/custom_logger_v2/hook_point_fields.py create mode 100644 litellm/integrations/custom_logger_v2/hooks.py create mode 100644 litellm/integrations/custom_logger_v2/host.py create mode 100644 litellm/integrations/custom_logger_v2/legacy.py create mode 100644 litellm/integrations/custom_logger_v2/logger.py create mode 100644 litellm/integrations/custom_logger_v2/manifest.py create mode 100644 litellm/integrations/custom_logger_v2/offpath.py create mode 100644 litellm/integrations/custom_logger_v2/outcomes.py create mode 100644 litellm/integrations/custom_logger_v2/patches.py create mode 100644 litellm/integrations/custom_logger_v2/planner.py create mode 100644 litellm/integrations/custom_logger_v2/registry.py create mode 100644 litellm/integrations/custom_logger_v2/selection.py create mode 100644 tests/unit/integrations/custom_logger_v2/__init__.py create mode 100644 tests/unit/integrations/custom_logger_v2/test_dispatcher.py create mode 100644 tests/unit/integrations/custom_logger_v2/test_enums.py create mode 100644 tests/unit/integrations/custom_logger_v2/test_hook_point_fields.py create mode 100644 tests/unit/integrations/custom_logger_v2/test_hooks.py create mode 100644 tests/unit/integrations/custom_logger_v2/test_host.py create mode 100644 tests/unit/integrations/custom_logger_v2/test_legacy.py create mode 100644 tests/unit/integrations/custom_logger_v2/test_logger.py create mode 100644 tests/unit/integrations/custom_logger_v2/test_manifest.py create mode 100644 tests/unit/integrations/custom_logger_v2/test_outcomes.py create mode 100644 tests/unit/integrations/custom_logger_v2/test_patches.py create mode 100644 tests/unit/integrations/custom_logger_v2/test_planner.py create mode 100644 tests/unit/integrations/custom_logger_v2/test_selection.py create mode 100644 tests/unit/proxy/test_v2_hook_routing.py create mode 100644 tests/unit/router_utils/test_fallback_event_v2_dispatch.py create mode 100644 tests/unit/test_router/core/test_callback_filters_v2_dispatch.py diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 4206620d813..16a68bc52e0 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -3,6 +3,8 @@ import re import traceback from collections.abc import AsyncGenerator, Mapping, Sequence +from functools import cache +from itertools import takewhile from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional @@ -1120,3 +1122,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac msg[key] = self._redact_base64(value=val, max_depth=max_depth) filtered_messages.append(msg) return filtered_messages + + +@cache +def subclass_overrides_hook(cls: type, hook_name: str, base: type = CustomLogger) -> bool: + below_base: Final = takewhile(lambda klass: klass is not base, cls.__mro__) + return any(hook_name in vars(klass) for klass in below_base) diff --git a/litellm/integrations/custom_logger_v2/__init__.py b/litellm/integrations/custom_logger_v2/__init__.py new file mode 100644 index 00000000000..91896830abd --- /dev/null +++ b/litellm/integrations/custom_logger_v2/__init__.py @@ -0,0 +1,74 @@ +from litellm.integrations.custom_logger_v2.call import CallContext, CallFacts, CallIds, CallKind, CallTarget, CallTiming +from litellm.integrations.custom_logger_v2.dispatcher import Dispatcher, HookResult, RequestScope +from litellm.integrations.custom_logger_v2.enforcement import ContractViolation, UndeclaredReadError +from litellm.integrations.custom_logger_v2.enums import Activation, Field, FieldTier, Hook, OptOutBy +from litellm.integrations.custom_logger_v2.hooks import ( + HOOK_SPECS, + HookSpec, + InvalidDeclarationError, + Subscription, +) +from litellm.integrations.custom_logger_v2.logger import CustomLoggerV2, on +from litellm.integrations.custom_logger_v2.manifest import CallbackManifest, build_manifest +from litellm.integrations.custom_logger_v2.offpath import InProcessSink, OffPathSink +from litellm.integrations.custom_logger_v2.outcomes import ( + Abort, + InvalidPatchError, + JsonValue, + MergeField, + Outcome, + Patch, + Reroute, + Respond, + SetField, +) +from litellm.integrations.custom_logger_v2.patches import apply_patch, apply_patches +from litellm.integrations.custom_logger_v2.planner import HookPlan, plan_hook +from litellm.integrations.custom_logger_v2.registry import CallbackRegistry +from litellm.integrations.custom_logger_v2.selection import RequestSelection, Selection, is_selected + +__all__ = [ + "HOOK_SPECS", + "Abort", + "Activation", + "CallContext", + "CallFacts", + "CallIds", + "CallKind", + "CallTarget", + "CallTiming", + "CallbackManifest", + "CallbackRegistry", + "ContractViolation", + "CustomLoggerV2", + "Dispatcher", + "Field", + "FieldTier", + "Hook", + "HookPlan", + "HookResult", + "HookSpec", + "InProcessSink", + "InvalidDeclarationError", + "InvalidPatchError", + "JsonValue", + "MergeField", + "OffPathSink", + "OptOutBy", + "Outcome", + "Patch", + "RequestScope", + "RequestSelection", + "Reroute", + "Respond", + "Selection", + "SetField", + "Subscription", + "UndeclaredReadError", + "apply_patch", + "apply_patches", + "build_manifest", + "is_selected", + "on", + "plan_hook", +] diff --git a/litellm/integrations/custom_logger_v2/call.py b/litellm/integrations/custom_logger_v2/call.py new file mode 100644 index 00000000000..0d96a86faa3 --- /dev/null +++ b/litellm/integrations/custom_logger_v2/call.py @@ -0,0 +1,55 @@ +from dataclasses import dataclass +from typing import Final, Generic, TypeVar + +from litellm.integrations.custom_logger_v2.enums import Hook + +StateT: Final = TypeVar("StateT") + + +@dataclass(frozen=True, slots=True) +class CallIds: + call_id: str + trace_id: str | None = None + session_id: str | None = None + parent_call_id: str | None = None + + +@dataclass(frozen=True, slots=True) +class CallKind: + call_type: str + route: str | None = None + stream: bool = False + + +@dataclass(frozen=True, slots=True) +class CallTarget: + model: str + model_group: str | None = None + provider: str | None = None + api_base: str | None = None + deployment_id: str | None = None + attempt: int | None = None + + +@dataclass(frozen=True, slots=True) +class CallTiming: + start: float + provider_start: float | None = None + first_token: float | None = None + end: float | None = None + + +@dataclass(frozen=True, slots=True) +class CallFacts: + ids: CallIds + kind: CallKind + target: CallTarget + timing: CallTiming + cache_hit: bool | None = None + + +@dataclass(frozen=True, slots=True) +class CallContext(Generic[StateT]): + hook: Hook + call: CallFacts | None + state: StateT diff --git a/litellm/integrations/custom_logger_v2/dispatcher.py b/litellm/integrations/custom_logger_v2/dispatcher.py new file mode 100644 index 00000000000..058a30efe9e --- /dev/null +++ b/litellm/integrations/custom_logger_v2/dispatcher.py @@ -0,0 +1,400 @@ +import asyncio +import inspect +import logging +from collections import Counter +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain +from operator import itemgetter +from types import TracebackType +from typing import Final, Protocol, TypeAlias, TypeVar + +from typing_extensions import Self + +from litellm.integrations.custom_logger_v2.call import CallContext, CallFacts +from litellm.integrations.custom_logger_v2.enforcement import ContractViolation, Enforced, Policy, Unchecked +from litellm.integrations.custom_logger_v2.enums import Field, Hook +from litellm.integrations.custom_logger_v2.hooks import HOOK_SPECS +from litellm.integrations.custom_logger_v2.manifest import LegacyBinding +from litellm.integrations.custom_logger_v2.offpath import InProcessSink, OffPathSink +from litellm.integrations.custom_logger_v2.outcomes import Abort, Outcome, Patch, Reroute, Respond +from litellm.integrations.custom_logger_v2.patches import apply_patch, apply_patches +from litellm.integrations.custom_logger_v2.planner import EMPTY_PLAN, Entry, HookPlan +from litellm.integrations.custom_logger_v2.registry import CallbackRegistry, RegistrySnapshot +from litellm.integrations.custom_logger_v2.selection import RequestSelection, is_selected + +LegacyT: Final = TypeVar("LegacyT") + +Between: TypeAlias = Callable[[Mapping[Field, object]], Awaitable[Mapping[Field, object]]] +SyncBinding: TypeAlias = Callable[[object, object], object] + +SWALLOWING_HOOKS: Final = frozenset({Hook.response_headers, Hook.post_call_failure}) + +_LOGGER: Final = logging.getLogger("LiteLLM") + + +class _Handler(Protocol): + def __call__(self, view: object, ctx: CallContext[object], /) -> object: ... + + +@dataclass(frozen=True, slots=True) +class HookResult: + fields: Mapping[Field, object] + patches: tuple[Patch, ...] = () + outcome: Outcome | None = None + + +@dataclass(frozen=True, slots=True) +class _Chain: + fields: Mapping[Field, object] + patches: tuple[tuple[int, Patch], ...] = () + reroute: Reroute | None = None + terminal: Abort | Respond | None = None + + @property + def ordered_patches(self) -> tuple[Patch, ...]: + return tuple(patch for _, patch in sorted(self.patches, key=itemgetter(0))) + + +def _handler(logger: object, name: str) -> _Handler: + found: Final[object] = getattr(logger, name, None) + if not callable(found): + raise ContractViolation(f"{name} is not callable") + return found + + +async def _awaited(returned: object) -> object: + return await returned if inspect.isawaitable(returned) else returned + + +def _decides(hook: Hook, result: object) -> bool: + match result: + case Abort() | Respond(): + return True + case ContractViolation(): + return True + case Exception(): + return hook not in SWALLOWING_HOOKS + case BaseException(): + return True + case _: + return False + + +def _fatal(error: BaseException) -> bool: + return isinstance(error, ContractViolation) or not isinstance(error, Exception) + + +def _split( + stages: tuple[tuple[Entry, ...], ...], +) -> tuple[tuple[tuple[Entry, ...], ...], tuple[tuple[Entry, ...], ...]]: + if stages and stages[-1][0].subscription.parallel: + return stages[:-1], stages[-1:] + return stages, () + + +def _legacy_start( + entry: Entry, binding: Callable[[object, object], LegacyT] | None, fields: Mapping[Field, object] +) -> LegacyT | None: + call: Final = entry.manifest.legacy_call + if call is None or binding is None: + raise ContractViolation( + f"{entry.manifest.selection.name} is a legacy entry, so the hook point must pass a binding" + ) + return call(binding, fields) + + +async def _run_legacy(entry: Entry, binding: LegacyBinding | None, fields: Mapping[Field, object]) -> object: + started: Final = _legacy_start(entry, binding, fields) + return None if started is None else await started + + +async def _gather_started(started: tuple[Awaitable[object] | None, ...], fail_fast: bool) -> tuple[object, ...]: + pending: Final = tuple(run for run in started if run is not None) + if not pending: + return tuple(None for _ in started) + done: Final = iter( + await asyncio.gather(*pending) if fail_fast else await asyncio.gather(*pending, return_exceptions=True) + ) + return tuple(None if run is None else next(done) for run in started) + + +def _fail_fast(hook: Hook, stage: tuple[Entry, ...]) -> bool: + return HOOK_SPECS[hook].fail_fast_parallel and all(entry.subscription.parallel for entry in stage) + + +async def run_legacy_plan( + hook: Hook, + plan: HookPlan, + fields: Mapping[Field, object], + binding: LegacyBinding, + between: Between | None = None, +) -> Mapping[Field, object]: + serial, parallel = _split(plan.stages) + current: Mapping[Field, object] = fields # rebind-ok: each legacy entry may replace the kwargs + for stage in serial: + for entry in stage: + result = await _run_legacy(entry, binding, current) + if isinstance(result, Patch): + current = apply_patch(current, result) + if between is not None: + current = await between(current) + for stage in parallel: + started = tuple(_legacy_start(entry, binding, current) for entry in stage) + results = await _gather_started(started, _fail_fast(hook, stage)) + failure = next((result for result in results if isinstance(result, BaseException)), None) + if failure is not None: + raise failure + for entry in plan.inline: + _ = await _run_legacy(entry, binding, current) + return current + + +def run_legacy_plan_sync( + plan: HookPlan, fields: Mapping[Field, object], binding: SyncBinding +) -> Mapping[Field, object]: + staged: Final = tuple(chain.from_iterable(plan.stages)) + if any(entry.subscription.parallel for entry in staged): + raise ContractViolation("a synchronous hook cannot run parallel legacy entries") + current: Mapping[Field, object] = fields # rebind-ok: each legacy entry may replace the kwargs + for entry in (*staged, *plan.inline): + result = _legacy_start(entry, binding, current) + if isinstance(result, Patch): + current = apply_patch(current, result) + return current + + +class Dispatcher: + def __init__( + self, + registry: CallbackRegistry, + sink: OffPathSink | None = None, + enforce: bool = False, + ) -> None: + self.registry: Final = registry + self.sink: Final[OffPathSink] = sink or InProcessSink() + self.policy: Final[Policy] = Enforced() if enforce else Unchecked() + self.failures: Final[Counter[tuple[str, Hook]]] = Counter() # mutable-ok: process-wide failure counts + + def open_scope(self, selection: RequestSelection) -> "RequestScope": + return RequestScope(self, self.registry.snapshot, selection) + + def record_failure(self, entry: Entry, hook: Hook, error: BaseException) -> None: + self.failures[(entry.manifest.selection.name, hook)] += 1 + _LOGGER.warning("custom logger v2 %s failed on %s: %r", entry.manifest.selection.name, hook.value, error) + + +class RequestScope: + def __init__(self, dispatcher: Dispatcher, snapshot: RegistrySnapshot, selection: RequestSelection) -> None: + self._dispatcher: Final = dispatcher + self._snapshot: Final = snapshot + self._selection: Final = selection + self._states: dict[int, object] = {} # mutable-ok: request_state instances are created lazily per logger + self.reroute: Reroute | None = None + + def __enter__(self) -> Self: + return self + + def __exit__( + self, kind: type[BaseException] | None, error: BaseException | None, trace: TracebackType | None + ) -> None: + self.close() + + def close(self) -> None: + self._states.clear() + + def selects(self, entry: Entry) -> bool: + return is_selected(entry.manifest.selection, entry.activation, self._selection) + + def plan(self, hook: Hook) -> HookPlan: + entries: Final = self._snapshot.entries[hook] + if not entries: + return EMPTY_PLAN + selected: Final = tuple(index for index, entry in enumerate(entries) if self.selects(entry)) + return self._snapshot.plan(hook, selected) if selected else EMPTY_PLAN + + def reads(self, hook: Hook) -> frozenset[Field]: + return self.plan(hook).reads + + async def run_hook( + self, + hook: Hook, + fields: Mapping[Field, object], + call: CallFacts | None, + ) -> HookResult: + plan: Final = self.plan(hook) + if plan is EMPTY_PLAN: + return HookResult(fields=fields) + chain: Final = await self._run_stages(hook, plan.stages, _Chain(fields=fields), call, None) + patches: Final = chain.ordered_patches + if chain.terminal is not None: + return HookResult(fields=chain.fields, patches=patches, outcome=chain.terminal) + await self._run_awaited(hook, plan.awaited, chain.fields, call) + if plan.off_path: + self._dispatcher.sink.submit(self._deliver_off_path(hook, plan.off_path, chain.fields, call)) + if chain.reroute is not None: + self.reroute = chain.reroute + return HookResult(fields=chain.fields, patches=patches, outcome=chain.reroute) + + async def run_plan( + self, + hook: Hook, + plan: HookPlan, + fields: Mapping[Field, object], + call: CallFacts | None, + binding: LegacyBinding, + between: Between | None = None, + ) -> HookResult: + serial, parallel = _split(plan.stages) + chain: Final = await self._run_stages(hook, serial, _Chain(fields=fields), call, binding) + if chain.terminal is not None: + return HookResult(fields=chain.fields, patches=chain.ordered_patches, outcome=chain.terminal) + joined: Final = ( + chain + if between is None + else _Chain(fields=await between(chain.fields), patches=chain.patches, reroute=chain.reroute) + ) + done: Final = await self._run_stages(hook, parallel, joined, call, binding) + if done.terminal is not None: + return HookResult(fields=done.fields, patches=done.ordered_patches, outcome=done.terminal) + for entry in plan.inline: + _ = await _run_legacy(entry, binding, done.fields) + await self._run_awaited(hook, plan.awaited, done.fields, call) + if plan.off_path: + self._dispatcher.sink.submit(self._deliver_off_path(hook, plan.off_path, done.fields, call)) + if done.reroute is not None: + self.reroute = done.reroute + return HookResult(fields=done.fields, patches=done.ordered_patches, outcome=done.reroute) + + def run_plan_sync( + self, + hook: Hook, + plan: HookPlan, + fields: Mapping[Field, object], + call: CallFacts | None, + binding: SyncBinding, + ) -> Mapping[Field, object]: + v2_entries: Final = tuple(entry for entry in chain.from_iterable(plan.stages) if not entry.subscription.legacy) + if v2_entries or plan.awaited: + raise ContractViolation( + f"{hook.value} runs synchronously here, so only off-path V2 callbacks may subscribe" + ) + result: Final = run_legacy_plan_sync(plan, fields, binding) + if plan.off_path: + self._dispatcher.sink.submit(self._deliver_off_path(hook, plan.off_path, result, call)) + return result + + def _state(self, entry: Entry) -> object: + factory: Final = entry.manifest.request_state + if factory is None: + return None + key: Final = id(entry.manifest.logger) + if key not in self._states: + self._states[key] = factory() + return self._states[key] + + async def _invoke( + self, hook: Hook, entry: Entry, fields: Mapping[Field, object], call: CallFacts | None, state: object + ) -> object: + policy: Final = self._dispatcher.policy + handler: Final = _handler(entry.manifest.logger, entry.subscription.handler) + context: Final = CallContext[object](hook=hook, call=call, state=state) + view: Final = policy.view(entry.subscription, fields) + result: Final = ( + await _awaited(handler(view, context)) + if inspect.iscoroutinefunction(handler) + else await asyncio.to_thread(handler, view, context) + ) + policy.check(entry.subscription, result) + return result + + def _start( + self, + hook: Hook, + entry: Entry, + fields: Mapping[Field, object], + call: CallFacts | None, + live: bool, + binding: LegacyBinding | None, + ) -> Awaitable[object] | None: + if entry.subscription.legacy: + return _legacy_start(entry, binding, fields) + return self._invoke(hook, entry, fields, call, self._state(entry) if live else None) + + async def _gather( + self, + hook: Hook, + entries: Sequence[Entry], + fields: Mapping[Field, object], + call: CallFacts | None, + live: bool, + binding: LegacyBinding | None = None, + ) -> tuple[object, ...]: + if len(entries) == 1 and entries[0].subscription.legacy: + return (await _run_legacy(entries[0], binding, fields),) + started: Final = tuple(self._start(hook, entry, fields, call, live, binding) for entry in entries) + return await _gather_started(started, _fail_fast(hook, tuple(entries))) + + async def _run_stages( + self, + hook: Hook, + stages: tuple[tuple[Entry, ...], ...], + chain: _Chain, + call: CallFacts | None, + binding: LegacyBinding | None, + ) -> _Chain: + if not stages: + return chain + stage: Final = stages[0] + results: Final = await self._gather(hook, stage, chain.fields, call, live=True, binding=binding) + folded: Final = self._fold(hook, stage, results, chain) + if folded.terminal is not None: + return folded + return await self._run_stages(hook, stages[1:], folded, call, binding) + + def _fold(self, hook: Hook, stage: tuple[Entry, ...], results: tuple[object, ...], chain: _Chain) -> _Chain: + pairs: Final = tuple(zip(stage, results, strict=True)) + decisive: Final = next((pair for pair in pairs if _decides(hook, pair[1])), None) + self._count_failures(hook, tuple(pair for pair in pairs if pair is not decisive)) + match decisive: + case (_, BaseException() as error): + raise error + case (_, Abort() | Respond() as outcome): + return _Chain(fields=chain.fields, patches=chain.patches, reroute=chain.reroute, terminal=outcome) + case _: + return self._merge(pairs, chain) + + def _merge(self, pairs: tuple[tuple[Entry, object], ...], chain: _Chain) -> _Chain: + stage_patches: Final = tuple((entry.position, result) for entry, result in pairs if isinstance(result, Patch)) + reroutes: Final = tuple(result for _, result in pairs if isinstance(result, Reroute)) + return _Chain( + fields=apply_patches(chain.fields, (patch for _, patch in stage_patches)), + patches=(*chain.patches, *stage_patches), + reroute=chain.reroute or next(iter(reroutes), None), + ) + + def _count_failures(self, hook: Hook, pairs: Sequence[tuple[Entry, object]]) -> None: + for entry, result in pairs: + if isinstance(result, BaseException) and not entry.subscription.legacy: + self._dispatcher.record_failure(entry, hook, result) + + def _swallow(self, hook: Hook, entries: Sequence[Entry], results: tuple[object, ...]) -> None: + pairs: Final = tuple(zip(entries, results, strict=True)) + fatal: Final = next( + (result for _, result in pairs if isinstance(result, BaseException) and _fatal(result)), + None, + ) + self._count_failures(hook, tuple(pair for pair in pairs if pair[1] is not fatal)) + if fatal is not None: + raise fatal + + async def _run_awaited( + self, hook: Hook, entries: tuple[Entry, ...], fields: Mapping[Field, object], call: CallFacts | None + ) -> None: + if entries: + self._swallow(hook, entries, await self._gather(hook, entries, fields, call, live=True)) + + async def _deliver_off_path( + self, hook: Hook, entries: tuple[Entry, ...], fields: Mapping[Field, object], call: CallFacts | None + ) -> None: + self._swallow(hook, entries, await self._gather(hook, entries, fields, call, live=False)) diff --git a/litellm/integrations/custom_logger_v2/enforcement.py b/litellm/integrations/custom_logger_v2/enforcement.py new file mode 100644 index 00000000000..8e1c709c741 --- /dev/null +++ b/litellm/integrations/custom_logger_v2/enforcement.py @@ -0,0 +1,95 @@ +from collections.abc import Mapping +from typing import Final, Protocol + +from litellm.integrations.custom_logger_v2.enums import Field, Hook +from litellm.integrations.custom_logger_v2.hooks import HOOK_SPECS, Subscription +from litellm.integrations.custom_logger_v2.outcomes import Abort, Patch, Reroute, Respond + + +class ContractViolation(Exception): + pass + + +class UndeclaredReadError(ContractViolation, AttributeError): + pass + + +def _field_named(name: str) -> Field | None: + member: Final = Field._value2member_map_.get(name) + return member if isinstance(member, Field) else None + + +class FieldView: + __slots__ = ("_fields",) + + def __init__(self, fields: Mapping[Field, object]) -> None: + self._fields = fields + + def __getattr__(self, name: str) -> object: + field: Final = _field_named(name) + if field is None or field not in self._fields: + raise AttributeError(name) + return self._fields[field] + + +class DeclaredView: + __slots__ = ("_fields", "_reads") + + def __init__(self, fields: Mapping[Field, object], reads: frozenset[Field]) -> None: + self._fields = fields + self._reads = reads + + def __getattr__(self, name: str) -> object: + field: Final = _field_named(name) + if field is None or field not in self._reads: + raise UndeclaredReadError(f"{name} is not among the declared reads") + if field not in self._fields: + raise AttributeError(name) + return self._fields[field] + + +class Policy(Protocol): + def view(self, subscription: Subscription, fields: Mapping[Field, object]) -> object: ... + + def check(self, subscription: Subscription, result: object) -> None: ... + + +class Unchecked: + def view(self, subscription: Subscription, fields: Mapping[Field, object]) -> object: + return FieldView(fields) + + def check(self, subscription: Subscription, result: object) -> None: + return None + + +def _outcome_problem(hook: Hook, subscription: Subscription, outcome: Abort | Respond | Reroute) -> str | None: + if not subscription.may_abort: + return f"returned {type(outcome).__name__} without may_abort" + if type(outcome) not in HOOK_SPECS[hook].outcomes: + return f"returned {type(outcome).__name__}, which {hook.value} does not allow" + return None + + +def _result_problem(subscription: Subscription, result: object) -> str | None: + match result: + case None: + return None + case Patch(): + undeclared: Final = result.fields - subscription.writes + if not undeclared: + return None + return f"patched {', '.join(sorted(field.value for field in undeclared))} outside its declared writes" + case Abort() | Respond() | Reroute(): + return _outcome_problem(subscription.hook, subscription, result) + case _: + return f"returned {type(result).__name__}, which is neither None, a Patch nor an outcome" + + +class Enforced: + def view(self, subscription: Subscription, fields: Mapping[Field, object]) -> object: + return DeclaredView(fields, subscription.reads) + + def check(self, subscription: Subscription, result: object) -> None: + problem: Final = _result_problem(subscription, result) + if problem is not None: + raise ContractViolation(f"{subscription.handler} on {subscription.hook.value} {problem}") diff --git a/litellm/integrations/custom_logger_v2/enums.py b/litellm/integrations/custom_logger_v2/enums.py new file mode 100644 index 00000000000..fcc63ecc4db --- /dev/null +++ b/litellm/integrations/custom_logger_v2/enums.py @@ -0,0 +1,97 @@ +from enum import Enum, IntEnum +from typing import Final + + +class Hook(str, Enum): + prompt_resolve = "prompt_resolve" + pre_route = "pre_route" + filter_deployments = "filter_deployments" + pre_call = "pre_call" + during_call = "during_call" + pre_deployment_call = "pre_deployment_call" + post_deployment_success = "post_deployment_success" + post_deployment_failure = "post_deployment_failure" + post_call_success = "post_call_success" + post_call_failure = "post_call_failure" + response_headers = "response_headers" + stream_chunk = "stream_chunk" + list_models = "list_models" + mcp_tool_result = "mcp_tool_result" + log_prepare = "log_prepare" + provider_request_sent = "provider_request_sent" + provider_response_received = "provider_response_received" + attempt_ended = "attempt_ended" + call_ended = "call_ended" + router_event = "router_event" + audit_event = "audit_event" + + +class FieldTier(IntEnum): + small_record = 1 + content = 2 + wire = 3 + compatibility = 4 + + +class Field(str, Enum): + _value_: str + bit: int + tier: FieldTier + + def __new__(cls, value: str, bit: int, tier: FieldTier) -> "Field": + member: Final = str.__new__(cls, value) + member._value_ = value + member.bit = bit + member.tier = tier + return member + + model = ("model", 0, FieldTier.small_record) + key_identity = ("key_identity", 1, FieldTier.small_record) + deployment = ("deployment", 2, FieldTier.small_record) + candidate_deployments = ("candidate_deployments", 3, FieldTier.small_record) + request_tags = ("request_tags", 4, FieldTier.small_record) + dynamic_params = ("dynamic_params", 5, FieldTier.small_record) + guardrail_info = ("guardrail_info", 6, FieldTier.small_record) + outcome = ("outcome", 7, FieldTier.small_record) + error = ("error", 8, FieldTier.small_record) + usage = ("usage", 9, FieldTier.small_record) + cost = ("cost", 10, FieldTier.small_record) + metadata = ("metadata", 11, FieldTier.content) + messages = ("messages", 12, FieldTier.content) + input = ("input", 13, FieldTier.content) + tools = ("tools", 14, FieldTier.content) + optional_params = ("optional_params", 15, FieldTier.content) + request_headers = ("request_headers", 16, FieldTier.content) + response = ("response", 17, FieldTier.content) + response_headers = ("response_headers", 18, FieldTier.content) + stream_chunk = ("stream_chunk", 19, FieldTier.content) + assembled_stream = ("assembled_stream", 20, FieldTier.content) + tool_result = ("tool_result", 21, FieldTier.content) + model_list = ("model_list", 22, FieldTier.content) + request_body = ("request_body", 23, FieldTier.wire) + provider_request_body = ("provider_request_body", 24, FieldTier.wire) + provider_response_raw = ("provider_response_raw", 25, FieldTier.wire) + provider_response_headers = ("provider_response_headers", 26, FieldTier.wire) + standard_logging_payload = ("standard_logging_payload", 27, FieldTier.compatibility) + deprecated_kwargs = ("deprecated_kwargs", 28, FieldTier.compatibility) + deprecated_response_obj = ("deprecated_response_obj", 29, FieldTier.compatibility) + + +class Activation(str, Enum): + required = "required" + default_on = "default_on" + opt_in = "opt_in" + + +class OptOutBy(str, Enum): + admin = "admin" + caller = "caller" + + +ALL_FIELDS_MASK: Final = sum(1 << field.bit for field in Field) + + +def field_mask(fields: frozenset[Field]) -> int: + if Field.deprecated_kwargs in fields: + return ALL_FIELDS_MASK + return sum(1 << field.bit for field in fields) diff --git a/litellm/integrations/custom_logger_v2/hook_point_fields.py b/litellm/integrations/custom_logger_v2/hook_point_fields.py new file mode 100644 index 00000000000..67f987df83d --- /dev/null +++ b/litellm/integrations/custom_logger_v2/hook_point_fields.py @@ -0,0 +1,98 @@ +from collections.abc import Iterator, Mapping +from functools import reduce +from typing import Final + +from litellm.integrations.custom_logger_v2.enums import Field +from litellm.integrations.custom_logger_v2.legacy import KwargsFields +from litellm.integrations.custom_logger_v2.outcomes import InvalidPatchError, Patch, PatchOp, SetField +from litellm.integrations.custom_logger_v2.patches import LiveFields + +_CANDIDATE_FIELDS: Final = (Field.deprecated_kwargs, Field.candidate_deployments, Field.model) + + +class CandidateFields(LiveFields): + __slots__ = ("_deployments", "_model") + + def __init__(self, deployments: object, model: str) -> None: + self._deployments: Final = deployments + self._model: Final = model + + @property + def deployments(self) -> object: + return self._deployments + + def __getitem__(self, field: Field) -> object: + if field is Field.deprecated_kwargs: + return self._deployments + if field is Field.model: + return self._model + if field is Field.candidate_deployments and isinstance(self._deployments, list | tuple): + return tuple(self._deployments) # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType, reportUnknownVariableType] # untyped deployment dicts + raise KeyError(field) + + def __iter__(self) -> Iterator[Field]: + return iter(tuple(field for field in _CANDIDATE_FIELDS if self._has(field))) + + def __len__(self) -> int: + return sum(1 for field in _CANDIDATE_FIELDS if self._has(field)) + + def _has(self, field: Field) -> bool: + return field is not Field.candidate_deployments or isinstance(self._deployments, list | tuple) + + def apply(self, patch: Patch) -> Mapping[Field, object]: + return reduce(_narrow, patch.ops, self) + + +def _narrow(fields: Mapping[Field, object], op: PatchOp) -> Mapping[Field, object]: + if not isinstance(fields, CandidateFields) or not isinstance(op, SetField): + raise InvalidPatchError("candidate deployments are replaced whole, never merged") + if op.field is Field.deprecated_kwargs: + return CandidateFields(op.value, str(fields[Field.model])) + if op.field is not Field.candidate_deployments: + raise InvalidPatchError(f"{op.field.value} has no place in filter_deployments") + return CandidateFields(_narrowed(fields.deployments, op.value), str(fields[Field.model])) + + +def _narrowed(current: object, proposed: object) -> list[object]: # mutable-ok: legacy filters take a list + if not isinstance(current, list | tuple) or not isinstance(proposed, list | tuple): + raise InvalidPatchError("candidate_deployments must stay a sequence of the current candidates") + allowed: Final = frozenset(map(id, current)) # pyright: ignore[reportUnknownArgumentType] # untyped deployment dicts + picked: Final = tuple(map(id, proposed)) # pyright: ignore[reportUnknownArgumentType] # see above + if not allowed.issuperset(picked) or len(frozenset(picked)) != len(picked): + raise InvalidPatchError("filter_deployments may only narrow the candidates, never add or repeat one") + return list(proposed) # pyright: ignore[reportUnknownArgumentType] # see above + + +class ResultFields(LiveFields): + __slots__ = ("_request", "_slot", "value") + + def __init__(self, request: object, keys: Mapping[Field, str], slot: Field, value: object) -> None: + self._request: Final = KwargsFields(request, keys) + self._slot: Final = slot + self.value = value + + def __getitem__(self, field: Field) -> object: + if field is self._slot or field is Field.deprecated_response_obj: + return self.value + return self._request[field] + + def __iter__(self) -> Iterator[Field]: + return iter((*self._request, self._slot, Field.deprecated_response_obj)) + + def __len__(self) -> int: + return len(self._request) + 2 + + def apply(self, patch: Patch) -> Mapping[Field, object]: + values: Final = tuple(_slot_value(op, self._slot) for op in patch.ops) + if values: + self.value = values[-1] + return self + + +def _slot_value(op: PatchOp, slot: Field) -> object: + if not isinstance(op, SetField) or op.field not in (slot, Field.deprecated_response_obj): + raise InvalidPatchError(f"{op.field.value} cannot be patched here, only {slot.value}") + return op.value + + +__all__ = ["CandidateFields", "ResultFields"] diff --git a/litellm/integrations/custom_logger_v2/hooks.py b/litellm/integrations/custom_logger_v2/hooks.py new file mode 100644 index 00000000000..b3fe1b18671 --- /dev/null +++ b/litellm/integrations/custom_logger_v2/hooks.py @@ -0,0 +1,211 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Literal, TypeAlias + +from litellm.integrations.custom_logger_v2.enums import Activation, Field, Hook +from litellm.integrations.custom_logger_v2.outcomes import Abort, Reroute, Respond + +LegacyGuardrailEvent: TypeAlias = Literal["pre_call", "during_call", "post_call", "post_mcp_call"] +OutcomeType: TypeAlias = type[Abort] | type[Respond] | type[Reroute] + + +class InvalidDeclarationError(TypeError): + pass + + +@dataclass(frozen=True, slots=True) +class HookSpec: + legacy_methods: tuple[str, ...] + request_path: bool + writes: frozenset[Field] = frozenset() + outcomes: frozenset[OutcomeType] = frozenset() + legacy_guardrail_events: frozenset[LegacyGuardrailEvent] = frozenset() + has_call: bool = True + subscribable: bool = True + fail_fast_parallel: bool = False + + +_NO_OUTCOME: Final[frozenset[OutcomeType]] = frozenset() +_ABORT: Final[frozenset[OutcomeType]] = frozenset({Abort}) +_ABORT_OR_RESPOND: Final[frozenset[OutcomeType]] = frozenset({Abort, Respond}) +_ANY_OUTCOME: Final[frozenset[OutcomeType]] = frozenset({Abort, Respond, Reroute}) + +HOOK_SPECS: Final[Mapping[Hook, HookSpec]] = MappingProxyType( + { + Hook.prompt_resolve: HookSpec( + legacy_methods=("get_chat_completion_prompt", "async_get_chat_completion_prompt"), + request_path=True, + writes=frozenset({Field.messages, Field.optional_params, Field.model}), + outcomes=_ABORT_OR_RESPOND, + ), + Hook.pre_route: HookSpec( + legacy_methods=("async_pre_routing_hook",), + request_path=True, + writes=frozenset({Field.model, Field.messages}), + outcomes=_ABORT_OR_RESPOND, + ), + Hook.filter_deployments: HookSpec( + legacy_methods=("async_filter_deployments",), + request_path=True, + writes=frozenset({Field.candidate_deployments}), + outcomes=_ABORT, + ), + Hook.pre_call: HookSpec( + legacy_methods=("async_pre_call_hook",), + request_path=True, + writes=frozenset( + { + Field.messages, + Field.optional_params, + Field.metadata, + Field.request_headers, + Field.guardrail_info, + Field.tools, + Field.model, + Field.input, + } + ), + outcomes=_ANY_OUTCOME, + legacy_guardrail_events=frozenset({"pre_call"}), + ), + Hook.during_call: HookSpec( + legacy_methods=("async_moderation_hook",), + request_path=True, + outcomes=_ABORT_OR_RESPOND, + legacy_guardrail_events=frozenset({"during_call"}), + fail_fast_parallel=True, + ), + Hook.pre_deployment_call: HookSpec( + legacy_methods=("async_pre_call_deployment_hook", "async_pre_request_hook"), + request_path=True, + writes=frozenset({Field.provider_request_body, Field.request_headers}), + outcomes=_ABORT_OR_RESPOND, + ), + Hook.post_deployment_success: HookSpec( + legacy_methods=("async_post_call_success_deployment_hook", "async_post_call_streaming_deployment_hook"), + request_path=True, + writes=frozenset({Field.response}), + outcomes=_ABORT_OR_RESPOND, + ), + Hook.post_deployment_failure: HookSpec( + legacy_methods=("async_post_call_failure_deployment_hook",), + request_path=True, + writes=frozenset({Field.error}), + outcomes=_ABORT, + ), + Hook.post_call_success: HookSpec( + legacy_methods=("async_post_call_success_hook",), + request_path=True, + writes=frozenset({Field.response, Field.guardrail_info}), + outcomes=_ABORT_OR_RESPOND, + legacy_guardrail_events=frozenset({"post_call"}), + ), + Hook.post_call_failure: HookSpec( + legacy_methods=("async_post_call_failure_hook",), + request_path=True, + writes=frozenset({Field.error}), + outcomes=_ABORT, + ), + Hook.response_headers: HookSpec( + legacy_methods=("async_post_call_response_headers_hook",), + request_path=True, + writes=frozenset({Field.response_headers}), + outcomes=_NO_OUTCOME, + ), + Hook.stream_chunk: HookSpec( + legacy_methods=( + "async_post_call_streaming_hook", + "async_post_call_streaming_iterator_hook", + "log_stream_event", + ), + request_path=True, + writes=frozenset({Field.stream_chunk}), + outcomes=_ABORT, + subscribable=False, + ), + Hook.list_models: HookSpec( + legacy_methods=("async_filter_listed_models",), + request_path=True, + writes=frozenset({Field.model_list}), + outcomes=_ABORT, + has_call=False, + ), + Hook.mcp_tool_result: HookSpec( + legacy_methods=("async_post_mcp_tool_call_hook",), + request_path=True, + writes=frozenset({Field.tool_result}), + outcomes=_ABORT_OR_RESPOND, + legacy_guardrail_events=frozenset({"post_mcp_call"}), + ), + Hook.log_prepare: HookSpec( + legacy_methods=( + "logging_hook", + "async_logging_hook", + "redact_standard_logging_payload_from_model_call_details", + ), + request_path=False, + writes=frozenset({Field.standard_logging_payload}), + ), + Hook.provider_request_sent: HookSpec(legacy_methods=("log_pre_api_call",), request_path=False), + Hook.provider_response_received: HookSpec(legacy_methods=("log_post_api_call",), request_path=False), + Hook.attempt_ended: HookSpec( + legacy_methods=( + "log_success_event", + "log_failure_event", + "async_log_success_event", + "async_log_failure_event", + ), + request_path=False, + ), + Hook.call_ended: HookSpec(legacy_methods=(), request_path=False), + Hook.router_event: HookSpec( + legacy_methods=( + "log_success_fallback_event", + "log_failure_fallback_event", + "log_model_group_rate_limit_error", + ), + request_path=False, + ), + Hook.audit_event: HookSpec( + legacy_methods=("async_log_audit_log_event", "async_log_batch_line_items"), + request_path=False, + has_call=False, + ), + } +) + + +@dataclass(frozen=True, slots=True) +class Subscription: + hook: Hook + handler: str + reads: frozenset[Field] = frozenset() + writes: frozenset[Field] = frozenset() + may_abort: bool = False + parallel: bool = False + awaited: bool = False + activation: Activation | None = None + legacy: bool = False + + def __post_init__(self) -> None: + problem: Final = subscription_problem(self) + if problem is not None: + raise InvalidDeclarationError(f"{self.handler} on {self.hook.value}: {problem}") + + +def subscription_problem(subscription: Subscription) -> str | None: + spec: Final = HOOK_SPECS[subscription.hook] + if not spec.subscribable: + return "this hook point has no V2 subscription support yet" + legacy_writes: Final = frozenset({Field.deprecated_kwargs}) if subscription.legacy else frozenset[Field]() + forbidden: Final = subscription.writes - spec.writes - legacy_writes + if forbidden: + return f"writes {', '.join(sorted(field.value for field in forbidden))}, which this hook point cannot apply" + if subscription.may_abort and not spec.outcomes: + return "may_abort is set, but this hook point allows no outcome" + if subscription.awaited and (subscription.writes or subscription.may_abort): + return "an awaited subscription must neither write nor abort" + if subscription.parallel and (subscription.writes or not subscription.may_abort): + return "parallel is only for subscriptions that may abort and write nothing" + return None diff --git a/litellm/integrations/custom_logger_v2/host.py b/litellm/integrations/custom_logger_v2/host.py new file mode 100644 index 00000000000..efa0f44d026 --- /dev/null +++ b/litellm/integrations/custom_logger_v2/host.py @@ -0,0 +1,62 @@ +from collections.abc import Generator +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Final + +from litellm.integrations.custom_logger_v2.dispatcher import Dispatcher, RequestScope +from litellm.integrations.custom_logger_v2.outcomes import Outcome, Patch +from litellm.integrations.custom_logger_v2.registry import CallbackRegistry +from litellm.integrations.custom_logger_v2.selection import RequestSelection + +_NO_SELECTION: Final = RequestSelection() +_CURRENT_SCOPE: Final[ContextVar[RequestScope | None]] = ContextVar("custom_logger_v2_scope", default=None) + + +class V2Outcome(Exception): + def __init__(self, outcome: Outcome | None, patches: tuple[Patch, ...]) -> None: + super().__init__(f"a V2 callback ended the hook with {type(outcome).__name__}") + self.outcome: Final = outcome + self.patches: Final = patches + + +class _Host: + def __init__(self) -> None: + self._dispatcher: Dispatcher | None = None + + def dispatcher(self) -> Dispatcher: + if self._dispatcher is None: + self._dispatcher = Dispatcher(CallbackRegistry()) + return self._dispatcher + + def reset(self) -> None: + self._dispatcher = None + + +_HOST: Final = _Host() + + +def get_dispatcher() -> Dispatcher: + return _HOST.dispatcher() + + +def reset_host_for_tests() -> None: + _HOST.reset() + + +def current_scope() -> RequestScope | None: + return _CURRENT_SCOPE.get() + + +@contextmanager +def bind_scope(scope: RequestScope) -> Generator[RequestScope]: + token: Final = _CURRENT_SCOPE.set(scope) + try: + yield scope + finally: + _CURRENT_SCOPE.reset(token) + + +@contextmanager +def request_scope(selection: RequestSelection = _NO_SELECTION) -> Generator[RequestScope]: + with _HOST.dispatcher().open_scope(selection) as scope, bind_scope(scope): + yield scope diff --git a/litellm/integrations/custom_logger_v2/legacy.py b/litellm/integrations/custom_logger_v2/legacy.py new file mode 100644 index 00000000000..242301eb432 --- /dev/null +++ b/litellm/integrations/custom_logger_v2/legacy.py @@ -0,0 +1,434 @@ +import weakref +from collections.abc import Callable, Iterator, Mapping, MutableMapping, Sequence +from dataclasses import dataclass +from enum import Enum +from functools import lru_cache, reduce +from itertools import count +from types import MappingProxyType +from typing import Final, TypeAlias, TypeVar + +from typing_extensions import TypeIs # noqa: TID251 # narrows the untyped legacy kwargs without copying them + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger, subclass_overrides_hook +from litellm.integrations.custom_logger_v2.call import CallFacts +from litellm.integrations.custom_logger_v2.dispatcher import ( + Between, + HookResult, + RequestScope, + run_legacy_plan, + run_legacy_plan_sync, +) +from litellm.integrations.custom_logger_v2.enums import Activation, Field, Hook +from litellm.integrations.custom_logger_v2.hooks import HOOK_SPECS, Subscription +from litellm.integrations.custom_logger_v2.host import V2Outcome, current_scope, get_dispatcher +from litellm.integrations.custom_logger_v2.logger import CustomLoggerV2 +from litellm.integrations.custom_logger_v2.manifest import CallbackManifest, LegacyBinding, build_manifest +from litellm.integrations.custom_logger_v2.outcomes import Abort, InvalidPatchError, MergeField, Patch, PatchOp, Respond +from litellm.integrations.custom_logger_v2.patches import LiveFields +from litellm.integrations.custom_logger_v2.planner import Entry, HookPlan, plan_hook +from litellm.integrations.custom_logger_v2.selection import RequestSelection, Selection + +ResultT: Final = TypeVar("ResultT") + +SyncLegacyBinding: TypeAlias = Callable[[object, object], object] + +# These hook points are the only ones where v1 custom logger subclasses are +# checked to see if they run in parallel +PARALLEL_GUARDRAIL_HOOKS: Final = frozenset({Hook.pre_call, Hook.post_call_success}) +LEGACY_READS: Final = frozenset({Field.deprecated_kwargs, Field.deprecated_response_obj}) +PROXY_REQUEST_KEYS: Final[Mapping[Field, str]] = MappingProxyType( + { + Field.model: "model", + Field.messages: "messages", + Field.input: "input", + Field.tools: "tools", + } +) +SUBSCRIBERS_CACHE_SIZE: Final = 128 +PLAN_CACHE_SIZE: Final = 64 + +_SERIALS: Final = count() + + +class Admit(str, Enum): + overriders = "overriders" + every_logger = "every_logger" + every_entry = "every_entry" + + +_ABSENT: Final = object() +_NO_SELECTION: Final = RequestSelection() + + +@lru_cache(maxsize=None) +def legacy_subscription(hook: Hook, parallel: bool) -> Subscription: + spec: Final = HOOK_SPECS[hook] + may_abort: Final = bool(spec.outcomes) + gathered: Final = may_abort and (parallel or hook is Hook.during_call) + return Subscription( + hook=hook, + handler="call", + reads=LEGACY_READS, + writes=frozenset({Field.deprecated_kwargs}) if spec.writes and not gathered else frozenset(), + may_abort=may_abort, + parallel=gathered, + legacy=True, + ) + + +def _subscribes(cls: type, guardrail: bool, hook: Hook) -> bool: + spec: Final = HOOK_SPECS[hook] + if not spec.subscribable: + return False + if guardrail and spec.legacy_guardrail_events: + return True + return any(subclass_overrides_hook(cls, method) for method in spec.legacy_methods) + + +@lru_cache(maxsize=512) +def _derived_subscriptions(cls: type, guardrail: bool, parallel: bool) -> Mapping[Hook, Subscription]: + return MappingProxyType( + { + hook: legacy_subscription(hook, parallel and hook in PARALLEL_GUARDRAIL_HOOKS) + for hook in Hook + if _subscribes(cls, guardrail, hook) + } + ) + + +class LegacyLogger(CustomLoggerV2): + def __init__(self, inner: CustomLogger) -> None: + guardrail: Final = isinstance(inner, CustomGuardrail) + self.parallel = guardrail and bool(getattr(inner, "run_in_parallel", False)) + self.serial: Final = next(_SERIALS) + self._inner: Final = weakref.ref(inner) + self.manifest: Final = CallbackManifest( + logger=self, + selection=Selection(name=type(inner).__name__, activation=Activation.required), + subscriptions=_derived_subscriptions(type(inner), guardrail, self.parallel), + request_state=None, + legacy_call=self.call, + ) + + @property + def inner(self) -> CustomLogger | None: + return self._inner() + + def subscription(self, hook: Hook, admit: Admit) -> Subscription | None: + derived: Final = self.manifest.subscriptions.get(hook) + if derived is not None or admit is Admit.overriders or not HOOK_SPECS[hook].subscribable: + return derived + return legacy_subscription(hook, self.parallel and hook in PARALLEL_GUARDRAIL_HOOKS) + + def call(self, binding: Callable[[object, object], ResultT], fields: Mapping[Field, object]) -> ResultT | None: + inner: Final = self._inner() + return None if inner is None else binding(inner, fields[Field.deprecated_kwargs]) + + +_ADAPTERS: Final[dict[int, LegacyLogger]] = {} # mutable-ok: adapters live as long as the logger they wrap + + +def adapter_for(inner: CustomLogger) -> LegacyLogger: + key: Final = id(inner) + cached: Final = _ADAPTERS.get(key) + if cached is not None and cached.inner is inner: + return cached + adapter: Final = LegacyLogger(inner) + _ADAPTERS[key] = adapter + _ = weakref.finalize(inner, _ADAPTERS.pop, key, None) + return adapter + + +def _readable(value: object) -> TypeIs[Mapping[str, object]]: # guard-ok: legacy kwargs are an untyped dict + return isinstance(value, Mapping) + + +def _writable( + value: object, +) -> TypeIs[MutableMapping[str, object]]: # guard-ok: legacy kwargs are an untyped dict # mutable-ok: written in place + return isinstance(value, MutableMapping) + + +class OpaqueLegacy(CustomLoggerV2): + """Holds a callback list entry that is not a CustomLogger (a string name or a plain function) at its list + position, for hook points that admit every entry. There is no class to derive subscriptions from, so it + subscribes to the one hook it was built for and runs serially.""" + + def __init__(self, inner: object, hook: Hook) -> None: + self.inner: Final = inner + self.serial: Final = next(_SERIALS) + self.manifest: Final = CallbackManifest( + logger=self, + selection=Selection(name=type(inner).__name__, activation=Activation.required), + subscriptions=MappingProxyType({hook: legacy_subscription(hook, False)}), + request_state=None, + legacy_call=self.call, + ) + + def call(self, binding: Callable[[object, object], ResultT], fields: Mapping[Field, object]) -> ResultT | None: + return binding(self.inner, fields[Field.deprecated_kwargs]) + + +class KwargsFields(LiveFields): + """A live Field view over a legacy hook's untyped kwargs dict, so the dispatcher can run over it without copying. + + Two fields cover the whole call: deprecated_kwargs is the dict itself, and deprecated_response_obj is the + response, when the hook point passes one. Patching either one swaps in the new value. + + Every other field maps to one key in the dict through the hook point's key table (PROXY_REQUEST_KEYS or + LOGGING_KEYS), and is present only when that key is. Patching one writes straight into the dict, so later + callbacks and the request see the change.""" + + __slots__ = ("_keys", "_kwargs", "_response") + + def __init__(self, kwargs: object, keys: Mapping[Field, str], response: object = _ABSENT) -> None: + self._kwargs: Final = kwargs + self._keys: Final = keys + self._response: Final = response + + def _key(self, field: Field) -> str | None: + key: Final = self._keys.get(field) + return key if key is not None and _readable(self._kwargs) and key in self._kwargs else None + + def __getitem__(self, field: Field) -> object: + if field is Field.deprecated_kwargs: + return self._kwargs + if field is Field.deprecated_response_obj and self._response is not _ABSENT: + return self._response + key: Final = self._key(field) + if key is None or not _readable(self._kwargs): + raise KeyError(field) + return self._kwargs[key] + + def __iter__(self) -> Iterator[Field]: + return iter(tuple(field for field in Field if self._has(field))) + + def __len__(self) -> int: + return sum(1 for field in Field if self._has(field)) + + def _has(self, field: Field) -> bool: + if field is Field.deprecated_kwargs: + return True + if field is Field.deprecated_response_obj: + return self._response is not _ABSENT + return self._key(field) is not None + + def apply(self, patch: Patch) -> Mapping[Field, object]: + return reduce(_apply_op, patch.ops, self) + + def replaced(self, field: Field, value: object) -> "KwargsFields": + if field is Field.deprecated_kwargs: + return KwargsFields(value, self._keys, self._response) + return KwargsFields(self._kwargs, self._keys, value) + + def write(self, field: Field, value: object) -> None: + target: Final = self._target(field) + target[self._keys[field]] = value + + def merge(self, field: Field, keys: Mapping[str, object]) -> None: + target: Final = self._target(field) + current: Final = target.get(self._keys[field]) + if current is None: + target[self._keys[field]] = dict(keys) + return + if not _writable(current): + raise InvalidPatchError(f"{field.value} holds no mapping here, so keys cannot be merged into it") + current.update(keys) + + def _target(self, field: Field) -> MutableMapping[str, object]: # mutable-ok: legacy kwargs are the live request + if field not in self._keys: + raise InvalidPatchError(f"{field.value} has no place in the legacy kwargs of this hook") + if not _writable(self._kwargs): + raise InvalidPatchError(f"{field.value} cannot be written because the legacy kwargs are not a mapping") + return self._kwargs + + +def _apply_op(fields: Mapping[Field, object], op: PatchOp) -> Mapping[Field, object]: + if not isinstance(fields, KwargsFields): + raise InvalidPatchError("legacy kwargs patches need KwargsFields") + if op.field in (Field.deprecated_kwargs, Field.deprecated_response_obj): + if isinstance(op, MergeField): + raise InvalidPatchError(f"{op.field.value} is replaced whole, never merged") + return fields.replaced(op.field, op.value) + if isinstance(op, MergeField): + fields.merge(op.field, op.keys) + return fields + fields.write(op.field, op.value) + return fields + + +def kwargs_patch(before: object, after: object) -> Patch | None: + return None if after is before else Patch().set(Field.deprecated_kwargs, after) + + +@dataclass(frozen=True, slots=True, eq=False) +class HookSubscribers: + hook: Hook + entries: tuple[Entry, ...] + plan: HookPlan + has_v2: bool + plans: Callable[[tuple[int, ...]], HookPlan] + + @property + def legacy(self) -> tuple[object, ...]: + return tuple(inner for entry in self.entries if (inner := _wrapped(entry.manifest.logger)) is not None) + + def plan_for(self, scope: RequestScope) -> HookPlan: + selected: Final = tuple( + index for index, entry in enumerate(self.entries) if entry.subscription.legacy or scope.selects(entry) + ) + return self.plan if len(selected) == len(self.entries) else self.plans(selected) + + +def _wrapped(logger: CustomLoggerV2) -> object: + if isinstance(logger, LegacyLogger | OpaqueLegacy): + return logger.inner + return None + + +def _v2_entry(v2_callback: CustomLoggerV2, hook: Hook) -> tuple[CallbackManifest, Subscription] | None: + manifest: Final = build_manifest(v2_callback) + subscription: Final = manifest.subscriptions.get(hook) + return None if subscription is None else (manifest, subscription) + + +def _legacy_entry( + inner: CustomLogger, hook: Hook, admit: Admit, parallel: bool +) -> tuple[CallbackManifest, Subscription] | None: + adapter: Final = adapter_for(inner) + subscription: Final = adapter.subscription(hook, admit) + if subscription is None: + return None + if subscription.parallel and not parallel: + return adapter.manifest, legacy_subscription(hook, False) + return adapter.manifest, subscription + + +def _opaque_entry(item: object, hook: Hook) -> tuple[CallbackManifest, Subscription] | None: + if not HOOK_SPECS[hook].subscribable: + return None + opaque: Final = OpaqueLegacy(item, hook) + return opaque.manifest, opaque.manifest.subscriptions[hook] + + +def _subscriber( + item: object, hook: Hook, admit: Admit, include_v2: bool, parallel: bool +) -> tuple[CallbackManifest, Subscription] | None: + if isinstance(item, CustomLoggerV2): + return _v2_entry(item, hook) if include_v2 else None + if isinstance(item, CustomLogger): + return _legacy_entry(item, hook, admit, parallel) + return _opaque_entry(item, hook) if admit is Admit.every_entry else None + + +def _build_subscribers( + hook: Hook, callbacks: Sequence[object], admit: Admit, include_v2: bool, parallel: bool +) -> HookSubscribers: + members: Final = tuple( + member for item in callbacks if (member := _subscriber(item, hook, admit, include_v2, parallel)) is not None + ) + entries: Final = tuple( + Entry(position, manifest, subscription) for position, (manifest, subscription) in enumerate(members) + ) + + def subset_plan(selected: tuple[int, ...]) -> HookPlan: + return plan_hook(tuple(entries[index] for index in selected)) + + return HookSubscribers( + hook=hook, + entries=entries, + plan=plan_hook(entries), + has_v2=any(not entry.subscription.legacy for entry in entries), + plans=lru_cache(maxsize=PLAN_CACHE_SIZE)(subset_plan), + ) + + +@dataclass(frozen=True, slots=True, eq=False) +class _CachedSubscribers: + subscribers: HookSubscribers + pinned: tuple[object, ...] + + +_SUBSCRIBERS: Final[dict[tuple[object, ...], _CachedSubscribers]] = {} # mutable-ok: bounded subscribers cache + + +def _pinned(item: object) -> bool: + return not isinstance(item, CustomLogger) or isinstance(item, CustomLoggerV2) + + +def _cache_subscribers(key: tuple[object, ...], subscribers: HookSubscribers, callbacks: Sequence[object]) -> None: + if len(_SUBSCRIBERS) >= SUBSCRIBERS_CACHE_SIZE: + _SUBSCRIBERS.clear() + _SUBSCRIBERS[key] = _CachedSubscribers( + subscribers=subscribers, pinned=tuple(item for item in callbacks if _pinned(item)) + ) + for item in callbacks: + if not _pinned(item): + _ = weakref.finalize(item, _SUBSCRIBERS.pop, key, None) + + +def hook_subscribers( + hook: Hook, + callbacks: Sequence[object], + *, + admit: Admit = Admit.overriders, + include_v2: bool = True, + parallel: bool = True, +) -> HookSubscribers: + key: Final = (hook, admit, include_v2, parallel, *map(id, callbacks)) + cached: Final = _SUBSCRIBERS.get(key) + if cached is not None: + return cached.subscribers + subscribers: Final = _build_subscribers(hook, callbacks, admit, include_v2, parallel) + _cache_subscribers(key, subscribers, callbacks) + return subscribers + + +def clear_subscribers_cache_for_tests() -> None: + _SUBSCRIBERS.clear() + + +async def run_subscribers( + subscribers: HookSubscribers, + fields: Mapping[Field, object], + binding: LegacyBinding, + *, + between: Between | None = None, + call: CallFacts | None = None, + scope: RequestScope | None = None, +) -> HookResult: + if not subscribers.has_v2: + return HookResult(fields=await run_legacy_plan(subscribers.hook, subscribers.plan, fields, binding, between)) + bound: Final = scope or current_scope() + if bound is not None: + return await bound.run_plan(subscribers.hook, subscribers.plan_for(bound), fields, call, binding, between) + with get_dispatcher().open_scope(_NO_SELECTION) as transient: + return await transient.run_plan( + subscribers.hook, subscribers.plan_for(transient), fields, call, binding, between + ) + + +def run_subscribers_sync( + subscribers: HookSubscribers, + fields: Mapping[Field, object], + binding: SyncLegacyBinding, + *, + call: CallFacts | None = None, + scope: RequestScope | None = None, +) -> Mapping[Field, object]: + if not subscribers.has_v2: + return run_legacy_plan_sync(subscribers.plan, fields, binding) + bound: Final = scope or current_scope() + if bound is not None: + return bound.run_plan_sync(subscribers.hook, subscribers.plan_for(bound), fields, call, binding) + with get_dispatcher().open_scope(_NO_SELECTION) as transient: + return transient.run_plan_sync(subscribers.hook, subscribers.plan_for(transient), fields, call, binding) + + +def kwargs_after(result: HookResult) -> object: + match result.outcome: + case Abort() | Respond() as terminal: + raise V2Outcome(terminal, result.patches) + case _: + return result.fields[Field.deprecated_kwargs] diff --git a/litellm/integrations/custom_logger_v2/logger.py b/litellm/integrations/custom_logger_v2/logger.py new file mode 100644 index 00000000000..ce8b9f163e0 --- /dev/null +++ b/litellm/integrations/custom_logger_v2/logger.py @@ -0,0 +1,109 @@ +import inspect +from collections.abc import Callable, Mapping, Sequence, Set +from dataclasses import replace +from types import MappingProxyType +from typing import ClassVar, Final, TypeVar + +from litellm.integrations.custom_logger_v2.enums import Activation, Field, Hook +from litellm.integrations.custom_logger_v2.hooks import InvalidDeclarationError, Subscription +from litellm.integrations.custom_logger_v2.selection import Selection + +HandlerT: Final = TypeVar("HandlerT", bound=Callable[..., object]) + +_SUBSCRIPTION_ATTR: Final = "__custom_logger_v2_subscription__" + + +def on( + hook: Hook, + *, + reads: Set[Field] = frozenset(), + writes: Set[Field] = frozenset(), + may_abort: bool = False, + parallel: bool | None = None, + awaited: bool = False, +) -> Callable[[HandlerT], HandlerT]: + def decorate(handler: HandlerT) -> HandlerT: + subscription: Final = Subscription( + hook=hook, + handler=handler.__name__, + reads=frozenset(reads), + writes=frozenset(writes), + may_abort=may_abort, + parallel=(hook is Hook.during_call and may_abort) if parallel is None else parallel, + awaited=awaited, + ) + setattr(handler, _SUBSCRIPTION_ATTR, subscription) + return handler + + return decorate + + +def _declared_subscription(value: object) -> Subscription | None: + marker: Final[object] = getattr(value, _SUBSCRIPTION_ATTR, None) + return marker if isinstance(marker, Subscription) else None + + +def _namespace(cls: type) -> Mapping[str, object]: + return vars(cls) + + +def _declared_in(cls: type) -> Mapping[Hook, Subscription]: + declared: Final = tuple( + replace(subscription, handler=name) + for name, value in _namespace(cls).items() + if (subscription := _declared_subscription(value)) is not None + ) + hooks: Final = tuple(subscription.hook for subscription in declared) + duplicated: Final = sorted({hook.value for hook in hooks if hooks.count(hook) > 1}) + if duplicated: + raise InvalidDeclarationError(f"{cls.__name__} declares more than one handler for {', '.join(duplicated)}") + return MappingProxyType({subscription.hook: subscription for subscription in declared}) + + +def _nearest(layers: Sequence[Mapping[Hook, Subscription]], hook: Hook) -> Subscription | None: + return next((layer[hook] for layer in layers if hook in layer), None) + + +def _handler_problem(cls: type, subscription: Subscription) -> str | None: + handler: Final[object] = getattr(cls, subscription.handler, None) + if not callable(handler): + return f"{subscription.handler} handles {subscription.hook.value} but is not callable" + redeclared: Final = _declared_subscription(handler) + if redeclared is not None and redeclared.hook is not subscription.hook: + return ( + f"{subscription.handler} is decorated for {redeclared.hook.value} but also inherits " + f"the {subscription.hook.value} subscription" + ) + blocking: Final = bool(subscription.writes) or subscription.may_abort or subscription.awaited + if blocking and not inspect.iscoroutinefunction(handler): + return f"{subscription.handler} must be async because it writes, may abort or is awaited" + return None + + +class CustomLoggerV2: + _own_subscriptions: ClassVar[Mapping[Hook, Subscription]] = MappingProxyType({}) + subscriptions: ClassVar[Mapping[Hook, Subscription]] = MappingProxyType({}) + request_state: ClassVar[type | None] = None + declared_hooks: ClassVar[frozenset[Hook]] = frozenset() + + selection: Selection | None = None + enabled_hooks: frozenset[Hook] | None = None + parallel: bool = False + activation_overrides: Mapping[Hook, Activation] = MappingProxyType({}) + + def __init_subclass__(cls) -> None: + super().__init_subclass__() + cls._own_subscriptions = _declared_in(cls) + layers: Final = tuple(base._own_subscriptions for base in cls.__mro__ if issubclass(base, CustomLoggerV2)) + resolved: Final = MappingProxyType( + {hook: subscription for hook in Hook if (subscription := _nearest(layers, hook)) is not None} + ) + problems: Final = tuple( + problem + for subscription in resolved.values() + if (problem := _handler_problem(cls, subscription)) is not None + ) + if problems: + raise InvalidDeclarationError(f"{cls.__name__}: {'; '.join(problems)}") + cls.subscriptions = resolved + CustomLoggerV2.declared_hooks = CustomLoggerV2.declared_hooks | resolved.keys() diff --git a/litellm/integrations/custom_logger_v2/manifest.py b/litellm/integrations/custom_logger_v2/manifest.py new file mode 100644 index 00000000000..5984628feeb --- /dev/null +++ b/litellm/integrations/custom_logger_v2/manifest.py @@ -0,0 +1,73 @@ +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass, replace +from types import MappingProxyType +from typing import Final, Protocol, TypeAlias, TypeVar + +from litellm.integrations.custom_logger_v2.enums import Activation, Field, Hook, OptOutBy +from litellm.integrations.custom_logger_v2.hooks import InvalidDeclarationError, Subscription +from litellm.integrations.custom_logger_v2.logger import CustomLoggerV2 +from litellm.integrations.custom_logger_v2.selection import Selection + +ResultT: Final = TypeVar("ResultT") + +LegacyBinding: TypeAlias = Callable[[object, object], Awaitable[object] | None] + + +class LegacyCall(Protocol): + def __call__( + self, binding: Callable[[object, object], ResultT], fields: Mapping[Field, object], / + ) -> ResultT | None: ... + + +@dataclass(frozen=True, slots=True) +class CallbackManifest: + logger: CustomLoggerV2 + selection: Selection + subscriptions: Mapping[Hook, Subscription] + request_state: type | None + legacy_call: LegacyCall | None = None + + +def _instance_subscription(subscription: Subscription, parallel: bool, activation: Activation | None) -> Subscription: + return replace( + subscription, + parallel=subscription.parallel or (parallel and subscription.may_abort), + activation=activation, + ) + + +def _caller_can_remove(selection: Selection) -> bool: + return selection.honors_no_log or OptOutBy.caller in selection.opt_out_by + + +def build_manifest(logger: CustomLoggerV2) -> CallbackManifest: + name: Final = type(logger).__name__ + declared: Final = type(logger).subscriptions + enabled: Final = frozenset(declared) if logger.enabled_hooks is None else logger.enabled_hooks + not_implemented: Final = enabled - declared.keys() + if not_implemented: + hooks: Final = ", ".join(sorted(hook.value for hook in not_implemented)) + raise InvalidDeclarationError(f"{name} enables {hooks}, which it does not implement") + stray_overrides: Final = logger.activation_overrides.keys() - enabled + if stray_overrides: + overridden: Final = ", ".join(sorted(hook.value for hook in stray_overrides)) + raise InvalidDeclarationError(f"{name} overrides activation for {overridden}, which it does not enable") + subscriptions: Final = MappingProxyType( + { + hook: _instance_subscription(declared[hook], logger.parallel, logger.activation_overrides.get(hook)) + for hook in Hook + if hook in enabled + } + ) + selection: Final = logger.selection or Selection(name=name) + blocking: Final = any(subscription.writes or subscription.may_abort for subscription in subscriptions.values()) + if blocking and _caller_can_remove(selection): + raise InvalidDeclarationError( + f"{selection.name} writes or may abort, so neither no-log nor a caller opt-out may remove it" + ) + return CallbackManifest( + logger=logger, + selection=selection, + subscriptions=subscriptions, + request_state=type(logger).request_state, + ) diff --git a/litellm/integrations/custom_logger_v2/offpath.py b/litellm/integrations/custom_logger_v2/offpath.py new file mode 100644 index 00000000000..e8f309d4878 --- /dev/null +++ b/litellm/integrations/custom_logger_v2/offpath.py @@ -0,0 +1,27 @@ +import asyncio +from collections.abc import Coroutine +from concurrent.futures import ThreadPoolExecutor +from typing import Final, Protocol + + +class OffPathSink(Protocol): + def submit(self, delivery: Coroutine[object, object, None]) -> None: ... + + +class InProcessSink: + def __init__(self) -> None: + self._pending: set[asyncio.Task[None]] = set() # mutable-ok: strong references keep running tasks alive + self._loopless: Final = ThreadPoolExecutor(max_workers=1, thread_name_prefix="custom-logger-v2-off-path") + + def submit(self, delivery: Coroutine[object, object, None]) -> None: + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + _ = self._loopless.submit(asyncio.run, delivery) + return + task: Final[asyncio.Task[None]] = loop.create_task(delivery) + self._pending.add(task) + task.add_done_callback(self._pending.discard) + + async def drain(self) -> None: + await asyncio.gather(*self._pending) diff --git a/litellm/integrations/custom_logger_v2/outcomes.py b/litellm/integrations/custom_logger_v2/outcomes.py new file mode 100644 index 00000000000..aed1355096d --- /dev/null +++ b/litellm/integrations/custom_logger_v2/outcomes.py @@ -0,0 +1,106 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import ClassVar, Final, TypeAlias + +from pydantic import BaseModel + +from litellm.integrations.custom_logger_v2.enums import Field + +JsonValue: TypeAlias = "str | int | float | bool | None | Sequence[JsonValue] | Mapping[str, JsonValue]" + +MERGEABLE_FIELDS: Final = frozenset( + { + Field.dynamic_params, + Field.guardrail_info, + Field.metadata, + Field.optional_params, + Field.request_headers, + Field.response_headers, + Field.provider_request_body, + Field.standard_logging_payload, + } +) +OUTCOME_PATCH_FIELDS: Final = frozenset({Field.guardrail_info, Field.metadata}) + + +class InvalidPatchError(ValueError): + pass + + +@dataclass(frozen=True, slots=True) +class SetField: + field: Field + value: object + + +@dataclass(frozen=True, slots=True) +class MergeField: + field: Field + keys: Mapping[str, object] + + +PatchOp: TypeAlias = SetField | MergeField + + +@dataclass(frozen=True, slots=True) +class Patch: + ops: tuple[PatchOp, ...] = () + + def set(self, field: Field, value: object) -> "Patch": + return Patch((*self.ops, SetField(field, value))) + + def merge(self, field: Field, keys: Mapping[str, object]) -> "Patch": + if field not in MERGEABLE_FIELDS: + raise InvalidPatchError(f"{field.value} is not a mapping field, so it cannot be merged by key") + return Patch((*self.ops, MergeField(field, MappingProxyType(dict(keys))))) + + @property + def fields(self) -> frozenset[Field]: + return frozenset(op.field for op in self.ops) + + +def _require_outcome_patch(patch: Patch | None) -> None: + if patch is None: + return + outside: Final = patch.fields - OUTCOME_PATCH_FIELDS + if outside: + names: Final = ", ".join(sorted(field.value for field in outside)) + raise InvalidPatchError(f"an outcome may only patch guardrail_info and metadata, not {names}") + + +@dataclass(frozen=True, slots=True) +class Abort: + terminal: ClassVar[bool] = True + + status: int + message: str + code: str | None = None + headers: Mapping[str, str] | None = None + body: JsonValue = None + patch: Patch | None = None + + def __post_init__(self) -> None: + _require_outcome_patch(self.patch) + + +@dataclass(frozen=True, slots=True) +class Respond: + terminal: ClassVar[bool] = True + + response: "JsonValue | BaseModel" + status: int = 200 + patch: Patch | None = None + + def __post_init__(self) -> None: + _require_outcome_patch(self.patch) + + +@dataclass(frozen=True, slots=True) +class Reroute: + terminal: ClassVar[bool] = False + + model: str + + +Outcome: TypeAlias = Abort | Respond | Reroute diff --git a/litellm/integrations/custom_logger_v2/patches.py b/litellm/integrations/custom_logger_v2/patches.py new file mode 100644 index 00000000000..b5784295f87 --- /dev/null +++ b/litellm/integrations/custom_logger_v2/patches.py @@ -0,0 +1,70 @@ +from abc import abstractmethod +from collections.abc import Iterable, Iterator, Mapping +from functools import reduce +from types import MappingProxyType +from typing import Final + +from pydantic import TypeAdapter + +from litellm.integrations.custom_logger_v2.enums import Field +from litellm.integrations.custom_logger_v2.outcomes import InvalidPatchError, MergeField, Patch, PatchOp, SetField + +_STRING_KEYED: Final = TypeAdapter(Mapping[str, object]) + + +class LiveFields(Mapping[Field, object]): + __slots__ = () + + @abstractmethod + def apply(self, patch: Patch) -> Mapping[Field, object]: ... + + +class PatchedFields(Mapping[Field, object]): + __slots__ = ("_base", "_changed") + + def __init__(self, base: Mapping[Field, object], changed: Mapping[Field, object]) -> None: + self._base = base + self._changed = changed + + def with_value(self, field: Field, value: object) -> "PatchedFields": + return PatchedFields(self._base, MappingProxyType({**self._changed, field: value})) + + def __getitem__(self, field: Field) -> object: + return self._changed[field] if field in self._changed else self._base[field] + + def __iter__(self) -> Iterator[Field]: + return iter(self._changed.keys() | self._base.keys()) + + def __len__(self) -> int: + return len(self._changed.keys() | self._base.keys()) + + +def _with_value(fields: Mapping[Field, object], field: Field, value: object) -> PatchedFields: + patched: Final = fields if isinstance(fields, PatchedFields) else PatchedFields(fields, MappingProxyType({})) + return patched.with_value(field, value) + + +def _merged(current: object, field: Field, keys: Mapping[str, object]) -> Mapping[str, object]: + if current is None: + return MappingProxyType(dict(keys)) + if not isinstance(current, Mapping): + raise InvalidPatchError(f"{field.value} holds no mapping, so keys cannot be merged into it") + return MappingProxyType({**_STRING_KEYED.validate_python(current), **keys}) + + +def _apply_op(fields: Mapping[Field, object], op: PatchOp) -> Mapping[Field, object]: + match op: + case SetField(field=field, value=value): + return _with_value(fields, field, value) + case MergeField(field=field, keys=keys): + return _with_value(fields, field, _merged(fields.get(field), field, keys)) + + +def apply_patch(fields: Mapping[Field, object], patch: Patch) -> Mapping[Field, object]: + if isinstance(fields, LiveFields): + return fields.apply(patch) + return reduce(_apply_op, patch.ops, fields) + + +def apply_patches(fields: Mapping[Field, object], patches: Iterable[Patch]) -> Mapping[Field, object]: + return reduce(apply_patch, patches, fields) diff --git a/litellm/integrations/custom_logger_v2/planner.py b/litellm/integrations/custom_logger_v2/planner.py new file mode 100644 index 00000000000..bbd2187eca2 --- /dev/null +++ b/litellm/integrations/custom_logger_v2/planner.py @@ -0,0 +1,77 @@ +from collections.abc import Sequence +from dataclasses import dataclass +from functools import reduce +from typing import Final + +from litellm.integrations.custom_logger_v2.enums import Activation, Field, field_mask +from litellm.integrations.custom_logger_v2.hooks import Subscription +from litellm.integrations.custom_logger_v2.manifest import CallbackManifest + + +@dataclass(frozen=True, slots=True, eq=False) +class Entry: + position: int + manifest: CallbackManifest + subscription: Subscription + + @property + def activation(self) -> Activation: + return self.subscription.activation or self.manifest.selection.activation + + +@dataclass(frozen=True, slots=True) +class HookPlan: + stages: tuple[tuple[Entry, ...], ...] = () + inline: tuple[Entry, ...] = () + awaited: tuple[Entry, ...] = () + off_path: tuple[Entry, ...] = () + reads: frozenset[Field] = frozenset() + + +EMPTY_PLAN: Final = HookPlan() + + +def _blocks(subscription: Subscription) -> bool: + return bool(subscription.writes) or subscription.may_abort + + +def _level(placed: tuple[tuple[Subscription, int], ...], subscription: Subscription) -> int: + if subscription.may_abort: + return max((level for _, level in placed), default=-1) + 1 + touched: Final = field_mask(subscription.reads | subscription.writes) + writes: Final = field_mask(subscription.writes) + after_writers: Final = tuple(level + 1 for prior, level in placed if field_mask(prior.writes) & touched) + not_before_readers: Final = tuple(level for prior, level in placed if field_mask(prior.reads) & writes) + return max((*after_writers, *not_before_readers), default=0) + + +def _place( + placed: tuple[tuple[Subscription, int], ...], subscription: Subscription +) -> tuple[tuple[Subscription, int], ...]: + return (*placed, (subscription, _level(placed, subscription))) + + +def _levels(entries: Sequence[Entry]) -> tuple[tuple[Entry, ...], ...]: + placed: Final = reduce(_place, (entry.subscription for entry in entries), ()) + levels: Final = tuple(level for _, level in placed) + return tuple(_at_level(entries, levels, level) for level in range(max(levels, default=-1) + 1)) + + +def _at_level(entries: Sequence[Entry], levels: tuple[int, ...], level: int) -> tuple[Entry, ...]: + return tuple(entry for entry, placed_at in zip(entries, levels, strict=True) if placed_at == level) + + +def plan_hook(entries: Sequence[Entry]) -> HookPlan: + if not entries: + return EMPTY_PLAN + blocking: Final = tuple(entry for entry in entries if _blocks(entry.subscription)) + serial: Final = tuple(entry for entry in blocking if not entry.subscription.parallel) + parallel: Final = tuple(entry for entry in blocking if entry.subscription.parallel) + passive: Final = tuple(entry for entry in entries if not _blocks(entry.subscription)) + return HookPlan( + stages=(*_levels(serial), *((parallel,) if parallel else ())), + inline=tuple(entry for entry in passive if entry.subscription.legacy), + awaited=tuple(entry for entry in passive if entry.subscription.awaited and not entry.subscription.legacy), + off_path=tuple(entry for entry in passive if not (entry.subscription.awaited or entry.subscription.legacy)), + reads=frozenset[Field]().union(*(entry.subscription.reads for entry in entries)), + ) diff --git a/litellm/integrations/custom_logger_v2/registry.py b/litellm/integrations/custom_logger_v2/registry.py new file mode 100644 index 00000000000..ef094005545 --- /dev/null +++ b/litellm/integrations/custom_logger_v2/registry.py @@ -0,0 +1,51 @@ +from collections.abc import Callable, Mapping, Sequence +from functools import lru_cache +from types import MappingProxyType +from typing import Final + +from litellm.integrations.custom_logger_v2.enums import Hook +from litellm.integrations.custom_logger_v2.logger import CustomLoggerV2 +from litellm.integrations.custom_logger_v2.manifest import CallbackManifest +from litellm.integrations.custom_logger_v2.planner import Entry, HookPlan, plan_hook + +DEFAULT_PLAN_CACHE_SIZE: Final = 256 + + +def _hook_entries(manifests: tuple[CallbackManifest, ...], hook: Hook) -> tuple[Entry, ...]: + return tuple( + Entry(position, manifest, manifest.subscriptions[hook]) + for position, manifest in enumerate(manifests) + if hook in manifest.subscriptions + ) + + +class RegistrySnapshot: + __slots__ = ("entries", "manifests", "plan") + + def __init__(self, manifests: tuple[CallbackManifest, ...], plan_cache_size: int) -> None: + self.manifests: Final = manifests + self.entries: Final[Mapping[Hook, tuple[Entry, ...]]] = MappingProxyType( + {hook: _hook_entries(manifests, hook) for hook in Hook} + ) + self.plan: Final[Callable[[Hook, tuple[int, ...]], HookPlan]] = lru_cache(maxsize=plan_cache_size)( + self._build_plan + ) + + def _build_plan(self, hook: Hook, selected: tuple[int, ...]) -> HookPlan: + hook_entries: Final = self.entries[hook] + return plan_hook(tuple(hook_entries[index] for index in selected)) + + +class CallbackRegistry: + def __init__( + self, manifests: Sequence[CallbackManifest] = (), plan_cache_size: int = DEFAULT_PLAN_CACHE_SIZE + ) -> None: + self._plan_cache_size: Final = plan_cache_size + self.snapshot: RegistrySnapshot = RegistrySnapshot(tuple(manifests), plan_cache_size) + + def register(self, manifest: CallbackManifest) -> None: + self.snapshot = RegistrySnapshot((*self.snapshot.manifests, manifest), self._plan_cache_size) + + def unregister(self, logger: CustomLoggerV2) -> None: + kept: Final = tuple(manifest for manifest in self.snapshot.manifests if manifest.logger is not logger) + self.snapshot = RegistrySnapshot(kept, self._plan_cache_size) diff --git a/litellm/integrations/custom_logger_v2/selection.py b/litellm/integrations/custom_logger_v2/selection.py new file mode 100644 index 00000000000..ce80bcf12ee --- /dev/null +++ b/litellm/integrations/custom_logger_v2/selection.py @@ -0,0 +1,36 @@ +from dataclasses import dataclass +from typing import Final + +from litellm.integrations.custom_logger_v2.enums import Activation, OptOutBy + + +@dataclass(frozen=True, slots=True) +class Selection: + name: str + aliases: frozenset[str] = frozenset() + activation: Activation = Activation.default_on + opt_out_by: frozenset[OptOutBy] = frozenset({OptOutBy.admin}) + honors_no_log: bool = False + + +@dataclass(frozen=True, slots=True) +class RequestSelection: + requested: frozenset[str] = frozenset() + admin_opt_outs: frozenset[str] = frozenset() + caller_opt_outs: frozenset[str] = frozenset() + no_log: bool = False + + +def _opted_out(selection: Selection, names: frozenset[str], request: RequestSelection) -> bool: + by_admin: Final = OptOutBy.admin in selection.opt_out_by and not names.isdisjoint(request.admin_opt_outs) + by_caller: Final = OptOutBy.caller in selection.opt_out_by and not names.isdisjoint(request.caller_opt_outs) + return by_admin or by_caller or (request.no_log and selection.honors_no_log) + + +def is_selected(selection: Selection, activation: Activation, request: RequestSelection) -> bool: + if activation is Activation.required: + return True + names: Final = selection.aliases | {selection.name} + if _opted_out(selection, names, request): + return False + return activation is Activation.default_on or not names.isdisjoint(request.requested) diff --git a/litellm/integrations/otel/langfuse_logger.py b/litellm/integrations/otel/langfuse_logger.py index 15b0f3f0cc2..8cb2e75ebca 100644 --- a/litellm/integrations/otel/langfuse_logger.py +++ b/litellm/integrations/otel/langfuse_logger.py @@ -20,14 +20,27 @@ from litellm.integrations.otel.plumbing.context import request_root_span if TYPE_CHECKING: from opentelemetry.trace import Span + from litellm.caching.dual_cache import DualCache from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth - from litellm.types.utils import ModelResponseStream + from litellm.types.utils import CallTypesLiteral, ModelResponseStream class LangfuseOpenTelemetryV2(OpenTelemetryV2): """Stamps the caller's trace controls (name, user, session, tags) on the request. Langfuse reads them off the root observation, and the proxy's root span is still recording when the LLM call starts.""" + async def async_pre_call_hook( + self, + user_api_key_dict: "UserAPIKeyAuth", + cache: "DualCache", + data: dict, + call_type: "CallTypesLiteral", + ) -> dict: + # Through v1.111.0 the proxy's leaf-only override check never ran OpenTelemetryV2's pre-call identity + # seeding for the Langfuse subclasses. The MRO-aware check that replaced it would, so this no-op keeps them + # as they were. + return data + def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None: root: Final = request_root_span() if root is not None and root.is_recording(): diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 35821e856d7..4443a500603 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -82,6 +82,20 @@ from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheCont from litellm.integrations.arize.arize import ArizeLogger from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_logger_v2.enums import Field, Hook +from litellm.integrations.custom_logger_v2.legacy import ( + Admit, + KwargsFields, + SyncLegacyBinding, + hook_subscribers, + kwargs_after, + kwargs_patch, + run_subscribers, + run_subscribers_sync, +) +from litellm.integrations.custom_logger_v2.logger import CustomLoggerV2 +from litellm.integrations.custom_logger_v2.manifest import LegacyBinding +from litellm.integrations.custom_logger_v2.outcomes import Patch from litellm.integrations.deepeval.deepeval import DeepEvalLogger from litellm.integrations.langtrace import langtrace_trace_endpoint from litellm.integrations.mlflow import MlflowLogger @@ -306,6 +320,7 @@ if TYPE_CHECKING: from litellm.llms.base_llm.passthrough.transformation import PassthroughStreamCollector from litellm.proxy.hooks.autorouter_baseline_cache import BaselineCacheContext, CapturedBaselineObservation from litellm.proxy.hooks.search_free_tier import SearchAllowanceAllocation + from litellm.types.llms.base import HiddenParams try: from litellm_enterprise.enterprise_callbacks.callback_controls import ( EnterpriseCallbackControls, @@ -740,6 +755,97 @@ def _callback_control_params(litellm_params: object) -> dict[str, object]: return litellm_params.model_dump() if isinstance(litellm_params, BaseModel) else {} +LOGGING_KEYS: Final[Mapping[Field, str]] = MappingProxyType( + { + Field.model: "model", + Field.messages: "messages", + Field.input: "input", + Field.optional_params: "optional_params", + Field.standard_logging_payload: "standard_logging_object", + } +) + + +def _holds_v2(callbacks: Iterable[object]) -> bool: + return any(isinstance(callback, CustomLoggerV2) for callback in callbacks) + + +# include_v2=False runs only the legacy callbacks in the list. Sync handlers pass it when they also run for an async +# request (the async handler delivers the same event) or for one chunk of a stream (the complete response is +# delivered once at the end), so a V2 callback isn't sent the same event twice or once per chunk. Legacy callbacks +# keep their old behaviour there. +def _run_callbacks_sync( + hook: Hook, callbacks: Sequence[object], kwargs: object, binding: SyncLegacyBinding, include_v2: bool = True +) -> None: + if include_v2 and _holds_v2(callbacks): + _ = run_subscribers_sync( + hook_subscribers(hook, callbacks, admit=Admit.every_entry), KwargsFields(kwargs, LOGGING_KEYS), binding + ) + return + for callback in callbacks: + if not isinstance(callback, CustomLoggerV2): + _ = binding(callback, kwargs) + + +async def _run_callbacks( + hook: Hook, callbacks: Sequence[object], kwargs: object, binding: LegacyBinding, include_v2: bool = True +) -> None: + if include_v2 and _holds_v2(callbacks): + _ = kwargs_after( + await run_subscribers( + hook_subscribers(hook, callbacks, admit=Admit.every_entry), KwargsFields(kwargs, LOGGING_KEYS), binding + ) + ) + return + for callback in callbacks: + if isinstance(callback, CustomLoggerV2): + continue + started = binding(callback, kwargs) + if started is not None: + _ = await started + + +@dataclass(slots=True, eq=False) +class _LogHookRun: + details: dict[str, object] # mutable-ok: legacy logging hooks take and return the live call details + result: object + + +@dataclass(slots=True, eq=False) +class _SyncSuccessRun: + call_details: dict[str, object] # mutable-ok: legacy loggers take the live call details + result: object + kwargs: Mapping[str, object] + litellm_params: object + start_time: datetime.datetime + end_time: datetime.datetime + complete_streaming_response: object + is_sync_request: bool + for_async_request: bool + call_details_view: Mapping[str, object] + + +@dataclass(slots=True, eq=False) +class _McpToolCallRun: + kwargs: dict[str, object] # mutable-ok: forwarded unchanged to legacy hooks that take a dict + response_obj: "CallToolResult" + start_time: datetime.datetime + end_time: datetime.datetime + hidden_params: "HiddenParams" + + +@dataclass(frozen=True, slots=True) +class _FailureRun: + exception: Exception + traceback_exception: str + result: object + start_time: datetime.datetime + end_time: datetime.datetime | None + litellm_params: object + is_sync_request: bool + for_async_request: bool + + class Logging(LiteLLMLoggingBaseClass): global \ supabaseClient, \ @@ -1796,62 +1902,12 @@ class Logging(LiteLLMLoggingBaseClass): if self.model_call_details.get("first_api_call_start_time") is None: self.model_call_details["first_api_call_start_time"] = api_call_start_time # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made - callbacks: Final = ( - *self._without_replaced_global_callbacks(litellm.input_callback), - *(self.dynamic_input_callbacks or ()), + _run_callbacks_sync( + Hook.provider_request_sent, + self._input_callbacks(), + self.model_call_details, + self._log_provider_request_sent, ) - for callback in callbacks: - try: - if callback == "supabase" and supabaseClient is not None: - verbose_logger.debug("reaches supabase for logging!") - model = self._model_call_details_view["model"] - messages = self._model_call_details_view["input"] - verbose_logger.debug("supabaseClient: %s", supabaseClient) - supabaseClient.input_log_event( - model=model, - messages=messages, - end_user=self._model_call_details_view.get("user", "default"), - litellm_call_id=self.litellm_params["litellm_call_id"], - print_verbose=print_verbose, - ) - elif callback == "sentry" and add_breadcrumb: - try: - details_to_log = copy.deepcopy(self.model_call_details) - except Exception: - details_to_log = self.model_call_details - if litellm.turn_off_message_logging: - # make a copy of the _model_Call_details and log it - details_to_log.pop("messages", None) - details_to_log.pop("input", None) - details_to_log.pop("prompt", None) - - add_breadcrumb( - category="litellm.llm_call", - message=f"Model Call Details pre-call: {details_to_log}", - level="info", - ) - - elif isinstance(callback, CustomLogger): # custom logger class - callback.log_pre_api_call( - model=self.model, - messages=self.messages, - kwargs=self.model_call_details, - ) - elif callable(callback) and customLogger is not None: # custom logger functions - customLogger.log_input_event( - model=self.model, - messages=self.messages, - kwargs=self.model_call_details, - print_verbose=print_verbose, - callback_func=callback, - ) - except Exception as e: - verbose_logger.exception("litellm.Logging.pre_call(): Exception occured - %s", e) - verbose_logger.debug( - "LiteLLM.Logging: is sentry capture exception initialized %s", capture_exception - ) - if capture_exception: # log this error to sentry for debugging - capture_exception(e) except Exception as e: verbose_logger.exception("LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e) verbose_logger.error("LiteLLM.Logging: is sentry capture exception initialized %s", capture_exception) @@ -1861,6 +1917,66 @@ class Logging(LiteLLMLoggingBaseClass): if self.raw_request_only: raise RawRequestCaptured() + def _input_callbacks(self) -> tuple[str | Callable[..., object] | CustomLogger, ...]: + return ( + *self._without_replaced_global_callbacks(litellm.input_callback), + *(self.dynamic_input_callbacks or ()), + ) + + def _breadcrumb_details(self) -> dict[str, object]: # mutable-ok: the breadcrumb pops message keys from it + try: + return copy.deepcopy(self.model_call_details) + except Exception: + return self.model_call_details + + def _log_provider_request_sent(self, callback: object, _kwargs: object) -> None: + try: + if callback == "supabase" and supabaseClient is not None: + verbose_logger.debug("reaches supabase for logging!") + model: Final = self._model_call_details_view["model"] + messages: Final = self._model_call_details_view["input"] + verbose_logger.debug("supabaseClient: %s", supabaseClient) + supabaseClient.input_log_event( + model=model, + messages=messages, + end_user=self._model_call_details_view.get("user", "default"), + litellm_call_id=self.litellm_params["litellm_call_id"], + print_verbose=print_verbose, + ) + elif callback == "sentry" and add_breadcrumb: + details_to_log: Final = self._breadcrumb_details() + if litellm.turn_off_message_logging: + # make a copy of the _model_Call_details and log it + details_to_log.pop("messages", None) + details_to_log.pop("input", None) + details_to_log.pop("prompt", None) + + add_breadcrumb( + category="litellm.llm_call", + message=f"Model Call Details pre-call: {details_to_log}", + level="info", + ) + + elif isinstance(callback, CustomLogger): # custom logger class + callback.log_pre_api_call( + model=self.model, + messages=self.messages, + kwargs=self.model_call_details, + ) + elif callable(callback) and customLogger is not None: # custom logger functions + customLogger.log_input_event( + model=self.model, + messages=self.messages, + kwargs=self.model_call_details, + print_verbose=print_verbose, + callback_func=callback, + ) + except Exception as e: + verbose_logger.exception("litellm.Logging.pre_call(): Exception occured - %s", e) + verbose_logger.debug("LiteLLM.Logging: is sentry capture exception initialized %s", capture_exception) + if capture_exception: # log this error to sentry for debugging + capture_exception(e) + def _print_llm_call_debugging_log( self, api_base: str, @@ -1972,49 +2088,47 @@ class Logging(LiteLLMLoggingBaseClass): ) # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made - callbacks: Final = ( - *self._without_replaced_global_callbacks(litellm.input_callback), - *(self.dynamic_input_callbacks or ()), + _run_callbacks_sync( + Hook.provider_response_received, + self._input_callbacks(), + self.model_call_details, + self._log_provider_response_received, ) - for callback in callbacks: - try: - if callback == "sentry" and add_breadcrumb: - verbose_logger.debug("reaches sentry breadcrumbing") - try: - details_to_log = copy.deepcopy(self.model_call_details) - except Exception: - details_to_log = self.model_call_details - if litellm.turn_off_message_logging: - # make a copy of the _model_Call_details and log it - details_to_log.pop("messages", None) - details_to_log.pop("input", None) - details_to_log.pop("prompt", None) - - add_breadcrumb( - category="litellm.llm_call", - message=f"Model Call Details post-call: {details_to_log}", - level="info", - ) - elif isinstance(callback, CustomLogger): # custom logger class - callback.log_post_api_call( - kwargs=self.model_call_details, - response_obj=None, - start_time=self.start_time, - end_time=None, - ) - except Exception as e: - verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while post-call logging with integrations %s", - e, - ) - verbose_logger.debug( - "LiteLLM.Logging: is sentry capture exception initialized %s", capture_exception - ) - if capture_exception: # log this error to sentry for debugging - capture_exception(e) except Exception as e: verbose_logger.exception("LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e) + def _log_provider_response_received(self, callback: object, _kwargs: object) -> None: + try: + if callback == "sentry" and add_breadcrumb: + verbose_logger.debug("reaches sentry breadcrumbing") + details_to_log: Final = self._breadcrumb_details() + if litellm.turn_off_message_logging: + # make a copy of the _model_Call_details and log it + details_to_log.pop("messages", None) + details_to_log.pop("input", None) + details_to_log.pop("prompt", None) + + add_breadcrumb( + category="litellm.llm_call", + message=f"Model Call Details post-call: {details_to_log}", + level="info", + ) + elif isinstance(callback, CustomLogger): # custom logger class + callback.log_post_api_call( + kwargs=self.model_call_details, + response_obj=None, + start_time=self.start_time, + end_time=None, + ) + except Exception as e: + verbose_logger.exception( + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while post-call logging with integrations %s", + e, + ) + verbose_logger.debug("LiteLLM.Logging: is sentry capture exception initialized %s", capture_exception) + if capture_exception: # log this error to sentry for debugging + capture_exception(e) + async def async_post_mcp_tool_call_hook( self, kwargs: dict, @@ -2024,61 +2138,75 @@ class Logging(LiteLLMLoggingBaseClass): ) -> "CallToolResult": """Apply ordered MCP content callbacks to the result returned to the caller.""" from litellm.types.llms.base import HiddenParams - from litellm.types.mcp import MCPPostCallResponseObject + run: Final = _McpToolCallRun( + kwargs=kwargs, # pyright: ignore[reportUnknownArgumentType] # kwargs arrives as a bare dict + response_obj=response_obj, + start_time=start_time, + end_time=end_time, + hidden_params=HiddenParams(), + ) callbacks: Final = self.get_combined_callback_list( dynamic_success_callbacks=self.dynamic_success_callbacks, global_callbacks=litellm.success_callback, ) - hidden_params = HiddenParams() - for callback in callbacks: - try: - if isinstance(callback, CustomLogger): - original_content = copy.deepcopy(response_obj.content) - original_structured_content = copy.deepcopy(response_obj.structured_content) - callback_response = MCPPostCallResponseObject( - mcp_tool_call_response=copy.deepcopy(original_content), hidden_params=hidden_params - ) - try: - response = await callback.async_post_mcp_tool_call_hook( - kwargs=kwargs, - response_obj=callback_response, - start_time=start_time, - end_time=end_time, - ) - hook_content = ( - self._parse_post_mcp_call_hook_response(response=response) - if response is not None - else callback_response.mcp_tool_call_response - ) - if response is not None: - hidden_params = response.hidden_params - except Exception as e: - verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e - ) - hook_content = None - structured_replacement_matches = ( - response_obj.structured_content != original_structured_content - and ( - hook_content is None - or hook_content == original_content - or response_obj.content == hook_content - ) - ) - if hook_content is not None and hook_content != original_content: - response_obj.content[:] = hook_content - if ( - response_obj.content != original_content - and response_obj.structured_content is not None - and not structured_replacement_matches - ): - response_obj.structured_content = None - response_obj.is_error = True - except Exception as e: - verbose_logger.exception("LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e) + await _run_callbacks( + Hook.mcp_tool_result, callbacks, run.kwargs, functools.partial(self._post_mcp_tool_call_start, run) + ) return response_obj + def _post_mcp_tool_call_start( + self, run: "_McpToolCallRun", callback: object, _kwargs: object + ) -> Awaitable[None] | None: + return self._post_mcp_tool_call_one(run, callback) if isinstance(callback, CustomLogger) else None + + async def _post_mcp_tool_call_one(self, run: "_McpToolCallRun", callback: CustomLogger) -> None: + response_obj: Final = run.response_obj + try: + original_content: Final = copy.deepcopy(response_obj.content) + original_structured_content: Final = copy.deepcopy(response_obj.structured_content) + callback_response: Final = MCPPostCallResponseObject( + mcp_tool_call_response=copy.deepcopy(original_content), hidden_params=run.hidden_params + ) + hook_content: Final = await self._post_mcp_tool_call_content(run, callback, callback_response) + structured_replacement_matches: Final = response_obj.structured_content != original_structured_content and ( + hook_content is None or hook_content == original_content or response_obj.content == hook_content + ) + if hook_content is not None and hook_content != original_content: + response_obj.content[:] = hook_content + if ( + response_obj.content != original_content + and response_obj.structured_content is not None + and not structured_replacement_matches + ): + response_obj.structured_content = None + response_obj.is_error = True + except Exception as e: + verbose_logger.exception("LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e) + + async def _post_mcp_tool_call_content( + self, run: "_McpToolCallRun", callback: CustomLogger, callback_response: MCPPostCallResponseObject + ) -> "Sequence[TextContent | ImageContent | EmbeddedResource] | None": + try: + response: Final = await callback.async_post_mcp_tool_call_hook( + kwargs=run.kwargs, + response_obj=callback_response, + start_time=run.start_time, + end_time=run.end_time, + ) + hook_content: Final = ( + self._parse_post_mcp_call_hook_response(response=response) + if response is not None + else callback_response.mcp_tool_call_response + ) + if response is not None: + run.hidden_params = response.hidden_params + except Exception as e: + verbose_logger.exception("LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e) + return None + else: + return hook_content + def _parse_post_mcp_call_hook_response( self, response: MCPPostCallResponseObject | None ) -> "Sequence[TextContent | ImageContent | EmbeddedResource] | None": @@ -3059,7 +3187,7 @@ class Logging(LiteLLMLoggingBaseClass): """Deprecated alias of mark_logging_complete, kept for callers of the old name""" self.mark_logging_complete(event_type=event_type) - def should_run_callback(self, callback: litellm.CALLBACK_TYPES, litellm_params: object, event_hook: str) -> bool: + def should_run_callback(self, callback: object, litellm_params: object, event_hook: str) -> bool: if litellm.global_disable_no_log_param: return True @@ -3867,8 +3995,7 @@ class Logging(LiteLLMLoggingBaseClass): if callback_logging_payload is not None or shared_logging_payload is None else shared_logging_payload.copy() ) - call_details: dict[str, object] # mutable-ok: legacy loggers take a plain dict - call_details = ( # rebind-ok: logging hooks may return a new dict + call_details: Final[dict[str, object]] = ( # mutable-ok: legacy loggers take a plain dict self.model_call_details if not for_async_request else ( @@ -3896,378 +4023,423 @@ class Logging(LiteLLMLoggingBaseClass): ) ## LOGGING HOOK ## hook_callbacks: Final = tuple(dict.fromkeys((*callbacks, *delivery_callbacks))) - for callback in hook_callbacks: - if isinstance(callback, CustomGuardrail): - from litellm.types.guardrails import GuardrailEventHooks + include_v2: Final = not for_async_request and not ( + bool(self.stream) and complete_streaming_response is None + ) + prepared: Final = _LogHookRun(details=call_details, result=result) + _run_callbacks_sync( + Hook.log_prepare, + hook_callbacks, + call_details, + functools.partial(self._sync_logging_hook, prepared, for_async_request), + include_v2=include_v2, + ) - if ( - callback.should_run_guardrail( - data=call_details, - event_type=GuardrailEventHooks.logging_only, - ) - is not True - ): - continue - - self.shadow_eval_request_snapshot = None - call_details, result = callback.logging_hook( # pyright: ignore[reportUnknownVariableType] # logging_hook is typed with a bare dict - kwargs=call_details, - result=result, - call_type=self.call_type, - ) - elif isinstance(callback, CustomLogger): - call_details, result = callback.logging_hook( # pyright: ignore[reportUnknownVariableType] # logging_hook is typed with a bare dict - kwargs=call_details, - result=result, - call_type=self.call_type, - ) - if not for_async_request: - self.model_call_details = call_details - - call_details_view: Final = MappingProxyType[str, object](call_details) self.mark_logging_complete(event_type="sync_success") - for callback in delivery_callbacks: - try: - should_run = self.should_run_callback( - callback=callback, - litellm_params=litellm_params, - event_hook="success_handler", - ) - if not should_run: - continue - if callback == "promptlayer" and promptLayerLogger is not None: - print_verbose("reaches promptlayer for logging!") - promptLayerLogger.log_event( - kwargs=call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - if callback == "supabase" and supabaseClient is not None: - print_verbose("reaches supabase for logging!") - kwargs = call_details - - # this only logs streaming once, complete_streaming_response exists i.e when stream ends - if self.stream: - if "complete_streaming_response" not in kwargs: - continue - else: - print_verbose("reaches supabase for streaming logging!") - result = call_details_view["complete_streaming_response"] - - model = call_details_view["model"] - messages = call_details_view["messages"] - optional_params = call_details_view.get("optional_params", {}) - litellm_params = call_details_view.get("litellm_params", {}) - supabaseClient.log_event( - model=model, - messages=messages, - end_user=require_attributes(optional_params, _ParamsLookup, "get").get("user", "default"), - response_obj=result, - start_time=start_time, - end_time=end_time, - litellm_call_id=( - current_call_id - if ( - current_call_id := require_attributes(litellm_params, _ParamsLookup, "get").get( - "litellm_call_id", None - ) - ) - is not None - else str(uuid.uuid4()) - ), - print_verbose=print_verbose, - ) - if callback == "wandb" and weightsBiasesLogger is not None: - print_verbose("reaches wandb for logging!") - weightsBiasesLogger.log_event( - kwargs=call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - if callback == "logfire" and logfireLogger is not None: - verbose_logger.debug("reaches logfire for success logging!") - kwargs = {} - for k, v in call_details_view.items(): - if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine - kwargs[k] = v - - # this only logs streaming once, complete_streaming_response exists i.e when stream ends - if self.stream: - if "complete_streaming_response" not in kwargs: - continue - else: - print_verbose("reaches logfire for streaming logging!") - result = call_details_view["complete_streaming_response"] - - logfireLogger.log_event( - kwargs=call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - level=LogfireLevel.INFO, - ) - - if callback == "lunary" and lunaryLogger is not None: - print_verbose("reaches lunary for logging!") - model = self.model - kwargs = call_details - - input = kwargs.get("messages", kwargs.get("input", None)) - - type = "embed" if self.call_type == CallTypes.embedding.value else "llm" - - # this only logs streaming once, complete_streaming_response exists i.e when stream ends - if self.stream: - if "complete_streaming_response" not in kwargs: - continue - else: - result = call_details_view["complete_streaming_response"] - - lunaryLogger.log_event( - type=type, - kwargs=kwargs, - event="end", - model=model, - input=input, - user_id=kwargs.get("user", None), - # user_props=self.model_call_details.get("user_props", None), - extra=call_details_view.get("optional_params", {}), - response_obj=result, - start_time=start_time, - end_time=end_time, - run_id=self.litellm_call_id, - print_verbose=print_verbose, - ) - if callback == "helicone" and heliconeLogger is not None: - print_verbose("reaches helicone for logging!") - model = self.model - messages = call_details_view["input"] - kwargs = call_details - - # this only logs streaming once, complete_streaming_response exists i.e when stream ends - if self.stream: - if "complete_streaming_response" not in kwargs: - continue - else: - print_verbose("reaches helicone for streaming logging!") - result = call_details_view["complete_streaming_response"] - - heliconeLogger.log_success( - model=model, - messages=messages, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - kwargs=kwargs, - ) - if callback == "langfuse": - global langFuseLogger - print_verbose("reaches langfuse for success logging!") - kwargs = {} - for k, v in call_details_view.items(): - if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine - kwargs[k] = v - # this only logs streaming once, complete_streaming_response exists i.e when stream ends - if self.stream: - verbose_logger.debug( - "is complete_streaming_response in kwargs: %s", - kwargs.get("complete_streaming_response", None), - ) - if complete_streaming_response is None: - continue - else: - print_verbose("reaches langfuse for streaming logging!") - result = call_details_view["complete_streaming_response"] - - langfuse_logger_to_use = LangFuseHandler.get_langfuse_logger_for_request( - globalLangfuseLogger=langFuseLogger, - standard_callback_dynamic_params=self.standard_callback_dynamic_params, - in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, - ) - if langfuse_logger_to_use is not None: - _response = langfuse_logger_to_use.log_event_on_langfuse( - kwargs=kwargs, - response_obj=result, - start_time=start_time, - end_time=end_time, - user_id=kwargs.get("user", None), - ) - if _response is not None and isinstance(_response, dict): - _trace_id = _response.get("trace_id", None) - if _trace_id is not None: - in_memory_trace_id_cache.set_cache( - litellm_call_id=self.litellm_call_id, - service_name="langfuse", - trace_id=_trace_id, - ) - if callback == "greenscale" and greenscaleLogger is not None: - kwargs = {} - for k, v in call_details_view.items(): - if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine - kwargs[k] = v - # this only logs streaming once, complete_streaming_response exists i.e when stream ends - if self.stream: - verbose_logger.debug( - "is complete_streaming_response in kwargs: %s", - kwargs.get("complete_streaming_response", None), - ) - if complete_streaming_response is None: - continue - else: - print_verbose("reaches greenscale for streaming logging!") - result = call_details_view["complete_streaming_response"] - - greenscaleLogger.log_event( - kwargs=kwargs, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - if callback == "athina" and athinaLogger is not None: - deep_copy = {} - for k, v in call_details_view.items(): - deep_copy[k] = v - athinaLogger.log_event( - kwargs=deep_copy, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - if callback == "traceloop": - deep_copy = {} - for k, v in call_details_view.items(): - if k != "original_response": - deep_copy[k] = v - traceloopLogger.log_event( - kwargs=deep_copy, - response_obj=result, - start_time=start_time, - end_time=end_time, - user_id=kwargs.get("user", None), - print_verbose=print_verbose, - ) - if callback == "s3": - global s3Logger - if s3Logger is None: - s3Logger = S3Logger() - if self.stream: - if "complete_streaming_response" in call_details: - print_verbose("S3Logger Logger: Got Stream Event - Completed Stream Response") - s3Logger.log_event( - kwargs=call_details, - response_obj=call_details_view["complete_streaming_response"], - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - else: - print_verbose("S3Logger Logger: Got Stream Event - No complete stream response as yet") - else: - s3Logger.log_event( - kwargs=call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - - if callback == "openmeter" and is_sync_request: - global openMeterLogger - if openMeterLogger is None: - print_verbose("Instantiates openmeter client") - openMeterLogger = OpenMeterLogger() - if self.stream and complete_streaming_response is None: - openMeterLogger.log_stream_event( - kwargs=call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - ) - else: - if self.stream and complete_streaming_response: - call_details["complete_response"] = call_details.get("complete_streaming_response", {}) - result = call_details_view["complete_response"] - openMeterLogger.log_success_event( - kwargs=call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - ) - if ( - isinstance(callback, CustomLogger) - and (is_sync_request or (for_async_request and not self.stream)) - and not ( - for_async_request - and self.stream - and callback in (self.dynamic_async_success_callbacks or ()) - ) - and self.call_type - != CallTypes.pass_through.value # pass-through endpoints call async_log_success_event - ): # custom logger class - if self.stream and complete_streaming_response is None: - callback.log_stream_event( - kwargs=redacted_model_call_details_for_custom_logger( - model_call_details=call_details, custom_logger=callback - ), - response_obj=redacted_response_obj_for_custom_logger(result, callback), - start_time=start_time, - end_time=end_time, - ) - else: - if self.stream and complete_streaming_response: - call_details["complete_response"] = call_details.get("complete_streaming_response", {}) - result = call_details_view["complete_response"] - - callback.log_success_event( - kwargs=redacted_model_call_details_for_custom_logger( - model_call_details=callback.redact_standard_logging_payload_from_model_call_details( - call_details - ), - custom_logger=callback, - ), - response_obj=redacted_response_obj_for_custom_logger(result, callback), - start_time=start_time, - end_time=end_time, - ) - if callable(callback) is True and is_sync_request: # custom logger functions - function_logger = customLogger if customLogger is not None else CustomLogger() - print_verbose(f"success callbacks: Running Custom Callback Function - {callback}") - - function_logger.log_event( - kwargs=call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - callback_func=callback, - ) - - except Exception as e: - print_verbose( - f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging with integrations {traceback.format_exc()}" - ) - print_verbose(f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}") - if capture_exception: # log this error to sentry for debugging - capture_exception(e) - # Track callback logging failures in Prometheus - try: - self._handle_callback_failure(callback=callback) - except Exception: - pass + run: Final = _SyncSuccessRun( + call_details=prepared.details, + result=prepared.result, + kwargs=kwargs, + litellm_params=litellm_params, + start_time=start_time, + end_time=end_time, + complete_streaming_response=complete_streaming_response, + is_sync_request=is_sync_request, + for_async_request=for_async_request, + call_details_view=MappingProxyType(prepared.details), + ) + _run_callbacks_sync( + Hook.attempt_ended, + delivery_callbacks, + prepared.details, + functools.partial(self._deliver_sync_success, run), + include_v2=include_v2, + ) except Exception as e: verbose_logger.exception( "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging %s", e, ) + def _sync_logging_hook( + self, prepared: _LogHookRun, for_async_request: bool, callback: object, _kwargs: object + ) -> Patch | None: + before: Final = prepared.details + if isinstance(callback, CustomGuardrail): + from litellm.types.guardrails import GuardrailEventHooks + + if ( + callback.should_run_guardrail( + data=before, + event_type=GuardrailEventHooks.logging_only, + ) + is not True + ): + return None + + self.shadow_eval_request_snapshot = None + prepared.details, prepared.result = callback.logging_hook( # pyright: ignore[reportUnknownMemberType] # logging_hook is typed with a bare dict + kwargs=before, + result=prepared.result, + call_type=self.call_type, + ) + elif isinstance(callback, CustomLogger): + prepared.details, prepared.result = callback.logging_hook( # pyright: ignore[reportUnknownMemberType] # logging_hook is typed with a bare dict + kwargs=before, + result=prepared.result, + call_type=self.call_type, + ) + if not for_async_request: + self.model_call_details = prepared.details + return kwargs_patch(before, prepared.details) + + def _deliver_sync_success(self, run: _SyncSuccessRun, callback: object, _kwargs: object) -> None: + call_details: Final = run.call_details + start_time: Final = run.start_time + end_time: Final = run.end_time + complete_streaming_response: Final = run.complete_streaming_response + is_sync_request: Final = run.is_sync_request + for_async_request: Final = run.for_async_request + call_details_view: Final = run.call_details_view + result = run.result # rebind-ok: integrations swap in the assembled stream for later callbacks + kwargs: Mapping[str, object] = run.kwargs # rebind-ok: integrations leave their kwargs to later callbacks + litellm_params = run.litellm_params # rebind-ok: supabase leaves its params to later callbacks + try: + should_run: Final = self.should_run_callback( + callback=callback, + litellm_params=litellm_params, + event_hook="success_handler", + ) + if not should_run: + return + if callback == "promptlayer" and promptLayerLogger is not None: + print_verbose("reaches promptlayer for logging!") + promptLayerLogger.log_event( + kwargs=call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) + if callback == "supabase" and supabaseClient is not None: + print_verbose("reaches supabase for logging!") + kwargs = call_details # rebind-ok: later callbacks see these kwargs + + # this only logs streaming once, complete_streaming_response exists i.e when stream ends + if self.stream: + if "complete_streaming_response" not in kwargs: + return + else: + print_verbose("reaches supabase for streaming logging!") + result = call_details_view["complete_streaming_response"] # rebind-ok: seen by later callbacks + + model = call_details_view["model"] # rebind-ok: each integration branch binds its own + messages = call_details_view["messages"] # rebind-ok: each integration branch binds its own + optional_params: Final = call_details_view.get("optional_params", {}) + litellm_params = call_details_view.get( # rebind-ok: later callbacks see the rebound params + "litellm_params", {} + ) + supabaseClient.log_event( + model=model, + messages=messages, + end_user=require_attributes(optional_params, _ParamsLookup, "get").get("user", "default"), + response_obj=result, + start_time=start_time, + end_time=end_time, + litellm_call_id=( + current_call_id + if ( + current_call_id := require_attributes(litellm_params, _ParamsLookup, "get").get( + "litellm_call_id", None + ) + ) + is not None + else str(uuid.uuid4()) + ), + print_verbose=print_verbose, + ) + if callback == "wandb" and weightsBiasesLogger is not None: + print_verbose("reaches wandb for logging!") + weightsBiasesLogger.log_event( + kwargs=call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) + if callback == "logfire" and logfireLogger is not None: + verbose_logger.debug("reaches logfire for success logging!") + kwargs = {} # rebind-ok: later callbacks see these kwargs + for k, v in call_details_view.items(): + if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine + kwargs[k] = v + + # this only logs streaming once, complete_streaming_response exists i.e when stream ends + if self.stream: + if "complete_streaming_response" not in kwargs: + return + else: + print_verbose("reaches logfire for streaming logging!") + result = call_details_view["complete_streaming_response"] # rebind-ok: seen by later callbacks + + logfireLogger.log_event( + kwargs=call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + level=LogfireLevel.INFO, + ) + + if callback == "lunary" and lunaryLogger is not None: + print_verbose("reaches lunary for logging!") + model = self.model # rebind-ok: each integration branch binds its own + kwargs = call_details # rebind-ok: later callbacks see these kwargs + + input: Final = kwargs.get("messages", kwargs.get("input", None)) + + type: Final = "embed" if self.call_type == CallTypes.embedding.value else "llm" + + # this only logs streaming once, complete_streaming_response exists i.e when stream ends + if self.stream: + if "complete_streaming_response" not in kwargs: + return + else: + result = call_details_view["complete_streaming_response"] # rebind-ok: seen by later callbacks + + lunaryLogger.log_event( + type=type, + kwargs=kwargs, + event="end", + model=model, + input=input, + user_id=kwargs.get("user", None), + # user_props=self.model_call_details.get("user_props", None), + extra=call_details_view.get("optional_params", {}), + response_obj=result, + start_time=start_time, + end_time=end_time, + run_id=self.litellm_call_id, + print_verbose=print_verbose, + ) + if callback == "helicone" and heliconeLogger is not None: + print_verbose("reaches helicone for logging!") + model = self.model # rebind-ok: each integration branch binds its own + messages = call_details_view["input"] # rebind-ok: each integration branch binds its own + kwargs = call_details # rebind-ok: later callbacks see these kwargs + + # this only logs streaming once, complete_streaming_response exists i.e when stream ends + if self.stream: + if "complete_streaming_response" not in kwargs: + return + else: + print_verbose("reaches helicone for streaming logging!") + result = call_details_view["complete_streaming_response"] # rebind-ok: seen by later callbacks + + heliconeLogger.log_success( + model=model, + messages=messages, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + kwargs=kwargs, + ) + if callback == "langfuse": + global langFuseLogger + print_verbose("reaches langfuse for success logging!") + kwargs = {} # rebind-ok: later callbacks see these kwargs + for k, v in call_details_view.items(): + if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine + kwargs[k] = v + # this only logs streaming once, complete_streaming_response exists i.e when stream ends + if self.stream: + verbose_logger.debug( + "is complete_streaming_response in kwargs: %s", + kwargs.get("complete_streaming_response", None), + ) + if complete_streaming_response is None: + return + else: + print_verbose("reaches langfuse for streaming logging!") + result = call_details_view["complete_streaming_response"] # rebind-ok: seen by later callbacks + + langfuse_logger_to_use: Final = LangFuseHandler.get_langfuse_logger_for_request( + globalLangfuseLogger=langFuseLogger, + standard_callback_dynamic_params=self.standard_callback_dynamic_params, + in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, + ) + if langfuse_logger_to_use is not None: + _response: Final = langfuse_logger_to_use.log_event_on_langfuse( + kwargs=kwargs, + response_obj=result, + start_time=start_time, + end_time=end_time, + user_id=kwargs.get("user", None), # pyright: ignore[reportArgumentType] # logger takes the user id as given + ) + if _response is not None and isinstance(_response, dict): + _trace_id: Final = _response.get("trace_id", None) + if _trace_id is not None: + in_memory_trace_id_cache.set_cache( + litellm_call_id=self.litellm_call_id, + service_name="langfuse", + trace_id=_trace_id, + ) + if callback == "greenscale" and greenscaleLogger is not None: + kwargs = {} # rebind-ok: later callbacks see these kwargs + for k, v in call_details_view.items(): + if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine + kwargs[k] = v + # this only logs streaming once, complete_streaming_response exists i.e when stream ends + if self.stream: + verbose_logger.debug( + "is complete_streaming_response in kwargs: %s", + kwargs.get("complete_streaming_response", None), + ) + if complete_streaming_response is None: + return + else: + print_verbose("reaches greenscale for streaming logging!") + result = call_details_view["complete_streaming_response"] # rebind-ok: seen by later callbacks + + greenscaleLogger.log_event( + kwargs=kwargs, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) + if callback == "athina" and athinaLogger is not None: + deep_copy = {} # rebind-ok: each integration branch binds its own + for k, v in call_details_view.items(): + deep_copy[k] = v + athinaLogger.log_event( + kwargs=deep_copy, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) + if callback == "traceloop": + deep_copy = {} # rebind-ok: each integration branch binds its own + for k, v in call_details_view.items(): + if k != "original_response": + deep_copy[k] = v + traceloopLogger.log_event( + kwargs=deep_copy, + response_obj=result, + start_time=start_time, + end_time=end_time, + user_id=kwargs.get("user", None), + print_verbose=print_verbose, + ) + if callback == "s3": + global s3Logger + if s3Logger is None: + s3Logger = S3Logger() # rebind-ok: lazily built module logger + if self.stream: + if "complete_streaming_response" in call_details: + print_verbose("S3Logger Logger: Got Stream Event - Completed Stream Response") + s3Logger.log_event( + kwargs=call_details, + response_obj=call_details_view["complete_streaming_response"], + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) + else: + print_verbose("S3Logger Logger: Got Stream Event - No complete stream response as yet") + else: + s3Logger.log_event( + kwargs=call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) + + if callback == "openmeter" and is_sync_request: + global openMeterLogger + if openMeterLogger is None: + print_verbose("Instantiates openmeter client") + openMeterLogger = OpenMeterLogger() # rebind-ok: lazily built module logger + if self.stream and complete_streaming_response is None: + openMeterLogger.log_stream_event( + kwargs=call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + ) + else: + if self.stream and complete_streaming_response: + call_details["complete_response"] = call_details.get("complete_streaming_response", {}) + result = call_details_view["complete_response"] # rebind-ok: seen by later callbacks + openMeterLogger.log_success_event( + kwargs=call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + ) + if ( + isinstance(callback, CustomLogger) + and (is_sync_request or (for_async_request and not self.stream)) + and not (for_async_request and self.stream and callback in (self.dynamic_async_success_callbacks or ())) + and self.call_type + != CallTypes.pass_through.value # pass-through endpoints call async_log_success_event + ): # custom logger class + if self.stream and complete_streaming_response is None: + callback.log_stream_event( + kwargs=redacted_model_call_details_for_custom_logger( + model_call_details=call_details, custom_logger=callback + ), + response_obj=redacted_response_obj_for_custom_logger(result, callback), + start_time=start_time, + end_time=end_time, + ) + else: + if self.stream and complete_streaming_response: + call_details["complete_response"] = call_details.get("complete_streaming_response", {}) + result = call_details_view["complete_response"] # rebind-ok: seen by later callbacks + + callback.log_success_event( + kwargs=redacted_model_call_details_for_custom_logger( + model_call_details=callback.redact_standard_logging_payload_from_model_call_details( + call_details + ), + custom_logger=callback, + ), + response_obj=redacted_response_obj_for_custom_logger(result, callback), + start_time=start_time, + end_time=end_time, + ) + if callable(callback) is True and is_sync_request: # custom logger functions + function_logger: Final = customLogger if customLogger is not None else CustomLogger() + print_verbose(f"success callbacks: Running Custom Callback Function - {callback}") + + function_logger.log_event( + kwargs=call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + callback_func=callback, + ) + + except Exception as e: + print_verbose( + f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging with integrations {traceback.format_exc()}" + ) + print_verbose(f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}") + if capture_exception: # log this error to sentry for debugging + capture_exception(e) + # Track callback logging failures in Prometheus + try: + self._handle_callback_failure(callback=callback) + except Exception: + pass + finally: + run.result, run.kwargs, run.litellm_params = result, kwargs, litellm_params + async def async_success_handler( self, result: object = None, # heterogeneous response object; varies by call type (ANN401 ignored, see ruff-strict.toml) @@ -4492,183 +4664,26 @@ class Logging(LiteLLMLoggingBaseClass): ## LOGGING HOOK ## hook_callbacks: Final = tuple(dict.fromkeys((*callbacks, *delivery_callbacks))) - for callback in hook_callbacks: - try: - if isinstance(callback, CustomGuardrail): - from litellm.types.guardrails import GuardrailEventHooks - - if ( - callback.should_run_guardrail( - data=self.model_call_details, - event_type=GuardrailEventHooks.logging_only, - ) - is not True - ): - continue - - self.shadow_eval_request_snapshot = None - self.model_call_details, result = await callback.async_logging_hook( - kwargs=self.model_call_details, - result=result, - call_type=self.call_type, - ) - elif isinstance(callback, CustomLogger): - result = redact_message_input_output_from_custom_logger( - result=result, litellm_logging_obj=self, custom_logger=callback - ) - self.model_call_details, result = await callback.async_logging_hook( - kwargs=self.model_call_details, - result=result, - call_type=self.call_type, - ) - except Exception: # noqa: BLE001 # one failing hook must not skip later callbacks (slot release) - verbose_logger.error( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred in async_logging_hook %s", - traceback.format_exc(), - ) - self._handle_callback_failure(callback=callback) + include_v2: Final = self.stream is not True or "async_complete_streaming_response" in self.model_call_details + prepared: Final = _LogHookRun(details=self.model_call_details, result=result) + await _run_callbacks( + Hook.log_prepare, + hook_callbacks, + self.model_call_details, + functools.partial(self._async_logging_hook_start, prepared), + include_v2=include_v2, + ) + result = prepared.result # rebind-ok: logging hooks may replace the result self.mark_logging_complete(event_type="async_success") - for callback in delivery_callbacks: - # check if callback can run for this request - litellm_params = self._model_call_details_view.get("litellm_params", {}) - should_run = self.should_run_callback( - callback=callback, - litellm_params=litellm_params, - event_hook="async_success_handler", - ) - if not should_run: - continue - try: - if callback == "openmeter" and openMeterLogger is not None: - if self.stream is True: - if "async_complete_streaming_response" in self.model_call_details: - await openMeterLogger.async_log_success_event( - kwargs=self.model_call_details, - response_obj=self.model_call_details["async_complete_streaming_response"], - start_time=start_time, - end_time=end_time, - ) - else: - await openMeterLogger.async_log_stream_event( # [TODO]: move this to being an async log stream event function - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - ) - else: - await openMeterLogger.async_log_success_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - ) - - if isinstance(callback, CustomLogger): # custom logger class - from litellm.integrations.shadow_eval_logger import ShadowEvalLogger - - model_call_details: dict = ( - self._model_call_details_with_request_messages() - if callback.exempt_from_message_redaction() - else self.model_call_details - ) - ################################## - # call redaction hook for custom logger - model_call_details = callback.redact_standard_logging_payload_from_model_call_details( - model_call_details=model_call_details - ) - model_call_details = redacted_model_call_details_for_custom_logger( - model_call_details=model_call_details, custom_logger=callback - ) - ################################## - if isinstance(callback, ShadowEvalLogger) and ( - not self.stream or "async_complete_streaming_response" in model_call_details - ): - await callback.async_log_success_event( - kwargs=model_call_details, - response_obj=model_call_details["async_complete_streaming_response"] - if self.stream - else redacted_response_obj_for_custom_logger(result, callback), - start_time=start_time, - end_time=end_time, - guardrail_snapshot=self.shadow_eval_request_snapshot, - ) - elif self.stream is True: - if "async_complete_streaming_response" in model_call_details: - await callback.async_log_success_event( - kwargs=model_call_details, - response_obj=model_call_details["async_complete_streaming_response"], - start_time=start_time, - end_time=end_time, - ) - else: - await callback.async_log_stream_event( # [TODO]: move this to being an async log stream event function - kwargs=model_call_details, - response_obj=redacted_response_obj_for_custom_logger(result, callback), - start_time=start_time, - end_time=end_time, - ) - else: - await callback.async_log_success_event( - kwargs=model_call_details, - response_obj=redacted_response_obj_for_custom_logger(result, callback), - start_time=start_time, - end_time=end_time, - ) - if callable(callback): # custom logger functions - global customLogger - if customLogger is None: - customLogger = CustomLogger() - if self.stream: - if "async_complete_streaming_response" in self.model_call_details: - await customLogger.async_log_event( - kwargs=self.model_call_details, - response_obj=self.model_call_details["async_complete_streaming_response"], - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - callback_func=callback, - ) - else: - await customLogger.async_log_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - callback_func=callback, - ) - if callback == "dynamodb": - global dynamoLogger - if dynamoLogger is None: - dynamoLogger = DyanmoDBLogger() - if self.stream: - if "async_complete_streaming_response" in self.model_call_details: - print_verbose("DynamoDB Logger: Got Stream Event - Completed Stream Response") - await dynamoLogger.async_log_event( - kwargs=self.model_call_details, - response_obj=self.model_call_details["async_complete_streaming_response"], - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - else: - print_verbose("DynamoDB Logger: Got Stream Event - No complete stream response as yet") - else: - await dynamoLogger.async_log_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - except Exception: - verbose_logger.error( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging %s", - traceback.format_exc(), - ) - self._handle_callback_failure(callback=callback) + await _run_callbacks( + Hook.attempt_ended, + delivery_callbacks, + self.model_call_details, + functools.partial(self._deliver_async_success, result, start_time, end_time), + include_v2=include_v2, + ) if not self.stream: try: self.handle_sync_success_callbacks_for_async_calls( @@ -4678,6 +4693,199 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.exception("Could not schedule synchronous success callbacks") await flush_post_call_redis_batches() + def _async_logging_hook_start( + self, prepared: _LogHookRun, callback: object, _kwargs: object + ) -> Awaitable[Patch | None] | None: + return self._async_logging_hook(prepared, callback) if isinstance(callback, CustomLogger) else None + + async def _async_logging_hook(self, prepared: _LogHookRun, callback: CustomLogger) -> Patch | None: + before: Final = self.model_call_details + try: + if isinstance(callback, CustomGuardrail): + from litellm.types.guardrails import GuardrailEventHooks + + if ( + callback.should_run_guardrail( + data=self.model_call_details, + event_type=GuardrailEventHooks.logging_only, + ) + is not True + ): + return None + + self.shadow_eval_request_snapshot = None + self.model_call_details, prepared.result = await callback.async_logging_hook( + kwargs=self.model_call_details, + result=prepared.result, + call_type=self.call_type, + ) + else: + prepared.result = redact_message_input_output_from_custom_logger( + result=prepared.result, litellm_logging_obj=self, custom_logger=callback + ) + self.model_call_details, prepared.result = await callback.async_logging_hook( + kwargs=self.model_call_details, + result=prepared.result, + call_type=self.call_type, + ) + except Exception: # noqa: BLE001 # one failing hook must not skip later callbacks (slot release) + verbose_logger.error( + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred in async_logging_hook %s", + traceback.format_exc(), + ) + self._handle_callback_failure(callback=callback) + prepared.details = self.model_call_details + return kwargs_patch(before, self.model_call_details) + + def _deliver_async_success( + self, + result: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + callback: object, + _kwargs: object, + ) -> Awaitable[None] | None: + litellm_params: Final = self._model_call_details_view.get("litellm_params", {}) + should_run: Final = self.should_run_callback( + callback=callback, + litellm_params=litellm_params, + event_hook="async_success_handler", + ) + return self._deliver_async_success_to(callback, result, start_time, end_time) if should_run else None + + async def _deliver_async_success_to( + self, + callback: object, + result: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + try: + if callback == "openmeter" and openMeterLogger is not None: + if self.stream is True: + if "async_complete_streaming_response" in self.model_call_details: + await openMeterLogger.async_log_success_event( + kwargs=self.model_call_details, + response_obj=self.model_call_details["async_complete_streaming_response"], + start_time=start_time, + end_time=end_time, + ) + else: + await openMeterLogger.async_log_stream_event( # [TODO]: move this to being an async log stream event function + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + ) + else: + await openMeterLogger.async_log_success_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + ) + + if isinstance(callback, CustomLogger): # custom logger class + from litellm.integrations.shadow_eval_logger import ShadowEvalLogger + + model_call_details: Final[dict] = redacted_model_call_details_for_custom_logger( + model_call_details=callback.redact_standard_logging_payload_from_model_call_details( + model_call_details=( + self._model_call_details_with_request_messages() + if callback.exempt_from_message_redaction() + else self.model_call_details + ) + ), + custom_logger=callback, + ) + if isinstance(callback, ShadowEvalLogger) and ( + not self.stream or "async_complete_streaming_response" in model_call_details + ): + await callback.async_log_success_event( + kwargs=model_call_details, + response_obj=model_call_details["async_complete_streaming_response"] + if self.stream + else redacted_response_obj_for_custom_logger(result, callback), + start_time=start_time, + end_time=end_time, + guardrail_snapshot=self.shadow_eval_request_snapshot, + ) + elif self.stream is True: + if "async_complete_streaming_response" in model_call_details: + await callback.async_log_success_event( + kwargs=model_call_details, + response_obj=model_call_details["async_complete_streaming_response"], + start_time=start_time, + end_time=end_time, + ) + else: + await callback.async_log_stream_event( # [TODO]: move this to being an async log stream event function + kwargs=model_call_details, + response_obj=redacted_response_obj_for_custom_logger(result, callback), + start_time=start_time, + end_time=end_time, + ) + else: + await callback.async_log_success_event( + kwargs=model_call_details, + response_obj=redacted_response_obj_for_custom_logger(result, callback), + start_time=start_time, + end_time=end_time, + ) + if callable(callback): # custom logger functions + global customLogger + if customLogger is None: + customLogger = CustomLogger() # rebind-ok: lazily built module logger + if self.stream: + if "async_complete_streaming_response" in self.model_call_details: + await customLogger.async_log_event( + kwargs=self.model_call_details, + response_obj=self.model_call_details["async_complete_streaming_response"], + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + callback_func=callback, + ) + else: + await customLogger.async_log_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + callback_func=callback, + ) + if callback == "dynamodb": + global dynamoLogger + if dynamoLogger is None: + dynamoLogger = DyanmoDBLogger() # rebind-ok: lazily built module logger + if self.stream: + if "async_complete_streaming_response" in self.model_call_details: + print_verbose("DynamoDB Logger: Got Stream Event - Completed Stream Response") + await dynamoLogger.async_log_event( + kwargs=self.model_call_details, + response_obj=self.model_call_details["async_complete_streaming_response"], + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) + else: + print_verbose("DynamoDB Logger: Got Stream Event - No complete stream response as yet") + else: + await dynamoLogger.async_log_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) + except Exception: + verbose_logger.error( + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging %s", + traceback.format_exc(), + ) + self._handle_callback_failure(callback=callback) + def _handle_callback_failure(self, callback: object): """ Handle callback logging failures by incrementing Prometheus metrics. @@ -4900,150 +5108,176 @@ class Logging(LiteLLMLoggingBaseClass): ) ) self.mark_logging_complete(event_type="sync_failure") - for callback in callbacks: - try: - should_run = self.should_run_callback( - callback=callback, - litellm_params=litellm_params, - event_hook="failure_handler", - ) - if not should_run: - continue - if callback == "lunary" and lunaryLogger is not None: - print_verbose("reaches lunary for logging error!") - - model = self.model - - input = self._model_call_details_view["input"] - - _type = "embed" if self.call_type == CallTypes.embedding.value else "llm" - - lunaryLogger.log_event( - kwargs=self.model_call_details, - type=_type, - event="error", - user_id=self._model_call_details_view.get("user", "default"), - model=model, - input=input, - error=traceback_exception, - run_id=self.litellm_call_id, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - ) - if callback == "sentry": - print_verbose("sending exception to sentry") - if capture_exception: - capture_exception(exception) - else: - print_verbose(f"capture exception not initialized: {capture_exception}") - elif callback == "supabase" and supabaseClient is not None: - print_verbose("reaches supabase for logging!") - print_verbose(f"supabaseClient: {supabaseClient}") - supabaseClient.log_event( - model=self.model if hasattr(self, "model") else "", - messages=self.messages, - end_user=self._model_call_details_view.get("user", "default"), - response_obj=result, - start_time=start_time, - end_time=end_time, - litellm_call_id=self._model_call_details_view["litellm_call_id"], - print_verbose=print_verbose, - ) - if callable(callback): # custom logger functions - global customLogger - if customLogger is None: - customLogger = CustomLogger() - customLogger.log_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - callback_func=callback, - ) - if ( - 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, - end_time=end_time, - response_obj=result, - kwargs=self.model_call_details, - ) - if callback == "langfuse": - global langFuseLogger - verbose_logger.debug("reaches langfuse for logging failure") - kwargs = {} - for k, v in self._model_call_details_view.items(): - if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine - kwargs[k] = v - # this only logs streaming once, complete_streaming_response exists i.e when stream ends - langfuse_logger_to_use = LangFuseHandler.get_langfuse_logger_for_request( - globalLangfuseLogger=langFuseLogger, - standard_callback_dynamic_params=self.standard_callback_dynamic_params, - in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, - ) - _response = langfuse_logger_to_use.log_event_on_langfuse( - start_time=start_time, - end_time=end_time, - response_obj=None, - user_id=kwargs.get("user", None), - status_message=str(exception), - level="ERROR", - kwargs=kwargs, - ) - if _response is not None and isinstance(_response, dict): - _trace_id = _response.get("trace_id", None) - if _trace_id is not None: - in_memory_trace_id_cache.set_cache( - litellm_call_id=self.litellm_call_id, - service_name="langfuse", - trace_id=_trace_id, - ) - if callback == "traceloop": - traceloopLogger.log_event( - start_time=start_time, - end_time=end_time, - response_obj=None, - user_id=self.model_call_details.get("user", None), - print_verbose=print_verbose, - status_message=str(exception), - level="ERROR", - kwargs=self.model_call_details, - ) - if callback == "logfire" and logfireLogger is not None: - verbose_logger.debug("reaches logfire for failure logging!") - kwargs = {} - for k, v in self._model_call_details_view.items(): - if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine - kwargs[k] = v - kwargs["exception"] = exception - - logfireLogger.log_event( - kwargs=kwargs, - response_obj=result, - start_time=start_time, - end_time=end_time, - level=LogfireLevel.ERROR, - print_verbose=print_verbose, - ) - - except Exception as e: - print_verbose( - f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging with integrations {e}" - ) - print_verbose(f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}") - if capture_exception: # log this error to sentry for debugging - capture_exception(e) + run: Final = _FailureRun( + exception=exception, + traceback_exception=traceback_exception, + result=result, + start_time=start_time, # pyright: ignore[reportUnknownArgumentType] # untyped failure helper + end_time=end_time, # pyright: ignore[reportUnknownArgumentType] # untyped failure helper + litellm_params=litellm_params, + is_sync_request=is_sync_request, + for_async_request=for_async_request, + ) + _run_callbacks_sync( + Hook.attempt_ended, + callbacks, + self.model_call_details, + functools.partial(self._deliver_sync_failure, run), + include_v2=not for_async_request, + ) except Exception as e: verbose_logger.exception( "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging %s", e ) + def _deliver_sync_failure(self, run: _FailureRun, callback: object, _kwargs: object) -> None: + exception: Final = run.exception + traceback_exception: Final = run.traceback_exception + result: Final = run.result + start_time: Final = run.start_time + end_time: Final = run.end_time + litellm_params: Final = run.litellm_params + is_sync_request: Final = run.is_sync_request + for_async_request: Final = run.for_async_request + try: + should_run: Final = self.should_run_callback( + callback=callback, + litellm_params=litellm_params, + event_hook="failure_handler", + ) + if not should_run: + return + if callback == "lunary" and lunaryLogger is not None: + print_verbose("reaches lunary for logging error!") + + model: Final = self.model + + input: Final = self._model_call_details_view["input"] + + _type: Final = "embed" if self.call_type == CallTypes.embedding.value else "llm" + + lunaryLogger.log_event( + kwargs=self.model_call_details, + type=_type, + event="error", + user_id=self._model_call_details_view.get("user", "default"), + model=model, + input=input, + error=traceback_exception, + run_id=self.litellm_call_id, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) + if callback == "sentry": + print_verbose("sending exception to sentry") + if capture_exception: + capture_exception(exception) + else: + print_verbose(f"capture exception not initialized: {capture_exception}") + elif callback == "supabase" and supabaseClient is not None: + print_verbose("reaches supabase for logging!") + print_verbose(f"supabaseClient: {supabaseClient}") + supabaseClient.log_event( + model=self.model if hasattr(self, "model") else "", + messages=self.messages, + end_user=self._model_call_details_view.get("user", "default"), + response_obj=result, + start_time=start_time, + end_time=end_time, + litellm_call_id=self._model_call_details_view["litellm_call_id"], + print_verbose=print_verbose, + ) + if callable(callback): # custom logger functions + global customLogger + if customLogger is None: + customLogger = CustomLogger() # rebind-ok: lazily built module logger + customLogger.log_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + callback_func=callback, + ) + if ( + 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, + end_time=end_time, + response_obj=result, + kwargs=self.model_call_details, + ) + if callback == "langfuse": + global langFuseLogger + verbose_logger.debug("reaches langfuse for logging failure") + kwargs = {} # rebind-ok: each integration branch binds its own + for k, v in self._model_call_details_view.items(): + if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine + kwargs[k] = v + # this only logs streaming once, complete_streaming_response exists i.e when stream ends + langfuse_logger_to_use: Final = LangFuseHandler.get_langfuse_logger_for_request( + globalLangfuseLogger=langFuseLogger, + standard_callback_dynamic_params=self.standard_callback_dynamic_params, + in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, + ) + _response: Final = langfuse_logger_to_use.log_event_on_langfuse( + start_time=start_time, + end_time=end_time, + response_obj=None, + user_id=kwargs.get("user", None), + status_message=str(exception), + level="ERROR", + kwargs=kwargs, + ) + if _response is not None and isinstance(_response, dict): + _trace_id: Final = _response.get("trace_id", None) + if _trace_id is not None: + in_memory_trace_id_cache.set_cache( + litellm_call_id=self.litellm_call_id, + service_name="langfuse", + trace_id=_trace_id, + ) + if callback == "traceloop": + traceloopLogger.log_event( + start_time=start_time, + end_time=end_time, + response_obj=None, + user_id=self.model_call_details.get("user", None), + print_verbose=print_verbose, + status_message=str(exception), + level="ERROR", + kwargs=self.model_call_details, + ) + if callback == "logfire" and logfireLogger is not None: + verbose_logger.debug("reaches logfire for failure logging!") + kwargs = {} # rebind-ok: each integration branch binds its own + for k, v in self._model_call_details_view.items(): + if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine + kwargs[k] = v + kwargs["exception"] = exception + + logfireLogger.log_event( + kwargs=kwargs, + response_obj=result, + start_time=start_time, + end_time=end_time, + level=LogfireLevel.ERROR, + print_verbose=print_verbose, + ) + + except Exception as e: + print_verbose( + f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging with integrations {e}" + ) + print_verbose(f"LiteLLM.Logging: is sentry capture exception initialized {capture_exception}") + if capture_exception: # log this error to sentry for debugging + capture_exception(e) + async def async_failure_handler( self, exception: Exception, @@ -5095,48 +5329,64 @@ class Logging(LiteLLMLoggingBaseClass): global_callbacks=litellm._async_failure_callback, ) - result: Final = None # result sent to all loggers, init this to None incase it's not created - self.mark_logging_complete(event_type="async_failure") - for callback in callbacks: - try: - litellm_params = self._model_call_details_view.get("litellm_params", {}) - should_run = self.should_run_callback( - callback=callback, - litellm_params=litellm_params, - event_hook="async_failure_handler", - ) - if not should_run: - continue - if isinstance(callback, CustomLogger): # custom logger class - await callback.async_log_failure_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - ) - if callable(callback): # custom logger functions - global customLogger - if customLogger is None: - customLogger = CustomLogger() - await customLogger.async_log_event( - kwargs=self.model_call_details, - response_obj=result, - start_time=start_time, - end_time=end_time, - print_verbose=print_verbose, - callback_func=callback, - ) - except Exception as e: - verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging %s\nCallback=%s", - e, - callback, - ) - # Track callback logging failures in Prometheus - self._handle_callback_failure(callback=callback) + await _run_callbacks( + Hook.attempt_ended, + callbacks, + self.model_call_details, + functools.partial( + self._deliver_async_failure, + start_time, # pyright: ignore[reportUnknownArgumentType] # untyped failure helper + end_time, # pyright: ignore[reportUnknownArgumentType] # untyped failure helper + ), + ) await flush_post_call_redis_batches() + async def _deliver_async_failure( + self, + start_time: datetime.datetime, + end_time: datetime.datetime | None, + callback: object, + _kwargs: object, + ) -> None: + result: Final = None # result sent to all loggers, init this to None incase it's not created + try: + litellm_params: Final = self._model_call_details_view.get("litellm_params", {}) + should_run: Final = self.should_run_callback( + callback=callback, + litellm_params=litellm_params, + event_hook="async_failure_handler", + ) + if not should_run: + return + if isinstance(callback, CustomLogger): # custom logger class + await callback.async_log_failure_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + ) + if callable(callback): # custom logger functions + global customLogger + if customLogger is None: + customLogger = CustomLogger() # rebind-ok: lazily built module logger + await customLogger.async_log_event( + kwargs=self.model_call_details, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + callback_func=callback, + ) + except Exception as e: + verbose_logger.exception( + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging %s\nCallback=%s", + e, + callback, + ) + # Track callback logging failures in Prometheus + self._handle_callback_failure(callback=callback) + def get_trace_id(self, service_name: Literal["langfuse"]) -> str | None: """ For the given service (e.g. langfuse), return the trace_id actually logged. diff --git a/litellm/llms/anthropic/messages/native/handler.py b/litellm/llms/anthropic/messages/native/handler.py index 16b03a711ac..a81ba087f30 100644 --- a/litellm/llms/anthropic/messages/native/handler.py +++ b/litellm/llms/anthropic/messages/native/handler.py @@ -19,6 +19,19 @@ from litellm._logging import verbose_logger from litellm.anthropic_beta_headers_manager import aload_beta_headers_config from litellm.constants import PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_logger_v2.enums import Hook +from litellm.integrations.custom_logger_v2.legacy import ( + PROXY_REQUEST_KEYS, + Admit, + KwargsFields, + hook_subscribers, + kwargs_after, + kwargs_patch, + run_subscribers, +) +from litellm.integrations.custom_logger_v2.logger import CustomLoggerV2 +from litellm.integrations.custom_logger_v2.outcomes import Patch from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params from litellm.litellm_core_utils.exception_mapping_utils import exception_type from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -374,7 +387,7 @@ async def _execute_pre_request_hooks( pass # Build complete request kwargs dict - request_kwargs = { + request_kwargs: dict[str, object] = { # mutable-ok: the hooks take and may replace the request dict "tools": tools, "stream": stream, "litellm_params": { @@ -386,10 +399,21 @@ async def _execute_pre_request_hooks( if not litellm.callbacks: return request_kwargs - from litellm.integrations.custom_logger import CustomLogger as _CustomLogger + callbacks: Final = litellm.callbacks + if ( + Hook.pre_deployment_call in CustomLoggerV2.declared_hooks + and (subscribers := hook_subscribers(Hook.pre_deployment_call, callbacks, admit=Admit.every_logger)).has_v2 + ): + return kwargs_after( # pyright: ignore[reportReturnType] # whatever the last hook returned, as before + await run_subscribers( + subscribers, + KwargsFields(request_kwargs, PROXY_REQUEST_KEYS), + partial(_start_pre_request_hook, model, messages), + ) + ) - for callback in litellm.callbacks: - if not isinstance(callback, _CustomLogger): + for callback in callbacks: + if not isinstance(callback, CustomLogger): continue # Call the pre-request hook @@ -402,6 +426,27 @@ async def _execute_pre_request_hooks( return request_kwargs +def _start_pre_request_hook( + model: str, + messages: list[dict[str, object]], # mutable-ok: the hook takes the messages list + callback: object, + request_kwargs: object, +) -> Awaitable[Patch | None] | None: + if not isinstance(callback, CustomLogger): + return None + return _pre_request_hook(callback, model, messages, request_kwargs) + + +async def _pre_request_hook( + callback: CustomLogger, + model: str, + messages: list[dict[str, object]], # mutable-ok: see above + request_kwargs: object, +) -> Patch | None: + modified_kwargs: Final = await callback.async_pre_request_hook(model, messages, request_kwargs) # pyright: ignore[reportArgumentType, reportUnknownMemberType, reportUnknownVariableType] # the request dict as the previous hook left it + return None if modified_kwargs is None else kwargs_patch(request_kwargs, modified_kwargs) # pyright: ignore[reportUnknownArgumentType] # see above + + async def _try_websearch_short_circuit( model: str, messages: list[dict], diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 8213f9dc779..9c3d428d2f6 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -45,6 +45,9 @@ from litellm.anthropic_beta_headers_manager import ( ) from litellm.constants import MAX_FILE_LIST_LIMIT, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.files.types import FileContentStreamingResult +from litellm.integrations.custom_logger_v2.enums import Hook +from litellm.integrations.custom_logger_v2.legacy import Admit, hook_subscribers +from litellm.integrations.custom_logger_v2.logger import CustomLoggerV2 from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs from litellm.litellm_core_utils.agentic_loop_refusal import ( called_tool_names, @@ -562,6 +565,14 @@ def _run_responses_deployment_hook( "custom_llm_provider": custom_llm_provider, }, CallTypes.responses.value, + include_v2=True, + ) + + +def _has_v2_deployment_hook() -> bool: + return ( + Hook.pre_deployment_call in CustomLoggerV2.declared_hooks + and hook_subscribers(Hook.pre_deployment_call, litellm.callbacks, admit=Admit.every_logger).has_v2 ) @@ -2940,7 +2951,7 @@ class BaseLLMHTTPHandler: dict[str, object], GenericLiteLLMParams, ]: - if not _has_pre_call_deployment_hook(logging_obj): + if not _has_pre_call_deployment_hook(logging_obj) and not _has_v2_deployment_hook(): return ( model, input, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 7fd43ee54a0..9a002fff375 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -18,7 +18,7 @@ from collections.abc import ( ) from dataclasses import dataclass from datetime import datetime, timezone -from functools import lru_cache +from functools import lru_cache, partial from http import HTTPStatus from itertools import chain from types import MappingProxyType @@ -66,6 +66,9 @@ from litellm.constants import ( UNSAFE_PROXY_RESPONSE_HEADERS, ) from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException +from litellm.integrations.custom_logger import subclass_overrides_hook +from litellm.integrations.custom_logger_v2.enums import Field, Hook +from litellm.integrations.custom_logger_v2.legacy import hook_subscribers from litellm.integrations.otel.runtime import phase_event from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure @@ -177,9 +180,13 @@ from litellm.proxy.route_llm_request import ( from litellm.proxy.spend_tracking.cost_finalization import finalize_proxy_response_cost, prepare_response_cost from litellm.proxy.telemetry.request_context import bind_call_id as bind_telemetry_call_id from litellm.proxy.utils import ( # noqa: F401 # legacy module exports + HookRun, ProxyLogging, _check_and_merge_model_level_guardrails, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export check_and_merge_model_level_guardrails, + hook_fields, + routes_through_apply_guardrail, + run_hook_subscribers, ) from litellm.responses.mcp.parked_stream_dispatch import in_flight_round_biller, parked_success from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task @@ -343,6 +350,7 @@ from litellm.proxy.litellm_pre_call_utils import ( from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.utils import ( + LLMResponseTypes, ModelResponse, ModelResponseStream, StandardLoggingPayloadErrorInformation, @@ -4496,11 +4504,9 @@ class ProxyBaseLLMRequestProcessing: is not True ): continue - leaf_attrs = type(cb).__dict__ - holds_until_moderated = "async_post_call_streaming_iterator_hook" in leaf_attrs + holds_until_moderated = subclass_overrides_hook(type(cb), "async_post_call_streaming_iterator_hook") if ( - "apply_guardrail" not in leaf_attrs - or cb.use_native_lifecycle_hooks + not routes_through_apply_guardrail(cb) or cb.mask_response_content or cb.run_in_parallel or ProxyBaseLLMRequestProcessing._releases_held_stream_on_scan(cb, holds_until_moderated) @@ -4946,7 +4952,7 @@ class ProxyBaseLLMRequestProcessing: Extracted as a static method so tests can call the production implementation directly rather than reimplementing the closure. """ - _response = assembled_response + run: Final[HookRun[LLMResponseTypes]] = HookRun(response=assembled_response) try: from litellm.proxy.proxy_server import llm_router as _global_llm_router from litellm.proxy.utils import ( @@ -4958,47 +4964,18 @@ class ProxyBaseLLMRequestProcessing: data=captured_data, llm_router=_global_llm_router ) stream_gated: Final = stream_gated_guardrail_names(captured_data, captured_user_api_key_dict) - for cb in litellm.callbacks: - if not isinstance(cb, CustomGuardrail): - continue - if cb.guardrail_name in stream_gated: - continue - if not cb.should_run_guardrail( - data=guardrail_data, - event_type=GuardrailEventHooks.post_call, - ): - continue - try: - guardrail_result = None - if "apply_guardrail" in type(cb).__dict__ and not cb.use_native_lifecycle_hooks: - # Skip — unified-routed guardrails already ran via - # unified_guardrail's end-of-stream block in the - # streaming iterator pipeline. Running them again - # here would duplicate the guardrail API call - # (e.g. double OpenAI Moderation charges). - continue - if "async_post_call_streaming_iterator_hook" in type(cb).__dict__: - # Skip — the guardrail already scanned the assembled - # response via its own streaming iterator hook in the - # streaming pipeline. re running this function async_post_call_success_hook - # here would duplicate the scan and can spuriously block the guardrail that already passed / failed. - continue - else: - guardrail_result = await cb.async_post_call_success_hook( - user_api_key_dict=captured_user_api_key_dict, - data=guardrail_data, - response=_response, - ) - if guardrail_result is not None: - _response = guardrail_result - except Exception as e: - verbose_proxy_logger.exception( - "Error running post-call guardrail %s on streaming response: %s", - getattr(cb, "guardrail_name", type(cb).__name__), - e, - ) - if isinstance(e, HTTPException) and hasattr(captured_logging_obj, "model_call_details"): - captured_logging_obj.model_call_details.setdefault("metadata", {})["guardrail_blocked"] = True + await run_hook_subscribers( + hook_subscribers(Hook.post_call_success, litellm.callbacks, parallel=False), + hook_fields(guardrail_data, run, Field.response), + partial( + ProxyBaseLLMRequestProcessing._deferred_stream_guardrail, + run, + guardrail_data, + stream_gated, + captured_user_api_key_dict, + captured_logging_obj, + ), + ) except Exception as e: verbose_proxy_logger.exception( "Error in deferred streaming guardrail initialization: %s", @@ -5013,7 +4990,7 @@ class ProxyBaseLLMRequestProcessing: # async markers stored in litellm_params). asyncio.create_task( _as_success_dispatcher(captured_logging_obj).dispatch_success_handlers( - _response, + run.response, cache_hit=cache_hit, start_time=None, end_time=None, @@ -5026,6 +5003,55 @@ class ProxyBaseLLMRequestProcessing: e, ) + @staticmethod + async def _deferred_stream_guardrail( + run: HookRun[LLMResponseTypes], + guardrail_data: dict[str, object], # mutable-ok: the captured request merged with model-level guardrails + stream_gated: frozenset[str], + captured_user_api_key_dict: "UserAPIKeyAuth", + captured_logging_obj: LiteLLMLoggingObj, + cb: object, + request: object, + ) -> None: + if not isinstance(cb, CustomGuardrail): + return + if cb.guardrail_name in stream_gated: + return + if not cb.should_run_guardrail( + data=guardrail_data, + event_type=GuardrailEventHooks.post_call, + ): + return + try: + if routes_through_apply_guardrail(cb): + # Skip — unified-routed guardrails already ran via + # unified_guardrail's end-of-stream block in the + # streaming iterator pipeline. Running them again + # here would duplicate the guardrail API call + # (e.g. double OpenAI Moderation charges). + return + if subclass_overrides_hook(type(cb), "async_post_call_streaming_iterator_hook"): + # Skip — the guardrail already scanned the assembled + # response via its own streaming iterator hook in the + # streaming pipeline. re running this function async_post_call_success_hook + # here would duplicate the scan and can spuriously block the guardrail that already passed / failed. + return + run.adopt( + await cb.async_post_call_success_hook( + user_api_key_dict=captured_user_api_key_dict, + data=guardrail_data, + response=run.response, + ) + ) + except Exception as e: + verbose_proxy_logger.exception( + "Error running post-call guardrail %s on streaming response: %s", + getattr(cb, "guardrail_name", type(cb).__name__), + e, + ) + if isinstance(e, HTTPException) and hasattr(captured_logging_obj, "model_call_details"): + captured_logging_obj.model_call_details.setdefault("metadata", {})["guardrail_blocked"] = True + def _apply_router_cooldown_retry_after(self, headers: dict, e: Exception) -> None: if isinstance(e, RouterRateLimitError) and e.cooldown_time > 0: headers["retry-after"] = str(math.ceil(e.cooldown_time)) diff --git a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py index 7cd36d51511..1effbbd8e0c 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py @@ -28,6 +28,7 @@ from litellm.constants import ( MAX_COMPETITOR_NAMES, ) from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import subclass_overrides_hook from litellm.litellm_core_utils.offline_mode import is_offline_mode from litellm.llms.openai.chat.guardrail_translation.handler import ( OpenAIChatCompletionsHandler, @@ -181,7 +182,7 @@ async def apply_policies( continue if not isinstance(callback, CustomGuardrail): continue - if "apply_guardrail" not in type(callback).__dict__: + if not subclass_overrides_hook(type(callback), "apply_guardrail", CustomGuardrail): verbose_proxy_logger.debug( "apply_policies: guardrail '%s' has no apply_guardrail, skipping", guardrail_name, diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index b9a60d9e37f..5d2477f77b4 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -20,7 +20,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, ) -from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_logger import CustomLogger, subclass_overrides_hook from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, @@ -699,7 +699,10 @@ class PipelineExecutor: @staticmethod def supports_unified_execution(callback: CustomGuardrail) -> bool: """Whether this guardrail runs through the unified apply_guardrail path.""" - return "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks + return ( + subclass_overrides_hook(type(callback), "apply_guardrail", CustomGuardrail) + and not callback.use_native_lifecycle_hooks + ) @staticmethod def supports_streaming_execution(callback: CustomGuardrail) -> bool: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 76a7046beed..084022236c8 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -32,7 +32,6 @@ from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText from email.utils import formatdate from functools import partial, reduce -from itertools import takewhile from operator import attrgetter from types import MappingProxyType from typing import ( @@ -160,7 +159,28 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, merged_guardrail_entries, ) -from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_logger import CustomLogger, subclass_overrides_hook +from litellm.integrations.custom_logger_v2.enums import Field, Hook +from litellm.integrations.custom_logger_v2.legacy import ( + PROXY_REQUEST_KEYS, + Admit, + HookSubscribers, + KwargsFields, + adapter_for, + hook_subscribers, + kwargs_after, + kwargs_patch, + run_subscribers, +) +from litellm.integrations.custom_logger_v2.logger import CustomLoggerV2 +from litellm.integrations.custom_logger_v2.outcomes import ( + MERGEABLE_FIELDS, + InvalidPatchError, + MergeField, + Patch, + PatchOp, +) +from litellm.integrations.custom_logger_v2.patches import LiveFields, apply_patch from litellm.integrations.prometheus import PrometheusLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import add_langfuse_trace_id_to_alert @@ -1914,12 +1934,19 @@ class _CallbackCapabilities: # avoids the per-request ``get_custom_logger_compatible_class`` walk for # every string entry in ``litellm.callbacks``. resolved_callbacks: tuple[object, ...] = field(default_factory=tuple) + hook_callbacks: tuple[object, ...] = field(default_factory=tuple) + has_v2: bool = False listed_models_filters: tuple[CustomLogger, ...] = field(default_factory=tuple) -def _overrides_hook(callback: CustomLogger, hook_name: str) -> bool: - leaf_to_base: Final = takewhile(lambda klass: klass is not CustomLogger, type(callback).__mro__) - return any(hook_name in klass.__dict__ for klass in leaf_to_base) +def _overrides_hook(callback: object, hook_name: str) -> bool: + return subclass_overrides_hook(type(callback), hook_name) + + +def routes_through_apply_guardrail(callback: CustomGuardrail) -> bool: + return subclass_overrides_hook(type(callback), "apply_guardrail", CustomGuardrail) and not getattr( + callback, "use_native_lifecycle_hooks", False + ) def _overrides_moderation_hook(callback: CustomLogger) -> bool: @@ -1937,23 +1964,6 @@ class MalformedListingFilterReturn: tag: Literal["malformed_listing_filter_return"] = "malformed_listing_filter_return" -async def _names_kept_by_listing_callbacks( - callbacks: Sequence[CustomLogger], - user_api_key_dict: UserAPIKeyAuth, - model_names: tuple[str, ...], -) -> tuple[str, ...] | MalformedListingFilterReturn: - if not callbacks or not model_names: - return model_names - returned: Final = await callbacks[0].async_filter_listed_models(user_api_key_dict, model_names) - try: - kept: Final = frozenset(_LISTED_MODEL_NAMES.validate_python(returned)) - except ValidationError: - return MalformedListingFilterReturn(callback=type(callbacks[0]).__name__) - return await _names_kept_by_listing_callbacks( - callbacks[1:], user_api_key_dict, tuple(name for name in model_names if name in kept) - ) - - def _raise_malformed_listing_filter_return(error: MalformedListingFilterReturn) -> NoReturn: raise ProxyException( message=f"{error.callback}.async_filter_listed_models must return a sequence of model names", @@ -1967,6 +1977,174 @@ def _proxy_hook_factory(hook_name: str) -> Callable[..., CustomLogger]: return get_proxy_hook(hook_name) +def _request_dict(value: object) -> dict[str, object] | None: # mutable-ok: the live proxy request dict + return cast("dict[str, object] | None", value) # cast-ok: legacy hooks carry the untyped request dict + + +def _parallel_input( + callback: CustomGuardrail, + data: dict[str, object], # mutable-ok: the live proxy request dict + raw_request_snapshot: dict[str, object] | None, # mutable-ok: same request-payload shape as data +) -> dict[str, object]: # mutable-ok: same request-payload shape as data + if not callback.scan_raw_request or raw_request_snapshot is None: + return data + return independent_snapshot(raw_request_snapshot) + + +@dataclass(slots=True, eq=False) +class _PreCallRun: + user_api_key_dict: UserAPIKeyAuth + call_type: CallTypesLiteral + guardrails_only: bool + skip_guardrails: bool + endpoint_type: EndpointType + raw_request_snapshot: dict[str, object] | None # mutable-ok: same request-payload shape as latest + latest: dict[str, object] | None # mutable-ok: the request as the callbacks have left it so far + pipeline_managed: frozenset[str] = frozenset() + deferred: SensitiveDataRouteException | None = None + parallel: list[tuple[CustomGuardrail, list[object]]] = field( # mutable-ok: parallel results in list order + default_factory=list + ) + + def settle_parallel(self) -> None: + data: Final = self.latest + if data is None: + return + for callback, slot in self.parallel: + if not slot: + continue + result = slot[0] + if callback.scan_raw_request and not isinstance(result, BaseException) and result is not None: + callback.mark_pre_call_hook_ran(data) + if isinstance(result, BaseException) and not isinstance(result, SensitiveDataRouteException): + _record_raising_guardrail(data, callback) + + def raise_parallel_flow_change(self) -> None: + raised: Final = next( + (slot[0] for _, slot in self.parallel if slot and isinstance(slot[0], BaseException)), None + ) + if raised is not None: + raise raised + + +@dataclass(slots=True, eq=False) +class HookRun(Generic[_T]): + response: _T + flow_changes: list[list[BaseException]] = field( # mutable-ok: one slot per parallel guardrail in list order + default_factory=list + ) + transformed: HTTPException | None = None + + def raise_flow_change(self) -> None: + raised: Final = next((slot[0] for slot in self.flow_changes if slot), None) + if raised is not None: + raise raised + + def hold(self, response: _T) -> None: + self.response = response + + def adopt(self, reply: _T | None) -> None: + if reply is not None: + self.response = reply + + def keep_first(self, error: HTTPException) -> None: + if self.transformed is None: + self.transformed = error + + def parallel_slot(self) -> list[BaseException]: # mutable-ok: filled with one guardrail's flow change + slot: Final[list[BaseException]] = [] # mutable-ok: see above + self.flow_changes.append(slot) + return slot + + +@dataclass(slots=True, eq=False) +class _PostCallGuardrails: + data: dict[str, object] # mutable-ok: guardrail_to_apply is written into the live request + guardrail_data: dict[str, object] # mutable-ok: the request merged with model-level guardrails + request_view: Mapping[str, object] + pipeline_managed: frozenset[str] + user_api_key_dict: UserAPIKeyAuth + parallel_view: Mapping[str, object] | None = None + + def parallel_request_view(self) -> Mapping[str, object]: + if self.parallel_view is None: + self.parallel_view = _REQUEST_FIELDS.validate_python(self.data) + return self.parallel_view + + +class HookRunFields(LiveFields, Generic[_T]): + __slots__ = ("_kwargs", "_run", "_slot") + + def __init__(self, kwargs: KwargsFields, run: HookRun[_T], slot: Field) -> None: + self._kwargs: Final = kwargs + self._run: Final = run + self._slot: Final = slot + + def _held(self, field: Field) -> bool: + return field is self._slot or field is Field.deprecated_response_obj + + def __getitem__(self, field: Field) -> object: + return self._run.response if self._held(field) else self._kwargs[field] + + def __iter__(self) -> Iterator[Field]: + return iter((*self._kwargs, Field.deprecated_response_obj, self._slot)) + + def __len__(self) -> int: + return len(self._kwargs) + 2 + + def apply(self, patch: Patch) -> Mapping[Field, object]: + for op in patch.ops: + self._apply_op(op) + return self + + def _apply_op(self, op: PatchOp) -> None: + if not self._held(op.field): + if op.field is Field.deprecated_kwargs: + raise InvalidPatchError("this hook keeps the request it was given") + _ = self._kwargs.apply(Patch((op,))) + return + if isinstance(op, MergeField): + self._merge(op.keys) + return + if self._slot in MERGEABLE_FIELDS: + raise InvalidPatchError(f"{self._slot.value} is add-only here, so merge keys into it") + self._run.hold(op.value) # pyright: ignore[reportArgumentType] # a V2 patch carries an untyped value + + def _merge(self, keys: Mapping[str, object]) -> None: + current: Final = self._run.response + if not isinstance(current, dict): + raise InvalidPatchError(f"{self._slot.value} holds no mapping here, so keys cannot be merged into it") + cast("dict[str, object]", current).update(keys) # cast-ok: merged headers are str-keyed + + +_NO_KWARGS: Final[Mapping[str, object]] = MappingProxyType({}) +_NO_KEYS: Final[Mapping[Field, str]] = MappingProxyType({}) + + +def _resolved_logger(callback: object) -> CustomLogger | None: + resolved: Final = ( + litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( + cast(_custom_logger_compatible_callbacks_literal, callback) # cast-ok: unknown names resolve to None + ) + if isinstance(callback, str) + else callback + ) + return resolved if isinstance(resolved, CustomLogger) else None + + +def hook_fields(kwargs: object, run: HookRun[_T], slot: Field) -> HookRunFields[_T]: + return HookRunFields(KwargsFields(kwargs, PROXY_REQUEST_KEYS), run, slot) + + +async def run_hook_subscribers( + subscribers: HookSubscribers, + fields: Mapping[Field, object], + binding: Callable[[object, object], Awaitable[None] | None], +) -> None: + if subscribers.entries: + _ = kwargs_after(await run_subscribers(subscribers, fields, binding)) + + class ProxyLogging: """ Logging/Custom Handlers for proxy. @@ -2660,9 +2838,7 @@ class ProxyLogging: Result from the guardrail execution """ # Use unified_guardrail if callback has apply_guardrail method - has_apply_guardrail: Final = "apply_guardrail" in type(callback).__dict__ and not getattr( - callback, "use_native_lifecycle_hooks", False - ) + has_apply_guardrail: Final = routes_through_apply_guardrail(callback) use_unified: Final = has_apply_guardrail and not ( hook_type == "during_call" and getattr(callback, "use_native_during_call_hook", False) ) @@ -3318,214 +3494,187 @@ class ProxyLogging: independent_snapshot(data) if needs_raw_request_snapshot else None ) - try: - if not skip_guardrails: - data, _ = await self._maybe_execute_pipelines( # rebind-ok: pipeline edits feed the callback loop below - data=data, - user_api_key_dict=user_api_key_dict, - call_type=call_type, - event_hook="pre_call", - raw_request_snapshot=raw_request_snapshot, - ) - - # Get pipeline-managed guardrails to skip in normal loop - pipeline_managed: Final[frozenset[str]] = ( - frozenset() if skip_guardrails else pipeline_managed_guardrail_names(data, "pre_call") + return await self._dispatch_pre_call( + _PreCallRun( + user_api_key_dict=user_api_key_dict, + call_type=call_type, + guardrails_only=guardrails_only, + skip_guardrails=skip_guardrails, + endpoint_type=endpoint_type, + raw_request_snapshot=raw_request_snapshot, + latest=data, ) + ) + + async def _dispatch_pre_call(self, run: "_PreCallRun") -> dict[str, object] | None: # mutable-ok: request dict + try: + if not run.skip_guardrails and run.latest is not None: + run.latest, _ = await self._maybe_execute_pipelines( + data=run.latest, + user_api_key_dict=run.user_api_key_dict, + call_type=run.call_type, + event_hook="pre_call", + raw_request_snapshot=run.raw_request_snapshot, + ) + if not run.skip_guardrails and run.latest is not None: + run.pipeline_managed = pipeline_managed_guardrail_names(run.latest, "pre_call") caps: Final = ProxyLogging.callback_capabilities() - # Skip the per-request callback walk entirely when nothing in - # ``litellm.callbacks`` overrides ``async_pre_call_hook`` and no - # CustomGuardrail is configured. Saves the loop overhead + - # ``time.time()`` x2 per registered callback for the common - # "callbacks=[]" case on small / dev deployments. if ( - (skip_guardrails or not caps.has_guardrail) + (run.skip_guardrails or not caps.has_guardrail) and not caps.has_content_enforcer - and (guardrails_only or not caps.has_pre_call_override) + and (run.guardrails_only or not caps.has_pre_call_override) + and not caps.has_v2 ): - if data is not None: - self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails) - return data + if run.latest is not None: + self._process_guardrail_metadata(run.latest, skip_guardrails=run.skip_guardrails) + return run.latest - parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = ( - () - if skip_guardrails - else tuple( - cb - for cb in caps.resolved_callbacks - if isinstance(cb, CustomGuardrail) - and getattr(cb, "run_in_parallel", False) - and not (cb.guardrail_name and cb.guardrail_name in pipeline_managed) - ) + subscribers: Final = hook_subscribers( + Hook.pre_call, caps.hook_callbacks, include_v2=not run.guardrails_only ) - - deferred_route_exc: SensitiveDataRouteException | None = None - for _callback in caps.resolved_callbacks: - start_time = time.time() - try: - if isinstance(_callback, CustomGuardrail) and data is not None: - if skip_guardrails: - continue - - # Skip guardrails managed by a pipeline - if _callback.guardrail_name and _callback.guardrail_name in pipeline_managed: - continue - - if getattr(_callback, "run_in_parallel", False): - continue - - data = await self._run_sequential_guardrail_callback( - callback=_callback, - data=data, - raw_request_snapshot=raw_request_snapshot, - user_api_key_dict=user_api_key_dict, - call_type=call_type, - ) - - elif ( - _callback is not None - and isinstance(_callback, CustomLogger) - and (not guardrails_only or _callback.enforces_request_content) - and "async_pre_call_hook" in vars(_callback.__class__) - and _callback.__class__.async_pre_call_hook != CustomLogger.async_pre_call_hook - ): - if call_type in MCP_GUARDRAIL_CALL_TYPES and user_api_key_dict is None: - continue - - response: Exception | str | Mapping[str, object] | None = ( - await _callback.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=self.call_details["user_api_key_cache"], - data=data, - call_type=call_type, - endpoint_type=endpoint_type, - ) - if isinstance(_callback, PROXY_MaxParallelRequestsHandler_v3) - else await _callback.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=self.call_details["user_api_key_cache"], - data=data, - call_type=call_type, - ) - ) - if response is not None: - data = await self.process_pre_call_hook_response( - response=response, data=data, call_type=call_type - ) - except SensitiveDataRouteException as e: - # Defer the reroute until remaining guardrails have run so later - # security checks are not skipped; the first reroute wins and a - # later guardrail that blocks still propagates. Fall through to the - # service-span recording below so the triggering guardrail is still - # timed like every other callback. - if deferred_route_exc is None: - deferred_route_exc = e - - end_time = time.time() - duration = end_time - start_time - if ( - hasattr(self, "service_logging_obj") and duration > 0.01 - ): # only if duration is non-negligible - don't spam the logs - await self.service_logging_obj.async_service_success_hook( - service=ServiceTypes.PROXY_PRE_CALL, - duration=duration, - call_type=f"{_callback.__class__.__name__}", - parent_otel_span=getattr(user_api_key_dict, "parent_otel_span", None), - start_time=start_time, - end_time=end_time, - ) - - if deferred_route_exc is not None and data is not None: - data = await self._handle_sensitive_data_route_exception(deferred_route_exc, data, user_api_key_dict) - - if parallel_guardrails and data is not None: - await self._run_parallel_pre_call_guardrails( - guardrails=parallel_guardrails, - data=data, - raw_request_snapshot=raw_request_snapshot, - user_api_key_dict=user_api_key_dict, - call_type=call_type, + try: + result: Final = await run_subscribers( + subscribers, + KwargsFields(run.latest, PROXY_REQUEST_KEYS), + partial(self._pre_call_start, run), + between=partial(self._pre_call_deferred_reroute, run), ) + finally: + run.settle_parallel() + run.latest = _request_dict(kwargs_after(result)) + run.raise_parallel_flow_change() - if data is not None: - self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails) - - return data + if run.latest is not None: + self._process_guardrail_metadata(run.latest, skip_guardrails=run.skip_guardrails) + return run.latest except SensitiveDataRouteException as e: - data = await self._handle_sensitive_data_route_exception(e, data, user_api_key_dict) - if data is not None: - self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails) - return data + rerouted: Final = await self._handle_sensitive_data_route_exception(e, run.latest, run.user_api_key_dict) + if rerouted is not None: + self._process_guardrail_metadata(rerouted, skip_guardrails=run.skip_guardrails) + return rerouted except Exception: - if data is not None: - self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails) + if run.latest is not None: + self._process_guardrail_metadata(run.latest, skip_guardrails=run.skip_guardrails) raise - async def _run_parallel_pre_call_guardrails( + def _pre_call_start(self, run: _PreCallRun, callback: object, request: object) -> Awaitable[Patch | None] | None: + data: Final = _request_dict(request) + if isinstance(callback, CustomGuardrail) and adapter_for(callback).parallel: + return self._pre_call_start_parallel(run, callback, data) + return self._pre_call_sequential(run, callback, data) + + async def _pre_call_sequential( self, - guardrails: tuple[CustomGuardrail, ...], - data: dict, - raw_request_snapshot: dict | None, # mutable-ok: same request-payload shape as data - user_api_key_dict: UserAPIKeyAuth, - call_type: CallTypesLiteral, - ) -> None: - """ - Run opted-in pre_call guardrails concurrently against one shared payload - snapshot. These guardrails are declared block-only, so any modified data - they return is discarded; they run for their blocking side effect (raising - to reject the request before it reaches the LLM). Every guardrail is - awaited to completion (``return_exceptions=True``) so a raise by one never - leaves the others running as unobserved background tasks. A guardrail that - blocks (any exception other than a reroute or passthrough) takes precedence - over one that only changes the request flow, so a fast reroute can never - let a slower block be bypassed; the request is rejected before it reaches - the LLM, preserving the pre-call barrier that ``during_call`` guardrails - cannot provide. Per-guardrail latency is recorded by - ``_process_guardrail_callback``'s own metrics. - - A guardrail that also opted into ``scan_raw_request`` evaluates - ``raw_request_snapshot`` (taken before the sequential loop ran) instead - of ``data`` (the sequential loop's output), for the same reason the - sequential branch does: its block decision must not depend on what a - sequential guardrail already masked or rewrote. - """ - - def _input_for(callback: CustomGuardrail) -> dict: # mutable-ok: same request-payload shape as data - if not callback.scan_raw_request or raw_request_snapshot is None: - return data - return independent_snapshot(raw_request_snapshot) - - results: Final = await asyncio.gather( - *( - self._process_guardrail_callback( + run: _PreCallRun, + callback: object, + data: dict[str, object] | None, # mutable-ok: the live proxy request dict + ) -> Patch | None: + start_time: Final = time.time() + updated: dict[str, object] | None = data # rebind-ok: a hook may replace the request # mutable-ok: request + try: + if isinstance(callback, CustomGuardrail) and data is not None: + if run.skip_guardrails: + return None + if callback.guardrail_name and callback.guardrail_name in run.pipeline_managed: + return None + updated = await self._run_sequential_guardrail_callback( # rebind-ok: see above callback=callback, - data=_input_for(callback), - user_api_key_dict=user_api_key_dict, - call_type=call_type, - event_type=GuardrailEventHooks.pre_call, + data=data, + raw_request_snapshot=run.raw_request_snapshot, + user_api_key_dict=run.user_api_key_dict, + call_type=run.call_type, ) - for callback in guardrails - ), - return_exceptions=True, - ) - for callback, result in zip(guardrails, results, strict=True): - # _process_guardrail_callback stamped mark_pre_call_hook_ran on - # _input_for's throwaway snapshot copy for a scan_raw_request - # guardrail, never on the live, shared `data` -- without this, a - # deployment-level guardrail sharing this name would see no marker - # via _pre_call_hook_already_ran and re-run it a second time on - # live kwargs. - if callback.scan_raw_request and not isinstance(result, BaseException) and result is not None: - callback.mark_pre_call_hook_ran(data) - if isinstance(result, BaseException) and not isinstance(result, SensitiveDataRouteException): - _record_raising_guardrail(data, callback) - raised: Final = tuple(result for result in results if isinstance(result, BaseException)) - blocking: Final = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None) - if blocking is not None: - raise blocking - if raised: - raise raised[0] + elif ( + isinstance(callback, CustomLogger) + and (not run.guardrails_only or callback.enforces_request_content) + and _overrides_hook(callback, "async_pre_call_hook") + ): + if run.call_type in MCP_GUARDRAIL_CALL_TYPES and run.user_api_key_dict is None: + return None + response: Final[Exception | str | Mapping[str, object] | None] = ( + await callback.async_pre_call_hook( + user_api_key_dict=run.user_api_key_dict, + cache=self.call_details["user_api_key_cache"], + data=data, + call_type=run.call_type, + endpoint_type=run.endpoint_type, + ) + if isinstance(callback, PROXY_MaxParallelRequestsHandler_v3) + else await callback.async_pre_call_hook( + user_api_key_dict=run.user_api_key_dict, + cache=self.call_details["user_api_key_cache"], + data=data, + call_type=run.call_type, + ) + ) + if response is not None: + updated = await self.process_pre_call_hook_response( # rebind-ok: see above + response=response, data=data, call_type=run.call_type + ) + except SensitiveDataRouteException as e: + if run.deferred is None: + run.deferred = e + + end_time: Final = time.time() + duration: Final = end_time - start_time + if hasattr(self, "service_logging_obj") and duration > 0.01: + await self.service_logging_obj.async_service_success_hook( + service=ServiceTypes.PROXY_PRE_CALL, + duration=duration, + call_type=f"{callback.__class__.__name__}", + parent_otel_span=getattr(run.user_api_key_dict, "parent_otel_span", None), + start_time=start_time, + end_time=end_time, + ) + run.latest = updated + return kwargs_patch(data, updated) + + async def _pre_call_deferred_reroute( + self, run: _PreCallRun, fields: Mapping[Field, object] + ) -> Mapping[Field, object]: + data: Final = _request_dict(fields[Field.deprecated_kwargs]) + if run.deferred is None or data is None: + return fields + rerouted: Final = await self._handle_sensitive_data_route_exception(run.deferred, data, run.user_api_key_dict) + run.latest = rerouted + return apply_patch(fields, Patch().set(Field.deprecated_kwargs, rerouted)) + + def _pre_call_start_parallel( + self, + run: _PreCallRun, + callback: CustomGuardrail, + data: dict[str, object] | None, # mutable-ok: the live proxy request dict + ) -> Awaitable[None] | None: + if run.skip_guardrails or data is None: + return None + if callback.guardrail_name and callback.guardrail_name in run.pipeline_managed: + return None + slot: Final[list[object]] = [] # mutable-ok: filled with this guardrail's result or exception + run.parallel.append((callback, slot)) + return self._pre_call_parallel(run, callback, _parallel_input(callback, data, run.raw_request_snapshot), slot) + + async def _pre_call_parallel( + self, + run: _PreCallRun, + callback: CustomGuardrail, + data: dict[str, object], # mutable-ok: the live request or a raw-request snapshot + slot: list[object], # mutable-ok: see _pre_call_start_parallel + ) -> None: + try: + result: Final = await self._process_guardrail_callback( + callback=callback, + data=data, + user_api_key_dict=run.user_api_key_dict, + call_type=run.call_type, + event_type=GuardrailEventHooks.pre_call, + ) + except BaseException as e: + slot.append(e) + if _exception_changes_request_flow(e): + return + raise + slot.append(result) async def _handle_sensitive_data_route_exception( self, @@ -3697,8 +3846,12 @@ class ProxyLogging: has_moderation_override = False iterator_overrides: Final[list[tuple[CustomLogger, str]]] = [] # (callback, kind) resolved_callbacks: Final[list[CustomLogger]] = [] + hook_callbacks: Final[list[object]] = [] # mutable-ok: built once per callback-list signature for callback in callbacks: + if isinstance(callback, CustomLoggerV2): + hook_callbacks.append(callback) + continue if isinstance(callback, str): resolved = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( cast(_custom_logger_compatible_callbacks_literal, callback) @@ -3708,6 +3861,7 @@ class ProxyLogging: if resolved is None or not isinstance(resolved, CustomLogger): continue resolved_callbacks.append(resolved) + hook_callbacks.append(resolved) cls = type(resolved) if cls is CustomLogger: continue @@ -3715,38 +3869,16 @@ class ProxyLogging: has_guardrail = True elif _overrides_moderation_hook(resolved): has_moderation_override = True - # Use the same leaf-class ``__dict__`` check as the other hook - # capabilities: only callbacks that actually override the hook - # contribute to the flag. Setting this for every ``CustomLogger`` - # instance (the prior behaviour) forced the full - # ``post_call_response_headers_hook`` body to run on every request - # even when no registered callback customized response headers. - cls_attrs = cls.__dict__ - if "async_post_call_response_headers_hook" in cls_attrs: + if _overrides_hook(resolved, "async_post_call_response_headers_hook"): has_post_call_response_headers = True - if "async_post_call_streaming_iterator_hook" in cls_attrs: + if _overrides_hook(resolved, "async_post_call_streaming_iterator_hook"): has_iterator_override = True iterator_overrides.append((resolved, "override")) - elif "apply_guardrail" in cls_attrs and not getattr(resolved, "use_native_lifecycle_hooks", False): + elif isinstance(resolved, CustomGuardrail) and routes_through_apply_guardrail(resolved): iterator_overrides.append((resolved, "apply_guardrail")) - # Walk the MRO for ``async_post_call_streaming_hook`` rather than - # using the leaf-class ``__dict__`` check used by the other flags: - # before this PR the hook was unconditionally invoked, so a - # callback that inherits an override from an intermediate parent - # (e.g. a vendor base class providing the override, with the - # registered class adding nothing else) MUST still be detected. - # A leaf-class miss here would silently drop the inherited hook. - base_streaming_hook = CustomLogger.async_post_call_streaming_hook - cls_streaming_hook = getattr( - cls, - "async_post_call_streaming_hook", - base_streaming_hook, - ) - if getattr(cls_streaming_hook, "__func__", cls_streaming_hook) is not getattr( - base_streaming_hook, "__func__", base_streaming_hook - ): + if _overrides_hook(resolved, "async_post_call_streaming_hook"): has_streaming_chunk_override = True - if "async_pre_call_hook" in cls_attrs: + if _overrides_hook(resolved, "async_pre_call_hook"): has_pre_call_override = True if resolved.enforces_request_content is True: has_content_enforcer = True @@ -3762,6 +3894,8 @@ class ProxyLogging: has_moderation_override=has_moderation_override, iterator_overrides=tuple(iterator_overrides), resolved_callbacks=tuple(resolved_callbacks), + hook_callbacks=tuple(hook_callbacks), + has_v2=len(hook_callbacks) > len(resolved_callbacks), listed_models_filters=tuple( callback for callback in resolved_callbacks if _overrides_hook(callback, "async_filter_listed_models") ), @@ -3828,70 +3962,63 @@ class ProxyLogging: guardrails_only: bool = False, ): caps: Final = ProxyLogging.callback_capabilities() - if not caps.has_guardrail and not caps.has_moderation_override: + if not caps.has_guardrail and not caps.has_moderation_override and not caps.has_v2: return data - # Step 1: Collect all guardrail tasks to run in parallel - guardrail_tasks: Final = [] + subscribers: Final = hook_subscribers(Hook.during_call, litellm.callbacks, include_v2=not guardrails_only) - for callback in litellm.callbacks: - if ( - isinstance(callback, CustomLogger) - and not isinstance(callback, CustomGuardrail) - and (not guardrails_only or callback.enforces_request_content) - and _overrides_moderation_hook(callback) - and user_api_key_dict is not None - ): - guardrail_tasks.append( - callback.async_moderation_hook( - data=data, - user_api_key_dict=user_api_key_dict, - call_type=call_type, - ) - ) - elif isinstance(callback, CustomGuardrail): - ################################################################ - # Check if guardrail should be run for GuardrailEventHooks.during_call hook - ################################################################ - - # V1 implementation - backwards compatibility - if callback.event_hook is None and hasattr(callback, "moderation_check"): - if callback.moderation_check == "pre_call": - continue - else: - # Main - V2 Guardrails implementation - from litellm.types.guardrails import GuardrailEventHooks - - event_type = GuardrailEventHooks.during_call - if call_type == CallTypes.call_mcp_tool.value: - event_type = GuardrailEventHooks.during_mcp_call - - if callback.should_run_guardrail(data=data, event_type=event_type) is not True: - continue - # Convert user_api_key_dict to proper format for async_moderation_hook - if call_type == CallTypes.call_mcp_tool.value: - user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict(user_api_key_dict) - else: - user_api_key_auth_dict = user_api_key_dict - guardrail_tasks.append( - self._run_during_call_guardrail( - callback=callback, - data=data, - user_api_key_dict=user_api_key_dict, - user_api_key_auth_dict=user_api_key_auth_dict, - call_type=call_type, - ) - ) - - # Step 2: Run all guardrail tasks in parallel - if guardrail_tasks: - try: - await asyncio.gather(*guardrail_tasks) - except Exception as e: - # If any guardrail raises an exception, it will propagate here - raise e + def start(callback: object, request: object) -> Awaitable[object] | None: + return self._start_during_call(callback, request, user_api_key_dict, call_type, guardrails_only) + _ = kwargs_after(await run_subscribers(subscribers, KwargsFields(data, PROXY_REQUEST_KEYS), start)) return data + def _start_during_call( + self, + callback: object, + request: object, + user_api_key_dict: UserAPIKeyAuth | None, + call_type: CallTypesLiteral, + guardrails_only: bool, + ) -> Awaitable[object] | None: + data: Final = _request_dict(request) + if data is None: + return None + if ( + isinstance(callback, CustomLogger) + and not isinstance(callback, CustomGuardrail) + and (not guardrails_only or callback.enforces_request_content) + and _overrides_moderation_hook(callback) + and user_api_key_dict is not None + ): + return callback.async_moderation_hook(data=data, user_api_key_dict=user_api_key_dict, call_type=call_type) + if not isinstance(callback, CustomGuardrail): + return None + if callback.event_hook is None and hasattr(callback, "moderation_check"): + if callback.moderation_check == "pre_call": + return None + elif ( + callback.should_run_guardrail( + data=data, + event_type=GuardrailEventHooks.during_mcp_call + if call_type == CallTypes.call_mcp_tool.value + else GuardrailEventHooks.during_call, + ) + is not True + ): + return None + user_api_key_auth_dict: Final = ( + self._convert_user_api_key_auth_to_dict(user_api_key_dict) + if call_type == CallTypes.call_mcp_tool.value + else user_api_key_dict + ) + return self._run_during_call_guardrail( + callback=callback, + data=data, + user_api_key_dict=user_api_key_dict, + user_api_key_auth_dict=user_api_key_auth_dict, + call_type=call_type, + ) + async def _run_during_call_guardrail( self, callback: CustomGuardrail, @@ -3901,8 +4028,7 @@ class ProxyLogging: call_type: CallTypesLiteral, ) -> None: if ( - "apply_guardrail" in type(callback).__dict__ - and not callback.use_native_lifecycle_hooks + routes_through_apply_guardrail(callback) and user_api_key_dict is not None and not callback.use_native_during_call_hook ): @@ -4193,42 +4319,49 @@ class ProxyLogging: redacted_traceback_str: Final = redact_string(traceback_str) if traceback_str is not None else None - # Track the first HTTPException returned or raised by any callback - transformed_exception: HTTPException | None = None + callbacks: Final = litellm.callbacks + run: Final[HookRun[Exception]] = HookRun(response=original_exception) + if not ProxyLogging.callback_capabilities().has_v2: + for callback in callbacks: + await self._post_call_failure_one( + run, request_data, user_api_key_dict, redacted_traceback_str, callback, request_data + ) + return run.transformed + await run_hook_subscribers( + hook_subscribers(Hook.post_call_failure, callbacks, admit=Admit.every_entry), + hook_fields(request_data, run, Field.error), + partial(self._post_call_failure_one, run, request_data, user_api_key_dict, redacted_traceback_str), + ) + return run.transformed - for callback in litellm.callbacks: + async def _post_call_failure_one( + self, + run: HookRun[Exception], + request_data: dict[str, object], # mutable-ok: same request-payload shape as post_call_failure_hook's + user_api_key_dict: UserAPIKeyAuth, + redacted_traceback_str: str | None, + callback: object, + request: object, + ) -> None: + try: + resolved: Final = _resolved_logger(callback) + if resolved is None: + return try: - _callback: CustomLogger | None = None - if isinstance(callback, str): - _callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( - cast(_custom_logger_compatible_callbacks_literal, callback) - ) - else: - _callback = callback - if _callback is not None and isinstance(_callback, CustomLogger): - try: - hook_result = await _callback.async_post_call_failure_hook( - request_data=request_data, - user_api_key_dict=user_api_key_dict, - original_exception=original_exception, - traceback_str=redacted_traceback_str, - ) - # If callback returned an HTTPException, use it (first one wins) - if isinstance(hook_result, HTTPException) and transformed_exception is None: - transformed_exception = hook_result - except HTTPException as e: - # If callback raised an HTTPException, use it (first one wins) - if transformed_exception is None: - transformed_exception = e - except Exception as e: - # Log non-HTTPException errors from callbacks but don't break the flow - verbose_proxy_logger.exception( - "[Non-Blocking] Error in async_post_call_failure_hook callback: %s", e - ) + hook_result: Final = await resolved.async_post_call_failure_hook( + request_data=request_data, + user_api_key_dict=user_api_key_dict, + original_exception=run.response, + traceback_str=redacted_traceback_str, + ) + if isinstance(hook_result, HTTPException): + run.keep_first(hook_result) + except HTTPException as e: + run.keep_first(e) except Exception as e: - verbose_proxy_logger.exception("[Non-Blocking] Error setting up post_call_failure_hook callback: %s", e) - - return transformed_exception + verbose_proxy_logger.exception("[Non-Blocking] Error in async_post_call_failure_hook callback: %s", e) + except Exception as e: + verbose_proxy_logger.exception("[Non-Blocking] Error setting up post_call_failure_hook callback: %s", e) def _is_proxy_only_llm_api_error( self, @@ -4469,165 +4602,166 @@ class ProxyLogging: """ from litellm.proxy.proxy_server import llm_router - from litellm.types.guardrails import GuardrailEventHooks + run: Final[HookRun[LLMResponseTypes]] = HookRun(response=response) pipeline_response: Final = await self._run_post_call_pipelines( data=data, user_api_key_dict=user_api_key_dict, response=response, ) - if pipeline_response is not None: - response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below + run.adopt(pipeline_response) pipeline_managed: Final = pipeline_managed_guardrail_names(data, "post_call") request_view: Final = _REQUEST_FIELDS.validate_python(data) guardrail_callbacks, other_callbacks = _partition_post_call_callbacks() - try: - # Merge model-level guardrails before checking which guardrails to run - guardrail_data: Final = check_and_merge_model_level_guardrails(data=data, llm_router=llm_router) - - parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = tuple( - callback - for callback in guardrail_callbacks - if getattr(callback, "run_in_parallel", False) - and not (callback.guardrail_name and callback.guardrail_name in pipeline_managed) - ) - - for callback in guardrail_callbacks: - # Main - V2 Guardrails implementation - - if callback.guardrail_name and callback.guardrail_name in pipeline_managed: - continue - - if getattr(callback, "run_in_parallel", False): - continue - - if ( - callback.should_run_guardrail( - data=guardrail_data, - event_type=GuardrailEventHooks.post_call, - ) - is not True - ): - continue - - if callback.post_call_already_checked(request_view, response) is True: - continue - - guardrail_response: Any | None = None - - if "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks: - data["guardrail_to_apply"] = callback - guardrail_response = await self._run_guardrail_with_metrics( - callback, - unified_guardrail.async_post_call_success_hook( - user_api_key_dict=user_api_key_dict, - data=data, - response=response, - ), - "post_call", - request_data=data, - ) - else: - guardrail_response = await self._run_guardrail_with_metrics( - callback, - callback.async_post_call_success_hook( - user_api_key_dict=user_api_key_dict, - data=data, - response=response, - ), - "post_call", - request_data=data, - ) - - if guardrail_response is not None: - response = guardrail_response - - if parallel_guardrails: - await self._run_parallel_post_call_guardrails( - guardrails=parallel_guardrails, - data=data, - guardrail_data=guardrail_data, - response=response, - user_api_key_dict=user_api_key_dict, - ) - - ############ Handle CustomLogger ############################### - ################################################################# - - for callback in other_callbacks: - callback_response: LLMResponseTypes | None = await callback.async_post_call_success_hook( - user_api_key_dict=user_api_key_dict, data=data, response=response - ) - if callback_response is not None: - response = callback_response - except Exception as e: - raise e - return response - - async def _run_parallel_post_call_guardrails( - self, - guardrails: tuple[CustomGuardrail, ...], - data: dict, - guardrail_data: dict, - response: LLMResponseTypes, - user_api_key_dict: UserAPIKeyAuth, - ) -> None: - """ - Run opted-in post_call guardrails concurrently against the response - produced by the sequential guardrails. These guardrails are declared - block-only, so any modified response they return is discarded; they run - for their blocking side effect (raising to reject the response before it - reaches the client). Every guardrail is awaited to completion - (``return_exceptions=True``) so a raise by one never leaves the others - running as unobserved background tasks. A guardrail that blocks (any - exception other than a passthrough) takes precedence over one that only - changes the response flow, so a fast passthrough can never let a slower - block be bypassed. Each per-guardrail coroutine sets ``guardrail_to_apply`` - immediately before awaiting, and the unified hook pops it before its first - suspension point, so concurrent guardrails never race on that key. - """ - request_view: Final = _REQUEST_FIELDS.validate_python(data) - - async def _run_one(callback: CustomGuardrail) -> None: - if callback.should_run_guardrail(data=guardrail_data, event_type=GuardrailEventHooks.post_call) is not True: - return - if callback.post_call_already_checked(request_view, response) is True: - return - if "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks: - data["guardrail_to_apply"] = callback - await self._run_guardrail_with_metrics( - callback, - unified_guardrail.async_post_call_success_hook( - user_api_key_dict=user_api_key_dict, - data=data, - response=response, - ), - "post_call", - request_data=data, - ) - else: - await self._run_guardrail_with_metrics( - callback, - callback.async_post_call_success_hook( - user_api_key_dict=user_api_key_dict, - data=data, - response=response, - ), - "post_call", - request_data=data, - ) - - results: Final = await asyncio.gather( - *(_run_one(callback) for callback in guardrails), - return_exceptions=True, + guardrails: Final = _PostCallGuardrails( + data=data, + guardrail_data=check_and_merge_model_level_guardrails(data=data, llm_router=llm_router), + request_view=request_view, + pipeline_managed=pipeline_managed, + user_api_key_dict=user_api_key_dict, + ) + if not any(isinstance(callback, CustomLoggerV2) for callback in other_callbacks): + await self._post_call_without_v2(run, guardrails, guardrail_callbacks, other_callbacks) + return run.response + fields: Final = hook_fields(data, run, Field.response) + await run_hook_subscribers( + hook_subscribers(Hook.post_call_success, guardrail_callbacks), + fields, + partial(self._post_call_guardrail_start, run, guardrails), + ) + run.raise_flow_change() + await run_hook_subscribers( + hook_subscribers(Hook.post_call_success, other_callbacks), + fields, + partial(self._post_call_logger, run, data, user_api_key_dict), + ) + return run.response + + async def _post_call_without_v2( + self, + run: HookRun[LLMResponseTypes], + guardrails: "_PostCallGuardrails", + guardrail_callbacks: tuple[CustomGuardrail, ...], + other_callbacks: tuple[CustomLogger, ...], + ) -> None: + flagged: Final = tuple((callback, adapter_for(callback).parallel) for callback in guardrail_callbacks) + for callback, is_parallel in flagged: + if not is_parallel: + await self._post_call_guardrail_begin(run, guardrails, callback, parallel=False) + parallel: Final = tuple(callback for callback, is_parallel in flagged if is_parallel) + if parallel: + started: Final = tuple( + self._post_call_guardrail_begin(run, guardrails, callback, parallel=True) for callback in parallel + ) + results: Final = await asyncio.gather(*started, return_exceptions=True) + raised: Final = next((result for result in results if isinstance(result, BaseException)), None) + if raised is not None: + raise raised + run.raise_flow_change() + for callback in other_callbacks: + if _overrides_hook(callback, "async_post_call_success_hook"): + await self._post_call_logger(run, guardrails.data, guardrails.user_api_key_dict, callback, None) + + def _post_call_guardrail_start( + self, run: HookRun[LLMResponseTypes], guardrails: "_PostCallGuardrails", callback: object, request: object + ) -> Awaitable[None] | None: + if not isinstance(callback, CustomGuardrail): + return None + return self._post_call_guardrail_begin(run, guardrails, callback, parallel=adapter_for(callback).parallel) + + async def _post_call_guardrail_begin( + self, + run: HookRun[LLMResponseTypes], + guardrails: "_PostCallGuardrails", + callback: CustomGuardrail, + parallel: bool, + ) -> None: + if callback.guardrail_name and callback.guardrail_name in guardrails.pipeline_managed: + return + if not parallel: + await self._post_call_sequential_guardrail(run, guardrails, callback) + return + await self._post_call_parallel_guardrail( + run, guardrails, callback, guardrails.parallel_request_view(), run.parallel_slot() + ) + + async def _post_call_sequential_guardrail( + self, run: HookRun[LLMResponseTypes], guardrails: "_PostCallGuardrails", callback: CustomGuardrail + ) -> None: + if ( + callback.should_run_guardrail(data=guardrails.guardrail_data, event_type=GuardrailEventHooks.post_call) + is not True + ): + return + if callback.post_call_already_checked(guardrails.request_view, run.response) is True: + return + run.adopt(await self._post_call_guardrail(guardrails, callback, run.response)) + + async def _post_call_parallel_guardrail( + self, + run: HookRun[LLMResponseTypes], + guardrails: "_PostCallGuardrails", + callback: CustomGuardrail, + request_view: Mapping[str, object], + slot: list[BaseException], # mutable-ok: see _post_call_guardrail_start + ) -> None: + try: + if ( + callback.should_run_guardrail(data=guardrails.guardrail_data, event_type=GuardrailEventHooks.post_call) + is not True + ): + return + if callback.post_call_already_checked(request_view, run.response) is True: + return + _ = await self._post_call_guardrail(guardrails, callback, run.response) + except BaseException as e: + if not _exception_changes_request_flow(e): + raise + slot.append(e) + + async def _post_call_guardrail( + self, guardrails: "_PostCallGuardrails", callback: CustomGuardrail, response: LLMResponseTypes + ) -> LLMResponseTypes | None: + data: Final = guardrails.data + if routes_through_apply_guardrail(callback): + data["guardrail_to_apply"] = callback + return await self._run_guardrail_with_metrics( # pyright: ignore[reportReturnType] # unified hook is untyped + callback, + unified_guardrail.async_post_call_success_hook( + user_api_key_dict=guardrails.user_api_key_dict, + data=data, + response=response, + ), + "post_call", + request_data=data, + ) + return await self._run_guardrail_with_metrics( + callback, + callback.async_post_call_success_hook( + user_api_key_dict=guardrails.user_api_key_dict, + data=data, + response=response, + ), + "post_call", + request_data=data, + ) + + async def _post_call_logger( + self, + run: HookRun[LLMResponseTypes], + data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data + user_api_key_dict: UserAPIKeyAuth, + callback: object, + request: object, + ) -> None: + if not isinstance(callback, CustomLogger): + return + run.adopt( + await callback.async_post_call_success_hook( + user_api_key_dict=user_api_key_dict, data=data, response=run.response + ) ) - raised: Final = tuple(result for result in results if isinstance(result, BaseException)) - blocking: Final = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None) - if blocking is not None: - raise blocking - if raised: - raise raised[0] async def post_mcp_call_hook( self, @@ -4651,7 +4785,7 @@ class ProxyLogging: swallowed into an unguarded result. """ caps: Final = ProxyLogging.callback_capabilities() - if not caps.has_guardrail: + if not caps.has_guardrail and not caps.has_v2: return response handler_cls: Final = load_guardrail_translation_mappings().get(CallTypes.call_mcp_tool) @@ -4659,20 +4793,34 @@ class ProxyLogging: verbose_proxy_logger.debug("MCP guardrail translation handler unavailable; skipping post_mcp_call hook") return response - for callback in caps.resolved_callbacks: - if not isinstance(callback, CustomGuardrail): - continue - if "apply_guardrail" not in type(callback).__dict__ or callback.use_native_lifecycle_hooks: - continue - if ( - callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_mcp_call) - is not True - ): - continue - response = await self._run_guardrail_with_metrics( + run: Final[HookRun[CallToolResult]] = HookRun(response=response) + await run_hook_subscribers( + hook_subscribers(Hook.mcp_tool_result, caps.hook_callbacks), + hook_fields(request_data, run, Field.tool_result), + partial(self._post_mcp_call_one, run, handler_cls, request_data, user_api_key_dict), + ) + return run.response + + async def _post_mcp_call_one( + self, + run: HookRun["CallToolResult"], + handler_cls: type["BaseTranslation"], + request_data: Mapping[str, Any], + user_api_key_dict: UserAPIKeyAuth | None, + callback: object, + request: object, + ) -> None: + if not isinstance(callback, CustomGuardrail): + return + if not routes_through_apply_guardrail(callback): + return + if callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_mcp_call) is not True: + return + run.hold( + await self._run_guardrail_with_metrics( callback, handler_cls().process_output_response( - response=response, + response=run.response, guardrail_to_apply=callback, litellm_logging_obj=request_data.get("litellm_logging_obj"), user_api_key_dict=user_api_key_dict, @@ -4681,7 +4829,7 @@ class ProxyLogging: "post_mcp_call", request_data=request_data, ) - return response + ) async def post_call_response_headers_hook( self, @@ -4697,62 +4845,91 @@ class ProxyLogging: Returns: Dict[str, str]: Merged headers from all callbacks. """ - merged_headers: Final[dict[str, str]] = {} - # Outer call sites in common_request_processing.py already gate this - # call with ``has_post_call_response_headers_callbacks()``. The - # cached detection makes the redundant interior guard cheap, but the - # guard would still iterate every code path through this function so - # keep it cheap and rely on the cached capability lookup. - if not ProxyLogging.callback_capabilities().has_post_call_response_headers: + merged_headers: Final[dict[str, str]] = {} # mutable-ok: callbacks merge their headers into it + caps: Final = ProxyLogging.callback_capabilities() + if not caps.has_post_call_response_headers and not caps.has_v2: return merged_headers try: - # Build litellm_call_info — normalized routing metadata for callbacks litellm_call_info: Final = self._build_litellm_call_info(data=data, response=response) - - for callback in litellm.callbacks: - _callback: CustomLogger | None = None - if isinstance(callback, str): - _callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( - cast(_custom_logger_compatible_callbacks_literal, callback) - ) - else: - _callback = callback - - if _callback is not None and isinstance(_callback, CustomLogger): - if _accepts_litellm_call_info(_callback): - result = await _callback.async_post_call_response_headers_hook( - data=data, - user_api_key_dict=user_api_key_dict, - response=response, - request_headers=request_headers, - litellm_call_info=litellm_call_info, - ) - else: - # Backwards compat: callback doesn't accept litellm_call_info - result = await _callback.async_post_call_response_headers_hook( - data=data, - user_api_key_dict=user_api_key_dict, - response=response, - request_headers=request_headers, - ) - if result is not None: - merged_headers.update(result) + await run_hook_subscribers( + hook_subscribers(Hook.response_headers, litellm.callbacks, admit=Admit.every_entry), + hook_fields(data, HookRun(response=merged_headers), Field.response_headers), + partial( + self._response_headers_one, + merged_headers, + data, + user_api_key_dict, + response, + request_headers, + litellm_call_info, + ), + ) except Exception as e: verbose_proxy_logger.exception("Error in post_call_response_headers_hook: %s", str(e)) return merged_headers + async def _response_headers_one( + self, + merged_headers: dict[str, str], # mutable-ok: the merged headers this hook returns + data: dict[str, object], # mutable-ok: same request-payload shape as post_call_response_headers_hook's + user_api_key_dict: UserAPIKeyAuth, + response: object, + request_headers: dict[str, str] | None, # mutable-ok: same shape as post_call_response_headers_hook's + litellm_call_info: dict[str, object], # mutable-ok: same shape as _build_litellm_call_info's + callback: object, + request: object, + ) -> None: + resolved: Final = _resolved_logger(callback) + if resolved is None or not _overrides_hook(resolved, "async_post_call_response_headers_hook"): + return + result: Final = ( + await resolved.async_post_call_response_headers_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=response, + request_headers=request_headers, + litellm_call_info=litellm_call_info, + ) + if _accepts_litellm_call_info(resolved) + else await resolved.async_post_call_response_headers_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=response, + request_headers=request_headers, + ) + ) + if result is not None: + merged_headers.update(result) + async def hidden_by_listing_callbacks( self, user_api_key_dict: UserAPIKeyAuth, model_names: Sequence[str] ) -> frozenset[str]: - filters: Final = ProxyLogging.callback_capabilities().listed_models_filters - if not filters: + caps: Final = ProxyLogging.callback_capabilities() + if not caps.listed_models_filters and not caps.has_v2: return frozenset() candidates: Final = tuple(model_names) - kept: Final = await _names_kept_by_listing_callbacks(filters, user_api_key_dict, candidates) - if isinstance(kept, MalformedListingFilterReturn): - _raise_malformed_listing_filter_return(kept) - return frozenset(candidates).difference(kept) + run: Final[HookRun[tuple[str, ...]]] = HookRun(response=candidates) + await run_hook_subscribers( + hook_subscribers(Hook.list_models, caps.hook_callbacks), + HookRunFields(KwargsFields(_NO_KWARGS, _NO_KEYS), run, Field.model_list), + partial(self._filter_listed_models, run, user_api_key_dict), + ) + return frozenset(candidates).difference(run.response) + + @staticmethod + async def _filter_listed_models( + run: HookRun[tuple[str, ...]], user_api_key_dict: UserAPIKeyAuth, callback: object, request: object + ) -> None: + names: Final = run.response + if not names or not isinstance(callback, CustomLogger): + return + returned: Final = await callback.async_filter_listed_models(user_api_key_dict, names) + try: + kept: Final = frozenset(_LISTED_MODEL_NAMES.validate_python(returned)) + except ValidationError: + _raise_malformed_listing_filter_return(MalformedListingFilterReturn(callback=type(callback).__name__)) + run.hold(tuple(name for name in names if name in kept)) @staticmethod def _build_litellm_call_info(data: dict, response: object) -> dict[str, object]: diff --git a/litellm/router_rust/steps/callbacks.py b/litellm/router_rust/steps/callbacks.py index e924193db78..5cda8675612 100644 --- a/litellm/router_rust/steps/callbacks.py +++ b/litellm/router_rust/steps/callbacks.py @@ -17,7 +17,8 @@ which outcome callbacks the Python router would still have running for a request router core. """ -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Awaitable, Iterable, Mapping, Sequence +from functools import partial from itertools import chain from typing import TYPE_CHECKING, Final from weakref import WeakSet @@ -27,6 +28,17 @@ from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import PROMPT_CACHE_PIN_CACHE_MAX_SIZE from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_logger_v2.enums import Hook +from litellm.integrations.custom_logger_v2.hook_point_fields import CandidateFields +from litellm.integrations.custom_logger_v2.legacy import ( + Admit, + hook_subscribers, + kwargs_after, + kwargs_patch, + run_subscribers, +) +from litellm.integrations.custom_logger_v2.logger import CustomLoggerV2 +from litellm.integrations.custom_logger_v2.outcomes import Patch from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.routing_common_utils import get_request_team_id from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy @@ -460,7 +472,11 @@ class OptionalChecks: def has_filters(self) -> bool: """Whether ``filter_deployments`` would run any callback: without one it keeps every deployment.""" - return any(filters_deployments(callback) for callback in self.routing_callbacks()) + callbacks: Final = tuple(self.routing_callbacks()) + return any(filters_deployments(callback) for callback in callbacks) or ( + Hook.filter_deployments in CustomLoggerV2.declared_hooks + and hook_subscribers(Hook.filter_deployments, callbacks, admit=Admit.every_logger).has_v2 + ) async def filter_deployments( self, @@ -473,25 +489,25 @@ class OptionalChecks: turn, then the prompt caching check names the deployment it prefers among the rest, which the Rust core takes only if its later filters keep it; a filter's error is raised for the Rust core to decide on.""" - remaining: list[object] = list(deployments) # mutable-ok: filter contract # rebind-ok: narrowed + remaining: object = list(deployments) # rebind-ok: narrowed by each filter request_kwargs: Final = dict(request) callbacks: Final = tuple(self.routing_callbacks()) - for callback in callbacks: - if not isinstance(callback, CustomLogger) or not filters_deployments(callback): - continue - if isinstance(callback, PromptCachingDeploymentCheck): - continue - remaining = await callback.async_filter_deployments( # pyright: ignore[reportUnknownMemberType, reportAssignmentType, reportUnknownVariableType] # see above - model=model, - healthy_deployments=remaining, - messages=messages, # pyright: ignore[reportArgumentType] # the request's messages, as given - request_kwargs=request_kwargs, - parent_otel_span=None, - ) + if ( + Hook.filter_deployments in CustomLoggerV2.declared_hooks + and (subscribers := hook_subscribers(Hook.filter_deployments, callbacks, admit=Admit.every_logger)).has_v2 + ): + filters: Final = partial(_start_filter, model, messages, request_kwargs) + narrowed: Final = await run_subscribers(subscribers, CandidateFields(remaining, model), filters) + remaining = kwargs_after(narrowed) # rebind-ok: narrowed by the hook's subscribers + else: + for callback in callbacks: + started = _filter_call(callback, model, messages, request_kwargs, remaining) + if started is not None: + remaining = await started prompt_caching: Final = next((cb for cb in callbacks if isinstance(cb, PromptCachingDeploymentCheck)), None) if prompt_caching is None: - return remaining, None - return remaining, await prompt_caching.async_preferred_deployment_id( + return remaining, None # pyright: ignore[reportReturnType] # whatever the last filter returned + return remaining, await prompt_caching.async_preferred_deployment_id( # pyright: ignore[reportReturnType] # see above model, remaining, # pyright: ignore[reportArgumentType] # RustCore's deployment dicts messages, # pyright: ignore[reportArgumentType] # the request's messages, as given @@ -580,6 +596,41 @@ def _ids_of(value: object) -> Sequence[object]: return value if isinstance(value, list | tuple) else () # pyright: ignore[reportUnknownVariableType] # the step's id list +def _filter_call( + callback: object, + model: str, + messages: object, + request_kwargs: dict[str, object], # mutable-ok: the filters take the request dict + remaining: object, +) -> Awaitable[object] | None: + if not isinstance(callback, CustomLogger) or not filters_deployments(callback): + return None + if isinstance(callback, PromptCachingDeploymentCheck): + return None + return callback.async_filter_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped deployment dicts + model=model, + healthy_deployments=remaining, # pyright: ignore[reportArgumentType] # whatever the previous filter returned + messages=messages, # pyright: ignore[reportArgumentType] # the request's messages, as given + request_kwargs=request_kwargs, + parent_otel_span=None, + ) + + +def _start_filter( + model: str, + messages: object, + request_kwargs: dict[str, object], # mutable-ok: the filters take the request dict + callback: object, + remaining: object, +) -> Awaitable[Patch | None] | None: + started: Final = _filter_call(callback, model, messages, request_kwargs, remaining) + return None if started is None else _narrowed_by(remaining, started) + + +async def _narrowed_by(remaining: object, started: Awaitable[object]) -> Patch | None: + return kwargs_patch(remaining, await started) + + def filters_deployments(callback: object) -> bool: """A routing callback that overrides ``async_filter_deployments``: the base filter returns its input unchanged, so only overriding filters are run.""" diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 88687269d4b..2f8cddf7740 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -1,11 +1,12 @@ import hashlib import json -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass, field from datetime import datetime from enum import Enum +from functools import partial, reduce from types import MappingProxyType -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, cast # noqa: TID251 # the router passes the request kwargs as an untyped dict import httpx from pydantic import TypeAdapter @@ -13,6 +14,9 @@ from pydantic import TypeAdapter import litellm from litellm._logging import verbose_router_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_logger_v2.enums import Field, Hook +from litellm.integrations.custom_logger_v2.legacy import Admit, hook_subscribers, kwargs_after, run_subscribers +from litellm.integrations.custom_logger_v2.logger import CustomLoggerV2 from litellm.litellm_core_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, get_fallback_error_info, @@ -975,60 +979,99 @@ async def run_async_fallback( raise error_from_fallbacks -async def log_success_fallback_event(original_model_group: str, kwargs: dict, original_exception: Exception): - """ - Log a successful fallback event to all registered callbacks. +async def log_success_fallback_event( + original_model_group: str, + kwargs: dict, # mutable-ok: the request kwargs the custom loggers receive + original_exception: Exception, +): + await _log_fallback_event( + True, + original_model_group, + cast( # cast-ok: the request kwargs dict, keyed by parameter name + dict[str, object], kwargs + ), + original_exception, + ) - Uses LoggingCallbackManager.get_custom_loggers_for_type() to get deduplicated - CustomLogger instances from all callback lists. - Args: - original_model_group (str): The original model group before fallback. - kwargs (dict): kwargs for the request +async def log_failure_fallback_event( + original_model_group: str, + kwargs: dict, # mutable-ok: the request kwargs the custom loggers receive + original_exception: Exception, +): + await _log_fallback_event( + False, + original_model_group, + cast( # cast-ok: the request kwargs dict, keyed by parameter name + dict[str, object], kwargs + ), + original_exception, + ) - Note: - Errors during logging are caught and reported but do not interrupt the process. - """ - # Get deduplicated CustomLogger instances from all callback lists - custom_loggers: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(CustomLogger) - for _callback_custom_logger in custom_loggers: - try: - await _callback_custom_logger.log_success_fallback_event( - original_model_group=original_model_group, - kwargs=kwargs, - original_exception=original_exception, +async def _log_fallback_event( + succeeded: bool, + original_model_group: str, + kwargs: dict[str, object], # mutable-ok: the request kwargs the custom loggers receive + original_exception: Exception, +) -> None: + callbacks: Final = reduce(_with_event_logger, litellm.logging_callback_manager.get_all_callbacks(), ()) + fields: Final = MappingProxyType( + { + Field.deprecated_kwargs: kwargs, + Field.model: original_model_group, + Field.error: original_exception, + Field.outcome: "fallback_success" if succeeded else "fallback_failure", + } + ) + subscribers: Final = hook_subscribers(Hook.router_event, callbacks, admit=Admit.every_logger) + _ = kwargs_after( + await run_subscribers( + subscribers, fields, partial(_start_fallback_event, succeeded, original_model_group, original_exception) + ) + ) + + +def _with_event_logger(kept: tuple[object, ...], callback: object) -> tuple[object, ...]: + if not isinstance(callback, CustomLogger | CustomLoggerV2) or callback in kept: + return kept + return (*kept, callback) + + +def _start_fallback_event( + succeeded: bool, original_model_group: str, original_exception: Exception, callback: object, kwargs: object +) -> Awaitable[None] | None: + if not isinstance(callback, CustomLogger): + return None + return _deliver_fallback_event( + callback, + succeeded, + original_model_group, + kwargs, # pyright: ignore[reportArgumentType] # the request kwargs as the router passed them + original_exception, + ) + + +async def _deliver_fallback_event( + callback: CustomLogger, + succeeded: bool, + original_model_group: str, + kwargs: dict[str, object], # mutable-ok: the request kwargs the custom loggers receive + original_exception: Exception, +) -> None: + try: + if succeeded: + await callback.log_success_fallback_event( + original_model_group=original_model_group, kwargs=kwargs, original_exception=original_exception ) - except Exception as e: - verbose_router_logger.error("Error in log_success_fallback_event: %s", e) - - -async def log_failure_fallback_event(original_model_group: str, kwargs: dict, original_exception: Exception): - """ - Log a failed fallback event to all registered callbacks. - - Uses LoggingCallbackManager.get_custom_loggers_for_type() to get deduplicated - CustomLogger instances from all callback lists. - - Args: - original_model_group (str): The original model group before fallback. - kwargs (dict): kwargs for the request - - Note: - Errors during logging are caught and reported but do not interrupt the process. - """ - # Get deduplicated CustomLogger instances from all callback lists - custom_loggers: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(CustomLogger) - - for _callback_custom_logger in custom_loggers: - try: - await _callback_custom_logger.log_failure_fallback_event( - original_model_group=original_model_group, - kwargs=kwargs, - original_exception=original_exception, + else: + await callback.log_failure_fallback_event( + original_model_group=original_model_group, kwargs=kwargs, original_exception=original_exception ) - except Exception as e: - verbose_router_logger.error("Error in log_failure_fallback_event: %s", e) + except Exception as e: + verbose_router_logger.error( + "Error in log_success_fallback_event: %s" if succeeded else "Error in log_failure_fallback_event: %s", e + ) def check_non_standard_fallback_format(fallbacks: Sequence[object] | None) -> bool: diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index 44cac1a0eb2..7147e99ac27 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -229,15 +229,19 @@ class StreamingLogBuilder(Protocol): class DeploymentHook(Protocol): - def __call__(self, kwargs: dict[str, object], call_type: str) -> Awaitable[object]: ... + def __call__(self, kwargs: dict[str, object], call_type: str, *, include_v2: bool) -> Awaitable[object]: ... class DeploymentSuccessHook(Protocol): - def __call__(self, request_data: dict[str, object], response: object, call_type: object) -> Awaitable[object]: ... + def __call__( + self, request_data: dict[str, object], response: object, call_type: object, *, include_v2: bool + ) -> Awaitable[object]: ... class DeploymentFailureHook(Protocol): - def __call__(self, request_data: Mapping[str, object], exception: Exception, call_type: str) -> Awaitable[None]: ... + def __call__( + self, request_data: Mapping[str, object], exception: Exception, call_type: str, *, include_v2: bool + ) -> Awaitable[None]: ... def update_logging( @@ -517,7 +521,7 @@ def before_deployment_call(kwargs: dict[str, object], call_type: str) -> Awaitab hook: Final = cast( # cast-ok: bounded adapter for the untyped deployment hook DeploymentHook, utils.async_pre_call_deployment_hook ) - return hook(kwargs, call_type) + return hook(kwargs, call_type, include_v2=True) def after_deployment_success(kwargs: dict[str, object], response: object, call_type: str) -> Awaitable[object]: @@ -528,7 +532,12 @@ def after_deployment_success(kwargs: dict[str, object], response: object, call_t DeploymentSuccessHook, utils.async_post_call_success_deployment_hook ) typed: Final = CallTypes(call_type) - return hook(kwargs, _messages_stream(response) if typed == CallTypes.anthropic_messages else response, typed) + return hook( + kwargs, + _messages_stream(response) if typed == CallTypes.anthropic_messages else response, + typed, + include_v2=True, + ) def _messages_stream(response: object) -> object: @@ -558,7 +567,7 @@ def after_deployment_failure(kwargs: dict[str, object], error: Exception, call_t hook: Final = cast( # cast-ok: bounded adapter for the untyped deployment hook DeploymentFailureHook, utils.async_post_call_failure_deployment_hook ) - return hook(kwargs, error, call_type) + return hook(kwargs, error, call_type, include_v2=True) _STREAM_HEADERS: Final = TypeAdapter(Mapping[str, str]) diff --git a/litellm/utils.py b/litellm/utils.py index b7610bcfaae..7d495ee60a5 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -83,6 +83,7 @@ from litellm.constants import ( PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO, TOOL_CHOICE_OBJECT_TOKEN_COUNT, ) +from litellm.integrations.custom_logger_v2.enums import Field, Hook from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure from litellm.litellm_core_utils.allowed_tools_tool_choice import ( HOSTED_TOOL_CHOICE_TYPES, @@ -134,21 +135,6 @@ from litellm.rust_bridge.configuration import Decision _CachingHandlerResponse = None _LLMCachingHandler = None _CustomGuardrail = None -_CustomLogger = None - - -def _get_cached_custom_logger(): - """ - Get cached CustomLogger class. - Lazy imports on first call to avoid loading custom_logger at import time. - Subsequent calls use cached class for better performance. - """ - global _CustomLogger - if _CustomLogger is None: - from litellm.integrations.custom_logger import CustomLogger - - _CustomLogger = CustomLogger - return _CustomLogger @lru_cache(maxsize=None) @@ -335,7 +321,7 @@ except (ImportError, AttributeError, TypeError): # Convert to str (if necessary) claude_json_str = json.dumps(json_data) import importlib.metadata -from collections.abc import AsyncIterator, Callable, Collection, Iterable, Iterator, Mapping, Sequence, Sized +from collections.abc import AsyncIterator, Awaitable, Callable, Collection, Iterable, Iterator, Mapping, Sequence, Sized from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast, runtime_checkable from typing_extensions import assert_never @@ -551,6 +537,21 @@ from .exceptions import ( if TYPE_CHECKING: from litellm import MockException +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_logger_v2.hook_point_fields import ResultFields +from litellm.integrations.custom_logger_v2.legacy import ( + PROXY_REQUEST_KEYS, + Admit, + KwargsFields, + hook_subscribers, + kwargs_after, + kwargs_patch, + run_subscribers, +) +from litellm.integrations.custom_logger_v2.logger import CustomLoggerV2 +from litellm.integrations.custom_logger_v2.outcomes import Patch + ####### ENVIRONMENT VARIABLES #################### # Adjust to your specific application needs / system capabilities. sentry_sdk_instance: Final = None @@ -1842,7 +1843,9 @@ def _deployment_provider(kwargs: Mapping[str, object]) -> str: return "" -async def async_pre_call_deployment_hook(kwargs: Mapping[str, object], call_type: str) -> dict[str, Any]: +async def async_pre_call_deployment_hook( + kwargs: Mapping[str, object], call_type: str, *, include_v2: bool = False +) -> dict[str, Any]: # mutable-ok: callers merge the returned kwargs """ Allow modifying the request just before it's sent to the deployment. @@ -1863,15 +1866,40 @@ async def async_pre_call_deployment_hook(kwargs: Mapping[str, object], call_type body={"type": "invalid_request_error", "param": resolved.field}, ) modified_kwargs = dict(resolved) + callbacks: Final = litellm.callbacks + if ( + include_v2 + and Hook.pre_deployment_call in CustomLoggerV2.declared_hooks + and (subscribers := hook_subscribers(Hook.pre_deployment_call, callbacks, admit=Admit.every_logger)).has_v2 + ): + modified_kwargs = kwargs_after( # rebind-ok: whatever the last hook returned + await run_subscribers( + subscribers, + KwargsFields(modified_kwargs, PROXY_REQUEST_KEYS), + partial(_start_pre_deployment, typed_call_type), + ) + ) + else: + for callback in callbacks: + if isinstance(callback, CustomLogger): + result = await callback.async_pre_call_deployment_hook(modified_kwargs, typed_call_type) + if result is not None: + modified_kwargs = result - custom_logger_class: Final = _get_cached_custom_logger() - for callback in litellm.callbacks: - if isinstance(callback, custom_logger_class): - result = await callback.async_pre_call_deployment_hook(modified_kwargs, typed_call_type) - if result is not None: - modified_kwargs = result + return _unstreamed_when_post_call_guardrails_apply(modified_kwargs, typed_call_type) # pyright: ignore[reportArgumentType] # whatever the last hook returned, as before - return _unstreamed_when_post_call_guardrails_apply(modified_kwargs, typed_call_type) + +def _start_pre_deployment( + call_type: CallTypes | None, callback: object, kwargs: object +) -> Awaitable[Patch | None] | None: + if not isinstance(callback, CustomLogger): + return None + return _pre_deployment(callback, kwargs, call_type) + + +async def _pre_deployment(callback: CustomLogger, kwargs: object, call_type: CallTypes | None) -> Patch | None: + result: Final = await callback.async_pre_call_deployment_hook(kwargs, call_type) # pyright: ignore[reportArgumentType, reportUnknownMemberType, reportUnknownVariableType] # the request dict as the previous hook left it + return None if result is None else kwargs_patch(kwargs, result) # pyright: ignore[reportUnknownArgumentType] # see above def _unstreamed_when_post_call_guardrails_apply( @@ -1896,7 +1924,11 @@ def _unstreamed_when_post_call_guardrails_apply( async def async_post_call_success_deployment_hook( - request_data: dict, response: object, call_type: CallTypes | None + request_data: dict, # mutable-ok: the attempt's request dict the hooks receive + response: object, + call_type: CallTypes | None, + *, + include_v2: bool = False, ) -> object: """ Allow modifying / reviewing the response just after it's received from the deployment. @@ -1906,33 +1938,65 @@ async def async_post_call_success_deployment_hook( except ValueError: typed_call_type = None # unknown call type - modified_response = response + callbacks: Final = litellm.callbacks + if ( + include_v2 + and Hook.post_deployment_success in CustomLoggerV2.declared_hooks + and (subscribers := hook_subscribers(Hook.post_deployment_success, callbacks, admit=Admit.every_logger)).has_v2 + ): + fields: Final = ResultFields(request_data, PROXY_REQUEST_KEYS, Field.response, response) # pyright: ignore[reportUnknownArgumentType] # the attempt's request dict + _ = kwargs_after( + await run_subscribers(subscribers, fields, partial(_start_post_deployment_success, fields, typed_call_type)) + ) + return fields.value - custom_logger_class: Final = _get_cached_custom_logger() - CustomGuardrail: Final = _get_cached_custom_guardrail() - for callback in litellm.callbacks: - if isinstance(callback, custom_logger_class): - try: - result = await callback.async_post_call_success_deployment_hook( - request_data, cast(LLMResponseTypes, modified_response), typed_call_type - ) - except Exception: # noqa: BLE001 # a broken callback must not fail a completed request - if isinstance(callback, CustomGuardrail): - raise - verbose_logger.exception( - "async_post_call_success_deployment_hook error in %s for call_type=%s", - type(callback).__name__, - typed_call_type, - ) - continue + modified_response = response # rebind-ok: each hook may replace it + for callback in callbacks: + if isinstance(callback, CustomLogger): + result = await _post_deployment_success(callback, request_data, modified_response, typed_call_type) if result is not None: modified_response = result return modified_response +def _start_post_deployment_success( + fields: ResultFields, call_type: CallTypes | None, callback: object, request_data: object +) -> Awaitable[Patch | None] | None: + if not isinstance(callback, CustomLogger): + return None + return _post_deployment_success_patch(callback, request_data, fields.value, call_type) + + +async def _post_deployment_success_patch( + callback: CustomLogger, request_data: object, response: object, call_type: CallTypes | None +) -> Patch | None: + result: Final = await _post_deployment_success(callback, request_data, response, call_type) + return None if result is None else Patch().set(Field.deprecated_response_obj, result) + + +async def _post_deployment_success( + callback: CustomLogger, request_data: object, response: object, call_type: CallTypes | None +) -> object: + try: + return await callback.async_post_call_success_deployment_hook( + request_data, # pyright: ignore[reportArgumentType] # the attempt's request dict + cast(LLMResponseTypes, response), + call_type, + ) + except Exception: # noqa: BLE001 # a broken callback must not fail a completed request + if isinstance(callback, CustomGuardrail): + raise + verbose_logger.exception( + "async_post_call_success_deployment_hook error in %s for call_type=%s", + type(callback).__name__, + call_type, + ) + return None + + async def async_post_call_failure_deployment_hook( - request_data: Mapping[str, object], exception: Exception, call_type: str + request_data: Mapping[str, object], exception: Exception, call_type: str, *, include_v2: bool = False ) -> None: """ Notify CustomLogger callbacks that a deployment attempt failed. @@ -1964,24 +2028,63 @@ async def async_post_call_failure_deployment_hook( safe_request_data: Final = MappingProxyType({k: v for k, v in request_data.items() if k != "attempted_targets"}) safe_exception: Final = _snapshot_exception_for_hook(exception) - custom_logger_class: Final = _get_cached_custom_logger() - for callback in litellm.callbacks: - if isinstance(callback, custom_logger_class): - try: - if _accepts_fallback_depth_kwarg_for_class(type(callback)): - await callback.async_post_call_failure_deployment_hook( - safe_request_data, safe_exception, typed_call_type, fallback_depth=fallback_depth - ) - else: - await callback.async_post_call_failure_deployment_hook( - safe_request_data, safe_exception, typed_call_type - ) - except Exception as callback_error: # noqa: BLE001 # a broken callback must not mask the real failure - verbose_logger.debug( - "async_post_call_failure_deployment_hook error in %s: %s", - type(callback).__name__, - callback_error, - ) + callbacks: Final = litellm.callbacks + if ( + include_v2 + and Hook.post_deployment_failure in CustomLoggerV2.declared_hooks + and (subscribers := hook_subscribers(Hook.post_deployment_failure, callbacks, admit=Admit.every_logger)).has_v2 + ): + fields: Final = ResultFields(safe_request_data, PROXY_REQUEST_KEYS, Field.error, safe_exception) + _ = kwargs_after( + await run_subscribers( + subscribers, fields, partial(_start_post_deployment_failure, fields, typed_call_type, fallback_depth) + ) + ) + return + for callback in callbacks: + if isinstance(callback, CustomLogger): + await _post_deployment_failure(callback, safe_request_data, safe_exception, typed_call_type, fallback_depth) + + +def _start_post_deployment_failure( + fields: ResultFields, + call_type: CallTypes | None, + fallback_depth: int | None, + callback: object, + request_data: object, +) -> Awaitable[None] | None: + if not isinstance(callback, CustomLogger): + return None + return _post_deployment_failure(callback, request_data, fields.value, call_type, fallback_depth) + + +async def _post_deployment_failure( + callback: CustomLogger, + request_data: object, + exception: object, + call_type: CallTypes | None, + fallback_depth: int | None, +) -> None: + try: + if _accepts_fallback_depth_kwarg_for_class(type(callback)): + await callback.async_post_call_failure_deployment_hook( + request_data, # pyright: ignore[reportArgumentType] # the attempt's read-only request view + exception, # pyright: ignore[reportArgumentType] # the exception snapshot + call_type, + fallback_depth=fallback_depth, + ) + else: + await callback.async_post_call_failure_deployment_hook( + request_data, # pyright: ignore[reportArgumentType] # see above + exception, # pyright: ignore[reportArgumentType] # see above + call_type, + ) + except Exception as callback_error: # noqa: BLE001 # a broken callback must not mask the real failure + verbose_logger.debug( + "async_post_call_failure_deployment_hook error in %s: %s", + type(callback).__name__, + callback_error, + ) def post_call_processing( @@ -2439,7 +2542,7 @@ def client(original_function): bind_budget_reservation_to_callbacks(logging_obj.litellm_params) kwargs["litellm_logging_obj"] = logging_obj - modified_kwargs: Final = await async_pre_call_deployment_hook(kwargs, call_type) + modified_kwargs: Final = await async_pre_call_deployment_hook(kwargs, call_type, include_v2=True) transcription_runs_unstreamed: Final = _transcription_runs_unstreamed(call_type, kwargs, modified_kwargs) if modified_kwargs is not None: kwargs = modified_kwargs @@ -2519,6 +2622,7 @@ def client(original_function): request_data=call_kwargs, exception=deployment_error, call_type=call_type, + include_v2=True, ) except asyncio.CancelledError: raise @@ -2582,6 +2686,7 @@ def client(original_function): else call_kwargs, response=result, call_type=_call_type_enum, + include_v2=True, ) ## Add response to cache diff --git a/tests/unit/integrations/custom_logger_v2/__init__.py b/tests/unit/integrations/custom_logger_v2/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/custom_logger_v2/test_dispatcher.py b/tests/unit/integrations/custom_logger_v2/test_dispatcher.py new file mode 100644 index 00000000000..394aeaf57fd --- /dev/null +++ b/tests/unit/integrations/custom_logger_v2/test_dispatcher.py @@ -0,0 +1,547 @@ +import asyncio +import gc +import threading +import weakref +from collections.abc import Awaitable, Callable, Coroutine, Mapping, Set +from dataclasses import dataclass, field +from types import MappingProxyType + +import pytest + +from litellm.integrations.custom_logger_v2 import ( + Abort, + Activation, + CallbackManifest, + CallbackRegistry, + CallContext, + ContractViolation, + CustomLoggerV2, + Dispatcher, + Field, + Hook, + Patch, + RequestSelection, + Reroute, + Respond, + Selection, + UndeclaredReadError, + build_manifest, + on, +) + +F = Field +Behaviour = Callable[[object, CallContext[object]], Awaitable[object]] + + +@dataclass +class Calls: + calls: int = 0 + + +async def _nothing(view: object, ctx: CallContext[object]) -> object: + return None + + +def v2_callback( + name: str, + hook: Hook = Hook.pre_call, + behaviour: Behaviour = _nothing, + *, + reads: Set[Field] = frozenset(), + writes: Set[Field] = frozenset(), + may_abort: bool = False, + parallel: bool | None = None, + awaited: bool = False, + selection: Selection | None = None, + state: type | None = None, +) -> CallbackManifest: + async def handle(self: CustomLoggerV2, view: object, ctx: CallContext[object]) -> object: + return await behaviour(view, ctx) + + decorated = on(hook, reads=reads, writes=writes, may_abort=may_abort, parallel=parallel, awaited=awaited)(handle) + namespace = {"handle": decorated, "selection": selection or Selection(name), "request_state": state} + return build_manifest(type(name, (CustomLoggerV2,), namespace)()) + + +def returning(result: object, journal: list[str] | None = None, name: str = "") -> Behaviour: + async def behaviour(view: object, ctx: CallContext[object]) -> object: + if journal is not None: + journal.append(f"{name}+") + await asyncio.sleep(0) + journal.append(f"{name}-") + return result + + return behaviour + + +class RecordingSink: + def __init__(self) -> None: + self.deliveries: list[Coroutine[object, object, None]] = [] + + def submit(self, delivery: Coroutine[object, object, None]) -> None: + self.deliveries.append(delivery) + + async def run_all(self) -> None: + for delivery in self.deliveries: + await delivery + self.deliveries.clear() + + +def dispatcher(*manifests: CallbackManifest, sink: RecordingSink | None = None, enforce: bool = False) -> Dispatcher: + return Dispatcher(CallbackRegistry(manifests), sink=sink or RecordingSink(), enforce=enforce) + + +MESSAGES = (MappingProxyType({"role": "user", "content": "my card is 4111"}),) +FIELDS: Mapping[Field, object] = MappingProxyType({F.messages: MESSAGES, F.metadata: MappingProxyType({"a": 1})}) + + +async def test_v2_callbacks_on_other_hooks_leave_this_hook_a_passthrough() -> None: + elsewhere = v2_callback("elsewhere", Hook.post_call_success, returning(Patch()), writes={F.response}) + + result = await dispatcher(elsewhere).open_scope(RequestSelection()).run_hook(Hook.pre_call, FIELDS, None) + + assert result.fields is FIELDS + assert result.outcome is None + assert result.patches == () + + +async def test_design_doc_worked_example_runs_in_its_stages_and_audit_sees_the_final_messages() -> None: + journal: list[str] = [] + masked = (MappingProxyType({"role": "user", "content": "my card is [CARD]"}),) + scanned = (MappingProxyType({"role": "user", "content": "[SCANNED]"}),) + seen_by_audit: list[object] = [] + + async def audit(view: object, ctx: CallContext[object]) -> None: + seen_by_audit.append(getattr(view, "messages")) + + async def secret_scan(view: object, ctx: CallContext[object]) -> Patch: + journal.append("secret_scan+") + assert getattr(view, "messages") == masked + return Patch().set(F.messages, scanned) + + sink = RecordingSink() + scope = dispatcher( + v2_callback( + "pii_mask", + behaviour=returning(Patch().set(F.messages, masked), journal, "pii_mask"), + reads={F.messages}, + writes={F.messages}, + may_abort=True, + ), + v2_callback( + "tagger", + behaviour=returning(Patch().merge(F.metadata, {"tag": "x"}), journal, "tagger"), + writes={F.metadata}, + ), + v2_callback("limiter", behaviour=returning(None, journal, "limiter"), reads={F.key_identity}, may_abort=True), + v2_callback("secret_scan", behaviour=secret_scan, reads={F.messages}, writes={F.messages}, may_abort=True), + v2_callback("audit", behaviour=audit, reads={F.messages}), + sink=sink, + ).open_scope(RequestSelection()) + + result = await scope.run_hook(Hook.pre_call, FIELDS, None) + + assert journal == ["pii_mask+", "tagger+", "pii_mask-", "tagger-", "limiter+", "limiter-", "secret_scan+"] + assert result.fields[F.messages] == scanned + assert result.fields[F.metadata] == {"a": 1, "tag": "x"} + assert seen_by_audit == [] + await sink.run_all() + assert seen_by_audit == [scanned] + + +async def test_patches_come_back_in_list_order_even_when_stages_run_out_of_it() -> None: + first = Patch().set(F.messages, ()) + second = Patch().set(F.optional_params, {}) + third = Patch().merge(F.metadata, {"b": 2}) + scope = dispatcher( + v2_callback("first", behaviour=returning(first), writes={F.messages}, may_abort=True), + v2_callback("second", behaviour=returning(second), writes={F.optional_params}, may_abort=True), + v2_callback("third", behaviour=returning(third), writes={F.metadata}), + ).open_scope(RequestSelection()) + + result = await scope.run_hook(Hook.pre_call, FIELDS, None) + + assert result.patches == (first, second, third) + + +async def test_a_later_writer_never_shows_an_earlier_reader_its_patch() -> None: + seen: list[object] = [] + + async def scanner(view: object, ctx: CallContext[object]) -> None: + seen.append(getattr(view, "messages")) + await asyncio.sleep(0) + + scope = dispatcher( + v2_callback("tagger", behaviour=returning(Patch().merge(F.metadata, {"t": 1})), writes={F.metadata}), + v2_callback("scanner", behaviour=scanner, reads={F.messages}, may_abort=True), + v2_callback("rewriter", behaviour=returning(Patch().set(F.messages, ())), writes={F.messages}), + ).open_scope(RequestSelection()) + + result = await scope.run_hook(Hook.pre_call, FIELDS, None) + + assert seen == [MESSAGES] + assert result.fields[F.messages] == () + + +@pytest.mark.parametrize("outcome", [Abort(status=400, message="blocked"), Respond(response={"ok": True})]) +async def test_terminal_outcome_stops_new_starts_and_discards_in_flight_patches( + outcome: Abort | Respond, +) -> None: + journal: list[str] = [] + scope = dispatcher( + v2_callback("guard", behaviour=returning(outcome, journal, "guard"), reads={F.messages}, may_abort=True), + v2_callback( + "tagger", behaviour=returning(Patch().merge(F.metadata, {"t": 1}), journal, "tagger"), writes={F.metadata} + ), + v2_callback("later", behaviour=returning(None, journal, "later"), may_abort=True), + ).open_scope(RequestSelection()) + + result = await scope.run_hook(Hook.pre_call, FIELDS, None) + + assert result.outcome is outcome + assert result.patches == () + assert journal == ["guard+", "tagger+", "guard-", "tagger-"] + + +async def test_patches_from_finished_stages_survive_a_later_abort() -> None: + patch = Patch().set(F.messages, ()) + abort = Abort(status=429, message="slow down") + scope = dispatcher( + v2_callback("masker", behaviour=returning(patch), writes={F.messages}, may_abort=True), + v2_callback("limiter", behaviour=returning(abort), may_abort=True), + ).open_scope(RequestSelection()) + + result = await scope.run_hook(Hook.pre_call, FIELDS, None) + + assert result.outcome is abort + assert result.patches == (patch,) + + +async def test_reroute_is_recorded_and_the_rest_of_the_list_still_runs() -> None: + journal: list[str] = [] + + scope = dispatcher( + v2_callback("router", behaviour=returning(Reroute(model="safe-model")), may_abort=True), + v2_callback("second_router", behaviour=returning(Reroute(model="other")), may_abort=True), + v2_callback("scanner", behaviour=returning(None, journal, "scanner"), may_abort=True), + ).open_scope(RequestSelection()) + + result = await scope.run_hook(Hook.pre_call, FIELDS, None) + + assert result.outcome == Reroute(model="safe-model") + assert scope.reroute == Reroute(model="safe-model") + assert journal == ["scanner+", "scanner-"] + + +async def test_a_later_abort_beats_an_earlier_reroute() -> None: + abort = Abort(status=400, message="blocked") + scope = dispatcher( + v2_callback("router", behaviour=returning(Reroute(model="safe-model")), may_abort=True), + v2_callback("guard", behaviour=returning(abort), may_abort=True), + ).open_scope(RequestSelection()) + + result = await scope.run_hook(Hook.pre_call, FIELDS, None) + + assert result.outcome is abort + assert scope.reroute is None + + +async def test_unexpected_exception_from_a_blocking_v2_callback_fails_the_request_closed() -> None: + error = ValueError("vendor down") + later = Calls() + + async def broken(view: object, ctx: CallContext[object]) -> None: + raise error + + async def counted(view: object, ctx: CallContext[object]) -> None: + later.calls += 1 + + scope = dispatcher( + v2_callback("guard", behaviour=broken, may_abort=True), + v2_callback("later", behaviour=counted, may_abort=True), + ).open_scope(RequestSelection()) + + with pytest.raises(ValueError, match="vendor down") as raised: + await scope.run_hook(Hook.pre_call, FIELDS, None) + + assert raised.value is error + assert later.calls == 0 + + +async def test_swallowing_hooks_count_the_failure_and_keep_the_other_patches() -> None: + async def broken(view: object, ctx: CallContext[object]) -> None: + raise ValueError("header lookup failed") + + patch = Patch().merge(F.response_headers, {"x-limit": "10"}) + owner = dispatcher( + v2_callback("broken", Hook.response_headers, broken, writes={F.response_headers}), + v2_callback("limits", Hook.response_headers, returning(patch), writes={F.response_headers}), + ) + + result = await owner.open_scope(RequestSelection()).run_hook(Hook.response_headers, {}, None) + + assert result.patches == (patch,) + assert owner.failures[("broken", Hook.response_headers)] == 1 + + +async def test_awaited_readers_run_after_blocking_work_in_parallel_and_never_fail_the_request() -> None: + journal: list[str] = [] + + async def failing(view: object, ctx: CallContext[object]) -> None: + journal.append("failing+") + await asyncio.sleep(0) + raise RuntimeError("redis down") + + owner = dispatcher( + v2_callback("release", behaviour=returning(None, journal, "release"), reads={F.key_identity}, awaited=True), + v2_callback("failing", behaviour=failing, awaited=True), + v2_callback("writer", behaviour=returning(Patch().set(F.messages, ()), journal, "writer"), writes={F.messages}), + ) + + result = await owner.open_scope(RequestSelection()).run_hook(Hook.pre_call, FIELDS, None) + + assert journal == ["writer+", "writer-", "release+", "failing+", "release-"] + assert result.fields[F.messages] == () + assert owner.failures[("failing", Hook.pre_call)] == 1 + + +async def test_during_call_abortables_run_in_parallel_and_list_order_picks_the_winner() -> None: + journal: list[str] = [] + first = Abort(status=400, message="first") + second = Abort(status=400, message="second") + scope = dispatcher( + v2_callback("slow", Hook.during_call, returning(None, journal, "slow"), may_abort=True), + v2_callback("first", Hook.during_call, returning(first, journal, "first"), may_abort=True), + v2_callback("second", Hook.during_call, returning(second, journal, "second"), may_abort=True), + ).open_scope(RequestSelection()) + + result = await scope.run_hook(Hook.during_call, FIELDS, None) + + assert journal == ["slow+", "first+", "second+", "slow-", "first-", "second-"] + assert result.outcome is first + + +async def test_parallel_pre_call_checks_see_the_patches_of_the_serial_chain() -> None: + seen: list[object] = [] + + async def check(view: object, ctx: CallContext[object]) -> None: + seen.append(getattr(view, "messages")) + + scope = dispatcher( + v2_callback("check", behaviour=check, reads={F.messages}, may_abort=True, parallel=True), + v2_callback("masker", behaviour=returning(Patch().set(F.messages, ())), writes={F.messages}, may_abort=True), + ).open_scope(RequestSelection()) + + await scope.run_hook(Hook.pre_call, FIELDS, None) + + assert seen == [()] + + +async def test_selection_mask_removes_opted_out_v2_callbacks_and_their_reads() -> None: + called: list[str] = [] + + async def note(view: object, ctx: CallContext[object]) -> None: + called.append("guard") + + owner = dispatcher( + v2_callback( + "guard", + behaviour=note, + reads={F.messages}, + may_abort=True, + selection=Selection("guard", activation=Activation.opt_in), + ), + v2_callback("limiter", reads={F.key_identity}, may_abort=True), + ) + + unrequested = owner.open_scope(RequestSelection()) + requested = owner.open_scope(RequestSelection(requested=frozenset({"guard"}))) + + assert unrequested.reads(Hook.pre_call) == {F.key_identity} + assert requested.reads(Hook.pre_call) == {F.key_identity, F.messages} + await unrequested.run_hook(Hook.pre_call, FIELDS, None) + assert called == [] + await requested.run_hook(Hook.pre_call, FIELDS, None) + assert called == ["guard"] + + +async def test_a_subscription_activation_override_beats_the_instance_activation() -> None: + class Masker(CustomLoggerV2): + selection = Selection("masker", activation=Activation.opt_in) + activation_overrides = MappingProxyType({Hook.log_prepare: Activation.required}) + + @on(Hook.pre_call, reads={F.messages}, may_abort=True) + async def check(self, view: object, ctx: CallContext[object]) -> None: ... + + @on(Hook.log_prepare, writes={F.standard_logging_payload}) + async def redact(self, view: object, ctx: CallContext[object]) -> None: ... + + scope = dispatcher(build_manifest(Masker())).open_scope(RequestSelection(admin_opt_outs=frozenset({"masker"}))) + + assert scope.plan(Hook.pre_call).stages == () + assert len(scope.plan(Hook.log_prepare).stages) == 1 + + +async def test_off_path_v2_callbacks_are_delivered_once_per_hook_when_it_finishes_without_state() -> None: + @dataclass + class Spans: + open: list[str] = field(default_factory=list) + + states: list[object] = [] + + async def remember(view: object, ctx: CallContext[object]) -> None: + states.append(ctx.state) + + sink = RecordingSink() + scope = dispatcher( + v2_callback("otel_pre", Hook.pre_call, remember, reads={F.messages}, state=Spans), + v2_callback("otel_awaited", Hook.pre_call, remember, awaited=True, state=Spans), + sink=sink, + ).open_scope(RequestSelection()) + + await scope.run_hook(Hook.pre_call, FIELDS, None) + + assert len(sink.deliveries) == 1 + assert len(states) == 1 + assert isinstance(states[0], Spans) + await sink.run_all() + assert states[1] is None + + +def test_sync_off_path_handler_runs_on_a_worker_thread() -> None: + threads: list[int] = [] + + class SyncLogger(CustomLoggerV2): + @on(Hook.attempt_ended, reads={F.usage}) + def ended(self, view: object, ctx: CallContext[object]) -> None: + threads.append(threading.get_ident()) + + sink = RecordingSink() + scope = dispatcher(build_manifest(SyncLogger()), sink=sink).open_scope(RequestSelection()) + + async def run() -> None: + await scope.run_hook(Hook.attempt_ended, {F.usage: 3}, None) + await sink.run_all() + + asyncio.run(run()) + + assert len(threads) == 1 + assert threads[0] != threading.get_ident() + + +async def test_request_state_is_shared_across_hooks_per_request_and_freed_on_close() -> None: + @dataclass + class Tokens: + tokens: dict[str, str] = field(default_factory=dict) + + states: list[object] = [] + + async def remember(view: object, ctx: CallContext[object]) -> None: + states.append(ctx.state) + + class PiiMask(CustomLoggerV2): + request_state = Tokens + + @on(Hook.pre_call, writes={F.messages}) + async def mask(self, view: object, ctx: CallContext[object]) -> None: + await remember(view, ctx) + + @on(Hook.post_call_success, writes={F.response}) + async def unmask(self, view: object, ctx: CallContext[object]) -> None: + await remember(view, ctx) + + owner = dispatcher(build_manifest(PiiMask())) + with owner.open_scope(RequestSelection()) as first: + await first.run_hook(Hook.pre_call, FIELDS, None) + await first.run_hook(Hook.post_call_success, {}, None) + with owner.open_scope(RequestSelection()) as second: + await second.run_hook(Hook.pre_call, FIELDS, None) + + assert isinstance(states[0], Tokens) + assert states[0] is states[1] + assert states[2] is not states[0] + freed = weakref.ref(states[0]) + states.clear() + gc.collect() + assert freed() is None + + +async def test_registry_change_reaches_new_scopes_and_leaves_open_ones_alone() -> None: + registry = CallbackRegistry() + owner = Dispatcher(registry, sink=RecordingSink()) + before = owner.open_scope(RequestSelection()) + late = v2_callback("late", reads={F.messages}, may_abort=True) + + registry.register(late) + after = owner.open_scope(RequestSelection()) + + assert before.plan(Hook.pre_call).stages == () + assert len(after.plan(Hook.pre_call).stages) == 1 + registry.unregister(late.logger) + assert owner.open_scope(RequestSelection()).plan(Hook.pre_call).stages == () + + +async def test_production_mode_applies_patches_without_checking_declarations() -> None: + sneaky = Patch().set(F.model, "other") + scope = dispatcher(v2_callback("sneaky", behaviour=returning(sneaky), writes={F.messages})).open_scope( + RequestSelection() + ) + + result = await scope.run_hook(Hook.pre_call, FIELDS, None) + + assert result.fields[F.model] == "other" + + +async def test_enforced_mode_restricts_views_to_declared_reads() -> None: + async def peek(view: object, ctx: CallContext[object]) -> None: + getattr(view, "metadata") + + scope = dispatcher( + v2_callback("peek", behaviour=peek, reads={F.messages}, may_abort=True), enforce=True + ).open_scope(RequestSelection()) + + with pytest.raises(UndeclaredReadError): + await scope.run_hook(Hook.pre_call, FIELDS, None) + + +@pytest.mark.parametrize( + ("result", "may_abort", "hook"), + [ + (Patch().set(F.model, "other"), False, Hook.pre_call), + (Abort(status=400, message="no"), False, Hook.pre_call), + (Reroute(model="other"), True, Hook.during_call), + ("not an outcome", True, Hook.pre_call), + ], +) +async def test_enforced_mode_rejects_results_outside_the_declaration( + result: object, may_abort: bool, hook: Hook +) -> None: + writes = frozenset({F.messages}) if hook is Hook.pre_call else frozenset() + scope = dispatcher( + v2_callback("liar", hook, returning(result), writes=writes, may_abort=may_abort, parallel=False), enforce=True + ).open_scope(RequestSelection()) + + with pytest.raises(ContractViolation): + await scope.run_hook(hook, FIELDS, None) + + +async def test_enforced_mode_violations_escape_the_awaited_swallow() -> None: + scope = dispatcher( + v2_callback("awaited", behaviour=returning(Patch().set(F.messages, ())), awaited=True), enforce=True + ).open_scope(RequestSelection()) + + with pytest.raises(ContractViolation): + await scope.run_hook(Hook.pre_call, FIELDS, None) + + +async def test_enforced_mode_accepts_declared_behaviour() -> None: + async def masker(view: object, ctx: CallContext[object]) -> Patch: + return Patch().set(F.messages, (*getattr(view, "messages"), "extra")) + + scope = dispatcher( + v2_callback("masker", behaviour=masker, reads={F.messages}, writes={F.messages}, may_abort=True), enforce=True + ).open_scope(RequestSelection()) + + result = await scope.run_hook(Hook.pre_call, FIELDS, None) + + assert result.fields[F.messages] == (*MESSAGES, "extra") diff --git a/tests/unit/integrations/custom_logger_v2/test_enums.py b/tests/unit/integrations/custom_logger_v2/test_enums.py new file mode 100644 index 00000000000..edbaef629f2 --- /dev/null +++ b/tests/unit/integrations/custom_logger_v2/test_enums.py @@ -0,0 +1,73 @@ +from collections import Counter +from enum import Enum + +import pytest + +from litellm.integrations.custom_logger_v2 import Activation, Field, FieldTier, Hook, OptOutBy + + +@pytest.mark.parametrize("enum_type", [Hook, Field, Activation, OptOutBy]) +def test_enum_values_are_their_lowercase_names_so_config_and_manifest_share_one_string(enum_type: type[Enum]) -> None: + mismatched = [(member.name, member.value) for member in enum_type if member.value != member.name.lower()] + assert mismatched == [] + assert all(isinstance(member, str) for member in enum_type) + + +def test_field_round_trips_from_its_config_string() -> None: + assert Field("provider_request_body") is Field.provider_request_body + assert Field.messages == "messages" + + +def test_field_bits_are_distinct_and_fit_a_u64_bitset() -> None: + bits = [member.bit for member in Field] + assert len(set(bits)) == len(bits) == 30 + assert all(0 <= bit < 64 for bit in bits) + + +def test_fields_fall_in_the_four_cost_tiers_of_the_design() -> None: + assert Counter(member.tier for member in Field) == { + FieldTier.small_record: 11, + FieldTier.content: 12, + FieldTier.wire: 4, + FieldTier.compatibility: 3, + } + assert {member for member in Field if member.tier is FieldTier.compatibility} == { + Field.standard_logging_payload, + Field.deprecated_kwargs, + Field.deprecated_response_obj, + } + + +def test_field_bits_are_the_wire_contract_shared_with_the_rust_manifest() -> None: + assert {member.value: member.bit for member in Field} == { + "model": 0, + "key_identity": 1, + "deployment": 2, + "candidate_deployments": 3, + "request_tags": 4, + "dynamic_params": 5, + "guardrail_info": 6, + "outcome": 7, + "error": 8, + "usage": 9, + "cost": 10, + "metadata": 11, + "messages": 12, + "input": 13, + "tools": 14, + "optional_params": 15, + "request_headers": 16, + "response": 17, + "response_headers": 18, + "stream_chunk": 19, + "assembled_stream": 20, + "tool_result": 21, + "model_list": 22, + "request_body": 23, + "provider_request_body": 24, + "provider_response_raw": 25, + "provider_response_headers": 26, + "standard_logging_payload": 27, + "deprecated_kwargs": 28, + "deprecated_response_obj": 29, + } diff --git a/tests/unit/integrations/custom_logger_v2/test_hook_point_fields.py b/tests/unit/integrations/custom_logger_v2/test_hook_point_fields.py new file mode 100644 index 00000000000..96afd48fb44 --- /dev/null +++ b/tests/unit/integrations/custom_logger_v2/test_hook_point_fields.py @@ -0,0 +1,76 @@ +from typing import Final + +import pytest + +from litellm.integrations.custom_logger_v2 import Field, InvalidPatchError, Patch, apply_patch +from litellm.integrations.custom_logger_v2.legacy import PROXY_REQUEST_KEYS +from litellm.integrations.custom_logger_v2.hook_point_fields import CandidateFields, ResultFields + + +def test_candidate_fields_show_legacy_the_list_and_v2_callbacks_a_tuple_and_keep_only_a_narrowing() -> None: + first: Final[dict[str, object]] = {"id": "a"} + second: Final[dict[str, object]] = {"id": "b"} + deployments: Final[list[object]] = [first, second] + fields: Final = CandidateFields(deployments, "group") + + narrowed: Final = apply_patch(fields, Patch().set(Field.candidate_deployments, (second,))) + + assert fields[Field.deprecated_kwargs] is deployments + assert fields[Field.candidate_deployments] == (first, second) + assert fields[Field.model] == "group" + assert narrowed[Field.deprecated_kwargs] == [second] + assert narrowed[Field.deprecated_kwargs][0] is second + assert set(narrowed) == {Field.deprecated_kwargs, Field.candidate_deployments, Field.model} + + +def test_candidate_fields_carry_whatever_a_legacy_filter_returned() -> None: + fields: Final = apply_patch(CandidateFields([{"id": "a"}], "group"), Patch().set(Field.deprecated_kwargs, None)) + + assert fields[Field.deprecated_kwargs] is None + assert Field.candidate_deployments not in fields + with pytest.raises(InvalidPatchError, match="stay a sequence"): + apply_patch(fields, Patch().set(Field.candidate_deployments, [])) + + +@pytest.mark.parametrize( + "patch", + [ + pytest.param(Patch().set(Field.candidate_deployments, [{"id": "a"}]), id="an equal copy"), + pytest.param(Patch().set(Field.model, "other"), id="another field"), + ], +) +def test_candidate_fields_refuse_anything_but_a_narrowing(patch: Patch) -> None: + with pytest.raises(InvalidPatchError): + apply_patch(CandidateFields([{"id": "a"}], "group"), patch) + + +def test_result_fields_keep_one_live_slot_that_legacy_and_v2_callback_patches_both_replace() -> None: + request: Final[dict[str, object]] = {"model": "m", "messages": ["hi"]} + fields: Final = ResultFields(request, PROXY_REQUEST_KEYS, Field.response, "original") + + after_legacy: Final = apply_patch(fields, Patch().set(Field.deprecated_response_obj, "from legacy")) + seen_by_v2: Final = fields[Field.response] + after_v2: Final = apply_patch(after_legacy, Patch().set(Field.response, "from v2")) + + assert after_legacy is fields + assert after_v2 is fields + assert seen_by_v2 == "from legacy" + assert fields.value == "from v2" + assert fields[Field.deprecated_kwargs] is request + assert fields[Field.model] == "m" + assert set(fields) == { + Field.deprecated_kwargs, + Field.model, + Field.messages, + Field.response, + Field.deprecated_response_obj, + } + + +def test_result_fields_refuse_a_patch_outside_the_slot_and_leave_it_unchanged() -> None: + fields: Final = ResultFields({"model": "m"}, PROXY_REQUEST_KEYS, Field.error, "original") + + with pytest.raises(InvalidPatchError, match="only error"): + apply_patch(fields, Patch().set(Field.error, "changed").set(Field.model, "other")) + + assert fields.value == "original" diff --git a/tests/unit/integrations/custom_logger_v2/test_hooks.py b/tests/unit/integrations/custom_logger_v2/test_hooks.py new file mode 100644 index 00000000000..fd18531cd40 --- /dev/null +++ b/tests/unit/integrations/custom_logger_v2/test_hooks.py @@ -0,0 +1,129 @@ +from collections import Counter + +import pytest + +from litellm.integrations.custom_logger_v2 import ( + HOOK_SPECS, + Abort, + Field, + Hook, + InvalidDeclarationError, + Reroute, + Respond, + Subscription, +) + +F = Field + +_LIFECYCLE_HOOKS = [hook for hook in Hook if not HOOK_SPECS[hook].request_path] + + +def test_every_hook_has_exactly_one_table_row() -> None: + assert set(HOOK_SPECS) == set(Hook) + + +def test_each_legacy_method_maps_onto_one_hook_point() -> None: + counts = Counter(method for spec in HOOK_SPECS.values() for method in spec.legacy_methods) + assert [method for method, count in counts.items() if count > 1] == [] + + +def test_legacy_methods_are_real_custom_logger_methods() -> None: + from litellm.integrations.custom_logger import CustomLogger + + missing = sorted( + method + for spec in HOOK_SPECS.values() + for method in spec.legacy_methods + if not callable(getattr(CustomLogger, method, None)) + ) + assert missing == [] + + +def test_legacy_guardrail_events_are_real_guardrail_event_hooks() -> None: + from litellm.types.guardrails import GuardrailEventHooks + + events = {event for spec in HOOK_SPECS.values() for event in spec.legacy_guardrail_events} + assert events <= {member.value for member in GuardrailEventHooks} + assert events == {"pre_call", "during_call", "post_call", "post_mcp_call"} + + +def test_outcomes_are_allowed_only_on_request_path_hooks() -> None: + assert {hook for hook in Hook if HOOK_SPECS[hook].outcomes} <= { + hook for hook in Hook if HOOK_SPECS[hook].request_path + } + + +def test_reroute_is_allowed_only_before_the_provider_call_on_pre_call() -> None: + assert {hook for hook in Hook if Reroute in HOOK_SPECS[hook].outcomes} == {Hook.pre_call} + + +def test_reroute_is_the_only_non_terminal_outcome() -> None: + assert (Abort.terminal, Respond.terminal, Reroute.terminal) == (True, True, False) + + +@pytest.mark.parametrize("hook", _LIFECYCLE_HOOKS) +def test_abort_on_a_lifecycle_hook_is_rejected(hook: Hook) -> None: + with pytest.raises(InvalidDeclarationError, match="allows no outcome"): + Subscription(hook=hook, handler="h", may_abort=True) + + +def test_abort_on_response_headers_is_rejected_because_its_errors_are_swallowed() -> None: + with pytest.raises(InvalidDeclarationError, match="allows no outcome"): + Subscription(hook=Hook.response_headers, handler="h", writes=frozenset({F.response_headers}), may_abort=True) + + +@pytest.mark.parametrize( + ("hook", "field"), + [ + (Hook.during_call, F.messages), + (Hook.post_call_success, F.messages), + (Hook.pre_call, F.response), + (Hook.pre_call, F.deprecated_kwargs), + (Hook.attempt_ended, F.metadata), + (Hook.provider_request_sent, F.provider_request_body), + (Hook.log_prepare, F.messages), + ], +) +def test_write_the_hook_cannot_apply_is_rejected(hook: Hook, field: Field) -> None: + with pytest.raises(InvalidDeclarationError, match=f"writes {field.value}, which this hook point cannot apply"): + Subscription(hook=hook, handler="h", writes=frozenset({field})) + + +def test_pre_call_may_write_tools_model_and_input() -> None: + writes = frozenset({F.tools, F.model, F.input, F.messages, F.metadata, F.guardrail_info}) + assert Subscription(hook=Hook.pre_call, handler="h", writes=writes, may_abort=True).writes == writes + + +def test_log_prepare_may_rewrite_the_standard_logging_payload() -> None: + writes = frozenset({F.standard_logging_payload}) + assert Subscription(hook=Hook.log_prepare, handler="h", writes=writes).writes == writes + + +def test_stream_chunk_cannot_be_subscribed_while_streams_are_out_of_scope() -> None: + with pytest.raises(InvalidDeclarationError, match="no V2 subscription support"): + Subscription(hook=Hook.stream_chunk, handler="h", reads=frozenset({F.stream_chunk})) + + +@pytest.mark.parametrize( + ("writes", "may_abort"), + [(frozenset({F.metadata}), False), (frozenset[Field](), True)], +) +def test_awaited_subscription_that_writes_or_aborts_is_rejected(writes: frozenset[Field], may_abort: bool) -> None: + with pytest.raises(InvalidDeclarationError, match="awaited subscription must neither write nor abort"): + Subscription(hook=Hook.pre_call, handler="h", writes=writes, may_abort=may_abort, awaited=True) + + +def test_awaited_read_only_subscription_is_accepted_on_a_lifecycle_hook() -> None: + subscription = Subscription( + hook=Hook.call_ended, handler="release", reads=frozenset({F.key_identity}), awaited=True + ) + assert (subscription.awaited, subscription.writes, subscription.may_abort) == (True, frozenset(), False) + + +@pytest.mark.parametrize( + ("writes", "may_abort"), + [(frozenset[Field](), False), (frozenset({F.messages}), True)], +) +def test_parallel_is_only_for_abort_only_subscriptions(writes: frozenset[Field], may_abort: bool) -> None: + with pytest.raises(InvalidDeclarationError, match="parallel is only for"): + Subscription(hook=Hook.pre_call, handler="h", writes=writes, may_abort=may_abort, parallel=True) diff --git a/tests/unit/integrations/custom_logger_v2/test_host.py b/tests/unit/integrations/custom_logger_v2/test_host.py new file mode 100644 index 00000000000..81ce45f427d --- /dev/null +++ b/tests/unit/integrations/custom_logger_v2/test_host.py @@ -0,0 +1,145 @@ +import asyncio +import threading +from collections.abc import Awaitable, Callable, Generator +from dataclasses import dataclass, field +from typing import Final + +import pytest + +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_logger_v2 import ( + Abort, + CallContext, + CustomLoggerV2, + Field, + Hook, + Patch, + on, +) +from litellm.integrations.custom_logger_v2.dispatcher import RequestScope +from litellm.integrations.custom_logger_v2.host import ( + V2Outcome, + bind_scope, + current_scope, + request_scope, + reset_host_for_tests, +) +from litellm.integrations.custom_logger_v2.legacy import ( + PROXY_REQUEST_KEYS, + KwargsFields, + clear_subscribers_cache_for_tests, + hook_subscribers, + kwargs_after, + run_subscribers, +) + +F = Field +ABORT = Abort(status=403, message="blocked") + + +@pytest.fixture(autouse=True) +def fresh_host() -> Generator[None]: + reset_host_for_tests() + clear_subscribers_cache_for_tests() + yield + reset_host_for_tests() + clear_subscribers_cache_for_tests() + + +@dataclass +class Seen: + states: list[object] = field(default_factory=list) + + +class Masker(CustomLoggerV2): + @dataclass + class request_state: + hits: int = 0 + + def __init__(self, seen: Seen, outcome: Abort | None = None) -> None: + self.seen = seen + self.outcome = outcome + + @on(Hook.pre_call, reads={F.messages}, writes={F.messages}, may_abort=True) + async def mask(self, view: object, ctx: CallContext[object]) -> Patch | Abort: + self.seen.states.append(ctx.state) + return self.outcome or Patch().set(F.messages, ("masked",)) + + +class PreCallOnly(CustomLogger): + async def async_pre_call_hook(self, *args: object, **kwargs: object) -> None: + return None + + +def _recording(ran: list[object]) -> Callable[[object, object], Awaitable[None]]: + def start(logger: object, data: object) -> Awaitable[None]: + async def run() -> None: + ran.append(logger) + + return run() + + return start + + +async def _run_pre_call(callbacks: list[object], ran: list[object], scope: RequestScope | None = None) -> object: + subscribers: Final = hook_subscribers(Hook.pre_call, callbacks) + fields: Final = KwargsFields({"messages": ("raw",)}, PROXY_REQUEST_KEYS) + return kwargs_after(await run_subscribers(subscribers, fields, _recording(ran), scope=scope)) + + +async def test_a_v2_callback_abort_raises_v2_callback_outcome_and_skips_the_legacy_entry_after_it() -> None: + ran: Final[list[object]] = [] + callbacks: Final[list[object]] = [Masker(Seen(), ABORT), PreCallOnly()] + + with pytest.raises(V2Outcome) as raised: + await _run_pre_call(callbacks, ran) + + assert raised.value.outcome is ABORT + assert ran == [] + + +async def test_request_scope_is_current_inside_and_shares_state_across_hook_runs() -> None: + seen: Final = Seen() + ran: Final[list[object]] = [] + callbacks: Final[list[object]] = [Masker(seen), PreCallOnly()] + + with request_scope() as scope: + assert current_scope() is scope + await _run_pre_call(callbacks, ran) + await _run_pre_call(callbacks, ran) + await _run_pre_call(callbacks, ran) + + assert current_scope() is None + assert len(ran) == 3 + assert seen.states[0] is seen.states[1] + assert seen.states[2] is not seen.states[0] + + +async def test_an_explicit_scope_is_used_where_no_scope_is_current() -> None: + seen: Final = Seen() + callbacks: Final[list[object]] = [Masker(seen)] + + with request_scope() as scope: + pass + await _run_pre_call(callbacks, [], scope=scope) + await _run_pre_call(callbacks, [], scope=scope) + + assert seen.states[0] is seen.states[1] + + +async def test_bind_scope_carries_a_scope_into_a_thread_that_did_not_inherit_it() -> None: + found: list[object] = [] + + def worker(scope: object) -> None: + found.append(current_scope()) + with request_scope() as other, bind_scope(other): + found.append(current_scope() is other) + + with request_scope() as scope: + thread = threading.Thread(target=worker, args=(scope,)) + thread.start() + await asyncio.to_thread(thread.join) + with bind_scope(scope): + assert current_scope() is scope + + assert found == [None, True] diff --git a/tests/unit/integrations/custom_logger_v2/test_legacy.py b/tests/unit/integrations/custom_logger_v2/test_legacy.py new file mode 100644 index 00000000000..2e058564980 --- /dev/null +++ b/tests/unit/integrations/custom_logger_v2/test_legacy.py @@ -0,0 +1,332 @@ +import asyncio +import gc +from collections.abc import Awaitable, Generator +from typing import Final + +import pytest + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_logger_v2 import ( + CallContext, + ContractViolation, + CustomLoggerV2, + Field, + Hook, + InvalidPatchError, + Patch, + on, +) +from litellm.integrations.custom_logger_v2.host import reset_host_for_tests +from litellm.integrations.custom_logger_v2.legacy import ( + PROXY_REQUEST_KEYS, + Admit, + KwargsFields, + adapter_for, + clear_subscribers_cache_for_tests, + hook_subscribers, + kwargs_after, + kwargs_patch, + run_subscribers, + run_subscribers_sync, +) +from litellm.types.guardrails import GuardrailEventHooks + +F = Field + + +@pytest.fixture(autouse=True) +def fresh() -> Generator[None]: + reset_host_for_tests() + clear_subscribers_cache_for_tests() + yield + reset_host_for_tests() + clear_subscribers_cache_for_tests() + + +class PreCallOnly(CustomLogger): + async def async_pre_call_hook(self, *args: object, **kwargs: object) -> None: + return None + + +class InheritsPreCall(PreCallOnly): + pass + + +class Moderator(CustomLogger): + async def async_moderation_hook(self, *args: object, **kwargs: object) -> None: + return None + + +class Exporter(CustomLogger): + def log_success_event(self, *args: object, **kwargs: object) -> None: + return None + + +class Masking(CustomLogger): + def logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + return kwargs, result + + +class Writer(CustomLoggerV2): + def __init__(self, journal: list[object], replacement: object) -> None: + self.journal = journal + self.replacement = replacement + + @on(Hook.pre_call, reads={F.messages}, writes={F.messages}, may_abort=True) + async def pre_call(self, view: object, ctx: CallContext[object]) -> Patch: + self.journal.append(("v2", getattr(view, "messages"))) + return Patch().set(F.messages, self.replacement) + + +class Tagger(CustomLoggerV2): + def __init__(self, journal: list[object]) -> None: + self.journal = journal + + @on(Hook.pre_call, writes={F.model}) + async def pre_call(self, view: object, ctx: CallContext[object]) -> Patch: + self.journal.append("tagger") + return Patch().set(F.model, "tagged") + + +class Observer(CustomLoggerV2): + def __init__(self, journal: list[object]) -> None: + self.journal = journal + + @on(Hook.attempt_ended, reads={F.deprecated_kwargs}) + async def ended(self, view: object, ctx: CallContext[object]) -> None: + self.journal.append("observer") + + +class BlockingAttemptEnded(CustomLoggerV2): + @on(Hook.log_prepare, writes={F.standard_logging_payload}) + async def prepare(self, view: object, ctx: CallContext[object]) -> None: + return None + + +def test_manifest_subscribes_only_the_hooks_the_class_overrides_through_its_mro() -> None: + direct: Final = adapter_for(PreCallOnly()).manifest.subscriptions + inherited: Final = adapter_for(InheritsPreCall()).manifest.subscriptions + + assert set(direct) == {Hook.pre_call} + assert set(inherited) == {Hook.pre_call} + assert direct[Hook.pre_call].reads == {F.deprecated_kwargs, F.deprecated_response_obj} + assert direct[Hook.pre_call].writes == {F.deprecated_kwargs} + assert (direct[Hook.pre_call].may_abort, direct[Hook.pre_call].parallel) == (True, False) + + +def test_during_call_and_parallel_guardrails_gather_and_write_nothing() -> None: + moderation: Final = adapter_for(Moderator()).manifest.subscriptions[Hook.during_call] + guardrail: Final = adapter_for( + CustomGuardrail(guardrail_name="g", event_hook=GuardrailEventHooks.pre_call, run_in_parallel=True) + ).manifest.subscriptions + + assert (moderation.writes, moderation.may_abort, moderation.parallel) == (frozenset(), True, True) + assert (guardrail[Hook.pre_call].writes, guardrail[Hook.pre_call].parallel) == (frozenset(), True) + assert guardrail[Hook.during_call].parallel + assert guardrail[Hook.mcp_tool_result].writes == {F.deprecated_kwargs} + + +def test_log_hooks_neither_write_nor_abort_but_log_prepare_keeps_its_kwargs_rewrite() -> None: + attempt: Final = adapter_for(Exporter()).manifest.subscriptions[Hook.attempt_ended] + prepare: Final = adapter_for(Masking()).manifest.subscriptions[Hook.log_prepare] + + assert (attempt.writes, attempt.may_abort) == (frozenset(), False) + assert (prepare.writes, prepare.may_abort) == ({F.deprecated_kwargs}, False) + + +def test_every_logger_admits_loggers_that_do_not_override_the_hook() -> None: + exporter: Final = Exporter() + + assert hook_subscribers(Hook.pre_call, [exporter]).entries == () + assert len(hook_subscribers(Hook.pre_call, [exporter], admit=Admit.every_logger).entries) == 1 + + +def test_adapters_are_cached_per_instance_and_dropped_with_it() -> None: + logger = PreCallOnly() + adapter = adapter_for(logger) + + assert adapter_for(logger) is adapter + assert adapter.inner is logger + del logger + gc.collect() + assert adapter.inner is None + + +def test_subscribers_are_cached_per_list_content_and_rebuilt_when_it_changes() -> None: + first: Final = PreCallOnly() + callbacks: Final[list[object]] = ["a-string-callback", first] + subscribers: Final = hook_subscribers(Hook.pre_call, callbacks) + + assert hook_subscribers(Hook.pre_call, callbacks) is subscribers + callbacks.append(PreCallOnly()) + grown: Final = hook_subscribers(Hook.pre_call, callbacks) + assert grown is not subscribers + assert len(grown.entries) == 2 + + +def test_a_cached_subscribers_is_dropped_when_a_legacy_logger_in_it_dies() -> None: + for _ in range(5): + dying = PreCallOnly() + _ = hook_subscribers(Hook.pre_call, [dying]) + del dying + gc.collect() + replacement = PreCallOnly() + + assert hook_subscribers(Hook.pre_call, [replacement]).legacy == (replacement,) + + +def test_legacy_kwargs_conflict_with_every_field_so_list_order_is_kept() -> None: + journal: Final[list[object]] = [] + subscribers: Final = hook_subscribers(Hook.pre_call, [Tagger(journal), PreCallOnly(), Tagger(journal)]) + + assert [len(stage) for stage in subscribers.plan.stages] == [1, 1, 1] + assert [entry.position for stage in subscribers.plan.stages for entry in stage] == [0, 1, 2] + + +def _recording(journal: list[object], name: str, replacement: object = None): + def start(logger: object, data: object) -> Awaitable[Patch | None]: + async def run() -> Patch | None: + journal.append((name, KwargsFields(data, PROXY_REQUEST_KEYS)[F.messages])) + return None if replacement is None else kwargs_patch(data, replacement) + + return run() + + return start + + +async def test_a_v2_callback_between_two_legacy_entries_sees_the_first_edit_and_patches_what_the_second_sees() -> None: + journal: Final[list[object]] = [] + first: Final = PreCallOnly() + second: Final = PreCallOnly() + subscribers: Final = hook_subscribers(Hook.pre_call, [first, Writer(journal, ["v2"]), second]) + replaced: Final[dict[str, object]] = {"messages": ["first"]} + + def start(logger: object, data: object) -> Awaitable[Patch | None]: + async def run() -> Patch | None: + journal.append( + ("first" if logger is first else "second", KwargsFields(data, PROXY_REQUEST_KEYS)[F.messages]) + ) + return kwargs_patch(data, replaced) if logger is first else None + + return run() + + result: Final = await run_subscribers( + subscribers, KwargsFields({"messages": ["original"]}, PROXY_REQUEST_KEYS), start + ) + + assert journal == [("first", ["original"]), ("v2", ["first"]), ("second", ["v2"])] + assert kwargs_after(result) is replaced + assert replaced["messages"] == ["v2"] + + +async def test_a_v2_callback_writer_without_abort_still_waits_for_the_legacy_entry_before_it() -> None: + journal: Final[list[object]] = [] + callbacks: Final[list[object]] = [PreCallOnly(), Tagger(journal)] + subscribers: Final = hook_subscribers(Hook.pre_call, callbacks) + data: Final[dict[str, object]] = {"model": "m", "messages": []} + + def start(logger: object, kwargs: object) -> Awaitable[None]: + async def run() -> None: + journal.append("legacy") + await asyncio.sleep(0) + + return run() + + await run_subscribers(subscribers, KwargsFields(data, PROXY_REQUEST_KEYS), start) + + assert journal == ["legacy", "tagger"] + assert data["model"] == "tagged" + + +async def test_legacy_only_and_dispatched_subscribers_give_the_same_result() -> None: + first: Final = PreCallOnly() + second: Final = PreCallOnly() + replacement: Final[dict[str, object]] = {"messages": ["replaced"]} + + async def run(callbacks: list[object]) -> tuple[list[object], object]: + journal: list[object] = [] + subscribers = hook_subscribers(Hook.pre_call, callbacks) + result = await run_subscribers( + subscribers, KwargsFields({"messages": ["original"]}, PROXY_REQUEST_KEYS), _recording(journal, "x", replacement) + ) + return journal, kwargs_after(result) + + plain: Final = await run([first, second]) + dispatched: Final = await run([first, second, Observer([])]) + + assert plain == dispatched + assert plain[1] is replacement + + +async def test_a_legacy_exception_propagates_unchanged_through_both_paths() -> None: + error: Final = RuntimeError("legacy said no") + + def start(logger: object, data: object) -> Awaitable[None]: + async def run() -> None: + raise error + + return run() + + for callbacks in ([PreCallOnly()], [PreCallOnly(), Writer([], [])]): + with pytest.raises(RuntimeError) as raised: + await run_subscribers( + hook_subscribers(Hook.pre_call, callbacks), KwargsFields({"messages": []}, PROXY_REQUEST_KEYS), start + ) + assert raised.value is error + + +def test_kwargs_fields_write_through_to_the_live_request() -> None: + metadata: Final[dict[str, object]] = {"kept": 1} + data: Final[dict[str, object]] = {"model": "m", "metadata": metadata} + fields: Final = KwargsFields(data, {**PROXY_REQUEST_KEYS, F.metadata: "metadata"}) + + patched: Final = fields.apply(Patch().set(F.model, "n").merge(F.metadata, {"added": 2})) + + assert patched[F.deprecated_kwargs] is data + assert data == {"model": "n", "metadata": {"kept": 1, "added": 2}} + assert data["metadata"] is metadata + assert F.messages not in fields + with pytest.raises(InvalidPatchError, match="no place"): + fields.apply(Patch().set(F.request_headers, {})) + + +def test_sync_subscribers_run_legacy_inline_in_order_and_hand_v2_callbacks_to_the_off_path() -> None: + journal: Final[list[object]] = [] + first: Final = Exporter() + second: Final = Exporter() + subscribers: Final = hook_subscribers(Hook.attempt_ended, [first, Observer(journal), second]) + + def start(logger: object, data: object) -> None: + journal.append("first" if logger is first else "second") + + run_subscribers_sync(subscribers, KwargsFields({}, {}), start) + + assert journal[:2] == ["first", "second"] + + +def test_sync_subscribers_reject_a_blocking_v2_callback() -> None: + subscribers: Final = hook_subscribers(Hook.log_prepare, [Masking(), BlockingAttemptEnded()]) + + with pytest.raises(ContractViolation, match="synchronously"): + run_subscribers_sync(subscribers, KwargsFields({}, {}), lambda logger, data: None) + + +def test_every_entry_keeps_string_and_callable_callbacks_in_list_order() -> None: + journal: Final[list[object]] = [] + + def plain_callback() -> None: + return None + + exporter: Final = Exporter() + callbacks: Final[list[object]] = ["langfuse", plain_callback, exporter] + + def start(logger: object, data: object) -> None: + journal.append(logger) + + run_subscribers_sync(hook_subscribers(Hook.attempt_ended, callbacks, admit=Admit.every_entry), KwargsFields({}, {}), start) + + assert journal == ["langfuse", plain_callback, exporter] + assert hook_subscribers(Hook.attempt_ended, callbacks).legacy == (exporter,) diff --git a/tests/unit/integrations/custom_logger_v2/test_logger.py b/tests/unit/integrations/custom_logger_v2/test_logger.py new file mode 100644 index 00000000000..3dfcc1bf433 --- /dev/null +++ b/tests/unit/integrations/custom_logger_v2/test_logger.py @@ -0,0 +1,242 @@ +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field + +import pytest + +from litellm.integrations.custom_logger_v2 import ( + Abort, + Activation, + CallContext, + CustomLoggerV2, + Field, + Hook, + InvalidDeclarationError, + Patch, + Selection, + Subscription, + build_manifest, + on, +) + +F = Field + + +class Prometheus(CustomLoggerV2): + @on(Hook.attempt_ended, reads={F.outcome, F.deployment, F.key_identity, F.usage, F.cost, F.request_tags}) + async def attempt_ended(self, view: object, ctx: CallContext[None]) -> None: ... + + @on(Hook.call_ended, reads={F.outcome, F.deployment, F.key_identity, F.error}) + async def call_ended(self, view: object, ctx: CallContext[None]) -> None: ... + + +class PiiMask(CustomLoggerV2): + @dataclass + class request_state: + tokens: dict[str, str] = field(default_factory=dict) + + @on(Hook.pre_call, reads={F.messages}, writes={F.messages}, may_abort=True) + async def mask(self, view: object, ctx: CallContext[object]) -> Patch | None: ... + + @on(Hook.post_call_success, reads={F.response}, writes={F.response}) + async def unmask(self, view: object, ctx: CallContext[object]) -> Patch | None: ... + + +class ParallelRequestLimiter(CustomLoggerV2): + selection = Selection("parallel_request_limiter", activation=Activation.required) + + @on(Hook.pre_call, reads={F.key_identity}, may_abort=True) + async def acquire(self, view: object, ctx: CallContext[None]) -> Abort | None: ... + + @on(Hook.response_headers, reads={F.key_identity}, writes={F.response_headers}) + async def headers(self, view: object, ctx: CallContext[None]) -> Patch: ... + + @on(Hook.call_ended, reads={F.key_identity, F.usage, F.outcome}) + async def release(self, view: object, ctx: CallContext[None]) -> None: ... + + +def test_prometheus_example_declares_two_read_only_lifecycle_subscriptions() -> None: + assert dict(Prometheus.subscriptions) == { + Hook.attempt_ended: Subscription( + hook=Hook.attempt_ended, + handler="attempt_ended", + reads=frozenset({F.outcome, F.deployment, F.key_identity, F.usage, F.cost, F.request_tags}), + ), + Hook.call_ended: Subscription( + hook=Hook.call_ended, + handler="call_ended", + reads=frozenset({F.outcome, F.deployment, F.key_identity, F.error}), + ), + } + + +def test_pii_mask_example_declares_its_writes_and_request_state() -> None: + assert dict(PiiMask.subscriptions) == { + Hook.pre_call: Subscription( + hook=Hook.pre_call, + handler="mask", + reads=frozenset({F.messages}), + writes=frozenset({F.messages}), + may_abort=True, + ), + Hook.post_call_success: Subscription( + hook=Hook.post_call_success, + handler="unmask", + reads=frozenset({F.response}), + writes=frozenset({F.response}), + ), + } + assert build_manifest(PiiMask()).request_state is PiiMask.request_state + + +def test_parallel_request_limiter_example_is_required_and_aborts_without_writing() -> None: + manifest = build_manifest(ParallelRequestLimiter()) + assert manifest.selection == Selection("parallel_request_limiter", activation=Activation.required) + assert dict(manifest.subscriptions) == { + Hook.pre_call: Subscription( + hook=Hook.pre_call, handler="acquire", reads=frozenset({F.key_identity}), may_abort=True + ), + Hook.response_headers: Subscription( + hook=Hook.response_headers, + handler="headers", + reads=frozenset({F.key_identity}), + writes=frozenset({F.response_headers}), + ), + Hook.call_ended: Subscription( + hook=Hook.call_ended, handler="release", reads=frozenset({F.key_identity, F.usage, F.outcome}) + ), + } + + +def test_a_declaration_the_hook_rejects_fails_when_the_class_is_defined() -> None: + with pytest.raises(InvalidDeclarationError, match="writes response, which this hook point cannot apply"): + + class Broken(CustomLoggerV2): + @on(Hook.pre_call, writes={F.response}) + async def handle(self, view: object, ctx: CallContext[None]) -> None: ... + + +def test_during_call_abort_subscriptions_are_parallel_unless_they_opt_out() -> None: + class Moderation(CustomLoggerV2): + @on(Hook.during_call, reads={F.messages}, may_abort=True) + async def moderate(self, view: object, ctx: CallContext[None]) -> Abort | None: ... + + class SerialModeration(CustomLoggerV2): + @on(Hook.during_call, reads={F.messages}, may_abort=True, parallel=False) + async def moderate(self, view: object, ctx: CallContext[None]) -> Abort | None: ... + + class Watcher(CustomLoggerV2): + @on(Hook.during_call, reads={F.messages}) + async def watch(self, view: object, ctx: CallContext[None]) -> None: ... + + assert Moderation.subscriptions[Hook.during_call].parallel is True + assert SerialModeration.subscriptions[Hook.during_call].parallel is False + assert Watcher.subscriptions[Hook.during_call].parallel is False + + +def test_pre_call_abort_subscriptions_stay_serial_by_default() -> None: + assert ParallelRequestLimiter.subscriptions[Hook.pre_call].parallel is False + + +def test_two_handlers_for_one_hook_in_one_class_are_rejected() -> None: + with pytest.raises(InvalidDeclarationError, match="more than one handler for pre_call"): + + class Twice(CustomLoggerV2): + @on(Hook.pre_call, reads={F.messages}) + async def first(self, view: object, ctx: CallContext[None]) -> None: ... + + @on(Hook.pre_call, reads={F.metadata}) + async def second(self, view: object, ctx: CallContext[None]) -> None: ... + + +@pytest.mark.parametrize( + ("hook", "writes", "may_abort", "awaited"), + [ + (Hook.pre_call, frozenset({F.metadata}), False, False), + (Hook.pre_call, frozenset[Field](), True, False), + (Hook.call_ended, frozenset[Field](), False, True), + ], +) +def test_sync_handler_that_blocks_the_request_is_rejected( + hook: Hook, writes: frozenset[Field], may_abort: bool, awaited: bool +) -> None: + with pytest.raises(InvalidDeclarationError, match="handle must be async"): + + class Sync(CustomLoggerV2): + @on(hook, writes=writes, may_abort=may_abort, awaited=awaited) + def handle(self, view: object, ctx: CallContext[None]) -> None: ... + + +def test_sync_handler_is_allowed_for_off_path_reads() -> None: + class SyncLogger(CustomLoggerV2): + @on(Hook.attempt_ended, reads={F.usage}) + def log(self, view: object, ctx: CallContext[None]) -> None: ... + + assert SyncLogger.subscriptions[Hook.attempt_ended].handler == "log" + + +async def test_undecorated_override_keeps_the_parent_subscription_and_runs_the_override() -> None: + override_patch = Patch().set(F.messages, ["masked"]) + + class Masked(PiiMask): + async def mask(self, view: object, ctx: CallContext[object]) -> Patch | None: + return override_patch + + subscription = Masked.subscriptions[Hook.pre_call] + handler: Callable[[object, CallContext[object]], Awaitable[Patch | None]] = getattr(Masked(), subscription.handler) + assert subscription is PiiMask.subscriptions[Hook.pre_call] + assert await handler(None, CallContext(hook=Hook.pre_call, call=None, state=None)) is override_patch + + +def test_redecorated_override_replaces_the_parent_subscription() -> None: + class WiderMask(PiiMask): + @on(Hook.pre_call, reads={F.messages, F.tools}, writes={F.messages, F.tools}, may_abort=True) + async def mask(self, view: object, ctx: CallContext[object]) -> Patch | None: ... + + assert WiderMask.subscriptions[Hook.pre_call].writes == frozenset({F.messages, F.tools}) + assert WiderMask.subscriptions[Hook.post_call_success] is PiiMask.subscriptions[Hook.post_call_success] + + +def test_new_handler_for_an_inherited_hook_replaces_the_parent_handler() -> None: + class Renamed(PiiMask): + @on(Hook.pre_call, reads={F.messages}) + async def inspect_only(self, view: object, ctx: CallContext[object]) -> None: ... + + assert Renamed.subscriptions[Hook.pre_call] == Subscription( + hook=Hook.pre_call, handler="inspect_only", reads=frozenset({F.messages}) + ) + + +def test_redecorating_an_inherited_handler_for_another_hook_is_rejected() -> None: + with pytest.raises(InvalidDeclarationError, match="mask is decorated for pre_route but also inherits"): + + class Moved(PiiMask): + @on(Hook.pre_route, reads={F.messages}) + async def mask(self, view: object, ctx: CallContext[object]) -> Patch | None: ... + + +def test_replacing_an_inherited_handler_with_a_non_callable_is_rejected() -> None: + with pytest.raises(InvalidDeclarationError, match="unmask handles post_call_success but is not callable"): + + class Removed(PiiMask): + unmask = None # pyright: ignore[reportAssignmentType] # the test removes the handler on purpose + + +def test_multiple_inheritance_merges_hooks_and_the_nearest_class_wins() -> None: + class Masking(CustomLoggerV2): + @on(Hook.pre_call, reads={F.messages}, writes={F.messages}) + async def mask(self, view: object, ctx: CallContext[None]) -> Patch | None: ... + + class Auditing(CustomLoggerV2): + @on(Hook.call_ended, reads={F.outcome}) + async def audit(self, view: object, ctx: CallContext[None]) -> None: ... + + @on(Hook.pre_call, reads={F.metadata}) + async def peek(self, view: object, ctx: CallContext[None]) -> None: ... + + class MaskingAudit(Masking, Auditing): + pass + + assert MaskingAudit.subscriptions == { + Hook.pre_call: Masking.subscriptions[Hook.pre_call], + Hook.call_ended: Auditing.subscriptions[Hook.call_ended], + } diff --git a/tests/unit/integrations/custom_logger_v2/test_manifest.py b/tests/unit/integrations/custom_logger_v2/test_manifest.py new file mode 100644 index 00000000000..25bef9f115f --- /dev/null +++ b/tests/unit/integrations/custom_logger_v2/test_manifest.py @@ -0,0 +1,133 @@ +import pytest + +from litellm.integrations.custom_logger_v2 import ( + Abort, + Activation, + CallContext, + CustomLoggerV2, + Field, + Hook, + InvalidDeclarationError, + OptOutBy, + Patch, + Selection, + build_manifest, + on, +) + +F = Field + + +class ContentGuardrail(CustomLoggerV2): + def __init__( + self, + selection: Selection, + enabled_hooks: frozenset[Hook] | None = None, + parallel: bool = False, + ) -> None: + self.selection = selection + self.enabled_hooks = enabled_hooks + self.parallel = parallel + + @on(Hook.pre_call, reads={F.messages}, may_abort=True) + async def check_input(self, view: object, ctx: CallContext[None]) -> Abort | None: ... + + @on(Hook.post_call_success, reads={F.response}, writes={F.response, F.guardrail_info}, may_abort=True) + async def check_output(self, view: object, ctx: CallContext[None]) -> Patch | Abort | None: ... + + @on(Hook.log_prepare, reads={F.standard_logging_payload}, writes={F.standard_logging_payload}) + async def mask_logs(self, view: object, ctx: CallContext[None]) -> Patch | None: ... + + +class PureLogger(CustomLoggerV2): + @on(Hook.attempt_ended, reads={F.standard_logging_payload}) + async def export(self, view: object, ctx: CallContext[None]) -> None: ... + + @on(Hook.call_ended, reads={F.key_identity}, awaited=True) + async def settle(self, view: object, ctx: CallContext[None]) -> None: ... + + +def test_manifest_keeps_only_the_hooks_the_instance_enables() -> None: + manifest = build_manifest(ContentGuardrail(Selection("input-only"), enabled_hooks=frozenset({Hook.pre_call}))) + assert dict(manifest.subscriptions) == {Hook.pre_call: ContentGuardrail.subscriptions[Hook.pre_call]} + + +def test_manifest_defaults_to_every_declared_hook_and_the_class_name() -> None: + manifest = build_manifest(PureLogger()) + assert set(manifest.subscriptions) == {Hook.attempt_ended, Hook.call_ended} + assert manifest.selection == Selection("PureLogger") + + +def test_two_instances_of_one_class_get_independent_manifests() -> None: + first = build_manifest(ContentGuardrail(Selection("a", activation=Activation.opt_in))) + second = build_manifest(ContentGuardrail(Selection("b"), enabled_hooks=frozenset({Hook.log_prepare}))) + assert (first.selection.name, first.selection.activation) == ("a", Activation.opt_in) + assert (second.selection.name, set(second.subscriptions)) == ("b", {Hook.log_prepare}) + + +def test_manifest_is_frozen_against_later_instance_changes() -> None: + guardrail = ContentGuardrail(Selection("before")) + manifest = build_manifest(guardrail) + guardrail.selection = Selection("after") + guardrail.enabled_hooks = frozenset() + assert manifest.selection.name == "before" + assert set(manifest.subscriptions) == {Hook.pre_call, Hook.post_call_success, Hook.log_prepare} + + +def test_enabling_a_hook_the_class_does_not_implement_is_rejected() -> None: + guardrail = ContentGuardrail(Selection("g"), enabled_hooks=frozenset({Hook.pre_call, Hook.during_call})) + with pytest.raises(InvalidDeclarationError, match="enables during_call, which it does not implement"): + build_manifest(guardrail) + + +def test_instance_parallel_flag_applies_to_its_abort_only_subscriptions() -> None: + guardrail = ContentGuardrail( + Selection("g"), enabled_hooks=frozenset({Hook.pre_call, Hook.log_prepare}), parallel=True + ) + manifest = build_manifest(guardrail) + assert manifest.subscriptions[Hook.pre_call].parallel is True + assert manifest.subscriptions[Hook.log_prepare].parallel is False + assert ContentGuardrail.subscriptions[Hook.pre_call].parallel is False + + +def test_instance_parallel_flag_on_an_abortable_writer_is_rejected() -> None: + guardrail = ContentGuardrail(Selection("g"), enabled_hooks=frozenset({Hook.post_call_success}), parallel=True) + with pytest.raises(InvalidDeclarationError, match="parallel is only for"): + build_manifest(guardrail) + + +def test_logging_only_step_can_be_required_while_the_instance_is_opt_in() -> None: + guardrail = ContentGuardrail(Selection("g", activation=Activation.opt_in)) + guardrail.activation_overrides = {Hook.log_prepare: Activation.required} + manifest = build_manifest(guardrail) + assert manifest.subscriptions[Hook.log_prepare].activation is Activation.required + assert manifest.subscriptions[Hook.pre_call].activation is None + assert manifest.selection.activation is Activation.opt_in + + +def test_activation_override_for_a_hook_the_instance_does_not_enable_is_rejected() -> None: + guardrail = ContentGuardrail(Selection("g"), enabled_hooks=frozenset({Hook.pre_call})) + guardrail.activation_overrides = {Hook.log_prepare: Activation.required} + with pytest.raises(InvalidDeclarationError, match="overrides activation for log_prepare"): + build_manifest(guardrail) + + +@pytest.mark.parametrize( + ("enabled_hooks", "selection"), + [ + (frozenset({Hook.pre_call}), Selection("g", honors_no_log=True)), + (frozenset({Hook.pre_call}), Selection("g", opt_out_by=frozenset({OptOutBy.admin, OptOutBy.caller}))), + (frozenset({Hook.log_prepare}), Selection("g", honors_no_log=True)), + ], +) +def test_callers_cannot_switch_off_a_callback_that_writes_or_aborts( + enabled_hooks: frozenset[Hook], selection: Selection +) -> None: + with pytest.raises(InvalidDeclarationError, match="neither no-log nor a caller opt-out may remove it"): + build_manifest(ContentGuardrail(selection, enabled_hooks=enabled_hooks)) + + +def test_read_only_logger_may_honor_no_log_and_caller_opt_out() -> None: + logger = PureLogger() + logger.selection = Selection("exporter", opt_out_by=frozenset(OptOutBy), honors_no_log=True) + assert build_manifest(logger).selection.honors_no_log is True diff --git a/tests/unit/integrations/custom_logger_v2/test_outcomes.py b/tests/unit/integrations/custom_logger_v2/test_outcomes.py new file mode 100644 index 00000000000..f8577c33f88 --- /dev/null +++ b/tests/unit/integrations/custom_logger_v2/test_outcomes.py @@ -0,0 +1,52 @@ +import pytest + +from litellm.integrations.custom_logger_v2 import ( + Abort, + Field, + InvalidPatchError, + MergeField, + Patch, + Respond, + SetField, +) + +F = Field + + +def test_patch_builders_return_new_patches_in_call_order() -> None: + empty = Patch() + patch = empty.set(F.messages, ["masked"]).merge(F.metadata, {"pii": True}) + assert empty.ops == () + assert patch.ops == (SetField(F.messages, ["masked"]), MergeField(F.metadata, {"pii": True})) + assert patch.fields == frozenset({F.messages, F.metadata}) + + +def test_merge_keeps_its_own_copy_of_the_keys() -> None: + keys = {"tag": "a"} + patch = Patch().merge(F.metadata, keys) + keys["tag"] = "b" + assert patch.ops == (MergeField(F.metadata, {"tag": "a"}),) + + +@pytest.mark.parametrize("field", [F.messages, F.response, F.model]) +def test_merge_on_a_non_mapping_field_is_rejected(field: Field) -> None: + with pytest.raises(InvalidPatchError, match=f"{field.value} is not a mapping field"): + Patch().merge(field, {"k": "v"}) + + +def test_abort_carries_a_json_body_and_a_guardrail_info_patch() -> None: + body = {"error": {"message": "blocked", "violations": ["pii"]}} + patch = Patch().merge(F.guardrail_info, {"pii_guard": "intervened"}).merge(F.metadata, {"blocked": True}) + abort = Abort(status=400, message="blocked", body=body, patch=patch) + assert (abort.body, abort.patch) == (body, patch) + + +@pytest.mark.parametrize( + "patch", + [Patch().set(F.messages, []), Patch().merge(F.guardrail_info, {"g": 1}).set(F.response, None)], +) +def test_outcome_patch_outside_guardrail_info_and_metadata_is_rejected(patch: Patch) -> None: + with pytest.raises(InvalidPatchError, match="may only patch guardrail_info and metadata"): + Abort(status=400, message="blocked", patch=patch) + with pytest.raises(InvalidPatchError, match="may only patch guardrail_info and metadata"): + Respond(response={"choices": []}, patch=patch) diff --git a/tests/unit/integrations/custom_logger_v2/test_patches.py b/tests/unit/integrations/custom_logger_v2/test_patches.py new file mode 100644 index 00000000000..d473318e779 --- /dev/null +++ b/tests/unit/integrations/custom_logger_v2/test_patches.py @@ -0,0 +1,39 @@ +from types import MappingProxyType + +import pytest + +from litellm.integrations.custom_logger_v2 import Field, InvalidPatchError, Patch, apply_patch, apply_patches + +F = Field +BASE = MappingProxyType({F.model: "gpt", F.metadata: MappingProxyType({"keep": 1, "replace": 1})}) + + +def test_set_replaces_one_field_and_leaves_the_base_untouched() -> None: + patched = apply_patch(BASE, Patch().set(F.model, "claude")) + + assert patched == {F.model: "claude", F.metadata: {"keep": 1, "replace": 1}} + assert BASE[F.model] == "gpt" + + +def test_merge_keeps_keys_it_does_not_name() -> None: + patched = apply_patch(BASE, Patch().merge(F.metadata, {"replace": 2, "new": 3})) + + assert patched[F.metadata] == {"keep": 1, "replace": 2, "new": 3} + + +def test_merge_into_an_absent_field_starts_from_its_keys() -> None: + patched = apply_patch(BASE, Patch().merge(F.guardrail_info, {"blocked": True})) + + assert patched[F.guardrail_info] == {"blocked": True} + assert len(patched) == 3 + + +def test_patches_apply_in_order_so_the_later_set_wins() -> None: + patched = apply_patches(BASE, (Patch().set(F.model, "first"), Patch().set(F.model, "second"))) + + assert patched[F.model] == "second" + + +def test_merge_into_a_non_mapping_value_is_rejected() -> None: + with pytest.raises(InvalidPatchError): + apply_patch({F.metadata: "text"}, Patch().merge(F.metadata, {"a": 1})) diff --git a/tests/unit/integrations/custom_logger_v2/test_planner.py b/tests/unit/integrations/custom_logger_v2/test_planner.py new file mode 100644 index 00000000000..7c797341022 --- /dev/null +++ b/tests/unit/integrations/custom_logger_v2/test_planner.py @@ -0,0 +1,142 @@ +from collections.abc import Set + +from litellm.integrations.custom_logger_v2 import ( + CallbackManifest, + CustomLoggerV2, + Field, + Hook, + HookPlan, + Selection, + Subscription, + plan_hook, +) +from litellm.integrations.custom_logger_v2.planner import Entry + +F = Field + + +def _entry( + position: int, + name: str, + *, + reads: Set[Field] = frozenset(), + writes: Set[Field] = frozenset(), + may_abort: bool = False, + parallel: bool = False, + awaited: bool = False, + hook: Hook = Hook.pre_call, +) -> Entry: + subscription = Subscription( + hook=hook, + handler=name, + reads=frozenset(reads), + writes=frozenset(writes), + may_abort=may_abort, + parallel=parallel, + awaited=awaited, + ) + manifest = CallbackManifest( + logger=CustomLoggerV2(), + selection=Selection(name), + subscriptions={hook: subscription}, + request_state=None, + ) + return Entry(position, manifest, subscription) + + +def _names(plan: HookPlan) -> tuple[tuple[str, ...], ...]: + return tuple(tuple(entry.subscription.handler for entry in stage) for stage in plan.stages) + + +def test_design_doc_worked_example() -> None: + entries = ( + _entry(0, "pii_mask", reads={F.messages}, writes={F.messages}, may_abort=True), + _entry(1, "tagger", writes={F.metadata}), + _entry(2, "limiter", reads={F.key_identity}, may_abort=True), + _entry(3, "secret_scan", reads={F.messages}, writes={F.messages}, may_abort=True), + _entry(4, "audit", reads={F.messages}), + ) + + plan = plan_hook(entries) + + assert _names(plan) == (("pii_mask", "tagger"), ("limiter",), ("secret_scan",)) + assert plan.off_path == (entries[4],) + assert plan.awaited == () + assert plan.reads == {F.messages, F.key_identity} + + +def test_rule_one_only_writers_and_abortables_block() -> None: + reader = _entry(0, "reader", reads={F.messages}) + awaited = _entry(1, "awaited", reads={F.key_identity}, awaited=True) + writer = _entry(2, "writer", writes={F.metadata}) + + plan = plan_hook((reader, awaited, writer)) + + assert _names(plan) == (("writer",),) + assert plan.off_path == (reader,) + assert plan.awaited == (awaited,) + + +def test_rule_two_every_abortable_waits_for_all_earlier_blocking_entries() -> None: + plan = plan_hook( + ( + _entry(0, "tagger", writes={F.metadata}), + _entry(1, "first", reads={F.key_identity}, may_abort=True), + _entry(2, "second", reads={F.key_identity}, may_abort=True), + ) + ) + + assert _names(plan) == (("tagger",), ("first",), ("second",)) + + +def test_rule_three_writers_wait_only_on_conflicting_earlier_writers() -> None: + plan = plan_hook( + ( + _entry(0, "slow_messages", writes={F.messages}), + _entry(1, "metadata", writes={F.metadata}), + _entry(2, "reads_metadata", reads={F.metadata}, writes={F.guardrail_info}), + _entry(3, "rewrites_messages", reads={F.messages}, writes={F.messages}), + ) + ) + + assert _names(plan) == (("slow_messages", "metadata"), ("reads_metadata", "rewrites_messages")) + + +def test_rule_three_writer_never_lands_before_an_earlier_reader_of_its_field() -> None: + plan = plan_hook( + ( + _entry(0, "tagger", writes={F.metadata}), + _entry(1, "scanner", reads={F.messages}, may_abort=True), + _entry(2, "rewriter", writes={F.messages}), + ) + ) + + assert _names(plan) == (("tagger",), ("scanner", "rewriter")) + + +def test_rule_three_writer_without_conflicts_joins_the_first_stage_past_abortables() -> None: + plan = plan_hook( + ( + _entry(0, "limiter", reads={F.key_identity}, may_abort=True), + _entry(1, "scanner", reads={F.key_identity}, may_abort=True), + _entry(2, "tagger", writes={F.metadata}), + ) + ) + + assert _names(plan) == (("limiter", "tagger"), ("scanner",)) + + +def test_parallel_abort_only_entries_run_together_after_the_serial_chain() -> None: + plan = plan_hook( + ( + _entry(0, "fast_check", reads={F.messages}, may_abort=True, parallel=True), + _entry(1, "masker", reads={F.messages}, writes={F.messages}, may_abort=True), + _entry(2, "slow_check", reads={F.messages}, may_abort=True, parallel=True), + ) + ) + + assert _names(plan) == (("masker",), ("fast_check", "slow_check")) + + +def test_no_entries_plan_nothing() -> None: + assert plan_hook(()) == HookPlan() diff --git a/tests/unit/integrations/custom_logger_v2/test_selection.py b/tests/unit/integrations/custom_logger_v2/test_selection.py new file mode 100644 index 00000000000..092ffc9f57c --- /dev/null +++ b/tests/unit/integrations/custom_logger_v2/test_selection.py @@ -0,0 +1,55 @@ +import pytest + +from litellm.integrations.custom_logger_v2 import Activation, OptOutBy, RequestSelection, Selection, is_selected + +ADMIN = frozenset({OptOutBy.admin}) +BOTH = frozenset({OptOutBy.admin, OptOutBy.caller}) + + +@pytest.mark.parametrize( + ("selection", "activation", "request_selection", "expected"), + [ + (Selection("limiter"), Activation.required, RequestSelection(admin_opt_outs=frozenset({"limiter"})), True), + ( + Selection("limiter", honors_no_log=True, opt_out_by=BOTH), + Activation.required, + RequestSelection(caller_opt_outs=frozenset({"limiter"}), no_log=True), + True, + ), + (Selection("langfuse"), Activation.default_on, RequestSelection(), True), + (Selection("langfuse"), Activation.default_on, RequestSelection(admin_opt_outs=frozenset({"langfuse"})), False), + (Selection("langfuse"), Activation.default_on, RequestSelection(caller_opt_outs=frozenset({"langfuse"})), True), + ( + Selection("langfuse", opt_out_by=BOTH), + Activation.default_on, + RequestSelection(caller_opt_outs=frozenset({"langfuse"})), + False, + ), + ( + Selection("v2_langfuse", aliases=frozenset({"langfuse"})), + Activation.default_on, + RequestSelection(admin_opt_outs=frozenset({"langfuse"})), + False, + ), + (Selection("guard"), Activation.opt_in, RequestSelection(), False), + (Selection("guard"), Activation.opt_in, RequestSelection(requested=frozenset({"guard"})), True), + ( + Selection("guard", opt_out_by=ADMIN), + Activation.opt_in, + RequestSelection(requested=frozenset({"guard"}), admin_opt_outs=frozenset({"guard"})), + False, + ), + (Selection("langfuse", honors_no_log=True), Activation.default_on, RequestSelection(no_log=True), False), + (Selection("langfuse"), Activation.default_on, RequestSelection(no_log=True), True), + ( + Selection("langfuse", honors_no_log=True, opt_out_by=frozenset()), + Activation.default_on, + RequestSelection(no_log=True), + False, + ), + ], +) +def test_is_selected( + selection: Selection, activation: Activation, request_selection: RequestSelection, expected: bool +) -> None: + assert is_selected(selection, activation, request_selection) is expected diff --git a/tests/unit/integrations/test_custom_guardrail_post_call_once.py b/tests/unit/integrations/test_custom_guardrail_post_call_once.py index fb823156929..ff93da50db5 100644 --- a/tests/unit/integrations/test_custom_guardrail_post_call_once.py +++ b/tests/unit/integrations/test_custom_guardrail_post_call_once.py @@ -310,7 +310,6 @@ PROXY_CASES: Final = ( pytest.param( ProxyCase(primary=(), request=("G",), style=InheritedApplyGuardrailG), id="K-I-request-chat-inherited-apply-guardrail", - marks=_known_gap("AI-2"), ), pytest.param(ProxyCase(primary=("other",), default_on=True), id="L-default-on"), ) diff --git a/tests/unit/integrations/test_custom_logger.py b/tests/unit/integrations/test_custom_logger.py index 0f1d5c79bcc..19ff7482d3c 100644 --- a/tests/unit/integrations/test_custom_logger.py +++ b/tests/unit/integrations/test_custom_logger.py @@ -25,7 +25,7 @@ import pytest_asyncio import litellm from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE -from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_logger import CustomLogger, subclass_overrides_hook from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -497,3 +497,48 @@ def setup_and_teardown(): if hasattr(litellm, "in_memory_llm_clients_cache"): litellm.in_memory_llm_clients_cache.flush_cache() yield + + +class _DefinesPreCallHook(CustomLogger): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + return data + + +class _InheritsPreCallHookFromBase(_DefinesPreCallHook): + pass + + +class _PreCallHookMixin: + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + return data + + +class _LoggerBeforeMixin(CustomLogger, _PreCallHookMixin): + pass + + +@pytest.mark.parametrize( + ("cls", "expected"), + [ + (CustomLogger, False), + (_DefinesPreCallHook, True), + (_InheritsPreCallHookFromBase, True), + (_LoggerBeforeMixin, False), + ], +) +def test_subclass_overrides_hook_walks_mro_up_to_custom_logger(cls, expected): + assert subclass_overrides_hook(cls, "async_pre_call_hook") is expected + + +def test_subclass_overrides_hook_stops_at_the_given_base(): + from litellm.integrations.custom_guardrail import CustomGuardrail + + class _AppliesGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + class _InheritsApplyGuardrail(_AppliesGuardrail): + pass + + assert subclass_overrides_hook(CustomGuardrail, "apply_guardrail", CustomGuardrail) is False + assert subclass_overrides_hook(_InheritsApplyGuardrail, "apply_guardrail", CustomGuardrail) is True diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 81166caa094..806e8ecd7f8 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -14,7 +14,7 @@ from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence from datetime import datetime as datetime_standard_logging, datetime as datetime_unit_test, datetime as dt_object from importlib.machinery import ModuleSpec from types import MappingProxyType, ModuleType -from typing import Final, Literal +from typing import Final, Literal, TypeAlias from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -40,6 +40,8 @@ from litellm.constants import ( ) from litellm.cost_calculator import ocr_batch_cost from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_logger_v2 import Abort, CallContext, CustomLoggerV2, Field, Hook, Patch, on +from litellm.integrations.custom_logger_v2.host import V2Outcome from litellm.llms.custom_llm import CustomLLM, CustomLLMError from litellm.litellm_core_utils import litellm_logging from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging @@ -61,7 +63,9 @@ from litellm.proxy.hooks.cache_control_check import PROXY_CacheControlCheck from litellm.proxy.hooks.max_iterations_limiter import PROXY_MaxIterationsHandler from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.utils import InternalUsageCache -from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse +from litellm.types.mcp import MCPPostCallResponseObject +from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse +from litellm.types.prompts.init_prompts import PromptSpec from litellm.types.rerank import RerankResponse from litellm.types.utils import ( CallTypes, @@ -69,6 +73,7 @@ from litellm.types.utils import ( ImageResponse, LiteLLMRealtimeStreamLoggingObject, ModelResponse, + StandardCallbackDynamicParams, StandardLoggingHiddenParams, StandardLoggingMetadata, StandardLoggingModelInformation, @@ -15195,3 +15200,492 @@ async def test_every_callback_runs_under_the_span_an_sdk_caller_passed( await asyncio.wait_for(asyncio.to_thread(logger.logged.wait, 10), timeout=10) assert (logger.seen.get(hook), current_request_parent_otel_span()) == (span, None) + + +_PromptResult: TypeAlias = tuple[str, list[AllMessageValues], dict[str, object]] + + +class _Inert(CustomLoggerV2): + @on(Hook.provider_request_sent) + def provider_request_sent(self, view: object, ctx: CallContext[object]) -> None: + return None + + @on(Hook.provider_response_received) + def provider_response_received(self, view: object, ctx: CallContext[object]) -> None: + return None + + @on(Hook.log_prepare) + def log_prepare(self, view: object, ctx: CallContext[object]) -> None: + return None + + @on(Hook.attempt_ended) + def attempt_ended(self, view: object, ctx: CallContext[object]) -> None: + return None + + @on(Hook.mcp_tool_result) + def mcp_tool_result(self, view: object, ctx: CallContext[object]) -> None: + return None + + +class _Observer(CustomLoggerV2): + def __init__(self, journal: list[object]) -> None: + self.journal = journal + self.done = threading.Event() + + def _note(self) -> None: + self.journal.append("v2") + self.done.set() + + @on(Hook.provider_request_sent) + def provider_request_sent(self, view: object, ctx: CallContext[object]) -> None: + self._note() + + @on(Hook.provider_response_received) + def provider_response_received(self, view: object, ctx: CallContext[object]) -> None: + self._note() + + @on(Hook.attempt_ended, awaited=True) + async def attempt_ended(self, view: object, ctx: CallContext[object]) -> None: + self._note() + + +class _PayloadStamp(CustomLoggerV2): + def __init__(self) -> None: + self.seen: list[object] = [] + + @on(Hook.log_prepare, reads={Field.standard_logging_payload}, writes={Field.standard_logging_payload}) + async def log_prepare(self, view: object, ctx: CallContext[object]) -> Patch | None: + payload: Final[object] = getattr(view, "standard_logging_payload", None) + if not isinstance(payload, dict): + return None + self.seen.append(payload.get("stamped_by")) + return Patch().set(Field.standard_logging_payload, {**payload, "stamped_by": "v2"}) + + +class _ToolGate(CustomLoggerV2): + def __init__(self, journal: list[object], outcome: Abort | None = None) -> None: + self.journal = journal + self.outcome = outcome + + @on(Hook.mcp_tool_result, may_abort=True) + async def mcp_tool_result(self, view: object, ctx: CallContext[object]) -> Abort | None: + self.journal.append("v2") + return self.outcome + + +class _Entry(CustomLogger): + def __init__(self, name: str, journal: list[object], error: Exception | None = None) -> None: + super().__init__() + self.name = name + self.journal = journal + self.error = error + self.kwargs: list[object] = [] + + def _note(self, event: str, detail: object, kwargs: object) -> None: + self.journal.append((self.name, event, detail)) + self.kwargs.append(kwargs) + if self.error is not None: + raise self.error + + def log_pre_api_call(self, model: object, messages: object, kwargs: object) -> None: + self._note("pre", None, kwargs) + + def log_post_api_call(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self._note("post", None, kwargs) + + def logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + self.journal.append((self.name, "hook", getattr(result, "id", None))) + return {**kwargs, f"hooked_by_{self.name}": True}, ModelResponse(id=self.name) + + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + payload: Final = kwargs.get("standard_logging_object") + self.journal.append((self.name, "hook", payload.get("stamped_by") if isinstance(payload, dict) else None)) + if self.error is not None: + raise self.error + if isinstance(payload, dict): + payload["stamped_by"] = self.name + return {**kwargs, f"hooked_by_{self.name}": True}, ModelResponse(id=self.name) + + def log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._note("success", getattr(response_obj, "id", None), kwargs) + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._note("success", getattr(response_obj, "id", None), kwargs) + + def log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._note("failure", None, kwargs) + + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._note("failure", None, kwargs) + + async def async_post_mcp_tool_call_hook( + self, + kwargs: dict[str, object], + response_obj: MCPPostCallResponseObject, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> MCPPostCallResponseObject | None: + content: Final = response_obj.mcp_tool_call_response + self.journal.append((self.name, "mcp", [getattr(item, "text", None) for item in content])) + response_obj.mcp_tool_call_response = [TextContent(type="text", text=f"{self.name} edit")] + return response_obj + + +def _plane_lists(monkeypatch: pytest.MonkeyPatch, **lists: Sequence[object]) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + for name in ("input_callback", "success_callback", "failure_callback"): + monkeypatch.setattr(litellm, name, list(lists.get(name, ()))) + for name in ("_async_success_callback", "_async_failure_callback"): + monkeypatch.setattr(litellm, name, list(lists.get(name, ()))) + + +def _with_inert(callbacks: Sequence[object], inert: bool) -> list[object]: + return [*callbacks, _Inert()] if inert else list(callbacks) + + +def _ran(journal: Sequence[object], event: str) -> list[object]: + return [entry for entry in journal if isinstance(entry, tuple) and entry[1] == event] + + +def _fire_provider_event(logging_obj: LitellmLogging, hook: Hook) -> None: + logging_obj.update_environment_variables(litellm_params={}, optional_params={}) + if hook is Hook.provider_request_sent: + logging_obj.pre_call(input="", api_key="") + else: + logging_obj.post_call(original_response="ok", input="", api_key="") + + +@pytest.mark.parametrize("inert", [False, True]) +@pytest.mark.parametrize( + ("hook", "event"), [(Hook.provider_request_sent, "pre"), (Hook.provider_response_received, "post")] +) +def test_provider_events_reach_every_input_callback_in_order_past_a_failing_one( + logging_obj: LitellmLogging, monkeypatch: pytest.MonkeyPatch, hook: Hook, event: str, inert: bool +) -> None: + journal: Final[list[object]] = [] + failing: Final = _Entry("failing", journal, error=RuntimeError("logger down")) + second: Final = _Entry("second", journal) + _plane_lists(monkeypatch, input_callback=_with_inert([failing, second], inert)) + + _fire_provider_event(logging_obj, hook) + + assert journal == [("failing", event, None), ("second", event, None)] + assert failing.kwargs[0] is logging_obj.model_call_details + assert second.kwargs[0] is logging_obj.model_call_details + + +@pytest.mark.parametrize("hook", [Hook.provider_request_sent, Hook.provider_response_received]) +def test_a_v2_callback_input_callback_observes_provider_events_after_the_legacy_ones( + logging_obj: LitellmLogging, monkeypatch: pytest.MonkeyPatch, hook: Hook +) -> None: + journal: Final[list[object]] = [] + observer: Final = _Observer(journal) + _plane_lists(monkeypatch, input_callback=[_Entry("first", journal), observer, _Entry("second", journal)]) + + _fire_provider_event(logging_obj, hook) + + assert observer.done.wait(timeout=30) + assert [entry[0] if isinstance(entry, tuple) else entry for entry in journal] == ["first", "second", "v2"] + + +@pytest.mark.parametrize("inert", [False, True]) +def test_sync_success_chains_logging_hooks_and_delivers_past_a_failing_logger( + logging_obj: LitellmLogging, monkeypatch: pytest.MonkeyPatch, inert: bool +) -> None: + journal: Final[list[object]] = [] + function_calls: Final[list[object]] = [] + + def function_callback(kwargs: object, response: object, start_time: object, end_time: object) -> None: + function_calls.append(getattr(response, "id", None)) + journal.append(("function", "success", getattr(response, "id", None))) + + first: Final = _Entry("first", journal) + failing: Final = _Entry("failing", journal, error=RuntimeError("logger down")) + last: Final = _Entry("last", journal) + _plane_lists(monkeypatch, success_callback=_with_inert([first, function_callback, failing, last], inert)) + logging_obj.stream = False + logging_obj.update_environment_variables(litellm_params={}, optional_params={}) + now: Final = datetime.datetime.now() + + logging_obj.success_handler(result=ModelResponse(id="served"), start_time=now, end_time=now) + + assert _ran(journal, "hook") == [ + ("first", "hook", "served"), + ("failing", "hook", "first"), + ("last", "hook", "failing"), + ] + assert _ran(journal, "success") == [ + ("first", "success", "last"), + ("function", "success", "last"), + ("failing", "success", "last"), + ("last", "success", "last"), + ] + delivered: Final = last.kwargs[0] + assert isinstance(delivered, dict) + assert all(delivered.get(f"hooked_by_{name}") for name in ("first", "failing", "last")) + assert logging_obj.model_call_details.get("hooked_by_last") is True + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_async_success_chains_logging_hooks_past_a_failing_hook_and_delivers_in_order( + logging_obj: LitellmLogging, monkeypatch: pytest.MonkeyPatch, inert: bool +) -> None: + journal: Final[list[object]] = [] + + async def function_callback(kwargs: object, response: object, start_time: object, end_time: object) -> None: + journal.append(("function", "success", getattr(response, "id", None))) + + first: Final = _Entry("first", journal) + failing: Final = _Entry("failing", journal, error=RuntimeError("hook down")) + last: Final = _Entry("last", journal) + _plane_lists(monkeypatch, _async_success_callback=_with_inert([first, failing, function_callback, last], inert)) + logging_obj.stream = False + logging_obj.update_environment_variables(litellm_params={"acompletion": True}, optional_params={}) + now: Final = datetime.datetime.now() + + await logging_obj.async_success_handler(result=ModelResponse(id="served"), start_time=now, end_time=now) + + assert [entry[0] for entry in _ran(journal, "hook")] == ["first", "failing", "last"] + assert _ran(journal, "success") == [ + ("first", "success", "last"), + ("failing", "success", "last"), + ("function", "success", "last"), + ("last", "success", "last"), + ] + assert logging_obj.model_call_details.get("hooked_by_first") is True + assert "hooked_by_failing" not in logging_obj.model_call_details + assert logging_obj.model_call_details.get("hooked_by_last") is True + + +async def test_a_v2_callback_between_two_async_logging_hooks_sees_the_first_edit_and_the_second_sees_its_patch( + logging_obj: LitellmLogging, monkeypatch: pytest.MonkeyPatch +) -> None: + journal: Final[list[object]] = [] + stamp: Final = _PayloadStamp() + _plane_lists(monkeypatch, _async_success_callback=[_Entry("first", journal), stamp, _Entry("second", journal)]) + logging_obj.stream = False + logging_obj.update_environment_variables(litellm_params={"acompletion": True}, optional_params={}) + now: Final = datetime.datetime.now() + + await logging_obj.async_success_handler(result=ModelResponse(id="served"), start_time=now, end_time=now) + + assert stamp.seen == ["first"] + assert _ran(journal, "hook") == [("first", "hook", None), ("second", "hook", "v2")] + payload: Final[object] = logging_obj.model_call_details["standard_logging_object"] + assert isinstance(payload, dict) + assert payload["stamped_by"] == "second" + + +async def test_an_attempt_ended_v2_callback_between_two_async_loggers_runs_after_both_deliveries( + logging_obj: LitellmLogging, monkeypatch: pytest.MonkeyPatch +) -> None: + journal: Final[list[object]] = [] + observer: Final = _Observer(journal) + _plane_lists(monkeypatch, _async_success_callback=[_Entry("first", journal), observer, _Entry("second", journal)]) + logging_obj.stream = False + logging_obj.update_environment_variables(litellm_params={"acompletion": True}, optional_params={}) + now: Final = datetime.datetime.now() + + await logging_obj.async_success_handler(result=ModelResponse(id="served"), start_time=now, end_time=now) + + assert [ + entry[0] if isinstance(entry, tuple) else entry for entry in journal if entry not in _ran(journal, "hook") + ] == [ + "first", + "second", + "v2", + ] + + +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("inert", [False, True]) +async def test_failure_handlers_deliver_to_every_callback_in_order_past_a_failing_one( + logging_obj: LitellmLogging, monkeypatch: pytest.MonkeyPatch, use_async: bool, inert: bool +) -> None: + journal: Final[list[object]] = [] + failing: Final = _Entry("failing", journal, error=RuntimeError("logger down")) + second: Final = _Entry("second", journal) + callbacks: Final = _with_inert([failing, second], inert) + if use_async: + _plane_lists(monkeypatch, _async_failure_callback=callbacks) + else: + _plane_lists(monkeypatch, failure_callback=callbacks) + logging_obj.stream = False + logging_obj.update_environment_variables(litellm_params={}, optional_params={}) + error: Final = ValueError("provider failed") + + if use_async: + await logging_obj.async_failure_handler(error, "trace") + else: + await asyncio.to_thread(logging_obj.failure_handler, error, "trace") + + assert journal == [("failing", "failure", None), ("second", "failure", None)] + assert second.kwargs[0] is logging_obj.model_call_details + assert logging_obj.model_call_details["exception"] is error + + +async def test_an_attempt_ended_v2_callback_between_two_async_failure_loggers_runs_after_both( + logging_obj: LitellmLogging, monkeypatch: pytest.MonkeyPatch +) -> None: + journal: Final[list[object]] = [] + observer: Final = _Observer(journal) + _plane_lists(monkeypatch, _async_failure_callback=[_Entry("first", journal), observer, _Entry("second", journal)]) + logging_obj.update_environment_variables(litellm_params={}, optional_params={}) + + await logging_obj.async_failure_handler(ValueError("provider failed"), "trace") + + assert journal == [("first", "failure", None), ("second", "failure", None), "v2"] + + +def _tool_result() -> CallToolResult: + return CallToolResult(content=[TextContent(type="text", text="original")]) + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_mcp_tool_result_callbacks_each_see_the_previous_edit( + logging_obj: LitellmLogging, monkeypatch: pytest.MonkeyPatch, inert: bool +) -> None: + journal: Final[list[object]] = [] + _plane_lists( + monkeypatch, success_callback=_with_inert([_Entry("first", journal), _Entry("second", journal)], inert) + ) + result: Final = _tool_result() + now: Final = datetime.datetime.now() + + returned: Final = await logging_obj.async_post_mcp_tool_call_hook({}, result, now, now) + + assert returned is result + assert journal == [("first", "mcp", ["original"]), ("second", "mcp", ["first edit"])] + assert [getattr(item, "text", None) for item in result.content] == ["second edit"] + + +async def test_an_mcp_tool_result_v2_callback_between_two_loggers_runs_between_them_and_its_abort_stops_the_rest( + logging_obj: LitellmLogging, monkeypatch: pytest.MonkeyPatch +) -> None: + journal: Final[list[object]] = [] + _plane_lists( + monkeypatch, success_callback=[_Entry("first", journal), _ToolGate(journal), _Entry("second", journal)] + ) + now: Final = datetime.datetime.now() + + await logging_obj.async_post_mcp_tool_call_hook({}, _tool_result(), now, now) + + assert journal == [("first", "mcp", ["original"]), "v2", ("second", "mcp", ["first edit"])] + abort: Final = Abort(status=403, message="tool result blocked") + journal.clear() + _plane_lists( + monkeypatch, success_callback=[_Entry("first", journal), _ToolGate(journal, abort), _Entry("second", journal)] + ) + + with pytest.raises(V2Outcome) as raised: + await logging_obj.async_post_mcp_tool_call_hook({}, _tool_result(), now, now) + + assert raised.value.outcome is abort + assert journal == [("first", "mcp", ["original"]), "v2"] + + +class _PromptManager(CustomLogger): + def __init__(self, resolved: _PromptResult, error: Exception | None = None) -> None: + super().__init__() + self.resolved = resolved + self.error = error + self.prompt_ids: list[str | None] = [] + + def get_chat_completion_prompt( + self, + model: str, + messages: list[AllMessageValues], + non_default_params: dict[str, object], + prompt_id: str | None, + prompt_variables: dict[str, object] | None, + dynamic_callback_params: StandardCallbackDynamicParams, + prompt_spec: PromptSpec | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, + ignore_prompt_manager_model: bool | None = False, + ignore_prompt_manager_optional_params: bool | None = False, + ) -> _PromptResult: + self.prompt_ids.append(prompt_id) + if self.error is not None: + raise self.error + return self.resolved + + async def async_get_chat_completion_prompt( + self, + model: str, + messages: list[AllMessageValues], + non_default_params: dict[str, object], + prompt_id: str | None, + prompt_variables: dict[str, object] | None, + dynamic_callback_params: StandardCallbackDynamicParams, + litellm_logging_obj: LitellmLogging, + prompt_spec: PromptSpec | None = None, + tools: list[dict[str, object]] | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, + ignore_prompt_manager_model: bool | None = False, + ignore_prompt_manager_optional_params: bool | None = False, + ) -> _PromptResult: + return self.get_chat_completion_prompt( + model, messages, non_default_params, prompt_id, prompt_variables, dynamic_callback_params + ) + + +async def _resolve_prompt(logging_obj: LitellmLogging, manager: CustomLogger, use_async: bool) -> _PromptResult: + messages: Final[list[AllMessageValues]] = [{"role": "user", "content": "raw"}] + if use_async: + return await logging_obj.async_get_chat_completion_prompt( + model="requested-model", + messages=messages, + non_default_params={}, + prompt_variables=None, + prompt_id="greeting", + prompt_management_logger=manager, + ) + return await asyncio.to_thread( + logging_obj.get_chat_completion_prompt, + model="requested-model", + messages=messages, + non_default_params={}, + prompt_variables=None, + prompt_id="greeting", + prompt_management_logger=manager, + ) + + +@pytest.mark.parametrize("use_async", [False, True]) +async def test_prompt_resolve_returns_the_managers_prompt_unchanged( + logging_obj: LitellmLogging, use_async: bool +) -> None: + resolved: Final[_PromptResult] = ("rendered-model", [{"role": "system", "content": "rendered"}], {"seed": 1}) + manager: Final = _PromptManager(resolved) + + model, messages, params = await _resolve_prompt(logging_obj, manager, use_async) + + assert model == resolved[0] + assert messages is resolved[1] + assert params is resolved[2] + assert logging_obj.messages is resolved[1] + assert manager.prompt_ids == ["greeting"] + + +@pytest.mark.parametrize("use_async", [False, True]) +async def test_prompt_resolve_propagates_the_managers_error(logging_obj: LitellmLogging, use_async: bool) -> None: + error: Final = LookupError("unknown prompt") + + with pytest.raises(LookupError) as raised: + await _resolve_prompt(logging_obj, _PromptManager(("m", [], {}), error), use_async) + + assert raised.value is error diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py index 54b76756975..a17e74a6eef 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py @@ -3,7 +3,7 @@ import copy import json import os import uuid -from collections.abc import Mapping +from collections.abc import Iterator, Mapping from dataclasses import dataclass from pathlib import Path from typing import Any, Dict, Final, List @@ -23,6 +23,8 @@ import litellm from litellm.llms.anthropic.messages.native import handler as python_messages from litellm.anthropic_interface import messages from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_logger_v2 import CallContext, CustomLoggerV2, Field, Hook, on +from litellm.llms.anthropic.messages.native.handler import _execute_pre_request_hooks from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -2996,3 +2998,136 @@ async def test_azure_messages_replay_encrypted_reasoning_only_to_responses( else: assert route.calls.last.request.url.path.endswith("/chat/completions") assert "qa-opaque-reasoning" not in json.dumps(body) + + +@pytest.mark.parametrize("with_v2", [False, True]) +async def test_messages_send_the_request_the_pre_request_hook_returned( + monkeypatch: pytest.MonkeyPatch, with_v2: bool +) -> None: + seen: Final[list[str]] = [] + + class MockRewriter(CustomLogger): + async def async_pre_request_hook( + self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object] + ) -> dict[str, object] | None: + seen.append("legacy") + return {**kwargs, "mock_response": "rewritten by the pre-request hook"} + + class Gate(CustomLoggerV2): + @on(Hook.pre_deployment_call, may_abort=True) + async def gate(self, view: object, ctx: CallContext[object]) -> None: + seen.append("v2") + + monkeypatch.setattr(litellm, "callbacks", [*([Gate()] if with_v2 else []), MockRewriter()]) + + response: Final = await messages.acreate( + model="anthropic/claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + max_tokens=10, + mock_response="original", + ) + + assert isinstance(response, dict) + assert response["content"][0]["text"] == "rewritten by the pre-request hook" + assert seen == (["v2", "v2", "legacy"] if with_v2 else ["legacy"]) + + +class _PreRequestHook(CustomLogger): + def __init__(self, name: str, journal: list[object], reply: object = None, error: Exception | None = None) -> None: + super().__init__() + self.name = name + self.journal = journal + self.reply = reply + self.error = error + self.seen: list[object] = [] + + async def async_pre_request_hook( + self, model: str, messages: list[dict[str, object]], kwargs: dict[str, object] + ) -> dict[str, object] | None: + self.journal.append(self.name) + self.seen.append(kwargs) + if self.error is not None: + raise self.error + return self.reply if not callable(self.reply) else self.reply(kwargs) + + +class _ToolsWatcher(CustomLoggerV2): + def __init__(self, journal: list[object]) -> None: + self.journal = journal + + @on(Hook.pre_deployment_call, reads={Field.tools}, may_abort=True) + async def watch(self, view: object, ctx: CallContext[object]) -> None: + self.journal.append(("v2", getattr(view, "tools"))) # noqa: B009 # the view's declared field + + +async def _pre_request_hooks() -> dict[str, object]: + return await _execute_pre_request_hooks( + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + tools=[{"name": "original"}], + stream=False, + custom_llm_provider="anthropic", + ) + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_pre_request_hooks_run_in_order_and_a_returned_dict_replaces_the_request( + monkeypatch: pytest.MonkeyPatch, inert: bool +) -> None: + journal: Final[list[object]] = [] + replaced: Final[dict[str, object]] = {"tools": [{"name": "replaced"}]} + first: Final = _PreRequestHook("first", journal, reply=replaced) + keeper: Final = _PreRequestHook("keeper", journal) + last: Final = _PreRequestHook("last", journal) + monkeypatch.setattr(litellm, "callbacks", [first, keeper, last, *([_ToolsWatcher([])] if inert else [])]) + + result: Final = await _pre_request_hooks() + + assert journal == ["first", "keeper", "last"] + assert first.seen[0] == { + "tools": [{"name": "original"}], + "stream": False, + "litellm_params": {"custom_llm_provider": "anthropic"}, + } + assert keeper.seen[0] is replaced + assert result is replaced + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_a_pre_request_hook_error_propagates_unchanged_and_stops_the_hook( + monkeypatch: pytest.MonkeyPatch, inert: bool +) -> None: + journal: Final[list[object]] = [] + error: Final = ValueError("request refused") + monkeypatch.setattr( + litellm, + "callbacks", + [ + _PreRequestHook("first", journal, error=error), + _PreRequestHook("next", journal), + *([_ToolsWatcher([])] if inert else []), + ], + ) + + with pytest.raises(ValueError, match="request refused") as raised: + await _pre_request_hooks() + + assert raised.value is error + assert journal == ["first"] + + +async def test_a_v2_callback_between_pre_request_hooks_sees_the_first_rewrite(monkeypatch: pytest.MonkeyPatch) -> None: + journal: Final[list[object]] = [] + monkeypatch.setattr( + litellm, + "callbacks", + [ + _PreRequestHook("first", journal, reply=lambda kwargs: {**kwargs, "tools": [{"name": "rewritten"}]}), + _ToolsWatcher(journal), + _PreRequestHook("second", journal), + ], + ) + + await _pre_request_hooks() + + assert journal == ["first", ("v2", [{"name": "rewritten"}]), "second"] diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index b51abfa1ad0..b15bd9326a7 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -5,7 +5,7 @@ import json import logging import threading import time -from collections.abc import Awaitable, Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, Mock, patch @@ -26,6 +26,7 @@ from litellm.integrations.code_interpreter_interception.handler import ( LITELLM_CODE_EXECUTION_TOOL_NAME, ) from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_logger_v2 import CallContext, CustomLoggerV2, Field, Hook, on from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig from litellm.llms.azure.vector_stores.transformation import AzureOpenAIVectorStoreConfig @@ -73,6 +74,7 @@ from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoo from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( + CallTypes, ImageObject, ImageResponse, ModelResponse, @@ -6902,3 +6904,67 @@ async def test_a_stream_answered_with_one_json_body_is_not_replayed_for_a_provid ) with pytest.raises(litellm.exceptions.MidStreamFallbackError, match="before any terminal event"): _ = tuple([event async for event in stream]) + + +@pytest.mark.parametrize("with_v2", [False, True]) +def test_sync_responses_send_the_params_the_pre_deployment_hook_returned( + monkeypatch: pytest.MonkeyPatch, with_v2: bool +) -> None: + seen: Final[list[str]] = [] + + class TemperatureRewriter(CustomLogger): + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict[str, object] | None: + seen.append("legacy") + return {**kwargs, "temperature": 0.25} + + class Gate(CustomLoggerV2): + @on(Hook.pre_deployment_call, may_abort=True) + async def gate(self, view: object, ctx: CallContext[object]) -> None: + seen.append("v2") + + monkeypatch.setattr(litellm, "callbacks", [*([Gate()] if with_v2 else []), TemperatureRewriter()]) + logging_obj: Final = Mock() + logging_obj.dynamic_success_callbacks = [] + + _, _, _, response_params, _ = BaseLLMHTTPHandler()._run_sync_responses_pre_call_deployment_hook( + model="gpt-5.5", + input="hi", + custom_llm_provider="openai", + response_api_optional_request_params={"temperature": 1.0}, + litellm_params=GenericLiteLLMParams(), + logging_obj=logging_obj, + ) + + assert response_params["temperature"] == 0.25 + assert seen == (["v2", "legacy"] if with_v2 else ["legacy"]) + + +@pytest.mark.parametrize("with_v2", [False, True]) +def test_sync_responses_run_a_v2_callback_pre_deployment_hook_without_any_legacy_override( + monkeypatch: pytest.MonkeyPatch, with_v2: bool +) -> None: + seen: Final[list[object]] = [] + + class Gate(CustomLoggerV2): + @on(Hook.pre_deployment_call, reads={Field.model}, may_abort=True) + async def gate(self, view: object, ctx: CallContext[object]) -> None: + seen.append(getattr(view, "model")) # noqa: B009 # the view's declared field + + monkeypatch.setattr(litellm, "callbacks", [CustomLogger(), *([Gate()] if with_v2 else [])]) + logging_obj: Final = Mock() + logging_obj.dynamic_success_callbacks = [] + params: Final[dict[str, object]] = {"temperature": 1.0} + + _, _, _, response_params, _ = BaseLLMHTTPHandler()._run_sync_responses_pre_call_deployment_hook( + model="gpt-5.5", + input="hi", + custom_llm_provider="openai", + response_api_optional_request_params=params, + litellm_params=GenericLiteLLMParams(), + logging_obj=logging_obj, + ) + + assert seen == (["gpt-5.5"] if with_v2 else []) + assert response_params["temperature"] == 1.0 diff --git a/tests/unit/proxy/test_proxy_logging_hook_detection.py b/tests/unit/proxy/test_proxy_logging_hook_detection.py index afcb89010fe..7e34bf6741f 100644 --- a/tests/unit/proxy/test_proxy_logging_hook_detection.py +++ b/tests/unit/proxy/test_proxy_logging_hook_detection.py @@ -1295,3 +1295,60 @@ async def test_parallel_post_call_guardrails_keep_native_hook_when_opted_out(mon assert opted_out.native_hooks_ran == ["post_call"] assert routed.native_hooks_ran == [] + + +class _TagsRequestBeforeCall(CustomLogger): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + data["tagged_by"] = type(self).__name__ + return data + + +class _InheritsPreCallHook(_TagsRequestBeforeCall): + pass + + +@pytest.mark.asyncio +async def test_pre_call_hook_runs_hook_inherited_from_intermediate_base(monkeypatch): + monkeypatch.setattr(litellm, "callbacks", [_InheritsPreCallHook()]) + + result = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), + data={"model": "gpt-5", "messages": [{"role": "user", "content": "hi"}]}, + call_type="acompletion", + ) + + assert result["tagged_by"] == "_InheritsPreCallHook" + + +class _RewritesSuccessResponse(CustomLogger): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + return {"rewritten_by": type(self).__name__} + + +class _InheritsSuccessHook(_RewritesSuccessResponse): + pass + + +@pytest.mark.asyncio +async def test_post_call_success_hook_awaits_only_loggers_whose_class_overrides_it(monkeypatch): + from litellm.types.utils import Choices, Message, ModelResponse + + calls: list[str] = [] + + async def _instance_level_hook(data, user_api_key_dict, response): + calls.append("instance") + return {"rewritten_by": "instance"} + + default_logger = CustomLogger() + default_logger.async_post_call_success_hook = _instance_level_hook + monkeypatch.setattr(litellm, "callbacks", [_InheritsSuccessHook(), default_logger]) + response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]) + + result = await ProxyLogging(user_api_key_cache=DualCache()).post_call_success_hook( + data={"messages": [{"role": "user", "content": "hi"}]}, + response=response, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), + ) + + assert calls == [] + assert result == {"rewritten_by": "_InheritsSuccessHook"} diff --git a/tests/unit/proxy/test_v2_hook_routing.py b/tests/unit/proxy/test_v2_hook_routing.py new file mode 100644 index 00000000000..b904ba754a3 --- /dev/null +++ b/tests/unit/proxy/test_v2_hook_routing.py @@ -0,0 +1,1212 @@ +import asyncio +from collections.abc import Awaitable, Callable, Generator, Mapping, Sequence +from typing import Final + +import pytest +from fastapi import HTTPException +from mcp.types import CallToolResult, TextContent + +import litellm +from litellm.caching import DualCache +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger +from litellm.exceptions import RejectedRequestError, SensitiveDataRouteException +from litellm.integrations.custom_logger_v2 import ( + Abort, + CallContext, + CustomLoggerV2, + Field, + Hook, + Patch, + on, +) +from litellm.integrations.custom_logger_v2.host import V2Outcome, reset_host_for_tests +from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep +from litellm.types.utils import CallTypesLiteral, GenericGuardrailAPIInputs, ModelResponse + +AUTH: Final = UserAPIKeyAuth(api_key="sk-test") + + +@pytest.fixture(autouse=True) +def fresh_host() -> Generator[None]: + reset_host_for_tests() + yield + reset_host_for_tests() + + +@pytest.fixture +def proxy_logging() -> ProxyLogging: + return ProxyLogging(user_api_key_cache=DualCache()) + + +class _Legacy(CustomLogger): + def __init__(self, order: list[str], reply: object = None, error: Exception | None = None) -> None: + super().__init__() + self.order = order + self.reply = reply + self.error = error + self.seen: list[object] = [] + + def _record(self, seen: object) -> object: + self.order.append("legacy") + self.seen.append(seen) + if self.error is not None: + raise self.error + return self.reply + + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral + ) -> object: + return self._record(data) + + async def async_moderation_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral + ) -> object: + return self._record(data) + + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: object + ) -> object: + return self._record(response) + + async def async_post_call_failure_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> object: + return self._record(original_exception) + + async def async_post_call_response_headers_hook( + self, + data: dict[str, object], + user_api_key_dict: UserAPIKeyAuth, + response: object, + request_headers: dict[str, str] | None = None, + litellm_call_info: dict[str, object] | None = None, + ) -> object: + return self._record(response) + + async def async_filter_listed_models( + self, user_api_key_dict: UserAPIKeyAuth, model_names: Sequence[str] + ) -> Sequence[str]: + self._record(tuple(model_names)) + return ("kept",) + + +class _StreamGuardrail(CustomGuardrail): + def __init__(self, order: list[str], reply: object) -> None: + super().__init__(guardrail_name="stream-witness", default_on=True, event_hook=GuardrailEventHooks.post_call) + self.order = order + self.reply = reply + + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: object + ) -> object: + self.order.append("legacy") + return self.reply + + +class _SuccessDispatch: + def __init__(self) -> None: + self.model_call_details: dict[str, object] = {"metadata": {}} + self.dispatched: list[object] = [] + + async def dispatch_success_handlers(self, result: object, **kwargs: object) -> None: + self.dispatched.append(result) + + +async def test_pre_call_passthrough_hands_the_request_to_legacy_once_and_returns_its_reply( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + order: Final[list[str]] = [] + reply: Final[dict[str, object]] = {"model": "replaced", "metadata": {}} + legacy: Final = _Legacy(order, reply=reply) + monkeypatch.setattr(litellm, "callbacks", [legacy]) + data: Final[dict[str, object]] = {"model": "gpt-4o", "metadata": {}} + + result: Final = await proxy_logging.pre_call_hook(user_api_key_dict=AUTH, data=data, call_type="completion") + + assert result is reply + assert legacy.seen == [data] + assert legacy.seen[0] is data + + +async def test_pre_call_passthrough_propagates_the_legacy_exception( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + error: Final = ValueError("rejected by a legacy hook") + monkeypatch.setattr(litellm, "callbacks", [_Legacy([], error=error)]) + + with pytest.raises(ValueError, match="rejected by a legacy hook") as raised: + await proxy_logging.pre_call_hook(user_api_key_dict=AUTH, data={"model": "gpt-4o"}, call_type="completion") + + assert raised.value is error + + +async def test_during_call_passthrough_returns_the_request_and_propagates_the_legacy_exception( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + order: Final[list[str]] = [] + legacy: Final = _Legacy(order) + monkeypatch.setattr(litellm, "callbacks", [legacy]) + data: Final = {"model": "gpt-4o"} + + result: Final = await proxy_logging.during_call_hook(data=data, user_api_key_dict=AUTH, call_type="completion") + + assert result is data + assert legacy.seen == [data] + error: Final = RuntimeError("moderation failed") + monkeypatch.setattr(litellm, "callbacks", [_Legacy([], error=error)]) + with pytest.raises(RuntimeError) as raised: + await proxy_logging.during_call_hook(data=data, user_api_key_dict=AUTH, call_type="completion") + assert raised.value is error + + +async def test_post_call_success_passthrough_returns_the_legacy_replacement( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + replacement: Final = ModelResponse() + legacy: Final = _Legacy([], reply=replacement) + monkeypatch.setattr(litellm, "callbacks", [legacy]) + response: Final = ModelResponse() + + result: Final = await proxy_logging.post_call_success_hook( + data={"model": "gpt-4o"}, response=response, user_api_key_dict=AUTH + ) + + assert result is replacement + assert legacy.seen == [response] + assert legacy.seen[0] is response + + +async def test_post_call_success_passthrough_propagates_the_legacy_exception( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + error: Final = ValueError("response rejected") + monkeypatch.setattr(litellm, "callbacks", [_Legacy([], error=error)]) + + with pytest.raises(ValueError, match="response rejected") as raised: + await proxy_logging.post_call_success_hook( + data={"model": "gpt-4o"}, response=ModelResponse(), user_api_key_dict=AUTH + ) + + assert raised.value is error + + +async def test_post_call_failure_passthrough_returns_the_first_legacy_http_exception( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + transformed: Final = HTTPException(status_code=418, detail="rewritten") + legacy: Final = _Legacy([], reply=transformed) + monkeypatch.setattr(litellm, "callbacks", [legacy, _Legacy([], reply=HTTPException(status_code=500))]) + original: Final = RuntimeError("upstream failed") + + result: Final = await proxy_logging.post_call_failure_hook( + request_data={"model": "gpt-4o"}, original_exception=original, user_api_key_dict=AUTH + ) + + assert result is transformed + assert legacy.seen == [original] + assert legacy.seen[0] is original + + +async def test_response_headers_passthrough_returns_the_legacy_headers( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + legacy: Final = _Legacy([], reply={"x-legacy": "1"}) + monkeypatch.setattr(litellm, "callbacks", [legacy]) + response: Final = ModelResponse() + + result: Final = await proxy_logging.post_call_response_headers_hook( + data={"model": "gpt-4o"}, user_api_key_dict=AUTH, response=response + ) + + assert result == {"x-legacy": "1"} + assert legacy.seen == [response] + + +async def test_list_models_passthrough_returns_the_names_legacy_hides( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + legacy: Final = _Legacy([]) + monkeypatch.setattr(litellm, "callbacks", [legacy]) + + result: Final = await proxy_logging.hidden_by_listing_callbacks(AUTH, ["kept", "hidden"]) + + assert result == frozenset({"hidden"}) + assert legacy.seen == [("kept", "hidden")] + + +async def test_post_mcp_call_passthrough_returns_the_result_untouched_without_guardrails( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + response: Final = CallToolResult(content=[TextContent(type="text", text="tool output")]) + + result: Final = await proxy_logging.post_mcp_call_hook(response=response, request_data={}, user_api_key_dict=AUTH) + + assert result is response + + +async def test_deferred_stream_guardrails_passthrough_dispatches_the_guardrail_replacement( + monkeypatch: pytest.MonkeyPatch, +) -> None: + order: Final[list[str]] = [] + replacement: Final = ModelResponse() + monkeypatch.setattr(litellm, "callbacks", [_StreamGuardrail(order, replacement)]) + dispatch: Final = _SuccessDispatch() + + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "gpt-4o", "metadata": {}}, + captured_user_api_key_dict=AUTH, + captured_logging_obj=dispatch, + assembled_response=ModelResponse(), + cache_hit=None, + ) + await _drain() + + assert order == ["legacy"] + assert dispatch.dispatched == [replacement] + assert dispatch.dispatched[0] is replacement + + +async def _drain() -> None: + for _ in range(3): + await asyncio.sleep(0) + + +Action = Callable[[dict[str, object]], Awaitable[None] | Exception | str | dict[str, object] | None] + + +def _passes(data: dict[str, object]) -> None: + return None + + +class _PreCall(CustomLogger): + def __init__(self, name: str, journal: list[object], action: Action = _passes, enforces: bool = False) -> None: + super().__init__() + self.name = name + self.journal = journal + self.action = action + self.enforces_request_content = enforces + self.seen: list[dict[str, object]] = [] + + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral + ) -> Exception | str | dict[str, object] | None: + self.journal.append(self.name) + self.seen.append(data) + return await _settled(self.action(data)) + + async def async_moderation_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral + ) -> object: + self.journal.append(self.name) + await asyncio.sleep(0) + return await _settled(self.action(data)) + + +class _Guard(CustomGuardrail): + def __init__( + self, + name: str, + journal: list[object], + action: Action = _passes, + parallel: bool = False, + delay: float = 0.0, + event: GuardrailEventHooks = GuardrailEventHooks.pre_call, + ) -> None: + super().__init__(guardrail_name=name, default_on=True, event_hook=event, run_in_parallel=parallel) + self.name = name + self.journal = journal + self.action = action + self.delay = delay + self.seen_models: list[object] = [] + + async def _run(self, data: dict[str, object]) -> Exception | str | dict[str, object] | None: + self.journal.append(self.name) + self.seen_models.append(data.get("model")) + await asyncio.sleep(self.delay) + return await _settled(self.action(data)) + + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral + ) -> Exception | str | dict[str, object] | None: + return await self._run(data) + + async def async_moderation_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral + ) -> object: + return await self._run(data) + + +class _V2Callback(CustomLoggerV2): + def __init__(self, journal: list[object], replacement: object = None, outcome: Abort | None = None) -> None: + self.journal = journal + self.replacement = replacement + self.outcome = outcome + + @on(Hook.pre_call, reads={Field.messages}, writes={Field.messages}, may_abort=True) + async def pre_call(self, view: object, ctx: CallContext[object]) -> Patch | Abort | None: + self.journal.append(("v2", getattr(view, "messages", None))) + if self.outcome is not None: + return self.outcome + return None if self.replacement is None else Patch().set(Field.messages, self.replacement) + + @on(Hook.during_call, reads={Field.messages}, may_abort=True) + async def during_call(self, view: object, ctx: CallContext[object]) -> Abort | None: + self.journal.append("v2") + await asyncio.sleep(0) + return self.outcome + + +async def _settled( + value: Awaitable[None] | Exception | str | dict[str, object] | None, +) -> Exception | str | dict[str, object] | None: + return await value if isinstance(value, Awaitable) else value + + +def _reroute(name: str) -> Action: + def action(data: dict[str, object]) -> None: + raise SensitiveDataRouteException( + route_to_model=name, session_id="s", guardrail_name=name, sticky_session_routing=False + ) + + return action + + +def _raise(error: BaseException) -> Action: + def action(data: dict[str, object]) -> None: + raise error + + return action + + +def _returning(value: str | dict[str, object] | None) -> Action: + def action(data: dict[str, object]) -> str | dict[str, object] | None: + return value + + return action + + +def _with_inert_v2(callbacks: list[object], inert: bool) -> list[object]: + return [*callbacks, _V2Callback([])] if inert else callbacks + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_pre_call_runs_sequential_entries_in_list_order_then_the_parallel_guardrails( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging, inert: bool +) -> None: + journal: Final[list[object]] = [] + callbacks: Final[list[object]] = [ + _Guard("parallel-1", journal, parallel=True), + _PreCall("logger-1", journal), + _Guard("sequential", journal), + _Guard("parallel-2", journal, parallel=True), + _PreCall("logger-2", journal), + ] + monkeypatch.setattr(litellm, "callbacks", _with_inert_v2(callbacks, inert)) + + await proxy_logging.pre_call_hook(user_api_key_dict=AUTH, data={"model": "m"}, call_type="completion") + + assert journal == ["logger-1", "sequential", "logger-2", "parallel-1", "parallel-2"] + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_pre_call_dict_replaces_the_request_and_none_keeps_it( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging, inert: bool +) -> None: + replacement: Final[dict[str, object]] = {"model": "replaced"} + first: Final = _PreCall("first", [], _returning(replacement)) + keeper: Final = _PreCall("keeper", []) + last: Final = _PreCall("last", []) + monkeypatch.setattr(litellm, "callbacks", _with_inert_v2([first, keeper, last], inert)) + data: Final[dict[str, object]] = {"model": "m"} + + result: Final = await proxy_logging.pre_call_hook(user_api_key_dict=AUTH, data=data, call_type="completion") + + assert first.seen[0] is data + assert keeper.seen[0] is replacement + assert last.seen[0] is replacement + assert result is replacement + + +@pytest.mark.parametrize(("call_type", "error"), [("completion", RejectedRequestError), ("embeddings", HTTPException)]) +async def test_pre_call_string_rejects_and_stops_the_hook( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging, call_type: CallTypesLiteral, error: type[Exception] +) -> None: + journal: Final[list[object]] = [] + monkeypatch.setattr( + litellm, "callbacks", [_PreCall("rejecter", journal, _returning("no thanks")), _PreCall("after", journal)] + ) + + with pytest.raises(error): + await proxy_logging.pre_call_hook(user_api_key_dict=AUTH, data={"model": "m"}, call_type=call_type) + + assert journal == ["rejecter"] + + +async def test_pre_call_sequential_exception_propagates_unchanged_and_stops_the_hook( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + journal: Final[list[object]] = [] + error: Final = HTTPException(status_code=403, detail="blocked") + monkeypatch.setattr(litellm, "callbacks", [_Guard("blocker", journal, _raise(error)), _PreCall("after", journal)]) + + with pytest.raises(HTTPException) as raised: + await proxy_logging.pre_call_hook(user_api_key_dict=AUTH, data={"model": "m"}, call_type="completion") + + assert raised.value is error + assert journal == ["blocker"] + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_parallel_guardrails_pick_the_first_block_in_list_order_and_a_block_beats_a_reroute( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging, inert: bool +) -> None: + journal: Final[list[object]] = [] + slow_error: Final = ValueError("slow but first") + callbacks: Final[list[object]] = [ + _Guard("reroute", journal, _reroute("onprem"), parallel=True), + _Guard("slow", journal, _raise(slow_error), parallel=True, delay=0.02), + _Guard("fast", journal, _raise(HTTPException(status_code=400)), parallel=True), + ] + monkeypatch.setattr(litellm, "callbacks", _with_inert_v2(callbacks, inert)) + metadata: Final[dict[str, object]] = {} + data: Final[dict[str, object]] = {"model": "m", "metadata": metadata} + + with pytest.raises(ValueError, match="slow but first") as raised: + await proxy_logging.pre_call_hook(user_api_key_dict=AUTH, data=data, call_type="completion") + + assert raised.value is slow_error + assert data["model"] == "m" + assert metadata["applied_guardrails"] == ["slow", "fast"] + + +async def test_parallel_reroute_alone_reroutes_the_request( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + journal: Final[list[object]] = [] + monkeypatch.setattr(litellm, "callbacks", [_Guard("reroute", journal, _reroute("onprem"), parallel=True)]) + + result: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=AUTH, data={"model": "m", "metadata": {}}, call_type="completion" + ) + + assert result["model"] == "onprem" + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_deferred_reroute_lets_later_guardrails_run_and_the_first_reroute_wins( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging, inert: bool +) -> None: + journal: Final[list[object]] = [] + later: Final = _Guard("later", journal) + parallel: Final = _Guard("parallel", journal, parallel=True) + callbacks: Final[list[object]] = [ + _Guard("first-reroute", journal, _reroute("first")), + _Guard("second-reroute", journal, _reroute("second")), + later, + parallel, + ] + monkeypatch.setattr(litellm, "callbacks", _with_inert_v2(callbacks, inert)) + + result: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=AUTH, data={"model": "m", "metadata": {}}, call_type="completion" + ) + + assert journal == ["first-reroute", "second-reroute", "later", "parallel"] + assert later.seen_models == ["m"] + assert parallel.seen_models == ["first"] + assert result["model"] == "first" + assert result["metadata"]["sensitive_data_routing_original_model"] == "m" + + +async def test_a_block_after_a_deferred_reroute_still_propagates( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + error: Final = HTTPException(status_code=400, detail="blocked") + monkeypatch.setattr( + litellm, + "callbacks", + [_Guard("reroute", [], _reroute("onprem")), _Guard("blocker", [], _raise(error))], + ) + data: Final[dict[str, object]] = {"model": "m", "metadata": {}} + + with pytest.raises(HTTPException) as raised: + await proxy_logging.pre_call_hook(user_api_key_dict=AUTH, data=data, call_type="completion") + + assert raised.value is error + assert data["model"] == "m" + + +async def test_a_pipeline_managed_guardrail_runs_in_its_pipeline_and_not_again_in_the_hook( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + journal: Final[list[object]] = [] + monkeypatch.setattr(litellm, "callbacks", [_Guard("managed", journal), _Guard("unmanaged", journal)]) + pipeline: Final = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="managed", on_pass="next")]) + data: Final[dict[str, object]] = { + "model": "m", + "messages": [], + "metadata": {"_guardrail_pipelines": [("p", pipeline)]}, + } + + await proxy_logging.pre_call_hook(user_api_key_dict=AUTH, data=data, call_type="completion") + + assert journal == ["managed", "unmanaged"] + + +@pytest.mark.parametrize( + ("skip_guardrails", "guardrails_only", "expected"), + [ + (False, False, ["logger", "guard", ("v2", None), "enforcer"]), + (True, False, ["logger", ("v2", None), "enforcer"]), + (False, True, ["guard", "enforcer"]), + ], +) +async def test_hook_point_modes_choose_which_entries_run( + monkeypatch: pytest.MonkeyPatch, + proxy_logging: ProxyLogging, + skip_guardrails: bool, + guardrails_only: bool, + expected: list[object], +) -> None: + journal: Final[list[object]] = [] + callbacks: Final[list[object]] = [ + _PreCall("logger", journal), + _Guard("guard", journal), + _V2Callback(journal), + _PreCall("enforcer", journal, enforces=True), + ] + monkeypatch.setattr(litellm, "callbacks", callbacks) + + await proxy_logging.pre_call_hook( + user_api_key_dict=AUTH, + data={"model": "m"}, + call_type="completion", + skip_guardrails=skip_guardrails, + guardrails_only=guardrails_only, + ) + + assert journal == expected + + +async def test_a_runtime_change_to_the_callback_list_reaches_the_next_request( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + journal: Final[list[object]] = [] + first: Final = _PreCall("first", journal) + callbacks: Final[list[object]] = [first] + monkeypatch.setattr(litellm, "callbacks", callbacks) + + await proxy_logging.pre_call_hook(user_api_key_dict=AUTH, data={"model": "m"}, call_type="completion") + callbacks.append(_PreCall("added", journal)) + await proxy_logging.pre_call_hook(user_api_key_dict=AUTH, data={"model": "m"}, call_type="completion") + callbacks.remove(first) + await proxy_logging.pre_call_hook(user_api_key_dict=AUTH, data={"model": "m"}, call_type="completion") + + assert journal == ["first", "first", "added", "added"] + + +async def test_a_v2_callback_between_two_legacy_loggers_sees_the_first_edit_and_the_second_sees_its_patch( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + journal: Final[list[object]] = [] + + def edit(data: dict[str, object]) -> None: + data["messages"] = ["edited by the first"] + return None + + seen_by_second: Final[list[object]] = [] + + def record(data: dict[str, object]) -> None: + seen_by_second.append(data["messages"]) + + first: Final = _PreCall("first", [], edit) + second: Final = _PreCall("second", [], record) + monkeypatch.setattr( + litellm, "callbacks", [first, _V2Callback(journal, replacement=["patched by the v2 callback"]), second] + ) + data: Final[dict[str, object]] = {"model": "m", "messages": ["original"]} + + result: Final = await proxy_logging.pre_call_hook(user_api_key_dict=AUTH, data=data, call_type="completion") + + assert journal == [("v2", ["edited by the first"])] + assert seen_by_second == [["patched by the v2 callback"]] + assert result is data + assert data["messages"] == ["patched by the v2 callback"] + + +async def test_a_v2_callback_abort_stops_the_hook_before_later_legacy_loggers( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + journal: Final[list[object]] = [] + abort: Final = Abort(status=429, message="slow down") + monkeypatch.setattr( + litellm, + "callbacks", + [_PreCall("before", journal), _V2Callback(journal, outcome=abort), _PreCall("after", journal)], + ) + + with pytest.raises(V2Outcome) as raised: + await proxy_logging.pre_call_hook(user_api_key_dict=AUTH, data={"model": "m"}, call_type="completion") + + assert raised.value.outcome is abort + assert journal == ["before", ("v2", None)] + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_during_call_starts_every_entry_in_list_order_and_returns_on_the_first_failure( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging, inert: bool +) -> None: + journal: Final[list[object]] = [] + release: Final = asyncio.Event() + error: Final = RuntimeError("moderation failed") + + async def hold(data: dict[str, object]) -> None: + await release.wait() + + slow: Final = _Guard("slow", journal, hold, event=GuardrailEventHooks.during_call) + callbacks: Final[list[object]] = [ + slow, + _PreCall("moderator", journal, _raise(error)), + _Guard("guard", journal, event=GuardrailEventHooks.during_call), + ] + monkeypatch.setattr(litellm, "callbacks", _with_inert_v2(callbacks, inert)) + + with pytest.raises(RuntimeError) as raised: + await asyncio.wait_for( + proxy_logging.during_call_hook(data={"model": "m"}, user_api_key_dict=AUTH, call_type="completion"), 1 + ) + + release.set() + await _drain() + assert raised.value is error + assert journal[:3] == ["slow", "moderator", "guard"] + + +async def test_during_call_v2_callback_between_legacy_runs_with_them_and_its_abort_wins_by_list_order( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + journal: Final[list[object]] = [] + abort: Final = Abort(status=400, message="flagged") + callbacks: Final[list[object]] = [ + _PreCall("first", journal), + _V2Callback(journal, outcome=abort), + _PreCall("second", journal), + ] + monkeypatch.setattr(litellm, "callbacks", callbacks) + + with pytest.raises(V2Outcome) as raised: + await proxy_logging.during_call_hook(data={"model": "m"}, user_api_key_dict=AUTH, call_type="completion") + + assert raised.value.outcome is abort + assert journal == ["first", "v2", "second"] + + +class _Inert(CustomLoggerV2): + @on(Hook.post_call_success, may_abort=True) + async def post_call_success(self, view: object, ctx: CallContext[object]) -> None: + return None + + @on(Hook.post_call_failure, may_abort=True) + async def post_call_failure(self, view: object, ctx: CallContext[object]) -> None: + return None + + @on(Hook.response_headers, writes={Field.response_headers}) + async def response_headers(self, view: object, ctx: CallContext[object]) -> None: + return None + + @on(Hook.list_models, may_abort=True) + async def list_models(self, view: object, ctx: CallContext[object]) -> None: + return None + + @on(Hook.mcp_tool_result, may_abort=True) + async def mcp_tool_result(self, view: object, ctx: CallContext[object]) -> None: + return None + + +def _with_inert(callbacks: list[object], inert: bool) -> list[object]: + return [*callbacks[:1], _Inert(), *callbacks[1:]] if inert else callbacks + + +def _read(view: object, name: str) -> object: + return getattr(view, name) # pyright: ignore[reportAny] # a view exposes each declared field as an attribute + + +class _Patcher(CustomLoggerV2): + def __init__(self, replacement: object = None, outcome: Abort | None = None) -> None: + self.replacement = replacement + self.outcome = outcome + self.seen: list[object] = [] + + def _answer(self, field: Field, current: object) -> Patch | Abort | None: + self.seen.append(current) + if self.outcome is not None: + return self.outcome + return None if self.replacement is None else Patch().set(field, self.replacement) + + @on(Hook.post_call_success, reads={Field.response}, writes={Field.response}, may_abort=True) + async def post_call_success(self, view: object, ctx: CallContext[object]) -> Patch | Abort | None: + return self._answer(Field.response, _read(view, "response")) + + @on(Hook.post_call_failure, reads={Field.error}, writes={Field.error}, may_abort=True) + async def post_call_failure(self, view: object, ctx: CallContext[object]) -> Patch | Abort | None: + return self._answer(Field.error, _read(view, "error")) + + @on(Hook.list_models, reads={Field.model_list}, writes={Field.model_list}, may_abort=True) + async def list_models(self, view: object, ctx: CallContext[object]) -> Patch | Abort | None: + return self._answer(Field.model_list, _read(view, "model_list")) + + @on(Hook.mcp_tool_result, reads={Field.tool_result}, writes={Field.tool_result}, may_abort=True) + async def mcp_tool_result(self, view: object, ctx: CallContext[object]) -> Patch | Abort | None: + return self._answer(Field.tool_result, _read(view, "tool_result")) + + +class _HeaderMerger(CustomLoggerV2): + def __init__(self, headers: dict[str, str]) -> None: + self.headers = headers + self.seen: list[dict[str, object]] = [] + + @on(Hook.response_headers, reads={Field.response_headers}, writes={Field.response_headers}) + async def response_headers(self, view: object, ctx: CallContext[object]) -> Patch: + headers: Final = _read(view, "response_headers") + assert isinstance(headers, Mapping) + self.seen.append(dict(headers)) + return Patch().merge(Field.response_headers, self.headers) + + +class _PostGuard(CustomGuardrail): + def __init__( + self, + name: str, + journal: list[object], + reply: object = None, + error: BaseException | None = None, + parallel: bool = False, + delay: float = 0.0, + ) -> None: + super().__init__( + guardrail_name=name, default_on=True, event_hook=GuardrailEventHooks.post_call, run_in_parallel=parallel + ) + self.name = name + self.journal = journal + self.reply = reply + self.error = error + self.delay = delay + self.seen: list[object] = [] + + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: object + ) -> object: + self.journal.append(self.name) + self.seen.append(response) + await asyncio.sleep(self.delay) + if self.error is not None: + raise self.error + self.journal.append(f"{self.name} done") + return self.reply + + +class _PostLogger(CustomLogger): + def __init__(self, name: str, journal: list[object], reply: object = None) -> None: + super().__init__() + self.name = name + self.journal = journal + self.reply = reply + self.seen: list[object] = [] + + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: object + ) -> object: + self.journal.append(self.name) + self.seen.append(response) + return self.reply + + +def _flow_change(name: str) -> SensitiveDataRouteException: + return SensitiveDataRouteException( + route_to_model=name, session_id="s", guardrail_name=name, sticky_session_routing=False + ) + + +async def _post_call(proxy_logging: ProxyLogging, response: ModelResponse) -> object: + return await proxy_logging.post_call_success_hook( + data={"model": "m", "metadata": {}}, response=response, user_api_key_dict=AUTH + ) + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_post_call_success_runs_guardrails_then_parallel_guardrails_then_loggers_and_chains_replies( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging, inert: bool +) -> None: + journal: Final[list[object]] = [] + from_guardrail: Final = ModelResponse() + from_logger: Final = ModelResponse() + parallel: Final = _PostGuard("parallel", journal, reply=ModelResponse(), parallel=True) + second_logger: Final = _PostLogger("logger-2", journal) + callbacks: Final[list[object]] = [ + _PostLogger("logger-1", journal, reply=from_logger), + _PostGuard("guard-1", journal, reply=from_guardrail), + parallel, + CustomLogger(), + second_logger, + _PostGuard("guard-2", journal), + ] + monkeypatch.setattr(litellm, "callbacks", _with_inert(callbacks, inert)) + + result: Final = await _post_call(proxy_logging, ModelResponse()) + + assert [entry for entry in journal if isinstance(entry, str) and not entry.endswith("done")] == [ + "guard-1", + "guard-2", + "parallel", + "logger-1", + "logger-2", + ] + assert parallel.seen[0] is from_guardrail + assert second_logger.seen[0] is from_logger + assert result is from_logger + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_post_call_success_parallel_guardrails_wait_for_each_other_and_the_first_block_in_list_order_wins( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging, inert: bool +) -> None: + journal: Final[list[object]] = [] + slow_block: Final = ValueError("slow but first") + callbacks: Final[list[object]] = [ + _PostGuard("flow", journal, error=_flow_change("flow"), parallel=True), + _PostGuard("slow", journal, error=slow_block, parallel=True, delay=0.02), + _PostGuard("fast", journal, error=HTTPException(status_code=400), parallel=True), + _PostGuard("passes", journal, parallel=True, delay=0.03), + _PostLogger("logger", journal), + ] + monkeypatch.setattr(litellm, "callbacks", _with_inert(callbacks, inert)) + + with pytest.raises(ValueError, match="slow but first") as raised: + await _post_call(proxy_logging, ModelResponse()) + + assert raised.value is slow_block + assert "passes done" in journal + assert "logger" not in journal + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_post_call_success_parallel_flow_change_propagates_only_when_nothing_blocked( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging, inert: bool +) -> None: + journal: Final[list[object]] = [] + first: Final = _flow_change("first") + callbacks: Final[list[object]] = [ + _PostGuard("first", journal, error=first, parallel=True, delay=0.02), + _PostGuard("second", journal, error=_flow_change("second"), parallel=True), + _PostLogger("logger", journal), + ] + monkeypatch.setattr(litellm, "callbacks", _with_inert(callbacks, inert)) + + with pytest.raises(SensitiveDataRouteException) as raised: + await _post_call(proxy_logging, ModelResponse()) + + assert raised.value is first + assert "logger" not in journal + + +async def test_post_call_success_v2_callback_between_two_loggers_sees_the_first_reply_and_the_second_sees_its_patch( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + from_first: Final = ModelResponse() + from_v2: Final = ModelResponse() + v2_callback: Final = _Patcher(replacement=from_v2) + second: Final = _PostLogger("second", []) + monkeypatch.setattr(litellm, "callbacks", [_PostLogger("first", [], reply=from_first), v2_callback, second]) + + result: Final = await _post_call(proxy_logging, ModelResponse()) + + assert v2_callback.seen[0] is from_first + assert second.seen[0] is from_v2 + assert result is from_v2 + + +async def test_post_call_success_v2_callback_abort_stops_the_loggers_after_it( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + journal: Final[list[object]] = [] + abort: Final = Abort(status=400, message="unsafe output") + monkeypatch.setattr( + litellm, + "callbacks", + [_PostLogger("before", journal), _Patcher(outcome=abort), _PostLogger("after", journal)], + ) + + with pytest.raises(V2Outcome) as raised: + await _post_call(proxy_logging, ModelResponse()) + + assert raised.value.outcome is abort + assert journal == ["before"] + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_post_call_failure_calls_every_logger_and_the_first_http_exception_wins( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging, inert: bool +) -> None: + order: Final[list[str]] = [] + raised_first: Final = HTTPException(status_code=429, detail="raised first") + callbacks: Final[list[object]] = [ + _Legacy(order, error=RuntimeError("swallowed")), + "not-a-known-integration", + _Legacy(order, reply="not an exception"), + _Legacy(order, error=raised_first), + _passes, + _Legacy(order, reply=HTTPException(status_code=418, detail="returned later")), + ] + monkeypatch.setattr(litellm, "callbacks", _with_inert(callbacks, inert)) + + result: Final = await proxy_logging.post_call_failure_hook( + request_data={"model": "m"}, original_exception=RuntimeError("upstream"), user_api_key_dict=AUTH + ) + + assert result is raised_first + assert order == ["legacy", "legacy", "legacy", "legacy"] + + +async def test_post_call_failure_v2_callback_between_two_loggers_rewrites_the_error_the_second_sees( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + original: Final = RuntimeError("upstream") + rewritten: Final = RuntimeError("rewritten by the v2 callback") + first: Final = _Legacy([]) + v2_callback: Final = _Patcher(replacement=rewritten) + second: Final = _Legacy([]) + monkeypatch.setattr(litellm, "callbacks", [first, v2_callback, second]) + + result: Final = await proxy_logging.post_call_failure_hook( + request_data={"model": "m"}, original_exception=original, user_api_key_dict=AUTH + ) + + assert result is None + assert first.seen[0] is original + assert v2_callback.seen[0] is original + assert second.seen[0] is rewritten + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_response_headers_merge_in_list_order_and_a_failure_keeps_what_was_merged( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging, inert: bool +) -> None: + order: Final[list[str]] = [] + callbacks: Final[list[object]] = [ + _Legacy(order, reply={"x-a": "1", "x-b": "1"}), + CustomLogger(), + _Legacy(order, reply={"x-b": "2"}), + _Legacy(order, error=RuntimeError("header hook failed")), + _Legacy(order, reply={"x-c": "never"}), + ] + monkeypatch.setattr(litellm, "callbacks", _with_inert(callbacks, inert)) + + result: Final = await proxy_logging.post_call_response_headers_hook( + data={"model": "m"}, user_api_key_dict=AUTH, response=ModelResponse() + ) + + assert result == {"x-a": "1", "x-b": "2"} + assert order == ["legacy", "legacy", "legacy"] + + +async def test_response_headers_v2_callback_between_two_loggers_adds_headers_and_later_entries_win( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + v2_callback: Final = _HeaderMerger({"x-b": "v2", "x-n": "v2"}) + monkeypatch.setattr( + litellm, + "callbacks", + [_Legacy([], reply={"x-a": "1", "x-b": "1"}), v2_callback, _Legacy([], reply={"x-n": "2"})], + ) + + result: Final = await proxy_logging.post_call_response_headers_hook( + data={"model": "m"}, user_api_key_dict=AUTH, response=ModelResponse() + ) + + assert v2_callback.seen == [{"x-a": "1", "x-b": "1"}] + assert result == {"x-a": "1", "x-b": "v2", "x-n": "2"} + + +class _Lister(CustomLogger): + def __init__(self, journal: list[object], hides: frozenset[str] = frozenset(), reply: object = None) -> None: + super().__init__() + self.journal = journal + self.hides = hides + self.reply = reply + + async def async_filter_listed_models( + self, user_api_key_dict: UserAPIKeyAuth, model_names: Sequence[str] + ) -> Sequence[str]: + self.journal.append(tuple(model_names)) + if self.reply is not None: + return self.reply # pyright: ignore[reportReturnType] # a malformed return is the case under test + return tuple(name for name in model_names if name not in self.hides) + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_list_models_filters_see_the_previous_output_and_stop_once_nothing_is_left( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging, inert: bool +) -> None: + journal: Final[list[object]] = [] + callbacks: Final[list[object]] = [ + _Lister(journal, hides=frozenset({"b"})), + CustomLogger(), + _Lister(journal, hides=frozenset({"a", "c"})), + _Lister(journal), + ] + monkeypatch.setattr(litellm, "callbacks", _with_inert(callbacks, inert)) + + hidden: Final = await proxy_logging.hidden_by_listing_callbacks(AUTH, ["a", "b", "c"]) + + assert hidden == frozenset({"a", "b", "c"}) + assert journal == [("a", "b", "c"), ("a", "c")] + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_list_models_malformed_return_raises_500_and_stops_the_filters( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging, inert: bool +) -> None: + journal: Final[list[object]] = [] + monkeypatch.setattr(litellm, "callbacks", _with_inert([_Lister(journal, reply=42), _Lister(journal)], inert)) + + with pytest.raises(ProxyException) as raised: + await proxy_logging.hidden_by_listing_callbacks(AUTH, ["a"]) + + assert raised.value.code == "500" + assert journal == [("a",)] + + +async def test_list_models_v2_callback_between_two_filters_narrows_what_the_second_sees( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + journal: Final[list[object]] = [] + v2_callback: Final = _Patcher(replacement=("c",)) + monkeypatch.setattr(litellm, "callbacks", [_Lister(journal, hides=frozenset({"a"})), v2_callback, _Lister(journal)]) + + hidden: Final = await proxy_logging.hidden_by_listing_callbacks(AUTH, ["a", "b", "c"]) + + assert v2_callback.seen == [("b", "c")] + assert journal == [("a", "b", "c"), ("c",)] + assert hidden == frozenset({"a", "b"}) + + +class _McpSuffix(CustomGuardrail): + def __init__(self, suffix: str) -> None: + super().__init__(guardrail_name=f"mcp{suffix}", default_on=True, event_hook=GuardrailEventHooks.post_mcp_call) + self.suffix = suffix + self.seen: list[list[str]] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: str, + logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + texts: Final = list(inputs.get("texts") or []) + self.seen.append(texts) + return GenericGuardrailAPIInputs(texts=[f"{text}{self.suffix}" for text in texts]) + + +def _tool_result(text: str) -> CallToolResult: + return CallToolResult(content=[TextContent(type="text", text=text)]) + + +def _texts(result: CallToolResult) -> list[str]: + return [item.text for item in result.content if isinstance(item, TextContent)] + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_post_mcp_call_guardrails_run_in_list_order_on_the_previous_result( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging, inert: bool +) -> None: + second: Final = _McpSuffix("-2") + monkeypatch.setattr(litellm, "callbacks", _with_inert([_McpSuffix("-1"), CustomLogger(), second], inert)) + + result: Final = await proxy_logging.post_mcp_call_hook( + response=_tool_result("out"), request_data={"metadata": {}}, user_api_key_dict=AUTH + ) + + assert second.seen == [["out-1"]] + assert _texts(result) == ["out-1-2"] + + +async def test_post_mcp_call_v2_callback_between_two_guardrails_replaces_the_result_the_second_scans( + monkeypatch: pytest.MonkeyPatch, proxy_logging: ProxyLogging +) -> None: + replacement: Final = _tool_result("from v2") + v2_callback: Final = _Patcher(replacement=replacement) + second: Final = _McpSuffix("-2") + monkeypatch.setattr(litellm, "callbacks", [_McpSuffix("-1"), v2_callback, second]) + + result: Final = await proxy_logging.post_mcp_call_hook( + response=_tool_result("out"), request_data={"metadata": {}}, user_api_key_dict=AUTH + ) + + seen_by_v2: Final = v2_callback.seen[0] + assert isinstance(seen_by_v2, CallToolResult) + assert _texts(seen_by_v2) == ["out-1"] + assert second.seen == [["from v2"]] + assert result is replacement + assert _texts(result) == ["from v2-2"] + + +async def _deferred(dispatch: _SuccessDispatch) -> None: + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "m", "metadata": {}}, + captured_user_api_key_dict=AUTH, + captured_logging_obj=dispatch, + assembled_response=ModelResponse(), + cache_hit=None, + ) + await _drain() + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_deferred_stream_guardrails_run_one_at_a_time_and_a_block_is_logged_then_success_is_dispatched( + monkeypatch: pytest.MonkeyPatch, inert: bool +) -> None: + journal: Final[list[object]] = [] + replacement: Final = ModelResponse() + last: Final = _PostGuard("last", journal) + callbacks: Final[list[object]] = [ + _PostGuard("parallel", journal, reply=replacement, parallel=True, delay=0.01), + _PostGuard("blocker", journal, error=HTTPException(status_code=400, detail="blocked")), + last, + ] + monkeypatch.setattr(litellm, "callbacks", _with_inert(callbacks, inert)) + dispatch: Final = _SuccessDispatch() + + await _deferred(dispatch) + + assert journal == ["parallel", "parallel done", "blocker", "last", "last done"] + assert last.seen[0] is replacement + assert dispatch.model_call_details["metadata"] == {"guardrail_blocked": True} + assert dispatch.dispatched == [replacement] + + +async def test_deferred_stream_v2_callback_between_two_guardrails_patches_what_the_second_sees_and_what_is_logged( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from_first: Final = ModelResponse() + from_v2: Final = ModelResponse() + v2_callback: Final = _Patcher(replacement=from_v2) + second: Final = _PostGuard("second", []) + monkeypatch.setattr(litellm, "callbacks", [_PostGuard("first", [], reply=from_first), v2_callback, second]) + dispatch: Final = _SuccessDispatch() + + await _deferred(dispatch) + + assert v2_callback.seen[0] is from_first + assert second.seen[0] is from_v2 + assert dispatch.dispatched == [from_v2] diff --git a/tests/unit/router_utils/test_fallback_event_v2_dispatch.py b/tests/unit/router_utils/test_fallback_event_v2_dispatch.py new file mode 100644 index 00000000000..7c6b81f71b0 --- /dev/null +++ b/tests/unit/router_utils/test_fallback_event_v2_dispatch.py @@ -0,0 +1,141 @@ +import asyncio +import threading +from collections.abc import Awaitable, Callable, Generator +from typing import Final + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_logger_v2 import CallContext, CustomLoggerV2, Field, Hook, on +from litellm.router_utils import fallback_event_handlers as handlers + +EventLogger = Callable[[str, dict[str, object], Exception], Awaitable[None]] +EVENTS: Final = ( + ("log_success_fallback_event", "fallback_success"), + ("log_failure_fallback_event", "fallback_failure"), +) +CALLBACK_LISTS: Final = ( + "callbacks", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", +) + + +@pytest.fixture(autouse=True) +def empty_callback_lists(monkeypatch: pytest.MonkeyPatch) -> Generator[None]: + for name in CALLBACK_LISTS: + monkeypatch.setattr(litellm, name, []) + yield + + +class Recorder(CustomLogger): + def __init__( + self, name: str, journal: list[object], error: Exception | None = None + ) -> None: # mutable-ok: the shared call log + super().__init__() # pyright: ignore[reportUnknownMemberType] # CustomLogger's untyped **kwargs + self.name = name + self.journal = journal + self.error = error + self.calls: list[tuple[str, str, dict[str, object], Exception]] = [] # mutable-ok: each call's arguments + + async def _record( + self, event: str, group: str, kwargs: dict[str, object], exception: Exception + ) -> None: # mutable-ok: the handler's kwargs dict + self.journal.append(self.name) + self.calls.append((event, group, kwargs, exception)) + await asyncio.sleep(0) + if self.error is not None: + raise self.error + + async def log_success_fallback_event( + self, original_model_group: str, kwargs: dict[str, object], original_exception: Exception + ) -> None: # mutable-ok: the handler's kwargs dict + await self._record("success", original_model_group, kwargs, original_exception) + + async def log_failure_fallback_event( + self, original_model_group: str, kwargs: dict[str, object], original_exception: Exception + ) -> None: # mutable-ok: the handler's kwargs dict + await self._record("failure", original_model_group, kwargs, original_exception) + + +class Watcher(CustomLoggerV2): + def __init__(self, journal: list[object], done: threading.Event) -> None: # mutable-ok: the shared call log + self.journal = journal + self.done = done + + @on(Hook.router_event, reads={Field.model, Field.outcome, Field.error}) + async def watch(self, view: object, ctx: CallContext[object]) -> None: + self.journal.append( + ("v2", getattr(view, "model"), getattr(view, "outcome"), str(getattr(view, "error"))) # noqa: B009 # the view's declared fields + ) + self.done.set() + + +def _public(name: str) -> EventLogger: + found: Final[EventLogger] = getattr(handlers, name) # pyright: ignore[reportAny] # the module's async handlers + return found + + +def _with_inert(callbacks: list[object], inert: bool) -> list[object]: # mutable-ok: a process list + return [*callbacks, Watcher([], threading.Event())] if inert else callbacks + + +@pytest.mark.parametrize("inert", [False, True]) +@pytest.mark.parametrize(("public", "_outcome"), EVENTS) +async def test_every_logger_in_the_merged_lists_runs_once_in_order_with_the_same_arguments( + public: str, _outcome: str, inert: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + journal: Final[list[object]] = [] # mutable-ok: the shared call log + first: Final = Recorder("first", journal) + second: Final = Recorder("second", journal) + third: Final = Recorder("third", journal) + monkeypatch.setattr(litellm, "callbacks", _with_inert([first, "langfuse", second], inert)) + monkeypatch.setattr(litellm, "_async_success_callback", [second, third, first]) + kwargs: Final[dict[str, object]] = {"model": "fallback"} # mutable-ok: the handler's kwargs dict + error: Final = RuntimeError("original") + + result: Final = await _public(public)("primary", kwargs, error) + + assert result is None + assert journal == ["first", "second", "third"] + event: Final = "success" if public.startswith("log_success") else "failure" + assert first.calls == [(event, "primary", kwargs, error)] + assert first.calls[0][2] is kwargs + assert first.calls[0][3] is error + + +@pytest.mark.parametrize("inert", [False, True]) +@pytest.mark.parametrize(("public", "_outcome"), EVENTS) +async def test_a_logger_error_is_logged_and_the_next_logger_still_runs( + public: str, _outcome: str, inert: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + journal: Final[list[object]] = [] # mutable-ok: the shared call log + logged: Final[list[tuple[object, ...]]] = [] # mutable-ok: the router logger's error calls + monkeypatch.setattr(handlers.verbose_router_logger, "error", lambda *args: logged.append(args)) + error: Final = ValueError("logger raised") + monkeypatch.setattr( + litellm, "callbacks", _with_inert([Recorder("broken", journal, error), Recorder("next", journal)], inert) + ) + + await _public(public)("primary", {}, RuntimeError("original")) + + assert journal == ["broken", "next"] + assert logged == [(f"Error in {public}: %s", error)] + + +@pytest.mark.parametrize(("public", "outcome"), EVENTS) +async def test_a_v2_callback_in_the_lists_is_delivered_off_path_after_the_legacy_loggers( + public: str, outcome: str, monkeypatch: pytest.MonkeyPatch +) -> None: + journal: Final[list[object]] = [] # mutable-ok: the shared call log + done: Final = threading.Event() + monkeypatch.setattr(litellm, "callbacks", [Recorder("before", journal), Watcher(journal, done)]) + monkeypatch.setattr(litellm, "_async_failure_callback", [Recorder("after", journal)]) + + await _public(public)("primary", {"model": "fallback"}, RuntimeError("original")) + + assert await asyncio.to_thread(done.wait, 30) + assert journal == ["before", "after", ("v2", "primary", outcome, "original")] diff --git a/tests/unit/test_router/core/test_callback_filters_v2_dispatch.py b/tests/unit/test_router/core/test_callback_filters_v2_dispatch.py new file mode 100644 index 00000000000..142834caeff --- /dev/null +++ b/tests/unit/test_router/core/test_callback_filters_v2_dispatch.py @@ -0,0 +1,232 @@ +from collections.abc import Callable, Sequence +from typing import Final + +import pytest + +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_logger_v2 import ( + Abort, + CallContext, + CustomLoggerV2, + Field, + Hook, + InvalidPatchError, + Patch, + on, +) +from litellm.integrations.custom_logger_v2.host import V2Outcome +from litellm.router_rust.steps.callbacks import OptionalChecks +from tests.unit.test_router.core.registry_helpers import MakeCore +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment + +GROUP: Final = [deployment("g", "a"), deployment("g", "b"), deployment("g", "c")] + +Deployments = list[dict[str, object]] # mutable-ok: the filter contract + + +def _id(deployment: object) -> object: + return deployment["model_info"]["id"] if isinstance(deployment, dict) else None # pyright: ignore[reportUnknownVariableType] # a deployment dict + + +def _ids(deployments: object) -> tuple[object, ...]: + return tuple(map(_id, deployments)) if isinstance(deployments, list | tuple) else () # pyright: ignore[reportUnknownArgumentType] # see above + + +class Narrow(CustomLogger): + def __init__( + self, name: str, journal: list[object], dropped: str = "", error: Exception | None = None + ) -> None: # mutable-ok: the shared call log + super().__init__() # pyright: ignore[reportUnknownMemberType] # CustomLogger's untyped **kwargs + self.name = name + self.journal = journal + self.dropped = dropped + self.error = error + self.seen: list[object] = [] # mutable-ok: each call's deployments + + async def async_filter_deployments( + self, + model: str, + healthy_deployments: Deployments, + messages: object, + request_kwargs: object = None, + parent_otel_span: object = None, + ) -> Deployments: + self.journal.append(self.name) + self.seen.append(healthy_deployments) + if self.error is not None: + raise self.error + return [d for d in healthy_deployments if _id(d) != self.dropped] + + +class Bystander(CustomLogger): + def __init__(self, journal: list[object]) -> None: # mutable-ok: the shared call log + super().__init__() # pyright: ignore[reportUnknownMemberType] # see above + self.journal = journal + + +class Picker(CustomLoggerV2): + def __init__( + self, + journal: list[object], + keep: Callable[[Sequence[object]], object] | None = None, + outcome: Abort | None = None, + ) -> None: # mutable-ok: the shared call log + self.journal = journal + self.keep = keep + self.outcome = outcome + + @on( + Hook.filter_deployments, + reads={Field.candidate_deployments, Field.model}, + writes={Field.candidate_deployments}, + may_abort=True, + ) + async def pick(self, view: object, ctx: CallContext[object]) -> Patch | Abort | None: + candidates: Final = getattr(view, "candidate_deployments") # noqa: B009 # the view's declared field + self.journal.append(("v2", getattr(view, "model"), _ids(candidates))) # noqa: B009 # see above + if self.outcome is not None: + return self.outcome + return None if self.keep is None else Patch().set(Field.candidate_deployments, self.keep(candidates)) + + +class Inert(CustomLoggerV2): + @on(Hook.pre_call) + async def ignore(self, view: object, ctx: CallContext[object]) -> None: + return None + + +def _checks(make_core: MakeCore) -> OptionalChecks: + core: Final = make_core(model_list=list(GROUP)) + checks: Final[OptionalChecks] = core.__dict__["_checks"] # pyright: ignore[reportAny] # RustCore's own attribute + return checks + + +def _install(host: ScriptedHost, callbacks: Sequence[object], inert: bool) -> None: + host.lists["callbacks"].extend((*callbacks, Picker([])) if inert else callbacks) + + +def _deployments() -> Deployments: + return [{"model_info": {"id": name}} for name in ("a", "b", "c")] + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_filters_run_in_list_order_and_each_sees_what_the_previous_one_kept( + make_core: MakeCore, host: ScriptedHost, inert: bool +) -> None: + journal: Final[list[object]] = [] # mutable-ok: the shared call log + first: Final = Narrow("first", journal, dropped="a") + second: Final = Narrow("second", journal, dropped="c") + checks: Final = _checks(make_core) + _install(host, [first, Bystander(journal), "a-string-callback", second], inert) + deployments: Final = _deployments() + + remaining, preferred = await checks.filter_deployments("g", deployments, None, {"model": "g"}) + + assert journal == ["first", "second"] + assert first.seen[0] is not deployments + assert _ids(first.seen[0]) == ("a", "b", "c") + assert _ids(second.seen[0]) == ("b", "c") + assert _ids(remaining) == ("b",) + assert preferred is None + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_a_filter_error_reaches_the_caller_unchanged_and_stops_later_filters( + make_core: MakeCore, host: ScriptedHost, inert: bool +) -> None: + journal: Final[list[object]] = [] # mutable-ok: the shared call log + error: Final = RuntimeError("a filter raised") + checks: Final = _checks(make_core) + _install(host, [Narrow("first", journal, error=error), Narrow("second", journal)], inert) + + with pytest.raises(RuntimeError) as raised: + await checks.filter_deployments("g", _deployments(), None, {}) + + assert raised.value is error + assert journal == ["first"] + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_an_empty_result_is_returned_for_the_core_to_fail_the_pick( + make_core: MakeCore, host: ScriptedHost, inert: bool +) -> None: + checks: Final = _checks(make_core) + _install( + host, [Narrow("only", [], dropped="a"), Narrow("all", [], dropped="b"), Narrow("rest", [], dropped="c")], inert + ) + + remaining, _ = await checks.filter_deployments("g", _deployments(), None, {}) + + assert remaining == [] + + +async def test_a_v2_callback_between_two_filters_sees_the_first_cut_and_the_second_sees_its_own( + make_core: MakeCore, host: ScriptedHost +) -> None: + journal: Final[list[object]] = [] # mutable-ok: the shared call log + second: Final = Narrow("second", journal) + checks: Final = _checks(make_core) + v2_callback: Final = Picker(journal, keep=lambda candidates: [d for d in candidates if _id(d) != "b"]) + _install(host, [Narrow("first", journal, dropped="a"), v2_callback, second], inert=False) + + remaining, _ = await checks.filter_deployments("g", _deployments(), None, {}) + + assert journal == ["first", ("v2", "g", ("b", "c")), "second"] + assert _ids(second.seen[0]) == ("c",) + assert _ids(remaining) == ("c",) + + +@pytest.mark.parametrize( + "keep", + [ + pytest.param(lambda candidates: [*candidates, {"model_info": {"id": "z"}}], id="adds"), + pytest.param(lambda candidates: [candidates[0], candidates[0]], id="repeats"), + ], +) +async def test_a_v2_callback_may_only_narrow_the_candidates( + make_core: MakeCore, host: ScriptedHost, keep: Callable[[Sequence[object]], object] +) -> None: + journal: Final[list[object]] = [] # mutable-ok: the shared call log + checks: Final = _checks(make_core) + _install(host, [Picker(journal, keep=keep), Narrow("after", journal)], inert=False) + + with pytest.raises(InvalidPatchError, match="only narrow"): + await checks.filter_deployments("g", _deployments(), None, {}) + + assert journal == [("v2", "g", ("a", "b", "c"))] + + +async def test_a_v2_callback_abort_ends_the_filters_as_an_error(make_core: MakeCore, host: ScriptedHost) -> None: + journal: Final[list[object]] = [] # mutable-ok: the shared call log + abort: Final = Abort(status=429, message="no capacity") + checks: Final = _checks(make_core) + _install(host, [Narrow("before", journal), Picker(journal, outcome=abort), Narrow("after", journal)], inert=False) + + with pytest.raises(V2Outcome) as raised: + await checks.filter_deployments("g", _deployments(), None, {}) + + assert raised.value.outcome is abort + assert journal == ["before", ("v2", "g", ("a", "b", "c"))] + + +async def test_only_a_v2_callback_that_filters_makes_the_core_ask_for_filters( + make_core: MakeCore, host: ScriptedHost +) -> None: + checks: Final = _checks(make_core) + host.lists["callbacks"].append(Inert()) + without_filter: Final = checks.has_filters() + host.lists["callbacks"].append(Picker([])) + + assert (without_filter, checks.has_filters()) == (False, True) + + +async def test_the_pick_honours_a_v2_callback_between_legacy_filters(make_core: MakeCore, host: ScriptedHost) -> None: + journal: Final[list[object]] = [] # mutable-ok: the shared call log + v2_callback: Final = Picker(journal, keep=lambda candidates: [d for d in candidates if _id(d) != "c"]) + host.lists["callbacks"].extend((Narrow("first", journal, dropped="a"), v2_callback, Narrow("second", journal))) + core: Final = make_core(model_list=list(GROUP)) + + picked: Final = {_id(await core.select("g")) for _ in range(10)} + + assert picked == {"b"} + assert journal[:3] == ["first", ("v2", "g", ("b", "c")), "second"] diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 7fb1e1d5244..8b83581e49b 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -13,7 +13,7 @@ import queue import threading import time import uuid -from collections.abc import Callable, Iterator, Mapping +from collections.abc import Awaitable, Callable, Iterator, Mapping from concurrent.futures import Future, ThreadPoolExecutor from datetime import datetime, timedelta, timezone from functools import partial @@ -44,6 +44,8 @@ from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_logger_v2 import Abort, CallContext, CustomLoggerV2, Field, Hook, Patch, on +from litellm.integrations.custom_logger_v2.host import V2Outcome from litellm.integrations.humanloop import HumanloopLogger from litellm.llms.base_llm.audio_transcription.transformation import BaseAudioTranscriptionConfig from litellm.integrations.langfuse.langfuse_prompt_management import LangfusePromptManagement @@ -11910,3 +11912,390 @@ async def test_pre_call_deployment_hook_hands_responses_compact_its_call_type(mo await async_pre_call_deployment_hook({"model": "gpt-5.4-mini"}, "acompact_responses") assert hook.seen_call_types == (CallTypes.acompact_responses,) + + +_REWRITTEN_RESPONSE: Final = ModelResponse(id="rewritten-by-the-deployment-hook") + + +class _DeploymentLogger(CustomLogger): + def __init__( + self, + name: str, + journal: list[object], + replacement: object = None, + error: Exception | None = None, + ) -> None: + super().__init__() + self.name = name + self.journal = journal + self.replacement = replacement + self.error = error + self.seen: list[tuple[object, ...]] = [] + + def _record(self, *seen: object) -> object: + self.journal.append(self.name) + self.seen.append(seen) + if self.error is not None: + raise self.error + return self.replacement + + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict[str, object] | None: + return cast("dict[str, object] | None", self._record(kwargs, call_type)) + + async def async_post_call_success_deployment_hook( + self, request_data: dict[str, object], response: object, call_type: CallTypes | None + ) -> object: + return self._record(request_data, response, call_type) + + async def async_post_call_failure_deployment_hook( + self, + request_data: Mapping[str, object], + exception: Exception, + call_type: CallTypes | None, + fallback_depth: int | None = None, + ) -> None: + self._record(request_data, exception, call_type, fallback_depth) + + +class _DeploymentGuardrail(CustomGuardrail): + def __init__(self, error: Exception) -> None: + super().__init__(guardrail_name="deployment-guard", default_on=True) + self.error = error + + async def async_post_call_success_deployment_hook( + self, request_data: dict[str, object], response: object, call_type: CallTypes | None + ) -> object: + raise self.error + + +class _DeploymentV2(CustomLoggerV2): + def __init__(self, journal: list[object], replacement: object = None, outcome: Abort | None = None) -> None: + self.journal = journal + self.replacement = replacement + self.outcome = outcome + + @on(Hook.pre_deployment_call, reads={Field.model}, may_abort=True) + async def pre_deployment(self, view: object, ctx: CallContext[object]) -> Abort | None: + self.journal.append(("v2", getattr(view, "model"))) # noqa: B009 # the view's declared field + return self.outcome + + @on(Hook.post_deployment_success, reads={Field.response}, writes={Field.response}, may_abort=True) + async def post_success(self, view: object, ctx: CallContext[object]) -> Patch | Abort | None: + self.journal.append(("v2", getattr(view, "response"))) # noqa: B009 # see above + if self.outcome is not None: + return self.outcome + return None if self.replacement is None else Patch().set(Field.response, self.replacement) + + @on(Hook.post_deployment_failure, reads={Field.error}, may_abort=True) + async def post_failure(self, view: object, ctx: CallContext[object]) -> Abort | None: + self.journal.append(("v2", str(getattr(view, "error")))) # noqa: B009 # see above + return self.outcome + + +def _with_inert_v2(callbacks: list[object], inert: bool) -> list[object]: + return [*callbacks, _DeploymentV2([])] if inert else callbacks + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_pre_deployment_hooks_run_in_order_and_a_returned_dict_replaces_the_kwargs( + monkeypatch: pytest.MonkeyPatch, inert: bool +) -> None: + journal: Final[list[object]] = [] + replaced: Final[dict[str, object]] = {"model": "replaced"} + first: Final = _DeploymentLogger("first", journal, replacement=replaced) + keeper: Final = _DeploymentLogger("keeper", journal) + last: Final = _DeploymentLogger("last", journal) + monkeypatch.setattr(litellm, "callbacks", _with_inert_v2([first, "a-string-callback", keeper, last], inert)) + + result: Final = await async_pre_call_deployment_hook({"model": "original"}, "acompletion", include_v2=True) + + assert journal == ["first", "keeper", "last"] + assert first.seen[0][0] == {"model": "original"} + assert keeper.seen[0][0] is replaced + assert last.seen[0][0] is replaced + assert result is replaced + assert first.seen[0][1] is CallTypes.acompletion + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_a_pre_deployment_hook_error_propagates_unchanged_and_stops_the_hook( + monkeypatch: pytest.MonkeyPatch, inert: bool +) -> None: + journal: Final[list[object]] = [] + error: Final = ValueError("rejected before the deployment") + monkeypatch.setattr( + litellm, + "callbacks", + _with_inert_v2([_DeploymentLogger("first", journal, error=error), _DeploymentLogger("next", journal)], inert), + ) + + with pytest.raises(ValueError, match="rejected before the deployment") as raised: + await async_pre_call_deployment_hook({"model": "m"}, "acompletion", include_v2=True) + + assert raised.value is error + assert journal == ["first"] + + +async def test_a_pre_deployment_v2_callback_between_legacy_hooks_sees_the_first_rewrite_and_its_abort_stops_the_rest( + monkeypatch: pytest.MonkeyPatch, +) -> None: + journal: Final[list[object]] = [] + abort: Final = Abort(status=403, message="deployment not allowed") + monkeypatch.setattr( + litellm, + "callbacks", + [ + _DeploymentLogger("first", journal, replacement={"model": "rewritten"}), + _DeploymentV2(journal, outcome=abort), + _DeploymentLogger("after", journal), + ], + ) + + with pytest.raises(V2Outcome) as raised: + await async_pre_call_deployment_hook({"model": "m"}, "acompletion", include_v2=True) + + assert raised.value.outcome is abort + assert journal == ["first", ("v2", "rewritten")] + + +@pytest.mark.parametrize( + "hook_call", + [ + pytest.param(lambda: async_pre_call_deployment_hook({"model": "m"}, "acompletion"), id="pre"), + pytest.param( + lambda: async_post_call_success_deployment_hook({"model": "m"}, ModelResponse(), CallTypes.acompletion), + id="success", + ), + pytest.param( + lambda: async_post_call_failure_deployment_hook({"model": "m"}, ValueError("x"), "acompletion"), + id="failure", + ), + ], +) +async def test_callers_that_do_not_ask_for_v2_callbacks_run_only_the_legacy_hooks( + monkeypatch: pytest.MonkeyPatch, hook_call: Callable[[], object] +) -> None: + journal: Final[list[object]] = [] + monkeypatch.setattr(litellm, "callbacks", [_DeploymentV2(journal), _DeploymentLogger("legacy", journal)]) + + await cast("Awaitable[object]", hook_call()) + + assert journal == ["legacy"] + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_post_deployment_success_chains_replacements_and_swallows_a_logger_error( + monkeypatch: pytest.MonkeyPatch, inert: bool +) -> None: + journal: Final[list[object]] = [] + replacement: Final = ModelResponse(id="replaced") + first: Final = _DeploymentLogger("first", journal, replacement=replacement) + broken: Final = _DeploymentLogger("broken", journal, error=RuntimeError("logger bug")) + last: Final = _DeploymentLogger("last", journal) + monkeypatch.setattr(litellm, "callbacks", _with_inert_v2([first, broken, last], inert)) + request: Final[dict[str, object]] = {"model": "m"} + response: Final = ModelResponse(id="original") + + result: Final = await async_post_call_success_deployment_hook( + request, response, CallTypes.acompletion, include_v2=True + ) + + assert journal == ["first", "broken", "last"] + assert first.seen[0][0] is request + assert first.seen[0][1] is response + assert broken.seen[0][1] is replacement + assert last.seen[0][1] is replacement + assert result is replacement + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_post_deployment_success_reraises_a_guardrail_error_and_stops_the_hook( + monkeypatch: pytest.MonkeyPatch, inert: bool +) -> None: + journal: Final[list[object]] = [] + error: Final = ValueError("response blocked") + monkeypatch.setattr( + litellm, + "callbacks", + _with_inert_v2([_DeploymentGuardrail(error), _DeploymentLogger("next", journal)], inert), + ) + + with pytest.raises(ValueError, match="response blocked") as raised: + await async_post_call_success_deployment_hook({}, ModelResponse(), CallTypes.acompletion, include_v2=True) + + assert raised.value is error + assert journal == [] + + +async def test_a_post_deployment_success_v2_callback_between_legacy_hooks_sees_and_replaces_the_response( + monkeypatch: pytest.MonkeyPatch, +) -> None: + journal: Final[list[object]] = [] + first_reply: Final = ModelResponse(id="first") + native_reply: Final = ModelResponse(id="v2") + second: Final = _DeploymentLogger("second", journal) + monkeypatch.setattr( + litellm, + "callbacks", + [ + _DeploymentLogger("first", journal, replacement=first_reply), + _DeploymentV2(journal, replacement=native_reply), + second, + ], + ) + + result: Final = await async_post_call_success_deployment_hook( + {}, ModelResponse(id="original"), CallTypes.acompletion, include_v2=True + ) + + assert journal == ["first", ("v2", first_reply), "second"] + assert second.seen[0][1] is native_reply + assert result is native_reply + + +@pytest.mark.parametrize("inert", [False, True]) +async def test_post_deployment_failure_hands_every_logger_a_snapshot_and_swallows_their_errors( + monkeypatch: pytest.MonkeyPatch, inert: bool +) -> None: + journal: Final[list[object]] = [] + broken: Final = _DeploymentLogger("broken", journal, error=RuntimeError("logger bug")) + last: Final = _DeploymentLogger("last", journal) + monkeypatch.setattr(litellm, "callbacks", _with_inert_v2([broken, last], inert)) + error: Final = ValueError("deployment failed") + + await async_post_call_failure_deployment_hook( + {"model": "m", "fallback_depth": 2, "attempted_targets": object()}, error, "acompletion", include_v2=True + ) + + assert journal == ["broken", "last"] + request, seen_error, call_type, depth = last.seen[0] + assert dict(cast("Mapping[str, object]", request)) == {"model": "m", "fallback_depth": 2} + assert type(seen_error) is ValueError + assert seen_error is not error + assert broken.seen[0][1] is seen_error + assert (call_type, depth) == (CallTypes.acompletion, 2) + + +async def test_a_post_deployment_failure_v2_callback_between_legacy_hooks_runs_in_list_order( + monkeypatch: pytest.MonkeyPatch, +) -> None: + journal: Final[list[object]] = [] + monkeypatch.setattr( + litellm, + "callbacks", + [_DeploymentLogger("first", journal), _DeploymentV2(journal), _DeploymentLogger("second", journal)], + ) + + await async_post_call_failure_deployment_hook({}, ValueError("deployment failed"), "acompletion", include_v2=True) + + assert journal == ["first", ("v2", "deployment failed"), "second"] + + +class _PreDeploymentRewriter(CustomLogger): + def __init__(self, seen: list[str]) -> None: + super().__init__() + self.seen = seen + + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict[str, object] | None: + self.seen.append("legacy") + return {**kwargs, "mock_response": "rewritten by the pre-deployment hook"} + + +class _PostDeploymentSuccessRewriter(CustomLogger): + def __init__(self, seen: list[str]) -> None: + super().__init__() + self.seen = seen + + async def async_post_call_success_deployment_hook( + self, request_data: dict[str, object], response: object, call_type: CallTypes | None + ) -> object: + self.seen.append("legacy") + return _REWRITTEN_RESPONSE + + +class _PostDeploymentFailureRecorder(CustomLogger): + def __init__(self, seen: list[str]) -> None: + super().__init__() + self.seen = seen + self.failures: tuple[Exception, ...] = () + + async def async_post_call_failure_deployment_hook( + self, request_data: Mapping[str, object], exception: Exception, call_type: CallTypes | None + ) -> None: + self.seen.append("legacy") + self.failures = (*self.failures, exception) + + +def _deployment_gate(hook: Hook, seen: list[str]) -> CustomLoggerV2: + class Gate(CustomLoggerV2): + @on(hook, may_abort=True) + async def gate(self, view: object, ctx: CallContext[object]) -> None: + seen.append("v2") + + return Gate() + + +def _install_deployment_hook( + monkeypatch: pytest.MonkeyPatch, legacy: CustomLogger, v2_hook: Hook | None, seen: list[str] +) -> None: + gate: Final = [] if v2_hook is None else [_deployment_gate(v2_hook, seen)] + monkeypatch.setattr(litellm, "callbacks", [*gate, legacy]) + + +@pytest.mark.parametrize("with_v2", [False, True]) +async def test_pre_deployment_call_sends_the_kwargs_the_legacy_hook_returned( + monkeypatch: pytest.MonkeyPatch, with_v2: bool +) -> None: + seen: Final[list[str]] = [] + _install_deployment_hook( + monkeypatch, _PreDeploymentRewriter(seen), Hook.pre_deployment_call if with_v2 else None, seen + ) + + response: Final = await litellm.acompletion( + model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], mock_response="original", num_retries=0 + ) + + assert isinstance(response, ModelResponse) + assert response.choices[0].message.content == "rewritten by the pre-deployment hook" + assert seen == (["v2", "legacy"] if with_v2 else ["legacy"]) + + +@pytest.mark.parametrize("with_v2", [False, True]) +async def test_post_deployment_success_returns_what_the_legacy_hook_returned( + monkeypatch: pytest.MonkeyPatch, with_v2: bool +) -> None: + seen: Final[list[str]] = [] + _install_deployment_hook( + monkeypatch, _PostDeploymentSuccessRewriter(seen), Hook.post_deployment_success if with_v2 else None, seen + ) + + response: Final = await litellm.acompletion( + model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], mock_response="original", num_retries=0 + ) + + assert response is _REWRITTEN_RESPONSE + assert seen == (["v2", "legacy"] if with_v2 else ["legacy"]) + + +@pytest.mark.parametrize("with_v2", [False, True]) +async def test_post_deployment_failure_reaches_the_legacy_hook_and_keeps_the_real_error( + monkeypatch: pytest.MonkeyPatch, with_v2: bool +) -> None: + seen: Final[list[str]] = [] + recorder: Final = _PostDeploymentFailureRecorder(seen) + _install_deployment_hook(monkeypatch, recorder, Hook.post_deployment_failure if with_v2 else None, seen) + error: Final = litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o-mini") + + with pytest.raises(litellm.AuthenticationError, match="bad key") as raised: + await litellm.acompletion( + model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], mock_response=error, num_retries=0 + ) + + assert len(recorder.failures) == 1 + assert type(recorder.failures[0]) is type(raised.value) + assert recorder.failures[0] is not raised.value + assert seen == (["v2", "legacy"] if with_v2 else ["legacy"])