mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
fix(guardrails): scan empty top-level system text blocks too
The hoisted structured row keeps every text block of the top-level system prompt, empty ones included, while the scanned texts dropped the empty ones. Guardrails that count one text per slot then came back with more texts than the handler could place, so their rewrite was rejected. User text blocks were already scanned empty or not; the system prompt now matches.
This commit is contained in:
parent
04410967ff
commit
1b594fc935
2 changed files with 43 additions and 1 deletions
|
|
@ -721,7 +721,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
return tuple(
|
||||
ScannedText(text_str, SystemBlockTextTarget(block_idx))
|
||||
for block_idx, block in enumerate(content)
|
||||
if isinstance(block, dict) and isinstance(text_str := block.get("text"), str) and text_str
|
||||
if isinstance(block, dict) and isinstance(text_str := block.get("text"), str)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -2499,6 +2499,29 @@ class PerRowTextGuardrail(CustomGuardrail):
|
|||
return {**inputs, "texts": [str(row.get("content")).replace("123-45-6789", "<US_SSN>") for row in rows]}
|
||||
|
||||
|
||||
class PerSlotTextGuardrail(CustomGuardrail):
|
||||
"""Answers one redacted text per text slot of every chat row it was shown, the
|
||||
way a guardrail that counts slots per message does, and hands back only texts."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(guardrail_name="per-slot-redactor")
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import message_slot_texts
|
||||
|
||||
rows = inputs.get("structured_messages") or []
|
||||
return {
|
||||
**inputs,
|
||||
"texts": [text.replace("123-45-6789", "<US_SSN>") for row in rows for text in message_slot_texts(row)],
|
||||
}
|
||||
|
||||
|
||||
class TestPerMessageTextWriteBack:
|
||||
"""Texts that no longer pair one-to-one with what the handler extracted must be
|
||||
rejected by name instead of sliding onto the wrong messages."""
|
||||
|
|
@ -2537,6 +2560,25 @@ class TestPerMessageTextWriteBack:
|
|||
assert data["system"] == original["system"], "a rejected rewrite must leave the request untouched"
|
||||
assert data["messages"] == original["messages"], "a rejected rewrite must leave the request untouched"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_one_text_per_slot_over_a_system_prompt_with_an_empty_block_is_applied(self):
|
||||
data = {
|
||||
"model": "claude-sonnet-4-5",
|
||||
"system": [
|
||||
{"type": "text", "text": ""},
|
||||
{"type": "text", "text": "Reply with exactly the SSN you were given."},
|
||||
],
|
||||
"messages": [{"role": "user", "content": "My SSN is 123-45-6789."}],
|
||||
}
|
||||
|
||||
await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=PerSlotTextGuardrail())
|
||||
|
||||
assert data["system"] == [
|
||||
{"type": "text", "text": ""},
|
||||
{"type": "text", "text": "Reply with exactly the SSN you were given."},
|
||||
]
|
||||
assert data["messages"] == [{"role": "user", "content": "My SSN is <US_SSN>."}]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_one_text_per_row_without_a_system_prompt_is_applied(self):
|
||||
data = {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue