mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
Emit native web_search_tool_result blocks for Anthropic clients (Claude Desktop / Cowork citations) (#27886)
* feat(custom_logger): add async_post_agentic_loop_response_hook Lets a CustomLogger shape the response returned by the agentic-loop follow-up call without bypassing the loop's safety / observability machinery (depth tracking, fingerprinting, etc.). Default returns the response unchanged. Used by websearch_interception to inject Anthropic-native web_search_tool_result blocks when the originating client requested a native web_search_* tool. * feat(llm_http_handler): call post-agentic-loop hook on the originating callback In _execute_anthropic_agentic_plan, after anthropic_messages.acreate returns, call the originating callback's async_post_agentic_loop_response_hook so it can mutate the final response (e.g. inject native tool_result blocks). Pass the callback through from _call_agentic_completion_hooks. Exceptions in the post-hook are caught and logged so a buggy callback can't kill the request. * feat(websearch_interception): add is_anthropic_native_web_search_tool Identifies tools the Anthropic-native clients (Claude Desktop, the Anthropic SDK, the Anthropic Console) use to request native search: type starts with "web_search_" (e.g. web_search_20250305). Rejects the LiteLLM standard tool, the OpenAI-function variant, the bare "WebSearch" legacy name, and the bare "web_search" Claude Code shape. This lets us decide per-request whether the client expects web_search_tool_result content blocks in the response, without renaming any existing constants or touching native-provider skip logic. * feat(websearch_interception): add build_web_search_tool_result_block Produces the Anthropic-native web_search_tool_result content block from a structured SearchResponse. Anthropic-native clients use this block to populate citations / source links — the existing text-blob flatten path only feeds readable evidence to the model and discards the structure, so this builder gives us the missing piece. Shape matches https://docs.anthropic.com/en/api/web-search-tool — web_search_result items carry url, title, page_age, encrypted_content (empty string when the search provider doesn't supply one). * feat(websearch_interception): emit native web_search_tool_result blocks When the originating client request carried a native Anthropic web_search_* tool, the final response now also carries web_search_tool_result content blocks alongside the model's text answer — so Claude Desktop / Anthropic SDK clients can populate the citations panel and replay conversation history with structured search evidence. Wiring: - Pre-request hooks (both deployment + Anthropic path) set a flag on kwargs when they see a native web_search_* tool, so the signal survives the conversion-to-litellm_web_search step regardless of which hook fires first. - _execute_search now returns (text, SearchResponse) so the structured results aren't lost when the text is flattened for the follow-up model call. - _build_anthropic_request_patch returns the parallel list of SearchResponse objects. - async_build_agentic_loop_plan pre-builds the web_search_tool_result blocks (one per tool_use_id) and stashes them on plan.metadata when the flag is set. - async_post_agentic_loop_response_hook reads the metadata and prepends the blocks to response.content. - _execute_agentic_loop mirrors the injection for the legacy path so both paths behave identically. Clients that send the LiteLLM standard tool keep the existing text-only behavior — no regression. * test(websearch_interception): cover native web_search_tool_result emission 18 tests across: - detector branches (native vs litellm-standard, OpenAI-function shape, Claude Desktop builtin WebSearch, bare web_search, missing type) - block-builder shape (results, none, empty) - pre-request hook flag-setting (native sets, standard does not) - async_build_agentic_loop_plan attaches blocks to plan.metadata when the flag is present, leaves metadata untouched when absent - post-hook injection into dict and object responses - legacy _execute_agentic_loop mirrors the injection so both paths return the same shape * test(websearch_short_circuit): keep _execute_search mocks in sync with new tuple return * test(websearch_thinking_constraint): keep _execute_search mocks in sync with new tuple return * feat(websearch_interception): emit native blocks from try_short_circuit_search The agentic-loop post-hook only fires when the model returns a tool_use block. Cowork / Claude Desktop on Bedrock actually make TWO requests per user turn: the main /v1/messages with their builtin tool, and a separate standalone /v1/messages whose only tool is web_search_20250305. That second request hits try_short_circuit_search — no agentic loop, no post-hook — and was returning text-only, leaving the citations panel empty. When the short-circuit input carries a native web_search_* tool, build a synthetic server_tool_use + web_search_tool_result pair (using the structured SearchResponse already returned by _execute_search) so the client gets the native shape it expects. The legacy text block is preserved so non-native short-circuit callers (Claude Code, github_copilot, etc.) see the same payload as before. Failure path still emits the native block pair (with empty results) plus the text-error block, so the client gets a well-formed response rather than a malformed half-shape. * test(websearch_native_blocks): cover short-circuit native-block emission Three new cases on top of the existing 18: - native web_search_20250305 short-circuit → [server_tool_use, web_search_tool_result, text], ids paired, urls/titles carried. - litellm_web_search short-circuit → text-only (no regression). - native short-circuit on search failure → still emits the native block pair (empty results) plus the text-error block, so the client never sees a malformed half-shape. * test(websearch_short_circuit): index assertions by block type, not by position Native short-circuit responses now have [server_tool_use, web_search_tool_result, text] when the input carries web_search_20250305 — find the text block by type rather than relying on content[0]. * fix(websearch_interception): gate legacy WebSearch name on schema absence Clients like Cowork / Claude Desktop ship a client-side tool named "WebSearch" with a full input_schema — they handle it themselves and expect to make a separate native web_search_20250305 sub-request for the actual search. Today is_web_search_tool matches the bare name regardless of other fields, which hijacks the client's tool server-side. The agentic loop fires on the main request, the model never gets to emit the client-side tool_use, and the separate native sub-request (where citation data flows) is never made. Net: citations panel empty. Real Anthropic client tools always carry 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: legacy callers (if any) keep working, real client-side WebSearch tools pass through untouched. * fix(websearch_interception): drop "WebSearch" from response-detection lists Post-conversion the model always sees ``litellm_web_search``, so the "WebSearch" entry in the response-side tool_use detection lists was dead at best. If a model ever did return ``tool_use(name="WebSearch")`` it would now (incorrectly) hijack the client's own ``WebSearch`` tool again — same Cowork problem we just fixed on the input side. Drop it. * test(websearch_native_blocks): cover the WebSearch legacy-name schema gate Three new cases: - {name: "WebSearch"} (bare interception marker) → still matched - {name: "WebSearch", input_schema: {...}} (Cowork client tool) → passes through untouched - {name: "WebSearch", description: "..."} (no schema) → still matched on the assumption it's a legacy marker rather than a malformed real client tool. --------- Co-authored-by: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
This commit is contained in:
ishaan-berri
GitHub
parent
9b6ab55c5f
commit
39e1831e84
8 changed files with 900 additions and 56 deletions
|
|
@ -697,6 +697,27 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
"""
|
||||
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,
|
||||
|
|
|
|||
|
|
@ -19,12 +19,14 @@ from litellm.integrations.custom_logger import CustomLogger
|
|||
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.llms.openai import AllMessageValues
|
|||
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 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
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 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
|
||||
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 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
"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 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
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 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
) -> 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 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
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 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
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 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
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 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
**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 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
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 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
)
|
||||
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 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
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 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
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 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
)
|
||||
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 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
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 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -100,11 +100,14 @@ class WebSearchTransformation:
|
|||
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 @@ class WebSearchTransformation:
|
|||
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 @@ class WebSearchTransformation:
|
|||
|
||||
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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -4634,6 +4634,7 @@ class BaseLLMHTTPHandler:
|
|||
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 @@ class BaseLLMHTTPHandler:
|
|||
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 @@ class BaseLLMHTTPHandler:
|
|||
}
|
||||
)
|
||||
|
||||
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 @@ class BaseLLMHTTPHandler:
|
|||
fingerprints=fingerprints,
|
||||
fingerprint=fingerprint,
|
||||
stream=stream,
|
||||
callback=callback,
|
||||
)
|
||||
except Exception as e:
|
||||
_call_id = getattr(logging_obj, "litellm_call_id", "unknown")
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
@ -30,7 +30,8 @@ class TestTryShortCircuitSearch:
|
|||
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 @@ 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 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 @@ class TestTryShortCircuitSearch:
|
|||
)
|
||||
|
||||
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 @@ 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 = ("search results here", None)
|
||||
|
||||
result = await logger.try_short_circuit_search(
|
||||
model="github_copilot/claude-sonnet-4",
|
||||
|
|
@ -246,7 +254,7 @@ class TestShortCircuitEntryPoint:
|
|||
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 @@ class TestShortCircuitEntryPoint:
|
|||
)
|
||||
|
||||
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 @@ class TestShortCircuitEntryPoint:
|
|||
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 @@ class TestShortCircuitEntryPoint:
|
|||
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 @@ class TestShortCircuitEntryPoint:
|
|||
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 @@ class TestShortCircuitEntryPoint:
|
|||
)
|
||||
|
||||
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"
|
||||
|
|
|
|||
|
|
@ -68,7 +68,9 @@ class TestThinkingBudgetTokensConstraint:
|
|||
"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 @@ class TestThinkingBudgetTokensConstraint:
|
|||
"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 @@ class TestThinkingBudgetTokensConstraint:
|
|||
"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 @@ class TestThinkingBudgetTokensConstraint:
|
|||
"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 @@ class TestThinkingBudgetTokensConstraint:
|
|||
"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 @@ class TestLoggingObjExcludedFromFollowUp:
|
|||
"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 @@ class TestLoggingObjExcludedFromFollowUp:
|
|||
"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 @@ class TestFollowUpErrorScenarios:
|
|||
"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 @@ class TestFollowUpErrorScenarios:
|
|||
"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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue