diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 502c3cc95e0..5220dc576fc 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -145,6 +145,16 @@ def _unified_hook_fields(guardrail: "CustomGuardrail", request_data: Mapping[str } +def _inspected_choice_index(request_data: Mapping[str, object]) -> int | None: + """The choice this round inspected, as the OpenAI guardrail handler labelled it. + + Read behind a type guard: the key rides on the request body, so anything reaching + here that is not a plain int was not written by that handler. + """ + labelled: Final = request_data.get("guardrail_choice_index") + return labelled if isinstance(labelled, int) and not isinstance(labelled, bool) else None + + def _instance_attributes(value: object) -> Mapping[str, object]: return vars(value) @@ -1414,7 +1424,7 @@ class CustomGuardrail(CustomLogger): end_time=end_time, duration=duration, masked_entity_count=masked_entity_count, - choice_index=request_data.get("guardrail_choice_index"), + choice_index=_inspected_choice_index(request_data), **(tracing_detail or {}), ) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 8dc7501bc3b..ee072fd5f80 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -794,6 +794,26 @@ def _strip_client_pricing_overrides(data: dict[str, object]) -> None: ) +# ``guardrail_choice_index`` is written by the OpenAI guardrail handler around each +# per-choice inspection round and read by guardrail spend logs as fact; a client value +# has no legitimate meaning and no key or team setting keeps it. +_GUARDRAIL_RESERVED_REQUEST_FIELDS: Final = frozenset({"guardrail_choice_index"}) + + +def _strip_guardrail_reserved_fields( + data: dict[str, object], # mutable-ok: strips in place on the request body the pre-call pipeline threads through +) -> None: + """Drop the guardrail-owned per-choice label from the client request body.""" + stripped: Final = tuple(sorted(field for field in _GUARDRAIL_RESERVED_REQUEST_FIELDS if field in data)) + for field in stripped: + data.pop(field, None) + if stripped: + verbose_proxy_logger.debug( + "Stripped guardrail-owned fields from request body: %s. The proxy writes these per round.", + ", ".join(stripped), + ) + + def _strip_router_reserved_metadata( data: dict[str, object], # mutable-ok: strips in place on the request body the pre-call pipeline threads through ) -> None: @@ -2650,6 +2670,7 @@ async def add_litellm_data_to_request( if not _key_or_team_allows_client_pricing_override(user_api_key_dict): _strip_client_pricing_overrides(data) _strip_router_reserved_metadata(data) + _strip_guardrail_reserved_fields(data) # Same reason as the strips above: runs after the metadata string-to-dict parse # so JSON-string metadata cannot smuggle callback credentials past the dict guard. diff --git a/tests/unit/proxy/guardrails/test_guardrail_choice_index_label.py b/tests/unit/proxy/guardrails/test_guardrail_choice_index_label.py new file mode 100644 index 00000000000..45f69ad03d8 --- /dev/null +++ b/tests/unit/proxy/guardrails/test_guardrail_choice_index_label.py @@ -0,0 +1,73 @@ +"""The per-choice label on guardrail spend logs is proxy-owned, never caller-supplied. + +`guardrail_choice_index` names which choice a guardrail round inspected. The OpenAI +handler writes it around each round; `StandardLoggingGuardrailInformation` reads it and +spend logs persist it. A caller putting the key in its own request body would otherwise +have that value stored verbatim, with any type or size, on every guardrail row of the +request. +""" + +from typing import Final +from unittest.mock import MagicMock + +import pytest +from fastapi import Request + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + +def _request_mock() -> Request: + request_mock = MagicMock(spec=Request) + request_mock.url = MagicMock() + request_mock.url.path = "/v1/chat/completions" + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + return request_mock + + +def _logged_choice_index(request_data: dict[str, object]) -> object: + guardrail: Final = CustomGuardrail(guardrail_name="choice-label-test") + guardrail.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={}, + request_data=request_data, + guardrail_status="success", + ) + metadata: Final = request_data["metadata"] + assert isinstance(metadata, dict) + rows: Final = metadata["standard_logging_guardrail_information"] + assert isinstance(rows, list) + return rows[0]["choice_index"] + + +@pytest.mark.asyncio +async def test_proxy_drops_a_caller_supplied_choice_index() -> None: + data: Final[dict[str, object]] = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "guardrail_choice_index": "x" * 4096, + } + + updated: Final = await add_litellm_data_to_request( + data=data, + request=_request_mock(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert "guardrail_choice_index" not in updated + + +def test_a_non_integer_label_is_not_logged() -> None: + assert _logged_choice_index({"guardrail_choice_index": "1; drop"}) is None + + +def test_an_integer_label_is_logged() -> None: + assert _logged_choice_index({"guardrail_choice_index": 1}) == 1