mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-29 01:42:19 +00:00
fix(responses): stream guardrail pre-call block as SSE with a typed output item (#42507)
* fix(responses): stream guardrail pre-call block as SSE with a typed output item Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(responses): import blocked usage helper from the guardrail utils module Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(e2e): drop narrating docstrings and poll without rebinding Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * 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> * test(integration): audit cells for responses guardrail block contract Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): observe upstream on the recorded chat route for responses denial cells Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): tidy responses denial audit cells Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): wait for worker count to recover after SIGKILL Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): require a replacement worker after SIGKILL Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(responses): type the blocked response test helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yucheng <yucheng@berri.ai>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
e73abe6c72
commit
eea1d0f269
8 changed files with 1022 additions and 63 deletions
|
|
@ -1540,7 +1540,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
from litellm.responses.streaming_iterator import build_synthetic_response_events
|
||||
|
||||
return build_synthetic_response_events(
|
||||
transformed=_blocked_response(exc, response_id=f"resp_{uuid.uuid4()}", model=exc.model),
|
||||
transformed=build_blocked_response(exc),
|
||||
logging_obj=None,
|
||||
chunk_size=max(len(exc.message), 1),
|
||||
)
|
||||
|
|
@ -1648,6 +1648,10 @@ def _blocked_output_item(exc: "ModifyResponseException") -> GenericResponseOutpu
|
|||
return GenericResponseOutputItem.model_validate(payload)
|
||||
|
||||
|
||||
def build_blocked_response(exc: "ModifyResponseException") -> ResponsesAPIResponse:
|
||||
return _blocked_response(exc, response_id=f"resp_{uuid.uuid4()}", model=exc.model)
|
||||
|
||||
|
||||
def _blocked_response(
|
||||
exc: "ModifyResponseException",
|
||||
response_id: str,
|
||||
|
|
|
|||
|
|
@ -1,13 +1,11 @@
|
|||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import time
|
||||
from collections.abc import AsyncIterator, Awaitable, Mapping, Sequence
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Mapping, Sequence
|
||||
from enum import Enum
|
||||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, TypeAlias, cast, get_args
|
||||
from uuid import uuid4
|
||||
|
||||
import fastapi
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, Response
|
||||
|
|
@ -21,8 +19,9 @@ from typing_extensions import ReadOnly, TypedDict
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import EMPTY_MAPPING
|
||||
from litellm.integrations.custom_guardrail import ModifyResponseException
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import (
|
||||
blocked_responses_api_usage as _blocked_responses_api_usage,
|
||||
from litellm.llms.openai.responses.guardrail_translation.handler import (
|
||||
OpenAIResponsesHandler,
|
||||
build_blocked_response,
|
||||
)
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
|
|
@ -30,7 +29,7 @@ from litellm.proxy.auth.user_api_key_auth import (
|
|||
user_api_key_auth,
|
||||
user_api_key_auth_websocket,
|
||||
)
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing, create_response
|
||||
from litellm.proxy.common_utils.http_parsing_utils import (
|
||||
_read_request_body,
|
||||
_safe_set_request_parsed_body,
|
||||
|
|
@ -440,17 +439,16 @@ async def responses_api(
|
|||
request_data=_data,
|
||||
)
|
||||
|
||||
violation_text: Final = e.message
|
||||
response_obj: Final = ResponsesAPIResponse(
|
||||
id=f"resp_{uuid4()}",
|
||||
object="response",
|
||||
created_at=int(time.time()),
|
||||
model=e.model or data.get("model"),
|
||||
output=cast(Any, [{"content": [{"type": "text", "text": violation_text}]}]),
|
||||
status="completed",
|
||||
usage=_blocked_responses_api_usage(e.original_response),
|
||||
)
|
||||
return response_obj
|
||||
if data.get("stream") is True:
|
||||
block_chunks: Final = OpenAIResponsesHandler().build_block_sse_chunks(e)
|
||||
|
||||
async def _blocked_stream() -> AsyncGenerator[str, None]:
|
||||
for chunk in block_chunks:
|
||||
yield chunk.decode()
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
return await create_response(generator=_blocked_stream(), media_type="text/event-stream", headers={})
|
||||
return build_blocked_response(e)
|
||||
except Exception as e:
|
||||
raise await processor._handle_llm_api_exception(
|
||||
e=e,
|
||||
|
|
|
|||
|
|
@ -37,4 +37,5 @@
|
|||
- {id: guardrail.mcp_security.pre_call.blocks, module: guardrail, tier: P2, hook_point: pre_call, assertions: [blocks], exercised_on: [mcp_operations], source: "guardrail_hooks/mcp_security", rationale: "MCP protocol security"}
|
||||
- {id: guardrail.llm_as_a_judge.pre_call.blocks, module: guardrail, tier: P2, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/llm_as_a_judge", rationale: "LLM-based judgment guardrail"}
|
||||
- {id: guardrail.litellm_content_filter.pre_mcp_call.blocks, module: guardrail, tier: P1, hook_point: pre_mcp_call, assertions: [blocks], exercised_on: [mcp_operations], source: "guardrail_hooks/litellm_content_filter/content_filter.py:_scan_mcp_tool_call_arguments", rationale: "A general content-filter guardrail configured mode=pre_mcp_call blocks a banned keyword in an MCP tool call's arguments before it reaches the upstream MCP server; a clean argument passes"}
|
||||
- {id: guardrail.custom_code.pre_call.blocks, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [responses], source: "response_api_endpoints/endpoints.py ModifyResponseException handler", rationale: "A pre_call custom_code block on /v1/responses must answer in the requested shape: SSE response.completed with a completed assistant output_text message item when stream=true, schema-valid JSON when not, both with zero usage"}
|
||||
- {id: guardrail.dispatch.pre_call.rejects_unknown_name, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "proxy guardrail dispatch (per-request `guardrails` selector)", rationale: "A request naming a guardrail this proxy does not serve must fail closed with a 4xx; today it is silently served unguarded, so a typo'd name drops the protection the caller asked for"}
|
||||
|
|
|
|||
|
|
@ -111,6 +111,15 @@ class ToolPermissionParamsBody(GuardrailParamsBase):
|
|||
on_disallowed_action: Literal["block", "rewrite"] = "block"
|
||||
|
||||
|
||||
class CustomCodeParamsBody(GuardrailParamsBase):
|
||||
"""Custom-code guardrail params: `custom_code` is the sandboxed source the
|
||||
proxy compiles, which must define `apply_guardrail(inputs, request_data,
|
||||
input_type)` returning `allow()` or `block(reason)`."""
|
||||
|
||||
guardrail: Literal["custom_code"] = "custom_code"
|
||||
custom_code: str
|
||||
|
||||
|
||||
GuardrailParamsBody = (
|
||||
ContentFilterParamsBody
|
||||
| BedrockGuardrailParamsBody
|
||||
|
|
@ -118,6 +127,7 @@ GuardrailParamsBody = (
|
|||
| BlockCodeExecutionParamsBody
|
||||
| PresidioParamsBody
|
||||
| ToolPermissionParamsBody
|
||||
| CustomCodeParamsBody
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -174,6 +184,7 @@ class _ResponsesGuardrailBody(BaseModel):
|
|||
model: str
|
||||
input: str
|
||||
guardrails: list[str] | None = None
|
||||
stream: bool | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -509,6 +520,24 @@ class GuardrailsClient:
|
|||
json=_ResponsesGuardrailBody(model=model, input=text, guardrails=guardrails),
|
||||
)
|
||||
|
||||
def responses_stream_raw(
|
||||
self,
|
||||
key: str,
|
||||
model: str,
|
||||
text: str,
|
||||
*,
|
||||
guardrails: list[str] | None = None,
|
||||
) -> StreamingResponse:
|
||||
"""Drive /v1/responses with stream=true, returning the raw HTTP outcome:
|
||||
a streamed block is judged on status, content-type, and the SSE event
|
||||
sequence, not a typed JSON body."""
|
||||
return self.proxy.transport.send(
|
||||
"/v1/responses",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=_ResponsesGuardrailBody(model=model, input=text, guardrails=guardrails, stream=True),
|
||||
stream=True,
|
||||
)
|
||||
|
||||
def apply_guardrail(self, key: str, *, name: str, text: str) -> Result[ApplyGuardrailResponse]:
|
||||
return self.proxy.transport.post(
|
||||
"/guardrails/apply_guardrail",
|
||||
|
|
|
|||
154
tests/e2e/guardrails/test_responses_pre_call_block_stream_e2e.py
Normal file
154
tests/e2e/guardrails/test_responses_pre_call_block_stream_e2e.py
Normal file
|
|
@ -0,0 +1,154 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, unique_marker
|
||||
from e2e_http import StreamingResponse
|
||||
from guardrails_client import CustomCodeParamsBody, GuardrailsClient
|
||||
from lifecycle import ResourceManager
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
DENIAL: Final = "This model is not currently available. Please contact support if you think this is a mistake."
|
||||
|
||||
CUSTOM_CODE: Final = f'''
|
||||
def apply_guardrail(inputs, request_data, input_type):
|
||||
return block("{DENIAL}")
|
||||
'''
|
||||
|
||||
|
||||
class _ContentPart(BaseModel):
|
||||
type: str
|
||||
text: str | None = None
|
||||
|
||||
|
||||
class _OutputItem(BaseModel):
|
||||
type: str | None = None
|
||||
id: str | None = None
|
||||
role: str | None = None
|
||||
status: str | None = None
|
||||
content: list[_ContentPart] = []
|
||||
|
||||
|
||||
class _Usage(BaseModel):
|
||||
total_tokens: int = 0
|
||||
|
||||
|
||||
class _ResponseBody(BaseModel):
|
||||
output: list[_OutputItem] = []
|
||||
usage: _Usage | None = None
|
||||
|
||||
|
||||
class _EventHead(BaseModel):
|
||||
type: str
|
||||
|
||||
|
||||
class _CompletedEvent(BaseModel):
|
||||
type: str
|
||||
response: _ResponseBody
|
||||
|
||||
|
||||
_EVENT_HEAD: Final = TypeAdapter(_EventHead)
|
||||
|
||||
|
||||
def _denial_delivered(result: StreamingResponse) -> bool:
|
||||
if not result.ok:
|
||||
return False
|
||||
if DENIAL in result.body:
|
||||
return True
|
||||
return any(DENIAL in event for event in result.stream_events)
|
||||
|
||||
|
||||
def _poll_terminal(result: StreamingResponse) -> bool:
|
||||
if _denial_delivered(result):
|
||||
return True
|
||||
if result.ok:
|
||||
return False
|
||||
return "Guardrail not found" not in result.body and result.status_code not in (-1, 401, 429)
|
||||
|
||||
|
||||
def _poll_attempt(call: Callable[[], StreamingResponse], deadline: float) -> StreamingResponse:
|
||||
result: Final = call()
|
||||
if _poll_terminal(result) or time.monotonic() >= deadline:
|
||||
return result
|
||||
time.sleep(POLL_INTERVAL)
|
||||
return _poll_attempt(call, deadline)
|
||||
|
||||
|
||||
def _poll_for_block(call: Callable[[], StreamingResponse]) -> StreamingResponse:
|
||||
return _poll_attempt(call, time.monotonic() + POLL_TIMEOUT)
|
||||
|
||||
|
||||
def _assert_blocked_response(response: _ResponseBody) -> None:
|
||||
item = next(iter(response.output), None)
|
||||
assert item is not None, f"blocked response carried no output item: {response.output!r}"
|
||||
assert item.type == "message", f"output[0] must be a message item, got {item.type!r}: {item!r}"
|
||||
assert item.role == "assistant", f"output[0] role must be assistant, got {item.role!r}"
|
||||
assert item.status == "completed", f"output[0] status must be completed, got {item.status!r}"
|
||||
part = next(iter(item.content), None)
|
||||
assert part is not None, f"output[0] carried no content part: {item!r}"
|
||||
assert part.type == "output_text", f"content[0] must be output_text, got {part.type!r}"
|
||||
assert part.text == DENIAL, f"content[0] text must be the denial, got {part.text!r}"
|
||||
assert response.usage is not None and response.usage.total_tokens == 0, (
|
||||
f"a blocked response never reached a provider, usage must be zero: {response.usage!r}"
|
||||
)
|
||||
|
||||
|
||||
class TestResponsesPreCallBlock:
|
||||
def _register_block(self, client: GuardrailsClient, resources: ResourceManager) -> str:
|
||||
name: Final = f"e2e-custom-code-responses-block-{unique_marker()}"
|
||||
guardrail_id: Final = client.register(
|
||||
name,
|
||||
CustomCodeParamsBody(mode="pre_call", default_on=False, custom_code=CUSTOM_CODE),
|
||||
)
|
||||
resources.defer(lambda: client.delete_guardrail(guardrail_id))
|
||||
return name
|
||||
|
||||
@pytest.mark.covers("guardrail.custom_code.pre_call.blocks", exercised_on=["responses"])
|
||||
def test_stream_block_is_sse_with_completed_assistant_message(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
name: Final = self._register_block(client, resources)
|
||||
model: Final = client.create_backend_model(
|
||||
resources, prefix="e2e-responses-block", backend="openai/gpt-4.1-mini", api_key="os.environ/OPENAI_API_KEY"
|
||||
)
|
||||
|
||||
result: Final = _poll_for_block(
|
||||
lambda: client.responses_stream_raw(scoped_key, model, "say hi", guardrails=[name])
|
||||
)
|
||||
|
||||
assert result.status_code == 200, f"a pre_call block answers 200, got {result.status_code}: {result.body[:400]}"
|
||||
assert (result.content_type or "").startswith("text/event-stream"), (
|
||||
f"stream=true must answer SSE, got content-type {result.content_type!r}: {result.body[:400]}"
|
||||
)
|
||||
events: Final = tuple(_EVENT_HEAD.validate_json(payload).type for payload in result.stream_events)
|
||||
completed: Final = tuple(
|
||||
_CompletedEvent.model_validate_json(payload)
|
||||
for payload, event_type in zip(result.stream_events, events)
|
||||
if event_type == "response.completed"
|
||||
)
|
||||
assert len(completed) == 1, (
|
||||
f"the denial stream must end in exactly one response.completed event, got events {events!r}"
|
||||
)
|
||||
_assert_blocked_response(completed[0].response)
|
||||
|
||||
@pytest.mark.covers("guardrail.custom_code.pre_call.blocks", exercised_on=["responses"])
|
||||
def test_non_stream_block_is_schema_valid_json(
|
||||
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
|
||||
) -> None:
|
||||
name: Final = self._register_block(client, resources)
|
||||
model: Final = client.create_backend_model(
|
||||
resources, prefix="e2e-responses-block", backend="openai/gpt-4.1-mini", api_key="os.environ/OPENAI_API_KEY"
|
||||
)
|
||||
|
||||
result: Final = _poll_for_block(lambda: client.responses(scoped_key, model, "say hi", guardrails=[name]))
|
||||
|
||||
assert result.status_code == 200, f"a pre_call block answers 200, got {result.status_code}: {result.body[:400]}"
|
||||
assert (result.content_type or "").startswith("application/json"), (
|
||||
f"a non-streaming block answers JSON, got content-type {result.content_type!r}"
|
||||
)
|
||||
_assert_blocked_response(_ResponseBody.model_validate_json(result.body))
|
||||
|
|
@ -1,15 +1,22 @@
|
|||
import json
|
||||
import os
|
||||
import signal
|
||||
import socket
|
||||
import uuid
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import psutil
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually, object_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.mcp import mcp_peer, register_mcp, tool_names
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.process import group_members, owned_proxy, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from openai import AsyncOpenAI, OpenAI
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.rewrite_reaches_correct_anthropic_positions")
|
||||
|
|
@ -208,8 +215,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(
|
||||
|
|
@ -500,3 +505,687 @@ 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 _deny_guardrail(name: str, denial: str = _RESPONSES_DENIAL) -> dict[str, object]:
|
||||
return {
|
||||
"guardrail_name": name,
|
||||
"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({denial!r})\n"),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _responses_denial_config(tmp_path: Path, identity: str, denial: str = _RESPONSES_DENIAL) -> Path:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["guardrails"] = [_deny_guardrail(identity, denial)]
|
||||
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
|
||||
|
||||
|
||||
def _response_id(index: int, response: httpx.Response) -> str:
|
||||
assert response.status_code == 200, (index, response.text)
|
||||
if index % 3 == 0:
|
||||
assert response.headers["content-type"].startswith("text/event-stream"), response.text
|
||||
return str(_blocked_stream_events(response.text)[-1]["response"]["id"])
|
||||
if index % 3 == 1:
|
||||
assert response.headers["content-type"].startswith("text/event-stream"), response.text
|
||||
blocked: Final = _blocked_stream_events(response.text)[-1]["response"]
|
||||
_assert_blocked_message_item(blocked["output"][0], blocked)
|
||||
return str(blocked["id"])
|
||||
assert response.headers["content-type"].startswith("application/json"), response.text
|
||||
body: Final = response.json()
|
||||
_assert_blocked_message_item(body["output"][0], body)
|
||||
return str(body["id"])
|
||||
|
||||
|
||||
@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(use_chat_completions_api=True)
|
||||
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(use_chat_completions_api=True)
|
||||
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"] == []
|
||||
|
||||
|
||||
_RESPONSES_OUTPUT_DENIAL: Final = "Output withheld by policy."
|
||||
_UPSTREAM_INPUT_TOKENS: Final = 20
|
||||
_UPSTREAM_OUTPUT_TOKENS: Final = 20
|
||||
_UPSTREAM_TOTAL_TOKENS: Final = 40
|
||||
|
||||
|
||||
def _responses_output_denial_config(tmp_path: Path, identity: str, model: 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": "post_call",
|
||||
"default_on": False,
|
||||
"custom_code": (
|
||||
"def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
f" return block({_RESPONSES_OUTPUT_DENIAL!r})\n"
|
||||
),
|
||||
},
|
||||
}
|
||||
]
|
||||
config["policies"] = {
|
||||
f"{identity}-pipeline": {
|
||||
"guardrails": {"add": [identity]},
|
||||
"pipeline": {
|
||||
"mode": "post_call",
|
||||
"steps": [
|
||||
{
|
||||
"guardrail": identity,
|
||||
"on_pass": "allow",
|
||||
"on_fail": "modify_response",
|
||||
"modify_response_message": _RESPONSES_OUTPUT_DENIAL,
|
||||
}
|
||||
],
|
||||
},
|
||||
}
|
||||
}
|
||||
config["policy_attachments"] = [{"policy": f"{identity}-pipeline", "models": [model]}]
|
||||
path: Final = tmp_path / "responses-output-deny.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
def _blocked_stream_events(text: str) -> tuple[dict[str, object], ...]:
|
||||
lines: Final = tuple(line for line in text.split("\n") if line.startswith("data: "))
|
||||
assert lines[-1] == "data: [DONE]", text
|
||||
return tuple(json.loads(line.removeprefix("data: ")) for line in lines[:-1])
|
||||
|
||||
|
||||
def _dead_api_base() -> str:
|
||||
with socket.socket() as reserve:
|
||||
reserve.bind(("127.0.0.1", 0))
|
||||
port: Final = reserve.getsockname()[1]
|
||||
return f"http://127.0.0.1:{port}/v1"
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_openai_sdk_streams_typed_message")
|
||||
def test_responses_pre_call_denial_openai_sdk_streams_typed_message(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(use_chat_completions_api=True)
|
||||
client: Final = OpenAI(
|
||||
base_url=f"{candidate.client.base_url}/v1", api_key=candidate.key, max_retries=0, timeout=15
|
||||
)
|
||||
events: Final = tuple(
|
||||
client.responses.create(model=model, input="say hi", stream=True, extra_body={"guardrails": [identity]})
|
||||
)
|
||||
assert events[-1].type == "response.completed", [event.type for event in events]
|
||||
completed: Final = events[-1].response
|
||||
assert completed is not None and len(completed.output) == 1, completed
|
||||
item: Final = completed.output[0]
|
||||
assert item.type == "message", item
|
||||
assert item.role == "assistant" and item.status == "completed", item
|
||||
assert item.content[0].type == "output_text" and item.content[0].text == _RESPONSES_DENIAL, item.content
|
||||
assert completed.usage is not None and completed.usage.total_tokens == 0, completed.usage
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_openai_async_sdk_streams_typed_message")
|
||||
async def test_responses_pre_call_denial_openai_async_sdk_streams_typed_message(
|
||||
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(use_chat_completions_api=True)
|
||||
client: Final = AsyncOpenAI(
|
||||
base_url=f"{candidate.client.base_url}/v1", api_key=candidate.key, max_retries=0, timeout=15
|
||||
)
|
||||
stream: Final = await client.responses.create(
|
||||
model=model, input="say hi", stream=True, extra_body={"guardrails": [identity]}
|
||||
)
|
||||
kinds: Final = [event.type async for event in stream]
|
||||
assert kinds[-1] == "response.completed", kinds
|
||||
assert "response.output_text.delta" in kinds, kinds
|
||||
assert "response.in_progress" in kinds, kinds
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_openai_sdk_returns_typed_message")
|
||||
def test_responses_pre_call_denial_openai_sdk_returns_typed_message(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(use_chat_completions_api=True)
|
||||
client: Final = OpenAI(
|
||||
base_url=f"{candidate.client.base_url}/v1", api_key=candidate.key, max_retries=0, timeout=15
|
||||
)
|
||||
body: Final = client.responses.create(model=model, input="say hi", extra_body={"guardrails": [identity]})
|
||||
assert body.object == "response" and body.status == "completed", body
|
||||
assert len(body.output) == 1, body.output
|
||||
item: Final = body.output[0]
|
||||
assert item.type == "message" and item.role == "assistant", item
|
||||
assert item.content[0].type == "output_text" and item.content[0].text == _RESPONSES_DENIAL, item.content
|
||||
assert body.output_text == _RESPONSES_DENIAL, body
|
||||
assert body.usage is not None and body.usage.total_tokens == 0, body.usage
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_stream_false_returns_json")
|
||||
def test_responses_pre_call_denial_stream_false_returns_json(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(use_chat_completions_api=True)
|
||||
response: Final = candidate.request(
|
||||
"POST", "/v1/responses", {"model": model, "input": "say hi", "stream": False, "guardrails": [identity]}
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("application/json"), response.text
|
||||
body: Final = response.json()
|
||||
_assert_blocked_message_item(body["output"][0], body)
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_stream_string_true_returns_json")
|
||||
def test_responses_pre_call_denial_stream_string_true_returns_json(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(use_chat_completions_api=True)
|
||||
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("application/json"), (
|
||||
response.headers["content-type"],
|
||||
response.text,
|
||||
)
|
||||
body: Final = response.json()
|
||||
_assert_blocked_message_item(body["output"][0], body)
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_stream_event_vocabulary")
|
||||
def test_responses_pre_call_denial_stream_event_vocabulary(gateway: Gateway, tmp_path: Path) -> None:
|
||||
identity: Final = "guardrail" + uuid.uuid4().hex
|
||||
second: Final = "guardrail-2-" + uuid.uuid4().hex
|
||||
config: Final = _responses_denial_config(tmp_path, identity)
|
||||
loaded: Final = yaml.safe_load(config.read_text())
|
||||
loaded["guardrails"].append(_deny_guardrail(second))
|
||||
config.write_text(yaml.safe_dump(loaded))
|
||||
with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario:
|
||||
model: Final = scenario.model(use_chat_completions_api=True)
|
||||
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, second]},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("text/event-stream"), response.text
|
||||
events: Final = _blocked_stream_events(response.text)
|
||||
kinds: Final = {event["type"] for event in events}
|
||||
assert kinds == {
|
||||
"response.created",
|
||||
"response.in_progress",
|
||||
"response.output_item.added",
|
||||
"response.content_part.added",
|
||||
"response.output_text.delta",
|
||||
"response.output_text.done",
|
||||
"response.content_part.done",
|
||||
"response.output_item.done",
|
||||
"response.completed",
|
||||
}, kinds
|
||||
item_done: Final = tuple(event for event in events if event["type"] == "response.output_item.done")
|
||||
assert len(item_done) == 1, events
|
||||
assert len(events[-1]["response"]["output"]) == 1, events[-1]
|
||||
assert observed.get("/__observations").json()["requests"] == []
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_stream_large_denial_text")
|
||||
def test_responses_pre_call_denial_stream_large_denial_text(gateway: Gateway, tmp_path: Path) -> None:
|
||||
identity: Final = "guardrail" + uuid.uuid4().hex
|
||||
denial: Final = ("Denied: " + "mixed ascii and unicode text " * 200 + "fin")[:5000]
|
||||
config: Final = _responses_denial_config(tmp_path, identity, denial)
|
||||
with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario:
|
||||
model: Final = scenario.model(use_chat_completions_api=True)
|
||||
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.text
|
||||
events: Final = _blocked_stream_events(response.text)
|
||||
assert "".join(event["delta"] for event in events if event["type"] == "response.output_text.delta") == denial
|
||||
done: Final = next(event for event in events if event["type"] == "response.output_text.done")
|
||||
assert done["text"] == denial, done
|
||||
completed: Final = events[-1]["response"]
|
||||
assert completed["output"][0]["content"][0]["text"] == denial, completed
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_stream_requests_have_distinct_ids")
|
||||
def test_responses_pre_call_denial_stream_requests_have_distinct_ids(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(use_chat_completions_api=True)
|
||||
with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as observed:
|
||||
observed.get("/__observations")
|
||||
responses: Final = tuple(
|
||||
candidate.request(
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
{"model": model, "input": "say hi", "stream": True, "guardrails": [identity]},
|
||||
)
|
||||
for _ in range(2)
|
||||
)
|
||||
completed: Final = tuple(_blocked_stream_events(response.text)[-1]["response"] for response in responses)
|
||||
for response in responses:
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("text/event-stream"), response.text
|
||||
assert completed[0]["id"] != completed[1]["id"], completed
|
||||
assert completed[0]["output"][0]["id"] != completed[1]["output"][0]["id"], completed
|
||||
assert observed.get("/__observations").json()["requests"] == []
|
||||
|
||||
|
||||
def _register_named_model(candidate: Gateway, name: str, api_base: str | None = None, **parameters: object) -> str:
|
||||
created: Final = candidate.post(
|
||||
"/model/new",
|
||||
{
|
||||
"model_name": name,
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "integration-provider-key",
|
||||
"api_base": api_base or f"{candidate.upstream_url}/v1",
|
||||
**parameters,
|
||||
},
|
||||
},
|
||||
)
|
||||
return str(created["model_info"]["id"])
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.responses_post_call_pipeline_denial_streams_real_usage")
|
||||
def test_responses_post_call_pipeline_denial_streams_real_usage(gateway: Gateway, tmp_path: Path) -> None:
|
||||
identity: Final = "guardrail" + uuid.uuid4().hex
|
||||
model: Final = f"integration-{uuid.uuid4().hex}"
|
||||
config: Final = _responses_output_denial_config(tmp_path, identity, model)
|
||||
with owned_proxy(gateway, tmp_path, {}, config=config) as candidate:
|
||||
model_id: Final = _register_named_model(candidate, model, use_chat_completions_api=True)
|
||||
try:
|
||||
response: Final = candidate.request(
|
||||
"POST", "/v1/responses", {"model": model, "input": f"say hi {uuid.uuid4().hex}", "stream": True}
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("text/event-stream"), response.text
|
||||
events: Final = _blocked_stream_events(response.text)
|
||||
assert events[-1]["type"] == "response.completed", events
|
||||
completed: Final = events[-1]["response"]
|
||||
item: Final = completed["output"][0]
|
||||
assert item["type"] == "message" and item["role"] == "assistant", item
|
||||
assert item["content"][0]["type"] == "output_text", item
|
||||
assert item["content"][0]["text"] == _RESPONSES_OUTPUT_DENIAL, item
|
||||
usage: Final = completed["usage"]
|
||||
assert (
|
||||
usage["input_tokens"],
|
||||
usage["output_tokens"],
|
||||
usage["total_tokens"],
|
||||
) == (_UPSTREAM_INPUT_TOKENS, _UPSTREAM_OUTPUT_TOKENS, _UPSTREAM_TOTAL_TOKENS), usage
|
||||
finally:
|
||||
candidate.post("/model/delete", {"id": model_id})
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.responses_post_call_pipeline_denial_returns_real_usage")
|
||||
def test_responses_post_call_pipeline_denial_returns_real_usage(gateway: Gateway, tmp_path: Path) -> None:
|
||||
identity: Final = "guardrail" + uuid.uuid4().hex
|
||||
model: Final = f"integration-{uuid.uuid4().hex}"
|
||||
config: Final = _responses_output_denial_config(tmp_path, identity, model)
|
||||
with owned_proxy(gateway, tmp_path, {}, config=config) as candidate:
|
||||
model_id: Final = _register_named_model(candidate, model, use_chat_completions_api=True)
|
||||
try:
|
||||
response: Final = candidate.request(
|
||||
"POST", "/v1/responses", {"model": model, "input": f"say hi {uuid.uuid4().hex}"}
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = response.json()
|
||||
item: Final = body["output"][0]
|
||||
assert item["type"] == "message" and item["role"] == "assistant", item
|
||||
assert item["content"][0]["type"] == "output_text", item
|
||||
assert item["content"][0]["text"] == _RESPONSES_OUTPUT_DENIAL, item
|
||||
usage: Final = body["usage"]
|
||||
assert (
|
||||
usage["input_tokens"],
|
||||
usage["output_tokens"],
|
||||
usage["total_tokens"],
|
||||
) == (_UPSTREAM_INPUT_TOKENS, _UPSTREAM_OUTPUT_TOKENS, _UPSTREAM_TOTAL_TOKENS), usage
|
||||
finally:
|
||||
candidate.post("/model/delete", {"id": model_id})
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.responses_denial_requires_authentication")
|
||||
def test_responses_denial_requires_authentication(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()
|
||||
response: Final = candidate.request(
|
||||
"POST", "/v1/responses", {"model": model, "input": "say hi", "guardrails": [identity]}, key="sk-invalid"
|
||||
)
|
||||
assert response.status_code == 401, (response.status_code, response.text)
|
||||
assert response.json()["error"]["type"] == "token_not_found_in_db", response.text
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_stream_does_not_reach_upstream")
|
||||
def test_responses_pre_call_denial_stream_does_not_reach_upstream(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(api_base=_dead_api_base())
|
||||
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.text
|
||||
events: Final = _blocked_stream_events(response.text)
|
||||
completed: Final = events[-1]["response"]
|
||||
_assert_blocked_message_item(completed["output"][0], completed)
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.responses_unguarded_stream_reaches_upstream")
|
||||
def test_responses_unguarded_stream_reaches_upstream(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:
|
||||
dead: Final = scenario.model(api_base=_dead_api_base())
|
||||
denied: Final = candidate.request(
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
{"model": dead, "input": "say hi", "stream": True, "guardrails": [identity]},
|
||||
)
|
||||
assert denied.status_code == 200, denied.text
|
||||
model: Final = scenario.model(use_chat_completions_api=True)
|
||||
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": f"say hi {uuid.uuid4().hex}", "stream": True, "guardrails": []},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("text/event-stream"), response.text
|
||||
assert "response.completed" in response.text, response.text
|
||||
requests: Final = eventually(
|
||||
lambda: observed.get("/__observations").json()["requests"],
|
||||
lambda values: len(values) >= 1,
|
||||
seconds=30,
|
||||
)
|
||||
assert len(requests) == 1, requests
|
||||
assert requests[0]["path"] == "/v1/chat/completions", requests
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.chat_pre_call_denial_streams_content_filter")
|
||||
def test_chat_pre_call_denial_streams_content_filter(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()
|
||||
response: Final = candidate.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": "say hi"}],
|
||||
"stream": True,
|
||||
"guardrails": [identity],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("text/event-stream"), response.text
|
||||
lines: Final = tuple(line for line in response.text.split("\n") if line.startswith("data: "))
|
||||
assert lines[-1] == "data: [DONE]", response.text
|
||||
chunks: Final = tuple(json.loads(line.removeprefix("data: ")) for line in lines[:-1])
|
||||
assert chunks[0]["choices"][0]["delta"]["content"] == _RESPONSES_DENIAL, chunks
|
||||
assert chunks[-1]["choices"][0]["finish_reason"] == "stop", chunks
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.chat_pre_call_denial_returns_content_filter")
|
||||
def test_chat_pre_call_denial_returns_content_filter(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()
|
||||
response: Final = candidate.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": "say hi"}],
|
||||
"guardrails": [identity],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = response.json()
|
||||
choice: Final = body["choices"][0]
|
||||
assert choice["finish_reason"] == "content_filter", body
|
||||
assert choice["message"]["content"] == _RESPONSES_DENIAL, body
|
||||
assert (
|
||||
body["usage"]["prompt_tokens"],
|
||||
body["usage"]["completion_tokens"],
|
||||
body["usage"]["total_tokens"],
|
||||
) == (0, 0, 0), body["usage"]
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.messages_pre_call_denial_returns_message")
|
||||
def test_messages_pre_call_denial_returns_message(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()
|
||||
response: Final = candidate.request(
|
||||
"POST",
|
||||
"/v1/messages",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": "say hi"}],
|
||||
"max_tokens": 16,
|
||||
"guardrails": [identity],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = response.json()
|
||||
assert body["type"] == "message" and body["role"] == "assistant", body
|
||||
assert body["content"] == [{"type": "text", "text": _RESPONSES_DENIAL}], body
|
||||
assert body["stop_reason"] == "end_turn", body
|
||||
assert (body["usage"]["input_tokens"], body["usage"]["output_tokens"]) == (0, 0), body
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.messages_pre_call_denial_streams_message")
|
||||
def test_messages_pre_call_denial_streams_message(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()
|
||||
response: Final = candidate.request(
|
||||
"POST",
|
||||
"/v1/messages",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": "say hi"}],
|
||||
"max_tokens": 16,
|
||||
"stream": True,
|
||||
"guardrails": [identity],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("text/event-stream"), response.text
|
||||
lines: Final = tuple(line for line in response.text.split("\n") if line.startswith("data: "))
|
||||
assert len(lines) == 1, response.text
|
||||
body: Final = json.loads(lines[0].removeprefix("data: "))
|
||||
assert body["type"] == "message" and body["role"] == "assistant", body
|
||||
assert body["content"] == [{"type": "text", "text": _RESPONSES_DENIAL}], body
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_writes_zero_spend_row")
|
||||
def test_responses_pre_call_denial_writes_zero_spend_row(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()
|
||||
response: Final = candidate.request(
|
||||
"POST", "/v1/responses", {"model": model, "input": "say hi", "guardrails": [identity]}
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT spend, total_tokens FROM "LiteLLM_SpendLogs" WHERE model=%s AND call_type=%s',
|
||||
(model, "aresponses"),
|
||||
),
|
||||
lambda values: len(values) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert float(rows[0]["spend"]) == 0, rows
|
||||
assert rows[0]["total_tokens"] == 0, rows
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_stream_survives_worker_burst")
|
||||
def test_responses_pre_call_denial_stream_survives_worker_burst(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, workers=2) as candidate, candidate.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=_dead_api_base())
|
||||
healthy: Final = scenario.model(use_chat_completions_api=True)
|
||||
with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as observed:
|
||||
observed.get("/__observations")
|
||||
|
||||
def burst(index: int) -> httpx.Response:
|
||||
if index % 3 == 0:
|
||||
return candidate.request(
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
{"model": healthy, "input": f"say hi {uuid.uuid4().hex} {index}", "stream": True},
|
||||
)
|
||||
stream: Final = index % 3 == 1
|
||||
return candidate.request(
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
{"model": model, "input": f"say hi {index}", "stream": stream, "guardrails": [identity]},
|
||||
)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
responses: Final = tuple(pool.map(burst, range(30)))
|
||||
response_ids: Final = frozenset(_response_id(index, response) for index, response in enumerate(responses))
|
||||
assert len(response_ids) == 30, response_ids
|
||||
assert len(observed.get("/__observations").json()["requests"]) == 10
|
||||
|
||||
|
||||
@pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_stream_survives_worker_kill")
|
||||
def test_responses_pre_call_denial_stream_survives_worker_kill(gateway: Gateway, tmp_path: Path) -> None:
|
||||
identity: Final = "guardrail" + uuid.uuid4().hex
|
||||
config: Final = _responses_denial_config(tmp_path, identity)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
with candidate.scenario() as scenario:
|
||||
model: Final = scenario.model(api_base=_dead_api_base())
|
||||
members: Final = tuple(
|
||||
member for member in group_members(owned.process.pid) if member.pid != owned.process.pid
|
||||
)
|
||||
children: Final = tuple(member.pid for member in members)
|
||||
workers: Final = tuple(
|
||||
member.pid for member in members if any("spawn_main" in part for part in member.cmdline())
|
||||
)
|
||||
assert len(workers) >= 2, workers
|
||||
os.kill(workers[0], signal.SIGKILL)
|
||||
expected: Final = len(children)
|
||||
eventually(
|
||||
lambda: tuple(
|
||||
member.pid
|
||||
for member in group_members(owned.process.pid)
|
||||
if member.pid != owned.process.pid
|
||||
and member.is_running()
|
||||
and member.status() != psutil.STATUS_ZOMBIE
|
||||
),
|
||||
lambda pids: len(pids) >= expected and any(pid not in children for pid in pids),
|
||||
seconds=30,
|
||||
)
|
||||
|
||||
def burst(index: int) -> httpx.Response:
|
||||
return candidate.request(
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
{"model": model, "input": f"say hi {index}", "stream": True, "guardrails": [identity]},
|
||||
)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=5) as pool:
|
||||
responses: Final = tuple(pool.map(burst, range(10)))
|
||||
for response in responses:
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("text/event-stream"), response.text
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ Test for response_api_endpoints/endpoints.py
|
|||
"""
|
||||
|
||||
import unittest
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Final, Literal
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
|
@ -14,6 +15,7 @@ from httpx import Response
|
|||
|
||||
import litellm
|
||||
from litellm.proxy.proxy_server import app
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -2193,6 +2195,59 @@ class TestCursorGateRecognizesRoutingGroups:
|
|||
assert "reasoning_effort" not in resolved
|
||||
|
||||
|
||||
BLOCK_MESSAGE = "Content flagged by policy, response withheld"
|
||||
|
||||
|
||||
def _post_blocked_responses(
|
||||
original_response: ResponsesAPIResponse | litellm.ModelResponse | None,
|
||||
payload: Mapping[str, object] | None = None,
|
||||
) -> httpx.Response:
|
||||
from litellm.integrations.custom_guardrail import ModifyResponseException
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
exc = ModifyResponseException(
|
||||
message=BLOCK_MESSAGE,
|
||||
model="gpt-4o-mini",
|
||||
request_data={"model": "gpt-4o-mini", "input": "hi"},
|
||||
guardrail_name="zero-usage-regression",
|
||||
original_response=original_response,
|
||||
)
|
||||
mock_proxy_logging = MagicMock()
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
api_key="sk-test", request_route="/v1/responses"
|
||||
)
|
||||
body = {"model": "gpt-4o-mini", "input": "Write a haiku about token accounting"}
|
||||
if payload:
|
||||
body.update(payload)
|
||||
try:
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing.base_process_llm_request",
|
||||
new=AsyncMock(side_effect=exc),
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging),
|
||||
):
|
||||
client = TestClient(app)
|
||||
return client.post("/v1/responses", json=body, headers={"Authorization": "Bearer sk-1234"})
|
||||
finally:
|
||||
app.dependency_overrides.pop(user_api_key_auth, None)
|
||||
|
||||
|
||||
def _assert_blocked_output_item(item: Mapping[str, object], text: str) -> None:
|
||||
assert item["type"] == "message"
|
||||
assert item["id"].startswith("msg_")
|
||||
assert item["role"] == "assistant"
|
||||
assert item["status"] == "completed"
|
||||
assert item["content"][0]["type"] == "output_text"
|
||||
assert item["content"][0]["text"] == text
|
||||
|
||||
|
||||
def _sse_data_frames(text: str) -> list[str]:
|
||||
return [line.removeprefix("data: ").strip() for line in text.splitlines() if line.startswith("data: ")]
|
||||
|
||||
|
||||
class TestGuardrailBlockedResponsesUsage:
|
||||
"""Regression tests for https://github.com/BerriAI/litellm/issues/36880.
|
||||
|
||||
|
|
@ -2202,38 +2257,7 @@ class TestGuardrailBlockedResponsesUsage:
|
|||
e.original_response, exactly like /v1/chat/completions already does."""
|
||||
|
||||
def _post_blocked_responses(self, original_response):
|
||||
from litellm.integrations.custom_guardrail import ModifyResponseException
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
exc = ModifyResponseException(
|
||||
message="Content flagged by policy, response withheld",
|
||||
model="gpt-4o-mini",
|
||||
request_data={"model": "gpt-4o-mini", "input": "hi"},
|
||||
guardrail_name="zero-usage-regression",
|
||||
original_response=original_response,
|
||||
)
|
||||
mock_proxy_logging = MagicMock()
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
api_key="sk-test", request_route="/v1/responses"
|
||||
)
|
||||
try:
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing.base_process_llm_request",
|
||||
new=AsyncMock(side_effect=exc),
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging),
|
||||
):
|
||||
client = TestClient(app)
|
||||
return client.post(
|
||||
"/v1/responses",
|
||||
json={"model": "gpt-4o-mini", "input": "Write a haiku about token accounting"},
|
||||
headers={"Authorization": "Bearer sk-1234"},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(user_api_key_auth, None)
|
||||
return _post_blocked_responses(original_response)
|
||||
|
||||
def test_post_call_block_reports_real_upstream_usage(self):
|
||||
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
|
||||
|
|
@ -2429,6 +2453,66 @@ class TestResponsesInputTokens:
|
|||
assert response.json()["error"]["message"] == "rate limited"
|
||||
|
||||
|
||||
class TestGuardrailBlockedResponsesShape:
|
||||
"""A pre_call block raises ModifyResponseException before any provider call.
|
||||
|
||||
The reply must satisfy the Responses API contract the request selected:
|
||||
stream=true answers SSE ending in one response.completed whose output[0] is
|
||||
a completed assistant message item with output_text content, and a plain
|
||||
POST answers JSON with the same item, both with the usage the blocked call
|
||||
consumed (zero for pre_call)."""
|
||||
|
||||
def test_non_stream_block_is_a_completed_assistant_message(self):
|
||||
response = _post_blocked_responses(None)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("application/json")
|
||||
body = response.json()
|
||||
_assert_blocked_output_item(body["output"][0], BLOCK_MESSAGE)
|
||||
assert body["usage"]["total_tokens"] == 0
|
||||
|
||||
def test_stream_block_answers_sse_with_completed_event(self):
|
||||
response = _post_blocked_responses(None, payload={"stream": True})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("text/event-stream")
|
||||
frames = _sse_data_frames(response.text)
|
||||
assert frames[-1] == "[DONE]"
|
||||
events = [json.loads(frame) for frame in frames[:-1]]
|
||||
types = [event["type"] for event in events]
|
||||
assert "response.created" in types
|
||||
completed = [event for event in events if event["type"] == "response.completed"]
|
||||
assert len(completed) == 1
|
||||
completed_response = completed[0]["response"]
|
||||
_assert_blocked_output_item(completed_response["output"][0], BLOCK_MESSAGE)
|
||||
assert completed_response["usage"]["total_tokens"] == 0
|
||||
delta_text = "".join(event["delta"] for event in events if event["type"] == "response.output_text.delta")
|
||||
assert delta_text == BLOCK_MESSAGE
|
||||
|
||||
def test_stream_block_keeps_upstream_usage(self):
|
||||
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
|
||||
|
||||
original = ResponsesAPIResponse(
|
||||
id="resp_upstream",
|
||||
created_at=1,
|
||||
model="gpt-4o-mini",
|
||||
object="response",
|
||||
output=[],
|
||||
status="completed",
|
||||
usage=ResponseAPIUsage(input_tokens=14, output_tokens=20, total_tokens=34),
|
||||
)
|
||||
|
||||
response = _post_blocked_responses(original, payload={"stream": True})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
frames = _sse_data_frames(response.text)
|
||||
completed = [json.loads(frame) for frame in frames[:-1] if json.loads(frame)["type"] == "response.completed"]
|
||||
usage = completed[0]["response"]["usage"]
|
||||
assert usage["input_tokens"] == 14
|
||||
assert usage["output_tokens"] == 20
|
||||
assert usage["total_tokens"] == 34
|
||||
|
||||
|
||||
def test_responses_routes_document_response_models_in_openapi_schema():
|
||||
from typing import cast
|
||||
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ proxy endpoints (/v1/chat/completions, /v1/completions, and /v1/responses).
|
|||
|
||||
A post-call block replaces the LLM response with the violation message, but the
|
||||
upstream call already consumed tokens. `_blocked_response_usage` (and its
|
||||
Responses API counterpart `_blocked_responses_api_usage`) reports that real
|
||||
Responses API counterpart `blocked_responses_api_usage`) reports that real
|
||||
usage (carried on `ModifyResponseException.original_response`) rather than
|
||||
zero; a pre-call block never invoked the LLM, so usage is zero.
|
||||
"""
|
||||
|
|
@ -91,8 +91,8 @@ def test_responses_api_blocked_reply_carries_real_usage():
|
|||
"""
|
||||
import time
|
||||
|
||||
from litellm.proxy.response_api_endpoints.endpoints import (
|
||||
_blocked_responses_api_usage,
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import (
|
||||
blocked_responses_api_usage,
|
||||
)
|
||||
|
||||
original_response = ResponsesAPIResponse(
|
||||
|
|
@ -105,7 +105,7 @@ def test_responses_api_blocked_reply_carries_real_usage():
|
|||
usage=ResponseAPIUsage(input_tokens=14, output_tokens=20, total_tokens=34),
|
||||
)
|
||||
|
||||
usage = _blocked_responses_api_usage(original_response)
|
||||
usage = blocked_responses_api_usage(original_response)
|
||||
|
||||
assert usage.input_tokens == 14
|
||||
assert usage.output_tokens == 20
|
||||
|
|
@ -114,11 +114,11 @@ def test_responses_api_blocked_reply_carries_real_usage():
|
|||
|
||||
def test_responses_api_blocked_reply_zero_usage_when_no_original_response():
|
||||
"""Pre-call block has no original_response, so usage must be zero."""
|
||||
from litellm.proxy.response_api_endpoints.endpoints import (
|
||||
_blocked_responses_api_usage,
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import (
|
||||
blocked_responses_api_usage,
|
||||
)
|
||||
|
||||
usage = _blocked_responses_api_usage(None)
|
||||
usage = blocked_responses_api_usage(None)
|
||||
|
||||
assert usage.input_tokens == 0
|
||||
assert usage.output_tokens == 0
|
||||
|
|
@ -128,14 +128,14 @@ def test_responses_api_blocked_reply_zero_usage_when_no_original_response():
|
|||
def test_responses_api_blocked_reply_maps_bridged_chat_usage():
|
||||
"""A chat model bridged through /v1/responses blocks with a ModelResponse whose
|
||||
Usage fields must map prompt_tokens -> input_tokens and completion_tokens -> output_tokens."""
|
||||
from litellm.proxy.response_api_endpoints.endpoints import (
|
||||
_blocked_responses_api_usage,
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import (
|
||||
blocked_responses_api_usage,
|
||||
)
|
||||
|
||||
resp = litellm.ModelResponse()
|
||||
resp.usage = litellm.Usage(prompt_tokens=14, completion_tokens=18, total_tokens=32)
|
||||
|
||||
usage = _blocked_responses_api_usage(resp)
|
||||
usage = blocked_responses_api_usage(resp)
|
||||
|
||||
assert usage.input_tokens == 14
|
||||
assert usage.output_tokens == 18
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue