mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
feat(guardrails): expose streaming knobs on generic_guardrail_api (#31730)
* feat(guardrails): expose streaming knobs on generic_guardrail_api Wire streaming_end_of_stream_only and streaming_sampling_rate through optional params, initialize_guardrail, and get_config_model so the generic guardrail API participates in UnifiedLLMGuardrails streaming checks with configurable cadence and end-of-stream-only mode. * fix(guardrails): use builtin type[] in get_config_model return Avoids a new UP006 violation that tripped the ruff strict-rule budget gate on the PR lint job. * fix(guardrails): default optional streaming knobs to None Non-None Pydantic defaults on GenericGuardrailAPIOptionalParams made _get_config_value treat unset nested fields as explicit values, which shadowed top-level litellm_params streaming flags whenever any other optional_params key was present. Real defaults stay in the constructor. * fix(guardrails): address review nits on generic_guardrail_api streaming Validate streaming_sampling_rate >= 1 in the constructor and Pydantic optional_params (ge=1), and add /v1/responses streaming coverage through the unified post-call hook so Responses API usage is exercised alongside chat completions. * fix(guardrails): read nested streaming config from dict optional_params Guardrail API/UI delivers optional_params as a plain dict, so getattr was silently ignoring streaming_sampling_rate and streaming_end_of_stream_only. Handle both dict and model shapes in _get_config_value with regression tests. * fix(guardrails): clear ruff findings in generic_guardrail_api tests/types * style(guardrails): ruff format generic_guardrail_api modules --------- Co-authored-by: Marton Schneider <marton@schneider.co.nl>
This commit is contained in:
parent
59f51b2d72
commit
1815636e1c
4 changed files with 775 additions and 5 deletions
|
|
@ -1,4 +1,4 @@
|
|||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, Any, Optional
|
||||
|
||||
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||
|
||||
|
|
@ -8,9 +8,23 @@ if TYPE_CHECKING:
|
|||
from litellm.types.guardrails import Guardrail, LitellmParams
|
||||
|
||||
|
||||
def _get_config_value(litellm_params: Any, optional_params: Any, attribute_name: str) -> Optional[Any]:
|
||||
if optional_params is not None:
|
||||
value = (
|
||||
optional_params.get(attribute_name)
|
||||
if isinstance(optional_params, dict)
|
||||
else getattr(optional_params, attribute_name, None)
|
||||
)
|
||||
if value is not None:
|
||||
return value
|
||||
return getattr(litellm_params, attribute_name, None)
|
||||
|
||||
|
||||
def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"):
|
||||
import litellm
|
||||
|
||||
optional_params = getattr(litellm_params, "optional_params", None)
|
||||
|
||||
_generic_guardrail_api_callback = GenericGuardrailAPI(
|
||||
api_base=litellm_params.api_base,
|
||||
api_key=litellm_params.api_key,
|
||||
|
|
@ -22,6 +36,8 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
|
|||
guardrail_name=guardrail.get("guardrail_name", ""),
|
||||
event_hook=litellm_params.mode,
|
||||
default_on=litellm_params.default_on,
|
||||
streaming_end_of_stream_only=_get_config_value(litellm_params, optional_params, "streaming_end_of_stream_only"),
|
||||
streaming_sampling_rate=_get_config_value(litellm_params, optional_params, "streaming_sampling_rate"),
|
||||
)
|
||||
|
||||
litellm.logging_callback_manager.add_litellm_callback(_generic_guardrail_api_callback)
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
|
||||
GUARDRAIL_NAME = "generic_guardrail_api"
|
||||
|
||||
|
|
@ -178,6 +179,8 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed",
|
||||
fail_on_error: Optional[bool] = True,
|
||||
extra_headers: Optional[list] = None,
|
||||
streaming_end_of_stream_only: Optional[bool] = None,
|
||||
streaming_sampling_rate: Optional[int] = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
|
||||
|
|
@ -209,6 +212,15 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
|
||||
self.fail_on_error: bool = True if fail_on_error is None else fail_on_error
|
||||
|
||||
# Read by UnifiedLLMGuardrails.async_post_call_streaming_iterator_hook
|
||||
# via getattr(guardrail_to_apply, "streaming_*", default).
|
||||
self.streaming_end_of_stream_only: bool = (
|
||||
False if streaming_end_of_stream_only is None else streaming_end_of_stream_only
|
||||
)
|
||||
if streaming_sampling_rate is not None and streaming_sampling_rate < 1:
|
||||
raise ValueError(f"streaming_sampling_rate must be >= 1 (got {streaming_sampling_rate})")
|
||||
self.streaming_sampling_rate: int = 5 if streaming_sampling_rate is None else streaming_sampling_rate
|
||||
|
||||
# Set supported event hooks
|
||||
if "supported_event_hooks" not in kwargs:
|
||||
kwargs["supported_event_hooks"] = [
|
||||
|
|
@ -470,3 +482,11 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
return self._handle_guardrail_request_error(e, inputs, input_type, logging_obj)
|
||||
except Exception as e:
|
||||
return self._handle_guardrail_request_error(e, inputs, input_type, logging_obj, is_unreachable=False)
|
||||
|
||||
@staticmethod
|
||||
def get_config_model() -> Optional[type["GuardrailConfigModel"]]:
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
GenericGuardrailAPIConfigModel,
|
||||
)
|
||||
|
||||
return GenericGuardrailAPIConfigModel
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
from typing import Any, Dict, List, Literal, Optional, Union
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
from typing_extensions import TYPE_CHECKING, TypedDict
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
|
|
@ -60,6 +60,30 @@ class GenericGuardrailAPIOptionalParams(BaseModel):
|
|||
),
|
||||
)
|
||||
|
||||
streaming_end_of_stream_only: Optional[bool] = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"If False (default when unset), the guardrail runs on sampled chunks during "
|
||||
"the stream at the cadence set by streaming_sampling_rate, and an in-flight "
|
||||
"BLOCKED stops further chunks from streaming. If True, the guardrail runs "
|
||||
"once at end of stream over the assembled response; lower cost and latency, "
|
||||
"but flagged content has already streamed to the client before the terminal "
|
||||
"block. Defaults are applied in GenericGuardrailAPI.__init__ when None so "
|
||||
"unset optional_params does not shadow top-level litellm_params."
|
||||
),
|
||||
)
|
||||
|
||||
streaming_sampling_rate: Optional[int] = Field(
|
||||
default=None,
|
||||
ge=1,
|
||||
description=(
|
||||
"When streaming_end_of_stream_only is False, the guardrail runs every Nth "
|
||||
"streamed chunk. Ignored when streaming_end_of_stream_only is True. "
|
||||
"Must be >= 1 when set. Defaults to 5 in GenericGuardrailAPI.__init__ "
|
||||
"when None so unset optional_params does not shadow top-level litellm_params."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class GenericGuardrailAPIConfigModel(
|
||||
GuardrailConfigModel[GenericGuardrailAPIOptionalParams],
|
||||
|
|
|
|||
|
|
@ -609,7 +609,6 @@ class TestImageSupport:
|
|||
request_data=mock_request_data_input,
|
||||
input_type="request",
|
||||
)
|
||||
result_texts = guardrailed_inputs.get("texts", [])
|
||||
result_images = guardrailed_inputs.get("images", None)
|
||||
|
||||
# Verify API was called with images
|
||||
|
|
@ -943,7 +942,7 @@ class TestMultimodalSupport:
|
|||
guardrail.async_handler, "post", return_value=mock_response
|
||||
) as mock_post:
|
||||
# This should not raise SerializationIterator error
|
||||
result = await guardrail.apply_guardrail(
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={
|
||||
"texts": ["What's in this image?"],
|
||||
"images": ["https://example.com/image.jpg"],
|
||||
|
|
@ -1006,7 +1005,7 @@ class TestMultimodalSupport:
|
|||
with patch.object(
|
||||
guardrail.async_handler, "post", return_value=mock_response
|
||||
) as mock_post:
|
||||
result = await guardrail.apply_guardrail(
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={
|
||||
"texts": ["Hello", "World"],
|
||||
"structured_messages": messages_with_iterable,
|
||||
|
|
@ -1023,6 +1022,717 @@ class TestMultimodalSupport:
|
|||
assert isinstance(json_payload["structured_messages"], list)
|
||||
|
||||
|
||||
def _make_stream_chunk(content: str, finish_reason=None):
|
||||
"""Build a real ModelResponseStream so the handler's isinstance checks pass."""
|
||||
from litellm.types.utils import Delta, ModelResponseStream
|
||||
|
||||
return ModelResponseStream(
|
||||
model="gpt-4",
|
||||
choices=[
|
||||
litellm.StreamingChoices(
|
||||
index=0,
|
||||
delta=Delta(role="assistant", content=content),
|
||||
finish_reason=finish_reason,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _make_assembled_model_response(content: str) -> ModelResponse:
|
||||
return ModelResponse(
|
||||
id="mock-response",
|
||||
model="gpt-4",
|
||||
choices=[
|
||||
litellm.Choices(
|
||||
index=0,
|
||||
message=litellm.Message(role="assistant", content=content),
|
||||
finish_reason="stop",
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _mock_guardrail_post_response(action: str = "NONE", texts=None, blocked_reason=None):
|
||||
mock_response = MagicMock()
|
||||
payload = {"action": action}
|
||||
if texts is not None:
|
||||
payload["texts"] = texts
|
||||
if blocked_reason is not None:
|
||||
payload["blocked_reason"] = blocked_reason
|
||||
mock_response.json.return_value = payload
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
return mock_response
|
||||
|
||||
|
||||
def _make_responses_stream_events(text: str):
|
||||
"""Minimal /v1/responses SSE event sequence ending in response.completed."""
|
||||
return (
|
||||
{"type": "response.created", "response": {"id": "resp_test"}},
|
||||
{
|
||||
"type": "response.output_item.added",
|
||||
"item": {"type": "message", "id": "msg_test"},
|
||||
},
|
||||
{
|
||||
"type": "response.content_part.added",
|
||||
"part": {"type": "output_text", "text": ""},
|
||||
},
|
||||
{"type": "response.output_text.delta", "delta": text},
|
||||
{
|
||||
"type": "response.output_text.done",
|
||||
"text": text,
|
||||
},
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "resp_test",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_test",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": text}],
|
||||
}
|
||||
],
|
||||
"status": "completed",
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class TestGenericGuardrailAPIStreamingConfig:
|
||||
"""Streaming knobs on GenericGuardrailAPI and initialize_guardrail plumbing."""
|
||||
|
||||
def test_streaming_defaults(self):
|
||||
guardrail = GenericGuardrailAPI(
|
||||
api_base="https://api.test.guardrail.com",
|
||||
guardrail_name="test-generic-guardrail",
|
||||
event_hook="post_call",
|
||||
)
|
||||
assert guardrail.streaming_end_of_stream_only is False
|
||||
assert guardrail.streaming_sampling_rate == 5
|
||||
|
||||
def test_streaming_overrides(self):
|
||||
guardrail = GenericGuardrailAPI(
|
||||
api_base="https://api.test.guardrail.com",
|
||||
guardrail_name="test-generic-guardrail",
|
||||
event_hook="post_call",
|
||||
streaming_end_of_stream_only=True,
|
||||
streaming_sampling_rate=2,
|
||||
)
|
||||
assert guardrail.streaming_end_of_stream_only is True
|
||||
assert guardrail.streaming_sampling_rate == 2
|
||||
|
||||
@pytest.mark.parametrize("invalid_rate", [0, -1, -5])
|
||||
def test_streaming_sampling_rate_rejects_non_positive(self, invalid_rate):
|
||||
with pytest.raises(ValueError, match="streaming_sampling_rate must be >= 1"):
|
||||
GenericGuardrailAPI(
|
||||
api_base="https://api.test.guardrail.com",
|
||||
guardrail_name="test-generic-guardrail",
|
||||
event_hook="post_call",
|
||||
streaming_sampling_rate=invalid_rate,
|
||||
)
|
||||
|
||||
def test_optional_params_streaming_sampling_rate_ge_one(self):
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
GenericGuardrailAPIOptionalParams,
|
||||
)
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
GenericGuardrailAPIOptionalParams(streaming_sampling_rate=0)
|
||||
|
||||
def test_get_config_model(self):
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
GenericGuardrailAPIConfigModel,
|
||||
)
|
||||
|
||||
assert GenericGuardrailAPI.get_config_model() is GenericGuardrailAPIConfigModel
|
||||
|
||||
def test_initialize_guardrail_forwards_streaming_flags(self):
|
||||
from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
initialize_guardrail,
|
||||
)
|
||||
from litellm.types.guardrails import LitellmParams
|
||||
|
||||
litellm_params = LitellmParams(
|
||||
guardrail="generic_guardrail_api",
|
||||
mode="post_call",
|
||||
api_base="https://api.test.guardrail.com",
|
||||
default_on=False,
|
||||
)
|
||||
# LitellmParams uses extra="allow" on the base; set streaming knobs dynamically
|
||||
litellm_params.streaming_end_of_stream_only = False # type: ignore[attr-defined]
|
||||
litellm_params.streaming_sampling_rate = 3 # type: ignore[attr-defined]
|
||||
|
||||
guardrail_config = {"guardrail_name": "test-generic-streaming"}
|
||||
|
||||
with patch(
|
||||
"litellm.logging_callback_manager.add_litellm_callback"
|
||||
):
|
||||
guardrail = initialize_guardrail(litellm_params, guardrail_config)
|
||||
|
||||
assert guardrail.streaming_end_of_stream_only is False
|
||||
assert guardrail.streaming_sampling_rate == 3
|
||||
|
||||
def test_initialize_guardrail_optional_params_defaults_do_not_shadow_top_level(
|
||||
self,
|
||||
):
|
||||
"""Top-level streaming knobs win when optional_params only carries siblings."""
|
||||
from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
initialize_guardrail,
|
||||
)
|
||||
from litellm.types.guardrails import LitellmParams
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
GenericGuardrailAPIOptionalParams,
|
||||
)
|
||||
|
||||
litellm_params = LitellmParams(
|
||||
guardrail="generic_guardrail_api",
|
||||
mode="post_call",
|
||||
api_base="https://api.test.guardrail.com",
|
||||
default_on=False,
|
||||
)
|
||||
litellm_params.streaming_end_of_stream_only = True # type: ignore[attr-defined]
|
||||
litellm_params.streaming_sampling_rate = 2 # type: ignore[attr-defined]
|
||||
# Sibling optional_params only; streaming fields stay at Pydantic default None.
|
||||
litellm_params.optional_params = GenericGuardrailAPIOptionalParams( # type: ignore[attr-defined]
|
||||
additional_provider_specific_params={"tenant": "acme"},
|
||||
)
|
||||
|
||||
guardrail_config = {"guardrail_name": "test-generic-streaming-mixed"}
|
||||
|
||||
with patch(
|
||||
"litellm.logging_callback_manager.add_litellm_callback"
|
||||
):
|
||||
guardrail = initialize_guardrail(litellm_params, guardrail_config)
|
||||
|
||||
assert guardrail.streaming_end_of_stream_only is True
|
||||
assert guardrail.streaming_sampling_rate == 2
|
||||
|
||||
def test_initialize_guardrail_explicit_optional_params_streaming_wins(self):
|
||||
from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
initialize_guardrail,
|
||||
)
|
||||
from litellm.types.guardrails import LitellmParams
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
GenericGuardrailAPIOptionalParams,
|
||||
)
|
||||
|
||||
litellm_params = LitellmParams(
|
||||
guardrail="generic_guardrail_api",
|
||||
mode="post_call",
|
||||
api_base="https://api.test.guardrail.com",
|
||||
default_on=False,
|
||||
)
|
||||
litellm_params.streaming_end_of_stream_only = False # type: ignore[attr-defined]
|
||||
litellm_params.streaming_sampling_rate = 9 # type: ignore[attr-defined]
|
||||
litellm_params.optional_params = GenericGuardrailAPIOptionalParams( # type: ignore[attr-defined]
|
||||
streaming_end_of_stream_only=True,
|
||||
streaming_sampling_rate=1,
|
||||
)
|
||||
|
||||
guardrail_config = {"guardrail_name": "test-generic-streaming-nested-wins"}
|
||||
|
||||
with patch(
|
||||
"litellm.logging_callback_manager.add_litellm_callback"
|
||||
):
|
||||
guardrail = initialize_guardrail(litellm_params, guardrail_config)
|
||||
|
||||
assert guardrail.streaming_end_of_stream_only is True
|
||||
assert guardrail.streaming_sampling_rate == 1
|
||||
|
||||
def test_initialize_guardrail_dict_optional_params_streaming_wins(self):
|
||||
"""Guardrail API/UI delivers optional_params as a plain dict, not a model."""
|
||||
from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
initialize_guardrail,
|
||||
)
|
||||
from litellm.types.guardrails import LitellmParams
|
||||
|
||||
litellm_params = LitellmParams(
|
||||
guardrail="generic_guardrail_api",
|
||||
mode="post_call",
|
||||
api_base="https://api.test.guardrail.com",
|
||||
default_on=False,
|
||||
)
|
||||
litellm_params.streaming_end_of_stream_only = False # type: ignore[attr-defined]
|
||||
litellm_params.streaming_sampling_rate = 9 # type: ignore[attr-defined]
|
||||
# Plain dict mirrors how configs arrive from the guardrail API/UI.
|
||||
litellm_params.optional_params = { # type: ignore[attr-defined]
|
||||
"streaming_end_of_stream_only": True,
|
||||
"streaming_sampling_rate": 1,
|
||||
}
|
||||
|
||||
guardrail_config = {"guardrail_name": "test-generic-streaming-dict-optional"}
|
||||
|
||||
with patch(
|
||||
"litellm.logging_callback_manager.add_litellm_callback"
|
||||
):
|
||||
guardrail = initialize_guardrail(litellm_params, guardrail_config)
|
||||
|
||||
assert guardrail.streaming_end_of_stream_only is True
|
||||
assert guardrail.streaming_sampling_rate == 1
|
||||
|
||||
def test_initialize_guardrail_dict_optional_params_sibling_only_falls_through(
|
||||
self,
|
||||
):
|
||||
"""Dict optional_params without streaming keys must not shadow top-level knobs."""
|
||||
from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
initialize_guardrail,
|
||||
)
|
||||
from litellm.types.guardrails import LitellmParams
|
||||
|
||||
litellm_params = LitellmParams(
|
||||
guardrail="generic_guardrail_api",
|
||||
mode="post_call",
|
||||
api_base="https://api.test.guardrail.com",
|
||||
default_on=False,
|
||||
)
|
||||
litellm_params.streaming_end_of_stream_only = True # type: ignore[attr-defined]
|
||||
litellm_params.streaming_sampling_rate = 2 # type: ignore[attr-defined]
|
||||
litellm_params.optional_params = { # type: ignore[attr-defined]
|
||||
"additional_provider_specific_params": {"tenant": "acme"},
|
||||
}
|
||||
|
||||
guardrail_config = {"guardrail_name": "test-generic-streaming-dict-sibling"}
|
||||
|
||||
with patch(
|
||||
"litellm.logging_callback_manager.add_litellm_callback"
|
||||
):
|
||||
guardrail = initialize_guardrail(litellm_params, guardrail_config)
|
||||
|
||||
assert guardrail.streaming_end_of_stream_only is True
|
||||
assert guardrail.streaming_sampling_rate == 2
|
||||
|
||||
|
||||
class TestGenericGuardrailAPIStreamingViaUnified:
|
||||
"""Streaming output checks routed through UnifiedLLMGuardrails."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_safe_content_yields_all_chunks(self):
|
||||
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
|
||||
UnifiedLLMGuardrails,
|
||||
)
|
||||
|
||||
guardrail = GenericGuardrailAPI(
|
||||
api_base="https://api.test.guardrail.com",
|
||||
guardrail_name="test-generic-guardrail",
|
||||
event_hook="post_call",
|
||||
)
|
||||
unified_guardrail = UnifiedLLMGuardrails()
|
||||
|
||||
async def mock_stream():
|
||||
chunks_data = ["Hello", " ", "world", "!", " Goodbye"]
|
||||
for i, content in enumerate(chunks_data):
|
||||
yield _make_stream_chunk(
|
||||
content,
|
||||
finish_reason="stop" if i == len(chunks_data) - 1 else None,
|
||||
)
|
||||
|
||||
mock_post = AsyncMock(
|
||||
return_value=_mock_guardrail_post_response(
|
||||
action="NONE", texts=["Hello world! Goodbye"]
|
||||
)
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(guardrail.async_handler, "post", mock_post),
|
||||
patch(
|
||||
"litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder",
|
||||
return_value=_make_assembled_model_response("Hello world! Goodbye"),
|
||||
),
|
||||
):
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test", request_route="/chat/completions"
|
||||
)
|
||||
request_data = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"guardrail_to_apply": guardrail,
|
||||
"metadata": {"guardrails": ["test-generic-guardrail"]},
|
||||
}
|
||||
|
||||
chunks_received = 0
|
||||
async for _ in unified_guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=mock_stream(),
|
||||
request_data=request_data,
|
||||
):
|
||||
chunks_received += 1
|
||||
|
||||
assert chunks_received == 5
|
||||
assert mock_post.await_count >= 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_blocked_content_raises(self):
|
||||
from litellm.exceptions import GuardrailRaisedException
|
||||
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
|
||||
UnifiedLLMGuardrails,
|
||||
)
|
||||
|
||||
guardrail = GenericGuardrailAPI(
|
||||
api_base="https://api.test.guardrail.com",
|
||||
guardrail_name="test-generic-guardrail",
|
||||
event_hook="post_call",
|
||||
streaming_sampling_rate=1,
|
||||
)
|
||||
unified_guardrail = UnifiedLLMGuardrails()
|
||||
|
||||
async def mock_stream():
|
||||
chunks_data = ["Hello", " ishaan", " here"]
|
||||
for i, content in enumerate(chunks_data):
|
||||
yield _make_stream_chunk(
|
||||
content,
|
||||
finish_reason="stop" if i == len(chunks_data) - 1 else None,
|
||||
)
|
||||
|
||||
mock_post = AsyncMock(
|
||||
return_value=_mock_guardrail_post_response(
|
||||
action="BLOCKED", blocked_reason="Ishaan is not allowed"
|
||||
)
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(guardrail.async_handler, "post", mock_post),
|
||||
patch(
|
||||
"litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder",
|
||||
return_value=_make_assembled_model_response("Hello ishaan here"),
|
||||
),
|
||||
):
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test", request_route="/chat/completions"
|
||||
)
|
||||
request_data = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"guardrail_to_apply": guardrail,
|
||||
"metadata": {"guardrails": ["test-generic-guardrail"]},
|
||||
}
|
||||
|
||||
with pytest.raises(GuardrailRaisedException) as exc_info:
|
||||
async for _ in unified_guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=mock_stream(),
|
||||
request_data=request_data,
|
||||
):
|
||||
pass
|
||||
|
||||
assert "Ishaan is not allowed" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_default_uses_sampled_cadence(self):
|
||||
"""Default samples every 5th chunk + final pass: 10 chunks → calls at 5, 10, and final = 3."""
|
||||
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
|
||||
UnifiedLLMGuardrails,
|
||||
)
|
||||
|
||||
guardrail = GenericGuardrailAPI(
|
||||
api_base="https://api.test.guardrail.com",
|
||||
guardrail_name="test-generic-guardrail",
|
||||
event_hook="post_call",
|
||||
)
|
||||
unified_guardrail = UnifiedLLMGuardrails()
|
||||
|
||||
async def mock_stream():
|
||||
chunks_data = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J"]
|
||||
for i, content in enumerate(chunks_data):
|
||||
yield _make_stream_chunk(
|
||||
content,
|
||||
finish_reason="stop" if i == len(chunks_data) - 1 else None,
|
||||
)
|
||||
|
||||
mock_post = AsyncMock(
|
||||
return_value=_mock_guardrail_post_response(
|
||||
action="NONE", texts=["ABCDEFGHIJ"]
|
||||
)
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(guardrail.async_handler, "post", mock_post),
|
||||
patch(
|
||||
"litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder",
|
||||
return_value=_make_assembled_model_response("ABCDEFGHIJ"),
|
||||
),
|
||||
):
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test", request_route="/chat/completions"
|
||||
)
|
||||
request_data = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"guardrail_to_apply": guardrail,
|
||||
"metadata": {"guardrails": ["test-generic-guardrail"]},
|
||||
}
|
||||
|
||||
async for _ in unified_guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=mock_stream(),
|
||||
request_data=request_data,
|
||||
):
|
||||
pass
|
||||
|
||||
assert mock_post.await_count == 3, (
|
||||
f"Expected 3 guardrail calls (2 sampled at chunks 5 / 10 + 1 final), "
|
||||
f"got {mock_post.await_count}"
|
||||
)
|
||||
for call in mock_post.await_args_list:
|
||||
assert call.kwargs["json"]["input_type"] == "response"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_end_of_stream_only_calls_guardrail_once(self):
|
||||
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
|
||||
UnifiedLLMGuardrails,
|
||||
)
|
||||
|
||||
guardrail = GenericGuardrailAPI(
|
||||
api_base="https://api.test.guardrail.com",
|
||||
guardrail_name="test-generic-guardrail",
|
||||
event_hook="post_call",
|
||||
streaming_end_of_stream_only=True,
|
||||
)
|
||||
unified_guardrail = UnifiedLLMGuardrails()
|
||||
|
||||
async def mock_stream():
|
||||
chunks_data = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J"]
|
||||
for i, content in enumerate(chunks_data):
|
||||
yield _make_stream_chunk(
|
||||
content,
|
||||
finish_reason="stop" if i == len(chunks_data) - 1 else None,
|
||||
)
|
||||
|
||||
mock_post = AsyncMock(
|
||||
return_value=_mock_guardrail_post_response(
|
||||
action="NONE", texts=["ABCDEFGHIJ"]
|
||||
)
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(guardrail.async_handler, "post", mock_post),
|
||||
patch(
|
||||
"litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder",
|
||||
return_value=_make_assembled_model_response("ABCDEFGHIJ"),
|
||||
),
|
||||
):
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test", request_route="/chat/completions"
|
||||
)
|
||||
request_data = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"guardrail_to_apply": guardrail,
|
||||
"metadata": {"guardrails": ["test-generic-guardrail"]},
|
||||
}
|
||||
|
||||
async for _ in unified_guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=mock_stream(),
|
||||
request_data=request_data,
|
||||
):
|
||||
pass
|
||||
|
||||
assert mock_post.await_count == 1, (
|
||||
f"Expected exactly one guardrail call at end of stream, "
|
||||
f"got {mock_post.await_count}"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_sampling_rate_override(self):
|
||||
"""sampling_rate=2 on 6 chunks → in-stream at 2,4,6 plus final = 4 calls."""
|
||||
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
|
||||
UnifiedLLMGuardrails,
|
||||
)
|
||||
|
||||
guardrail = GenericGuardrailAPI(
|
||||
api_base="https://api.test.guardrail.com",
|
||||
guardrail_name="test-generic-guardrail",
|
||||
event_hook="post_call",
|
||||
streaming_end_of_stream_only=False,
|
||||
streaming_sampling_rate=2,
|
||||
)
|
||||
unified_guardrail = UnifiedLLMGuardrails()
|
||||
|
||||
async def mock_stream():
|
||||
chunks_data = ["A", "B", "C", "D", "E", "F"]
|
||||
for i, content in enumerate(chunks_data):
|
||||
yield _make_stream_chunk(
|
||||
content,
|
||||
finish_reason="stop" if i == len(chunks_data) - 1 else None,
|
||||
)
|
||||
|
||||
mock_post = AsyncMock(
|
||||
return_value=_mock_guardrail_post_response(action="NONE", texts=["ABCDEF"])
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(guardrail.async_handler, "post", mock_post),
|
||||
patch(
|
||||
"litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder",
|
||||
return_value=_make_assembled_model_response("ABCDEF"),
|
||||
),
|
||||
):
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test", request_route="/chat/completions"
|
||||
)
|
||||
request_data = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"guardrail_to_apply": guardrail,
|
||||
"metadata": {"guardrails": ["test-generic-guardrail"]},
|
||||
}
|
||||
|
||||
async for _ in unified_guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=mock_stream(),
|
||||
request_data=request_data,
|
||||
):
|
||||
pass
|
||||
|
||||
assert mock_post.await_count == 4, (
|
||||
f"Expected 4 guardrail calls (3 sampled + 1 final aggregate), "
|
||||
f"got {mock_post.await_count}"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_fail_open_on_unreachable_continues_stream(self):
|
||||
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
|
||||
UnifiedLLMGuardrails,
|
||||
)
|
||||
|
||||
guardrail = GenericGuardrailAPI(
|
||||
api_base="https://api.test.guardrail.com",
|
||||
guardrail_name="test-generic-guardrail",
|
||||
event_hook="post_call",
|
||||
unreachable_fallback="fail_open",
|
||||
streaming_end_of_stream_only=True,
|
||||
)
|
||||
unified_guardrail = UnifiedLLMGuardrails()
|
||||
|
||||
async def mock_stream():
|
||||
for i, content in enumerate(["A", "B", "C"]):
|
||||
yield _make_stream_chunk(
|
||||
content, finish_reason="stop" if i == 2 else None
|
||||
)
|
||||
|
||||
mock_post = AsyncMock(side_effect=httpx.ConnectError("connection refused"))
|
||||
|
||||
with (
|
||||
patch.object(guardrail.async_handler, "post", mock_post),
|
||||
patch(
|
||||
"litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder",
|
||||
return_value=_make_assembled_model_response("ABC"),
|
||||
),
|
||||
):
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test", request_route="/chat/completions"
|
||||
)
|
||||
request_data = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"guardrail_to_apply": guardrail,
|
||||
"metadata": {"guardrails": ["test-generic-guardrail"]},
|
||||
}
|
||||
|
||||
chunks_received = 0
|
||||
async for _ in unified_guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=mock_stream(),
|
||||
request_data=request_data,
|
||||
):
|
||||
chunks_received += 1
|
||||
|
||||
assert chunks_received == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_responses_api_streaming_end_of_stream_only_calls_guardrail_once(self):
|
||||
"""/v1/responses path through unified hook; end-of-stream-only = one call."""
|
||||
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
|
||||
UnifiedLLMGuardrails,
|
||||
)
|
||||
|
||||
guardrail = GenericGuardrailAPI(
|
||||
api_base="https://api.test.guardrail.com",
|
||||
guardrail_name="test-generic-guardrail",
|
||||
event_hook="post_call",
|
||||
streaming_end_of_stream_only=True,
|
||||
)
|
||||
unified_guardrail = UnifiedLLMGuardrails()
|
||||
|
||||
async def mock_responses_stream():
|
||||
for event in _make_responses_stream_events("Hello world"):
|
||||
yield event
|
||||
|
||||
mock_post = AsyncMock(
|
||||
return_value=_mock_guardrail_post_response(
|
||||
action="NONE", texts=["Hello world"]
|
||||
)
|
||||
)
|
||||
|
||||
with patch.object(guardrail.async_handler, "post", mock_post):
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test", request_route="/v1/responses"
|
||||
)
|
||||
request_data = {
|
||||
"input": "hi",
|
||||
"guardrail_to_apply": guardrail,
|
||||
"metadata": {"guardrails": ["test-generic-guardrail"]},
|
||||
}
|
||||
|
||||
events_received = 0
|
||||
async for _ in unified_guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=mock_responses_stream(),
|
||||
request_data=request_data,
|
||||
):
|
||||
events_received += 1
|
||||
|
||||
assert events_received == 6
|
||||
assert mock_post.await_count == 1, (
|
||||
f"Expected exactly one guardrail call at end of /v1/responses stream, "
|
||||
f"got {mock_post.await_count}"
|
||||
)
|
||||
assert mock_post.await_args.kwargs["json"]["input_type"] == "response"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_responses_api_streaming_blocked_raises(self):
|
||||
"""Mid-stream BLOCKED on /v1/responses surfaces GuardrailRaisedException."""
|
||||
from litellm.exceptions import GuardrailRaisedException
|
||||
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
|
||||
UnifiedLLMGuardrails,
|
||||
)
|
||||
|
||||
guardrail = GenericGuardrailAPI(
|
||||
api_base="https://api.test.guardrail.com",
|
||||
guardrail_name="test-generic-guardrail",
|
||||
event_hook="post_call",
|
||||
streaming_sampling_rate=1,
|
||||
)
|
||||
unified_guardrail = UnifiedLLMGuardrails()
|
||||
|
||||
async def mock_responses_stream():
|
||||
for event in _make_responses_stream_events("blocked content"):
|
||||
yield event
|
||||
|
||||
mock_post = AsyncMock(
|
||||
return_value=_mock_guardrail_post_response(
|
||||
action="BLOCKED", blocked_reason="Responses content not allowed"
|
||||
)
|
||||
)
|
||||
|
||||
with patch.object(guardrail.async_handler, "post", mock_post):
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test", request_route="/v1/responses"
|
||||
)
|
||||
request_data = {
|
||||
"input": "hi",
|
||||
"guardrail_to_apply": guardrail,
|
||||
"metadata": {"guardrails": ["test-generic-guardrail"]},
|
||||
}
|
||||
|
||||
with pytest.raises(GuardrailRaisedException) as exc_info:
|
||||
async for _ in unified_guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=mock_responses_stream(),
|
||||
request_data=request_data,
|
||||
):
|
||||
pass
|
||||
|
||||
assert "Responses content not allowed" in str(exc_info.value)
|
||||
|
||||
class TestToolSupport:
|
||||
"""Test tool handling in guardrail requests"""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue