scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 201

test(guardrails): cover streamed /v1/completions under post_call guardrails

Browse source
This commit is contained in:
mateo-berri 2026-10-20 19:53:59 -07:00
parent cc156a73b8
commit db4f91fb16
3 changed files with 278 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

95
tests/integration/streaming/test_text_completion_stream_guardrail_rebuild.py
View file

@ -27,7 +27,12 @@ _PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
_PROVIDER_KEY: Final = "synthetic-provider-key"
_GUARDRAIL: Final = "audit-text-rebuild-tool-permission"
_TEXT_MODEL: Final = f"audit-text-instruct-{uuid.uuid4().hex}"
_GOVERNED_MODEL: Final = f"audit-text-governed-{uuid.uuid4().hex}"
_KEYWORD_FILTER: Final = "audit-text-keyword-filter"
_POLICY: Final = "audit-text-post-call-governance"
_PIECES: Final = ("Hello", " streamed", " world")
_MASKED_PIECES: Final = ("A pers", "im", "mon tart")
_BLOCKED_PIECES: Final = ("A kum", "quat tart")
_PROMPT_TOKENS: Final = 11
_COMPLETION_TOKENS: Final = 3
_USAGE: Final[dict[str, JsonValue]] = {
@ -97,6 +102,14 @@ def _text_choice(text: str, finish_reason: str | None) -> dict[str, JsonValue]:
return {"text": text, "index": 0, "logprobs": None, "finish_reason": finish_reason}
def _pieces_for(identity: str) -> tuple[str, ...]:
if identity.startswith("cmpl-mask-"):
return _MASKED_PIECES
if identity.startswith("cmpl-block-"):
return _BLOCKED_PIECES
return _PIECES
def _respond(request: Request) -> Reply:
body: Final = JSON_OBJECT.validate_json(request.body)
assert body.get("stream") is True, body
@ -106,7 +119,7 @@ def _respond(request: Request) -> Reply:
return Reply(
content_type="text/event-stream",
chunks=(
*(_text_frame(identity, [_text_choice(piece, None)]) for piece in _PIECES),
*(_text_frame(identity, [_text_choice(piece, None)]) for piece in _pieces_for(identity)),
_text_frame(identity, [_text_choice("", "length")]),
_text_frame(identity, [], usage=_USAGE),
b"data: [DONE]\n\n",
@ -114,20 +127,25 @@ def _respond(request: Request) -> Reply:
)
def _text_deployment(model_name: str, wire: Wire) -> dict[str, JsonValue]:
return {
"model_name": model_name,
"litellm_params": {
"model": "text-completion-openai/gpt-3.5-turbo-instruct",
"api_base": f"{wire.url}/v1",
"api_key": _PROVIDER_KEY,
},
}
def _config(directory: Path, wire: Wire) -> Path:
base: Final = object_value(yaml.safe_load(_PROXY_CONFIG.read_text()))
config: Final = {
**base,
"model_list": [
*list_value(base["model_list"]),
{
"model_name": _TEXT_MODEL,
"litellm_params": {
"model": "text-completion-openai/gpt-3.5-turbo-instruct",
"api_base": f"{wire.url}/v1",
"api_key": _PROVIDER_KEY,
},
},
_text_deployment(_TEXT_MODEL, wire),
_text_deployment(_GOVERNED_MODEL, wire),
],
"guardrails": [
{
@ -139,8 +157,30 @@ def _config(directory: Path, wire: Wire) -> Path:
"default_action": "allow",
"rules": [],
},
}
},
{
"guardrail_name": _KEYWORD_FILTER,
"litellm_params": {
"guardrail": "litellm_content_filter",
"mode": "post_call",
"default_on": False,
"blocked_words": [
{"keyword": "persimmon", "action": "MASK"},
{"keyword": "kumquat", "action": "BLOCK"},
],
},
},
],
"policies": {
_POLICY: {
"guardrails": {"add": [_KEYWORD_FILTER]},
"pipeline": {
"mode": "post_call",
"steps": [{"guardrail": _KEYWORD_FILTER, "on_pass": "allow", "on_fail": "block"}],
},
}
},
"policy_attachments": [{"policy": _POLICY, "models": [_GOVERNED_MODEL]}],
}
path: Final = directory / f"text-rebuild-{uuid.uuid4().hex}.yaml"
path.write_text(yaml.safe_dump(config))
@ -163,14 +203,14 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
yield _Rig(wire, owned)
def _stream(rig: _Rig, identity: str) -> _Stream:
def _stream(rig: _Rig, identity: str, model: str = _TEXT_MODEL) -> _Stream:
gateway: Final = rig.proxy.gateway
with gateway.client.stream(
"POST",
"/v1/completions",
headers={"Authorization": f"Bearer {gateway.key}"},
json={
"model": _TEXT_MODEL,
"model": model,
"prompt": identity,
"max_tokens": 16,
"stream": True,
@ -212,6 +252,37 @@ def test_guarded_text_completion_stream_reaches_the_caller_whole_with_its_usage(
assert _GUARDRAIL in list_value(object_value(row["metadata"])["applied_guardrails"]), row
def test_post_call_pipeline_masks_a_streamed_text_completion(rig: _Rig) -> None:
identity: Final = f"cmpl-mask-{uuid.uuid4().hex}"
stream: Final = _stream(rig, identity, _GOVERNED_MODEL)
assert stream.status == 200, stream.body
assert stream.body.rstrip().endswith(b"data: [DONE]"), stream.body[-400:]
assert stream.text == "A [KEYWORD_REDACTED] tart", stream.frames
assert b"persimmon" not in stream.body, stream.body
assert stream.finish_reasons == ("length",), stream.frames
assert [(usage["prompt_tokens"], usage["completion_tokens"]) for usage in stream.usages] == [
(_PROMPT_TOKENS, _COMPLETION_TOKENS)
], stream.frames
(row,) = eventually(
lambda: rig.proxy.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (identity,)),
lambda rows: len(rows) == 1,
seconds=70,
)
assert _KEYWORD_FILTER in list_value(object_value(row["metadata"])["applied_guardrails"]), row
def test_post_call_pipeline_blocks_a_streamed_text_completion(rig: _Rig) -> None:
identity: Final = f"cmpl-block-{uuid.uuid4().hex}"
stream: Final = _stream(rig, identity, _GOVERNED_MODEL)
assert stream.status == 400, stream.body
assert b"Content blocked: keyword 'kumquat' detected" in stream.body, stream.body
assert b"quat tart" not in stream.body, stream.body
def _together_chunk(choices: list[JsonValue], **extra: JsonValue) -> bytes:
return _frame(
{

147
tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py
View file

@ -6,6 +6,7 @@ from typing import List, Optional, Tuple
from unittest.mock import MagicMock
import pytest
from fastapi import HTTPException
from litellm.integrations.custom_guardrail import CustomGuardrail
@ -13,7 +14,11 @@ from litellm.llms import get_guardrail_translation_mapping
from litellm.llms.openai.completion.guardrail_translation.handler import (
OpenAITextCompletionHandler,
)
from litellm.types.utils import CallTypes, TextChoices, TextCompletionResponse
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails
from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks
from litellm.types.utils import CallTypes, TextChoices, TextCompletionResponse, Usage
class MockGuardrail(CustomGuardrail):
@ -341,3 +346,143 @@ class TestPIIMaskingScenario:
assert "[EMAIL_REDACTED]" in result["prompt"][1]
assert "alice@example.com" not in result["prompt"][0]
assert "bob@test.org" not in result["prompt"][1]
class _RecordingKeywordFilter(ContentFilterGuardrail):
def __init__(self):
super().__init__(
guardrail_name="keyword-filter",
blocked_words=[
BlockedWord(keyword="persimmon", action=ContentFilterAction.MASK),
BlockedWord(keyword="kumquat", action=ContentFilterAction.BLOCK),
],
event_hook=GuardrailEventHooks.post_call,
default_on=True,
)
self.streaming_sampling_rate = 1
self.scanned_texts: tuple[tuple[str, ...], ...] = ()
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
self.scanned_texts = (*self.scanned_texts, tuple(inputs.get("texts") or ()))
return await super().apply_guardrail(inputs, request_data, input_type, logging_obj)
def _text_chunk(index: int, text: str | None, finish_reason: str | None = None) -> TextCompletionResponse:
return TextCompletionResponse(
id="cmpl-stream",
created=1589478378,
model="gpt-3.5-turbo-instruct",
object="text_completion",
choices=[TextChoices(index=index, text=text, finish_reason=finish_reason)],
)
def _usage_chunk() -> TextCompletionResponse:
return TextCompletionResponse(
id="cmpl-stream",
created=1589478378,
model="gpt-3.5-turbo-instruct",
object="text_completion",
choices=[],
usage=Usage(prompt_tokens=3, completion_tokens=5, total_tokens=8),
)
def _delivered(chunks: list[TextCompletionResponse]) -> list[tuple[int, str | None, str | None]]:
return [(choice.index, choice.text, choice.finish_reason) for chunk in chunks for choice in chunk.choices]
async def _drive_flat_stream(guardrail: CustomGuardrail, chunks: list[TextCompletionResponse]) -> list[object]:
async def upstream():
for chunk in chunks:
yield chunk
return [
item
async for item in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/completions"),
response=upstream(),
request_data={"guardrail_to_apply": guardrail, "model": "gpt-3.5-turbo-instruct"},
)
]
class TestStreamingOutputProcessing:
@pytest.mark.asyncio
async def test_ended_stream_delivers_each_choice_masked_text(self):
guardrail = _RecordingKeywordFilter()
chunks = [
_text_chunk(0, "A pers"),
_text_chunk(1, "kiwi "),
_text_chunk(0, "im"),
_text_chunk(1, "pie", "stop"),
_text_chunk(0, "mon tart"),
_text_chunk(0, None, "length"),
_usage_chunk(),
]
await OpenAITextCompletionHandler().process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=guardrail,
request_data={},
deliver_ended_stream_rewrites=True,
)
assert guardrail.scanned_texts == (("A persimmon tart", "kiwi pie"),)
assert _delivered(chunks) == [
(0, "A [KEYWORD_REDACTED] tart", None),
(1, "kiwi ", None),
(0, "", None),
(1, "pie", "stop"),
(0, "", None),
(0, None, "length"),
]
assert chunks[-1].usage.total_tokens == 8
@pytest.mark.asyncio
async def test_ended_stream_scanned_without_delivery_keeps_chunks_and_stores_the_assembled_response(self):
chunks = [_text_chunk(0, "A pers"), _text_chunk(0, "immon tart"), _text_chunk(0, "", "stop"), _usage_chunk()]
request_data: dict[str, object] = {}
await OpenAITextCompletionHandler().process_output_streaming_response(
responses_so_far=chunks, guardrail_to_apply=_RecordingKeywordFilter(), request_data=request_data
)
assert _delivered(chunks) == [(0, "A pers", None), (0, "immon tart", None), (0, "", "stop")]
assembled = request_data["response"]
assert isinstance(assembled, TextCompletionResponse)
assert [(c.text, c.finish_reason) for c in assembled.choices] == [("A [KEYWORD_REDACTED] tart", "stop")]
assert assembled.usage.total_tokens == 8
@pytest.mark.asyncio
async def test_ended_stream_block_raises_the_guardrail_error(self):
chunks = [_text_chunk(0, "A kum"), _text_chunk(0, "quat tart"), _text_chunk(0, "", "stop")]
with pytest.raises(HTTPException) as exc_info:
await OpenAITextCompletionHandler().process_output_streaming_response(
responses_so_far=chunks, guardrail_to_apply=_RecordingKeywordFilter(), deliver_ended_stream_rewrites=True
)
assert exc_info.value.status_code == 400
assert "kumquat" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_flat_stream_is_scanned_once_after_it_ends(self):
guardrail = _RecordingKeywordFilter()
chunks = [_text_chunk(0, "A pers"), _text_chunk(0, "immon tart"), _text_chunk(0, "", "stop"), _usage_chunk()]
delivered = await _drive_flat_stream(guardrail, chunks)
assert guardrail.scanned_texts == (("A persimmon tart",),)
assert len(delivered) == len(chunks)
@pytest.mark.asyncio
async def test_flat_stream_block_ends_with_an_in_band_error_frame(self):
chunks = [_text_chunk(0, "A kum"), _text_chunk(0, "quat tart"), _text_chunk(0, "", "stop"), _usage_chunk()]
delivered = await _drive_flat_stream(_RecordingKeywordFilter(), chunks)
error_frame = delivered[-1]
assert isinstance(error_frame, bytes)
assert error_frame.startswith(b'data: {"error": ')
assert b"kumquat" in error_frame

50
tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py
View file

@ -41,6 +41,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import (
GuardrailPipeline,
PipelineStep,
)
from litellm.types.utils import TextChoices, TextCompletionResponse
@pytest.fixture(autouse=True)
@ -1886,7 +1887,7 @@ def test_streamable_post_call_pipelines_keeps_hook_guardrails_and_drops_iterator
@pytest.mark.parametrize(
"request_route",
["/v1/completions", "/v1beta/models/gemini-2.5-flash:streamGenerateContent", "/a2a/agent"],
["/v1beta/models/gemini-2.5-flash:streamGenerateContent", "/a2a/agent"],
)
def test_streamable_post_call_pipelines_keeps_legacy_hooks_off_routes_that_assemble_no_response(
make_user_api_key_auth, monkeypatch, caplog, request_route
@ -2036,6 +2037,53 @@ async def test_streaming_iterator_hook_delivers_legacy_hook_rewrite_on_every_cho
assert _warnings(caplog) == []
def _text_completion_chunk(text: str, finish_reason: str | None = None) -> TextCompletionResponse:
return TextCompletionResponse(
id="cmpl-1", model="m", choices=[TextChoices(index=0, text=text, finish_reason=finish_reason)]
)
def _rewritten_text_completion(response: Any) -> TextCompletionResponse:
return TextCompletionResponse(
id=response.id,
model=response.model,
choices=[TextChoices(index=0, text="[REWRITTEN] " + response.choices[0].text, finish_reason="stop")],
)
@pytest.mark.asyncio
async def test_streaming_iterator_hook_delivers_legacy_hook_rewrite_on_a_text_completion_stream(
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
):
seen: Dict[str, Any] = {}
monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen, rewrite=_rewritten_text_completion)])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = {**_post_call_pipeline_data(stream=True), "prompt": "hi"}
chunks = [_text_completion_chunk("hello "), _text_completion_chunk("world"), _text_completion_chunk("", "stop")]
auth = make_user_api_key_auth(request_route="/v1/completions")
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
await proxy_logging.pre_call_hook(
user_api_key_dict=auth, data=data, call_type="text_completion", guardrails_only=True
)
delivered = [
item
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data
)
]
assert isinstance(seen["response"], TextCompletionResponse)
assert seen["response"].choices[0].text == "hello world"
assert [(item.choices[0].text, item.choices[0].finish_reason) for item in delivered] == [
("[REWRITTEN] hello world", None),
("", None),
("", "stop"),
]
assert data["metadata"]["applied_guardrails"] == ["gr-post"]
assert _warnings(caplog) == []
@pytest.mark.asyncio
async def test_streaming_iterator_hook_releases_stream_untouched_when_legacy_hook_returns_none(
proxy_logging, make_user_api_key_auth, monkeypatch