diff --git a/tests/integration/authorization/test_guardrail_submissions.py b/tests/integration/authorization/test_guardrail_submissions.py new file mode 100644 index 00000000000..f120028f46d --- /dev/null +++ b/tests/integration/authorization/test_guardrail_submissions.py @@ -0,0 +1,521 @@ +import json +import uuid +from contextlib import ExitStack +from datetime import datetime +from hashlib import sha256 +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, object_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_PROMPT_PENDING: Final = "pending submission must not run" +_PROMPT_APPROVED: Final = "approved submission runs" +_PROMPT_REJECTED: Final = "rejected submission must not run" +_APPLY_TEXT: Final = "apply guardrail after approval" +_PROVIDER_RESPONSE: Final = { + "id": "chatcmpl-guardrail-submission", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "provider response"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +class _RegisterResponse(BaseModel): + model_config = ConfigDict(extra="forbid") + + guardrail_id: str + guardrail_name: str + status: str + submitted_at: datetime + + +class _GuardrailRow(BaseModel): + guardrail_id: str + status: str + team_id: str + + +class _ProxyError(BaseModel): + model_config = ConfigDict(extra="forbid") + + message: str + type: str + param: str | None + code: str + + +class _ErrorResponse(BaseModel): + model_config = ConfigDict(extra="forbid") + + error: _ProxyError + + +class _SubmissionForbiddenResponse(BaseModel): + model_config = ConfigDict(extra="forbid") + + detail: str + + +class _AdminActionResponse(BaseModel): + model_config = ConfigDict(extra="forbid") + + guardrail_id: str + status: str + message: str + + +class _ApplyResponse(BaseModel): + model_config = ConfigDict(extra="forbid") + + response_text: str + + +class _ChatMessage(BaseModel): + role: str + content: str | None + + +class _ChatChoice(BaseModel): + index: int + message: _ChatMessage + finish_reason: str | None + + +class _ChatResponse(BaseModel): + choices: list[_ChatChoice] + + +class _SubmissionItem(BaseModel): + guardrail_id: str + + +class _SubmissionsResponse(BaseModel): + submissions: list[_SubmissionItem] + summary: dict[str, int] + + +def _delete_guardrail(gateway: Gateway, guardrail_id: str) -> None: + deleted: Final = gateway.request("DELETE", f"/guardrails/{guardrail_id}") + assert deleted.status_code == 200, deleted.text + assert read_rows('SELECT guardrail_id FROM "LiteLLM_GuardrailsTable" WHERE guardrail_id=%s', (guardrail_id,)) == [] + + +def _assert_submission_status(guardrail_id: str, status: str, team_id: str) -> None: + rows: Final = read_rows( + 'SELECT guardrail_id, status, team_id FROM "LiteLLM_GuardrailsTable" WHERE guardrail_id=%s', + (guardrail_id,), + ) + assert len(rows) == 1, rows + assert _GuardrailRow.model_validate(rows[0]) == _GuardrailRow( + guardrail_id=guardrail_id, status=status, team_id=team_id + ), rows + + +def _assert_submission_actions_are_forbidden(responses: tuple[httpx.Response, ...]) -> None: + assert tuple(response.status_code for response in responses) == (403, 403, 403, 403), tuple( + response.text for response in responses + ) + expected_error: Final = {"detail": "Admin access required"} + assert tuple( + _SubmissionForbiddenResponse.model_validate_json(response.content).model_dump() for response in responses + ) == (expected_error,) * 4, tuple(response.text for response in responses) + + +def _mask_user_id(user_id: str) -> str: + return f"{user_id[:6]}{'*' * (len(user_id) - 8)}{user_id[-2:]}" + + +def _register( + gateway: Gateway, + scenario_cleanups: ExitStack, + name: str, + team_id: str | None, + key: str, + api_base: str, +) -> str: + params: Final[dict[str, JsonValue]] = { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "api_base": api_base, + "api_key": "synthetic-submission-guardrail-key", + } + request_body: Final[dict[str, JsonValue]] = {"guardrail_name": name, "litellm_params": params} + complete_body: Final[dict[str, JsonValue]] = { + **request_body, + **({"team_id": team_id} if team_id is not None else {}), + } + response: Final = gateway.request("POST", "/guardrails/register", complete_body, key=key) + assert response.status_code == 200, response.text + created: Final = _RegisterResponse.model_validate_json(response.content) + assert (created.guardrail_name, created.status) == (name, "pending_review"), response.text + scenario_cleanups.callback(_delete_guardrail, gateway, created.guardrail_id) + return created.guardrail_id + + +def _chat(gateway: Gateway, model: str, key: str, text: str, guardrail: str) -> _ChatResponse: + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": text}], "guardrails": [guardrail]}, + key=key, + ) + assert response.status_code == 200, response.text + parsed: Final = _ChatResponse.model_validate_json(response.content) + assert parsed.choices[0].message.content == "provider response", response.text + return parsed + + +def _assert_provider_request(request: Request, text: str) -> None: + assert (request.method, request.target) == ("POST", "/v1/chat/completions") + payload: Final = _JSON_OBJECT.validate_json(request.body) + assert payload == {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": text}]}, payload + + +def _assert_guardrail_request( + request: Request, + text: str, + model: str | None, + key: str, + *, + key_alias: str | None = None, +) -> None: + assert (request.method, request.target) == ("POST", "/beta/litellm_basic_guardrail_api") + payload: Final = _JSON_OBJECT.validate_json(request.body) + assert set(payload) == { + "input_type", + "litellm_call_id", + "litellm_trace_id", + "structured_messages", + "images", + "tools", + "texts", + "request_data", + "request_headers", + "litellm_version", + "additional_provider_specific_params", + "tool_calls", + "model", + }, payload + assert { + field: payload[field] + for field in ( + "input_type", + "structured_messages", + "images", + "tools", + "texts", + "additional_provider_specific_params", + "tool_calls", + "model", + ) + } == { + "input_type": "request", + "structured_messages": [{"role": "user", "content": text}] if model is not None else None, + "images": None, + "tools": None, + "texts": [text], + "additional_provider_specific_params": {}, + "tool_calls": None, + "model": model, + }, payload + assert payload["litellm_call_id"] is None or isinstance(payload["litellm_call_id"], str), payload + assert payload["litellm_trace_id"] is None or isinstance(payload["litellm_trace_id"], str), payload + assert payload["litellm_version"] is None or isinstance(payload["litellm_version"], str), payload + request_data: Final = object_value(payload["request_data"]) + if model is None: + assert request_data == {}, payload + else: + assert set(request_data).issubset( + { + "user_api_key_hash", + "user_api_key_alias", + "user_api_key_user_id", + "user_api_key_user_email", + "user_api_key_team_id", + "user_api_key_team_alias", + "user_api_key_end_user_id", + "user_api_key_org_id", + } + ), request_data + assert request_data["user_api_key_hash"] == sha256(key.encode()).hexdigest(), request_data + assert request_data.get("user_api_key_alias") == key_alias, request_data + assert all(value is None or isinstance(value, str) for value in request_data.values()), request_data + request_headers: Final = payload["request_headers"] + assert request_headers is None or isinstance(request_headers, dict), payload + if isinstance(request_headers, dict): + assert "authorization" not in request_headers + assert request.headers["x-api-key"] == "synthetic-submission-guardrail-key" + + +def test_non_admin_keys_scoped_to_the_submission_routes_get_403_from_the_admin_check(gateway: Gateway) -> None: + with ( + wire_server(lambda _request: Reply(body=json.dumps(_PROVIDER_RESPONSE).encode())) as provider, + wire_server(lambda _request: Reply(body=b'{"action":"NONE"}')) as guardrail, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=f"{provider.url}/v1", api_key="synthetic-provider-key") + team_id: Final = scenario.team(models=[model]) + member: Final = scenario.member(team_id) + team_admin: Final = scenario.member(team_id, role="admin") + registration_key: Final = scenario.key(user_id=member, team_id=team_id, models=[model]) + plain_member_key: Final = scenario.key(user_id=member, team_id=team_id, models=[model]) + plain_team_admin_key: Final = scenario.key(user_id=team_admin, team_id=team_id, models=[model]) + member_key: Final = scenario.key( + user_id=member, + team_id=team_id, + models=[model], + allowed_routes=["/guardrails/submissions"], + ) + team_admin_key: Final = scenario.key( + user_id=team_admin, + team_id=team_id, + models=[model], + allowed_routes=["/guardrails/submissions"], + ) + name: Final = f"submitted-auth-{uuid.uuid4().hex}" + guardrail_id: Final = _register( + gateway, + scenario.cleanups, + name, + team_id, + registration_key, + f"{guardrail.url}/beta/litellm_basic_guardrail_api", + ) + plain_responses: Final = ( + gateway.request("POST", f"/guardrails/submissions/{guardrail_id}/approve", key=plain_team_admin_key), + gateway.request("POST", f"/guardrails/submissions/{guardrail_id}/reject", key=plain_team_admin_key), + gateway.request("POST", f"/guardrails/submissions/{guardrail_id}/approve", key=plain_member_key), + gateway.request("POST", f"/guardrails/submissions/{guardrail_id}/reject", key=plain_member_key), + ) + expected_routes: Final = tuple( + f"/guardrails/submissions/{guardrail_id}/{action}" for action in ("approve", "reject", "approve", "reject") + ) + expected_user_ids: Final = (team_admin, team_admin, member, member) + expected_plain_errors: Final = tuple( + _ErrorResponse( + error=_ProxyError( + message=( + "Authentication Error, Only proxy admin can be used to generate, delete, update info for " + f"new keys/users/teams. Route={route}. Your role=internal_user. " + f"Your user_id={_mask_user_id(user_id)}" + ), + type="auth_error", + param="None", + code="401", + ) + ) + for route, user_id in zip(expected_routes, expected_user_ids, strict=True) + ) + route_check_refusals: Final = tuple((401, error.model_dump()) for error in expected_plain_errors) + handler_forbidden: Final = _SubmissionForbiddenResponse(detail="Admin access required").model_dump() + handler_refusals: Final = ((403, handler_forbidden),) * 4 + response_models: Final = {401: _ErrorResponse, 403: _SubmissionForbiddenResponse} + observed: Final = tuple( + ( + response.status_code, + ( + response_models[response.status_code].model_validate_json(response.content).model_dump() + if response.status_code in response_models + else response.text + ), + ) + for response in plain_responses + ) + assert observed in (route_check_refusals, handler_refusals), tuple( + response.text for response in plain_responses + ) + _assert_submission_status(guardrail_id, "pending_review", team_id) + responses: Final = ( + gateway.request("POST", f"/guardrails/submissions/{guardrail_id}/approve", key=team_admin_key), + gateway.request("POST", f"/guardrails/submissions/{guardrail_id}/reject", key=team_admin_key), + gateway.request("POST", f"/guardrails/submissions/{guardrail_id}/approve", key=member_key), + gateway.request("POST", f"/guardrails/submissions/{guardrail_id}/reject", key=member_key), + ) + _assert_submission_actions_are_forbidden(responses) + assert guardrail.drain() == () + response: Final = _chat(gateway, model, registration_key, _PROMPT_PENDING, name) + assert response.choices[0].message.content == "provider response" + assert guardrail.drain() == () + + +def test_team_admins_and_internal_users_get_403_on_guardrail_submission_actions(gateway: Gateway) -> None: + pytest.skip( + "BUG: team admins and internal users get 401 from the route check instead of the " + "documented 403 on guardrail submission approve and reject" + ) + with ( + wire_server(lambda _request: Reply(body=json.dumps(_PROVIDER_RESPONSE).encode())) as provider, + wire_server(lambda _request: Reply(body=b'{"action":"NONE"}')) as guardrail, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=f"{provider.url}/v1", api_key="synthetic-provider-key") + team_id: Final = scenario.team(models=[model]) + member: Final = scenario.member(team_id) + team_admin: Final = scenario.member(team_id, role="admin") + member_key: Final = scenario.key(user_id=member, team_id=team_id, models=[model]) + team_admin_key: Final = scenario.key(user_id=team_admin, team_id=team_id, models=[model]) + name: Final = f"submitted-auth-{uuid.uuid4().hex}" + guardrail_id: Final = _register( + gateway, + scenario.cleanups, + name, + team_id, + member_key, + f"{guardrail.url}/beta/litellm_basic_guardrail_api", + ) + responses: Final = ( + gateway.request("POST", f"/guardrails/submissions/{guardrail_id}/approve", key=team_admin_key), + gateway.request("POST", f"/guardrails/submissions/{guardrail_id}/reject", key=team_admin_key), + gateway.request("POST", f"/guardrails/submissions/{guardrail_id}/approve", key=member_key), + gateway.request("POST", f"/guardrails/submissions/{guardrail_id}/reject", key=member_key), + ) + _assert_submission_actions_are_forbidden(responses) + _assert_submission_status(guardrail_id, "pending_review", team_id) + assert guardrail.drain() == () + + +def test_team_guardrail_submissions_require_admin_approval( + gateway: Gateway, +) -> None: + with ( + wire_server(lambda _request: Reply(body=json.dumps(_PROVIDER_RESPONSE).encode())) as provider, + wire_server(lambda _request: Reply(body=b'{"action":"NONE"}')) as guardrail, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=f"{provider.url}/v1", api_key="synthetic-provider-key") + team_id: Final = scenario.team(models=[model]) + member: Final = scenario.member(team_id) + team_admin: Final = scenario.member(team_id, role="admin") + member_key_alias: Final = f"submission-member-{uuid.uuid4().hex}" + member_key: Final = scenario.key( + user_id=member, + team_id=team_id, + models=[model], + key_alias=member_key_alias, + ) + team_admin_key: Final = scenario.key( + user_id=team_admin, + team_id=team_id, + models=[model], + key_alias=f"submission-team-admin-{uuid.uuid4().hex}", + ) + other_team_id: Final = scenario.team(models=[model]) + api_base: Final = f"{guardrail.url}/beta/litellm_basic_guardrail_api" + first_name: Final = f"submitted-member-{uuid.uuid4().hex}" + second_name: Final = f"submitted-admin-{uuid.uuid4().hex}" + other_name: Final = f"submitted-other-team-{uuid.uuid4().hex}" + first_id: Final = _register(gateway, scenario.cleanups, first_name, team_id, member_key, api_base) + second_id: Final = _register(gateway, scenario.cleanups, second_name, None, team_admin_key, api_base) + other_id: Final = _register(gateway, scenario.cleanups, other_name, other_team_id, gateway.key, api_base) + + _assert_submission_status(first_id, "pending_review", team_id) + _assert_submission_status(second_id, "pending_review", team_id) + _assert_submission_status(other_id, "pending_review", other_team_id) + + pending_apply: Final = gateway.request( + "POST", "/apply_guardrail", {"guardrail_name": first_name, "text": _APPLY_TEXT} + ) + assert pending_apply.status_code == 404, pending_apply.text + assert _ErrorResponse.model_validate_json(pending_apply.content).model_dump() == { + "error": { + "message": ( + f"Guardrail '{first_name}' not found. " + "Please ensure the guardrail is configured in your LiteLLM proxy." + ), + "type": "internal_server_error", + "param": None, + "code": "404", + } + }, pending_apply.text + + pending_chat: Final = _chat(gateway, model, member_key, _PROMPT_PENDING, first_name) + assert len(pending_chat.choices) == 1 + assert guardrail.drain() == () + + _assert_submission_status(first_id, "pending_review", team_id) + + approved: Final = gateway.request("POST", f"/guardrails/submissions/{first_id}/approve") + assert approved.status_code == 200, approved.text + assert _AdminActionResponse.model_validate_json(approved.content).model_dump() == { + "guardrail_id": first_id, + "status": "active", + "message": "Guardrail approved", + }, approved.text + _assert_submission_status(first_id, "active", team_id) + + approved_chat: Final = _chat(gateway, model, member_key, _PROMPT_APPROVED, first_name) + assert len(approved_chat.choices) == 1 + active_apply: Final = gateway.request( + "POST", "/apply_guardrail", {"guardrail_name": first_name, "text": _APPLY_TEXT} + ) + assert active_apply.status_code == 200, active_apply.text + assert _ApplyResponse.model_validate_json(active_apply.content).response_text == _APPLY_TEXT, active_apply.text + + rejected: Final = gateway.request("POST", f"/guardrails/submissions/{second_id}/reject") + assert rejected.status_code == 200, rejected.text + assert _AdminActionResponse.model_validate_json(rejected.content).model_dump() == { + "guardrail_id": second_id, + "status": "rejected", + "message": "Guardrail rejected", + }, rejected.text + _assert_submission_status(second_id, "rejected", team_id) + rejected_chat: Final = _chat(gateway, model, member_key, _PROMPT_REJECTED, second_name) + assert len(rejected_chat.choices) == 1 + rejected_apply: Final = gateway.request( + "POST", "/apply_guardrail", {"guardrail_name": second_name, "text": _APPLY_TEXT} + ) + assert rejected_apply.status_code == 404, rejected_apply.text + assert _ErrorResponse.model_validate_json(rejected_apply.content).model_dump() == { + "error": { + "message": ( + f"Guardrail '{second_name}' not found. " + "Please ensure the guardrail is configured in your LiteLLM proxy." + ), + "type": "internal_server_error", + "param": None, + "code": "404", + } + }, rejected_apply.text + + team_list: Final = gateway.request("GET", "/guardrails/submissions", key=member_key) + assert team_list.status_code == 200, team_list.text + team_submissions: Final = _SubmissionsResponse.model_validate_json(team_list.content) + assert {item.guardrail_id for item in team_submissions.submissions} == {first_id, second_id}, team_list.text + admin_list: Final = gateway.request("GET", "/guardrails/submissions") + assert admin_list.status_code == 200, admin_list.text + admin_submissions: Final = _SubmissionsResponse.model_validate_json(admin_list.content) + assert {first_id, second_id, other_id} <= {item.guardrail_id for item in admin_submissions.submissions}, ( + admin_list.text + ) + + provider_requests: Final = provider.drain() + assert [(request.method, request.target) for request in provider_requests] == [ + ("POST", "/v1/chat/completions"), + ("POST", "/v1/chat/completions"), + ("POST", "/v1/chat/completions"), + ] + for request, text in zip(provider_requests, (_PROMPT_PENDING, _PROMPT_APPROVED, _PROMPT_REJECTED), strict=True): + _assert_provider_request(request, text) + guardrail_requests: Final = guardrail.drain() + assert [(request.method, request.target) for request in guardrail_requests] == [ + ("POST", "/beta/litellm_basic_guardrail_api"), + ("POST", "/beta/litellm_basic_guardrail_api"), + ] + _assert_guardrail_request( + guardrail_requests[0], + _PROMPT_APPROVED, + model, + member_key, + key_alias=member_key_alias, + ) + _assert_guardrail_request(guardrail_requests[1], _APPLY_TEXT, None, gateway.key) diff --git a/tests/integration/management/test_policy_attachment_runtime.py b/tests/integration/management/test_policy_attachment_runtime.py new file mode 100644 index 00000000000..adf5b57dc9c --- /dev/null +++ b/tests/integration/management/test_policy_attachment_runtime.py @@ -0,0 +1,1297 @@ +import json +import uuid +from contextlib import ExitStack +from hashlib import sha256 +from typing import Final, Literal + +import httpx +import pytest +from integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from openai import OpenAI +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MASKED: Final = "masked integration prompt" +_MASKED_FIRST: Final = "masked original policy prompt" +_MASKED_SECOND: Final = "masked published policy prompt" +_PROMPT: Final = "integration prompt containing a secret" +_PROVIDER_RESPONSE: Final = { + "id": "chatcmpl-policy-attachment", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "provider response"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +class _GuardrailCreated(BaseModel): + guardrail_id: str + guardrail_name: str + + +class _PolicyResponse(BaseModel): + policy_id: str + policy_name: str + version_number: int + version_status: str + guardrails_add: list[str] + guardrails_remove: list[str] + + +class _PolicyVersionsResponse(BaseModel): + policy_name: str + versions: list[_PolicyResponse] + total_count: int + + +class _AttachmentResponse(BaseModel): + attachment_id: str + policy_name: str + scope: str | None + teams: list[str] + keys: list[str] + models: list[str] + tags: list[str] + priority: int | None + default: bool + + +class _ChatMessage(BaseModel): + role: str + content: str | None + + +class _ChatChoice(BaseModel): + index: int + message: _ChatMessage + finish_reason: str | None + + +class _ChatResponse(BaseModel): + choices: list[_ChatChoice] + + +def _delete_guardrail(gateway: Gateway, guardrail_id: str) -> None: + deleted: Final = gateway.request("DELETE", f"/guardrails/{guardrail_id}") + assert deleted.status_code == 200, deleted.text + assert read_rows('SELECT guardrail_id FROM "LiteLLM_GuardrailsTable" WHERE guardrail_id=%s', (guardrail_id,)) == [] + + +def _delete_policy(gateway: Gateway, policy_id: str) -> None: + deleted: Final = gateway.request("DELETE", f"/policies/{policy_id}") + assert deleted.status_code == 200, deleted.text + assert read_rows('SELECT policy_id FROM "LiteLLM_PolicyTable" WHERE policy_id=%s', (policy_id,)) == [] + + +def _create_guardrail( + gateway: Gateway, + cleanups: ExitStack, + guardrail_name: str, + api_base: str, + api_key: str, +) -> str: + created: Final = gateway.request( + "POST", + "/guardrails", + { + "guardrail": { + "guardrail_name": guardrail_name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "api_base": api_base, + "api_key": api_key, + }, + } + }, + ) + assert created.status_code == 200, created.text + payload: Final = _GuardrailCreated.model_validate_json(created.content) + cleanups.callback(_delete_guardrail, gateway, payload.guardrail_id) + return payload.guardrail_name + + +def _create_policy(gateway: Gateway, cleanups: ExitStack, policy_name: str, guardrail_name: str) -> str: + created: Final = gateway.request( + "POST", + "/policies", + {"policy_name": policy_name, "guardrails_add": [guardrail_name]}, + ) + assert created.status_code == 200, created.text + payload: Final = _PolicyResponse.model_validate_json(created.content) + cleanups.callback(_delete_policy, gateway, payload.policy_id) + assert (payload.policy_name, payload.version_status, payload.guardrails_add) == ( + policy_name, + "production", + [guardrail_name], + ), created.text + return payload.policy_id + + +def _create_published_version( + gateway: Gateway, + cleanups: ExitStack, + policy_name: str, + source_policy_id: str, + guardrail_name: str, +) -> str: + create_version: Final = gateway.request( + "POST", f"/policies/name/{policy_name}/versions", {"source_policy_id": source_policy_id} + ) + assert create_version.status_code == 200, create_version.text + draft: Final = _PolicyResponse.model_validate_json(create_version.content) + cleanups.callback(_delete_policy, gateway, draft.policy_id) + update_version: Final = gateway.request( + "PUT", + f"/policies/{draft.policy_id}", + {"guardrails_add": [guardrail_name], "guardrails_remove": []}, + ) + assert update_version.status_code == 200, update_version.text + assert _PolicyResponse.model_validate_json(update_version.content).guardrails_add == [guardrail_name], ( + update_version.text + ) + publish_version: Final = gateway.request( + "PUT", f"/policies/{draft.policy_id}/status", {"version_status": "published"} + ) + assert publish_version.status_code == 200, publish_version.text + assert _PolicyResponse.model_validate_json(publish_version.content).version_status == "published", ( + publish_version.text + ) + versions_response: Final = gateway.request("GET", f"/policies/name/{policy_name}/versions") + assert versions_response.status_code == 200, versions_response.text + versions: Final = _PolicyVersionsResponse.model_validate_json(versions_response.content) + assert versions.policy_name == policy_name, versions_response.text + second_version: Final = next(version for version in versions.versions if version.policy_id == draft.policy_id) + assert second_version.guardrails_add == [guardrail_name], versions_response.text + assert second_version.version_status == "published", versions_response.text + return second_version.policy_id + + +def _assert_guardrail_request( + request: Request, + text: str, + model: str, + key: str, + *, + key_alias: str | None = None, + response: httpx.Response | None = None, +) -> None: + assert request.method == "POST" + assert request.target == "/beta/litellm_basic_guardrail_api" + payload: Final = _JSON_OBJECT.validate_json(request.body) + assert set(payload) == { + "input_type", + "litellm_call_id", + "litellm_trace_id", + "structured_messages", + "images", + "tools", + "texts", + "request_data", + "request_headers", + "litellm_version", + "additional_provider_specific_params", + "tool_calls", + "model", + }, payload + assert { + key: payload[key] + for key in ( + "input_type", + "structured_messages", + "images", + "tools", + "texts", + "additional_provider_specific_params", + "tool_calls", + "model", + ) + } == { + "input_type": "request", + "structured_messages": [{"role": "user", "content": text}], + "images": None, + "tools": None, + "texts": [text], + "additional_provider_specific_params": {}, + "tool_calls": None, + "model": model, + }, payload + assert payload["litellm_call_id"] is None or isinstance(payload["litellm_call_id"], str), payload + assert payload["litellm_trace_id"] is None or isinstance(payload["litellm_trace_id"], str), payload + assert payload["litellm_version"] is None or isinstance(payload["litellm_version"], str), payload + if response is not None and response.headers.get("x-litellm-call-id") is not None: + assert payload["litellm_call_id"] == response.headers["x-litellm-call-id"], payload + request_data: Final = object_value(payload["request_data"]) + assert set(request_data).issubset( + { + "user_api_key_hash", + "user_api_key_alias", + "user_api_key_user_id", + "user_api_key_user_email", + "user_api_key_team_id", + "user_api_key_team_alias", + "user_api_key_end_user_id", + "user_api_key_org_id", + } + ), request_data + assert request_data["user_api_key_hash"] == sha256(key.encode()).hexdigest(), request_data + assert request_data.get("user_api_key_alias") == key_alias, request_data + assert all(value is None or isinstance(value, str) for value in request_data.values()), request_data + request_headers: Final = payload["request_headers"] + assert request_headers is None or isinstance(request_headers, dict), payload + if isinstance(request_headers, dict): + assert "authorization" not in request_headers + + +def _assert_provider_request(request: Request, text: str) -> None: + assert request.method == "POST" + assert request.target == "/v1/chat/completions" + payload: Final = _JSON_OBJECT.validate_json(request.body) + assert payload == { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": text}], + }, payload + + +def _mask_policy_prompt(request: Request, mask: str) -> str: + payload: Final = _JSON_OBJECT.validate_json(request.body) + texts: Final = payload["texts"] + assert isinstance(texts, list) and len(texts) == 1, payload + return f"{mask}: {string_value(texts[0])}" + + +def _policy_guardrail_reply(request: Request) -> Reply: + body: Final = { + "action": "GUARDRAIL_INTERVENED", + "texts": [_mask_policy_prompt(request, _MASKED)], + } + return Reply(body=json.dumps(body).encode()) + + +def _chat( + gateway: Gateway, + model: str, + key: str, + text: str, + *, + policies: list[str] | None = None, +) -> httpx.Response: + body: Final = _JSON_OBJECT.validate_python( + { + "model": model, + "messages": [{"role": "user", "content": text}], + **({} if policies is None else {"policies": policies}), + } + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + body, + key=key, + ) + assert response.status_code == 200, response.text + _ChatResponse.model_validate_json(response.content) + return response + + +def _sse(events: tuple[dict[str, JsonValue], ...]) -> tuple[bytes, ...]: + return tuple(f"data: {json.dumps(event)}\n\n".encode() for event in events) + (b"data: [DONE]\n\n",) + + +def _streaming_chat_reply(_request: Request) -> Reply: + identity: Final = f"chatcmpl-policy-{uuid.uuid4().hex}" + chunk: Final[dict[str, JsonValue]] = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + } + events: Final[tuple[dict[str, JsonValue], ...]] = ( + {**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "provider "}}]}, + {**chunk, "choices": [{"index": 0, "delta": {"content": "response"}}]}, + {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + ) + return Reply(content_type="text/event-stream", chunks=_sse(events)) + + +def _stream_chat( + gateway: Gateway, + model: str, + key: str, + text: str, + *, + policy_name: str | None = None, +) -> tuple[httpx.Headers, str]: + with httpx.Client(trust_env=False, timeout=15) as http_client: + with OpenAI( + base_url=f"{str(gateway.client.base_url).rstrip('/')}/v1", + api_key=key, + http_client=http_client, + ) as client: + raw_response: Final = client.chat.completions.with_raw_response.create( + model=model, + messages=[{"role": "user", "content": text}], + stream=True, + extra_body=None if policy_name is None else {"policies": [policy_name]}, + ) + assert raw_response.status_code == 200, raw_response.text + with raw_response.parse() as stream: + response_text: Final = "".join( + chunk.choices[0].delta.content or "" for chunk in stream if chunk.choices + ) + return raw_response.headers, response_text + + +def _assert_streaming_provider_request(request: Request, text: str) -> None: + assert request.method == "POST" + assert request.target == "/v1/chat/completions" + payload: Final = _JSON_OBJECT.validate_json(request.body) + assert payload == { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": text}], + "stream": True, + "stream_options": {"include_usage": True}, + }, payload + + +def _attachment_scope( + spelling: str, + policy_name: str, + matching_model: str, + prefix: str, +) -> dict[str, JsonValue]: + return { + "policy_name": policy_name, + "models": [matching_model] if spelling == "models" else [], + "teams": [f"{prefix}-*"] if spelling == "teams" else [], + "keys": [f"{prefix}-*"] if spelling == "keys" else [], + "scope": "*" if spelling == "*" else None, + } + + +def _matching_and_nonmatching_keys( + scenario: Scenario, + spelling: str, + prefix: str, + model: str, + second_model: str, +) -> tuple[str, str | None]: + if spelling == "teams": + matching_team: Final = scenario.team(team_alias=f"{prefix}-eng", models=[model]) + nonmatching_team: Final = scenario.team(team_alias=f"other-{prefix}-team", models=[model]) + return ( + scenario.key(team_id=matching_team, models=[model], key_alias=f"{prefix}-key"), + scenario.key(team_id=nonmatching_team, models=[model], key_alias=f"other-{prefix}-key"), + ) + if spelling == "keys": + team: Final = scenario.team(models=[model]) + return ( + scenario.key(team_id=team, models=[model], key_alias=f"{prefix}-key"), + scenario.key(team_id=team, models=[model], key_alias=f"other-{prefix}-key"), + ) + return scenario.key(models=[model, second_model]), None + + +@pytest.mark.parametrize("spelling", ("models", "teams", "keys", "*")) +def test_policy_attachments_match_live_models_teams_keys_and_global_scope(gateway: Gateway, spelling: str) -> None: + with ( + wire_server(lambda _request: Reply(body=json.dumps(_PROVIDER_RESPONSE).encode())) as provider, + wire_server( + lambda _request: Reply(body=b'{"action":"GUARDRAIL_INTERVENED","texts":["masked integration prompt"]}') + ) as guardrail, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=f"{provider.url}/v1", api_key="synthetic-provider-key") + second_model: Final = scenario.model(api_base=f"{provider.url}/v1", api_key="synthetic-provider-key") + prefix: Final = f"attachment-{uuid.uuid4().hex}" + guardrail_name: Final = _create_guardrail( + gateway, + scenario.cleanups, + f"guardrail-{prefix}", + f"{guardrail.url}/beta/litellm_basic_guardrail_api", + "synthetic-guardrail-key", + ) + policy_name: Final = f"policy-{prefix}" + _create_policy(gateway, scenario.cleanups, policy_name, guardrail_name) + active_scope: Final = _attachment_scope(spelling, policy_name, model, prefix) + matching_key, nonmatching_key = _matching_and_nonmatching_keys(scenario, spelling, prefix, model, second_model) + attachment: Final = gateway.request("POST", "/policies/attachments", active_scope) + assert attachment.status_code == 200, attachment.text + created_attachment: Final = _AttachmentResponse.model_validate_json(attachment.content) + scenario.cleanups.callback(_delete_attachment, gateway, created_attachment.attachment_id) + expected_scope: Final = { + "attachment_id": created_attachment.attachment_id, + "policy_name": policy_name, + "scope": "*" if spelling == "*" else None, + "teams": [f"{prefix}-*"] if spelling == "teams" else [], + "keys": [f"{prefix}-*"] if spelling == "keys" else [], + "models": [model] if spelling == "models" else [], + "tags": [], + "priority": None, + "default": False, + } + assert created_attachment.model_dump() == expected_scope, attachment.text + + readback: Final = gateway.request("GET", f"/policies/attachments/{created_attachment.attachment_id}") + assert readback.status_code == 200, readback.text + assert _AttachmentResponse.model_validate_json(readback.content).model_dump() == expected_scope, readback.text + rows: Final = read_rows( + 'SELECT policy_name, scope, teams, keys, models FROM "LiteLLM_PolicyAttachmentTable" ' + "WHERE attachment_id=%s", + (created_attachment.attachment_id,), + ) + assert rows == [ + { + "policy_name": policy_name, + "scope": "*" if spelling == "*" else None, + "teams": [f"{prefix}-*"] if spelling == "teams" else [], + "keys": [f"{prefix}-*"] if spelling == "keys" else [], + "models": [model] if spelling == "models" else [], + } + ] + + matching_response: Final = _chat(gateway, model, matching_key, _PROMPT) + assert matching_response.headers.get("x-litellm-applied-policies") == policy_name, matching_response.text + global_response: Final = _chat(gateway, second_model, matching_key, _PROMPT) if spelling == "*" else None + if global_response is not None: + assert global_response.headers.get("x-litellm-applied-policies") == policy_name, global_response.text + elif spelling == "models": + nonmatching_model_response: Final = _chat(gateway, second_model, matching_key, _PROMPT) + assert nonmatching_model_response.headers.get("x-litellm-applied-policies") is None, ( + nonmatching_model_response.text + ) + elif nonmatching_key is not None: + nonmatching_key_response: Final = _chat(gateway, model, nonmatching_key, _PROMPT) + assert nonmatching_key_response.headers.get("x-litellm-applied-policies") is None, ( + nonmatching_key_response.text + ) + global_responses: Final[tuple[httpx.Response, ...]] = () if global_response is None else (global_response,) + + guardrail_requests: Final = guardrail.drain() + provider_requests: Final = provider.drain() + expected_matches: Final = 2 if spelling == "*" else 1 + assert len(guardrail_requests) == expected_matches + assert len(provider_requests) == 2 + expected_guardrail_models: Final = (model, second_model) if spelling == "*" else (model,) + matching_responses: Final = (matching_response,) + global_responses + expected_key_alias: Final = f"{prefix}-key" if spelling in ("teams", "keys") else None + for request, expected_model, response in zip( + guardrail_requests, + expected_guardrail_models, + matching_responses, + strict=True, + ): + _assert_guardrail_request( + request, + _PROMPT, + expected_model, + matching_key, + key_alias=expected_key_alias, + response=response, + ) + assert [(request.method, request.target) for request in guardrail_requests] == [ + ("POST", "/beta/litellm_basic_guardrail_api") + ] * expected_matches + expected_provider_texts: Final = ( + (_MASKED, _MASKED) + if spelling == "*" + else (_MASKED, _PROMPT) + if spelling in ("models", "teams", "keys") + else (_MASKED,) + ) + assert [(request.method, request.target) for request in provider_requests] == [ + ("POST", "/v1/chat/completions") + ] * len(expected_provider_texts) + for request, expected_text in zip(provider_requests, expected_provider_texts, strict=True): + _assert_provider_request(request, expected_text) + + +@pytest.mark.parametrize( + ("model_priority", "global_priority"), + ( + pytest.param(10, 20, id="model-priority-first"), + pytest.param(20, 10, id="global-priority-first"), + ), +) +def test_matching_model_and_global_attachments_apply_by_priority( + gateway: Gateway, + model_priority: int, + global_priority: int, +) -> None: + with ( + wire_server(lambda _request: Reply(body=json.dumps(_PROVIDER_RESPONSE).encode())) as provider, + wire_server(lambda _request: Reply(body=b'{"action":"NONE"}')) as model_guardrail, + wire_server(lambda _request: Reply(body=b'{"action":"NONE"}')) as global_guardrail, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=f"{provider.url}/v1", api_key="synthetic-provider-key") + other_model: Final = scenario.model(api_base=f"{provider.url}/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model, other_model]) + prefix: Final = f"attachment-{uuid.uuid4().hex}" + model_guardrail_key: Final = "synthetic-model-guardrail-key" + global_guardrail_key: Final = "synthetic-global-guardrail-key" + model_guardrail_name: Final = _create_guardrail( + gateway, + scenario.cleanups, + f"guardrail-model-{prefix}", + f"{model_guardrail.url}/beta/litellm_basic_guardrail_api", + model_guardrail_key, + ) + global_guardrail_name: Final = _create_guardrail( + gateway, + scenario.cleanups, + f"guardrail-global-{prefix}", + f"{global_guardrail.url}/beta/litellm_basic_guardrail_api", + global_guardrail_key, + ) + model_policy_name: Final = f"policy-model-{prefix}" + global_policy_name: Final = f"policy-global-{prefix}" + _create_policy(gateway, scenario.cleanups, model_policy_name, model_guardrail_name) + _create_policy(gateway, scenario.cleanups, global_policy_name, global_guardrail_name) + model_attachment_response: Final = gateway.request( + "POST", + "/policies/attachments", + {"policy_name": model_policy_name, "models": [model], "priority": model_priority}, + ) + assert model_attachment_response.status_code == 200, model_attachment_response.text + model_attachment: Final = _AttachmentResponse.model_validate_json(model_attachment_response.content) + scenario.cleanups.callback(_delete_attachment, gateway, model_attachment.attachment_id) + global_attachment_response: Final = gateway.request( + "POST", + "/policies/attachments", + {"policy_name": global_policy_name, "scope": "*", "priority": global_priority}, + ) + assert global_attachment_response.status_code == 200, global_attachment_response.text + global_attachment: Final = _AttachmentResponse.model_validate_json(global_attachment_response.content) + scenario.cleanups.callback(_delete_attachment, gateway, global_attachment.attachment_id) + + expected_model_attachment: Final = { + "attachment_id": model_attachment.attachment_id, + "policy_name": model_policy_name, + "scope": None, + "teams": [], + "keys": [], + "models": [model], + "tags": [], + "priority": model_priority, + "default": False, + } + expected_global_attachment: Final = { + "attachment_id": global_attachment.attachment_id, + "policy_name": global_policy_name, + "scope": "*", + "teams": [], + "keys": [], + "models": [], + "tags": [], + "priority": global_priority, + "default": False, + } + assert model_attachment.model_dump() == expected_model_attachment, model_attachment_response.text + assert global_attachment.model_dump() == expected_global_attachment, global_attachment_response.text + for attachment_id, expected_attachment in ( + (model_attachment.attachment_id, expected_model_attachment), + (global_attachment.attachment_id, expected_global_attachment), + ): + readback: Final = gateway.request("GET", f"/policies/attachments/{attachment_id}") + assert readback.status_code == 200, readback.text + assert _AttachmentResponse.model_validate_json(readback.content).model_dump() == expected_attachment, ( + readback.text + ) + + expected_order: Final = ( + (model_policy_name, global_policy_name) + if model_priority < global_priority + else (global_policy_name, model_policy_name) + ) + matched: Final = _chat(gateway, model, key, _PROMPT) + assert matched.headers.get("x-litellm-applied-policies") == ",".join(expected_order), matched.text + matched_model_guardrail_requests: Final = model_guardrail.drain() + matched_global_guardrail_requests: Final = global_guardrail.drain() + assert len(matched_model_guardrail_requests) == 1 + assert len(matched_global_guardrail_requests) == 1 + _assert_guardrail_request(matched_model_guardrail_requests[0], _PROMPT, model, key, response=matched) + _assert_guardrail_request(matched_global_guardrail_requests[0], _PROMPT, model, key, response=matched) + assert matched_model_guardrail_requests[0].headers["x-api-key"] == model_guardrail_key + assert matched_global_guardrail_requests[0].headers["x-api-key"] == global_guardrail_key + matched_provider_requests: Final = provider.drain() + assert len(matched_provider_requests) == 1 + _assert_provider_request(matched_provider_requests[0], _PROMPT) + + unmatched: Final = _chat(gateway, other_model, key, _PROMPT) + assert unmatched.headers.get("x-litellm-applied-policies") == global_policy_name, unmatched.text + assert model_guardrail.drain() == () + unmatched_global_guardrail_requests: Final = global_guardrail.drain() + assert len(unmatched_global_guardrail_requests) == 1 + _assert_guardrail_request(unmatched_global_guardrail_requests[0], _PROMPT, other_model, key, response=unmatched) + assert unmatched_global_guardrail_requests[0].headers["x-api-key"] == global_guardrail_key + unmatched_provider_requests: Final = provider.drain() + assert len(unmatched_provider_requests) == 1 + _assert_provider_request(unmatched_provider_requests[0], _PROMPT) + + +def _delete_attachment(gateway: Gateway, attachment_id: str) -> None: + deleted: Final = gateway.request("DELETE", f"/policies/attachments/{attachment_id}") + assert deleted.status_code == 200, deleted.text + assert ( + read_rows('SELECT attachment_id FROM "LiteLLM_PolicyAttachmentTable" WHERE attachment_id=%s', (attachment_id,)) + == [] + ) + + +def test_request_policy_name_selects_its_production_guardrail_and_is_removed_from_provider_body( + gateway: Gateway, +) -> None: + with ( + wire_server(lambda _request: Reply(body=json.dumps(_PROVIDER_RESPONSE).encode())) as provider, + wire_server( + lambda request: Reply( + body=json.dumps( + { + "action": "GUARDRAIL_INTERVENED", + "texts": [_mask_policy_prompt(request, _MASKED_FIRST)], + } + ).encode() + ) + ) as guardrail, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=f"{provider.url}/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + policy_name: Final = f"policy-{uuid.uuid4().hex}" + first_guardrail: Final = _create_guardrail( + gateway, + scenario.cleanups, + f"guardrail-one-{uuid.uuid4().hex}", + f"{guardrail.url}/beta/litellm_basic_guardrail_api", + "guardrail-one", + ) + second_guardrail: Final = _create_guardrail( + gateway, + scenario.cleanups, + f"guardrail-two-{uuid.uuid4().hex}", + f"{guardrail.url}/beta/litellm_basic_guardrail_api", + "guardrail-two", + ) + policy_id: Final = _create_policy(gateway, scenario.cleanups, policy_name, first_guardrail) + _create_published_version(gateway, scenario.cleanups, policy_name, policy_id, second_guardrail) + + response: Final = _chat(gateway, model, key, _PROMPT, policies=[policy_name]) + assert response.headers.get("x-litellm-applied-policies") == policy_name, response.text + + sdk_prompt: Final = f"{_PROMPT} via SDK {uuid.uuid4().hex}" + with httpx.Client(trust_env=False, timeout=15) as client: + sdk: Final = OpenAI( + base_url=f"{str(gateway.client.base_url).rstrip('/')}/v1", + api_key=key, + http_client=client, + ) + sdk_response: Final = sdk.chat.completions.create( + model=model, + messages=[{"role": "user", "content": sdk_prompt}], + extra_body={"policies": [policy_name]}, + ) + assert sdk_response.choices[0].message.content == "provider response" + guardrail_requests: Final = guardrail.drain() + provider_requests: Final = provider.drain() + assert len(guardrail_requests) == 2 + assert len(provider_requests) == 2 + assert [request.headers["x-api-key"] for request in guardrail_requests] == ["guardrail-one"] * 2 + _assert_guardrail_request( + guardrail_requests[0], + _PROMPT, + model, + key, + response=response, + ) + _assert_guardrail_request(guardrail_requests[1], sdk_prompt, model, key) + expected_provider_texts: Final = (f"{_MASKED_FIRST}: {_PROMPT}", f"{_MASKED_FIRST}: {sdk_prompt}") + for request, expected_text in zip(provider_requests, expected_provider_texts, strict=True): + _assert_provider_request(request, expected_text) + assert [(request.method, request.target) for request in guardrail_requests] == [ + ("POST", "/beta/litellm_basic_guardrail_api") + ] * len(guardrail_requests) + assert [(request.method, request.target) for request in provider_requests] == [ + ("POST", "/v1/chat/completions") + ] * len(provider_requests) + + +def test_request_policy_version_id_selects_its_published_guardrail(gateway: Gateway) -> None: + pytest.skip( + "BUG: a just-published policy version id in request-body policies applies no guardrail until a DB resync" + ) + with ( + wire_server(lambda _request: Reply(body=json.dumps(_PROVIDER_RESPONSE).encode())) as provider, + wire_server( + lambda request: Reply( + body=json.dumps( + { + "action": "GUARDRAIL_INTERVENED", + "texts": [ + _MASKED_FIRST if request.headers.get("x-api-key") == "guardrail-one" else _MASKED_SECOND + ], + } + ).encode() + ) + ) as guardrail, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=f"{provider.url}/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + policy_name: Final = f"policy-{uuid.uuid4().hex}" + first_guardrail: Final = _create_guardrail( + gateway, + scenario.cleanups, + f"guardrail-one-{uuid.uuid4().hex}", + f"{guardrail.url}/beta/litellm_basic_guardrail_api", + "guardrail-one", + ) + second_guardrail: Final = _create_guardrail( + gateway, + scenario.cleanups, + f"guardrail-two-{uuid.uuid4().hex}", + f"{guardrail.url}/beta/litellm_basic_guardrail_api", + "guardrail-two", + ) + policy_id: Final = _create_policy(gateway, scenario.cleanups, policy_name, first_guardrail) + version_id: Final = _create_published_version( + gateway, scenario.cleanups, policy_name, policy_id, second_guardrail + ) + + response: Final = _chat(gateway, model, key, _PROMPT, policies=[f"policy_{version_id}"]) + assert response.headers.get("x-litellm-applied-policies") == policy_name, response.text + guardrail_requests: Final = guardrail.drain() + assert len(guardrail_requests) == 1 + assert guardrail_requests[0].headers["x-api-key"] == "guardrail-two" + _assert_guardrail_request(guardrail_requests[0], _PROMPT, model, key, response=response) + provider_requests: Final = provider.drain() + assert len(provider_requests) == 1 + _assert_provider_request(provider_requests[0], _MASKED_SECOND) + assert [(request.method, request.target) for request in guardrail_requests] == [ + ("POST", "/beta/litellm_basic_guardrail_api") + ] + assert [(request.method, request.target) for request in provider_requests] == [("POST", "/v1/chat/completions")] + + +def test_request_published_policy_version_id_resolves_after_force_sync(gateway: Gateway) -> None: + with ( + wire_server(lambda _request: Reply(body=json.dumps(_PROVIDER_RESPONSE).encode())) as provider, + wire_server( + lambda request: Reply( + body=json.dumps( + { + "action": "GUARDRAIL_INTERVENED", + "texts": [ + _MASKED_FIRST if request.headers.get("x-api-key") == "guardrail-one" else _MASKED_SECOND + ], + } + ).encode() + ) + ) as guardrail, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=f"{provider.url}/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + policy_name: Final = f"policy-{uuid.uuid4().hex}" + first_guardrail: Final = _create_guardrail( + gateway, + scenario.cleanups, + f"guardrail-one-{uuid.uuid4().hex}", + f"{guardrail.url}/beta/litellm_basic_guardrail_api", + "guardrail-one", + ) + second_guardrail: Final = _create_guardrail( + gateway, + scenario.cleanups, + f"guardrail-two-{uuid.uuid4().hex}", + f"{guardrail.url}/beta/litellm_basic_guardrail_api", + "guardrail-two", + ) + policy_id: Final = _create_policy(gateway, scenario.cleanups, policy_name, first_guardrail) + version_id: Final = _create_published_version( + gateway, scenario.cleanups, policy_name, policy_id, second_guardrail + ) + + synchronized: Final = gateway.request( + "POST", + "/policies/resolve?force_sync=true", + {"model": model}, + ) + assert synchronized.status_code == 200, synchronized.text + + response: Final = _chat(gateway, model, key, _PROMPT, policies=[f"policy_{version_id}"]) + assert response.headers.get("x-litellm-applied-policies") == policy_name, response.text + guardrail_requests: Final = guardrail.drain() + assert [(request.method, request.target) for request in guardrail_requests] == [ + ("POST", "/beta/litellm_basic_guardrail_api") + ] + assert [request.headers["x-api-key"] for request in guardrail_requests] == ["guardrail-two"] + _assert_guardrail_request(guardrail_requests[0], _PROMPT, model, key, response=response) + provider_requests: Final = provider.drain() + assert [(request.method, request.target) for request in provider_requests] == [("POST", "/v1/chat/completions")] + _assert_provider_request(provider_requests[0], _MASKED_SECOND) + + +class _ChatErrorBody(BaseModel): + model_config = ConfigDict(extra="forbid") + + message: str + type: str | None + param: str | None + code: JsonValue + + +class _ChatError(BaseModel): + model_config = ConfigDict(extra="forbid") + + error: _ChatErrorBody + + +def test_policy_attachment_with_a_blocking_guardrail_stops_the_request_before_the_provider( + gateway: Gateway, +) -> None: + blocked_reason: Final = f"policy attachment blocked {uuid.uuid4().hex}" + with ( + wire_server(lambda _request: Reply(body=json.dumps(_PROVIDER_RESPONSE).encode())) as provider, + wire_server( + lambda _request: Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": blocked_reason}).encode()) + ) as guardrail, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=f"{provider.url}/v1", api_key="synthetic-provider-key") + other_model: Final = scenario.model(api_base=f"{provider.url}/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model, other_model]) + prefix: Final = f"attachment-{uuid.uuid4().hex}" + guardrail_name: Final = _create_guardrail( + gateway, + scenario.cleanups, + f"guardrail-{prefix}", + f"{guardrail.url}/beta/litellm_basic_guardrail_api", + "synthetic-guardrail-key", + ) + policy_name: Final = f"policy-{prefix}" + _create_policy(gateway, scenario.cleanups, policy_name, guardrail_name) + attachment: Final = gateway.request( + "POST", "/policies/attachments", {"policy_name": policy_name, "models": [model]} + ) + assert attachment.status_code == 200, attachment.text + created_attachment: Final = _AttachmentResponse.model_validate_json(attachment.content) + scenario.cleanups.callback(_delete_attachment, gateway, created_attachment.attachment_id) + + blocked: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": _PROMPT}]}, + key=key, + ) + assert blocked.status_code == 400, blocked.text + error: Final = _ChatError.model_validate_json(blocked.content) + assert error.model_dump() == { + "error": { + "message": blocked_reason, + "type": "invalid_request_error", + "param": None, + "code": "400", + } + }, blocked.text + + allowed: Final = _chat(gateway, other_model, key, _PROMPT) + assert allowed.headers.get("x-litellm-applied-policies") is None, allowed.text + + guardrail_requests: Final = guardrail.drain() + assert len(guardrail_requests) == 1 + _assert_guardrail_request(guardrail_requests[0], _PROMPT, model, key, response=blocked) + provider_requests: Final = provider.drain() + assert [(request.method, request.target) for request in provider_requests] == [("POST", "/v1/chat/completions")] + _assert_provider_request(provider_requests[0], _PROMPT) + + +@pytest.mark.parametrize( + "tag_spelling", + ( + pytest.param("metadata_tags", id="metadata_tags"), + pytest.param("header", id="header"), + pytest.param("top_level_tags", id="top_level_tags"), + ), +) +def test_policy_attachment_tag_scope_matches_only_tagged_requests( + gateway: Gateway, + tag_spelling: Literal["metadata_tags", "header", "top_level_tags"], +) -> None: + with ( + wire_server( + lambda _request: Reply( + body=json.dumps({**_PROVIDER_RESPONSE, "id": f"chatcmpl-{uuid.uuid4().hex}"}).encode() + ) + ) as provider, + wire_server( + lambda _request: Reply(body=b'{"action":"GUARDRAIL_INTERVENED","texts":["masked integration prompt"]}') + ) as guardrail, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=f"{provider.url}/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + prefix: Final = f"attachment-{uuid.uuid4().hex}" + guardrail_name: Final = _create_guardrail( + gateway, + scenario.cleanups, + f"guardrail-{prefix}", + f"{guardrail.url}/beta/litellm_basic_guardrail_api", + "synthetic-guardrail-key", + ) + policy_name: Final = f"policy-{prefix}" + _create_policy(gateway, scenario.cleanups, policy_name, guardrail_name) + attachment: Final = gateway.request( + "POST", "/policies/attachments", {"policy_name": policy_name, "tags": [f"{prefix}-*"]} + ) + assert attachment.status_code == 200, attachment.text + created_attachment: Final = _AttachmentResponse.model_validate_json(attachment.content) + scenario.cleanups.callback(_delete_attachment, gateway, created_attachment.attachment_id) + expected_scope: Final = { + "attachment_id": created_attachment.attachment_id, + "policy_name": policy_name, + "scope": None, + "teams": [], + "keys": [], + "models": [], + "tags": [f"{prefix}-*"], + "priority": None, + "default": False, + } + assert created_attachment.model_dump() == expected_scope, attachment.text + readback: Final = gateway.request("GET", f"/policies/attachments/{created_attachment.attachment_id}") + assert readback.status_code == 200, readback.text + assert _AttachmentResponse.model_validate_json(readback.content).model_dump() == expected_scope, readback.text + rows: Final = read_rows( + 'SELECT policy_name, tags, is_default FROM "LiteLLM_PolicyAttachmentTable" WHERE attachment_id=%s', + (created_attachment.attachment_id,), + ) + assert rows == [{"policy_name": policy_name, "tags": [f"{prefix}-*"], "is_default": False}] + + tag: Final = f"{prefix}-prod" + unmatched_tag: Final = f"unmatched-{uuid.uuid4().hex}" + header_tags: Final = [unmatched_tag, tag] + tagged_body: Final[dict[str, JsonValue]] = { + "model": model, + "messages": [{"role": "user", "content": _PROMPT}], + **({"metadata": {"tags": [tag]}} if tag_spelling == "metadata_tags" else {}), + **({"tags": [tag]} if tag_spelling == "top_level_tags" else {}), + } + tagged: Final = gateway.request( + "POST", + "/v1/chat/completions", + tagged_body, + key=key, + headers={"x-litellm-tags": ", ".join(header_tags)} if tag_spelling == "header" else None, + ) + assert tagged.status_code == 200, tagged.text + _ChatResponse.model_validate_json(tagged.content) + tagged_id: Final = string_value(object_value(tagged.json())["id"]) + assert tagged.headers.get("x-litellm-applied-policies") == policy_name, tagged.text + spend_rows: Final = eventually( + lambda: read_rows( + 'SELECT request_tags FROM "LiteLLM_SpendLogs" WHERE request_id=%s AND status=%s', + (tagged_id, "success"), + ), + lambda rows: len(rows) == 1, + seconds=60, + ) + expected_request_tags: Final = header_tags if tag_spelling == "header" else [tag] + persisted_caller_tags: Final = [ + request_tag for request_tag in spend_rows[0]["request_tags"] if not request_tag.startswith("User-Agent: ") + ] + assert persisted_caller_tags == expected_request_tags, ( + f"{tag_spelling} request tags should be persisted in caller order, excluding proxy-added User-Agent tags" + ) + tagged_guardrail_requests: Final = guardrail.drain() + assert len(tagged_guardrail_requests) == 1 + _assert_guardrail_request(tagged_guardrail_requests[0], _PROMPT, model, key, response=tagged) + tagged_provider_requests: Final = provider.drain() + assert len(tagged_provider_requests) == 1 + assert [(request.method, request.target) for request in tagged_provider_requests] == [ + ("POST", "/v1/chat/completions") + ] + _assert_provider_request(tagged_provider_requests[0], _MASKED) + + untagged: Final = _chat(gateway, model, key, _PROMPT) + assert untagged.headers.get("x-litellm-applied-policies") is None, untagged.text + assert guardrail.drain() == () + untagged_provider_requests: Final = provider.drain() + assert len(untagged_provider_requests) == 1 + assert [(request.method, request.target) for request in untagged_provider_requests] == [ + ("POST", "/v1/chat/completions") + ] + _assert_provider_request(untagged_provider_requests[0], _PROMPT) + + +def test_default_policy_attachment_applies_only_when_no_other_attachment_matches(gateway: Gateway) -> None: + with ( + wire_server(lambda _request: Reply(body=json.dumps(_PROVIDER_RESPONSE).encode())) as provider, + wire_server( + lambda _request: Reply(body=b'{"action":"GUARDRAIL_INTERVENED","texts":["masked default policy prompt"]}') + ) as default_guardrail, + wire_server( + lambda _request: Reply(body=b'{"action":"GUARDRAIL_INTERVENED","texts":["masked scoped policy prompt"]}') + ) as scoped_guardrail, + gateway.scenario() as scenario, + ): + first_model: Final = scenario.model(api_base=f"{provider.url}/v1", api_key="synthetic-provider-key") + second_model: Final = scenario.model(api_base=f"{provider.url}/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[first_model, second_model]) + prefix: Final = f"attachment-{uuid.uuid4().hex}" + default_guardrail_name: Final = _create_guardrail( + gateway, + scenario.cleanups, + f"guardrail-default-{prefix}", + f"{default_guardrail.url}/beta/litellm_basic_guardrail_api", + "synthetic-guardrail-key", + ) + scoped_guardrail_name: Final = _create_guardrail( + gateway, + scenario.cleanups, + f"guardrail-scoped-{prefix}", + f"{scoped_guardrail.url}/beta/litellm_basic_guardrail_api", + "synthetic-guardrail-key", + ) + default_policy: Final = f"policy-default-{prefix}" + _create_policy(gateway, scenario.cleanups, default_policy, default_guardrail_name) + scoped_policy: Final = f"policy-scoped-{prefix}" + _create_policy(gateway, scenario.cleanups, scoped_policy, scoped_guardrail_name) + + default_attachment: Final = gateway.request( + "POST", + "/policies/attachments", + {"policy_name": default_policy, "scope": "*", "default": True}, + ) + assert default_attachment.status_code == 200, default_attachment.text + created_default: Final = _AttachmentResponse.model_validate_json(default_attachment.content) + scenario.cleanups.callback(_delete_attachment, gateway, created_default.attachment_id) + assert created_default.model_dump() == { + "attachment_id": created_default.attachment_id, + "policy_name": default_policy, + "scope": "*", + "teams": [], + "keys": [], + "models": [], + "tags": [], + "priority": None, + "default": True, + }, default_attachment.text + scoped_attachment: Final = gateway.request( + "POST", + "/policies/attachments", + {"policy_name": scoped_policy, "models": [second_model]}, + ) + assert scoped_attachment.status_code == 200, scoped_attachment.text + created_scoped: Final = _AttachmentResponse.model_validate_json(scoped_attachment.content) + scenario.cleanups.callback(_delete_attachment, gateway, created_scoped.attachment_id) + assert created_scoped.model_dump() == { + "attachment_id": created_scoped.attachment_id, + "policy_name": scoped_policy, + "scope": None, + "teams": [], + "keys": [], + "models": [second_model], + "tags": [], + "priority": None, + "default": False, + }, scoped_attachment.text + rows: Final = read_rows( + 'SELECT policy_name, is_default FROM "LiteLLM_PolicyAttachmentTable" WHERE attachment_id=%s', + (created_default.attachment_id,), + ) + assert rows == [{"policy_name": default_policy, "is_default": True}] + + first_response: Final = _chat(gateway, first_model, key, _PROMPT) + assert first_response.headers.get("x-litellm-applied-policies") == default_policy, first_response.text + second_response: Final = _chat(gateway, second_model, key, _PROMPT) + assert second_response.headers.get("x-litellm-applied-policies") == scoped_policy, second_response.text + + default_requests: Final = default_guardrail.drain() + assert [(request.method, request.target) for request in default_requests] == [ + ("POST", "/beta/litellm_basic_guardrail_api") + ] + _assert_guardrail_request(default_requests[0], _PROMPT, first_model, key, response=first_response) + scoped_requests: Final = scoped_guardrail.drain() + assert [(request.method, request.target) for request in scoped_requests] == [ + ("POST", "/beta/litellm_basic_guardrail_api") + ] + _assert_guardrail_request(scoped_requests[0], _PROMPT, second_model, key, response=second_response) + provider_requests: Final = provider.drain() + assert [(request.method, request.target) for request in provider_requests] == [ + ("POST", "/v1/chat/completions"), + ("POST", "/v1/chat/completions"), + ] + _assert_provider_request(provider_requests[0], "masked default policy prompt") + _assert_provider_request(provider_requests[1], "masked scoped policy prompt") + + +def test_policy_attachment_tag_scope_matches_tags_from_the_key_and_team(gateway: Gateway) -> None: + with ( + wire_server(lambda _request: Reply(body=json.dumps(_PROVIDER_RESPONSE).encode())) as provider, + wire_server(_policy_guardrail_reply) as guardrail, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=f"{provider.url}/v1", api_key="synthetic-provider-key") + prefix: Final = f"attachment-{uuid.uuid4().hex}" + tag: Final = f"{prefix}-prod" + guardrail_name: Final = _create_guardrail( + gateway, + scenario.cleanups, + f"guardrail-{prefix}", + f"{guardrail.url}/beta/litellm_basic_guardrail_api", + "synthetic-guardrail-key", + ) + policy_name: Final = f"policy-{prefix}" + _create_policy(gateway, scenario.cleanups, policy_name, guardrail_name) + attachment: Final = gateway.request( + "POST", + "/policies/attachments", + {"policy_name": policy_name, "tags": [tag]}, + ) + assert attachment.status_code == 200, attachment.text + created_attachment: Final = _AttachmentResponse.model_validate_json(attachment.content) + scenario.cleanups.callback(_delete_attachment, gateway, created_attachment.attachment_id) + assert created_attachment.model_dump() == { + "attachment_id": created_attachment.attachment_id, + "policy_name": policy_name, + "scope": None, + "teams": [], + "keys": [], + "models": [], + "tags": [tag], + "priority": None, + "default": False, + }, attachment.text + + key_tagged: Final = scenario.key(models=[model], metadata={"tags": [tag]}) + team_tagged: Final = scenario.team(models=[model], metadata={"tags": [tag]}) + team_key: Final = scenario.key(team_id=team_tagged, models=[model]) + untagged_team: Final = scenario.team(models=[model]) + untagged_key: Final = scenario.key(team_id=untagged_team, models=[model]) + assert guardrail.drain() == () + assert provider.drain() == () + + key_prompt: Final = f"{_PROMPT} key metadata" + team_prompt: Final = f"{_PROMPT} team metadata" + control_prompt: Final = f"{_PROMPT} untagged control" + key_response: Final = _chat(gateway, model, key_tagged, key_prompt) + assert key_response.headers.get("x-litellm-applied-policies") == policy_name, key_response.text + team_response: Final = _chat(gateway, model, team_key, team_prompt) + assert team_response.headers.get("x-litellm-applied-policies") == policy_name, team_response.text + guardrail_requests: Final = guardrail.drain() + assert len(guardrail_requests) == 2 + _assert_guardrail_request(guardrail_requests[0], key_prompt, model, key_tagged, response=key_response) + _assert_guardrail_request(guardrail_requests[1], team_prompt, model, team_key, response=team_response) + provider_requests: Final = provider.drain() + assert len(provider_requests) == 2 + _assert_provider_request(provider_requests[0], f"{_MASKED}: {key_prompt}") + _assert_provider_request(provider_requests[1], f"{_MASKED}: {team_prompt}") + + control_response: Final = _chat(gateway, model, untagged_key, control_prompt) + assert control_response.headers.get("x-litellm-applied-policies") is None, control_response.text + assert guardrail.drain() == () + control_provider_requests: Final = provider.drain() + assert len(control_provider_requests) == 1 + _assert_provider_request(control_provider_requests[0], control_prompt) + + +def test_streaming_chat_on_an_attached_model_runs_the_policy_guardrail(gateway: Gateway) -> None: + with ( + wire_server(_streaming_chat_reply) as provider, + wire_server(_policy_guardrail_reply) as guardrail, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=f"{provider.url}/v1", api_key="synthetic-provider-key") + second_model: Final = scenario.model(api_base=f"{provider.url}/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model, second_model]) + prefix: Final = f"attachment-{uuid.uuid4().hex}" + guardrail_name: Final = _create_guardrail( + gateway, + scenario.cleanups, + f"guardrail-{prefix}", + f"{guardrail.url}/beta/litellm_basic_guardrail_api", + "synthetic-guardrail-key", + ) + policy_name: Final = f"policy-{prefix}" + _create_policy(gateway, scenario.cleanups, policy_name, guardrail_name) + attachment: Final = gateway.request( + "POST", + "/policies/attachments", + {"policy_name": policy_name, "models": [model]}, + ) + assert attachment.status_code == 200, attachment.text + created_attachment: Final = _AttachmentResponse.model_validate_json(attachment.content) + scenario.cleanups.callback(_delete_attachment, gateway, created_attachment.attachment_id) + assert created_attachment.model_dump() == { + "attachment_id": created_attachment.attachment_id, + "policy_name": policy_name, + "scope": None, + "teams": [], + "keys": [], + "models": [model], + "tags": [], + "priority": None, + "default": False, + }, attachment.text + assert guardrail.drain() == () + assert provider.drain() == () + + attached_headers, attached_text = _stream_chat(gateway, model, key, _PROMPT) + assert attached_headers.get("x-litellm-applied-policies") == policy_name, dict(attached_headers) + assert attached_text == "provider response" + guardrail_requests: Final = guardrail.drain() + assert len(guardrail_requests) == 1 + _assert_guardrail_request(guardrail_requests[0], _PROMPT, model, key) + provider_requests: Final = provider.drain() + assert len(provider_requests) == 1 + _assert_streaming_provider_request(provider_requests[0], f"{_MASKED}: {_PROMPT}") + + assert guardrail.drain() == () + unattached_headers, unattached_text = _stream_chat(gateway, second_model, key, _PROMPT) + assert unattached_headers.get("x-litellm-applied-policies") is None, dict(unattached_headers) + assert unattached_text == "provider response" + assert guardrail.drain() == () + unattached_provider_requests: Final = provider.drain() + assert len(unattached_provider_requests) == 1 + _assert_streaming_provider_request(unattached_provider_requests[0], _PROMPT) + + +def test_streaming_request_body_policies_run_the_guardrail_and_are_removed_from_the_provider_body( + gateway: Gateway, +) -> None: + with ( + wire_server(_streaming_chat_reply) as provider, + wire_server(_policy_guardrail_reply) as guardrail, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=f"{provider.url}/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + prefix: Final = f"attachment-{uuid.uuid4().hex}" + guardrail_name: Final = _create_guardrail( + gateway, + scenario.cleanups, + f"guardrail-{prefix}", + f"{guardrail.url}/beta/litellm_basic_guardrail_api", + "synthetic-guardrail-key", + ) + policy_name: Final = f"policy-{prefix}" + _create_policy(gateway, scenario.cleanups, policy_name, guardrail_name) + assert guardrail.drain() == () + assert provider.drain() == () + + response_headers, response_text = _stream_chat(gateway, model, key, _PROMPT, policy_name=policy_name) + assert response_headers.get("x-litellm-applied-policies") == policy_name, dict(response_headers) + assert response_text == "provider response" + guardrail_requests: Final = guardrail.drain() + assert len(guardrail_requests) == 1 + _assert_guardrail_request(guardrail_requests[0], _PROMPT, model, key) + provider_requests: Final = provider.drain() + assert len(provider_requests) == 1 + _assert_streaming_provider_request(provider_requests[0], f"{_MASKED}: {_PROMPT}") diff --git a/tests/integration/observability/test_apply_guardrail_endpoint.py b/tests/integration/observability/test_apply_guardrail_endpoint.py new file mode 100644 index 00000000000..4ff11a50667 --- /dev/null +++ b/tests/integration/observability/test_apply_guardrail_endpoint.py @@ -0,0 +1,313 @@ +import json +import uuid +from contextlib import ExitStack +from pathlib import Path +from typing import Final + +import httpx +import yaml +from integration._support.client import Gateway, string_value +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_BLOCK_REASON: Final = "synthetic blocked by apply endpoint" +_BLOCKED_TEXT: Final = "forbidden integration prompt" +_METADATA_TEXT: Final = "metadata forwarding prompt" +_PROVIDER_RESPONSE: Final = { + "id": "chatcmpl-apply-guardrail", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "provider response"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +class _CreatedGuardrail(BaseModel): + guardrail_id: str + guardrail_name: str + + +class _ApplyResponse(BaseModel): + model_config = ConfigDict(extra="forbid") + + response_text: str + + +class _ApplyErrorBody(BaseModel): + model_config = ConfigDict(extra="forbid") + + message: str + type: str + param: str | None + code: int + + +class _ApplyError(BaseModel): + model_config = ConfigDict(extra="forbid") + + error: _ApplyErrorBody + + +def _delete_guardrail(gateway: Gateway, guardrail_id: str) -> None: + deleted: Final = gateway.request("DELETE", f"/guardrails/{guardrail_id}") + assert deleted.status_code == 200, deleted.text + + +def _create_guardrail( + gateway: Gateway, + cleanups: ExitStack, + name: str, + params: dict[str, JsonValue], +) -> str: + response: Final = gateway.request( + "POST", + "/guardrails", + {"guardrail": {"guardrail_name": name, "litellm_params": params}}, + ) + assert response.status_code == 200, response.text + created: Final = _CreatedGuardrail.model_validate_json(response.content) + assert created.guardrail_name == name, response.text + cleanups.callback(_delete_guardrail, gateway, created.guardrail_id) + return created.guardrail_name + + +def _assert_guardrail_request(request: Request, text: str, response: httpx.Response) -> None: + assert (request.method, request.target) == ("POST", "/beta/litellm_basic_guardrail_api") + payload: Final = _JSON_OBJECT.validate_json(request.body) + assert set(payload) == { + "input_type", + "litellm_call_id", + "litellm_trace_id", + "structured_messages", + "images", + "tools", + "texts", + "request_data", + "request_headers", + "litellm_version", + "additional_provider_specific_params", + "tool_calls", + "model", + }, response.text + assert { + field: payload[field] + for field in ( + "input_type", + "structured_messages", + "images", + "tools", + "texts", + "additional_provider_specific_params", + "tool_calls", + "model", + ) + } == { + "input_type": "request", + "structured_messages": None, + "images": None, + "tools": None, + "texts": [text], + "additional_provider_specific_params": {}, + "tool_calls": None, + "model": None, + }, response.text + assert payload["request_data"] == {}, response.text + request_headers: Final = payload["request_headers"] + assert request_headers is None or isinstance(request_headers, dict), response.text + if isinstance(request_headers, dict): + assert "authorization" not in request_headers, response.text + assert payload["litellm_call_id"] is None or isinstance(payload["litellm_call_id"], str), response.text + assert payload["litellm_trace_id"] is None or isinstance(payload["litellm_trace_id"], str), response.text + assert payload["litellm_version"] is None or isinstance(payload["litellm_version"], str), response.text + if response.headers.get("x-litellm-call-id") is not None: + assert payload["litellm_call_id"] == response.headers["x-litellm-call-id"], response.text + assert request.headers["x-api-key"] == "synthetic-apply-guardrail-key", response.text + + +def test_apply_guardrail_returns_a_block_error_for_generic_guardrails(gateway: Gateway) -> None: + def block(_request: Request) -> Reply: + return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": _BLOCK_REASON}).encode()) + + with wire_server(block) as guardrail, gateway.scenario() as scenario: + name: Final = _create_guardrail( + gateway, + scenario.cleanups, + f"blocked-apply-{uuid.uuid4().hex}", + { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "api_base": f"{guardrail.url}/beta/litellm_basic_guardrail_api", + "api_key": "synthetic-apply-guardrail-key", + }, + ) + response: Final = gateway.request( + "POST", + "/guardrails/apply_guardrail", + {"guardrail_name": name, "text": _BLOCKED_TEXT}, + ) + assert response.status_code == 400, response.text + error: Final = _ApplyError.model_validate_json(response.content) + assert error.model_dump() == { + "error": { + "message": _BLOCK_REASON, + "type": "internal_server_error", + "param": None, + "code": 400, + } + }, response.text + requests: Final = guardrail.drain() + assert [(request.method, request.target) for request in requests] == [ + ("POST", "/beta/litellm_basic_guardrail_api") + ] + request: Final = requests[0] + _assert_guardrail_request(request, _BLOCKED_TEXT, response) + + +def test_apply_guardrail_forwards_messages_to_the_guardrail(gateway: Gateway) -> None: + def allow(_request: Request) -> Reply: + return Reply(body=json.dumps({"data": {"guardrailsResult": {"Allowed": True, "Reason": ""}}}).encode()) + + text: Final = f"akto apply text {uuid.uuid4().hex}" + messages: Final = [ + {"role": "system", "content": "akto system prompt"}, + {"role": "user", "content": "akto user content different from the text field"}, + ] + with wire_server(allow) as akto, gateway.scenario() as scenario: + name: Final = _create_guardrail( + gateway, + scenario.cleanups, + f"akto-apply-{uuid.uuid4().hex}", + { + "guardrail": "akto", + "mode": "pre_call", + "default_on": False, + "akto_base_url": akto.url, + "akto_api_key": "synthetic-akto-apply-key", + }, + ) + with_messages: Final = gateway.request( + "POST", + "/guardrails/apply_guardrail", + {"guardrail_name": name, "text": text, "messages": messages}, + ) + assert with_messages.status_code == 200, with_messages.text + assert _ApplyResponse.model_validate_json(with_messages.content).model_dump() == {"response_text": text}, ( + with_messages.text + ) + without_messages: Final = gateway.request( + "POST", + "/guardrails/apply_guardrail", + {"guardrail_name": name, "text": text}, + ) + assert without_messages.status_code == 200, without_messages.text + assert _ApplyResponse.model_validate_json(without_messages.content).model_dump() == {"response_text": text}, ( + without_messages.text + ) + + requests: Final = akto.drain() + assert [(request.method, request.target) for request in requests] == [ + ("POST", "/api/http-proxy?akto_connector=litellm&guardrails=true"), + ("POST", "/api/http-proxy?akto_connector=litellm&guardrails=true"), + ] + assert [request.headers["authorization"] for request in requests] == ["synthetic-akto-apply-key"] * 2 + bodies: Final = tuple( + _JSON_OBJECT.validate_json( + string_value( + json.loads(string_value(_JSON_OBJECT.validate_json(request.body)["requestPayload"]))["body"] + ) + ) + for request in requests + ) + assert bodies[0]["messages"] == messages, bodies[0] + assert bodies[1]["messages"] == [{"role": "user", "content": text}], bodies[1] + + +def test_apply_guardrail_forwards_metadata_to_custom_code( + gateway: Gateway, + tmp_path: Path, +) -> None: + def record(request: Request) -> Reply: + assert request.target == "/metadata" + return Reply(body=b'{"received":true}') + + with wire_server(record) as recorder: + config: Final = _JSON_OBJECT.validate_python( + yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + ) + settings: Final = _JSON_OBJECT.validate_python(config["litellm_settings"]) + config_path: Final = tmp_path / "apply-guardrail-owned-proxy.yaml" + config_path.write_text( + yaml.safe_dump( + { + **config, + "litellm_settings": {**settings, "user_url_allowed_hosts": ["127.0.0.1"]}, + } + ) + ) + custom_code: Final = f""" +async def apply_guardrail(inputs, request_data, input_type): + await http_post( + "{recorder.url}/metadata", + body={{ + "texts": inputs.get("texts", []), + "metadata": request_data.get("metadata", {{}}), + "input_type": input_type, + }}, + ) + return allow() +""" + with owned_proxy(gateway, tmp_path, {}, config=config_path) as candidate: + with candidate.scenario() as scenario: + name: Final = _create_guardrail( + candidate, + scenario.cleanups, + f"metadata-apply-{uuid.uuid4().hex}", + { + "guardrail": "custom_code", + "mode": "pre_call", + "default_on": False, + "custom_code": custom_code, + }, + ) + with_metadata: Final = candidate.request( + "POST", + "/guardrails/apply_guardrail", + { + "guardrail_name": name, + "text": _METADATA_TEXT, + "metadata": {"forbidden_topics": ["tax"]}, + }, + ) + assert with_metadata.status_code == 200, with_metadata.text + assert _ApplyResponse.model_validate_json(with_metadata.content).model_dump() == { + "response_text": _METADATA_TEXT + }, with_metadata.text + without_metadata: Final = candidate.request( + "POST", + "/guardrails/apply_guardrail", + {"guardrail_name": name, "text": _METADATA_TEXT}, + ) + assert without_metadata.status_code == 200, without_metadata.text + assert _ApplyResponse.model_validate_json(without_metadata.content).model_dump() == { + "response_text": _METADATA_TEXT + }, without_metadata.text + + requests: Final = recorder.drain() + assert [(request.method, request.target) for request in requests] == [ + ("POST", "/metadata"), + ("POST", "/metadata"), + ] + assert tuple(_JSON_OBJECT.validate_json(request.body) for request in requests) == ( + { + "texts": [_METADATA_TEXT], + "metadata": {"forbidden_topics": ["tax"]}, + "input_type": "request", + }, + {"texts": [_METADATA_TEXT], "metadata": {}, "input_type": "request"}, + ) diff --git a/tests/integration/observability/test_presidio_entity_masking.py b/tests/integration/observability/test_presidio_entity_masking.py index 03ac65532e4..843afb51c5c 100644 --- a/tests/integration/observability/test_presidio_entity_masking.py +++ b/tests/integration/observability/test_presidio_entity_masking.py @@ -11,7 +11,7 @@ import httpx import pytest from integration._support.client import Gateway, string_value from integration._support.wire import Reply, Request, Wire, wire_server -from pydantic import JsonValue +from pydantic import BaseModel, ConfigDict, JsonValue CARD: Final = "4111-1111-1111-1111" EMAIL: Final = "jane.doe@example.com" @@ -24,6 +24,27 @@ RECOGNIZERS: Final = { } +class _ApplyGuardrailResponse(BaseModel): + model_config = ConfigDict(extra="forbid") + + response_text: str + + +class _ApplyGuardrailErrorBody(BaseModel): + model_config = ConfigDict(extra="forbid") + + message: str + type: str + param: str | None + code: str + + +class _ApplyGuardrailError(BaseModel): + model_config = ConfigDict(extra="forbid") + + error: _ApplyGuardrailErrorBody + + def _detect(entity: str, text: str) -> tuple[dict[str, JsonValue], ...]: return tuple( {"entity_type": entity, "start": match.start(), "end": match.end(), "score": 0.95} @@ -91,6 +112,56 @@ def _requested_entities(analyzer: Wire) -> list[JsonValue]: return [json.loads(request.body).get("entities") for request in analyzer.drain()] +def _assert_apply_guardrail_successes( + presidio: Presidio, + text: str, + responses: tuple[httpx.Response, ...], +) -> None: + assert tuple(response.status_code for response in responses) == (200,) * len(responses), tuple( + response.text for response in responses + ) + expected_response: Final = {"response_text": "alias contract and "} + assert tuple( + _ApplyGuardrailResponse.model_validate_json(response.content).model_dump() for response in responses + ) == (expected_response,) * len(responses), tuple(response.text for response in responses) + + analyzer_requests: Final = presidio.analyzer.drain() + assert len(analyzer_requests) == len(responses) + assert tuple((request.method, request.target) for request in analyzer_requests) == ( + ("POST", "/analyze"), + ) * len(responses) + expected_analyzer_body: Final = {"text": text, "language": "en"} + assert tuple(json.loads(request.body) for request in analyzer_requests) == ( + expected_analyzer_body, + ) * len(responses) + + anonymizer_requests: Final = presidio.anonymizer.drain() + assert len(anonymizer_requests) == len(responses) + assert tuple((request.method, request.target) for request in anonymizer_requests) == ( + ("POST", "/anonymize"), + ) * len(responses) + expected_anonymizer_body: Final = { + "text": text, + "analyzer_results": [ + { + "entity_type": "CREDIT_CARD", + "start": text.index(CARD), + "end": text.index(CARD) + len(CARD), + "score": 0.95, + }, + { + "entity_type": "EMAIL_ADDRESS", + "start": text.index(EMAIL), + "end": text.index(EMAIL) + len(EMAIL), + "score": 0.95, + }, + ], + } + assert tuple(json.loads(request.body) for request in anonymizer_requests) == ( + expected_anonymizer_body, + ) * len(responses) + + def test_pre_call_masks_only_the_configured_entities_before_the_provider_sees_the_prompt(gateway: Gateway) -> None: with ( _presidio(gateway, "pre_call", {"CREDIT_CARD": "MASK", "EMAIL_ADDRESS": "MASK"}) as presidio, @@ -139,3 +210,192 @@ def test_apply_guardrail_with_the_default_config_masks_every_detected_entity( assert masked == "card and email ", masked assert _requested_entities(presidio.analyzer) == [None] assert len(presidio.anonymizer.drain()) == 1 + + +def test_apply_guardrail_forwards_documented_language_and_entities_to_presidio(gateway: Gateway) -> None: + pytest.skip("BUG: /guardrails/apply_guardrail drops the documented language and entities fields before Presidio") + with _presidio(gateway, "pre_call", None) as presidio: + text: Final = f"apply language contract {CARD} and {EMAIL}" + response: Final = gateway.request( + "POST", + "/guardrails/apply_guardrail", + { + "guardrail_name": presidio.name, + "text": text, + "language": "fr", + "entities": ["EMAIL_ADDRESS"], + }, + ) + assert response.status_code == 200, response.text + masked: Final = string_value(response.json()["response_text"]) + assert masked == f"apply language contract {CARD} and ", masked + analyzer_requests: Final = presidio.analyzer.drain() + assert len(analyzer_requests) == 1 + assert json.loads(analyzer_requests[0].body) == { + "text": text, + "language": "fr", + "entities": ["EMAIL_ADDRESS"], + }, analyzer_requests[0].body + assert len(presidio.anonymizer.drain()) == 1 + + +def test_apply_guardrail_aliases_accept_admin_auth(gateway: Gateway) -> None: + with _presidio(gateway, "pre_call", None) as presidio: + text: Final = f"alias contract {CARD} and {EMAIL}" + body: Final = {"guardrail_name": presidio.name, "text": text} + paths: Final = ("/apply_guardrail", "/guardrails/apply_guardrail") + + unauthenticated_responses: Final = tuple(gateway.client.post(path, json=body) for path in paths) + assert tuple(response.status_code for response in unauthenticated_responses) == (401, 401), tuple( + response.text for response in unauthenticated_responses + ) + assert presidio.analyzer.drain() == () + assert presidio.anonymizer.drain() == () + + admin_requests: Final = tuple( + gateway.client.post(path, json=body, headers=headers) + for path, headers in ( + (paths[0], {"Authorization": f"Bearer {gateway.key}"}), + (paths[1], {"Authorization": f"Bearer {gateway.key}"}), + (paths[0], {"x-litellm-api-key": gateway.key}), + (paths[1], {"x-litellm-api-key": gateway.key}), + ) + ) + assert tuple(response.status_code for response in admin_requests) == (200, 200, 200, 200), tuple( + response.text for response in admin_requests + ) + for response in admin_requests: + assert _ApplyGuardrailResponse.model_validate_json(response.content).model_dump() == { + "response_text": "alias contract and " + }, response.text + + analyzer_requests: Final = presidio.analyzer.drain() + assert len(analyzer_requests) == 4 + assert all(request.method == "POST" and request.target == "/analyze" for request in analyzer_requests) + assert len({request.body for request in analyzer_requests}) == 1 + assert json.loads(analyzer_requests[0].body) == {"text": text, "language": "en"}, analyzer_requests[0].body + anonymizer_requests: Final = presidio.anonymizer.drain() + assert len(anonymizer_requests) == 4 + assert all(request.method == "POST" and request.target == "/anonymize" for request in anonymizer_requests) + assert len({request.body for request in anonymizer_requests}) == 1 + assert json.loads(anonymizer_requests[0].body) == { + "text": text, + "analyzer_results": [ + { + "entity_type": "CREDIT_CARD", + "start": text.index(CARD), + "end": text.index(CARD) + len(CARD), + "score": 0.95, + }, + { + "entity_type": "EMAIL_ADDRESS", + "start": text.index(EMAIL), + "end": text.index(EMAIL) + len(EMAIL), + "score": 0.95, + }, + ], + }, anonymizer_requests[0].body + + +def test_long_apply_guardrail_alias_accepts_an_llm_api_scoped_key(gateway: Gateway) -> None: + with ( + _presidio(gateway, "pre_call", None) as presidio, + gateway.scenario() as scenario, + ): + proxy_admin: Final = scenario.user(user_role="proxy_admin") + key: Final = scenario.key(user_id=proxy_admin, key_type="llm_api") + text: Final = f"alias contract {CARD} and {EMAIL}" + body: Final = {"guardrail_name": presidio.name, "text": text} + path: Final = "/guardrails/apply_guardrail" + + assert presidio.analyzer.drain() == () + assert presidio.anonymizer.drain() == () + responses: Final = tuple( + gateway.client.post(path, json=body, headers=headers) + for headers in ( + {"Authorization": f"Bearer {key}"}, + {"x-litellm-api-key": key}, + ) + ) + _assert_apply_guardrail_successes(presidio, text, responses) + + unauthenticated: Final = gateway.client.post(path, json=body) + assert unauthenticated.status_code == 401, unauthenticated.text + assert _ApplyGuardrailError.model_validate_json(unauthenticated.content).model_dump() == { + "error": { + "message": "Authentication Error, No api key passed in.", + "type": "auth_error", + "param": "None", + "code": "401", + } + }, unauthenticated.text + assert presidio.analyzer.drain() == () + assert presidio.anonymizer.drain() == () + + +def test_short_apply_guardrail_alias_accepts_an_llm_api_scoped_key(gateway: Gateway) -> None: + pytest.skip( + "BUG: llm_api scoped keys are refused on /apply_guardrail while " + "/guardrails/apply_guardrail serves them" + ) + with ( + _presidio(gateway, "pre_call", None) as presidio, + gateway.scenario() as scenario, + ): + proxy_admin: Final = scenario.user(user_role="proxy_admin") + key: Final = scenario.key(user_id=proxy_admin, key_type="llm_api") + text: Final = f"alias contract {CARD} and {EMAIL}" + body: Final = {"guardrail_name": presidio.name, "text": text} + path: Final = "/apply_guardrail" + + assert presidio.analyzer.drain() == () + assert presidio.anonymizer.drain() == () + responses: Final = tuple( + gateway.client.post(path, json=body, headers=headers) + for headers in ( + {"Authorization": f"Bearer {key}"}, + {"x-litellm-api-key": key}, + ) + ) + _assert_apply_guardrail_successes(presidio, text, responses) + + +def test_apply_guardrail_aliases_accept_internal_user_auth(gateway: Gateway) -> None: + pytest.skip("BUG: non-admin virtual keys get 401 on /apply_guardrail and /guardrails/apply_guardrail") + with ( + _presidio(gateway, "pre_call", None) as presidio, + gateway.scenario() as scenario, + ): + internal_user: Final = scenario.user(user_role="internal_user") + internal_key: Final = scenario.key( + user_id=internal_user, + key_alias=f"presidio-apply-short-{uuid.uuid4().hex}", + ) + text: Final = f"alias contract {CARD} and {EMAIL}" + body: Final = {"guardrail_name": presidio.name, "text": text} + paths: Final = ("/apply_guardrail", "/guardrails/apply_guardrail") + internal_auth_requests: Final = ( + (paths[0], {"Authorization": f"Bearer {internal_key}"}), + (paths[1], {"Authorization": f"Bearer {internal_key}"}), + (paths[0], {"x-litellm-api-key": internal_key}), + (paths[1], {"x-litellm-api-key": internal_key}), + ) + responses: Final = tuple( + gateway.client.post(path, json=body, headers=headers) for path, headers in internal_auth_requests + ) + assert tuple(response.status_code for response in responses) == (200, 200, 200, 200), tuple( + response.text for response in responses + ) + assert ( + tuple(_ApplyGuardrailResponse.model_validate_json(response.content).model_dump() for response in responses) + == ({"response_text": "alias contract and "},) * 4 + ) + analyzer_requests: Final = presidio.analyzer.drain() + assert len(analyzer_requests) == 4 + assert all(request.method == "POST" and request.target == "/analyze" for request in analyzer_requests) + assert len({request.body for request in analyzer_requests}) == 1 + assert json.loads(analyzer_requests[0].body)["text"] == text + anonymizer_requests: Final = presidio.anonymizer.drain() + assert len(anonymizer_requests) == 4 + assert all(request.method == "POST" and request.target == "/anonymize" for request in anonymizer_requests) + assert len({request.body for request in anonymizer_requests}) == 1 diff --git a/tests/integration/security/test_stored_config_slots.py b/tests/integration/security/test_stored_config_slots.py index b2386306e42..cb788617f28 100644 --- a/tests/integration/security/test_stored_config_slots.py +++ b/tests/integration/security/test_stored_config_slots.py @@ -64,6 +64,10 @@ from integration.security._sweeps import ( sweep_all, sweep_sink, ) +from litellm.types.guardrails import ( + GuardrailInfoResponse, + ListGuardrailsResponse, +) OUTCOMES: Final = ("success", "provider_4xx") BEDROCK_MODEL: Final = "bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0" @@ -81,6 +85,11 @@ VERTEX_MODEL_PATH: Final = ( ) +class _CreatedGuardrail(BaseModel): + guardrail_id: str + guardrail_name: str + + @pytest.fixture(scope="module") def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: """Shared proxy for slots created through the API, with team model_config overrides on.""" @@ -785,7 +794,7 @@ def _guardrail(request: Request) -> Reply: return Reply(body=json.dumps({"action": "NONE"}).encode()) -def _guardrail_params(url: str, secret: Canary) -> dict[str, object]: +def _guardrail_params(url: str, secret: Canary) -> dict[str, JsonValue]: return { "guardrail": "generic_guardrail_api", "mode": "pre_call", @@ -849,6 +858,121 @@ def test_config_guardrail_api_key_reaches_only_the_guardrail( ) +@pytest.mark.timeout(240) +@pytest.mark.parametrize("outcome", OUTCOMES) +def test_api_created_guardrail_api_key_is_encrypted_at_rest_and_reaches_only_the_guardrail( + rig: Rig, outcome: str, request: pytest.FixtureRequest +) -> None: + started: Final = datetime.now(UTC) + secret: Final = canary("E1") + marker: Final = canary(MARKER) + name: Final = f"canary-api-guardrail-{uuid.uuid4().hex}" + with wire_server(_guardrail) as wire, rig.proxy.scenario() as scenario: + guardrail: Final = Recorder(wire) + created_response: Final = rig.proxy.request( + "POST", + "/guardrails", + {"guardrail": {"guardrail_name": name, "litellm_params": _guardrail_params(wire.url, secret)}}, + ) + assert created_response.status_code == 200, created_response.text + created: Final = _CreatedGuardrail.model_validate_json(created_response.content) + assert created.guardrail_name == name, created_response.text + guardrail_id: Final = created.guardrail_id + + def delete_guardrail() -> None: + deleted: Final = rig.proxy.request("DELETE", f"/guardrails/{guardrail_id}") + assert deleted.status_code == 200, deleted.text + assert ( + read_rows( + 'SELECT guardrail_id FROM "LiteLLM_GuardrailsTable" WHERE guardrail_id=%s', + (guardrail_id,), + ) + == [] + ) + + scenario.cleanups.callback(delete_guardrail) + caller: Final = _caller(scenario, models=[CONFIG_MODEL]) + response, request_id = _chat(rig.proxy, caller.key, CONFIG_MODEL, "E1", marker, outcome) + _guardrail_delivered(guardrail, marker, secret) + guardrail_requests: Final = guardrail.requests() + assert [(item.method, item.target) for item in guardrail_requests] == [("POST", GUARDRAIL_PATH)] + + _assert_stored_without_canary( + "SELECT litellm_params->>'api_key' FROM \"LiteLLM_GuardrailsTable\" WHERE guardrail_id=%s", + (guardrail_id,), + secret, + ) + params_rows: Final = read_rows( + 'SELECT litellm_params::text AS params FROM "LiteLLM_GuardrailsTable" WHERE guardrail_id=%s', + (guardrail_id,), + ) + assert len(params_rows) == 1, params_rows + stored_params: Final = string_value(params_rows[0]["params"]) + assert find_canary(stored_params, (secret,)) == (), f"Slot {secret.slot} stored in plaintext JSON" + + masked_secret: Final = f"{secret.value[:2]}{'*' * 4}{secret.value[-2:]}" + detail_routes: Final = (f"/guardrails/{guardrail_id}", f"/guardrails/{guardrail_id}/info") + detail_responses: Final = tuple(rig.proxy.request("GET", path) for path in detail_routes) + assert tuple(item.status_code for item in detail_responses) == (200, 200), tuple( + item.text for item in detail_responses + ) + details: Final = tuple(GuardrailInfoResponse.model_validate_json(item.content) for item in detail_responses) + assert all(item.created_at is not None and item.updated_at is not None for item in details), details + assert ( + tuple( + ( + item.guardrail_id, + item.guardrail_name, + item.guardrail_definition_location.value, + ) + for item in details + ) + == ((guardrail_id, name, "db"),) * 2 + ), tuple(item.text for item in detail_responses) + assert ( + tuple((item.created_at, item.updated_at) for item in details) + == ((details[0].created_at, details[0].updated_at),) * 2 + ) + assert all(item.litellm_params is not None for item in details), tuple(item.text for item in detail_responses) + assert ( + tuple( + ( + item.litellm_params.default_on, + item.litellm_params.api_base, + item.litellm_params.api_key, + ) + for item in details + if item.litellm_params is not None + ) + == ((True, wire.url, masked_secret),) * 2 + ), tuple(item.text for item in detail_responses) + + list_response: Final = rig.proxy.request("GET", "/v2/guardrails/list") + assert list_response.status_code == 200, list_response.text + listed_guardrails: Final = ListGuardrailsResponse.model_validate_json(list_response.content).guardrails + matching_entries: Final = tuple(item for item in listed_guardrails if item.guardrail_id == guardrail_id) + assert len(matching_entries) == 1, list_response.text + assert find_canary(list_response.content, (secret,)) == (), list_response.text + + _finish( + rig, + rig.proxy, + request, + secrets=(secret,), + marker=marker, + response=response, + request_id=request_id, + caller=caller, + ids={"model": CONFIG_MODEL, "guardrail_id": guardrail_id}, + detail_routes=detail_routes, + reads=_reads(rig.proxy, {"/guardrails/list": {}, "/v2/guardrails/list": {}}), + extra_sinks={GUARDRAIL_SINK: guardrail.requests}, + own_headers={GUARDRAIL_SINK: ("x-api-key", "E1")}, + context=f"slot E1 (API-created DB guardrail), {outcome}", + since=started, + ) + + def _langfuse(request: Request) -> Reply: if request.method == "GET" and request.target.startswith("/api/public/projects"): return Reply(body=json.dumps({"data": [{"id": "canary-project", "name": "canary"}]}).encode())