From a8ebcf646d7c0255a4ad9e0fe040310fd4d250b3 Mon Sep 17 00:00:00 2001 From: Joseph Barker Date: Tue, 7 Apr 2026 12:29:30 -0700 Subject: [PATCH] Add Rubrik as officially-supported guardrail plugin Adds tool blocking and batch logging integration with an external Rubrik webhook service. The plugin validates LLM tool calls against a policy service (fail-open on errors) and batch-logs all requests/responses. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../docs/proxy/guardrails/rubrik.md | 140 ++++ docs/my-website/sidebars.js | 1 + litellm/integrations/rubrik.py | 459 +++++++++++ .../guardrail_hooks/rubrik/__init__.py | 35 + litellm/types/guardrails.py | 1 + .../integrations/rubrik_test_helpers.py | 23 + .../test_litellm/integrations/test_rubrik.py | 732 ++++++++++++++++++ 7 files changed, 1391 insertions(+) create mode 100644 docs/my-website/docs/proxy/guardrails/rubrik.md create mode 100644 litellm/integrations/rubrik.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py create mode 100644 tests/test_litellm/integrations/rubrik_test_helpers.py create mode 100644 tests/test_litellm/integrations/test_rubrik.py diff --git a/docs/my-website/docs/proxy/guardrails/rubrik.md b/docs/my-website/docs/proxy/guardrails/rubrik.md new file mode 100644 index 00000000000..3e0d439f1ed --- /dev/null +++ b/docs/my-website/docs/proxy/guardrails/rubrik.md @@ -0,0 +1,140 @@ +# Rubrik Guardrail + +Use Rubrik's tool blocking and logging integration to validate LLM tool calls against an external policy service and batch-log all LLM requests/responses. + +**Key features:** +- **Tool blocking**: Validates tool calls against an external Rubrik service after LLM completion. Blocked tool calls trigger a policy violation response. +- **Batch logging**: Logs all LLM requests and responses to Rubrik with configurable sampling and batching. +- **Fail-open**: If the tool blocking service is unavailable, requests are allowed through unchanged. + +--- + +## Quick Start + +### 1. Set Environment Variables + +```bash +export RUBRIK_WEBHOOK_URL="https://your-rubrik-service.example.com" +export RUBRIK_API_KEY="your-rubrik-api-key" # optional +``` + +### 2. Configure `config.yaml` + +```yaml +model_list: + - model_name: gpt-4 + litellm_params: + model: openai/gpt-4 + api_key: os.environ/OPENAI_API_KEY + +guardrails: + - guardrail_name: "rubrik" + litellm_params: + guardrail: rubrik + mode: "post_call" + api_key: os.environ/RUBRIK_API_KEY + api_base: os.environ/RUBRIK_WEBHOOK_URL + default_on: true +``` + +### 3. Launch the Proxy + +```bash +litellm --config config.yaml --port 4000 +``` + +### 4. Test It + +```bash +curl -X POST http://localhost:4000/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-1234" \ + -d '{ + "model": "gpt-4", + "messages": [{"role": "user", "content": "What is the weather in SF?"}], + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string"} + }, + "required": ["location"] + } + } + } + ] + }' +``` + +--- + +## Configuration Reference + +### Environment Variables + +| Variable | Required | Default | Description | +|----------|----------|---------|-------------| +| `RUBRIK_WEBHOOK_URL` | Yes (or `api_base` in config) | — | Base URL of the Rubrik webhook service | +| `RUBRIK_API_KEY` | No | — | Bearer token for authenticating with the Rubrik service | +| `RUBRIK_SAMPLING_RATE` | No | `1.0` | Fraction of requests to log (0.0 to 1.0). Set to `0.5` to log ~50% of requests. | +| `RUBRIK_BATCH_SIZE` | No | `512` | Number of log entries to buffer before flushing. Logs are also flushed on a periodic interval. | + +### YAML Config Parameters + +| Parameter | Description | +|-----------|-------------| +| `guardrail: rubrik` | Selects the Rubrik guardrail integration | +| `mode: "post_call"` | Run after the LLM response is received | +| `api_key` | Rubrik API key (can use `os.environ/RUBRIK_API_KEY`) | +| `api_base` | Rubrik webhook base URL (can use `os.environ/RUBRIK_WEBHOOK_URL`) | +| `default_on` | When `true`, the guardrail runs on all requests without needing per-request opt-in | + +--- + +## How Tool Blocking Works + +1. After the LLM returns a response with tool calls, the Rubrik guardrail sends them to the blocking service at `{api_base}/v1/after_completion/openai/v1`. +2. The service evaluates each tool call against configured policies and returns the set of **allowed** tool calls. +3. If any tool calls are blocked, the proxy returns a `200` response with the policy violation explanation instead of the original LLM response. +4. If the blocking service is unreachable or returns an error, the guardrail **fails open** — the original response is returned unchanged. + +### Request/Response format + +The guardrail sends a JSON envelope to the blocking service: + +```json +{ + "request": { + "messages": [...], + "model": "gpt-4", + "proxy_server_request": {...} + }, + "response": { + "id": "chatcmpl-...", + "object": "chat.completion", + "choices": [{ + "message": { + "role": "assistant", + "tool_calls": [...] + } + }] + } +} +``` + +The service should return an OpenAI chat completion format response containing only the **allowed** tool calls and an optional `content` field with the blocking explanation. + +--- + +## How Batch Logging Works + +All LLM requests (successes and failures) are queued and sent in batches to `{api_base}/v1/litellm/batch`. + +- Logs are flushed when the queue reaches `RUBRIK_BATCH_SIZE` (default 512) or on a periodic interval (default 5 seconds). +- Use `RUBRIK_SAMPLING_RATE` to reduce logging volume in high-traffic deployments. +- For Anthropic `/v1/messages` requests, the log ID is normalized to `litellm_call_id` for consistency across tool blocking and logging. diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index ab8f257c7d5..4ca963c81e4 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -90,6 +90,7 @@ const sidebars = { "proxy/guardrails/custom_code_guardrail", "proxy/guardrails/prompt_injection", "proxy/guardrails/tool_permission", + "proxy/guardrails/rubrik", "proxy/guardrails/zscaler_ai_guard", "proxy/guardrails/javelin" ].sort(), diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py new file mode 100644 index 00000000000..ae525f6b1ee --- /dev/null +++ b/litellm/integrations/rubrik.py @@ -0,0 +1,459 @@ +"""Rubrik LiteLLM Plugin for tool blocking and batch logging.""" + +import asyncio +import copy +import os +import random +import time +import urllib.parse +import uuid +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal, Optional + +import httpx +from litellm._logging import verbose_logger +from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + ModifyResponseException, +) +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import ( + ChatCompletionMessageToolCall, + Function, + GenericGuardrailAPIInputs, + StandardLoggingPayload, +) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObj, + ) + +_ENDPOINT_ANTHROPIC_MESSAGES = "/messages" +_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): + def __init__( + self, + api_key: str | None = None, + api_base: str | None = None, + **kwargs, + ): + self.flush_lock = asyncio.Lock() + kwargs.setdefault("guardrail_name", "rubrik") + kwargs.setdefault("event_hook", GuardrailEventHooks.post_call) + kwargs.setdefault("default_on", True) + super().__init__( + flush_lock=self.flush_lock, + **kwargs, + ) + + verbose_logger.debug("initializing rubrik logger") + + self.sampling_rate = 1.0 + rbrk_sampling_rate = os.getenv("RUBRIK_SAMPLING_RATE") + if rbrk_sampling_rate is not None: + try: + parsed_rate = float(rbrk_sampling_rate.strip()) + self.sampling_rate = max(0.0, min(1.0, parsed_rate)) + if parsed_rate != self.sampling_rate: + verbose_logger.warning( + f"RUBRIK_SAMPLING_RATE={parsed_rate} clamped to " + f"{self.sampling_rate}" + ) + except ValueError: + verbose_logger.warning( + f"Invalid RUBRIK_SAMPLING_RATE: {rbrk_sampling_rate!r}, using 1.0" + ) + + self.key = api_key or os.getenv("RUBRIK_API_KEY") + if not self.key: + verbose_logger.warning( + "Rubrik: No API key configured. Requests will be unauthenticated." + ) + _batch_size = os.getenv("RUBRIK_BATCH_SIZE") + + if _batch_size: + try: + self.batch_size = int(_batch_size) + except ValueError: + verbose_logger.warning( + f"Invalid RUBRIK_BATCH_SIZE: {_batch_size!r}, using default" + ) + + _webhook_url = api_base or os.getenv("RUBRIK_WEBHOOK_URL") + + if _webhook_url is None: + raise ValueError( + "Rubrik webhook URL not configured. " + "Set RUBRIK_WEBHOOK_URL or pass api_base." + ) + + _webhook_url = _webhook_url.rstrip("/").removesuffix("/v1") + self.tool_blocking_endpoint = f"{_webhook_url}{_WEBHOOK_PATH_TOOL_BLOCKING}" + self.logging_endpoint = f"{_webhook_url}{_WEBHOOK_PATH_LOGGING_BATCH}" + + self.async_httpx_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.LoggingCallback + ) + + # Dedicated client to avoid connection pooling issues with LiteLLM's shared client + self.tool_blocking_client = httpx.AsyncClient( + timeout=httpx.Timeout(5.0, connect=2.0), + limits=httpx.Limits(max_connections=10, max_keepalive_connections=5), + ) + + self._headers: dict[str, str] = {"Content-Type": "application/json"} + if self.key: + self._headers["Authorization"] = f"Bearer {self.key}" + + asyncio.create_task(self.periodic_flush()) + + async def aclose(self): + """Close the dedicated tool blocking HTTP client.""" + await self.tool_blocking_client.aclose() + + # -- Guardrail hook -------------------------------------------------------- + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + """Validate tool calls against the blocking service (fail-open).""" + if input_type != "response": + return inputs + + tool_calls = inputs.get("tool_calls") + if not tool_calls: + return inputs + + try: + return await self._check_tool_calls( + inputs, tool_calls, request_data, logging_obj + ) + except ModifyResponseException: + raise + except Exception as e: + verbose_logger.error( + f"Tool blocking hook failed: {e}. " + "Returning original response unchanged.", + exc_info=True, + ) + return inputs + + async def _check_tool_calls( + self, + inputs: GenericGuardrailAPIInputs, + tool_calls: Any, + request_data: dict, + logging_obj: Optional["LiteLLMLoggingObj"], + ) -> GenericGuardrailAPIInputs: + """Send tool calls to blocking service, raise if any are blocked.""" + message_tool_calls = self._normalize_tool_calls(tool_calls) + + call_details = ( + getattr(logging_obj, "model_call_details", {}) if logging_obj else {} + ) + response = request_data.get("response") + request_id = getattr(response, "id", None) if response else None + if logging_obj and not call_details: + verbose_logger.warning( + "Rubrik: logging_obj present but model_call_details is empty " + "-- request context will be missing" + ) + + response_data = self._build_tool_call_payload(message_tool_calls, request_id) + req_data = self._extract_request_data(call_details) + + service_response = await self._post_to_tool_blocking_service( + response_data, req_data + ) + blocked = self._extract_blocked_tools(service_response, message_tool_calls) + + if blocked: + model = self._resolve_model(request_data, call_details) + raise ModifyResponseException( + message=blocked.explanation, + model=model, + request_data=request_data, + guardrail_name=self.guardrail_name, + ) + + return inputs + + @staticmethod + def _normalize_tool_calls(tool_calls: Any) -> list[ChatCompletionMessageToolCall]: + """Convert tool_calls from inputs to ChatCompletionMessageToolCall objects.""" + result = [] + for tc in tool_calls: + if isinstance(tc, ChatCompletionMessageToolCall): + result.append(tc) + elif isinstance(tc, dict): + func = tc.get("function", {}) + result.append( + ChatCompletionMessageToolCall( + id=tc.get("id", ""), + type=tc.get("type", "function"), + function=Function( + name=func.get("name", ""), + arguments=func.get("arguments", ""), + ), + ) + ) + elif hasattr(tc, "id") and hasattr(tc, "function"): + result.append( + ChatCompletionMessageToolCall( + id=tc.id or "", + type=getattr(tc, "type", None) or "function", + function=tc.function, + ) + ) + else: + raise TypeError( + f"Cannot normalize tool_call of type {type(tc).__name__}: {tc!r}" + ) + return result + + @staticmethod + def _build_tool_call_payload( + tool_calls: list[ChatCompletionMessageToolCall], + request_id: str | None, + ) -> dict[str, Any]: + """Build a full OpenAI ChatCompletion-format dict for the blocking service.""" + return { + "id": request_id or f"chatcmpl-{uuid.uuid4()}", + "object": "chat.completion", + "created": int(time.time()), + "model": "", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + tc.model_dump(exclude_none=True) for tc in tool_calls + ], + }, + "finish_reason": "tool_calls", + } + ], + } + + @staticmethod + def _extract_request_data(call_details: dict[str, Any]) -> dict[str, Any]: + """Extract original request data from model_call_details.""" + if not call_details: + return {} + litellm_params = call_details.get("litellm_params", {}) or {} + return { + "messages": call_details.get("messages"), + "model": call_details.get("model"), + "proxy_server_request": litellm_params.get("proxy_server_request"), + } + + @staticmethod + def _resolve_model( + request_data: dict[str, Any], call_details: dict[str, Any] + ) -> str: + """Get the model name for the ModifyResponseException.""" + response = request_data.get("response") + if response and hasattr(response, "model"): + return response.model or "unknown" + return call_details.get("model", "unknown") + + # -- Logging hooks --------------------------------------------------------- + + async def _prepare_log_payload( + self, kwargs: dict, event_type: str + ) -> StandardLoggingPayload | None: + """Shared logic for success and failure logging.""" + if random.random() > self.sampling_rate: + verbose_logger.debug( + f"Skipping Rubrik {event_type} logging " + f"(sampling_rate={self.sampling_rate})" + ) + return None + + # Deep-copy so mutations don't affect other callbacks sharing this object + standard_logging_payload: StandardLoggingPayload = copy.deepcopy( + kwargs["standard_logging_object"] + ) + + # For Anthropic /v1/messages requests, LiteLLM creates a separate + # ModelResponse (with a generated chatcmpl-* id) for logging, which + # differs from the original Anthropic msg-* id on the response dict. + # Normalize to litellm_call_id so that the logging and tool-blocking + # endpoints see the same request identifier. + litellm_params = kwargs.get("litellm_params", {}) or {} + proxy_request = litellm_params.get("proxy_server_request", {}) or {} + url_path = urllib.parse.urlparse(proxy_request.get("url", "")).path + if url_path.endswith(_ENDPOINT_ANTHROPIC_MESSAGES): + _litellm_call_id = kwargs.get("litellm_call_id") + if _litellm_call_id: + standard_logging_payload["id"] = _litellm_call_id # type: ignore[literal-required] + + if "system" in kwargs: + system_prompt_msg_list = kwargs["system"] + try: + if system_prompt_msg_list: + system_scaffold = { + "role": "system", + "content": system_prompt_msg_list, + } + if isinstance(standard_logging_payload["messages"], list): + standard_logging_payload["messages"].insert( + 0, system_scaffold + ) + elif isinstance( + standard_logging_payload["messages"], (dict, str) + ): + standard_logging_payload["messages"] = [ + system_scaffold, + standard_logging_payload["messages"], + ] + except Exception as e: + verbose_logger.warning( + f"Rubrik: failed to prepend system prompt: {e}", + exc_info=True, + ) + + return standard_logging_payload + + async def _enqueue_log_event(self, kwargs: dict, event_type: str): + try: + payload = await self._prepare_log_payload(kwargs, event_type) + if payload is None: + return + + self.log_queue.append(payload) + + if len(self.log_queue) >= self.batch_size: + await self.flush_queue() + except Exception as e: + verbose_logger.error( + f"Rubrik {event_type} logging hook failed: {e}. " + "Skipping logging for this event.", + exc_info=True, + ) + + async def async_log_success_event( + self, kwargs, response_obj, start_time, end_time + ): + await self._enqueue_log_event(kwargs, "success") + + async def async_log_failure_event( + self, kwargs, response_obj, start_time, end_time + ): + await self._enqueue_log_event(kwargs, "failure") + + # -- Batch logging --------------------------------------------------------- + + async def _log_batch_to_rubrik(self, data): + try: + response = await self.async_httpx_client.post( + url=self.logging_endpoint, + json=data, + headers=self._headers, + ) + response.raise_for_status() + except httpx.HTTPStatusError as e: + verbose_logger.exception( + f"Rubrik HTTP Error: {e.response.status_code} - {e.response.text}" + ) + except Exception: + verbose_logger.exception("Rubrik Layer Error") + + async def async_send_batch(self): + """Handles sending batches of responses to Rubrik.""" + if not self.log_queue: + return + + await self._log_batch_to_rubrik( + data=self.log_queue, + ) + + # -- Tool blocking service ------------------------------------------------- + + async def _post_to_tool_blocking_service( + self, + response_data: dict[str, Any], + request_data: dict[str, Any], + ) -> dict[str, Any]: + """Post a payload to the tool blocking service and return the response. + + Args: + response_data: The OpenAI-formatted response payload to send. + request_data: Original LLM request data to include alongside + the response for additional context. Empty dict if unavailable. + + Raises: + Exception: If the service is unavailable or returns an error. + """ + envelope = { + "request": request_data, + "response": response_data, + } + verbose_logger.debug( + f"Sending request to tool blocking service: " + f"{self.tool_blocking_endpoint}" + ) + http_response = await self.tool_blocking_client.post( + self.tool_blocking_endpoint, + json=envelope, + headers=self._headers, + ) + http_response.raise_for_status() + result: dict[str, Any] = http_response.json() + return result + + @staticmethod + 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. + + 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. + + Expects service_response in OpenAI chat completion format: + {"choices": [{"message": {"tool_calls": [...], "content": "..."}}]} + """ + choices = service_response.get("choices", []) + if not choices: + raise Exception("Tool blocking service returned empty response") + + message = choices[0].get("message", {}) + returned_tool_calls = message.get("tool_calls", []) + 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] + + if len(allowed_tools) == len(all_tool_calls): + return None + + return BlockedToolsResult( + allowed_tools=allowed_tools, + explanation=f"\n\n{blocking_explanation}", + ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py new file mode 100644 index 00000000000..ab347130a30 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py @@ -0,0 +1,35 @@ +"""Rubrik guardrail integration for LiteLLM.""" + +from typing import TYPE_CHECKING + +from litellm.integrations.rubrik import RubrikLogger +from litellm.types.guardrails import SupportedGuardrailIntegrations + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail( + litellm_params: "LitellmParams", guardrail: "Guardrail" +) -> RubrikLogger: + import litellm + + rubrik_callback = RubrikLogger( + api_key=litellm_params.api_key, + api_base=litellm_params.api_base, + guardrail_name=guardrail.get("guardrail_name", ""), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + + litellm.logging_callback_manager.add_litellm_callback(rubrik_callback) + return rubrik_callback + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.RUBRIK.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.RUBRIK.value: RubrikLogger, +} diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index cfec0398c81..e78cfc5fd92 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -84,6 +84,7 @@ class SupportedGuardrailIntegrations(Enum): BLOCK_CODE_EXECUTION = "block_code_execution" AKTO = "akto" MCP_JWT_SIGNER = "mcp_jwt_signer" + RUBRIK = "rubrik" class Role(Enum): diff --git a/tests/test_litellm/integrations/rubrik_test_helpers.py b/tests/test_litellm/integrations/rubrik_test_helpers.py new file mode 100644 index 00000000000..1bdb8cb247b --- /dev/null +++ b/tests/test_litellm/integrations/rubrik_test_helpers.py @@ -0,0 +1,23 @@ +"""Shared helpers for Rubrik plugin tests.""" + +from typing import Any, Dict + +from litellm.types.utils import GenericGuardrailAPIInputs + + +def make_tool_call_dict( + tc_id: str, name: str, arguments: str = "{}" +) -> Dict[str, Any]: + """Create a tool call dict matching the ChatCompletionMessageToolCall schema.""" + return { + "id": tc_id, + "type": "function", + "function": {"name": name, "arguments": arguments}, + } + + +def make_inputs_with_tools( + tool_calls: list, texts: list | None = None +) -> GenericGuardrailAPIInputs: + """Create GenericGuardrailAPIInputs with tool_calls.""" + return GenericGuardrailAPIInputs(texts=texts or [], tool_calls=tool_calls) diff --git a/tests/test_litellm/integrations/test_rubrik.py b/tests/test_litellm/integrations/test_rubrik.py new file mode 100644 index 00000000000..d543516f183 --- /dev/null +++ b/tests/test_litellm/integrations/test_rubrik.py @@ -0,0 +1,732 @@ +""" +Tests for the Rubrik LiteLLM plugin. + +Covers initialization, apply_guardrail tool blocking (all allowed, all blocked, +partial blocking, fail-open), batch logging, and Anthropic format handling. +""" + +import os +from typing import Any, Dict +from unittest.mock import AsyncMock, Mock, patch + +import httpx +import pytest + +from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.integrations.rubrik import RubrikLogger + +from tests.test_litellm.integrations.rubrik_test_helpers import ( + make_inputs_with_tools, + make_tool_call_dict, +) + + +@pytest.fixture +def mock_env(): + """Set up environment variables for testing.""" + with patch.dict( + os.environ, + { + "RUBRIK_WEBHOOK_URL": "http://localhost:8080", + "RUBRIK_API_KEY": "test-api-key", + }, + ): + yield + + +@pytest.fixture +def handler(mock_env): + """Create a RubrikLogger instance for testing.""" + with patch("asyncio.create_task", Mock()): + return RubrikLogger() + + +# -- Initialization ----------------------------------------------------------- + + +class TestInitialization: + def test_init_success(self, mock_env): + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger() + assert ( + handler.tool_blocking_endpoint + == "http://localhost:8080/v1/after_completion/openai/v1" + ) + assert handler.logging_endpoint == "http://localhost:8080/v1/litellm/batch" + assert handler.key == "test-api-key" + assert isinstance(handler.tool_blocking_client, httpx.AsyncClient) + + def test_init_with_constructor_params(self): + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger( + api_key="ctor-key", api_base="http://ctor-host:9090" + ) + assert handler.key == "ctor-key" + assert ( + handler.tool_blocking_endpoint + == "http://ctor-host:9090/v1/after_completion/openai/v1" + ) + + def test_init_without_url(self): + with patch.dict(os.environ, {}, clear=True): + with pytest.raises(ValueError, match="Rubrik webhook URL not configured"): + RubrikLogger() + + def test_init_without_api_key(self): + with patch.dict( + os.environ, {"RUBRIK_WEBHOOK_URL": "http://localhost:8080"}, clear=True + ): + with patch("asyncio.create_task", Mock()): + assert RubrikLogger().key is None + + def test_trailing_slash_removed(self): + with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://localhost:8080/"}): + with patch("asyncio.create_task", Mock()): + assert ( + RubrikLogger().tool_blocking_endpoint + == "http://localhost:8080/v1/after_completion/openai/v1" + ) + + def test_v1_suffix_stripped_as_substring_not_charset(self): + with patch("asyncio.create_task", Mock()): + with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://host/v1"}): + assert ( + RubrikLogger().tool_blocking_endpoint + == "http://host/v1/after_completion/openai/v1" + ) + + with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://host/v11"}): + assert ( + RubrikLogger().tool_blocking_endpoint + == "http://host/v11/v1/after_completion/openai/v1" + ) + + def test_sampling_rate_fractional(self): + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_SAMPLING_RATE": "0.5"}, + ): + assert RubrikLogger().sampling_rate == 0.5 + + def test_sampling_rate_invalid_ignored(self): + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_SAMPLING_RATE": "abc"}, + ): + assert RubrikLogger().sampling_rate == 1.0 + + def test_sampling_rate_clamped(self): + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_SAMPLING_RATE": "2.0"}, + ): + assert RubrikLogger().sampling_rate == 1.0 + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_SAMPLING_RATE": "-0.5"}, + ): + assert RubrikLogger().sampling_rate == 0.0 + + def test_batch_size_invalid_ignored(self): + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_BATCH_SIZE": "abc"}, + ): + # Should use default without crashing + assert isinstance(RubrikLogger().batch_size, int) + + def test_batch_size_valid(self): + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_BATCH_SIZE": "256"}, + ): + assert RubrikLogger().batch_size == 256 + + def test_headers_with_api_key(self, handler): + assert handler._headers["Authorization"] == "Bearer test-api-key" + assert handler._headers["Content-Type"] == "application/json" + + def test_headers_without_api_key(self): + with patch.dict( + os.environ, {"RUBRIK_WEBHOOK_URL": "http://host"}, clear=True + ): + with patch("asyncio.create_task", Mock()): + h = RubrikLogger() + assert "Authorization" not in h._headers + + +# -- Batch Logging ------------------------------------------------------------ + + +@pytest.mark.asyncio +class TestBatchLogging: + async def test_log_success_event_appends_to_queue(self, handler): + kwargs = { + "standard_logging_object": { + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + }, + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert len(handler.log_queue) == 1 + + async def test_log_failure_event_appends_to_queue(self, handler): + kwargs = { + "standard_logging_object": { + "messages": [{"role": "user", "content": "hi"}], + "response": "error", + }, + } + await handler.async_log_failure_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert len(handler.log_queue) == 1 + + async def test_log_success_event_sampling_skips(self, handler): + handler.sampling_rate = 0.0 + kwargs = { + "standard_logging_object": { + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + }, + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert len(handler.log_queue) == 0 + + async def test_flush_queue_sends_batch(self, handler): + handler.log_queue = [{"msg": "a"}, {"msg": "b"}] + mock_response = Mock() + mock_response.status_code = 200 + handler.async_httpx_client = AsyncMock() + handler.async_httpx_client.post = AsyncMock(return_value=mock_response) + await handler.flush_queue() + handler.async_httpx_client.post.assert_called_once() + assert len(handler.log_queue) == 0 + + async def test_log_batch_error_does_not_crash(self, handler): + handler.log_queue = [{"msg": "a"}] + mock_response = Mock() + mock_response.status_code = 500 + mock_response.text = "Internal Server Error" + mock_response.raise_for_status = Mock( + side_effect=httpx.HTTPStatusError( + "err", request=Mock(), response=mock_response + ) + ) + handler.async_httpx_client = AsyncMock() + handler.async_httpx_client.post = AsyncMock(return_value=mock_response) + await handler.flush_queue() + assert len(handler.log_queue) == 0 + + async def test_system_prompt_prepended_to_messages(self, handler): + kwargs = { + "standard_logging_object": { + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + }, + "system": "You are a helpful assistant.", + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert len(handler.log_queue) == 1 + msgs = handler.log_queue[0]["messages"] + assert msgs[0]["role"] == "system" + assert msgs[0]["content"] == "You are a helpful assistant." + + async def test_system_prompt_with_dict_messages(self, handler): + kwargs = { + "standard_logging_object": { + "messages": {"role": "user", "content": "hi"}, + "response": "hello", + }, + "system": "Be concise.", + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert len(handler.log_queue) == 1 + msgs = handler.log_queue[0]["messages"] + assert isinstance(msgs, list) + assert msgs[0]["role"] == "system" + assert msgs[1] == {"role": "user", "content": "hi"} + + async def test_anthropic_id_normalization(self, handler): + kwargs = { + "standard_logging_object": { + "id": "chatcmpl-original", + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + }, + "litellm_params": { + "proxy_server_request": { + "url": "http://proxy/v1/messages", + }, + }, + "litellm_call_id": "litellm-call-123", + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert handler.log_queue[0]["id"] == "litellm-call-123" + + async def test_non_anthropic_id_unchanged(self, handler): + kwargs = { + "standard_logging_object": { + "id": "chatcmpl-original", + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + }, + "litellm_params": { + "proxy_server_request": { + "url": "http://proxy/v1/chat/completions", + }, + }, + "litellm_call_id": "litellm-call-123", + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert handler.log_queue[0]["id"] == "chatcmpl-original" + + async def test_payload_deep_copied_not_mutated(self, handler): + """Verify the shared standard_logging_object is not mutated.""" + original_payload = { + "id": "original-id", + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + } + kwargs = { + "standard_logging_object": original_payload, + "system": "System prompt.", + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + # Original payload should NOT have been mutated + assert original_payload["id"] == "original-id" + assert len(original_payload["messages"]) == 1 + + +# -- Tool Blocking (apply_guardrail) ------------------------------------------ + + +def _mock_service_response(response_json): + """Create a mock tool blocking client that returns the given JSON.""" + + async def mock_post(*_args, **kwargs): + mock_resp = Mock() + mock_resp.json.return_value = response_json + mock_resp.raise_for_status = Mock() + return mock_resp + + mock_client = AsyncMock() + mock_client.post = mock_post + return mock_client + + +def _echo_service(): + """Create a mock tool blocking client that echoes the payload back.""" + + async def mock_post(*_args, **kwargs): + mock_resp = Mock() + mock_resp.json.return_value = kwargs.get("json", {}).get("response", {}) + mock_resp.raise_for_status = Mock() + return mock_resp + + mock_client = AsyncMock() + mock_client.post = mock_post + return mock_client + + +@pytest.mark.asyncio +class TestApplyGuardrail: + async def test_skips_requests(self, handler): + inputs = make_inputs_with_tools( + [make_tool_call_dict("call_1", "test_tool")] + ) + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + async def test_no_tool_calls(self, handler): + from litellm.types.utils import GenericGuardrailAPIInputs + + inputs = GenericGuardrailAPIInputs(texts=["hello"]) + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + async def test_all_allowed(self, handler): + tc1 = make_tool_call_dict("call_1", "get_weather") + tc2 = make_tool_call_dict("call_2", "get_time") + inputs = make_inputs_with_tools([tc1, tc2]) + + handler.tool_blocking_client = _echo_service() + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + async def test_all_blocked(self, handler): + tc1 = make_tool_call_dict("call_1", "delete_table") + tc2 = make_tool_call_dict("call_2", "drop_database") + inputs = make_inputs_with_tools([tc1, tc2]) + + handler.tool_blocking_client = _mock_service_response( + { + "choices": [ + { + "message": { + "role": "assistant", + "content": "Tool blocked by policy", + "tool_calls": [], + } + } + ], + } + ) + + with pytest.raises(ModifyResponseException) as exc_info: + await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert "Tool blocked by policy" in exc_info.value.message + + async def test_partial_blocking(self, handler): + tc_blocked = make_tool_call_dict("call_A", "blocked_tool") + tc_allowed = make_tool_call_dict("call_B", "allowed_tool") + inputs = make_inputs_with_tools([tc_blocked, tc_allowed]) + + async def mock_post(*_args, **kwargs): + payload = kwargs.get("json", {}).get("response", {}) + all_tcs = payload["choices"][0]["message"]["tool_calls"] + allowed = [tc for tc in all_tcs if tc.get("id") == "call_B"] + mock_resp = Mock() + mock_resp.json.return_value = { + "choices": [ + { + "message": { + "role": "assistant", + "content": "blocked", + "tool_calls": allowed, + } + } + ], + } + mock_resp.raise_for_status = Mock() + return mock_resp + + mock_client = AsyncMock() + mock_client.post = mock_post + handler.tool_blocking_client = mock_client + + with pytest.raises(ModifyResponseException): + await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + + async def test_service_failure_fail_open(self, handler): + tc1 = make_tool_call_dict("call_1", "test_tool") + inputs = make_inputs_with_tools([tc1]) + + mock_client = AsyncMock() + mock_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout")) + handler.tool_blocking_client = mock_client + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + async def test_service_empty_choices_fail_open(self, handler): + tc1 = make_tool_call_dict("call_1", "test_tool") + inputs = make_inputs_with_tools([tc1]) + + handler.tool_blocking_client = _mock_service_response({"choices": []}) + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + async def test_blocking_service_payload_format(self, handler): + tc1 = make_tool_call_dict("call_1", "get_weather", '{"location": "SF"}') + tc2 = make_tool_call_dict( + "call_2", "send_email", '{"to": "user@example.com"}' + ) + inputs = make_inputs_with_tools([tc1, tc2]) + + captured_payload: Dict[str, Any] = {} + + async def mock_post(*_args, **kwargs): + captured_payload.update(kwargs.get("json", {})) + mock_resp = Mock() + mock_resp.json.return_value = captured_payload.get("response", {}) + mock_resp.raise_for_status = Mock() + return mock_resp + + mock_client = AsyncMock() + mock_client.post = mock_post + handler.tool_blocking_client = mock_client + + await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + + # Verify envelope structure + assert "request" in captured_payload + assert "response" in captured_payload + + response_data = captured_payload["response"] + message = response_data["choices"][0]["message"] + assert message["role"] == "assistant" + assert len(message["tool_calls"]) == 2 + assert message["tool_calls"][0]["id"] == "call_1" + assert message["tool_calls"][0]["function"]["name"] == "get_weather" + assert message["tool_calls"][1]["id"] == "call_2" + assert message["tool_calls"][1]["function"]["name"] == "send_email" + + async def test_request_data_included_in_envelope(self, handler): + tc = make_tool_call_dict("call_1", "test_tool") + inputs = make_inputs_with_tools([tc]) + + captured_payload: Dict[str, Any] = {} + + async def mock_post(*_args, **kwargs): + captured_payload.update(kwargs.get("json", {})) + mock_resp = Mock() + mock_resp.json.return_value = captured_payload.get("response", {}) + mock_resp.raise_for_status = Mock() + return mock_resp + + mock_client = AsyncMock() + mock_client.post = mock_post + handler.tool_blocking_client = mock_client + + logging_obj = Mock() + logging_obj.model_call_details = { + "messages": [{"role": "user", "content": "hi"}], + "model": "gpt-4", + "litellm_params": { + "proxy_server_request": {"url": "/chat/completions"}, + }, + } + + await handler.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="response", + logging_obj=logging_obj, + ) + + req = captured_payload["request"] + assert req["model"] == "gpt-4" + assert req["messages"] == [{"role": "user", "content": "hi"}] + + +# -- Anthropic format ---------------------------------------------------------- + + +@pytest.mark.asyncio +class TestApplyGuardrailAnthropicFormat: + """Verify blocking works correctly regardless of original provider format. + + The framework converts Anthropic tool_use blocks to OpenAI-format + tool_calls before calling apply_guardrail. + """ + + async def test_single_tool_allowed(self, handler): + tc = make_tool_call_dict( + "toolu_123", "get_weather", '{"location": "Portland, OR"}' + ) + inputs = make_inputs_with_tools([tc], texts=["I'll check the weather."]) + + handler.tool_blocking_client = _echo_service() + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + async def test_single_tool_blocked(self, handler): + tc = make_tool_call_dict("toolu_123", "dangerous_tool", '{"arg": "value"}') + inputs = make_inputs_with_tools([tc]) + + handler.tool_blocking_client = _mock_service_response( + { + "choices": [ + { + "message": { + "role": "assistant", + "content": "blocked", + "tool_calls": [], + } + } + ], + } + ) + + with pytest.raises(ModifyResponseException): + await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + + async def test_text_only_response_no_blocking(self, handler): + from litellm.types.utils import GenericGuardrailAPIInputs + + inputs = GenericGuardrailAPIInputs(texts=["Hello! I'm Claude."]) + + mock_client = AsyncMock() + mock_client.post = AsyncMock() + handler.tool_blocking_client = mock_client + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + + assert result is inputs + mock_client.post.assert_not_called() + + async def test_service_failure_preserves_tools(self, handler): + tc = make_tool_call_dict("toolu_123", "get_weather", '{"location": "SF"}') + inputs = make_inputs_with_tools([tc]) + + mock_client = AsyncMock() + mock_client.post = AsyncMock( + side_effect=httpx.TimeoutException("Timeout") + ) + handler.tool_blocking_client = mock_client + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + +# -- Normalize tool calls ------------------------------------------------------ + + +class TestNormalizeToolCalls: + def test_dict_input(self): + tc = make_tool_call_dict("call_1", "test", '{"a": 1}') + result = RubrikLogger._normalize_tool_calls([tc]) + assert len(result) == 1 + assert result[0].id == "call_1" + assert result[0].function.name == "test" + assert result[0].function.arguments == '{"a": 1}' + + def test_typed_object_input(self): + from litellm.types.utils import ChatCompletionMessageToolCall, Function + + tc = ChatCompletionMessageToolCall( + id="call_2", + type="function", + function=Function(name="fn", arguments="{}"), + ) + result = RubrikLogger._normalize_tool_calls([tc]) + assert len(result) == 1 + assert result[0].id == "call_2" + assert result[0].function.name == "fn" + + def test_unsupported_type_raises(self): + with pytest.raises(TypeError, match="Cannot normalize"): + RubrikLogger._normalize_tool_calls(["not_a_tool_call"]) + + +# -- Extract blocked tools ----------------------------------------------------- + + +class TestExtractBlockedTools: + def test_all_allowed_returns_none(self): + from litellm.types.utils import ChatCompletionMessageToolCall, Function + + tc = ChatCompletionMessageToolCall( + id="call_1", type="function", function=Function(name="fn", arguments="{}") + ) + service_resp = { + "choices": [ + { + "message": { + "tool_calls": [{"id": "call_1"}], + "content": "", + } + } + ] + } + result = RubrikLogger._extract_blocked_tools(service_resp, [tc]) + assert result is None + + def test_some_blocked_returns_result(self): + from litellm.types.utils import ChatCompletionMessageToolCall, Function + + tc1 = ChatCompletionMessageToolCall( + id="call_1", + type="function", + function=Function(name="fn1", arguments="{}"), + ) + tc2 = ChatCompletionMessageToolCall( + id="call_2", + type="function", + function=Function(name="fn2", arguments="{}"), + ) + service_resp = { + "choices": [ + { + "message": { + "tool_calls": [{"id": "call_1"}], + "content": "blocked fn2", + } + } + ] + } + 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 + + def test_empty_choices_raises(self): + with pytest.raises(Exception, match="empty response"): + RubrikLogger._extract_blocked_tools({"choices": []}, []) + + +# -- Resolve model ------------------------------------------------------------- + + +class TestResolveModel: + def test_model_from_response(self): + from unittest.mock import Mock + + response = Mock() + response.model = "gpt-4" + result = RubrikLogger._resolve_model({"response": response}, {}) + assert result == "gpt-4" + + def test_model_from_call_details(self): + result = RubrikLogger._resolve_model({}, {"model": "claude-3"}) + assert result == "claude-3" + + def test_fallback_to_unknown(self): + result = RubrikLogger._resolve_model({}, {}) + assert result == "unknown" + + def test_empty_model_on_response_returns_unknown(self): + from unittest.mock import Mock + + response = Mock() + response.model = "" + result = RubrikLogger._resolve_model({"response": response}, {"model": "fallback"}) + assert result == "unknown"