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"
|
_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
|
_DEFAULT_TIMEOUT: Final = 30.0
|
||||||
_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({})
|
_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({})
|
||||||
|
_MCP_MODEL_PREFIX: Final = "MCP:"
|
||||||
|
|
||||||
|
|
||||||
class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object):
|
class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object):
|
||||||
|
|
@ -375,6 +376,7 @@ class SingulrGuardrail(CustomGuardrail):
|
||||||
singulr_resp_obj = SingulrGuardrailPayload(
|
singulr_resp_obj = SingulrGuardrailPayload(
|
||||||
correlation_id=request_data.get("litellm_call_id"),
|
correlation_id=request_data.get("litellm_call_id"),
|
||||||
guardrail_scope="response",
|
guardrail_scope="response",
|
||||||
|
model_name=request_data.get("model"),
|
||||||
messages=request_data.get("messages"),
|
messages=request_data.get("messages"),
|
||||||
images=inputs.get("images"),
|
images=inputs.get("images"),
|
||||||
response=assistant_message,
|
response=assistant_message,
|
||||||
|
|
@ -405,6 +407,7 @@ class SingulrGuardrail(CustomGuardrail):
|
||||||
try:
|
try:
|
||||||
return SingulrGuardrailPayload(
|
return SingulrGuardrailPayload(
|
||||||
correlation_id=kwargs.get("litellm_call_id"),
|
correlation_id=kwargs.get("litellm_call_id"),
|
||||||
|
model_name=kwargs.get("model"),
|
||||||
guardrail_scope="response",
|
guardrail_scope="response",
|
||||||
response=result,
|
response=result,
|
||||||
metadata=metadata,
|
metadata=metadata,
|
||||||
|
|
@ -459,12 +462,21 @@ class SingulrGuardrail(CustomGuardrail):
|
||||||
return "guardrail_intervened"
|
return "guardrail_intervened"
|
||||||
return "success"
|
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(
|
async def async_logging_hook(
|
||||||
self,
|
self,
|
||||||
kwargs: dict, # mutable-ok: matches CustomLogger override; mutated via setdefault
|
kwargs: dict, # mutable-ok: matches CustomLogger override; mutated via setdefault
|
||||||
result: Any, # noqa: ANN401 # required by CustomLogger.async_logging_hook override signature
|
result: Any, # noqa: ANN401 # required by CustomLogger.async_logging_hook override signature
|
||||||
call_type: str,
|
call_type: str,
|
||||||
) -> tuple[dict, Any]:
|
) -> 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)
|
start_time: Final = datetime.now(timezone.utc)
|
||||||
guardrail_status: Final = await self._logging_only_guardrail_status(kwargs=kwargs, result=result)
|
guardrail_status: Final = await self._logging_only_guardrail_status(kwargs=kwargs, result=result)
|
||||||
if guardrail_status is None:
|
if guardrail_status is None:
|
||||||
|
|
|
||||||
|
|
@ -935,6 +935,41 @@ class TestSingulrLoggingHook:
|
||||||
guardrail_information = returned_kwargs["standard_logging_object"]["guardrail_information"]
|
guardrail_information = returned_kwargs["standard_logging_object"]["guardrail_information"]
|
||||||
assert guardrail_information[0]["guardrail_status"] == "success"
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_records_standard_logging_guardrail_information(self, singulr_guardrail):
|
async def test_records_standard_logging_guardrail_information(self, singulr_guardrail):
|
||||||
resp = _make_response({"should_block": False})
|
resp = _make_response({"should_block": False})
|
||||||
|
|
@ -1112,7 +1147,7 @@ class TestSingulrRequestWiring:
|
||||||
)
|
)
|
||||||
call_kwargs = mock_post.call_args.kwargs
|
call_kwargs = mock_post.call_args.kwargs
|
||||||
assert call_kwargs["timeout"] == 5.0
|
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:
|
class TestSingulrBuildHeaders:
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue