mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
feat(guardrails): skip MCP tool calls in the Singulr logging hook
MCP traffic already gets reported through the pre_mcp_call and post_mcp_call hooks with the richer mcp_request / mcp_response payloads. The logging_only hook sees the same call again as a completion whose model is "MCP: <tool_name>", so it was sending Singulr a duplicate, lower-fidelity report of every tool call. Also refreshes the endpoint assertion that the litellm-v2 guard endpoint bump left stale. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
171c9c0d0f
commit
da298ca74b
2 changed files with 49 additions and 2 deletions
|
|
@ -44,9 +44,10 @@ from litellm.types.utils import (
|
|||
)
|
||||
|
||||
_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):
|
||||
|
|
@ -375,6 +376,7 @@ class SingulrGuardrail(CustomGuardrail):
|
|||
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,
|
||||
|
|
@ -405,6 +407,7 @@ class SingulrGuardrail(CustomGuardrail):
|
|||
try:
|
||||
return SingulrGuardrailPayload(
|
||||
correlation_id=kwargs.get("litellm_call_id"),
|
||||
model_name=kwargs.get("model"),
|
||||
guardrail_scope="response",
|
||||
response=result,
|
||||
metadata=metadata,
|
||||
|
|
@ -459,12 +462,21 @@ class SingulrGuardrail(CustomGuardrail):
|
|||
return "guardrail_intervened"
|
||||
return "success"
|
||||
|
||||
@staticmethod
|
||||
def _is_mcp_call(kwargs: Mapping[str, Any]) -> bool:
|
||||
model: Final = kwargs.get("model")
|
||||
return isinstance(model, str) and model.startswith(_MCP_MODEL_PREFIX)
|
||||
|
||||
async def async_logging_hook(
|
||||
self,
|
||||
kwargs: dict, # mutable-ok: matches CustomLogger override; mutated via setdefault
|
||||
result: Any, # noqa: ANN401 # required by CustomLogger.async_logging_hook override signature
|
||||
call_type: str,
|
||||
) -> tuple[dict, Any]:
|
||||
if self._is_mcp_call(kwargs):
|
||||
verbose_proxy_logger.debug("Singulr: skipping logging_only report for MCP call %s", kwargs.get("model"))
|
||||
return kwargs, result
|
||||
|
||||
start_time: Final = datetime.now(timezone.utc)
|
||||
guardrail_status: Final = await self._logging_only_guardrail_status(kwargs=kwargs, result=result)
|
||||
if guardrail_status is None:
|
||||
|
|
|
|||
|
|
@ -935,6 +935,41 @@ class TestSingulrLoggingHook:
|
|||
guardrail_information = returned_kwargs["standard_logging_object"]["guardrail_information"]
|
||||
assert guardrail_information[0]["guardrail_status"] == "success"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_tool_call_is_not_reported(self, singulr_guardrail):
|
||||
"""MCP traffic is already covered by the pre/post_mcp_call hooks, which send
|
||||
the richer mcp_request/mcp_response payloads. The logging_only hook sees the
|
||||
same call again with model="MCP: <tool_name>" and must skip it so Singulr
|
||||
doesn't get a duplicate, lower-fidelity report of every tool call."""
|
||||
kwargs = {"model": "MCP: get_weather", "messages": [{"role": "user", "content": "hi"}]}
|
||||
with patch.object(singulr_guardrail.async_handler, "post") as mock_post:
|
||||
updated_kwargs, result = await singulr_guardrail.async_logging_hook(
|
||||
kwargs=kwargs, result={"choices": []}, call_type="acompletion"
|
||||
)
|
||||
mock_post.assert_not_called()
|
||||
assert "standard_logging_object" not in updated_kwargs
|
||||
assert result == {"choices": []}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_list_tools_call_is_not_reported(self, singulr_guardrail):
|
||||
kwargs = {"model": "MCP: list_tools", "messages": [{"role": "user", "content": "hi"}]}
|
||||
with patch.object(singulr_guardrail.async_handler, "post") as mock_post:
|
||||
await singulr_guardrail.async_logging_hook(kwargs=kwargs, result=None, call_type="acompletion")
|
||||
mock_post.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_mcp_model_is_still_reported(self, singulr_guardrail):
|
||||
"""Guard against the skip being too broad: a normal LLM call whose model
|
||||
merely mentions MCP later in the name must still be reported."""
|
||||
resp = _make_response({"should_block": False})
|
||||
kwargs = {"model": "gpt-4o-mcp", "messages": [{"role": "user", "content": "hi"}]}
|
||||
with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post:
|
||||
updated_kwargs, _ = await singulr_guardrail.async_logging_hook(
|
||||
kwargs=kwargs, result=None, call_type="acompletion"
|
||||
)
|
||||
mock_post.assert_called_once()
|
||||
assert updated_kwargs["standard_logging_object"]["guardrail_information"][0]["guardrail_status"] == "success"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_records_standard_logging_guardrail_information(self, singulr_guardrail):
|
||||
resp = _make_response({"should_block": False})
|
||||
|
|
@ -1112,7 +1147,7 @@ class TestSingulrRequestWiring:
|
|||
)
|
||||
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"
|
||||
assert call_kwargs["url"] == "https://api.test.singulr.ai/api/v1/ai-gateway/litellm-v2"
|
||||
|
||||
|
||||
class TestSingulrBuildHeaders:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue