mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
test(integration): cover pre-call guardrail block on /v1/responses stream and json
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng
parent
5025bff4b2
commit
e4af26ed3f
2 changed files with 105 additions and 5 deletions
|
|
@ -1737,6 +1737,12 @@
|
|||
],
|
||||
"tests/integration/mcp/test_mcp_lifecycle.py::test_same_url_server_grants_scope_discovery_and_direct_or_virtual_execution[bearer]": [
|
||||
"other.mcp.permissions.same_url_servers_enforce_discovery_and_execution"
|
||||
],
|
||||
"tests/integration/observability/test_guardrail_effects.py::test_responses_pre_call_denial_streams_sse_with_typed_message_item": [
|
||||
"other.observability.guardrails.responses_pre_call_denial_streams_typed_message"
|
||||
],
|
||||
"tests/integration/observability/test_guardrail_effects.py::test_responses_pre_call_denial_returns_json_with_typed_message_item": [
|
||||
"other.observability.guardrails.responses_pre_call_denial_returns_typed_message"
|
||||
]
|
||||
},
|
||||
"browser": {
|
||||
|
|
|
|||
|
|
@ -3,11 +3,10 @@ import uuid
|
|||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.mcp import mcp_peer, register_mcp, tool_names
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
|
|
@ -124,8 +123,6 @@ def test_guardrail_denial_prevents_provider_and_preserves_allowed_control(gatewa
|
|||
with owned_proxy(gateway, tmp_path, {}, config=path) as candidate, candidate.scenario() as scenario:
|
||||
model: Final = scenario.model()
|
||||
key: Final = scenario.key(models=[model])
|
||||
import httpx
|
||||
|
||||
with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as observed:
|
||||
observed.get("/__observations")
|
||||
denied: Final = candidate.request(
|
||||
|
|
@ -214,3 +211,100 @@ def test_request_selected_mcp_guardrail_blocks_direct_and_virtual_calls(gateway:
|
|||
assert len(calls) == 1
|
||||
assert calls[0]["body"]["params"]["name"] == tool
|
||||
assert calls[0]["body"]["params"]["arguments"] == arguments
|
||||
|
||||
|
||||
_RESPONSES_DENIAL: Final = "This model is not currently available."
|
||||
|
||||
|
||||
def _responses_denial_config(tmp_path: Path, identity: str) -> Path:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["guardrails"] = [
|
||||
{
|
||||
"guardrail_name": identity,
|
||||
"litellm_params": {
|
||||
"guardrail": "custom_code",
|
||||
"mode": "pre_call",
|
||||
"default_on": False,
|
||||
"custom_code": (
|
||||
f"def apply_guardrail(inputs, request_data, input_type):\n return block({_RESPONSES_DENIAL!r})\n"
|
||||
),
|
||||
},
|
||||
}
|
||||
]
|
||||
path: Final = tmp_path / "responses-deny.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
def _assert_blocked_message_item(item: dict[str, object], response: dict[str, object]) -> None:
|
||||
assert item["type"] == "message", item
|
||||
assert item["role"] == "assistant", item
|
||||
assert item["status"] == "completed", item
|
||||
assert str(item["id"]).startswith("msg_"), item
|
||||
assert item["content"] == [{"type": "output_text", "text": _RESPONSES_DENIAL, "annotations": []}], item
|
||||
assert response["status"] == "completed", response
|
||||
usage: Final = response["usage"]
|
||||
assert isinstance(usage, dict), response
|
||||
assert (usage["input_tokens"], usage["output_tokens"], usage["total_tokens"]) == (0, 0, 0), usage
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_streams_typed_message")
|
||||
def test_responses_pre_call_denial_streams_sse_with_typed_message_item(gateway: Gateway, tmp_path: Path) -> None:
|
||||
identity: Final = "guardrail" + uuid.uuid4().hex
|
||||
config: Final = _responses_denial_config(tmp_path, identity)
|
||||
with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario:
|
||||
model: Final = scenario.model()
|
||||
with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as observed:
|
||||
observed.get("/__observations")
|
||||
response: Final = candidate.request(
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
{"model": model, "input": "say hi", "stream": True, "guardrails": [identity]},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("text/event-stream"), (
|
||||
response.headers["content-type"],
|
||||
response.text,
|
||||
)
|
||||
lines: Final = tuple(line for line in response.text.split("\n") if line.startswith("data: "))
|
||||
assert lines[-1] == "data: [DONE]", response.text
|
||||
events: Final = tuple(json.loads(line.removeprefix("data: ")) for line in lines[:-1])
|
||||
kinds: Final = tuple(event["type"] for event in events)
|
||||
assert tuple(kind for kind in kinds if kind != "response.output_text.delta") == (
|
||||
"response.created",
|
||||
"response.in_progress",
|
||||
"response.output_item.added",
|
||||
"response.content_part.added",
|
||||
"response.output_text.done",
|
||||
"response.content_part.done",
|
||||
"response.output_item.done",
|
||||
"response.completed",
|
||||
), kinds
|
||||
assert kinds.index("response.output_text.delta") == kinds.index("response.content_part.added") + 1, kinds
|
||||
assert "".join(event["delta"] for event in events if event["type"] == "response.output_text.delta") == (
|
||||
_RESPONSES_DENIAL
|
||||
)
|
||||
completed: Final = events[-1]["response"]
|
||||
assert completed["output"] == [events[-2]["item"]], (completed, events[-2])
|
||||
_assert_blocked_message_item(completed["output"][0], completed)
|
||||
assert observed.get("/__observations").json()["requests"] == []
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_returns_typed_message")
|
||||
def test_responses_pre_call_denial_returns_json_with_typed_message_item(gateway: Gateway, tmp_path: Path) -> None:
|
||||
identity: Final = "guardrail" + uuid.uuid4().hex
|
||||
config: Final = _responses_denial_config(tmp_path, identity)
|
||||
with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario:
|
||||
model: Final = scenario.model()
|
||||
with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as observed:
|
||||
observed.get("/__observations")
|
||||
response: Final = candidate.request(
|
||||
"POST", "/v1/responses", {"model": model, "input": "say hi", "guardrails": [identity]}
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("application/json"), response.headers["content-type"]
|
||||
body: Final = response.json()
|
||||
assert body["object"] == "response", body
|
||||
assert len(body["output"]) == 1, body
|
||||
_assert_blocked_message_item(body["output"][0], body)
|
||||
assert observed.get("/__observations").json()["requests"] == []
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue