From 0379a0fe9b25d72d4bb320e70266b5f12e5a055c Mon Sep 17 00:00:00 2001 From: 0xxmemo Date: Mon, 6 Apr 2026 16:09:30 -0500 Subject: [PATCH] Update tests and fix FakeAnthropicMessagesStreamIterator for native format - 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 --- .../messages/fake_stream_iterator.py | 48 +--- .../test_websearch_short_circuit.py | 211 +++++++++++------- 2 files changed, 147 insertions(+), 112 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py index 9dde357f4ee..12fd5ba6f91 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py @@ -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. diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_short_circuit.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_short_circuit.py index 82c1c9839e7..8a81d736f16 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_short_circuit.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_short_circuit.py @@ -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"