diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index f2b24984ccf..721a6653597 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -499,6 +499,39 @@ class AnthropicMessagesHandler(BaseTranslation): {"role": "system", "content": blocks} if blocks else None # mutable-ok: API message payload ) # mutable-ok: API message payload + @staticmethod + def _fold_leading_systems_into_top_level( + data: dict[str, object], # mutable-ok: API message payload + leading_systems: Sequence[object], + include_existing_system: bool, + ) -> None: + """Deliver leading system rows through Anthropic's top-level system param, which rejects them in messages.""" + existing: Final = data.get("system") if include_existing_system else None + existing_blocks: Final[list[object]] = ( # mutable-ok: API message payload + [{"type": "text", "text": existing}] + if isinstance(existing, str) and existing + else list(existing) + if isinstance(existing, list) + else [] + ) + converted_rows: Final = tuple( + AnthropicMessagesHandler._openai_system_message_to_anthropic(message) + for message in leading_systems + if isinstance(message, dict) + ) + folded: Final[list[object]] = existing_blocks + [ # mutable-ok: API message payload + block + for row in converted_rows + if row is not None + for block in ( + [{"type": "text", "text": row["content"]}] if isinstance(row["content"], str) else row["content"] + ) + ] + if folded: + data["system"] = folded # rebind-ok: write-back mutates the request payload in place + else: + data.pop("system", None) + @staticmethod def _is_hoisted_top_level_system(message: object, hoisted_system_message: object) -> bool: """Match the hoisted prompt by identity, or by value after serialization.""" @@ -575,9 +608,24 @@ class AnthropicMessagesHandler(BaseTranslation): ) ordered: Final = AnthropicMessagesHandler._defer_systems_inside_tool_exchanges(structured_messages) + leading_count: Final = next( + (index for index, message in enumerate(ordered) if not _is_system(message)), + len(ordered), + ) + leading_systems: Final = ordered[:leading_count] + hoisted_in_leading: Final = any( + AnthropicMessagesHandler._is_hoisted_top_level_system(message, hoisted_system_message) + for message in leading_systems + ) + if leading_systems and not (leading_count == 1 and hoisted_in_leading): + AnthropicMessagesHandler._fold_leading_systems_into_top_level( + data, + leading_systems, + include_existing_system=hoisted_system_message is None, + ) run: Final[list] = [] # mutable-ok: API message payload - hoisted_dropped = False # rebind-ok: flips once the hoisted prompt is dropped - for message in ordered: + hoisted_dropped = hoisted_in_leading # rebind-ok: flips once the hoisted prompt is dropped + for message in ordered[leading_count:]: if not _is_system(message): run.append(message) continue diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index cefbaf17d57..b219dcba491 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -121,6 +121,45 @@ class MockCompactingGuardrail(CustomGuardrail): return rewritten +class MockStructuredMaskingGuardrail(CustomGuardrail): + """Mask an email in texts and in a rebuilt structured view, like a PII-masking guardrail (LIT-5696).""" + + def __init__(self): + super().__init__(guardrail_name="structured-masking-test") + + @staticmethod + def _mask(text: str) -> str: + return text.replace("bob@example.com", "") + + def _mask_content(self, content: object) -> object: + if isinstance(content, str): + return self._mask(content) + if not isinstance(content, list): + return content + return [ + {**block, "text": self._mask(block["text"])} + if isinstance(block, dict) and isinstance(block.get("text"), str) + else block + for block in content + ] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + masked = inputs.copy() + masked["texts"] = [self._mask(text) for text in inputs.get("texts", [])] + structured = inputs.get("structured_messages") + if structured is not None: + masked["structured_messages"] = [ + {**message, "content": self._mask_content(message.get("content"))} for message in structured + ] + return masked + + class TestAnthropicMessagesHandlerStreamingRequestData: """Post-call guardrails on streaming /v1/messages receive the response and identity metadata""" @@ -602,7 +641,7 @@ class TestAnthropicMessagesHandlerInputProcessing: assert data["system"] == "trusted top-level system prompt" @pytest.mark.asyncio - async def test_compaction_rewrite_keeps_leading_midturn_system_when_system_is_skipped( + async def test_leading_system_row_appends_to_skipped_top_level_system( self, ): handler = AnthropicMessagesHandler() @@ -624,11 +663,14 @@ class TestAnthropicMessagesHandlerInputProcessing: await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) - assert [m["role"] for m in data["messages"]] == ["system", "user"] - assert data["messages"][0]["content"] == "use the corrected result" + assert [m["role"] for m in data["messages"]] == ["user"] + assert data["system"] == [ + {"type": "text", "text": "trusted top-level system prompt"}, + {"type": "text", "text": "use the corrected result"}, + ] @pytest.mark.asyncio - async def test_compaction_rewrite_keeps_leading_correction_when_top_level_system_hoists_nothing( + async def test_leading_correction_appends_when_top_level_system_hoists_nothing( self, ): handler = AnthropicMessagesHandler() @@ -650,11 +692,14 @@ class TestAnthropicMessagesHandlerInputProcessing: await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) - assert [m["role"] for m in data["messages"]] == ["system", "user"] - assert data["messages"][0]["content"] == "use the corrected result" + assert [m["role"] for m in data["messages"]] == ["user"] + assert data["system"] == [ + {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}, + {"type": "text", "text": "use the corrected result"}, + ] @pytest.mark.asyncio - async def test_compaction_rewrite_keeps_leading_correction_when_hoisted_prompt_is_dropped( + async def test_leading_correction_replaces_top_level_system_when_hoisted_prompt_is_dropped( self, ): handler = AnthropicMessagesHandler() @@ -681,9 +726,81 @@ class TestAnthropicMessagesHandlerInputProcessing: "role": "system", "content": "TRUSTED", } - assert [m["role"] for m in data["messages"]] == ["system", "user"] - assert data["messages"][0]["content"] == "CLIENT CORRECTION" - assert data["system"] == "TRUSTED" + assert [m["role"] for m in data["messages"]] == ["user"] + assert data["system"] == [{"type": "text", "text": "CLIENT CORRECTION"}] + + @pytest.mark.asyncio + async def test_masked_hoisted_system_folds_into_top_level_system(self): + """LIT-5696: a guardrail-modified top-level prompt must go back through the system + param; emitting it as messages[0] is rejected by Anthropic, dropping it leaks the + unmasked original.""" + handler = AnthropicMessagesHandler() + guardrail = MockStructuredMaskingGuardrail() + data = { + "model": "claude-3-5-sonnet-20241022", + "system": [{"type": "text", "text": "You are helpful. The admin is bob@example.com."}], + "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data["system"] == [{"type": "text", "text": "You are helpful. The admin is ."}] + assert [m["role"] for m in data["messages"]] == ["user"] + + @pytest.mark.asyncio + async def test_client_leading_system_row_folds_into_top_level_system(self): + """LIT-5696: a client-sent leading system row folds into the system param instead of + being sent back as messages[0], which Anthropic rejects.""" + handler = AnthropicMessagesHandler() + guardrail = MockStructuredMaskingGuardrail() + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + {"role": "system", "content": [{"type": "text", "text": "You are helpful."}]}, + {"role": "user", "content": [{"type": "text", "text": "hi"}]}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data["system"] == [{"type": "text", "text": "You are helpful."}] + assert [m["role"] for m in data["messages"]] == ["user"] + + @pytest.mark.asyncio + async def test_masked_midturn_system_after_user_stays_in_messages(self): + handler = AnthropicMessagesHandler() + guardrail = MockStructuredMaskingGuardrail() + data = { + "model": "claude-3-5-sonnet-20241022", + "system": [{"type": "text", "text": "You are helpful. The admin is bob@example.com."}], + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "hi"}]}, + {"role": "assistant", "content": [{"type": "text", "text": "hello"}]}, + {"role": "system", "content": [{"type": "text", "text": "Mid-turn: admin bob@example.com"}]}, + {"role": "user", "content": [{"type": "text", "text": "next"}]}, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data["system"] == [{"type": "text", "text": "You are helpful. The admin is ."}] + assert [m["role"] for m in data["messages"]] == ["user", "assistant", "system", "user"] + assert data["messages"][2]["content"] == [{"type": "text", "text": "Mid-turn: admin "}] + + @pytest.mark.asyncio + async def test_unmodified_structured_copy_leaves_top_level_system_untouched(self): + handler = AnthropicMessagesHandler() + guardrail = MockStructuredMaskingGuardrail() + data = { + "model": "claude-3-5-sonnet-20241022", + "system": "You are helpful.", + "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data["system"] == "You are helpful." + assert [m["role"] for m in data["messages"]] == ["user"] @pytest.mark.asyncio async def test_compaction_rewrite_drops_hoisted_prompt_matched_by_content_copy(self): @@ -934,8 +1051,8 @@ class TestAnthropicMessagesHandlerInputProcessing: with patch.object(litellm, "modify_params", True): await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) - assert [m["role"] for m in data["messages"]] == ["system", "user"] - assert data["messages"][0]["content"] == "use the corrected result" + assert data["messages"] == [{"role": "user", "content": [{"type": "text", "text": "Please continue."}]}] + assert data["system"] == [{"type": "text", "text": "use the corrected result"}] @pytest.mark.asyncio async def test_compaction_rewrite_without_system_messages_is_unchanged(self):