From c05907f867b3048244a31f2f099ade810d861b38 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 27 Oct 2026 01:06:46 -0700 Subject: [PATCH] fix(proxy): run per-chunk hooks set on the instance beside a class-level hook --- litellm/proxy/utils.py | 11 +++-- ...request_only_guardrail_stream_fast_path.py | 41 +++++++++++++++++ .../test_callback_capabilities_class.py | 9 ++-- .../proxy_logging/test_streaming_hooks.py | 46 +++++++++++++------ 4 files changed, 87 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 53a48a2731c..9a7ef76fede 100644 --- a/litellm/proxy/utils.py +++ b/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, 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 4f728d44aac..7df3b5f2dc2 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 @@ -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) 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 389a53c6bc6..24ca87be156 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 @@ -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 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 49de5fd21c7..1b48c62f1be 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/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):