diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index 57044b96f44..ac91b7bd260 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -63,25 +63,21 @@ _REPLY_NESTED_FIELDS: Final = MappingProxyType( } ) _REPLY_MODEL_OUTPUT_FIELDS: Final = frozenset(("args", "arguments", "input", "parameters")) -_REPLY_MARKER_FIELDS: Final = frozenset( - ( - "candidates", - "choices", - "finish_reason", - "object", - "output", - "predictions", - "stop_reason", - "usage", - "usageMetadata", - ) -) +_REPLY_LIST_FIELDS: Final = ("candidates", "choices", "output", "predictions") _IMAGE_REPLY_FIELDS: Final = frozenset(("created", "data")) _NO_FIELDS: Final[frozenset[str]] = frozenset() def _is_provider_reply(response: Mapping[str, object]) -> bool: - return bool(_REPLY_MARKER_FIELDS & response.keys()) or _IMAGE_REPLY_FIELDS.issubset(response) + """Whether a JSON response has the shape of a provider's reply, by the type of the field that names it, since a + custom API's JSON can use the same field names.""" + return ( + any(isinstance(response.get(field), list) for field in _REPLY_LIST_FIELDS) + or isinstance(response.get("object"), str) + or (response.get("type") == "message" and "stop_reason" in response) + or ("finish_reason" in response and "message" in response) + or _IMAGE_REPLY_FIELDS.issubset(response) + ) def _reply_field_position(name: str, skipped: frozenset[str] | None) -> frozenset[str] | None: diff --git a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py index 43c8ce9cc1a..e028ff2a7bd 100644 --- a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py @@ -286,6 +286,9 @@ async def test_post_call_scan_reads_a_field_named_like_the_envelope_in_what_the_ "4111111111111111", ), ({"choices": [{"message": {"tool_calls": [{"function": {"arguments": {"n": 5551234567}}}]}}]}, "5551234567"), + ({"usage": {}, "customer": {"phone": 5551234567}}, "5551234567"), + ({"output": {"card": 4111111111111111}}, "4111111111111111"), + ({"type": "message", "phone": 5551234567}, "5551234567"), ], ) async def test_post_call_scan_reads_an_integer_in_what_the_model_or_api_wrote(response, expected): @@ -311,6 +314,7 @@ async def test_post_call_scan_reads_a_field_name_in_what_the_model_or_api_wrote( {"created": 1793600123, "data": [{"url": "https://x.example/a.png"}]}, {"id": "batch_1", "object": "batch", "created_at": 1793600123, "request_counts": {"total": 5551234567}}, {"predictions": [{"embeddings": {"values": [0.0123456789]}, "total": 5551234567}]}, + {"finish_reason": "COMPLETE", "message": {"role": "assistant"}, "usage": {"total": 5551234567}}, ], ) async def test_post_call_scan_leaves_out_the_numbers_of_a_provider_reply(response):