mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
fix(proxy): log blocked streaming guardrail responses as failures, not success
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
jesus
parent
7da6fe54b5
commit
0f9dc5cbcb
3 changed files with 184 additions and 6 deletions
|
|
@ -82,7 +82,11 @@ from litellm._logging import _redact_string, verbose_proxy_logger
|
|||
from litellm._service_logger import ServiceLogging, ServiceTypes
|
||||
from litellm.caching.caching import DualCache, RedisCache
|
||||
from litellm.caching.dual_cache import LimitedSizeOrderedDict
|
||||
from litellm.exceptions import RejectedRequestError, SensitiveDataRouteException
|
||||
from litellm.exceptions import (
|
||||
GuardrailRaisedException,
|
||||
RejectedRequestError,
|
||||
SensitiveDataRouteException,
|
||||
)
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
ModifyResponseException,
|
||||
|
|
@ -2628,6 +2632,7 @@ class ProxyLogging:
|
|||
- Authentication Errors from user_api_key_auth
|
||||
- HTTP HTTPException (rate limit errors)
|
||||
- ProxyException (guardrail blocks, budget / rate-limit errors)
|
||||
- GuardrailRaisedException (guardrail blocks / guardrail failures)
|
||||
"""
|
||||
|
||||
#########################################################
|
||||
|
|
@ -2642,9 +2647,9 @@ class ProxyLogging:
|
|||
if not (RouteChecks.is_llm_api_route(route) or RouteChecks.is_info_route(route)):
|
||||
return False
|
||||
|
||||
return isinstance(original_exception, (HTTPException, ProxyException)) or (
|
||||
error_type == ProxyErrorTypes.auth_error
|
||||
)
|
||||
return isinstance(
|
||||
original_exception, (HTTPException, ProxyException, GuardrailRaisedException)
|
||||
) or (error_type == ProxyErrorTypes.auth_error)
|
||||
|
||||
async def _handle_logging_proxy_only_error(
|
||||
self,
|
||||
|
|
@ -3181,7 +3186,7 @@ class ProxyLogging:
|
|||
except (GeneratorExit, asyncio.CancelledError):
|
||||
raise
|
||||
except Exception:
|
||||
ProxyLogging._fire_deferred_stream_logging(request_data)
|
||||
ProxyLogging._discard_deferred_stream_logging_for_failure(request_data)
|
||||
raise
|
||||
ProxyLogging._fire_deferred_stream_logging(request_data)
|
||||
return
|
||||
|
|
@ -3241,7 +3246,7 @@ class ProxyLogging:
|
|||
except (GeneratorExit, asyncio.CancelledError):
|
||||
raise
|
||||
except Exception:
|
||||
ProxyLogging._fire_deferred_stream_logging(request_data)
|
||||
ProxyLogging._discard_deferred_stream_logging_for_failure(request_data)
|
||||
raise
|
||||
|
||||
# Fire deferred logging AFTER all guardrail end-of-stream blocks
|
||||
|
|
@ -3272,6 +3277,44 @@ class ProxyLogging:
|
|||
logging_obj._deferred_stream_complete_args = None
|
||||
asyncio.create_task(_deferred_cb(*_args))
|
||||
|
||||
@staticmethod
|
||||
def _discard_deferred_stream_logging_for_failure(request_data: dict) -> None:
|
||||
"""Discard the deferred stream-complete dispatch when the stream ends in
|
||||
a failure (e.g. an end-of-stream guardrail block raising out of the
|
||||
callback chain). The deferred dispatch is the success logging path —
|
||||
firing it here would record the blocked request as a success callback
|
||||
and a ``status=success`` spend row before the outer generator's
|
||||
``post_call_failure_hook`` writes the failure row. The CSW shape parks
|
||||
``(assembled ModelResponse, cache_hit)``; record its partial usage so
|
||||
the failure row bills what the stream consumed instead of zero. The
|
||||
native /v1/messages and responses shapes park ``(coroutine,)`` and
|
||||
still need the flush (no success row is produced without it), so they
|
||||
keep the existing fire behaviour.
|
||||
"""
|
||||
logging_obj: Final = request_data.get("litellm_logging_obj")
|
||||
if logging_obj is None:
|
||||
return
|
||||
_deferred_cb: Final[Callable[..., Coroutine[object, object, object]] | None] = getattr(
|
||||
logging_obj, "_on_deferred_stream_complete", None
|
||||
)
|
||||
_args: Final[tuple[object, ...] | None] = getattr(
|
||||
logging_obj, "_deferred_stream_complete_args", None
|
||||
)
|
||||
if _deferred_cb is None or _args is None:
|
||||
return
|
||||
assembled: Final = _args[0]
|
||||
if not isinstance(assembled, ModelResponse):
|
||||
ProxyLogging._fire_deferred_stream_logging(request_data)
|
||||
return
|
||||
logging_obj._on_deferred_stream_complete = None
|
||||
logging_obj._deferred_stream_complete_args = None
|
||||
usage: Final[Usage | None] = getattr(assembled, "usage", None)
|
||||
if isinstance(usage, Usage):
|
||||
logging_obj.record_partial_usage_for_failure(
|
||||
usage,
|
||||
logging_obj._response_cost_calculator(result=assembled) or 0.0,
|
||||
)
|
||||
|
||||
async def _arelease_max_parallel_requests_on_disconnect(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import pytest
|
|||
from fastapi import HTTPException
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import GuardrailRaisedException
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.proxy._types import AlertType, ProxyErrorTypes
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
|
@ -49,12 +50,17 @@ def test_is_proxy_only_llm_api_truth_table(proxy_logging):
|
|||
error_type=ProxyErrorTypes.auth_error,
|
||||
route="/chat/completions",
|
||||
),
|
||||
"guardrail_raised_on_llm_route": proxy_logging._is_proxy_only_llm_api_error(
|
||||
original_exception=GuardrailRaisedException(guardrail_name="g", message="blocked"),
|
||||
route="/chat/completions",
|
||||
),
|
||||
}
|
||||
assert snapshot == {
|
||||
"no_route": False,
|
||||
"non_llm_route": False,
|
||||
"http_on_llm_route": True,
|
||||
"auth_short_circuit": True,
|
||||
"guardrail_raised_on_llm_route": True,
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -267,3 +273,49 @@ async def test_handle_logging_proxy_only_path_propagates_async_failure_raises(
|
|||
route="/chat/completions",
|
||||
original_exception=Exception("x"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_failure_hook_guardrail_block_fires_failure_callback(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
):
|
||||
"""A ``GuardrailRaisedException`` on an LLM route must reach the logging
|
||||
object's ``async_failure_handler`` so custom loggers see a ``failure``
|
||||
status - without this, guardrail blocks produce only
|
||||
``post_call_failure_hook`` and no failure logging event."""
|
||||
from datetime import datetime
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
recorded: dict[str, Any] = {}
|
||||
|
||||
class _StatusRecorder(CustomLogger):
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
recorded["status"] = (kwargs.get("standard_logging_object") or {}).get("status")
|
||||
|
||||
monkeypatch.setattr(litellm, "_async_failure_callback", [_StatusRecorder()])
|
||||
logging_obj = LiteLLMLoggingObj(
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=False,
|
||||
call_type="acompletion",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="test_guardrail_block_failure_cb",
|
||||
function_id="test_guardrail_block_failure_cb",
|
||||
)
|
||||
request_data = {
|
||||
"litellm_logging_obj": logging_obj,
|
||||
"litellm_call_id": "test_guardrail_block_failure_cb",
|
||||
"model": "gpt-4o-mini",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {},
|
||||
}
|
||||
proxy_logging.alert_types = []
|
||||
await proxy_logging.post_call_failure_hook(
|
||||
request_data=request_data,
|
||||
original_exception=GuardrailRaisedException(guardrail_name="g", message="blocked"),
|
||||
user_api_key_dict=make_user_api_key_auth(request_route="/chat/completions"),
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
await asyncio.sleep(0)
|
||||
assert recorded["status"] == "failure"
|
||||
|
|
|
|||
|
|
@ -479,6 +479,89 @@ async def test_native_messages_stream_logging_fires_when_guardrail_blocks_after_
|
|||
assert logging_obj._deferred_stream_complete_args is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_stream_guardrail_block_after_stream_end_logs_failure_not_success(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
):
|
||||
"""
|
||||
On /chat/completions streams the CSW shape parks
|
||||
``(assembled ModelResponse, cache_hit)`` as the deferred args. A guardrail
|
||||
that raises ``GuardrailRaisedException`` at end of stream must NOT dispatch
|
||||
that deferred success logging - the request is logged via the failure path
|
||||
instead, with the consumed usage carried over so the failure row bills
|
||||
correctly.
|
||||
"""
|
||||
from litellm.exceptions import GuardrailRaisedException
|
||||
from litellm.types.utils import Usage
|
||||
|
||||
events: List[Any] = []
|
||||
request_data: Dict[str, Any] = {"metadata": {}}
|
||||
logging_obj = LiteLLMLoggingObj(
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
call_type="acompletion",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="test_chat_stream_guardrail_block",
|
||||
function_id="test_chat_stream_guardrail_block",
|
||||
)
|
||||
logging_obj.optional_params = {}
|
||||
logging_obj.litellm_params = {}
|
||||
logging_obj.standard_built_in_tools_params = None
|
||||
|
||||
async def _dispatch_deferred_logging(*args):
|
||||
events.append("success_dispatched")
|
||||
|
||||
logging_obj._on_deferred_stream_complete = _dispatch_deferred_logging
|
||||
request_data["litellm_logging_obj"] = logging_obj
|
||||
|
||||
assembled = litellm.ModelResponse(
|
||||
model="gpt-4o-mini",
|
||||
choices=[{"index": 0, "message": {"role": "assistant", "content": "BANANA"}}],
|
||||
usage=Usage(prompt_tokens=3, completion_tokens=5, total_tokens=8),
|
||||
)
|
||||
|
||||
async def _upstream():
|
||||
yield {"id": "c1", "choices": [{"index": 0, "delta": {"content": "BAN"}}]}
|
||||
yield {"id": "c1", "choices": [{"index": 0, "delta": {"content": "ANA"}}]}
|
||||
logging_obj._deferred_stream_complete_args = (assembled, False)
|
||||
|
||||
class _BlockingGuardrail(CustomLogger):
|
||||
async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):
|
||||
async for chunk in response:
|
||||
yield chunk
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name="g", message="blocked", blocked_content=True
|
||||
)
|
||||
|
||||
monkeypatch.setattr(litellm, "callbacks", [_BlockingGuardrail()])
|
||||
|
||||
with pytest.raises(GuardrailRaisedException):
|
||||
async for _ in proxy_logging.async_post_call_streaming_iterator_hook(
|
||||
response=_upstream(),
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
request_data=request_data,
|
||||
):
|
||||
pass
|
||||
await asyncio.sleep(0)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
snapshot = {
|
||||
"events": events,
|
||||
"callback_cleared": logging_obj._on_deferred_stream_complete is None,
|
||||
"args_cleared": logging_obj._deferred_stream_complete_args is None,
|
||||
"combined_usage_total_tokens": logging_obj.model_call_details["combined_usage_object"].total_tokens,
|
||||
"response_cost_positive": logging_obj.model_call_details["response_cost"] > 0,
|
||||
}
|
||||
assert snapshot == {
|
||||
"events": [],
|
||||
"callback_cleared": True,
|
||||
"args_cleared": True,
|
||||
"combined_usage_total_tokens": 8,
|
||||
"response_cost_positive": True,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _fire_deferred_stream_logging
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue