tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-14 02:22:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #24135 from michelligabriele/fix/deferred-guardrail-logging-v2

Browse Source 
fix(proxy): defer logging until post-call guardrails complete
This commit is contained in:
Krish Dholakia
2026-03-19 17:57:00 -07:00
committed by GitHub
parent b20c448188 573f6b78ea
commit d4857f682d
10 changed files with 1498 additions and 162 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/proxy/guardrails/custom_guardrail.md
+10 -2
View File
@@ -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.
:::
<details>
<summary>Advanced: Multiple modes with individual event hooks</summary>
@@ -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
docs/my-website/sidebars.js
+14
View File
@@ -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",
litellm/litellm_core_utils/litellm_logging.py
+14
View File
@@ -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=(
litellm/litellm_core_utils/streaming_handler.py
+25 -11
View File
@@ -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:
litellm/proxy/_new_secret_config.yaml
+27 -36
View File
@@ -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)
litellm/proxy/common_request_processing.py
+356 -104
View File
@@ -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,
litellm/utils.py
+31 -8
View File
@@ -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,
poetry.lock
Generated
+1 -1
View File
@@ -8018,4 +8018,4 @@ utils = ["numpydoc"]
[metadata]
lock-version = "2.1"
python-versions = ">=3.9,<4.0"
content-hash = "eda34dfd8b35474beffee18893d6782c7b3d0d3d2c610f66237eb97176f43527"
content-hash = "2cf958f1a04fd5f1ab0e5cfc33bdbf441b518ed6c82d0f2546bf64cd3d2f89be"
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
+76
View File
@@ -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."""
tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py
+944
View File
@@ -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"