mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
fix(websearch-interception)!: stream native server-tool blocks to clients (#40967)
* fix(websearch-interception): stream native server-tool blocks to clients Streaming native clients (Anthropic SDK, Claude Desktop, Claude Code) never received the server_tool_use and web_search_tool_result blocks the interception layer builds for them, so citation panels stayed empty and result counts read zero. Two gaps in the same flow: FakeAnthropicMessagesStreamIterator._create_content_block_chunks only knew text, thinking, redacted_thinking and tool_use; any other block type emitted a bare content_block_stop with no start and no payload, silently dropping the block from the fake stream. server_tool_use now streams as content_block_start plus input_json_delta and any other server-generated block rides whole in content_block_start, matching Anthropic's native streaming of these types. try_short_circuit_search only recognized a native client by finding a web_search_* tool in the tools list, but the pre-request and deployment hooks convert that tool to the LiteLLM standard shape before the short-circuit runs, recording the conversion in WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY. The flag now counts as native too, so the synthetic response carries the native blocks on the converted path instead of only when hook ordering happens to leave the tool untouched. Fixes #35333 * fix(anthropic): accept server_tool_use and web_search_tool_result blocks when costing a /v1/messages response The short-circuit now returns native web search blocks, and the cost path validated the served body against a content union that only allowed text, tool_use and thinking blocks, so non-streaming requests answered 500 * fix(anthropic): keep every server_tool_use field when re-streaming a /v1/messages reply * fix(anthropic): keep tool_use fields, redacted thinking data and text citations when re-streaming a /v1/messages reply * test(websearch-interception): build the cost test's logging object with a datetime start time * fix(websearch-interception): send search snippets once in the direct search answer * fix(guardrails): mask web search queries, titles and urls on /v1/messages and deliver stream rewrites past server_tool_use blocks Output guardrails now scan the native web search blocks the short-circuit emits, ahead of the reply text so guardrails that read the last text still scan the reply. A guardrail reply of a different length falls back to the reply-only write-back. Ended-stream tool call rewrites land on server_tool_use blocks too, the ended-stream scan key covers the result strings, and legacy-hook pipeline steps pass through the result strings their response view lacks. The wire test reads the streamed server_tool_use input from its input_json_delta. * fix(guardrails): deliver a legacy post-call hook's tool-call rewrites on streams instead of discarding the stream rewrite A legacy hook that masked a tool call's arguments on a buffered stream made the whole rewrite undeliverable, text included, so the original stream went out. With web search blocks in the stream that now hit any hook masking the search query. The adapter hands same-count tool-call rewrites to the translation's write-back, as a unified guardrail's already are, and still discards a rewrite that adds or drops a tool call. * fix(guardrails): show a legacy post-call hook the web search result blocks of a /v1/messages stream The legacy hook's view of a buffered Anthropic stream is rebuilt from the assembled ModelResponse, which keeps web_search_tool_result blocks only in provider-specific fields, so a hook masking result titles on a stream never saw them. The hook's view now carries those blocks, its rewrite of them goes back through the ended-stream write-back, and the adapter no longer needs to pass texts outside the hook's view through untouched. * test(guardrails): type the legacy hook stream test helpers * fix(guardrails): scan a streamed web search query as text like the non-streaming reply * test(guardrails): type the masking guardrail test helpers * fix(guardrails): keep the block id when scanning a streamed web search query --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Devon Krisman
GitHub
parent
4a5c5665fa
commit
953750ed41
11 changed files with 1191 additions and 119 deletions
|
|
@ -398,14 +398,17 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
(t for t in tools if is_anthropic_native_web_search_tool(t)),
|
||||
None,
|
||||
)
|
||||
emit_native_blocks: Final = native_tool is not None or (
|
||||
kwargs is not None and bool(kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY))
|
||||
)
|
||||
|
||||
outcome: Final = await self._short_circuit_search_outcome(query, kwargs=kwargs)
|
||||
search_result_text: Final = WebSearchTransformation.search_outcome_text(outcome)
|
||||
|
||||
content: Final[list[dict[str, object]]] = []
|
||||
if native_tool is not None:
|
||||
if emit_native_blocks:
|
||||
tool_use_id: Final = f"srvtoolu_{uuid.uuid4().hex}"
|
||||
tool_name: Final = native_tool.get("name") or "web_search"
|
||||
tool_name: Final = (native_tool.get("name") if native_tool is not None else None) or "web_search"
|
||||
content.append(
|
||||
{
|
||||
"type": "server_tool_use",
|
||||
|
|
@ -415,10 +418,10 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
}
|
||||
)
|
||||
content.append(
|
||||
WebSearchTransformation.build_web_search_outcome_block(tool_use_id=tool_use_id, outcome=outcome)
|
||||
WebSearchTransformation.build_web_search_outcome_block(
|
||||
tool_use_id=tool_use_id, outcome=outcome, with_snippets=False
|
||||
)
|
||||
)
|
||||
# Keep the text block so non-native short-circuit callers (Claude Code,
|
||||
# github_copilot, etc.) see the same payload they always have.
|
||||
content.append({"type": "text", "text": search_result_text})
|
||||
|
||||
response: Final[dict[str, object]] = {
|
||||
|
|
@ -435,7 +438,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
verbose_logger.debug(
|
||||
"WebSearchInterception: Short-circuit search completed, returning synthetic response (%s chars, native_blocks=%s)",
|
||||
len(search_result_text),
|
||||
native_tool is not None,
|
||||
emit_native_blocks,
|
||||
)
|
||||
return response
|
||||
|
||||
|
|
|
|||
|
|
@ -412,6 +412,7 @@ class WebSearchTransformation:
|
|||
def build_web_search_tool_result_block(
|
||||
tool_use_id: str,
|
||||
search_response: SearchResponse | None,
|
||||
with_snippets: bool = True,
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Build an Anthropic-native ``web_search_tool_result`` content block.
|
||||
|
|
@ -445,6 +446,8 @@ class WebSearchTransformation:
|
|||
``litellm.asearch()``. If None or empty, the block is still
|
||||
emitted with an empty result list (signals "search ran, no
|
||||
results" rather than "search did not run").
|
||||
with_snippets: False when a text block in the same reply already
|
||||
carries the snippets.
|
||||
"""
|
||||
items: Final[list[dict[str, object]]] = []
|
||||
if search_response is not None:
|
||||
|
|
@ -460,7 +463,7 @@ class WebSearchTransformation:
|
|||
"title": title,
|
||||
"page_age": page_age,
|
||||
"encrypted_content": "",
|
||||
"snippet": getattr(r, "snippet", "") or "",
|
||||
**({"snippet": getattr(r, "snippet", "") or ""} if with_snippets else {}),
|
||||
}
|
||||
)
|
||||
return {
|
||||
|
|
@ -481,12 +484,15 @@ class WebSearchTransformation:
|
|||
}
|
||||
|
||||
@staticmethod
|
||||
def build_web_search_outcome_block(tool_use_id: str, outcome: SearchOutcome) -> dict[str, object]:
|
||||
def build_web_search_outcome_block(
|
||||
tool_use_id: str, outcome: SearchOutcome, with_snippets: bool = True
|
||||
) -> dict[str, object]:
|
||||
match outcome:
|
||||
case SearchSucceeded(response=response):
|
||||
return WebSearchTransformation.build_web_search_tool_result_block(
|
||||
tool_use_id=tool_use_id,
|
||||
search_response=response,
|
||||
with_snippets=with_snippets,
|
||||
)
|
||||
case SearchFailed(error_code=error_code):
|
||||
return WebSearchTransformation.build_web_search_tool_result_error_block(
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config
|
|||
from typing_extensions import ReadOnly, TypedDict, assert_never
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
from litellm.llms.anthropic.messages.adapters.transformation import (
|
||||
OPENAI_MAX_TOOL_NAME_LENGTH,
|
||||
|
|
@ -61,6 +61,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
|
|||
unappliable_request_rewrite,
|
||||
)
|
||||
from litellm.types.llms.anthropic import (
|
||||
ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX,
|
||||
AllAnthropicToolsValues,
|
||||
AnthropicMessagesRequest,
|
||||
)
|
||||
|
|
@ -153,6 +154,7 @@ class _TextAttrBlock(Protocol):
|
|||
text: str
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _WritableMessage(Protocol):
|
||||
@overload
|
||||
def get(self, key: str, /) -> object | None: ...
|
||||
|
|
@ -195,6 +197,140 @@ class _ToolCallShape:
|
|||
arguments: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _WebSearchString:
|
||||
block_idx: int
|
||||
result_idx: int | None
|
||||
field: str
|
||||
value: str
|
||||
|
||||
|
||||
def _streamed_web_search_result_blocks(response: ModelResponse) -> tuple[Mapping[str, object], ...]:
|
||||
choice: Final = response.choices[0] if response.choices else None
|
||||
fields: Final = as_str_mapping(choice.message.provider_specific_fields) if isinstance(choice, Choices) else None
|
||||
blocks: Final = (
|
||||
(as_str_mapping(item) for item in iter_items(fields.get("web_search_results")))
|
||||
if fields is not None and isinstance(fields.get("web_search_results"), list)
|
||||
else ()
|
||||
)
|
||||
return tuple(
|
||||
deepcopy(block) for block in blocks if block is not None and block.get("type") == "web_search_tool_result"
|
||||
)
|
||||
|
||||
|
||||
def _web_search_results(block: Mapping[str, object]) -> Iterable[object]:
|
||||
return iter_items(block.get("content")) if isinstance(block.get("content"), list) else ()
|
||||
|
||||
|
||||
def _web_search_result_strings(block_idx: int, result_idx: int, result: object) -> tuple[_WebSearchString, ...]:
|
||||
fields: Final = as_str_mapping(result)
|
||||
if fields is None or fields.get("type") != "web_search_result":
|
||||
return ()
|
||||
values: Final = ((field, fields.get(field)) for field in ("title", "url"))
|
||||
return tuple(
|
||||
_WebSearchString(block_idx, result_idx, field, value)
|
||||
for field, value in values
|
||||
if isinstance(value, str) and value
|
||||
)
|
||||
|
||||
|
||||
def _is_server_tool_use(block: Mapping[str, object]) -> bool:
|
||||
block_id: Final = block.get("id")
|
||||
match block.get("type"):
|
||||
case "server_tool_use":
|
||||
return True
|
||||
case "tool_use":
|
||||
return isinstance(block_id, str) and block_id.startswith(ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX)
|
||||
case _:
|
||||
return False
|
||||
|
||||
|
||||
def _web_search_block_strings(
|
||||
block_idx: int, block: Mapping[str, object], *, with_query: bool
|
||||
) -> tuple[_WebSearchString, ...]:
|
||||
"""The strings a client reads off a native web search block: the ``server_tool_use`` query and each
|
||||
``web_search_tool_result`` result's title and url."""
|
||||
match block.get("type"):
|
||||
case "server_tool_use" | "tool_use" if with_query and _is_server_tool_use(block):
|
||||
tool_input: Final = as_str_mapping(block.get("input"))
|
||||
query: Final = None if tool_input is None else tool_input.get("query")
|
||||
return (_WebSearchString(block_idx, None, "query", query),) if isinstance(query, str) and query else ()
|
||||
case "web_search_tool_result":
|
||||
return tuple(
|
||||
chain.from_iterable(
|
||||
_web_search_result_strings(block_idx, result_idx, result)
|
||||
for result_idx, result in enumerate(_web_search_results(block))
|
||||
)
|
||||
)
|
||||
case _:
|
||||
return ()
|
||||
|
||||
|
||||
def _rewritten_web_search_result(
|
||||
result_idx: int, result: object, rewrites: Mapping[tuple[int | None, str], str]
|
||||
) -> object:
|
||||
fields: Final = as_str_mapping(result)
|
||||
if fields is None:
|
||||
return result
|
||||
return {
|
||||
**fields,
|
||||
**{field: rewrites[(result_idx, field)] for field in ("title", "url") if (result_idx, field) in rewrites},
|
||||
}
|
||||
|
||||
|
||||
def _rewritten_web_search_block(
|
||||
block: Mapping[str, object], rewrites: Sequence[tuple[_WebSearchString, str]]
|
||||
) -> Mapping[str, object]:
|
||||
by_location: Final = MappingProxyType({(target.result_idx, target.field): value for target, value in rewrites})
|
||||
if _is_server_tool_use(block):
|
||||
query: Final = by_location.get((None, "query"))
|
||||
tool_input: Final = as_str_mapping(block.get("input")) or MappingProxyType({})
|
||||
return block if query is None else {**block, "input": {**tool_input, "query": query}}
|
||||
if not isinstance(block.get("content"), list):
|
||||
return block
|
||||
return {
|
||||
**block,
|
||||
"content": [
|
||||
_rewritten_web_search_result(result_idx, result, by_location)
|
||||
for result_idx, result in enumerate(_web_search_results(block))
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _write_back_web_search_block(content_block: object, rewrites: Sequence[tuple[_WebSearchString, str]]) -> None:
|
||||
block: Final = as_str_mapping(content_block)
|
||||
if block is None or not isinstance(content_block, _WritableMessage):
|
||||
return
|
||||
rewritten: Final = _rewritten_web_search_block(block, rewrites)
|
||||
writable: Final = _as_writable(content_block)
|
||||
for key in ("input", "content"):
|
||||
if key in rewritten:
|
||||
writable[key] = rewritten[key]
|
||||
|
||||
|
||||
def _split_web_search_texts(
|
||||
guardrailed: Sequence[str], web_search_count: int, reply_count: int
|
||||
) -> tuple[Sequence[str], Sequence[str]]:
|
||||
split_at: Final = web_search_count if len(guardrailed) == web_search_count + reply_count else 0
|
||||
return guardrailed[:split_at], guardrailed[split_at:]
|
||||
|
||||
|
||||
def _web_search_rewrites_by_block(
|
||||
strings: Sequence[_WebSearchString], guardrailed: Sequence[object]
|
||||
) -> Mapping[int, tuple[tuple[_WebSearchString, str], ...]]:
|
||||
rewrites: Final = tuple(
|
||||
(target, value)
|
||||
for target, value in zip(strings, guardrailed)
|
||||
if isinstance(value, str) and value != target.value
|
||||
)
|
||||
return MappingProxyType(
|
||||
{
|
||||
block_idx: tuple(rewrite for rewrite in rewrites if rewrite[0].block_idx == block_idx)
|
||||
for block_idx in {target.block_idx for target, _ in rewrites}
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _is_client_tool_use(block: Mapping[str, object]) -> bool:
|
||||
return (
|
||||
block.get("type") == "tool_use"
|
||||
|
|
@ -289,6 +425,7 @@ def _write_back_message_text(message: _WritableMessage, target: MessageTextTarge
|
|||
|
||||
_TOOL_USE_INPUT_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||
_DECODED_SSE_DATA: Final = TypeAdapter(object)
|
||||
_REBUILT_TOOL_CALL_BLOCK_TYPES: Final = frozenset({"tool_use", "server_tool_use"})
|
||||
_RELEASED_TOOL_USE_STOP: Final = (
|
||||
b"event: message_delta\n"
|
||||
b'data: {"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": null}, '
|
||||
|
|
@ -315,6 +452,51 @@ def _write_back_tool_use(
|
|||
block["name"] = shape.name
|
||||
|
||||
|
||||
def _streamed_query_strings(
|
||||
tool_call_blocks: Sequence[tuple[int, Mapping[str, object]]], tool_calls: Sequence[_ToolCallShape]
|
||||
) -> tuple[_WebSearchString, ...]:
|
||||
if len(tool_call_blocks) != len(tool_calls):
|
||||
return ()
|
||||
return tuple(
|
||||
chain.from_iterable(
|
||||
_web_search_block_strings(
|
||||
index, {**block, "input": _rewritten_tool_use_input(tool_call.arguments)}, with_query=True
|
||||
)
|
||||
for (index, block), tool_call in zip(tool_call_blocks, tool_calls)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _query_rewrites(strings: Sequence[_WebSearchString], guardrailed: Sequence[object]) -> Mapping[int, str]:
|
||||
return MappingProxyType(
|
||||
{
|
||||
target.block_idx: value
|
||||
for target, value in zip(strings, guardrailed)
|
||||
if isinstance(value, str) and value != target.value
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _with_query(tool_call: _ToolCallShape, query: str) -> _ToolCallShape:
|
||||
arguments: Final = _rewritten_tool_use_input(tool_call.arguments)
|
||||
if arguments is None:
|
||||
return tool_call
|
||||
return _ToolCallShape(name=tool_call.name, arguments=json.dumps({**arguments, "query": query}))
|
||||
|
||||
|
||||
def _with_query_rewrites(
|
||||
tool_calls: tuple[_ToolCallShape, ...],
|
||||
tool_call_blocks: Sequence[tuple[int, Mapping[str, object]]],
|
||||
query_rewrites: Mapping[int, str],
|
||||
) -> tuple[_ToolCallShape, ...]:
|
||||
if not query_rewrites or len(tool_call_blocks) != len(tool_calls):
|
||||
return tool_calls
|
||||
return tuple(
|
||||
_with_query(tool_call, query_rewrites[index]) if index in query_rewrites else tool_call
|
||||
for (index, _), tool_call in zip(tool_call_blocks, tool_calls)
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _SSEFieldRewrite:
|
||||
"""One field of one nested section of a buffered SSE event, rewritten."""
|
||||
|
|
@ -402,7 +584,12 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
def post_call_hook_response(self, response: object) -> object:
|
||||
if not isinstance(response, ModelResponse):
|
||||
return response
|
||||
return self.adapter.translate_openai_response_to_anthropic(response)
|
||||
translated: Final = self.adapter.translate_openai_response_to_anthropic(response)
|
||||
result_blocks: Final = _streamed_web_search_result_blocks(response)
|
||||
if not result_blocks:
|
||||
return translated
|
||||
content: Final = translated.get("content") or ()
|
||||
return {**translated, "content": [*result_blocks, *content]}
|
||||
|
||||
@staticmethod
|
||||
def _build_streaming_usage_response(
|
||||
|
|
@ -1273,9 +1460,17 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
task_mappings,
|
||||
tool_calls_to_check,
|
||||
)
|
||||
block_fields: Final = tuple(map(self._output_block_fields, response_content))
|
||||
web_search_strings: Final = tuple(
|
||||
chain.from_iterable(
|
||||
_web_search_block_strings(content_idx, fields[1], with_query=True)
|
||||
for content_idx, fields in enumerate(block_fields)
|
||||
if fields is not None
|
||||
)
|
||||
)
|
||||
|
||||
# Step 2: Apply guardrail to all texts in batch
|
||||
if texts_to_check or tool_calls_to_check:
|
||||
if texts_to_check or tool_calls_to_check or web_search_strings:
|
||||
request_data = self._prepare_request_data(
|
||||
request_data,
|
||||
response,
|
||||
|
|
@ -1284,7 +1479,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
)
|
||||
|
||||
inputs: Final = self._build_guardrail_inputs(
|
||||
texts_to_check,
|
||||
[*(target.value for target in web_search_strings), *texts_to_check],
|
||||
images_to_check,
|
||||
tool_calls_to_check,
|
||||
response,
|
||||
|
|
@ -1297,14 +1492,19 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
|
||||
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
|
||||
web_search_rewrites, reply_rewrites = _split_web_search_texts(
|
||||
guardrailed_inputs.get("texts", []), len(web_search_strings), len(texts_to_check)
|
||||
)
|
||||
|
||||
# Step 3: Map guardrail responses back to original response structure
|
||||
await self._apply_guardrail_responses_to_output(
|
||||
response=response,
|
||||
responses=guardrailed_texts,
|
||||
responses=list(reply_rewrites),
|
||||
task_mappings=task_mappings,
|
||||
)
|
||||
self._write_back_web_search_strings(
|
||||
response_content, _web_search_rewrites_by_block(web_search_strings, web_search_rewrites)
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug("Anthropic Messages: Processed output response: %s", response)
|
||||
|
||||
|
|
@ -1350,8 +1550,14 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
string_so_far = first_choice.message.content
|
||||
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_list or ())
|
||||
guardrail_inputs: Final = GenericGuardrailAPIInputs()
|
||||
if string_so_far:
|
||||
guardrail_inputs["texts"] = [string_so_far]
|
||||
leading_texts: Final = (string_so_far,) if isinstance(string_so_far, str) and string_so_far else ()
|
||||
tool_call_blocks: Final = self._streamed_tool_call_blocks(responses_so_far)
|
||||
streamed_query_strings: Final = _streamed_query_strings(tool_call_blocks, pre_guardrail_tool_calls)
|
||||
streamed_result_strings: Final = self._streamed_web_search_result_strings(responses_so_far)
|
||||
streamed_web_search_strings: Final = (*streamed_result_strings, *streamed_query_strings)
|
||||
texts_to_check: Final = [*(target.value for target in streamed_web_search_strings), *leading_texts]
|
||||
if texts_to_check:
|
||||
guardrail_inputs["texts"] = texts_to_check
|
||||
if tool_calls_list:
|
||||
guardrail_inputs["tool_calls"] = tool_calls_list
|
||||
|
||||
|
|
@ -1374,29 +1580,42 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
responses_so_far, request_data
|
||||
)
|
||||
raise
|
||||
guardrailed_texts: Final = _guardrailed_inputs.get("texts")
|
||||
web_search_rewrites, reply_rewrites = _split_web_search_texts(
|
||||
_guardrailed_inputs.get("texts") or [], len(streamed_web_search_strings), len(leading_texts)
|
||||
)
|
||||
if (
|
||||
deliver_ended_stream_rewrites
|
||||
and isinstance(string_so_far, str)
|
||||
and string_so_far
|
||||
and guardrailed_texts
|
||||
and guardrailed_texts[0] != string_so_far
|
||||
and reply_rewrites
|
||||
and reply_rewrites[0] != string_so_far
|
||||
):
|
||||
self._write_ended_stream_text_rewrite(
|
||||
responses_so_far,
|
||||
guardrailed_texts[0],
|
||||
reply_rewrites[0],
|
||||
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
|
||||
)
|
||||
if deliver_ended_stream_rewrites:
|
||||
self._write_ended_stream_web_search_rewrites(
|
||||
responses_so_far,
|
||||
_web_search_rewrites_by_block(streamed_result_strings, web_search_rewrites),
|
||||
)
|
||||
if deliver_ended_stream_rewrites:
|
||||
returned_tool_calls: Final = _guardrailed_inputs.get("tool_calls")
|
||||
self._write_ended_stream_tool_call_rewrites(
|
||||
responses_so_far,
|
||||
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
|
||||
post_guardrail_tool_calls=_tool_call_shapes(
|
||||
returned_tool_calls
|
||||
if isinstance(returned_tool_calls, list)
|
||||
and len(returned_tool_calls) == len(pre_guardrail_tool_calls)
|
||||
else tool_calls_list or ()
|
||||
post_guardrail_tool_calls=_with_query_rewrites(
|
||||
_tool_call_shapes(
|
||||
returned_tool_calls
|
||||
if isinstance(returned_tool_calls, list)
|
||||
and len(returned_tool_calls) == len(pre_guardrail_tool_calls)
|
||||
else tool_calls_list or ()
|
||||
),
|
||||
tool_call_blocks,
|
||||
_query_rewrites(
|
||||
streamed_query_strings, web_search_rewrites[len(streamed_result_strings) :]
|
||||
),
|
||||
),
|
||||
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
|
||||
)
|
||||
|
|
@ -1451,6 +1670,14 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
request_data["litellm_metadata"] = user_metadata
|
||||
return request_data
|
||||
|
||||
@staticmethod
|
||||
def _write_back_web_search_strings(
|
||||
response_content: Sequence[object],
|
||||
rewrites_by_block: Mapping[int, tuple[tuple[_WebSearchString, str], ...]],
|
||||
) -> None:
|
||||
for block_idx, rewrites in rewrites_by_block.items():
|
||||
_write_back_web_search_block(_content_block_at(response_content, block_idx), rewrites)
|
||||
|
||||
@staticmethod
|
||||
def _get_response_content(response: object) -> Sequence[object]:
|
||||
"""Extract content list from a dict or object response."""
|
||||
|
|
@ -1573,30 +1800,22 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
) -> None:
|
||||
"""Deliver ended-stream guardrail tool-call rewrites by rewriting the
|
||||
buffered chunks in place: the rebuilt response lists tool calls in the
|
||||
order of the stream's ``tool_use`` blocks, so the nth rewritten call lands
|
||||
on the nth block, its first ``input_json_delta`` carrying the full rewritten
|
||||
order of the stream's ``tool_use`` and ``server_tool_use`` blocks, so the nth
|
||||
rewritten call lands on the nth block, its first ``input_json_delta`` carrying the full rewritten
|
||||
arguments, every later one blanked, and ``content_block_start`` carrying the
|
||||
rewritten name. Blocks that do not line up with the rebuilt tool calls make
|
||||
the rewrite undeliverable, so the pipeline executor discards it and releases
|
||||
the original chunks."""
|
||||
if post_guardrail_tool_calls == pre_guardrail_tool_calls:
|
||||
return
|
||||
block_indices: Final = tuple(
|
||||
index
|
||||
for item in responses_so_far
|
||||
for event in cls._iter_sse_events(item)
|
||||
if event.get("type") == "content_block_start"
|
||||
and isinstance(block := event.get("content_block"), Mapping)
|
||||
and block.get("type") == "tool_use"
|
||||
and isinstance(index := event.get("index"), int)
|
||||
)
|
||||
block_indices: Final = tuple(index for index, _ in cls._streamed_tool_call_blocks(responses_so_far))
|
||||
if len(block_indices) != len(post_guardrail_tool_calls):
|
||||
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
|
||||
|
||||
raise UndeliverableStreamRewrite(
|
||||
guardrail_name,
|
||||
f"the guardrail returned {len(post_guardrail_tool_calls)} tool calls for a stream that carried "
|
||||
f"{len(block_indices)} tool_use blocks",
|
||||
f"{len(block_indices)} tool_use or server_tool_use blocks",
|
||||
)
|
||||
rewrites_by_block: Final = MappingProxyType(
|
||||
{
|
||||
|
|
@ -1636,6 +1855,60 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _streamed_web_search_result_strings(cls, responses_so_far: Sequence[object]) -> tuple[_WebSearchString, ...]:
|
||||
events: Final = chain.from_iterable(cls._iter_sse_events(item) for item in responses_so_far)
|
||||
result_blocks: Final = (
|
||||
(index, block)
|
||||
for event in events
|
||||
if event.get("type") == "content_block_start"
|
||||
and isinstance(index := event.get("index"), int)
|
||||
and (block := as_str_mapping(event.get("content_block"))) is not None
|
||||
)
|
||||
return tuple(
|
||||
chain.from_iterable(
|
||||
_web_search_block_strings(index, block, with_query=False) for index, block in result_blocks
|
||||
)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _write_ended_stream_web_search_rewrites(
|
||||
cls,
|
||||
responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place
|
||||
rewrites_by_block: Mapping[int, tuple[tuple[_WebSearchString, str], ...]],
|
||||
) -> None:
|
||||
"""Deliver ended-stream guardrail rewrites of web search result titles and urls into the
|
||||
``content_block_start`` that carries each whole ``web_search_tool_result`` block."""
|
||||
if not rewrites_by_block:
|
||||
return
|
||||
|
||||
def rewrite_result_block(event: Mapping[str, object]) -> _SSEFieldRewrite | None:
|
||||
index: Final = event.get("index")
|
||||
block: Final = as_str_mapping(event.get("content_block"))
|
||||
rewrites: Final = rewrites_by_block.get(index) if isinstance(index, int) else None
|
||||
if event.get("type") != "content_block_start" or block is None or rewrites is None:
|
||||
return None
|
||||
return _SSEFieldRewrite(
|
||||
"content_block", "content", _rewritten_web_search_block(block, rewrites).get("content")
|
||||
)
|
||||
|
||||
cls._rewrite_ended_stream_events(responses_so_far, rewrite_result_block)
|
||||
|
||||
@classmethod
|
||||
def _streamed_tool_call_blocks(
|
||||
cls, responses_so_far: Sequence[object]
|
||||
) -> tuple[tuple[int, Mapping[str, object]], ...]:
|
||||
events: Final = chain.from_iterable(cls._iter_sse_events(item) for item in responses_so_far)
|
||||
return tuple(
|
||||
(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 isinstance(block_type := block.get("type"), str)
|
||||
and block_type in _REBUILT_TOOL_CALL_BLOCK_TYPES
|
||||
and isinstance(index := event.get("index"), int)
|
||||
)
|
||||
|
||||
@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."""
|
||||
|
|
@ -1738,8 +2011,10 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
|
||||
stream_ended: Final = self._check_streaming_has_ended(responses_so_far)
|
||||
tool_use_fingerprints: Final = self._streamed_tool_use_fingerprints(responses_so_far)
|
||||
streamed_text: Final = self.get_streaming_string_so_far(responses_so_far)
|
||||
result_strings: Final = self._streamed_web_search_result_strings(responses_so_far) if stream_ended else ()
|
||||
return StreamingScanKey(
|
||||
texts=(self.get_streaming_string_so_far(responses_so_far),),
|
||||
texts=(*(target.value for target in result_strings), streamed_text),
|
||||
tool_calls=tool_use_fingerprints if stream_ended else (),
|
||||
stream_ended=stream_ended,
|
||||
tool_calls_in_flight=bool(tool_use_fingerprints) and not stream_ended,
|
||||
|
|
|
|||
|
|
@ -12,11 +12,20 @@ import json
|
|||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final, cast
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.llms.anthropic.messages.native.utils import anthropic_message_delta
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import (
|
||||
AnthropicMessagesResponse,
|
||||
)
|
||||
|
||||
_CITATIONS: Final = TypeAdapter(tuple[object, ...])
|
||||
|
||||
|
||||
def _citations_of(block: Mapping[str, object]) -> tuple[object, ...]:
|
||||
citations: Final = block.get("citations")
|
||||
return _CITATIONS.validate_python(citations) if isinstance(citations, list) else ()
|
||||
|
||||
|
||||
class FakeAnthropicMessagesStreamIterator:
|
||||
"""
|
||||
|
|
@ -47,16 +56,30 @@ class FakeAnthropicMessagesStreamIterator:
|
|||
|
||||
def _create_content_block_chunks(self, block_dict: Mapping[str, object], index: int) -> list[bytes]:
|
||||
"""Build SSE chunks for a single content block."""
|
||||
chunks: Final = []
|
||||
chunks: Final = list[bytes]()
|
||||
block_type: Final = block_dict.get("type")
|
||||
|
||||
if block_type == "text":
|
||||
citations: Final = _citations_of(block_dict)
|
||||
content_block_start = {
|
||||
"type": "content_block_start",
|
||||
"index": index,
|
||||
"content_block": {"type": "text", "text": ""},
|
||||
"content_block": {"type": "text", "text": "", "citations": ()}
|
||||
if citations
|
||||
else {"type": "text", "text": ""},
|
||||
}
|
||||
chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode())
|
||||
citation_deltas: Final = (
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": index,
|
||||
"delta": {"type": "citations_delta", "citation": citation},
|
||||
}
|
||||
for citation in citations
|
||||
)
|
||||
chunks.extend(
|
||||
f"event: content_block_delta\ndata: {json.dumps(delta)}\n\n".encode() for delta in citation_deltas
|
||||
)
|
||||
text: Final = block_dict.get("text", "")
|
||||
content_block_delta = {
|
||||
"type": "content_block_delta",
|
||||
|
|
@ -89,36 +112,23 @@ class FakeAnthropicMessagesStreamIterator:
|
|||
}
|
||||
chunks.append(f"event: content_block_delta\ndata: {json.dumps(signature_delta)}\n\n".encode())
|
||||
|
||||
elif block_type == "redacted_thinking":
|
||||
content_block_start = {
|
||||
elif block_type in ("tool_use", "server_tool_use"):
|
||||
empty_input: Final[Mapping[str, object]] = {}
|
||||
tool_use_start: Final = {
|
||||
"type": "content_block_start",
|
||||
"index": index,
|
||||
"content_block": {"type": "redacted_thinking"},
|
||||
"content_block": {**block_dict, "input": empty_input},
|
||||
}
|
||||
chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode())
|
||||
|
||||
elif block_type == "tool_use":
|
||||
content_block_start = {
|
||||
"type": "content_block_start",
|
||||
"index": index,
|
||||
"content_block": {
|
||||
"type": "tool_use",
|
||||
"id": block_dict.get("id"),
|
||||
"name": block_dict.get("name"),
|
||||
"input": {},
|
||||
},
|
||||
}
|
||||
chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode())
|
||||
input_data: Final = block_dict.get("input", {})
|
||||
content_block_delta = {
|
||||
chunks.append(f"event: content_block_start\ndata: {json.dumps(tool_use_start)}\n\n".encode())
|
||||
tool_use_delta: Final = {
|
||||
"type": "content_block_delta",
|
||||
"index": index,
|
||||
"delta": {
|
||||
"type": "input_json_delta",
|
||||
"partial_json": json.dumps(input_data),
|
||||
"partial_json": json.dumps(block_dict.get("input", {})),
|
||||
},
|
||||
}
|
||||
chunks.append(f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n".encode())
|
||||
chunks.append(f"event: content_block_delta\ndata: {json.dumps(tool_use_delta)}\n\n".encode())
|
||||
|
||||
else:
|
||||
passthrough_start: Final = {
|
||||
|
|
|
|||
|
|
@ -98,15 +98,11 @@ def _count_change(sent: tuple[object, ...] | None, returned: tuple[object, ...]
|
|||
return (len(sent), len(returned))
|
||||
|
||||
|
||||
def _tool_call_mismatch_reason(
|
||||
sent: tuple[tuple[object, object], ...] | None, returned: tuple[tuple[object, object], ...] | None
|
||||
) -> str | None:
|
||||
if sent == returned:
|
||||
return None
|
||||
def _tool_call_count_mismatch_reason(sent: Sequence[object] | None, returned: Sequence[object] | None) -> str | None:
|
||||
sent_count: Final = len(sent or ())
|
||||
returned_count: Final = len(returned or ())
|
||||
if sent_count == returned_count:
|
||||
return "the legacy hook changed a tool call's name or arguments, which this path cannot write back"
|
||||
return None
|
||||
return f"the legacy hook returned {returned_count} tool calls for a stream that carried {sent_count}"
|
||||
|
||||
|
||||
|
|
@ -205,11 +201,10 @@ class _LegacyHookStreamAdapter(CustomGuardrail):
|
|||
non-streaming hooks, an exception it raises ends the stream through the executor's
|
||||
fail/error classification, and the response it hands back, or the one it changed in place
|
||||
and returned ``None`` for, is re-scanned by the same translation so its texts reach the
|
||||
client through the translation's ended-stream write-back. A
|
||||
replacement whose scanned texts do not line up with the originals, or whose tool calls
|
||||
differ from them, is undeliverable, so the executor releases the original chunks. A stream
|
||||
that carried no text to scan, such as a tool-only Anthropic message, stays deliverable as
|
||||
long as the hook left the tool calls alone."""
|
||||
client through the translation's ended-stream write-back. Tool calls the hook rewrote go to the translation's
|
||||
tool-call write-back like a unified guardrail's. A replacement whose scanned texts or tool
|
||||
calls do not line up with the originals by count is undeliverable, so the executor releases
|
||||
the original chunks."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -258,14 +253,19 @@ class _LegacyHookStreamAdapter(CustomGuardrail):
|
|||
guardrail_name,
|
||||
f"the legacy hook returned {returned_text_count} texts for a stream that carried {sent_text_count}",
|
||||
)
|
||||
tool_call_mismatch: Final = _tool_call_mismatch_reason(
|
||||
_tool_call_shapes(inputs.get("tool_calls")), _tool_call_shapes(rescanned.get("tool_calls"))
|
||||
sent_tool_calls: Final = inputs.get("tool_calls")
|
||||
rewritten_tool_calls: Final = rescanned.get("tool_calls")
|
||||
tool_call_count_mismatch: Final = _tool_call_count_mismatch_reason(sent_tool_calls, rewritten_tool_calls)
|
||||
if tool_call_count_mismatch is not None:
|
||||
raise UndeliverableStreamRewrite(guardrail_name, tool_call_count_mismatch)
|
||||
texts_update: Final[GenericGuardrailAPIInputs] = {"texts": rewritten} if rewritten else {}
|
||||
tool_calls_update: Final[GenericGuardrailAPIInputs] = (
|
||||
{"tool_calls": rewritten_tool_calls}
|
||||
if rewritten_tool_calls is not None
|
||||
and _rewrote(_tool_call_shapes(sent_tool_calls), _tool_call_shapes(rewritten_tool_calls))
|
||||
else {}
|
||||
)
|
||||
if tool_call_mismatch is not None:
|
||||
raise UndeliverableStreamRewrite(guardrail_name, tool_call_mismatch)
|
||||
if not rewritten:
|
||||
return inputs
|
||||
rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten}
|
||||
rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, **texts_update, **tool_calls_update}
|
||||
return rewritten_inputs
|
||||
|
||||
async def _rescan(
|
||||
|
|
|
|||
|
|
@ -713,6 +713,23 @@ class AnthropicResponseContentBlockRedactedThinking(LiteLLMBaseModel):
|
|||
data: str
|
||||
|
||||
|
||||
class AnthropicResponseContentBlockServerToolUse(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
type: Literal["server_tool_use"]
|
||||
id: str
|
||||
name: str
|
||||
input: Mapping[str, object]
|
||||
|
||||
|
||||
class AnthropicResponseContentBlockWebSearchToolResult(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
type: Literal["web_search_tool_result"]
|
||||
tool_use_id: str
|
||||
content: Sequence[Mapping[str, object]] | Mapping[str, object]
|
||||
|
||||
|
||||
class AnthropicResponseUsageBlock(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
|
|
@ -744,6 +761,8 @@ class AnthropicResponse(LiteLLMBaseModel):
|
|||
| AnthropicResponseContentBlockToolUse
|
||||
| AnthropicResponseContentBlockThinking
|
||||
| AnthropicResponseContentBlockRedactedThinking
|
||||
| AnthropicResponseContentBlockServerToolUse
|
||||
| AnthropicResponseContentBlockWebSearchToolResult
|
||||
]
|
||||
"""Content generated by the model."""
|
||||
|
||||
|
|
|
|||
|
|
@ -29,6 +29,24 @@ def sse_events(text: str) -> tuple[tuple[str, dict[str, object]], ...]:
|
|||
)
|
||||
|
||||
|
||||
def streamed_blocks(events: tuple[tuple[str, dict[str, object]], ...]) -> tuple[dict[str, object], ...]:
|
||||
input_json: Final = tuple(
|
||||
(event["index"], event["delta"]["partial_json"])
|
||||
for name, event in events
|
||||
if name == "content_block_delta" and event["delta"]["type"] == "input_json_delta"
|
||||
)
|
||||
return tuple(
|
||||
{
|
||||
**event["content_block"],
|
||||
"input": json.loads("".join(part for index, part in input_json if index == event["index"])),
|
||||
}
|
||||
if any(index == event["index"] for index, _ in input_json)
|
||||
else event["content_block"]
|
||||
for name, event in events
|
||||
if name == "content_block_start"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.covers("other.provider_wire.bedrock.websearch_interception_streamed_capped_turn_ends_with_native_results")
|
||||
def test_streamed_web_search_turn_capped_by_max_agentic_loops_ends_turn_with_snippets_and_ordered_blocks(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
|
|
@ -136,7 +154,7 @@ def test_streamed_web_search_turn_capped_by_max_agentic_loops_ends_turn_with_sni
|
|||
for earlier, earlier_event in events[:position]
|
||||
if earlier == "content_block_start"
|
||||
}, response.text
|
||||
started: Final = tuple(event["content_block"] for name, event in events if name == "content_block_start")
|
||||
started: Final = streamed_blocks(events)
|
||||
search_ids: Final = tuple(block["id"] for block in started if block["type"] == "server_tool_use")
|
||||
assert search_ids and all(search_id.startswith("srvtoolu_") for search_id in search_ids), response.text
|
||||
assert started[-1] == {"type": "text", "text": ""}, response.text
|
||||
|
|
@ -162,7 +180,12 @@ def test_streamed_web_search_turn_capped_by_max_agentic_loops_ends_turn_with_sni
|
|||
)
|
||||
), response.text
|
||||
assert (
|
||||
"".join(event["delta"]["text"] for name, event in events if name == "content_block_delta") == "turn-1"
|
||||
"".join(
|
||||
event["delta"]["text"]
|
||||
for name, event in events
|
||||
if name == "content_block_delta" and event["delta"]["type"] == "text_delta"
|
||||
)
|
||||
== "turn-1"
|
||||
), response.text
|
||||
assert [event["delta"]["stop_reason"] for name, event in events if name == "message_delta"] == [
|
||||
"end_turn"
|
||||
|
|
|
|||
|
|
@ -5,7 +5,9 @@ Tests the short-circuit path that detects web-search-only /v1/messages requests
|
|||
and executes the search directly without routing through the backend LLM.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import json
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
|
|
@ -16,6 +18,7 @@ import litellm
|
|||
from litellm.integrations.websearch_interception.handler import (
|
||||
WebSearchInterceptionLogger,
|
||||
)
|
||||
from litellm.messages.dispatch import MessagesResult
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Detection tests
|
||||
|
|
@ -73,6 +76,42 @@ class TestTryShortCircuitSearch:
|
|||
assert "Result" in text_block["text"]
|
||||
mock_search.assert_called_once_with("Search for Claude Code releases")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_emits_native_blocks_when_tools_already_converted(self, perplexity_search: respx.Route):
|
||||
"""Pre-request/deployment hooks convert the native tool to the LiteLLM
|
||||
standard shape before the short-circuit runs and record that in
|
||||
WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY; the flag alone must be enough for
|
||||
native blocks to be emitted."""
|
||||
from litellm.integrations.websearch_interception.handler import (
|
||||
WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY,
|
||||
)
|
||||
|
||||
logger: Final = WebSearchInterceptionLogger(enabled_providers=["hosted_vllm"])
|
||||
|
||||
result: Final = await logger.try_short_circuit_search(
|
||||
model="hosted_vllm/some-model",
|
||||
messages=[{"role": "user", "content": "Search for something"}],
|
||||
tools=[
|
||||
{
|
||||
"name": "litellm_web_search",
|
||||
"description": "Search the web for information.",
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {"query": {"type": "string"}},
|
||||
"required": ["query"],
|
||||
},
|
||||
}
|
||||
],
|
||||
custom_llm_provider="hosted_vllm",
|
||||
kwargs={WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True},
|
||||
)
|
||||
|
||||
assert perplexity_search.call_count == 1
|
||||
assert result is not None
|
||||
assert [b["type"] for b in result["content"]] == ["server_tool_use", "web_search_tool_result", "text"]
|
||||
assert result["content"][0]["name"] == "web_search"
|
||||
assert [r["url"] for r in result["content"][1]["content"]] == ["https://example.com/releases"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_does_not_short_circuit_mixed_tools(self):
|
||||
"""Mix of web_search and other tools → NOT short-circuited"""
|
||||
|
|
@ -468,6 +507,89 @@ class TestShortCircuitEntryPoint:
|
|||
assert text_block["text"] == "results"
|
||||
|
||||
|
||||
class TestShortCircuitThroughMessagesEntryPoint:
|
||||
"""The pre-request hook renames the native tool before the short-circuit
|
||||
runs, so only a request through anthropic_messages shows whether native
|
||||
clients still get their server_tool_use and web_search_tool_result blocks."""
|
||||
|
||||
NATIVE_TOOL: Final = {"type": "web_search_20250305", "name": "web_search"}
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _intercept_web_search(self, perplexity_search: respx.Route) -> Iterator[respx.Route]:
|
||||
with patch("litellm.callbacks", [WebSearchInterceptionLogger(enabled_providers=["bedrock"])]):
|
||||
yield perplexity_search
|
||||
|
||||
@staticmethod
|
||||
async def _send(*, tools: Sequence[Mapping[str, object]], stream: bool, **kwargs: object) -> MessagesResult:
|
||||
return await litellm.anthropic_messages(
|
||||
model="bedrock/converse/us.anthropic.claude-sonnet-5-5",
|
||||
messages=[{"role": "user", "content": "latest python release"}],
|
||||
tools=tools,
|
||||
max_tokens=64,
|
||||
stream=stream,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_tool_gets_native_blocks(self, perplexity_search: respx.Route):
|
||||
response: Final = await self._send(tools=[self.NATIVE_TOOL], stream=False)
|
||||
|
||||
assert perplexity_search.call_count == 1
|
||||
assert [block["type"] for block in response["content"]] == [
|
||||
"server_tool_use",
|
||||
"web_search_tool_result",
|
||||
"text",
|
||||
]
|
||||
server_tool_use, tool_result, text = response["content"]
|
||||
assert server_tool_use["name"] == "web_search"
|
||||
assert server_tool_use["input"] == {"query": "latest python release"}
|
||||
assert tool_result["tool_use_id"] == server_tool_use["id"]
|
||||
assert [result["url"] for result in tool_result["content"]] == ["https://example.com/releases"]
|
||||
assert "Release notes" in text["text"]
|
||||
assert "Release notes" not in json.dumps(tool_result)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_blocks_reply_can_be_costed(self):
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
||||
logging_obj: Final = Logging(
|
||||
model="bedrock/converse/us.anthropic.claude-sonnet-5-5",
|
||||
messages=[{"role": "user", "content": "latest python release"}],
|
||||
stream=False,
|
||||
call_type="anthropic_messages",
|
||||
start_time=datetime.datetime(2026, 1, 1),
|
||||
litellm_call_id="short-circuit-cost",
|
||||
function_id="short-circuit-cost",
|
||||
)
|
||||
logging_obj.update_environment_variables(litellm_params={}, optional_params={})
|
||||
response: Final = await self._send(tools=[self.NATIVE_TOOL], stream=False, litellm_logging_obj=logging_obj)
|
||||
|
||||
assert logging_obj.response_cost_calculator(result=response) == 0.0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_tool_streams_native_blocks(self):
|
||||
stream: Final = await self._send(tools=[self.NATIVE_TOOL], stream=True)
|
||||
body: Final = b"".join([chunk async for chunk in stream])
|
||||
|
||||
starts: Final = [
|
||||
json.loads(line.removeprefix(b"data: "))["content_block"]
|
||||
for line in body.splitlines()
|
||||
if line.startswith(b"data: ") and b'"content_block_start"' in line
|
||||
]
|
||||
assert [block["type"] for block in starts] == ["server_tool_use", "web_search_tool_result", "text"]
|
||||
assert starts[1]["tool_use_id"] == starts[0]["id"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_litellm_standard_tool_stays_text_only(self):
|
||||
from litellm.integrations.websearch_interception.tools import (
|
||||
get_litellm_web_search_tool,
|
||||
)
|
||||
|
||||
response: Final = await self._send(tools=[get_litellm_web_search_tool()], stream=False)
|
||||
|
||||
assert [block["type"] for block in response["content"]] == ["text"]
|
||||
|
||||
|
||||
class TestShortCircuitSafeguardResults:
|
||||
"""A short-circuited reply never reaches the backend, so the proxy answers Claude Code's safeguards itself."""
|
||||
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ with guardrail transformations, specifically testing edge cases with empty choic
|
|||
"""
|
||||
|
||||
import json
|
||||
import re
|
||||
from typing import Any, Literal, Optional, Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
|
@ -552,13 +553,25 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
|
|||
assert chunks == original
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_ended_stream_tool_use_rewrite_with_server_tool_use_block_fails_closed(self):
|
||||
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
|
||||
|
||||
async def test_deliver_ended_stream_tool_use_rewrite_lands_on_server_tool_use_and_tool_use_blocks(self):
|
||||
handler = AnthropicMessagesHandler()
|
||||
server_tool_use = [
|
||||
("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}}}),
|
||||
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"query": "fruit"}'}}),
|
||||
(
|
||||
"content_block_start",
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}},
|
||||
},
|
||||
),
|
||||
(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "input_json_delta", "partial_json": '{"query": "fruit"}'},
|
||||
},
|
||||
),
|
||||
("content_block_stop", {"type": "content_block_stop", "index": 0}),
|
||||
]
|
||||
tool_use = self._ended_tool_use_sse_chunks()
|
||||
|
|
@ -568,6 +581,26 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
|
|||
+ [chunk.replace(b'"index": 0', b'"index": 1') for chunk in tool_use[1:]]
|
||||
)
|
||||
|
||||
await handler.process_output_streaming_response(
|
||||
responses_so_far=chunks,
|
||||
guardrail_to_apply=self._argument_masking_guardrail(),
|
||||
litellm_logging_obj=MagicMock(),
|
||||
deliver_ended_stream_rewrites=True,
|
||||
)
|
||||
|
||||
assert self._partial_jsons(chunks) == ['{"fruit": "[MASKED]"}', '{"fruit": "[MASKED]"}', "", ""]
|
||||
raw = b"".join(chunks).decode()
|
||||
assert '"id": "srvtoolu_1"' in raw and '"id": "toolu_1"' in raw
|
||||
assert "persim" not in raw and '"query"' not in raw
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_ended_stream_tool_use_rewrite_on_an_unindexed_tool_block_fails_closed(self):
|
||||
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
|
||||
|
||||
handler = AnthropicMessagesHandler()
|
||||
tool_use = self._ended_tool_use_sse_chunks()
|
||||
chunks = tool_use[:1] + [tool_use[1].replace(b'"index": 0', b'"index": "0"')] + tool_use[2:]
|
||||
|
||||
with pytest.raises(UndeliverableStreamRewrite):
|
||||
await handler.process_output_streaming_response(
|
||||
responses_so_far=chunks,
|
||||
|
|
@ -576,6 +609,270 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
|
|||
deliver_ended_stream_rewrites=True,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _python_masking_guardrail() -> CustomGuardrail:
|
||||
class MaskPython(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):
|
||||
tool_call.function.arguments = re.sub(
|
||||
"python", "[MASKED]", tool_call.function.arguments, flags=re.I
|
||||
)
|
||||
return {
|
||||
**inputs,
|
||||
"texts": [re.sub("python", "[MASKED]", text, flags=re.I) for text in inputs.get("texts", [])],
|
||||
}
|
||||
|
||||
return MaskPython(guardrail_name="test")
|
||||
|
||||
@staticmethod
|
||||
def _python_texts_masking_guardrail() -> CustomGuardrail:
|
||||
class MaskPythonTexts(CustomGuardrail):
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict[str, object],
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[LiteLLMLoggingObj] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
return {
|
||||
**inputs,
|
||||
"texts": [re.sub("python", "[MASKED]", text, flags=re.I) for text in inputs.get("texts", [])],
|
||||
}
|
||||
|
||||
return MaskPythonTexts(guardrail_name="test")
|
||||
|
||||
@staticmethod
|
||||
def _web_search_result_block() -> dict[str, object]:
|
||||
return {
|
||||
"type": "web_search_tool_result",
|
||||
"tool_use_id": "srvtoolu_1",
|
||||
"content": [
|
||||
{
|
||||
"type": "web_search_result",
|
||||
"url": "https://docs.python.org/3.14/",
|
||||
"title": "Python 3.14 docs",
|
||||
"encrypted_content": "",
|
||||
"page_age": None,
|
||||
},
|
||||
{
|
||||
"type": "web_search_result",
|
||||
"url": "https://peps.example/745",
|
||||
"title": "PEP 745",
|
||||
"encrypted_content": "",
|
||||
"page_age": None,
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def _web_search_stream_chunks(cls, server_tool_block_type: str = "server_tool_use") -> list[bytes]:
|
||||
events = [
|
||||
(
|
||||
"message_start",
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-5-5",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"usage": {"input_tokens": 1, "output_tokens": 0},
|
||||
},
|
||||
},
|
||||
),
|
||||
(
|
||||
"content_block_start",
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {
|
||||
"type": server_tool_block_type,
|
||||
"id": "srvtoolu_1",
|
||||
"name": "web_search",
|
||||
"input": {},
|
||||
},
|
||||
},
|
||||
),
|
||||
(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "input_json_delta", "partial_json": '{"query": "python 3.14 release date"}'},
|
||||
},
|
||||
),
|
||||
("content_block_stop", {"type": "content_block_stop", "index": 0}),
|
||||
(
|
||||
"content_block_start",
|
||||
{"type": "content_block_start", "index": 1, "content_block": cls._web_search_result_block()},
|
||||
),
|
||||
("content_block_stop", {"type": "content_block_stop", "index": 1}),
|
||||
(
|
||||
"content_block_start",
|
||||
{"type": "content_block_start", "index": 2, "content_block": {"type": "text", "text": ""}},
|
||||
),
|
||||
(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 2,
|
||||
"delta": {"type": "text_delta", "text": "Python 3.14 shipped"},
|
||||
},
|
||||
),
|
||||
("content_block_stop", {"type": "content_block_stop", "index": 2}),
|
||||
(
|
||||
"message_delta",
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
||||
"usage": {"output_tokens": 2},
|
||||
},
|
||||
),
|
||||
("message_stop", {"type": "message_stop"}),
|
||||
]
|
||||
return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events]
|
||||
|
||||
@pytest.mark.parametrize("server_tool_block_type", ["server_tool_use", "tool_use"])
|
||||
@pytest.mark.parametrize(
|
||||
"guardrail_factory",
|
||||
[_python_masking_guardrail, _python_texts_masking_guardrail],
|
||||
ids=["texts_and_tool_calls", "texts_only"],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_ended_stream_rewrites_mask_the_web_search_query_results_and_text(
|
||||
self, guardrail_factory, server_tool_block_type
|
||||
):
|
||||
chunks = self._web_search_stream_chunks(server_tool_block_type)
|
||||
|
||||
await AnthropicMessagesHandler().process_output_streaming_response(
|
||||
responses_so_far=chunks,
|
||||
guardrail_to_apply=guardrail_factory(),
|
||||
litellm_logging_obj=MagicMock(),
|
||||
deliver_ended_stream_rewrites=True,
|
||||
)
|
||||
|
||||
raw = b"".join(chunks).decode()
|
||||
assert "python" not in raw.lower()
|
||||
result_block = next(
|
||||
json.loads(line[len("data:") :])["content_block"]
|
||||
for line in raw.split("\n")
|
||||
if line.startswith("data:") and '"web_search_tool_result"' in line
|
||||
)
|
||||
assert [(result["title"], result["url"]) for result in result_block["content"]] == [
|
||||
("[MASKED] 3.14 docs", "https://docs.[MASKED].org/3.14/"),
|
||||
("PEP 745", "https://peps.example/745"),
|
||||
]
|
||||
assert '"id": "srvtoolu_1"' in raw and '"tool_use_id": "srvtoolu_1"' in raw
|
||||
|
||||
@staticmethod
|
||||
def _last_text_redacting_guardrail() -> CustomGuardrail:
|
||||
class RedactLastText(CustomGuardrail):
|
||||
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
|
||||
return {**inputs, "texts": ["[REDACTED] " + inputs["texts"][-1]]}
|
||||
|
||||
return RedactLastText(guardrail_name="test")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_output_response_keeps_the_reply_as_the_last_text_a_guardrail_scans(self):
|
||||
response = {
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-5-5",
|
||||
"content": [
|
||||
{
|
||||
"type": "server_tool_use",
|
||||
"id": "srvtoolu_1",
|
||||
"name": "web_search",
|
||||
"input": {"query": "python 3.14 release date"},
|
||||
},
|
||||
self._web_search_result_block(),
|
||||
{"type": "text", "text": "Python 3.14 shipped"},
|
||||
],
|
||||
"stop_reason": "end_turn",
|
||||
}
|
||||
|
||||
await AnthropicMessagesHandler().process_output_response(
|
||||
response=response,
|
||||
guardrail_to_apply=self._last_text_redacting_guardrail(),
|
||||
litellm_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert response["content"][0]["input"] == {"query": "python 3.14 release date"}
|
||||
assert response["content"][1] == self._web_search_result_block()
|
||||
assert response["content"][2] == {"type": "text", "text": "[REDACTED] Python 3.14 shipped"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_ended_stream_keeps_the_reply_as_the_last_text_a_guardrail_scans(self):
|
||||
chunks = self._web_search_stream_chunks()
|
||||
|
||||
await AnthropicMessagesHandler().process_output_streaming_response(
|
||||
responses_so_far=chunks,
|
||||
guardrail_to_apply=self._last_text_redacting_guardrail(),
|
||||
litellm_logging_obj=MagicMock(),
|
||||
deliver_ended_stream_rewrites=True,
|
||||
)
|
||||
|
||||
raw = b"".join(chunks).decode()
|
||||
assert "[REDACTED] Python 3.14 shipped" in raw
|
||||
assert raw.count("[REDACTED]") == 1
|
||||
assert json.dumps(self._web_search_result_block()) in raw
|
||||
|
||||
def test_ended_stream_scan_key_differs_from_a_mid_stream_key_by_the_web_search_results(self):
|
||||
handler = AnthropicMessagesHandler()
|
||||
chunks = self._web_search_stream_chunks()
|
||||
|
||||
mid_stream_key = handler.get_streaming_scan_key(chunks[:-2])
|
||||
ended_key = handler.get_streaming_scan_key(chunks)
|
||||
|
||||
assert mid_stream_key.texts == ("Python 3.14 shipped",)
|
||||
assert ended_key != mid_stream_key
|
||||
assert ended_key.texts[-1] == "Python 3.14 shipped"
|
||||
assert {"Python 3.14 docs", "https://docs.python.org/3.14/"} <= set(ended_key.texts)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_output_response_masks_the_web_search_query_results_and_text(self):
|
||||
response = {
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-5-5",
|
||||
"content": [
|
||||
{
|
||||
"type": "server_tool_use",
|
||||
"id": "srvtoolu_1",
|
||||
"name": "web_search",
|
||||
"input": {"query": "python 3.14 release date"},
|
||||
},
|
||||
self._web_search_result_block(),
|
||||
{"type": "text", "text": "Python 3.14 shipped"},
|
||||
],
|
||||
"stop_reason": "end_turn",
|
||||
}
|
||||
|
||||
await AnthropicMessagesHandler().process_output_response(
|
||||
response=response,
|
||||
guardrail_to_apply=self._python_masking_guardrail(),
|
||||
litellm_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert "python" not in json.dumps(response).lower()
|
||||
assert response["content"][0]["input"] == {"query": "[MASKED] 3.14 release date"}
|
||||
assert [(result["title"], result["url"]) for result in response["content"][1]["content"]] == [
|
||||
("[MASKED] 3.14 docs", "https://docs.[MASKED].org/3.14/"),
|
||||
("PEP 745", "https://peps.example/745"),
|
||||
]
|
||||
assert response["content"][2] == {"type": "text", "text": "[MASKED] 3.14 shipped"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self):
|
||||
handler = AnthropicMessagesHandler()
|
||||
|
|
|
|||
|
|
@ -1,6 +1,10 @@
|
|||
import json
|
||||
from collections.abc import Iterable, Mapping
|
||||
from typing import Final
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from typing import Final, cast
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from litellm.llms.anthropic.messages.native.fake_stream_iterator import (
|
||||
FakeAnthropicMessagesStreamIterator,
|
||||
|
|
@ -91,3 +95,140 @@ def test_message_delta_omits_absent_optional_fields():
|
|||
"usage": {"input_tokens": 10, "output_tokens": 2},
|
||||
}
|
||||
assert _message_delta(response) == {"stop_reason": "end_turn", "stop_sequence": None}
|
||||
|
||||
|
||||
_WEB_SEARCH_RESULT_BLOCK: Final = {
|
||||
"type": "web_search_tool_result",
|
||||
"tool_use_id": "srvtoolu_01",
|
||||
"content": [
|
||||
{
|
||||
"type": "web_search_result",
|
||||
"url": "https://example.com/a",
|
||||
"title": "Example A",
|
||||
"encrypted_content": "abc123",
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _content_events(content: Sequence[Mapping[str, object]]) -> tuple[tuple[str, dict[str, object]], ...]:
|
||||
response: Final = {
|
||||
"id": "msg_ws",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"content": list(content),
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5},
|
||||
}
|
||||
return tuple(
|
||||
event for event in _decoded_events(FakeAnthropicMessagesStreamIterator(response=response))
|
||||
if event[0].startswith("content_block_")
|
||||
)
|
||||
|
||||
|
||||
def test_server_tool_use_streams_empty_start_then_input_json_delta():
|
||||
"""A real Anthropic stream opens server_tool_use with an empty input and
|
||||
sends the query as an input_json_delta, which SDK accumulators read."""
|
||||
caller: Final = {"type": "code_execution_20260120", "tool_id": "srvtoolu_00"}
|
||||
events: Final = _content_events(
|
||||
[
|
||||
{
|
||||
"type": "server_tool_use",
|
||||
"id": "srvtoolu_01",
|
||||
"name": "web_search",
|
||||
"input": {"query": "example search"},
|
||||
"caller": caller,
|
||||
}
|
||||
]
|
||||
)
|
||||
assert [event_type for event_type, _ in events] == [
|
||||
"content_block_start",
|
||||
"content_block_delta",
|
||||
"content_block_stop",
|
||||
]
|
||||
assert events[0][1]["content_block"] == {
|
||||
"type": "server_tool_use",
|
||||
"id": "srvtoolu_01",
|
||||
"name": "web_search",
|
||||
"input": {},
|
||||
"caller": caller,
|
||||
}
|
||||
delta: Final = events[1][1]["delta"]
|
||||
assert isinstance(delta, dict)
|
||||
assert delta["type"] == "input_json_delta"
|
||||
assert json.loads(delta["partial_json"]) == {"query": "example search"}
|
||||
|
||||
|
||||
def test_web_search_tool_result_rides_whole_in_content_block_start():
|
||||
events: Final = _content_events([_WEB_SEARCH_RESULT_BLOCK])
|
||||
assert [event_type for event_type, _ in events] == ["content_block_start", "content_block_stop"]
|
||||
assert events[0][1]["content_block"] == _WEB_SEARCH_RESULT_BLOCK
|
||||
|
||||
|
||||
def test_native_search_blocks_keep_indices_aligned_with_text():
|
||||
events: Final = _content_events(
|
||||
[
|
||||
{"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "q"}},
|
||||
_WEB_SEARCH_RESULT_BLOCK,
|
||||
{"type": "text", "text": "Here is what I found."},
|
||||
]
|
||||
)
|
||||
starts: Final = tuple(payload for event_type, payload in events if event_type == "content_block_start")
|
||||
stops: Final = tuple(payload["index"] for event_type, payload in events if event_type == "content_block_stop")
|
||||
assert [(start["index"], cast(dict[str, object], start["content_block"])["type"]) for start in starts] == [
|
||||
(0, "server_tool_use"),
|
||||
(1, "web_search_tool_result"),
|
||||
(2, "text"),
|
||||
]
|
||||
assert stops == (0, 1, 2)
|
||||
|
||||
|
||||
_CALLER: Final = {"type": "code_execution_20260120", "tool_id": "srvtoolu_00"}
|
||||
|
||||
_REPLAYED_BLOCKS: Final = (
|
||||
{"type": "thinking", "thinking": "Look up the release page.", "signature": "sig_01"},
|
||||
{"type": "redacted_thinking", "data": "EncryptedThinking=="},
|
||||
{"type": "tool_use", "id": "toolu_01", "name": "get_weather", "input": {"city": "Paris"}, "caller": _CALLER},
|
||||
{"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_fetch", "input": {"url": "https://example.com/a"}, "caller": _CALLER},
|
||||
_WEB_SEARCH_RESULT_BLOCK,
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Example A is the newest release.",
|
||||
"citations": [
|
||||
{
|
||||
"type": "web_search_result_location",
|
||||
"cited_text": "Example A was released today.",
|
||||
"url": "https://example.com/a",
|
||||
"title": "Example A",
|
||||
"encrypted_index": "idx_01",
|
||||
}
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _sdk_rebuilt_content(block: Mapping[str, object]) -> list[dict[str, object]]:
|
||||
response: Final = {
|
||||
"id": "msg_replay",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-5-5",
|
||||
"content": [block],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5},
|
||||
}
|
||||
body: Final = b"".join(FakeAnthropicMessagesStreamIterator(response=response))
|
||||
transport: Final = httpx.MockTransport(
|
||||
lambda _request: httpx.Response(200, content=body, headers={"content-type": "text/event-stream"})
|
||||
)
|
||||
client: Final = anthropic.Anthropic(api_key="sk-ant-test", http_client=httpx.Client(transport=transport))
|
||||
with client.messages.stream(
|
||||
model="claude-sonnet-5-5", max_tokens=64, messages=[{"role": "user", "content": "hi"}]
|
||||
) as stream:
|
||||
return [content_block.model_dump(exclude_none=True) for content_block in stream.get_final_message().content]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("block", _REPLAYED_BLOCKS, ids=lambda block: str(block["type"]))
|
||||
def test_anthropic_sdk_rebuilds_every_block_field_from_the_stream(block: Mapping[str, object]):
|
||||
assert _sdk_rebuilt_content(block) == [block]
|
||||
|
|
|
|||
|
|
@ -5,9 +5,12 @@ Uses mock guardrails to validate pipeline execution without external services.
|
|||
"""
|
||||
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
import pickle
|
||||
from typing import Literal
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
from itertools import chain
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
|
@ -15,6 +18,7 @@ import pytest
|
|||
import litellm
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import (
|
||||
|
|
@ -1726,21 +1730,6 @@ async def test_streaming_step_discards_legacy_rewrite_whose_texts_do_not_line_up
|
|||
assert chunks == [_chunk()]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_discards_legacy_rewrite_that_changes_a_tool_call(monkeypatch, caplog):
|
||||
masked_tool_call = {"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}}
|
||||
guardrail = _LegacyHookGuardrail(
|
||||
replacement=_legacy_replacement("[REWRITTEN] hello world", tool_calls=[masked_tool_call])
|
||||
)
|
||||
chunks = [_chunk()]
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
result = await _run_legacy_streaming_step(monkeypatch, guardrail, chunks)
|
||||
|
||||
_assert_passed_with_discard_warning(result, caplog, "the legacy hook changed a tool call's name or arguments")
|
||||
assert chunks == [_chunk()]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_discards_legacy_rewrite_that_drops_the_tool_calls(monkeypatch, caplog):
|
||||
guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world", tool_calls=[]))
|
||||
|
|
@ -1772,21 +1761,6 @@ async def test_streaming_step_passes_a_tool_only_stream_the_legacy_hook_left_alo
|
|||
assert not any("discarded" in record.getMessage() for record in caplog.records)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_discards_a_legacy_tool_call_rewrite_on_a_tool_only_stream(monkeypatch, caplog):
|
||||
masked_tool_call = {"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}}
|
||||
guardrail = _LegacyHookGuardrail(replacement=_legacy_replacement(tool_calls=[masked_tool_call]))
|
||||
chunks = [_tool_only_chunk()]
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
result = await _run_legacy_streaming_step(
|
||||
monkeypatch, guardrail, chunks, translation=_ToolOnlyLegacyScanningTranslation()
|
||||
)
|
||||
|
||||
_assert_passed_with_discard_warning(result, caplog, "the legacy hook changed a tool call's name or arguments")
|
||||
assert chunks == [_tool_only_chunk()]
|
||||
|
||||
|
||||
class _NoHooksGuardrail(CustomGuardrail):
|
||||
pass
|
||||
|
||||
|
|
@ -1825,6 +1799,208 @@ async def test_streaming_step_discards_a_legacy_rewrite_the_translation_cannot_r
|
|||
assert chunks == [_chunk()]
|
||||
|
||||
|
||||
_AnthropicHookResponse = dict[str, list[dict[str, object]]]
|
||||
|
||||
|
||||
class _AnthropicTextRewritingLegacyGuardrail(_LegacyHookGuardrail):
|
||||
async def async_post_call_success_hook(
|
||||
self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: _AnthropicHookResponse
|
||||
) -> None:
|
||||
self.calls.append({"response": copy.deepcopy(response)})
|
||||
for block in response["content"]:
|
||||
if block["type"] == "text" and isinstance(text := block.get("text"), str):
|
||||
block["text"] = "[REWRITTEN] " + text
|
||||
|
||||
|
||||
def _anthropic_web_search_stream() -> list[bytes]:
|
||||
result_block: Final = {
|
||||
"type": "web_search_tool_result",
|
||||
"tool_use_id": "srvtoolu_1",
|
||||
"content": [
|
||||
{"type": "web_search_result", "url": "https://example.test/a", "title": "A", "encrypted_content": ""}
|
||||
],
|
||||
}
|
||||
events: Final = (
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "m",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"usage": {"input_tokens": 1, "output_tokens": 0},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}},
|
||||
},
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "input_json_delta", "partial_json": '{"query": "a"}'},
|
||||
},
|
||||
{"type": "content_block_stop", "index": 0},
|
||||
{"type": "content_block_start", "index": 1, "content_block": result_block},
|
||||
{"type": "content_block_stop", "index": 1},
|
||||
{"type": "content_block_start", "index": 2, "content_block": {"type": "text", "text": ""}},
|
||||
{"type": "content_block_delta", "index": 2, "delta": {"type": "text_delta", "text": "hello world"}},
|
||||
{"type": "content_block_stop", "index": 2},
|
||||
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 2}},
|
||||
{"type": "message_stop"},
|
||||
)
|
||||
return _anthropic_sse(events)
|
||||
|
||||
|
||||
def _anthropic_sse(events: Iterable[Mapping[str, object]]) -> list[bytes]:
|
||||
return [f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events]
|
||||
|
||||
|
||||
def _anthropic_tool_only_stream() -> list[bytes]:
|
||||
return _anthropic_sse(
|
||||
(
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "m",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"usage": {"input_tokens": 1, "output_tokens": 0},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": {}},
|
||||
},
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "input_json_delta", "partial_json": '{"ssn"'},
|
||||
},
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "input_json_delta", "partial_json": ': "123"}'},
|
||||
},
|
||||
{"type": "content_block_stop", "index": 0},
|
||||
{"type": "message_delta", "delta": {"stop_reason": "tool_use"}, "usage": {"output_tokens": 2}},
|
||||
{"type": "message_stop"},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _streamed_events(chunks: Sequence[bytes]) -> list[dict[str, object]]:
|
||||
return [
|
||||
json.loads(line[len("data: ") :])
|
||||
for line in b"".join(chunks).decode().splitlines()
|
||||
if line.startswith("data: ")
|
||||
]
|
||||
|
||||
|
||||
def _streamed_block_deltas(chunks: Sequence[bytes], index: int, field: str) -> str:
|
||||
deltas: Final = (
|
||||
event["delta"]
|
||||
for event in _streamed_events(chunks)
|
||||
if event.get("index") == index and event["type"] == "content_block_delta"
|
||||
)
|
||||
return "".join(str(delta[field]) for delta in deltas if isinstance(delta, dict))
|
||||
|
||||
|
||||
def _streamed_web_search_results(chunks: Sequence[bytes]) -> list[dict[str, str]]:
|
||||
blocks: Final = (
|
||||
event["content_block"] for event in _streamed_events(chunks) if event["type"] == "content_block_start"
|
||||
)
|
||||
results: Final = (
|
||||
block["content"] for block in blocks if isinstance(block, dict) and block["type"] == "web_search_tool_result"
|
||||
)
|
||||
return list(chain.from_iterable(results))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_delivers_a_legacy_hook_mask_of_web_search_result_titles(
|
||||
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
chunks: Final = _anthropic_web_search_stream()
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
result: Final = await _run_legacy_streaming_step(
|
||||
monkeypatch, _AnthropicMaskingLegacyGuardrail(), chunks, translation=AnthropicMessagesHandler()
|
||||
)
|
||||
|
||||
assert [step.outcome for step in result.step_results] == ["pass"]
|
||||
assert [(item["title"], item["url"]) for item in _streamed_web_search_results(chunks)] == [
|
||||
("[MASKED]", "https://example.test/a")
|
||||
]
|
||||
assert _streamed_block_deltas(chunks, 2, "text") == "[MASKED] world"
|
||||
assert not any("discarded" in record.getMessage() for record in caplog.records)
|
||||
|
||||
|
||||
class _AnthropicMaskingLegacyGuardrail(_LegacyHookGuardrail):
|
||||
async def async_post_call_success_hook(
|
||||
self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: _AnthropicHookResponse
|
||||
) -> None:
|
||||
for block in response["content"]:
|
||||
if block["type"] == "text" and isinstance(text := block.get("text"), str):
|
||||
block["text"] = text.replace("hello", "[MASKED]")
|
||||
if block["type"] == "tool_use" and isinstance(tool_input := block.get("input"), dict):
|
||||
block["input"] = {key: "[MASKED]" for key in tool_input}
|
||||
if block["type"] == "web_search_tool_result" and isinstance(results := block.get("content"), list):
|
||||
block["content"] = [{**result, "title": "[MASKED]"} for result in results]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"stream, masked_key",
|
||||
[(_anthropic_web_search_stream, "query"), (_anthropic_tool_only_stream, "ssn")],
|
||||
ids=["web_search_query", "tool_only"],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_delivers_a_legacy_hook_mask_of_a_streamed_tool_input(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
stream: Callable[[], list[bytes]],
|
||||
masked_key: str,
|
||||
) -> None:
|
||||
chunks: Final = stream()
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
result: Final = await _run_legacy_streaming_step(
|
||||
monkeypatch, _AnthropicMaskingLegacyGuardrail(), chunks, translation=AnthropicMessagesHandler()
|
||||
)
|
||||
|
||||
assert result.terminal_action == "allow"
|
||||
assert [step.outcome for step in result.step_results] == ["pass"]
|
||||
assert json.loads(_streamed_block_deltas(chunks, 0, "partial_json")) == {masked_key: "[MASKED]"}
|
||||
assert "hello" not in b"".join(chunks).decode()
|
||||
assert not any("discarded" in record.getMessage() for record in caplog.records)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_step_delivers_a_legacy_text_rewrite_on_a_web_search_stream(
|
||||
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
guardrail: Final = _AnthropicTextRewritingLegacyGuardrail()
|
||||
chunks: Final = _anthropic_web_search_stream()
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
result: Final = await _run_legacy_streaming_step(
|
||||
monkeypatch, guardrail, chunks, translation=AnthropicMessagesHandler()
|
||||
)
|
||||
|
||||
assert result.terminal_action == "allow"
|
||||
assert [step.outcome for step in result.step_results] == ["pass"]
|
||||
raw: Final = b"".join(chunks).decode()
|
||||
assert "[REWRITTEN] hello world" in raw
|
||||
assert '"title": "A"' in raw and '"url": "https://example.test/a"' in raw
|
||||
assert not any("discarded" in record.getMessage() for record in caplog.records)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_later_legacy_step_sees_the_stream_as_the_earlier_step_left_it(monkeypatch):
|
||||
masker = _LegacyHookGuardrail(replacement=_legacy_replacement("[REWRITTEN] hello world"))
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue