mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(anthropic): emit native content blocks when faking a stream
The fake stream iterator only knew text/thinking/redacted_thinking/tool_use, so websearch interception's injected server_tool_use and web_search_tool_result blocks were reduced to a bare content_block_stop on the streaming path
This commit is contained in:
parent
71b825a7f0
commit
a1bc77d045
2 changed files with 111 additions and 2 deletions
|
|
@ -90,12 +90,12 @@ class FakeAnthropicMessagesStreamIterator:
|
|||
}
|
||||
chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode())
|
||||
|
||||
elif block_type == "tool_use":
|
||||
elif block_type in ("tool_use", "server_tool_use", "mcp_tool_use"):
|
||||
content_block_start = {
|
||||
"type": "content_block_start",
|
||||
"index": index,
|
||||
"content_block": {
|
||||
"type": "tool_use",
|
||||
"type": block_type,
|
||||
"id": block_dict.get("id"),
|
||||
"name": block_dict.get("name"),
|
||||
"input": {},
|
||||
|
|
@ -113,6 +113,14 @@ class FakeAnthropicMessagesStreamIterator:
|
|||
}
|
||||
chunks.append(f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n".encode())
|
||||
|
||||
else:
|
||||
content_block_start = {
|
||||
"type": "content_block_start",
|
||||
"index": index,
|
||||
"content_block": block_dict,
|
||||
}
|
||||
chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode())
|
||||
|
||||
content_block_stop = {"type": "content_block_stop", "index": index}
|
||||
chunks.append(f"event: content_block_stop\ndata: {json.dumps(content_block_stop)}\n\n".encode())
|
||||
return chunks
|
||||
|
|
|
|||
|
|
@ -0,0 +1,101 @@
|
|||
"""
|
||||
Unit tests for FakeAnthropicMessagesStreamIterator.
|
||||
|
||||
Regression coverage for websearch interception: the native
|
||||
``server_tool_use`` / ``web_search_tool_result`` blocks that
|
||||
``_inject_native_blocks`` adds must survive the non-streaming -> streaming
|
||||
re-wrap, not collapse into a bare ``content_block_stop``.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, List, cast
|
||||
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
|
||||
FakeAnthropicMessagesStreamIterator,
|
||||
)
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import (
|
||||
AnthropicMessagesResponse,
|
||||
)
|
||||
|
||||
|
||||
def _events(response: Dict[str, Any]) -> List[Dict[str, Any]]:
|
||||
return [
|
||||
json.loads(line.removeprefix("data: "))
|
||||
for chunk in FakeAnthropicMessagesStreamIterator(response=cast(AnthropicMessagesResponse, response))
|
||||
for line in chunk.decode().splitlines()
|
||||
if line.startswith("data: ")
|
||||
]
|
||||
|
||||
|
||||
def _response(content: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||||
return {
|
||||
"id": "msg_123",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": content,
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5},
|
||||
}
|
||||
|
||||
|
||||
def test_web_search_tool_result_block_is_emitted_with_payload():
|
||||
search_result_block = {
|
||||
"type": "web_search_tool_result",
|
||||
"tool_use_id": "srvtoolu_1",
|
||||
"content": [
|
||||
{
|
||||
"type": "web_search_result",
|
||||
"url": "https://docs.litellm.ai",
|
||||
"title": "LiteLLM docs",
|
||||
"page_age": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
events = _events(_response([search_result_block, {"type": "text", "text": "hi"}]))
|
||||
|
||||
starts = [e for e in events if e["type"] == "content_block_start"]
|
||||
assert {"index": 0, "type": "content_block_start", "content_block": search_result_block} in starts
|
||||
assert [e["index"] for e in events if e["type"] == "content_block_stop"] == [0, 1]
|
||||
|
||||
|
||||
def test_server_tool_use_block_keeps_type_and_input():
|
||||
events = _events(
|
||||
_response(
|
||||
[
|
||||
{
|
||||
"type": "server_tool_use",
|
||||
"id": "srvtoolu_1",
|
||||
"name": "web_search",
|
||||
"input": {"query": "litellm"},
|
||||
}
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
start = next(e for e in events if e["type"] == "content_block_start")
|
||||
assert start["content_block"] == {
|
||||
"type": "server_tool_use",
|
||||
"id": "srvtoolu_1",
|
||||
"name": "web_search",
|
||||
"input": {},
|
||||
}
|
||||
delta = next(e for e in events if e["type"] == "content_block_delta")
|
||||
assert json.loads(delta["delta"]["partial_json"]) == {"query": "litellm"}
|
||||
|
||||
|
||||
def test_unknown_block_type_is_passed_through_verbatim():
|
||||
block = {"type": "code_execution_tool_result", "tool_use_id": "srvtoolu_2", "content": {"stdout": "42"}}
|
||||
events = _events(_response([block]))
|
||||
|
||||
start = next(e for e in events if e["type"] == "content_block_start")
|
||||
assert start["content_block"] == block
|
||||
|
||||
|
||||
def test_text_block_still_streams_as_delta():
|
||||
events = _events(_response([{"type": "text", "text": "hello"}]))
|
||||
|
||||
start = next(e for e in events if e["type"] == "content_block_start")
|
||||
assert start["content_block"] == {"type": "text", "text": ""}
|
||||
delta = next(e for e in events if e["type"] == "content_block_delta")
|
||||
assert delta["delta"] == {"type": "text_delta", "text": "hello"}
|
||||
Loading…
Add table
Reference in a new issue