fix(proxy): keep per-chunk hooks set on the instance running beside a guardrail

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,

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)

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,
     }
 
 

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):