diff --git a/litellm/main.py b/litellm/main.py index 78b6f8e4153..873aef483bf 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6517,6 +6517,7 @@ _CHAT_SURFACE_KWARGS: Final = ( "mock_delay", "mock_timeout", ) +_SKIP_RESPONSES_API_BRIDGE: Final = "_skip_responses_api_bridge" _TEXT_COMPLETION_PROVIDERS: Final = frozenset( ("text-completion-openai", "text-completion-codestral", "text-completion-inception") ) @@ -6530,6 +6531,7 @@ class PlannedChat: kwargs: Mapping[str, object] non_default_params: Mapping[str, object] control_options: ControlOptions + skip_responses_api_bridge: bool def _optional_str(value: object) -> str | None: @@ -6599,7 +6601,11 @@ def plan_chat(request: Mapping[str, object], logging_obj: LiteLLMLoggingObj) -> ) if resolved_provider in _TEXT_COMPLETION_PROVIDERS: return None - kwargs: Final = {key: value for key, value in request.items() if key not in _EXECUTE_CHAT_PARAMS} + kwargs: Final = { + key: value + for key, value in request.items() + if key not in _EXECUTE_CHAT_PARAMS and key != _SKIP_RESPONSES_API_BRIDGE + } tools: Final = validate_and_fix_openai_tools(tools=_object_list(request.get("tools"))) messages: Final = get_completion_messages( messages=validate_and_fix_openai_messages(messages=_object_list(request.get("messages")) or []), @@ -6639,6 +6645,7 @@ def plan_chat(request: Mapping[str, object], logging_obj: LiteLLMLoggingObj) -> kwargs=MappingProxyType(kwargs), non_default_params=MappingProxyType(non_default_params), control_options=_resolve_control_options(kwargs, model), + skip_responses_api_bridge=request.get(_SKIP_RESPONSES_API_BRIDGE) is True, ) @@ -6648,7 +6655,7 @@ def _run_planned_chat(plan: PlannedChat, is_async: bool) -> ModelResponse | Cust kwargs={**plan.kwargs, "acompletion": is_async}, non_default_params=dict(plan.non_default_params), control_options=plan.control_options, - skip_responses_api_bridge=False, + skip_responses_api_bridge=plan.skip_responses_api_bridge, prompt_id=None, prompt_variables=None, ) diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 5489d2223ac..0b148985b9a 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -3,7 +3,7 @@ Handler for transforming responses api requests to litellm.completion requests """ from collections.abc import Coroutine, Mapping -from typing import Final +from typing import TYPE_CHECKING, Final from pydantic import TypeAdapter @@ -23,9 +23,35 @@ from litellm.types.llms.openai import ( ) from litellm.types.utils import ModelResponse +if TYPE_CHECKING: + from litellm.main import PlannedChat + _REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(dict[str, object]) +def _plan_bridged_chat(request: Mapping[str, object]) -> "PlannedChat | None": + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.main import plan_chat + + logging_obj: Final = request.get("litellm_logging_obj") + return plan_chat(request, logging_obj) if isinstance(logging_obj, Logging) else None + + +def _price_planned_chat_response( + request: Mapping[str, object], response: ModelResponse | litellm.CustomStreamWrapper +) -> None: + """Prices the chat response as `acompletion` would have, before it becomes a Responses object.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj: Final = request.get("litellm_logging_obj") + if not isinstance(logging_obj, Logging): + return + if isinstance(response, ModelResponse): + logging_obj.price_bridged_provider_response(response) + else: + logging_obj.price_bridged_provider_stream(response) + + class LiteLLMCompletionTransformationHandler: def response_api_handler( self, @@ -71,15 +97,18 @@ class LiteLLMCompletionTransformationHandler: **kwargs, ) - completion_args: Final = {} - completion_args.update(kwargs) - completion_args.update(litellm_completion_request) - completion_args["_skip_responses_api_bridge"] = True + completion_args: Final = {**kwargs, **litellm_completion_request, "_skip_responses_api_bridge": True} - litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = litellm.completion( - **completion_args, + from litellm.main import execute_planned_chat + + planned_chat: Final = _plan_bridged_chat(_REQUEST_KWARGS_ADAPTER.validate_python(completion_args)) + litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = ( + litellm.completion(**completion_args) if planned_chat is None else execute_planned_chat(planned_chat) ) + if planned_chat is not None: + _price_planned_chat_response(completion_args, litellm_completion_response) + if isinstance(litellm_completion_response, ModelResponse): responses_api_response: Final[ResponsesAPIResponse] = ( LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( @@ -122,15 +151,20 @@ class LiteLLMCompletionTransformationHandler: previous_response_client_ids=previous_response_client_ids, ) - acompletion_args: Final = {} - acompletion_args.update(kwargs) - acompletion_args.update(litellm_completion_request) - acompletion_args["_skip_responses_api_bridge"] = True + acompletion_args: Final = {**kwargs, **litellm_completion_request, "_skip_responses_api_bridge": True} - litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = await litellm.acompletion( - **acompletion_args, + from litellm.main import aexecute_planned_chat + + planned_chat: Final = _plan_bridged_chat(_REQUEST_KWARGS_ADAPTER.validate_python(acompletion_args)) + litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = ( + await litellm.acompletion(**acompletion_args) + if planned_chat is None + else await aexecute_planned_chat(planned_chat) ) + if planned_chat is not None: + _price_planned_chat_response(acompletion_args, litellm_completion_response) + if isinstance(litellm_completion_response, ModelResponse): responses_api_response: Final[ResponsesAPIResponse] = ( LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 615bdb5bbfb..b2aadc18198 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -1,8 +1,11 @@ import asyncio import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from itertools import groupby from typing import TYPE_CHECKING, Any, Final, cast +from pydantic import TypeAdapter + import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER @@ -28,6 +31,7 @@ else: # Cold Storage Handler ######################################################## COLD_STORAGE_HANDLER: Final = ColdStorageHandler() +_STORED_REQUEST: Final = TypeAdapter(dict[str, object]) ######################################################## @@ -66,6 +70,76 @@ def _stored_responses_api_output(stored_response: object) -> ResponseInputParam return None +def _stored_web_search_results(stored_response: object) -> Sequence[Mapping[str, object]]: + """The Anthropic ``web_search_tool_result`` blocks a bridged response kept in its ``provider_specific_fields``.""" + fields: Final = stored_response.get("provider_specific_fields") if isinstance(stored_response, Mapping) else None + results: Final = fields.get("web_search_results") if isinstance(fields, Mapping) else None + return tuple(result for result in results if isinstance(result, Mapping)) if isinstance(results, list) else () + + +def _server_web_search_call_id(item: object, web_search_results: Sequence[Mapping[str, object]]) -> str | None: + """The ``server_tool_use`` id behind a stored ``web_search_call`` item, when its result block was kept.""" + if not isinstance(item, Mapping) or item.get("type") != "web_search_call": + return None + item_id: Final = item.get("id") + call_id: Final = item_id.removeprefix("ws_") if isinstance(item_id, str) else None + if call_id and any(result.get("tool_use_id") == call_id for result in web_search_results): + return call_id + return None + + +def _server_web_search_message(items: Sequence[object], web_search_results: Sequence[Mapping[str, object]]) -> Message: + """Replay stored web search calls as the assistant tool calls the chat completion carried, + so an Anthropic target rebuilds its ``server_tool_use`` / ``web_search_tool_result`` blocks.""" + calls: Final = tuple( + (call_id, item) + for item in items + if isinstance(item, Mapping) and (call_id := _server_web_search_call_id(item, web_search_results)) is not None + ) + call_ids: Final = frozenset(call_id for call_id, _ in calls) + return Message( + role="assistant", + content=None, + tool_calls=[ + { + "id": call_id, + "type": "function", + "function": {"name": "web_search", "arguments": json.dumps({"query": _web_search_query(item)})}, + } + for call_id, item in calls + ], + provider_specific_fields={ + "web_search_results": [result for result in web_search_results if result.get("tool_use_id") in call_ids] + }, + ) + + +def _web_search_query(item: Mapping[str, object]) -> str: + action: Final = item.get("action") + query: Final = action.get("query") if isinstance(action, Mapping) else None + return query if isinstance(query, str) else "" + + +_CHAT_MESSAGE_CONTENT_ITEM_TYPES: Final = ("reasoning", "message") + + +def _in_chat_message_order( + output: ResponseInputParam, web_search_results: Sequence[Mapping[str, object]] +) -> ResponseInputParam: + """A turn with server web searches in the order its chat completion message replayed it: reasoning, then text, + then the tool calls as they happened. Any other turn keeps its output order.""" + if not any(_server_web_search_call_id(item, web_search_results) is not None for item in output): + return output + return sorted( + output, + key=lambda item: ( + _CHAT_MESSAGE_CONTENT_ITEM_TYPES.index(item_type) + if isinstance(item, Mapping) and (item_type := item.get("type")) in _CHAT_MESSAGE_CONTENT_ITEM_TYPES + else len(_CHAT_MESSAGE_CONTENT_ITEM_TYPES) + ), + ) + + def _stored_instructions(proxy_server_request: Mapping[str, object] | None) -> str | None: instructions: Final = None if proxy_server_request is None else proxy_server_request.get("instructions") return instructions if isinstance(instructions, str) and instructions else None @@ -189,14 +263,22 @@ class ResponsesSessionHandler: ############################################################ _response_output: Final[object] = spend_log.get("response", "{}") if (responses_api_output := _stored_responses_api_output(_response_output)) is not None: - chat_completion_message_history.extend( - LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( - input=responses_api_output, - responses_api_request={}, - replay_reasoning=True, - item_cache_control=item_cache_control, + web_search_results: Final = _stored_web_search_results(_response_output) + for is_server_web_search, items in groupby( + _in_chat_message_order(responses_api_output, web_search_results), + key=lambda item: _server_web_search_call_id(item, web_search_results) is not None, + ): + if is_server_web_search: + chat_completion_message_history.append(_server_web_search_message(tuple(items), web_search_results)) + continue + chat_completion_message_history.extend( + LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=list(items), + responses_api_request={}, + replay_reasoning=True, + item_cache_control=item_cache_control, + ) ) - ) elif isinstance(_response_output, dict) and _response_output and _response_output != {}: # transform `ChatCompletion Response` to `ResponsesAPIResponse` model_response: Final = ModelResponse(**_response_output) @@ -214,30 +296,23 @@ class ResponsesSessionHandler: Get the parsed proxy server request from the spend log """ proxy_server_request: Final[str | dict] = spend_log.get("proxy_server_request") or "{}" - proxy_server_request_dict: dict[str, object] | None = None - if isinstance(proxy_server_request, dict): - proxy_server_request_dict = proxy_server_request - else: - proxy_server_request_dict = json.loads(proxy_server_request) + proxy_server_request_dict: Final = _STORED_REQUEST.validate_python( + proxy_server_request if isinstance(proxy_server_request, dict) else json.loads(proxy_server_request) + ) ############################################################ # Check if user has setup cold storage for session handling ############################################################ - if ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_server_request_dict): - # Try to get cold storage object key from spend log metadata - _proxy_server_request_dict: dict | None = None - cold_storage_object_key = ResponsesSessionHandler._get_cold_storage_object_key_from_spend_log(spend_log) - if cold_storage_object_key: - # Use the object key directly from metadata - _proxy_server_request_dict = ( - await ResponsesSessionHandler.get_proxy_server_request_from_cold_storage_with_object_key( - object_key=cold_storage_object_key, - ) - ) - if _proxy_server_request_dict: - proxy_server_request_dict = _proxy_server_request_dict - - return proxy_server_request_dict + if not ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_server_request_dict): + return proxy_server_request_dict + cold_storage_object_key: Final = ResponsesSessionHandler._get_cold_storage_object_key_from_spend_log(spend_log) + if not cold_storage_object_key: + return proxy_server_request_dict + return ( + await ResponsesSessionHandler.get_proxy_server_request_from_cold_storage_with_object_key( + object_key=cold_storage_object_key, + ) + ) or proxy_server_request_dict @staticmethod def _get_cold_storage_object_key_from_spend_log( @@ -279,14 +354,10 @@ class ResponsesSessionHandler: """ verbose_proxy_logger.debug("inside get_proxy_server_request_from_cold_storage_with_object_key...") - proxy_server_request_dict: Final = ( - await COLD_STORAGE_HANDLER.get_proxy_server_request_from_cold_storage_with_object_key( - object_key=object_key, - ) + return await COLD_STORAGE_HANDLER.get_proxy_server_request_from_cold_storage_with_object_key( + object_key=object_key, ) - return proxy_server_request_dict - @staticmethod def _should_check_cold_storage_for_full_payload( proxy_server_request_dict: Mapping[str, object] | None, diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index 025f1315aab..19749d51061 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -1,3 +1,4 @@ +import datetime import json from collections.abc import Mapping, Sequence from itertools import chain @@ -10,6 +11,7 @@ import respx from pydantic import BaseModel import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, @@ -86,8 +88,8 @@ async def test_aresponses_forwards_timeout_to_acompletion(): and similar providers, with calls falling back to the provider SDK default (~600s for Anthropic). """ - with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion: - mock_acompletion.return_value = ModelResponse( + with patch("litellm.main.aexecute_planned_chat", new_callable=AsyncMock) as mock_chat: + mock_chat.return_value = ModelResponse( id="id", created=0, model="anthropic/claude-sonnet-4-5", @@ -102,8 +104,8 @@ async def test_aresponses_forwards_timeout_to_acompletion(): api_key="sk-ant-fake", ) - assert mock_acompletion.call_count == 1 - forwarded_timeout = mock_acompletion.call_args.kwargs.get("timeout") + assert mock_chat.call_count == 1 + forwarded_timeout = mock_chat.call_args.args[0].named.get("timeout") assert forwarded_timeout == 42, ( f"timeout was not forwarded to acompletion (got {forwarded_timeout!r}); " "this means Router(timeout=N) silently fails for providers on the " @@ -541,3 +543,153 @@ def test_streamed_responses_over_anthropic_emit_the_responses_event_sequence(res assert deltas == "Argentina won." assert completed.response.status == "completed" assert completed.response.output[0].content[0].text == "Argentina won." + + +class _DeploymentHookCounter(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.calls = 0 + + async def async_pre_call_deployment_hook(self, kwargs: Mapping[str, object], call_type: object) -> None: + self.calls += 1 + + +@pytest.mark.asyncio +async def test_bridged_responses_call_runs_deployment_hooks_once(monkeypatch: pytest.MonkeyPatch): + counter: Final = _DeploymentHookCounter() + monkeypatch.setattr(litellm, "callbacks", [counter]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock(assert_all_called=False) as upstream: + messages_route: Final = upstream.post("http://anthropic.test/v1/messages").mock( + return_value=httpx.Response( + 200, + json={ + "id": "msg_bridge", + "type": "message", + "role": "assistant", + "model": "claude-bridge", + "content": [{"type": "text", "text": "from upstream"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 5}, + }, + ) + ) + response: Final = await litellm.aresponses( + model="anthropic/claude-bridge", input="hi", api_base="http://anthropic.test", api_key="sk-test" + ) + + assert response.output[0].content[0].text == "from upstream" + assert messages_route.call_count == 1 + assert counter.calls == 1 + + +@pytest.mark.asyncio +async def test_bridged_responses_call_is_billed_like_the_chat_call_for_one_hour_cache_writes_and_web_searches( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + usage: Final = { + "input_tokens": 100, + "output_tokens": 50, + "cache_creation_input_tokens": 2000, + "cache_read_input_tokens": 3000, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 2000}, + "server_tool_use": {"web_search_requests": 2}, + } + upstream_message: Final = { + "id": "msg_bridge_cost", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "searched"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": usage, + } + credentials: Final = {"api_base": "http://anthropic.test", "api_key": "sk-test"} + + with respx.mock(assert_all_called=False) as upstream: + upstream.post("http://anthropic.test/v1/messages").mock(return_value=httpx.Response(200, json=upstream_message)) + chat: Final = await litellm.acompletion( + model="anthropic/claude-sonnet-4-5", messages=[{"role": "user", "content": "hi"}], **credentials + ) + bridged: Final = await litellm.aresponses(model="anthropic/claude-sonnet-4-5", input="hi", **credentials) + + chat_cost: Final = chat._hidden_params["response_cost"] + assert isinstance(chat_cost, float) and chat_cost > 0 + assert bridged._hidden_params["response_cost"] == pytest.approx(chat_cost) + + +_ANTHROPIC_MESSAGE: Final = { + "id": "msg_bridge_headers", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 5}, +} + + +@pytest.mark.asyncio +async def test_bridged_responses_call_names_the_provider_request_id_in_its_headers(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock(assert_all_called=False) as upstream: + upstream.post("http://anthropic.test/v1/messages").mock( + return_value=httpx.Response(200, json=_ANTHROPIC_MESSAGE, headers={"x-request-id": "req-upstream-1"}) + ) + bridged: Final = await litellm.aresponses( + model="anthropic/claude-sonnet-4-5", input="hi", api_base="http://anthropic.test", api_key="sk-test" + ) + + assert bridged._hidden_params["additional_headers"]["llm_provider-request-id"] == "req-upstream-1" + + +@pytest.mark.asyncio +async def test_bridged_responses_stream_opens_with_the_cost_breakdown_headers(monkeypatch: pytest.MonkeyPatch): + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy.common_request_processing import cost_headers + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + logging_obj: Final = Logging( + model="anthropic/claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="aresponses", + start_time=datetime.datetime.now(datetime.timezone.utc), + litellm_call_id="bridge-stream-headers-0000-4000-8000-000000000001", + function_id="1234", + ) + events: Final = ( + {"type": "message_start", "message": {**_ANTHROPIC_MESSAGE, "content": [], "stop_reason": None}}, + {"type": "message_stop"}, + ) + sse: Final = "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events) + + with respx.mock(assert_all_called=False) as upstream: + upstream.post("http://anthropic.test/v1/messages").mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=sse) + ) + stream: Final = await litellm.aresponses( + model="anthropic/claude-sonnet-4-5", + input="hi", + stream=True, + api_base="http://anthropic.test", + api_key="sk-test", + litellm_logging_obj=logging_obj, + ) + headers: Final = cost_headers(None, logging_obj) + async for _ in stream: + pass + + assert headers["x-litellm-response-cost-output"] == "0.0" + assert headers["x-litellm-response-cost-discount-amount"] == "0.0" + assert headers["x-litellm-response-cost-margin-amount"] == "0.0" diff --git a/tests/unit/responses/litellm_completion_transformation/test_google_ai_studio_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_google_ai_studio_responses_bridge.py index 24db797f53b..8aa03d7629a 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_google_ai_studio_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_google_ai_studio_responses_bridge.py @@ -31,8 +31,8 @@ async def test_mock_basic_google_ai_studio_responses_api_with_tools(): ], ) - with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion: - mock_acompletion.return_value = mock_response + with patch("litellm.main.aexecute_planned_chat", new_callable=AsyncMock) as mock_chat: + mock_chat.return_value = mock_response request_model = "gemini/gemini-2.5-flash" await litellm.aresponses( @@ -41,11 +41,10 @@ async def test_mock_basic_google_ai_studio_responses_api_with_tools(): tools=[{"type": "web_search_preview", "search_context_size": "low"}], ) - # Verify that acompletion was called - assert mock_acompletion.called + assert mock_chat.called - # Get the call arguments - call_args, call_kwargs = mock_acompletion.call_args + plan = mock_chat.call_args.args[0] + call_kwargs = {**plan.kwargs, **plan.named} # Verify the expected parameters were passed print( @@ -65,5 +64,5 @@ async def test_mock_basic_google_ai_studio_responses_api_with_tools(): call_kwargs["messages"][0]["content"] == "what is the latest version of supabase python package and when was it released?" ) - assert "tools" not in call_kwargs - assert "tool_choice" not in call_kwargs + assert call_kwargs.get("tools") is None + assert call_kwargs.get("tool_choice") is None diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py index 7a146c1b665..6738d6c5c99 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -233,11 +233,10 @@ async def test_e2e_cold_storage_successful_retrieval(): patch.object(session_handler, "COLD_STORAGE_HANDLER") as mock_cold_storage, patch("litellm.cold_storage_custom_logger", return_value="s3"), ): - # Setup mocks mock_get_spend_logs.return_value = mock_spend_logs - mock_cold_storage.get_proxy_server_request_from_cold_storage_with_object_key = ( - AsyncMock(return_value=full_proxy_request) + mock_cold_storage.get_proxy_server_request_from_cold_storage_with_object_key = AsyncMock( + return_value=full_proxy_request ) # Call the main function @@ -291,7 +290,6 @@ async def test_e2e_cold_storage_fallback_to_truncated_payload(): ) as mock_get_spend_logs, patch.object(session_handler, "COLD_STORAGE_HANDLER") as mock_cold_storage, ): - # Setup mocks mock_get_spend_logs.return_value = mock_spend_logs @@ -358,41 +356,25 @@ async def test_should_check_cold_storage_for_full_payload(): with patch("litellm.cold_storage_custom_logger", return_value="s3"): # Test case 1: Should return True for truncated content - result1 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_with_truncated_pdf - ) - assert ( - result1 == True - ), "Should return True for proxy request with truncated PDF content" + result1 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_with_truncated_pdf) + assert result1 == True, "Should return True for proxy request with truncated PDF content" # Test case 2: Should return False for regular content - result2 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_regular - ) - assert ( - result2 == False - ), "Should return False for regular proxy request without truncation" + result2 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_regular) + assert result2 == False, "Should return False for regular proxy request without truncation" # Test case 3: Should return True for empty request - result3 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_empty - ) + result3 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_empty) assert result3 == True, "Should return True for empty proxy request" # Test case 4: Should return True for None request - result4 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_none - ) + result4 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_none) assert result4 == True, "Should return True for None proxy request" # Test case 5: Should return False when cold storage is not configured with patch.object(litellm, "cold_storage_custom_logger", None): - result5 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_with_truncated_pdf - ) - assert ( - result5 == False - ), "Should return False when cold storage is not configured, even with truncated content" + result5 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_with_truncated_pdf) + assert result5 == False, "Should return False when cold storage is not configured, even with truncated content" @pytest.mark.asyncio @@ -422,9 +404,7 @@ async def test_get_chat_completion_message_history_empty_response_dict(): } ] - with patch.object( - ResponsesSessionHandler, "get_all_spend_logs_for_previous_response_id" - ) as mock_get_spend_logs: + with patch.object(ResponsesSessionHandler, "get_all_spend_logs_for_previous_response_id") as mock_get_spend_logs: mock_get_spend_logs.return_value = mock_spend_logs # Call the function @@ -518,9 +498,7 @@ async def test_message_history_reconstructs_list_shaped_input(): ) as mock_get_spend_logs: mock_get_spend_logs.return_value = mock_spend_logs - result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( - request_id - ) + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(request_id) messages = result["messages"] assert [(message.get("role"), message.get("content")) for message in messages] == [ @@ -549,9 +527,7 @@ async def test_message_history_retries_a_spend_log_the_batch_writer_has_not_flus fake_prisma_client = _FakePrismaClient(results=[[], [spend_log]]) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): - result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( - request_id - ) + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(request_id) messages = result["messages"] assert [(message.get("role"), message.get("content")) for message in messages] == [ @@ -618,14 +594,10 @@ async def test_message_history_looks_up_the_decoded_chat_completion_id(): model_id="e0f302a1412e78470ebb28cbed01fff5f88c0d331c667e9f2ba4b413c6fbd282", response_id=request_id, ) - fake_prisma_client = _FakePrismaClient( - results=[[_spend_log(request_id, "session-a", "Hello.", "Hi.")]] - ) + fake_prisma_client = _FakePrismaClient(results=[[_spend_log(request_id, "session-a", "Hello.", "Hi.")]]) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): - await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( - encoded_response_id - ) + await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(encoded_response_id) assert len(fake_prisma_client.db.calls) == 1 assert request_id in fake_prisma_client.db.calls[0] @@ -724,8 +696,9 @@ async def test_session_lookup_does_not_retry_when_spend_logs_are_disabled( """ fake_prisma_client = _FakePrismaClient(results=[]) - with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client), patch( - "litellm.proxy.proxy_server.disable_spend_logs", True + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client), + patch("litellm.proxy.proxy_server.disable_spend_logs", True), ): spend_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( "chatcmpl-does-not-exist" @@ -741,7 +714,11 @@ def test_normalize_redacted_arguments_skips_custom_tool_calls(): content=None, tool_calls=[ {"id": "call_c", "type": "custom", "custom": {"name": "run_code", "input": "print(1)"}}, - {"id": "call_f", "type": "function", "function": {"name": "get_weather", "arguments": "redacted-by-litellm"}}, + { + "id": "call_f", + "type": "function", + "function": {"name": "get_weather", "arguments": "redacted-by-litellm"}, + }, ], ) @@ -796,11 +773,13 @@ async def test_message_history_normalizes_redacted_tool_call_arguments(): } ] - with patch.object( # test-quality-ok: the handler has no DI seam for the spend-log fetch; every test in this file stubs this same boundary - ResponsesSessionHandler, - "get_all_spend_logs_for_previous_response_id", - new_callable=AsyncMock, - ) as mock_get_spend_logs: + with ( + patch.object( # test-quality-ok: the handler has no DI seam for the spend-log fetch; every test in this file stubs this same boundary + ResponsesSessionHandler, + "get_all_spend_logs_for_previous_response_id", + new_callable=AsyncMock, + ) as mock_get_spend_logs + ): mock_get_spend_logs.return_value = mock_spend_logs result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( @@ -889,7 +868,9 @@ class _SqliteBackedPrismaDB: self._connection.executemany('INSERT INTO "LiteLLM_SpendLogs" VALUES (?, ?, ?, ?, ?, ?)', rows) async def query_raw(self, query: str, *args: object) -> list[dict[str, object]]: - sqlite_query: Final = query.replace("= ANY($1::text[])", "IN (SELECT value FROM json_each(?))").replace("$1", "?") + sqlite_query: Final = query.replace("= ANY($1::text[])", "IN (SELECT value FROM json_each(?))").replace( + "$1", "?" + ) sqlite_args: Final = tuple(json.dumps(arg) if isinstance(arg, list) else arg for arg in args) return [dict(row) for row in self._connection.execute(sqlite_query, sqlite_args)] @@ -955,3 +936,134 @@ async def test_session_lookup_follows_the_owner_across_keys_and_keeps_ownerless_ assert await _history_request_ids(sqlite_db, "alice-new-key") == ["alice-old-key", "alice-new-key"] assert await _history_request_ids(sqlite_db, "ownerless-2") == ["ownerless-2"] + + +def _text(content: object) -> object: + if isinstance(content, list): + return "".join(part["text"] for part in content) + return content + + +def _responses_spend_log(request_id: str, session_id: str, prompt: str, output: list[dict[str, object]]) -> dict: + return { + "request_id": request_id, + "call_type": "aresponses", + "session_id": session_id, + "proxy_server_request": {"input": [{"role": "user", "content": prompt}], "model": "claude-bridge"}, + "response": {"id": request_id, "object": "response", "status": "completed", "output": output}, + } + + +@pytest.mark.asyncio +async def test_message_history_reads_responses_shaped_rows_next_to_chat_shaped_rows(): + session_id = "0f6f2d8e-6c8a-4b55-9a39-3d0b1f3c9e11" + fake_prisma_client = _FakePrismaClient( + results=[ + [ + _spend_log("chatcmpl-3333", session_id, "My favorite color is chartreuse.", "Got it."), + _responses_spend_log( + "resp_4444", + session_id, + "What is the weather in Lisbon?", + [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Let me check.", "annotations": []}], + }, + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_weather", + "name": "get_weather", + "arguments": '{"city": "Lisbon"}', + }, + ], + ), + ] + ] + ) + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id("resp_4444") + + messages: Final = result["messages"] + assert [(message.get("role"), _text(message.get("content"))) for message in messages[:4]] == [ + ("user", "My favorite color is chartreuse."), + ("assistant", "Got it."), + ("user", "What is the weather in Lisbon?"), + ("assistant", "Let me check."), + ] + tool_calls: Final = [call for message in messages[3:] for call in message.get("tool_calls") or []] + assert [(call["id"], call["function"]["name"], call["function"]["arguments"]) for call in tool_calls] == [ + ("call_weather", "get_weather", '{"city": "Lisbon"}') + ] + + +_WEB_SEARCH_RESULT: Final = { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": [{"type": "web_search_result", "url": "https://x.com", "title": "X", "encrypted_content": "EC"}], +} + + +def _bridged_anthropic_web_search_turn() -> dict[str, object]: + """A Claude turn served through the Responses to chat bridge, stored as its Responses object.""" + return { + "id": "resp_1", + "object": "response", + "output": [ + { + "id": "ws_srvtoolu_1", + "type": "web_search_call", + "status": "completed", + "action": { + "type": "search", + "query": "weather sf", + "sources": [{"type": "url", "url": "https://x.com"}], + }, + }, + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "It is sunny.", "annotations": []}], + }, + ], + "provider_specific_fields": {"web_search_results": [_WEB_SEARCH_RESULT]}, + } + + +def test_message_history_replays_a_bridged_anthropic_web_search_as_server_tool_blocks(): + from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt + + history: Final = ResponsesSessionHandler.extend_chat_completion_message_with_spend_log_payload( + spend_log={"request_id": "resp_1", "response": _bridged_anthropic_web_search_turn()}, + proxy_server_request_dict={"input": "weather in sf?"}, + chat_completion_message_history=[], + ) + anthropic_messages: Final = anthropic_messages_pt( + messages=history, model="claude-sonnet-4-5", llm_provider="anthropic" + ) + + assert [message["role"] for message in anthropic_messages] == ["user", "assistant"] + assistant_content: Final = anthropic_messages[1]["content"] + assert [block["type"] for block in assistant_content] == ["text", "server_tool_use", "web_search_tool_result"] + assert assistant_content[1]["id"] == "srvtoolu_1" + assert assistant_content[1]["input"] == {"query": "weather sf"} + assert assistant_content[2]["content"] == _WEB_SEARCH_RESULT["content"] + + +def test_message_history_keeps_a_web_search_without_a_stored_result_as_text(): + turn: Final = {**_bridged_anthropic_web_search_turn(), "provider_specific_fields": None} + + history: Final = ResponsesSessionHandler.extend_chat_completion_message_with_spend_log_payload( + spend_log={"request_id": "resp_1", "response": turn}, + proxy_server_request_dict=None, + chat_completion_message_history=[], + ) + + assert str(_text_of(history[0]["content"])).startswith("Hosted web search: ")