diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index fdba9f82f10..2f0cbc0d474 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -6,6 +6,7 @@ Route to previously cached model id, if valid import secrets from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from itertools import chain from types import MappingProxyType from typing import Final, NamedTuple, Protocol, TypedDict, cast @@ -29,6 +30,7 @@ from litellm.responses.litellm_completion_transformation.transformation import ( RESPONSES_TOOL_TYPES_WITHOUT_CHAT_FORM, RESPONSES_WEB_SEARCH_TOOL_TYPES, LiteLLMCompletionResponsesConfig, + ResponsesToolChatForm, ResponseTools, ) from litellm.responses.utils import ResponsesAPIRequestUtils @@ -82,6 +84,20 @@ def new_responses_pins_in_flight() -> ResponsesPinsInFlight: ) +def _key_chat_tool(source: Mapping[str, object], chat_tool: Mapping[str, object]) -> Mapping[str, object]: + function: Final = as_str_mapping(chat_tool.get("function")) + if source.get("type") != "function" or function is None: + return chat_tool + return {**chat_tool, "function": {key: source[key] for key in function if key in source}} + + +def _key_chat_tools( + sources: Sequence[Mapping[str, object]], forms: Sequence[ResponsesToolChatForm] +) -> Iterator[Mapping[str, object]]: + for source, form in zip(sources, forms, strict=True): + yield from (_key_chat_tool(source, chat_tool) for chat_tool in form.chat_tools) + + def _responses_prompt(request_data: Mapping[str, object]) -> _ResponsesPrompt | None: request_input: Final = request_data.get("input") request_tools: Final = request_data.get("tools") @@ -107,13 +123,12 @@ def _responses_prompt(request_data: Mapping[str, object]) -> _ResponsesPrompt | item_cache_control=True, only_since_compaction=True, ) - tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( - convertible_tools - ) + forms: Final = LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(convertible_tools) except Exception as error: # noqa: BLE001 # the provider judges the input, a routing preference never does verbose_logger.debug("Prompt Caching: skipping an input with no chat form: %s", type(error).__name__) return None - key_tools: Final = [*tools, *tools_kept_as_sent] + tools: Final = list(chain.from_iterable(form.chat_tools for form in forms)) + key_tools: Final = [*_key_chat_tools(convertible_tools or (), forms), *tools_kept_as_sent] return _ResponsesPrompt( messages=cast(list[AllMessageValues], messages), # cast-ok: the bridge's chat messages, read as plain JSON tools=cast(list[AllToolParamValues], tools) or None, # cast-ok: the bridge's chat tools, read as plain JSON @@ -465,10 +480,12 @@ class PromptCachingDeploymentCheck(CustomLogger): def _take_responses_pin(self, kwargs: _PromptCachingSuccessEventKwargs) -> object: litellm_params: Final = kwargs.get("litellm_params") or {} - metadata: Final = as_str_mapping(litellm_params.get("litellm_metadata")) or {} - pin_token: Final = metadata.get(_RESPONSES_PIN_METADATA_KEY) + metadata: Final = litellm_params.get("litellm_metadata") + pin_token: Final = (as_str_mapping(metadata) or {}).get(_RESPONSES_PIN_METADATA_KEY) if not isinstance(pin_token, str): return None + if isinstance(metadata, dict): + metadata.pop(_RESPONSES_PIN_METADATA_KEY, None) # pyright: ignore[reportUnknownMemberType] # the caller's metadata dict recorded_pin: Final = self.responses_pins_in_flight.get_cache(pin_token) self.responses_pins_in_flight.delete_cache(pin_token) return recorded_pin diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 34cb1811488..568cc4a618c 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -2096,6 +2096,35 @@ async def test_responses_with_another_tool_set_are_a_different_prefix(): assert len(pinned) == 2 +@pytest.mark.asyncio +async def test_chat_and_responses_requests_sending_the_same_tool_share_one_prefix(): + described_tool: Final = {**_RESPONSES_TOOL, "description": "Get the time"} + chat_tool: Final = { + "type": "function", + "function": {key: value for key, value in described_tool.items() if key != "type"}, + } + router: Final = _responses_router(RESPONSES_BRIDGED_MODEL) + await router.acompletion( + model=MODEL_GROUP_ALIAS, + messages=[ + {"role": "system", "content": [{**_marked_input_text(LONG_PROMPT), "type": "text"}]}, + {"role": "user", "content": "hello"}, + ], + tools=[chat_tool], + mock_response="ok", + ) + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + chat_pin_keys: Final = set(router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + router.prompt_caching_pin_cache.in_memory_cache.flush_cache() + + await router.aresponses( + model=MODEL_GROUP_ALIAS, input=_marked_responses_input(), tools=[described_tool], mock_response="ok" + ) + + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + assert set(router.prompt_caching_pin_cache.in_memory_cache.cache_dict) == chat_pin_keys + + @pytest.mark.asyncio async def test_responses_marked_tool_set_long_enough_to_cache_pins_on_its_own(): marked_tool = {**_RESPONSES_TOOL, "description": LONG_PROMPT, "cache_control": {"type": "ephemeral"}} @@ -2413,6 +2442,22 @@ async def test_logger_serializing_the_request_metadata_never_sees_the_prompt_cac assert not [pin_key for pin_key in pin_keys if pin_key.split(":")[1] in logged_metadata] +@pytest.mark.asyncio +async def test_responses_call_hands_back_the_callers_metadata_without_the_pin_token(): + router: Final = _responses_router(RESPONSES_BRIDGED_MODEL) + caller_metadata: Final[dict[str, object]] = {"tags": ["team-a"]} + + await router.aresponses( + model=MODEL_GROUP_ALIAS, + input=_marked_responses_input(), + litellm_metadata=caller_metadata, + mock_response="ok", + ) + + assert await _eventually(lambda: router.prompt_caching_pin_cache.in_memory_cache.cache_dict) + assert "responses_prompt_cache_pin" not in caller_metadata + + @pytest.mark.asyncio @pytest.mark.parametrize( "failed_in_request,expected_deployments",