diff --git a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py index 5109f09d9c2..a91812bb474 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py @@ -1,4 +1,7 @@ +import json import os +from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import Any, Final from urllib.parse import urlparse @@ -19,20 +22,26 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.base import ( GuardrailConfigModel, ) from litellm.types.proxy.guardrails.guardrail_hooks.singulr import ( + AssistantMessage, SingulrGuardrailPayload, - SingulrGuardrailRequest, SingulrGuardrailResponse, + SingulrMcpGuardrailPayload, + ToolCall, + ToolCallFunction, ) -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs _DEFAULT_API_BASE: Final = "http://localhost:8003" -_GUARD_ENDPOINT: Final = "/api/v1/ai-gateway/litellm" +_GUARD_ENDPOINT: Final = "/api/v1/ai-gateway/litellm-v2" _DEFAULT_TIMEOUT: Final = 30.0 +_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({}) +_MCP_MODEL_PREFIX: Final = "MCP:" class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): @@ -51,8 +60,8 @@ class SingulrGuardrail(CustomGuardrail): **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: self.singulr_api_key = singulr_api_key or os.environ.get("SINGULR_API_KEY") - self.singulr_api_base = (singulr_api_base or os.environ.get("SINGULR_API_BASE") or _DEFAULT_API_BASE).rstrip( - "/" + self.singulr_api_base = ( + (singulr_api_base or os.environ.get("SINGULR_API_BASE") or _DEFAULT_API_BASE).strip().rstrip("/") ) parsed: Final = urlparse(self.singulr_api_base) if parsed.scheme == "http" and parsed.hostname not in ( @@ -85,6 +94,9 @@ class SingulrGuardrail(CustomGuardrail): kwargs["supported_event_hooks"] = [ GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call, + GuardrailEventHooks.logging_only, + GuardrailEventHooks.pre_mcp_call, + GuardrailEventHooks.post_mcp_call, ] super().__init__(**kwargs) @@ -97,52 +109,70 @@ class SingulrGuardrail(CustomGuardrail): return SingulrGuardrailConfigModel - def _build_payload( - self, - request_data: dict[str, Any], - inputs: GenericGuardrailAPIInputs, - input_type: str, - ) -> dict[str, object]: - if not request_data: - texts: Final = inputs.get("texts", []) - - payload = SingulrGuardrailPayload( - input_type=input_type, - is_playground_request=True, - playground_text=texts[0] if texts else None, + @staticmethod + def _metadata_containers(request_data: Mapping[str, Any]) -> tuple[Mapping[str, Any], ...]: + litellm_params: Final = request_data.get("litellm_params") or _EMPTY_MAPPING + return tuple( + container + for container in ( + request_data.get("litellm_metadata"), + request_data.get("metadata"), + litellm_params.get("litellm_metadata") if litellm_params else None, + litellm_params.get("metadata") if litellm_params else None, ) - else: - response: Final = request_data.get("response") - singulr_req_object: Final = SingulrGuardrailRequest( - model=request_data.get("model"), - messages=request_data.get("messages"), - tools=request_data.get("tools"), - model_response=response.model_dump(mode="json") if input_type == "response" and response else None, - litellm_metadata=request_data.get("litellm_metadata"), - ) - payload = SingulrGuardrailPayload( - litellm_call_id=request_data.get("litellm_call_id"), - request_data=singulr_req_object, - input_type=input_type, - ) - - return payload.model_dump(mode="json") - - def _build_headers(self) -> dict[str, str]: - return dict( - (header, value) - for header, value in ( - ("Content-Type", "application/json"), - ("X-Singulr-Gateway-Token", self.singulr_api_key), - ( - "X-Singulr-Enforcement-Entity-Id", - self.singulr_application_id or "", - ), - ("X-Singulr-Guardrail-Id", self.singulr_guardrail_id or ""), - ) - if value + if container ) + @classmethod + def _resolve_metadata_value(cls, request_data: Mapping[str, Any], key: str) -> str | None: + for container in cls._metadata_containers(request_data=request_data): + value = container.get(key) + if value: + return value + return None + + @classmethod + def _resolve_user_role_from_request_data(cls, request_data: Mapping[str, Any]) -> str | None: + for container in cls._metadata_containers(request_data=request_data): + auth = container.get("user_api_key_auth") + if isinstance(auth, UserAPIKeyAuth) and auth.user_role: + return auth.user_role.value + return None + + @classmethod + def _build_metadata(cls, request_data: Mapping[str, Any]) -> Mapping[str, str] | None: + fields: Final = ( + "user_api_key_alias", + "user_api_key_user_id", + "user_api_key_user_email", + "user_api_key_org_id", + "user_api_key_org_alias", + "user_api_key_team_id", + "user_api_key_team_alias", + ) + resolved: Final = ( + *((field, cls._resolve_metadata_value(request_data=request_data, key=field)) for field in fields), + ("user_api_key_user_role", cls._resolve_user_role_from_request_data(request_data=request_data)), + ) + if not any(value for _, value in resolved): + return None + return {key: value for key, value in resolved if value} # mutable-ok: short-lived JSON payload dict + + @staticmethod + def _build_user_message(text: str) -> Mapping[str, Any]: + return {"role": "user", "content": text} # mutable-ok: short-lived JSON payload dict + + def _build_headers(self) -> Mapping[str, str]: + all_headers: Final = MappingProxyType( + { + "Content-Type": "application/json", + "X-Singulr-Gateway-Token": self.singulr_api_key, + "X-Singulr-Enforcement-Entity-Id": self.singulr_application_id, + "X-Singulr-Guardrail-Id": self.singulr_guardrail_id, + } + ) + return MappingProxyType({header: value for header, value in all_headers.items() if value}) + async def _call_api(self, payload: dict[str, object]) -> SingulrGuardrailResponse | None: endpoint: Final = f"{self.singulr_api_base}{_GUARD_ENDPOINT}" verbose_proxy_logger.debug("Singulr: %s", endpoint) @@ -168,7 +198,7 @@ class SingulrGuardrail(CustomGuardrail): if self.block_on_error: raise GuardrailRaisedException( guardrail_name=self.guardrail_name, - message=(f"Singulr API returned HTTP {exc.response.status_code}: {exc.response.text}"), + message=f"Singulr API returned HTTP {exc.response.status_code}: {exc.response.text}", ) from exc return None @@ -190,33 +220,218 @@ class SingulrGuardrail(CustomGuardrail): ) from exc return None - @log_guardrail_information - async def apply_guardrail( + async def _apply_guardrail_on_request( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, - input_type: str, - logging_obj: "LiteLLMLoggingObj | None" = None, + texts: Sequence[str], + structured_messages: Sequence[Any], + request_data: Mapping[str, Any], ) -> GenericGuardrailAPIInputs: - payload: Final = self._build_payload(request_data, inputs, input_type) - if not payload: - return inputs - - result: Final = await self._call_api(payload) - if result is None: - return inputs - - verbose_proxy_logger.debug( - "Singulr: should_block=%s blocking_due_to=%s", - result.should_block, - result.blocking_due_to, + messages: Final = ( + tuple(structured_messages) + if structured_messages + else tuple(self._build_user_message(text) for text in texts) ) - if result.should_block: + images: Final = inputs.get("images") + tools: Final = inputs.get("tools") + + if not messages and not images and not tools: + verbose_proxy_logger.debug("Singulr: No messages, images, or tools to check after filtering") + return inputs + + metadata: Final = self._build_metadata(request_data=request_data) + + singulr_req_obj = SingulrGuardrailPayload( + correlation_id=request_data.get("litellm_call_id"), + model_name=inputs.get("model"), + guardrail_scope="request", + messages=messages, + images=images, + tools=tools, + metadata=metadata, + ) + payload = singulr_req_obj.model_dump(mode="json") + guardrail_resp = await self._call_api(payload) + + if guardrail_resp is None: + return inputs + + if guardrail_resp.should_block: raise GuardrailRaisedException( guardrail_name=self.guardrail_name, - message=f"Blocked by Singulr: {result.blocking_due_to or 'unknown'}", + status_code=400, + message=f"Blocked by Singulr, Blocking due to {guardrail_resp.blocking_due_to or 'unknown'}", + blocked_content=True, + ) + return inputs + + @staticmethod + def _mcp_tool_name(request_data: Mapping[str, Any]) -> str | None: + return request_data.get("mcp_tool_name") or request_data.get("name") + + @staticmethod + def _mcp_arguments(request_data: Mapping[str, Any]) -> object: + arguments: Final = request_data.get("mcp_arguments") + return arguments if arguments is not None else request_data.get("arguments") + + @staticmethod + def _is_mcp_call(request_data: Mapping[str, Any], logging_obj: LiteLLMLoggingObj | None) -> bool: + call_type: Final = logging_obj.call_type if logging_obj is not None else request_data.get("call_type") + if call_type is not None: + return call_type == CallTypes.call_mcp_tool.value + model: Final = request_data.get("model") + return "mcp_tool_name" in request_data or (isinstance(model, str) and model.startswith(_MCP_MODEL_PREFIX)) + + async def _apply_guardrail_on_mcp_request(self, request_data: Mapping[str, Any]) -> None: + metadata: Final = self._build_metadata(request_data=request_data) + + singulr_mcp_obj = SingulrMcpGuardrailPayload( + guardrail_scope="mcp_request", + tool_name=self._mcp_tool_name(request_data), + tool_arguments=self._mcp_arguments(request_data), + mcp_server_name=request_data.get("mcp_server_name"), + metadata=metadata, + ) + payload = singulr_mcp_obj.model_dump(mode="json") + guardrail_resp = await self._call_api(payload) + + if guardrail_resp is None: + return + + if guardrail_resp.should_block: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + status_code=400, + message=f"Blocked by Singulr, Blocking due to {guardrail_resp.blocking_due_to or 'unknown'}", + blocked_content=True, + ) + + async def _apply_guardrail_on_mcp_response( + self, inputs: GenericGuardrailAPIInputs, texts: Sequence[str], request_data: Mapping[str, Any] + ) -> GenericGuardrailAPIInputs: + if not texts: + return inputs + + metadata: Final = self._build_metadata(request_data=request_data) + + singulr_mcp_obj = SingulrMcpGuardrailPayload( + model_name=request_data.get("model"), + guardrail_scope="mcp_response", + tool_result=texts, + metadata=metadata, + ) + payload = singulr_mcp_obj.model_dump(mode="json") + guardrail_resp = await self._call_api(payload) + + if guardrail_resp is None: + return inputs + + if guardrail_resp.should_block: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + status_code=400, + message=f"Blocked by Singulr, Blocking due to {guardrail_resp.blocking_due_to or 'unknown'}", blocked_content=True, ) return inputs + + @staticmethod + def _build_tool_call(tool_call: Mapping[str, Any]) -> "ToolCall | None": + tool_call_id: Final = tool_call.get("id") + fun: Final = tool_call.get("function") + if not tool_call_id or not fun: + return None + func_name: Final = fun.get("name") + args: Final = fun.get("arguments") + if not func_name or args is None: + return None + call_type: Final = tool_call.get("type") + return ToolCall( + id=tool_call_id, + type=call_type if isinstance(call_type, str) and call_type else "function", + function=ToolCallFunction( + name=func_name, + arguments=args if isinstance(args, str) else json.dumps(args, default=str), + ), + ) + + async def _apply_guardrail_on_response( + self, inputs: GenericGuardrailAPIInputs, texts: Sequence[str], request_data: Mapping[str, Any] + ) -> GenericGuardrailAPIInputs: + combined_texts: Final = "\n".join(texts) if texts else None + + tool_calls: Final = inputs.get("tool_calls", ()) + tool_calls_res: Final = tuple( + tool_call_res + for tool_call_res in (self._build_tool_call(tool_call) for tool_call in tool_calls) + if tool_call_res is not None + ) + + assistant_message: Final = AssistantMessage( + role="assistant", + content=combined_texts, + tool_calls=tool_calls_res, + ) + + metadata: Final = self._build_metadata(request_data=request_data) + + singulr_resp_obj = SingulrGuardrailPayload( + correlation_id=request_data.get("litellm_call_id"), + guardrail_scope="response", + model_name=request_data.get("model"), + messages=request_data.get("messages"), + images=inputs.get("images"), + response=assistant_message, + metadata=metadata, + ) + + payload = singulr_resp_obj.model_dump(mode="json") + guardrail_resp = await self._call_api(payload) + + if guardrail_resp is None: + return inputs + + if guardrail_resp.should_block: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + status_code=400, + message=f"Blocked by Singulr, Blocking due to {guardrail_resp.blocking_due_to or 'unknown'}", + blocked_content=True, + ) + return inputs + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, # mutable-ok: required by CustomGuardrail.apply_guardrail override signature + input_type: str, + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: + texts: Final = inputs.get("texts", ()) + structured_messages: Final = inputs.get("structured_messages", ()) + + verbose_proxy_logger.debug( + "Singulr Guardrail: apply_guardrail called with input_type=%s, texts=%d, structured_messages=%d", + input_type, + len(texts), + len(structured_messages), + ) + + is_mcp_call: Final = self._is_mcp_call(request_data, logging_obj) + if input_type == "request": + if is_mcp_call: + await self._apply_guardrail_on_mcp_request(request_data=request_data) + return inputs + return await self._apply_guardrail_on_request( + inputs=inputs, texts=texts, structured_messages=structured_messages, request_data=request_data + ) + elif input_type == "response": + if is_mcp_call: + return await self._apply_guardrail_on_mcp_response( + inputs=inputs, texts=texts, request_data=request_data + ) + return await self._apply_guardrail_on_response(inputs=inputs, texts=texts, request_data=request_data) + return inputs diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/singulr.py b/litellm/types/proxy/guardrails/guardrail_hooks/singulr.py index d0d19d191c1..ea1e6238181 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/singulr.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/singulr.py @@ -1,24 +1,53 @@ -from typing import Any +from collections.abc import Mapping, Sequence +from typing import Literal from pydantic import BaseModel, Field from .base import GuardrailConfigModel -class SingulrGuardrailRequest(BaseModel): - model: str | None = None - messages: list[dict[str, Any]] | None = None - tools: list[dict[str, Any]] | None = None - model_response: dict[str, Any] | None = None - litellm_metadata: dict[str, Any] | None = None +class ContentBlock(BaseModel): + type: str | None = None + text: str | None = None + + +class ToolCallFunction(BaseModel): + name: str + arguments: str + + +class ToolCall(BaseModel): + id: str + type: str = "function" + function: ToolCallFunction + + +class AssistantMessage(BaseModel): + role: Literal["assistant"] = "assistant" + content: str | Sequence[ContentBlock] | None = None + tool_calls: Sequence[ToolCall] | None = None class SingulrGuardrailPayload(BaseModel): - litellm_call_id: str | None = None - request_data: SingulrGuardrailRequest | None = None - input_type: str - is_playground_request: bool | None = None - playground_text: str | None = None + correlation_id: str | None = None + model_name: str | None = None + model_provider_name: str | None = None + guardrail_scope: str | None = None + messages: Sequence[Mapping[str, object]] | None = None + images: Sequence[str] | None = None + tools: Sequence[Mapping[str, object]] | None = None + response: AssistantMessage | None = None + metadata: Mapping[str, str] | None = None + + +class SingulrMcpGuardrailPayload(BaseModel): + model_name: str | None = None + guardrail_scope: str | None = None + tool_name: str | None = None + tool_arguments: object = None + mcp_server_name: str | None = None + tool_result: Sequence[str] | None = None + metadata: Mapping[str, str] | None = None class SingulrGuardrailResponse(BaseModel): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_singulr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_singulr.py index 14d8e90e027..b775d399b86 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_singulr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_singulr.py @@ -1,22 +1,22 @@ +import json from unittest.mock import MagicMock, patch import httpx import pytest +import litellm from litellm.exceptions import GuardrailRaisedException +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.singulr.singulr import SingulrGuardrail from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.singulr import ( SingulrGuardrailConfigModel, ) +from litellm.types.utils import ModelResponse -# --------------------------------------------------------------------------- -# Fixtures -# --------------------------------------------------------------------------- @pytest.fixture def singulr_guardrail(): - """Create a SingulrGuardrail instance with test credentials.""" return SingulrGuardrail( singulr_api_base="https://api.test.singulr.ai", singulr_api_key="test_token_1234", @@ -28,8 +28,26 @@ def singulr_guardrail(): ) +@pytest.fixture +def logging_only_guardrail(): + return SingulrGuardrail( + singulr_api_base="https://api.test.singulr.ai", + singulr_api_key="test_token_1234", + singulr_guardrail_id="test_guardrail_id", + singulr_application_id="test_enforcement_entity", + guardrail_name="test-singulr", + event_hook="logging_only", + default_on=True, + ) + + +def _logging_obj(call_type: str) -> MagicMock: + logging_obj = MagicMock() + logging_obj.call_type = call_type + return logging_obj + + def _make_response(body: dict) -> MagicMock: - """Build a mock httpx response with the given JSON body.""" mock = MagicMock() mock.json.return_value = body mock.raise_for_status = MagicMock() @@ -37,11 +55,6 @@ def _make_response(body: dict) -> MagicMock: return mock -# --------------------------------------------------------------------------- -# Configuration -# --------------------------------------------------------------------------- - - class TestSingulrConfiguration: def test_init_with_explicit_credentials(self): guardrail = SingulrGuardrail( @@ -55,6 +68,25 @@ class TestSingulrConfiguration: assert guardrail.singulr_guardrail_id == "id123" assert guardrail.singulr_application_id == "entity123" + def test_api_base_strips_surrounding_whitespace(self): + guardrail = SingulrGuardrail( + singulr_api_key="test_key", + singulr_api_base=" https://custom.api.local ", + ) + assert guardrail.singulr_api_base == "https://custom.api.local" + + def test_api_base_strips_trailing_slash(self): + guardrail = SingulrGuardrail(singulr_api_key="test_key", singulr_api_base="https://custom.api.local/") + assert guardrail.singulr_api_base == "https://custom.api.local" + + def test_non_local_http_api_base_raises(self): + with pytest.raises(ValueError, match="HTTPS"): + SingulrGuardrail(singulr_api_key="test_key", singulr_api_base="http://guardrails.singulr.ai") + + def test_localhost_http_api_base_is_allowed(self): + guardrail = SingulrGuardrail(singulr_api_key="test_key", singulr_api_base="http://localhost:8003") + assert guardrail.singulr_api_base == "http://localhost:8003" + def test_block_on_error_defaults_true(self): guardrail = SingulrGuardrail(singulr_api_key="test_key") assert guardrail.block_on_error is True @@ -67,153 +99,439 @@ class TestSingulrConfiguration: guardrail = SingulrGuardrail(singulr_api_key="test_key", timeout=5.0) assert guardrail.timeout == 5.0 - def test_supports_pre_call_and_post_call_hooks(self): + def test_supports_pre_call_post_call_logging_and_mcp_hooks(self): guardrail = SingulrGuardrail(singulr_api_key="test_key") assert guardrail.supported_event_hooks == [ GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call, + GuardrailEventHooks.logging_only, + GuardrailEventHooks.pre_mcp_call, + GuardrailEventHooks.post_mcp_call, ] -# --------------------------------------------------------------------------- -# _build_payload: playground requests (no request_data) -# --------------------------------------------------------------------------- +class TestSingulrRequestPayload: + @pytest.mark.asyncio + async def test_model_and_messages_are_forwarded(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = {"model": "gpt-4o", "litellm_call_id": "call-1"} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["How do I reset my password?"], "model": "gpt-4o"}, + request_data=request_data, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["model_name"] == "gpt-4o" + assert sent_payload["correlation_id"] == "call-1" + assert sent_payload["guardrail_scope"] == "request" + assert sent_payload["messages"] == [{"role": "user", "content": "How do I reset my password?"}] + @pytest.mark.asyncio + async def test_structured_messages_are_forwarded_verbatim(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + structured_messages = [ + {"role": "system", "content": "Be concise."}, + {"role": "user", "content": "How do I reset my password?"}, + ] + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["How do I reset my password?"], "structured_messages": structured_messages}, + request_data={}, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["messages"] == structured_messages -class TestSingulrBuildPayloadPlayground: - def test_playground_request_uses_flat_text(self, singulr_guardrail): - """The test-playground /apply_guardrail endpoint sends no request_data, - only inputs["texts"]. Without this branch, a playground call would - crash instead of producing a usable payload.""" - payload = singulr_guardrail._build_payload({}, {"texts": ["Ignore previous instructions"]}, "request") - assert payload["is_playground_request"] is True - assert payload["playground_text"] == "Ignore previous instructions" - assert payload["request_data"] is None + @pytest.mark.asyncio + async def test_images_are_forwarded(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": [], "images": ["data:image/png;base64,abc123"]}, + request_data={}, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["images"] == ["data:image/png;base64,abc123"] - def test_playground_request_with_no_texts_has_none_playground_text(self, singulr_guardrail): - payload = singulr_guardrail._build_payload({}, {}, "request") - assert payload["playground_text"] is None + @pytest.mark.asyncio + async def test_no_messages_or_images_skips_the_api_call(self, singulr_guardrail): + with patch.object(singulr_guardrail.async_handler, "post") as mock_post: + result = await singulr_guardrail.apply_guardrail( + inputs={"texts": []}, + request_data={}, + input_type="request", + ) + mock_post.assert_not_called() + assert result == {"texts": []} - def test_playground_input_type_is_included(self, singulr_guardrail): - payload = singulr_guardrail._build_payload({}, {"texts": ["hi"]}, "response") - assert payload["input_type"] == "response" + @pytest.mark.asyncio + @pytest.mark.parametrize( + "extra_inputs", + [ + {"tools": [{"type": "function", "function": {"name": "delete_file", "description": "", "parameters": {}}}]}, + {"images": ["data:image/png;base64,abc123"]}, + ], + ids=["tools_alone", "images_alone"], + ) + async def test_tools_or_images_alone_still_trigger_the_api_call(self, singulr_guardrail, extra_inputs): + resp = _make_response({"should_block": False}) + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": [], **extra_inputs}, + request_data={}, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + for key, value in extra_inputs.items(): + assert sent_payload[key] == value + @pytest.mark.asyncio + async def test_tools_are_forwarded(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + tools = [ + { + "type": "function", + "function": {"name": "search_docs", "description": "Search internal docs", "parameters": {}}, + } + ] + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["How do I reset my password?"], "tools": tools}, + request_data={}, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["tools"] == tools -# --------------------------------------------------------------------------- -# _build_payload: real proxy requests (request_data present) -# --------------------------------------------------------------------------- + @pytest.mark.asyncio + async def test_responses_api_mcp_tools_are_forwarded(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + tools = [ + { + "type": "mcp", + "server_label": "docs-server", + "server_url": "https://mcp.example.com", + "allowed_tools": ["search_docs"], + } + ] + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["How do I reset my password?"], "tools": tools}, + request_data={}, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["tools"] == tools + @pytest.mark.asyncio + async def test_user_api_key_alias_is_forwarded_in_metadata(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = {"litellm_metadata": {"user_api_key_alias": "my-key-alias"}} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=request_data, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["metadata"] == {"user_api_key_alias": "my-key-alias"} -class TestSingulrBuildPayloadRequestData: - def test_model_messages_and_tools_are_forwarded(self, singulr_guardrail): + @pytest.mark.asyncio + async def test_falls_back_to_regular_metadata_for_key_alias(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = {"metadata": {"user_api_key_alias": "fallback-alias"}} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=request_data, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["metadata"] == {"user_api_key_alias": "fallback-alias"} + + @pytest.mark.asyncio + async def test_user_api_key_user_id_is_forwarded_in_metadata(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = {"litellm_metadata": {"user_api_key_user_id": "my-user-id"}} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=request_data, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["metadata"] == {"user_api_key_user_id": "my-user-id"} + + @pytest.mark.asyncio + async def test_falls_back_to_regular_metadata_for_user_id(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = {"metadata": {"user_api_key_user_id": "fallback-user-id"}} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=request_data, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["metadata"] == {"user_api_key_user_id": "fallback-user-id"} + + @pytest.mark.asyncio + async def test_user_api_key_user_email_is_forwarded_in_metadata(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = {"litellm_metadata": {"user_api_key_user_email": "user@example.com"}} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=request_data, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["metadata"] == {"user_api_key_user_email": "user@example.com"} + + @pytest.mark.asyncio + async def test_user_api_key_organization_alias_is_forwarded_in_metadata(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = {"litellm_metadata": {"user_api_key_org_alias": "Acme Org"}} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=request_data, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["metadata"] == {"user_api_key_org_alias": "Acme Org"} + + @pytest.mark.asyncio + async def test_user_api_key_team_alias_is_forwarded_in_metadata(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = {"litellm_metadata": {"user_api_key_team_alias": "AI Content Security Team"}} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=request_data, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["metadata"] == {"user_api_key_team_alias": "AI Content Security Team"} + + @pytest.mark.asyncio + async def test_user_api_key_org_id_is_forwarded_in_metadata(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = {"litellm_metadata": {"user_api_key_org_id": "org-123"}} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=request_data, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["metadata"] == {"user_api_key_org_id": "org-123"} + + @pytest.mark.asyncio + async def test_user_api_key_team_id_is_forwarded_in_metadata(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = {"litellm_metadata": {"user_api_key_team_id": "team-456"}} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=request_data, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["metadata"] == {"user_api_key_team_id": "team-456"} + + @pytest.mark.asyncio + async def test_user_api_key_user_role_is_forwarded_in_metadata(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY) + request_data = {"litellm_metadata": {"user_api_key_auth": auth}} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=request_data, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["metadata"] == {"user_api_key_user_role": LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value} + + @pytest.mark.asyncio + async def test_no_user_role_available_omits_role_from_metadata(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = {"litellm_metadata": {"user_api_key_alias": "my-key-alias"}} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=request_data, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert "user_api_key_user_role" not in sent_payload["metadata"] + + @pytest.mark.asyncio + async def test_all_user_metadata_fields_forwarded_together(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY) request_data = { - "model": "gpt-4o", - "messages": [{"role": "user", "content": "How do I reset my password?"}], - "tools": [{"type": "function", "function": {"name": "get_weather"}}], + "litellm_metadata": { + "user_api_key_alias": "my-key-alias", + "user_api_key_user_id": "my-user-id", + "user_api_key_user_email": "user@example.com", + "user_api_key_org_id": "org-123", + "user_api_key_org_alias": "Acme Org", + "user_api_key_team_id": "team-456", + "user_api_key_team_alias": "AI Content Security Team", + "user_api_key_auth": auth, + } + } + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=request_data, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["metadata"] == { + "user_api_key_alias": "my-key-alias", + "user_api_key_user_id": "my-user-id", + "user_api_key_user_email": "user@example.com", + "user_api_key_org_id": "org-123", + "user_api_key_org_alias": "Acme Org", + "user_api_key_team_id": "team-456", + "user_api_key_team_alias": "AI Content Security Team", + "user_api_key_user_role": LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, } - payload = singulr_guardrail._build_payload(request_data, {"texts": []}, "request") - assert payload["request_data"]["model"] == "gpt-4o" - assert payload["request_data"]["messages"] == request_data["messages"] - assert payload["request_data"]["tools"] == request_data["tools"] - assert payload["is_playground_request"] is None - def test_model_response_absent_on_request_side(self, singulr_guardrail): - """The response hasn't happened yet at request time, so model_response - must not be forwarded even if request_data carries a stale response - object from a previous call.""" - from litellm.types.utils import ModelResponse + @pytest.mark.asyncio + async def test_no_key_alias_available_sends_no_metadata(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data={}, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["metadata"] is None - request_data = {"model": "gpt-4o", "response": ModelResponse()} - payload = singulr_guardrail._build_payload(request_data, {"texts": []}, "request") - assert payload["request_data"]["model_response"] is None - def test_model_response_is_forwarded_and_json_serializable(self, singulr_guardrail): - """Regression: request_data["response"] is a ModelResponse (pydantic) - object containing nested non-JSON-safe values (e.g. a `created` - unix timestamp is fine, but nested pydantic submodels are not plain - dicts). Without mode="json" on both the inner and outer dumps, this - payload cannot be sent via httpx's json= kwarg.""" - import json as _json - - from litellm.types.utils import Choices, Message, ModelResponse, Usage - - response = ModelResponse( - choices=[Choices(message=Message(role="assistant", content="Go to settings."))], - usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), - ) - request_data = {"model": "gpt-4o", "response": response} - payload = singulr_guardrail._build_payload(request_data, {"texts": ["Go to settings."]}, "response") - - # Must not raise - this is what httpx's json= kwarg effectively does. - serialized = _json.dumps(payload) - assert "Go to settings." in serialized - assert payload["request_data"]["model_response"]["choices"][0]["message"]["content"] == "Go to settings." - - def test_model_requested_tool_calls_are_forwarded_in_model_response(self, singulr_guardrail): - """Tool calls the model requests arrive inside response.choices[].message.tool_calls. - They must survive the dump so Singulr can inspect what tools the - model is trying to invoke.""" - from litellm.types.utils import Choices, Message, ModelResponse - - response = ModelResponse( - choices=[ - Choices( - message=Message( - role="assistant", - content=None, - tool_calls=[ - { - "id": "call_1", - "type": "function", - "function": {"name": "get_current_time", "arguments": "{}"}, - } - ], - ) - ) +class TestSingulrResponsePayload: + @pytest.mark.asyncio + async def test_assistant_text_and_tool_calls_are_forwarded(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + inputs = { + "texts": ["Go to settings."], + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_current_time", "arguments": "{}"}, + } ], - ) - request_data = {"model": "gpt-4o", "response": response} - payload = singulr_guardrail._build_payload(request_data, {"texts": []}, "response") - - tool_calls = payload["request_data"]["model_response"]["choices"][0]["message"]["tool_calls"] - assert tool_calls[0]["function"]["name"] == "get_current_time" - - def test_litellm_metadata_is_forwarded(self, singulr_guardrail): - request_data = {"model": "gpt-4o", "litellm_metadata": {"user_api_key_hash": "abc123"}} - payload = singulr_guardrail._build_payload(request_data, {"texts": []}, "request") - assert payload["request_data"]["litellm_metadata"] == {"user_api_key_hash": "abc123"} - - def test_internal_logging_object_is_not_forwarded(self, singulr_guardrail): - """Regression: request_data can carry internal proxy objects (e.g. the - Logging instance) that aren't JSON-serializable at all. _build_payload - must only pull known request/response fields out of request_data, - not dump it wholesale, or this crashes on every real proxy call.""" - import json as _json - - class _NotSerializable: - pass - - request_data = { - "model": "gpt-4o", - "messages": [{"role": "user", "content": "hi"}], - "litellm_logging_obj": _NotSerializable(), } - payload = singulr_guardrail._build_payload(request_data, {"texts": ["hi"]}, "request") + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="response", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["guardrail_scope"] == "response" + assert sent_payload["response"]["content"] == "Go to settings." + assert sent_payload["response"]["tool_calls"][0]["function"]["name"] == "get_current_time" - # Must not raise. - _json.dumps(payload) - assert "litellm_logging_obj" not in payload["request_data"] + @pytest.mark.asyncio + async def test_response_images_are_forwarded(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + inputs = {"texts": ["ok"], "images": ["data:image/png;base64,xyz"]} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="response", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["images"] == ["data:image/png;base64,xyz"] + @pytest.mark.asyncio + async def test_incomplete_tool_calls_are_dropped(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + inputs = { + "texts": [], + "tool_calls": [ + {"id": None, "type": "function", "function": {"name": "f", "arguments": "{}"}}, + {"id": "call_2", "type": "function", "function": None}, + {"id": "call_3", "type": "function", "function": {"name": None, "arguments": "{}"}}, + {"id": "call_4", "type": "function", "function": {"name": "f", "arguments": None}}, + ], + } + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="response", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["response"]["tool_calls"] == [] -# --------------------------------------------------------------------------- -# Allow / block decisions -# --------------------------------------------------------------------------- + @pytest.mark.asyncio + @pytest.mark.parametrize( + "raw_type, expected_type", + [(None, "function"), ("custom", "custom")], + ids=["type_missing", "type_not_function"], + ) + async def test_tool_call_type_other_than_function_is_still_scanned( + self, singulr_guardrail, raw_type, expected_type + ): + resp = _make_response({"should_block": False}) + tool_call = {"id": "call_1", "function": {"name": "get_current_time", "arguments": "{}"}} + inputs = { + "texts": [], + "tool_calls": [tool_call if raw_type is None else {**tool_call, "type": raw_type}], + } + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="response") + sent_tool_calls = mock_post.call_args.kwargs["json"]["response"]["tool_calls"] + assert [call["type"] for call in sent_tool_calls] == [expected_type] + assert sent_tool_calls[0]["function"]["name"] == "get_current_time" + + @pytest.mark.asyncio + async def test_non_string_tool_call_arguments_are_serialized(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + inputs = { + "texts": [], + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "rm", "arguments": {"path": "/etc/passwd"}}} + ], + } + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="response") + sent_tool_calls = mock_post.call_args.kwargs["json"]["response"]["tool_calls"] + assert json.loads(sent_tool_calls[0]["function"]["arguments"]) == {"path": "/etc/passwd"} + + @pytest.mark.asyncio + async def test_block_verdict_still_raises_for_a_non_function_tool_call(self, singulr_guardrail): + resp = _make_response({"should_block": True, "blocking_due_to": "dangerous_tool"}) + inputs = { + "texts": [], + "tool_calls": [{"id": "call_1", "function": {"name": "rm", "arguments": "{}"}}], + } + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp): + with pytest.raises(GuardrailRaisedException) as exc_info: + await singulr_guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="response") + assert "dangerous_tool" in str(exc_info.value) class TestSingulrAllowAction: @pytest.mark.asyncio - async def test_allow_returns_inputs_unchanged(self, singulr_guardrail): - resp = _make_response({"should_block": False}) + @pytest.mark.parametrize( + "guard_response", + [{"should_block": False}, {}], + ids=["should_block_false", "should_block_omitted"], + ) + async def test_should_block_falsy_returns_inputs_unchanged_on_request(self, singulr_guardrail, guard_response): + resp = _make_response(guard_response) inputs = {"texts": ["How do I reset my password?"]} with patch.object(singulr_guardrail.async_handler, "post", return_value=resp): result = await singulr_guardrail.apply_guardrail( @@ -223,18 +541,68 @@ class TestSingulrAllowAction: ) assert result is inputs + @pytest.mark.asyncio + async def test_should_block_false_returns_inputs_unchanged_on_response(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + inputs = {"texts": ["Here is your answer."]} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp): + result = await singulr_guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="response", + ) + assert result is inputs + + @pytest.mark.asyncio + async def test_response_returns_inputs_unchanged_when_api_unreachable_and_block_on_error_false(self): + guardrail = SingulrGuardrail( + singulr_api_base="https://api.test.singulr.ai", + singulr_api_key="test_token_1234", + guardrail_name="test-singulr", + block_on_error=False, + ) + inputs = {"texts": ["Here is your answer."]} + with patch.object(guardrail.async_handler, "post", side_effect=httpx.TransportError("unreachable")): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="response", + ) + assert result is inputs + + @pytest.mark.asyncio + async def test_explicit_null_verdict_fails_closed_by_default(self, singulr_guardrail): + resp = _make_response({"should_block": None}) + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp): + with pytest.raises(GuardrailRaisedException, match="invalid response"): + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + @pytest.mark.asyncio + async def test_explicit_null_verdict_fails_open_when_block_on_error_false(self): + guardrail = SingulrGuardrail( + singulr_api_base="https://api.test.singulr.ai", + singulr_api_key="test_token_1234", + guardrail_name="test-singulr", + block_on_error=False, + ) + resp = _make_response({"should_block": None}) + inputs = {"texts": ["hi"]} + with patch.object(guardrail.async_handler, "post", return_value=resp): + assert await guardrail._call_api({"guardrail_scope": "request"}) is None + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={"model": "gpt-4o"}, input_type="request" + ) + assert result is inputs + class TestSingulrBlockAction: @pytest.mark.asyncio - async def test_block_raises_guardrail_exception(self, singulr_guardrail): - """Regression: a should_block=True response must stop the request - instead of silently letting it through.""" - resp = _make_response( - { - "should_block": True, - "blocking_due_to": "PII Information detected", - } - ) + async def test_should_block_true_raises_on_request(self, singulr_guardrail): + resp = _make_response({"should_block": True, "blocking_due_to": "PII Information detected"}) with patch.object(singulr_guardrail.async_handler, "post", return_value=resp): with pytest.raises(GuardrailRaisedException) as exc_info: await singulr_guardrail.apply_guardrail( @@ -243,6 +611,20 @@ class TestSingulrBlockAction: input_type="request", ) assert "PII Information detected" in str(exc_info.value) + assert exc_info.value.blocked_content is True + + @pytest.mark.asyncio + async def test_should_block_true_raises_on_response(self, singulr_guardrail): + resp = _make_response({"should_block": True, "blocking_due_to": "Toxic content detected"}) + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp): + with pytest.raises(GuardrailRaisedException) as exc_info: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["Here is something toxic."]}, + request_data={}, + input_type="response", + ) + assert "Toxic content detected" in str(exc_info.value) + assert exc_info.value.blocked_content is True @pytest.mark.asyncio async def test_block_without_reason_uses_unknown_placeholder(self, singulr_guardrail): @@ -256,17 +638,409 @@ class TestSingulrBlockAction: ) -# --------------------------------------------------------------------------- -# HTTP call wiring (endpoint, timeout, headers) -# --------------------------------------------------------------------------- +class TestSingulrMcpRequest: + @pytest.mark.asyncio + async def test_mcp_tool_name_routes_to_mcp_request_payload(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = { + "mcp_tool_name": "search_docs", + "mcp_arguments": {"query": "reset password"}, + "mcp_server_name": "docs-server", + } + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + result = await singulr_guardrail.apply_guardrail( + inputs={"texts": []}, + request_data=request_data, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["guardrail_scope"] == "mcp_request" + assert sent_payload["tool_name"] == "search_docs" + assert sent_payload["tool_arguments"] == {"query": "reset password"} + assert sent_payload["mcp_server_name"] == "docs-server" + assert result == {"texts": []} + + @pytest.mark.asyncio + async def test_mcp_request_should_block_true_raises(self, singulr_guardrail): + resp = _make_response({"should_block": True, "blocking_due_to": "Disallowed tool"}) + request_data = {"mcp_tool_name": "delete_file", "mcp_arguments": {"path": "/etc/passwd"}} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp): + with pytest.raises(GuardrailRaisedException, match="Disallowed tool") as exc_info: + await singulr_guardrail.apply_guardrail( + inputs={"texts": []}, + request_data=request_data, + input_type="request", + ) + assert exc_info.value.blocked_content is True + + @pytest.mark.asyncio + async def test_mcp_request_is_a_noop_when_api_unreachable_and_block_on_error_false(self): + guardrail = SingulrGuardrail( + singulr_api_base="https://api.test.singulr.ai", + singulr_api_key="test_token_1234", + guardrail_name="test-singulr", + block_on_error=False, + ) + request_data = {"mcp_tool_name": "search_docs", "mcp_arguments": {"query": "reset password"}} + with patch.object(guardrail.async_handler, "post", side_effect=httpx.TransportError("unreachable")): + result = await guardrail.apply_guardrail( + inputs={"texts": []}, + request_data=request_data, + input_type="request", + ) + assert result == {"texts": []} + + @pytest.mark.asyncio + async def test_mcp_rest_body_shape_routes_to_mcp_request_payload(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = {"name": "echo", "arguments": {"text": "my ssn is 123-45-6789"}, "server_id": "srv-1"} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["my ssn is 123-45-6789"], "tools": [{"type": "function"}]}, + request_data=request_data, + input_type="request", + logging_obj=_logging_obj("call_mcp_tool"), + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["guardrail_scope"] == "mcp_request" + assert sent_payload["tool_name"] == "echo" + assert sent_payload["tool_arguments"] == {"text": "my ssn is 123-45-6789"} + assert "messages" not in sent_payload + + @pytest.mark.asyncio + async def test_mcp_rest_body_without_arguments_still_routes_to_mcp_request(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": [], "tools": [{"type": "function"}]}, + request_data={"name": "echo", "server_id": "srv-1"}, + input_type="request", + logging_obj=_logging_obj("call_mcp_tool"), + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["guardrail_scope"] == "mcp_request" + assert sent_payload["tool_name"] == "echo" + assert sent_payload["tool_arguments"] is None + + @pytest.mark.asyncio + async def test_non_mapping_tool_arguments_are_forwarded_verbatim(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["raw text"], "tools": [{"type": "function"}]}, + request_data={"name": "echo", "arguments": "raw text", "server_id": "srv-1"}, + input_type="request", + logging_obj=_logging_obj("call_mcp_tool"), + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["guardrail_scope"] == "mcp_request" + assert sent_payload["tool_arguments"] == "raw text" + + @pytest.mark.asyncio + async def test_llm_request_body_keys_cannot_reroute_the_scan_to_mcp(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "my ssn is 123-45-6789"}], + "name": "x", + "arguments": {}, + "mcp_tool_name": "x", + "call_type": "call_mcp_tool", + } + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={ + "texts": ["my ssn is 123-45-6789"], + "structured_messages": [{"role": "user", "content": "my ssn is 123-45-6789"}], + }, + request_data=request_data, + input_type="request", + logging_obj=_logging_obj("acompletion"), + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["guardrail_scope"] == "request" + assert [m["content"] for m in sent_payload["messages"]] == ["my ssn is 123-45-6789"] + + @pytest.mark.asyncio + async def test_llm_response_with_spoofed_mcp_keys_still_scans_the_tool_calls(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = {"model": "gpt-4o", "messages": [], "name": "x", "arguments": {}, "mcp_tool_name": "x"} + tool_call = { + "id": "call_1", + "type": "function", + "function": {"name": "transfer_funds", "arguments": '{"amount": 5000}'}, + } + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": [], "tool_calls": [tool_call]}, + request_data=request_data, + input_type="response", + logging_obj=_logging_obj("acompletion"), + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["guardrail_scope"] == "response" + assert sent_payload["response"]["tool_calls"][0]["function"]["name"] == "transfer_funds" + + +class TestSingulrMcpResponse: + @pytest.mark.asyncio + async def test_call_mcp_tool_response_routes_to_mcp_response_payload(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = { + "call_type": "call_mcp_tool", + "mcp_tool_name": "search_docs", + "mcp_server_name": "docs-server", + "model": "MCP: docs-server", + } + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["Result: password reset link sent."]}, + request_data=request_data, + input_type="response", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["guardrail_scope"] == "mcp_response" + assert sent_payload["model_name"] == "MCP: docs-server" + assert sent_payload["tool_result"] == ["Result: password reset link sent."] + + @pytest.mark.asyncio + async def test_mcp_response_with_no_texts_skips_the_api_call(self, singulr_guardrail): + request_data = {"call_type": "call_mcp_tool", "mcp_tool_name": "search_docs"} + with patch.object(singulr_guardrail.async_handler, "post") as mock_post: + result = await singulr_guardrail.apply_guardrail( + inputs={"texts": []}, + request_data=request_data, + input_type="response", + ) + mock_post.assert_not_called() + assert result == {"texts": []} + + @pytest.mark.asyncio + async def test_mcp_response_should_block_true_raises(self, singulr_guardrail): + resp = _make_response({"should_block": True, "blocking_due_to": "Sensitive tool output"}) + request_data = {"call_type": "call_mcp_tool", "mcp_tool_name": "search_docs"} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp): + with pytest.raises(GuardrailRaisedException, match="Sensitive tool output") as exc_info: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["leaked secret"]}, + request_data=request_data, + input_type="response", + ) + assert exc_info.value.blocked_content is True + + @pytest.mark.asyncio + async def test_mcp_response_resolves_metadata_from_nested_litellm_params(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY) + request_data = { + "call_type": "call_mcp_tool", + "mcp_tool_name": "search_docs", + "litellm_params": { + "metadata": { + "user_api_key_alias": "my-key-alias", + "user_api_key_user_id": "my-user-id", + "user_api_key_user_email": "user@example.com", + "user_api_key_org_id": "org-123", + "user_api_key_org_alias": "Acme Org", + "user_api_key_team_id": "team-456", + "user_api_key_team_alias": "AI Content Security Team", + "user_api_key_auth": auth, + } + }, + } + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["Result: password reset link sent."]}, + request_data=request_data, + input_type="response", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["metadata"] == { + "user_api_key_alias": "my-key-alias", + "user_api_key_user_id": "my-user-id", + "user_api_key_user_email": "user@example.com", + "user_api_key_org_id": "org-123", + "user_api_key_org_alias": "Acme Org", + "user_api_key_team_id": "team-456", + "user_api_key_team_alias": "AI Content Security Team", + "user_api_key_user_role": LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, + } + + @pytest.mark.asyncio + async def test_mcp_response_prefers_top_level_metadata_over_nested_litellm_params(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = { + "call_type": "call_mcp_tool", + "mcp_tool_name": "search_docs", + "litellm_metadata": {"user_api_key_alias": "top-level-alias"}, + "litellm_params": {"metadata": {"user_api_key_alias": "nested-alias"}}, + } + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=request_data, + input_type="response", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["metadata"] == {"user_api_key_alias": "top-level-alias"} + + @pytest.mark.asyncio + async def test_mcp_response_returns_inputs_unchanged_when_api_unreachable_and_block_on_error_false(self): + guardrail = SingulrGuardrail( + singulr_api_base="https://api.test.singulr.ai", + singulr_api_key="test_token_1234", + guardrail_name="test-singulr", + block_on_error=False, + ) + request_data = {"call_type": "call_mcp_tool", "mcp_tool_name": "search_docs"} + inputs = {"texts": ["leaked secret"]} + with patch.object(guardrail.async_handler, "post", side_effect=httpx.TransportError("unreachable")): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + ) + assert result is inputs + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("request_data", "logging_obj"), + [ + ({"call_type": "call_mcp_tool", "model": "MCP: echo"}, None), + ({"model": "MCP: echo"}, None), + ({"name": "echo", "arguments": {"text": "hi"}}, _logging_obj("call_mcp_tool")), + ], + ids=["post_mcp_call_model_call_details", "logging_only_scratch_request", "rest_pre_call_logger"], + ) + async def test_mcp_response_is_detected_from_each_producer(self, singulr_guardrail, request_data, logging_obj): + resp = _make_response({"should_block": False}) + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["tool output"]}, + request_data=request_data, + input_type="response", + logging_obj=logging_obj, + ) + assert mock_post.call_args.kwargs["json"]["guardrail_scope"] == "mcp_response" + + +class TestSingulrApplyGuardrailDispatch: + @pytest.mark.asyncio + async def test_unknown_input_type_returns_inputs_unchanged(self, singulr_guardrail): + with patch.object(singulr_guardrail.async_handler, "post") as mock_post: + inputs = {"texts": ["hi"]} + result = await singulr_guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="unsupported", + ) + mock_post.assert_not_called() + assert result is inputs + + +class TestSingulrLoggingHook: + @staticmethod + def _logged_call(**overrides): + kwargs = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "litellm_call_id": "call-1", + "litellm_params": {"metadata": {"user_api_key_alias": "my-key-alias", "user_api_key_org_id": "org-123"}}, + "standard_logging_object": {"guardrail_information": []}, + } + return {**kwargs, **overrides} + + @pytest.mark.asyncio + async def test_scans_request_then_response_as_an_assistant_message(self, logging_only_guardrail): + resp = _make_response({"should_block": False}) + result = ModelResponse( + choices=[{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hello there"}}] + ) + with patch.object(logging_only_guardrail.async_handler, "post", return_value=resp) as mock_post: + updated_kwargs, returned = await logging_only_guardrail.async_logging_hook( + kwargs=self._logged_call(), result=result, call_type="acompletion" + ) + + assert returned is result + scopes = [call.kwargs["json"]["guardrail_scope"] for call in mock_post.call_args_list] + assert scopes == ["request", "response"] + request_payload = mock_post.call_args_list[0].kwargs["json"] + response_payload = mock_post.call_args_list[1].kwargs["json"] + assert request_payload["messages"] == [{"role": "user", "content": "hi"}] + assert request_payload["correlation_id"] == "call-1" + assert response_payload["response"] == {"role": "assistant", "content": "hello there", "tool_calls": []} + expected_metadata = {"user_api_key_alias": "my-key-alias", "user_api_key_org_id": "org-123"} + assert request_payload["metadata"] == expected_metadata + assert response_payload["metadata"] == expected_metadata + statuses = [ + entry["guardrail_status"] for entry in updated_kwargs["standard_logging_object"]["guardrail_information"] + ] + assert statuses == ["success", "success"] + + @pytest.mark.asyncio + async def test_block_verdict_is_recorded_as_intervened_without_failing_the_call(self, logging_only_guardrail): + resp = _make_response({"should_block": True, "blocking_due_to": "pii"}) + with patch.object(logging_only_guardrail.async_handler, "post", return_value=resp): + updated_kwargs, returned = await logging_only_guardrail.async_logging_hook( + kwargs=self._logged_call(messages=[{"role": "user", "content": "my ssn is 123-45-6789"}]), + result=None, + call_type="acompletion", + ) + assert returned is None + entries = updated_kwargs["standard_logging_object"]["guardrail_information"] + assert entries[0]["guardrail_status"] == "guardrail_intervened" + assert entries[0]["guardrail_mode"] == "logging_only" + assert "Blocking due to pii" in str(entries[0]["guardrail_response"]) + + @pytest.mark.asyncio + async def test_vendor_timeout_is_recorded_as_failed_to_respond(self, logging_only_guardrail): + timeout = litellm.Timeout("Singulr timed out", model="gpt-4o", llm_provider="singulr") + with patch.object(logging_only_guardrail.async_handler, "post", side_effect=timeout): + updated_kwargs, returned = await logging_only_guardrail.async_logging_hook( + kwargs=self._logged_call(), result=None, call_type="acompletion" + ) + assert returned is None + entries = updated_kwargs["standard_logging_object"]["guardrail_information"] + assert [entry["guardrail_status"] for entry in entries] == ["guardrail_failed_to_respond"] + assert "timed out" in str(entries[0]["guardrail_response"]) + + @pytest.mark.asyncio + async def test_mcp_tool_result_is_scanned_as_mcp_response(self, logging_only_guardrail): + from mcp.types import CallToolResult, TextContent + + resp = _make_response({"should_block": False}) + result = CallToolResult(content=[TextContent(type="text", text="ssn 123-45-6789")]) + with patch.object(logging_only_guardrail.async_handler, "post", return_value=resp) as mock_post: + await logging_only_guardrail.async_logging_hook( + kwargs=self._logged_call(model="MCP: get_customer_record", messages=None), + result=result, + call_type="call_mcp_tool", + ) + payloads = [call.kwargs["json"] for call in mock_post.call_args_list] + assert [payload["guardrail_scope"] for payload in payloads] == ["mcp_response"] + assert payloads[0]["tool_result"] == ["ssn 123-45-6789"] + assert payloads[0]["model_name"] == "MCP: get_customer_record" + + def test_sync_logging_hook_never_calls_singulr(self, logging_only_guardrail): + from concurrent.futures import ThreadPoolExecutor + + kwargs = {"messages": [{"role": "user", "content": "hi"}], "standard_logging_object": {}} + + def _run(): + with patch.object(logging_only_guardrail.async_handler, "post") as mock_post: + returned = logging_only_guardrail.logging_hook(kwargs=kwargs, result=None, call_type="acompletion") + mock_post.assert_not_called() + return returned + + with ThreadPoolExecutor(max_workers=1) as pool: + returned_kwargs, returned_result = pool.submit(_run).result() + assert returned_result is None + assert returned_kwargs == {"messages": [{"role": "user", "content": "hi"}], "standard_logging_object": {}} class TestSingulrRequestWiring: @pytest.mark.asyncio - async def test_sends_configured_timeout(self): - """litellm_params.timeout must reach the httpx call so operators can - tighten or loosen the latency budget instead of being stuck with a - hardcoded 30s regardless of configuration.""" + async def test_sends_configured_timeout_and_calls_the_guard_endpoint(self): guardrail = SingulrGuardrail( singulr_api_key="test_key", singulr_api_base="https://api.test.singulr.ai", @@ -279,7 +1053,9 @@ class TestSingulrRequestWiring: request_data={}, input_type="request", ) - assert mock_post.call_args.kwargs["timeout"] == 5.0 + call_kwargs = mock_post.call_args.kwargs + assert call_kwargs["timeout"] == 5.0 + assert call_kwargs["url"] == "https://api.test.singulr.ai/api/v1/ai-gateway/litellm-v2" class TestSingulrBuildHeaders: @@ -300,11 +1076,6 @@ class TestSingulrBuildHeaders: assert "X-Singulr-Guardrail-Id" not in headers -# --------------------------------------------------------------------------- -# Non-JSON / malformed response handling -# --------------------------------------------------------------------------- - - class TestSingulrInvalidResponse: @pytest.mark.asyncio async def test_non_json_response_block_on_error_false_returns_inputs(self): @@ -349,18 +1120,17 @@ class TestSingulrInvalidResponse: @pytest.mark.asyncio async def test_response_missing_expected_fields_block_on_error_true_raises(self): - """Regression: a response body that fails SingulrGuardrailResponse - validation (e.g. should_block is a string, not a bool) must raise - GuardrailRaisedException instead of letting pydantic.ValidationError - propagate unhandled.""" guardrail = SingulrGuardrail( singulr_api_base="https://api.test.singulr.ai", singulr_api_key="test_token_1234", guardrail_name="test-singulr", block_on_error=True, ) - resp = _make_response({"should_block": "not-a-bool"}) - with patch.object(guardrail.async_handler, "post", return_value=resp): + mock_resp = MagicMock() + mock_resp.raise_for_status = MagicMock() + mock_resp.json.side_effect = ValueError("not valid json") + + with patch.object(guardrail.async_handler, "post", return_value=mock_resp): with pytest.raises(GuardrailRaisedException): await guardrail.apply_guardrail( inputs={"texts": ["test"]}, @@ -369,11 +1139,6 @@ class TestSingulrInvalidResponse: ) -# --------------------------------------------------------------------------- -# Transport error handling -# --------------------------------------------------------------------------- - - class TestSingulrTransportError: @pytest.mark.asyncio async def test_remote_protocol_error_block_on_error_false_returns_inputs(self): @@ -417,11 +1182,6 @@ class TestSingulrTransportError: ) -# --------------------------------------------------------------------------- -# HTTP status error handling -# --------------------------------------------------------------------------- - - class TestSingulrHttpStatusError: @pytest.mark.asyncio async def test_http_error_message_names_status_code_not_unreachable(self): @@ -472,19 +1232,12 @@ class TestSingulrHttpStatusError: assert result is inputs -# --------------------------------------------------------------------------- -# Config model -# --------------------------------------------------------------------------- - - class TestSingulrConfigModel: def test_ui_friendly_name(self): assert SingulrGuardrailConfigModel.ui_friendly_name() == "Singulr" - -# --------------------------------------------------------------------------- -# Initializer and registry -# --------------------------------------------------------------------------- + def test_get_config_model_returns_singulr_config_model(self): + assert SingulrGuardrail.get_config_model() is SingulrGuardrailConfigModel class TestSingulrInitializer: @@ -496,11 +1249,6 @@ class TestSingulrInitializer: assert callable(initialize_guardrail) def test_initialize_guardrail_reads_singulr_prefixed_fields(self): - """Regression: the UI config form (and YAML config) populate the - singulr_-prefixed fields declared on SingulrGuardrailConfigModel, not - the generic api_base/api_key fields. initialize_guardrail must read - those, or a UI-configured singulr_api_base is silently ignored and - the guardrail falls back to the localhost default.""" from litellm.proxy.guardrails.guardrail_hooks.singulr import ( initialize_guardrail, ) @@ -525,10 +1273,6 @@ class TestSingulrInitializer: assert cb.singulr_guardrail_id == "configured_guardrail_id" def test_initialize_guardrail_wires_timeout(self): - """BaseLitellmParams.timeout exists so operators can override the - per-request latency budget. initialize_guardrail must forward it to - SingulrGuardrail instead of leaving every deployment stuck on the - hardcoded default regardless of configuration.""" from litellm.proxy.guardrails.guardrail_hooks.singulr import ( initialize_guardrail, )