fix(rubrik): collapse BlockedToolsResult dead-code into Optional[str]

The `allowed_tools` field on `BlockedToolsResult` was computed in
`_extract_blocked_tools` but never read by the only caller — when any
tool was blocked the integration unconditionally raised
`ModifyResponseException` to reject the full response, never doing
partial filtering. Drop the dataclass and return the blocking
explanation directly as `Optional[str]` so there's no misleading shape
hinting at unused partial-filter capability.

Co-authored-by: Greptile <greptile-apps[bot]@users.noreply.github.com>
This commit is contained in:
Claude 2026-05-20 19:15:38 +00:00
parent b8050a19ba
commit 5862f45a6c
No known key found for this signature in database
2 changed files with 15 additions and 28 deletions

View file

@ -6,7 +6,6 @@ import random
import time
import urllib.parse
import uuid
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Literal, Optional
import httpx
@ -39,14 +38,6 @@ _WEBHOOK_PATH_TOOL_BLOCKING = "/v1/after_completion/openai/v1"
_WEBHOOK_PATH_LOGGING_BATCH = "/v1/litellm/batch"
@dataclass
class BlockedToolsResult:
"""Returned by _extract_blocked_tools when at least one tool was blocked."""
allowed_tools: list
explanation: str
class RubrikLogger(CustomGuardrail, CustomBatchLogger):
preserve_events_added_during_flush = True
@ -210,12 +201,14 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
service_response = await self._post_to_tool_blocking_service(
response_data, req_data
)
blocked = self._extract_blocked_tools(service_response, message_tool_calls)
blocked_explanation = self._extract_blocked_tools(
service_response, message_tool_calls
)
if blocked:
if blocked_explanation is not None:
model = self._resolve_model(request_data, call_details)
raise ModifyResponseException(
message=blocked.explanation,
message=blocked_explanation,
model=model,
request_data=request_data,
guardrail_name=self.guardrail_name,
@ -504,12 +497,12 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
def _extract_blocked_tools(
service_response: dict[str, Any],
all_tool_calls: list[ChatCompletionMessageToolCall],
) -> BlockedToolsResult | None:
"""Determine whether any tool calls were blocked by the service.
) -> Optional[str]:
"""Return the blocking explanation if any tool calls were blocked.
Compares the service response (which contains only allowed tools) against
the full set of tool calls. Returns None if all tools are allowed, or a
BlockedToolsResult.
the full set of tool calls. Returns ``None`` if all tools are allowed, or
the explanation string (prefixed with newlines) otherwise.
Expects service_response in OpenAI chat completion format:
{"choices": [{"message": {"tool_calls": [...], "content": "..."}}]}
@ -523,13 +516,10 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
blocking_explanation = message.get("content", "")
allowed_ids = {tc["id"] for tc in returned_tool_calls if tc.get("id")}
allowed_tools = [tc for tc in all_tool_calls if tc.id in allowed_ids]
allowed_count = sum(1 for tc in all_tool_calls if tc.id in allowed_ids)
if len(allowed_tools) == len(all_tool_calls):
if allowed_count == len(all_tool_calls):
return None
explanation = blocking_explanation or "Tool call blocked by policy."
return BlockedToolsResult(
allowed_tools=allowed_tools,
explanation=f"\n\n{explanation}",
)
return f"\n\n{explanation}"

View file

@ -830,7 +830,7 @@ class TestExtractBlockedTools:
result = RubrikLogger._extract_blocked_tools(service_resp, [tc])
assert result is None
def test_some_blocked_returns_result(self):
def test_some_blocked_returns_explanation(self):
from litellm.types.utils import ChatCompletionMessageToolCall, Function
tc1 = ChatCompletionMessageToolCall(
@ -855,9 +855,7 @@ class TestExtractBlockedTools:
}
result = RubrikLogger._extract_blocked_tools(service_resp, [tc1, tc2])
assert result is not None
assert len(result.allowed_tools) == 1
assert result.allowed_tools[0].id == "call_1"
assert "blocked fn2" in result.explanation
assert "blocked fn2" in result
def test_empty_choices_raises(self):
with pytest.raises(Exception, match="empty response"):
@ -881,8 +879,7 @@ class TestExtractBlockedTools:
}
result = RubrikLogger._extract_blocked_tools(service_resp, [tc])
assert result is not None
assert result.allowed_tools == []
assert "blocked everything" in result.explanation
assert "blocked everything" in result
# -- Sanitize proxy server request -------------------------------------------