diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index a38b8bed595..fd25d58642d 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -37,6 +37,7 @@ from litellm.integrations.websearch_interception.transformation import ( from litellm.litellm_core_utils.agentic_loop_settings import ( validated_max_agentic_loops, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.types.integrations.custom_logger import ( CHAT_COMPLETION_AGENTIC_SURFACE, @@ -514,7 +515,13 @@ class WebSearchInterceptionLogger(CustomLogger): get_litellm_web_search_tool_responses() if is_web_search_tool_responses(tool) else tool for tool in tools ] - converted_kwargs: Final = {**kwargs, "tools": converted_tools} + tool_choice: Final = as_str_mapping(kwargs.get("tool_choice")) + choice_patch: Final = ( + {"tool_choice": {"type": "function", "name": LITELLM_WEB_SEARCH_TOOL_NAME}} + if tool_choice is not None and is_web_search_tool_responses(tool_choice) + else {} + ) + converted_kwargs: Final = {**kwargs, "tools": converted_tools, **choice_patch} if kwargs.get("stream"): verbose_logger.debug("WebSearchInterception: deployment hook converting stream=True to stream=False") @@ -1266,7 +1273,7 @@ class WebSearchInterceptionLogger(CustomLogger): model=full_model_name, messages=input_list, tools=tools_param if isinstance(tools_param, list) else None, - optional_params=optional_params_clean, + optional_params={**optional_params_clean, "tool_choice": "auto"}, kwargs=kwargs_for_followup, ) diff --git a/litellm/llms/anthropic/messages/native/fake_stream_iterator.py b/litellm/llms/anthropic/messages/native/fake_stream_iterator.py index 39bdfbbc53a..1b2333adc0c 100644 --- a/litellm/llms/anthropic/messages/native/fake_stream_iterator.py +++ b/litellm/llms/anthropic/messages/native/fake_stream_iterator.py @@ -93,7 +93,7 @@ class FakeAnthropicMessagesStreamIterator: content_block_start = { "type": "content_block_start", "index": index, - "content_block": {"type": "redacted_thinking"}, + "content_block": block_dict, } chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode()) @@ -102,10 +102,8 @@ class FakeAnthropicMessagesStreamIterator: "type": "content_block_start", "index": index, "content_block": { - "type": "tool_use", - "id": block_dict.get("id"), - "name": block_dict.get("name"), - "input": {}, + **block_dict, + "input": dict[str, object](), }, } chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode()) diff --git a/litellm/utils.py b/litellm/utils.py index 6f0b1ca2e3e..f74b4d53a3b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1014,6 +1014,33 @@ def _is_converted_stream_result(result: object) -> bool: return isinstance(result, (CustomStreamWrapper, BaseResponsesAPIStreamingIterator)) +_MESSAGES_REPLAY_REQUEST: Final = TypeAdapter(Mapping[str, object]) +_MESSAGES_REPLAY_RESPONSE: Final = TypeAdapter(object) + + +def _is_buffered_messages_stream(result: object, request_data: Mapping[str, object], call_type: str) -> bool: + from litellm.types.integrations.custom_logger import converted_stream_requested + + return ( + call_type == CallTypes.anthropic_messages.value + and isinstance(result, dict) + and (bool(request_data.get("stream")) or converted_stream_requested(request_data)) + ) + + +def _restore_buffered_messages_stream(result: object, request_data: Mapping[str, object], call_type: str) -> object: + if not _is_buffered_messages_stream(result, request_data, call_type): + return result + from litellm.llms.anthropic.messages.native.fake_stream_iterator import FakeAnthropicMessagesStreamIterator + from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse + + return FakeAnthropicMessagesStreamIterator( + response=cast( # cast-ok: adapter-validated Messages body; preserve provider fields + AnthropicMessagesResponse, result + ) + ) + + async def _run_success_deployment_hook_on_converted_chat_stream( result: object, request_data: dict[str, object], call_type: str ) -> None: @@ -2205,6 +2232,12 @@ def client(original_function): if modified_kwargs is not None: kwargs = modified_kwargs + messages_replay_request: Final[Mapping[str, object]] = ( + _MESSAGES_REPLAY_REQUEST.validate_python(kwargs) + if call_type == CallTypes.anthropic_messages.value + else {} + ) + # Sync logging_obj.stream after deployment hooks (they may convert it). _hook_stream: Final = kwargs.get("stream") if ( @@ -2252,7 +2285,9 @@ def client(original_function): if _is_converted_stream_result(_caching_handler_response.cached_result): logging_obj.stream = True logging_obj.model_call_details["stream"] = True - return _caching_handler_response.cached_result + return _restore_buffered_messages_stream( + _caching_handler_response.cached_result, messages_replay_request, call_type + ) elif _caching_handler_response.embedding_all_elements_cache_hit is True: return _caching_handler_response.final_embedding_cached_response @@ -2316,7 +2351,10 @@ def client(original_function): end_time = datetime.datetime.now() streaming_requested: Final = _is_streaming_request(kwargs=kwargs, call_type=call_type) - if streaming_requested or _is_converted_stream_result(result): + buffered_messages_stream: Final = _is_buffered_messages_stream( + _MESSAGES_REPLAY_RESPONSE.validate_python(result), messages_replay_request, call_type + ) + if (streaming_requested and not buffered_messages_stream) or _is_converted_stream_result(result): logging_obj.stream = True logging_obj.model_call_details["stream"] = True if not streaming_requested: @@ -2406,7 +2444,9 @@ def client(original_function): is_litellm_internal_call=_is_litellm_internal_call, ) - return result + return _restore_buffered_messages_stream( + _MESSAGES_REPLAY_RESPONSE.validate_python(result), messages_replay_request, call_type + ) except Exception as e: traceback_exception: Final = traceback.format_exc() # Reuse the timestamp taken right when the deployment call itself failed, before diff --git a/tests/unit/integrations/websearch_interception/test_websearch_responses.py b/tests/unit/integrations/websearch_interception/test_websearch_responses.py index 00b7299cc31..d1ffe057616 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_responses.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_responses.py @@ -6,11 +6,16 @@ tool calls returned by /v1/responses, executes the search server-side, and builds a Responses-format follow-up request. """ +import json from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx +import litellm from litellm.integrations.websearch_interception.handler import ( WebSearchInterceptionLogger, ) @@ -247,7 +252,7 @@ async def test_build_responses_plan_produces_responses_input(): } assert patch_obj.tools == [{"type": "function", "name": "litellm_web_search"}] - assert "tool_choice" not in patch_obj.optional_params + assert patch_obj.optional_params["tool_choice"] == "auto" assert "_agentic_loop_api_surface" not in patch_obj.kwargs assert patch_obj.model == "openai/gpt-4o" @@ -316,3 +321,169 @@ async def test_deployment_hook_responses_converts_stream_to_non_stream(): assert result is not None assert result["stream"] is False assert result["_websearch_interception_converted_stream"] is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("search_type", ["web_search", "web_search_preview", "web_search_preview_2025_03_11"]) +async def test_forced_responses_search_names_the_converted_function(search_type: str, stream: bool) -> None: + logger: Final = WebSearchInterceptionLogger(enabled_providers=[LlmProviders.BEDROCK]) + choice: Final = {"type": search_type} + result: Final = await logger.async_pre_call_deployment_hook( + kwargs={ + "model": "us.openai.gpt-6.1-sol", + "custom_llm_provider": "bedrock", + "tools": [{"type": search_type}], + "tool_choice": choice, + "stream": stream, + }, + call_type=CallTypes.aresponses, + ) + + assert result is not None + assert result["tool_choice"] == {"type": "function", "name": result["tools"][0]["name"]} + assert result["stream"] is False + assert choice == {"type": search_type} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "choice", + [ + None, + "auto", + "none", + "required", + {"type": "function", "name": "get_weather"}, + {"type": "function", "name": "litellm_web_search"}, + ], +) +async def test_responses_search_conversion_preserves_other_tool_choices(choice: object) -> None: + logger: Final = WebSearchInterceptionLogger(enabled_providers=[LlmProviders.BEDROCK]) + weather: Final = {"type": "function", "name": "get_weather", "parameters": {"type": "object"}} + result: Final = await logger.async_pre_call_deployment_hook( + kwargs={ + "model": "us.openai.gpt-6.1-sol", + "custom_llm_provider": "bedrock", + "tools": [{"type": "web_search"}, weather], + "tool_choice": choice, + }, + call_type=CallTypes.aresponses, + ) + + assert result is not None + assert result["tool_choice"] == choice + assert result["tools"][1] is weather + + +@pytest.mark.asyncio +async def test_forced_responses_search_executes_once_then_allows_the_answer(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("TAVILY_API_BASE", "https://search.example.test") + monkeypatch.setattr(litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["bedrock"])]) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "search-canary", + "litellm_params": { + "model": "bedrock/us.openai.gpt-6.1-sol", + "api_key": "synthetic-bedrock-bearer", + "aws_region_name": "us-east-1", + }, + } + ], + search_tools=[ + { + "search_tool_name": "fixture-search", + "litellm_params": { + "search_provider": "tavily", + "api_key": "synthetic-search-key", + "api_base": "https://search.example.test", + }, + } + ], + num_retries=0, + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + common: Final = { + "object": "response", + "created_at": 1, + "status": "completed", + "model": "us.openai.gpt-6.1-sol", + "usage": {"input_tokens": 11, "output_tokens": 5, "total_tokens": 16}, + } + with respx.mock as upstream: + search: Final = upstream.post("https://search.example.test/search").respond( + 200, + json={ + "results": [ + { + "title": "Synthetic search", + "url": "https://example.test/search", + "content": "Fixture evidence", + } + ] + }, + ) + responses_route: Final = upstream.post( + "https://bedrock-runtime.us-east-1.amazonaws.com/openai/v1/responses" + ).mock( + side_effect=[ + httpx.Response( + 200, + json={ + **common, + "id": "resp_search", + "output": [ + { + "type": "function_call", + "id": "fc_search", + "call_id": "call_search", + "name": "litellm_web_search", + "arguments": '{"query":"synthetic query"}', + "status": "completed", + } + ], + }, + ), + httpx.Response( + 200, + json={ + **common, + "id": "resp_answer", + "output": [ + { + "type": "message", + "id": "msg_answer", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Answer from search", "annotations": []}], + } + ], + }, + ), + ] + ) + result: Final = await router.aresponses( + model="search-canary", + input="Search for synthetic evidence", + max_output_tokens=100, + tools=[{"type": "web_search"}], + tool_choice={"type": "web_search"}, + ) + + assert result.output[0].content[0].text == "Answer from search" + assert len(responses_route.calls) == 2 + first: Final = json.loads(responses_route.calls[0].request.content) + followup: Final = json.loads(responses_route.calls[1].request.content) + assert first["tool_choice"] == {"type": "function", "name": first["tools"][0]["name"]} + assert followup["tool_choice"] == "auto" + assert followup["input"][-1] == { + "type": "function_call_output", + "call_id": "call_search", + "output": "Title: Synthetic search\nURL: https://example.test/search\nSnippet: Fixture evidence", + } + assert len(search.calls) == 1 + assert json.loads(search.calls[0].request.content)["query"] == "synthetic query" diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py index c11f58faec4..2075f7814c9 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py @@ -2,34 +2,38 @@ import asyncio import copy import json import os +import unittest.mock import uuid +from itertools import chain from typing import Any, Dict, Final, List +from unittest.mock import AsyncMock, MagicMock, patch import httpx +import pytest import responses import respx -import pytest from fastapi.testclient import TestClient from pydantic import ValidationError - -from unittest.mock import AsyncMock, MagicMock, patch - import litellm from litellm.anthropic_interface import messages +from litellm.caching.caching import Cache +from litellm.caching.caching_handler import _PENDING_CACHE_WRITES from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.router import Router from litellm.types.utils import ( + CallTypes, Delta, + LLMResponseTypes, ModelResponse, StandardLoggingPayloadErrorInformation, StreamingChoices, ) from litellm.utils import _invalidate_model_cost_lowercase_map -import unittest.mock @pytest.mark.asyncio @@ -2396,3 +2400,145 @@ def test_native_messages_tool_config_point_marks_last_tool_and_drops_the_key(mon assert "cache_control" not in captured["params"]["tools"][0] assert captured["messages"][0]["content"][-1]["cache_control"] == {"type": "ephemeral"} assert tools == tools_before + + +class _MessagesAnswerRewrite(CustomLogger): + async def async_post_call_success_deployment_hook( + self, request_data: dict, response: LLMResponseTypes, call_type: CallTypes | None + ) -> LLMResponseTypes | None: + if call_type == CallTypes.anthropic_messages and isinstance(response, dict): + return {**response, "content": [{"type": "text", "text": "Reviewed answer"}]} + return None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("tool_call", [False, True], ids=["text", "ordinary-tool"]) +@pytest.mark.parametrize("rewrite", [False, True], ids=["original", "success-hook"]) +@pytest.mark.parametrize("cache_enabled", [False, True], ids=["uncached", "response-cache"]) +async def test_bedrock_messages_mixed_search_tools_preserve_the_client_stream( + monkeypatch: pytest.MonkeyPatch, stream: bool, tool_call: bool, rewrite: bool, cache_enabled: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr( + litellm, + "callbacks", + [WebSearchInterceptionLogger(enabled_providers=["bedrock"])] + ([_MessagesAnswerRewrite()] if rewrite else []), + ) + if cache_enabled: + monkeypatch.setattr(litellm, "cache", Cache(type="local")) + content: Final = ( + [{"toolUse": {"toolUseId": "call_weather", "name": "get_weather", "input": {"city": "Paris"}}}] + if tool_call + else [{"text": "No search needed"}] + ) + stop_reason: Final = "tool_use" if tool_call else "end_turn" + with respx.mock as upstream: + converse: Final = upstream.post( + "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.openai.gpt-6.1-sol/converse" + ).respond( + 200, + json={ + "output": {"message": {"role": "assistant", "content": content}}, + "stopReason": stop_reason, + "usage": {"inputTokens": 11, "outputTokens": 5, "totalTokens": 16}, + "metrics": {"latencyMs": 1}, + }, + ) + result: Final = await litellm.anthropic.messages.acreate( + model="bedrock/us.openai.gpt-6.1-sol", + api_key="synthetic-bedrock-bearer", + aws_region_name="us-east-1", + max_tokens=100, + messages=[{"role": "user", "content": "Answer or call get_weather for Paris without searching."}], + tools=[ + {"type": "web_search_20250305", "name": "web_search"}, + { + "name": "get_weather", + "description": "Get weather", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + ], + stream=stream, + ) + assert len(converse.calls) == 1 + assert "stream" not in json.loads(converse.calls.last.request.content) + if cache_enabled: + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) + replay: Final = await litellm.anthropic.messages.acreate( + model="bedrock/us.openai.gpt-6.1-sol", + api_key="synthetic-bedrock-bearer", + aws_region_name="us-east-1", + max_tokens=100, + messages=[{"role": "user", "content": "Answer or call get_weather for Paris without searching."}], + tools=[ + {"type": "web_search_20250305", "name": "web_search"}, + { + "name": "get_weather", + "description": "Get weather", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + ], + stream=stream, + ) + if stream: + original_chunks: Final = tuple([chunk async for chunk in result]) + replay_chunks: Final = tuple([chunk async for chunk in replay]) + assert replay_chunks == original_chunks + else: + assert replay["content"] == result["content"] + assert len(converse.calls) == 1 + + if not stream: + assert result["stop_reason"] == stop_reason + assert result["content"] == ( + [{"type": "text", "text": "Reviewed answer"}] + if rewrite + else ( + [ + {"type": "text", "text": ""}, + {"type": "tool_use", "id": "call_weather", "name": "get_weather", "input": {"city": "Paris"}}, + ] + if tool_call + else [{"type": "text", "text": "No search needed"}] + ) + ) + return + + assert hasattr(result, "__aiter__"), result + chunks: Final = original_chunks if cache_enabled else tuple([chunk async for chunk in result]) + lines: Final = chain.from_iterable(chunk.decode().splitlines() for chunk in chunks) + events: Final = [json.loads(line.removeprefix("data: ")) for line in lines if line.startswith("data: ")] + assert events[0]["type"] == "message_start" + assert events[-1] == {"type": "message_stop", "usage": {"input_tokens": 11, "output_tokens": 5}} + assert events[0]["message"]["usage"]["input_tokens"] == 11 + assert events[-2]["usage"]["output_tokens"] == 5 + assert events[-2]["delta"]["stop_reason"] == stop_reason + starts: Final = [event["content_block"] for event in events if event["type"] == "content_block_start"] + deltas: Final = [event["delta"] for event in events if event["type"] == "content_block_delta"] + assert starts[-1] == ( + {"type": "text", "text": ""} + if rewrite + else ( + {"type": "tool_use", "id": "call_weather", "name": "get_weather", "input": {}} + if tool_call + else {"type": "text", "text": ""} + ) + ) + assert deltas[-1] == ( + {"type": "text_delta", "text": "Reviewed answer"} + if rewrite + else ( + {"type": "input_json_delta", "partial_json": json.dumps({"city": "Paris"})} + if tool_call + else {"type": "text_delta", "text": "No search needed"} + ) + ) diff --git a/tests/unit/llms/anthropic/messages/native/test_fake_stream_iterator.py b/tests/unit/llms/anthropic/messages/native/test_fake_stream_iterator.py index 233145f2706..8201c5523a5 100644 --- a/tests/unit/llms/anthropic/messages/native/test_fake_stream_iterator.py +++ b/tests/unit/llms/anthropic/messages/native/test_fake_stream_iterator.py @@ -91,3 +91,25 @@ def test_message_delta_omits_absent_optional_fields(): "usage": {"input_tokens": 10, "output_tokens": 2}, } assert _message_delta(response) == {"stop_reason": "end_turn", "stop_sequence": None} + + +def test_replayed_stream_preserves_redacted_reasoning_and_tool_signatures(): + redacted: Final = {"type": "redacted_thinking", "data": "opaque-continuation-data"} + tool: Final = { + "type": "tool_use", + "id": "call_signed", + "name": "get_weather", + "input": {"city": "Paris"}, + "provider_specific_fields": {"thought_signature": "opaque-tool-signature"}, + } + response: Final = { + "id": "msg_signed", + "content": [redacted, tool], + "stop_reason": "tool_use", + "usage": {"input_tokens": 11, "output_tokens": 5}, + } + events: Final = _decoded_events(FakeAnthropicMessagesStreamIterator(response=response)) + starts: Final = tuple(payload["content_block"] for kind, payload in events if kind == "content_block_start") + assert starts == (redacted, {**tool, "input": {}}) + deltas: Final = tuple(payload["delta"] for kind, payload in events if kind == "content_block_delta") + assert deltas == ({"type": "input_json_delta", "partial_json": json.dumps(tool["input"])},)