mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
fix(websearch_interception): surface a failed search as a web_search_tool_result_error block and end the turn
This commit is contained in:
parent
bc6b540205
commit
cdc0e57e93
8 changed files with 492 additions and 93 deletions
|
|
@ -44,6 +44,8 @@ from litellm.types.integrations.custom_logger import (
|
|||
from litellm.types.integrations.websearch_interception import (
|
||||
AnthropicSearchQuery,
|
||||
AnthropicServerToolUseBlock,
|
||||
SearchFailed,
|
||||
SearchOutcome,
|
||||
WebSearchInterceptionConfig,
|
||||
)
|
||||
from litellm.types.llms.anthropic import AnthropicThinkingParam
|
||||
|
|
@ -332,16 +334,8 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
None,
|
||||
)
|
||||
|
||||
# Execute search — keep the structured SearchResponse so the native
|
||||
# block can carry per-result url/title/page_age.
|
||||
try:
|
||||
if kwargs is None:
|
||||
search_result_text, structured = await self._execute_search(query)
|
||||
else:
|
||||
search_result_text, structured = await self._execute_search(query, kwargs=kwargs)
|
||||
except Exception as e:
|
||||
verbose_logger.error("WebSearchInterception: Short-circuit search failed: %s", e)
|
||||
search_result_text, structured = f"Search failed: {e}", None
|
||||
outcome: Final = await self._short_circuit_search_outcome(query, kwargs=kwargs)
|
||||
search_result_text: Final = WebSearchTransformation.search_outcome_text(outcome)
|
||||
|
||||
content: Final[list[dict[str, object]]] = []
|
||||
if native_tool is not None:
|
||||
|
|
@ -355,12 +349,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
"input": {"query": query},
|
||||
}
|
||||
)
|
||||
content.append(
|
||||
WebSearchTransformation.build_web_search_tool_result_block(
|
||||
tool_use_id=tool_use_id,
|
||||
search_response=structured,
|
||||
)
|
||||
)
|
||||
content.append(WebSearchTransformation.build_web_search_outcome_block(tool_use_id=tool_use_id, outcome=outcome))
|
||||
# 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})
|
||||
|
|
@ -934,7 +923,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
|
||||
tool_calls: Final = tools["tool_calls"]
|
||||
thinking_blocks: Final = tools.get("thinking_blocks", [])
|
||||
request_patch, structured_results = await self._build_anthropic_request_patch(
|
||||
request_patch, search_outcomes = await self._build_anthropic_request_patch(
|
||||
model=model,
|
||||
messages=messages,
|
||||
tool_calls=tool_calls,
|
||||
|
|
@ -953,17 +942,19 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
# 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,
|
||||
)
|
||||
if not kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY):
|
||||
return AgenticLoopPlan(run_agentic_loop=True, request_patch=request_patch, metadata=metadata)
|
||||
|
||||
return AgenticLoopPlan(
|
||||
run_agentic_loop=True,
|
||||
request_patch=request_patch,
|
||||
metadata=metadata,
|
||||
metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY] = self._build_native_result_blocks(
|
||||
tool_calls=tool_calls,
|
||||
search_outcomes=search_outcomes,
|
||||
)
|
||||
every_search_failed: Final = bool(search_outcomes) and all(
|
||||
isinstance(outcome, SearchFailed) for outcome in search_outcomes
|
||||
)
|
||||
if every_search_failed:
|
||||
return AgenticLoopPlan(run_agentic_loop=False, terminate=True, stop_reason="web_search_failed", metadata=metadata)
|
||||
return AgenticLoopPlan(run_agentic_loop=True, request_patch=request_patch, metadata=metadata)
|
||||
|
||||
async def async_post_agentic_loop_response_hook(
|
||||
self,
|
||||
|
|
@ -992,7 +983,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
@staticmethod
|
||||
def _build_native_result_blocks(
|
||||
tool_calls: list[dict],
|
||||
structured_results: list[SearchResponse | None],
|
||||
search_outcomes: Sequence[SearchOutcome],
|
||||
) -> tuple[Mapping[str, object], ...]:
|
||||
"""
|
||||
Build a ``server_tool_use`` + ``web_search_tool_result`` pair per tool_call.
|
||||
|
|
@ -1004,10 +995,10 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
"""
|
||||
return tuple(
|
||||
block
|
||||
for i, tool_call in enumerate(tool_calls)
|
||||
for tool_call, outcome in zip(tool_calls, search_outcomes, strict=True)
|
||||
for block in WebSearchInterceptionLogger._native_result_pair(
|
||||
query=WebSearchInterceptionLogger._tool_call_query(tool_call),
|
||||
search_response=structured_results[i] if i < len(structured_results) else None,
|
||||
outcome=outcome,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -1022,15 +1013,12 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
@staticmethod
|
||||
def _native_result_pair(
|
||||
query: str,
|
||||
search_response: SearchResponse | None,
|
||||
outcome: SearchOutcome,
|
||||
) -> tuple[Mapping[str, object], Mapping[str, object]]:
|
||||
tool_use_id: Final = f"srvtoolu_{uuid.uuid4().hex}"
|
||||
return (
|
||||
AnthropicServerToolUseBlock(id=tool_use_id, input=AnthropicSearchQuery(query=query)).model_dump(),
|
||||
WebSearchTransformation.build_web_search_tool_result_block(
|
||||
tool_use_id=tool_use_id,
|
||||
search_response=search_response,
|
||||
),
|
||||
WebSearchTransformation.build_web_search_outcome_block(tool_use_id=tool_use_id, outcome=outcome),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -1306,7 +1294,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
kwargs: Mapping[str, object],
|
||||
) -> "AnthropicMessagesResponse | AsyncIterator[object]":
|
||||
"""Legacy path: execute search + build patch + run follow-up call."""
|
||||
request_patch, structured_results = await self._build_anthropic_request_patch(
|
||||
request_patch, search_outcomes = await self._build_anthropic_request_patch(
|
||||
model=model,
|
||||
messages=messages,
|
||||
tool_calls=tool_calls,
|
||||
|
|
@ -1344,7 +1332,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
if kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY):
|
||||
native_blocks: Final = self._build_native_result_blocks(
|
||||
tool_calls=tool_calls,
|
||||
structured_results=structured_results,
|
||||
search_outcomes=search_outcomes,
|
||||
)
|
||||
response = self._inject_native_blocks(response, native_blocks)
|
||||
|
||||
|
|
@ -1359,15 +1347,14 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
anthropic_messages_optional_request_params: dict,
|
||||
logging_obj: "LiteLLMLoggingObj | None",
|
||||
kwargs: dict,
|
||||
) -> tuple[AgenticLoopRequestPatch, list[SearchResponse | None]]:
|
||||
) -> tuple[AgenticLoopRequestPatch, tuple[SearchOutcome, ...]]:
|
||||
"""
|
||||
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.
|
||||
Returns the patch alongside the parallel tuple of search outcomes (one
|
||||
per tool_call). The caller uses these to optionally build
|
||||
Anthropic-native ``web_search_tool_result`` content blocks for the
|
||||
final response and to decide whether a follow-up call is worth making.
|
||||
"""
|
||||
|
||||
# Extract search queries from tool_use blocks
|
||||
|
|
@ -1385,27 +1372,10 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
# Execute searches in parallel
|
||||
verbose_logger.debug("WebSearchInterception: Executing %s search(es) in parallel", len(search_tasks))
|
||||
search_results: Final = await asyncio.gather(*search_tasks, return_exceptions=True)
|
||||
|
||||
# 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: Final[list[str]] = []
|
||||
structured_results: Final[list[SearchResponse | None]] = []
|
||||
for i, result in enumerate(search_results):
|
||||
if isinstance(result, Exception):
|
||||
verbose_logger.error("WebSearchInterception: Search %s failed with error: %s", i, result)
|
||||
final_search_results.append(f"Search failed: {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:
|
||||
# Defensive: legacy callers / unexpected shape — preserve text,
|
||||
# drop structure.
|
||||
verbose_logger.debug("WebSearchInterception: Unexpected result type %s at index %s", type(result), i)
|
||||
final_search_results.append(str(result))
|
||||
structured_results.append(None)
|
||||
search_outcomes: Final = tuple(WebSearchTransformation.search_outcome(result) for result in search_results)
|
||||
final_search_results: Final = tuple(
|
||||
WebSearchTransformation.search_outcome_text(outcome) for outcome in search_outcomes
|
||||
)
|
||||
|
||||
# Build assistant and user messages using transformation
|
||||
assistant_message, user_message = WebSearchTransformation.transform_response(
|
||||
|
|
@ -1449,7 +1419,16 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
optional_params=optional_params_without_max_tokens,
|
||||
kwargs=kwargs_for_followup,
|
||||
)
|
||||
return patch, structured_results
|
||||
return patch, search_outcomes
|
||||
|
||||
async def _short_circuit_search_outcome(self, query: str, kwargs: Mapping[str, object] | None) -> SearchOutcome:
|
||||
try:
|
||||
result: Final = (
|
||||
await self._execute_search(query) if kwargs is None else await self._execute_search(query, kwargs=kwargs)
|
||||
)
|
||||
except Exception as e:
|
||||
return WebSearchTransformation.search_outcome(e)
|
||||
return WebSearchTransformation.search_outcome(result)
|
||||
|
||||
async def _execute_search(
|
||||
self, query: str, kwargs: Mapping[str, object] | None = None
|
||||
|
|
|
|||
|
|
@ -5,11 +5,21 @@ Transforms between Anthropic/OpenAI tool_use format and LiteLLM search format.
|
|||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, Final
|
||||
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import LITELLM_WEB_SEARCH_TOOL_NAME
|
||||
from litellm.exceptions import BadRequestError, RateLimitError
|
||||
from litellm.llms.base_llm.search.transformation import SearchResponse
|
||||
from litellm.types.integrations.websearch_interception import (
|
||||
SearchFailed,
|
||||
SearchOutcome,
|
||||
SearchSucceeded,
|
||||
WebSearchToolResultErrorCode,
|
||||
)
|
||||
|
||||
|
||||
class WebSearchTransformation:
|
||||
|
|
@ -280,7 +290,7 @@ class WebSearchTransformation:
|
|||
@staticmethod
|
||||
def transform_response(
|
||||
tool_calls: list[dict],
|
||||
search_results: list[str],
|
||||
search_results: Sequence[str],
|
||||
response_format: str = "anthropic",
|
||||
thinking_blocks: list[dict] | None = None,
|
||||
) -> tuple[dict, dict | list[dict]]:
|
||||
|
|
@ -314,7 +324,7 @@ class WebSearchTransformation:
|
|||
@staticmethod
|
||||
def _transform_response_anthropic(
|
||||
tool_calls: list[dict],
|
||||
search_results: list[str],
|
||||
search_results: Sequence[str],
|
||||
thinking_blocks: list[dict] | None = None,
|
||||
) -> tuple[dict, dict]:
|
||||
"""Transform to Anthropic format (single user message with tool_result blocks)"""
|
||||
|
|
@ -364,7 +374,7 @@ class WebSearchTransformation:
|
|||
@staticmethod
|
||||
def _transform_response_openai(
|
||||
tool_calls: list[dict],
|
||||
search_results: list[str],
|
||||
search_results: Sequence[str],
|
||||
) -> tuple[dict, list[dict]]:
|
||||
"""Transform to OpenAI format (assistant with tool_calls, separate tool messages)"""
|
||||
# Build assistant message with tool_calls
|
||||
|
|
@ -456,6 +466,67 @@ class WebSearchTransformation:
|
|||
"content": items,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def build_web_search_tool_result_error_block(
|
||||
tool_use_id: str,
|
||||
error_code: WebSearchToolResultErrorCode,
|
||||
) -> dict[str, object]:
|
||||
return {
|
||||
"type": "web_search_tool_result",
|
||||
"tool_use_id": tool_use_id,
|
||||
"content": {"type": "web_search_tool_result_error", "error_code": error_code},
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def build_web_search_outcome_block(tool_use_id: str, outcome: SearchOutcome) -> dict[str, object]:
|
||||
match outcome:
|
||||
case SearchSucceeded(response=response):
|
||||
return WebSearchTransformation.build_web_search_tool_result_block(
|
||||
tool_use_id=tool_use_id,
|
||||
search_response=response,
|
||||
)
|
||||
case SearchFailed(error_code=error_code):
|
||||
return WebSearchTransformation.build_web_search_tool_result_error_block(
|
||||
tool_use_id=tool_use_id,
|
||||
error_code=error_code,
|
||||
)
|
||||
case _:
|
||||
assert_never(outcome)
|
||||
|
||||
@staticmethod
|
||||
def search_error_code(error: BaseException) -> WebSearchToolResultErrorCode:
|
||||
match error:
|
||||
case RateLimitError():
|
||||
return "too_many_requests"
|
||||
case BadRequestError():
|
||||
return "invalid_tool_input"
|
||||
case _:
|
||||
return "unavailable"
|
||||
|
||||
@staticmethod
|
||||
def search_outcome(result: object) -> SearchOutcome:
|
||||
match result:
|
||||
case BaseException():
|
||||
verbose_logger.error("WebSearchInterception: Search failed with error: %s", result)
|
||||
return SearchFailed(error_code=WebSearchTransformation.search_error_code(result), message=str(result))
|
||||
case (str() as text, SearchResponse() as response):
|
||||
return SearchSucceeded(text=text, response=response)
|
||||
case (str() as text, None):
|
||||
return SearchSucceeded(text=text, response=None)
|
||||
case _:
|
||||
verbose_logger.debug("WebSearchInterception: Unexpected search result type %s", type(result))
|
||||
return SearchSucceeded(text=str(result), response=None)
|
||||
|
||||
@staticmethod
|
||||
def search_outcome_text(outcome: SearchOutcome) -> str:
|
||||
match outcome:
|
||||
case SearchSucceeded(text=text):
|
||||
return text
|
||||
case SearchFailed(message=message):
|
||||
return f"Search failed: {message}"
|
||||
case _:
|
||||
assert_never(outcome)
|
||||
|
||||
@staticmethod
|
||||
def format_search_response(result: SearchResponse) -> str:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1410,12 +1410,19 @@ class _ReplayedWebSearchResult(BaseModel):
|
|||
encrypted_content: str = ""
|
||||
|
||||
|
||||
class _ReplayedWebSearchToolResultError(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
type: Literal["web_search_tool_result_error"]
|
||||
error_code: str = ""
|
||||
|
||||
|
||||
class _ReplayedWebSearchToolResult(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
type: Literal["web_search_tool_result"]
|
||||
tool_use_id: str
|
||||
content: tuple[_ReplayedWebSearchResult, ...]
|
||||
content: tuple[_ReplayedWebSearchResult, ...] | _ReplayedWebSearchToolResultError
|
||||
|
||||
|
||||
class _ReplayedServerToolUse(BaseModel):
|
||||
|
|
@ -1441,15 +1448,17 @@ def _flattenable_web_search_tool_result(block: object) -> _ReplayedWebSearchTool
|
|||
``encrypted_content``, else None for anything Anthropic itself issued.
|
||||
|
||||
An empty ``content`` list is flattenable too. It is what the interceptor emits
|
||||
when a search legitimately returns nothing and when a search raises, and it
|
||||
carries neither evidence to preserve nor an ``encrypted_content`` to respect,
|
||||
so leaving it in place only buys the 400 this whole function exists to avoid.
|
||||
when a search legitimately returns nothing, and it carries neither evidence to
|
||||
preserve nor an ``encrypted_content`` to respect, so leaving it in place only
|
||||
buys the 400 this whole function exists to avoid. The same goes for the
|
||||
``web_search_tool_result_error`` object the interceptor emits when a search
|
||||
raises: it never carries ``encrypted_content``, so it is flattened as well.
|
||||
"""
|
||||
try:
|
||||
parsed: Final = _WEB_SEARCH_TOOL_RESULT_ADAPTER.validate_python(block)
|
||||
except ValidationError:
|
||||
return None
|
||||
if any(result.encrypted_content for result in parsed.content):
|
||||
if isinstance(parsed.content, tuple) and any(result.encrypted_content for result in parsed.content):
|
||||
return None
|
||||
return parsed
|
||||
|
||||
|
|
@ -1461,8 +1470,12 @@ def _replayed_server_tool_use(block: object) -> _ReplayedServerToolUse | None:
|
|||
return None
|
||||
|
||||
|
||||
def _render_web_search_results(query: str, results: tuple[_ReplayedWebSearchResult, ...]) -> str:
|
||||
def _render_web_search_results(
|
||||
query: str, results: tuple[_ReplayedWebSearchResult, ...] | _ReplayedWebSearchToolResultError
|
||||
) -> str:
|
||||
header: Final = f"Web search results for '{query}':" if query else "Web search results:"
|
||||
if isinstance(results, _ReplayedWebSearchToolResultError):
|
||||
return f"{header}\n\nSearch failed: {results.error_code or 'unavailable'}"
|
||||
if not results:
|
||||
return f"{header}\n\nNo results were returned."
|
||||
body: Final = "\n\n".join(
|
||||
|
|
|
|||
|
|
@ -5905,7 +5905,15 @@ class BaseLLMHTTPHandler:
|
|||
callback.__class__.__name__,
|
||||
plan.stop_reason,
|
||||
)
|
||||
return self._maybe_wrap_in_fake_stream(response, logging_obj, api_surface)
|
||||
return self._maybe_wrap_in_fake_stream(
|
||||
await callback.async_post_agentic_loop_response_hook(
|
||||
response=self._finalize_refused_agentic_response(response=response, tool_calls=tool_calls),
|
||||
plan=plan,
|
||||
kwargs=kwargs_with_provider,
|
||||
),
|
||||
logging_obj,
|
||||
api_surface,
|
||||
)
|
||||
if not plan.run_agentic_loop:
|
||||
continue
|
||||
|
||||
|
|
|
|||
|
|
@ -2,11 +2,15 @@
|
|||
Type definitions for WebSearch Interception integration.
|
||||
"""
|
||||
|
||||
from typing import Literal, TypedDict
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Literal, TypeAlias, TypedDict
|
||||
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.llms.base_llm.search.transformation import SearchResponse
|
||||
|
||||
|
||||
class AnthropicSearchQuery(BaseModel):
|
||||
"""``input`` of an Anthropic ``server_tool_use`` block for a web search."""
|
||||
|
|
@ -27,6 +31,31 @@ class AnthropicServerToolUseBlock(BaseModel):
|
|||
input: AnthropicSearchQuery
|
||||
|
||||
|
||||
WebSearchToolResultErrorCode: TypeAlias = Literal[
|
||||
"invalid_tool_input",
|
||||
"unavailable",
|
||||
"max_uses_exceeded",
|
||||
"too_many_requests",
|
||||
"query_too_long",
|
||||
"request_too_large",
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SearchSucceeded:
|
||||
text: str
|
||||
response: "SearchResponse | None"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SearchFailed:
|
||||
error_code: WebSearchToolResultErrorCode
|
||||
message: str
|
||||
|
||||
|
||||
SearchOutcome: TypeAlias = SearchSucceeded | SearchFailed
|
||||
|
||||
|
||||
class WebSearchInterceptionConfig(TypedDict, total=False):
|
||||
"""
|
||||
Configuration parameters for WebSearchInterceptionLogger.
|
||||
|
|
|
|||
|
|
@ -12,19 +12,23 @@ config.yaml through to the settings the loop actually reads.
|
|||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import AuthenticationError, RateLimitError
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.websearch_interception.handler import (
|
||||
WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY,
|
||||
WebSearchInterceptionLogger,
|
||||
)
|
||||
from litellm.integrations.websearch_interception.tools import get_litellm_web_search_tool
|
||||
from litellm.litellm_core_utils.agentic_loop_settings import DEFAULT_MAX_AGENTIC_LOOPS
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
|
||||
FakeAnthropicMessagesStreamIterator,
|
||||
)
|
||||
from litellm.litellm_core_utils.agentic_loop_settings import DEFAULT_MAX_AGENTIC_LOOPS
|
||||
from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
|
||||
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
|
||||
from litellm.secret_managers.main import get_secret
|
||||
from litellm.types.integrations.custom_logger import (
|
||||
|
|
@ -490,6 +494,135 @@ class TestOuterFramePostHookStillRuns:
|
|||
assert result["stop_reason"] == "end_turn"
|
||||
|
||||
|
||||
def _response_asking_for_searches(*queries: str) -> dict:
|
||||
return {
|
||||
"id": "msg_123",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [
|
||||
{"id": f"toolu_internal_{index}", "type": "tool_use", "name": INTERNAL_TOOL_NAME, "input": {"query": query}}
|
||||
for index, query in enumerate(queries, start=1)
|
||||
],
|
||||
"stop_reason": "tool_use",
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5},
|
||||
}
|
||||
|
||||
|
||||
class TestFailedSearchEndsTheTurn:
|
||||
"""
|
||||
A search that failed used to come back to the client as an empty successful
|
||||
``web_search_tool_result`` while the model was re-asked the same query until
|
||||
the loop cap tripped. When the client sent a native web search tool, the
|
||||
turn now ends after the first failed search, with Anthropic's
|
||||
``web_search_tool_result_error`` object in the tool result and no follow-up
|
||||
model call. An iteration where some search still succeeded keeps its
|
||||
follow-up call.
|
||||
"""
|
||||
|
||||
def setup_method(self):
|
||||
self.handler = BaseLLMHTTPHandler()
|
||||
self.logger = WebSearchInterceptionLogger(enabled_providers=["anthropic"])
|
||||
self.followup_calls: list[dict] = []
|
||||
|
||||
async def _fake_acreate(self, **call_kwargs):
|
||||
self.followup_calls.append(call_kwargs)
|
||||
return {
|
||||
"id": "msg_followup",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [{"type": "text", "text": "final answer"}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 20, "output_tokens": 5},
|
||||
}
|
||||
|
||||
async def _run(self, response: dict, converted_stream: bool = False):
|
||||
return await self.handler._call_agentic_completion_hooks(
|
||||
response=response,
|
||||
model="claude-sonnet-4-5",
|
||||
messages=[{"role": "user", "content": "who won the world cup"}],
|
||||
anthropic_messages_provider_config=MagicMock(),
|
||||
anthropic_messages_optional_request_params={"tools": [get_litellm_web_search_tool()]},
|
||||
logging_obj=_logging_obj(self.logger, converted_stream=converted_stream),
|
||||
stream=False,
|
||||
custom_llm_provider="anthropic",
|
||||
kwargs={"_agentic_loop_depth": 0, "max_agentic_loops": 3, WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True},
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_all_failed_iteration_ends_the_turn_without_a_follow_up_call(self, monkeypatch):
|
||||
monkeypatch.setattr("litellm.anthropic_interface.messages.acreate", self._fake_acreate)
|
||||
|
||||
with patch.object(
|
||||
self.logger,
|
||||
"_execute_search",
|
||||
side_effect=AuthenticationError("401 Unauthorized", llm_provider="tavily", model="tavily"),
|
||||
):
|
||||
result = await self._run(_response_asking_for_searches("who won the world cup"))
|
||||
|
||||
assert self.followup_calls == []
|
||||
assert result["stop_reason"] == "end_turn"
|
||||
assert INTERNAL_TOOL_NAME not in _tool_use_names(result)
|
||||
assert _block_types(result) == ["server_tool_use", "web_search_tool_result"]
|
||||
server_tool_use, tool_result = result["content"]
|
||||
assert server_tool_use["id"].startswith("srvtoolu_")
|
||||
assert server_tool_use["input"] == {"query": "who won the world cup"}
|
||||
assert tool_result["tool_use_id"] == server_tool_use["id"]
|
||||
assert tool_result["content"] == {"type": "web_search_tool_result_error", "error_code": "unavailable"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_all_failed_iteration_streams_the_error_block(self, monkeypatch):
|
||||
monkeypatch.setattr("litellm.anthropic_interface.messages.acreate", self._fake_acreate)
|
||||
|
||||
with patch.object(
|
||||
self.logger,
|
||||
"_execute_search",
|
||||
side_effect=AuthenticationError("401 Unauthorized", llm_provider="tavily", model="tavily"),
|
||||
):
|
||||
result = await self._run(_response_asking_for_searches("who won the world cup"), converted_stream=True)
|
||||
|
||||
assert self.followup_calls == []
|
||||
assert isinstance(result, FakeAnthropicMessagesStreamIterator)
|
||||
events = _stream_events(result.response)
|
||||
started = [event["content_block"] for event in events if event["type"] == "content_block_start"]
|
||||
assert [block["type"] for block in started] == ["server_tool_use", "web_search_tool_result"]
|
||||
assert started[1]["tool_use_id"] == started[0]["id"]
|
||||
assert started[1]["content"] == {"type": "web_search_tool_result_error", "error_code": "unavailable"}
|
||||
assert [event["delta"]["stop_reason"] for event in events if event["type"] == "message_delta"] == ["end_turn"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mixed_iteration_keeps_the_follow_up_call(self, monkeypatch):
|
||||
monkeypatch.setattr("litellm.anthropic_interface.messages.acreate", self._fake_acreate)
|
||||
|
||||
async def search(query, kwargs=None):
|
||||
if query == "fails":
|
||||
raise RateLimitError("slow down", llm_provider="tavily", model="tavily")
|
||||
found = SearchResult(title="Result", url="https://example.com", snippet="A result.", date=None)
|
||||
return ("Title: Result\nURL: https://example.com", SearchResponse(results=[found]))
|
||||
|
||||
with patch.object(self.logger, "_execute_search", side_effect=search):
|
||||
result = await self._run(_response_asking_for_searches("fails", "works"))
|
||||
|
||||
assert len(self.followup_calls) == 1
|
||||
tool_results = self.followup_calls[0]["messages"][-1]["content"]
|
||||
assert [block["type"] for block in tool_results] == ["tool_result", "tool_result"]
|
||||
assert tool_results[0]["content"] == "Search failed: litellm.RateLimitError: slow down"
|
||||
assert tool_results[1]["content"] == "Title: Result\nURL: https://example.com"
|
||||
assert result["stop_reason"] == "end_turn"
|
||||
assert _block_types(result) == [
|
||||
"server_tool_use",
|
||||
"web_search_tool_result",
|
||||
"server_tool_use",
|
||||
"web_search_tool_result",
|
||||
"text",
|
||||
]
|
||||
assert result["content"][0]["input"] == {"query": "fails"}
|
||||
assert result["content"][1]["content"] == {"type": "web_search_tool_result_error", "error_code": "too_many_requests"}
|
||||
assert result["content"][2]["input"] == {"query": "works"}
|
||||
assert result["content"][3]["content"][0]["url"] == "https://example.com"
|
||||
|
||||
|
||||
class TestMaxAgenticLoopsConfigKnob:
|
||||
def test_from_config_yaml_reads_the_knob(self):
|
||||
logger = WebSearchInterceptionLogger.from_config_yaml(
|
||||
|
|
|
|||
|
|
@ -10,6 +10,13 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
|
||||
import pytest
|
||||
|
||||
from litellm.exceptions import (
|
||||
APIConnectionError,
|
||||
AuthenticationError,
|
||||
BadRequestError,
|
||||
RateLimitError,
|
||||
Timeout,
|
||||
)
|
||||
from litellm.integrations.websearch_interception.handler import (
|
||||
WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY,
|
||||
WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY,
|
||||
|
|
@ -27,6 +34,10 @@ from litellm.types.integrations.custom_logger import (
|
|||
AgenticLoopPlan,
|
||||
AgenticLoopRequestPatch,
|
||||
)
|
||||
from litellm.types.integrations.websearch_interception import (
|
||||
SearchFailed,
|
||||
SearchSucceeded,
|
||||
)
|
||||
|
||||
|
||||
def _make_search_response() -> SearchResponse:
|
||||
|
|
@ -48,6 +59,10 @@ def _make_search_response() -> SearchResponse:
|
|||
)
|
||||
|
||||
|
||||
def _succeeded_outcome() -> SearchSucceeded:
|
||||
return SearchSucceeded(text="Title: LiteLLM Docs\nURL: https://docs.litellm.ai/", response=_make_search_response())
|
||||
|
||||
|
||||
class TestIsAnthropicNativeWebSearchTool:
|
||||
"""The detector must match native tools without catching look-alikes."""
|
||||
|
||||
|
|
@ -227,12 +242,10 @@ class TestBuildPlanAttachesBlocks:
|
|||
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)),
|
||||
new=AsyncMock(return_value=(patch_obj, (_succeeded_outcome(),))),
|
||||
):
|
||||
plan = await logger.async_build_agentic_loop_plan(
|
||||
tools={"tool_calls": tool_calls, "thinking_blocks": []},
|
||||
|
|
@ -277,7 +290,7 @@ class TestBuildPlanAttachesBlocks:
|
|||
with patch.object(
|
||||
logger,
|
||||
"_build_anthropic_request_patch",
|
||||
new=AsyncMock(return_value=(patch_obj, [_make_search_response()])),
|
||||
new=AsyncMock(return_value=(patch_obj, (_succeeded_outcome(),))),
|
||||
):
|
||||
plan = await logger.async_build_agentic_loop_plan(
|
||||
tools={"tool_calls": tool_calls, "thinking_blocks": []},
|
||||
|
|
@ -294,6 +307,145 @@ class TestBuildPlanAttachesBlocks:
|
|||
assert WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY not in plan.metadata
|
||||
|
||||
|
||||
class TestFailedSearchOutcome:
|
||||
"""A search that raises becomes a ``web_search_tool_result_error`` block, coded by exception type."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("error", "expected_code"),
|
||||
[
|
||||
(RateLimitError("slow down", llm_provider="tavily", model="tavily"), "too_many_requests"),
|
||||
(BadRequestError("bad query", model="tavily", llm_provider="tavily"), "invalid_tool_input"),
|
||||
(AuthenticationError("401 Unauthorized", llm_provider="tavily", model="tavily"), "unavailable"),
|
||||
(APIConnectionError("connection refused", llm_provider="tavily", model="tavily"), "unavailable"),
|
||||
(Timeout("timed out", model="tavily", llm_provider="tavily"), "unavailable"),
|
||||
(RuntimeError("boom"), "unavailable"),
|
||||
],
|
||||
)
|
||||
def test_error_block_carries_the_mapped_error_code(self, error, expected_code):
|
||||
outcome = WebSearchTransformation.search_outcome(error)
|
||||
|
||||
assert outcome == SearchFailed(error_code=expected_code, message=str(error))
|
||||
assert WebSearchTransformation.build_web_search_outcome_block("srvtoolu_x", outcome) == {
|
||||
"type": "web_search_tool_result",
|
||||
"tool_use_id": "srvtoolu_x",
|
||||
"content": {"type": "web_search_tool_result_error", "error_code": expected_code},
|
||||
}
|
||||
assert WebSearchTransformation.search_outcome_text(outcome) == f"Search failed: {error}"
|
||||
|
||||
def test_succeeded_outcome_still_yields_result_items(self):
|
||||
outcome = WebSearchTransformation.search_outcome(("Title: x", _make_search_response()))
|
||||
|
||||
assert outcome == SearchSucceeded(text="Title: x", response=_make_search_response())
|
||||
block = WebSearchTransformation.build_web_search_outcome_block("srvtoolu_x", outcome)
|
||||
assert [item["type"] for item in block["content"]] == ["web_search_result", "web_search_result"]
|
||||
assert block["content"][0]["url"] == "https://docs.litellm.ai/"
|
||||
assert WebSearchTransformation.search_outcome_text(outcome) == "Title: x"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_all_failed_iteration_terminates_when_native_blocks_are_emitted(self):
|
||||
logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"])
|
||||
tool_calls = [
|
||||
{"id": "toolu_one", "type": "tool_use", "name": "litellm_web_search", "input": {"query": "q1"}},
|
||||
{"id": "toolu_two", "type": "tool_use", "name": "litellm_web_search", "input": {"query": "q2"}},
|
||||
]
|
||||
|
||||
with patch.object(
|
||||
logger,
|
||||
"_execute_search",
|
||||
side_effect=AuthenticationError("401 Unauthorized", llm_provider="tavily", model="tavily"),
|
||||
):
|
||||
plan = await logger.async_build_agentic_loop_plan(
|
||||
tools={"tool_calls": tool_calls, "thinking_blocks": []},
|
||||
model="bedrock/claude",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
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},
|
||||
)
|
||||
|
||||
assert plan.run_agentic_loop is False
|
||||
assert plan.terminate is True
|
||||
assert plan.stop_reason == "web_search_failed"
|
||||
blocks = plan.metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY]
|
||||
assert [b["type"] for b in blocks] == [
|
||||
"server_tool_use",
|
||||
"web_search_tool_result",
|
||||
"server_tool_use",
|
||||
"web_search_tool_result",
|
||||
]
|
||||
assert blocks[1]["tool_use_id"] == blocks[0]["id"]
|
||||
assert blocks[1]["content"] == {"type": "web_search_tool_result_error", "error_code": "unavailable"}
|
||||
assert blocks[3]["tool_use_id"] == blocks[2]["id"]
|
||||
assert blocks[3]["content"] == {"type": "web_search_tool_result_error", "error_code": "unavailable"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_all_failed_iteration_keeps_the_follow_up_without_native_blocks(self):
|
||||
logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"])
|
||||
tool_calls = [
|
||||
{"id": "toolu_one", "type": "tool_use", "name": "litellm_web_search", "input": {"query": "q1"}},
|
||||
]
|
||||
|
||||
with patch.object(
|
||||
logger,
|
||||
"_execute_search",
|
||||
side_effect=AuthenticationError("401 Unauthorized", llm_provider="tavily", model="tavily"),
|
||||
):
|
||||
plan = await logger.async_build_agentic_loop_plan(
|
||||
tools={"tool_calls": tool_calls, "thinking_blocks": []},
|
||||
model="bedrock/claude",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
response=MagicMock(),
|
||||
anthropic_messages_provider_config=None,
|
||||
anthropic_messages_optional_request_params={},
|
||||
logging_obj=MagicMock(model_call_details={}),
|
||||
stream=False,
|
||||
kwargs={},
|
||||
)
|
||||
|
||||
assert plan.run_agentic_loop is True
|
||||
assert plan.terminate is False
|
||||
assert plan.request_patch is not None
|
||||
tool_results = plan.request_patch.messages[-1]["content"]
|
||||
assert "Search failed: litellm.AuthenticationError: 401 Unauthorized" in tool_results[0]["content"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mixed_iteration_keeps_the_follow_up_and_pairs_each_block(self):
|
||||
logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"])
|
||||
tool_calls = [
|
||||
{"id": "toolu_one", "type": "tool_use", "name": "litellm_web_search", "input": {"query": "fails"}},
|
||||
{"id": "toolu_two", "type": "tool_use", "name": "litellm_web_search", "input": {"query": "works"}},
|
||||
]
|
||||
|
||||
async def search(query, kwargs=None):
|
||||
if query == "fails":
|
||||
raise RateLimitError("slow down", llm_provider="tavily", model="tavily")
|
||||
return ("Title: x", _make_search_response())
|
||||
|
||||
with patch.object(logger, "_execute_search", side_effect=search):
|
||||
plan = await logger.async_build_agentic_loop_plan(
|
||||
tools={"tool_calls": tool_calls, "thinking_blocks": []},
|
||||
model="bedrock/claude",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
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},
|
||||
)
|
||||
|
||||
assert plan.run_agentic_loop is True
|
||||
assert plan.terminate is False
|
||||
blocks = plan.metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY]
|
||||
assert blocks[0]["input"] == {"query": "fails"}
|
||||
assert blocks[1]["content"] == {"type": "web_search_tool_result_error", "error_code": "too_many_requests"}
|
||||
assert blocks[2]["input"] == {"query": "works"}
|
||||
assert blocks[3]["content"][0]["url"] == "https://docs.litellm.ai/"
|
||||
|
||||
|
||||
class TestPostHookInjectsBlocks:
|
||||
"""The post-hook must prepend blocks; absent metadata is a no-op."""
|
||||
|
||||
|
|
@ -437,13 +589,17 @@ class TestShortCircuitEmitsNativeBlocks:
|
|||
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."""
|
||||
async def test_native_short_circuit_failure_emits_the_error_block(self):
|
||||
"""Search failure on native path: the tool result carries Anthropic's
|
||||
error object (rendered as "Web search error: <code>" by the client)
|
||||
next to the legacy text-error block."""
|
||||
logger = WebSearchInterceptionLogger(enabled_providers=["github_copilot"])
|
||||
|
||||
with patch.object(logger, "_execute_search", side_effect=RuntimeError("boom")):
|
||||
with patch.object(
|
||||
logger,
|
||||
"_execute_search",
|
||||
side_effect=RateLimitError("slow down", llm_provider="tavily", model="tavily"),
|
||||
):
|
||||
result = await logger.try_short_circuit_search(
|
||||
model="github_copilot/claude-sonnet-4",
|
||||
messages=[{"role": "user", "content": "search query"}],
|
||||
|
|
@ -455,9 +611,10 @@ class TestShortCircuitEmitsNativeBlocks:
|
|||
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"] == []
|
||||
assert tool_result["tool_use_id"] == result["content"][0]["id"]
|
||||
assert tool_result["content"] == {"type": "web_search_tool_result_error", "error_code": "too_many_requests"}
|
||||
text_block = result["content"][2]
|
||||
assert "Search failed" in text_block["text"]
|
||||
assert text_block["text"] == "Search failed: litellm.RateLimitError: slow down"
|
||||
|
||||
|
||||
class TestLegacyPathMatchesNewPath:
|
||||
|
|
@ -489,7 +646,7 @@ class TestLegacyPathMatchesNewPath:
|
|||
patch.object(
|
||||
logger,
|
||||
"_build_anthropic_request_patch",
|
||||
new=AsyncMock(return_value=(patch_obj, [_make_search_response()])),
|
||||
new=AsyncMock(return_value=(patch_obj, (_succeeded_outcome(),))),
|
||||
),
|
||||
patch(
|
||||
"litellm.integrations.websearch_interception.handler.anthropic_messages.acreate",
|
||||
|
|
|
|||
|
|
@ -1828,7 +1828,11 @@ class TestAnthropicThinkingSignatureSelfHeal:
|
|||
|
||||
assert out[0] is msgs[0]
|
||||
|
||||
def test_flatten_unencrypted_web_search_results_leaves_error_blocks_alone(self):
|
||||
def test_flatten_unencrypted_web_search_results_flattens_error_blocks(self):
|
||||
"""A failed intercepted search is replayed by the client as the error
|
||||
object LiteLLM emitted. Anthropic rejects a replayed ``server_tool_use``
|
||||
it never issued, so the pair is flattened to text the same way a
|
||||
successful unencrypted result is."""
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
flatten_unencrypted_web_search_results_in_anthropic_messages,
|
||||
)
|
||||
|
|
@ -1837,6 +1841,7 @@ class TestAnthropicThinkingSignatureSelfHeal:
|
|||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "q"}},
|
||||
{
|
||||
"type": "web_search_tool_result",
|
||||
"tool_use_id": "srvtoolu_1",
|
||||
|
|
@ -1844,14 +1849,18 @@ class TestAnthropicThinkingSignatureSelfHeal:
|
|||
"type": "web_search_tool_result_error",
|
||||
"error_code": "max_uses_exceeded",
|
||||
},
|
||||
}
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
out = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs)
|
||||
once = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs)
|
||||
twice = flatten_unencrypted_web_search_results_in_anthropic_messages(once)
|
||||
|
||||
assert out[0] is msgs[0]
|
||||
assert once[0]["content"] == [
|
||||
{"type": "text", "text": "Web search results for 'q':\n\nSearch failed: max_uses_exceeded"}
|
||||
]
|
||||
assert json.dumps(twice) == json.dumps(once)
|
||||
|
||||
def test_sanitize_tool_use_ids_in_anthropic_messages(self):
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue