fix(proxy): run per-chunk hooks set on the instance beside a class-level hook

litellm/proxy/utils.py

@@ -3256,12 +3256,17 @@ 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.
+        # A hook set on the instance runs per chunk only when a guardrail or a
+        # class-level override is present, the cases where the per-chunk loop
+        # used to reach it.
+        reaches_instance_hooks: Final = has_guardrail or any(
+            _overrides_streaming_chunk_hook(callback) for callback in resolved_callbacks
+        )
         streaming_chunk_overrides: Final = tuple(
             callback
             for callback in resolved_callbacks
-            if _overrides_streaming_chunk_hook(callback) or (has_guardrail and _assigns_streaming_chunk_hook(callback))
+            if _overrides_streaming_chunk_hook(callback)
+            or (reaches_instance_hooks 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

@@ -36,6 +36,8 @@ _OPT_IN: Final = "opt-in-chunk-check"
 _INHERITED: Final = "inherited-chunk-check"
 _RAISING: Final = "raising-chunk-check"
 _INSTANCE: Final = "instance-chunk-check"
+_CLASS_LOGGER: Final = "class-chunk-logger"
+_INSTANCE_LOGGER: Final = "instance-chunk-logger"
 _CHUNK_ERROR: Final = "Chunk check failed"
 _MODULE: Final = "request_only_guardrail"
 _LOG_ENV: Final = "REQUEST_ONLY_GUARDRAIL_LOG"
@@ -55,6 +57,7 @@ import re
 from fastapi import HTTPException
 
 from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.integrations.custom_logger import CustomLogger
 
 MARKER = re.compile("marker-([0-9a-f]{{32}})")
 
@@ -105,6 +108,26 @@ class RaisingPerChunkGuardrail(CustomGuardrail):
         if "word3 " in response:
             raise HTTPException(status_code=400, detail={{"error": "{_CHUNK_ERROR}"}})
         return None
+
+
+class ClassChunkLogger(CustomLogger):
+    async def async_post_call_streaming_hook(self, user_api_key_dict, response):
+        _log(_marker_in(response), "chunk:{_CLASS_LOGGER}")
+        return None
+
+
+class InstanceChunkLogger(CustomLogger):
+    def __init__(self):
+        super().__init__()
+        self.async_post_call_streaming_hook = self._see_chunk
+
+    async def _see_chunk(self, user_api_key_dict, response):
+        _log(_marker_in(response), "chunk:{_INSTANCE_LOGGER}")
+        return None
+
+
+class_chunk_logger = ClassChunkLogger()
+instance_chunk_logger = InstanceChunkLogger()
 '''
 
 
@@ -360,6 +383,17 @@ def instance_rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFact
         yield built
 
 
+@pytest.fixture(scope="module")
+def logger_rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
+    with _rig(
+        session_gateway,
+        tmp_path_factory.mktemp("per-chunk-loggers-without-guardrails"),
+        [],
+        {"callbacks": [f"{_MODULE}.class_chunk_logger", f"{_MODULE}.instance_chunk_logger"]},
+    ) 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:
@@ -619,6 +653,13 @@ def test_a_per_chunk_hook_set_on_the_instance_still_sees_every_chunk(instance_ri
     assert _events(instance_rig, marker).count("post_call") <= 1, _events(instance_rig, marker)
 
 
+def test_a_per_chunk_hook_set_on_the_instance_runs_beside_a_class_hook_without_guardrails(logger_rig: Rig) -> None:
+    marker: Final = uuid.uuid4().hex
+    text: Final = _stream(logger_rig, "chat", marker)
+    _assert_every_chunk_seen(logger_rig, marker, _CLASS_LOGGER, text)
+    _assert_every_chunk_seen(logger_rig, marker, _INSTANCE_LOGGER, text)
+
+
 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

@@ -9,13 +9,14 @@ Covers ``_callback_capabilities`` (the cached deriver),
 
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, Final
 
 import pytest
 
 import litellm
 from litellm.integrations.custom_guardrail import CustomGuardrail
 from litellm.integrations.custom_logger import CustomLogger
+from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.utils import ProxyLogging, _CallbackCapabilities
 
 
@@ -49,12 +50,12 @@ class _OverridesPreCall(CustomLogger):
 
 
 def _instance_per_chunk_logger() -> CustomLogger:
-    logger = _PlainLogger()
+    logger: Final = _PlainLogger()
 
-    async def per_chunk(user_api_key_dict, response):
+    async def per_chunk(user_api_key_dict: UserAPIKeyAuth, response: str) -> None:
         return None
 
-    logger.async_post_call_streaming_hook = per_chunk  # type: ignore[method-assign]
+    vars(logger)["async_post_call_streaming_hook"] = per_chunk
     return logger
 
 

tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py

@@ -442,34 +442,54 @@ 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()
+def _instance_rewriter() -> CustomLogger:
+    rewriter: Final = CustomLogger()
 
-    async def per_chunk(user_api_key_dict, response):
+    async def per_chunk(user_api_key_dict: UserAPIKeyAuth, response: str) -> str:
         return f"modified-{response}"
 
-    rewriter.async_post_call_streaming_hook = per_chunk  # type: ignore[method-assign]
-    monkeypatch.setattr(litellm, "callbacks", [guardrail, rewriter])
-    chunk = ModelResponseStream(
+    vars(rewriter)["async_post_call_streaming_hook"] = per_chunk
+    return rewriter
+
+
+class _ClassChunkLogger(CustomLogger):
+    async def async_post_call_streaming_hook(self, user_api_key_dict: UserAPIKeyAuth, response: str) -> None:
+        return None
+
+
+async def _say_hi_through_chunk_hooks(proxy_logging: ProxyLogging, user_api_key_dict: UserAPIKeyAuth) -> object:
+    chunk: Final = 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(
+    return 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(),
+        user_api_key_dict=user_api_key_dict,
         str_so_far="say ",
     )
-    assert out == "modified-say hi"
+
+
+@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: Final = _RequestOnlyGuardrail()
+    monkeypatch.setattr(litellm, "callbacks", [guardrail, _instance_rewriter()])
+    assert await _say_hi_through_chunk_hooks(proxy_logging, make_user_api_key_auth()) == "modified-say hi"
     assert guardrail.asked == []
 
 
+@pytest.mark.asyncio
+async def test_async_post_call_streaming_hook_runs_an_instance_hook_beside_a_class_hook(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    monkeypatch.setattr(litellm, "callbacks", [_ClassChunkLogger(), _instance_rewriter()])
+    assert await _say_hi_through_chunk_hooks(proxy_logging, make_user_api_key_auth()) == "modified-say hi"
+
+
 @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):