diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5323f692b80d..a5262505983e 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -354,9 +354,9 @@ def __init__( ) self.function_id = function_id self.streaming_chunks: List[Any] = [] # for generating complete stream response - self.sync_streaming_chunks: List[ - Any - ] = [] # for generating complete stream response + self.sync_streaming_chunks: List[Any] = ( + [] + ) # for generating complete stream response self.log_raw_request_response = log_raw_request_response # Initialize dynamic callbacks @@ -801,9 +801,9 @@ def _auto_detect_prompt_management_logger( prompt_spec=prompt_spec, dynamic_callback_params=dynamic_callback_params, ): - self.model_call_details[ - "prompt_integration" - ] = logger.__class__.__name__ + self.model_call_details["prompt_integration"] = ( + logger.__class__.__name__ + ) return logger except Exception: # If check fails, continue to next logger @@ -871,9 +871,9 @@ def get_custom_logger_for_prompt_management( if anthropic_cache_control_logger := AnthropicCacheControlHook.get_custom_logger_for_anthropic_cache_control_hook( non_default_params ): - self.model_call_details[ - "prompt_integration" - ] = anthropic_cache_control_logger.__class__.__name__ + self.model_call_details["prompt_integration"] = ( + anthropic_cache_control_logger.__class__.__name__ + ) return anthropic_cache_control_logger ######################################################### @@ -885,9 +885,9 @@ def get_custom_logger_for_prompt_management( internal_usage_cache=None, llm_router=None, ) - self.model_call_details[ - "prompt_integration" - ] = vector_store_custom_logger.__class__.__name__ + self.model_call_details["prompt_integration"] = ( + vector_store_custom_logger.__class__.__name__ + ) # Add to global callbacks so post-call hooks are invoked if ( vector_store_custom_logger @@ -947,9 +947,9 @@ def _pre_call(self, input, api_key, model=None, additional_args={}): model ): # if model name was changes pre-call, overwrite the initial model call name with the new one self.model_call_details["model"] = model - self.model_call_details["litellm_params"][ - "api_base" - ] = self._get_masked_api_base(additional_args.get("api_base", "")) + self.model_call_details["litellm_params"]["api_base"] = ( + self._get_masked_api_base(additional_args.get("api_base", "")) + ) def pre_call(self, input, api_key, model=None, additional_args={}): # noqa: PLR0915 # Log the exact input to the LLM API @@ -978,9 +978,7 @@ def pre_call(self, input, api_key, model=None, additional_args={}): # noqa: PLR try: # [Non-blocking Extra Debug Information in metadata] if turn_off_message_logging is True: - _metadata[ - "raw_request" - ] = "redacted by litellm. \ + _metadata["raw_request"] = "redacted by litellm. \ 'litellm.turn_off_message_logging=True'" else: curl_command = self._get_request_curl_command( @@ -992,35 +990,31 @@ def pre_call(self, input, api_key, model=None, additional_args={}): # noqa: PLR _metadata["raw_request"] = str(curl_command) # split up, so it's easier to parse in the UI - self.model_call_details[ - "raw_request_typed_dict" - ] = RawRequestTypedDict( - raw_request_api_base=str( - additional_args.get("api_base") or "" - ), - raw_request_body=self._get_raw_request_body( - additional_args.get("complete_input_dict", {}) - ), - # NOTE: setting ignore_sensitive_headers to True will cause - # the Authorization header to be leaked when calls to the health - # endpoint are made and fail. - raw_request_headers=self._get_masked_headers( - additional_args.get("headers", {}) or {}, - ), - error=None, + self.model_call_details["raw_request_typed_dict"] = ( + RawRequestTypedDict( + raw_request_api_base=str( + additional_args.get("api_base") or "" + ), + raw_request_body=self._get_raw_request_body( + additional_args.get("complete_input_dict", {}) + ), + # NOTE: setting ignore_sensitive_headers to True will cause + # the Authorization header to be leaked when calls to the health + # endpoint are made and fail. + raw_request_headers=self._get_masked_headers( + additional_args.get("headers", {}) or {}, + ), + error=None, + ) ) except Exception as e: - self.model_call_details[ - "raw_request_typed_dict" - ] = RawRequestTypedDict( - error=str(e), - ) - _metadata[ - "raw_request" - ] = "Unable to Log \ - raw request: {}".format( - str(e) + self.model_call_details["raw_request_typed_dict"] = ( + RawRequestTypedDict( + error=str(e), + ) ) + _metadata["raw_request"] = "Unable to Log \ + raw request: {}".format(str(e)) if getattr(self, "logger_fn", None) and callable(self.logger_fn): try: self.logger_fn( @@ -1320,13 +1314,13 @@ async def async_post_mcp_tool_call_hook( for callback in callbacks: try: if isinstance(callback, CustomLogger): - response: Optional[ - MCPPostCallResponseObject - ] = await callback.async_post_mcp_tool_call_hook( - kwargs=kwargs, - response_obj=post_mcp_tool_call_response_obj, - start_time=start_time, - end_time=end_time, + response: Optional[MCPPostCallResponseObject] = ( + await callback.async_post_mcp_tool_call_hook( + kwargs=kwargs, + response_obj=post_mcp_tool_call_response_obj, + start_time=start_time, + end_time=end_time, + ) ) ###################################################################### # if any of the callbacks modify the response, use the modified response @@ -1527,9 +1521,9 @@ def _response_cost_calculator( verbose_logger.debug( f"response_cost_failure_debug_information: {debug_info}" ) - self.model_call_details[ - "response_cost_failure_debug_information" - ] = debug_info + self.model_call_details["response_cost_failure_debug_information"] = ( + debug_info + ) return None try: @@ -1555,9 +1549,9 @@ def _response_cost_calculator( verbose_logger.debug( f"response_cost_failure_debug_information: {debug_info}" ) - self.model_call_details[ - "response_cost_failure_debug_information" - ] = debug_info + self.model_call_details["response_cost_failure_debug_information"] = ( + debug_info + ) return None @@ -1706,9 +1700,9 @@ def _merge_hidden_params_from_response_into_metadata( self.model_call_details["litellm_params"].setdefault("metadata", {}) if self.model_call_details["litellm_params"]["metadata"] is None: self.model_call_details["litellm_params"]["metadata"] = {} - self.model_call_details["litellm_params"]["metadata"][ - "hidden_params" - ] = getattr(logging_result, "_hidden_params", {}) + self.model_call_details["litellm_params"]["metadata"]["hidden_params"] = ( + getattr(logging_result, "_hidden_params", {}) + ) def _process_hidden_params_and_response_cost( self, @@ -1737,9 +1731,9 @@ def _process_hidden_params_and_response_cost( result=logging_result ) - self.model_call_details[ - "standard_logging_object" - ] = self._build_standard_logging_payload(logging_result, start_time, end_time) + self.model_call_details["standard_logging_object"] = ( + self._build_standard_logging_payload(logging_result, start_time, end_time) + ) if ( standard_logging_payload := self.model_call_details.get( @@ -1817,9 +1811,9 @@ def _success_handler_helper_fn( end_time = datetime.datetime.now() if self.completion_start_time is None: self.completion_start_time = end_time - self.model_call_details[ - "completion_start_time" - ] = self.completion_start_time + self.model_call_details["completion_start_time"] = ( + self.completion_start_time + ) self.model_call_details["log_event_type"] = "successful_api_call" self.model_call_details["end_time"] = end_time @@ -1856,10 +1850,10 @@ def _success_handler_helper_fn( end_time=end_time, ) elif isinstance(result, dict) or isinstance(result, list): - self.model_call_details[ - "standard_logging_object" - ] = self._build_standard_logging_payload( - result, start_time, end_time + self.model_call_details["standard_logging_object"] = ( + self._build_standard_logging_payload( + result, start_time, end_time + ) ) if ( standard_logging_payload := self.model_call_details.get( @@ -1868,9 +1862,9 @@ def _success_handler_helper_fn( ) is not None: emit_standard_logging_payload(standard_logging_payload) elif standard_logging_object is not None: - self.model_call_details[ - "standard_logging_object" - ] = standard_logging_object + self.model_call_details["standard_logging_object"] = ( + standard_logging_object + ) else: self.model_call_details["response_cost"] = None @@ -2028,20 +2022,20 @@ def success_handler( # noqa: PLR0915 verbose_logger.debug( "Logging Details LiteLLM-Success Call streaming complete" ) - self.model_call_details[ - "complete_streaming_response" - ] = complete_streaming_response - self.model_call_details[ - "response_cost" - ] = self._response_cost_calculator(result=complete_streaming_response) + self.model_call_details["complete_streaming_response"] = ( + complete_streaming_response + ) + self.model_call_details["response_cost"] = ( + self._response_cost_calculator(result=complete_streaming_response) + ) self._merge_hidden_params_from_response_into_metadata( complete_streaming_response ) ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details[ - "standard_logging_object" - ] = self._build_standard_logging_payload( - complete_streaming_response, start_time, end_time + self.model_call_details["standard_logging_object"] = ( + self._build_standard_logging_payload( + complete_streaming_response, start_time, end_time + ) ) if ( standard_logging_payload := self.model_call_details.get( @@ -2375,10 +2369,10 @@ def success_handler( # noqa: PLR0915 ) else: if self.stream and complete_streaming_response: - self.model_call_details[ - "complete_response" - ] = self.model_call_details.get( - "complete_streaming_response", {} + self.model_call_details["complete_response"] = ( + self.model_call_details.get( + "complete_streaming_response", {} + ) ) result = self.model_call_details["complete_response"] openMeterLogger.log_success_event( @@ -2402,10 +2396,10 @@ def success_handler( # noqa: PLR0915 ) else: if self.stream and complete_streaming_response: - self.model_call_details[ - "complete_response" - ] = self.model_call_details.get( - "complete_streaming_response", {} + self.model_call_details["complete_response"] = ( + self.model_call_details.get( + "complete_streaming_response", {} + ) ) result = self.model_call_details["complete_response"] @@ -2544,9 +2538,9 @@ async def async_success_handler( # noqa: PLR0915 if complete_streaming_response is not None: print_verbose("Async success callbacks: Got a complete streaming response") - self.model_call_details[ - "async_complete_streaming_response" - ] = complete_streaming_response + self.model_call_details["async_complete_streaming_response"] = ( + complete_streaming_response + ) try: if self.model_call_details.get("cache_hit", False) is True: @@ -2557,10 +2551,10 @@ async def async_success_handler( # noqa: PLR0915 model_call_details=self.model_call_details ) # base_model defaults to None if not set on model_info - self.model_call_details[ - "response_cost" - ] = self._response_cost_calculator( - result=complete_streaming_response + self.model_call_details["response_cost"] = ( + self._response_cost_calculator( + result=complete_streaming_response + ) ) verbose_logger.debug( @@ -2577,10 +2571,10 @@ async def async_success_handler( # noqa: PLR0915 ) ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details[ - "standard_logging_object" - ] = self._build_standard_logging_payload( - complete_streaming_response, start_time, end_time + self.model_call_details["standard_logging_object"] = ( + self._build_standard_logging_payload( + complete_streaming_response, start_time, end_time + ) ) # print standard logging payload @@ -2607,9 +2601,9 @@ async def async_success_handler( # noqa: PLR0915 # _success_handler_helper_fn if self.model_call_details.get("standard_logging_object") is None: ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details[ - "standard_logging_object" - ] = self._build_standard_logging_payload(result, start_time, end_time) + self.model_call_details["standard_logging_object"] = ( + self._build_standard_logging_payload(result, start_time, end_time) + ) # print standard logging payload if ( @@ -2852,18 +2846,18 @@ def _failure_handler_helper_fn( ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details[ - "standard_logging_object" - ] = get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj={}, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="failure", - error_str=str(exception), - original_exception=exception, - standard_built_in_tools_params=self.standard_built_in_tools_params, + self.model_call_details["standard_logging_object"] = ( + get_standard_logging_object_payload( + kwargs=self.model_call_details, + init_response_obj={}, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="failure", + error_str=str(exception), + original_exception=exception, + standard_built_in_tools_params=self.standard_built_in_tools_params, + ) ) return start_time, end_time @@ -3831,9 +3825,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 service_name=arize_config.project_name, ) - os.environ[ - "OTEL_EXPORTER_OTLP_TRACES_HEADERS" - ] = f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}" + os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( + f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}" + ) for callback in _in_memory_loggers: if ( isinstance(callback, ArizeLogger) @@ -3859,13 +3853,13 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "") # Add openinference.project.name attribute if existing_attrs: - os.environ[ - "OTEL_RESOURCE_ATTRIBUTES" - ] = f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}" + os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( + f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}" + ) else: - os.environ[ - "OTEL_RESOURCE_ATTRIBUTES" - ] = f"openinference.project.name={arize_phoenix_config.project_name}" + os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( + f"openinference.project.name={arize_phoenix_config.project_name}" + ) # Set Phoenix project name from environment variable phoenix_project_name = os.environ.get("PHOENIX_PROJECT_NAME", None) @@ -3873,19 +3867,19 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "") # Add openinference.project.name attribute if existing_attrs: - os.environ[ - "OTEL_RESOURCE_ATTRIBUTES" - ] = f"{existing_attrs},openinference.project.name={phoenix_project_name}" + os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( + f"{existing_attrs},openinference.project.name={phoenix_project_name}" + ) else: - os.environ[ - "OTEL_RESOURCE_ATTRIBUTES" - ] = f"openinference.project.name={phoenix_project_name}" + os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( + f"openinference.project.name={phoenix_project_name}" + ) # auth can be disabled on local deployments of arize phoenix if arize_phoenix_config.otlp_auth_headers is not None: - os.environ[ - "OTEL_EXPORTER_OTLP_TRACES_HEADERS" - ] = arize_phoenix_config.otlp_auth_headers + os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( + arize_phoenix_config.otlp_auth_headers + ) for callback in _in_memory_loggers: if ( @@ -4072,9 +4066,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 exporter="otlp_http", endpoint="https://langtrace.ai/api/trace", ) - os.environ[ - "OTEL_EXPORTER_OTLP_TRACES_HEADERS" - ] = f"api_key={os.getenv('LANGTRACE_API_KEY')}" + os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( + f"api_key={os.getenv('LANGTRACE_API_KEY')}" + ) for callback in _in_memory_loggers: if ( isinstance(callback, OpenTelemetry) @@ -4998,10 +4992,10 @@ def get_hidden_params( for key in StandardLoggingHiddenParams.__annotations__.keys(): if key in hidden_params: if key == "additional_headers": - clean_hidden_params[ - "additional_headers" - ] = StandardLoggingPayloadSetup.get_additional_headers( - hidden_params[key] + clean_hidden_params["additional_headers"] = ( + StandardLoggingPayloadSetup.get_additional_headers( + hidden_params[key] + ) ) else: clean_hidden_params[key] = hidden_params[key] # type: ignore @@ -5640,9 +5634,9 @@ def scrub_sensitive_keys_in_metadata(litellm_params: Optional[dict]): ): for k, v in metadata["user_api_key_metadata"].items(): if k == "logging": # prevent logging user logging keys - cleaned_user_api_key_metadata[ - k - ] = "scrubbed_by_litellm_for_sensitive_keys" + cleaned_user_api_key_metadata[k] = ( + "scrubbed_by_litellm_for_sensitive_keys" + ) else: cleaned_user_api_key_metadata[k] = v diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py index fbd1da749c27..91e7e9d55504 100644 --- a/litellm/llms/a2a/chat/guardrail_translation/handler.py +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -111,6 +111,7 @@ async def process_output_response( guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + request_data: Optional[dict] = None, ) -> Any: """ Process A2A output response by applying guardrails to text content. @@ -166,13 +167,21 @@ async def process_output_response( return response # Step 2: Apply guardrail to all texts in batch - # Create a request_data dict with response info and user API key metadata - request_data: dict = {"response": response_dict} + # Use the real request_data if provided (proxy path), otherwise + # create a standalone dict (SDK / direct-call path). + if request_data is None: + request_data = {"response": response_dict} + else: + if "response" not in request_data: + request_data["response"] = response_dict # Add user API key metadata with prefixed keys - user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + if "litellm_metadata" not in request_data: + user_metadata = self.transform_user_api_key_dict_to_metadata( + user_api_key_dict + ) + if user_metadata: + request_data["litellm_metadata"] = user_metadata inputs = GenericGuardrailAPIInputs(texts=texts_to_check) @@ -213,6 +222,7 @@ async def process_output_streaming_response( guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + request_data: Optional[dict] = None, ) -> List[Any]: """ Process A2A streaming output by applying guardrails to accumulated text. @@ -258,10 +268,18 @@ async def process_output_streaming_response( if not combined_text: return responses_so_far - request_data: dict = {"responses_so_far": responses_so_far} - user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + if request_data is None: + request_data = {"responses_so_far": responses_so_far} + else: + if "responses_so_far" not in request_data: + request_data["responses_so_far"] = responses_so_far + + if "litellm_metadata" not in request_data: + user_metadata = self.transform_user_api_key_dict_to_metadata( + user_api_key_dict + ) + if user_metadata: + request_data["litellm_metadata"] = user_metadata inputs = GenericGuardrailAPIInputs(texts=[combined_text]) guardrailed_inputs = await guardrail_to_apply.apply_guardrail( diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 5372757cbb67..b6139a2f84d6 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -252,6 +252,7 @@ async def process_output_response( guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional[Any] = None, user_api_key_dict: Optional[Any] = None, + request_data: Optional[dict] = None, ) -> Any: """ Process output response by applying guardrails to text content and tool calls. @@ -323,15 +324,21 @@ async def process_output_response( # Step 2: Apply guardrail to all texts in batch if texts_to_check or tool_calls_to_check: - # Create a request_data dict with response info and user API key metadata - request_data: dict = {"response": response} + # Use the real request_data if provided (proxy path), otherwise + # create a standalone dict (SDK / direct-call path). + if request_data is None: + request_data = {"response": response} + else: + if "response" not in request_data: + request_data["response"] = response # Add user API key metadata with prefixed keys - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + if "litellm_metadata" not in request_data: + user_metadata = self.transform_user_api_key_dict_to_metadata( + user_api_key_dict + ) + if user_metadata: + request_data["litellm_metadata"] = user_metadata inputs = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: @@ -375,6 +382,7 @@ async def process_output_streaming_response( guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional[Any] = None, user_api_key_dict: Optional[Any] = None, + request_data: Optional[dict] = None, ) -> List[Any]: """ Process output streaming response by applying guardrails to text content. @@ -413,7 +421,7 @@ async def process_output_streaming_response( _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( # allow rejecting the response, if invalid inputs=guardrail_inputs, - request_data={}, + request_data=request_data if request_data is not None else {}, input_type="response", logging_obj=litellm_logging_obj, ) @@ -426,7 +434,7 @@ async def process_output_streaming_response( string_so_far = self.get_streaming_string_so_far(responses_so_far) _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( # allow rejecting the response, if invalid inputs={"texts": [string_so_far]}, - request_data={}, + request_data=request_data if request_data is not None else {}, input_type="response", logging_obj=litellm_logging_obj, ) diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index a7982cb606e3..e1da0dfa29e8 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -73,6 +73,7 @@ async def process_output_response( guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + request_data: Optional[dict] = None, ) -> Any: """ Process output response with guardrails. @@ -91,6 +92,7 @@ async def process_output_streaming_response( guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + request_data: Optional[dict] = None, ) -> Any: """ Process output streaming response with guardrails. diff --git a/litellm/llms/cohere/rerank/guardrail_translation/handler.py b/litellm/llms/cohere/rerank/guardrail_translation/handler.py index b8133c59f7d0..e9a5823d2b82 100644 --- a/litellm/llms/cohere/rerank/guardrail_translation/handler.py +++ b/litellm/llms/cohere/rerank/guardrail_translation/handler.py @@ -83,6 +83,7 @@ async def process_output_response( guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional[Any] = None, user_api_key_dict: Optional[Any] = None, + request_data: Optional[dict] = None, ) -> Any: """ Process output response - not applicable for rerank. diff --git a/litellm/llms/mistral/ocr/guardrail_translation/handler.py b/litellm/llms/mistral/ocr/guardrail_translation/handler.py index 697bd2daa3da..7d3797a1dbe6 100644 --- a/litellm/llms/mistral/ocr/guardrail_translation/handler.py +++ b/litellm/llms/mistral/ocr/guardrail_translation/handler.py @@ -91,6 +91,7 @@ async def process_output_response( guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional[Any] = None, user_api_key_dict: Optional[Any] = None, + request_data: Optional[dict] = None, ) -> Any: """ Process OCR output by applying guardrails to extracted page text. @@ -127,14 +128,27 @@ async def process_output_response( if model: inputs["model"] = model + # Use the real request_data if provided (proxy path), otherwise + # create a standalone dict (SDK / direct-call path). + if request_data is None: + request_data = {} + # Add user metadata if available if user_api_key_dict is not None: - metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) - inputs.update(metadata) # type: ignore + user_metadata = self.transform_user_api_key_dict_to_metadata( + user_api_key_dict + ) + if user_metadata: + # Preserve original behavior: inject metadata into inputs for + # third-party guardrail providers that read it from there + inputs.update(user_metadata) # type: ignore + # Also store in request_data for the logging pipeline + if "litellm_metadata" not in request_data: + request_data["litellm_metadata"] = user_metadata guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, - request_data={}, + request_data=request_data, input_type="response", logging_obj=litellm_logging_obj, ) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index bab4c3b5eb75..0ce9d57a65f7 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -86,9 +86,9 @@ async def process_input_messages( if tool_calls_to_check: inputs["tool_calls"] = tool_calls_to_check # type: ignore if messages: - inputs[ - "structured_messages" - ] = messages # pass the openai /chat/completions messages to the guardrail, as-is + inputs["structured_messages"] = ( + messages # pass the openai /chat/completions messages to the guardrail, as-is + ) # Pass tools (function definitions) to the guardrail tools = data.get("tools") if tools: @@ -260,6 +260,7 @@ async def process_output_response( guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional[Any] = None, user_api_key_dict: Optional[Any] = None, + request_data: Optional[dict] = None, ) -> Any: """ Process output response by applying guardrails to text content. @@ -308,15 +309,21 @@ async def process_output_response( # Step 2: Apply guardrail to all texts and tool calls in batch if texts_to_check or tool_calls_to_check: - # Create a request_data dict with response info and user API key metadata - request_data: dict = {"response": response} + # Use the real request_data if provided (proxy path), otherwise + # create a standalone dict (SDK / direct-call path). + if request_data is None: + request_data = {"response": response} + else: + if "response" not in request_data: + request_data["response"] = response # Add user API key metadata with prefixed keys - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + if "litellm_metadata" not in request_data: + user_metadata = self.transform_user_api_key_dict_to_metadata( + user_api_key_dict + ) + if user_metadata: + request_data["litellm_metadata"] = user_metadata inputs = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: @@ -364,6 +371,7 @@ async def process_output_streaming_response( guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional[Any] = None, user_api_key_dict: Optional[Any] = None, + request_data: Optional[dict] = None, ) -> List["ModelResponseStream"]: """ Process output streaming responses by applying guardrails to text content. @@ -402,6 +410,7 @@ async def process_output_streaming_response( guardrail_to_apply=guardrail_to_apply, litellm_logging_obj=litellm_logging_obj, user_api_key_dict=user_api_key_dict, + request_data=request_data, ) return responses_so_far @@ -436,15 +445,21 @@ async def process_output_streaming_response( # Step 3: Apply guardrail to all combined texts in batch if texts_to_check: - # Create a request_data dict with response info and user API key metadata - request_data: dict = {"responses": responses_so_far} + # Use the real request_data if provided (proxy path), otherwise + # create a standalone dict (SDK / direct-call path). + if request_data is None: + request_data = {"responses": responses_so_far} + else: + if "responses" not in request_data: + request_data["responses"] = responses_so_far # Add user API key metadata with prefixed keys - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + if "litellm_metadata" not in request_data: + user_metadata = self.transform_user_api_key_dict_to_metadata( + user_api_key_dict + ) + if user_metadata: + request_data["litellm_metadata"] = user_metadata inputs = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: diff --git a/litellm/llms/openai/completion/guardrail_translation/handler.py b/litellm/llms/openai/completion/guardrail_translation/handler.py index 1f8c6159da03..593ab0ed2e5a 100644 --- a/litellm/llms/openai/completion/guardrail_translation/handler.py +++ b/litellm/llms/openai/completion/guardrail_translation/handler.py @@ -125,6 +125,7 @@ async def process_output_response( guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional[Any] = None, user_api_key_dict: Optional[Any] = None, + request_data: Optional[dict] = None, ) -> Any: """ Process output response by applying guardrails to completion text. @@ -155,15 +156,21 @@ async def process_output_response( # Apply guardrails in batch if texts_to_check: - # Create a request_data dict with response info and user API key metadata - request_data: dict = {"response": response} + # Use the real request_data if provided (proxy path), otherwise + # create a standalone dict (SDK / direct-call path). + if request_data is None: + request_data = {"response": response} + else: + if "response" not in request_data: + request_data["response"] = response # Add user API key metadata with prefixed keys - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + if "litellm_metadata" not in request_data: + user_metadata = self.transform_user_api_key_dict_to_metadata( + user_api_key_dict + ) + if user_metadata: + request_data["litellm_metadata"] = user_metadata inputs = GenericGuardrailAPIInputs(texts=texts_to_check) # Include model information from the response if available diff --git a/litellm/llms/openai/embeddings/guardrail_translation/handler.py b/litellm/llms/openai/embeddings/guardrail_translation/handler.py index 7458020e109b..ff5021b8ce07 100644 --- a/litellm/llms/openai/embeddings/guardrail_translation/handler.py +++ b/litellm/llms/openai/embeddings/guardrail_translation/handler.py @@ -155,6 +155,7 @@ async def process_output_response( guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional[Any] = None, user_api_key_dict: Optional[Any] = None, + request_data: Optional[dict] = None, ) -> Any: """ Process output response - embeddings responses contain vectors, not text. diff --git a/litellm/llms/openai/image_generation/guardrail_translation/handler.py b/litellm/llms/openai/image_generation/guardrail_translation/handler.py index e6340ba4705a..76610088d0cd 100644 --- a/litellm/llms/openai/image_generation/guardrail_translation/handler.py +++ b/litellm/llms/openai/image_generation/guardrail_translation/handler.py @@ -87,6 +87,7 @@ async def process_output_response( guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional[Any] = None, user_api_key_dict: Optional[Any] = None, + request_data: Optional[dict] = None, ) -> Any: """ Process output response - typically not needed for image generation. diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 466e2e76f183..76f40eed71fe 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -347,6 +347,7 @@ async def process_output_response( guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional[Any] = None, user_api_key_dict: Optional[Any] = None, + request_data: Optional[dict] = None, ) -> Any: """ Process output response by applying guardrails to text content and tool calls. @@ -402,15 +403,21 @@ async def process_output_response( # Step 2: Apply guardrail to all texts in batch if texts_to_check or tool_calls_to_check: - # Create a request_data dict with response info and user API key metadata - request_data: dict = {"response": response} + # Use the real request_data if provided (proxy path), otherwise + # create a standalone dict (SDK / direct-call path). + if request_data is None: + request_data = {"response": response} + else: + if "response" not in request_data: + request_data["response"] = response # Add user API key metadata with prefixed keys - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + if "litellm_metadata" not in request_data: + user_metadata = self.transform_user_api_key_dict_to_metadata( + user_api_key_dict + ) + if user_metadata: + request_data["litellm_metadata"] = user_metadata inputs = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: @@ -454,6 +461,7 @@ async def process_output_streaming_response( guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional[Any] = None, user_api_key_dict: Optional[Any] = None, + request_data: Optional[dict] = None, ) -> List[Any]: """ Process output streaming response by applying guardrails to text content. @@ -481,7 +489,7 @@ async def process_output_streaming_response( inputs["model"] = model_response_stream.model _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, - request_data={}, + request_data=request_data if request_data is not None else {}, input_type="response", logging_obj=litellm_logging_obj, ) @@ -512,7 +520,7 @@ async def process_output_streaming_response( if tool_calls or text: _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=guardrail_inputs, - request_data={}, + request_data=request_data if request_data is not None else {}, input_type="response", logging_obj=litellm_logging_obj, ) @@ -537,7 +545,7 @@ async def process_output_streaming_response( inputs["model"] = response_model _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, - request_data={}, + request_data=request_data if request_data is not None else {}, input_type="response", logging_obj=litellm_logging_obj, ) diff --git a/litellm/llms/openai/speech/guardrail_translation/handler.py b/litellm/llms/openai/speech/guardrail_translation/handler.py index e6796fbac2a4..f0c3149d0aee 100644 --- a/litellm/llms/openai/speech/guardrail_translation/handler.py +++ b/litellm/llms/openai/speech/guardrail_translation/handler.py @@ -85,6 +85,7 @@ async def process_output_response( guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional[Any] = None, user_api_key_dict: Optional[Any] = None, + request_data: Optional[dict] = None, ) -> Any: """ Process output - not applicable for text-to-speech. diff --git a/litellm/llms/openai/transcriptions/guardrail_translation/handler.py b/litellm/llms/openai/transcriptions/guardrail_translation/handler.py index 3d76a21c3898..92cf4398f053 100644 --- a/litellm/llms/openai/transcriptions/guardrail_translation/handler.py +++ b/litellm/llms/openai/transcriptions/guardrail_translation/handler.py @@ -58,6 +58,7 @@ async def process_output_response( guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional[Any] = None, user_api_key_dict: Optional[Any] = None, + request_data: Optional[dict] = None, ) -> Any: """ Process output transcription by applying guardrails to transcribed text. @@ -79,15 +80,21 @@ async def process_output_response( if isinstance(response.text, str): original_text = response.text - # Create a request_data dict with response info and user API key metadata - request_data: dict = {"response": response} + # Use the real request_data if provided (proxy path), otherwise + # create a standalone dict (SDK / direct-call path). + if request_data is None: + request_data = {"response": response} + else: + if "response" not in request_data: + request_data["response"] = response # Add user API key metadata with prefixed keys - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + if "litellm_metadata" not in request_data: + user_metadata = self.transform_user_api_key_dict_to_metadata( + user_api_key_dict + ) + if user_metadata: + request_data["litellm_metadata"] = user_metadata inputs = GenericGuardrailAPIInputs(texts=[original_text]) # Include model information from the response if available diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index 40433d534137..a8cc42d7c548 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -139,6 +139,7 @@ async def process_output_response( guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, user_api_key_dict: Optional[Any] = None, + request_data: Optional[dict] = None, ) -> Any: """ Process output response by applying guardrails to targeted fields. @@ -171,17 +172,27 @@ async def process_output_response( if not text_to_check: return response - # Create a request_data dict with response info and user API key metadata - request_data: dict = ( - {"response": response} - if not isinstance(response, dict) - else response.copy() - ) + # Use the real request_data if provided (proxy path), otherwise + # create a standalone dict (SDK / direct-call path). + if request_data is None: + request_data = ( + {"response": response} + if not isinstance(response, dict) + else response.copy() + ) + else: + if "response" not in request_data: + request_data["response"] = ( + response if not isinstance(response, dict) else response.copy() + ) # Add user API key metadata with prefixed keys - user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + if "litellm_metadata" not in request_data: + user_metadata = self.transform_user_api_key_dict_to_metadata( + user_api_key_dict + ) + if user_metadata: + request_data["litellm_metadata"] = user_metadata # Apply guardrail (pass-through doesn't modify the text, just checks it) inputs = GenericGuardrailAPIInputs(texts=[text_to_check]) diff --git a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py index 14bbb82808df..6997f5241de9 100644 --- a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py +++ b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py @@ -92,6 +92,7 @@ async def process_output_response( guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional[Any] = None, user_api_key_dict: Optional[Any] = None, + request_data: Optional[dict] = None, ) -> Any: verbose_proxy_logger.debug( "MCP Guardrail: Output processing not implemented for MCP tools", diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py index 4bd94345727a..4ddeac9a2084 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py @@ -5,9 +5,11 @@ from typing import ( TYPE_CHECKING, + Dict, Literal, Optional, Type, + Union, ) from fastapi import HTTPException @@ -22,7 +24,8 @@ get_async_httpx_client, httpxSpecialProvider, ) -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus from .base import OpenAIGuardrailBase @@ -223,12 +226,98 @@ async def apply_guardrail( # Make moderation request moderation_response = await self.async_make_request(input_text=text_to_moderate) + # Stash full moderation response in request_data for logging + # (Model Armor pattern — per-request dict avoids race conditions) + if isinstance(request_data, dict): + metadata = request_data.get("metadata") or {} + request_data["metadata"] = metadata + metadata["_openai_moderation_response"] = moderation_response.model_dump() + # Check if content is flagged and raise exception if needed self._check_moderation_result(moderation_response) # Moderation doesn't modify content, just blocks - return inputs unchanged return inputs + def _process_response( + self, + response: Optional[Dict], + request_data: dict, + start_time: Optional[float] = None, + end_time: Optional[float] = None, + duration: Optional[float] = None, + event_type: Optional[GuardrailEventHooks] = None, + original_inputs: Optional[Dict] = None, + ): + """ + Override to log the full OpenAI Moderation API response instead of + the decorator's simplified "allow"/"mask" string. + + Follows the Model Armor pattern (model_armor.py:325-360). + """ + if isinstance(request_data, dict): + metadata = request_data.get("metadata") or {} + request_data["metadata"] = metadata # anchor so pop() mutates the real dict + else: + metadata = {} + + # .pop() cleans up the internal key so it doesn't leak to downstream + # loggers. Falls back to "allow" when no moderation call was made + # (e.g. no text to moderate — early return in apply_guardrail). + guardrail_response = metadata.pop("_openai_moderation_response", "allow") + + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=guardrail_response, + request_data=request_data, + guardrail_status="success", + duration=duration, + start_time=start_time, + end_time=end_time, + event_type=event_type, + ) + return response + + def _process_error( + self, + e: Exception, + request_data: dict, + start_time: Optional[float] = None, + end_time: Optional[float] = None, + duration: Optional[float] = None, + event_type: Optional[GuardrailEventHooks] = None, + ): + """ + Override to log the full OpenAI Moderation API response on error + instead of the stringified exception. + """ + guardrail_status: GuardrailStatus = ( + "guardrail_intervened" + if self._is_guardrail_intervention(e) + else "guardrail_failed_to_respond" + ) + + if isinstance(request_data, dict): + metadata = request_data.get("metadata") or {} + request_data["metadata"] = metadata # anchor so pop() mutates the real dict + else: + metadata = {} + + # Use the stashed moderation response if available, fall back to exception + guardrail_response: Union[dict, Exception, str] = metadata.pop( + "_openai_moderation_response", e + ) + + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=guardrail_response, + request_data=request_data, + guardrail_status=guardrail_status, + duration=duration, + start_time=start_time, + end_time=end_time, + event_type=event_type, + ) + raise e + @staticmethod def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 84bbf6d20e13..a1623121da58 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -247,6 +247,7 @@ async def async_post_call_success_hook( guardrail_to_apply=guardrail_to_apply, litellm_logging_obj=data.get("litellm_logging_obj"), user_api_key_dict=user_api_key_dict, + request_data=data, ) # Add guardrail to applied guardrails header add_guardrail_to_applied_guardrails_header( @@ -397,6 +398,7 @@ async def async_post_call_streaming_iterator_hook( # noqa: PLR0915 guardrail_to_apply=guardrail_to_apply, litellm_logging_obj=request_data.get("litellm_logging_obj"), user_api_key_dict=user_api_key_dict, + request_data=request_data, ) except HTTPException as e: # Response already started (we already yielded chunks); cannot send 400. @@ -457,6 +459,7 @@ async def async_post_call_streaming_iterator_hook( # noqa: PLR0915 guardrail_to_apply=guardrail_to_apply, litellm_logging_obj=request_data.get("litellm_logging_obj"), user_api_key_dict=user_api_key_dict, + request_data=request_data, ) except HTTPException as e: if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 7954f0b64606..b9c1e2f2bc32 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1898,9 +1898,9 @@ async def _handle_logging_proxy_only_error( normalized_call_type = CallTypes.aembedding.value if normalized_call_type is not None: litellm_logging_obj.call_type = normalized_call_type - litellm_logging_obj.model_call_details[ - "call_type" - ] = normalized_call_type + litellm_logging_obj.model_call_details["call_type"] = ( + normalized_call_type + ) # Pass-through endpoints are logged via the callback loop's # async_post_call_failure_hook — skip pre_call and failure handlers. if litellm_logging_obj.call_type == CallTypes.pass_through.value: @@ -2498,8 +2498,7 @@ async def check_view_exists(self): required_view = "LiteLLM_VerificationTokenView" expected_views_str = ", ".join(f"'{view}'" for view in expected_views) pg_schema = os.getenv("DATABASE_SCHEMA", "public") - ret = await self.db.query_raw( - f""" + ret = await self.db.query_raw(f""" WITH existing_views AS ( SELECT viewname FROM pg_views @@ -2511,8 +2510,7 @@ async def check_view_exists(self): (SELECT COUNT(*) FROM existing_views) AS view_count, ARRAY_AGG(viewname) AS view_names FROM existing_views - """ - ) + """) expected_total_views = len(expected_views) if ret[0]["view_count"] == expected_total_views: verbose_proxy_logger.info("All necessary views exist!") @@ -2521,8 +2519,7 @@ async def check_view_exists(self): ## check if required view exists ## if ret[0]["view_names"] and required_view not in ret[0]["view_names"]: await self.health_check() # make sure we can connect to db - await self.db.execute_raw( - """ + await self.db.execute_raw(""" CREATE VIEW "LiteLLM_VerificationTokenView" AS SELECT v.*, @@ -2532,8 +2529,7 @@ async def check_view_exists(self): t.rpm_limit AS team_rpm_limit FROM "LiteLLM_VerificationToken" v LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id; - """ - ) + """) verbose_proxy_logger.info( "LiteLLM_VerificationTokenView Created in DB!" @@ -2759,7 +2755,7 @@ async def get_data( # noqa: PLR0915 and reset_at is not None ): response = await self.db.litellm_verificationtoken.find_many( - where={ # type:ignore + where={ # type: ignore "OR": [ {"expires": None}, {"expires": {"gt": expires}}, @@ -2819,7 +2815,7 @@ async def get_data( # noqa: PLR0915 ) # type: ignore elif query_type == "find_all" and reset_at is not None: response = await self.db.litellm_usertable.find_many( - where={ # type:ignore + where={ # type: ignore "budget_reset_at": {"lt": reset_at}, } ) @@ -2831,10 +2827,10 @@ async def get_data( # noqa: PLR0915 if expires is not None: response = await self.db.litellm_usertable.find_many( # type: ignore order={"spend": "desc"}, - where={ # type:ignore + where={ # type: ignore "OR": [ - {"expires": None}, # type:ignore - {"expires": {"gt": expires}}, # type:ignore + {"expires": None}, # type: ignore + {"expires": {"gt": expires}}, # type: ignore ], }, ) @@ -2881,7 +2877,7 @@ async def get_data( # noqa: PLR0915 elif table_name == "budget" and reset_at is not None: if query_type == "find_all": response = await self.db.litellm_budgettable.find_many( - where={ # type:ignore + where={ # type: ignore "OR": [ { "AND": [ @@ -2909,7 +2905,7 @@ async def get_data( # noqa: PLR0915 ) elif query_type == "find_all" and reset_at is not None: response = await self.db.litellm_teamtable.find_many( - where={ # type:ignore + where={ # type: ignore "budget_reset_at": {"lt": reset_at}, } ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py index 3a17bbd0025c..5c19e7189e06 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py @@ -488,3 +488,330 @@ async def mock_stream(): assert exc_info.value.status_code == 400 assert "Violated OpenAI moderation policy" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_openai_moderation_guardrail_logs_full_response_safe_content(): + """Test that safe content logs the full moderation response (categories, scores) + in StandardLoggingGuardrailInformation, not just 'allow'.""" + from litellm.types.utils import GenericGuardrailAPIInputs + + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + guardrail = OpenAIModerationGuardrail( + guardrail_name="test-openai-moderation", + ) + + mock_response = OpenAIModerationResponse( + id="modr-123", + model="omni-moderation-latest", + results=[ + OpenAIModerationResult( + flagged=False, + categories={ + "sexual": False, + "hate": False, + "harassment": False, + "self-harm": False, + "violence": False, + }, + category_scores={ + "sexual": 0.001, + "hate": 0.002, + "harassment": 0.001, + "self-harm": 0.001, + "violence": 0.003, + }, + category_applied_input_types={ + "sexual": [], + "hate": [], + "harassment": [], + "self-harm": [], + "violence": [], + }, + ) + ], + ) + + with patch.object(guardrail, "async_make_request", return_value=mock_response): + inputs = GenericGuardrailAPIInputs( + structured_messages=[ + {"role": "user", "content": "Hello, how are you?"} + ] + ) + request_data = {"metadata": {}} + + await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + guardrail_info_list = request_data["metadata"][ + "standard_logging_guardrail_information" + ] + assert len(guardrail_info_list) == 1 + + info = guardrail_info_list[0] + assert info["guardrail_name"] == "test-openai-moderation" + assert info["guardrail_status"] == "success" + + # Full moderation response, NOT "allow" + guardrail_resp = info["guardrail_response"] + assert isinstance(guardrail_resp, dict) + assert guardrail_resp["results"][0]["flagged"] is False + assert "category_scores" in guardrail_resp["results"][0] + + # Internal key cleaned up (.pop()) + assert "_openai_moderation_response" not in request_data["metadata"] + + +@pytest.mark.asyncio +async def test_openai_moderation_guardrail_logs_full_response_harmful_content(): + """Test that harmful content logs guardrail_intervened status with the full + moderation response, not just the exception string.""" + from litellm.types.utils import GenericGuardrailAPIInputs + + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + guardrail = OpenAIModerationGuardrail( + guardrail_name="test-openai-moderation", + ) + + mock_response = OpenAIModerationResponse( + id="modr-456", + model="omni-moderation-latest", + results=[ + OpenAIModerationResult( + flagged=True, + categories={ + "sexual": False, + "hate": True, + "harassment": False, + "self-harm": False, + "violence": False, + }, + category_scores={ + "sexual": 0.001, + "hate": 0.95, + "harassment": 0.001, + "self-harm": 0.001, + "violence": 0.001, + }, + category_applied_input_types={ + "sexual": [], + "hate": ["text"], + "harassment": [], + "self-harm": [], + "violence": [], + }, + ) + ], + ) + + with patch.object(guardrail, "async_make_request", return_value=mock_response): + inputs = GenericGuardrailAPIInputs( + structured_messages=[ + {"role": "user", "content": "Hateful content"} + ] + ) + request_data = {"metadata": {}} + + from fastapi import HTTPException + + with pytest.raises(HTTPException): + await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + guardrail_info_list = request_data["metadata"][ + "standard_logging_guardrail_information" + ] + info = guardrail_info_list[0] + assert info["guardrail_status"] == "guardrail_intervened" + + # Full moderation response, NOT stringified exception + guardrail_resp = info["guardrail_response"] + assert isinstance(guardrail_resp, dict) + assert guardrail_resp["results"][0]["flagged"] is True + assert guardrail_resp["results"][0]["category_scores"]["hate"] == 0.95 + + # Internal key cleaned up by _process_error (.pop()) + assert "_openai_moderation_response" not in request_data["metadata"] + + +@pytest.mark.asyncio +async def test_openai_moderation_post_call_request_data_passthrough(): + """Test that post-call guardrail info flows through to the real request_data + via the unified guardrail dispatcher (Bug 1 fix).""" + from unittest.mock import AsyncMock + + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + from litellm.types.utils import ModelResponse + + import litellm + + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + guardrail = OpenAIModerationGuardrail( + guardrail_name="test-openai-moderation", + event_hook="post_call", + ) + unified_guardrail = UnifiedLLMGuardrails() + + mock_mod_response = OpenAIModerationResponse( + id="modr-789", + model="omni-moderation-latest", + results=[ + OpenAIModerationResult( + flagged=False, + categories={ + "sexual": False, + "hate": False, + "harassment": False, + "self-harm": False, + "violence": False, + }, + category_scores={ + "sexual": 0.001, + "hate": 0.002, + "harassment": 0.001, + "self-harm": 0.001, + "violence": 0.001, + }, + category_applied_input_types={ + "sexual": [], + "hate": [], + "harassment": [], + "self-harm": [], + "violence": [], + }, + ) + ], + ) + + llm_response = ModelResponse( + id="chatcmpl-test", + model="gpt-4", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message( + role="assistant", content="Hello world" + ), + finish_reason="stop", + ) + ], + ) + + request_data = { + "messages": [{"role": "user", "content": "Hello"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": ["test-openai-moderation"]}, + } + + mock_make_request = AsyncMock(return_value=mock_mod_response) + with patch.object(guardrail, "async_make_request", mock_make_request): + await unified_guardrail.async_post_call_success_hook( + data=request_data, + user_api_key_dict=UserAPIKeyAuth( + api_key="test", request_route="/chat/completions" + ), + response=llm_response, + ) + + mock_make_request.assert_called_once() + + # Guardrail info in the REAL request_data (not a throwaway) + guardrail_info_list = request_data["metadata"].get( + "standard_logging_guardrail_information" + ) + assert guardrail_info_list is not None + assert isinstance(guardrail_info_list[0]["guardrail_response"], dict) + assert "results" in guardrail_info_list[0]["guardrail_response"] + + +def test_openai_moderation_process_response_metadata_none_edge_case(): + """ + Test that _process_response anchors the metadata dict back into + request_data when metadata is None, so pop() doesn't operate on a + temporary and the moderation response is correctly logged. + """ + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + guardrail = OpenAIModerationGuardrail( + guardrail_name="test-openai-moderation", + event_hook="post_call", + ) + + mod_dict = {"id": "modr-test", "model": "omni-moderation-latest", "results": []} + + # Simulate apply_guardrail having stashed the response but metadata + # was None initially — apply_guardrail anchors it, so metadata is a + # real dict with the stashed key by the time _process_response runs. + request_data = {"metadata": {"_openai_moderation_response": mod_dict}} + + guardrail._process_response( + response={"inputs": {}}, + request_data=request_data, + ) + + # Full moderation dict should be logged, not "allow" + info_list = request_data["metadata"].get( + "standard_logging_guardrail_information" + ) + assert info_list is not None + assert info_list[0]["guardrail_response"] == mod_dict + + # Internal key should have been cleaned up by pop() + assert "_openai_moderation_response" not in request_data["metadata"] + + +def test_openai_moderation_process_error_metadata_none_edge_case(): + """ + Test that _process_error anchors the metadata dict back into + request_data when metadata starts as None (or {}), so pop() doesn't + operate on a temporary. + """ + from fastapi import HTTPException + + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + guardrail = OpenAIModerationGuardrail( + guardrail_name="test-openai-moderation", + event_hook="post_call", + ) + + mod_dict = { + "id": "modr-test", + "model": "omni-moderation-latest", + "results": [{"flagged": True, "categories": {"hate": True}}], + } + + # metadata is None — exercises the or {} anchor + request_data: dict = {"metadata": None} + + # Simulate stashing the response then calling _process_error + # (normally apply_guardrail stashes, then the decorator calls + # _process_error on HTTPException) + # First anchor metadata like apply_guardrail does: + metadata = request_data.get("metadata") or {} + request_data["metadata"] = metadata + metadata["_openai_moderation_response"] = mod_dict + + exc = HTTPException(status_code=400, detail="Violated policy") + with pytest.raises(HTTPException): + guardrail._process_error( + e=exc, + request_data=request_data, + ) + + # Full moderation dict should be logged, not the exception + info_list = request_data["metadata"].get( + "standard_logging_guardrail_information" + ) + assert info_list is not None + assert info_list[0]["guardrail_response"] == mod_dict + assert info_list[0]["guardrail_status"] == "guardrail_intervened" + + # Internal key cleaned up + assert "_openai_moderation_response" not in request_data["metadata"] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py index c77a5d07b3b3..2595a1df7f1e 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py @@ -170,3 +170,106 @@ async def mock_stream(): assert exc_info.value.status_code == 400 assert "Violated OpenAI moderation policy" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_openai_moderation_streaming_end_of_stream_request_data_passthrough(): + """Test that streaming end-of-stream guardrail info flows through to the + real request_data (Bug 1 fix for streaming path).""" + from litellm.types.llms.openai import ( + OpenAIModerationResponse, + OpenAIModerationResult, + ) + + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + openai_guardrail = OpenAIModerationGuardrail( + guardrail_name="test-openai-moderation", + event_hook="post_call", + ) + unified_guardrail = UnifiedLLMGuardrails() + + mock_mod_response = OpenAIModerationResponse( + id="modr-stream-test", + model="omni-moderation-latest", + results=[ + OpenAIModerationResult( + flagged=False, + categories={"hate": False, "violence": False}, + category_scores={"hate": 0.001, "violence": 0.002}, + category_applied_input_types={"hate": [], "violence": []}, + ) + ], + ) + + async def mock_stream(): + import litellm + + chunks_data = ["Hello", " world"] + for i, content in enumerate(chunks_data): + chunk = MagicMock(spec=ModelResponseStream) + chunk.model = "gpt-4" + choice = MagicMock() + choice.delta = MagicMock() + choice.delta.content = content + choice.finish_reason = "stop" if i == len(chunks_data) - 1 else None + chunk.choices = [choice] + yield chunk + + import litellm + + mock_model_response = ModelResponse( + id="mock-stream-response", + model="gpt-4", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message( + role="assistant", content="Hello world" + ), + finish_reason="stop", + ) + ], + ) + + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": openai_guardrail, + "metadata": { + "guardrails": ["test-openai-moderation"], + "guardrail_config": {"streaming_sampling_rate": 1}, + }, + } + + with patch.object( + openai_guardrail, "async_make_request", return_value=mock_mod_response + ), patch( + "litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder", + return_value=mock_model_response, + ): + user_api_key_dict = UserAPIKeyAuth( + api_key="test", request_route="/chat/completions" + ) + + async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + pass + + # Verify guardrail info reached the REAL request_data (not a throwaway) + guardrail_info_list = request_data["metadata"].get( + "standard_logging_guardrail_information" + ) + assert guardrail_info_list is not None, ( + "Guardrail info should be in request_data after streaming" + ) + info = guardrail_info_list[0] + assert info["guardrail_status"] == "success" + + # Full moderation response dict, NOT the simplified "allow" string + guardrail_resp = info["guardrail_response"] + assert isinstance(guardrail_resp, dict), ( + f"Expected full moderation response dict, got {type(guardrail_resp)}: {guardrail_resp}" + ) + assert "results" in guardrail_resp diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 7c29c8161db6..bbba8e4d03dd 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -78,9 +78,7 @@ async def test_uses_mcp_event_type(self): data = { "guardrail_to_apply": guardrail, - "messages": [ - {"role": "user", "content": "Tool: test\nArguments: {}"} - ], + "messages": [{"role": "user", "content": "Tool: test\nArguments: {}"}], "model": "mcp-tool-call", } @@ -102,9 +100,7 @@ async def test_uses_mcp_event_type(self): data = { "guardrail_to_apply": guardrail, - "messages": [ - {"role": "user", "content": "Tool: test\nArguments: {}"} - ], + "messages": [{"role": "user", "content": "Tool: test\nArguments: {}"}], "model": "mcp-tool-call", } @@ -168,6 +164,7 @@ async def process_output_streaming_response( guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, + request_data=None, ): # Simulate what the real handler does: # put combined text in first chunk, clear the rest @@ -200,10 +197,12 @@ async def process_output_streaming_response( chunks = [] for i in range(10): chunk = ModelResponseStream( - choices=[StreamingChoices( - delta=Delta(content=f"word{i} ", role="assistant"), - finish_reason=None, - )], + choices=[ + StreamingChoices( + delta=Delta(content=f"word{i} ", role="assistant"), + finish_reason=None, + ) + ], ) chunks.append(chunk) @@ -228,7 +227,9 @@ async def mock_stream(): response=mock_stream(), request_data=request_data, ): - content = item.choices[0].delta.content if item.choices[0].delta else None + content = ( + item.choices[0].delta.content if item.choices[0].delta else None + ) yielded_contents.append(content) # Every chunk should have non-empty content @@ -271,10 +272,15 @@ async def test_pre_call_hook_invokes_ocr_handler_for_input(self): assert guardrail.event_history == [GuardrailEventHooks.pre_call] assert len(guardrail.apply_calls) == 1 assert guardrail.apply_calls[0]["input_type"] == "request" - assert "https://arxiv.org/pdf/2201.04234" in guardrail.apply_calls[0]["inputs"]["texts"] + assert ( + "https://arxiv.org/pdf/2201.04234" + in guardrail.apply_calls[0]["inputs"]["texts"] + ) # Data should be returned with document intact - assert result["document"]["document_url"] == "https://arxiv.org/pdf/2201.04234" + assert ( + result["document"]["document_url"] == "https://arxiv.org/pdf/2201.04234" + ) @pytest.mark.asyncio async def test_moderation_hook_invokes_ocr_handler(self): @@ -302,7 +308,10 @@ async def test_moderation_hook_invokes_ocr_handler(self): assert guardrail.event_history == [GuardrailEventHooks.during_call] assert len(guardrail.apply_calls) == 1 - assert "https://example.com/scan.png" in guardrail.apply_calls[0]["inputs"]["texts"] + assert ( + "https://example.com/scan.png" + in guardrail.apply_calls[0]["inputs"]["texts"] + ) @pytest.mark.asyncio async def test_post_call_success_hook_guardrails_ocr_output(self): @@ -318,7 +327,9 @@ def __init__(self): def should_run_guardrail(self, data, event_type): # type: ignore[override] return True - async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + async def apply_guardrail( + self, inputs, request_data, input_type, **kwargs + ): texts = inputs.get("texts", []) return {"texts": [t.replace("SECRET", "[REDACTED]") for t in texts]}