fix(guardrails): pydantic type for payload

This commit is contained in:
aniket-kardile 2026-08-24 20:58:15 +05:30
parent 4315d8818f
commit 74c7523fe6
No known key found for this signature in database
2 changed files with 29 additions and 2 deletions

View file

@ -3,7 +3,10 @@ from typing import Any, Literal
from pydantic import BaseModel, Field
from litellm.types.llms.openai import ChatCompletionToolParam
from litellm.types.llms.openai import (
ALL_RESPONSES_API_TOOL_PARAMS,
ChatCompletionToolParam,
)
from .base import GuardrailConfigModel
@ -37,7 +40,7 @@ class SingulrGuardrailPayload(BaseModel):
guardrail_scope: str | None = None
messages: Sequence[Any] | None = None
images: Sequence[str] | None = None
tools: Sequence[ChatCompletionToolParam] | None = None
tools: Sequence[ChatCompletionToolParam | ALL_RESPONSES_API_TOOL_PARAMS] | None = None
response: Any = None # pyright: ignore[reportExplicitAny] # logging_only reports raw litellm callback results (ModelResponse, EmbeddingResponse, etc.)
metadata: Mapping[str, Any] | None = None

View file

@ -212,6 +212,30 @@ class TestSingulrRequestPayload:
sent_payload = mock_post.call_args.kwargs["json"]
assert sent_payload["tools"] == tools
@pytest.mark.asyncio
async def test_responses_api_mcp_tools_are_forwarded(self, singulr_guardrail):
"""Regression: Responses API tools (e.g. {"type": "mcp", "server_label": ...})
have no "function" key, unlike Chat Completions tools. SingulrGuardrailPayload
rejected them with a pydantic ValidationError, turning every Responses API
request carrying an MCP tool into a 500."""
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):
"""Regression: the alias must be sent as {"user_api_key_alias": <alias>},