scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-07 07:58:09 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 150

refactor(proxy): reuse streaming guardrail call type

Browse source
This commit is contained in:
zouxiaoyue 2026-08-10 10:21:34 +08:00
parent 9fb6f38034
commit 2ded3d4cf4
3 changed files with 20 additions and 48 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/integrations/custom_guardrail.py
View file

@ -626,9 +626,6 @@ class CustomGuardrail(CustomLogger):
def uses_apply_guardrail_interface(self) -> bool:
return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail
def supports_streaming_call_type(self, call_type: CallTypes) -> bool:
return call_type in self.supported_streaming_call_types
def _deployment_pre_call_target(self) -> "CustomLogger":
if not self.uses_apply_guardrail_interface():
return self

37
litellm/proxy/utils.py
View file

@ -1753,35 +1753,23 @@ class ProxyLogging:
return caps
@staticmethod
def _stream_requires_guardrail_translation(user_api_key_dict: UserAPIKeyAuth) -> bool:
from litellm.litellm_core_utils.api_route_to_call_types import (
get_call_types_for_route,
)
route: Final = user_api_key_dict.request_route
if not route:
return False
call_types: Final = get_call_types_for_route(route)
if not call_types:
return False
return call_types[0] in NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES
@staticmethod
def _guardrail_supports_streaming_route(
guardrail: CustomGuardrail,
def _stream_guardrail_translation_call_type(
user_api_key_dict: UserAPIKeyAuth,
) -> bool:
) -> CallTypes | None:
from litellm.litellm_core_utils.api_route_to_call_types import (
get_call_types_for_route,
)
route: Final = user_api_key_dict.request_route
if not route:
return False
return None
call_types: Final = get_call_types_for_route(route)
if not call_types:
return False
return any(guardrail.supports_streaming_call_type(call_type) for call_type in call_types)
return None
call_type: Final = call_types[0]
if call_type not in NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES:
return None
return call_type
@staticmethod
def has_post_call_response_headers_callbacks() -> bool:
@ -2744,7 +2732,7 @@ class ProxyLogging:
request_data = _check_and_merge_model_level_guardrails(data=request_data, llm_router=llm_router)
current_response = response
stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict)
stream_translation_call_type: Final = ProxyLogging._stream_guardrail_translation_call_type(user_api_key_dict)
for resolved_callback, kind in caps.iterator_overrides:
if isinstance(resolved_callback, CustomGuardrail):
@ -2757,14 +2745,11 @@ class ProxyLogging:
"apply_guardrail"
if (
kind == "override"
and stream_needs_translation
and stream_translation_call_type is not None
and isinstance(resolved_callback, CustomGuardrail)
and resolved_callback.uses_apply_guardrail_interface()
and not resolved_callback.mask_response_content
and not ProxyLogging._guardrail_supports_streaming_route(
resolved_callback,
user_api_key_dict,
)
and stream_translation_call_type not in resolved_callback.supported_streaming_call_types
)
else kind
)

28
tests/test_litellm/proxy/test_proxy_logging_hook_detection.py
View file

@ -51,17 +51,6 @@ def test_has_streaming_callbacks_detects_guardrails(monkeypatch):
assert ProxyLogging.has_streaming_callbacks() is True
def test_custom_guardrail_streaming_call_type_capability_is_opt_in():
from litellm.types.utils import CallTypes
class _ProtocolAwareGuardrail(CustomGuardrail):
supported_streaming_call_types = frozenset({CallTypes.anthropic_messages})
assert CustomGuardrail().supports_streaming_call_type(CallTypes.anthropic_messages) is False
assert _ProtocolAwareGuardrail().supports_streaming_call_type(CallTypes.anthropic_messages) is True
assert _ProtocolAwareGuardrail().supports_streaming_call_type(CallTypes.aresponses) is False
@pytest.mark.asyncio
async def test_post_call_response_headers_hook_returns_early_without_callbacks(
monkeypatch,
@ -244,27 +233,28 @@ def _streaming_logging_obj():
)
def test_stream_requires_guardrail_translation_route_detection():
def test_stream_guardrail_translation_call_type_route_detection():
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import CallTypes
assert (
ProxyLogging._stream_requires_guardrail_translation(
ProxyLogging._stream_guardrail_translation_call_type(
UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages")
)
is True
is CallTypes.anthropic_messages
)
assert (
ProxyLogging._stream_requires_guardrail_translation(
ProxyLogging._stream_guardrail_translation_call_type(
UserAPIKeyAuth(api_key="sk-1234", request_route="/chat/completions")
)
is False
is None
)
assert ProxyLogging._stream_requires_guardrail_translation(UserAPIKeyAuth(api_key="sk-1234")) is False
assert ProxyLogging._stream_guardrail_translation_call_type(UserAPIKeyAuth(api_key="sk-1234")) is None
assert (
ProxyLogging._stream_requires_guardrail_translation(
ProxyLogging._stream_guardrail_translation_call_type(
UserAPIKeyAuth(api_key="sk-1234", request_route="/route/without/call/types")
)
is False
is None
)