diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index cdc108a6b4e..84106b77a7d 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -11,6 +11,7 @@ across pods or stop races; the hook reads active jobs through a short-TTL cache. import asyncio import hashlib +import json import random import traceback from collections.abc import Awaitable, Callable, Mapping, Sequence @@ -28,7 +29,7 @@ from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.websearch_interception.tools import is_web_search_tool_responses -from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs +from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs, independent_snapshot from litellm.litellm_core_utils.internal_call_metadata import sanitized_forwardable_call_metadata from litellm.litellm_core_utils.llm_judge import ( default_router_provider, @@ -281,10 +282,8 @@ class _SurfaceOps: request (messages plus translated generation params) and how its response yields the judgeable final text. Membership in this table IS the sampling allowlist; unknown call types fail closed. ``wire_params`` marks the surfaces whose params - come from the proxy's wire-body snapshot, which is taken before the guardrail - pre-call hook: those rows must not sample a request a pre-call guardrail rewrote, - or the shadow call would replay content (tools, unmasked entities) the guardrail - removed.""" + come from the proxy's native request snapshot. Requests rewritten by guardrails + require a post-hook snapshot whose guardrail history is still current.""" __slots__ = ("chat_request", "final_text", "wire_params") @@ -311,19 +310,85 @@ _NON_MUTATING_GUARDRAIL_MODES: Final = frozenset( ) +def _guardrail_is_non_mutating(entry: Mapping[str, object], allowed_modes: frozenset[str]) -> bool: + modes: Final = entry.get("guardrail_mode") + return all( + isinstance(mode, str) and mode in allowed_modes + for mode in (modes if isinstance(modes, list | tuple) else (modes,)) + ) + + def _request_mutating_guardrail_ran(request_metadata: Mapping[str, object]) -> bool: - """Whether a guardrail that can rewrite the outbound request ran on this one, read - from the same guardrail-information entries spend logging uses. str-enum modes - compare equal to their plain-string values, and an entry whose mode is missing or - unrecognized counts as mutating.""" raw: Final = request_metadata.get("standard_logging_guardrail_information") entries: Final = raw if isinstance(raw, Sequence) else () - modes_per_entry: Final = tuple(entry.get("guardrail_mode") for entry in entries if isinstance(entry, Mapping)) return any( - not all( - mode in _NON_MUTATING_GUARDRAIL_MODES for mode in (modes if isinstance(modes, list | tuple) else (modes,)) + not _guardrail_is_non_mutating(entry, _NON_MUTATING_GUARDRAIL_MODES) + for entry in entries + if isinstance(entry, Mapping) + ) + + +def request_guardrail_fingerprint(request_metadata: Mapping[str, object]) -> str | None: + raw: Final = request_metadata.get("standard_logging_guardrail_information") + entries: Final = raw if isinstance(raw, Sequence) else () + replay_safe_modes: Final = _NON_MUTATING_GUARDRAIL_MODES - frozenset(("logging_only",)) + relevant: Final = tuple( + entry + for entry in entries + if isinstance(entry, Mapping) and not _guardrail_is_non_mutating(entry, replay_safe_modes) + ) + try: + serialized: Final = json.dumps(relevant, sort_keys=True, default=str) + except (TypeError, ValueError): + return None + return hashlib.sha256(serialized.encode()).hexdigest() + + +@dataclass(frozen=True, slots=True) +class GuardrailRequestSnapshot: + body: Mapping[str, object] + fingerprint: str + + @staticmethod + def capture(body: Mapping[str, object], metadata: Mapping[str, object]) -> "GuardrailRequestSnapshot | None": + if not _request_mutating_guardrail_ran(metadata): + return None + fingerprint: Final = request_guardrail_fingerprint(metadata) + if fingerprint is None: + return None + return GuardrailRequestSnapshot( + body=MappingProxyType( + _CHAT_REQUEST_ADAPTER.validate_python( + independent_snapshot(dict(body)) # mutable-ok: snapshot helper requires a plain dictionary + ) + ), + fingerprint=fingerprint, ) - for modes in modes_per_entry + + +def _post_guardrail_kwargs( + kwargs: Mapping[str, object], + request_metadata: Mapping[str, object], + ops: _SurfaceOps, + guardrail_snapshot: GuardrailRequestSnapshot | None, +) -> Mapping[str, object] | None: + if guardrail_snapshot is None or guardrail_snapshot.fingerprint != request_guardrail_fingerprint(request_metadata): + return None + raw_params: Final = kwargs.get("litellm_params") + litellm_params: Final = raw_params if isinstance(raw_params, Mapping) else _EMPTY_METADATA + raw_request: Final = litellm_params.get("proxy_server_request") + request: Final = raw_request if isinstance(raw_request, Mapping) else _EMPTY_METADATA + body: Final = guardrail_snapshot.body + return MappingProxyType( + { + **kwargs, + "messages": body.get("input" if ops is _RESPONSES_OPS else "messages"), + "system": body.get("system"), + "instructions": body.get("instructions"), + "litellm_params": MappingProxyType( + {**litellm_params, "proxy_server_request": MappingProxyType({**request, "body": body})} + ), + } ) @@ -881,6 +946,8 @@ class ShadowEvalLogger(CustomLogger): response_obj: object, start_time: object, end_time: object, + *, + guardrail_snapshot: GuardrailRequestSnapshot | None = None, ) -> None: try: payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") # pyright: ignore[reportAssignmentType] # untyped callback kwargs @@ -914,8 +981,13 @@ class ShadowEvalLogger(CustomLogger): ops: Final = _SURFACE_OPS.get(str(payload.get("call_type") or "")) if ops is None: return # only surfaces this table can normalize are comparable; unknown types fail closed - if ops.wire_params and _request_mutating_guardrail_ran(request_metadata): - return # the wire-body snapshot predates the rewrite; replaying it would resurrect stripped content + sample_kwargs: Final = ( + _post_guardrail_kwargs(kwargs, request_metadata, ops, guardrail_snapshot) + if ops.wire_params and _request_mutating_guardrail_ran(request_metadata) + else kwargs + ) + if sample_kwargs is None: + return active_jobs: Final = await self._active_jobs() eligible: Final = self._sampled_jobs( tuple(job for target in targets for job in active_jobs.get(target, ())), @@ -927,7 +999,7 @@ class ShadowEvalLogger(CustomLogger): return sample: Final = _judgeable_sample( ops, - kwargs, + sample_kwargs, MappingProxyType(dict(payload.get("model_parameters") or {})), # mutable-ok: frozen snapshot response_obj, ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 8e28a0d543d..b038762a476 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -226,6 +226,7 @@ if TYPE_CHECKING: from litellm.integrations.otel.logger import OpenTelemetryV2 from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config + from litellm.integrations.shadow_eval_logger import GuardrailRequestSnapshot from litellm.litellm_core_utils.llm_cost_calc.utils import BilledTokenRates from litellm.llms.base_llm.passthrough.transformation import PassthroughStreamCollector from litellm.proxy.hooks.autorouter_baseline_cache import BaselineCacheContext, CapturedBaselineObservation @@ -714,6 +715,7 @@ class Logging(LiteLLMLoggingBaseClass): self._defer_async_logging: bool = False self._enqueue_deferred_logging: Callable[[], None] | None = None self._on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None + self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None: """Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``.""" @@ -2825,6 +2827,7 @@ class Logging(LiteLLMLoggingBaseClass): ): continue + self.shadow_eval_request_snapshot = None self.model_call_details, result = callback.logging_hook( kwargs=self.model_call_details, result=result, @@ -3391,6 +3394,7 @@ class Logging(LiteLLMLoggingBaseClass): ): continue + self.shadow_eval_request_snapshot = None self.model_call_details, result = await callback.async_logging_hook( kwargs=self.model_call_details, result=result, @@ -3450,6 +3454,8 @@ class Logging(LiteLLMLoggingBaseClass): ) if isinstance(callback, CustomLogger): # custom logger class + from litellm.integrations.shadow_eval_logger import ShadowEvalLogger + model_call_details: dict = self.model_call_details ################################## # call redaction hook for custom logger @@ -3460,7 +3466,19 @@ class Logging(LiteLLMLoggingBaseClass): model_call_details=model_call_details, custom_logger=callback ) ################################## - if self.stream is True: + 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 result, + 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, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index cbf4786affe..2e43c6b0d22 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2208,7 +2208,7 @@ class ProxyBaseLLMRequestProcessing: # Refresh AFTER pre_call_hook: guardrails (e.g. Presidio PII masking) may # have mutated `self.data` in place, and the audit-trail snapshot taken in # add_litellm_data_to_request predates that mutation. - refresh_proxy_server_request_body_snapshot(self.data) + refresh_proxy_server_request_body_snapshot(self.data, guardrails_applied=True) verbose_proxy_logger.debug("receiving data: %s", self.data) if "messages" in self.data and self.data["messages"]: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 90bba82aa84..3bf6fea62f7 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1923,6 +1923,8 @@ class LiteLLMProxyRequestSetup: def refresh_proxy_server_request_body_snapshot( data: MutableMapping[str, object], + *, + guardrails_applied: bool = False, ) -> None: """ Re-snapshot ``data["proxy_server_request"]["body"]`` from the current state of ``data``. @@ -1938,13 +1940,27 @@ def refresh_proxy_server_request_body_snapshot( ``Logging`` instance, so it must be excluded here the same way ``secret_fields`` and ``proxy_server_request`` are. """ - proxy_server_request = data.get("proxy_server_request") + from litellm.integrations.shadow_eval_logger import GuardrailRequestSnapshot + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj: Final = data.get("litellm_logging_obj") + if isinstance(logging_obj, Logging): + logging_obj.shadow_eval_request_snapshot = None + proxy_server_request: Final = data.get("proxy_server_request") if not isinstance(proxy_server_request, dict): return - _body_snapshot_exclude = ( + _body_snapshot_exclude: Final = ( frozenset({"secret_fields", "proxy_server_request", "litellm_logging_obj"}) | _TRANSPORT_ONLY_CREDENTIAL_KEYS ) - proxy_server_request["body"] = {k: v for k, v in data.items() if k not in _body_snapshot_exclude} + body: Final = { # mutable-ok: audit JSON serialization requires a dict with shared nested messages + k: v for k, v in data.items() if k not in _body_snapshot_exclude + } + proxy_server_request["body"] = body + if guardrails_applied and isinstance(logging_obj, Logging): + metadata: Final = data.get(get_metadata_variable_name_from_kwargs(data)) + logging_obj.shadow_eval_request_snapshot = GuardrailRequestSnapshot.capture( + body, metadata if isinstance(metadata, Mapping) else MappingProxyType({}) + ) async def add_litellm_data_to_request( diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index 76efe9c8576..e3f059a7941 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -4,7 +4,7 @@ the detached pipeline's single attempt-row write, and the cache-first job lookup import asyncio from collections.abc import Mapping from datetime import datetime, timedelta, timezone -from typing import Final +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock import pytest @@ -19,11 +19,13 @@ from litellm.integrations.shadow_eval_logger import ( JUDGE_MAX_OUTPUT_TOKENS, PAIRWISE_JUDGE_RESPONSE_FORMAT, ActiveShadowEvalJob, + GuardrailRequestSnapshot, ShadowEvalLogger, _failure_detail, _judge_user_prompt, _sample_hits, _unmask_preference, + request_guardrail_fingerprint, ) from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( @@ -35,6 +37,15 @@ from litellm.types.utils import ( ) +def test_guardrail_fingerprint_excludes_auth_metadata() -> None: + history: Final = [{"guardrail_name": "mask", "guardrail_mode": "pre_call"}] + metadata: Final = {"standard_logging_guardrail_information": history} + fingerprint: Final = request_guardrail_fingerprint(metadata) + assert fingerprint == request_guardrail_fingerprint({**metadata, "user_api_key": "first-test-credential"}) + assert fingerprint == request_guardrail_fingerprint({**metadata, "user_api_key": "second-test-credential"}) + assert fingerprint != request_guardrail_fingerprint({"standard_logging_guardrail_information": []}) + + def _job(**overrides) -> ActiveShadowEvalJob: defaults = dict( id="job-1", @@ -312,11 +323,19 @@ class TestSurfaceNormalization: """/v1/messages and /v1/responses arms: the hook normalizes each surface's logged request through litellm's own transformations and judges only text-final turns.""" - async def _drive(self, hook_kwargs, response_obj): - prisma = _prisma() - router = _router() - logger = _logger(router=router, prisma=prisma, jobs=(_job(),)) - await logger.async_log_success_event(hook_kwargs, response_obj, None, None) + async def _drive( + self, + hook_kwargs: Mapping[str, object], + response_obj: object, + *, + guardrail_snapshot: GuardrailRequestSnapshot | None = None, + ) -> tuple[MagicMock, MagicMock]: + prisma: Final = _prisma() + router: Final = _router() + logger: Final = _logger(router=router, prisma=prisma, jobs=(_job(),)) + await logger.async_log_success_event( + hook_kwargs, response_obj, None, None, guardrail_snapshot=guardrail_snapshot + ) await _drain(logger) return prisma, router @@ -711,41 +730,142 @@ class TestSurfaceNormalization: prisma.db.litellm_shadowevalattempt.create.assert_not_called() @pytest.mark.parametrize( - "call_type,guardrail_mode,sampled", + "call_type,guardrail_mode,checkpoint,later_mode,sampled", [ - ("anthropic_messages", ["logging_only", "pre_call"], False), - ("aresponses", GuardrailEventHooks.pre_call, False), - ("anthropic_messages", "post_call", True), - ("acompletion", "pre_call", True), + ("anthropic_messages", "pre_call", "absent", None, False), + ("aresponses", "pre_call", "corrupt", None, False), + ("anthropic_messages", ["logging_only", "pre_call"], "missing", None, False), + ("aresponses", GuardrailEventHooks.pre_call, "missing", None, False), + ("anthropic_messages", "pre_call", "unapproved", None, False), + ("aresponses", "pre_call", "unapproved", None, False), + ("anthropic_messages", ["logging_only", "pre_call"], "approved", None, True), + ("aresponses", GuardrailEventHooks.pre_call, "approved", None, True), + ("anthropic_messages", "pre_call", "approved", "pre_call", False), + ("aresponses", "pre_call", "approved", "pre_call", False), + ("anthropic_messages", "pre_call", "approved", "logging_only", False), + ("aresponses", "pre_call", "approved", "logging_only", False), + ("anthropic_messages", "pre_call", "approved", "post_call", True), + ("aresponses", "pre_call", "approved", "post_call", True), + ("anthropic_messages", "post_call", "missing", None, True), + ("acompletion", "pre_call", "missing", None, True), ], - ids=["anthropic-pre-call-list", "responses-pre-call-enum", "anthropic-post-call-only", "chat-pre-call"], ) - async def test_guardrail_rewritten_requests_never_replay_the_wire_body(self, call_type, guardrail_mode, sampled): - """The proxy snapshots the wire body before the guardrail pre-call hook, so the - wire-sourced surfaces skip requests a request-mutating guardrail ran on rather - than replay stripped tools or unmasked content; chat sources the dispatched - call and keeps sampling, as do requests only response-mode guardrails touched.""" - hook_kwargs = _success_kwargs( + async def test_guardrail_replay_requires_current_approved_snapshot( + self, + call_type: str, + guardrail_mode: str | list[str], + checkpoint: Literal["absent", "corrupt", "missing", "unapproved", "approved"], + later_mode: str | None, + sampled: bool, + ) -> None: + history: Final[list[dict[str, object]]] = [{"guardrail_name": "g", "guardrail_mode": guardrail_mode}] + if checkpoint == "corrupt": + history[0]["guardrail_response"] = history + body: Final[dict[str, object]] = { + "model": "model", + "messages": [{"role": "user", "content": "approved input"}], + "input": "approved input", + } + snapshot: Final = ( + GuardrailRequestSnapshot.capture(body, {"standard_logging_guardrail_information": history}) + if checkpoint in ("approved", "corrupt") else None + ) + if checkpoint == "corrupt": + assert snapshot is None + base_kwargs: Final = _success_kwargs( call_type=call_type, request_metadata={ - "standard_logging_guardrail_information": [{"guardrail_name": "g", "guardrail_mode": guardrail_mode}] + "standard_logging_guardrail_information": history + + ([{"guardrail_name": "g", "guardrail_mode": later_mode}] if later_mode else []) }, ) - response = RESPONSE - if call_type == "anthropic_messages": - hook_kwargs["messages"] = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] - elif call_type == "aresponses": - hook_kwargs["messages"] = "hi" - response = RESPONSES_API_RESPONSE + hook_kwargs: Final = { + **base_kwargs, + "messages": "hi" if call_type == "aresponses" else base_kwargs["messages"], + "litellm_params": { + **base_kwargs["litellm_params"], + "proxy_server_request": None if checkpoint == "absent" else {"body": body}, + }, + } - prisma, router = await self._drive(hook_kwargs, response) + prisma, router = await self._drive( + hook_kwargs, + RESPONSES_API_RESPONSE if call_type == "aresponses" else RESPONSE, + guardrail_snapshot=snapshot, + ) if sampled: + assert router.acompletion.call_count == 2 prisma.db.litellm_shadowevalattempt.create.assert_called_once() else: router.acompletion.assert_not_called() prisma.db.litellm_shadowevalattempt.create.assert_not_called() + @pytest.mark.parametrize("call_type", ["anthropic_messages", "aresponses"]) + @pytest.mark.parametrize("remove_optional_fields", [False, True]) + async def test_approved_guardrail_snapshot_replays_independent_native_input( + self, call_type: str, remove_optional_fields: bool + ) -> None: + is_responses: Final = call_type == "aresponses" + metadata: Final = { + "standard_logging_guardrail_information": [{"guardrail_name": "g", "guardrail_mode": "pre_call"}] + } + live_message: Final = {"role": "user", "content": "approved input"} + live_tool: Final = { + "name": "approved_tool", + "description": "approved tool", + "strict": False, + "parameters" if is_responses else "input_schema": {"type": "object", "properties": {}}, + **({"type": "function"} if is_responses else {}), + } + data: Final[dict[str, object]] = { + "model": "model", + "input" if is_responses else "messages": [live_message], + "max_output_tokens" if is_responses else "max_tokens": 123, + **({} if remove_optional_fields else { + "instructions" if is_responses else "system": "approved system", + "tools": [live_tool], + "temperature": 0.2, + }), + } + snapshot: Final = GuardrailRequestSnapshot.capture(data, metadata) + assert snapshot is not None + live_message["content"] = "changed after checkpoint" + live_tool["name"] = "changed_after_checkpoint" + base_kwargs: Final = _success_kwargs(call_type=call_type, request_metadata=metadata) + hook_kwargs: Final = { + **base_kwargs, + "messages": "stale input" if is_responses else [{"role": "user", "content": "stale input"}], + "system": "stale system", + "instructions": "stale system", + "standard_logging_object": { + **base_kwargs["standard_logging_object"], + "model_parameters": {"tools": [{"name": "stale_tool"}], "temperature": 0.9, "max_tokens": 999}, + }, + "litellm_params": {**base_kwargs["litellm_params"], "proxy_server_request": {"body": data}}, + } + + prisma, router = await self._drive( + hook_kwargs, RESPONSES_API_RESPONSE if is_responses else RESPONSE, guardrail_snapshot=snapshot + ) + + assert router.acompletion.call_count == 2 + shadow_call: Final = router.acompletion.call_args_list[0].kwargs + assert shadow_call["messages"] == ( + [] if remove_optional_fields else [{"role": "system", "content": "approved system"}] + ) + [{"role": "user", "content": "approved input"}] + assert shadow_call["max_tokens"] == 123 + assert {key: shadow_call[key] for key in ("tools", "temperature") if key in shadow_call} == ( + {} if remove_optional_fields else { + "temperature": 0.2, + "tools": [{"type": "function", "function": { + "name": "approved_tool", "description": "approved tool", "strict": False, + "parameters": {"type": "object", "properties": {}}, + }}], + } + ) + prisma.db.litellm_shadowevalattempt.create.assert_called_once() + @pytest.mark.parametrize( "call_type,messages,response_obj", [ diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 3d5c38c3acd..23c01841b1b 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -2473,6 +2473,7 @@ def test_success_handler_skips_guardrail_logging_hook_when_disabled(logging_obj) from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger + from litellm.integrations.shadow_eval_logger import GuardrailRequestSnapshot from litellm.types.guardrails import GuardrailEventHooks class DummyGuardrail(CustomGuardrail): @@ -2482,6 +2483,12 @@ def test_success_handler_skips_guardrail_logging_hook_when_disabled(logging_obj) pass logging_obj.stream = False + snapshot: Final = GuardrailRequestSnapshot.capture( + {"messages": [{"role": "user", "content": "approved"}]}, + {"standard_logging_guardrail_information": [{"guardrail_mode": "pre_call"}]}, + ) + assert snapshot is not None + logging_obj.shadow_eval_request_snapshot = snapshot model_response = ModelResponse( id="resp-guardrail-skip", @@ -2523,6 +2530,7 @@ def test_success_handler_skips_guardrail_logging_hook_when_disabled(logging_obj) assert guardrail_call_kwargs["event_type"] == GuardrailEventHooks.logging_only guardrail.logging_hook.assert_not_called() dummy_logger.logging_hook.assert_called_once() + assert logging_obj.shadow_eval_request_snapshot is snapshot def test_success_handler_runs_guardrail_logging_hook_when_enabled(logging_obj): @@ -2530,12 +2538,18 @@ def test_success_handler_runs_guardrail_logging_hook_when_enabled(logging_obj): import datetime from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.integrations.shadow_eval_logger import GuardrailRequestSnapshot from litellm.types.guardrails import GuardrailEventHooks class DummyGuardrail(CustomGuardrail): pass logging_obj.stream = False + logging_obj.shadow_eval_request_snapshot = GuardrailRequestSnapshot.capture( + {"messages": [{"role": "user", "content": "approved"}]}, + {"standard_logging_guardrail_information": [{"guardrail_mode": "pre_call"}]}, + ) + assert logging_obj.shadow_eval_request_snapshot is not None model_response = ModelResponse( id="resp-guardrail-run", @@ -2580,6 +2594,88 @@ def test_success_handler_runs_guardrail_logging_hook_when_enabled(logging_obj): assert guardrail_call_kwargs["event_type"] == GuardrailEventHooks.logging_only guardrail.logging_hook.assert_called_once() assert logging_obj.model_call_details.get("guardrail_hook_ran") is True + assert logging_obj.shadow_eval_request_snapshot is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hook_mode", ["disabled", "mask", "raises"]) +@pytest.mark.parametrize("stream", [False, True]) +async def test_shadow_snapshot_stays_private_and_is_invalidated_before_logging_guardrails( + monkeypatch: pytest.MonkeyPatch, hook_mode: Literal["disabled", "mask", "raises"], stream: bool +) -> None: + from litellm.caching.in_memory_cache import InMemoryCache + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.integrations.shadow_eval_logger import GuardrailRequestSnapshot, ShadowEvalLogger + from litellm.types.guardrails import GuardrailEventHooks + + shadow_snapshots: Final[list[GuardrailRequestSnapshot | None]] = [] + hook_snapshots: Final[list[GuardrailRequestSnapshot | None]] = [] + other_payloads: Final[list[Mapping[str, object]]] = [] + prisma_reads: Final[list[bool]] = [] + + def no_prisma() -> None: + prisma_reads.append(True) + + class RecordingShadowLogger(ShadowEvalLogger): + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, + end_time: object, *, guardrail_snapshot: GuardrailRequestSnapshot | None = None, + ) -> None: + shadow_snapshots.append(guardrail_snapshot) + await super().async_log_success_event( + kwargs, response_obj, start_time, end_time, guardrail_snapshot=guardrail_snapshot + ) + + class RecordingLogger(CustomLogger): + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object, + ) -> None: + other_payloads.append(kwargs) + + class LoggingGuardrail(CustomGuardrail): + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str, + ) -> tuple[dict[str, object], object]: + hook_snapshots.append(logging_obj.shadow_eval_request_snapshot) + if hook_mode == "raises": + raise RuntimeError("logging guardrail failed without recording history") + return {**kwargs, "messages": [{"role": "user", "content": "masked"}]}, result + + metadata: Final = { + "standard_logging_guardrail_information": [{"guardrail_mode": "pre_call"}], + "user_api_key_hash": "test-key", + } + snapshot: Final = GuardrailRequestSnapshot.capture( + {"messages": [{"role": "user", "content": "snapshot-only"}]}, metadata, + ) + assert snapshot is not None + shadow: Final = RecordingShadowLogger(prisma_provider=no_prisma, jobs_cache=InMemoryCache()) + guardrail: Final = LoggingGuardrail( + guardrail_name="late-mask", default_on=True, + event_hook=GuardrailEventHooks.pre_call if hook_mode == "disabled" else GuardrailEventHooks.logging_only, + ) + monkeypatch.setattr(litellm, "_async_success_callback", []) + logging_obj: Final = LitellmLogging( + model="test-model", messages=[], stream=stream, call_type="anthropic_messages", + start_time=datetime.datetime.now(), litellm_call_id="private-snapshot", function_id="private-snapshot", + dynamic_async_success_callbacks=[shadow, RecordingLogger(), guardrail], + ) + logging_obj.update_messages([{"role": "user", "content": "logged input"}]) + logging_obj.update_environment_variables(litellm_params={"metadata": metadata}, optional_params={}) + logging_obj.shadow_eval_request_snapshot = snapshot + payload: Final = { + "id": "private-snapshot", "call_type": "anthropic_messages", "metadata": metadata, + "model_group": "test-model", "model_parameters": {}, + } + + await logging_obj.async_success_handler(result=ModelResponse(), standard_logging_object=payload) + + assert shadow_snapshots == ([snapshot] if hook_mode == "disabled" else [None]) + assert hook_snapshots == ([] if hook_mode == "disabled" else [None]) + assert prisma_reads == ([True] if hook_mode == "disabled" else []) + assert len(other_payloads) == 1 + assert "snapshot-only" not in json.dumps(other_payloads[0], default=str) + assert "snapshot-only" not in json.dumps(logging_obj.model_call_details, default=str) def test_get_user_agent_tags(): diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 0ef02f76a4f..7b74e69685c 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -3,7 +3,7 @@ import copy import datetime import json from types import MappingProxyType, SimpleNamespace -from typing import AsyncGenerator, Callable, Final, Iterator, Optional, Sequence +from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence from urllib.parse import unquote_plus from unittest.mock import AsyncMock, MagicMock, patch @@ -495,60 +495,92 @@ class TestProxyBaseLLMRequestProcessing: add_litellm_data_to_request.assert_not_awaited() @pytest.mark.asyncio + @pytest.mark.parametrize("safe_memory_mode", [False, True]) + @pytest.mark.parametrize( + "route_type,input_key,system_key,token_key", + [ + ("acompletion", "messages", "system", "max_tokens"), + ("anthropic_messages", "messages", "system", "max_tokens"), + ("aresponses", "input", "instructions", "max_output_tokens"), + ], + ) async def test_common_processing_pre_call_logic_refreshes_proxy_server_request_body_after_guardrails( - self, monkeypatch - ): - """ - A guardrail (e.g. Presidio PII masking) mutates data["messages"] in place inside - pre_call_hook. The proxy_server_request.body snapshot is taken before that hook - runs, so it must be refreshed afterward or SpendLogs (when store_prompts_in_spend_logs - is enabled) persists the raw pre-guardrail body, bypassing the masking entirely. - """ - processing_obj = ProxyBaseLLMRequestProcessing(data={}) - mock_request = MagicMock(spec=Request) + self, + monkeypatch: pytest.MonkeyPatch, + safe_memory_mode: bool, + route_type: Literal["acompletion", "anthropic_messages", "aresponses"], + input_key: str, + system_key: str, + token_key: str, + ) -> None: + from litellm.integrations.shadow_eval_logger import request_guardrail_fingerprint + + monkeypatch.setattr(litellm, "safe_memory_mode", safe_memory_mode) + processing_obj: Final = ProxyBaseLLMRequestProcessing(data={}) + mock_request: Final = MagicMock(spec=Request) mock_request.headers = {} + metadata_key: Final = "metadata" if route_type == "acompletion" else "litellm_metadata" + raw_body: Final = { + input_key: [{"role": "user", "content": "private input"}], + system_key: "private system", + "tools": [{"name": "private", "description": "private tool"}], + "tool_choice": {"type": "tool", "name": "private"}, + token_key: 100, + } + approved_messages: Final = [{"role": "user", "content": ""}] + approved_tools: Final = [{"name": "allowed", "description": ""}] + approved_body: Final = {input_key: approved_messages, "tools": approved_tools, token_key: 64} + recorded: Final = [{"guardrail_name": "mask", "guardrail_mode": "pre_call", "guardrail_status": "success"}] - raw_messages = [{"role": "user", "content": "my ssn is 123-45-6789"}] - - async def mock_add_litellm_data_to_request(*args, **kwargs): + async def mock_pre_call_hook( + user_api_key_dict: UserAPIKeyAuth, + data: dict[str, object], + call_type: str, + skip_guardrails: bool = False, + ) -> dict[str, object]: + logging_obj: Final = data["litellm_logging_obj"] + assert isinstance(logging_obj, LiteLLMLoggingObj) + assert logging_obj.shadow_eval_request_snapshot is None return { - "messages": raw_messages, - "proxy_server_request": { - "url": "http://testserver/chat/completions", - "method": "POST", - "body": {"messages": raw_messages}, - }, + **{key: value for key, value in data.items() if key not in (system_key, "tool_choice")}, + **approved_body, + metadata_key: {"standard_logging_guardrail_information": recorded}, } - async def mock_pre_call_hook(user_api_key_dict, data, call_type, skip_guardrails=False): - data["messages"] = [{"role": "user", "content": "my ssn is "}] - return data - - mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj: Final = MagicMock(spec=ProxyLogging) mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=mock_pre_call_hook) monkeypatch.setattr( litellm.proxy.common_request_processing, "add_litellm_data_to_request", - mock_add_litellm_data_to_request, + AsyncMock(return_value={**raw_body, metadata_key: {}, "proxy_server_request": {"body": raw_body}}), ) - returned_data, _ = await processing_obj.common_processing_pre_call_logic( + returned_data, logging_obj = await processing_obj.common_processing_pre_call_logic( request=mock_request, general_settings={}, user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), proxy_logging_obj=mock_proxy_logging_obj, proxy_config=MagicMock(spec=ProxyConfig), - route_type="acompletion", + route_type=route_type, ) - persisted_body = returned_data["proxy_server_request"]["body"] - assert persisted_body["messages"] == returned_data["messages"] - assert "123-45-6789" not in json.dumps(persisted_body["messages"]) - # litellm_logging_obj is stamped onto `data` by function_setup between the - # initial snapshot and pre_call_hook; it must never leak into the persisted - # audit body, which needs to stay plain-JSON-serializable end to end. + proxy_request: Final = returned_data["proxy_server_request"] + persisted_body: Final = proxy_request["body"] + snapshot: Final = logging_obj.shadow_eval_request_snapshot + expected_content: Final = copy.deepcopy(approved_body) + assert snapshot is not None + assert {key: persisted_body[key] for key in raw_body if key in persisted_body} == expected_content + assert {key: snapshot.body[key] for key in raw_body if key in snapshot.body} == expected_content + assert snapshot.fingerprint == request_guardrail_fingerprint( + {"standard_logging_guardrail_information": recorded} + ) assert "litellm_logging_obj" not in persisted_body - json.dumps(persisted_body) + assert "private" not in json.dumps(persisted_body) + approved_messages[0]["content"] = "later input mutation" + approved_tools[0]["description"] = "later tool mutation" + assert {key: snapshot.body[key] for key in raw_body if key in snapshot.body} == expected_content + assert persisted_body[input_key][0]["content"] == "later input mutation" + assert persisted_body["tools"][0]["description"] == "later tool mutation" @staticmethod def _guardrail_tag_budget_harness( diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 0d4b9e8d21f..b2241191ced 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -5,6 +5,7 @@ import os import time from datetime import datetime, timezone from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -807,6 +808,9 @@ async def test_add_litellm_data_to_request_body_snapshot_excludes_proxy_server_r "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hello"}], "api_key": "request-key", + "proxy_server_request": { + "body": {"messages": [{"role": "user", "content": "forged"}]}, + }, } user_api_key_dict = UserAPIKeyAuth( @@ -836,6 +840,77 @@ async def test_add_litellm_data_to_request_body_snapshot_excludes_proxy_server_r ) assert "api_key" not in snapshot_body assert updated["proxy_server_request"]["credential_fields"] == ("api_key",) + assert snapshot_body["messages"] == [{"role": "user", "content": "hello"}] + + +def test_initial_snapshot_refresh_clears_a_previous_guardrail_checkpoint() -> None: + from litellm.integrations.shadow_eval_logger import GuardrailRequestSnapshot + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy.litellm_pre_call_utils import refresh_proxy_server_request_body_snapshot + + logging_obj: Final = Logging( + model="test-model", messages=[], stream=False, call_type="acompletion", + start_time=datetime.now(), litellm_call_id="new-request", function_id="new-request", + ) + logging_obj.shadow_eval_request_snapshot = GuardrailRequestSnapshot.capture( + {"messages": [{"role": "user", "content": "previous request"}]}, + {"standard_logging_guardrail_information": [{"guardrail_mode": "pre_call"}]}, + ) + assert logging_obj.shadow_eval_request_snapshot is not None + proxy_request: Final = {"body": {}} + data: Final = { + "messages": [{"role": "user", "content": "new request"}], + "proxy_server_request": proxy_request, + "litellm_logging_obj": logging_obj, + } + + refresh_proxy_server_request_body_snapshot(data) + + assert logging_obj.shadow_eval_request_snapshot is None + assert proxy_request == {"body": {"messages": [{"role": "user", "content": "new request"}]}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("pre_call_ran", [False, True]) +async def test_post_guardrail_snapshot_preserves_logging_only_masking_in_spend_logs( + monkeypatch: pytest.MonkeyPatch, pre_call_ran: bool +) -> None: + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy.guardrails.guardrail_hooks.presidio import _OPTIONAL_PresidioPIIMasking + from litellm.proxy.litellm_pre_call_utils import refresh_proxy_server_request_body_snapshot + from litellm.proxy.spend_tracking.spend_tracking_utils import _get_proxy_server_request_for_spend_logs_payload + + monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true") + messages: Final = [{"role": "user", "content": "email probe@example.invalid"}] + metadata: Final = { + "standard_logging_guardrail_information": [{"guardrail_mode": "pre_call"}] if pre_call_ran else [] + } + data: Final = {"messages": messages, "metadata": metadata, "proxy_server_request": {}} + logging_obj: Final = Logging( + model="test-model", messages=messages, stream=False, call_type="acompletion", + start_time=datetime.now(), litellm_call_id="mask-spend", function_id="mask-spend", kwargs=data, + ) + data["litellm_logging_obj"] = logging_obj + refresh_proxy_server_request_body_snapshot(data, guardrails_applied=True) + logging_obj.update_messages(messages) + snapshot: Final = logging_obj.shadow_eval_request_snapshot + assert (snapshot is not None) is pre_call_ran + guardrail: Final = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, logging_only=True, mock_redacted_text={"text": "email [EMAIL]", "items": []} + ) + + kwargs, _ = await guardrail.async_logging_hook( + kwargs=logging_obj.model_call_details, result=None, call_type="acompletion" + ) + stored: Final = json.loads(_get_proxy_server_request_for_spend_logs_payload( + metadata={}, litellm_params=kwargs["litellm_params"], kwargs=kwargs, + )) + + assert kwargs["messages"] == [{"role": "user", "content": "email [EMAIL]"}] + assert stored["messages"] == kwargs["messages"] + if snapshot is not None: + assert snapshot.body["messages"] == [{"role": "user", "content": "email probe@example.invalid"}] + assert "probe@example.invalid" not in json.dumps(stored) def test_refresh_proxy_server_request_body_snapshot_picks_up_guardrail_masking(): @@ -2850,7 +2925,7 @@ def test_add_headers_to_llm_call_by_model_group_existing_headers_in_data(): litellm.model_group_settings = original_model_group_settings -from typing import Final, Optional +from typing import Optional from fastapi.responses import Response