diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index 31e92201a41..57044b96f44 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -34,21 +34,26 @@ if TYPE_CHECKING: # under, and the fields that hold what the model wrote, which are read whole. _REPLY_TOP_LEVEL_FIELDS: Final = frozenset( ( + "finish_reason", "id", "model", "modelVersion", "object", "responseId", + "role", "service_tier", "status", + "stop_reason", + "stop_sequence", "system_fingerprint", + "type", ) ) _REPLY_NESTED_FIELDS: Final = MappingProxyType( { "candidates": frozenset(("finishReason",)), "choices": frozenset(("finish_reason",)), - "content": frozenset(("role", "type")), + "content": frozenset(("id", "role", "type")), "data": frozenset(("b64_json", "embedding", "object")), "message": frozenset(("role",)), "output": frozenset(("id", "role", "status", "type")), @@ -58,7 +63,19 @@ _REPLY_NESTED_FIELDS: Final = MappingProxyType( } ) _REPLY_MODEL_OUTPUT_FIELDS: Final = frozenset(("args", "arguments", "input", "parameters")) -_REPLY_MARKER_FIELDS: Final = frozenset(("candidates", "choices", "object", "output", "predictions")) +_REPLY_MARKER_FIELDS: Final = frozenset( + ( + "candidates", + "choices", + "finish_reason", + "object", + "output", + "predictions", + "stop_reason", + "usage", + "usageMetadata", + ) +) _IMAGE_REPLY_FIELDS: Final = frozenset(("created", "data")) _NO_FIELDS: Final[frozenset[str]] = frozenset() 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 e19939ddd3e..43c8ce9cc1a 100644 --- a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py @@ -319,6 +319,25 @@ async def test_post_call_scan_leaves_out_the_numbers_of_a_provider_reply(respons ) +@pytest.mark.asyncio +async def test_post_call_scan_reads_only_the_text_of_an_anthropic_message(): + response: Final = { + "id": "msg_5551234567", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [ + {"type": "text", "text": "Paris."}, + {"type": "tool_use", "id": "toolu_5551234567", "name": "lookup", "input": {"card": 4111111111111111}}, + ], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5551234567, "output_tokens": 1}, + } + + assert await _scanned_response_texts(response) == ("Paris.\nlookup\ncard\n4111111111111111",) + + @pytest.mark.asyncio async def test_post_call_scan_skips_a_reply_with_no_text(): response: Final = {"object": "list", "data": [{"object": "embedding", "embedding": [0.4111111111111111]}]}