diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8f41db94b98..53a48a2731c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1512,6 +1512,10 @@ def _overrides_streaming_chunk_hook(callback: CustomLogger) -> bool: ) +def _assigns_streaming_chunk_hook(callback: CustomLogger) -> bool: + return "async_post_call_streaming_hook" in vars(callback) + + def _overrides_moderation_hook(callback: CustomLogger) -> bool: return _overrides_hook(callback, "async_moderation_hook") @@ -3252,8 +3256,12 @@ class ProxyLogging: if resolved.enforces_request_content is True: has_content_enforcer = True + # A hook set on the instance runs per chunk only beside a guardrail, + # the one case where the per-chunk loop used to reach it. streaming_chunk_overrides: Final = tuple( - callback for callback in resolved_callbacks if _overrides_streaming_chunk_hook(callback) + callback + for callback in resolved_callbacks + if _overrides_streaming_chunk_hook(callback) or (has_guardrail and _assigns_streaming_chunk_hook(callback)) ) caps: Final = _CallbackCapabilities( has_post_call_response_headers=has_post_call_response_headers, diff --git a/tests/integration/streaming/test_request_only_guardrail_stream_fast_path.py b/tests/integration/streaming/test_request_only_guardrail_stream_fast_path.py index 01a0f728ff3..4f728d44aac 100644 --- a/tests/integration/streaming/test_request_only_guardrail_stream_fast_path.py +++ b/tests/integration/streaming/test_request_only_guardrail_stream_fast_path.py @@ -35,6 +35,7 @@ _CHUNK_GUARDRAIL: Final = "per-chunk-check" _OPT_IN: Final = "opt-in-chunk-check" _INHERITED: Final = "inherited-chunk-check" _RAISING: Final = "raising-chunk-check" +_INSTANCE: Final = "instance-chunk-check" _CHUNK_ERROR: Final = "Chunk check failed" _MODULE: Final = "request_only_guardrail" _LOG_ENV: Final = "REQUEST_ONLY_GUARDRAIL_LOG" @@ -89,6 +90,16 @@ class InheritedPerChunkGuardrail(PerChunkGuardrail): pass +class InstancePerChunkGuardrail(CustomGuardrail): + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.async_post_call_streaming_hook = self._see_chunk + + async def _see_chunk(self, user_api_key_dict, response): + _log(_marker_in(response), "chunk:" + self.guardrail_name) + return None + + class RaisingPerChunkGuardrail(CustomGuardrail): async def async_post_call_streaming_hook(self, user_api_key_dict, response): if "word3 " in response: @@ -339,6 +350,16 @@ def opt_in_rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactor yield built +@pytest.fixture(scope="module") +def instance_rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with _rig( + session_gateway, + tmp_path_factory.mktemp("request-only-guardrail-beside-instance-per-chunk-hook"), + [*_request_only(), _guardrail(_INSTANCE, "InstancePerChunkGuardrail", "post_call")], + ) as built: + yield built + + @pytest.fixture(scope="module") def chaos_rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: with _rig(session_gateway, tmp_path_factory.mktemp("request-only-guardrail-stream-chaos"), _mixed()) as built: @@ -591,6 +612,13 @@ def test_an_inherited_per_chunk_hook_still_sees_every_chunk(opt_in_rig: Rig) -> _assert_every_chunk_seen(opt_in_rig, marker, _INHERITED, text) +def test_a_per_chunk_hook_set_on_the_instance_still_sees_every_chunk(instance_rig: Rig) -> None: + marker: Final = uuid.uuid4().hex + text: Final = _stream(instance_rig, "chat", marker) + _assert_every_chunk_seen(instance_rig, marker, _INSTANCE, text) + assert _events(instance_rig, marker).count("post_call") <= 1, _events(instance_rig, marker) + + def test_a_model_level_per_chunk_guardrail_sees_every_chunk(opt_in_rig: Rig) -> None: marker: Final = uuid.uuid4().hex text: Final = _stream(opt_in_rig, "chat", marker, model=_GUARDED_MODEL) diff --git a/tests/unit/proxy/utils/proxy_logging/test_callback_capabilities_class.py b/tests/unit/proxy/utils/proxy_logging/test_callback_capabilities_class.py index 4d2bc6aa098..389a53c6bc6 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_callback_capabilities_class.py +++ b/tests/unit/proxy/utils/proxy_logging/test_callback_capabilities_class.py @@ -48,6 +48,16 @@ class _OverridesPreCall(CustomLogger): return None +def _instance_per_chunk_logger() -> CustomLogger: + logger = _PlainLogger() + + async def per_chunk(user_api_key_dict, response): + return None + + logger.async_post_call_streaming_hook = per_chunk # type: ignore[method-assign] + return logger + + @pytest.fixture(autouse=True) def _clear_caps_cache(): ProxyLogging._callback_capabilities_cache.clear() @@ -271,12 +281,24 @@ def test_needs_per_chunk_streaming_hook_truth_table(proxy_logging, monkeypatch, monkeypatch.setattr(litellm, "callbacks", [_PerChunkGuardrail(guardrail_name="g", event_hook="post_call")]) ProxyLogging._callback_capabilities_cache.clear() snapshot["per_chunk_guardrail_true"] = proxy_logging.needs_per_chunk_streaming_hook() + monkeypatch.setattr(litellm, "callbacks", [_instance_per_chunk_logger()]) + ProxyLogging._callback_capabilities_cache.clear() + snapshot["instance_hook_alone_false"] = proxy_logging.needs_per_chunk_streaming_hook() + monkeypatch.setattr( + litellm, + "callbacks", + [CustomGuardrail(guardrail_name="g", event_hook="pre_call"), _instance_per_chunk_logger()], + ) + ProxyLogging._callback_capabilities_cache.clear() + snapshot["instance_hook_beside_guardrail_true"] = proxy_logging.needs_per_chunk_streaming_hook() assert snapshot == { "empty_false": False, "per_chunk_override_true": True, "only_iter_override_false": False, "request_only_guardrail_false": False, "per_chunk_guardrail_true": True, + "instance_hook_alone_false": False, + "instance_hook_beside_guardrail_true": True, } diff --git a/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py index 06cfba33860..49de5fd21c7 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -442,6 +442,34 @@ async def test_async_post_call_streaming_hook_skips_a_request_only_guardrail_bes assert guardrail.asked == [] +@pytest.mark.asyncio +async def test_async_post_call_streaming_hook_runs_an_instance_hook_beside_a_guardrail( + proxy_logging, make_user_api_key_auth, monkeypatch +): + guardrail = _RequestOnlyGuardrail() + rewriter = CustomLogger() + + async def per_chunk(user_api_key_dict, response): + return f"modified-{response}" + + rewriter.async_post_call_streaming_hook = per_chunk # type: ignore[method-assign] + monkeypatch.setattr(litellm, "callbacks", [guardrail, rewriter]) + chunk = ModelResponseStream( + id="rid", + choices=[{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": None}], + created=0, + model="gpt-4o-mini", + ) + out = await proxy_logging.async_post_call_streaming_hook( + data={"messages": [{"role": "user", "content": "hi"}]}, + response=chunk, + user_api_key_dict=make_user_api_key_auth(), + str_so_far="say ", + ) + assert out == "modified-say hi" + assert guardrail.asked == [] + + @pytest.mark.asyncio async def test_async_post_call_streaming_hook_callback_error_raises(proxy_logging, make_user_api_key_auth, monkeypatch): class _Per(CustomLogger):