diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index cacf9d27d4b..ab64578cdd4 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -24,7 +24,13 @@ from litellm.types.llms.openai import ( BaseLiteLLMOpenAIResponseObject, ResponsesAPIResponse, ) -from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, CallTypesLiteral, LLMResponseTypes, SpecialEnums +from litellm.types.utils import ( + ADDRESSED_COMPARISON_RESPONSE_ID_FIELD, + ADDRESSED_RESPONSE_ID_FIELD, + CallTypesLiteral, + LLMResponseTypes, + SpecialEnums, +) if TYPE_CHECKING: from litellm.caching.caching import DualCache @@ -155,37 +161,29 @@ class ResponsesIDSecurity(CustomLogger): data: Mapping[str, object], user_api_key_dict: "UserAPIKeyAuth", ) -> Mapping[str, object]: - options: Final = self._authorized_prompt_cache_options(data.get("prompt_cache_options"), user_api_key_dict) extra_body: Final = as_str_mapping(data.get("extra_body")) - extra_body_options: Final = ( - None - if extra_body is None - else self._authorized_prompt_cache_options(extra_body.get("prompt_cache_options"), user_api_key_dict) + overriding_extra_body: Final = ( + extra_body if extra_body is not None and "prompt_cache_options" in extra_body else None ) - options_field: Final[Mapping[str, object]] = {} if options is None else {"prompt_cache_options": options} - extra_body_field: Final[Mapping[str, object]] = ( - {} - if extra_body is None or extra_body_options is None - else {"extra_body": {**extra_body, "prompt_cache_options": extra_body_options}} + options: Final = as_str_mapping( + data.get("prompt_cache_options") + if overriding_extra_body is None + else overriding_extra_body.get("prompt_cache_options") ) - return {**options_field, **extra_body_field} - - def _authorized_prompt_cache_options( - self, - prompt_cache_options: object, - user_api_key_dict: "UserAPIKeyAuth", - ) -> Mapping[str, object] | None: - options: Final = as_str_mapping(prompt_cache_options) - comparison_response_id: Final = None if options is None else options.get("comparison_response_id") - if ( - options is None - or not isinstance(comparison_response_id, str) - or not self._is_encrypted_response_id(comparison_response_id) - ): - return None - original_response_id, user_id, team_id = self._decrypt_response_id(comparison_response_id) + retained_id: Final = data.get(ADDRESSED_COMPARISON_RESPONSE_ID_FIELD) + requested_id: Final = None if options is None else options.get("comparison_response_id") + addressed_id: Final = retained_id if isinstance(retained_id, str) and retained_id else requested_id + if options is None or not isinstance(addressed_id, str) or not self._is_encrypted_response_id(addressed_id): + return {} + original_response_id, user_id, team_id = self._decrypt_response_id(addressed_id) self.check_user_access_to_response_id(user_id, team_id, user_api_key_dict) - return {**options, "comparison_response_id": original_response_id} + authorized_options: Final = {**options, "comparison_response_id": original_response_id} + authorized_field: Final[Mapping[str, object]] = ( + {"prompt_cache_options": authorized_options} + if overriding_extra_body is None + else {"extra_body": {**overriding_extra_body, "prompt_cache_options": authorized_options}} + ) + return {**authorized_field, ADDRESSED_COMPARISON_RESPONSE_ID_FIELD: addressed_id} def authorize_response_id( self, diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 9f1b35fdce3..b6b17a6b02e 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -55,6 +55,7 @@ RoutingStrategyName: TypeAlias = Literal[ TRUSTED_CALLBACK_VARS_FIELD: Final = "litellm_trusted_callback_vars" ADDRESSED_RESPONSE_ID_FIELD: Final = "_litellm_addressed_response_id" +ADDRESSED_COMPARISON_RESPONSE_ID_FIELD: Final = "_litellm_addressed_comparison_response_id" STREAM_USAGE_INJECTED_FIELD: Final = "_litellm_stream_usage_injected" WIRE_NAME: Final = "wire_name" @@ -406,6 +407,9 @@ class ProxyRequestState: secret_fields: "SecretFields | None" = None trusted_callback_vars: Mapping[str, str] | None = field(default=None, metadata=wire(TRUSTED_CALLBACK_VARS_FIELD)) addressed_response_id: str | None = field(default=None, metadata=wire(ADDRESSED_RESPONSE_ID_FIELD)) + addressed_comparison_response_id: str | None = field( + default=None, metadata=wire(ADDRESSED_COMPARISON_RESPONSE_ID_FIELD) + ) strip_stream_usage: bool | None = field(default=None, metadata=wire("_litellm_strip_stream_usage")) stream_usage_injected: bool | None = field(default=None, metadata=wire(STREAM_USAGE_INJECTED_FIELD)) client_side_timeout: bool | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 0c3440e88f7..a301a5fb2af 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4334,6 +4334,7 @@ bedrock_batch_litellm_params: Final = BEDROCK_BATCH_KWARG_NAMES TRUSTED_CALLBACK_VARS_FIELD: Final = _litellm_params.TRUSTED_CALLBACK_VARS_FIELD ADDRESSED_RESPONSE_ID_FIELD: Final = _litellm_params.ADDRESSED_RESPONSE_ID_FIELD +ADDRESSED_COMPARISON_RESPONSE_ID_FIELD: Final = _litellm_params.ADDRESSED_COMPARISON_RESPONSE_ID_FIELD anthropic_wif_litellm_params: Final = tuple(sorted(ANTHROPIC_WIF_KWARGS_KEYS)) openai_wif_litellm_params: Final = tuple(sorted(OPENAI_WIF_KWARGS_KEYS))