diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 48925c06408..b4a1ddb22c3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -237,13 +237,26 @@ def _bedrock_items(value: object) -> Sequence[object]: return _BEDROCK_ITEMS.validate_python(value) if isinstance(value, list) else () +def _bedrock_search_result_text(fields: Mapping[str, object]) -> str: + """The joined text of a Converse ``searchResult`` block or an Anthropic ``search_result`` block, else nothing.""" + search_result: Final = ( + fields if fields.get("type") == "search_result" else as_str_mapping(fields.get("searchResult")) + ) + entries: Final = _bedrock_items((search_result or {}).get("content")) + texts: Final = ((as_str_mapping(entry) or {}).get("text") for entry in entries) + return "\n".join(text for text in texts if isinstance(text, str) and text) + + def _bedrock_text_block(block: object) -> object: - """A Converse ``guardContent`` text block or ``json`` tool-result block as a plain ``{"text": ...}`` block, so its - text is inspected like any other; every other block is returned as given.""" + """A Converse ``guardContent`` text block, ``json`` tool-result block or search result as a plain ``{"text": ...}`` + block, so its text is inspected like any other; every other block is returned as given.""" fields: Final = as_str_mapping(block) or {} guard_text: Final = as_str_mapping((as_str_mapping(fields.get("guardContent")) or {}).get("text")) or {} if isinstance(guard_text.get("text"), str): return {"text": guard_text["text"]} + search_text: Final = _bedrock_search_result_text(fields) + if search_text: + return {"text": search_text} return {"text": json.dumps(fields["json"])} if "json" in fields else block diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py index 6f1fe3256f4..d002fec9f20 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py @@ -2587,6 +2587,26 @@ class TestBedrockPassthroughBody: {"role": "tool", "content": '{"card": "on file"}'}, ] + def test_reads_converse_and_anthropic_search_results(self): + converse_result: Final = {"source": "s", "title": "t", "content": [{"text": "first"}, {"text": "second"}]} + tool_result: Final = {"toolUseId": "t1", "content": [{"searchResult": converse_result}]} + anthropic_result: Final = { + "type": "search_result", + "source": "s", + "title": "t", + "content": [{"type": "text", "text": "third"}], + } + body: Final = { + "messages": [ + {"role": "user", "content": [{"toolResult": tool_result}]}, + {"role": "user", "content": [anthropic_result]}, + ], + } + assert _build_lakera_inspection_messages(_bedrock_data(body), call_type="allm_passthrough_route") == [ + {"role": "tool", "content": "first\nsecond"}, + {"role": "user", "content": "third"}, + ] + def test_a_nested_body_for_another_provider_is_not_read(self): data: Final = {**_bedrock_data(_invoke_body()), "custom_llm_provider": "vertex_ai"} assert _build_lakera_inspection_messages(data, call_type="allm_passthrough_route") == []