diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 48bf1d2675f..4b769325dab 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1,6 +1,7 @@ import contextvars import copy import hashlib +import json import os import secrets from collections.abc import Iterable, Mapping, Sequence @@ -10,7 +11,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args import httpx -from pydantic import TypeAdapter +from pydantic import BaseModel, TypeAdapter from litellm._internal_context import with_service_target from litellm._logging import verbose_logger @@ -940,14 +941,26 @@ class CustomGuardrail(CustomLogger): return _request_carries_marker(data, self._marker(_DEPLOYMENT_HOOK_RAN_TOKEN)) def _post_call_deployment_marker(self, response: object) -> str | None: - response_id: Final[object] = ( - cast("Mapping[object, object]", response).get("id") # cast-ok: only reading a mapping value - if isinstance(response, Mapping) - else getattr(response, "id", None) - ) + try: + response_data: Final[Mapping[str, object]] = ( + response.model_dump(mode="json") + if isinstance(response, BaseModel) + else _REQUEST_FIELDS.validate_python(response) + ) + output: Final = tuple( + (key, response_data[key]) for key in ("choices", "output", "content", "data") if key in response_data + ) + if not output: + return None + digest: Final = hashlib.sha256( + json.dumps(output, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() + ).hexdigest() + except (TypeError, ValueError): + return None + response_id: Final = response_data.get("id") if not isinstance(response_id, str) or not response_id.strip(): return None - return self._marker(f"{_POST_CALL_DEPLOYMENT_HOOK_RAN_TOKEN}:{response_id}") + return self._marker(f"{_POST_CALL_DEPLOYMENT_HOOK_RAN_TOKEN}:{response_id}:{digest}") def post_call_deployment_hook_ran(self, data: Mapping[str, object], response: object) -> bool: return _request_carries_marker(data, self._post_call_deployment_marker(response)) diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index ea9a8f707ed..f6b68eaf8e7 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -3133,9 +3133,32 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook: def test_post_call_marker_requires_usable_id(self, response_id: object) -> None: guardrail: Final = CustomGuardrail(guardrail_name="post-mod") data: Final[dict[str, object]] = {"metadata": {}} - guardrail._record_post_call_deployment_hook_run(data, {"id": response_id}) + guardrail._record_post_call_deployment_hook_run(data, {"id": response_id, "choices": []}) assert data == {"metadata": {}} - assert guardrail.post_call_deployment_hook_ran(data, {"id": response_id}) is False + assert guardrail.post_call_deployment_hook_ran(data, {"id": response_id, "choices": []}) is False + + @pytest.mark.parametrize("output", [None, object(), float("nan")]) + def test_post_call_marker_requires_serializable_output(self, output: object) -> None: + guardrail: Final = CustomGuardrail(guardrail_name="post-mod") + data: Final[dict[str, object]] = {"metadata": {}} + response: Final = {"id": "chatcmpl-same", **({"choices": output} if output is not None else {})} + guardrail._record_post_call_deployment_hook_run(data, response) + assert data == {"metadata": {}} + assert guardrail.post_call_deployment_hook_ran(data, response) is False + + def test_post_call_marker_tracks_tool_calls_but_not_routing_fields(self) -> None: + guardrail: Final = CustomGuardrail(guardrail_name="post-mod") + data: Final[dict[str, object]] = {"metadata": {}} + response: Final = { + "id": "chatcmpl-same", + "model": "deployment", + "choices": [{"message": {"tool_calls": [{"function": {"arguments": "checked"}}]}}], + } + guardrail._record_post_call_deployment_hook_run(data, response) + assert guardrail.post_call_deployment_hook_ran(data, {**response, "model": "alias"}) is True + assert guardrail.post_call_deployment_hook_ran( + data, {**response, "choices": [{"message": {"tool_calls": [{"function": {"arguments": "rewritten"}}]}}]} + ) is False @pytest.mark.asyncio async def test_apply_guardrail_retains_request_identity(self) -> None: diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py index f8fa32cf476..a171b14d6b4 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py @@ -14,7 +14,8 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import CallTypes, UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.utils import GenericGuardrailAPIInputs, ModelResponse +from litellm.types.utils import GenericGuardrailAPIInputs, LLMResponseTypes, ModelResponse +from litellm.utils import async_post_call_success_deployment_hook @pytest.fixture(autouse=True) @@ -86,6 +87,52 @@ async def test_post_call_moderates_returned_response_once( assert guardrail.evaluations == (2 if scenario in ("discarded", "no-id") else 1) +@pytest.mark.asyncio +@pytest.mark.parametrize("parallel", [False, True]) +@pytest.mark.parametrize("rewrite", [False, True]) +async def test_post_call_rechecks_output_after_later_deployment_callback( + proxy_logging: ProxyLogging, monkeypatch: pytest.MonkeyPatch, parallel: bool, rewrite: bool +) -> None: + class Guardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="post-mod", event_hook=GuardrailEventHooks.post_call) + self.evaluations = 0 + self.run_in_parallel = parallel + + async def async_post_call_success_hook( + self, data: dict[str, object], response: LLMResponseTypes, user_api_key_dict: UserAPIKeyAuth + ) -> None: + self.evaluations += 1 + if isinstance(response, ModelResponse) and "restored" in response.model_dump_json(): + raise HTTPException(status_code=400, detail="blocked restored output") + + class Rewriter(CustomLogger): + async def async_post_call_success_deployment_hook( + self, request_data: dict[str, object], response: LLMResponseTypes, call_type: CallTypes | None + ) -> ModelResponse | None: + if not rewrite or not isinstance(response, ModelResponse): + return None + return ModelResponse( + id=response.id, choices=[{"message": {"role": "assistant", "content": "restored"}}] + ) + + guardrail: Final = Guardrail() + monkeypatch.setattr(litellm, "callbacks", [guardrail, Rewriter()]) + data: Final = {"model": "m", "guardrails": ["post-mod"], "metadata": {}} + response: Final = ModelResponse( + id="chatcmpl-same", choices=[{"message": {"role": "assistant", "content": "placeholder"}}] + ) + returned: Final = await async_post_call_success_deployment_hook(data, response, CallTypes.acompletion) + if rewrite: + with pytest.raises(HTTPException, match="blocked restored output"): + await proxy_logging.post_call_success_hook(data=data, user_api_key_dict=UserAPIKeyAuth(), response=returned) + else: + assert await proxy_logging.post_call_success_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), response=returned + ) is returned + assert guardrail.evaluations == (2 if rewrite else 1) + + @pytest.mark.asyncio async def test_post_call_success_hook_returns_response_when_no_callbacks(proxy_logging, make_user_api_key_auth, mock_callbacks_disabled): response = {"original": True, "model": "m", "choices": []}