feat(guardrails): buffer TrendAI streams by default

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,

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

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
         )

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] = {