mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(websearch): address review - ReadOnly TypedDict fields, suppression reason, call-site coverage
- RichWebSearchInput fields are ReadOnly and constructed literally - the pyright suppression now states why the str provider name is safe - new tests drive _build_anthropic_request_patch and _build_chat_completion_request_patch end to end so the tool-call -> _rich_search_input wiring is covered, not just _execute_search Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
parent
5ef05b97c5
commit
05908bbe57
3 changed files with 85 additions and 8 deletions
|
|
@ -1646,18 +1646,25 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
"""
|
||||
if not isinstance(tool_input, Mapping):
|
||||
return None
|
||||
rich: RichWebSearchInput = {}
|
||||
objective = tool_input.get("objective")
|
||||
if isinstance(objective, str) and objective.strip():
|
||||
rich["objective"] = objective
|
||||
valid_objective = (
|
||||
objective if isinstance(objective, str) and objective.strip() else None
|
||||
)
|
||||
raw_queries = tool_input.get("search_queries")
|
||||
valid_queries: list[str] | None = None
|
||||
if isinstance(raw_queries, Sequence) and not isinstance(raw_queries, str):
|
||||
queries = [q for q in raw_queries if isinstance(q, str) and q.strip()]
|
||||
if queries:
|
||||
# Providers cap multi-query requests (Parallel drops queries
|
||||
# past the fifth); trim here so nothing is silently ignored.
|
||||
rich["search_queries"] = queries[:5]
|
||||
return rich or None
|
||||
valid_queries = queries[:5]
|
||||
if valid_objective is not None and valid_queries is not None:
|
||||
return {"objective": valid_objective, "search_queries": valid_queries}
|
||||
if valid_objective is not None:
|
||||
return {"objective": valid_objective}
|
||||
if valid_queries is not None:
|
||||
return {"search_queries": valid_queries}
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _provider_supports_rich_search(search_provider: str | None) -> bool:
|
||||
|
|
@ -1672,7 +1679,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
# misses the config map and returns None rather than raising.
|
||||
config = ProviderConfigManager.get_provider_search_config(
|
||||
search_provider
|
||||
) # pyright: ignore[reportArgumentType]
|
||||
) # pyright: ignore[reportArgumentType] -- SearchProviders is a str enum, so the router's provider string hashes to the matching member; unknown strings miss the map and yield None
|
||||
return config is not None and config.supports_rich_search_input()
|
||||
|
||||
async def _execute_search(
|
||||
|
|
|
|||
|
|
@ -36,10 +36,10 @@ class RichWebSearchInput(TypedDict, total=False):
|
|||
other provider keeps receiving the single ``query`` string.
|
||||
"""
|
||||
|
||||
objective: str
|
||||
objective: ReadOnly[str]
|
||||
"""Natural-language description of the goal behind the search."""
|
||||
|
||||
search_queries: list[str]
|
||||
search_queries: ReadOnly[list[str]]
|
||||
"""Two to five short keyword queries covering different angles."""
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -186,3 +186,73 @@ class TestExecuteSearchShape:
|
|||
|
||||
call_kwargs = mock_asearch.await_args.kwargs
|
||||
assert call_kwargs["objective"] == "configured objective"
|
||||
|
||||
|
||||
class TestCallSiteWiring:
|
||||
"""Drive the patch builders end to end so regressions in the tool-call ->
|
||||
_rich_search_input wiring are caught, not just _execute_search itself."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_tool_call_forwards_rich_shape(self, monkeypatch):
|
||||
import litellm
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
logger = WebSearchInterceptionLogger()
|
||||
mock_asearch = AsyncMock(return_value=_search_response())
|
||||
monkeypatch.setattr(proxy_server, "llm_router", _mock_router("parallel_ai"))
|
||||
monkeypatch.setattr(litellm, "asearch", mock_asearch)
|
||||
|
||||
tool_calls = [
|
||||
{"id": "toolu_1", "name": "litellm_web_search", "input": dict(RICH_INPUT)}
|
||||
]
|
||||
await logger._build_anthropic_request_patch(
|
||||
model="claude",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
tool_calls=tool_calls,
|
||||
thinking_blocks=[],
|
||||
anthropic_messages_optional_request_params={},
|
||||
logging_obj=None,
|
||||
kwargs={},
|
||||
)
|
||||
|
||||
call_kwargs = mock_asearch.await_args.kwargs
|
||||
assert call_kwargs["query"] == RICH_INPUT["search_queries"]
|
||||
assert call_kwargs["objective"] == RICH_INPUT["objective"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_completion_tool_call_forwards_rich_shape(self, monkeypatch):
|
||||
import json
|
||||
|
||||
import litellm
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
logger = WebSearchInterceptionLogger()
|
||||
mock_asearch = AsyncMock(return_value=_search_response())
|
||||
monkeypatch.setattr(proxy_server, "llm_router", _mock_router("parallel_ai"))
|
||||
monkeypatch.setattr(litellm, "asearch", mock_asearch)
|
||||
|
||||
# The normalized shape transform_request produces for OpenAI responses:
|
||||
# function.arguments (raw) plus top-level name/input (parsed).
|
||||
tool_calls = [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"name": "litellm_web_search",
|
||||
"function": {
|
||||
"name": "litellm_web_search",
|
||||
"arguments": json.dumps(RICH_INPUT),
|
||||
},
|
||||
"input": dict(RICH_INPUT),
|
||||
}
|
||||
]
|
||||
await logger._build_chat_completion_request_patch(
|
||||
model="claude",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
tool_calls=tool_calls,
|
||||
optional_params={},
|
||||
kwargs={},
|
||||
)
|
||||
|
||||
call_kwargs = mock_asearch.await_args.kwargs
|
||||
assert call_kwargs["query"] == RICH_INPUT["search_queries"]
|
||||
assert call_kwargs["objective"] == RICH_INPUT["objective"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue