From 882ec0f731bfea895b99d006c2acf5306188f026 Mon Sep 17 00:00:00 2001 From: Oliver Fei Date: Tue, 27 Oct 2026 16:21:52 -0400 Subject: [PATCH] feat(guardrails): buffer TrendAI streams by default --- .../guardrail_hooks/trendai/__init__.py | 8 +- .../guardrail_hooks/trendai/_models.py | 1 + .../guardrail_hooks/trendai/trendai.py | 2 + .../guardrail_hooks/trendai/test_trendai.py | 111 +++++++++++++++++- 4 files changed, 120 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/trendai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/trendai/__init__.py index 7976f4fc734..53cb68df8f2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/trendai/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/trendai/__init__.py @@ -1,5 +1,7 @@ from typing import TYPE_CHECKING, Final +import httpx + import litellm from litellm.types.guardrails import GuardrailEventHooks, Mode @@ -20,7 +22,9 @@ def _normalize_event_hook( return mode -def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> TrendAIGuardrail: +def initialize_guardrail( + litellm_params: "LitellmParams", guardrail: "Guardrail", *, async_handler: httpx.AsyncClient | None = None +) -> TrendAIGuardrail: settings: Final = TrendAISettings.model_validate(litellm_params.model_dump(mode="python", exclude_none=True)) guardrail_name: Final = guardrail["guardrail_name"] callback: Final = TrendAIGuardrail( @@ -31,6 +35,8 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" timeout=settings.timeout, stream_overlap_size=settings.stream_overlap_size, response_content_chunk_size_bytes=settings.response_content_chunk_size_bytes, + streaming_buffer_until_moderated=settings.streaming_buffer_until_moderated, + async_handler=async_handler, guardrail_name=guardrail_name, event_hook=_normalize_event_hook(litellm_params.mode), default_on=litellm_params.default_on is True, diff --git a/litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py b/litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py index 8172113ae7a..b1d09cb0377 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py +++ b/litellm/proxy/guardrails/guardrail_hooks/trendai/_models.py @@ -14,6 +14,7 @@ class TrendAISettings(BaseModel, frozen=True): timeout: float = 5.0 stream_overlap_size: int = 256 response_content_chunk_size_bytes: int = 49_500 + streaming_buffer_until_moderated: bool = True class TrendAIChatMessage(BaseModel, frozen=True): diff --git a/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py b/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py index 8dc82f37536..f21f18f2711 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py @@ -87,6 +87,7 @@ class TrendAIGuardrail(CustomGuardrail): GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None ) = None, default_on: bool = False, + streaming_buffer_until_moderated: bool = True, ) -> None: resolved_api_key: Final = api_key or os.environ.get("TMV1_API_KEY") if not resolved_api_key: @@ -116,6 +117,7 @@ class TrendAIGuardrail(CustomGuardrail): self.timeout = timeout self.stream_overlap_size: int = stream_overlap_size self.response_content_chunk_size_bytes: int = response_content_chunk_size_bytes + self.streaming_buffer_until_moderated: bool = streaming_buffer_until_moderated self.async_handler: _AsyncHTTPClient = async_handler or get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback ) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/trendai/test_trendai.py b/tests/unit/proxy/guardrails/guardrail_hooks/trendai/test_trendai.py index c886a0aff43..44fb52ee546 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/trendai/test_trendai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/trendai/test_trendai.py @@ -1,5 +1,7 @@ +import asyncio import json -from collections.abc import Callable, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Mapping, Sequence +from contextlib import aclosing from copy import deepcopy from functools import reduce from typing import Final, Literal @@ -26,14 +28,19 @@ from litellm.proxy.guardrails.guardrail_hooks.trendai.trendai import ( OPENAI_CHAT_COMPLETION_RESPONSE_V1, PLUGIN_VERSION, ) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler from litellm.types.guardrails import GuardrailEventHooks, LitellmParams from litellm.types.utils import ( CallTypes, Choices, + Delta, GenericGuardrailAPIInputs, Message, ModelResponse, + ModelResponseStream, + StreamingChoices, ) @@ -818,6 +825,108 @@ def test_utf8_windows_respect_byte_limits_and_character_overlap( assert all(content[window.start : window.start + len(window.text)] == window.text for window in windows) +async def _streaming_response(fragments: Sequence[str]) -> AsyncIterator[ModelResponseStream]: + for text in fragments: + yield ModelResponseStream( + id="chatcmpl-trendai-stream", + created=0, + model="test-model", + choices=[StreamingChoices(index=0, delta=Delta(content=text))], + ) + yield ModelResponseStream( + id="chatcmpl-trendai-stream", + created=0, + model="test-model", + choices=[StreamingChoices(index=0, delta=Delta(content=None), finish_reason="stop")], + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("config", "buffered"), + [ + ({}, True), + ({"streaming_buffer_until_moderated": None}, True), + ({"streaming_buffer_until_moderated": True}, True), + ({"streaming_buffer_until_moderated": False}, False), + ], + ids=["omitted", "null", "enabled", "disabled"], +) +async def test_initialized_stream_waits_for_moderation_unless_buffering_is_disabled( + config: Mapping[str, bool | None], buffered: bool +) -> None: + moderated: Final = asyncio.Event() + + def respond(request: httpx.Request) -> httpx.Response: + moderated.set() + return httpx.Response(200, json={"action": "allow"}) + + params: Final = LitellmParams.model_validate( + { + "guardrail": "trendai", + "mode": "post_call", + "api_key": "test-key", + "api_base": "https://guard.example.com", + **config, + } + ) + fragments: Final = ("first ", "second ", "last") + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + callback: Final = initialize_guardrail(params, {"guardrail_name": "trendai-buffering"}, async_handler=client) + try: + async with aclosing( + UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"), + response=_streaming_response(fragments), + request_data={"metadata": {"guardrails": ["trendai-buffering"]}}, + guardrail_to_apply=callback, + ) + ) as stream: + first: Final = ModelResponseStream.model_validate(await anext(stream)) + assert moderated.is_set() is buffered + remaining: Final = tuple([ModelResponseStream.model_validate(item) async for item in stream]) + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(callback) + + chunks: Final = (first, *remaining) + contents: Final = (chunk.choices[0].delta.content or "" for chunk in chunks) + assert "".join(contents) == "".join(fragments) + assert chunks[-1].choices[0].finish_reason == "stop" + assert moderated.is_set() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("verdict", ["block", "provider_failure"]) +async def test_default_buffering_withholds_all_chunks_when_a_later_window_fails( + verdict: Literal["block", "provider_failure"], +) -> None: + respond, scanned = _engine( + block_on="zzzzz" if verdict == "block" else None, + fail_on="zzzzz" if verdict == "provider_failure" else None, + ) + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + callback: Final = _guardrail( + async_handler=client, + event_hook=GuardrailEventHooks.post_call, + response_content_chunk_size_bytes=5, + stream_overlap_size=0, + ) + async with aclosing( + UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"), + response=_streaming_response(("aaaaa", "bbbbb", "zzzzz", "ccccc")), + request_data={"metadata": {"guardrails": ["trendai"]}}, + guardrail_to_apply=callback, + ) + ) as stream: + with pytest.raises(GuardrailRaisedException) as raised: + await anext(stream) + + assert raised.value.status_code == (400 if verdict == "block" else 503) + assert raised.value.blocked_content is (verdict == "block") + assert scanned == ["aaaaa", "bbbbb", "zzzzz"] + + def _logged_call(user_text: str, assistant_text: str) -> tuple[dict[str, object], ModelResponse]: response = ModelResponse(choices=[Choices(message=Message(role="assistant", content=assistant_text))]) kwargs: dict[str, object] = {