diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 300c311f36dd..481cf7fce8e5 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -697,6 +697,27 @@ async def async_build_agentic_loop_plan( """ return AgenticLoopPlan(run_agentic_loop=False) + async def async_post_agentic_loop_response_hook( + self, + response: Any, + plan: AgenticLoopPlan, + kwargs: Dict, + ) -> Any: + """ + Post-process the response returned by the agentic-loop follow-up call. + + Called after BaseLLMHTTPHandler executes ``AgenticLoopPlan.request_patch`` + and receives the final response from the provider. Lets callbacks shape + what the client sees without bypassing the loop's safety / observability + machinery (depth tracking, fingerprinting, etc.). + + Use ``plan.metadata`` to carry whatever the build step decided to expose + for post-processing (e.g. native tool_result blocks to inject). + + Default returns ``response`` unchanged. + """ + return response + async def async_should_run_chat_completion_agentic_loop( self, response: Any, diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 41618c726278..37528e7dcd50 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -19,12 +19,14 @@ from litellm.integrations.websearch_interception.tools import ( get_litellm_web_search_tool, get_litellm_web_search_tool_openai, + is_anthropic_native_web_search_tool, is_web_search_tool, is_web_search_tool_chat_completion, ) from litellm.integrations.websearch_interception.transformation import ( WebSearchTransformation, ) +from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.types.integrations.websearch_interception import ( WebSearchInterceptionConfig, ) @@ -36,6 +38,16 @@ from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager +# Key used to flag, on per-request kwargs, that the originating client sent +# an Anthropic-native ``web_search_*`` tool — meaning the final response +# should include ``web_search_tool_result`` content blocks so the client +# (e.g. Claude Desktop's citations panel) can render sources. +WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY = "_websearch_interception_emit_native_blocks" + +# Key on ``AgenticLoopPlan.metadata`` carrying the list of pre-built +# ``web_search_tool_result`` blocks to inject into the final response. +WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY = "websearch_native_blocks" + class WebSearchInterceptionLogger(CustomLogger): """ @@ -152,22 +164,55 @@ async def try_short_circuit_search( f"(provider={provider_str}, query='{query}')" ) - # Execute search + # Native clients (Claude Desktop / Cowork / Anthropic SDK) make a + # standalone /v1/messages sub-request just for the search, and they + # expect the response in native shape with server_tool_use + + # web_search_tool_result content blocks so the citations panel can + # render. The agentic-loop post-hook never fires on this path because + # there is no model call — emit the native blocks here instead. + native_tool = next( + (t for t in tools if is_anthropic_native_web_search_tool(t)), + None, + ) + + # Execute search — keep the structured SearchResponse so the native + # block can carry per-result url/title/page_age. try: - search_result_text = await self._execute_search(query) + search_result_text, structured = await self._execute_search(query) except Exception as e: verbose_logger.error( f"WebSearchInterception: Short-circuit search failed: {e}" ) - search_result_text = f"Search failed: {e}" + search_result_text, structured = f"Search failed: {e}", None + + content: List[Dict[str, Any]] = [] + if native_tool is not None: + tool_use_id = f"srvtoolu_{uuid.uuid4().hex}" + tool_name = native_tool.get("name") or "web_search" + content.append( + { + "type": "server_tool_use", + "id": tool_use_id, + "name": tool_name, + "input": {"query": query}, + } + ) + content.append( + WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id=tool_use_id, + search_response=structured, + ) + ) + # 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}) - # Build synthetic Anthropic response response: Dict[str, Any] = { "id": f"msg_{str(uuid.uuid4())}", "type": "message", "role": "assistant", "model": model, - "content": [{"type": "text", "text": search_result_text}], + "content": content, "stop_reason": "end_turn", "stop_sequence": None, "usage": {"input_tokens": 0, "output_tokens": 0}, @@ -175,7 +220,8 @@ async def try_short_circuit_search( verbose_logger.debug( "WebSearchInterception: Short-circuit search completed, " - f"returning synthetic response ({len(search_result_text)} chars)" + f"returning synthetic response ({len(search_result_text)} chars, " + f"native_blocks={native_tool is not None})" ) return response @@ -219,6 +265,14 @@ async def async_pre_call_deployment_hook( "WebSearchInterception: Converting native web_search tools to LiteLLM standard" ) + # If the client sent an Anthropic-native web_search_* tool, mark the + # request so the agentic loop emits native web_search_tool_result + # blocks in the final response (matches async_pre_request_hook). This + # deployment hook fires before async_pre_request_hook on some paths, + # so flagging here ensures the signal isn't lost regardless of order. + if any(is_anthropic_native_web_search_tool(t) for t in tools): + kwargs[WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY] = True + # Convert native/custom web_search tools to LiteLLM standard converted_tools = [] for tool in tools: @@ -342,6 +396,14 @@ async def async_pre_request_hook( f"WebSearchInterception: Pre-request hook triggered for provider={custom_llm_provider}" ) + # If the client sent an Anthropic-native web_search_* tool, mark the + # request so the agentic loop emits native web_search_tool_result + # blocks in the final response (for citations panels, etc.). The flag + # is read by async_build_agentic_loop_plan; the leading underscore + # prefix ensures it is stripped before the follow-up call kwargs. + if any(is_anthropic_native_web_search_tool(t) for t in tools): + kwargs[WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY] = True + # Convert native web search tools to LiteLLM standard converted_tools = [] for tool in tools: @@ -591,7 +653,7 @@ async def async_build_agentic_loop_plan( ) -> AgenticLoopPlan: tool_calls = tools["tool_calls"] thinking_blocks = tools.get("thinking_blocks", []) - request_patch = await self._build_anthropic_request_patch( + request_patch, structured_results = await self._build_anthropic_request_patch( model=model, messages=messages, tool_calls=tool_calls, @@ -600,12 +662,92 @@ async def async_build_agentic_loop_plan( logging_obj=logging_obj, kwargs=kwargs, ) + + metadata: Dict[str, Any] = { + "tool_type": "websearch", + "response_format": "anthropic", + } + + # If the client request originally carried a native web_search_* tool, + # pre-build the Anthropic-native ``web_search_tool_result`` blocks now + # (while we still have the structured SearchResponse list) and stash + # them on plan metadata for the post-hook to inject. + if kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY): + metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY] = ( + self._build_native_result_blocks( + tool_calls=tool_calls, + structured_results=structured_results, + ) + ) + return AgenticLoopPlan( run_agentic_loop=True, request_patch=request_patch, - metadata={"tool_type": "websearch", "response_format": "anthropic"}, + metadata=metadata, ) + async def async_post_agentic_loop_response_hook( + self, + response: Any, + plan: AgenticLoopPlan, + kwargs: Dict, + ) -> Any: + """ + Inject Anthropic-native ``web_search_tool_result`` blocks into the + final response when the originating client used a native + ``web_search_*`` tool. + + See ``WebSearchTransformation.build_web_search_tool_result_block`` for + the block shape. The blocks are prepended to ``response.content`` so + Anthropic-native clients (Claude Desktop, the Anthropic SDK) can + render citations / sources alongside the model's textual reply. + """ + native_blocks = plan.metadata.get(WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY) + if not native_blocks: + return response + return self._inject_native_blocks(response, native_blocks) + + @staticmethod + def _build_native_result_blocks( + tool_calls: List[Dict], + structured_results: List[Optional[SearchResponse]], + ) -> List[Dict[str, Any]]: + """Build one ``web_search_tool_result`` block per tool_call.""" + blocks: List[Dict[str, Any]] = [] + for i, tool_call in enumerate(tool_calls): + tool_use_id = tool_call.get("id") or "" + structured = structured_results[i] if i < len(structured_results) else None + blocks.append( + WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id=tool_use_id, + search_response=structured, + ) + ) + return blocks + + @staticmethod + def _inject_native_blocks( + response: Any, native_blocks: List[Dict[str, Any]] + ) -> Any: + """Prepend native blocks to response content, dict or object form.""" + if not native_blocks: + return response + if isinstance(response, dict): + existing = response.get("content") or [] + response["content"] = list(native_blocks) + list(existing) + return response + existing = getattr(response, "content", None) or [] + try: + response.content = list(native_blocks) + list(existing) + except (AttributeError, TypeError): + # Object refused write — fall through and leave the response + # untouched rather than crash the request. + verbose_logger.debug( + "WebSearchInterception: could not inject native blocks into " + f"response of type {type(response).__name__}" + ) + return response + async def async_run_chat_completion_agentic_loop( self, tools: Dict, @@ -733,7 +875,7 @@ async def _execute_agentic_loop( kwargs: Dict, ) -> Any: """Legacy path: execute search + build patch + run follow-up call.""" - request_patch = await self._build_anthropic_request_patch( + request_patch, structured_results = await self._build_anthropic_request_patch( model=model, messages=messages, tool_calls=tool_calls, @@ -755,7 +897,7 @@ async def _execute_agentic_loop( if max_tokens is None: max_tokens = cast(int, kwargs.get("max_tokens", 1024)) - return await anthropic_messages.acreate( + response = await anthropic_messages.acreate( max_tokens=max_tokens, messages=request_patch.messages, model=request_patch.model or model, @@ -763,6 +905,18 @@ async def _execute_agentic_loop( **request_patch.kwargs, ) + # Legacy path: the new path goes through the typed plan + core + # dispatcher which runs the post-hook automatically. Mirror the + # native-block injection here so both paths behave identically. + if kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY): + native_blocks = self._build_native_result_blocks( + tool_calls=tool_calls, + structured_results=structured_results, + ) + response = self._inject_native_blocks(response, native_blocks) + + return response + async def _build_anthropic_request_patch( self, model: str, @@ -772,8 +926,16 @@ async def _build_anthropic_request_patch( anthropic_messages_optional_request_params: Dict, logging_obj: Any, kwargs: Dict, - ) -> AgenticLoopRequestPatch: - """Execute litellm.search() and build follow-up request patch.""" + ) -> Tuple[AgenticLoopRequestPatch, List[Optional[SearchResponse]]]: + """ + Execute litellm.search() and build follow-up request patch. + + Returns the patch alongside the parallel list of structured + ``SearchResponse`` objects (one per tool_call, ``None`` when the + search failed or the tool_call had no query). The caller uses these + to optionally build Anthropic-native ``web_search_tool_result`` + content blocks for the final response. + """ # Extract search queries from tool_use blocks search_tasks = [] @@ -797,23 +959,38 @@ async def _build_anthropic_request_patch( ) search_results = await asyncio.gather(*search_tasks, return_exceptions=True) - # Handle any exceptions in search results + # Split the gathered (text, structured) tuples into two parallel lists. + # The text list feeds the follow-up model call; the structured list + # is returned to the caller for native-block emission. final_search_results: List[str] = [] + structured_results: List[Optional[SearchResponse]] = [] for i, result in enumerate(search_results): if isinstance(result, Exception): verbose_logger.error( f"WebSearchInterception: Search {i} failed with error: {str(result)}" ) final_search_results.append(f"Search failed: {str(result)}") - elif isinstance(result, str): - # Explicitly cast to str for type checker - final_search_results.append(cast(str, result)) + structured_results.append(None) + elif isinstance(result, tuple) and len(result) == 2: + text_value, structured_value = result + final_search_results.append( + cast(str, text_value) + if isinstance(text_value, str) + else str(text_value) + ) + structured_results.append( + structured_value + if isinstance(structured_value, SearchResponse) + else None + ) else: - # Should never happen, but handle for type safety + # Defensive: legacy callers / unexpected shape — preserve text, + # drop structure. verbose_logger.debug( f"WebSearchInterception: Unexpected result type {type(result)} at index {i}" ) final_search_results.append(str(result)) + structured_results.append(None) # Build assistant and user messages using transformation assistant_message, user_message = WebSearchTransformation.transform_response( @@ -859,16 +1036,26 @@ async def _build_anthropic_request_patch( len(follow_up_messages), len(final_search_results), ) - return AgenticLoopRequestPatch( + patch = AgenticLoopRequestPatch( model=full_model_name, messages=follow_up_messages, max_tokens=max_tokens, optional_params=optional_params_without_max_tokens, kwargs=kwargs_for_followup, ) + return patch, structured_results - async def _execute_search(self, query: str) -> str: - """Execute a single web search using router's search tools""" + async def _execute_search(self, query: str) -> Tuple[str, Optional[SearchResponse]]: + """ + Execute a single web search using router's search tools. + + Returns both the formatted text (fed back to the model in the follow-up + call) and the structured ``SearchResponse`` (preserved so callers can + build Anthropic-native ``web_search_tool_result`` blocks for clients + that requested a native ``web_search_*`` tool). The structured value + is None on the failure path so callers can still emit an empty result + block rather than dropping the search entirely. + """ try: # Import router from proxy_server try: @@ -934,7 +1121,7 @@ async def _execute_search(self, query: str) -> str: verbose_logger.debug( f"WebSearchInterception: Search completed for '{query}', got {len(search_result_text)} chars" ) - return search_result_text + return search_result_text, result except Exception as e: verbose_logger.error( f"WebSearchInterception: Search failed for '{query}': {str(e)}" @@ -1015,7 +1202,8 @@ async def _build_chat_completion_request_patch( # noqa: PLR0915 ) search_results = await asyncio.gather(*search_tasks, return_exceptions=True) - # Handle any exceptions in search results + # Chat-completion path only needs text — OpenAI tool_result format + # has no equivalent of Anthropic's web_search_tool_result block. final_search_results: List[str] = [] for i, result in enumerate(search_results): if isinstance(result, Exception): @@ -1023,8 +1211,13 @@ async def _build_chat_completion_request_patch( # noqa: PLR0915 f"WebSearchInterception: Search {i} failed with error: {str(result)}" ) final_search_results.append(f"Search failed: {str(result)}") - elif isinstance(result, str): - final_search_results.append(cast(str, result)) + elif isinstance(result, tuple) and len(result) == 2: + text_value, _ = result + final_search_results.append( + cast(str, text_value) + if isinstance(text_value, str) + else str(text_value) + ) else: verbose_logger.debug( f"WebSearchInterception: Unexpected result type {type(result)} at index {i}" @@ -1112,9 +1305,11 @@ async def _build_chat_completion_request_patch( # noqa: PLR0915 kwargs=kwargs_for_followup, ) - async def _create_empty_search_result(self) -> str: + async def _create_empty_search_result( + self, + ) -> Tuple[str, Optional[SearchResponse]]: """Create an empty search result for tool calls without queries""" - return "No search query provided" + return "No search query provided", None @staticmethod def initialize_from_proxy_config( diff --git a/litellm/integrations/websearch_interception/tools.py b/litellm/integrations/websearch_interception/tools.py index e373b64cdda0..b29372af9edc 100644 --- a/litellm/integrations/websearch_interception/tools.py +++ b/litellm/integrations/websearch_interception/tools.py @@ -126,6 +126,27 @@ def is_web_search_tool_chat_completion(tool: Dict[str, Any]) -> bool: return False +def is_anthropic_native_web_search_tool(tool: Dict[str, Any]) -> bool: + """ + Check if a tool is an Anthropic-native ``web_search_*`` tool. + + Native clients (Anthropic SDK, Claude Desktop, Anthropic Console) send + tools like ``{"type": "web_search_20250305", "name": "web_search"}`` and + expect the response to contain ``web_search_tool_result`` content blocks + so that citations can be rendered. This helper identifies that contract + so the agentic loop can emit native-format blocks for those clients + without affecting clients that send the LiteLLM standard tool. + + Returns False for the LiteLLM standard tool (``litellm_web_search``), + the OpenAI-shaped variant, the bare ``WebSearch`` legacy name, and the + bare ``web_search`` name (Claude Code style). + """ + tool_type = tool.get("type", "") + if not isinstance(tool_type, str): + return False + return tool_type.startswith("web_search_") and tool_type != "function" + + def is_web_search_tool(tool: Dict[str, Any]) -> bool: """ Check if a tool is a web search tool (native or LiteLLM standard). @@ -135,7 +156,22 @@ def is_web_search_tool(tool: Dict[str, Any]) -> bool: - OpenAI format: type == "function" with function.name == "litellm_web_search" - Anthropic native: type starts with "web_search_" (e.g., "web_search_20250305") - Claude Code: name == "web_search" with a type field - - Custom: name == "WebSearch" (legacy format) + - Custom: name == "WebSearch" (legacy interception marker — only matched + when input_schema is absent; see note below) + + Note on the legacy ``WebSearch`` name: + Clients like Claude Desktop / Cowork ship a *client-side* tool called + ``WebSearch`` (a fully-formed Anthropic client tool with its own + ``input_schema``) that they handle themselves. Treating that as our + interception marker hijacks it server-side and the client's own tool + handler never fires — which means Cowork's separate native + ``web_search_20250305`` sub-request (where citation data actually + flows) never gets made. + + Real Anthropic client tools always carry an ``input_schema`` (the API + rejects them otherwise), so a bare ``{name: "WebSearch"}`` with no + schema is the only thing that could be a legacy interception marker. + Gate the match on schema absence to keep both groups working. Args: tool: Tool dictionary to check @@ -152,6 +188,10 @@ def is_web_search_tool(tool: Dict[str, Any]) -> bool: True >>> is_web_search_tool({"name": "calculator"}) False + >>> is_web_search_tool({"name": "WebSearch"}) # legacy interception marker + True + >>> is_web_search_tool({"name": "WebSearch", "input_schema": {"type": "object"}}) # Cowork client tool + False """ tool_name = tool.get("name", "") tool_type = tool.get("type", "") @@ -175,8 +215,9 @@ def is_web_search_tool(tool: Dict[str, Any]) -> bool: if tool_name == "web_search" and tool_type: return True - # Check for legacy WebSearch format - if tool_name == "WebSearch": + # Legacy "WebSearch" interception marker — only when no schema is + # present, so real client-side WebSearch tools (Cowork) pass through. + if tool_name == "WebSearch" and "input_schema" not in tool: return True return False diff --git a/litellm/integrations/websearch_interception/transformation.py b/litellm/integrations/websearch_interception/transformation.py index 00d4829ad39f..9c20a3f6c775 100644 --- a/litellm/integrations/websearch_interception/transformation.py +++ b/litellm/integrations/websearch_interception/transformation.py @@ -100,11 +100,14 @@ def _detect_from_non_streaming_response( block_id = getattr(block, "id", None) block_input = getattr(block, "input", {}) - # Check for LiteLLM standard or legacy web search tools - # Handles: litellm_web_search, WebSearch, web_search + # Detect tool_use blocks that came from interception. After + # pre-request conversion the model always sees + # ``litellm_web_search``; the bare ``web_search`` entry handles + # callers that bypass our pre-request hooks (e.g. direct + # litellm.acompletion). "WebSearch" is intentionally omitted — + # see is_web_search_tool for the Cowork rationale. if block_type == "tool_use" and block_name in ( LITELLM_WEB_SEARCH_TOOL_NAME, - "WebSearch", "web_search", ): # Convert to dict for easier handling @@ -190,10 +193,12 @@ def _detect_from_openai_response( getattr(function, "arguments", None) if function else None ) - # Check for LiteLLM standard or legacy web search tools + # Detect function-style web search tool_calls. ``WebSearch`` is + # intentionally omitted — see is_web_search_tool for the Cowork + # rationale (clients ship their own client-side ``WebSearch`` and + # we must not hijack it). if tool_type == "function" and function_name in ( LITELLM_WEB_SEARCH_TOOL_NAME, - "WebSearch", "web_search", ): # Parse arguments (might be JSON string) @@ -350,6 +355,57 @@ def _transform_response_openai( return assistant_message, tool_messages + @staticmethod + def build_web_search_tool_result_block( + tool_use_id: str, + search_response: Optional[SearchResponse], + ) -> Dict[str, Any]: + """ + Build an Anthropic-native ``web_search_tool_result`` content block. + + Native Anthropic clients (Claude Desktop, the Anthropic SDK, the + Anthropic Console) expect search-tool results to be returned as + structured ``web_search_tool_result`` blocks so that citations and + source links can be rendered. The agentic loop currently feeds the + model a flat text blob in the follow-up call (which is correct — the + model needs readable evidence). This helper produces the *additional* + block that should accompany the model's text reply when the original + request used a native ``web_search_*`` tool. + + Spec reference: + https://docs.anthropic.com/en/api/web-search-tool + + Args: + tool_use_id: The ``tool_use_id`` the model emitted on the first + turn. Must match exactly so the client can pair the result + with its tool_use block. + search_response: Structured ``SearchResponse`` from + ``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"). + """ + items: List[Dict[str, Any]] = [] + if search_response is not None: + results = getattr(search_response, "results", None) or [] + for r in results: + url = getattr(r, "url", "") or "" + title = getattr(r, "title", "") or "" + page_age = getattr(r, "date", None) or getattr(r, "last_updated", None) + items.append( + { + "type": "web_search_result", + "url": url, + "title": title, + "page_age": page_age, + "encrypted_content": "", + } + ) + return { + "type": "web_search_tool_result", + "tool_use_id": tool_use_id, + "content": items, + } + @staticmethod def format_search_response(result: SearchResponse) -> str: """ diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index fa1253d90058..2ff63cc2d7f0 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -4634,6 +4634,7 @@ async def _execute_anthropic_agentic_plan( fingerprints: List[str], fingerprint: str, stream: bool = False, + callback: Optional[Any] = None, ) -> Any: from litellm.anthropic_interface import messages as anthropic_messages @@ -4675,7 +4676,7 @@ async def _execute_anthropic_agentic_plan( kwargs_for_followup["max_agentic_loops"] = max_loops kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint] - return await anthropic_messages.acreate( + response = await anthropic_messages.acreate( **{ "max_tokens": max_tokens, "messages": patch.messages, @@ -4686,6 +4687,23 @@ async def _execute_anthropic_agentic_plan( } ) + if callback is not None: + try: + response = await callback.async_post_agentic_loop_response_hook( + response=response, plan=plan, kwargs=kwargs + ) + except Exception as e: + _call_id = getattr(logging_obj, "litellm_call_id", "unknown") + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in " + "async_post_agentic_loop_response_hook [call_id=%s model=%s]: %s", + _call_id, + model, + str(e), + ) + + return response + async def _execute_chat_completion_agentic_plan( self, plan: AgenticLoopPlan, @@ -4869,6 +4887,7 @@ async def _call_agentic_completion_hooks( fingerprints=fingerprints, fingerprint=fingerprint, stream=stream, + callback=callback, ) except Exception as e: _call_id = getattr(logging_obj, "litellm_call_id", "unknown") diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py new file mode 100644 index 000000000000..544abab8dcf4 --- /dev/null +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py @@ -0,0 +1,484 @@ +""" +Tests for Anthropic-native ``web_search_tool_result`` block emission. + +Covers the path that lets Claude Desktop / Anthropic SDK clients render +citations when their request used a native ``web_search_*`` tool against a +provider (e.g. Bedrock) that can't run web search natively. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.integrations.websearch_interception.handler import ( + WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY, + WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY, + WebSearchInterceptionLogger, +) +from litellm.integrations.websearch_interception.tools import ( + is_anthropic_native_web_search_tool, + is_web_search_tool, +) +from litellm.integrations.websearch_interception.transformation import ( + WebSearchTransformation, +) +from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult +from litellm.types.integrations.custom_logger import ( + AgenticLoopPlan, + AgenticLoopRequestPatch, +) + + +def _make_search_response() -> SearchResponse: + return SearchResponse( + results=[ + SearchResult( + title="LiteLLM Docs", + url="https://docs.litellm.ai/", + snippet="Unified interface for LLMs.", + date="2025-01-15", + ), + SearchResult( + title="Bedrock Pricing", + url="https://aws.amazon.com/bedrock/pricing/", + snippet="Pay-per-use pricing model.", + date=None, + ), + ] + ) + + +class TestIsAnthropicNativeWebSearchTool: + """The detector must match native tools without catching look-alikes.""" + + def test_matches_web_search_20250305(self): + assert is_anthropic_native_web_search_tool( + {"type": "web_search_20250305", "name": "web_search", "max_uses": 5} + ) + + def test_matches_future_dated_variant(self): + assert is_anthropic_native_web_search_tool( + {"type": "web_search_20260101", "name": "web_search"} + ) + + def test_rejects_litellm_standard(self): + assert not is_anthropic_native_web_search_tool( + {"name": "litellm_web_search", "input_schema": {}} + ) + + def test_rejects_openai_function_shape(self): + assert not is_anthropic_native_web_search_tool( + {"type": "function", "function": {"name": "litellm_web_search"}} + ) + + def test_rejects_claude_desktop_builtin(self): + # Claude Desktop's builtin client-side ``WebSearch`` tool must not be + # misidentified — that's the collision PR #25242 introduced. + assert not is_anthropic_native_web_search_tool({"name": "WebSearch"}) + + def test_rejects_unrelated_tool(self): + assert not is_anthropic_native_web_search_tool( + {"type": "function", "function": {"name": "calculator"}} + ) + + def test_handles_missing_type(self): + assert not is_anthropic_native_web_search_tool({"name": "web_search"}) + + +class TestLegacyWebSearchNameGate: + """The bare ``WebSearch`` name is a legacy interception marker. Real + client-side ``WebSearch`` tools (Cowork, Claude Desktop) carry an + ``input_schema`` and must pass through untouched — otherwise the proxy + hijacks them server-side and the client's own tool handler never fires, + which means the separate ``web_search_20250305`` sub-request (where + citations actually flow) is never made.""" + + def test_bare_legacy_name_still_matched(self): + # Caller deliberately uses the bare-name interception marker — + # back-compat for anyone relying on the old shape. + assert is_web_search_tool({"name": "WebSearch"}) + + def test_real_client_tool_passes_through(self): + # Cowork's client-side WebSearch tool ships with input_schema. + cowork_tool = { + "name": "WebSearch", + "input_schema": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + } + assert not is_web_search_tool(cowork_tool) + + def test_real_client_tool_with_description_passes_through(self): + # description-only client tools (no schema) are not valid Anthropic + # tools; only the schema-bearing shape is the disambiguator. This + # case stays matched on the assumption it's a legacy marker. + assert is_web_search_tool({"name": "WebSearch", "description": "search"}) + + +class TestBuildWebSearchToolResultBlock: + """The block-builder must produce the Anthropic-native shape exactly.""" + + def test_shape_with_results(self): + block = WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="toolu_abc", + search_response=_make_search_response(), + ) + assert block["type"] == "web_search_tool_result" + assert block["tool_use_id"] == "toolu_abc" + assert len(block["content"]) == 2 + first = block["content"][0] + assert first["type"] == "web_search_result" + assert first["url"] == "https://docs.litellm.ai/" + assert first["title"] == "LiteLLM Docs" + assert first["page_age"] == "2025-01-15" + assert first["encrypted_content"] == "" + + def test_handles_none_search_response(self): + block = WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="toolu_abc", + search_response=None, + ) + assert block["type"] == "web_search_tool_result" + assert block["tool_use_id"] == "toolu_abc" + assert block["content"] == [] + + def test_handles_empty_results(self): + block = WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="toolu_xyz", + search_response=SearchResponse(results=[]), + ) + assert block["content"] == [] + + +class TestPreRequestHookFlagsNativeTools: + """The pre-request hook must mark the request when a native tool is used.""" + + @pytest.mark.asyncio + async def test_native_tool_sets_flag(self): + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + kwargs = { + "tools": [ + {"type": "web_search_20250305", "name": "web_search", "max_uses": 5} + ], + "litellm_params": {"custom_llm_provider": "bedrock"}, + } + out = await logger.async_pre_request_hook( + model="bedrock/claude", messages=[], kwargs=kwargs + ) + assert out is not None + assert out.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY) is True + + @pytest.mark.asyncio + async def test_litellm_standard_tool_does_not_set_flag(self): + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + kwargs = { + "tools": [{"name": "litellm_web_search", "input_schema": {}}], + "litellm_params": {"custom_llm_provider": "bedrock"}, + } + out = await logger.async_pre_request_hook( + model="bedrock/claude", messages=[], kwargs=kwargs + ) + assert out is not None + assert WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY not in out + + +class TestBuildPlanAttachesBlocks: + """async_build_agentic_loop_plan must put pre-built blocks on metadata.""" + + @pytest.mark.asyncio + async def test_metadata_carries_blocks_when_flag_set(self): + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + tool_calls = [ + { + "id": "toolu_one", + "type": "tool_use", + "name": "litellm_web_search", + "input": {"query": "what is litellm"}, + } + ] + patch_obj = AgenticLoopRequestPatch( + model="bedrock/claude", + messages=[{"role": "user", "content": "hi"}], + max_tokens=1024, + ) + structured = [_make_search_response()] + + with patch.object( + logger, + "_build_anthropic_request_patch", + new=AsyncMock(return_value=(patch_obj, structured)), + ): + plan = await logger.async_build_agentic_loop_plan( + tools={"tool_calls": tool_calls, "thinking_blocks": []}, + model="bedrock/claude", + messages=[], + response=MagicMock(), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(model_call_details={}), + stream=False, + kwargs={WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True}, + ) + + blocks = plan.metadata.get(WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY) + assert isinstance(blocks, list) + assert len(blocks) == 1 + assert blocks[0]["type"] == "web_search_tool_result" + assert blocks[0]["tool_use_id"] == "toolu_one" + assert blocks[0]["content"][0]["url"] == "https://docs.litellm.ai/" + + @pytest.mark.asyncio + async def test_metadata_does_not_carry_blocks_when_flag_absent(self): + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + tool_calls = [ + { + "id": "toolu_one", + "type": "tool_use", + "name": "litellm_web_search", + "input": {"query": "what is litellm"}, + } + ] + patch_obj = AgenticLoopRequestPatch( + model="bedrock/claude", + messages=[{"role": "user", "content": "hi"}], + max_tokens=1024, + ) + + with patch.object( + logger, + "_build_anthropic_request_patch", + new=AsyncMock(return_value=(patch_obj, [_make_search_response()])), + ): + plan = await logger.async_build_agentic_loop_plan( + tools={"tool_calls": tool_calls, "thinking_blocks": []}, + model="bedrock/claude", + messages=[], + response=MagicMock(), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(model_call_details={}), + stream=False, + kwargs={}, + ) + + assert WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY not in plan.metadata + + +class TestPostHookInjectsBlocks: + """The post-hook must prepend blocks; absent metadata is a no-op.""" + + @pytest.mark.asyncio + async def test_injects_when_metadata_present(self): + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + block = WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="toolu_abc", + search_response=_make_search_response(), + ) + plan = AgenticLoopPlan( + run_agentic_loop=True, + metadata={WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: [block]}, + ) + response = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "Based on the search..."}], + "stop_reason": "end_turn", + } + + out = await logger.async_post_agentic_loop_response_hook( + response=response, plan=plan, kwargs={} + ) + + # Native block must be first so the client can pair it with the + # tool_use before reading the assistant text. + assert out["content"][0]["type"] == "web_search_tool_result" + assert out["content"][0]["tool_use_id"] == "toolu_abc" + assert out["content"][1]["type"] == "text" + + @pytest.mark.asyncio + async def test_noop_when_metadata_absent(self): + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + plan = AgenticLoopPlan(run_agentic_loop=True, metadata={}) + response = { + "id": "msg_1", + "content": [{"type": "text", "text": "answer"}], + } + out = await logger.async_post_agentic_loop_response_hook( + response=response, plan=plan, kwargs={} + ) + assert out == response + + @pytest.mark.asyncio + async def test_handles_object_style_response(self): + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + block = WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="toolu_obj", + search_response=_make_search_response(), + ) + plan = AgenticLoopPlan( + run_agentic_loop=True, + metadata={WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: [block]}, + ) + + class _Resp: + def __init__(self): + self.content = [{"type": "text", "text": "ok"}] + + resp = _Resp() + out = await logger.async_post_agentic_loop_response_hook( + response=resp, plan=plan, kwargs={} + ) + assert out.content[0]["type"] == "web_search_tool_result" + assert out.content[1]["type"] == "text" + + +class TestShortCircuitEmitsNativeBlocks: + """Standalone /v1/messages sub-requests (Cowork's separate search call) + hit ``try_short_circuit_search``, which builds a synthetic response and + never enters the agentic loop. The native-block emission must happen + here too, otherwise the citations panel stays empty.""" + + @pytest.mark.asyncio + async def test_native_tool_short_circuit_emits_blocks(self): + logger = WebSearchInterceptionLogger(enabled_providers=["github_copilot"]) + + with patch.object( + logger, + "_execute_search", + new=AsyncMock(return_value=("Title: x\nURL: y", _make_search_response())), + ): + result = await logger.try_short_circuit_search( + model="github_copilot/claude-sonnet-4", + messages=[{"role": "user", "content": "search query"}], + tools=[ + { + "type": "web_search_20250305", + "name": "web_search", + "max_uses": 3, + } + ], + custom_llm_provider="github_copilot", + ) + + assert result is not None + block_types = [b["type"] for b in result["content"]] + # Order matters: native clients expect tool_use before tool_result. + assert block_types == ["server_tool_use", "web_search_tool_result", "text"] + server_use, tool_result, _ = result["content"] + assert server_use["name"] == "web_search" + assert server_use["input"] == {"query": "search query"} + # tool_use_id must match between the server_tool_use and the + # web_search_tool_result block so the client can pair them. + assert server_use["id"].startswith("srvtoolu_") + assert tool_result["tool_use_id"] == server_use["id"] + # The actual search results carry through (urls + titles). + assert len(tool_result["content"]) == 2 + assert tool_result["content"][0]["url"] == "https://docs.litellm.ai/" + + @pytest.mark.asyncio + async def test_litellm_standard_tool_short_circuit_stays_text_only(self): + """Non-native tool → existing text-only short-circuit, no regression.""" + logger = WebSearchInterceptionLogger(enabled_providers=["github_copilot"]) + + with patch.object( + logger, + "_execute_search", + new=AsyncMock(return_value=("Title: x\nURL: y", _make_search_response())), + ): + result = await logger.try_short_circuit_search( + model="github_copilot/claude-sonnet-4", + messages=[{"role": "user", "content": "search query"}], + tools=[ + { + "name": "litellm_web_search", + "input_schema": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + } + ], + custom_llm_provider="github_copilot", + ) + + assert result is not None + block_types = [b["type"] for b in result["content"]] + assert block_types == ["text"] + + @pytest.mark.asyncio + async def test_native_short_circuit_failure_still_emits_blocks(self): + """Search failure on native path: emit blocks with empty results + + the legacy text-error block, so the client gets a well-formed + response instead of a malformed half-shape.""" + logger = WebSearchInterceptionLogger(enabled_providers=["github_copilot"]) + + with patch.object(logger, "_execute_search", side_effect=RuntimeError("boom")): + result = await logger.try_short_circuit_search( + 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", + ) + + assert result is not None + block_types = [b["type"] for b in result["content"]] + assert block_types == ["server_tool_use", "web_search_tool_result", "text"] + tool_result = result["content"][1] + assert tool_result["content"] == [] + text_block = result["content"][2] + assert "Search failed" in text_block["text"] + + +class TestLegacyPathMatchesNewPath: + """The legacy ``_execute_agentic_loop`` must inject blocks too.""" + + @pytest.mark.asyncio + async def test_legacy_path_injects_when_flag_set(self): + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + tool_calls = [ + { + "id": "toolu_legacy", + "type": "tool_use", + "name": "litellm_web_search", + "input": {"query": "q"}, + } + ] + patch_obj = AgenticLoopRequestPatch( + model="bedrock/claude", + messages=[{"role": "user", "content": "hi"}], + max_tokens=1024, + optional_params={}, + ) + followup_response = { + "id": "msg_followup", + "content": [{"type": "text", "text": "final answer"}], + } + + with ( + patch.object( + logger, + "_build_anthropic_request_patch", + new=AsyncMock(return_value=(patch_obj, [_make_search_response()])), + ), + patch( + "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate", + new=AsyncMock(return_value=followup_response), + ), + ): + out = await logger._execute_agentic_loop( + model="bedrock/claude", + messages=[], + tool_calls=tool_calls, + thinking_blocks=[], + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(model_call_details={}), + stream=False, + kwargs={WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True}, + ) + + assert out["content"][0]["type"] == "web_search_tool_result" + assert out["content"][0]["tool_use_id"] == "toolu_legacy" + assert out["content"][1]["type"] == "text" 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 82c1c9839e79..7de8892b8fc0 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 @@ -30,7 +30,8 @@ async def test_short_circuits_single_web_search_tool(self): logger, "_execute_search", new_callable=AsyncMock ) as mock_search: mock_search.return_value = ( - "Title: Result\nURL: https://example.com\nSnippet: test" + "Title: Result\nURL: https://example.com\nSnippet: test", + None, ) result = await logger.try_short_circuit_search( @@ -48,9 +49,15 @@ async def test_short_circuits_single_web_search_tool(self): 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 web_search_20250305 client → short-circuit emits native + # blocks (server_tool_use + web_search_tool_result) plus the legacy + # text block so Cowork / Claude Desktop citations panels populate. + block_types = [b["type"] for b in result["content"]] + assert "server_tool_use" in block_types + assert "web_search_tool_result" in block_types + assert "text" in block_types + text_block = next(b for b in result["content"] if b["type"] == "text") + assert "Result" in text_block["text"] mock_search.assert_called_once_with("Search for Claude Code releases") @pytest.mark.asyncio @@ -173,7 +180,8 @@ async def test_search_failure_returns_error_text(self): ) assert result is not None - assert "Search failed" in result["content"][0]["text"] + text_block = next(b for b in result["content"] if b["type"] == "text") + assert "Search failed" in text_block["text"] @pytest.mark.asyncio async def test_response_has_valid_structure(self): @@ -183,7 +191,7 @@ async def test_response_has_valid_structure(self): with patch.object( logger, "_execute_search", new_callable=AsyncMock ) as mock_search: - mock_search.return_value = "search results here" + mock_search.return_value = ("search results here", None) result = await logger.try_short_circuit_search( model="github_copilot/claude-sonnet-4", @@ -246,7 +254,7 @@ async def test_returns_dict_when_not_streaming(self): with patch.object( logger, "_execute_search", new_callable=AsyncMock ) as mock_search: - mock_search.return_value = "results" + mock_search.return_value = ("results", None) with patch("litellm.callbacks", [logger]): result = await _try_websearch_short_circuit( model="github_copilot/claude-sonnet-4", @@ -257,7 +265,8 @@ async def test_returns_dict_when_not_streaming(self): ) assert isinstance(result, dict) - assert result["content"][0]["text"] == "results" + text_block = next(b for b in result["content"] if b["type"] == "text") + assert text_block["text"] == "results" @pytest.mark.asyncio async def test_returns_stream_iterator_when_streaming(self): @@ -273,7 +282,7 @@ async def test_returns_stream_iterator_when_streaming(self): with patch.object( logger, "_execute_search", new_callable=AsyncMock ) as mock_search: - mock_search.return_value = "streaming results" + mock_search.return_value = ("streaming results", None) with patch("litellm.callbacks", [logger]): result = await _try_websearch_short_circuit( model="github_copilot/claude-sonnet-4", @@ -338,7 +347,7 @@ async def test_uses_original_stream_not_hook_converted(self): with patch.object( logger, "_execute_search", new_callable=AsyncMock ) as mock_search: - mock_search.return_value = "streaming results" + mock_search.return_value = ("streaming results", None) 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 @@ -368,7 +377,7 @@ async def test_short_circuits_with_provider_from_model_string(self): with patch.object( logger, "_execute_search", new_callable=AsyncMock ) as mock_search: - mock_search.return_value = "results" + mock_search.return_value = ("results", None) with patch("litellm.callbacks", [logger]): # Simulate the caller having derived custom_llm_provider from # the model string before calling _try_websearch_short_circuit @@ -381,4 +390,5 @@ async def test_short_circuits_with_provider_from_model_string(self): ) assert result is not None - assert result["content"][0]["text"] == "results" + text_block = next(b for b in result["content"] if b["type"] == "text") + assert text_block["text"] == "results" diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_thinking_constraint.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_thinking_constraint.py index a939951c4309..b2d5225070c1 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_thinking_constraint.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_thinking_constraint.py @@ -68,7 +68,9 @@ async def _fake_acreate(**kw): "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate", side_effect=_fake_acreate, ), - patch.object(logger, "_execute_search", return_value="search result"), + patch.object( + logger, "_execute_search", return_value=("search result", None) + ), ): await logger._execute_agentic_loop( @@ -102,7 +104,9 @@ async def _fake_acreate(**kw): "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate", side_effect=_fake_acreate, ), - patch.object(logger, "_execute_search", return_value="search result"), + patch.object( + logger, "_execute_search", return_value=("search result", None) + ), ): await logger._execute_agentic_loop( @@ -136,7 +140,9 @@ async def _fake_acreate(**kw): "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate", side_effect=_fake_acreate, ), - patch.object(logger, "_execute_search", return_value="search result"), + patch.object( + logger, "_execute_search", return_value=("search result", None) + ), ): await logger._execute_agentic_loop( @@ -170,7 +176,9 @@ async def _fake_acreate(**kw): "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate", side_effect=_fake_acreate, ), - patch.object(logger, "_execute_search", return_value="search result"), + patch.object( + logger, "_execute_search", return_value=("search result", None) + ), ): await logger._execute_agentic_loop( @@ -201,7 +209,9 @@ async def _fake_acreate(**kw): "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate", side_effect=_fake_acreate, ), - patch.object(logger, "_execute_search", return_value="search result"), + patch.object( + logger, "_execute_search", return_value=("search result", None) + ), ): await logger._execute_agentic_loop( @@ -286,7 +296,9 @@ async def _fake_acreate(**kw): "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate", side_effect=_fake_acreate, ), - patch.object(logger, "_execute_search", return_value="search result"), + patch.object( + logger, "_execute_search", return_value=("search result", None) + ), ): await logger._execute_agentic_loop( @@ -325,7 +337,9 @@ async def _fake_acreate(**kw): "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate", side_effect=_fake_acreate, ), - patch.object(logger, "_execute_search", return_value="search result"), + patch.object( + logger, "_execute_search", return_value=("search result", None) + ), ): await logger._execute_agentic_loop( @@ -373,7 +387,9 @@ async def _fail_acreate(**kw): "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate", side_effect=_fail_acreate, ), - patch.object(logger, "_execute_search", return_value="search result"), + patch.object( + logger, "_execute_search", return_value=("search result", None) + ), ): with pytest.raises(Exception, match="max_tokens must be greater"): @@ -450,7 +466,9 @@ async def _fake_acreate(**kw): "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate", side_effect=_fake_acreate, ), - patch.object(logger, "_execute_search", return_value="search result"), + patch.object( + logger, "_execute_search", return_value=("search result", None) + ), ): await logger._execute_agentic_loop(