From c10ab571041096df5568c84217260d49ab2c4317 Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:55:55 -0700 Subject: [PATCH 1/6] fix(proxy): keep the streaming fast path when guardrails only check the request The per-chunk streaming gates also opened for any attached guardrail, so a guardrail running only before the call sent every streamed chunk through the per-chunk hook and was asked about the response once per chunk. Only callbacks that override the per-chunk hook ever act there, and that is already detected on its own When another callback does override the per-chunk hook, the per-chunk loop now visits only the overriding callbacks, so a request-only guardrail is not asked about each chunk in that config either --- litellm/proxy/common_request_processing.py | 14 +- litellm/proxy/proxy_server.py | 8 +- litellm/proxy/utils.py | 87 ++++---- ...request_only_guardrail_stream_fast_path.py | 189 ++++++++++++++++++ .../proxy/test_common_request_processing.py | 28 +++ .../test_callback_capabilities_class.py | 14 ++ .../proxy_logging/test_streaming_hooks.py | 60 +++++- 7 files changed, 341 insertions(+), 59 deletions(-) create mode 100644 tests/integration/streaming/test_request_only_guardrail_stream_fast_path.py diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 59db23820fc..bd6be2f4a79 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4391,15 +4391,15 @@ class ProxyBaseLLMRequestProcessing: """ verbose_proxy_logger.debug("inside generator") # Resolve per-stream (not per-chunk) whether the heavy per-chunk path - # is needed. When no callback overrides ``async_post_call_streaming_hook``, - # no CustomGuardrail is active, and cost injection is disabled, the - # per-chunk hook returns the chunk unchanged, ``str_so_far`` is never - # consumed, and cost injection is a no-op -- so the per-chunk coroutine - # await, response-string materialization, and cost-injection call are - # pure overhead on the streaming hot path (the default config). + # is needed. When no callback overrides ``async_post_call_streaming_hook`` + # and cost injection is disabled, the per-chunk hook returns the chunk + # unchanged, ``str_so_far`` is never consumed, and cost injection is a + # no-op -- so the per-chunk coroutine await, response-string + # materialization, and cost-injection call are pure overhead on the + # streaming hot path (the default config). caps: Final = ProxyLogging.callback_capabilities() cost_injection_enabled: Final = bool(getattr(litellm, "include_cost_in_streaming_usage", False)) - fast_path = not caps.has_streaming_chunk_override and not caps.has_guardrail and not cost_injection_enabled + fast_path = not caps.has_streaming_chunk_override and not cost_injection_enabled debug_enabled: Final = verbose_proxy_logger.isEnabledFor(logging.DEBUG) stream_completed = False client_disconnected = False diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 967cae27990..0efc100de1d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10420,10 +10420,10 @@ async def async_data_generator( # ``async_post_call_streaming_iterator_hook`` or has # ``apply_guardrail``; the per-chunk hook (which builds ``str_so_far`` # and calls ``async_post_call_streaming_hook``) is only needed when - # there is an active CustomGuardrail or a class that overrides the - # per-chunk hook. Coalescing them into a single flag forced wasted - # ``get_response_string`` work per chunk on every deployment that - # happened to ship a streaming-iterator override (the default). + # a class overrides the per-chunk hook. Coalescing them into a single + # flag forced wasted ``get_response_string`` work per chunk on every + # deployment that happened to ship a streaming-iterator override (the + # default). needs_per_chunk_hook: Final = proxy_logging_obj.needs_per_chunk_streaming_hook() is_raw_sse_stream: Final = bool(request_data.get("_litellm_raw_sse_stream")) strip_stream_usage: Final = bool(request_data.get("_litellm_strip_stream_usage")) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b5101898bb1..8edc0973826 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1496,6 +1496,10 @@ class _CallbackCapabilities: # every string entry in ``litellm.callbacks``. resolved_callbacks: tuple[object, ...] = field(default_factory=tuple) listed_models_filters: tuple[CustomLogger, ...] = field(default_factory=tuple) + # Resolved callbacks that override ``async_post_call_streaming_hook``, in + # ``litellm.callbacks`` order. The per-chunk loop only visits these, since + # every other callback's per-chunk hook is the base no-op. + streaming_chunk_overrides: tuple[CustomLogger, ...] = field(default_factory=tuple) def _overrides_hook(callback: CustomLogger, hook_name: str) -> bool: @@ -1503,6 +1507,21 @@ def _overrides_hook(callback: CustomLogger, hook_name: str) -> bool: return any(hook_name in klass.__dict__ for klass in leaf_to_base) +def _overrides_streaming_chunk_hook(callback: CustomLogger) -> bool: + # Walk the MRO for ``async_post_call_streaming_hook`` rather than + # using the leaf-class ``__dict__`` check used by the other flags: + # before this PR the hook was unconditionally invoked, so a + # callback that inherits an override from an intermediate parent + # (e.g. a vendor base class providing the override, with the + # registered class adding nothing else) MUST still be detected. + # A leaf-class miss here would silently drop the inherited hook. + base_streaming_hook: Final = CustomLogger.async_post_call_streaming_hook + cls_streaming_hook: Final = getattr(type(callback), "async_post_call_streaming_hook", base_streaming_hook) + return getattr(cls_streaming_hook, "__func__", cls_streaming_hook) is not getattr( + base_streaming_hook, "__func__", base_streaming_hook + ) + + def _overrides_moderation_hook(callback: CustomLogger) -> bool: return _overrides_hook(callback, "async_moderation_hook") @@ -3200,7 +3219,6 @@ class ProxyLogging: has_post_call_response_headers = False has_iterator_override = False - has_streaming_chunk_override = False has_guardrail = False has_pre_call_override = False has_content_enforcer = False @@ -3239,33 +3257,19 @@ class ProxyLogging: iterator_overrides.append((resolved, "override")) elif "apply_guardrail" in cls_attrs and not getattr(resolved, "use_native_lifecycle_hooks", False): iterator_overrides.append((resolved, "apply_guardrail")) - # Walk the MRO for ``async_post_call_streaming_hook`` rather than - # using the leaf-class ``__dict__`` check used by the other flags: - # before this PR the hook was unconditionally invoked, so a - # callback that inherits an override from an intermediate parent - # (e.g. a vendor base class providing the override, with the - # registered class adding nothing else) MUST still be detected. - # A leaf-class miss here would silently drop the inherited hook. - base_streaming_hook = CustomLogger.async_post_call_streaming_hook - cls_streaming_hook = getattr( - cls, - "async_post_call_streaming_hook", - base_streaming_hook, - ) - if getattr(cls_streaming_hook, "__func__", cls_streaming_hook) is not getattr( - base_streaming_hook, "__func__", base_streaming_hook - ): - has_streaming_chunk_override = True if "async_pre_call_hook" in cls_attrs: has_pre_call_override = True if resolved.enforces_request_content is True: has_content_enforcer = True + streaming_chunk_overrides: Final = tuple( + callback for callback in resolved_callbacks if _overrides_streaming_chunk_hook(callback) + ) caps: Final = _CallbackCapabilities( has_post_call_response_headers=has_post_call_response_headers, has_iterator_override=has_iterator_override or any(kind == "apply_guardrail" for _, kind in iterator_overrides), - has_streaming_chunk_override=has_streaming_chunk_override, + has_streaming_chunk_override=bool(streaming_chunk_overrides), has_guardrail=has_guardrail, has_pre_call_override=has_pre_call_override, has_content_enforcer=has_content_enforcer, @@ -3275,6 +3279,7 @@ class ProxyLogging: listed_models_filters=tuple( callback for callback in resolved_callbacks if _overrides_hook(callback, "async_filter_listed_models") ), + streaming_chunk_overrides=streaming_chunk_overrides, ) # Limit cache to handle test churn without leaking; production # callback lists are stable so this rarely grows past 1 entry. @@ -3323,8 +3328,7 @@ class ProxyLogging: ``_apply_streaming_chunk_hooks`` for every emitted chunk. Instance method for the same reason as :py:meth:`needs_iterator_wrap`. """ - caps: Final = ProxyLogging.callback_capabilities() - return caps.has_streaming_chunk_override or caps.has_guardrail + return ProxyLogging.callback_capabilities().has_streaming_chunk_override @staticmethod def has_during_call_guardrails() -> bool: @@ -4253,12 +4257,11 @@ class ProxyLogging: """ # Per-chunk fast path: skip the response-string materialization and # callback scan when no configured callback overrides - # ``async_post_call_streaming_hook`` AND no CustomGuardrail is - # active. ``get_response_string`` walks every choice/delta on the - # chunk so paying it per chunk for no-op callbacks dominated stream - # CPU time even after the iterator-chain fix. + # ``async_post_call_streaming_hook``. ``get_response_string`` walks + # every choice/delta on the chunk so paying it per chunk for no-op + # callbacks dominated stream CPU time even after the iterator-chain fix. caps: Final = ProxyLogging.callback_capabilities() - if not caps.has_streaming_chunk_override and not caps.has_guardrail: + if not caps.has_streaming_chunk_override: return response from litellm.proxy.proxy_server import llm_router @@ -4279,9 +4282,8 @@ class ProxyLogging: stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset() ) - for callback in litellm.callbacks: + for callback in caps.streaming_chunk_overrides: try: - _callback: CustomLogger | None = None if isinstance(callback, CustomGuardrail): if callback.guardrail_name in pipeline_gated: continue @@ -4303,26 +4305,17 @@ class ProxyLogging: is not True ): continue - if isinstance(callback, str): - _callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( - cast(_custom_logger_compatible_callbacks_literal, callback) - ) + if str_so_far is not None: + complete_response = str_so_far + response_str else: - _callback = callback - if _callback is not None and isinstance(_callback, CustomLogger): - if str_so_far is not None: - complete_response = str_so_far + response_str - else: - complete_response = response_str - callback_response: ( - ModelResponse | EmbeddingResponse | ImageResponse | ModelResponseStream | None - ) - callback_response = await _callback.async_post_call_streaming_hook( - user_api_key_dict=user_api_key_dict, - response=complete_response, - ) - if callback_response is not None: - response = callback_response + complete_response = response_str + callback_response: ModelResponse | EmbeddingResponse | ImageResponse | ModelResponseStream | None + callback_response = await callback.async_post_call_streaming_hook( + user_api_key_dict=user_api_key_dict, + response=complete_response, + ) + if callback_response is not None: + response = callback_response except Exception as e: raise e return response 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 new file mode 100644 index 00000000000..54ad304fed0 --- /dev/null +++ b/tests/integration/streaming/test_request_only_guardrail_stream_fast_path.py @@ -0,0 +1,189 @@ +import json +import re +import uuid +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, list_value, object_value, string_value +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +pytestmark: Final = [ + pytest.mark.timeout(2 * graceful_stop_seconds() + 120), + pytest.mark.xdist_group("request-only-guardrail-stream-fast-path"), +] + +_MODEL: Final = "request-only-guardrail-stream" +_GUARDRAIL: Final = "request-only-check" +_CHUNK_GUARDRAIL: Final = "per-chunk-check" +_MODULE: Final = "request_only_guardrail" +_LOG_ENV: Final = "REQUEST_ONLY_GUARDRAIL_LOG" +_PIECES: Final = tuple(f"word{index} " for index in range(24)) +_MARKER: Final = re.compile(rb"marker-([0-9a-f]{32})") +_MODULE_SOURCE: Final = f'''import json +import os +import re + +from litellm.integrations.custom_guardrail import CustomGuardrail + + +class RequestOnlyGuardrail(CustomGuardrail): + def should_run_guardrail(self, data, event_type): + found = re.search("marker-([0-9a-f]{{32}})", str(data.get("messages") or data.get("input"))) + if found is not None: + with open(os.environ["{_LOG_ENV}"], "a") as log: + log.write(json.dumps({{"marker": found.group(1), "event": str(getattr(event_type, "value", event_type))}}) + "\\n") + return super().should_run_guardrail(data=data, event_type=event_type) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + return data + + +class PerChunkGuardrail(CustomGuardrail): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + with open(os.environ["{_LOG_ENV}"], "a") as log: + log.write(json.dumps({{"marker": None, "event": "chunk"}}) + "\\n") + return None +''' + + +def _chat_chunk(identity: str, content: str | None, finish: str | None = None) -> bytes: + body: Final = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": {} if content is None else {"content": content}, "finish_reason": finish}], + } + return b"data: " + json.dumps(body).encode() + b"\n\n" + + +def _respond(request: Request) -> Reply: + found: Final = _MARKER.search(request.body) + assert found is not None, request.body + identity: Final = found.group(1).decode() + assert request.target.endswith("/chat/completions"), request.target + return Reply( + content_type="text/event-stream", + chunks=( + *(_chat_chunk(f"chatcmpl-{identity}", piece) for piece in _PIECES), + _chat_chunk(f"chatcmpl-{identity}", None, finish="stop"), + b"data: [DONE]\n\n", + ), + ) + + +def _guardrail(name: str, guardrail_class: str, mode: str) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": {"guardrail": f"{_MODULE}.{guardrail_class}", "mode": mode, "default_on": True}, + } + + +def _config(upstream_url: str, guardrails: list[JsonValue]) -> dict[str, JsonValue]: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + return { + **base, + "guardrails": guardrails, + "model_list": [ + { + "model_name": _MODEL, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{upstream_url}/v1", + "api_key": "synthetic-stream-key", + }, + } + ], + } + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: OwnedProxy + log: Path + + +@contextmanager +def _rig(gateway: Gateway, directory: Path, guardrails: list[JsonValue]) -> Iterator[Rig]: + with wire_server(_respond) as upstream: + (directory / f"{_MODULE}.py").write_text(_MODULE_SOURCE) + config_path: Final = directory / "request-only-guardrail-stream-fast-path.yaml" + config_path.write_text(yaml.safe_dump(_config(upstream.url, guardrails))) + log: Final = directory / "guardrail-asks.jsonl" + log.touch() + with owned_proxy_process(gateway, directory, {_LOG_ENV: str(log)}, config=config_path, workers=2) as proxy: + yield Rig(proxy, log) + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with _rig( + session_gateway, + tmp_path_factory.mktemp("request-only-guardrail-stream-fast-path"), + [_guardrail(_GUARDRAIL, "RequestOnlyGuardrail", "pre_call")], + ) as built: + yield built + + +@pytest.fixture(scope="module") +def mixed_rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with _rig( + session_gateway, + tmp_path_factory.mktemp("request-only-guardrail-beside-per-chunk-guardrail"), + [ + _guardrail(_GUARDRAIL, "RequestOnlyGuardrail", "pre_call"), + _guardrail(_CHUNK_GUARDRAIL, "PerChunkGuardrail", "post_call"), + ], + ) as built: + yield built + + +def _events(rig: Rig, marker: str | None) -> tuple[str, ...]: + entries: Final = (JSON_OBJECT.validate_json(line) for line in rig.log.read_text().splitlines() if line) + return tuple(string_value(entry["event"]) for entry in entries if entry["marker"] == marker) + + +def _choices(text: str) -> Iterator[JsonValue]: + for line in text.splitlines(): + if line.startswith("data: ") and line != "data: [DONE]": + yield from list_value(JSON_OBJECT.validate_json(line.removeprefix("data: ")).get("choices") or []) + + +def _chat_text(text: str) -> str: + return "".join( + string_value(object_value(object_value(choice)["delta"]).get("content") or "") for choice in _choices(text) + ) + + +def _stream(rig: Rig, marker: str) -> str: + body: Final = {"model": _MODEL, "stream": True, "messages": [{"role": "user", "content": f"marker-{marker}"}]} + response: Final = rig.proxy.gateway.request("POST", "/v1/chat/completions", body) + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("text/event-stream"), response.headers + assert _chat_text(response.text) == "".join(_PIECES), response.text + return response.text + + +def test_a_request_only_guardrail_is_not_asked_about_each_streamed_chunk(rig: Rig) -> None: + marker: Final = uuid.uuid4().hex + _stream(rig, marker) + asks: Final = _events(rig, marker) + assert "pre_call" in asks, asks + assert asks.count("post_call") == 0, f"asked about the response {asks.count('post_call')} times: {asks}" + + +def test_a_per_chunk_guardrail_still_sees_every_chunk_without_asking_the_request_only_one(mixed_rig: Rig) -> None: + marker: Final = uuid.uuid4().hex + text: Final = _stream(mixed_rig, marker) + asks: Final = _events(mixed_rig, marker) + assert "pre_call" in asks, asks + assert asks.count("post_call") <= 1, f"asked about the response {asks.count('post_call')} times: {asks}" + delivered: Final = sum(1 for line in text.splitlines() if line.startswith("data: ") and line != "data: [DONE]") + assert _events(mixed_rig, None).count("chunk") >= delivered, (delivered, _events(mixed_rig, None)) diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index d742956bb59..d8c061fc6b7 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -32,6 +32,7 @@ from litellm.constants import ( RETURN_RAW_MODEL_NAME_METADATA_KEY, STREAM_SSE_KEEPALIVE_PING_BYTES, ) +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.opentelemetry import UserAPIKeyAuth from litellm.proxy.common_request_processing import ( @@ -4971,6 +4972,33 @@ class TestAsyncStreamingDataGeneratorFastPath: assert out == chunks # bytes pass through return_sse_chunk untouched hook_spy.assert_not_awaited() + @pytest.mark.asyncio + async def test_fast_path_skips_per_chunk_hook_with_a_request_only_guardrail(self, monkeypatch): + monkeypatch.setattr(litellm, "callbacks", [CustomGuardrail(guardrail_name="g", event_hook="pre_call")]) + ProxyLogging._callback_capabilities_cache.clear() + + proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock()) + hook_spy = AsyncMock(side_effect=lambda **kw: kw["response"]) + monkeypatch.setattr(proxy_logging_obj, "async_post_call_streaming_hook", hook_spy) + + chunks = [b"event: a\ndata: {}\n\n", b"event: b\ndata: {}\n\n"] + out = [ + c + async for c in ProxyBaseLLMRequestProcessing.async_streaming_data_generator( + response=self._aiter(chunks), + user_api_key_dict=ProxyUserAPIKeyAuth(), + request_data={"model": "claude-x"}, + proxy_logging_obj=proxy_logging_obj, + serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk, + serialize_error=lambda e: "data: error\n\n", + ) + ] + + assert out == chunks + hook_spy.assert_not_awaited() + + ProxyLogging._callback_capabilities_cache.clear() + @pytest.mark.asyncio async def test_slow_path_runs_per_chunk_hook(self, monkeypatch): """A callback that overrides async_post_call_streaming_hook forces the 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 75a91177f00..903a6aac762 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 @@ -14,6 +14,7 @@ from typing import Any import pytest import litellm +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.proxy.utils import ProxyLogging, _CallbackCapabilities @@ -37,6 +38,11 @@ class _OverridesPerChunk(CustomLogger): return None +class _PerChunkGuardrail(CustomGuardrail): + async def async_post_call_streaming_hook(self, *args, **kwargs): + return None + + class _OverridesPreCall(CustomLogger): async def async_pre_call_hook(self, *args, **kwargs): # type: ignore[override] return None @@ -251,10 +257,18 @@ def test_needs_per_chunk_streaming_hook_truth_table(proxy_logging, monkeypatch, monkeypatch.setattr(litellm, "callbacks", [_OverridesIterator()]) ProxyLogging._callback_capabilities_cache.clear() snapshot["only_iter_override_false"] = proxy_logging.needs_per_chunk_streaming_hook() + monkeypatch.setattr(litellm, "callbacks", [CustomGuardrail(guardrail_name="g", event_hook="pre_call")]) + ProxyLogging._callback_capabilities_cache.clear() + snapshot["request_only_guardrail_false"] = proxy_logging.needs_per_chunk_streaming_hook() + 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() 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, } 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 50f50478ad3..06cfba33860 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -29,7 +29,7 @@ from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.utils import Usage +from litellm.types.utils import ModelResponseStream, Usage @pytest.fixture(autouse=True) @@ -384,6 +384,64 @@ async def test_async_post_call_streaming_hook_invokes_per_chunk_callback(proxy_l assert out.startswith("modified-") +class _RequestOnlyGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="request-only", event_hook=GuardrailEventHooks.pre_call, default_on=True) + self.asked: List[GuardrailEventHooks] = [] + + def should_run_guardrail(self, data, event_type: GuardrailEventHooks) -> bool: + self.asked.append(event_type) + return super().should_run_guardrail(data=data, event_type=event_type) + + +@pytest.mark.asyncio +async def test_async_post_call_streaming_hook_does_not_ask_a_request_only_guardrail( + proxy_logging, make_user_api_key_auth, monkeypatch +): + guardrail = _RequestOnlyGuardrail() + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + 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(), + ) + assert out is chunk + assert guardrail.asked == [] + + +class _PerChunkRewriter(CustomLogger): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): # type: ignore[override] + return f"modified-{response}" + + +@pytest.mark.asyncio +async def test_async_post_call_streaming_hook_skips_a_request_only_guardrail_beside_a_per_chunk_callback( + proxy_logging, make_user_api_key_auth, monkeypatch +): + guardrail = _RequestOnlyGuardrail() + monkeypatch.setattr(litellm, "callbacks", [guardrail, _PerChunkRewriter()]) + 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): From 66af9d4cf660b4ed4e0d0e2b335b15819e4bf807 Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 18:41:56 -0700 Subject: [PATCH 2/6] fix(proxy): align the per-chunk override probe with the streaming gate --- litellm/proxy/utils.py | 13 +------------ .../test_callback_capabilities_class.py | 8 ++++++++ 2 files changed, 9 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8edc0973826..8f41db94b98 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1496,9 +1496,6 @@ class _CallbackCapabilities: # every string entry in ``litellm.callbacks``. resolved_callbacks: tuple[object, ...] = field(default_factory=tuple) listed_models_filters: tuple[CustomLogger, ...] = field(default_factory=tuple) - # Resolved callbacks that override ``async_post_call_streaming_hook``, in - # ``litellm.callbacks`` order. The per-chunk loop only visits these, since - # every other callback's per-chunk hook is the base no-op. streaming_chunk_overrides: tuple[CustomLogger, ...] = field(default_factory=tuple) @@ -1508,13 +1505,6 @@ def _overrides_hook(callback: CustomLogger, hook_name: str) -> bool: def _overrides_streaming_chunk_hook(callback: CustomLogger) -> bool: - # Walk the MRO for ``async_post_call_streaming_hook`` rather than - # using the leaf-class ``__dict__`` check used by the other flags: - # before this PR the hook was unconditionally invoked, so a - # callback that inherits an override from an intermediate parent - # (e.g. a vendor base class providing the override, with the - # registered class adding nothing else) MUST still be detected. - # A leaf-class miss here would silently drop the inherited hook. base_streaming_hook: Final = CustomLogger.async_post_call_streaming_hook cls_streaming_hook: Final = getattr(type(callback), "async_post_call_streaming_hook", base_streaming_hook) return getattr(cls_streaming_hook, "__func__", cls_streaming_hook) is not getattr( @@ -3313,8 +3303,7 @@ class ProxyLogging: def has_streaming_chunk_hook_overrides() -> bool: """True iff any callback overrides ``async_post_call_streaming_hook`` (the per-chunk hook, distinct from the iterator wrapper).""" - caps: Final = ProxyLogging.callback_capabilities() - return caps.has_streaming_chunk_override or caps.has_guardrail + return ProxyLogging.callback_capabilities().has_streaming_chunk_override def needs_iterator_wrap(self) -> bool: """Whether ``async_data_generator`` needs to wrap the upstream stream 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 903a6aac762..4d2bc6aa098 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 @@ -201,10 +201,18 @@ def test_has_streaming_chunk_hook_overrides_truth_table(monkeypatch, mock_callba monkeypatch.setattr(litellm, "callbacks", [_OverridesIterator()]) ProxyLogging._callback_capabilities_cache.clear() snapshot["only_iterator_false"] = ProxyLogging.has_streaming_chunk_hook_overrides() + monkeypatch.setattr(litellm, "callbacks", [CustomGuardrail(guardrail_name="g", event_hook="pre_call")]) + ProxyLogging._callback_capabilities_cache.clear() + snapshot["request_only_guardrail_false"] = ProxyLogging.has_streaming_chunk_hook_overrides() + monkeypatch.setattr(litellm, "callbacks", [_PerChunkGuardrail(guardrail_name="g", event_hook="post_call")]) + ProxyLogging._callback_capabilities_cache.clear() + snapshot["per_chunk_guardrail_true"] = ProxyLogging.has_streaming_chunk_hook_overrides() assert snapshot == { "empty_false": False, "per_chunk_override_true": True, "only_iterator_false": False, + "request_only_guardrail_false": False, + "per_chunk_guardrail_true": True, } From c939dbfed384e54be8277fa5bc9e7d8cc1e8a622 Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:52:25 -0700 Subject: [PATCH 3/6] test(proxy): cover every streaming endpoint, client and per-chunk guardrail mode for the fast path --- ...request_only_guardrail_stream_fast_path.py | 604 ++++++++++++++++-- 1 file changed, 541 insertions(+), 63 deletions(-) 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 54ad304fed0..01a0f728ff3 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 @@ -1,15 +1,20 @@ +import asyncio import json import re import uuid -from collections.abc import Iterator +from collections.abc import Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path -from typing import Final +from typing import Final, Literal, TypeAlias +import anthropic +import httpx +import openai import pytest import yaml -from integration._support.client import JSON_OBJECT, Gateway, list_value, object_value, string_value +from integration._support.client import JSON_OBJECT, Gateway, eventually, list_value, object_value, string_value from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue @@ -19,26 +24,55 @@ pytestmark: Final = [ pytest.mark.xdist_group("request-only-guardrail-stream-fast-path"), ] +Endpoint: TypeAlias = Literal["chat", "messages", "responses"] +Client: TypeAlias = Literal["sync", "async"] + _MODEL: Final = "request-only-guardrail-stream" +_MESSAGES_MODEL: Final = "request-only-guardrail-messages" +_GUARDED_MODEL: Final = "request-only-guardrail-guarded" _GUARDRAIL: Final = "request-only-check" _CHUNK_GUARDRAIL: Final = "per-chunk-check" +_OPT_IN: Final = "opt-in-chunk-check" +_INHERITED: Final = "inherited-chunk-check" +_RAISING: Final = "raising-chunk-check" +_CHUNK_ERROR: Final = "Chunk check failed" _MODULE: Final = "request_only_guardrail" _LOG_ENV: Final = "REQUEST_ONLY_GUARDRAIL_LOG" +_ENDPOINTS: Final[tuple[Endpoint, ...]] = ("chat", "messages", "responses") +_PATHS: Final[Mapping[Endpoint, str]] = { + "chat": "/v1/chat/completions", + "messages": "/v1/messages", + "responses": "/v1/responses", +} _PIECES: Final = tuple(f"word{index} " for index in range(24)) _MARKER: Final = re.compile(rb"marker-([0-9a-f]{32})") +_DROP: Final = re.compile(rb"drop-after-([0-9]+)") _MODULE_SOURCE: Final = f'''import json import os import re +from fastapi import HTTPException + from litellm.integrations.custom_guardrail import CustomGuardrail +MARKER = re.compile("marker-([0-9a-f]{{32}})") + + +def _marker_in(value): + found = MARKER.search(str(value)) + return None if found is None else found.group(1) + + +def _log(marker, event): + with open(os.environ["{_LOG_ENV}"], "a") as log: + log.write(json.dumps({{"marker": marker, "event": event}}) + "\\n") + class RequestOnlyGuardrail(CustomGuardrail): def should_run_guardrail(self, data, event_type): - found = re.search("marker-([0-9a-f]{{32}})", str(data.get("messages") or data.get("input"))) - if found is not None: - with open(os.environ["{_LOG_ENV}"], "a") as log: - log.write(json.dumps({{"marker": found.group(1), "event": str(getattr(event_type, "value", event_type))}}) + "\\n") + marker = _marker_in(data.get("messages") or data.get("input")) + if marker is not None: + _log(marker, str(getattr(event_type, "value", event_type))) return super().should_run_guardrail(data=data, event_type=event_type) async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): @@ -47,59 +81,197 @@ class RequestOnlyGuardrail(CustomGuardrail): class PerChunkGuardrail(CustomGuardrail): async def async_post_call_streaming_hook(self, user_api_key_dict, response): - with open(os.environ["{_LOG_ENV}"], "a") as log: - log.write(json.dumps({{"marker": None, "event": "chunk"}}) + "\\n") + _log(_marker_in(response), "chunk:" + self.guardrail_name) + return None + + +class InheritedPerChunkGuardrail(PerChunkGuardrail): + pass + + +class RaisingPerChunkGuardrail(CustomGuardrail): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + if "word3 " in response: + raise HTTPException(status_code=400, detail={{"error": "{_CHUNK_ERROR}"}}) return None ''' -def _chat_chunk(identity: str, content: str | None, finish: str | None = None) -> bytes: - body: Final = { - "id": identity, - "object": "chat.completion.chunk", +def _text(marker: str) -> str: + return f"marker-{marker} " + "".join(_PIECES) + + +def _pieces(marker: str) -> tuple[str, ...]: + return (f"marker-{marker} ", *_PIECES) + + +def _sse(name: str | None, payload: Mapping[str, JsonValue]) -> bytes: + prefix: Final = f"event: {name}\n" if name is not None else "" + return f"{prefix}data: {json.dumps(payload)}\n\n".encode() + + +def _chat_chunk(identity: str, delta: Mapping[str, JsonValue], finish: str | None = None) -> bytes: + return _sse( + None, + { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish}], + }, + ) + + +def _chat_stream(marker: str) -> tuple[bytes, ...]: + identity: Final = f"chatcmpl-{marker}" + first, *rest = _pieces(marker) + return ( + _chat_chunk(identity, {"role": "assistant", "content": first}), + *(_chat_chunk(identity, {"content": piece}) for piece in rest), + _chat_chunk(identity, {}, finish="stop"), + b"data: [DONE]\n\n", + ) + + +def _message(marker: str) -> dict[str, JsonValue]: + return { + "id": f"msg_{marker}", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [{"type": "text", "text": _text(marker)}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 2, "output_tokens": 25}, + } + + +def _messages_stream(marker: str) -> tuple[bytes, ...]: + started: Final = {**_message(marker), "content": [], "stop_reason": None} + events: Final[tuple[dict[str, JsonValue], ...]] = ( + {"type": "message_start", "message": {**started, "usage": {"input_tokens": 2, "output_tokens": 0}}}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + *( + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": piece}} + for piece in _pieces(marker) + ), + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 25}}, + {"type": "message_stop"}, + ) + return tuple(_sse(string_value(event["type"]), event) for event in events) + + +def _response_output(marker: str) -> dict[str, JsonValue]: + return { + "type": "message", + "id": f"msg_{marker}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": _text(marker), "annotations": []}], + } + + +def _response(marker: str) -> dict[str, JsonValue]: + return { + "id": f"resp_{marker}", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [_response_output(marker)], + "usage": {"input_tokens": 2, "output_tokens": 25, "total_tokens": 27}, + } + + +def _responses_stream(marker: str) -> tuple[bytes, ...]: + response: Final = _response(marker) + output: Final = _response_output(marker) + part: Final[dict[str, JsonValue]] = {"type": "output_text", "text": _text(marker), "annotations": []} + position: Final[dict[str, JsonValue]] = {"item_id": f"msg_{marker}", "output_index": 0, "content_index": 0} + started: Final = {**response, "status": "in_progress", "output": []} + events: Final[tuple[dict[str, JsonValue], ...]] = ( + {"type": "response.created", "response": started}, + {"type": "response.in_progress", "response": started}, + { + "type": "response.output_item.added", + "item": {**output, "status": "in_progress", "content": []}, + "output_index": 0, + }, + {"type": "response.content_part.added", **position, "part": {**part, "text": ""}}, + *({"type": "response.output_text.delta", **position, "delta": piece} for piece in _pieces(marker)), + {"type": "response.output_text.done", **position, "text": _text(marker)}, + {"type": "response.content_part.done", **position, "part": part}, + {"type": "response.output_item.done", "item": output, "output_index": 0}, + {"type": "response.completed", "response": response}, + ) + return tuple( + _sse(string_value(event["type"]), {**event, "sequence_number": index}) for index, event in enumerate(events) + ) + + +def _chat_completion(marker: str) -> dict[str, JsonValue]: + return { + "id": f"chatcmpl-{marker}", + "object": "chat.completion", "created": 1, "model": "gpt-4o-mini", - "choices": [{"index": 0, "delta": {} if content is None else {"content": content}, "finish_reason": finish}], + "choices": [{"index": 0, "message": {"role": "assistant", "content": _text(marker)}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 2, "completion_tokens": 25, "total_tokens": 27}, } - return b"data: " + json.dumps(body).encode() + b"\n\n" def _respond(request: Request) -> Reply: found: Final = _MARKER.search(request.body) assert found is not None, request.body - identity: Final = found.group(1).decode() + marker: Final = found.group(1).decode() + streamed: Final = JSON_OBJECT.validate_json(request.body).get("stream") is True + if request.target.endswith("/messages"): + if not streamed: + return Reply(body=json.dumps(_message(marker)).encode()) + return Reply(content_type="text/event-stream", chunks=_messages_stream(marker)) + if request.target.endswith("/responses"): + if not streamed: + return Reply(body=json.dumps(_response(marker)).encode()) + return Reply(content_type="text/event-stream", chunks=_responses_stream(marker)) assert request.target.endswith("/chat/completions"), request.target - return Reply( - content_type="text/event-stream", - chunks=( - *(_chat_chunk(f"chatcmpl-{identity}", piece) for piece in _PIECES), - _chat_chunk(f"chatcmpl-{identity}", None, finish="stop"), - b"data: [DONE]\n\n", - ), - ) + if not streamed: + return Reply(body=json.dumps(_chat_completion(marker)).encode()) + drop: Final = _DROP.search(request.body) + if drop is not None: + return Reply( + content_type="text/event-stream", + chunks=_chat_stream(marker), + abort_after=int(drop.group(1)), + pause_between_chunks=0.02, + ) + return Reply(content_type="text/event-stream", chunks=_chat_stream(marker)) -def _guardrail(name: str, guardrail_class: str, mode: str) -> dict[str, JsonValue]: +def _guardrail(name: str, guardrail_class: str, mode: str, *, default_on: bool = True) -> dict[str, JsonValue]: return { "guardrail_name": name, - "litellm_params": {"guardrail": f"{_MODULE}.{guardrail_class}", "mode": mode, "default_on": True}, + "litellm_params": {"guardrail": f"{_MODULE}.{guardrail_class}", "mode": mode, "default_on": default_on}, } -def _config(upstream_url: str, guardrails: list[JsonValue]) -> dict[str, JsonValue]: +def _deployment(name: str, model: str, api_base: str, **extra: JsonValue) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": {"model": model, "api_base": api_base, "api_key": "synthetic-stream-key", **extra}, + } + + +def _config(upstream_url: str, guardrails: list[JsonValue], settings: Mapping[str, JsonValue]) -> dict[str, JsonValue]: base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) return { **base, + "litellm_settings": {**object_value(base["litellm_settings"]), **settings}, "guardrails": guardrails, "model_list": [ - { - "model_name": _MODEL, - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_base": f"{upstream_url}/v1", - "api_key": "synthetic-stream-key", - }, - } + _deployment(_MODEL, "openai/gpt-4o-mini", f"{upstream_url}/v1"), + _deployment(_MESSAGES_MODEL, "anthropic/claude-sonnet-4-5-20250929", upstream_url), + _deployment(_GUARDED_MODEL, "openai/gpt-4o-mini", f"{upstream_url}/v1", guardrails=[_OPT_IN]), ], } @@ -111,23 +283,34 @@ class Rig: @contextmanager -def _rig(gateway: Gateway, directory: Path, guardrails: list[JsonValue]) -> Iterator[Rig]: +def _rig( + gateway: Gateway, + directory: Path, + guardrails: list[JsonValue], + settings: Mapping[str, JsonValue] | None = None, +) -> Iterator[Rig]: with wire_server(_respond) as upstream: (directory / f"{_MODULE}.py").write_text(_MODULE_SOURCE) config_path: Final = directory / "request-only-guardrail-stream-fast-path.yaml" - config_path.write_text(yaml.safe_dump(_config(upstream.url, guardrails))) + config_path.write_text(yaml.safe_dump(_config(upstream.url, guardrails, settings or {}))) log: Final = directory / "guardrail-asks.jsonl" log.touch() with owned_proxy_process(gateway, directory, {_LOG_ENV: str(log)}, config=config_path, workers=2) as proxy: yield Rig(proxy, log) +def _request_only() -> list[JsonValue]: + return [_guardrail(_GUARDRAIL, "RequestOnlyGuardrail", "pre_call")] + + +def _mixed() -> list[JsonValue]: + return [*_request_only(), _guardrail(_CHUNK_GUARDRAIL, "PerChunkGuardrail", "post_call")] + + @pytest.fixture(scope="module") def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: with _rig( - session_gateway, - tmp_path_factory.mktemp("request-only-guardrail-stream-fast-path"), - [_guardrail(_GUARDRAIL, "RequestOnlyGuardrail", "pre_call")], + session_gateway, tmp_path_factory.mktemp("request-only-guardrail-stream-fast-path"), _request_only() ) as built: yield built @@ -135,55 +318,350 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I @pytest.fixture(scope="module") def mixed_rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: with _rig( - session_gateway, - tmp_path_factory.mktemp("request-only-guardrail-beside-per-chunk-guardrail"), - [ - _guardrail(_GUARDRAIL, "RequestOnlyGuardrail", "pre_call"), - _guardrail(_CHUNK_GUARDRAIL, "PerChunkGuardrail", "post_call"), - ], + session_gateway, tmp_path_factory.mktemp("request-only-guardrail-beside-per-chunk-guardrail"), _mixed() ) as built: yield built -def _events(rig: Rig, marker: str | None) -> tuple[str, ...]: +@pytest.fixture(scope="module") +def opt_in_rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with _rig( + session_gateway, + tmp_path_factory.mktemp("request-only-guardrail-beside-opt-in-per-chunk-guardrails"), + [ + *_request_only(), + _guardrail(_OPT_IN, "PerChunkGuardrail", "post_call", default_on=False), + _guardrail(_INHERITED, "InheritedPerChunkGuardrail", "post_call", default_on=False), + _guardrail(_RAISING, "RaisingPerChunkGuardrail", "post_call", default_on=False), + ], + {"include_cost_in_streaming_usage": True}, + ) 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: + yield built + + +def _events(rig: Rig, marker: str) -> tuple[str, ...]: entries: Final = (JSON_OBJECT.validate_json(line) for line in rig.log.read_text().splitlines() if line) return tuple(string_value(entry["event"]) for entry in entries if entry["marker"] == marker) -def _choices(text: str) -> Iterator[JsonValue]: - for line in text.splitlines(): - if line.startswith("data: ") and line != "data: [DONE]": - yield from list_value(JSON_OBJECT.validate_json(line.removeprefix("data: ")).get("choices") or []) +def _prompt(marker: str, directive: str) -> str: + return f"marker-{marker} {directive}".strip() -def _chat_text(text: str) -> str: - return "".join( - string_value(object_value(object_value(choice)["delta"]).get("content") or "") for choice in _choices(text) +def _body( + endpoint: Endpoint, + marker: str, + *, + stream: bool = True, + model: str = _MODEL, + guardrails: Sequence[str] = (), + directive: str = "", +) -> dict[str, JsonValue]: + requested: Final[dict[str, JsonValue]] = {"guardrails": list(guardrails)} if guardrails else {} + if endpoint == "responses": + return {"model": model, "stream": stream, "input": _prompt(marker, directive), **requested} + messages: Final[list[JsonValue]] = [{"role": "user", "content": _prompt(marker, directive)}] + if endpoint == "messages": + return {"model": _MESSAGES_MODEL, "stream": stream, "max_tokens": 64, "messages": messages, **requested} + return {"model": model, "stream": stream, "messages": messages, "num_retries": 0, **requested} + + +def _payloads(text: str) -> tuple[dict[str, JsonValue], ...]: + return tuple( + JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" ) -def _stream(rig: Rig, marker: str) -> str: - body: Final = {"model": _MODEL, "stream": True, "messages": [{"role": "user", "content": f"marker-{marker}"}]} - response: Final = rig.proxy.gateway.request("POST", "/v1/chat/completions", body) +def _event_text(event: Mapping[str, JsonValue], endpoint: Endpoint) -> str: + if endpoint == "responses": + return string_value(event.get("delta", "")) if event.get("type") == "response.output_text.delta" else "" + if endpoint == "messages": + if event.get("type") != "content_block_delta": + return "" + return string_value(object_value(event["delta"]).get("text", "")) + return "".join( + string_value(object_value(object_value(choice).get("delta", {})).get("content") or "") + for choice in list_value(event.get("choices") or []) + ) + + +def _streamed_text(text: str, endpoint: Endpoint) -> str: + return "".join(_event_text(event, endpoint) for event in _payloads(text)) + + +def _send( + rig: Rig, + endpoint: Endpoint, + marker: str, + *, + stream: bool = True, + model: str = _MODEL, + guardrails: Sequence[str] = (), + directive: str = "", +) -> httpx.Response: + body: Final = _body(endpoint, marker, stream=stream, model=model, guardrails=guardrails, directive=directive) + return rig.proxy.gateway.request("POST", _PATHS[endpoint], body) + + +def _stream(rig: Rig, endpoint: Endpoint, marker: str, *, model: str = _MODEL, guardrails: Sequence[str] = ()) -> str: + response: Final = _send(rig, endpoint, marker, model=model, guardrails=guardrails) assert response.status_code == 200, response.text assert response.headers["content-type"].startswith("text/event-stream"), response.headers - assert _chat_text(response.text) == "".join(_PIECES), response.text + assert _streamed_text(response.text, endpoint) == _text(marker), response.text return response.text -def test_a_request_only_guardrail_is_not_asked_about_each_streamed_chunk(rig: Rig) -> None: +def _delivered_chunks(text: str) -> int: + return sum(1 for payload in _payloads(text) if "choices" in payload) + + +def _assert_every_chunk_seen(rig: Rig, marker: str, guardrail: str, text: str) -> None: + seen: Final = _events(rig, marker).count(f"chunk:{guardrail}") + assert seen >= _delivered_chunks(text) > 0, (seen, _delivered_chunks(text), _events(rig, marker)) + + +@pytest.mark.parametrize("endpoint", _ENDPOINTS) +def test_a_request_only_guardrail_is_not_asked_about_each_streamed_chunk(rig: Rig, endpoint: Endpoint) -> None: marker: Final = uuid.uuid4().hex - _stream(rig, marker) + _stream(rig, endpoint, marker) asks: Final = _events(rig, marker) assert "pre_call" in asks, asks assert asks.count("post_call") == 0, f"asked about the response {asks.count('post_call')} times: {asks}" +@pytest.mark.parametrize("endpoint", _ENDPOINTS) +def test_a_request_only_guardrail_answers_non_streaming_calls_on_every_endpoint(rig: Rig, endpoint: Endpoint) -> None: + marker: Final = uuid.uuid4().hex + response: Final = _send(rig, endpoint, marker, stream=False) + assert response.status_code == 200, response.text + assert _text(marker) in response.text, response.text + asks: Final = _events(rig, marker) + assert "pre_call" in asks, asks + assert asks.count("post_call") <= 1, asks + + +def _openai_text(client: openai.OpenAI, endpoint: Endpoint, marker: str) -> str: + if endpoint == "responses": + events: Final = client.responses.create(model=_MODEL, input=_prompt(marker, ""), stream=True) + return "".join(event.delta for event in events if event.type == "response.output_text.delta") + chunks: Final = client.chat.completions.create( + model=_MODEL, messages=[{"role": "user", "content": _prompt(marker, "")}], stream=True + ) + return "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) + + +async def _async_openai_text(client: openai.AsyncOpenAI, endpoint: Endpoint, marker: str) -> str: + pieces: Final[list[str]] = [] + if endpoint == "responses": + async for event in await client.responses.create(model=_MODEL, input=_prompt(marker, ""), stream=True): + if event.type == "response.output_text.delta": + pieces.append(event.delta) + return "".join(pieces) + async for chunk in await client.chat.completions.create( + model=_MODEL, messages=[{"role": "user", "content": _prompt(marker, "")}], stream=True + ): + pieces.extend(choice.delta.content or "" for choice in chunk.choices) + return "".join(pieces) + + +def _anthropic_text(client: anthropic.Anthropic, marker: str) -> str: + events: Final = client.messages.create( + model=_MESSAGES_MODEL, max_tokens=64, messages=[{"role": "user", "content": _prompt(marker, "")}], stream=True + ) + return "".join( + event.delta.text for event in events if event.type == "content_block_delta" and event.delta.type == "text_delta" + ) + + +async def _async_anthropic_text(client: anthropic.AsyncAnthropic, marker: str) -> str: + pieces: Final[list[str]] = [] + async for event in await client.messages.create( + model=_MESSAGES_MODEL, max_tokens=64, messages=[{"role": "user", "content": _prompt(marker, "")}], stream=True + ): + if event.type == "content_block_delta" and event.delta.type == "text_delta": + pieces.append(event.delta.text) + return "".join(pieces) + + +def _sync_sdk_text(gateway: Gateway, endpoint: Endpoint, marker: str) -> str: + origin: Final = str(gateway.client.base_url).rstrip("/") + with httpx.Client(timeout=30, trust_env=False) as transport: + if endpoint == "messages": + with anthropic.Anthropic( + api_key=gateway.key, base_url=origin, max_retries=0, http_client=transport + ) as client: + return _anthropic_text(client, marker) + with openai.OpenAI( + api_key=gateway.key, base_url=f"{origin}/v1", max_retries=0, http_client=transport + ) as client: + return _openai_text(client, endpoint, marker) + + +def _async_sdk_text(gateway: Gateway, endpoint: Endpoint, marker: str) -> str: + origin: Final = str(gateway.client.base_url).rstrip("/") + + async def collect() -> str: + async with httpx.AsyncClient(timeout=30, trust_env=False) as transport: + if endpoint == "messages": + async with anthropic.AsyncAnthropic( + api_key=gateway.key, base_url=origin, max_retries=0, http_client=transport + ) as client: + return await _async_anthropic_text(client, marker) + async with openai.AsyncOpenAI( + api_key=gateway.key, base_url=f"{origin}/v1", max_retries=0, http_client=transport + ) as client: + return await _async_openai_text(client, endpoint, marker) + + return asyncio.run(collect()) + + +@pytest.mark.parametrize("client", ["sync", "async"]) +@pytest.mark.parametrize("endpoint", _ENDPOINTS) +def test_sdk_streams_through_a_request_only_guardrail_without_per_chunk_asks( + rig: Rig, endpoint: Endpoint, client: Client +) -> None: + marker: Final = uuid.uuid4().hex + read: Final = _sync_sdk_text if client == "sync" else _async_sdk_text + assert read(rig.proxy.gateway, endpoint, marker) == _text(marker) + asks: Final = _events(rig, marker) + assert "pre_call" in asks, asks + assert asks.count("post_call") == 0, asks + + +def test_a_request_only_guardrail_stream_writes_one_success_spend_row_naming_it(rig: Rig) -> None: + marker: Final = uuid.uuid4().hex + response: Final = _send(rig, "chat", marker) + assert response.status_code == 200, response.text + assert _streamed_text(response.text, "chat") == _text(marker), response.text + call_id: Final = response.headers["x-litellm-call-id"] + rows: Final = eventually( + lambda: rig.proxy.gateway.rows( + 'SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id = %s', (call_id,) + ), + lambda found: len(found) >= 1, + seconds=30, + ) + assert [row["status"] for row in rows] == ["success"], rows + assert _GUARDRAIL in list_value(object_value(rows[0]["metadata"]).get("applied_guardrails") or []), rows + + def test_a_per_chunk_guardrail_still_sees_every_chunk_without_asking_the_request_only_one(mixed_rig: Rig) -> None: marker: Final = uuid.uuid4().hex - text: Final = _stream(mixed_rig, marker) + text: Final = _stream(mixed_rig, "chat", marker) asks: Final = _events(mixed_rig, marker) assert "pre_call" in asks, asks assert asks.count("post_call") <= 1, f"asked about the response {asks.count('post_call')} times: {asks}" - delivered: Final = sum(1 for line in text.splitlines() if line.startswith("data: ") and line != "data: [DONE]") - assert _events(mixed_rig, None).count("chunk") >= delivered, (delivered, _events(mixed_rig, None)) + _assert_every_chunk_seen(mixed_rig, marker, _CHUNK_GUARDRAIL, text) + + +@pytest.mark.parametrize("endpoint", ["messages", "responses"]) +def test_a_per_chunk_guardrail_beside_a_request_only_one_keeps_other_endpoints_whole( + mixed_rig: Rig, endpoint: Endpoint +) -> None: + marker: Final = uuid.uuid4().hex + _stream(mixed_rig, endpoint, marker) + asks: Final = _events(mixed_rig, marker) + assert "pre_call" in asks, asks + assert asks.count("post_call") <= 1, asks + + +def test_an_opt_in_per_chunk_guardrail_sees_every_chunk_only_when_requested(opt_in_rig: Rig) -> None: + requested: Final = uuid.uuid4().hex + text: Final = _stream(opt_in_rig, "chat", requested, guardrails=[_OPT_IN]) + _assert_every_chunk_seen(opt_in_rig, requested, _OPT_IN, text) + assert _events(opt_in_rig, requested).count("post_call") <= 1, _events(opt_in_rig, requested) + + skipped: Final = uuid.uuid4().hex + _stream(opt_in_rig, "chat", skipped) + assert not any(event.startswith("chunk:") for event in _events(opt_in_rig, skipped)), _events(opt_in_rig, skipped) + assert _events(opt_in_rig, skipped).count("post_call") <= 1, _events(opt_in_rig, skipped) + + +def test_an_inherited_per_chunk_hook_still_sees_every_chunk(opt_in_rig: Rig) -> None: + marker: Final = uuid.uuid4().hex + text: Final = _stream(opt_in_rig, "chat", marker, guardrails=[_INHERITED]) + _assert_every_chunk_seen(opt_in_rig, marker, _INHERITED, 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) + _assert_every_chunk_seen(opt_in_rig, marker, _OPT_IN, text) + + +def test_a_raising_per_chunk_guardrail_ends_the_stream_with_its_error(opt_in_rig: Rig) -> None: + marker: Final = uuid.uuid4().hex + response: Final = _send(opt_in_rig, "chat", marker, guardrails=[_RAISING]) + assert response.status_code == 200, response.text + lines: Final = [line for line in response.text.splitlines() if line.startswith("data: ")] + assert lines[-1] == "data: [DONE]" and lines.count("data: [DONE]") == 1, response.text + payloads: Final = _payloads(response.text) + assert "error" in payloads[-1] and all("error" not in payload for payload in payloads[:-1]), response.text + assert object_value(payloads[-1]["error"])["message"] == _CHUNK_ERROR, payloads[-1] + delivered: Final = "".join(_event_text(payload, "chat") for payload in payloads[:-1]) + assert delivered and _text(marker).startswith(delivered) and "word3 " not in delivered, response.text + assert _events(opt_in_rig, marker).count("post_call") <= 1, _events(opt_in_rig, marker) + + +@pytest.mark.parametrize("endpoint", _ENDPOINTS) +def test_cost_injection_keeps_streams_whole_beside_a_request_only_guardrail( + opt_in_rig: Rig, endpoint: Endpoint +) -> None: + marker: Final = uuid.uuid4().hex + _stream(opt_in_rig, endpoint, marker) + assert _events(opt_in_rig, marker).count("post_call") <= 1, _events(opt_in_rig, marker) + + +@dataclass(frozen=True, slots=True) +class Sent: + endpoint: Endpoint + marker: str + dropped: bool + response: httpx.Response + + +def _burst_member(rig: Rig, endpoint: Endpoint, dropped: bool) -> Sent: + marker: Final = uuid.uuid4().hex + directive: Final = "drop-after-8" if dropped else "" + return Sent(endpoint, marker, dropped, _send(rig, endpoint, marker, directive=directive)) + + +def _assert_burst_member(rig: Rig, sent: Sent) -> None: + assert sent.response.status_code == 200, sent.response.text + text: Final = _streamed_text(sent.response.text, sent.endpoint) + if sent.dropped: + assert text and _text(sent.marker).startswith(text) and text != _text(sent.marker), sent.response.text + else: + assert text == _text(sent.marker), sent.response.text + asks: Final = _events(rig, sent.marker) + assert "pre_call" in asks and asks.count("post_call") <= 1, asks + if sent.endpoint == "chat": + _assert_every_chunk_seen(rig, sent.marker, _CHUNK_GUARDRAIL, sent.response.text) + + +def test_a_mixed_burst_with_dropped_upstreams_and_a_killed_worker_keeps_every_guardrail_contract( + chaos_rig: Rig, +) -> None: + plan: Final[tuple[tuple[Endpoint, bool], ...]] = ( + *(("chat", index % 2 == 0) for index in range(8)), + *(("messages", False) for _ in range(8)), + *(("responses", False) for _ in range(8)), + ) + with ThreadPoolExecutor(max_workers=len(plan)) as pool: + burst: Final = tuple(pool.map(lambda member: _burst_member(chaos_rig, *member), plan)) + for sent in burst: + _assert_burst_member(chaos_rig, sent) + + victim: Final = min(chaos_rig.proxy.worker_pids()) + replacement: Final = chaos_rig.proxy.kill_worker(victim) + assert replacement in chaos_rig.proxy.worker_pids(), (replacement, chaos_rig.proxy.worker_pids()) + for _ in range(8): + _assert_burst_member(chaos_rig, _burst_member(chaos_rig, "chat", False)) From 2db6641b87ace926446a9351de74b5136e607a9b Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 26 Oct 2026 18:05:50 -0700 Subject: [PATCH 4/6] fix(proxy): keep per-chunk hooks set on the instance running beside a guardrail --- litellm/proxy/utils.py | 10 ++++++- ...request_only_guardrail_stream_fast_path.py | 28 +++++++++++++++++++ .../test_callback_capabilities_class.py | 22 +++++++++++++++ .../proxy_logging/test_streaming_hooks.py | 28 +++++++++++++++++++ 4 files changed, 87 insertions(+), 1 deletion(-) 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): From c05907f867b3048244a31f2f099ade810d861b38 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 27 Oct 2026 01:06:46 -0700 Subject: [PATCH 5/6] 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): From 01ab5bb7a1b00124df57dd7152546cede6360094 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 27 Oct 2026 17:16:34 -0700 Subject: [PATCH 6/6] test(proxy): type the per-chunk streaming hook test helpers --- .../proxy/test_common_request_processing.py | 12 +++--- .../test_callback_capabilities_class.py | 2 +- .../proxy_logging/test_streaming_hooks.py | 42 +++++++++++-------- 3 files changed, 33 insertions(+), 23 deletions(-) diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index d8c061fc6b7..f7bb4f54f0a 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -4973,16 +4973,18 @@ class TestAsyncStreamingDataGeneratorFastPath: hook_spy.assert_not_awaited() @pytest.mark.asyncio - async def test_fast_path_skips_per_chunk_hook_with_a_request_only_guardrail(self, monkeypatch): + async def test_fast_path_skips_per_chunk_hook_with_a_request_only_guardrail( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: monkeypatch.setattr(litellm, "callbacks", [CustomGuardrail(guardrail_name="g", event_hook="pre_call")]) ProxyLogging._callback_capabilities_cache.clear() - proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock()) - hook_spy = AsyncMock(side_effect=lambda **kw: kw["response"]) + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=MagicMock()) + hook_spy: Final = AsyncMock(side_effect=lambda **kw: kw["response"]) monkeypatch.setattr(proxy_logging_obj, "async_post_call_streaming_hook", hook_spy) - chunks = [b"event: a\ndata: {}\n\n", b"event: b\ndata: {}\n\n"] - out = [ + chunks: Final = [b"event: a\ndata: {}\n\n", b"event: b\ndata: {}\n\n"] + out: Final = [ c async for c in ProxyBaseLLMRequestProcessing.async_streaming_data_generator( response=self._aiter(chunks), 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 24ca87be156..09cdd31598a 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 @@ -40,7 +40,7 @@ class _OverridesPerChunk(CustomLogger): class _PerChunkGuardrail(CustomGuardrail): - async def async_post_call_streaming_hook(self, *args, **kwargs): + async def async_post_call_streaming_hook(self, user_api_key_dict: UserAPIKeyAuth, response: str) -> None: return None 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 1b48c62f1be..b19b31413e0 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -10,7 +10,7 @@ Covers ``_wrap_streaming_iterator_with_enrichment``, from __future__ import annotations import asyncio -from collections.abc import AsyncGenerator, AsyncIterator +from collections.abc import AsyncGenerator, AsyncIterator, Callable from datetime import datetime from typing import Any, Dict, Final, List from unittest.mock import AsyncMock, MagicMock @@ -389,24 +389,26 @@ class _RequestOnlyGuardrail(CustomGuardrail): super().__init__(guardrail_name="request-only", event_hook=GuardrailEventHooks.pre_call, default_on=True) self.asked: List[GuardrailEventHooks] = [] - def should_run_guardrail(self, data, event_type: GuardrailEventHooks) -> bool: + def should_run_guardrail(self, data: dict[str, object], event_type: GuardrailEventHooks) -> bool: self.asked.append(event_type) return super().should_run_guardrail(data=data, event_type=event_type) @pytest.mark.asyncio async def test_async_post_call_streaming_hook_does_not_ask_a_request_only_guardrail( - proxy_logging, make_user_api_key_auth, monkeypatch -): - guardrail = _RequestOnlyGuardrail() + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, +) -> None: + guardrail: Final = _RequestOnlyGuardrail() monkeypatch.setattr(litellm, "callbacks", [guardrail]) - chunk = ModelResponseStream( + 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( + out: Final = 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(), @@ -416,23 +418,25 @@ async def test_async_post_call_streaming_hook_does_not_ask_a_request_only_guardr class _PerChunkRewriter(CustomLogger): - async def async_post_call_streaming_hook(self, user_api_key_dict, response): # type: ignore[override] + async def async_post_call_streaming_hook(self, user_api_key_dict: UserAPIKeyAuth, response: str) -> str: return f"modified-{response}" @pytest.mark.asyncio async def test_async_post_call_streaming_hook_skips_a_request_only_guardrail_beside_a_per_chunk_callback( - proxy_logging, make_user_api_key_auth, monkeypatch -): - guardrail = _RequestOnlyGuardrail() + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, +) -> None: + guardrail: Final = _RequestOnlyGuardrail() monkeypatch.setattr(litellm, "callbacks", [guardrail, _PerChunkRewriter()]) - chunk = ModelResponseStream( + 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( + out: Final = 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(), @@ -474,8 +478,10 @@ async def _say_hi_through_chunk_hooks(proxy_logging: ProxyLogging, user_api_key_ @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 -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, +) -> None: 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" @@ -484,8 +490,10 @@ async def test_async_post_call_streaming_hook_runs_an_instance_hook_beside_a_gua @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 -): + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.setattr(litellm, "callbacks", [_ClassChunkLogger(), _instance_rewriter()]) assert await _say_hi_through_chunk_hooks(proxy_logging, make_user_api_key_auth()) == "modified-say hi"