mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-18 04:44:37 +00:00
fix(anthropic): carry safeguard_results and stop_details in synthesized message_delta events (#46096)
* fix(anthropic): carry safeguard_results and stop_details in synthesized message_delta events When the proxy renders a complete Anthropic Messages response as an SSE stream for a client that asked to stream, the message_delta event carried only stop_reason and stop_sequence. A real Anthropic stream puts safeguard_results and stop_details inside message_delta.delta, so a client such as Claude Code auto mode lost the server-side classifier verdict and a refusal lost its stop_details. Both stream builders now copy those keys from the source response through one shared helper when they are present. * fix(guardrails): keep safeguard_results and stop_details when re-emitting a rewritten Anthropic stream A post-call guardrail that rewrites an Anthropic stream (presidio, bedrock, tool permission, model armor) now carries the upstream message_delta safeguard_results and stop_details into the re-emitted stream. The router fallback message_start no longer leaks those final fields early * fix(guardrails): keep stop_details and safeguard_results in rewritten Anthropic streams without explanations A guardrail that rewrites a /v1/messages stream now carries stop_details (type and category) and the dangerous_tool_use verdicts for the tool calls it still delivers, with every explanation dropped so a masked stream never echoes the unmasked text through those fields. * fix(guardrails): copy only upstream stop_details and forward-compatible safeguard_results into rewritten Anthropic streams A rewritten stream now takes stop_details and safeguard_results from the upstream message_delta alone, so a refusal the upstream sent without stop_details is never given one by the OpenAI round trip. Unknown safeguard types, statuses, reasons and extra fields pass through, a verdict follows a tool call only when its input is unchanged, and tool_permission keeps the explanations since it rewrites no text. * fix(guardrails): apply the safeguard_results and stop_details rules to policy pipeline stream rewrites * fix(router): deliver a refused stream built from a finished agentic turn instead of falling back Websearch, MCP and advisor streams replay a complete reply whose tools already ran, so a refusal fallback would run every tool again. Also give the recovered response's message_start a null stop_details as the API sends it, drop safeguard_results a masking rewrite cannot read, and add the audit cells for synthesized streams. * test(router): annotate the finished-turn refusal stream test locals with Final --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
22e17b69fa
commit
46ae68021e
21 changed files with 1621 additions and 32 deletions
|
|
@ -30,6 +30,7 @@ from litellm.llms.anthropic.pass_through.adapters.transformation import (
|
|||
LiteLLMAnthropicMessagesAdapter,
|
||||
is_provider_native_tool_dict,
|
||||
)
|
||||
from litellm.llms.anthropic.pass_through.messages.utils import filtered_safeguard_results, redacted_stop_details
|
||||
from litellm.llms.anthropic.pass_through.stream_assembly import (
|
||||
build_complete_streaming_response,
|
||||
build_usage_only_response_from_chunks,
|
||||
|
|
@ -1456,6 +1457,9 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
return _SSEFieldRewrite("delta", "text", next(replacements))
|
||||
|
||||
cls._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta)
|
||||
cls._redact_ended_stream_message_delta(
|
||||
responses_so_far, kept_tool_use_ids=frozenset(cls._streamed_tool_use_ids(responses_so_far).values())
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _write_ended_stream_tool_call_rewrites(
|
||||
|
|
@ -1523,6 +1527,64 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
return None
|
||||
|
||||
cls._rewrite_ended_stream_events(responses_so_far, rewrite_tool_use)
|
||||
streamed_tool_use_ids: Final = cls._streamed_tool_use_ids(responses_so_far)
|
||||
cls._redact_ended_stream_message_delta(
|
||||
responses_so_far,
|
||||
kept_tool_use_ids=frozenset(
|
||||
tool_use_id for index, tool_use_id in streamed_tool_use_ids.items() if index not in rewrites_by_block
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _streamed_tool_use_ids(cls, responses_so_far: Sequence[object]) -> Mapping[int, str]:
|
||||
"""The id of each buffered ``tool_use`` block, by its content block index."""
|
||||
events: Final = chain.from_iterable(cls._iter_sse_events(item) for item in responses_so_far)
|
||||
tool_use_blocks: Final = (
|
||||
(event.get("index"), block)
|
||||
for event in events
|
||||
if event.get("type") == "content_block_start"
|
||||
and (block := as_str_mapping(event.get("content_block"))) is not None
|
||||
and block.get("type") == "tool_use"
|
||||
)
|
||||
return MappingProxyType(
|
||||
{
|
||||
index: tool_use_id
|
||||
for index, block in tool_use_blocks
|
||||
if isinstance(index, int) and isinstance(tool_use_id := block.get("id"), str)
|
||||
}
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _redact_ended_stream_message_delta(
|
||||
cls,
|
||||
responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place
|
||||
*,
|
||||
kept_tool_use_ids: frozenset[str],
|
||||
) -> None:
|
||||
"""Hold the buffered ``message_delta`` to the rules a re-emitted rewrite follows: a verdict survives
|
||||
only for a tool call delivered unchanged, and no free-text ``explanation`` describes the text the
|
||||
guardrail rewrote."""
|
||||
|
||||
def rewrite_stop_details(event: Mapping[str, object]) -> _SSEFieldRewrite | None:
|
||||
delta: Final = as_str_mapping(event.get("delta")) if event.get("type") == "message_delta" else None
|
||||
if delta is None or "stop_details" not in delta:
|
||||
return None
|
||||
return _SSEFieldRewrite(
|
||||
"delta", "stop_details", redacted_stop_details(delta["stop_details"], keep_explanations=False)
|
||||
)
|
||||
|
||||
def rewrite_safeguard_results(event: Mapping[str, object]) -> _SSEFieldRewrite | None:
|
||||
delta: Final = as_str_mapping(event.get("delta")) if event.get("type") == "message_delta" else None
|
||||
if delta is None or "safeguard_results" not in delta:
|
||||
return None
|
||||
return _SSEFieldRewrite(
|
||||
"delta",
|
||||
"safeguard_results",
|
||||
filtered_safeguard_results(delta["safeguard_results"], kept_tool_use_ids, keep_explanations=False),
|
||||
)
|
||||
|
||||
cls._rewrite_ended_stream_events(responses_so_far, rewrite_stop_details)
|
||||
cls._rewrite_ended_stream_events(responses_so_far, rewrite_safeguard_results)
|
||||
|
||||
@staticmethod
|
||||
def _rewrite_ended_stream_events(
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import json
|
|||
from collections.abc import Mapping
|
||||
from typing import Any, Final, cast
|
||||
|
||||
from litellm.llms.anthropic.pass_through.messages.utils import anthropic_message_delta
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import (
|
||||
AnthropicMessagesResponse,
|
||||
)
|
||||
|
|
@ -34,8 +35,13 @@ class FakeAnthropicMessagesStreamIterator:
|
|||
- message_stop
|
||||
"""
|
||||
|
||||
def __init__(self, response: AnthropicMessagesResponse):
|
||||
def __init__(
|
||||
self,
|
||||
response: AnthropicMessagesResponse,
|
||||
message_delta_fields: Mapping[str, object] | None = None,
|
||||
):
|
||||
self.response = response
|
||||
self.message_delta_fields = message_delta_fields
|
||||
self.chunks = self._create_streaming_chunks()
|
||||
self.current_index = 0
|
||||
|
||||
|
|
@ -175,14 +181,9 @@ class FakeAnthropicMessagesStreamIterator:
|
|||
cache_read_input_tokens: Final = usage.get("cache_read_input_tokens")
|
||||
if cache_read_input_tokens is not None:
|
||||
delta_usage["cache_read_input_tokens"] = cache_read_input_tokens
|
||||
safeguard_results: Final = self.response.get("safeguard_results")
|
||||
message_delta: Final = {
|
||||
"type": "message_delta",
|
||||
"delta": {
|
||||
"stop_reason": self.response.get("stop_reason"),
|
||||
"stop_sequence": self.response.get("stop_sequence"),
|
||||
**({"safeguard_results": safeguard_results} if safeguard_results is not None else {}),
|
||||
},
|
||||
"delta": anthropic_message_delta(response_dict, self.message_delta_fields),
|
||||
"usage": delta_usage,
|
||||
}
|
||||
chunks.append(f"event: message_delta\ndata: {json.dumps(message_delta)}\n\n".encode())
|
||||
|
|
|
|||
|
|
@ -18,7 +18,11 @@ from litellm.litellm_core_utils.core_helpers import process_response_headers
|
|||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.llms.anthropic.common_utils import ANTHROPIC_ERROR_STATUS_CODE_MAP
|
||||
from litellm.llms.anthropic.pass_through.messages.utils import INCOMPLETE_STREAM_ERROR_MESSAGE
|
||||
from litellm.llms.anthropic.pass_through.messages.utils import (
|
||||
INCOMPLETE_STREAM_ERROR_MESSAGE,
|
||||
MESSAGE_DELTA_OPTIONAL_KEYS,
|
||||
anthropic_message_delta,
|
||||
)
|
||||
from litellm.llms.anthropic.pass_through.safeguards import StreamedSafeguardResults
|
||||
from litellm.llms.anthropic.pass_through.utils import anthropic_sse_frame
|
||||
from litellm.proxy.pass_through_endpoints.success_handler import (
|
||||
|
|
@ -280,30 +284,32 @@ def anthropic_messages_response_as_sse_events(response: AnthropicMessagesRespons
|
|||
content_events: Final = (
|
||||
event for index, block in enumerate(content_blocks) for event in _anthropic_content_block_events(index, block)
|
||||
)
|
||||
# A real message_start always carries a null stop_reason/stop_sequence and
|
||||
# a zero output_tokens - those are only known once generation finishes, so
|
||||
# copying the completed response's final values here would let a client
|
||||
# treat the message as already finished, or double-count output tokens.
|
||||
# A real message_start always carries a null stop_reason/stop_sequence/stop_details,
|
||||
# a zero output_tokens and no safeguard_results - those are only known once
|
||||
# generation finishes, so copying the completed response's final values here
|
||||
# would let a client treat the message as already finished, or double-count
|
||||
# output tokens.
|
||||
message_start_usage: Final = {
|
||||
**(response.get("usage") or {}),
|
||||
"output_tokens": 0,
|
||||
}
|
||||
message_start_fields: Final = {
|
||||
key: value for key, value in response.items() if key not in MESSAGE_DELTA_OPTIONAL_KEYS
|
||||
}
|
||||
message_start_payload: Final = {
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
**response,
|
||||
**message_start_fields,
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"stop_sequence": None,
|
||||
"stop_details": None,
|
||||
"usage": message_start_usage,
|
||||
},
|
||||
}
|
||||
message_delta_payload: Final = {
|
||||
"type": "message_delta",
|
||||
"delta": {
|
||||
"stop_reason": response.get("stop_reason"),
|
||||
"stop_sequence": response.get("stop_sequence"),
|
||||
},
|
||||
"delta": anthropic_message_delta(response),
|
||||
"usage": response.get("usage") or {},
|
||||
}
|
||||
return (
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ from collections.abc import Iterable, Mapping, Sequence
|
|||
from functools import lru_cache
|
||||
from typing import TYPE_CHECKING, Final, cast, get_type_hints
|
||||
|
||||
from pydantic import JsonValue
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
|
|
@ -49,6 +49,95 @@ def refusal_stop_details(explanation: str | None) -> AnthropicStopDetails:
|
|||
return AnthropicStopDetails(type="refusal", category=None, explanation=explanation)
|
||||
|
||||
|
||||
MESSAGE_DELTA_OPTIONAL_KEYS: Final = ("stop_details", "safeguard_results")
|
||||
|
||||
|
||||
def anthropic_message_delta(
|
||||
response: Mapping[str, object], optional_fields: Mapping[str, object] | None = None
|
||||
) -> Mapping[str, object]:
|
||||
"""
|
||||
The ``delta`` of the ``message_delta`` event synthesized from a complete
|
||||
Anthropic Messages response. A real stream carries ``stop_details`` and
|
||||
``safeguard_results`` there, so a client that asked to stream reads them
|
||||
from this event and nowhere else. ``optional_fields`` replaces the
|
||||
response's own when it was rebuilt from a stream that carried them.
|
||||
"""
|
||||
optional_source: Final = response if optional_fields is None else optional_fields
|
||||
present_optional: Final = {
|
||||
key: value for key in MESSAGE_DELTA_OPTIONAL_KEYS if (value := optional_source.get(key)) is not None
|
||||
}
|
||||
return {
|
||||
"stop_reason": response.get("stop_reason"),
|
||||
"stop_sequence": response.get("stop_sequence"),
|
||||
**present_optional,
|
||||
}
|
||||
|
||||
|
||||
_SAFEGUARD_RESULT_ITEMS: Final = TypeAdapter(tuple[object, ...])
|
||||
|
||||
|
||||
def redacted_stop_details(raw: object, keep_explanations: bool) -> object:
|
||||
"""``stop_details`` as a rewritten stream may deliver it: its free-text ``explanation`` can quote
|
||||
the text a guardrail rewrote, so it is nulled unless the rewrite changed no text."""
|
||||
stop_details: Final = as_str_mapping(raw)
|
||||
if keep_explanations or stop_details is None or "explanation" not in stop_details:
|
||||
return raw
|
||||
return {**stop_details, "explanation": None}
|
||||
|
||||
|
||||
def _without_explanation(raw: object) -> object:
|
||||
mapping: Final = as_str_mapping(raw)
|
||||
if mapping is None:
|
||||
return raw
|
||||
return {key: value for key, value in mapping.items() if key != "explanation"}
|
||||
|
||||
|
||||
def _status_without_explanations(raw: object) -> object:
|
||||
status: Final = as_str_mapping(_without_explanation(raw))
|
||||
if status is None:
|
||||
return raw
|
||||
tool_uses: Final = as_str_mapping(status.get("tool_uses"))
|
||||
if tool_uses is None:
|
||||
return status
|
||||
return {
|
||||
**status,
|
||||
"tool_uses": {tool_use_id: _without_explanation(verdict) for tool_use_id, verdict in tool_uses.items()},
|
||||
}
|
||||
|
||||
|
||||
def _result_without_explanations(raw: object) -> object:
|
||||
result: Final = as_str_mapping(_without_explanation(raw))
|
||||
if result is None:
|
||||
return raw
|
||||
if "status" not in result:
|
||||
return result
|
||||
return {**result, "status": _status_without_explanations(result["status"])}
|
||||
|
||||
|
||||
def _kept_tool_use_verdicts(raw: object, kept_tool_use_ids: frozenset[str]) -> object:
|
||||
result: Final = as_str_mapping(raw)
|
||||
status: Final = as_str_mapping(result.get("status")) if result is not None else None
|
||||
tool_uses: Final = as_str_mapping(status.get("tool_uses")) if status is not None else None
|
||||
if result is None or status is None or tool_uses is None:
|
||||
return raw
|
||||
kept_tool_uses: Final = {
|
||||
tool_use_id: verdict for tool_use_id, verdict in tool_uses.items() if tool_use_id in kept_tool_use_ids
|
||||
}
|
||||
return {**result, "status": {**status, "tool_uses": kept_tool_uses}}
|
||||
|
||||
|
||||
def filtered_safeguard_results(raw: object, kept_tool_use_ids: frozenset[str], keep_explanations: bool) -> object:
|
||||
"""``safeguard_results`` as a rewritten stream may deliver them: a ``tool_uses`` verdict survives only
|
||||
for a tool call the client gets with the input the classifier judged, and the free-text ``explanation`` of a
|
||||
result, its status and each verdict is dropped unless the rewrite changed no text."""
|
||||
try:
|
||||
results: Final = _SAFEGUARD_RESULT_ITEMS.validate_python(raw)
|
||||
except ValidationError:
|
||||
return raw if keep_explanations else None
|
||||
kept: Final = tuple(_kept_tool_use_verdicts(result, kept_tool_use_ids) for result in results)
|
||||
return kept if keep_explanations else tuple(_result_without_explanations(result) for result in kept)
|
||||
|
||||
|
||||
def _mapping_field(container: object, key: str) -> object | None:
|
||||
"""One key of a raw provider payload, or None when the payload is not a mapping."""
|
||||
if not isinstance(container, Mapping):
|
||||
|
|
|
|||
|
|
@ -1,16 +1,18 @@
|
|||
import codecs
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.core_helpers import map_finish_reason
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, map_finish_reason
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.anthropic.chat.handler import (
|
||||
ModelResponseIterator as AnthropicModelResponseIterator,
|
||||
)
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
from litellm.llms.anthropic.pass_through.messages.utils import MESSAGE_DELTA_OPTIONAL_KEYS
|
||||
from litellm.types.utils import (
|
||||
Choices,
|
||||
Message,
|
||||
|
|
@ -447,6 +449,27 @@ def _anthropic_message_start(sse_stream: str) -> Mapping[str, object] | None:
|
|||
)
|
||||
|
||||
|
||||
def anthropic_message_delta_fields(all_chunks: Sequence[object]) -> Mapping[str, object]:
|
||||
"""The ``stop_details`` and ``safeguard_results`` of the stream's ``message_delta``.
|
||||
|
||||
The assembled OpenAI shape has no slot for them, and a client that asked to stream reads them
|
||||
from that event and nowhere else (a Claude Code auto-mode session takes its classifier verdict
|
||||
from ``safeguard_results``), so a caller re-emitting a rewritten response merges them back in.
|
||||
"""
|
||||
sse_stream: Final = _joined_sse_stream(all_chunks)
|
||||
if sse_stream is None:
|
||||
return MappingProxyType({})
|
||||
message_delta: Final = next(
|
||||
(event for event in reversed(_parsed_sse_events(sse_stream)) if event.get("type") == "message_delta"), None
|
||||
)
|
||||
delta: Final = as_str_mapping(message_delta.get("delta")) if message_delta is not None else None
|
||||
if delta is None:
|
||||
return MappingProxyType({})
|
||||
return MappingProxyType(
|
||||
{key: value for key in MESSAGE_DELTA_OPTIONAL_KEYS if (value := delta.get(key)) is not None}
|
||||
)
|
||||
|
||||
|
||||
def is_anthropic_sse_stream(all_chunks: Sequence[object]) -> bool:
|
||||
"""Whether raw SSE frames are Anthropic Messages events.
|
||||
|
||||
|
|
|
|||
|
|
@ -8,8 +8,13 @@ hook scan such a stream, and re-emit it when the guardrail rewrote the response.
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.llms.anthropic.pass_through.messages.utils import filtered_safeguard_results, redacted_stop_details
|
||||
from litellm.llms.anthropic.pass_through.stream_assembly import anthropic_message_delta_fields
|
||||
from litellm.llms.anthropic.pass_through.stream_assembly import (
|
||||
assemble_anthropic_sse_stream as assemble_anthropic_sse_stream, # noqa: PLC0414 # public re-export
|
||||
)
|
||||
|
|
@ -50,7 +55,62 @@ from litellm.llms.anthropic.pass_through.stream_assembly import ( # noqa: E402
|
|||
)
|
||||
|
||||
|
||||
def anthropic_sse_chunks_from_response(assembled: ModelResponse) -> tuple[bytes, ...]:
|
||||
def _tool_use_inputs(content: Sequence[object]) -> Mapping[str, object]:
|
||||
blocks: Final = (block for raw in content if (block := as_str_mapping(raw)) is not None)
|
||||
return MappingProxyType(
|
||||
{
|
||||
tool_use_id: block.get("input")
|
||||
for block in blocks
|
||||
if block.get("type") == "tool_use" and isinstance(tool_use_id := block.get("id"), str)
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _unchanged_tool_use_ids(source_content: Sequence[object], delivered_content: Sequence[object]) -> frozenset[str]:
|
||||
source_inputs: Final = _tool_use_inputs(source_content)
|
||||
return frozenset(
|
||||
tool_use_id
|
||||
for tool_use_id, tool_input in _tool_use_inputs(delivered_content).items()
|
||||
if tool_use_id in source_inputs and source_inputs[tool_use_id] == tool_input
|
||||
)
|
||||
|
||||
|
||||
def _rewritten_message_delta_fields(
|
||||
source_chunks: Sequence[object], unchanged_tool_use_ids: frozenset[str], keep_explanations: bool
|
||||
) -> Mapping[str, object]:
|
||||
source_fields: Final = anthropic_message_delta_fields(source_chunks)
|
||||
stop_details: Final = source_fields.get("stop_details")
|
||||
safeguard_results: Final = source_fields.get("safeguard_results")
|
||||
return MappingProxyType(
|
||||
{
|
||||
**(
|
||||
{"stop_details": redacted_stop_details(stop_details, keep_explanations)}
|
||||
if stop_details is not None
|
||||
else {}
|
||||
),
|
||||
**(
|
||||
{
|
||||
"safeguard_results": filtered_safeguard_results(
|
||||
safeguard_results, unchanged_tool_use_ids, keep_explanations
|
||||
)
|
||||
}
|
||||
if safeguard_results is not None
|
||||
else {}
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def anthropic_sse_chunks_from_response(
|
||||
assembled: ModelResponse, source_chunks: Sequence[object] = (), *, keep_explanations: bool = False
|
||||
) -> tuple[bytes, ...]:
|
||||
"""Re-emit a rewritten response as the SSE events of the ``source_chunks`` stream it replaces.
|
||||
|
||||
``message_delta`` carries the source stream's ``stop_details`` and ``safeguard_results`` and
|
||||
nothing the OpenAI round trip made up, with a verdict kept only for a tool call delivered with
|
||||
the input the classifier judged. Their free-text ``explanation`` was written about the text the
|
||||
guardrail rewrote, so it is dropped unless the caller rewrites no text (``keep_explanations``).
|
||||
"""
|
||||
from litellm.llms.anthropic.pass_through.adapters.transformation import (
|
||||
LiteLLMAnthropicMessagesAdapter,
|
||||
)
|
||||
|
|
@ -58,7 +118,22 @@ def anthropic_sse_chunks_from_response(assembled: ModelResponse) -> tuple[bytes,
|
|||
FakeAnthropicMessagesStreamIterator,
|
||||
)
|
||||
|
||||
anthropic_response: Final = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(
|
||||
response=assembled
|
||||
adapter: Final = LiteLLMAnthropicMessagesAdapter()
|
||||
anthropic_response: Final = adapter.translate_openai_response_to_anthropic(response=assembled)
|
||||
source_response: Final = assemble_anthropic_sse_stream(source_chunks)
|
||||
source_content: Final = (
|
||||
adapter.translate_openai_response_to_anthropic(response=source_response).get("content")
|
||||
if source_response is not None
|
||||
else None
|
||||
)
|
||||
unchanged_tool_use_ids: Final = _unchanged_tool_use_ids(
|
||||
source_content or (), anthropic_response.get("content") or ()
|
||||
)
|
||||
return tuple(
|
||||
FakeAnthropicMessagesStreamIterator(
|
||||
response=anthropic_response,
|
||||
message_delta_fields=_rewritten_message_delta_fields(
|
||||
source_chunks, unchanged_tool_use_ids, keep_explanations
|
||||
),
|
||||
).chunks
|
||||
)
|
||||
return tuple(FakeAnthropicMessagesStreamIterator(response=anthropic_response).chunks)
|
||||
|
|
|
|||
|
|
@ -3058,7 +3058,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
#########################################################################
|
||||
if raw_sse:
|
||||
for sse_chunk in (
|
||||
anthropic_sse_chunks_from_response(assembled_model_response)
|
||||
anthropic_sse_chunks_from_response(assembled_model_response, all_chunks)
|
||||
if model_response_text(assembled_model_response) != pre_guardrail_text
|
||||
else all_chunks
|
||||
):
|
||||
|
|
|
|||
|
|
@ -1263,7 +1263,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
|
||||
# Return sanitized stream
|
||||
if surface is _StreamSurface.ANTHROPIC_MESSAGES:
|
||||
for sse_chunk in anthropic_sse_chunks_from_response(assembled_response):
|
||||
for sse_chunk in anthropic_sse_chunks_from_response(assembled_response, all_chunks):
|
||||
yield sse_chunk
|
||||
return
|
||||
mock_response: Final = MockResponseIterator(model_response=assembled_response)
|
||||
|
|
|
|||
|
|
@ -1532,7 +1532,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
|
|||
await self._process_response_for_pii(response=assembled, request_data=request_data, mode="mask")
|
||||
if model_response_text(assembled) == original_text:
|
||||
return chunks
|
||||
return anthropic_sse_chunks_from_response(assembled)
|
||||
return anthropic_sse_chunks_from_response(assembled, chunks)
|
||||
|
||||
@staticmethod
|
||||
def _unmask_sse_bytes_chunk(chunk: bytes, pii_tokens: dict[str, str]) -> bytes:
|
||||
|
|
|
|||
|
|
@ -971,7 +971,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
|
|||
return
|
||||
|
||||
self._modify_response_with_permission_errors(anthropic_response, anthropic_denials)
|
||||
for sse_chunk in anthropic_sse_chunks_from_response(anthropic_response):
|
||||
for sse_chunk in anthropic_sse_chunks_from_response(anthropic_response, all_chunks, keep_explanations=True):
|
||||
yield sse_chunk
|
||||
|
||||
def _check_assembled_stream(
|
||||
|
|
|
|||
|
|
@ -6055,6 +6055,9 @@ class Router:
|
|||
including pings and provider error frames, is forwarded live and
|
||||
verbatim instead.
|
||||
"""
|
||||
from litellm.llms.anthropic.pass_through.messages.fake_stream_iterator import (
|
||||
FakeAnthropicMessagesStreamIterator,
|
||||
)
|
||||
from litellm.llms.anthropic.pass_through.messages.streaming_iterator import (
|
||||
aclose_if_supported,
|
||||
aclose_superseded_attempt,
|
||||
|
|
@ -6066,6 +6069,7 @@ class Router:
|
|||
)
|
||||
|
||||
source_iterator: Final = response
|
||||
replays_a_finished_agentic_turn: Final = isinstance(source_iterator, FakeAnthropicMessagesStreamIterator)
|
||||
|
||||
async def close_source() -> None:
|
||||
with anyio.CancelScope(shield=True), contextlib.suppress(BaseException):
|
||||
|
|
@ -6126,7 +6130,11 @@ class Router:
|
|||
if not has_generated_content and error_event is None
|
||||
else None
|
||||
)
|
||||
if refusal_stop_details is not None and self._refusal_fallback_available(model, initial_kwargs):
|
||||
if (
|
||||
refusal_stop_details is not None
|
||||
and not replays_a_finished_agentic_turn
|
||||
and self._refusal_fallback_available(model, initial_kwargs)
|
||||
):
|
||||
refusal_error = safeguard_refusal_error(model=model, stop_details=refusal_stop_details)
|
||||
raise MidStreamFallbackError(
|
||||
message=refusal_error.message,
|
||||
|
|
|
|||
|
|
@ -4,28 +4,35 @@ with no verdicts."""
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from itertools import chain
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import object_value
|
||||
from integration._support.anthropic_refusal import DOCS_STOP_DETAILS
|
||||
from integration._support.client import list_value, object_value
|
||||
from integration._support.mcp import McpPeer, mcp_peer, tool_calls
|
||||
from integration._support.process import graceful_stop_seconds
|
||||
from integration._support.safeguards_rig import (
|
||||
BASH_TOOL,
|
||||
CHAT_MODEL,
|
||||
CLAUDE_MODEL,
|
||||
Responder,
|
||||
Rig,
|
||||
auto_mode_body,
|
||||
available,
|
||||
chat_backend,
|
||||
json_body,
|
||||
native_message,
|
||||
nonce,
|
||||
request_json,
|
||||
safeguards_rig,
|
||||
sse_events,
|
||||
stream_results,
|
||||
stream_tool_use_ids,
|
||||
tool_id,
|
||||
tool_use_block,
|
||||
tool_use_ids,
|
||||
)
|
||||
from integration._support.wire import Reply, Request
|
||||
|
|
@ -115,3 +122,45 @@ def test_auto_executed_mcp_tool_costs_no_classifier_call(rig: Rig, peer: McpPeer
|
|||
received: Final = rig.received(request_nonce)
|
||||
assert rig.classifier_calls(received) == (), [request.target for request in received]
|
||||
assert len(rig.backend_calls(received)) == 2, [request.target for request in received]
|
||||
|
||||
|
||||
def carries_a_tool_result_block(request: Request) -> bool:
|
||||
messages: Final = list_value(request_json(request).get("messages", []))
|
||||
contents: Final = tuple(object_value(message).get("content") for message in messages)
|
||||
blocks: Final = chain.from_iterable(content for content in contents if isinstance(content, list))
|
||||
return any(object_value(block).get("type") == "tool_result" for block in blocks)
|
||||
|
||||
|
||||
def claude_mcp_tool_then_refusal(request_nonce: str) -> Responder:
|
||||
identity: Final = tool_id(request_nonce)
|
||||
tool_turn: Final = native_message(
|
||||
identity,
|
||||
content=(tool_use_block(f"{identity}m", name=f"{MCP_ALIAS}-add", arguments={"a": 2, "b": 3}),),
|
||||
model="claude-sonnet-4-6",
|
||||
)
|
||||
refusal_turn: Final = {
|
||||
**native_message(f"{identity}r", content=(), stop_reason="refusal", model="claude-sonnet-4-6"),
|
||||
"stop_details": DOCS_STOP_DETAILS,
|
||||
}
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
turn: Final = refusal_turn if carries_a_tool_result_block(request) else tool_turn
|
||||
return Reply(body=json.dumps(turn).encode())
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def test_auto_executed_mcp_loop_streams_the_final_turn_refusal_stop_details(rig: Rig, peer: McpPeer) -> None:
|
||||
request_nonce: Final = nonce()
|
||||
rig.register(request_nonce, claude_mcp_tool_then_refusal(request_nonce))
|
||||
peer.drain()
|
||||
status, raw = rig.stream(auto_mode_body(CLAUDE_MODEL, request_nonce, tools=[dict(AUTO_EXECUTED), dict(BASH_TOOL)]))
|
||||
assert status == 200, raw
|
||||
assert executed_tools(peer) == ("add",)
|
||||
deltas: Final = tuple(
|
||||
object_value(event["delta"]) for event in sse_events(raw) if event.get("type") == "message_delta"
|
||||
)
|
||||
assert len(deltas) == 1, raw
|
||||
assert deltas[0]["stop_reason"] == "refusal", raw
|
||||
assert deltas[0].get("stop_details") == DOCS_STOP_DETAILS, raw
|
||||
assert len(rig.backend_calls(rig.received(request_nonce))) == 2
|
||||
|
|
|
|||
|
|
@ -0,0 +1,470 @@
|
|||
"""A /v1/messages stream LiteLLM builds from a complete reply (web search interception, a fallback's answer, a guardrail
|
||||
rewrite) carries the ``stop_details`` and ``safeguard_results`` a real Anthropic stream puts on its ``message_delta``."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from typing import Final
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import pytest
|
||||
from integration._support.client import eventually, object_value
|
||||
from integration._support.process import graceful_stop_seconds
|
||||
from integration._support.safeguards_rig import (
|
||||
BASH_TOOL,
|
||||
CLAUDE_MODEL,
|
||||
CLAUDE_TARGET,
|
||||
TAVILY_TARGET,
|
||||
Responder,
|
||||
Rig,
|
||||
auto_mode_body,
|
||||
available,
|
||||
deployment,
|
||||
native_message,
|
||||
native_stream_frames,
|
||||
nonce,
|
||||
request_json,
|
||||
safeguards_rig,
|
||||
sse_events,
|
||||
stream_results,
|
||||
stream_tool_use_ids,
|
||||
text_block,
|
||||
tool_id,
|
||||
tool_use_block,
|
||||
unavailable,
|
||||
)
|
||||
from integration._support.wire import Reply, Request
|
||||
from pydantic import JsonValue
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 240)
|
||||
|
||||
CLAUDE: Final = "anthropic/claude-sonnet-4-6"
|
||||
PRIMARY_MODEL: Final = "synthesized-primary"
|
||||
FALLBACK_MODEL: Final = "synthesized-fallback"
|
||||
REFUSED_MODEL: Final = "synthesized-refused"
|
||||
MOCK_MODEL: Final = "synthesized-mock"
|
||||
PIPELINE_MODEL: Final = "synthesized-pipeline"
|
||||
AZURE_REFUSED_MODEL: Final = "synthesized-azure-refused"
|
||||
CACHED_FALLBACK_MODEL: Final = "synthesized-cached-fallback"
|
||||
PRIMARY_TARGET: Final = "/primary/v1/messages"
|
||||
FALLBACK_TARGET: Final = "/fallback/v1/messages"
|
||||
REFUSED_TARGET: Final = "/refused/v1/messages"
|
||||
PIPELINE_TARGET: Final = "/pipeline/v1/messages"
|
||||
AZURE_REFUSED_TARGET: Final = "/azure/anthropic/v1/messages"
|
||||
CACHED_FALLBACK_TARGET: Final = "/cached-fallback/v1/messages"
|
||||
MOCK_ANSWER: Final = "mocked fallback answer"
|
||||
TOOL_PERMISSION: Final = "synthesized-tool-permission"
|
||||
PIPELINE_REWRITE: Final = "synthesized-pipeline-rewrite"
|
||||
SECRET: Final = "SECRET"
|
||||
MASKED: Final = "[MASKED]"
|
||||
REWRITE_CODE: Final = (
|
||||
"def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
" if input_type != 'response':\n"
|
||||
" return allow()\n"
|
||||
f" texts = [text.replace({SECRET!r}, {MASKED!r}) for text in inputs.get('texts') or []]\n"
|
||||
" calls = [\n"
|
||||
" {'id': call['id'], 'type': 'function', 'function': {'name': call['function']['name'],\n"
|
||||
f" 'arguments': call['function']['arguments'].replace({SECRET!r}, {MASKED!r})}}}}\n"
|
||||
" for call in inputs.get('tool_calls') or []\n"
|
||||
" ]\n"
|
||||
" return modify(texts=texts, tool_calls=calls)\n"
|
||||
)
|
||||
WEB_SEARCH_TOOL: Final[Mapping[str, JsonValue]] = {"type": "web_search_20250305", "name": "web_search", "max_uses": 1}
|
||||
READ_TOOL: Final[Mapping[str, JsonValue]] = {
|
||||
"name": "Read",
|
||||
"description": "Read a file",
|
||||
"input_schema": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]},
|
||||
}
|
||||
NO_CACHE: Final[Mapping[str, JsonValue]] = {"cache": {"no-cache": True}}
|
||||
REFUSAL_STOP_DETAILS: Final[Mapping[str, JsonValue]] = {
|
||||
"type": "refusal",
|
||||
"category": "cyber",
|
||||
"explanation": "This request was declined because it could enable cyber harm.",
|
||||
}
|
||||
_FRAME_BOUNDARY: Final = re.compile(rb"(?<=\n\n)")
|
||||
|
||||
LITELLM_SETTINGS: Final[Mapping[str, JsonValue]] = {
|
||||
"callbacks": ["websearch_interception"],
|
||||
"websearch_interception_params": {"enabled_providers": ["anthropic"], "search_tool_name": "tavily-search"},
|
||||
}
|
||||
|
||||
|
||||
def extra_models(wire_url: str) -> Sequence[Mapping[str, JsonValue]]:
|
||||
return (
|
||||
deployment(PRIMARY_MODEL, CLAUDE, f"{wire_url}/primary"),
|
||||
deployment(FALLBACK_MODEL, CLAUDE, f"{wire_url}/fallback"),
|
||||
deployment(REFUSED_MODEL, CLAUDE, f"{wire_url}/refused"),
|
||||
deployment(MOCK_MODEL, CLAUDE, f"{wire_url}/mock", mock_response=MOCK_ANSWER),
|
||||
deployment(PIPELINE_MODEL, CLAUDE, f"{wire_url}/pipeline"),
|
||||
deployment(AZURE_REFUSED_MODEL, "azure_ai/claude-sonnet-4-6", f"{wire_url}/azure/anthropic"),
|
||||
deployment(CACHED_FALLBACK_MODEL, CLAUDE, f"{wire_url}/cached-fallback"),
|
||||
)
|
||||
|
||||
|
||||
def top_level(wire_url: str) -> Mapping[str, JsonValue]:
|
||||
tavily: Final[Mapping[str, JsonValue]] = {
|
||||
"search_tool_name": "tavily-search",
|
||||
"litellm_params": {
|
||||
"search_provider": "tavily",
|
||||
"api_key": "synthetic-tavily-key",
|
||||
"api_base": f"{wire_url}/tavily",
|
||||
},
|
||||
}
|
||||
tool_permission: Final[Mapping[str, JsonValue]] = {
|
||||
"guardrail_name": TOOL_PERMISSION,
|
||||
"litellm_params": {
|
||||
"guardrail": "tool_permission",
|
||||
"mode": "post_call",
|
||||
"default_on": False,
|
||||
"default_action": "allow",
|
||||
"on_disallowed_action": "rewrite",
|
||||
"rules": [{"id": "deny_bash", "tool_name": "^Bash$", "decision": "deny"}],
|
||||
},
|
||||
}
|
||||
pipeline_rewrite: Final[Mapping[str, JsonValue]] = {
|
||||
"guardrail_name": PIPELINE_REWRITE,
|
||||
"litellm_params": {
|
||||
"guardrail": "custom_code",
|
||||
"mode": "post_call",
|
||||
"default_on": False,
|
||||
"custom_code": REWRITE_CODE,
|
||||
},
|
||||
}
|
||||
masking_policy: Final[Mapping[str, JsonValue]] = {
|
||||
"guardrails": {"add": [PIPELINE_REWRITE]},
|
||||
"pipeline": {
|
||||
"mode": "post_call",
|
||||
"steps": [{"guardrail": PIPELINE_REWRITE, "on_pass": "allow", "on_fail": "block"}],
|
||||
},
|
||||
}
|
||||
return {
|
||||
"search_tools": [dict(tavily)],
|
||||
"guardrails": [dict(tool_permission), dict(pipeline_rewrite)],
|
||||
"policies": {"synthesized-masking": dict(masking_policy)},
|
||||
"policy_attachments": [{"policy": "synthesized-masking", "models": [PIPELINE_MODEL]}],
|
||||
"router_settings": {
|
||||
"disable_cooldowns": True,
|
||||
"content_policy_fallbacks": [
|
||||
{PRIMARY_MODEL: [FALLBACK_MODEL]},
|
||||
{REFUSED_MODEL: [MOCK_MODEL]},
|
||||
{AZURE_REFUSED_MODEL: [CACHED_FALLBACK_MODEL]},
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def synthesizing(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp("synthesized-streams")
|
||||
with safeguards_rig(
|
||||
directory, litellm_settings=LITELLM_SETTINGS, top_level=top_level, extra_models=extra_models
|
||||
) as rig:
|
||||
yield rig
|
||||
|
||||
|
||||
def claude_backend(message: Mapping[str, JsonValue], delta_fields: Mapping[str, JsonValue]) -> Responder:
|
||||
def respond(request: Request) -> Reply:
|
||||
if request_json(request).get("stream") is True:
|
||||
return Reply(
|
||||
content_type="text/event-stream", chunks=native_stream_frames(message, delta_extra=delta_fields)
|
||||
)
|
||||
return Reply(body=json.dumps({**message, **delta_fields}).encode())
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def by_target(backends: Mapping[str, Responder]) -> Responder:
|
||||
def respond(request: Request) -> Reply:
|
||||
return backends[urlsplit(request.target).path](request)
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def refusal(identity: str) -> Mapping[str, JsonValue]:
|
||||
return native_message(identity, content=(), stop_reason="refusal", model="claude-sonnet-4-6")
|
||||
|
||||
|
||||
def answer(identity: str, text: str) -> Mapping[str, JsonValue]:
|
||||
return native_message(identity, content=(text_block(text),), stop_reason="end_turn", model="claude-sonnet-4-6")
|
||||
|
||||
|
||||
def refusal_fields() -> Mapping[str, JsonValue]:
|
||||
return {"stop_details": dict(REFUSAL_STOP_DETAILS), "safeguard_results": list(available({}))}
|
||||
|
||||
|
||||
def search_body(model: str, request_nonce: str) -> Mapping[str, JsonValue]:
|
||||
return auto_mode_body(model, request_nonce, tools=[dict(WEB_SEARCH_TOOL)], prompt="search the web for", **NO_CACHE)
|
||||
|
||||
|
||||
def message_delta(raw: bytes) -> Mapping[str, JsonValue]:
|
||||
deltas: Final = tuple(
|
||||
object_value(event["delta"]) for event in sse_events(raw) if event.get("type") == "message_delta"
|
||||
)
|
||||
assert len(deltas) == 1, raw
|
||||
return deltas[0]
|
||||
|
||||
|
||||
def message_start(raw: bytes) -> Mapping[str, JsonValue]:
|
||||
starts: Final = tuple(
|
||||
object_value(event["message"]) for event in sse_events(raw) if event.get("type") == "message_start"
|
||||
)
|
||||
assert len(starts) == 1, raw
|
||||
return starts[0]
|
||||
|
||||
|
||||
def streamed_text(raw: bytes) -> str:
|
||||
deltas: Final = tuple(
|
||||
object_value(event["delta"]) for event in sse_events(raw) if event.get("type") == "content_block_delta"
|
||||
)
|
||||
return "".join(str(delta["text"]) for delta in deltas if delta.get("type") == "text_delta")
|
||||
|
||||
|
||||
def delta_frame(raw: bytes) -> bytes:
|
||||
frames: Final = tuple(frame for frame in _FRAME_BOUNDARY.split(raw) if b'"message_delta"' in frame)
|
||||
assert len(frames) == 1, raw
|
||||
return frames[0]
|
||||
|
||||
|
||||
def targets(rig: Rig, request_nonce: str) -> tuple[str, ...]:
|
||||
return tuple(urlsplit(request.target).path for request in rig.received(request_nonce))
|
||||
|
||||
|
||||
def test_intercepted_refusal_stream_delivers_the_upstream_stop_details(synthesizing: Rig) -> None:
|
||||
request_nonce: Final = nonce()
|
||||
synthesizing.register(request_nonce, claude_backend(refusal(tool_id(request_nonce)), refusal_fields()))
|
||||
status, raw = synthesizing.stream(search_body(CLAUDE_MODEL, request_nonce))
|
||||
assert status == 200, raw
|
||||
assert message_delta(raw) == {
|
||||
"stop_reason": "refusal",
|
||||
"stop_sequence": None,
|
||||
"stop_details": dict(REFUSAL_STOP_DETAILS),
|
||||
"safeguard_results": list(available({})),
|
||||
}, raw
|
||||
assert targets(synthesizing, request_nonce) == (CLAUDE_TARGET,)
|
||||
|
||||
|
||||
def test_intercepted_refusal_stream_with_a_content_policy_fallback_delivers_the_refusal_without_falling_back(
|
||||
synthesizing: Rig,
|
||||
) -> None:
|
||||
request_nonce: Final = nonce()
|
||||
identity: Final = tool_id(request_nonce)
|
||||
synthesizing.register(
|
||||
request_nonce,
|
||||
by_target(
|
||||
{
|
||||
PRIMARY_TARGET: claude_backend(refusal(identity), refusal_fields()),
|
||||
FALLBACK_TARGET: claude_backend(answer(f"{identity}f", f"fallback answer {request_nonce}"), {}),
|
||||
}
|
||||
),
|
||||
)
|
||||
status, raw = synthesizing.stream(search_body(PRIMARY_MODEL, request_nonce))
|
||||
assert status == 200, raw
|
||||
assert message_delta(raw) == {
|
||||
"stop_reason": "refusal",
|
||||
"stop_sequence": None,
|
||||
"stop_details": dict(REFUSAL_STOP_DETAILS),
|
||||
"safeguard_results": list(available({})),
|
||||
}, raw
|
||||
assert targets(synthesizing, request_nonce) == (PRIMARY_TARGET,)
|
||||
|
||||
|
||||
def test_tool_permission_rewrite_keeps_only_the_delivered_tool_verdict_with_its_explanation(synthesizing: Rig) -> None:
|
||||
request_nonce: Final = nonce()
|
||||
denied: Final = tool_id(request_nonce)
|
||||
allowed: Final = f"{denied}r"
|
||||
allowed_verdict: Final[Mapping[str, JsonValue]] = {
|
||||
"type": "evaluated",
|
||||
"outcome": "flagged",
|
||||
"explanation": f"reads a secret {allowed}",
|
||||
}
|
||||
verdicts: Final[Mapping[str, JsonValue]] = {
|
||||
denied: {"type": "evaluated", "outcome": "flagged", "explanation": f"destructive {denied}"},
|
||||
allowed: dict(allowed_verdict),
|
||||
}
|
||||
message: Final = native_message(
|
||||
denied,
|
||||
content=(
|
||||
text_block("Cleaning up."),
|
||||
tool_use_block(denied),
|
||||
tool_use_block(allowed, name="Read", arguments={"path": "/work/repo/.env"}),
|
||||
),
|
||||
model="claude-sonnet-4-6",
|
||||
)
|
||||
synthesizing.register(request_nonce, claude_backend(message, {"safeguard_results": list(available(verdicts))}))
|
||||
body: Final = auto_mode_body(
|
||||
CLAUDE_MODEL, request_nonce, tools=[dict(BASH_TOOL), dict(READ_TOOL)], guardrails=[TOOL_PERMISSION], **NO_CACHE
|
||||
)
|
||||
status, raw = synthesizing.stream(body)
|
||||
assert status == 200, raw
|
||||
assert stream_tool_use_ids(raw) == (allowed,), raw
|
||||
assert stream_results(raw) == available({allowed: dict(allowed_verdict)}), raw
|
||||
assert targets(synthesizing, request_nonce) == (CLAUDE_TARGET,)
|
||||
|
||||
|
||||
def test_intercepted_answer_stream_keeps_its_safeguard_results(synthesizing: Rig) -> None:
|
||||
request_nonce: Final = nonce()
|
||||
text: Final = f"found it {request_nonce}"
|
||||
results: Final = list(available({}))
|
||||
synthesizing.register(
|
||||
request_nonce,
|
||||
claude_backend(answer(tool_id(request_nonce), text), {"stop_details": None, "safeguard_results": results}),
|
||||
)
|
||||
status, raw = synthesizing.stream(search_body(CLAUDE_MODEL, request_nonce))
|
||||
assert status == 200, raw
|
||||
assert streamed_text(raw) == text, raw
|
||||
assert message_delta(raw) == {"stop_reason": "end_turn", "stop_sequence": None, "safeguard_results": results}, raw
|
||||
received: Final = targets(synthesizing, request_nonce)
|
||||
assert received == (CLAUDE_TARGET,) and TAVILY_TARGET not in received, received
|
||||
|
||||
|
||||
def test_native_refusal_stream_without_a_guardrail_keeps_the_upstream_message_delta_bytes(synthesizing: Rig) -> None:
|
||||
request_nonce: Final = nonce()
|
||||
identity: Final = tool_id(request_nonce)
|
||||
message: Final = native_message(
|
||||
identity, content=(text_block("Partial."),), stop_reason="refusal", model="claude-sonnet-4-6"
|
||||
)
|
||||
fields: Final = {**refusal_fields(), "safeguard_results": list(unavailable((identity,), "error"))}
|
||||
sent: Final = native_stream_frames(message, delta_extra=fields)
|
||||
synthesizing.register(request_nonce, claude_backend(message, fields))
|
||||
status, raw = synthesizing.stream(auto_mode_body(CLAUDE_MODEL, request_nonce, **NO_CACHE))
|
||||
assert status == 200, raw
|
||||
assert delta_frame(raw) == delta_frame(b"".join(sent)), raw
|
||||
assert targets(synthesizing, request_nonce) == (CLAUDE_TARGET,)
|
||||
|
||||
|
||||
def test_native_refusal_falling_back_to_a_mock_reply_renders_a_fresh_message_lifecycle(synthesizing: Rig) -> None:
|
||||
request_nonce: Final = nonce()
|
||||
synthesizing.register(request_nonce, claude_backend(refusal(tool_id(request_nonce)), refusal_fields()))
|
||||
status, raw = synthesizing.stream(auto_mode_body(REFUSED_MODEL, request_nonce, **NO_CACHE))
|
||||
assert status == 200, raw
|
||||
started: Final = message_start(raw)
|
||||
assert started["stop_reason"] is None, raw
|
||||
assert started.get("stop_details") is None and "safeguard_results" not in started, raw
|
||||
assert message_delta(raw) == {"stop_reason": "end_turn", "stop_sequence": None}, raw
|
||||
assert streamed_text(raw) == MOCK_ANSWER, raw
|
||||
assert b'"refusal"' not in raw, raw
|
||||
assert targets(synthesizing, request_nonce) == (REFUSED_TARGET,)
|
||||
|
||||
|
||||
def test_native_refusal_falling_back_to_a_cached_intercepted_reply_streams_its_fields_in_message_delta(
|
||||
synthesizing: Rig,
|
||||
) -> None:
|
||||
request_nonce: Final = nonce()
|
||||
identity: Final = tool_id(request_nonce)
|
||||
text: Final = f"cached answer {request_nonce}"
|
||||
results: Final = list(available({}))
|
||||
synthesizing.register(
|
||||
request_nonce,
|
||||
by_target(
|
||||
{
|
||||
AZURE_REFUSED_TARGET: claude_backend(refusal(identity), refusal_fields()),
|
||||
CACHED_FALLBACK_TARGET: claude_backend(
|
||||
answer(f"{identity}c", text), {"stop_details": None, "safeguard_results": results}
|
||||
),
|
||||
}
|
||||
),
|
||||
)
|
||||
cacheable: Final = auto_mode_body(
|
||||
CACHED_FALLBACK_MODEL, request_nonce, tools=[dict(WEB_SEARCH_TOOL)], prompt="search the web for"
|
||||
)
|
||||
|
||||
def prime() -> tuple[str, ...]:
|
||||
primed: Final = synthesizing.post({**cacheable, "stream": False})
|
||||
assert primed.status_code == 200, primed.text
|
||||
return targets(synthesizing, request_nonce)
|
||||
|
||||
first_prime: Final = prime()
|
||||
assert first_prime == (CACHED_FALLBACK_TARGET,), first_prime
|
||||
eventually(prime, lambda received: received == ())
|
||||
status, raw = synthesizing.stream({**cacheable, "model": AZURE_REFUSED_MODEL})
|
||||
assert status == 200, raw
|
||||
assert streamed_text(raw) == text, raw
|
||||
started: Final = message_start(raw)
|
||||
assert started["stop_details"] is None and "safeguard_results" not in started, raw
|
||||
assert message_delta(raw) == {"stop_reason": "end_turn", "stop_sequence": None, "safeguard_results": results}, raw
|
||||
assert targets(synthesizing, request_nonce) == (AZURE_REFUSED_TARGET,)
|
||||
|
||||
|
||||
def flagged_verdict(explanation: str) -> Mapping[str, JsonValue]:
|
||||
return {"type": "evaluated", "outcome": "flagged", "explanation": explanation}
|
||||
|
||||
|
||||
FLAGGED_WITHOUT_EXPLANATION: Final[Mapping[str, JsonValue]] = {"type": "evaluated", "outcome": "flagged"}
|
||||
|
||||
|
||||
def masked_refusal(identity: str) -> tuple[Mapping[str, JsonValue], Mapping[str, JsonValue], Mapping[str, JsonValue]]:
|
||||
message: Final = native_message(
|
||||
identity, content=(text_block(f"Partial {SECRET}."),), stop_reason="refusal", model="claude-sonnet-4-6"
|
||||
)
|
||||
expected: Final[Mapping[str, JsonValue]] = {
|
||||
"stop_reason": "refusal",
|
||||
"stop_sequence": None,
|
||||
"stop_details": {**REFUSAL_STOP_DETAILS, "explanation": None},
|
||||
"safeguard_results": list(available({})),
|
||||
}
|
||||
return message, refusal_fields(), expected
|
||||
|
||||
|
||||
def masked_tool_turn(
|
||||
identity: str,
|
||||
) -> tuple[Mapping[str, JsonValue], Mapping[str, JsonValue], Mapping[str, JsonValue]]:
|
||||
message: Final = native_message(
|
||||
identity,
|
||||
content=(text_block(f"Reading {SECRET}."), tool_use_block(identity, name="Read", arguments={"path": "/work"})),
|
||||
model="claude-sonnet-4-6",
|
||||
)
|
||||
sent: Final[Mapping[str, JsonValue]] = {
|
||||
"safeguard_results": list(available({identity: dict(flagged_verdict(f"reads {identity}"))}))
|
||||
}
|
||||
expected: Final[Mapping[str, JsonValue]] = {
|
||||
"stop_reason": "tool_use",
|
||||
"stop_sequence": None,
|
||||
"safeguard_results": list(available({identity: dict(FLAGGED_WITHOUT_EXPLANATION)})),
|
||||
}
|
||||
return message, sent, expected
|
||||
|
||||
|
||||
MASKED_STREAMS: Final = {"refusal_stop_details": masked_refusal, "verdict_explanation": masked_tool_turn}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", tuple(MASKED_STREAMS))
|
||||
def test_pipeline_text_rewrite_drops_explanations_from_the_upstream_message_delta(
|
||||
synthesizing: Rig, shape: str
|
||||
) -> None:
|
||||
request_nonce: Final = nonce()
|
||||
message, sent, expected = MASKED_STREAMS[shape](tool_id(request_nonce))
|
||||
synthesizing.register(request_nonce, claude_backend(message, sent))
|
||||
body: Final = auto_mode_body(PIPELINE_MODEL, request_nonce, tools=[dict(READ_TOOL)], **NO_CACHE)
|
||||
status, raw = synthesizing.stream(body)
|
||||
assert status == 200, raw
|
||||
assert MASKED in streamed_text(raw) and SECRET not in streamed_text(raw), raw
|
||||
assert message_delta(raw) == expected, raw
|
||||
assert targets(synthesizing, request_nonce) == (PIPELINE_TARGET,)
|
||||
|
||||
|
||||
def test_pipeline_tool_call_rewrite_keeps_only_the_untouched_tool_verdict(synthesizing: Rig) -> None:
|
||||
request_nonce: Final = nonce()
|
||||
rewritten: Final = tool_id(request_nonce)
|
||||
untouched: Final = f"{rewritten}r"
|
||||
message: Final = native_message(
|
||||
rewritten,
|
||||
content=(
|
||||
tool_use_block(rewritten, arguments={"command": f"cat {SECRET}"}),
|
||||
tool_use_block(untouched, name="Read", arguments={"path": "/work"}),
|
||||
),
|
||||
model="claude-sonnet-4-6",
|
||||
)
|
||||
verdicts: Final[Mapping[str, JsonValue]] = {
|
||||
rewritten: dict(flagged_verdict(f"reads a file {rewritten}")),
|
||||
untouched: dict(flagged_verdict(f"reads a directory {untouched}")),
|
||||
}
|
||||
synthesizing.register(request_nonce, claude_backend(message, {"safeguard_results": list(available(verdicts))}))
|
||||
body: Final = auto_mode_body(PIPELINE_MODEL, request_nonce, tools=[dict(BASH_TOOL), dict(READ_TOOL)], **NO_CACHE)
|
||||
status, raw = synthesizing.stream(body)
|
||||
assert status == 200, raw
|
||||
assert stream_tool_use_ids(raw) == (rewritten, untouched), raw
|
||||
assert SECRET.encode() not in raw and MASKED.encode() in raw, raw
|
||||
assert stream_results(raw) == available({untouched: dict(FLAGGED_WITHOUT_EXPLANATION)}), raw
|
||||
assert targets(synthesizing, request_nonce) == (PIPELINE_TARGET,)
|
||||
|
|
@ -2,7 +2,7 @@ import json
|
|||
import re
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from dataclasses import dataclass
|
||||
|
|
@ -10,6 +10,7 @@ from pathlib import Path
|
|||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.process import OwnedProxy, owned_proxy_process
|
||||
|
|
@ -457,6 +458,96 @@ def test_anthropic_messages_per_request_guardrails_selects_masking(gateway: Gate
|
|||
assert len(rig.analyzer.drain()) == 1
|
||||
|
||||
|
||||
READ_TOOL_USE_ID: Final = "toolu_presidio_read"
|
||||
READ_INPUT: Final = {"path": "/work/report.txt"}
|
||||
|
||||
|
||||
def anthropic_tool_turn_stream(identity: str, text: str, delta: Mapping[str, object]) -> tuple[bytes, ...]:
|
||||
opening, text_start, text_delta, text_stop, _, closing = anthropic_stream(identity, text)
|
||||
tool_use: Final = {"type": "tool_use", "id": READ_TOOL_USE_ID, "name": "Read", "input": {}}
|
||||
return (
|
||||
opening,
|
||||
text_start,
|
||||
text_delta,
|
||||
text_stop,
|
||||
anthropic_frame("content_block_start", {"type": "content_block_start", "index": 1, "content_block": tool_use}),
|
||||
anthropic_frame(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 1,
|
||||
"delta": {"type": "input_json_delta", "partial_json": json.dumps(READ_INPUT)},
|
||||
},
|
||||
),
|
||||
anthropic_frame("content_block_stop", {"type": "content_block_stop", "index": 1}),
|
||||
anthropic_frame("message_delta", {"type": "message_delta", "delta": delta, "usage": {"output_tokens": 4}}),
|
||||
closing,
|
||||
)
|
||||
|
||||
|
||||
def read_verdicts(verdict: Mapping[str, object]) -> Sequence[Mapping[str, object]]:
|
||||
return [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": {READ_TOOL_USE_ID: verdict}}}]
|
||||
|
||||
|
||||
def anthropic_deltas(received: Received, event_type: str) -> tuple[Mapping[str, object], ...]:
|
||||
events: Final = tuple(
|
||||
json.loads(line.removeprefix("data: ")) for line in received.text.split("\n") if line.startswith("data: ")
|
||||
)
|
||||
return tuple(event["delta"] for event in events if event.get("type") == event_type)
|
||||
|
||||
|
||||
def anthropic_message_delta(received: Received) -> Mapping[str, object]:
|
||||
deltas: Final = anthropic_deltas(received, "message_delta")
|
||||
assert len(deltas) == 1, received.text
|
||||
return deltas[0]
|
||||
|
||||
|
||||
def anthropic_text_deltas(received: Received) -> str:
|
||||
deltas: Final = anthropic_deltas(received, "content_block_delta")
|
||||
return "".join(str(delta["text"]) for delta in deltas if delta.get("type") == "text_delta")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("delta", "field", "kept"),
|
||||
[
|
||||
(
|
||||
{
|
||||
"stop_reason": "tool_use",
|
||||
"stop_sequence": None,
|
||||
"safeguard_results": read_verdicts(
|
||||
{"type": "evaluated", "outcome": "flagged", "explanation": f"{PERSON} asked for the report"}
|
||||
),
|
||||
},
|
||||
"safeguard_results",
|
||||
read_verdicts({"type": "evaluated", "outcome": "flagged"}),
|
||||
),
|
||||
(
|
||||
{
|
||||
"stop_reason": "refusal",
|
||||
"stop_sequence": None,
|
||||
"stop_details": {"type": "refusal", "category": "cyber", "explanation": f"declined for {PERSON}"},
|
||||
},
|
||||
"stop_details",
|
||||
{"type": "refusal", "category": "cyber", "explanation": None},
|
||||
),
|
||||
],
|
||||
ids=["verdict_explanation", "stop_details_explanation"],
|
||||
)
|
||||
def test_anthropic_messages_masked_stream_keeps_upstream_safeguard_fields_without_explanations(
|
||||
gateway: Gateway, tmp_path: Path, delta: Mapping[str, object], field: str, kept: object
|
||||
) -> None:
|
||||
identity: Final = "msg_" + uuid.uuid4().hex
|
||||
provider: Final = anthropic_provider(anthropic_tool_turn_stream(identity, f"{PERSON} designed it.", delta))
|
||||
with presidio_rig(gateway, tmp_path, provider) as rig:
|
||||
received: Final = rig.stream("/v1/messages", rig.messages_body())
|
||||
assert received.status == 200, received.text
|
||||
assert anthropic_text_deltas(received) == f"{MASK} designed it.", received.text
|
||||
assert anthropic_message_id(received) == identity, received.text
|
||||
assert anthropic_message_delta(received).get(field) == kept, received.text
|
||||
assert PERSON not in received.text, received.text
|
||||
assert len(rig.upstream.drain()) == 1
|
||||
|
||||
|
||||
def openai_provider(identity: str) -> Callable[[Request], Reply]:
|
||||
def provider(request: Request) -> Reply:
|
||||
assert request.target == "/v1/chat/completions", request.target
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ with guardrail transformations, specifically testing edge cases with empty choic
|
|||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Literal, Optional
|
||||
from typing import Any, Literal, Optional, Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -18,7 +18,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import Streami
|
|||
from litellm.llms.anthropic.chat.guardrail_translation.handler import (
|
||||
AnthropicMessagesHandler,
|
||||
)
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall, GenericGuardrailAPIInputs
|
||||
|
||||
|
||||
class MockPassThroughGuardrail(CustomGuardrail):
|
||||
|
|
@ -383,6 +383,97 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
|
|||
assert '"name": "lookup_fruit"' not in raw
|
||||
assert json.loads("".join(self._partial_jsons(chunks))) == {"fruit": "persimmon"}
|
||||
|
||||
@staticmethod
|
||||
def _message_delta(chunks: list[bytes]) -> dict[str, object]:
|
||||
lines: Final = b"".join(chunks).decode().split("\n")
|
||||
events: Final = (json.loads(line[len("data:") :]) for line in lines if line.startswith("data:"))
|
||||
return next(event["delta"] for event in events if event.get("type") == "message_delta")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_ended_stream_text_rewrite_drops_the_upstream_explanations(self):
|
||||
delta_fields: Final = {
|
||||
"stop_details": {"type": "refusal", "category": "cyber", "explanation": "the reply says world"},
|
||||
"safeguard_results": [
|
||||
{"type": "prompt_injection", "status": {"type": "available", "explanation": "quotes world"}}
|
||||
],
|
||||
}
|
||||
source: Final[list[bytes]] = self._ended_sse_chunks()
|
||||
chunks: Final = [
|
||||
chunk.replace(
|
||||
b'"stop_sequence": null}', b'"stop_sequence": null, ' + json.dumps(delta_fields)[1:-1].encode() + b"}"
|
||||
)
|
||||
for chunk in source
|
||||
]
|
||||
|
||||
await AnthropicMessagesHandler().process_output_streaming_response(
|
||||
responses_so_far=chunks,
|
||||
guardrail_to_apply=self._masking_guardrail(),
|
||||
litellm_logging_obj=MagicMock(),
|
||||
deliver_ended_stream_rewrites=True,
|
||||
)
|
||||
|
||||
assert self._delta_texts(chunks) == ["hello [MASKED]", ""]
|
||||
assert self._message_delta(chunks) == {
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"stop_details": {"type": "refusal", "category": "cyber", "explanation": None},
|
||||
"safeguard_results": [{"type": "prompt_injection", "status": {"type": "available"}}],
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_ended_stream_tool_use_rewrite_drops_the_verdict_of_the_rewritten_call_only(self):
|
||||
class MaskPersimmon(CustomGuardrail):
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict[str, object],
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[LiteLLMLoggingObj] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
for tool_call in inputs.get("tool_calls", []):
|
||||
if isinstance(tool_call, ChatCompletionMessageToolCall) and "persimmon" in tool_call.function.arguments:
|
||||
tool_call.function.arguments = '{"fruit": "[MASKED]"}'
|
||||
return inputs
|
||||
|
||||
first_call: Final[list[bytes]] = self._ended_tool_use_sse_chunks()
|
||||
second_call: Final = [
|
||||
chunk.replace(b'"index": 0', b'"index": 1')
|
||||
.replace(b"toolu_1", b"toolu_2")
|
||||
.replace(b"persim", b"pear")
|
||||
.replace(b'mon\\"}', b'\\"}')
|
||||
for chunk in first_call[1:6]
|
||||
]
|
||||
verdicts: Final = {
|
||||
"toolu_1": {"type": "evaluated", "outcome": "flagged", "explanation": "looks up persimmon"},
|
||||
"toolu_2": {"type": "evaluated", "outcome": "not_flagged", "explanation": "looks up pear"},
|
||||
}
|
||||
safeguard_results: Final = [
|
||||
{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": verdicts}}
|
||||
]
|
||||
message_delta: Final = first_call[6].replace(
|
||||
b'"stop_sequence": null}',
|
||||
b'"stop_sequence": null, "safeguard_results": ' + json.dumps(safeguard_results).encode() + b"}",
|
||||
)
|
||||
chunks: Final = [*first_call[:6], *second_call, message_delta, first_call[7]]
|
||||
|
||||
await AnthropicMessagesHandler().process_output_streaming_response(
|
||||
responses_so_far=chunks,
|
||||
guardrail_to_apply=MaskPersimmon(guardrail_name="test"),
|
||||
litellm_logging_obj=MagicMock(),
|
||||
deliver_ended_stream_rewrites=True,
|
||||
)
|
||||
|
||||
assert self._partial_jsons(chunks)[:3] == ['{"fruit": "[MASKED]"}', "", ""]
|
||||
assert self._message_delta(chunks)["safeguard_results"] == [
|
||||
{
|
||||
"type": "dangerous_tool_use",
|
||||
"status": {
|
||||
"type": "available",
|
||||
"tool_uses": {"toolu_2": {"type": "evaluated", "outcome": "not_flagged"}},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ended_stream_tool_use_rewrite_leaves_chunks_untouched_by_default(self):
|
||||
handler = AnthropicMessagesHandler()
|
||||
|
|
|
|||
|
|
@ -0,0 +1,93 @@
|
|||
import json
|
||||
from collections.abc import Iterable, Mapping
|
||||
from typing import Final
|
||||
|
||||
from litellm.llms.anthropic.pass_through.messages.fake_stream_iterator import (
|
||||
FakeAnthropicMessagesStreamIterator,
|
||||
)
|
||||
|
||||
_CLAUDE_CODE_SAFEGUARD_RESULTS: Final = [
|
||||
{
|
||||
"type": "dangerous_tool_use",
|
||||
"status": {
|
||||
"type": "available",
|
||||
"tool_uses": {
|
||||
"toolu_01": {"type": "evaluated", "outcome": "flagged", "explanation": "deletes the filesystem root"}
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def _decoded_event(chunk: object) -> tuple[str, dict[str, object]]:
|
||||
assert isinstance(chunk, bytes)
|
||||
event_line, data_line = chunk.decode().split("\n")[:2]
|
||||
assert event_line.startswith("event: ")
|
||||
return event_line.removeprefix("event: "), json.loads(data_line.removeprefix("data: "))
|
||||
|
||||
|
||||
def _decoded_events(chunks: Iterable[object]) -> tuple[tuple[str, dict[str, object]], ...]:
|
||||
return tuple(_decoded_event(chunk) for chunk in chunks)
|
||||
|
||||
|
||||
def _message_delta(response: Mapping[str, object]) -> object:
|
||||
decoded: Final = _decoded_events(FakeAnthropicMessagesStreamIterator(response=response))
|
||||
assert tuple(event_type for event_type, _ in decoded)[-2:] == ("message_delta", "message_stop")
|
||||
deltas: Final = tuple(payload["delta"] for event_type, payload in decoded if event_type == "message_delta")
|
||||
assert len(deltas) == 1
|
||||
return deltas[0]
|
||||
|
||||
|
||||
def test_message_delta_carries_safeguard_results():
|
||||
"""Websearch interception and the MCP loop render a complete response as
|
||||
this stream for a client that asked to stream (the guardrail path does too,
|
||||
after merging the upstream fields back in, which
|
||||
tests/unit/proxy/guardrails/test_anthropic_sse.py covers); a Claude Code
|
||||
auto-mode client reads the classifier verdict from message_delta.delta,
|
||||
where a real Anthropic stream puts it."""
|
||||
response: Final = {
|
||||
"id": "msg_1",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"role": "assistant",
|
||||
"type": "message",
|
||||
"content": [{"type": "text", "text": "I won't run that."}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"stop_details": None,
|
||||
"safeguard_results": _CLAUDE_CODE_SAFEGUARD_RESULTS,
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5},
|
||||
}
|
||||
assert _message_delta(response) == {
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"safeguard_results": _CLAUDE_CODE_SAFEGUARD_RESULTS,
|
||||
}
|
||||
|
||||
|
||||
def test_message_delta_carries_refusal_stop_details():
|
||||
stop_details: Final = {"type": "refusal", "category": "cyber", "explanation": "declined"}
|
||||
response: Final = {
|
||||
"id": "msg_2",
|
||||
"content": [],
|
||||
"stop_reason": "refusal",
|
||||
"stop_sequence": None,
|
||||
"stop_details": stop_details,
|
||||
"usage": {"input_tokens": 10, "output_tokens": 0},
|
||||
}
|
||||
assert _message_delta(response) == {
|
||||
"stop_reason": "refusal",
|
||||
"stop_sequence": None,
|
||||
"stop_details": stop_details,
|
||||
}
|
||||
|
||||
|
||||
def test_message_delta_omits_absent_optional_fields():
|
||||
response: Final = {
|
||||
"id": "msg_3",
|
||||
"content": [{"type": "text", "text": "hello"}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"stop_details": None,
|
||||
"usage": {"input_tokens": 10, "output_tokens": 2},
|
||||
}
|
||||
assert _message_delta(response) == {"stop_reason": "end_turn", "stop_sequence": None}
|
||||
|
|
@ -1306,6 +1306,33 @@ def test_anthropic_messages_streaming_response_reports_withheld_output_of_its_st
|
|||
)
|
||||
|
||||
|
||||
def test_anthropic_messages_response_as_sse_events_keeps_message_delta_fields_out_of_message_start():
|
||||
safeguard_results = [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": {}}}]
|
||||
stop_details = {"type": "refusal", "category": "cyber", "explanation": "declined"}
|
||||
response = {
|
||||
"id": "msg_1",
|
||||
"model": "claude-haiku",
|
||||
"role": "assistant",
|
||||
"type": "message",
|
||||
"stop_reason": "refusal",
|
||||
"stop_sequence": None,
|
||||
"stop_details": stop_details,
|
||||
"safeguard_results": safeguard_results,
|
||||
"content": [{"type": "text", "text": "no"}],
|
||||
"usage": {"input_tokens": 3, "output_tokens": 1},
|
||||
}
|
||||
decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response))
|
||||
message_start = decoded[0][1]["message"]
|
||||
assert message_start["stop_details"] is None
|
||||
assert "safeguard_results" not in message_start
|
||||
assert decoded[-2][1]["delta"] == {
|
||||
"stop_reason": "refusal",
|
||||
"stop_sequence": None,
|
||||
"stop_details": stop_details,
|
||||
"safeguard_results": safeguard_results,
|
||||
}
|
||||
|
||||
|
||||
def test_anthropic_messages_response_as_sse_events_no_content_blocks():
|
||||
response = {"id": "msg_4", "content": [], "stop_reason": "end_turn"}
|
||||
decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response))
|
||||
|
|
@ -1530,6 +1557,72 @@ async def test_relay_teardown_dispatches_deferred_billing_when_sentinel_never_co
|
|||
assert deferred_fired.is_set()
|
||||
|
||||
|
||||
_CLAUDE_CODE_SAFEGUARD_RESULTS: Final = [
|
||||
{
|
||||
"type": "dangerous_tool_use",
|
||||
"status": {
|
||||
"type": "available",
|
||||
"tool_uses": {
|
||||
"toolu_01": {"type": "evaluated", "outcome": "flagged", "explanation": "deletes the filesystem root"}
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def _sse_events_message_delta(response: Mapping[str, object]) -> object:
|
||||
decoded: Final = _decode_sse_events(anthropic_messages_response_as_sse_events(response))
|
||||
deltas: Final = tuple(payload["delta"] for event_type, payload in decoded if event_type == "message_delta")
|
||||
assert len(deltas) == 1
|
||||
return deltas[0]
|
||||
|
||||
|
||||
def test_anthropic_messages_response_as_sse_events_message_delta_carries_safeguard_results():
|
||||
"""A Claude Code auto-mode request's classifier verdict rides in
|
||||
message_delta.delta on a real Anthropic stream; a synthesized stream that
|
||||
drops it leaves the client without the verdict."""
|
||||
response = {
|
||||
"id": "msg_7",
|
||||
"content": [{"type": "text", "text": "I won't run that."}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"stop_details": None,
|
||||
"safeguard_results": _CLAUDE_CODE_SAFEGUARD_RESULTS,
|
||||
}
|
||||
assert _sse_events_message_delta(response) == {
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"safeguard_results": _CLAUDE_CODE_SAFEGUARD_RESULTS,
|
||||
}
|
||||
|
||||
|
||||
def test_anthropic_messages_response_as_sse_events_message_delta_carries_refusal_stop_details():
|
||||
stop_details = {"type": "refusal", "category": "cyber", "explanation": "declined"}
|
||||
response = {
|
||||
"id": "msg_8",
|
||||
"content": [],
|
||||
"stop_reason": "refusal",
|
||||
"stop_sequence": None,
|
||||
"stop_details": stop_details,
|
||||
}
|
||||
assert _sse_events_message_delta(response) == {
|
||||
"stop_reason": "refusal",
|
||||
"stop_sequence": None,
|
||||
"stop_details": stop_details,
|
||||
}
|
||||
|
||||
|
||||
def test_anthropic_messages_response_as_sse_events_message_delta_omits_absent_optional_fields():
|
||||
response = {
|
||||
"id": "msg_9",
|
||||
"content": [{"type": "text", "text": "hello"}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"stop_details": None,
|
||||
}
|
||||
assert _sse_events_message_delta(response) == {"stop_reason": "end_turn", "stop_sequence": None}
|
||||
|
||||
|
||||
class _OneVerdictClassifier:
|
||||
def __init__(self, reply: str) -> None:
|
||||
self._reply = reply
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from typing import Final
|
|||
import pytest
|
||||
|
||||
from litellm.llms.anthropic.pass_through.stream_assembly import (
|
||||
anthropic_message_delta_fields,
|
||||
assemble_anthropic_sse_stream,
|
||||
is_anthropic_sse_stream,
|
||||
is_raw_sse_stream,
|
||||
|
|
@ -55,6 +56,35 @@ ANTHROPIC_STREAM: Final = (
|
|||
_frame("message_stop", {"type": "message_stop"}),
|
||||
)
|
||||
|
||||
SAFEGUARD_RESULTS: Final = [
|
||||
{
|
||||
"type": "dangerous_tool_use",
|
||||
"status": {
|
||||
"type": "available",
|
||||
"tool_uses": {
|
||||
"toolu_01": {"type": "evaluated", "outcome": "flagged", "explanation": "deletes the filesystem root"}
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
SAFEGUARDED_STREAM: Final = (
|
||||
*ANTHROPIC_STREAM[:5],
|
||||
_frame(
|
||||
"message_delta",
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"stop_details": None,
|
||||
"safeguard_results": SAFEGUARD_RESULTS,
|
||||
},
|
||||
"usage": {"output_tokens": 2},
|
||||
},
|
||||
),
|
||||
_frame("message_stop", {"type": "message_stop"}),
|
||||
)
|
||||
|
||||
GOOGLE_STREAM: Final = (b'data: {"candidates": [{"content": {"parts": [{"text": "Hello"}]}}]}\n\n',)
|
||||
ANTHROPIC_ERROR: Final = _frame("error", {"type": "error", "error": {"type": "overloaded_error", "message": "busy"}})
|
||||
CHAT_ERROR: Final = b'data: {"error": {"message": "blocked by guardrail"}}\n\n'
|
||||
|
|
@ -108,6 +138,12 @@ def test_assemble_anthropic_sse_stream_restores_upstream_identity_only_when_aske
|
|||
assert (restored.id, restored.model) == (MESSAGE_ID, MODEL)
|
||||
|
||||
|
||||
def test_anthropic_message_delta_fields_reads_the_fields_the_openai_shape_drops():
|
||||
assert dict(anthropic_message_delta_fields(ANTHROPIC_STREAM)) == {}
|
||||
assert dict(anthropic_message_delta_fields(SAFEGUARDED_STREAM)) == {"safeguard_results": SAFEGUARD_RESULTS}
|
||||
assert dict(anthropic_message_delta_fields(GOOGLE_STREAM)) == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"chunks",
|
||||
[GOOGLE_STREAM, INVALID_UTF8, (ANTHROPIC_ERROR,), ()],
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics)
|
|||
import json
|
||||
import logging
|
||||
import re
|
||||
from typing import Final
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -1379,6 +1380,54 @@ class TestToolPermissionGuardrailAnthropicMessages:
|
|||
assert message_start["id"].startswith("chatcmpl-"), "the rewritten stream must not adopt the upstream message id"
|
||||
assert message_start["model"] == "claude-sonnet-4-5", "the rewritten stream reports the model the upstream served"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rewrite_mode_keeps_the_allowed_tool_use_verdict_with_its_explanation(self):
|
||||
"""Claude Code auto mode reads the classifier verdict from message_delta. The rewrite changes
|
||||
no text the explanation could repeat, so the allowed call keeps its verdict whole."""
|
||||
default_on_rewriting: Final = ToolPermissionGuardrail(
|
||||
guardrail_name="anthropic-rewrite-default-on",
|
||||
rules=self.rules,
|
||||
default_action="deny",
|
||||
on_disallowed_action="rewrite",
|
||||
event_hook=GuardrailEventHooks.post_call,
|
||||
default_on=True,
|
||||
)
|
||||
bash_verdict: Final = {"type": "evaluated", "outcome": "flagged", "explanation": "lists the home directory"}
|
||||
safeguard_results: Final = [
|
||||
{
|
||||
"type": "dangerous_tool_use",
|
||||
"status": {
|
||||
"type": "available",
|
||||
"tool_uses": {"tu_ok": bash_verdict, "tu_bad": {"type": "evaluated", "outcome": "not_flagged"}},
|
||||
},
|
||||
}
|
||||
]
|
||||
bash_chunks: Final = self._sse_chunks("Bash", "tu_ok")
|
||||
read_chunks: Final = self._sse_chunks("Read", "tu_bad")
|
||||
read_tool_use: Final = [chunk.replace(b'"index": 1', b'"index": 2') for chunk in read_chunks[4:7]]
|
||||
message_delta: Final = {
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "tool_use", "safeguard_results": safeguard_results},
|
||||
"usage": {"output_tokens": 5},
|
||||
}
|
||||
chunks: Final = [
|
||||
*bash_chunks[:7],
|
||||
*read_tool_use,
|
||||
f"event: message_delta\ndata: {json.dumps(message_delta)}\n\n".encode(),
|
||||
bash_chunks[-1],
|
||||
]
|
||||
|
||||
body: Final = b"".join(await self._drain(default_on_rewriting, chunks)).decode()
|
||||
delta: Final = next(
|
||||
json.loads(line[6:])["delta"]
|
||||
for line in body.splitlines()
|
||||
if line.startswith("data: ") and json.loads(line[6:]).get("type") == "message_delta"
|
||||
)
|
||||
assert "Permission denied" in body
|
||||
assert delta["safeguard_results"] == [
|
||||
{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": {"tu_ok": bash_verdict}}}
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_message_start_without_a_dict_message_fails_closed(self):
|
||||
"""Malformed SSE must not be forwarded unscanned.
|
||||
|
|
|
|||
316
tests/unit/proxy/guardrails/test_anthropic_sse.py
Normal file
316
tests/unit/proxy/guardrails/test_anthropic_sse.py
Normal file
|
|
@ -0,0 +1,316 @@
|
|||
import json
|
||||
from collections.abc import Callable
|
||||
from itertools import chain
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.guardrails.anthropic_sse import (
|
||||
anthropic_sse_chunks_from_response,
|
||||
assemble_anthropic_sse_stream,
|
||||
)
|
||||
from litellm.types.utils import Choices, ModelResponse
|
||||
|
||||
|
||||
def _frame(event: str, data: dict[str, object]) -> bytes:
|
||||
return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode()
|
||||
|
||||
|
||||
def _decoded_event(chunk: bytes) -> tuple[str, dict[str, object]]:
|
||||
event_line, data_line = chunk.decode().split("\n")[:2]
|
||||
assert event_line.startswith("event: ")
|
||||
return event_line.removeprefix("event: "), json.loads(data_line.removeprefix("data: "))
|
||||
|
||||
|
||||
def _decode(chunks: tuple[bytes, ...]) -> tuple[tuple[str, dict[str, object]], ...]:
|
||||
return tuple(_decoded_event(chunk) for chunk in chunks)
|
||||
|
||||
|
||||
def _rewritten_delta(
|
||||
upstream: tuple[bytes, ...], rewrite: Callable[[ModelResponse], None], keep_explanations: bool = False
|
||||
) -> object:
|
||||
assembled: Final = assemble_anthropic_sse_stream(upstream, restore_identity=True)
|
||||
assert assembled is not None
|
||||
rewrite(assembled)
|
||||
rewritten: Final = anthropic_sse_chunks_from_response(assembled, upstream, keep_explanations=keep_explanations)
|
||||
return next(payload["delta"] for event, payload in _decode(rewritten) if event == "message_delta")
|
||||
|
||||
|
||||
def _set_text(text: str) -> Callable[[ModelResponse], None]:
|
||||
def rewrite(assembled: ModelResponse) -> None:
|
||||
choice: Final = assembled.choices[0]
|
||||
assert isinstance(choice, Choices)
|
||||
choice.message.content = text
|
||||
|
||||
return rewrite
|
||||
|
||||
|
||||
def _tool_use_frames(index: int, tool_use_id: str, name: str) -> tuple[bytes, ...]:
|
||||
return (
|
||||
_frame(
|
||||
"content_block_start",
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": index,
|
||||
"content_block": {"type": "tool_use", "id": tool_use_id, "name": name, "input": {}},
|
||||
},
|
||||
),
|
||||
_frame(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": index,
|
||||
"delta": {"type": "input_json_delta", "partial_json": '{"command": "ls"}'},
|
||||
},
|
||||
),
|
||||
_frame("content_block_stop", {"type": "content_block_stop", "index": index}),
|
||||
)
|
||||
|
||||
|
||||
def _upstream_stream(
|
||||
delta: dict[str, object], text: str = "call 555-0100", tool_uses: tuple[tuple[str, str], ...] = ()
|
||||
) -> tuple[bytes, ...]:
|
||||
tool_use_frames = tuple(
|
||||
chain.from_iterable(
|
||||
_tool_use_frames(index, tool_use_id, name) for index, (tool_use_id, name) in enumerate(tool_uses, start=1)
|
||||
)
|
||||
)
|
||||
return (
|
||||
_frame(
|
||||
"message_start",
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_upstream_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-upstream",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 3, "output_tokens": 1},
|
||||
},
|
||||
},
|
||||
),
|
||||
_frame(
|
||||
"content_block_start",
|
||||
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
||||
),
|
||||
_frame(
|
||||
"content_block_delta",
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}},
|
||||
),
|
||||
_frame("content_block_stop", {"type": "content_block_stop", "index": 0}),
|
||||
*tool_use_frames,
|
||||
_frame("message_delta", {"type": "message_delta", "delta": delta, "usage": {"output_tokens": 2}}),
|
||||
_frame("message_stop", {"type": "message_stop"}),
|
||||
)
|
||||
|
||||
|
||||
def _safeguard_results(tool_uses: dict[str, dict[str, str]]) -> list[dict[str, object]]:
|
||||
return [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": tool_uses}}]
|
||||
|
||||
|
||||
SAFEGUARD_RESULTS: Final = _safeguard_results({})
|
||||
REFUSAL_STOP_DETAILS: Final = {"type": "refusal", "category": "cyber", "explanation": None}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("stop_reason", "delta_fields"),
|
||||
[
|
||||
("end_turn", {"safeguard_results": SAFEGUARD_RESULTS}),
|
||||
("refusal", {"stop_details": REFUSAL_STOP_DETAILS}),
|
||||
],
|
||||
ids=["safeguard_results", "refusal_stop_details"],
|
||||
)
|
||||
def test_rewritten_stream_keeps_the_upstream_message_delta_fields(stop_reason: str, delta_fields: dict[str, object]):
|
||||
"""A guardrail that rewrites a streamed response assembles it into the OpenAI shape, which has no
|
||||
slot for ``safeguard_results`` or ``stop_details``, so the re-emitted ``message_delta`` used to
|
||||
drop the classifier verdict a Claude Code auto-mode session reads from it."""
|
||||
upstream_delta: dict[str, object] = {"stop_reason": stop_reason, "stop_sequence": None, **delta_fields}
|
||||
upstream = _upstream_stream(upstream_delta)
|
||||
assembled = assemble_anthropic_sse_stream(upstream, restore_identity=True)
|
||||
assert assembled is not None
|
||||
choice = assembled.choices[0]
|
||||
assert isinstance(choice, Choices)
|
||||
choice.message.content = "call <PHONE>"
|
||||
|
||||
decoded = _decode(anthropic_sse_chunks_from_response(assembled, upstream))
|
||||
payloads = {event: payload for event, payload in decoded}
|
||||
|
||||
assert payloads["content_block_delta"]["delta"] == {"type": "text_delta", "text": "call <PHONE>"}
|
||||
assert payloads["message_delta"]["delta"] == upstream_delta
|
||||
assert not set(delta_fields) & set(payloads["message_start"]["message"])
|
||||
|
||||
|
||||
SSN: Final = "123-45-6789"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("upstream_delta", "tool_uses", "expected_fields"),
|
||||
[
|
||||
(
|
||||
{
|
||||
"stop_reason": "refusal",
|
||||
"stop_sequence": None,
|
||||
"stop_details": {"type": "refusal", "category": "cyber", "explanation": f"Request involves SSN {SSN}"},
|
||||
},
|
||||
(),
|
||||
{"stop_details": {"type": "refusal", "category": "cyber", "explanation": None}},
|
||||
),
|
||||
(
|
||||
{
|
||||
"stop_reason": "tool_use",
|
||||
"stop_sequence": None,
|
||||
"safeguard_results": _safeguard_results(
|
||||
{"toolu_bash": {"type": "evaluated", "outcome": "flagged", "explanation": f"leaks SSN {SSN}"}}
|
||||
),
|
||||
},
|
||||
(("toolu_bash", "Bash"),),
|
||||
{"safeguard_results": _safeguard_results({"toolu_bash": {"type": "evaluated", "outcome": "flagged"}})},
|
||||
),
|
||||
(
|
||||
{
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"safeguard_results": [
|
||||
{
|
||||
"type": "prompt_injection",
|
||||
"status": {"type": "available", "explanation": f"quotes SSN {SSN}"},
|
||||
"explanation": f"saw SSN {SSN}",
|
||||
}
|
||||
],
|
||||
},
|
||||
(),
|
||||
{"safeguard_results": [{"type": "prompt_injection", "status": {"type": "available"}}]},
|
||||
),
|
||||
],
|
||||
ids=["stop_details_explanation", "verdict_explanation", "unknown_shape_explanation"],
|
||||
)
|
||||
def test_masked_stream_drops_the_unmasked_upstream_explanations(
|
||||
upstream_delta: dict[str, object], tool_uses: tuple[tuple[str, str], ...], expected_fields: dict[str, object]
|
||||
):
|
||||
upstream = _upstream_stream(upstream_delta, text=f"my SSN is {SSN}", tool_uses=tool_uses)
|
||||
assembled = assemble_anthropic_sse_stream(upstream, restore_identity=True)
|
||||
assert assembled is not None
|
||||
choice = assembled.choices[0]
|
||||
assert isinstance(choice, Choices)
|
||||
choice.message.content = "my SSN is <US_SSN>"
|
||||
|
||||
rewritten = anthropic_sse_chunks_from_response(assembled, upstream)
|
||||
message_delta = next(payload for event, payload in _decode(rewritten) if event == "message_delta")
|
||||
|
||||
assert not [chunk for chunk in rewritten if SSN.encode() in chunk]
|
||||
assert {key: message_delta["delta"][key] for key in expected_fields} == expected_fields
|
||||
|
||||
|
||||
def test_masked_stream_drops_safeguard_results_it_cannot_read():
|
||||
upstream_delta: Final = {
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"safeguard_results": {"type": "dangerous_tool_use", "explanation": f"saw SSN {SSN}"},
|
||||
}
|
||||
|
||||
assert _rewritten_delta(_upstream_stream(upstream_delta), _set_text("my SSN is <US_SSN>")) == {
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
}
|
||||
|
||||
|
||||
def test_rewritten_stream_keeps_only_the_verdicts_of_tool_uses_it_still_delivers():
|
||||
kept_verdict = {"type": "evaluated", "outcome": "not_flagged"}
|
||||
upstream = _upstream_stream(
|
||||
{
|
||||
"stop_reason": "tool_use",
|
||||
"stop_sequence": None,
|
||||
"safeguard_results": _safeguard_results(
|
||||
{"toolu_glob": {"type": "evaluated", "outcome": "flagged"}, "toolu_bash": kept_verdict}
|
||||
),
|
||||
},
|
||||
tool_uses=(("toolu_glob", "Glob"), ("toolu_bash", "Bash")),
|
||||
)
|
||||
assembled = assemble_anthropic_sse_stream(upstream)
|
||||
assert assembled is not None
|
||||
choice = assembled.choices[0]
|
||||
assert isinstance(choice, Choices)
|
||||
choice.message.tool_calls = [call for call in choice.message.tool_calls or [] if call.id == "toolu_bash"]
|
||||
choice.message.content = "Permission denied: Tool 'Glob' denied"
|
||||
|
||||
decoded = _decode(anthropic_sse_chunks_from_response(assembled, upstream))
|
||||
message_delta = next(payload for event, payload in decoded if event == "message_delta")
|
||||
delivered_tool_use_ids = [
|
||||
payload["content_block"]["id"]
|
||||
for event, payload in decoded
|
||||
if event == "content_block_start" and payload["content_block"]["type"] == "tool_use"
|
||||
]
|
||||
|
||||
assert delivered_tool_use_ids == ["toolu_bash"]
|
||||
assert message_delta["delta"]["safeguard_results"] == _safeguard_results({"toolu_bash": kept_verdict})
|
||||
|
||||
|
||||
def test_rewritten_refusal_gets_no_stop_details_the_upstream_did_not_send():
|
||||
"""The OpenAI round trip sets ``stop_details`` on every refusal, and the repo reads a refusal
|
||||
carrying one as a safeguard refusal, so a plain refusal must not come back as one."""
|
||||
upstream_delta: Final = {"stop_reason": "refusal", "stop_sequence": None}
|
||||
|
||||
assert _rewritten_delta(_upstream_stream(upstream_delta), _set_text("call <PHONE>")) == upstream_delta
|
||||
|
||||
|
||||
def test_rewritten_stream_keeps_safeguard_shapes_it_does_not_know():
|
||||
upstream_delta: Final = {
|
||||
"stop_reason": "tool_use",
|
||||
"stop_sequence": None,
|
||||
"stop_details": {"type": "pause_turn", "category": None, "explanation": None, "policy_url": "https://x.test"},
|
||||
"safeguard_results": [
|
||||
{
|
||||
"type": "dangerous_tool_use",
|
||||
"status": {
|
||||
"type": "available",
|
||||
"tool_uses": {"toolu_bash": {"type": "unavailable", "reason": "overloaded"}},
|
||||
},
|
||||
"classifier_version": "2026-10",
|
||||
},
|
||||
{"type": "prompt_injection", "status": {"type": "skipped"}},
|
||||
],
|
||||
}
|
||||
upstream: Final = _upstream_stream(upstream_delta, tool_uses=(("toolu_bash", "Bash"),))
|
||||
|
||||
assert _rewritten_delta(upstream, _set_text("call <PHONE>")) == upstream_delta
|
||||
|
||||
|
||||
def test_rewrite_that_changes_no_text_keeps_the_upstream_explanations():
|
||||
upstream_delta: Final = {
|
||||
"stop_reason": "tool_use",
|
||||
"stop_sequence": None,
|
||||
"safeguard_results": _safeguard_results(
|
||||
{"toolu_bash": {"type": "evaluated", "outcome": "flagged", "explanation": "lists the home directory"}}
|
||||
),
|
||||
}
|
||||
upstream: Final = _upstream_stream(upstream_delta, tool_uses=(("toolu_bash", "Bash"),))
|
||||
|
||||
assert _rewritten_delta(upstream, _set_text("Permission denied"), keep_explanations=True) == upstream_delta
|
||||
|
||||
|
||||
def test_rewritten_stream_drops_the_verdict_of_a_tool_use_whose_input_the_guardrail_changed():
|
||||
"""A verdict scored on the upstream input must not vouch for the input the client runs."""
|
||||
verdict: Final = {"type": "evaluated", "outcome": "not_flagged"}
|
||||
upstream: Final = _upstream_stream(
|
||||
{
|
||||
"stop_reason": "tool_use",
|
||||
"stop_sequence": None,
|
||||
"safeguard_results": _safeguard_results({"toolu_masked": verdict, "toolu_untouched": verdict}),
|
||||
},
|
||||
tool_uses=(("toolu_masked", "Bash"), ("toolu_untouched", "Bash")),
|
||||
)
|
||||
|
||||
def mask_first_tool_input(assembled: ModelResponse) -> None:
|
||||
choice: Final = assembled.choices[0]
|
||||
assert isinstance(choice, Choices)
|
||||
tool_calls: Final = choice.message.tool_calls or []
|
||||
assert [call.id for call in tool_calls] == ["toolu_masked", "toolu_untouched"]
|
||||
tool_calls[0].function.arguments = '{"command": "<REDACTED>"}'
|
||||
|
||||
rewritten_delta: Final = _rewritten_delta(upstream, mask_first_tool_input)
|
||||
|
||||
assert isinstance(rewritten_delta, dict)
|
||||
assert rewritten_delta["safeguard_results"] == _safeguard_results({"toolu_untouched": verdict})
|
||||
|
|
@ -14,6 +14,7 @@ import respx
|
|||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.llms.anthropic.pass_through.messages.fake_stream_iterator import FakeAnthropicMessagesStreamIterator
|
||||
from litellm.litellm_core_utils import get_llm_provider_logic
|
||||
from litellm.router_utils.cooldown_handlers import mark_advisor_orchestration_failure
|
||||
from litellm.router_utils.fallback_event_handlers import (
|
||||
|
|
@ -32,6 +33,7 @@ from litellm.router_utils.fallback_event_handlers import (
|
|||
is_bad_request_without_fallback,
|
||||
log_failure_fallback_event,
|
||||
log_success_fallback_event,
|
||||
mid_stream_fallback_hop_kwargs,
|
||||
mid_stream_fallback_snapshot_kwargs,
|
||||
mid_stream_retry_kwargs,
|
||||
record_pre_routing_selection,
|
||||
|
|
@ -1865,6 +1867,41 @@ async def test_streaming_refusal_without_fallback_row_passes_frames_through():
|
|||
assert len(fake.calls) == 1
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
|
||||
@pytest.mark.asyncio
|
||||
async def test_refused_stream_built_from_a_finished_agentic_turn_is_delivered_without_a_fallback():
|
||||
"""Websearch, MCP and advisor loops run their tools before streaming the final reply, so a
|
||||
fallback would run every tool call again for a reply the client then gets from another model."""
|
||||
fake: Final = FakeAnthropicUpstream()
|
||||
router: Final = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}])
|
||||
|
||||
with fake.install():
|
||||
stream: Final = await router._aanthropic_messages_streaming_iterator(
|
||||
response=FakeAnthropicMessagesStreamIterator(REFUSAL_RESPONSE),
|
||||
initial_kwargs=mid_stream_fallback_hop_kwargs(
|
||||
model="fable-tier",
|
||||
original_generic_function=litellm.anthropic_messages,
|
||||
controls=None,
|
||||
kwargs={
|
||||
"model": "fable-tier",
|
||||
"max_tokens": 16,
|
||||
"stream": True,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
},
|
||||
),
|
||||
)
|
||||
events: Final = [
|
||||
json.loads(line[len(b"data: ") :])
|
||||
for line in (await _collect(stream)).splitlines()
|
||||
if line.startswith(b"data: ")
|
||||
]
|
||||
|
||||
assert fake.calls == []
|
||||
message_delta: Final = next(event for event in events if event["type"] == "message_delta")
|
||||
assert message_delta["delta"]["stop_reason"] == "refusal"
|
||||
assert message_delta["delta"]["stop_details"] == REFUSAL_RESPONSE["stop_details"]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_refusal_on_routed_tier_matches_tier_keyed_row_without_inbound_metadata():
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue