scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

Update tests and fix FakeAnthropicMessagesStreamIterator for native format

Browse source 
- Update test_websearch_short_circuit.py to validate native Anthropic response
  format (server_tool_use + web_search_tool_result) instead of plain text
- Add tests for structured search hits, tool_use_id linking, usage tracking,
  and uniform provider handling
- Fix server_tool_use/web_search_tool_result block handling in
  _create_content_block_chunks method (was orphaned from cherry-pick)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
0xxmemo 2026-04-06 16:09:30 -05:00
parent cf6c277386
commit 0379a0fe9b
2 changed files with 147 additions and 112 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py
View file

@ -128,6 +128,18 @@ class FakeAnthropicMessagesStreamIterator:
f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n".encode()
)
elif block_type in ("server_tool_use", "web_search_tool_result"):
# Emit the full block as content_block_start — same as
# Anthropic's native streaming format for server-side tools.
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()
@ -169,42 +181,6 @@ class FakeAnthropicMessagesStreamIterator:
block_dict = cast(Dict[str, Any], block)
chunks.extend(self._create_content_block_chunks(block_dict, index))
elif block_type == "server_tool_use":
# Emit the full block as content_block_start (same as
# Anthropic's native streaming format).
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
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()
)
elif block_type == "web_search_tool_result":
# Emit the full block as content_block_start (same as
# Anthropic's native streaming format).
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
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()
)
# 5. message_delta event (with final usage and stop_reason)
# Include cache usage fields so clients that only read message_delta
# (like Claude Code's SDK) see the full input token breakdown.

211
tests/test_litellm/integrations/websearch_interception/test_websearch_short_circuit.py
View file

@ -3,6 +3,9 @@ Unit tests for WebSearch Short-Circuit
Tests the short-circuit path that detects web-search-only /v1/messages requests
and executes the search directly without routing through the backend LLM.
The response uses native Anthropic format (server_tool_use + web_search_tool_result)
so Claude Code's WebSearchTool parser works correctly.
"""
from unittest.mock import AsyncMock, patch
@ -48,11 +51,103 @@ class TestTryShortCircuitSearch:
assert result["type"] == "message"
assert result["role"] == "assistant"
assert result["stop_reason"] == "end_turn"
assert len(result["content"]) == 1
assert result["content"][0]["type"] == "text"
assert "Result" in result["content"][0]["text"]
# Native format: server_tool_use + web_search_tool_result + text
assert result["content"][0]["type"] == "server_tool_use"
assert result["content"][0]["name"] == "web_search"
assert result["content"][1]["type"] == "web_search_tool_result"
assert result["content"][2]["type"] == "text"
assert "Result" in result["content"][2]["text"]
mock_search.assert_called_once_with("Search for Claude Code releases")
@pytest.mark.asyncio
async def test_native_format_search_hits(self):
"""Search results are structured as web_search_result hits"""
logger = WebSearchInterceptionLogger(enabled_providers=["anthropic"])
with patch.object(
logger, "_execute_search", new_callable=AsyncMock
) as mock_search:
mock_search.return_value = (
"Title: First Result\nURL: https://example.com/1\nSnippet: first\n\n"
"Title: Second Result\nURL: https://example.com/2\nSnippet: second"
)
result = await logger.try_short_circuit_search(
model="anthropic/claude-sonnet-4",
messages=[{"role": "user", "content": "Search query"}],
tools=[{"type": "web_search_20250305", "name": "web_search"}],
custom_llm_provider="anthropic",
)
assert result is not None
hits = result["content"][1]["content"]
assert len(hits) == 2
assert hits[0]["type"] == "web_search_result"
assert hits[0]["url"] == "https://example.com/1"
assert hits[0]["title"] == "First Result"
assert hits[1]["url"] == "https://example.com/2"
@pytest.mark.asyncio
async def test_server_tool_use_has_query(self):
"""server_tool_use block contains the original search query"""
logger = WebSearchInterceptionLogger(enabled_providers=["anthropic"])
with patch.object(
logger, "_execute_search", new_callable=AsyncMock
) as mock_search:
mock_search.return_value = "Title: R\nURL: https://x.com\nSnippet: s"
result = await logger.try_short_circuit_search(
model="anthropic/claude-sonnet-4",
messages=[{"role": "user", "content": "trending AI topics"}],
tools=[{"type": "web_search_20250305", "name": "web_search"}],
custom_llm_provider="anthropic",
)
stu = result["content"][0]
assert stu["type"] == "server_tool_use"
assert stu["name"] == "web_search"
assert stu["input"]["query"] == "trending AI topics"
assert stu["id"].startswith("srvtoolu_")
@pytest.mark.asyncio
async def test_tool_use_id_links_blocks(self):
"""server_tool_use.id matches web_search_tool_result.tool_use_id"""
logger = WebSearchInterceptionLogger(enabled_providers=["anthropic"])
with patch.object(
logger, "_execute_search", new_callable=AsyncMock
) as mock_search:
mock_search.return_value = "Title: R\nURL: https://x.com\nSnippet: s"
result = await logger.try_short_circuit_search(
model="anthropic/claude-sonnet-4",
messages=[{"role": "user", "content": "query"}],
tools=[{"type": "web_search_20250305", "name": "web_search"}],
custom_llm_provider="anthropic",
)
assert result["content"][0]["id"] == result["content"][1]["tool_use_id"]
@pytest.mark.asyncio
async def test_usage_includes_web_search_requests(self):
"""Usage includes server_tool_use.web_search_requests count"""
logger = WebSearchInterceptionLogger(enabled_providers=["anthropic"])
with patch.object(
logger, "_execute_search", new_callable=AsyncMock
) as mock_search:
mock_search.return_value = "Title: R\nURL: https://x.com\nSnippet: s"
result = await logger.try_short_circuit_search(
model="anthropic/claude-sonnet-4",
messages=[{"role": "user", "content": "query"}],
tools=[{"type": "web_search_20250305", "name": "web_search"}],
custom_llm_provider="anthropic",
)
assert result["usage"]["server_tool_use"]["web_search_requests"] == 1
@pytest.mark.asyncio
async def test_does_not_short_circuit_mixed_tools(self):
"""Mix of web_search and other tools → NOT short-circuited"""
@ -115,27 +210,26 @@ class TestTryShortCircuitSearch:
assert result is None
@pytest.mark.asyncio
async def test_does_not_short_circuit_bedrock(self):
"""Bedrock has native agentic loop support → NOT short-circuited.
async def test_short_circuits_all_enabled_providers_uniformly(self):
"""All enabled providers (including anthropic, bedrock) go through
the same short-circuit funnel — no provider is skipped."""
for provider in ["anthropic", "bedrock", "github_copilot"]:
logger = WebSearchInterceptionLogger(enabled_providers=[provider])
Providers with a BaseAnthropicMessagesConfig (bedrock, vertex_ai, etc.)
use the agentic loop which includes a follow-up LLM synthesis step.
The short-circuit must not fire for them.
"""
logger = WebSearchInterceptionLogger(
enabled_providers=["bedrock", "github_copilot"]
)
with patch.object(
logger, "_execute_search", new_callable=AsyncMock
) as mock_search:
mock_search.return_value = "Title: R\nURL: https://x.com\nSnippet: s"
result = await logger.try_short_circuit_search(
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
messages=[{"role": "user", "content": "Search for something"}],
tools=[
{"type": "web_search_20250305", "name": "web_search", "max_uses": 8}
],
custom_llm_provider="bedrock",
)
result = await logger.try_short_circuit_search(
model=f"{provider}/claude-sonnet-4",
messages=[{"role": "user", "content": "search query"}],
tools=[{"type": "web_search_20250305", "name": "web_search"}],
custom_llm_provider=provider,
)
assert result is None
assert result is not None, f"Short-circuit should fire for {provider}"
assert result["content"][0]["type"] == "server_tool_use"
@pytest.mark.asyncio
async def test_does_not_short_circuit_no_messages(self):
@ -173,7 +267,10 @@ class TestTryShortCircuitSearch:
)
assert result is not None
assert "Search failed" in result["content"][0]["text"]
# Error text is in the last content block (text)
text_block = result["content"][-1]
assert text_block["type"] == "text"
assert "Search failed" in text_block["text"]
@pytest.mark.asyncio
async def test_response_has_valid_structure(self):
@ -183,7 +280,7 @@ class TestTryShortCircuitSearch:
with patch.object(
logger, "_execute_search", new_callable=AsyncMock
) as mock_search:
mock_search.return_value = "search results here"
mock_search.return_value = "Title: R\nURL: https://x.com\nSnippet: test"
result = await logger.try_short_circuit_search(
model="github_copilot/claude-sonnet-4",
@ -203,11 +300,8 @@ class TestTryShortCircuitSearch:
assert result["stop_sequence"] is None
assert "usage" in result
assert "content" in result
# ---------------------------------------------------------------------------
# Query extraction tests
# ---------------------------------------------------------------------------
# Content has 3 blocks: server_tool_use, web_search_tool_result, text
assert len(result["content"]) == 3
# ---------------------------------------------------------------------------
@ -237,7 +331,7 @@ class TestShortCircuitEntryPoint:
@pytest.mark.asyncio
async def test_returns_dict_when_not_streaming(self):
"""Non-streaming short-circuit → returns dict"""
"""Non-streaming short-circuit → returns dict with native format"""
from litellm.llms.anthropic.experimental_pass_through.messages.handler import (
_try_websearch_short_circuit,
)
@ -246,7 +340,7 @@ class TestShortCircuitEntryPoint:
with patch.object(
logger, "_execute_search", new_callable=AsyncMock
) as mock_search:
mock_search.return_value = "results"
mock_search.return_value = "Title: R\nURL: https://x.com\nSnippet: results"
with patch("litellm.callbacks", [logger]):
result = await _try_websearch_short_circuit(
model="github_copilot/claude-sonnet-4",
@ -257,7 +351,8 @@ class TestShortCircuitEntryPoint:
)
assert isinstance(result, dict)
assert result["content"][0]["text"] == "results"
assert result["content"][0]["type"] == "server_tool_use"
assert result["content"][1]["type"] == "web_search_tool_result"
@pytest.mark.asyncio
async def test_returns_stream_iterator_when_streaming(self):
@ -273,7 +368,9 @@ class TestShortCircuitEntryPoint:
with patch.object(
logger, "_execute_search", new_callable=AsyncMock
) as mock_search:
mock_search.return_value = "streaming results"
mock_search.return_value = (
"Title: Result\nURL: https://example.com\nSnippet: streaming results"
)
with patch("litellm.callbacks", [logger]):
result = await _try_websearch_short_circuit(
model="github_copilot/claude-sonnet-4",
@ -291,12 +388,12 @@ class TestShortCircuitEntryPoint:
chunks.append(chunk)
assert len(chunks) > 0
# First chunk should be message_start
assert b"event: message_start" in chunks[0]
# Last chunk should be message_stop
assert b"event: message_stop" in chunks[-1]
# Should contain the search results text
# Should contain server_tool_use and web_search_tool_result blocks
all_data = b"".join(chunks)
assert b"server_tool_use" in all_data
assert b"web_search_tool_result" in all_data
assert b"streaming results" in all_data
@pytest.mark.asyncio
@ -321,12 +418,7 @@ class TestShortCircuitEntryPoint:
@pytest.mark.asyncio
async def test_uses_original_stream_not_hook_converted(self):
"""Verify that the entry point passes original_stream to the short-circuit.
The pre-request hook converts stream=True → stream=False for the agentic
loop. The short-circuit must use the ORIGINAL stream value so streaming
callers get SSE events instead of a plain dict.
"""
"""Verify that the entry point passes original_stream to the short-circuit."""
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
FakeAnthropicMessagesStreamIterator,
)
@ -338,47 +430,14 @@ class TestShortCircuitEntryPoint:
with patch.object(
logger, "_execute_search", new_callable=AsyncMock
) as mock_search:
mock_search.return_value = "streaming results"
mock_search.return_value = "Title: R\nURL: https://x.com\nSnippet: s"
with patch("litellm.callbacks", [logger]):
# Simulate what anthropic_messages() does: original_stream=True
# is passed to the short-circuit, even though the hook would have
# already converted stream to False in request_kwargs.
result = await _try_websearch_short_circuit(
model="github_copilot/claude-sonnet-4",
messages=[{"role": "user", "content": "search query"}],
tools=[{"type": "web_search_20250305", "name": "web_search"}],
custom_llm_provider="github_copilot",
stream=True, # original_stream, NOT the hook-converted value
stream=True,
)
# Must return a stream iterator, not a plain dict
assert isinstance(result, FakeAnthropicMessagesStreamIterator)
@pytest.mark.asyncio
async def test_short_circuits_with_provider_from_model_string(self):
"""Provider embedded in model string (custom_llm_provider=None) should
still fire the short-circuit when the caller propagates the derived
provider.
"""
from litellm.llms.anthropic.experimental_pass_through.messages.handler import (
_try_websearch_short_circuit,
)
logger = WebSearchInterceptionLogger(enabled_providers=["github_copilot"])
with patch.object(
logger, "_execute_search", new_callable=AsyncMock
) as mock_search:
mock_search.return_value = "results"
with patch("litellm.callbacks", [logger]):
# Simulate the caller having derived custom_llm_provider from
# the model string before calling _try_websearch_short_circuit
result = await _try_websearch_short_circuit(
model="github_copilot/claude-sonnet-4",
messages=[{"role": "user", "content": "search query"}],
tools=[{"type": "web_search_20250305", "name": "web_search"}],
custom_llm_provider="github_copilot",
stream=False,
)
assert result is not None
assert result["content"][0]["text"] == "results"