diff --git a/litellm/constants.py b/litellm/constants.py index 2ef54e080a0..8a003d35e91 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -190,6 +190,10 @@ DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE = int( # Override with LITELLM_MAX_CALLBACKS env var for large deployments (e.g., many teams with guardrails) MAX_CALLBACKS = get_env_int("LITELLM_MAX_CALLBACKS", 100) +# Metadata key recording which pre_call guardrails the proxy loop already ran, +# so the deployment-level hook does not re-run them for the same request +PRE_CALL_EXECUTED_GUARDRAILS_KEY = "_pre_call_executed_guardrails" + # Generic fallback for unknown models DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET = int( os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128) diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 213622cb43a..296bfb6fc85 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -27,6 +27,11 @@ else: LiteLLMLoggingObj = Any +# Anthropic (and Bedrock Claude) reject requests with more than 4 cache_control +# breakpoints: "A maximum of 4 blocks with cache_control may be provided." +MAX_CACHE_CONTROL_BLOCKS = 4 + + class AnthropicCacheControlHook(CustomPromptManagement): def get_chat_completion_prompt( self, @@ -61,16 +66,30 @@ class AnthropicCacheControlHook(CustomPromptManagement): processed_messages = copy.deepcopy(messages) # Separate message-level and non-message-level injection points - remaining_points = [] + message_points: List[CacheControlMessageInjectionPoint] = [] + remaining_points: List[CacheControlInjectionPoint] = [] for point in injection_points: if point.get("location") == "message": - point = cast(CacheControlMessageInjectionPoint, point) - processed_messages = self._process_message_injection( - point=point, messages=processed_messages - ) + message_points.append(cast(CacheControlMessageInjectionPoint, point)) else: remaining_points.append(point) + # Non-message points (currently Bedrock tool_config) are handled in the + # provider transform, where each tool_config point appends at most one + # cachePoint to the tools. That block also counts toward Anthropic's + # limit, so reserve a slot for it here to leave room. + reserved_blocks = ( + 1 + if any(p.get("location") == "tool_config" for p in remaining_points) + else 0 + ) + + processed_messages = self._apply_message_injections( + points=message_points, + messages=processed_messages, + max_blocks=MAX_CACHE_CONTROL_BLOCKS - reserved_blocks, + ) + # Pass through non-message injection points for provider-specific handling if remaining_points: non_default_params["cache_control_injection_points"] = remaining_points @@ -78,14 +97,71 @@ class AnthropicCacheControlHook(CustomPromptManagement): return model, processed_messages, non_default_params @staticmethod - def _process_message_injection( - point: CacheControlMessageInjectionPoint, messages: List[AllMessageValues] + def _apply_message_injections( + points: List[CacheControlMessageInjectionPoint], + messages: List[AllMessageValues], + max_blocks: int, ) -> List[AllMessageValues]: - """Process message-level cache control injection.""" - control: ChatCompletionCachedContent = point.get( - "control", None - ) or ChatCompletionCachedContent(type="ephemeral") + """Apply message-level cache control injection points in order. + Anthropic allows at most ``MAX_CACHE_CONTROL_BLOCKS`` cache_control + breakpoints per request. Client-supplied breakpoints count toward that + limit, so we never inject onto a message that already carries + cache_control (preserving the client's TTL) and we stop injecting once + ``max_blocks`` is reached. Injection points are honored in config order, + so earlier points win when slots are scarce. + """ + used_blocks = sum( + AnthropicCacheControlHook._count_cache_control_blocks(msg) + for msg in messages + ) + + limit_reached = False + for point in points: + if used_blocks >= max_blocks: + limit_reached = True + break + + control: ChatCompletionCachedContent = point.get( + "control", None + ) or ChatCompletionCachedContent(type="ephemeral") + + for target_index in AnthropicCacheControlHook._resolve_target_indices( + point=point, messages=messages + ): + if used_blocks >= max_blocks: + limit_reached = True + break + + if AnthropicCacheControlHook._message_has_cache_control( + messages[target_index] + ): + # Client already marked this message; don't overwrite it. + continue + + messages[target_index] = ( + AnthropicCacheControlHook._safe_insert_cache_control_in_message( + messages[target_index], control + ) + ) + used_blocks += 1 + + if limit_reached: + break + + if limit_reached: + verbose_logger.warning( + f"AnthropicCacheControlHook: Reached the Anthropic limit of " + f"{MAX_CACHE_CONTROL_BLOCKS} cache_control blocks. Skipping further injection." + ) + + return messages + + @staticmethod + def _resolve_target_indices( + point: CacheControlMessageInjectionPoint, messages: List[AllMessageValues] + ) -> List[int]: + """Resolve which message indices an injection point targets.""" _targetted_index: Optional[Union[int, str]] = point.get("index", None) targetted_index: Optional[int] = None if isinstance(_targetted_index, str): @@ -96,36 +172,49 @@ class AnthropicCacheControlHook(CustomPromptManagement): else: targetted_index = _targetted_index - targetted_role = point.get("role", None) - # Case 1: Target by specific index if targetted_index is not None: original_index = targetted_index - # Handle negative indices (convert to positive) if targetted_index < 0: targetted_index += len(messages) if 0 <= targetted_index < len(messages): - messages[targetted_index] = ( - AnthropicCacheControlHook._safe_insert_cache_control_in_message( - messages[targetted_index], control - ) - ) - else: - verbose_logger.warning( - f"AnthropicCacheControlHook: Provided index {original_index} is out of bounds for message list of length {len(messages)}. " - f"Targeted index was {targetted_index}. Skipping cache control injection for this point." - ) + return [targetted_index] + + verbose_logger.warning( + f"AnthropicCacheControlHook: Provided index {original_index} is out of bounds for message list of length {len(messages)}. " + f"Targeted index was {targetted_index}. Skipping cache control injection for this point." + ) + return [] + # Case 2: Target by role - elif targetted_role is not None: - for msg in messages: - if msg.get("role") == targetted_role: - msg = ( - AnthropicCacheControlHook._safe_insert_cache_control_in_message( - message=msg, control=control - ) - ) - return messages + targetted_role = point.get("role", None) + if targetted_role is not None: + return [ + idx + for idx, msg in enumerate(messages) + if msg.get("role") == targetted_role + ] + + return [] + + @staticmethod + def _count_cache_control_blocks(message: AllMessageValues) -> int: + """Count cache_control breakpoints on a message (message + content level).""" + count = 0 + if message.get("cache_control") is not None: + count += 1 + content = message.get("content") + if isinstance(content, list): + for block in content: + if isinstance(block, dict) and block.get("cache_control") is not None: + count += 1 + return count + + @staticmethod + def _message_has_cache_control(message: AllMessageValues) -> bool: + """Return True if the message already carries any cache_control.""" + return AnthropicCacheControlHook._count_cache_control_blocks(message) > 0 @staticmethod def _safe_insert_cache_control_in_message( diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 82a35f2eedd..586831a078f 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1,3 +1,4 @@ +import secrets from datetime import datetime from typing import ( TYPE_CHECKING, @@ -43,12 +44,19 @@ if TYPE_CHECKING: dc = DualCache() +from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY from litellm.exceptions import ( BlockedPiiEntityError, GuardrailRaisedException, ModifyResponseException, ) +# Per-process secret tagging each recorded marker. The deployment hook only +# honors markers carrying this token, so a caller cannot forge the metadata +# field to suppress a guardrail on the direct-SDK path that never reaches the +# proxy's metadata sanitizer. +_PRE_CALL_EXECUTED_TOKEN = secrets.token_hex(16) + class CustomGuardrail(CustomLogger): # If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path. @@ -325,6 +333,49 @@ class CustomGuardrail(CustomLogger): return False + def _pre_call_marker(self) -> Optional[str]: + name = self.guardrail_name + if not name: + return None + return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}" + + def mark_pre_call_hook_ran(self, data: Dict[str, Any]) -> None: + """ + Record that this guardrail's ``async_pre_call_hook`` already ran for this + request, so the deployment-level hook does not run it a second time. + + The proxy runs pre-call guardrails in ``ProxyLogging.pre_call_hook``. The + router later spreads a deployment's model-level ``guardrails`` into the + top-level request kwargs, which would otherwise re-trigger the same hook + from ``async_pre_call_deployment_hook``. + """ + marker = self._pre_call_marker() + if marker is None: + return + for meta_key in ("metadata", "litellm_metadata"): + meta = data.get(meta_key) + if isinstance(meta, dict): + executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY) + if isinstance(executed, list): + if marker not in executed: + executed.append(marker) + else: + meta[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker] + return + data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]} + + def _pre_call_hook_already_ran(self, data: Dict[str, Any]) -> bool: + marker = self._pre_call_marker() + if marker is None: + return False + for meta_key in ("metadata", "litellm_metadata"): + meta = data.get(meta_key) + if isinstance(meta, dict): + executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY) + if isinstance(executed, list) and marker in executed: + return True + return False + async def async_pre_call_deployment_hook( self, kwargs: Dict[str, Any], call_type: Optional[CallTypes] ) -> Optional[dict]: @@ -335,6 +386,9 @@ class CustomGuardrail(CustomLogger): if litellm_guardrails is None or not isinstance(litellm_guardrails, list): return kwargs + if self._pre_call_hook_already_ran(kwargs): + return kwargs + if ( self.should_run_guardrail( data=kwargs, event_type=GuardrailEventHooks.pre_call diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index c3e555f6e89..79a9219a39c 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -41,6 +41,7 @@ from litellm.integrations.datadog.datadog_handler import ( ) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.llms.custom_httpx.http_handler import ( + MaskedHTTPStatusError, _get_httpx_client, get_async_httpx_client, httpxSpecialProvider, @@ -68,6 +69,22 @@ DD_LOGGED_SUCCESS_SERVICE_TYPES = [ ] +def _resolve_dd_batch_size() -> int: + raw = os.getenv("DD_BATCH_SIZE") + if raw is None: + return DD_MAX_BATCH_SIZE + try: + value = int(raw) + except ValueError: + verbose_logger.warning( + "Datadog: ignoring invalid DD_BATCH_SIZE=%r, using %s", + raw, + DD_MAX_BATCH_SIZE, + ) + return DD_MAX_BATCH_SIZE + return max(1, min(value, DD_MAX_BATCH_SIZE)) + + class DataDogLogger( CustomBatchLogger, AdditionalLoggingUtils, @@ -128,7 +145,9 @@ class DataDogLogger( asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() super().__init__( - **kwargs, flush_lock=self.flush_lock, batch_size=DD_MAX_BATCH_SIZE + **kwargs, + flush_lock=self.flush_lock, + batch_size=_resolve_dd_batch_size(), ) except Exception as e: verbose_logger.exception( @@ -339,28 +358,14 @@ class DataDogLogger( "[DATADOG MOCK] Mock mode enabled - API calls will be intercepted" ) - response = await self.async_send_compressed_data(batch_to_send) - if response.status_code == 413: - verbose_logger.exception(DD_ERRORS.DATADOG_413_ERROR.value) - self.log_queue = batch_to_send + self.log_queue - return - - response.raise_for_status() - if response.status_code != 202: - raise Exception( - f"Response from datadog API status_code: {response.status_code}, text: {response.text}" - ) + undelivered = await self._send_with_413_split(batch_to_send) + if undelivered: + self.log_queue = undelivered + self.log_queue if self.is_mock_mode: verbose_logger.debug( f"[DATADOG MOCK] Batch of {len(batch_to_send)} events successfully mocked" ) - else: - verbose_logger.debug( - "Datadog: Response from datadog API status_code: %s, text: %s", - response.status_code, - response.text, - ) except Exception as e: self.log_queue = batch_to_send + self.log_queue @@ -368,6 +373,62 @@ class DataDogLogger( f"Datadog Error sending batch API - {str(e)}\n{traceback.format_exc()}" ) + async def _send_with_413_split(self, batch: List) -> List: + """ + Send a batch, halving any sub-batch that 413s (payload too large) and retrying the + halves, since Datadog enforces a 5MB uncompressed limit per request. + + A 413 surfaces as a raised MaskedHTTPStatusError (httpx raise_for_status), not a + returned response, so both paths are handled. A lone event that still 413s is + dropped to avoid wedging the queue on an undeliverable payload. Returns the events + that could not be delivered because of a non-413 (transient) error, so the caller + re-queues only those and never the events already accepted by Datadog. + """ + pending: List[List] = [batch] + while pending: + chunk = pending.pop() + if not chunk: + continue + try: + response = await self.async_send_compressed_data(chunk) + except Exception as e: + if isinstance(e, MaskedHTTPStatusError) and e.status_code == 413: + response = e.response + else: + verbose_logger.exception( + f"Datadog Error sending batch API - {str(e)}" + ) + return self._undelivered(chunk, pending) + + if response.status_code == 413: + if len(chunk) == 1: + verbose_logger.error(DD_ERRORS.DATADOG_413_ERROR.value) + continue + mid = len(chunk) // 2 + pending.append(chunk[mid:]) + pending.append(chunk[:mid]) + continue + + if response.status_code != 202: + verbose_logger.error( + "Datadog: unexpected response status_code=%s, text=%s", + response.status_code, + response.text, + ) + return self._undelivered(chunk, pending) + + verbose_logger.debug( + "Datadog: delivered %s events, status_code=%s, text=%s", + len(chunk), + response.status_code, + response.text, + ) + return [] + + @staticmethod + def _undelivered(chunk: List, pending: List[List]) -> List: + return chunk + [event for remaining in reversed(pending) for event in remaining] + async def flush_queue(self): if self.flush_lock is None: return diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 63fa0e64695..18321c70f6e 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1607,6 +1607,90 @@ class Logging(LiteLLMLoggingBaseClass): ) -> Optional[float]: return self._response_cost_calculator(result=result, cache_hit=cache_hit) + @staticmethod + def _is_sync_litellm_request(litellm_params: dict) -> bool: + """True for sync SDK entrypoints (``completion``), false for async (``acompletion``, etc.).""" + return ( + litellm_params.get(CallTypes.acompletion.value, False) is not True + and litellm_params.get(CallTypes.aresponses.value, False) is not True + and litellm_params.get(CallTypes.aembedding.value, False) is not True + and litellm_params.get(CallTypes.aimage_generation.value, False) is not True + and litellm_params.get(CallTypes.atranscription.value, False) is not True + ) + + def _is_assembled_stream_success(self, result=None) -> bool: + """Final assembled stream export (not a per-chunk success call). + + Per-chunk callers pass a ``ModelResponseStream`` (or ``None``); the + final assembled response is any other non-``None`` value (typically a + ``ModelResponse``). Treating a chunk as the assembled response would + prematurely set the ``has_dispatched_final_stream_success`` dedup + guard and silently suppress the real final stream log. + """ + if self.stream is not True: + return False + if result is not None and not isinstance(result, ModelResponseStream): + return True + return ( + "async_complete_streaming_response" in self.model_call_details + or self.model_call_details.get("complete_streaming_response") is not None + ) + + async def dispatch_success_handlers( + self, + result=None, + start_time=None, + end_time=None, + cache_hit=None, + prefer_async_handlers: bool = False, + **kwargs, + ) -> None: + """Route success logging to async and/or sync handlers for this request. + + ``prefer_async_handlers`` only bypasses the sync-SDK-only shortcut (e.g. + ``async for`` on a stream from ``completion()``). Legacy string callbacks + still run via ``executor.submit(success_handler)`` when configured. + """ + from litellm.litellm_core_utils.thread_pool_executor import executor + + if self._is_assembled_stream_success(result): + if self.model_call_details.get("has_dispatched_final_stream_success"): + return + self.model_call_details["has_dispatched_final_stream_success"] = True + + litellm_params = self.model_call_details.get("litellm_params", {}) or {} + sync_sdk = self._is_sync_litellm_request(litellm_params) + passthrough = self.call_type == CallTypes.pass_through.value + if sync_sdk and not prefer_async_handlers and not passthrough: + self.success_handler( + result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + **kwargs, + ) + return + + await self.async_success_handler( + result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + **kwargs, + ) + + if not self._should_run_sync_callbacks_for_async_calls(): + return + + executor.submit( + self.success_handler, + result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + **kwargs, + ) + def should_run_logging( self, event_type: Literal[ @@ -2029,13 +2113,7 @@ class Logging(LiteLLMLoggingBaseClass): standard_logging_object=kwargs.get("standard_logging_object", None), ) litellm_params = self.model_call_details.get("litellm_params", {}) - is_sync_request = ( - litellm_params.get(CallTypes.acompletion.value, False) is not True - and litellm_params.get(CallTypes.aresponses.value, False) is not True - and litellm_params.get(CallTypes.aembedding.value, False) is not True - and litellm_params.get(CallTypes.aimage_generation.value, False) is not True - and litellm_params.get(CallTypes.atranscription.value, False) is not True - ) + is_sync_request = self._is_sync_litellm_request(litellm_params) try: ## BUILD COMPLETE STREAMED RESPONSE complete_streaming_response: Optional[ @@ -2491,9 +2569,11 @@ class Logging(LiteLLMLoggingBaseClass): print_verbose( "Logging Details LiteLLM-Async Success Call, cache_hit={}".format(cache_hit) ) - if not self.should_run_logging( + if not self._is_assembled_stream_success( + result + ) and not self.should_run_logging( event_type="async_success" - ): # prevent double logging + ): # prevent double logging (non-streaming) return ## CALCULATE COST FOR BATCH JOBS @@ -2943,13 +3023,7 @@ class Logging(LiteLLMLoggingBaseClass): ): # prevent double logging return litellm_params = self.model_call_details.get("litellm_params", {}) - is_sync_request = ( - litellm_params.get(CallTypes.acompletion.value, False) is not True - and litellm_params.get(CallTypes.aresponses.value, False) is not True - and litellm_params.get(CallTypes.aembedding.value, False) is not True - and litellm_params.get(CallTypes.aimage_generation.value, False) is not True - and litellm_params.get(CallTypes.atranscription.value, False) is not True - ) + is_sync_request = self._is_sync_litellm_request(litellm_params) try: start_time, end_time = self._failure_handler_helper_fn( diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index 6c749118dec..b7adda3a9a4 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -394,6 +394,22 @@ class LoggingCallbackManager: + litellm._async_failure_callback ) + def remove_callback_from_all_lists(self, obj, require_self=False) -> None: + """ + Remove a callback object from every callback list it may have been + promoted into, so a re-initialized callback leaves no stale instance behind. + """ + for callback_list in ( + litellm.callbacks, + litellm.success_callback, + litellm.failure_callback, + litellm._async_success_callback, + litellm._async_failure_callback, + ): + self.remove_callback_from_list_by_object( + callback_list, obj, require_self=require_self + ) + def get_active_additional_logging_utils_from_custom_logger( self, ) -> Set[AdditionalLoggingUtils]: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index fa7faf3035d..29c0d0629e8 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1808,8 +1808,10 @@ class CustomStreamWrapper: processed_chunk, None, None, cache_hit ) ) - ## SYNC LOGGING - self.logging_obj.success_handler(processed_chunk, None, None, cache_hit) + ## SYNC LOGGING — only for sync SDK entrypoints; async proxy paths export via async_success_handler + litellm_params = self.logging_obj.model_call_details.get("litellm_params", {}) + if self.logging_obj._is_sync_litellm_request(litellm_params): + self.logging_obj.success_handler(processed_chunk, None, None, cache_hit) def finish_reason_handler(self): model_response = self.model_response_creator() @@ -2206,23 +2208,19 @@ class CustomStreamWrapper: cache_hit, ) else: + # prefer_async_handlers routes CustomLogger to async_success_handler + # when consumers use ``async for`` on sync-SDK streams. Legacy string + # callbacks still run via executor.submit inside dispatch_success_handlers. asyncio.create_task( - self.logging_obj.async_success_handler( + self.logging_obj.dispatch_success_handlers( complete_streaming_response, cache_hit=cache_hit, start_time=None, end_time=None, + prefer_async_handlers=True, ) ) - executor.submit( - self.logging_obj.success_handler, - complete_streaming_response, - cache_hit=cache_hit, - start_time=None, - end_time=None, - ) - raise StopAsyncIteration # Re-raise StopIteration else: self.sent_last_chunk = True diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9a34e6359c5..09cf9656b93 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3683,6 +3683,7 @@ class ProxyException(Exception): provider_specific_fields: Optional[dict] = None, ): self.message = str(message) + super().__init__(self.message) self.type = type self.param = param self.openai_code = openai_code or code diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index ef1d64335b4..4e26af7fe03 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1290,7 +1290,7 @@ class ProxyBaseLLMRequestProcessing: # (ProxyLogging._fire_deferred_stream_logging) fires the # closure after the full streaming pipeline finishes. # The closure runs non-apply_guardrail hooks on the - # assembled response, then fires both logging handlers. + # assembled response, then fires success logging. # Only for CustomStreamWrapper — raw async generators from # passthrough routes bypass CSW and would orphan the closure. from litellm.litellm_core_utils.streaming_handler import ( @@ -1411,33 +1411,18 @@ class ProxyBaseLLMRequestProcessing: logging_obj._on_deferred_stream_complete = None # type: ignore[union-attr] try: asyncio.create_task( - logging_obj.async_success_handler( + logging_obj.dispatch_success_handlers( response, cache_hit=None, start_time=None, end_time=None, + prefer_async_handlers=True, ) ) except Exception as e: verbose_proxy_logger.exception( "Error in orphaned streaming async logging: %s", e ) - try: - from litellm.litellm_core_utils.thread_pool_executor import ( - executor as _exc, - ) - - _exc.submit( - logging_obj.success_handler, - response, - cache_hit=None, - start_time=None, - end_time=None, - ) - except Exception as e: - verbose_proxy_logger.exception( - "Error in orphaned streaming sync logging: %s", e - ) # Always return the client-requested model name (not provider-prefixed internal identifiers) # for OpenAI-compatible responses. @@ -1639,7 +1624,7 @@ class ProxyBaseLLMRequestProcessing: ) -> None: """ Run non-streaming post-call guardrail hooks on an assembled streaming - response, then fire both async and sync logging handlers. + response, then fire success logging via ``dispatch_success_handlers``. Called by ProxyLogging._fire_deferred_stream_logging after the full streaming pipeline (including unified_guardrail end-of-stream blocks) @@ -1655,8 +1640,6 @@ class ProxyBaseLLMRequestProcessing: Extracted as a static method so tests can call the production implementation directly rather than reimplementing the closure. """ - from litellm.litellm_core_utils.thread_pool_executor import executor - _response = assembled_response try: from litellm.proxy.proxy_server import llm_router as _global_llm_router @@ -1715,31 +1698,23 @@ class ProxyBaseLLMRequestProcessing: ) finally: try: + # Proxy streaming always runs in async context and proxy spend + # logging is async-only; force async dispatch so DB/spend + # callbacks fire regardless of the call-type heuristic in + # _is_sync_litellm_request (which only recognizes a subset of + # async markers stored in litellm_params). asyncio.create_task( - captured_logging_obj.async_success_handler( + captured_logging_obj.dispatch_success_handlers( _response, cache_hit=cache_hit, start_time=None, end_time=None, + prefer_async_handlers=True, ) ) except Exception as e: verbose_proxy_logger.exception( - "Error in deferred streaming async logging: %s", - e, - ) - - try: - executor.submit( - captured_logging_obj.success_handler, - _response, - cache_hit=cache_hit, - start_time=None, - end_time=None, - ) - except Exception as e: - verbose_proxy_logger.exception( - "Error in deferred streaming sync logging: %s", + "Error in deferred streaming success logging: %s", e, ) @@ -1824,6 +1799,13 @@ class ProxyBaseLLMRequestProcessing: except Exception: pass + if isinstance(e, ProxyException): + e.headers = { + **e.headers, + **{k: v if isinstance(v, str) else str(v) for k, v in headers.items()}, + } + raise e + if isinstance(e, HTTPException): raw_detail = getattr(e, "detail", str(e)) message, structured_fields = _serialize_http_exception_detail(raw_detail) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index c5b97db07a0..694198e86df 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Literal, import litellm from litellm import get_secret from litellm._logging import verbose_proxy_logger +from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy._types import CommonProxyErrors, LiteLLMPromptInjectionParams @@ -482,6 +483,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS = frozenset( "guardrail_config", "_guardrail_pipelines", "_pipeline_managed_guardrails", + PRE_CALL_EXECUTED_GUARDRAILS_KEY, "disable_global_guardrails", "disable_global_guardrail", "opted_out_global_guardrails", diff --git a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py index 5b5f91195e7..d70c8e4f310 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py @@ -9,7 +9,6 @@ import json import os from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional, Type, Union -from fastapi import HTTPException from pydantic import BaseModel from websockets.asyncio.client import ClientConnection, connect @@ -21,7 +20,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.guardrails._content_utils import ( apply_redacted_messages_back, build_inspection_messages, @@ -129,6 +128,16 @@ class AimGuardrail(CustomGuardrail): verbose_proxy_logger.error(f"Aim: {action_type} action") return data + @staticmethod + def _rejection(message: str, *, openai_code: str | None = None) -> ProxyException: + return ProxyException( + message=message, + type="invalid_request_error", + param=None, + code=400, + openai_code=openai_code, + ) + def _handle_block_action(self, analysis_result: Any, required_action: Any) -> None: detection_message = required_action.get("detection_message", None) verbose_proxy_logger.info( @@ -136,7 +145,7 @@ class AimGuardrail(CustomGuardrail): policies=list(analysis_result["policy_drill_down"].keys()), ), ) - raise HTTPException(status_code=400, detail=detection_message) + raise self._rejection(detection_message, openai_code="content_policy_violation") def _anonymize_request(self, res: Any, data: dict) -> dict: verbose_proxy_logger.info("Aim: anonymize action") @@ -148,14 +157,11 @@ class AimGuardrail(CustomGuardrail): # parts from a multimodal request — degrade to block so the # multimodal payload is never silently rewritten. if has_non_string_content(data): - raise HTTPException( - status_code=400, - detail=( - "Aim: anonymize action requested for multimodal input " - "but mask-in-place would drop non-text parts. Send the " - "request with plain string content to use anonymize, " - "or rely on block-mode policies." - ), + raise self._rejection( + "Aim: anonymize action requested for multimodal input " + "but mask-in-place would drop non-text parts. Send the " + "request with plain string content to use anonymize, " + "or rely on block-mode policies." ) redacted_messages = [ { @@ -287,9 +293,9 @@ class AimGuardrail(CustomGuardrail): if aim_output_guardrail_result and aim_output_guardrail_result.get( "detection_message" ): - raise HTTPException( - status_code=400, - detail=aim_output_guardrail_result.get("detection_message"), + raise self._rejection( + aim_output_guardrail_result.get("detection_message"), + openai_code="content_policy_violation", ) if aim_output_guardrail_result and aim_output_guardrail_result.get( "redacted_output" diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index aafcc5f1819..d589217a7ea 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -5,6 +5,8 @@ import os from datetime import datetime, timezone from typing import Any, Dict, List, Literal, Optional, Set, Type, cast +from pydantic import ValidationError + import litellm from litellm import Router from litellm._logging import verbose_proxy_logger @@ -598,21 +600,25 @@ class InMemoryGuardrailHandler: def delete_in_memory_guardrail(self, guardrail_id: str) -> None: """ Delete a guardrail in memory and remove from litellm callbacks. + + The callback is purged from every callback list, not just + litellm.callbacks: request handling promotes guardrail callbacks into the + success/failure/async lists, so removing it from only litellm.callbacks + leaves the old instance stranded in those lists on every re-initialization. """ # Remove from in-memory storage self.IN_MEMORY_GUARDRAILS.pop(guardrail_id, None) self._sources.pop(guardrail_id, None) - # Remove the callback from litellm.callbacks custom_guardrail_callback = self.guardrail_id_to_custom_guardrail.pop( guardrail_id, None ) - if custom_guardrail_callback: - litellm.logging_callback_manager.remove_callback_from_list_by_object( - callback_list=litellm.callbacks, - obj=custom_guardrail_callback, - require_self=False, - ) + if custom_guardrail_callback is None: + return + + litellm.logging_callback_manager.remove_callback_from_all_lists( + custom_guardrail_callback + ) def list_in_memory_guardrails(self) -> List[Guardrail]: """ @@ -654,6 +660,34 @@ class InMemoryGuardrailHandler: self.delete_in_memory_guardrail(guardrail_id) return stale_ids + @staticmethod + def _normalize_litellm_params_for_comparison( + params: Optional[Any], + ) -> Optional[Dict[str, Any]]: + """ + Render litellm_params to a canonical dict so an in-memory LitellmParams and + the raw dict loaded from the DB compare equal when they describe the same + config. The in-memory side is a LitellmParams whose model_dump() carries + every field default and coerces enums, while the DB side is the raw stored + dict holding only the keys originally provided. Comparing those two shapes + directly never matches, so each DB poll would re-initialize the guardrail + forever; normalizing both through LitellmParams keeps the diff meaningful. + """ + if params is None: + return None + if isinstance(params, LitellmParams): + return params.model_dump() + if isinstance(params, dict): + try: + return LitellmParams(**params).model_dump() + except ValidationError as e: + verbose_proxy_logger.warning( + f"Could not normalize guardrail litellm_params for comparison; " + f"treating the guardrail as changed. Error: {e}" + ) + return params + return params + def _has_guardrail_params_changed( self, guardrail_id: str, new_guardrail: Guardrail ) -> bool: @@ -670,19 +704,11 @@ class InMemoryGuardrailHandler: return True # Compare litellm_params - existing_params = existing.get("litellm_params") - new_params = new_guardrail.get("litellm_params") - - # Convert to dicts for comparison - existing_dict = ( - existing_params.model_dump() - if isinstance(existing_params, LitellmParams) - else existing_params + existing_dict = self._normalize_litellm_params_for_comparison( + existing.get("litellm_params") ) - new_dict = ( - new_params.model_dump() - if isinstance(new_params, LitellmParams) - else new_params + new_dict = self._normalize_litellm_params_for_comparison( + new_guardrail.get("litellm_params") ) # Compare and identify specific differences diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 0d27b283c47..ed1fc76cd63 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -13,6 +13,7 @@ from starlette.datastructures import Headers import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging +from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import is_url_destination_allowed_by_host @@ -161,6 +162,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS = ( "secret_fields", "_guardrail_pipelines", "_pipeline_managed_guardrails", + PRE_CALL_EXECUTED_GUARDRAILS_KEY, ) _UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS = frozenset( diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 3224cf4d39a..a912a88a993 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -150,6 +150,13 @@ class AnthropicPassthroughLoggingHandler: handles streaming and non-streaming responses """ + # Only record complete_streaming_response for actual streaming responses. + # perform_redaction scrubs this field only when stream is True, so setting + # it on a non-streaming response would bypass message redaction. + if logging_obj.model_call_details.get("stream") is True: + logging_obj.model_call_details["complete_streaming_response"] = ( + litellm_model_response + ) try: # Get custom_llm_provider from logging object if available (e.g., azure_ai for Azure Anthropic) custom_llm_provider = logging_obj.model_call_details.get( diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 26f89e6b315..9c196b9cbd3 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -871,6 +871,9 @@ async def pass_through_request( # noqa: PLR0915 ) if stream: + logging_obj.stream = True + logging_obj.model_call_details["stream"] = True + if is_multipart: response = ( await HttpPassThroughEndpointHelpers.make_multipart_http_request( @@ -931,6 +934,9 @@ async def pass_through_request( # noqa: PLR0915 verbose_proxy_logger.debug("response.headers= %s", response.headers) if _is_streaming_response(response) is True: + logging_obj.stream = True + logging_obj.model_call_details["stream"] = True + try: response.raise_for_status() except httpx.HTTPStatusError as e: diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 235a38b75f9..33a6b719280 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -7,7 +7,6 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.proxy._types import PassThroughEndpointLoggingResultValues from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType @@ -145,25 +144,16 @@ class PassThroughStreamingHandler: end_time=end_time, model=model, ) - await litellm_logging_obj.async_success_handler( + # Always reached from an async context (anthropic_messages, + # google_genai, and proxy pass-through stream tasks). prefer_async_handlers + # keeps async-only loggers running even when call_type isn't pass_through + # and litellm_params lacks an async flag (e.g. aanthropic_messages). + await litellm_logging_obj.dispatch_success_handlers( result=standard_logging_response_object, start_time=start_time, end_time=end_time, cache_hit=False, - **kwargs, - ) - if ( - litellm_logging_obj._should_run_sync_callbacks_for_async_calls() - is False - ): - return - - executor.submit( - litellm_logging_obj.success_handler, - result=standard_logging_response_object, - end_time=end_time, - cache_hit=False, - start_time=start_time, + prefer_async_handlers=True, **kwargs, ) except Exception as e: diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 0bc0183aa7c..292871bae67 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -11,7 +11,6 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( PassthroughStandardLoggingPayload, ) from litellm.types.utils import StandardPassThroughResponseObject -from litellm.utils import executor as thread_pool_executor from .llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, @@ -94,19 +93,15 @@ class PassThroughEndpointLogging: cache_hit: bool, **kwargs, ): - """Helper function to handle both sync and async logging operations""" - # Submit to thread pool for sync logging - thread_pool_executor.submit( - logging_obj.success_handler, - standard_logging_response_object, - start_time, - end_time, - cache_hit, - **kwargs, - ) - - # Handle async logging - await logging_obj.async_success_handler( + """Log pass-through success via the shared async dispatch path.""" + # Always reached from pass_through_async_success_handler, which runs in + # an async context. call_type is "pass_through_endpoint" here, so the + # passthrough guard in dispatch_success_handlers already forces the + # async handler to run; pass prefer_async_handlers explicitly to match + # the streaming sibling (_route_streaming_logging_to_handler) and keep + # async-only loggers (e.g. the proxy spend logger) firing regardless of + # how the call-type classification evolves. + await logging_obj.dispatch_success_handlers( result=( json.dumps(result) if isinstance(result, dict) @@ -115,6 +110,7 @@ class PassThroughEndpointLogging: start_time=start_time, end_time=end_time, cache_hit=False, + prefer_async_handlers=True, **kwargs, ) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 3c5a1d67be4..e46e3e1dc9f 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -171,6 +171,10 @@ class PipelineExecutor: data=data, call_type=call_type, # type: ignore ) + if isinstance(callback, CustomGuardrail): + callback.mark_pre_call_hook_ran(data) + if isinstance(response, dict): + callback.mark_pre_call_hook_ran(response) elif mode == "post_call": response = await target.async_post_call_success_hook( user_api_key_dict=user_api_key_dict, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c75a58bb912..effd62da717 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1126,6 +1126,8 @@ class ProxyLogging: response=response, data=data, call_type=call_type ) + callback.mark_pre_call_hook_ran(data) + except Exception as e: status = "error" error_type = type(e).__name__ @@ -1965,7 +1967,7 @@ class ProxyLogging: litellm_call_id=request_data.get("litellm_call_id", ""), status="fail" ) if AlertType.llm_exceptions in self.alert_types and not isinstance( - original_exception, HTTPException + original_exception, (HTTPException, ProxyException) ): """ Just alert on LLM API exceptions. Do not alert on user errors @@ -2069,6 +2071,7 @@ class ProxyLogging: e.g should only return True for: - Authentication Errors from user_api_key_auth - HTTP HTTPException (rate limit errors) + - ProxyException (guardrail blocks, budget / rate-limit errors) """ ######################################################### @@ -2085,7 +2088,7 @@ class ProxyLogging: ): return False - return isinstance(original_exception, HTTPException) or ( + return isinstance(original_exception, (HTTPException, ProxyException)) or ( error_type == ProxyErrorTypes.auth_error ) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 282baff07fe..c3bf729238f 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3171,6 +3171,7 @@ all_litellm_params = ( "allowed_openai_params", "litellm_session_id", "use_litellm_proxy", + "use_chat_completions_api", "prompt_label", "shared_session", "search_tool_name", diff --git a/pyproject.toml b/pyproject.toml index b480e80df9c..f1acc0d19a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.87.3" +version = "1.87.4" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.14" @@ -257,7 +257,7 @@ source-exclude = [ profile = "black" [tool.commitizen] -version = "1.87.3" +version = "1.87.4" version_files = [ "pyproject.toml:^version", ] diff --git a/tests/litellm_utils_tests/test_logging_callback_manager.py b/tests/litellm_utils_tests/test_logging_callback_manager.py index d9540f8f850..d9bfca425e4 100644 --- a/tests/litellm_utils_tests/test_logging_callback_manager.py +++ b/tests/litellm_utils_tests/test_logging_callback_manager.py @@ -192,6 +192,29 @@ def test_remove_callback_from_list_by_object(): assert len(litellm._async_failure_callback) == 0 +def test_remove_callback_from_all_lists(): + manager = LoggingCallbackManager() + manager._reset_all_callbacks() + + class TestLogger(CustomLogger): + pass + + obj = TestLogger() + manager.add_litellm_callback(obj) + manager.add_litellm_success_callback(obj) + manager.add_litellm_failure_callback(obj) + manager.add_litellm_async_success_callback(obj) + manager.add_litellm_async_failure_callback(obj) + + manager.remove_callback_from_all_lists(obj) + + assert obj not in litellm.callbacks + assert obj not in litellm.success_callback + assert obj not in litellm.failure_callback + assert obj not in litellm._async_success_callback + assert obj not in litellm._async_failure_callback + + def test_reset_callbacks(callback_manager): # Add various callbacks callback_manager.add_litellm_callback("test") diff --git a/tests/local_testing/test_aim_guardrails.py b/tests/local_testing/test_aim_guardrails.py index 31416c565c1..2cb7f9cd357 100644 --- a/tests/local_testing/test_aim_guardrails.py +++ b/tests/local_testing/test_aim_guardrails.py @@ -6,10 +6,10 @@ import sys from unittest.mock import AsyncMock, patch, call import pytest -from fastapi.exceptions import HTTPException from httpx import Request, Response from litellm import DualCache +from litellm.proxy._types import ProxyException from litellm.proxy.guardrails.guardrail_hooks.aim.aim import ( AimGuardrail, AimGuardrailMissingSecrets, @@ -101,7 +101,7 @@ async def test_block_callback(mode: str): ], } - with pytest.raises(HTTPException, match="Jailbreak detected"): + with pytest.raises(ProxyException, match="Jailbreak detected") as exc_info: with patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", return_value=Response( @@ -135,6 +135,137 @@ async def test_block_callback(mode: str): call_type="completion", ) + exc = exc_info.value + assert exc.code == "400" + assert exc.type == "invalid_request_error" + assert exc.param is None + assert exc.openai_code == "content_policy_violation" + + +@pytest.mark.asyncio +async def test_output_block_raises_proxy_exception(): + """An output-side block is a content-policy violation, like the input block: + it must surface a conformant ProxyException, not a bare HTTPException whose + type/param serialize as the literal string "None". Regression for LIT-3751.""" + init_guardrails_v2( + all_guardrails=[ + { + "guardrail_name": "gibberish-guard", + "litellm_params": { + "guardrail": "aim", + "mode": "post_call", + "api_key": "hs-aim-key", + }, + }, + ], + config_file_path="", + ) + aim_guardrails = [ + callback for callback in litellm.callbacks if isinstance(callback, AimGuardrail) + ] + assert len(aim_guardrails) == 1 + aim_guardrail = aim_guardrails[0] + + block_on_output = Response( + json={ + "analysis_result": {"policy_drill_down": {"PII": {}}}, + "required_action": { + "action_type": "block_action", + "detection_message": "Output blocked: leaked secret", + "policy_name": "blocking policy", + }, + }, + status_code=200, + request=Request(method="POST", url="http://aim"), + ) + response = ModelResponse( + choices=[ + { + "finish_reason": "stop", + "index": 0, + "message": {"content": "here is the secret", "role": "assistant"}, + } + ] + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=block_on_output, + ): + with pytest.raises(ProxyException, match="Output blocked") as exc_info: + await aim_guardrail.async_post_call_success_hook( + data={"messages": [{"role": "user", "content": "tell me a secret"}]}, + response=response, + user_api_key_dict=UserAPIKeyAuth(), + ) + + exc = exc_info.value + assert exc.code == "400" + assert exc.type == "invalid_request_error" + assert exc.param is None + assert exc.openai_code == "content_policy_violation" + + +@pytest.mark.asyncio +async def test_anonymize_multimodal_rejection_raises_proxy_exception(): + """Anonymize on multimodal input degrades to a 400 because mask-in-place would + drop non-text parts. That is a usage error, not a content-policy violation, so + it must raise a conformant ProxyException WITHOUT the content_policy_violation + code. Regression for LIT-3751.""" + init_guardrails_v2( + all_guardrails=[ + { + "guardrail_name": "gibberish-guard", + "litellm_params": { + "guardrail": "aim", + "mode": "pre_call", + "api_key": "hs-aim-key", + }, + }, + ], + config_file_path="", + ) + aim_guardrails = [ + callback for callback in litellm.callbacks if isinstance(callback, AimGuardrail) + ] + assert len(aim_guardrails) == 1 + aim_guardrail = aim_guardrails[0] + + data = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Hi my name is Brian"}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}, + }, + ], + }, + ], + } + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=response_with_detections, + ): + with pytest.raises( + ProxyException, match="anonymize action requested for multimodal" + ) as exc_info: + await aim_guardrail.async_pre_call_hook( + data=data, + cache=DualCache(), + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", + ) + + exc = exc_info.value + assert exc.code == "400" + assert exc.type == "invalid_request_error" + assert exc.param is None + assert exc.openai_code != "content_policy_violation" + @pytest.mark.asyncio @pytest.mark.parametrize("mode", ["pre_call", "during_call"]) diff --git a/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py b/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py index 455c72ff636..5ab0319da47 100644 --- a/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py +++ b/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py @@ -318,6 +318,7 @@ def test_handle_logging_anthropic_collected_chunks(all_chunks): from litellm.types.utils import ModelResponse litellm_logging_obj = Mock() + litellm_logging_obj.model_call_details = {} pass_through_logging_obj = Mock() sent_args = { diff --git a/tests/pass_through_unit_tests/test_unit_test_streaming.py b/tests/pass_through_unit_tests/test_unit_test_streaming.py index 38b650121bd..63965320f2b 100644 --- a/tests/pass_through_unit_tests/test_unit_test_streaming.py +++ b/tests/pass_through_unit_tests/test_unit_test_streaming.py @@ -97,6 +97,123 @@ async def test_chunk_processor_yields_raw_bytes(endpoint_type, url_route): ), "Collected chunks do not match raw chunks" +@pytest.mark.asyncio +async def test_route_streaming_logging_runs_async_handler_for_sdk_passthrough(): + """ + SDK pass-through streaming (anthropic_messages, google generate_content) must run + the async success handler so async-only loggers record the assembled stream. + + Regression for duplicate-trace dedupe: dispatch_success_handlers treated these as + sync SDK requests because call_type is not ``pass_through_endpoint`` and + litellm_params carries no ``acompletion`` flag, so only the sync success_handler + ran and CustomLogger.async_log_success_event never fired. + """ + import time + + from litellm.types.utils import CallTypes + + logging_obj = LiteLLMLoggingObj( + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type=CallTypes.anthropic_messages.value, + start_time=time.time(), + litellm_call_id="test-id", + function_id="fn", + ) + logging_obj.model_call_details["litellm_params"] = {"anthropic_messages": True} + + with ( + patch.object( + PassThroughStreamingHandler, + "_build_passthrough_logging_result", + return_value=({"id": "slp"}, {}), + ), + patch.object( + logging_obj, "async_success_handler", new_callable=AsyncMock + ) as mock_async, + patch.object( + logging_obj, "success_handler", new_callable=MagicMock + ) as mock_sync, + patch.object( + logging_obj, + "_should_run_sync_callbacks_for_async_calls", + return_value=False, + ), + ): + await PassThroughStreamingHandler._route_streaming_logging_to_handler( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + request_body={}, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + raw_bytes=[], + end_time=datetime.now(), + ) + + mock_async.assert_awaited_once() + mock_sync.assert_not_called() + + +@pytest.mark.asyncio +async def test_handle_logging_runs_async_handler_for_passthrough(): + """ + Non-streaming pass-through logging (_handle_logging) must always run the + async success handler so async-only loggers (e.g. the proxy spend logger) + record the request. + + _handle_logging is only ever reached from pass_through_async_success_handler + (an async context), so it forces async dispatch via prefer_async_handlers. + This pins that contract independent of the call-type classification: even a + call_type that _is_sync_litellm_request would classify as sync (here + "completion" with no async marker in litellm_params) must still reach + async_success_handler. Without prefer_async_handlers=True the sync-only + branch would return early and async_log_success_event would never fire. + """ + import time + + from litellm.types.utils import CallTypes + + logging_obj = LiteLLMLoggingObj( + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type=CallTypes.completion.value, + start_time=time.time(), + litellm_call_id="test-id", + function_id="fn", + ) + logging_obj.model_call_details["litellm_params"] = {} + + handler = PassThroughEndpointLogging() + + with ( + patch.object( + logging_obj, "async_success_handler", new_callable=AsyncMock + ) as mock_async, + patch.object( + logging_obj, "success_handler", new_callable=MagicMock + ) as mock_sync, + patch.object( + logging_obj, + "_should_run_sync_callbacks_for_async_calls", + return_value=False, + ), + ): + await handler._handle_logging( + logging_obj=logging_obj, + standard_logging_response_object={"id": "slp"}, + result="", + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + ) + + mock_async.assert_awaited_once() + mock_sync.assert_not_called() + + def test_convert_raw_bytes_to_str_lines(): """ Test that the _convert_raw_bytes_to_str_lines method correctly converts raw bytes to a list of strings diff --git a/tests/proxy_unit_tests/test_proxy_reject_logging.py b/tests/proxy_unit_tests/test_proxy_reject_logging.py index 51a92fa3b4b..e0b575f4a71 100644 --- a/tests/proxy_unit_tests/test_proxy_reject_logging.py +++ b/tests/proxy_unit_tests/test_proxy_reject_logging.py @@ -95,6 +95,21 @@ router = Router( ) +def _register_proxy_test_logger(callback_logger: testLogger) -> None: + """ + Register the test logger on global callback lists. + + ``function_setup`` dedupes by object identity; each parametrized case + constructs a new ``testLogger`` and must replace the global lists, not + only ``litellm.callbacks``. + """ + litellm.callbacks = [callback_logger] + litellm.success_callback = [callback_logger] + litellm.failure_callback = [callback_logger] + litellm._async_success_callback = [callback_logger] + litellm._async_failure_callback = [callback_logger] + + @pytest.mark.parametrize( "route, body", [ @@ -115,7 +130,7 @@ router = Router( "/v1/embeddings", { "input": "The food was delicious and the waiter...", - "model": "text-embedding-ada-002", + "model": "fake-model", "encoding_format": "float", }, ), @@ -133,7 +148,7 @@ async def test_chat_completion_request_with_redaction(route, body): setattr(proxy_server, "llm_router", router) _test_logger = testLogger() - litellm.callbacks = [_test_logger] + _register_proxy_test_logger(_test_logger) litellm.set_verbose = True # Prepare the query string diff --git a/tests/test_litellm/integrations/datadog/test_datadog_logger_batching.py b/tests/test_litellm/integrations/datadog/test_datadog_logger_batching.py index e4d7227cc88..d1c7a4032fb 100644 --- a/tests/test_litellm/integrations/datadog/test_datadog_logger_batching.py +++ b/tests/test_litellm/integrations/datadog/test_datadog_logger_batching.py @@ -1,10 +1,49 @@ from unittest.mock import AsyncMock, Mock, patch +import httpx import pytest from httpx import Request, Response from litellm.integrations.datadog.datadog import DataDogLogger -from litellm.types.integrations.datadog import DatadogPayload +from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError +from litellm.types.integrations.datadog import DD_MAX_BATCH_SIZE, DatadogPayload + + +def _payloads(n): + return [ + DatadogPayload( + ddsource="litellm", + ddtags="env:test", + hostname="host", + message=f'{{"event": {i}}}', + service="svc", + status="info", + ) + for i in range(n) + ] + + +def _raised_413(): + request = Request("POST", "https://example.com") + response = Response(413, request=request, text="Payload Too Large") + return MaskedHTTPStatusError( + httpx.HTTPStatusError("413", request=request, response=response) + ) + + +def _make_send(max_ok, delivered, *, raise_413=True): + """Datadog double: 413 batches larger than max_ok, 202 (recording delivery) otherwise.""" + + async def _send(data): + request = Request("POST", "https://example.com") + if len(data) > max_ok: + if raise_413: + raise _raised_413() + return Response(413, request=request, text="Payload Too Large") + delivered.extend(event["message"] for event in data) + return Response(202, request=request, text="Accepted") + + return _send @pytest.fixture @@ -75,40 +114,152 @@ async def test_failure_hook_threshold_flush_uses_flush_queue(datadog_env): @pytest.mark.asyncio -async def test_async_send_batch_requeues_events_on_413(datadog_env): +async def test_413_splits_oversized_batch_and_delivers_every_event(datadog_env): + """A raised 413 (the real httpx path) halves the batch until each piece is accepted.""" with patch("asyncio.create_task"): logger = DataDogLogger() - logger.log_queue = [ - DatadogPayload( - ddsource="litellm", - ddtags="env:test", - hostname="host", - message=f'{{"event": {i}}}', - service="svc", - status="info", + logger.log_queue = _payloads(4) + delivered: list = [] + logger.async_send_compressed_data = AsyncMock(side_effect=_make_send(1, delivered)) + + await logger.async_send_batch() + + assert sorted(delivered) == [f'{{"event": {i}}}' for i in range(4)] + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_413_does_not_requeue_oversized_batch(datadog_env): + """Regression for the infinite 413 loop: an undeliverable batch must not be re-queued.""" + with patch("asyncio.create_task"): + logger = DataDogLogger() + + logger.log_queue = _payloads(4) + logger.async_send_compressed_data = AsyncMock(side_effect=_make_send(0, [])) + + await logger.async_send_batch() + await logger.async_send_batch() + + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_413_drops_single_oversized_event(datadog_env): + with patch("asyncio.create_task"): + logger = DataDogLogger() + + logger.log_queue = _payloads(1) + send = AsyncMock(side_effect=_make_send(0, [])) + logger.async_send_compressed_data = send + + await logger.async_send_batch() + + assert send.await_count == 1 + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_413_returned_response_also_splits(datadog_env): + """Defensive path: a 413 returned (not raised) is handled the same way.""" + with patch("asyncio.create_task"): + logger = DataDogLogger() + + logger.log_queue = _payloads(4) + delivered: list = [] + logger.async_send_compressed_data = AsyncMock( + side_effect=_make_send(1, delivered, raise_413=False) + ) + + await logger.async_send_batch() + + assert sorted(delivered) == [f'{{"event": {i}}}' for i in range(4)] + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_partial_delivery_then_transient_error_requeues_only_undelivered( + datadog_env, +): + """A transient error after a partial split delivery must not duplicate delivered events.""" + with patch("asyncio.create_task"): + logger = DataDogLogger() + + logger.log_queue = _payloads(4) + delivered: list = [] + + async def _send(data): + messages = [event["message"] for event in data] + if len(data) > 2: + raise _raised_413() + if messages == ['{"event": 2}', '{"event": 3}']: + raise RuntimeError("transient network error") + delivered.extend(messages) + return Response( + 202, request=Request("POST", "https://example.com"), text="Accepted" ) - for i in range(2) + + logger.async_send_compressed_data = AsyncMock(side_effect=_send) + + await logger.async_send_batch() + + assert delivered == ['{"event": 0}', '{"event": 1}'] + assert [event["message"] for event in logger.log_queue] == [ + '{"event": 2}', + '{"event": 3}', ] + +@pytest.mark.asyncio +async def test_unexpected_non_202_status_requeues(datadog_env): + """A non-413, non-202 response is treated as undelivered and re-queued.""" + with patch("asyncio.create_task"): + logger = DataDogLogger() + + logger.log_queue = _payloads(2) logger.async_send_compressed_data = AsyncMock( return_value=Response( - 413, - request=Request("POST", "https://example.com"), - text="Payload Too Large", + 200, request=Request("POST", "https://example.com"), text="OK" ) ) await logger.async_send_batch() - assert logger.async_send_compressed_data.await_count == 1 - assert len(logger.log_queue) == 2 assert [event["message"] for event in logger.log_queue] == [ '{"event": 0}', '{"event": 1}', ] +@pytest.mark.parametrize( + "value, expected", + [ + ("50", 50), + ("1", 1), + ("0", 1), + ("-5", 1), + (str(DD_MAX_BATCH_SIZE + 100), DD_MAX_BATCH_SIZE), + ("not_an_int", DD_MAX_BATCH_SIZE), + ], +) +def test_dd_batch_size_env_resolution(monkeypatch, value, expected): + monkeypatch.setenv("DD_API_KEY", "test_api_key") + monkeypatch.setenv("DD_SITE", "test.datadoghq.com") + monkeypatch.setenv("DD_BATCH_SIZE", value) + with patch("asyncio.create_task"): + logger = DataDogLogger() + assert logger.batch_size == expected + + +def test_dd_batch_size_defaults_to_max(monkeypatch): + monkeypatch.setenv("DD_API_KEY", "test_api_key") + monkeypatch.setenv("DD_SITE", "test.datadoghq.com") + monkeypatch.delenv("DD_BATCH_SIZE", raising=False) + with patch("asyncio.create_task"): + logger = DataDogLogger() + assert logger.batch_size == DD_MAX_BATCH_SIZE + + @pytest.mark.asyncio async def test_async_send_batch_handles_empty_queue(datadog_env): with patch("asyncio.create_task"): diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index 1a4d03528e7..6afe5efc54d 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -1087,3 +1087,357 @@ async def test_anthropic_cache_control_hook_string_negative_index(): f"Expected cachePoint in last message content, got: {last_message_content}. " "String index '-1' was not parsed correctly (str.isdigit() returns False for negative strings)." ) + + +def _count_cache_control(messages: List[AllMessageValues]) -> int: + """Count cache_control breakpoints across messages (message + content level).""" + count = 0 + for message in messages: + if message.get("cache_control") is not None: + count += 1 + content = message.get("content") + if isinstance(content, list): + for block in content: + if isinstance(block, dict) and block.get("cache_control") is not None: + count += 1 + return count + + +def _build_injection_points(): + return [ + { + "location": "message", + "role": "system", + "control": {"type": "ephemeral", "ttl": "1h"}, + }, + { + "location": "message", + "index": -1, + "control": {"type": "ephemeral", "ttl": "5m"}, + }, + ] + + +def test_cache_control_hook_caps_at_four_blocks_with_client_cache_control(): + """Regression for LIT-3667 / Anthropic 'A maximum of 4 blocks ... Found 5'. + + A Hermes-style request already carries 4 client cache_control breakpoints on + its system messages. With both auto-inject points configured the hook must + NOT add a 5th breakpoint, and must NOT overwrite the client's existing + breakpoints (TTL must be preserved). + """ + hook = AnthropicCacheControlHook() + + messages: List[AllMessageValues] = [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": f"System block {i}", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + for i in range(4) + ] + messages.append({"role": "user", "content": "hello"}) + + _, processed, _ = hook.get_chat_completion_prompt( + model="bedrock/us.anthropic.claude-opus-4-6-v1:0", + messages=messages, + non_default_params={ + "cache_control_injection_points": _build_injection_points() + }, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + ) + + assert ( + _count_cache_control(processed) == 4 + ), "Hook must cap cache_control at Anthropic's limit of 4 blocks" + + # Client TTL on system blocks must be preserved (not overwritten by config). + for i in range(4): + assert processed[i]["content"][-1]["cache_control"] == { + "type": "ephemeral", + "ttl": "1h", + } + + # The last (user) message must not receive a 5th breakpoint. + user_message = processed[-1] + assert user_message.get("cache_control") is None + user_content = user_message.get("content") + if isinstance(user_content, list): + assert all( + block.get("cache_control") is None + for block in user_content + if isinstance(block, dict) + ) + + +def test_cache_control_hook_caps_at_four_blocks_without_client_cache_control(): + """Four plain system messages + role:system + index:-1 must stay at 4 blocks. + + role:system fills all four slots, so the index:-1 point is skipped. + """ + hook = AnthropicCacheControlHook() + + messages: List[AllMessageValues] = [ + {"role": "system", "content": f"System {i}"} for i in range(4) + ] + messages.append({"role": "user", "content": "hello"}) + + _, processed, _ = hook.get_chat_completion_prompt( + model="bedrock/us.anthropic.claude-opus-4-6-v1:0", + messages=messages, + non_default_params={ + "cache_control_injection_points": _build_injection_points() + }, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + ) + + assert _count_cache_control(processed) == 4 + # All four system messages cached; user message skipped (limit reached). + assert all(processed[i].get("cache_control") is not None for i in range(4)) + assert processed[-1].get("cache_control") is None + + +def test_cache_control_hook_does_not_overwrite_existing_cache_control(): + """If a targeted message already has client cache_control, do not inject.""" + hook = AnthropicCacheControlHook() + + messages: List[AllMessageValues] = [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "Cached by client", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + {"role": "user", "content": "hello"}, + ] + + _, processed, _ = hook.get_chat_completion_prompt( + model="bedrock/us.anthropic.claude-opus-4-6-v1:0", + messages=messages, + # Target the already-cached system message with a different TTL. + non_default_params={ + "cache_control_injection_points": [ + { + "location": "message", + "index": 0, + "control": {"type": "ephemeral", "ttl": "5m"}, + } + ] + }, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + ) + + # Client's 1h TTL must be preserved, not replaced by the config's 5m. + assert processed[0]["content"][-1]["cache_control"] == { + "type": "ephemeral", + "ttl": "1h", + } + assert _count_cache_control(processed) == 1 + + +@pytest.mark.asyncio +async def test_cache_control_hook_bedrock_payload_caps_cachepoints_at_four(): + """End-to-end: outgoing Bedrock payload must not exceed 4 cachePoint blocks. + + Reproduces the customer report where 4 client cache_control system blocks + plus auto-inject produced 5 cachePoint blocks and Bedrock returned 400. + """ + with patch.dict( + os.environ, + { + "AWS_ACCESS_KEY_ID": "fake_access_key_id", + "AWS_SECRET_ACCESS_KEY": "fake_secret_access_key", + "AWS_REGION_NAME": "us-east-1", + }, + ): + litellm.callbacks = [AnthropicCacheControlHook()] + + mock_response = MagicMock() + mock_response.json.return_value = { + "output": {"message": {"role": "assistant", "content": "ok"}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 100, "outputTokens": 4, "totalTokens": 104}, + } + mock_response.status_code = 200 + + client = AsyncHTTPHandler() + with patch.object(client, "post", return_value=mock_response) as mock_post: + messages = [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": f"System block {i}", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + for i in range(4) + ] + messages.append({"role": "user", "content": "hello"}) + + await litellm.acompletion( + model="bedrock/us.anthropic.claude-opus-4-6-v1:0", + messages=messages, + max_tokens=32, + cache_control_injection_points=_build_injection_points(), + client=client, + ) + + request_body = json.loads(mock_post.call_args.kwargs["data"]) + + cache_points = sum( + 1 + for block in request_body.get("system", []) + if isinstance(block, dict) and "cachePoint" in block + ) + for msg in request_body.get("messages", []): + content = msg.get("content", []) + if isinstance(content, list): + cache_points += sum( + 1 + for block in content + if isinstance(block, dict) and "cachePoint" in block + ) + + assert cache_points <= 4, ( + f"Bedrock payload exceeded Anthropic's 4 cache_control block limit: " + f"found {cache_points} cachePoint blocks" + ) + + +def test_cache_control_hook_reserves_slot_for_tool_config_point(): + """A tool_config injection point consumes one of the 4 slots downstream. + + With role:system targeting 4 system messages plus a tool_config point, the + hook must inject at most 3 message-level blocks so the tool_config cachePoint + appended by the Bedrock transform keeps the total at 4, not 5. + """ + hook = AnthropicCacheControlHook() + + messages: List[AllMessageValues] = [ + {"role": "system", "content": f"System {i}"} for i in range(4) + ] + messages.append({"role": "user", "content": "hello"}) + + _, processed, non_default_params = hook.get_chat_completion_prompt( + model="bedrock/us.anthropic.claude-opus-4-6-v1:0", + messages=messages, + non_default_params={ + "cache_control_injection_points": [ + { + "location": "message", + "role": "system", + "control": {"type": "ephemeral", "ttl": "1h"}, + }, + {"location": "tool_config"}, + ] + }, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + ) + + assert _count_cache_control(processed) == 3 + # The tool_config point is passed through for the provider transform. + assert non_default_params["cache_control_injection_points"] == [ + {"location": "tool_config"} + ] + + +@pytest.mark.asyncio +async def test_cache_control_hook_bedrock_payload_caps_with_tool_config_point(): + """End-to-end: message + tool_config injection must not exceed 4 cachePoints.""" + with patch.dict( + os.environ, + { + "AWS_ACCESS_KEY_ID": "fake_access_key_id", + "AWS_SECRET_ACCESS_KEY": "fake_secret_access_key", + "AWS_REGION_NAME": "us-east-1", + }, + ): + litellm.callbacks = [AnthropicCacheControlHook()] + + mock_response = MagicMock() + mock_response.json.return_value = { + "output": {"message": {"role": "assistant", "content": "ok"}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 100, "outputTokens": 4, "totalTokens": 104}, + } + mock_response.status_code = 200 + + client = AsyncHTTPHandler() + with patch.object(client, "post", return_value=mock_response) as mock_post: + messages = [ + {"role": "system", "content": f"System block {i}"} for i in range(4) + ] + messages.append({"role": "user", "content": "What is the weather?"}) + + await litellm.acompletion( + model="bedrock/us.anthropic.claude-opus-4-6-v1:0", + messages=messages, + max_tokens=32, + tools=[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a location", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + }, + } + ], + cache_control_injection_points=[ + { + "location": "message", + "role": "system", + "control": {"type": "ephemeral", "ttl": "1h"}, + }, + {"location": "tool_config"}, + ], + client=client, + ) + + request_body = json.loads(mock_post.call_args.kwargs["data"]) + + cache_points = sum( + 1 + for block in request_body.get("system", []) + if isinstance(block, dict) and "cachePoint" in block + ) + for msg in request_body.get("messages", []): + content = msg.get("content", []) + if isinstance(content, list): + cache_points += sum( + 1 + for block in content + if isinstance(block, dict) and "cachePoint" in block + ) + for tool in request_body.get("toolConfig", {}).get("tools", []): + if isinstance(tool, dict) and "cachePoint" in tool: + cache_points += 1 + + assert cache_points <= 4, ( + f"Bedrock payload exceeded Anthropic's 4 cache_control block limit " + f"when mixing message and tool_config injection: found {cache_points}" + ) diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index a881044dc18..09469cd33e5 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -84,6 +84,112 @@ class TestCustomGuardrailDeploymentHook: assert result["messages"] == mock_result["messages"] assert result["messages"] != original_messages + @pytest.mark.asyncio + async def test_deployment_hook_skips_when_pre_call_already_ran(self): + """The deployment hook must not re-run async_pre_call_hook once the proxy + pre-call loop has already run it for this request.""" + + class CountingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__(guardrail_name="g1", default_on=True) + self.pre_call_count = 0 + + async def async_pre_call_hook( + self, user_api_key_dict, cache, data, call_type + ): + self.pre_call_count += 1 + return data + + guardrail = CountingGuardrail() + kwargs = { + "messages": [{"role": "user", "content": "hi"}], + "model": "gpt-3.5-turbo", + "guardrails": ["g1"], + "metadata": {}, + } + + guardrail.mark_pre_call_hook_ran(kwargs) + await guardrail.async_pre_call_deployment_hook( + kwargs=kwargs, call_type=CallTypes.completion + ) + + assert guardrail.pre_call_count == 0 + + @pytest.mark.asyncio + async def test_deployment_hook_runs_when_not_marked(self): + """Without the proxy marker (direct-SDK usage) the deployment hook is the + only execution path and must still run the guardrail exactly once.""" + + class CountingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__(guardrail_name="g1", default_on=True) + self.pre_call_count = 0 + + async def async_pre_call_hook( + self, user_api_key_dict, cache, data, call_type + ): + self.pre_call_count += 1 + return data + + guardrail = CountingGuardrail() + kwargs = { + "messages": [{"role": "user", "content": "hi"}], + "model": "gpt-3.5-turbo", + "guardrails": ["g1"], + "metadata": {}, + } + + await guardrail.async_pre_call_deployment_hook( + kwargs=kwargs, call_type=CallTypes.completion + ) + + assert guardrail.pre_call_count == 1 + + def test_mark_pre_call_hook_ran_uses_litellm_metadata(self): + """The marker is recorded in litellm_metadata when that is the metadata + bucket in use, and is then visible to the skip check.""" + from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY + + guardrail = CustomGuardrail(guardrail_name="g1") + kwargs = {"litellm_metadata": {}} + + guardrail.mark_pre_call_hook_ran(kwargs) + + assert kwargs["litellm_metadata"][PRE_CALL_EXECUTED_GUARDRAILS_KEY] + assert guardrail._pre_call_hook_already_ran(kwargs) is True + + @pytest.mark.asyncio + async def test_deployment_hook_ignores_forged_caller_marker(self): + """A direct-SDK caller controls request metadata but cannot know the + per-process token, so a hand-crafted marker must not suppress a + requested guardrail in async_pre_call_deployment_hook.""" + from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY + + class CountingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__(guardrail_name="g1", default_on=True) + self.pre_call_count = 0 + + async def async_pre_call_hook( + self, user_api_key_dict, cache, data, call_type + ): + self.pre_call_count += 1 + return data + + guardrail = CountingGuardrail() + kwargs = { + "messages": [{"role": "user", "content": "hi"}], + "model": "gpt-3.5-turbo", + "guardrails": ["g1"], + "metadata": {PRE_CALL_EXECUTED_GUARDRAILS_KEY: ["g1"]}, + } + + await guardrail.async_pre_call_deployment_hook( + kwargs=kwargs, call_type=CallTypes.completion + ) + + assert guardrail.pre_call_count == 1 + class TestCustomGuardrailShouldRunGuardrail: diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 07ab29c5231..b64cb7c6905 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1,6 +1,7 @@ import os import sys -from unittest.mock import MagicMock, patch +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -786,6 +787,211 @@ def test_success_handler_runs_sync_callbacks_for_sync_requests(logging_obj, call dummy_logger.log_stream_event.assert_not_called() +def test_is_sync_litellm_request(): + assert LitellmLogging._is_sync_litellm_request({}) is True + assert LitellmLogging._is_sync_litellm_request({"acompletion": True}) is False + + +@pytest.mark.asyncio +async def test_dispatch_success_handlers_invokes_callbacks_once_for_final_stream( + logging_obj, +): + """Second final-stream dispatch must not re-export (CSW + deferred guardrail paths).""" + import litellm + from litellm.integrations.custom_logger import CustomLogger + + class MockCallback(CustomLogger): + pass + + mock_callback = MockCallback() + original_async_callbacks = list(litellm._async_success_callback or []) + litellm._async_success_callback = [mock_callback] + + result = ModelResponse( + id="resp-dedupe", + model="gpt-4o-mini", + choices=[ + { + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + "index": 0, + } + ], + usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + ) + + try: + logging_obj.stream = True + logging_obj.model_call_details["litellm_params"] = {"acompletion": True} + + with ( + patch.object( + mock_callback, "async_log_success_event", new_callable=AsyncMock + ) as mock_async_log, + patch.object(mock_callback, "log_success_event") as mock_sync_log, + patch.object( + logging_obj, + "_success_handler_helper_fn", + return_value=(time.time(), time.time(), result), + ), + patch.object( + logging_obj, + "_get_assembled_streaming_response", + return_value=result, + ), + patch.object( + logging_obj, + "_should_run_sync_callbacks_for_async_calls", + return_value=True, + ), + ): + await logging_obj.dispatch_success_handlers(result=result) + await logging_obj.dispatch_success_handlers(result=result) + + mock_async_log.assert_awaited_once() + mock_sync_log.assert_not_called() + finally: + litellm._async_success_callback = original_async_callbacks + + +@pytest.mark.asyncio +async def test_dispatch_success_handlers_sync_path_invokes_callback_once_for_final_stream( + logging_obj, +): + """Sync dispatch path must also dedupe when dispatch is called twice.""" + import litellm + from litellm.integrations.custom_logger import CustomLogger + + class MockCallback(CustomLogger): + pass + + mock_callback = MockCallback() + original_success_callbacks = list(litellm.success_callback or []) + litellm.success_callback = [mock_callback] + + result = ModelResponse( + id="resp-sync-dedupe", + model="gpt-4o-mini", + choices=[ + { + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + "index": 0, + } + ], + usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + ) + + try: + logging_obj.stream = True + logging_obj.model_call_details["litellm_params"] = {} + + with ( + patch.object(mock_callback, "log_success_event") as mock_sync_log, + patch.object( + mock_callback, "async_log_success_event", new_callable=AsyncMock + ) as mock_async_log, + patch.object( + logging_obj, + "_success_handler_helper_fn", + return_value=(time.time(), time.time(), result), + ), + patch.object( + logging_obj, + "_get_assembled_streaming_response", + return_value=result, + ), + ): + await logging_obj.dispatch_success_handlers(result=result) + await logging_obj.dispatch_success_handlers(result=result) + + mock_sync_log.assert_called_once() + mock_async_log.assert_not_awaited() + finally: + litellm.success_callback = original_success_callbacks + + +@pytest.mark.asyncio +async def test_dispatch_prefer_async_handlers_runs_legacy_callbacks( + logging_obj, +): + """``prefer_async_handlers`` must not skip executor.submit for string callbacks.""" + result = ModelResponse( + id="resp-prefer-async", + model="gpt-4o-mini", + choices=[ + { + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + "index": 0, + } + ], + ) + + logging_obj.stream = True + logging_obj.model_call_details["litellm_params"] = {} + + with ( + patch.object( + logging_obj, "async_success_handler", new_callable=AsyncMock + ) as mock_async, + patch.object( + logging_obj, "success_handler", new_callable=MagicMock + ) as mock_sync, + patch.object( + logging_obj, + "_should_run_sync_callbacks_for_async_calls", + return_value=True, + ), + patch( + "litellm.litellm_core_utils.litellm_logging.executor.submit" + ) as mock_submit, + ): + await logging_obj.dispatch_success_handlers( + result=result, + prefer_async_handlers=True, + ) + + mock_async.assert_awaited_once() + mock_sync.assert_not_called() + mock_submit.assert_called_once() + + +@pytest.mark.asyncio +async def test_dispatch_success_handlers_invokes_async_callback_for_pass_through( + logging_obj, +): + """Pass-through must use async_success_handler (CustomLogger skips sync success_handler).""" + import litellm + from litellm.integrations.custom_logger import CustomLogger + from litellm.types.utils import CallTypes + + class MockCallback(CustomLogger): + pass + + mock_callback = MockCallback() + original_async_callbacks = list(litellm._async_success_callback or []) + litellm._async_success_callback = [mock_callback] + + logging_obj.call_type = CallTypes.pass_through.value + logging_obj.stream = False + logging_obj.model_call_details["litellm_params"] = {} + + try: + with ( + patch.object( + mock_callback, "async_log_success_event", new_callable=AsyncMock + ) as mock_async_log, + patch.object(mock_callback, "log_success_event") as mock_sync_log, + ): + await logging_obj.dispatch_success_handlers(result={"id": "pt-1"}) + + mock_async_log.assert_awaited_once() + mock_sync_log.assert_not_called() + finally: + litellm._async_success_callback = original_async_callbacks + + def test_success_handler_skips_guardrail_logging_hook_when_disabled(logging_obj): """Ensure CustomGuardrail logging_hook is skipped when should_run_guardrail is False.""" import datetime @@ -1351,7 +1557,7 @@ async def test_e2e_generate_cold_storage_object_key_with_custom_logger_s3_path() Test that _generate_cold_storage_object_key uses s3_path from custom logger instance. """ from datetime import datetime, timezone - from unittest.mock import MagicMock, patch + from unittest.mock import AsyncMock, MagicMock, patch from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup @@ -1404,7 +1610,7 @@ async def test_e2e_generate_cold_storage_object_key_with_logger_no_s3_path(): Test that _generate_cold_storage_object_key falls back to empty s3_path when logger has no s3_path. """ from datetime import datetime, timezone - from unittest.mock import MagicMock, patch + from unittest.mock import AsyncMock, MagicMock, patch from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 49d3c51e340..63e2cb7f35c 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -569,8 +569,6 @@ async def test_streaming_with_usage_and_logging(sync_mode: bool): == final_usage_block ) - print(mock_log_success_event.call_args.kwargs.keys()) - def test_streaming_handler_with_stop_chunk( initialized_custom_stream_wrapper: CustomStreamWrapper, @@ -2036,23 +2034,19 @@ async def test_azure_streaming_role_preserved_with_include_usage(sync_mode: bool chunks.append(chunk) # The prompt_filter chunk should be forwarded with choices=[] - assert len(chunks[0].choices) == 0, ( - f"Expected prompt_filter chunk with choices=[], got {len(chunks[0].choices)} choices" - ) + assert ( + len(chunks[0].choices) == 0 + ), f"Expected prompt_filter chunk with choices=[], got {len(chunks[0].choices)} choices" # At least one chunk must have role='assistant' in its delta has_role = any( - len(c.choices) > 0 - and getattr(c.choices[0].delta, "role", None) == "assistant" + len(c.choices) > 0 and getattr(c.choices[0].delta, "role", None) == "assistant" for c in chunks ) assert has_role, ( "No chunk contained role='assistant' in delta (issue #24221). " "Chunk deltas: " - + str([ - c.choices[0].delta if c.choices else "no choices" - for c in chunks - ]) + + str([c.choices[0].delta if c.choices else "no choices" for c in chunks]) ) diff --git a/tests/test_litellm/llms/openai/test_use_chat_completions_api_no_leak.py b/tests/test_litellm/llms/openai/test_use_chat_completions_api_no_leak.py new file mode 100644 index 00000000000..9a266fca81f --- /dev/null +++ b/tests/test_litellm/llms/openai/test_use_chat_completions_api_no_leak.py @@ -0,0 +1,74 @@ +""" +Regression test for issue #28146. + +`use_chat_completions_api` is a LiteLLM-internal control flag (it forces the +/responses -> /chat/completions bridge). When set as a model-level param in the +proxy config, it must never be forwarded to the upstream provider's request +body. OpenAI/Anthropic reject unknown body params with HTTP 400. +""" + +import os +import sys +from unittest.mock import MagicMock + +sys.path.insert(0, os.path.abspath("../../../..")) + +import litellm +from litellm.types.utils import all_litellm_params +from litellm.utils import get_non_default_completion_params + + +def test_use_chat_completions_api_is_a_known_litellm_param(): + assert "use_chat_completions_api" in all_litellm_params + + +def test_use_chat_completions_api_not_forwarded_as_provider_param(): + forwarded = get_non_default_completion_params( + {"use_chat_completions_api": True, "temperature": 0.5} + ) + assert "use_chat_completions_api" not in forwarded + + +def test_completion_does_not_leak_flag_into_provider_request_body(): + mock_response = MagicMock() + mock_response.model_dump.return_value = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1234567890, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + } + + mock_raw_response = MagicMock() + mock_raw_response.headers = {} + mock_raw_response.parse.return_value = mock_response + + mock_client = MagicMock() + mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_raw_response + ) + + litellm.completion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + use_chat_completions_api=True, + api_key="sk-test", + client=mock_client, + ) + + create_kwargs = ( + mock_client.chat.completions.with_raw_response.create.call_args.kwargs + ) + assert "use_chat_completions_api" not in create_kwargs + assert "use_chat_completions_api" not in (create_kwargs.get("extra_body") or {}) diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index e10258c0829..e9ff193e044 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -18,7 +18,7 @@ import asyncio import os import sys from typing import Any -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -38,6 +38,24 @@ from litellm.types.guardrails import GuardrailEventHooks # --------------------------------------------------------------------------- +def _attach_mock_success_dispatch(mock_logging_obj, async_success_fn): + """Match production entrypoint: ``_run_deferred_stream_guardrails`` uses dispatch.""" + + async def dispatch_success_handlers( + result=None, start_time=None, end_time=None, cache_hit=None, **kwargs + ): + await async_success_fn( + result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + **kwargs, + ) + + mock_logging_obj.dispatch_success_handlers = dispatch_success_handlers + mock_logging_obj.async_success_handler = async_success_fn + + class PostCallGuardrail(CustomGuardrail): """A post-call guardrail.""" @@ -454,7 +472,7 @@ class TestDeferredStreamingClosure: async def track_async_success(*args, **kwargs): pass - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) tracking_guardrail = TrackingGuardrail() tracking_logger = TrackingLogger() @@ -511,7 +529,7 @@ class TestDeferredStreamingClosure: nonlocal logged_response logged_response = args[0] if args else None - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) class ModifyingGuardrail(CustomGuardrail): def __init__(self): @@ -573,7 +591,7 @@ class TestDeferredStreamingClosure: nonlocal logging_called logging_called = True - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) guardrail = BlockingGuardrail() @@ -621,7 +639,7 @@ class TestDeferredStreamingClosure: async def track_async_success(*args, **kwargs): pass - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) guardrail = TransientErrorGuardrail() @@ -656,7 +674,7 @@ class TestDeferredStreamingClosure: nonlocal logged_response logged_response = args[0] if args else None - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) class TestGuardrail(CustomGuardrail): def __init__(self): @@ -739,7 +757,7 @@ class TestDeferredStreamingClosure: async def track_async_success(*args, **kwargs): pass - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) guardrail = ApplyGuardrailType() @@ -792,7 +810,7 @@ class TestDeferredStreamingClosure: async def track_async_success(*args, **kwargs): pass - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) guardrail = IteratorHookGuardrail() @@ -847,7 +865,7 @@ class TestDeferredStreamingClosure: async def track_async_success(*args, **kwargs): pass - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) guardrail = InspectingGuardrail() @@ -914,7 +932,7 @@ class TestDeferredStreamingClosure: async def track_async_success(*args, **kwargs): pass - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) guardrail_a = TaggedGuardrail("guardrail-a") guardrail_b = TaggedGuardrail("guardrail-b") @@ -962,7 +980,7 @@ class TestDeferredStreamingClosure: nonlocal logging_called logging_called = True - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) def exploding_merge(data, llm_router): raise RuntimeError("Simulated init failure") @@ -986,6 +1004,67 @@ class TestDeferredStreamingClosure: logging_called is True ), "Logging must fire even when guardrail initialization raises" + @pytest.mark.asyncio + async def test_deferred_logging_forces_async_for_sync_classified_call_type(self): + """ + Regression: proxy deferred streaming logging must reach the async success + handler (which runs the async-only DB/spend logger) even when the call + type is classified as a sync SDK request by _is_sync_litellm_request. + + Without prefer_async_handlers=True, an async proxy stream whose + litellm_params lacks a recognized async marker would enter the sync + branch of dispatch_success_handlers and silently skip spend tracking. + + Uses the real dispatch_success_handlers via the production + _run_deferred_stream_guardrails entrypoint. + """ + import time + + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObj, + ) + + logging_obj = LiteLLMLoggingObj( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", # not pass_through_endpoint + start_time=time.time(), + litellm_call_id="test-id", + function_id="fn", + ) + # litellm_params with no recognized async marker -> classified sync. + logging_obj.model_call_details["litellm_params"] = {} + assert LiteLLMLoggingObj._is_sync_litellm_request({}) is True + + with ( + patch.object( + logging_obj, "async_success_handler", new_callable=AsyncMock + ) as mock_async, + patch.object( + logging_obj, "success_handler", new_callable=MagicMock + ) as mock_sync, + patch.object( + logging_obj, + "_should_run_sync_callbacks_for_async_calls", + return_value=False, + ), + patch("litellm.callbacks", [PostCallGuardrail()]), + ): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "gpt-4o-mini", "metadata": {}}, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=logging_obj, + assembled_response=MagicMock(), + cache_hit=False, + ) + + await asyncio.sleep(0) + await asyncio.sleep(0) + + mock_async.assert_awaited_once() + mock_sync.assert_not_called() + # --------------------------------------------------------------------------- # 7. _fire_deferred_stream_logging @@ -1054,7 +1133,7 @@ class TestFireDeferredStreamLogging: nonlocal logged_response logged_response = args[0] if args else None - mock_logging_obj.async_success_handler = track_async_success + _attach_mock_success_dispatch(mock_logging_obj, track_async_success) class InfoWritingGuardrail(CustomGuardrail): def __init__(self): diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 9f7173383b0..0ef9ad857f9 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -180,3 +180,172 @@ def test_sync_guardrail_from_db_marks_source_db_when_unchanged(): handler.sync_guardrail_from_db(g) assert handler.get_source("collide") == "db" + + +def _db_litellm_params() -> dict: + """ + Shape produced by GuardrailRegistry.get_all_guardrails_from_db: litellm_params + is a raw dict (not a LitellmParams), holding only the keys originally stored, + a non-schema extra key, and plain-string enum values. + """ + return { + "guardrail": "litellm_content_filter", + "mode": "pre_call", + "default_on": True, + "version": 2, + "blocked_words": [{"keyword": "secret", "action": "BLOCK"}], + } + + +def test_unchanged_db_params_do_not_register_as_changed(): + """ + A DB poll returns litellm_params as a raw dict while the in-memory copy is a + LitellmParams whose model_dump() fills every field default and coerces enums. + The two shapes must compare equal when the config is identical; otherwise + every poll cycle re-initializes the guardrail indefinitely. + """ + handler = InMemoryGuardrailHandler() + raw = _db_litellm_params() + gid = "11111111-1111-1111-1111-111111111111" + handler.IN_MEMORY_GUARDRAILS[gid] = Guardrail( + guardrail_id=gid, + guardrail_name="cf", + litellm_params=LitellmParams(**raw), + ) + + new = Guardrail(guardrail_id=gid, guardrail_name="cf", litellm_params=dict(raw)) + assert handler._has_guardrail_params_changed(gid, new) is False + + +def test_changed_db_params_register_as_changed(): + """Normalizing both sides must still surface a genuine config change.""" + handler = InMemoryGuardrailHandler() + raw = _db_litellm_params() + gid = "22222222-2222-2222-2222-222222222222" + handler.IN_MEMORY_GUARDRAILS[gid] = Guardrail( + guardrail_id=gid, + guardrail_name="cf", + litellm_params=LitellmParams(**raw), + ) + + changed = {**raw, "blocked_words": [{"keyword": "different", "action": "BLOCK"}]} + new = Guardrail(guardrail_id=gid, guardrail_name="cf", litellm_params=changed) + assert handler._has_guardrail_params_changed(gid, new) is True + + +def test_unnormalizable_db_params_register_as_changed_without_raising(): + """ + A DB row whose litellm_params fail LitellmParams validation must not crash the + poll loop. The comparison falls back to treating the guardrail as changed so it + re-initializes (and surfaces the bad row in logs) rather than propagating the + validation error up through the polling cycle. + """ + handler = InMemoryGuardrailHandler() + raw = _db_litellm_params() + gid = "55555555-5555-5555-5555-555555555555" + handler.IN_MEMORY_GUARDRAILS[gid] = Guardrail( + guardrail_id=gid, + guardrail_name="cf", + litellm_params=LitellmParams(**raw), + ) + + malformed = {**raw, "default_on": "not-a-bool-xyz"} + new = Guardrail(guardrail_id=gid, guardrail_name="cf", litellm_params=malformed) + assert handler._has_guardrail_params_changed(gid, new) is True + + +def _all_callback_lists(): + import litellm + + return [ + litellm.callbacks, + litellm.success_callback, + litellm.failure_callback, + litellm._async_success_callback, + litellm._async_failure_callback, + ] + + +def test_delete_in_memory_guardrail_removes_callback_from_all_lists(): + """ + Request handling promotes guardrail callbacks from litellm.callbacks into the + success/failure/async lists. delete_in_memory_guardrail must purge the callback + from every list, otherwise a re-initialized guardrail leaves its old instance + stranded in those lists and instances accumulate. + """ + handler = InMemoryGuardrailHandler() + callback = CustomGuardrail( + guardrail_name="cf-delete", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + ) + gid = "33333333-3333-3333-3333-333333333333" + handler.IN_MEMORY_GUARDRAILS[gid] = _make_guardrail(gid, "cf-delete") + handler._sources[gid] = "db" + handler.guardrail_id_to_custom_guardrail[gid] = callback + + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + for cb_list in lists: + cb_list.append(callback) + + handler.delete_in_memory_guardrail(gid) + + for cb_list in lists: + assert callback not in cb_list + finally: + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot + + +def test_repeated_db_sync_does_not_accumulate_runner_instances(): + """ + End-to-end regression for the OOM: across repeated DB polls (with the config + genuinely changing each cycle to force re-initialization), exactly one live + guardrail instance must exist across all callback lists. On the unfixed code + the stale instance lingers in the success/failure lists and the distinct count + climbs above one. + """ + import litellm + + handler = InMemoryGuardrailHandler() + gid = "44444444-4444-4444-4444-444444444444" + name = "cf-accum" + + def db_guardrail(word: str) -> Guardrail: + params = { + **_db_litellm_params(), + "blocked_words": [{"keyword": word, "action": "BLOCK"}], + } + return Guardrail(guardrail_id=gid, guardrail_name=name, litellm_params=params) + + def promote_into_request_lists() -> None: + manager = litellm.logging_callback_manager + for callback in list(litellm.callbacks): + manager.add_litellm_success_callback(callback) + manager.add_litellm_failure_callback(callback) + manager.add_litellm_async_success_callback(callback) + manager.add_litellm_async_failure_callback(callback) + + def distinct_runner_instances() -> int: + seen = set() + for callback in litellm.logging_callback_manager._get_all_callbacks(): + if ( + isinstance(callback, CustomGuardrail) + and getattr(callback, "guardrail_name", None) == name + ): + seen.add(id(callback)) + return len(seen) + + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + for cycle in range(5): + handler.sync_guardrail_from_db(db_guardrail(f"word-{cycle}")) + promote_into_request_lists() + + assert distinct_runner_instances() == 1 + finally: + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index ab533738cb7..1083c3e9252 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -1374,3 +1374,335 @@ class TestPureTextFastPathParity: AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks(all_chunks) is None ) + + +class TestStreamFalseDeduplication: + """ + Regression tests for the duplicate-callback bug where a streaming pass-through + request had stream=False hardcoded on its Logging object. + + Before the fix: + - logging_obj.stream was always False for pass-through requests + - _is_assembled_stream_success() checked `self.stream is not True` and returned + False immediately, so has_dispatched_final_stream_success was never set + - Any second dispatch_success_handlers call went through unchecked + + After the fix: + - pass_through_endpoints.py sets logging_obj.stream = True after detecting stream + - _create_anthropic_response_logging_payload sets complete_streaming_response on + model_call_details so callbacks see the correct assembled response state + - _is_assembled_stream_success returns True, dedup guard fires on first dispatch + """ + + @staticmethod + def _sse(event, data): + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + @staticmethod + def _make_logging_obj(stream: bool = False) -> LiteLLMLoggingObj: + logging_obj = LiteLLMLoggingObj( + model="claude-3-5-sonnet-20241022", + messages=[{"role": "user", "content": "hello"}], + stream=stream, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="1245", + ) + return logging_obj + + @staticmethod + def _build_chunks(): + frames = [ + TestStreamFalseDeduplication._sse( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_abc", + "type": "message", + "role": "assistant", + "model": "claude-3-5-sonnet-20241022", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 0}, + }, + }, + ), + TestStreamFalseDeduplication._sse( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ), + TestStreamFalseDeduplication._sse( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello"}, + }, + ), + TestStreamFalseDeduplication._sse( + "content_block_stop", {"type": "content_block_stop", "index": 0} + ), + TestStreamFalseDeduplication._sse( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 5}, + }, + ), + TestStreamFalseDeduplication._sse("message_stop", {"type": "message_stop"}), + ] + from litellm.proxy.pass_through_endpoints.streaming_handler import ( + PassThroughStreamingHandler, + ) + + return PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(frames) + + def test_complete_streaming_response_set_on_model_call_details(self): + """ + After the fix, _create_anthropic_response_logging_payload must set + complete_streaming_response on logging_obj.model_call_details so that + callbacks like _PROXY_track_cost_callback see the assembled response + instead of None. + + Before the fix: model_call_details had no complete_streaming_response key. + The log showed: "kwargs stream: True + complete streaming response: None" + """ + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + EndpointType, + ) + + # pass_through_request sets the stream flag before the streaming handler + # reconstructs the response; mirror that here. + logging_obj = self._make_logging_obj(stream=True) + logging_obj.model_call_details["stream"] = True + all_chunks = list(self._build_chunks()) + + result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/anthropic/v1/messages", + request_body={"model": "claude-3-5-sonnet-20241022", "stream": True}, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + all_chunks=all_chunks, + end_time=datetime.now(), + ) + + # The assembled response must be stored on model_call_details so callbacks + # can identify this as a completed streaming call, not an in-progress one. + assert ( + logging_obj.model_call_details.get("complete_streaming_response") + is not None + ), "complete_streaming_response must be set on model_call_details after assembly" + + # The returned result must match what was stored + assert result["result"] is logging_obj.model_call_details.get( + "complete_streaming_response" + ) + + def test_dedup_guard_fires_when_stream_true_on_logging_obj(self): + """ + When logging_obj.stream is True (set by pass_through_endpoints.py after + detecting a streaming request), dispatch_success_handlers must set + has_dispatched_final_stream_success=True on the first call so that any + second call is a no-op. + + This is the _is_assembled_stream_success gate: with stream=False it + always returned False and the guard was permanently disabled. + """ + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + EndpointType, + ) + from litellm.types.utils import ModelResponse + + # Simulate what pass_through_endpoints.py now does after stream detection + logging_obj = self._make_logging_obj(stream=False) + logging_obj.stream = True # fix applied + logging_obj.model_call_details["stream"] = True + + # Simulate what _create_anthropic_response_logging_payload now does + mock_response = ModelResponse(model="claude-3-5-sonnet-20241022") + logging_obj.model_call_details["complete_streaming_response"] = mock_response + + assert logging_obj._is_assembled_stream_success(result=mock_response) is True + + # First dispatch sets the flag + assert not logging_obj.model_call_details.get( + "has_dispatched_final_stream_success" + ) + logging_obj.model_call_details["has_dispatched_final_stream_success"] = True + + # Second dispatch would be blocked — simulate the guard check + would_skip = bool( + logging_obj._is_assembled_stream_success(result=mock_response) + and logging_obj.model_call_details.get( + "has_dispatched_final_stream_success" + ) + ) + assert would_skip is True, ( + "Dedup guard must block a second dispatch_success_handlers call for the " + "same assembled streaming response" + ) + + def test_sse_fallback_path_sets_stream_true_for_dedup(self): + """ + When a nominally non-streaming request receives an SSE response + (_is_streaming_response returns True), the fallback branch in + pass_through_endpoints.py must set logging_obj.stream = True so the + dedup guard activates. + + Before the fix the fallback path never set stream=True, so + _is_assembled_stream_success always returned False and duplicate + callback dispatches were never blocked. + """ + from litellm.types.utils import ModelResponse + + # logging_obj starts with stream=False, as created before the request + logging_obj = self._make_logging_obj(stream=False) + assert logging_obj._is_assembled_stream_success(result=MagicMock()) is False + + # Simulate what the SSE fallback branch in pass_through_endpoints.py now does + logging_obj.stream = True + logging_obj.model_call_details["stream"] = True + + mock_response = ModelResponse(model="claude-3-5-sonnet-20241022") + logging_obj.model_call_details["complete_streaming_response"] = mock_response + + # With stream=True the dedup guard must be active + assert logging_obj._is_assembled_stream_success(result=mock_response) is True + + logging_obj.model_call_details["has_dispatched_final_stream_success"] = True + + would_skip = bool( + logging_obj._is_assembled_stream_success(result=mock_response) + and logging_obj.model_call_details.get( + "has_dispatched_final_stream_success" + ) + ) + assert would_skip is True + + def test_stream_false_logging_obj_bypasses_dedup_guard(self): + """ + Demonstrates the pre-fix state: with stream=False on the logging object, + _is_assembled_stream_success always returns False regardless of whether + complete_streaming_response is set. This means the dedup guard can never + fire, so duplicate dispatches go through unchecked. + + This test documents the old broken behavior so the fix is clearly justified. + """ + from litellm.types.utils import ModelResponse + + logging_obj = self._make_logging_obj(stream=False) + mock_response = ModelResponse(model="claude-3-5-sonnet-20241022") + logging_obj.model_call_details["complete_streaming_response"] = mock_response + + # With stream=False, _is_assembled_stream_success returns False even though + # complete_streaming_response is present — the guard is permanently disabled. + assert logging_obj._is_assembled_stream_success(result=mock_response) is False + + +class TestNonStreamingResponseRedaction: + """ + Regression tests ensuring _create_anthropic_response_logging_payload only sets + complete_streaming_response for streaming responses. perform_redaction scrubs + that field exclusively when model_call_details["stream"] is True, so storing it + on a non-streaming response would deliver the unredacted response to logging + callbacks when message logging is disabled. + """ + + @staticmethod + def _make_logging_obj(stream: bool) -> LiteLLMLoggingObj: + logging_obj = LiteLLMLoggingObj( + model="claude-3-5-sonnet-20241022", + messages=[{"role": "user", "content": "hello"}], + stream=stream, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="1245", + ) + # pass_through_request mirrors the stream flag onto model_call_details, + # which is the key perform_redaction inspects. + logging_obj.model_call_details["stream"] = stream + return logging_obj + + def test_non_streaming_does_not_set_complete_streaming_response(self): + from litellm.types.utils import ModelResponse + + logging_obj = self._make_logging_obj(stream=False) + response = ModelResponse(model="claude-3-5-sonnet-20241022") + + AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=response, + model="claude-3-5-sonnet-20241022", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + assert ( + "complete_streaming_response" not in logging_obj.model_call_details + ), "non-streaming responses must not populate complete_streaming_response" + + def test_streaming_sets_complete_streaming_response(self): + from litellm.types.utils import ModelResponse + + logging_obj = self._make_logging_obj(stream=True) + response = ModelResponse(model="claude-3-5-sonnet-20241022") + + AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=response, + model="claude-3-5-sonnet-20241022", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + assert ( + logging_obj.model_call_details.get("complete_streaming_response") + is response + ) + + def test_non_streaming_response_is_redacted_when_message_logging_off(self): + from litellm.litellm_core_utils.redact_messages import ( + redact_message_input_output_from_logging, + ) + from litellm.types.utils import Choices, Message, ModelResponse + + logging_obj = self._make_logging_obj(stream=False) + response = ModelResponse( + model="claude-3-5-sonnet-20241022", + choices=[Choices(message=Message(role="assistant", content="secret"))], + ) + + AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=response, + model="claude-3-5-sonnet-20241022", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + logging_obj.model_call_details["litellm_params"] = { + "metadata": {"headers": {"x-litellm-enable-message-redaction": True}} + } + + redacted = redact_message_input_output_from_logging( + model_call_details=logging_obj.model_call_details, + result=response, + ) + + leaked = logging_obj.model_call_details.get("complete_streaming_response") + assert leaked is None + assert redacted.choices[0].message.content == "redacted-by-litellm" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 97a21136198..89b58973079 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -989,6 +989,131 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): assert metadata["user_api_key_user_id"] == "test-user-id" +@pytest.mark.asyncio +async def test_pass_through_request_streaming_marks_logging_obj_as_stream(): + """ + Regression: a streaming pass-through request must flag its logging object as + streaming (logging_obj.stream and model_call_details["stream"]) before the + response is dispatched, so cost/success callbacks treat it as a stream and the + streaming dedup guard fires instead of double-logging. + """ + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.PassThroughStreamingHandler.chunk_processor" + ) as mock_chunk_processor: + mock_proxy_logging.pre_call_hook = AsyncMock( + return_value={"model": "claude-3", "stream": True} + ) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + + upstream_response = MagicMock() + upstream_response.status_code = 200 + upstream_response.headers = {} + upstream_response.raise_for_status = MagicMock() + + async_client = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + async def _empty_chunks(*args, **kwargs): + return + yield # pragma: no cover + + mock_chunk_processor.return_value = _empty_chunks() + + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = "http://test-proxy.com/v1/messages" + mock_request.body = AsyncMock( + return_value=b'{"model": "claude-3", "stream": true}' + ) + mock_request.headers = Headers({}) + mock_request.query_params = QueryParams({}) + + await pass_through_request( + request=mock_request, + target="http://target-api.com/v1/messages", + custom_headers={}, + user_api_key_dict=MagicMock(), + stream=True, + ) + + async_client.send.assert_awaited_once() + assert async_client.send.call_args.kwargs["stream"] is True + + mock_chunk_processor.assert_called_once() + logging_obj = mock_chunk_processor.call_args.kwargs[ + "litellm_logging_obj" + ] + assert logging_obj.stream is True + assert logging_obj.model_call_details["stream"] is True + + +@pytest.mark.asyncio +async def test_pass_through_request_sse_response_marks_logging_obj_as_stream(): + """ + Regression: a request that is not flagged as streaming up front but whose + upstream response comes back as an SSE stream (content-type text/event-stream) + must still flag its logging object as streaming before dispatch. Otherwise the + cost/success callbacks treat the assembled stream as a non-stream and the dedup + guard never fires, double-logging the request. + """ + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.PassThroughStreamingHandler.chunk_processor" + ) as mock_chunk_processor: + mock_proxy_logging.pre_call_hook = AsyncMock( + return_value={"model": "claude-3"} + ) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + + upstream_response = MagicMock() + upstream_response.status_code = 200 + upstream_response.headers = {"content-type": "text/event-stream"} + upstream_response.raise_for_status = MagicMock() + + async_client = MagicMock() + async_client.request = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + async def _empty_chunks(*args, **kwargs): + return + yield # pragma: no cover + + mock_chunk_processor.return_value = _empty_chunks() + + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = "http://test-proxy.com/v1/messages" + mock_request.body = AsyncMock(return_value=b'{"model": "claude-3"}') + mock_request.headers = Headers({}) + mock_request.query_params = QueryParams({}) + + await pass_through_request( + request=mock_request, + target="http://target-api.com/v1/messages", + custom_headers={}, + user_api_key_dict=MagicMock(), + stream=False, + ) + + async_client.request.assert_awaited_once() + + mock_chunk_processor.assert_called_once() + logging_obj = mock_chunk_processor.call_args.kwargs[ + "litellm_logging_obj" + ] + assert logging_obj.stream is True + assert logging_obj.model_call_details["stream"] is True + + @pytest.mark.asyncio async def test_create_pass_through_endpoint(): """ diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 265a82d4a44..657529c85f1 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -2244,6 +2244,41 @@ class TestHandleLLMApiExceptionDictDetail: assert proxy_exc.message == "Content blocked by guardrail" assert proxy_exc.provider_specific_fields is None + async def test_already_normalized_proxy_exception_is_honored(self): + """A ProxyException raised mid-request (e.g. a guardrail block) is already + the OpenAI wire format. The funnel must re-raise it untouched instead of + re-deriving the status from a (nonexistent) status_code attribute and + defaulting to 500. Regression for LIT-3751.""" + from litellm.proxy._types import ProxyException + + exc = ProxyException( + message='"Leroy Jenkins" detected as name', + type="invalid_request_error", + param=None, + code=400, + openai_code="content_policy_violation", + ) + proxy_exc = await self._invoke(exc) + assert proxy_exc is exc + assert proxy_exc.code == "400" + assert proxy_exc.type == "invalid_request_error" + assert proxy_exc.param is None + assert proxy_exc.openai_code == "content_policy_violation" + assert proxy_exc.message == '"Leroy Jenkins" detected as name' + + # The body the OpenAI-SDK client actually receives. The HTTP status line + # comes from int(exc.code) == 400; the wire ``code`` stays the status + # string. ``openai_code`` ("content_policy_violation") is intentionally + # NOT serialized here - to_dict() emits only ``code`` - so this asserts + # the real contract rather than the write-only attribute. + assert int(proxy_exc.code) == 400 + assert proxy_exc.to_dict() == { + "message": '"Leroy Jenkins" detected as name', + "type": "invalid_request_error", + "param": None, + "code": "400", + } + class TestAsyncStreamingDataGeneratorFastPath: """Fast/slow path branching in async_streaming_data_generator.""" diff --git a/tests/test_litellm/proxy/test_model_level_guardrails.py b/tests/test_litellm/proxy/test_model_level_guardrails.py index 3d74edd772b..9a79fa7f496 100644 --- a/tests/test_litellm/proxy/test_model_level_guardrails.py +++ b/tests/test_litellm/proxy/test_model_level_guardrails.py @@ -19,7 +19,6 @@ from litellm.proxy.utils import ( _merge_guardrails_with_existing, ) - # --------------------------------------------------------------------------- # Unit tests for _check_and_merge_model_level_guardrails # --------------------------------------------------------------------------- @@ -159,6 +158,157 @@ class TestCheckAndMergeModelLevelGuardrails: assert "existing" in result["metadata"]["guardrails"] +# --------------------------------------------------------------------------- +# Regression test: pre_call hook must run exactly once with model-level guardrails +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_pre_call_hook_runs_once_with_model_level_guardrails(): + """ + A guardrail attached at the model level (litellm_params.guardrails) is + spread into the top-level request kwargs by the router. The proxy pre-call + loop (async_pre_call_hook) and the deployment-level hook + (async_pre_call_deployment_hook) must together invoke async_pre_call_hook + exactly once, not twice. + """ + from litellm.caching.caching import DualCache + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy._types import CallTypes, UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging + from litellm.types.guardrails import GuardrailEventHooks + + class CountingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="counting-guardrail", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + self.pre_call_count = 0 + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.pre_call_count += 1 + return data + + guardrail = CountingGuardrail() + + with patch("litellm.callbacks", [guardrail]): + ProxyLogging._callback_capabilities_cache.clear() + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello"}], + "metadata": {}, + } + + # Path A: proxy pre-call loop runs the guardrail and records that it ran + data = await proxy_logging.pre_call_hook( + user_api_key_dict=user_api_key_dict, + data=data, + call_type="acompletion", + ) + + # Path B: the router spreads the deployment's model-level guardrails into + # the top-level kwargs, then litellm.acompletion fires the deployment hook + data["guardrails"] = ["counting-guardrail"] + await guardrail.async_pre_call_deployment_hook(data, CallTypes.acompletion) + + assert guardrail.pre_call_count == 1 + + +@pytest.mark.asyncio +async def test_pre_call_hook_runs_once_when_hook_returns_fresh_dict(): + """ + async_pre_call_hook may return a brand-new request dict instead of mutating + or spreading the one it received. The exactly-once marker must live on the + data that flows downstream, so the deployment hook still skips the guardrail + even when the proxy loop swapped in a fresh dict that never carried it. + """ + from litellm.caching.caching import DualCache + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy._types import CallTypes, UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging + from litellm.types.guardrails import GuardrailEventHooks + + class FreshDictGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="counting-guardrail", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + self.pre_call_count = 0 + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.pre_call_count += 1 + return {"model": data["model"], "messages": data["messages"]} + + guardrail = FreshDictGuardrail() + + with patch("litellm.callbacks", [guardrail]): + ProxyLogging._callback_capabilities_cache.clear() + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello"}], + "metadata": {}, + } + + data = await proxy_logging.pre_call_hook( + user_api_key_dict=user_api_key_dict, + data=data, + call_type="acompletion", + ) + + data["guardrails"] = ["counting-guardrail"] + await guardrail.async_pre_call_deployment_hook(data, CallTypes.acompletion) + + assert guardrail.pre_call_count == 1 + + +@pytest.mark.asyncio +async def test_deployment_hook_runs_pre_call_without_proxy_loop(): + """ + Direct-SDK usage (litellm.acompletion(..., guardrails=[...]) without the + proxy) never runs the proxy pre-call loop, so the deployment hook is the + only place the guardrail executes and it must still run exactly once. + """ + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy._types import CallTypes + from litellm.types.guardrails import GuardrailEventHooks + + class CountingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="counting-guardrail", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + ) + self.pre_call_count = 0 + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.pre_call_count += 1 + return data + + guardrail = CountingGuardrail() + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello"}], + "guardrails": ["counting-guardrail"], + "metadata": {}, + } + + await guardrail.async_pre_call_deployment_hook(data, CallTypes.acompletion) + + assert guardrail.pre_call_count == 1 + + # --------------------------------------------------------------------------- # Integration test: post_call_success_hook with model-level guardrails # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 7a2b20bd8fb..539a32db57d 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -321,3 +321,117 @@ class TestPostCallFailureHookLiftsFirstApiCallStartTime: await self._run(request_data) assert "first_api_call_start_time" not in request_data assert "litellm_logging_obj" not in request_data + + +class TestPostCallFailureHookLLMExceptionAlerting: + """The llm_exceptions alert is for infra / LLM-API failures, not user + errors (https://github.com/BerriAI/litellm/issues/3395). Already-normalized + client errors must be excluded so a guardrail content-policy block never + pages on-call. ProxyException is such an error; before LIT-3751 only + HTTPException was excluded, so AIM blocks paged as if the LLM API failed.""" + + async def _alerted(self, exc) -> bool: + import asyncio + from unittest.mock import AsyncMock, patch + + from litellm.proxy._types import AlertType, UserAPIKeyAuth + + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.alert_types = [AlertType.llm_exceptions] + alerting_handler = AsyncMock() + with ( + patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()), + patch.object(proxy_logging_obj, "alerting_handler", new=alerting_handler), + ): + await proxy_logging_obj.post_call_failure_hook( + request_data={}, + original_exception=exc, + user_api_key_dict=UserAPIKeyAuth(), + ) + await asyncio.sleep(0) # let the fire-and-forget alert task run + return alerting_handler.called + + @pytest.mark.asyncio + async def test_proxy_exception_does_not_alert(self): + from litellm.proxy._types import ProxyException + + exc = ProxyException( + message="content blocked", + type="invalid_request_error", + param=None, + code=400, + openai_code="content_policy_violation", + ) + assert await self._alerted(exc) is False + + @pytest.mark.asyncio + async def test_http_exception_does_not_alert(self): + assert ( + await self._alerted(HTTPException(status_code=400, detail="blocked")) + is False + ) + + @pytest.mark.asyncio + async def test_genuine_llm_api_error_still_alerts(self): + assert await self._alerted(Exception("upstream 503")) is True + + +class TestPostCallFailureHookProxyExceptionLogging: + """A guardrail block raises a ProxyException; on an LLM route it must still + drive proxy-only failure logging (_handle_logging_proxy_only_error) so the + blocked request is recorded, exactly as the old HTTPException did. Before + LIT-3751 the classifier only matched HTTPException, so switching AIM to + ProxyException silently dropped the rejected prompt from failure logs.""" + + async def _logged(self, exc, *, request_route) -> bool: + from unittest.mock import AsyncMock, patch + + from litellm.proxy._types import UserAPIKeyAuth + + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.alert_types = [] + handle_mock = AsyncMock() + with ( + patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()), + patch.object( + proxy_logging_obj, + "_handle_logging_proxy_only_error", + new=handle_mock, + ), + ): + await proxy_logging_obj.post_call_failure_hook( + request_data={}, + original_exception=exc, + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-test", request_route=request_route + ), + ) + return handle_mock.await_count > 0 + + def _block(self): + from litellm.proxy._types import ProxyException + + return ProxyException( + message="content blocked", + type="invalid_request_error", + param=None, + code=400, + openai_code="content_policy_violation", + ) + + @pytest.mark.asyncio + async def test_proxy_exception_on_llm_route_is_logged(self): + assert ( + await self._logged(self._block(), request_route="/v1/chat/completions") + is True + ) + + @pytest.mark.asyncio + async def test_generic_exception_on_llm_route_is_not_logged(self): + # A raw provider/unknown exception is logged by the LLM call path, not here. + assert ( + await self._logged( + Exception("upstream 503"), request_route="/v1/chat/completions" + ) + is False + ) diff --git a/uv.lock b/uv.lock index 20a9e23a019..1ea60c7409f 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-06-10T23:28:01.952719Z" +exclude-newer = "2026-06-17T19:22:29.366621Z" exclude-newer-span = "P3D" [manifest] @@ -3273,7 +3273,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.87.3" +version = "1.87.4" source = { editable = "." } dependencies = [ { name = "aiohttp" },