fix(proxy): defer logging until post-call guardrails complete
guardrail_information is None in StandardLoggingPayload because logging fires before post-call guardrails write to metadata. Non-streaming: wrapper_async stores a closure instead of calling create_task immediately. The proxy fires it in a try/finally after post_call_success_hook so the SLP is built with guardrail info. Streaming: a closure on logging_obj is called by CSW.__anext__ at stream end. The closure runs only guardrail hooks (not all callbacks) on the assembled response, then fires both logging handlers. This avoids behavioral changes for non-guardrail callbacks on streaming.
This commit is contained in:
michelligabriele
parent
e5baa2232f
commit
001501fb31
@ -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
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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.
|
||||
_has_post_call_guardrails = 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 _has_post_call_guardrails 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,181 @@ class ProxyBaseLLMRequestProcessing:
|
||||
|
||||
response = responses[1]
|
||||
|
||||
hidden_params = getattr(response, "_hidden_params", {}) or {}
|
||||
model_id = self._get_model_id_from_response(hidden_params, self.data)
|
||||
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 _has_post_call_guardrails 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
|
||||
)
|
||||
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
|
||||
)
|
||||
|
||||
# Always return the client-requested model name (not provider-prefixed internal identifiers)
|
||||
# for OpenAI-compatible responses.
|
||||
@ -1217,6 +1296,126 @@ 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
|
||||
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
|
||||
UnifiedLLMGuardrails,
|
||||
)
|
||||
from litellm.proxy.proxy_server import llm_router as _global_llm_router
|
||||
from litellm.proxy.utils import _check_and_merge_model_level_guardrails
|
||||
|
||||
_response = assembled_response
|
||||
_unified_guardrail = UnifiedLLMGuardrails()
|
||||
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__:
|
||||
captured_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=captured_data,
|
||||
response=_response,
|
||||
)
|
||||
)
|
||||
else:
|
||||
guardrail_result = await cb.async_post_call_success_hook(
|
||||
user_api_key_dict=captured_user_api_key_dict,
|
||||
data=captured_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
|
||||
|
||||
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,
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -0,0 +1,687 @@
|
||||
"""
|
||||
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
|
||||
from starlette.exceptions import HTTPException
|
||||
|
||||
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).
|
||||
"""
|
||||
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 closure passes the guardrail-modified response to logging handlers."""
|
||||
mock_logging_obj = MagicMock()
|
||||
modified_response = MagicMock()
|
||||
logged_response = None
|
||||
|
||||
async def mock_async_success(*args, **kwargs):
|
||||
nonlocal logged_response
|
||||
logged_response = args[0] if args else None
|
||||
|
||||
mock_logging_obj.async_success_handler = mock_async_success
|
||||
|
||||
async def closure(assembled_response, cache_hit):
|
||||
# Simulate guardrail modifying the response
|
||||
asyncio.create_task(
|
||||
mock_logging_obj.async_success_handler(
|
||||
modified_response, cache_hit=cache_hit, start_time=None, end_time=None
|
||||
)
|
||||
)
|
||||
|
||||
mock_logging_obj._on_deferred_stream_complete = closure
|
||||
|
||||
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 logged_response is modified_response
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_closure_logs_even_on_guardrail_exception(self):
|
||||
"""If the guardrail raises HTTPException, logging still fires
|
||||
and guardrail_blocked is set in metadata."""
|
||||
logging_called = False
|
||||
|
||||
async def mock_async_success(*args, **kwargs):
|
||||
nonlocal logging_called
|
||||
logging_called = True
|
||||
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.model_call_details = {"metadata": {}}
|
||||
mock_logging_obj.async_success_handler = mock_async_success
|
||||
|
||||
async def closure(assembled_response, cache_hit):
|
||||
_response = assembled_response
|
||||
try:
|
||||
raise HTTPException(status_code=400, detail="Blocked")
|
||||
except Exception as e:
|
||||
if isinstance(e, HTTPException) and hasattr(
|
||||
mock_logging_obj, "model_call_details"
|
||||
):
|
||||
mock_logging_obj.model_call_details.setdefault(
|
||||
"metadata", {}
|
||||
)["guardrail_blocked"] = True
|
||||
|
||||
asyncio.create_task(
|
||||
mock_logging_obj.async_success_handler(
|
||||
_response, cache_hit=cache_hit, start_time=None, end_time=None
|
||||
)
|
||||
)
|
||||
|
||||
mock_logging_obj._on_deferred_stream_complete = closure
|
||||
|
||||
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 logging_called is True
|
||||
assert mock_logging_obj.model_call_details["metadata"].get(
|
||||
"guardrail_blocked"
|
||||
) is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_transient_error_does_not_set_guardrail_blocked(self):
|
||||
"""Transient errors (not HTTPException) should NOT set guardrail_blocked."""
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.model_call_details = {"metadata": {}}
|
||||
|
||||
async def mock_async_success(*args, **kwargs):
|
||||
pass
|
||||
|
||||
mock_logging_obj.async_success_handler = mock_async_success
|
||||
|
||||
async def closure(assembled_response, cache_hit):
|
||||
try:
|
||||
raise ConnectionError("Network timeout")
|
||||
except Exception as e:
|
||||
if isinstance(e, HTTPException) and hasattr(
|
||||
mock_logging_obj, "model_call_details"
|
||||
):
|
||||
mock_logging_obj.model_call_details.setdefault(
|
||||
"metadata", {}
|
||||
)["guardrail_blocked"] = True
|
||||
|
||||
asyncio.create_task(
|
||||
mock_logging_obj.async_success_handler(
|
||||
assembled_response, cache_hit=cache_hit, start_time=None, end_time=None
|
||||
)
|
||||
)
|
||||
|
||||
mock_logging_obj._on_deferred_stream_complete = closure
|
||||
|
||||
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 mock_logging_obj.model_call_details["metadata"].get(
|
||||
"guardrail_blocked"
|
||||
) is not True
|
||||
|
||||
@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"
|
||||
Loading…
Reference in New Issue
Block a user