scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

fix(guardrails): inspect Anthropic /v1/messages output for Akamai FAI

Browse source
This commit is contained in:
Scott Jacobsen 2026-07-28 09:16:26 -05:00
parent 1e08938b14
commit 720d6082b8
2 changed files with 116 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
litellm/proxy/guardrails/guardrail_hooks/akamai_firewall_for_ai/akamai_firewall_for_ai.py
View file

@ -184,6 +184,31 @@ def _iter_responses_api_output_text(response: ResponsesAPIResponse) -> Iterator[
yield from _iter_function_fragments(item)
def _iter_anthropic_output_text(content: Any) -> Iterator[str]:
"""Yield text and tool-call payloads from an Anthropic ``/v1/messages`` reply.
The non-streaming ``/v1/messages`` response reaches the hook as a native
dict whose generated text lives in ``content[].text`` and whose tool calls
live in ``content[].input`` (``type == "tool_use"``); neither is reachable
via the Chat-Completions ``choices`` or the Responses-API ``output`` shapes.
"""
if not isinstance(content, list):
return
for block in content:
block_type = _item_get(block, "type")
if block_type == "text":
text = _item_get(block, "text")
if isinstance(text, str) and text:
yield text
elif block_type == "tool_use":
name = _item_get(block, "name")
if isinstance(name, str) and name:
yield name
tool_input = _item_get(block, "input")
if isinstance(tool_input, dict) and tool_input:
yield json.dumps(tool_input, sort_keys=True)
class AkamaiRuleTriggered(TypedDict, total=False):
action: str
category: str
@ -276,6 +301,8 @@ class AkamaiFirewallForAIGuardrail(CustomGuardrail):
return get_content_from_model_response(response)
if isinstance(response, ResponsesAPIResponse):
return "\n".join(_iter_responses_api_output_text(response))
if isinstance(response, dict) and response.get("type") == "message":
return "\n".join(_iter_anthropic_output_text(response.get("content")))
return ""
async def _detect(
@ -390,14 +417,17 @@ class AkamaiFirewallForAIGuardrail(CustomGuardrail):
return response
@classmethod
def _streaming_output_text(cls, chunks: list) -> str:
def _streaming_output_text(cls, chunks: list, request_data: dict) -> str:
"""Extract inspectable output text from a fully buffered stream.
Chat streams (``ModelResponse`` / ``ModelResponseStream`` chunks) are
assembled with ``stream_chunk_builder``. Responses-API streams instead
emit events, the terminal one of which carries the complete
``ResponsesAPIResponse``; reuse ``_output_text`` on it so streamed
Responses output and tool calls are inspected as well.
Responses output and tool calls are inspected as well. Anthropic
``/v1/messages`` streams arrive as raw SSE ``bytes``; the shared
passthrough assembler rebuilds them into a ``ModelResponse`` so streamed
Anthropic text and tool calls are inspected through the same path.
"""
if isinstance(chunks[0], (ModelResponse, ModelResponseStream)):
from litellm.main import stream_chunk_builder
@ -405,6 +435,18 @@ class AkamaiFirewallForAIGuardrail(CustomGuardrail):
assembled = stream_chunk_builder(chunks=chunks)
return cls._output_text(assembled) if isinstance(assembled, ModelResponse) else ""
if isinstance(chunks[0], (bytes, str)):
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
)
assembled = AnthropicPassthroughLoggingHandler._build_complete_streaming_response(
all_chunks=chunks,
litellm_logging_obj=request_data.get("litellm_logging_obj"),
model=str(request_data.get("model") or ""),
)
return cls._output_text(assembled) if isinstance(assembled, ModelResponse) else ""
for chunk in reversed(chunks):
candidate = _item_get(chunk, "response")
if isinstance(candidate, ResponsesAPIResponse):
@ -429,7 +471,7 @@ class AkamaiFirewallForAIGuardrail(CustomGuardrail):
try:
await self._detect(
client_request_id=self._client_request_id(request_data),
llm_output=self._streaming_output_text(chunks),
llm_output=self._streaming_output_text(chunks, request_data),
)
except HTTPException as exc:
error_obj = dict(exc.detail) if isinstance(exc.detail, dict) else {"message": str(exc.detail)}

71
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_akamai_firewall_for_ai.py
View file

@ -736,3 +736,74 @@ async def test_streaming_hook_blocks_responses_api_stream():
# the Responses events are withheld; only the SSE block is emitted
assert all(not isinstance(chunk, (OutputTextDeltaEvent, ResponseCompletedEvent)) for chunk in yielded)
assert len(yielded) == 1 and "Blocked by Akamai Firewall for AI" in yielded[0]
@pytest.mark.asyncio
async def test_output_hook_inspects_anthropic_messages_response():
"""Regression: /v1/messages returns a native Anthropic dict, not a ModelResponse.
Before the fix ``_output_text`` returned "" for that shape, so the generated
text and tool_use arguments were released without a detect request. Both the
text block and the tool_use input must be sent to Akamai and blocked.
"""
guardrail = _init("post_call")
data = {"litellm_call_id": "req-1", "guardrails": ["akamai-guard"], "messages": [{"role": "user", "content": "hi"}]}
response = {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-6",
"content": [
{"type": "text", "text": "here is the plan"},
{"type": "tool_use", "id": "tu1", "name": "exfiltrate", "input": {"secret": "AKIA-super-secret"}},
],
"stop_reason": "end_turn",
}
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new=AsyncMock(return_value=_response(BLOCK_BODY)),
) as mock_post:
with pytest.raises(HTTPException):
await guardrail.async_post_call_success_hook(
data=data, user_api_key_dict=UserAPIKeyAuth(), response=response
)
llm_output = mock_post.call_args.kwargs["json"]["llmOutput"]
assert "here is the plan" in llm_output
assert "AKIA-super-secret" in llm_output
assert "exfiltrate" in llm_output
@pytest.mark.asyncio
async def test_streaming_hook_blocks_anthropic_messages_stream():
"""A streamed /v1/messages reply arrives as raw Anthropic SSE bytes.
Those bytes are not ModelResponse chunks nor Responses events, so before the
fix the stream was released uninspected. The shared passthrough assembler
must rebuild them into a ModelResponse, the generated text scanned, and a
blocking verdict withhold the bytes before delivery.
"""
guardrail = _init("post_call")
request_data = {"litellm_call_id": "req-1", "guardrails": ["akamai-guard"], "model": "claude-sonnet-4-6"}
events = [
b'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-6","content":[],"stop_reason":null,"usage":{"input_tokens":1,"output_tokens":1}}}\n\n',
b'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}\n\n',
b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"here is a SECRET_STREAM_PAYLOAD"}}\n\n',
b'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n',
b'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":5}}\n\n',
b'event: message_stop\ndata: {"type":"message_stop"}\n\n',
]
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new=AsyncMock(return_value=_response(BLOCK_BODY)),
) as mock_post:
yielded = [
chunk
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=UserAPIKeyAuth(), response=_aiter(events), request_data=request_data
)
]
assert "SECRET_STREAM_PAYLOAD" in mock_post.call_args.kwargs["json"]["llmOutput"]
# none of the raw Anthropic SSE bytes are delivered
assert all(not isinstance(chunk, (bytes, bytearray)) for chunk in yielded)
assert len(yielded) == 1 and "Blocked by Akamai Firewall for AI" in yielded[0]