diff --git a/docs/my-website/docs/proxy/guardrails/custom_guardrail.md b/docs/my-website/docs/proxy/guardrails/custom_guardrail.md index c9115cf826..638cae9c83 100644 --- a/docs/my-website/docs/proxy/guardrails/custom_guardrail.md +++ b/docs/my-website/docs/proxy/guardrails/custom_guardrail.md @@ -117,6 +117,14 @@ guardrails: ::: +:::note Streaming and post_call guardrails + +For **streaming responses**, `post_call` guardrails run on the fully assembled response **after** all chunks have been delivered to the client. This means `post_call` guardrails on streaming are **audit-only** — they can inspect and log the complete response, but cannot block content delivery. Guardrail results are recorded in `guardrail_information` within the logging payload for compliance and auditing. + +To filter or block streaming content in real-time, use `async_post_call_streaming_iterator_hook` instead, which processes chunks as they arrive. + +::: +
Advanced: Multiple modes with individual event hooks @@ -655,8 +663,8 @@ class myCustomGuardrail(CustomGuardrail): | `apply_guardrail` | Simple method to check and optionally modify text | ✅ | INPUT or OUTPUT | ✅ | ✅ | ✅ | | `async_pre_call_hook` | A hook that runs before the LLM API call | ✅ | INPUT | ✅ | ❌ | ✅ | | `async_moderation_hook` | A hook that runs during the LLM API call| ✅ | INPUT | ❌ | ❌ | ✅ | -| `async_post_call_success_hook` | A hook that runs after a successful LLM API call| ✅ | INPUT, OUTPUT | ❌ | ✅ | ✅ | -| `async_post_call_streaming_iterator_hook` | A hook that processes streaming responses | ✅ | OUTPUT | ❌ | ✅ | ✅ | +| `async_post_call_success_hook` | A hook that runs after a successful LLM API call. For streaming, runs on the assembled response after delivery (audit-only, cannot block). | ✅ | INPUT, OUTPUT | ❌ | ✅ | ✅ (non-streaming only) | +| `async_post_call_streaming_iterator_hook` | A hook that processes streaming responses in real-time (can filter/block chunks) | ✅ | OUTPUT | ❌ | ✅ | ✅ | ## Frequently Asked Questions diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 7db72da276..4c0471fb8f 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -48,6 +48,20 @@ const sidebars = { slug: "/guardrail_providers" }, items: [ + { + type: "category", + label: "Contributing to Guardrails", + items: [ + "adding_provider/generic_guardrail_api", + "adding_provider/simple_guardrail_tutorial", + "adding_provider/adding_guardrail_support", + ] + }, + { + type: "doc", + id: "proxy/guardrails/team_based_guardrails", + label: "Team Bring-Your-Own Guardrails", + }, ...[ "proxy/guardrails/qualifire", "proxy/guardrails/aim_security", diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 826396a70d..4e63dd7076 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1455,6 +1455,20 @@ class Logging(LiteLLMLoggingBaseClass): ): # use model_id if not already set router_model_id = hidden_params["model_id"] + # Fallback: extract router_model_id from litellm_params when not available + # from the result object. ResponsesAPIResponse objects (used by /v1/responses + # streaming) don't carry _hidden_params["model_id"] like ModelResponse does. + if router_model_id is None and hasattr(self, "litellm_params"): + for metadata_key in ("litellm_metadata", "metadata"): + _metadata: dict = ( + self.litellm_params.get(metadata_key, {}) or {} + ) + _model_info: dict = _metadata.get("model_info", {}) or {} + _model_id = _model_info.get("id") + if _model_id is not None: + router_model_id = _model_id + break + ## RESPONSE COST ## custom_pricing = use_custom_pricing_for_model( litellm_params=( diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 6e991e6911..ccba18088e 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2136,22 +2136,36 @@ class CustomStreamWrapper: self.sent_stream_usage = True return response - asyncio.create_task( - self.logging_obj.async_success_handler( + _deferred_cb = getattr( + self.logging_obj, + "_on_deferred_stream_complete", + None, + ) + if _deferred_cb is not None: + # Proxy has post-call guardrails — let the closure + # run guardrails on the assembled response, then + # fire logging with guardrail_information populated. + self.logging_obj._on_deferred_stream_complete = None # type: ignore[attr-defined] + asyncio.create_task( + _deferred_cb(complete_streaming_response, cache_hit) + ) + else: + asyncio.create_task( + self.logging_obj.async_success_handler( + complete_streaming_response, + cache_hit=cache_hit, + start_time=None, + end_time=None, + ) + ) + + executor.submit( + self.logging_obj.success_handler, complete_streaming_response, cache_hit=cache_hit, start_time=None, end_time=None, ) - ) - - 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: diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 508c1c9465..604e7d5f41 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -1,41 +1,32 @@ model_list: - - model_name: gpt-3.5-turbo - litellm_params: - model: openai/gpt-3.5-turbo - api_key: os.environ/OPENAI_API_KEY - - model_name: gpt-4o - litellm_params: - model: openai/gpt-4o - api_key: os.environ/OPENAI_API_KEY - - model_name: claude-sonnet-4-5-20250929 - litellm_params: - model: anthropic/claude-sonnet-4-5-20250929 - - model_name: gpt-4.1-mini + + # OpenAI model for /v1/chat/completions test — 200x custom pricing + - model_name: "gpt-4.1-mini" litellm_params: model: openai/gpt-4.1-mini - - model_name: gpt-5-mini + api_key: os.environ/OPENAI_API_KEY + model_info: + id: gpt-4.1-mini-custom-pricing + input_cost_per_token: 0.00004 # 100x standard ($0.40/1M = $0.0000004) + output_cost_per_token: 0.00016 # 100x standard ($1.60/1M = $0.0000016) + + # OpenAI model for /v1/responses test — 100x custom pricing + - model_name: "gpt-5" litellm_params: - model: openai/gpt-5-mini - - model_name: custom_litellm_model + model: openai/gpt-5 + api_key: os.environ/OPENAI_API_KEY + model_info: + id: gpt-5-custom-pricing + mode: "chat" + input_cost_per_token: 125 # 100x standard ($1.25/1M = $0.00000125) + output_cost_per_token: 10 # 100x standard ($10.00/1M = $0.00001) + + # Anthropic model for /v1/messages test — 100x custom pricing + - model_name: "claude-sonnet-4-20250514" litellm_params: - model: litellm_agent/claude-sonnet-4-5-20250929 - litellm_system_prompt: "Be a helpful assistant." - - -guardrails: - - guardrail_name: "tool_policy" - litellm_params: - guardrail: tool_policy - mode: [pre_call, post_call] - default_on: true - -mcp_servers: - my_http_server: - url: "http://0.0.0.0:8001/mcp" - transport: "http" - description: "My custom MCP server" - available_on_public_internet: true - -general_settings: - store_model_in_db: true - store_prompts_in_spend_logs: true + model: anthropic/claude-sonnet-4-20250514 + api_key: os.environ/ANTHROPIC_API_KEY + model_info: + id: claude-sonnet-4-custom-pricing + input_cost_per_token: 0.0003 # 100x standard ($0.000003) + output_cost_per_token: 0.0015 # 100x standard ($0.000015) \ No newline at end of file diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 72765aab7d..96fd219571 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -45,7 +45,9 @@ from litellm.proxy.common_utils.callback_utils import ( from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import ProxyLogging +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.router import Router +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ServerToolUse # Type alias for streaming chunk serializer (chunk after hooks + cost injection -> wire format) @@ -801,7 +803,7 @@ class ProxyBaseLLMRequestProcessing: json.dumps(self.data, indent=4, default=str), ) - async def base_process_llm_request( + async def base_process_llm_request( # noqa: PLR0915 self, request: Request, fastapi_response: Response, @@ -926,6 +928,26 @@ class ProxyBaseLLMRequestProcessing: llm_router=llm_router, ) + # Defer async logging when post-call guardrails are configured so the + # StandardLoggingPayload is built after guardrails write to metadata. + # Cache the result to avoid scanning litellm.callbacks twice. + _post_call_guardrails_active = self._has_post_call_guardrails() + + # Non-streaming: defer the create_task in wrapper_async so the + # SLP is built after guardrails write to metadata. Streaming + # uses a separate closure mechanism (see below). + # + # Edge case: if _is_streaming_request is False but the response + # turns out to be a CustomStreamWrapper (rare provider behavior), + # wrapper_async exits early before the _defer_async_logging block + # so _enqueue_deferred_logging is never stored — the finally + # block is a no-op. The CSW path handles this correctly via + # _on_deferred_stream_complete, which fires its own logging. + if _post_call_guardrails_active and not self._is_streaming_request( + data=self.data, is_streaming_request=is_streaming_request + ): + logging_obj._defer_async_logging = True # type: ignore + tasks = [] # Start the moderation check (during_call_hook) as early as possible # This gives it a head start to mask/validate input while the proxy handles routing @@ -962,124 +984,229 @@ class ProxyBaseLLMRequestProcessing: response = responses[1] - hidden_params = getattr(response, "_hidden_params", {}) or {} - model_id = self._get_model_id_from_response(hidden_params, self.data) + _exception_raised = False + try: + hidden_params = getattr(response, "_hidden_params", {}) or {} + model_id = self._get_model_id_from_response(hidden_params, self.data) - cache_key, api_base, response_cost = ( - hidden_params.get("cache_key", None) or "", - hidden_params.get("api_base", None) or "", - hidden_params.get("response_cost", None) or "", - ) - fastest_response_batch_completion, additional_headers = ( - hidden_params.get("fastest_response_batch_completion", None), - hidden_params.get("additional_headers", {}) or {}, - ) - - # Post Call Processing - if llm_router is not None: - self.data["deployment"] = llm_router.get_deployment(model_id=model_id) - asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=self.data.get("litellm_call_id", ""), status="success" + cache_key, api_base, response_cost = ( + hidden_params.get("cache_key", None) or "", + hidden_params.get("api_base", None) or "", + hidden_params.get("response_cost", None) or "", ) - ) - if self._is_streaming_request( - data=self.data, is_streaming_request=is_streaming_request - ) or self._is_streaming_response( - response - ): # use generate_responses to stream responses - custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( - user_api_key_dict=user_api_key_dict, - call_id=logging_obj.litellm_call_id, - model_id=model_id, - cache_key=cache_key, - api_base=api_base, - version=version, - response_cost=response_cost, - model_region=getattr(user_api_key_dict, "allowed_model_region", ""), - fastest_response_batch_completion=fastest_response_batch_completion, - request_data=self.data, - hidden_params=hidden_params, - litellm_logging_obj=logging_obj, - **additional_headers, + fastest_response_batch_completion, additional_headers = ( + hidden_params.get("fastest_response_batch_completion", None), + hidden_params.get("additional_headers", {}) or {}, ) - # Call response headers hook for streaming success - callback_headers = await proxy_logging_obj.post_call_response_headers_hook( - data=self.data, - user_api_key_dict=user_api_key_dict, - response=response, - request_headers=dict(request.headers), + # Post Call Processing + if llm_router is not None: + self.data["deployment"] = llm_router.get_deployment(model_id=model_id) + asyncio.create_task( + proxy_logging_obj.update_request_status( + litellm_call_id=self.data.get("litellm_call_id", ""), status="success" + ) ) - if callback_headers: - custom_headers.update(callback_headers) + if self._is_streaming_request( + data=self.data, is_streaming_request=is_streaming_request + ) or self._is_streaming_response( + response + ): # use generate_responses to stream responses + custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + call_id=logging_obj.litellm_call_id, + model_id=model_id, + cache_key=cache_key, + api_base=api_base, + version=version, + response_cost=response_cost, + model_region=getattr(user_api_key_dict, "allowed_model_region", ""), + fastest_response_batch_completion=fastest_response_batch_completion, + request_data=self.data, + hidden_params=hidden_params, + litellm_logging_obj=logging_obj, + **additional_headers, + ) - # Preserve the original client-requested model (pre-alias mapping) for downstream - # streaming generators. Pre-call processing can rewrite `self.data["model"]` for - # aliasing/routing, but the OpenAI-compatible response `model` field should reflect - # what the client sent. - if requested_model_from_client: - self.data[ - "_litellm_client_requested_model" - ] = requested_model_from_client - if route_type == "allm_passthrough_route": - # Check if response is an async generator - if self._is_streaming_response(response): - if asyncio.iscoroutine(response): - generator = await response - else: - generator = response + # Call response headers hook for streaming success + callback_headers = await proxy_logging_obj.post_call_response_headers_hook( + data=self.data, + user_api_key_dict=user_api_key_dict, + response=response, + request_headers=dict(request.headers), + ) + if callback_headers: + custom_headers.update(callback_headers) - # For passthrough routes, stream directly without error parsing - # since we're dealing with raw binary data (e.g., AWS event streams) - return StreamingResponse( - content=generator, - status_code=status.HTTP_200_OK, - headers=custom_headers, - ) - else: - # Traditional HTTP response with aiter_bytes - return StreamingResponse( - content=response.aiter_bytes(), - status_code=response.status_code, - headers=custom_headers, - ) - elif route_type == "anthropic_messages": - # Check if response is actually a streaming response (async generator) - # Non-streaming responses (dict) should be returned directly - # This handles cases like websearch_interception agentic loop - # which returns a non-streaming dict even for streaming requests - if self._is_streaming_response(response): - selected_data_generator = ( - ProxyBaseLLMRequestProcessing.async_sse_data_generator( - response=response, - user_api_key_dict=user_api_key_dict, - request_data=self.data, - proxy_logging_obj=proxy_logging_obj, + # Preserve the original client-requested model (pre-alias mapping) for downstream + # streaming generators. Pre-call processing can rewrite `self.data["model"]` for + # aliasing/routing, but the OpenAI-compatible response `model` field should reflect + # what the client sent. + if requested_model_from_client: + self.data[ + "_litellm_client_requested_model" + ] = requested_model_from_client + + # Streaming: attach a closure that CSW.__anext__ will call + # at stream end instead of firing logging directly. The + # closure runs ONLY guardrail hooks (not all callbacks) on + # the assembled response so guardrail_information is + # populated, then fires both logging handlers. + # Only for CustomStreamWrapper — raw async generators from + # passthrough routes bypass CSW and would orphan the closure. + from litellm.litellm_core_utils.streaming_handler import ( + CustomStreamWrapper, + ) + + if _post_call_guardrails_active and isinstance( + response, CustomStreamWrapper + ): + # Intentionally a live reference (not a copy) — mirrors + # ProxyLogging.post_call_success_hook which also mutates + # data["guardrail_to_apply"] during iteration. + _captured_data = self.data + _captured_user_api_key_dict = user_api_key_dict + _captured_logging_obj = logging_obj + + async def _on_deferred_stream_complete( + assembled_response, cache_hit + ): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data=_captured_data, + captured_user_api_key_dict=_captured_user_api_key_dict, + captured_logging_obj=_captured_logging_obj, + assembled_response=assembled_response, + cache_hit=cache_hit, ) + + logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete # type: ignore[attr-defined] + + if route_type == "allm_passthrough_route": + # Check if response is an async generator + if self._is_streaming_response(response): + if asyncio.iscoroutine(response): + generator = await response + else: + generator = response + + # For passthrough routes, stream directly without error parsing + # since we're dealing with raw binary data (e.g., AWS event streams) + return StreamingResponse( + content=generator, + status_code=status.HTTP_200_OK, + headers=custom_headers, + ) + else: + # Traditional HTTP response with aiter_bytes + return StreamingResponse( + content=response.aiter_bytes(), + status_code=response.status_code, + headers=custom_headers, + ) + elif route_type == "anthropic_messages": + # Check if response is actually a streaming response (async generator) + # Non-streaming responses (dict) should be returned directly + # This handles cases like websearch_interception agentic loop + # which returns a non-streaming dict even for streaming requests + if self._is_streaming_response(response): + selected_data_generator = ( + ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=response, + user_api_key_dict=user_api_key_dict, + request_data=self.data, + proxy_logging_obj=proxy_logging_obj, + ) + ) + return await create_response( + generator=selected_data_generator, + media_type="text/event-stream", + headers=custom_headers, + ) + # Non-streaming response - fall through to normal response handling + elif select_data_generator: + selected_data_generator = select_data_generator( + response=response, + user_api_key_dict=user_api_key_dict, + request_data=self.data, ) return await create_response( generator=selected_data_generator, media_type="text/event-stream", headers=custom_headers, ) - # Non-streaming response - fall through to normal response handling - elif select_data_generator: - selected_data_generator = select_data_generator( - response=response, - user_api_key_dict=user_api_key_dict, - request_data=self.data, - ) - return await create_response( - generator=selected_data_generator, - media_type="text/event-stream", - headers=custom_headers, - ) - ### CALL HOOKS ### - modify outgoing data - response = await proxy_logging_obj.post_call_success_hook( - data=self.data, user_api_key_dict=user_api_key_dict, response=response - ) + ### CALL HOOKS ### - modify outgoing data + # If we reach here with a streaming closure still set, it means + # no early-return route consumed the CSW (hypothetical fallthrough). + # Clear the closure so guardrails run inline as before — this + # preserves blocking behavior and avoids double invocation. + if getattr(logging_obj, "_on_deferred_stream_complete", None): + logging_obj._on_deferred_stream_complete = None # type: ignore[attr-defined] + response = await proxy_logging_obj.post_call_success_hook( + data=self.data, user_api_key_dict=user_api_key_dict, response=response + ) + except Exception: + _exception_raised = True + raise + finally: + # Enqueue deferred logging after post-call guardrails have written + # guardrail_information to metadata. The finally block ensures + # logging fires even if a guardrail raises. + # For streaming early-returns: no closure is stored (wrapper_async + # returns before the deferred block), so _enqueue_fn is None — no-op. + _enqueue_fn = getattr(logging_obj, "_enqueue_deferred_logging", None) + if _enqueue_fn is not None: + logging_obj._enqueue_deferred_logging = None # type: ignore[attr-defined] + try: + _enqueue_fn() + except Exception as e: + verbose_proxy_logger.exception( + "Error firing deferred logging: %s", e + ) + + # Streaming cleanup: if an exception occurred AND the deferred + # streaming closure is still set, no streaming route will + # consume the CSW — the closure is orphaned. Clear it and + # fire logging directly to avoid silent loss. + # + # On normal streaming returns the closure must stay: CSW calls + # it at stream end. _exception_raised is function-scoped and + # immune to outer exception context, avoiding false positives. + if _exception_raised: + _deferred_fn = getattr( + logging_obj, "_on_deferred_stream_complete", None + ) + if _deferred_fn is not None: + logging_obj._on_deferred_stream_complete = None # type: ignore[attr-defined] + try: + asyncio.create_task( + logging_obj.async_success_handler( + response, + cache_hit=None, + start_time=None, + end_time=None, + ) + ) + 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. @@ -1217,6 +1344,131 @@ class ProxyBaseLLMRequestProcessing: return True return False + @staticmethod + def _has_post_call_guardrails() -> bool: + """ + Check if any registered callback is a post-call guardrail. + + Uses the global litellm.callbacks list rather than per-request + should_run_guardrail() — intentionally conservative so that the + check is simple and stateless. The deferral path produces + identical logging output, just fires it slightly later, so + false-positives are harmless. + """ + for cb in litellm.callbacks: + if isinstance(cb, CustomGuardrail) and cb._event_hook_is_event_type( + GuardrailEventHooks.post_call + ): + return True + return False + + @staticmethod + async def _run_deferred_stream_guardrails( + captured_data: dict, + captured_user_api_key_dict: "UserAPIKeyAuth", + captured_logging_obj: Any, + assembled_response: Any, + cache_hit: Any, + ) -> None: + """ + Run only post-call guardrail hooks on an assembled streaming response, + then fire both async and sync logging handlers. + + Called by CSW.__anext__ at stream end via a closure stored on + logging_obj._on_deferred_stream_complete. + + This is audit-only — content has already been delivered to the client. + Blocking guardrails that raise HTTPException cannot prevent content + delivery for streaming. Per-chunk filtering should use + async_post_call_streaming_hook instead. + + 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 + from litellm.proxy.utils import ( + _check_and_merge_model_level_guardrails, + unified_guardrail as _unified_guardrail, + ) + + guardrail_data = _check_and_merge_model_level_guardrails( + data=captured_data, llm_router=_global_llm_router + ) + for cb in litellm.callbacks: + if not isinstance(cb, CustomGuardrail): + continue + if not cb.should_run_guardrail( + data=guardrail_data, + event_type=GuardrailEventHooks.post_call, + ): + continue + try: + guardrail_result = None + if "apply_guardrail" in type(cb).__dict__: + guardrail_data["guardrail_to_apply"] = cb + guardrail_result = ( + await _unified_guardrail.async_post_call_success_hook( + user_api_key_dict=captured_user_api_key_dict, + data=guardrail_data, + response=_response, + ) + ) + else: + guardrail_result = await cb.async_post_call_success_hook( + user_api_key_dict=captured_user_api_key_dict, + data=guardrail_data, + response=_response, + ) + if guardrail_result is not None: + _response = guardrail_result + except Exception as e: + verbose_proxy_logger.exception( + "Error running post-call guardrail %s on streaming response: %s", + getattr(cb, "guardrail_name", type(cb).__name__), + e, + ) + if isinstance(e, HTTPException) and hasattr( + captured_logging_obj, "model_call_details" + ): + captured_logging_obj.model_call_details.setdefault( + "metadata", {} + )["guardrail_blocked"] = True + except Exception as e: + verbose_proxy_logger.exception( + "Error in deferred streaming guardrail initialization: %s", e, + ) + finally: + try: + asyncio.create_task( + captured_logging_obj.async_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 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", e, + ) + async def _handle_llm_api_exception( self, e: Exception, diff --git a/litellm/utils.py b/litellm/utils.py index 81d749ab82..0fda994aea 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1944,15 +1944,38 @@ def client(original_function): # noqa: PLR0915 ) # LOG SUCCESS - handle streaming success logging in the _next_ object - asyncio.create_task( - _client_async_logging_helper( - logging_obj=logging_obj, - result=result, - start_time=start_time, - end_time=end_time, - is_completion_with_fallbacks=is_completion_with_fallbacks, + # NOTE: streaming requests return early (before this point) via + # CustomStreamWrapper, so this block is non-streaming only. + if getattr(logging_obj, "_defer_async_logging", False): + # Proxy has post-call guardrails that must complete before the + # SLP is built. Store a closure the proxy will call after + # post_call_success_hook so guardrail_information is in metadata. + # Only create_task is deferred; sync callbacks fire immediately + # (below, outside the if/else) for billing/rate-limiting. + def _enqueue_deferred_logging() -> None: + asyncio.create_task( + _client_async_logging_helper( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + ) + ) + + logging_obj._enqueue_deferred_logging = _enqueue_deferred_logging # type: ignore + else: + asyncio.create_task( + _client_async_logging_helper( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + ) ) - ) + + # Sync callbacks always fire immediately regardless of deferral logging_obj.handle_sync_success_callbacks_for_async_calls( result=result, start_time=start_time, diff --git a/poetry.lock b/poetry.lock index e3b083d778..dc25864442 100644 --- a/poetry.lock +++ b/poetry.lock @@ -8018,4 +8018,4 @@ utils = ["numpydoc"] [metadata] lock-version = "2.1" python-versions = ">=3.9,<4.0" -content-hash = "eda34dfd8b35474beffee18893d6782c7b3d0d3d2c610f66237eb97176f43527" +content-hash = "2cf958f1a04fd5f1ab0e5cfc33bdbf441b518ed6c82d0f2546bf64cd3d2f89be" 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 6e9b72e96c..fe4851283f 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -180,6 +180,82 @@ def test_use_custom_pricing_not_detected_litellm_metadata_no_pricing(): assert use_custom_pricing_for_model(litellm_params) is False +def test_response_cost_calculator_uses_router_model_id_from_litellm_metadata(): + """_response_cost_calculator should extract router_model_id from + litellm_params.litellm_metadata.model_info.id when the result object + does not carry _hidden_params (e.g. ResponsesAPIResponse from /v1/responses + streaming). Regression test for custom pricing on streaming responses.""" + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.llms.openai import ResponsesAPIResponse + + custom_model_id = "gpt-5-custom-pricing" + custom_input_cost = 125.0 + custom_output_cost = 10.0 + + litellm.register_model( + model_cost={ + custom_model_id: { + "input_cost_per_token": custom_input_cost, + "output_cost_per_token": custom_output_cost, + "max_tokens": 128000, + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "litellm_provider": "openai", + } + } + ) + + try: + logging_obj = LiteLLMLoggingObj( + model="gpt-5", + messages=[{"role": "user", "content": "Hi"}], + stream=True, + call_type="aresponses", + start_time=time.time(), + litellm_call_id="test-123", + function_id="test-fn", + ) + + logging_obj.update_environment_variables( + model="gpt-5", + user="", + optional_params={}, + litellm_params={ + "api_base": "", + "litellm_metadata": { + "model_info": { + "id": custom_model_id, + "input_cost_per_token": custom_input_cost, + "output_cost_per_token": custom_output_cost, + }, + }, + }, + ) + + response_obj = ResponsesAPIResponse( + id="resp_abc", + created_at=1234567890, + model="gpt-5", + output=[], + usage={ + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + }, + ) + + cost = logging_obj._response_cost_calculator(result=response_obj) + + assert cost is not None, "Cost should not be None" + expected_cost = (10 * custom_input_cost) + (5 * custom_output_cost) + assert cost == pytest.approx( + expected_cost + ), f"Expected {expected_cost}, got {cost}" + finally: + litellm.model_cost.pop(custom_model_id, None) + + class TestUpdateFromKwargs: """Tests for the update_from_kwargs convenience wrapper.""" diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py new file mode 100644 index 0000000000..c4d2dce587 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -0,0 +1,944 @@ +""" +Tests for deferred logging with post-call guardrails. + +When post-call guardrails are configured, the async logging task is deferred +until after guardrails complete. This ensures the StandardLoggingPayload +is built with guardrail_information populated. + +Non-streaming: create_task in wrapper_async is replaced by a closure that + the proxy fires in a try/finally after post_call_success_hook. + +Streaming: a closure on logging_obj is called by CSW.__anext__ at stream end. + The closure runs ONLY guardrail hooks (not all callbacks), then fires + both logging handlers. +""" + +import asyncio +import os +import sys +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +import litellm +from litellm.caching.caching import DualCache +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class PostCallGuardrail(CustomGuardrail): + """A post-call guardrail.""" + + def __init__(self): + super().__init__( + guardrail_name="post-call", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + return response + + +class PreCallGuardrail(CustomGuardrail): + """A pre-call-only guardrail — should NOT trigger deferral.""" + + def __init__(self): + super().__init__( + guardrail_name="pre-call", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + ) + + +class AllEventsGuardrail(CustomGuardrail): + """A guardrail with event_hook=None (runs on all events).""" + + def __init__(self): + super().__init__( + guardrail_name="all-events", + default_on=True, + event_hook=None, + ) + + +# --------------------------------------------------------------------------- +# 1. _has_post_call_guardrails detection +# --------------------------------------------------------------------------- + + +class TestHasPostCallGuardrails: + def test_returns_true_for_post_call_guardrail(self): + with patch("litellm.callbacks", [PostCallGuardrail()]): + assert ProxyBaseLLMRequestProcessing._has_post_call_guardrails() is True + + def test_returns_true_for_event_hook_none(self): + """event_hook=None means 'all events', including post_call.""" + with patch("litellm.callbacks", [AllEventsGuardrail()]): + assert ProxyBaseLLMRequestProcessing._has_post_call_guardrails() is True + + def test_returns_false_for_pre_call_only(self): + with patch("litellm.callbacks", [PreCallGuardrail()]): + assert ProxyBaseLLMRequestProcessing._has_post_call_guardrails() is False + + def test_returns_false_for_no_callbacks(self): + with patch("litellm.callbacks", []): + assert ProxyBaseLLMRequestProcessing._has_post_call_guardrails() is False + + def test_ignores_non_guardrail_callbacks(self): + """String callbacks and CustomLogger instances are not guardrails.""" + with patch("litellm.callbacks", ["langfuse", CustomLogger()]): + assert ProxyBaseLLMRequestProcessing._has_post_call_guardrails() is False + + def test_returns_true_for_list_with_post_call(self): + """event_hook as a list containing post_call should trigger deferral.""" + + class ListGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="list-post", + default_on=True, + event_hook=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call], + ) + + with patch("litellm.callbacks", [ListGuardrail()]): + assert ProxyBaseLLMRequestProcessing._has_post_call_guardrails() is True + + def test_returns_false_for_list_without_post_call(self): + """event_hook as a list without post_call should not trigger deferral.""" + + class ListGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="list-pre", + default_on=True, + event_hook=[GuardrailEventHooks.pre_call], + ) + + with patch("litellm.callbacks", [ListGuardrail()]): + assert ProxyBaseLLMRequestProcessing._has_post_call_guardrails() is False + + +# --------------------------------------------------------------------------- +# 2. Non-streaming: deferral flag → closure stored, create_task skipped +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_deferred_flag_stores_and_executes_closure(): + """ + When _defer_async_logging is True on logging_obj: + 1. wrapper_async stores a callable closure instead of calling create_task + 2. Calling the closure fires create_task + 3. Sync callbacks fire immediately (not deferred) + """ + mock_logging_obj = MagicMock() + mock_logging_obj._defer_async_logging = True + mock_logging_obj._enqueue_deferred_logging = None + + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello!", + litellm_logging_obj=mock_logging_obj, + ) + + # Closure was stored + enqueue_fn = mock_logging_obj._enqueue_deferred_logging + assert callable(enqueue_fn), "Closure should be stored on logging_obj" + + # Sync callbacks fired immediately + mock_logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once() + + # Calling the closure fires create_task + created_tasks = [] + real_create_task = asyncio.create_task + + def tracking_create_task(coro): + task = real_create_task(coro) + created_tasks.append(task) + return task + + with patch("asyncio.create_task", side_effect=tracking_create_task): + enqueue_fn() + + assert len(created_tasks) >= 1, "Closure should fire asyncio.create_task" + + for task in created_tasks: + if not task.done(): + task.cancel() + try: + await task + except (asyncio.CancelledError, Exception): + pass + + +# --------------------------------------------------------------------------- +# 3. Non-streaming regression: without flag, create_task fires normally +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_no_flag_fires_create_task_normally(): + """Without _defer_async_logging, wrapper_async calls create_task as before.""" + created_tasks = [] + real_create_task = asyncio.create_task + + def tracking_create_task(coro): + task = real_create_task(coro) + created_tasks.append(task) + return task + + with patch("asyncio.create_task", side_effect=tracking_create_task): + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello!", + ) + + assert len(created_tasks) >= 1 + + for task in created_tasks: + if not task.done(): + task.cancel() + try: + await task + except (asyncio.CancelledError, Exception): + pass + + +# --------------------------------------------------------------------------- +# 4. Non-streaming: deferred logging fires even if guardrail raises +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_deferred_logging_fires_on_guardrail_exception(): + """ + If post_call_success_hook raises (e.g., guardrail blocks content), + the deferred logging closure must still fire (via try/finally). + """ + from fastapi import HTTPException # noqa: local import for test isolation + + enqueue_called = False + + def mock_enqueue(): + nonlocal enqueue_called + enqueue_called = True + + class BlockingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="blocker", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + raise HTTPException(status_code=400, detail="Content blocked") + + guardrail = BlockingGuardrail() + + logging_obj = MagicMock() + logging_obj._enqueue_deferred_logging = mock_enqueue + + with patch("litellm.callbacks", [guardrail]): + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + with pytest.raises(HTTPException): + try: + await proxy_logging.post_call_success_hook( + data={"model": "gpt-4", "metadata": {}}, + response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + ) + finally: + # Mirrors the proxy's finally block + _enqueue_fn = getattr(logging_obj, "_enqueue_deferred_logging", None) + if _enqueue_fn is not None: + logging_obj._enqueue_deferred_logging = None + _enqueue_fn() + + assert enqueue_called is True + assert logging_obj._enqueue_deferred_logging is None + + +# --------------------------------------------------------------------------- +# 5. Streaming: closure defers logging at stream end +# --------------------------------------------------------------------------- + + +class TestDeferredStreamingClosure: + @pytest.mark.asyncio + async def test_streaming_closure_defers_logging(self): + """When _on_deferred_stream_complete is set, CSW calls the closure + instead of firing async_success_handler directly.""" + mock_logging_obj = MagicMock() + callback_called = False + callback_args = {} + + async def mock_callback(assembled_response, cache_hit): + nonlocal callback_called, callback_args + callback_called = True + callback_args = {"response": assembled_response, "cache_hit": cache_hit} + + mock_logging_obj._on_deferred_stream_complete = mock_callback + + resp = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello!", + stream=True, + litellm_logging_obj=mock_logging_obj, + ) + async for _ in resp: + pass + + await asyncio.sleep(0) + + assert callback_called is True, "Closure should be called at stream end" + assert callback_args["response"] is not None + assert mock_logging_obj._on_deferred_stream_complete is None + + @pytest.mark.asyncio + async def test_streaming_no_closure_fires_normally(self): + """Regression: without closure, CSW fires logging immediately.""" + created_tasks = [] + real_create_task = asyncio.create_task + + def tracking_create_task(coro): + task = real_create_task(coro) + created_tasks.append(task) + return task + + resp = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello!", + stream=True, + ) + with patch("asyncio.create_task", side_effect=tracking_create_task): + async for _ in resp: + pass + + assert len(created_tasks) >= 1 + for task in created_tasks: + if not task.done(): + task.cancel() + try: + await task + except (asyncio.CancelledError, Exception): + pass + + @pytest.mark.asyncio + async def test_closure_runs_only_guardrail_hooks(self): + """The closure must call only CustomGuardrail hooks, not all callbacks. + This is the key v2 change — PR #23929 called post_call_success_hook + which ran ALL callbacks, causing behavioral changes for streaming.""" + guardrail_called = False + logger_called = False + + class TrackingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="tracker", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + nonlocal guardrail_called + guardrail_called = True + return response + + class TrackingLogger(CustomLogger): + async def async_post_call_success_hook( + self, user_api_key_dict, data, response + ): + nonlocal logger_called + logger_called = True + return response + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): + pass + + mock_logging_obj.async_success_handler = track_async_success + + tracking_guardrail = TrackingGuardrail() + tracking_logger = TrackingLogger() + + # Use the real production static method via a thin closure + _captured_data = {"model": "gpt-4", "metadata": {}} + _captured_user_api_key_dict = UserAPIKeyAuth(api_key="test") + + async def _on_deferred_stream_complete(assembled_response, cache_hit): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data=_captured_data, + captured_user_api_key_dict=_captured_user_api_key_dict, + captured_logging_obj=mock_logging_obj, + assembled_response=assembled_response, + cache_hit=cache_hit, + ) + + mock_logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete + + with patch("litellm.callbacks", [tracking_guardrail, tracking_logger]): + resp = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello!", + stream=True, + litellm_logging_obj=mock_logging_obj, + ) + async for _ in resp: + pass + + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert guardrail_called is True, "Guardrail hook should be called" + assert logger_called is False, "Non-guardrail logger should NOT be called by closure" + + @pytest.mark.asyncio + async def test_closure_passes_guardrail_modified_response_to_logging(self): + """The production _run_deferred_stream_guardrails must pass the + guardrail-modified response to async_success_handler.""" + logged_response = None + modified_response = MagicMock() + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): + nonlocal logged_response + logged_response = args[0] if args else None + + mock_logging_obj.async_success_handler = track_async_success + + class ModifyingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="modifier", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + return modified_response + + guardrail = ModifyingGuardrail() + + with patch("litellm.callbacks", [guardrail]): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "gpt-4", "metadata": {}}, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=MagicMock(), + cache_hit=False, + ) + + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert logged_response is modified_response, \ + "Logging must receive the guardrail-modified response" + + @pytest.mark.asyncio + async def test_closure_logs_even_on_guardrail_exception(self): + """If a guardrail raises HTTPException, the production + _run_deferred_stream_guardrails must still fire logging + and set guardrail_blocked in metadata.""" + from fastapi import HTTPException # noqa: local import for test isolation + + logging_called = False + + class BlockingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="blocker", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + raise HTTPException(status_code=400, detail="Blocked") + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): + nonlocal logging_called + logging_called = True + + mock_logging_obj.async_success_handler = track_async_success + + guardrail = BlockingGuardrail() + + with patch("litellm.callbacks", [guardrail]): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "gpt-4", "metadata": {}}, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=MagicMock(), + cache_hit=False, + ) + + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert logging_called is True, \ + "Logging must fire even when guardrail raises HTTPException" + assert mock_logging_obj.model_call_details["metadata"].get( + "guardrail_blocked" + ) is True, "guardrail_blocked must be set for HTTPException" + + @pytest.mark.asyncio + async def test_transient_error_does_not_set_guardrail_blocked(self): + """Transient errors (not HTTPException) should NOT set + guardrail_blocked. Uses the production _run_deferred_stream_guardrails.""" + + class TransientErrorGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="transient", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + raise ConnectionError("Network timeout") + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): + pass + + mock_logging_obj.async_success_handler = track_async_success + + guardrail = TransientErrorGuardrail() + + with patch("litellm.callbacks", [guardrail]): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "gpt-4", "metadata": {}}, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=MagicMock(), + cache_hit=False, + ) + + await asyncio.sleep(0) + + assert mock_logging_obj.model_call_details["metadata"].get( + "guardrail_blocked" + ) is not True, "guardrail_blocked must NOT be set for transient errors" + + @pytest.mark.asyncio + async def test_production_closure_integration(self): + """Integration test: calls the real _run_deferred_stream_guardrails + static method and verifies it calls guardrail hooks and passes + the modified response to logging.""" + hook_called = False + logged_response = None + modified_response = MagicMock() + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): + nonlocal logged_response + logged_response = args[0] if args else None + + mock_logging_obj.async_success_handler = track_async_success + + class TestGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="test", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + nonlocal hook_called + hook_called = True + return modified_response + + guardrail = TestGuardrail() + + async def _on_deferred_stream_complete(assembled_response, cache_hit): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "gpt-4", "metadata": {}}, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=assembled_response, + cache_hit=cache_hit, + ) + + mock_logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete + + with patch("litellm.callbacks", [guardrail]): + resp = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello!", + stream=True, + litellm_logging_obj=mock_logging_obj, + ) + async for _ in resp: + pass + + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert hook_called is True, \ + "Production closure must call guardrail hook" + assert logged_response is modified_response, \ + "Production closure must pass guardrail-modified response to logging" + + @pytest.mark.asyncio + async def test_apply_guardrail_path_uses_unified_guardrail(self): + """Guardrails that define apply_guardrail should be dispatched through + UnifiedLLMGuardrails.async_post_call_success_hook via the real + _run_deferred_stream_guardrails static method.""" + from litellm.types.utils import GenericGuardrailAPIInputs + + unified_hook_called = False + + class ApplyGuardrailType(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="apply-type", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def apply_guardrail( + self, inputs, request_data, input_type, logging_obj=None + ) -> GenericGuardrailAPIInputs: + nonlocal unified_hook_called + unified_hook_called = True + return inputs + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + logged_response = None + + async def track_async_success(*args, **kwargs): + nonlocal logged_response + logged_response = args[0] if args else None + + mock_logging_obj.async_success_handler = track_async_success + + guardrail = ApplyGuardrailType() + + async def _on_deferred_stream_complete(assembled_response, cache_hit): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "gpt-4", "metadata": {}}, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=assembled_response, + cache_hit=cache_hit, + ) + + mock_logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete + + with patch("litellm.callbacks", [guardrail]): + resp = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello!", + stream=True, + litellm_logging_obj=mock_logging_obj, + ) + async for _ in resp: + pass + + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert unified_hook_called is True, \ + "apply_guardrail guardrails must be dispatched through UnifiedLLMGuardrails" + assert logged_response is not None, \ + "Logging must fire after unified guardrail path" + + @pytest.mark.asyncio + async def test_hooks_receive_merged_guardrail_data(self): + """Hooks must receive guardrail_data (the merged dict from + _check_and_merge_model_level_guardrails), not the original + captured_data. This ensures model-level non-default guardrails + are visible to any inner should_run_guardrail re-checks. + + Uses a deep-copy mock to break the shallow-copy side-effect that + would otherwise mask the bug — verifying the code is explicitly + correct, not correct-by-accident.""" + import copy + + hook_received_data = None + + class InspectingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="inspector", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + nonlocal hook_received_data + hook_received_data = data + return response + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): + pass + + mock_logging_obj.async_success_handler = track_async_success + + guardrail = InspectingGuardrail() + + captured_data = {"model": "gpt-4", "metadata": {"existing_key": "value"}} + + def mock_merge(data, llm_router): + """Return a fully independent dict (deep copy) so the original + captured_data is NOT mutated. This simulates a correct merge + implementation and proves _run_deferred_stream_guardrails uses + the return value, not the original data.""" + merged = copy.deepcopy(data) + merged["metadata"]["guardrails"] = ["model-guardrail"] + merged["_merged_marker"] = True + return merged + + with patch("litellm.callbacks", [guardrail]), \ + patch( + "litellm.proxy.utils._check_and_merge_model_level_guardrails", + side_effect=mock_merge, + ): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data=captured_data, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=MagicMock(), + cache_hit=False, + ) + + assert hook_received_data is not None, "Guardrail hook must be called" + assert hook_received_data.get("_merged_marker") is True, \ + "Hook must receive guardrail_data (merged), not original captured_data" + assert "model-guardrail" in hook_received_data.get("metadata", {}).get( + "guardrails", [] + ), "Hook data must contain model-level guardrails" + + @pytest.mark.asyncio + async def test_apply_guardrail_path_receives_merged_guardrail_data(self): + """The apply_guardrail path (through UnifiedLLMGuardrails) must also + receive guardrail_data so that the inner should_run_guardrail re-check + inside UnifiedLLMGuardrails sees model-level guardrails. + + This is the specific scenario Greptile flagged: a default_on=False + guardrail configured at the model level would pass the outer gate but + be silently skipped at execution time if captured_data (unmerged) were + passed instead of guardrail_data (merged).""" + import copy + from litellm.types.utils import GenericGuardrailAPIInputs + + unified_received_data = None + + class ModelLevelApplyGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="model-apply-guardrail", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def apply_guardrail( + self, inputs, request_data, input_type, logging_obj=None + ) -> GenericGuardrailAPIInputs: + return inputs + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): + pass + + mock_logging_obj.async_success_handler = track_async_success + + guardrail = ModelLevelApplyGuardrail() + captured_data = {"model": "gpt-4", "metadata": {}} + + def mock_merge(data, llm_router): + merged = copy.deepcopy(data) + merged["metadata"]["guardrails"] = ["model-apply-guardrail"] + merged["_merged_marker"] = True + return merged + + # Capture what UnifiedLLMGuardrails.async_post_call_success_hook receives + original_unified_hook = None + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + original_unified_hook = UnifiedLLMGuardrails.async_post_call_success_hook + + async def tracking_unified_hook(self, user_api_key_dict, data, response): + nonlocal unified_received_data + unified_received_data = data + return response + + with patch("litellm.callbacks", [guardrail]), \ + patch( + "litellm.proxy.utils._check_and_merge_model_level_guardrails", + side_effect=mock_merge, + ), \ + patch.object( + UnifiedLLMGuardrails, + "async_post_call_success_hook", + tracking_unified_hook, + ): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data=captured_data, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=MagicMock(), + cache_hit=False, + ) + + assert unified_received_data is not None, \ + "UnifiedLLMGuardrails must be called for apply_guardrail guardrails" + assert unified_received_data.get("_merged_marker") is True, \ + "UnifiedLLMGuardrails must receive guardrail_data (merged), not captured_data" + assert "model-apply-guardrail" in unified_received_data.get( + "metadata", {} + ).get("guardrails", []), \ + "UnifiedLLMGuardrails data must contain model-level guardrails" + + @pytest.mark.asyncio + async def test_multiple_guardrails_all_receive_merged_data(self): + """When multiple guardrails are configured, ALL of them must receive + guardrail_data (merged), not just the first one.""" + import copy + + received_data_per_guardrail = {} + + class TaggedGuardrail(CustomGuardrail): + def __init__(self, name): + super().__init__( + guardrail_name=name, + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + received_data_per_guardrail[self.guardrail_name] = data + return response + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): + pass + + mock_logging_obj.async_success_handler = track_async_success + + guardrail_a = TaggedGuardrail("guardrail-a") + guardrail_b = TaggedGuardrail("guardrail-b") + + captured_data = {"model": "gpt-4", "metadata": {}} + + def mock_merge(data, llm_router): + merged = copy.deepcopy(data) + merged["metadata"]["guardrails"] = ["guardrail-a", "guardrail-b"] + merged["_merged_marker"] = True + return merged + + with patch("litellm.callbacks", [guardrail_a, guardrail_b]), \ + patch( + "litellm.proxy.utils._check_and_merge_model_level_guardrails", + side_effect=mock_merge, + ): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data=captured_data, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=MagicMock(), + cache_hit=False, + ) + + for name in ("guardrail-a", "guardrail-b"): + assert name in received_data_per_guardrail, \ + f"{name} must be called" + assert received_data_per_guardrail[name].get("_merged_marker") is True, \ + f"{name} must receive guardrail_data (merged), not captured_data" + + @pytest.mark.asyncio + async def test_logging_fires_even_if_guardrail_init_raises(self): + """If _check_and_merge_model_level_guardrails raises during + initialization, logging must still fire via the try/finally guard. + This prevents silent logging loss on transient init errors.""" + logging_called = False + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): + nonlocal logging_called + logging_called = True + + mock_logging_obj.async_success_handler = track_async_success + + def exploding_merge(data, llm_router): + raise RuntimeError("Simulated init failure") + + with patch( + "litellm.proxy.utils._check_and_merge_model_level_guardrails", + side_effect=exploding_merge, + ): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "gpt-4", "metadata": {}}, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=MagicMock(), + cache_hit=False, + ) + + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert logging_called is True, \ + "Logging must fire even when guardrail initialization raises"