diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py index 2e49da45ce9..4bcbe8bae37 100644 --- a/litellm/integrations/rubrik.py +++ b/litellm/integrations/rubrik.py @@ -1,12 +1,14 @@ -"""Rubrik LiteLLM Plugin for tool blocking and batch logging.""" +"""Rubrik LiteLLM Plugin for prompt/response moderation and batch logging.""" import asyncio import os import random import time -import urllib.parse import uuid from collections import Counter +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Literal, Optional import httpx @@ -18,6 +20,10 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.litellm_core_utils.core_helpers import safe_deep_copy +from litellm.litellm_core_utils.model_param_helper import ModelParamHelper +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + convert_content_list_to_str, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -35,15 +41,16 @@ if TYPE_CHECKING: Logging as LiteLLMLoggingObj, ) -_ENDPOINT_ANTHROPIC_MESSAGES = "/v1/messages" -_WEBHOOK_PATH_TOOL_BLOCKING = "/v1/after_completion/openai/v1" +_WEBHOOK_PATH_RESPONSE_MODERATION = "/v1/after_completion/openai/v1" +_WEBHOOK_PATH_PROMPT_MODERATION = "/v1/before_prompt/openai/v1" _WEBHOOK_PATH_LOGGING_BATCH = "/v1/litellm/batch" _MAX_QUEUE_SIZE = 10_000 _DROP_WARNING_INTERVAL_SECONDS = 60.0 +_EMPTY_MAPPING: Mapping[str, Any] = MappingProxyType({}) class _MalformedToolBlockingResponseError(Exception): - """Raised when the tool blocking service returns a structurally invalid + """Raised when the response moderation service returns a structurally invalid response (e.g. empty ``choices``). Distinct from transient network/HTTP errors so callers can surface a @@ -52,11 +59,15 @@ class _MalformedToolBlockingResponseError(Exception): """ -class RubrikLogger(CustomGuardrail, CustomBatchLogger): - @classmethod - def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: - return [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call] +@dataclass +class BlockedResponseResult: + """Returned by _extract_response_block when the response was blocked + (response text replaced, or at least one tool call removed).""" + explanation: str + + +class RubrikLogger(CustomGuardrail, CustomBatchLogger): def __init__( self, api_key: str | None = None, @@ -67,21 +78,82 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): kwargs.setdefault("guardrail_name", "rubrik") # `initialize_guardrail` always passes these kwargs explicitly, with # value `None` when the user omits `mode` / `default_on` from the - # guardrail config. Coerce None (omitted) to the desired default - # while preserving any explicit value the caller did set -- - # in particular `default_on=False` if the user wants the guardrail - # off by default. + # guardrail config. Follow the standard litellm convention: omitted + # resolves to False (off by default, user must opt in explicitly). kwargs["event_hook"] = kwargs.get("event_hook") or GuardrailEventHooks.post_call if kwargs.get("default_on") is None: - kwargs["default_on"] = True - kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) + kwargs["default_on"] = False super().__init__( flush_lock=self.flush_lock, + supported_event_hooks=list(self.get_supported_event_hooks()), **kwargs, ) verbose_logger.debug("initializing rubrik logger") + # Defining ``apply_guardrail`` routes streaming responses through + # litellm's ``unified_guardrail.async_post_call_streaming_iterator_hook``. + # By default that hook samples intermediate chunks + # (``streaming_sampling_rate``, default 5) and also moderates at + # end-of-stream, so a streamed response costs ~ceil(N/5)+1 Rubrik + # webhook round-trips. litellm reads this attribute via + # ``getattr(guardrail, "streaming_end_of_stream_only", False)``; when + # True it yields chunks unprocessed and only moderates the fully + # assembled response once at end of stream. + self.streaming_end_of_stream_only = True + + # ``streaming_end_of_stream_only`` is detect-only: it releases every + # chunk to the client *before* moderating, so a block can only append a + # trailing message -- the original content has already been delivered. + # ``streaming_buffer_until_moderated`` (litellm >= BerriAI/litellm#31389) + # withholds all chunks until end-of-stream moderation passes, then + # releases the original response (clean) or only the block message + # (blocked). On older litellm this attribute is ignored and we fall + # back to the detect-only behavior above. + self.streaming_buffer_until_moderated = True + + self._parse_sampling_rate() + + 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.") + + self._parse_batch_size() + + # Cap the in-memory retry queue so a Rubrik webhook outage cannot let + # authenticated traffic accumulate prompt/response payloads until the + # proxy runs out of memory. Once the cap is reached, oldest events are + # dropped to make room for fresh ones (drop-oldest backpressure). + self.max_queue_size = _MAX_QUEUE_SIZE + self._dropped_since_warning = 0 + self._last_drop_warning_time = 0.0 + + _webhook_url = api_base or os.getenv("RUBRIK_WEBHOOK_URL") + if not _webhook_url: + raise ValueError("Rubrik webhook URL not configured. Set RUBRIK_WEBHOOK_URL or pass api_base.") + + _webhook_url = _webhook_url.rstrip("/").removesuffix("/v1") + self._setup_clients(_webhook_url) + + self._headers: Mapping[str, str] = MappingProxyType( + {"Content-Type": "application/json", "Authorization": f"Bearer {self.key}"} + if self.key + else {"Content-Type": "application/json"} + ) + + self._periodic_flush_task: asyncio.Task[Any] | None = self._start_periodic_flush_task() + + @classmethod + def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: + """Return the guardrail event hooks this integration supports. + + Prompt moderation (``pre_call``) evaluates the user's message before + the LLM is called. Response moderation (``post_call``) evaluates the + assistant's reply and tool calls after the LLM returns. + """ + return [GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call] + + def _parse_sampling_rate(self) -> None: self.sampling_rate = 1.0 rbrk_sampling_rate = os.getenv("RUBRIK_SAMPLING_RATE") if rbrk_sampling_rate is not None: @@ -93,80 +165,54 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): 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.") + def _parse_batch_size(self) -> None: _batch_size = os.getenv("RUBRIK_BATCH_SIZE") - if _batch_size: try: - self.batch_size = int(_batch_size) + parsed_size = int(_batch_size) + if parsed_size <= 0: + verbose_logger.warning(f"RUBRIK_BATCH_SIZE={_batch_size!r} must be > 0, using default") + else: + self.batch_size = parsed_size except ValueError: verbose_logger.warning(f"Invalid RUBRIK_BATCH_SIZE: {_batch_size!r}, using default") - # Cap the in-memory retry queue so a Rubrik webhook outage cannot let - # authenticated traffic accumulate prompt/response payloads until the - # proxy runs out of memory. Once the cap is reached, oldest events are - # dropped to make room for fresh ones (drop-oldest backpressure). - self.max_queue_size = _MAX_QUEUE_SIZE - self._dropped_since_warning = 0 - self._last_drop_warning_time = 0.0 - - _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}" + def _setup_clients(self, webhook_url: str) -> None: + self.response_moderation_endpoint = f"{webhook_url}{_WEBHOOK_PATH_RESPONSE_MODERATION}" + self.prompt_moderation_endpoint = f"{webhook_url}{_WEBHOOK_PATH_PROMPT_MODERATION}" + self.logging_endpoint = f"{webhook_url}{_WEBHOOK_PATH_LOGGING_BATCH}" self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) - self.tool_blocking_client = get_async_httpx_client( + self.moderation_client = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback, params={"timeout": httpx.Timeout(5.0, connect=2.0)}, ) - self._headers: dict[str, str] = {"Content-Type": "application/json"} - if self.key: - self._headers["Authorization"] = f"Bearer {self.key}" - - # Periodic flush is started lazily on the first log event so that - # low-traffic deployments still get their batches drained even when the - # logger is instantiated outside a running event loop (sync init). - self._flush_task: asyncio.Task[Any] | None = self._start_periodic_flush_task() - def _start_periodic_flush_task(self) -> asyncio.Task[Any] | None: """Start the periodic flush task only when an event loop is already running.""" try: loop = asyncio.get_running_loop() except RuntimeError: - verbose_logger.debug( - "Rubrik logger init: no running event loop, periodic flush will start on first log event." - ) return None return loop.create_task(self.periodic_flush()) def _ensure_periodic_flush_task(self) -> None: - # Synchronous helper: in asyncio's cooperative model there is no await - # between the check and assignment, so two callers cannot race here. - if self._flush_task is None or self._flush_task.done(): - self._flush_task = self._start_periodic_flush_task() + if self._periodic_flush_task is None or self._periodic_flush_task.done(): + self._periodic_flush_task = self._start_periodic_flush_task() async def aclose(self): - """Close the dedicated HTTP clients used by this logger.""" - # Cancel the periodic flush task before closing the HTTP clients so - # the loop doesn't wake up and try to POST via a closed client. - if self._flush_task is not None and not self._flush_task.done(): - self._flush_task.cancel() - try: - await self._flush_task - except (asyncio.CancelledError, Exception): - pass - self._flush_task = None - await self.tool_blocking_client.close() - await self.async_httpx_client.close() + """Cancel the periodic flush task. + + ``moderation_client`` and ``async_httpx_client`` are shared objects + from LiteLLM's global HTTP-client cache (``get_async_httpx_client`` + uses the same cache key for all instances with equal parameters). + Closing them here would close the shared connection pool for every + other logger instance; let LiteLLM manage their lifecycle instead. + """ + task = getattr(self, "_periodic_flush_task", None) + if task is not None: + task.cancel() # -- Guardrail hook -------------------------------------------------------- @@ -177,67 +223,104 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): 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 + """Moderate prompts (request) and responses (response); fail-open. - tool_calls = inputs.get("tool_calls") - if not tool_calls: - return inputs + - ``request``: evaluate the prompt via the before_prompt webhook and + block disallowed prompts before the model is called. + - ``response``: evaluate the assistant's response text and tool calls + via the after_completion webhook and block on a policy violation. + litellm's guardrail-translation layer normalizes Anthropic and OpenAI + requests/responses into ``inputs`` before this runs, so a single code + path covers both wire formats. The configured guardrail ``mode`` + selects which surface(s) run. + """ + if input_type == "request": + return await self._guarded( + self._moderate_prompt(inputs, request_data, logging_obj), + inputs, + "Prompt moderation", + ) + if input_type == "response": + return await self._guarded( + self._moderate_response(inputs, request_data, logging_obj), + inputs, + "Response moderation", + ) + return inputs + + @staticmethod + async def _guarded( + coro: Any, + inputs: GenericGuardrailAPIInputs, + label: str, + ) -> GenericGuardrailAPIInputs: + """Await a moderation coroutine fail-open: re-raise an intentional + block, log at critical for malformed service responses, and swallow + any other error returning ``inputs`` unchanged.""" try: - return await self._check_tool_calls(inputs, tool_calls, request_data, logging_obj) + return await coro except ModifyResponseException: raise except _MalformedToolBlockingResponseError as e: - # Distinct from transient errors: the service responded but the - # payload was structurally invalid, which usually indicates a - # misconfigured webhook or a breaking change in its response - # format. Log loudly so operators notice their tool-blocking - # policy is not actually being enforced. + # The service responded but the payload was structurally invalid, + # which usually indicates a misconfigured webhook or a breaking + # change in its response format. Log loudly so operators notice + # their moderation policy is not actually being enforced. verbose_logger.critical( - "Tool blocking service returned a malformed response: %s. " - "Tool calls are NOT being checked -- verify the webhook " - "configuration. Returning original response unchanged.", + "Response moderation service returned a malformed response: %s. " + "Requests are NOT being checked -- verify the webhook " + "configuration. Returning original inputs unchanged.", e, exc_info=True, ) return inputs except Exception as e: verbose_logger.error( - f"Tool blocking hook failed: {e}. Returning original response unchanged.", + f"{label} hook failed: {e}. Returning original inputs unchanged.", exc_info=True, ) return inputs - async def _check_tool_calls( + async def _moderate_response( 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) + """Send response text + tool calls to the after_completion webhook and + raise if either the response text or any tool call is blocked.""" + tool_calls = inputs.get("tool_calls") + texts = inputs.get("texts") + if not tool_calls and not texts: + return inputs - 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 + message_tool_calls = self._normalize_tool_calls(tool_calls or ()) + sent_content = self._join_texts(texts) + + call_details = getattr(logging_obj, "model_call_details", _EMPTY_MAPPING) if logging_obj else _EMPTY_MAPPING 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) + # The moderation payload's ``id`` becomes the tool-blocking log's + # correlation key (the S3 filename), so it must match the failure + # (response) log written for the same blocked request. Both use + # ``litellm_call_id`` -- see ``_correlation_id``. + request_id = self._correlation_id(call_details, request_data) - service_response = await self._post_to_tool_blocking_service(response_data, req_data) - blocked_explanation = self._extract_blocked_tools(service_response, message_tool_calls) + response_data = self._build_response_moderation_payload(message_tool_calls, sent_content, request_id) + req_data = self._extract_request_data(call_details, request_data) - if blocked_explanation is not None: + service_response = await self._post_to_response_moderation_endpoint(response_data, req_data) + blocked = self._extract_response_block(service_response, message_tool_calls, sent_content) + + if blocked: model = self._resolve_model(request_data, call_details) + self._stash_block_context(logging_obj, request_data) raise ModifyResponseException( - message=blocked_explanation, + message=blocked.explanation, model=model, request_data=request_data, guardrail_name=self.guardrail_name, @@ -245,43 +328,125 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): 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__}") - return result + async def _moderate_prompt( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + logging_obj: Optional["LiteLLMLoggingObj"], + ) -> GenericGuardrailAPIInputs: + """Send the (normalized) prompt to the before_prompt webhook and raise + if the prompt is blocked.""" + messages = inputs.get("structured_messages") + if not messages: + # For non-chat request types (e.g. /v1/completions), litellm + # supplies the prompt as ``texts`` with no structured_messages. + # Synthesise a user-message so the webhook can evaluate the prompt. + texts = inputs.get("texts") + if texts: + joined = "\n".join(t for t in texts if t) + if joined: + messages = [{"role": "user", "content": joined}] + if not messages: + return inputs + + payload = self._build_prompt_moderation_payload(inputs, request_data) + service_response = await self._post_to_prompt_moderation_endpoint(payload) + refusal = self._extract_prompt_refusal(service_response) + if refusal is None: + return inputs + + model = inputs.get("model") or request_data.get("model") or "unknown" + self._stash_block_context(logging_obj, request_data) + raise ModifyResponseException( + message=refusal, + model=model, + request_data=request_data, + guardrail_name=self.guardrail_name, + ) @staticmethod - def _build_tool_call_payload( - tool_calls: list[ChatCompletionMessageToolCall], + def _stash_block_context( + logging_obj: Optional["LiteLLMLoggingObj"], + request_data: dict, + ) -> None: + """Stash signals so the deferred success-event skips this request and + ``async_post_call_failure_hook`` can build the failure payload. + + - Sets a flag on ``logging_obj.model_call_details`` so the deferred + success-event handler short-circuits. + - Stashes a reference to ``logging_obj`` on ``request_data`` under a + custom key. ``ProxyLogging.post_call_failure_hook`` pops only + ``litellm_logging_obj`` before iterating callbacks, so this key + survives. + + When ``logging_obj`` is ``None`` the success-event has no way to + observe the block (the flag has nowhere to live), so we log an error + instead of silently dropping the signal. + """ + if logging_obj is None: + verbose_logger.error( + "Rubrik: moderation block fired with logging_obj=None for " + f"litellm_call_id={request_data.get('litellm_call_id')}; " + "cannot suppress success event or attach failure payload." + ) + request_data["_rubrik_logging_obj"] = None + return + logging_obj.model_call_details["_rubrik_blocked"] = True + request_data["_rubrik_logging_obj"] = logging_obj + + @staticmethod + def _normalize_tool_calls(tool_calls: Any) -> tuple[ChatCompletionMessageToolCall, ...]: + """Convert tool_calls from inputs to ChatCompletionMessageToolCall objects.""" + return tuple(RubrikLogger._normalize_tool_call(tc) for tc in tool_calls) + + @staticmethod + def _normalize_tool_call(tc: Any) -> ChatCompletionMessageToolCall: + if isinstance(tc, ChatCompletionMessageToolCall): + return tc + if isinstance(tc, dict): + func = tc.get("function") or _EMPTY_MAPPING + return ChatCompletionMessageToolCall( + id=tc.get("id", ""), + type=tc.get("type", "function"), + function=Function( + name=func.get("name", ""), + arguments=func.get("arguments", ""), + ), + ) + if hasattr(tc, "id") and hasattr(tc, "function"): + return ChatCompletionMessageToolCall( + id=tc.id or "", + type=getattr(tc, "type", None) or "function", + function=tc.function, + ) + raise TypeError(f"Cannot normalize tool_call of type {type(tc).__name__}: {tc!r}") + + @staticmethod + def _join_texts(texts: Any) -> str: + """Join response text segments into the single content string the + webhook evaluates. Empty when there is no assistant text.""" + if not texts: + return "" + return "\n".join(t for t in texts if t) + + @staticmethod + def _build_response_moderation_payload( + tool_calls: Sequence[ChatCompletionMessageToolCall], + content: str, request_id: str | None, - ) -> dict[str, Any]: - """Build a full OpenAI ChatCompletion-format dict for the blocking service.""" + ) -> Mapping[str, Any]: + """Build an OpenAI ChatCompletion-format dict (assistant text + tool + calls) for the after_completion webhook. + + ``content`` is sent so the webhook can moderate the response text; + ``None`` when the assistant produced no text (tool-call-only response). + """ + message: dict[str, Any] = { + "role": "assistant", + "content": content or None, + } + if tool_calls: + message["tool_calls"] = tuple(tc.model_dump(exclude_none=True) for tc in tool_calls) return { "id": request_id or f"chatcmpl-{uuid.uuid4()}", "object": "chat.completion", @@ -290,42 +455,133 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): "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", + "message": message, + "finish_reason": "tool_calls" if tool_calls else "stop", } ], } @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 {} + def _flatten_messages_for_moderation(messages: Any) -> tuple[Mapping[str, Any], ...]: + """Collapse each message's content to a plain string for the webhook. + + litellm normalizes Anthropic ``/v1/messages`` requests to OpenAI shape, + but a turn sent as content-parts (``[{"type": "text", ...}]``) stays a + list. The before_prompt webhook reads ``content`` as a string and drops + non-string content, so we flatten text parts here (images skipped, per + ``convert_content_list_to_str``) -- otherwise block-content prompts + would pass through unmoderated. Builds a new list; never mutates the + shared ``structured_messages``. + """ + return tuple( + { + "role": message.get("role"), + "content": "\n".join(p for p in RubrikLogger._moderation_text_parts(message) if p), + } + for message in messages or () + if isinstance(message, dict) + ) + + @staticmethod + def _moderation_text_parts(message: Mapping[str, Any]) -> tuple[str, ...]: + """Every attacker-controlled text segment of a message: its content plus + the arguments of any tool call or deprecated function call.""" + fc = message.get("function_call") + return ( + # Base text content (flattens Anthropic content-part arrays) + convert_content_list_to_str(message), # pyright: ignore[reportArgumentType] # dict[str,Any] is AllMessageValues at runtime + *( + str((tc.get("function") or _EMPTY_MAPPING).get("arguments") or "") + for tc in message.get("tool_calls") or () + if isinstance(tc, dict) + ), + str((fc.get("arguments") if isinstance(fc, dict) else None) or ""), + ) + + @staticmethod + def _build_prompt_moderation_payload( + inputs: GenericGuardrailAPIInputs, + request_data: Mapping[str, Any], + ) -> Mapping[str, Any]: + """Build the bare OpenAI request the before_prompt webhook consumes. + + Unlike the after_completion envelope, this endpoint takes a raw OpenAI + chat-completions request. ``structured_messages`` is litellm's + OpenAI-normalized view of the prompt, so this works for Anthropic + ``/v1/messages`` requests too. Optional fields are sent only when + present so the payload stays clean. + """ + payload: dict[str, Any] = { + "model": inputs.get("model") or request_data.get("model") or "", + "messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")), + } + tools = inputs.get("tools") + if tools is not None: + payload["tools"] = tools + user = request_data.get("user") + if user: + payload["user"] = user + # Fall back to litellm_call_id, the stable cross-provider join key the + # response/tool path uses (see _correlation_id). LiteLLM does not + # populate request_data["correlation_key"]; it carries litellm_call_id. + # The before_prompt webhook skips the *_prompt_moderation.json S3 write + # when correlation_key is empty, so without this the block fires but no + # log is ever written. An explicit correlation_key still wins. + correlation_key = request_data.get("correlation_key") or request_data.get("litellm_call_id") + if correlation_key: + payload["correlation_key"] = correlation_key + return payload + + @staticmethod + def _extract_request_data( + call_details: Mapping[str, Any], + request_data: Mapping[str, Any] | None, + ) -> Mapping[str, Any]: + """Extract original request data from model_call_details for the + response moderation service envelope. + + Includes the agent's declared ``tools`` (OpenAI-format) when available + so the webhook's hallucination evaluator can compare returned tool calls + against the declared tool list. + """ + if not call_details and not request_data: + return _EMPTY_MAPPING + call_details = call_details or _EMPTY_MAPPING + request_data = request_data or _EMPTY_MAPPING + optional_params = call_details.get("optional_params") or _EMPTY_MAPPING + + # Use ``in`` rather than truthy ``or`` so an explicit empty list + # (caller declared the agent has NO tools) is forwarded as-is. + # The response moderation service uses that signal to flag tool-call + # hallucinations -- ``or`` would mask it by falling through to + # optional_params. + if "tools" in request_data: + tools = request_data["tools"] + else: + tools = optional_params.get("tools") + + # The response moderation service consumes only messages/model/tools. + # Don't forward proxy_server_request -- in litellm >=1.83 its ``body`` + # snapshot carries a UserAPIKeyAuth instance that breaks json.dumps, + # silently fail-opening the guardrail. return { "messages": call_details.get("messages"), "model": call_details.get("model"), - "proxy_server_request": RubrikLogger._sanitize_proxy_server_request( - litellm_params.get("proxy_server_request") - ), + "tools": tools, } @staticmethod def _sanitize_proxy_server_request(proxy_server_request: Any) -> Any: """Allowlist only routing fields (``url``, ``method``) when forwarding - ``proxy_server_request`` to the external Rubrik webhook, dropping - inbound ``headers`` (Authorization, Cookie, x-api-key, ...) and the raw + ``proxy_server_request`` to an external webhook, dropping inbound + ``headers`` (Authorization, Cookie, x-api-key, ...) and the raw request ``body`` so proxy credentials are not exfiltrated.""" if not isinstance(proxy_server_request, dict): return proxy_server_request return {key: proxy_server_request[key] for key in ("url", "method") if key in proxy_server_request} @staticmethod - def _resolve_model(request_data: dict[str, Any], call_details: dict[str, Any]) -> str: + def _resolve_model(request_data: Mapping[str, Any], call_details: Mapping[str, Any]) -> str: """Get the model name for the ModifyResponseException.""" response = request_data.get("response") if response and hasattr(response, "model"): @@ -334,8 +590,70 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # -- Logging hooks --------------------------------------------------------- - async def _prepare_log_payload(self, kwargs: dict, event_type: str) -> StandardLoggingPayload | None: - """Shared logic for success and failure logging.""" + @staticmethod + def _correlation_id(call_details: Mapping[str, Any], request_data: Mapping[str, Any] | None = None) -> str | None: + """The id that joins a blocked request's two S3 logs by filename: the + moderation (``_blocking``) log and the failure (response) log. + + Always ``litellm_call_id``. It is assigned at request start and is + present identically in both the guardrail path (``model_call_details`` + / ``request_data``) and the failure-hook path. Unlike ``response.id`` + or ``standard_logging_object["id"]`` it is immune to the race where a + block fires before the response/logging object is populated, so the + two logs correlate for every provider (OpenAI and Anthropic alike). + """ + return call_details.get("litellm_call_id") or (request_data or _EMPTY_MAPPING).get("litellm_call_id") + + @classmethod + def _apply_correlation_id(cls, payload: dict[str, Any], source: Mapping[str, Any]) -> None: + """Pin ``payload["id"]`` to ``litellm_call_id`` in place so this log + shares its S3 filename id with the moderation (``_blocking``) and + failure logs for the same request -- for every provider. + + ``standard_logging_object["id"]`` is the provider response id + (``response_obj.get("id", litellm_call_id)``), a ``chatcmpl-*`` value + for OpenAI, which would not correlate. ``litellm_call_id`` is assigned + at request start and is identical across all log paths. Falls back to + the existing id when ``litellm_call_id`` is somehow absent rather than + writing a null filename key. + + ``source`` may be ``model_call_details`` directly or a ``kwargs`` dict + that aliases it -- same shape either way. + """ + correlated = cls._correlation_id(source) + if correlated: + payload["id"] = correlated + + @staticmethod + def _prepend_system_prompt(payload: dict[str, Any], source: Mapping[str, Any]) -> None: + """Prepend ``source["system"]`` onto ``payload["messages"]``. + + Builds a NEW messages list rather than mutating ``payload["messages"]`` + in place. The fallback branch of ``_prepare_block_failure_payload`` + aliases ``call_details["messages"]`` directly, so an in-place + ``list.insert(0, ...)`` would mutate the shared source dict. + + No-op if no system prompt is present. Tolerates list/dict/str + message shapes; on unexpected shape, leaves payload alone. + """ + system_prompt = source.get("system") + if not system_prompt: + return + try: + system_scaffold = {"role": "system", "content": system_prompt} + messages = payload.get("messages") + if isinstance(messages, list): + payload["messages"] = (system_scaffold, *messages) + elif isinstance(messages, (dict, str)): + payload["messages"] = (system_scaffold, messages) + except Exception as e: + verbose_logger.warning( + f"Rubrik: failed to prepend system prompt: {e}", + exc_info=True, + ) + + async def _prepare_log_payload(self, kwargs: Mapping[str, Any], event_type: str) -> StandardLoggingPayload | None: + """Shared logic for success logging (sampled).""" if random.random() > self.sampling_rate: verbose_logger.debug(f"Skipping Rubrik {event_type} logging (sampling_rate={self.sampling_rate})") return None @@ -343,59 +661,17 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # Deep-copy so mutations don't affect other callbacks sharing this object standard_logging_payload: StandardLoggingPayload = safe_deep_copy(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, - ) + self._apply_correlation_id(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime + self._prepend_system_prompt(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime return standard_logging_payload - async def _enqueue_log_event(self, kwargs: dict, event_type: str): - try: - self._ensure_periodic_flush_task() - payload = await self._prepare_log_payload(kwargs, event_type) - if payload is None: - return - - self.log_queue.append(payload) - self._enforce_max_queue_size() - - 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 _append_and_maybe_flush(self, payload) -> None: + self._ensure_periodic_flush_task() + self.log_queue.append(payload) + self._enforce_max_queue_size() + if len(self.log_queue) >= self.batch_size: + await self.flush_queue() def _enforce_max_queue_size(self) -> None: overflow = len(self.log_queue) - self.max_queue_size @@ -415,18 +691,213 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): self._dropped_since_warning = 0 self._last_drop_warning_time = now + async def _enqueue_log_event(self, kwargs: Mapping[str, Any], event_type: str): + try: + payload = await self._prepare_log_payload(kwargs, event_type) + if payload is None: + return + await self._append_and_maybe_flush(payload) + 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): + # Blocked requests are logged via async_post_call_failure_hook; + # skip here to avoid double-logging the pre-block response. + if kwargs.get("_rubrik_blocked"): + verbose_logger.debug( + f"Rubrik: skipping success event for blocked request litellm_call_id={kwargs.get('litellm_call_id')}" + ) + return await self._enqueue_log_event(kwargs, "success") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + # Log regular LLM failures (timeouts, upstream errors, etc.) to Rubrik. + # NOTE: ``ModifyResponseException`` blocks are NOT routed here; they + # bypass ``Logging.async_failure_handler`` entirely and reach + # ``async_post_call_failure_hook`` instead. So there is no risk of + # double-logging a block through this path. await self._enqueue_log_event(kwargs, "failure") + async def async_post_call_failure_hook( + self, + request_data: dict, + original_exception: Exception, + user_api_key_dict: Any, + traceback_str: str | None = None, + ) -> None: + """Log blocked requests signalled via ``ModifyResponseException`` + (prompt blocks, response/tool blocks, streaming blocks). + + Carries the stashed ``_rubrik_logging_obj``. For every other + exception we no-op; LiteLLM's standard failure plumbing handles those. + """ + if not isinstance(original_exception, ModifyResponseException): + return + + # Guard by guardrail_name so that when multiple Rubrik instances are + # registered, only the instance that raised the block handles it. + # The failure hook is called for every registered callback; without + # this check the first instance pops the stash and the originating + # instance finds None and silently skips logging. + if getattr(original_exception, "guardrail_name", None) != self.guardrail_name: + return + + logging_obj = request_data.pop("_rubrik_logging_obj", None) + if logging_obj is None: + # Legitimate when a non-Rubrik guardrail raised the block; + # problematic if Rubrik did and the stash was lost (e.g. + # ``_stash_block_context`` ran with ``logging_obj=None``). Either + # way we cannot build the payload. + verbose_logger.warning( + "Rubrik: block exception without stashed logging_obj. " + f"litellm_call_id={request_data.get('litellm_call_id')}, " + f"model={request_data.get('model')}, " + f"user_id={getattr(user_api_key_dict, 'user_id', None)}, " + f"raising_guardrail=" + f"{getattr(original_exception, 'guardrail_name', None)}" + ) + return + + call_id: str | None = None + await self._build_and_enqueue_block_event(logging_obj, original_exception, call_id) + + async def _build_and_enqueue_block_event( + self, + logging_obj: "LiteLLMLoggingObj", + exception: "ModifyResponseException", + call_id: str | None, + ) -> None: + try: + call_details = logging_obj.model_call_details + # Do NOT pop "_rubrik_blocked" here. The deferred success-handler + # task may still be iterating callbacks, and popping mid-iteration + # (between two awaited callback invocations) would cause this + # plugin's success-event callback to read the flag as absent and + # log the pre-block response -- the exact bug this hook exists to + # prevent. The flag dies with model_call_details when the request + # completes; there's nothing to clean up. + call_id = call_details.get("litellm_call_id") + payload = self._prepare_block_failure_payload(logging_obj, exception) + except (AttributeError, KeyError, TypeError) as e: + verbose_logger.error( + f"Rubrik: failed to build blocked-tool payload for " + f"litellm_call_id={call_id}: {e}. Event will NOT be logged.", + exc_info=True, + ) + return + + try: + await self._append_and_maybe_flush(payload) + except Exception as e: + verbose_logger.error( + f"Rubrik: failed to enqueue blocked-tool event for litellm_call_id={call_id}: {e}.", + exc_info=True, + ) + + def _prepare_block_failure_payload( + self, + logging_obj: "LiteLLMLoggingObj", + exception: "ModifyResponseException", + ) -> StandardLoggingPayload: + """Build a failure-style payload using the exception text as response. + + Blocked-tool events are security-relevant and **bypass sampling**: + every block is logged. + + The deferred success-handler runs as a separately-scheduled task and + races with this hook, so ``standard_logging_object`` on + ``model_call_details`` may not yet be populated. If present we reuse + it; otherwise we fall back to a best-effort payload built from the + fields available at block time. + + For prompt blocks the LLM is never called, so ``standard_logging_object`` + is never populated. The fallback therefore must carry enough fields to + pass the log processor's ``LogEntry`` schema (``BaseLogEntry`` requires + ``metadata``, ``model_id``, ``model_group``, ``model_parameters``, + ``startTime``, ``endTime``, and ``completionStartTime``). Without a + parseable payload the log processor discards the entry with a parse + error and no session is created, so prompt-moderation violations are + silently dropped even though the ``_prompt_moderation.json`` forensic + log is written correctly. + + Field sourcing for the fallback path: + - ``model`` / ``model_group``: ``call_details["model"]`` -- this is the + model-group name (e.g. "gpt-4o") set by the proxy before the guardrail + fires. The router writes ``metadata["model_group"]`` only inside + ``acompletion()``, which hasn't run yet for a prompt block. + - ``model_id``: not available before the LLM returns hidden_params; + defaults to empty string. + - ``user_api_key_hash``: ``call_details["metadata"]["user_api_key"]`` -- + the hashed token written by ``add_user_information_to_request_data`` + before ``pre_call_hook`` fires. + - time fields: ``call_details["start_time"]`` reused for all three; + end/completion times are meaningless for a prompt block. + """ + call_details = logging_obj.model_call_details + exception_text = f"{type(exception).__name__}: {exception.message}" + + base = call_details.get("standard_logging_object") + if base is not None: + payload: dict = safe_deep_copy(base) + else: + verbose_logger.debug( + "Rubrik: standard_logging_object not yet on model_call_details " + f"for litellm_call_id={call_details.get('litellm_call_id')}; " + "using best-effort fallback payload." + ) + payload = self._build_fallback_payload(call_details) + + payload["response"] = exception_text + + # Pin the correlation key to litellm_call_id so this failure log shares + # its S3 filename id with the moderation (``_blocking``) log for the + # same request. The copied ``standard_logging_object["id"]`` is + # ``response_obj.get("id", litellm_call_id)`` -- a provider ``chatcmpl-*`` + # value for OpenAI -- which would not correlate; overwrite it. + payload["id"] = self._correlation_id(call_details) or f"chatcmpl-{uuid.uuid4()}" + self._prepend_system_prompt(payload, call_details) + + return payload # type: ignore[return-value] + + @staticmethod + def _build_fallback_payload(call_details: Mapping[str, Any]) -> dict[str, Any]: + _metadata: Mapping[str, Any] = call_details.get("metadata") or _EMPTY_MAPPING + # Convert datetime to a Unix float so json.dumps can serialize it. + # httpx's json= parameter uses stdlib json.dumps with no custom encoder. + _raw_start = call_details.get("start_time") + _start = _raw_start.timestamp() if _raw_start is not None else None + return { + "id": call_details.get("litellm_call_id"), + "model": call_details.get("model") or "", + # model_group is set by the router inside acompletion(), which + # hasn't run for a prompt block; use the model name instead. + "model_group": call_details.get("model") or "", + # model_id comes from response.hidden_params -- unavailable here. + "model_id": "", + "model_parameters": ModelParamHelper.get_standard_logging_model_parameters( + call_details.get("optional_params") or _EMPTY_MAPPING # pyright: ignore[reportArgumentType] # helper only reads the mapping + ), + "startTime": _start, + "endTime": _start, + "completionStartTime": _start, + "messages": call_details.get("messages") or (), + "metadata": { + # "user_api_key" is the hashed token written by + # add_user_information_to_request_data before guardrails fire. + "user_api_key_hash": _metadata.get("user_api_key_hash") or _metadata.get("user_api_key") or "", + }, + "status": "failure", + } + # -- Batch logging --------------------------------------------------------- async def _log_batch_to_rubrik(self, data): - # NOTE: this method intentionally re-raises on failure so the parent - # CustomBatchLogger.flush_queue keeps the unsent events in the queue - # for the next flush attempt instead of silently dropping them. + # NOTE: this method intentionally re-raises on failure so flush_queue + # can preserve the unsent events for the next flush attempt instead of + # silently dropping them. try: response = await self.async_httpx_client.post( url=self.logging_endpoint, @@ -452,10 +923,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): if not self.log_queue: return - log_queue_snapshot = list(self.log_queue) - verbose_logger.debug("Rubrik: Flushing batch of %s events", len(log_queue_snapshot)) await self._log_batch_to_rubrik( - data=log_queue_snapshot, + data=self.log_queue, ) async def flush_queue(self): @@ -463,8 +932,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): Overrides the base implementation so the same snapshot drives both the HTTP send and the queue truncation. This avoids the subtle - coupling where the base class captures `len(self.log_queue)` - separately from the snapshot taken inside `async_send_batch`, + coupling where the base class captures ``len(self.log_queue)`` + separately from the snapshot taken inside ``async_send_batch``, which could otherwise drift in a future refactor and cause duplicate deliveries to Rubrik. """ @@ -485,70 +954,141 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): del self.log_queue[: len(snapshot)] self.last_flush_time = time.time() - # -- Tool blocking service ------------------------------------------------- + # -- Webhook services ------------------------------------------------------ - async def _post_to_tool_blocking_service( + async def _post_json(self, endpoint: str, payload: Mapping[str, Any], service_name: str) -> Mapping[str, Any]: + """POST ``payload`` to a Rubrik webhook and return its dict response. + + Raises: + Exception: If the service is unavailable or returns an error. + TypeError: If the response JSON is not a dict. + """ + verbose_logger.debug(f"Sending request to {service_name}: {endpoint}") + http_response = await self.moderation_client.post( + endpoint, + json=payload, + headers=self._headers, + ) + http_response.raise_for_status() + result = http_response.json() + if not isinstance(result, dict): + raise TypeError( + f"{service_name} returned non-dict JSON " + f"({type(result).__name__}); expected OpenAI chat completion " + "shape or empty object." + ) + return result + + async def _post_to_response_moderation_endpoint( 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. + response_data: Mapping[str, Any], + request_data: Mapping[str, Any], + ) -> Mapping[str, Any]: + """Post the ``{request, response}`` envelope to the after_completion + webhook and return its (possibly rewritten) 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: {self.tool_blocking_endpoint}") - http_response = await self.tool_blocking_client.post( - self.tool_blocking_endpoint, - json=envelope, - headers=self._headers, + envelope = {"request": request_data, "response": response_data} + return await self._post_json( + self.response_moderation_endpoint, + envelope, + "Response moderation service", ) - http_response.raise_for_status() - result: dict[str, Any] = http_response.json() - return result + + async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, Any]) -> Mapping[str, Any]: + """Post a bare OpenAI request to the before_prompt webhook. + + Returns ``{}`` (passthrough) or a synthetic chat.completion (block). + """ + return await self._post_json(self.prompt_moderation_endpoint, payload, "Prompt moderation service") @staticmethod - def _extract_blocked_tools( - service_response: dict[str, Any], - all_tool_calls: list[ChatCompletionMessageToolCall], - ) -> str | None: - """Return the blocking explanation if any tool calls were blocked. + def _extract_prompt_refusal(service_response: Mapping[str, Any]) -> str | None: + """Return the refusal text when the prompt was blocked, else None. - Compares the service response (which contains only allowed tools) against - the full set of tool calls. Returns ``None`` if all tools are allowed, or - the explanation string (prefixed with newlines) otherwise. + The before_prompt webhook returns ``{}`` (passthrough) or a synthetic + chat.completion whose ``choices[0].message.content`` is the refusal + explanation. + """ + choices = service_response.get("choices") + if not choices: + return None + message = choices[0].get("message") or _EMPTY_MAPPING + content = message.get("content") + return content or "Request blocked by policy." + + @staticmethod + def _extract_response_block( + service_response: Mapping[str, Any], + all_tool_calls: Sequence[ChatCompletionMessageToolCall], + sent_content: str, + ) -> BlockedResponseResult | None: + """Detect whether the webhook moderated the response text or tool calls. + + The after_completion webhook rewrites the response in place with no + explicit "blocked" flag, so we infer a block by diffing what we sent + against what came back: + + - Tool block: a tool call we sent is absent from the returned (allowed) + set. + - Text block: the returned content was REPLACED wholesale (a text + violation), as opposed to having a tool-block explanation APPENDED to + the original content. We tell them apart with ``startswith``, which + mirrors the webhook's own append-vs-replace behavior. + + Returns None when nothing was moderated. A text block supersedes a tool + block (mirroring the webhook, which drops tool calls on a text block). Expects service_response in OpenAI chat completion format: {"choices": [{"message": {"tool_calls": [...], "content": "..."}}]} """ - choices = service_response.get("choices", []) + choices = service_response.get("choices") or () if not choices: - raise _MalformedToolBlockingResponseError("Tool blocking service returned empty response") + raise _MalformedToolBlockingResponseError("Response moderation service returned empty response") - message = choices[0].get("message", {}) - returned_tool_calls = message.get("tool_calls") or [] - blocking_explanation = message.get("content", "") + message = choices[0].get("message") or _EMPTY_MAPPING + returned_tool_calls = message.get("tool_calls") or () + returned_content = message.get("content") or "" - allowed_id_counts: Counter = Counter( - tc["id"] for tc in returned_tool_calls if isinstance(tc, dict) and tc.get("id") - ) - required_id_counts: Counter = Counter(tc.id for tc in all_tool_calls if tc.id) - - all_allowed = len(returned_tool_calls) >= len(all_tool_calls) and all( - allowed_id_counts.get(tc_id, 0) >= count for tc_id, count in required_id_counts.items() + # Use Counter so duplicate IDs are handled correctly: if the model + # emits two calls with the same ID (one allowed, one prohibited) and + # the service returns only the allowed one, a set-based check would + # miss the block. Counter preserves multiplicity. + returned_id_counts: Counter[str] = Counter(tc["id"] for tc in returned_tool_calls if tc.get("id")) + required_id_counts: Counter[str] = Counter(tc.id for tc in all_tool_calls if tc.id) + # Cardinality check catches ID-less tool calls (not counted in + # required_id_counts because tc.id is falsy); Counter check catches + # duplicate-ID attacks where one occurrence is silently removed. + tools_blocked = len(returned_tool_calls) < len(all_tool_calls) or not all( + returned_id_counts.get(tc_id, 0) >= count for tc_id, count in required_id_counts.items() ) - if all_allowed: - return None + # The webhook either replaces content wholesale (text block) or appends + # a tool-block explanation to the original text. ``appended`` tells the + # two apart, and is reused below to recover just the explanation. A text + # block requires there to have been assistant text to block. + # Use the documented ``\n\n`` separator to distinguish a tool-block + # append from a text replacement that shares the original as a prefix. + # Without the separator, a replacement like "Hello, blocked." where the + # original was "Hello" would be classified as an append (not a text + # block) and silently pass through to the client. + appended = bool(sent_content) and returned_content.startswith(f"{sent_content}\n\n") + text_blocked = bool(sent_content) and returned_content != sent_content and not appended - explanation = blocking_explanation or "Tool call blocked by policy." - return f"\n\n{explanation}" + if text_blocked: + return BlockedResponseResult(explanation=returned_content or "Response blocked by policy.") + + if tools_blocked: + if appended: + # Recover just the appended explanation: drop the original text + # and the leading separator the webhook inserted before it. + explanation = returned_content[len(sent_content) :].lstrip("\n") + else: + explanation = returned_content + return BlockedResponseResult(explanation=explanation or "Tool call blocked by policy.") + + return None diff --git a/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py index 4ad29bbeae8..2f2228ae312 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py @@ -10,6 +10,18 @@ if TYPE_CHECKING: def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> RubrikLogger: + """Create and register a RubrikLogger instance. + + The ``mode`` field in the guardrail config controls which surfaces are + moderated: + - ``pre_call`` (or a mode that includes it): prompt moderation via the + ``/v1/before_prompt/openai/v1`` webhook. + - ``post_call`` (the default when ``mode`` is omitted): response and tool + call moderation via the ``/v1/after_completion/openai/v1`` webhook. + + Both hooks are active when ``mode`` covers both ``pre_call`` and + ``post_call``. + """ import litellm rubrik_callback = RubrikLogger( diff --git a/tests/test_litellm/integrations/test_rubrik.py b/tests/test_litellm/integrations/test_rubrik.py index 922d2fe8a15..7f589dc15bf 100644 --- a/tests/test_litellm/integrations/test_rubrik.py +++ b/tests/test_litellm/integrations/test_rubrik.py @@ -1,8 +1,8 @@ """ 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. +Covers initialization, apply_guardrail (prompt moderation + response/tool +blocking), batch logging, and Anthropic format handling. """ import os @@ -13,8 +13,10 @@ import httpx import pytest from litellm.integrations.custom_guardrail import ModifyResponseException -from litellm.integrations.rubrik import RubrikLogger -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.integrations.rubrik import ( + RubrikLogger, + _MalformedToolBlockingResponseError, +) from tests.test_litellm.integrations.rubrik_test_helpers import ( make_inputs_with_tools, @@ -50,19 +52,19 @@ class TestInitialization: with patch("asyncio.create_task", Mock()): handler = RubrikLogger() assert ( - handler.tool_blocking_endpoint + handler.response_moderation_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, AsyncHTTPHandler) + assert handler.moderation_client is not None 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 + handler.response_moderation_endpoint == "http://ctor-host:9090/v1/after_completion/openai/v1" ) @@ -82,7 +84,7 @@ class TestInitialization: with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://localhost:8080/"}): with patch("asyncio.create_task", Mock()): assert ( - RubrikLogger().tool_blocking_endpoint + RubrikLogger().response_moderation_endpoint == "http://localhost:8080/v1/after_completion/openai/v1" ) @@ -90,13 +92,13 @@ class TestInitialization: with patch("asyncio.create_task", Mock()): with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://host/v1"}): assert ( - RubrikLogger().tool_blocking_endpoint + RubrikLogger().response_moderation_endpoint == "http://host/v1/after_completion/openai/v1" ) with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://host/v11"}): assert ( - RubrikLogger().tool_blocking_endpoint + RubrikLogger().response_moderation_endpoint == "http://host/v11/v1/after_completion/openai/v1" ) @@ -155,10 +157,10 @@ class TestInitialization: # Do NOT patch asyncio.create_task — the real call should be # guarded and fall back gracefully when there is no event loop. handler = RubrikLogger() - assert handler.tool_blocking_endpoint.startswith("http://localhost:8080") + assert handler.response_moderation_endpoint.startswith("http://localhost:8080") # Without a running loop at init, the periodic flush task should be # deferred so batches still get drained once a log event arrives. - assert handler._flush_task is None + assert handler._periodic_flush_task is None @pytest.mark.asyncio async def test_periodic_flush_task_started_lazily_on_first_log(self, mock_env): @@ -170,7 +172,7 @@ class TestInitialization: side_effect=RuntimeError("no running loop"), ): handler = RubrikLogger() - assert handler._flush_task is None + assert handler._periodic_flush_task is None kwargs = { "standard_logging_object": { @@ -183,8 +185,8 @@ class TestInitialization: with patch.object(handler, "_log_batch_to_rubrik", AsyncMock()): await handler.async_log_success_event(kwargs, None, None, None) - assert handler._flush_task is not None - handler._flush_task.cancel() + assert handler._periodic_flush_task is not None + handler._periodic_flush_task.cancel() def test_event_hook_defaults_to_post_call_when_none_passed(self, mock_env): """`initialize_guardrail` always passes ``event_hook=litellm_params.mode`` @@ -204,14 +206,13 @@ class TestInitialization: handler = RubrikLogger(event_hook=GuardrailEventHooks.pre_call) assert handler.event_hook == GuardrailEventHooks.pre_call - def test_default_on_defaults_to_true_when_none_passed(self, mock_env): - """`initialize_guardrail` always passes ``default_on=litellm_params.default_on`` - (which is ``None`` when the user omits ``default_on``). The logger must - coerce a None ``default_on`` to True, otherwise ``should_run_guardrail`` - (which checks ``self.default_on is True``) silently skips the guardrail.""" + def test_default_on_defaults_to_false_when_none_passed(self, mock_env): + """Follows the standard litellm pattern: omitted ``default_on`` resolves + to ``False`` (off by default). Users must explicitly set + ``default_on: true`` to enable the guardrail for all requests.""" with patch("asyncio.create_task", Mock()): handler = RubrikLogger(default_on=None) - assert handler.default_on is True + assert handler.default_on is False def test_explicit_default_on_false_preserved(self, mock_env): """A user explicitly setting ``default_on: false`` in their guardrail @@ -421,7 +422,7 @@ class TestBatchLogging: ) assert len(handler.log_queue) == 1 msgs = handler.log_queue[0]["messages"] - assert isinstance(msgs, list) + assert isinstance(msgs, tuple) assert msgs[0]["role"] == "system" assert msgs[1] == {"role": "user", "content": "hi"} @@ -444,7 +445,10 @@ class TestBatchLogging: ) assert handler.log_queue[0]["id"] == "litellm-call-123" - async def test_non_anthropic_id_unchanged(self, handler): + async def test_litellm_call_id_always_used_as_correlation_key(self, handler): + """The merged plugin always uses litellm_call_id as the log ID for all + providers (not just Anthropic) so that logs correlate with the + moderation (_blocking) and failure logs for the same request.""" kwargs = { "standard_logging_object": { "id": "chatcmpl-original", @@ -461,7 +465,7 @@ class TestBatchLogging: 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" + assert handler.log_queue[0]["id"] == "litellm-call-123" async def test_payload_deep_copied_not_mutated(self, handler): """Verify the shared standard_logging_object is not mutated.""" @@ -536,7 +540,7 @@ class TestApplyGuardrail: tc2 = make_tool_call_dict("call_2", "get_time") inputs = make_inputs_with_tools([tc1, tc2]) - handler.tool_blocking_client = _echo_service() + handler.moderation_client = _echo_service() result = await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" @@ -548,7 +552,7 @@ class TestApplyGuardrail: tc2 = make_tool_call_dict("call_2", "drop_database") inputs = make_inputs_with_tools([tc1, tc2]) - handler.tool_blocking_client = _mock_service_response( + handler.moderation_client = _mock_service_response( { "choices": [ { @@ -594,7 +598,7 @@ class TestApplyGuardrail: mock_client = AsyncMock() mock_client.post = mock_post - handler.tool_blocking_client = mock_client + handler.moderation_client = mock_client with pytest.raises(ModifyResponseException): await handler.apply_guardrail( @@ -607,7 +611,7 @@ class TestApplyGuardrail: mock_client = AsyncMock() mock_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout")) - handler.tool_blocking_client = mock_client + handler.moderation_client = mock_client result = await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" @@ -618,7 +622,7 @@ class TestApplyGuardrail: tc1 = make_tool_call_dict("call_1", "test_tool") inputs = make_inputs_with_tools([tc1]) - handler.tool_blocking_client = _mock_service_response({"choices": []}) + handler.moderation_client = _mock_service_response({"choices": []}) result = await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" @@ -641,7 +645,7 @@ class TestApplyGuardrail: mock_client = AsyncMock() mock_client.post = mock_post - handler.tool_blocking_client = mock_client + handler.moderation_client = mock_client await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" @@ -675,7 +679,7 @@ class TestApplyGuardrail: mock_client = AsyncMock() mock_client.post = mock_post - handler.tool_blocking_client = mock_client + handler.moderation_client = mock_client logging_obj = Mock() logging_obj.model_call_details = { @@ -697,7 +701,10 @@ class TestApplyGuardrail: assert req["model"] == "gpt-4" assert req["messages"] == [{"role": "user", "content": "hi"}] - async def test_proxy_server_request_headers_stripped(self, handler): + async def test_proxy_server_request_not_forwarded(self, handler): + """proxy_server_request is intentionally NOT included in the request + envelope: in litellm >=1.83 its ``body`` carries a UserAPIKeyAuth + instance that breaks json.dumps, silently fail-opening the guardrail.""" tc = make_tool_call_dict("call_1", "test_tool") inputs = make_inputs_with_tools([tc]) @@ -712,7 +719,7 @@ class TestApplyGuardrail: mock_client = AsyncMock() mock_client.post = mock_post - handler.tool_blocking_client = mock_client + handler.moderation_client = mock_client logging_obj = Mock() logging_obj.model_call_details = { @@ -739,8 +746,8 @@ class TestApplyGuardrail: logging_obj=logging_obj, ) - forwarded = captured_payload["request"]["proxy_server_request"] - assert forwarded == {"url": "/chat/completions", "method": "POST"} + # proxy_server_request is deliberately excluded from the forwarded envelope + assert "proxy_server_request" not in captured_payload["request"] # -- Anthropic format ---------------------------------------------------------- @@ -760,7 +767,7 @@ class TestApplyGuardrailAnthropicFormat: ) inputs = make_inputs_with_tools([tc], texts=["I'll check the weather."]) - handler.tool_blocking_client = _echo_service() + handler.moderation_client = _echo_service() result = await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" @@ -771,7 +778,7 @@ class TestApplyGuardrailAnthropicFormat: tc = make_tool_call_dict("toolu_123", "dangerous_tool", '{"arg": "value"}') inputs = make_inputs_with_tools([tc]) - handler.tool_blocking_client = _mock_service_response( + handler.moderation_client = _mock_service_response( { "choices": [ { @@ -790,21 +797,21 @@ class TestApplyGuardrailAnthropicFormat: inputs=inputs, request_data={}, input_type="response" ) - async def test_text_only_response_no_blocking(self, handler): + async def test_text_only_response_sent_to_moderation(self, handler): + """Text-only responses (no tool calls) are sent to the response + moderation service to check the assistant's text content.""" 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 + # Service allows the response (returns the content unchanged) + handler.moderation_client = _echo_service() 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"}') @@ -812,7 +819,7 @@ class TestApplyGuardrailAnthropicFormat: mock_client = AsyncMock() mock_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout")) - handler.tool_blocking_client = mock_client + handler.moderation_client = mock_client result = await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="response" @@ -850,10 +857,13 @@ class TestNormalizeToolCalls: RubrikLogger._normalize_tool_calls(["not_a_tool_call"]) -# -- Extract blocked tools ----------------------------------------------------- +# -- Extract response block ---------------------------------------------------- -class TestExtractBlockedTools: +class TestExtractResponseBlock: + """Tests for _extract_response_block, which replaces the upstream + _extract_blocked_tools and handles both text blocks and tool blocks.""" + def test_all_allowed_returns_none(self): from litellm.types.utils import ChatCompletionMessageToolCall, Function @@ -870,7 +880,7 @@ class TestExtractBlockedTools: } ] } - result = RubrikLogger._extract_blocked_tools(service_resp, [tc]) + result = RubrikLogger._extract_response_block(service_resp, [tc], "") assert result is None def test_some_blocked_returns_explanation(self): @@ -896,13 +906,13 @@ class TestExtractBlockedTools: } ] } - result = RubrikLogger._extract_blocked_tools(service_resp, [tc1, tc2]) + result = RubrikLogger._extract_response_block(service_resp, [tc1, tc2], "") assert result is not None - assert "blocked fn2" in result + assert "blocked fn2" in result.explanation def test_empty_choices_raises(self): - with pytest.raises(Exception, match="empty response"): - RubrikLogger._extract_blocked_tools({"choices": []}, []) + with pytest.raises(_MalformedToolBlockingResponseError): + RubrikLogger._extract_response_block({"choices": []}, [], "") def test_null_tool_calls_treated_as_all_blocked(self): from litellm.types.utils import ChatCompletionMessageToolCall, Function @@ -920,36 +930,55 @@ class TestExtractBlockedTools: } ] } - result = RubrikLogger._extract_blocked_tools(service_resp, [tc]) + result = RubrikLogger._extract_response_block(service_resp, [tc], "") assert result is not None - assert "blocked everything" in result + assert "blocked everything" in result.explanation - def test_duplicate_ids_block_when_only_one_returned(self): + def test_text_block_detected(self): + """When the service replaces the response text wholesale, it's a text block.""" from litellm.types.utils import ChatCompletionMessageToolCall, Function - tc1 = ChatCompletionMessageToolCall( - id="call_dup", - type="function", - function=Function(name="fn", arguments="{}"), - ) - tc2 = ChatCompletionMessageToolCall( - id="call_dup", - type="function", - function=Function(name="fn", arguments="{}"), - ) service_resp = { "choices": [ { "message": { - "tool_calls": [{"id": "call_dup"}], - "content": "blocked duplicate", + "tool_calls": [], + "content": "This content violates policy.", } } ] } - result = RubrikLogger._extract_blocked_tools(service_resp, [tc1, tc2]) + result = RubrikLogger._extract_response_block( + service_resp, [], "Original assistant text." + ) assert result is not None - assert "blocked duplicate" in result + assert "violates policy" in result.explanation + + def test_tool_block_with_appended_explanation(self): + """When the service appends an explanation to the original text, only the + appended part is returned as the explanation.""" + from litellm.types.utils import ChatCompletionMessageToolCall, Function + + tc = ChatCompletionMessageToolCall( + id="call_1", type="function", function=Function(name="fn", arguments="{}") + ) + original_text = "Here is my response." + appended_explanation = "Tool call was blocked." + service_resp = { + "choices": [ + { + "message": { + "tool_calls": [], + "content": original_text + "\n\n" + appended_explanation, + } + } + ] + } + result = RubrikLogger._extract_response_block( + service_resp, [tc], original_text + ) + assert result is not None + assert appended_explanation in result.explanation # -- Sanitize proxy server request ------------------------------------------- @@ -1010,3 +1039,796 @@ class TestResolveModel: {"response": response}, {"model": "fallback"} ) assert result == "unknown" + + +# -- Additional Initialization edge cases ------------------------------------ + + +class TestInitializationEdgeCases: + def test_batch_size_zero_uses_default(self): + """RUBRIK_BATCH_SIZE=0 must warn and fall back to the default.""" + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_BATCH_SIZE": "0"}, + ): + h = RubrikLogger() + # Should use default, not 0 + assert h.batch_size > 0 + + def test_batch_size_negative_uses_default(self): + """RUBRIK_BATCH_SIZE=-1 must warn and fall back to the default.""" + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_BATCH_SIZE": "-5"}, + ): + h = RubrikLogger() + assert h.batch_size > 0 + + +# -- aclose() ----------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestAclose: + async def test_aclose_cancels_task_does_not_close_shared_client(self, mock_env): + """aclose() cancels the periodic flush task but does NOT close the shared + moderation_client — closing a shared cached client would break other + RubrikLogger instances that share the same connection pool.""" + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger() + + mock_task = Mock() + mock_task.cancel = Mock() + handler._periodic_flush_task = mock_task + + handler.moderation_client = AsyncMock() + handler.moderation_client.close = AsyncMock() + + await handler.aclose() + + mock_task.cancel.assert_called_once() + handler.moderation_client.close.assert_not_awaited() + + async def test_aclose_with_none_task_does_not_close_client(self, mock_env): + """aclose() with no flush task still does not close the shared client.""" + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger() + + handler._periodic_flush_task = None + handler.moderation_client = AsyncMock() + handler.moderation_client.close = AsyncMock() + + await handler.aclose() + + handler.moderation_client.close.assert_not_awaited() + + +# -- apply_guardrail edge cases ----------------------------------------------- + + +@pytest.mark.asyncio +class TestApplyGuardrailEdgeCases: + async def test_unknown_input_type_returns_inputs_unchanged(self, handler): + """When input_type is not 'request' or 'response', inputs are returned as-is.""" + inputs = make_inputs_with_tools([make_tool_call_dict("call_1", "tool")]) + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="unknown" + ) + assert result is inputs + + async def test_response_with_no_texts_and_no_tool_calls_returns_inputs(self, handler): + """_moderate_response early-returns when both texts and tool_calls are empty.""" + from litellm.types.utils import GenericGuardrailAPIInputs + + inputs = GenericGuardrailAPIInputs() + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + async def test_moderate_response_empty_call_details_emits_warning(self, handler): + """When logging_obj is present but model_call_details is empty, a warning is + logged and moderation proceeds (fail-open on HTTP error).""" + tc = make_tool_call_dict("call_1", "test_tool") + inputs = make_inputs_with_tools([tc]) + + logging_obj = Mock() + logging_obj.model_call_details = {} + + handler.moderation_client = _echo_service() + + result = await handler.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="response", + logging_obj=logging_obj, + ) + assert result is inputs + + +# -- Prompt moderation -------------------------------------------------------- + + +@pytest.mark.asyncio +class TestPromptModeration: + async def test_prompt_moderation_passthrough(self, handler): + """Webhook returns {} (empty dict) → inputs returned unchanged.""" + inputs = {"structured_messages": [{"role": "user", "content": "Hello"}]} + + handler.moderation_client = _mock_service_response({}) + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + async def test_prompt_moderation_blocked_raises(self, handler): + """Webhook returns synthetic chat.completion → raises ModifyResponseException.""" + inputs = { + "structured_messages": [{"role": "user", "content": "Harmful prompt"}], + "model": "gpt-4", + } + + handler.moderation_client = _mock_service_response( + { + "choices": [ + { + "message": { + "role": "assistant", + "content": "This request violates our policy.", + } + } + ] + } + ) + + with pytest.raises(ModifyResponseException) as exc_info: + await handler.apply_guardrail( + inputs=inputs, request_data={"model": "gpt-4"}, input_type="request" + ) + assert "violates our policy" in exc_info.value.message + + async def test_prompt_moderation_no_messages_skips_moderation(self, handler): + """When structured_messages is absent/empty, moderation is skipped.""" + inputs = {"model": "gpt-4"} + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + async def test_prompt_moderation_stashes_logging_obj_on_block(self, handler): + """On a prompt block, _stash_block_context must set the blocked flag.""" + inputs = { + "structured_messages": [{"role": "user", "content": "bad prompt"}], + } + + handler.moderation_client = _mock_service_response( + { + "choices": [ + {"message": {"role": "assistant", "content": "Blocked."}} + ] + } + ) + + logging_obj = Mock() + logging_obj.model_call_details = {} + request_data: dict = {} + + with pytest.raises(ModifyResponseException): + await handler.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + logging_obj=logging_obj, + ) + + assert logging_obj.model_call_details.get("_rubrik_blocked") is True + assert request_data.get("_rubrik_logging_obj") is logging_obj + + +# -- _stash_block_context ----------------------------------------------------- + + +class TestStashBlockContext: + def test_with_non_none_logging_obj_sets_flag_and_stashes(self): + """Sets _rubrik_blocked flag and stores logging_obj on request_data.""" + logging_obj = Mock() + logging_obj.model_call_details = {} + request_data: dict = {} + + RubrikLogger._stash_block_context(logging_obj, request_data) + + assert logging_obj.model_call_details["_rubrik_blocked"] is True + assert request_data["_rubrik_logging_obj"] is logging_obj + + def test_with_none_logging_obj_stores_none_on_request_data(self): + """When logging_obj is None, stores None on request_data (logged as error).""" + request_data: dict = {"litellm_call_id": "test-id"} + + RubrikLogger._stash_block_context(None, request_data) + + assert request_data["_rubrik_logging_obj"] is None + + +# -- _normalize_tool_calls duck-typed ----------------------------------------- + + +class TestNormalizeToolCallsDuckTyped: + def test_duck_typed_object_with_id_and_function_attrs(self): + """Objects that have .id and .function attrs but are not + ChatCompletionMessageToolCall are handled by the third branch.""" + from litellm.types.utils import Function + + tc = Mock() + tc.id = "call_duck" + tc.type = "function" + tc.function = Function(name="duck_tool", arguments='{"x": 1}') + # Make isinstance(..., ChatCompletionMessageToolCall) return False + # by using a plain Mock (not a ChatCompletionMessageToolCall subclass) + + result = RubrikLogger._normalize_tool_calls([tc]) + assert len(result) == 1 + assert result[0].id == "call_duck" + assert result[0].function.name == "duck_tool" + + def test_duck_typed_without_type_defaults_to_function(self): + """getattr(tc, "type", None) falls back to "function" when absent.""" + from litellm.types.utils import Function + + tc = Mock(spec=["id", "function"]) # no .type attr + tc.id = "call_no_type" + tc.function = Function(name="fn", arguments="{}") + + result = RubrikLogger._normalize_tool_calls([tc]) + assert result[0].type == "function" + + +# -- _flatten_messages_for_moderation ----------------------------------------- + + +class TestFlattenMessagesForModeration: + def test_plain_string_content_preserved(self): + messages = [{"role": "user", "content": "Hello world"}] + result = RubrikLogger._flatten_messages_for_moderation(messages) + assert len(result) == 1 + assert result[0]["role"] == "user" + assert result[0]["content"] == "Hello world" + + def test_content_list_flattened_to_string(self): + """Content as a list of parts (e.g. Anthropic multi-part) is flattened.""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Hello from parts"}, + ], + } + ] + result = RubrikLogger._flatten_messages_for_moderation(messages) + assert len(result) == 1 + assert result[0]["role"] == "user" + assert "Hello from parts" in result[0]["content"] + + def test_non_dict_messages_skipped(self): + """Non-dict entries in the messages list are silently skipped.""" + messages = [ + "raw string message", + {"role": "user", "content": "valid"}, + ] + result = RubrikLogger._flatten_messages_for_moderation(messages) + assert len(result) == 1 + assert result[0]["content"] == "valid" + + def test_none_messages_returns_empty(self): + result = RubrikLogger._flatten_messages_for_moderation(None) + assert result == () + + def test_multiple_messages_preserved_in_order(self): + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Question?"}, + ] + result = RubrikLogger._flatten_messages_for_moderation(messages) + assert len(result) == 2 + assert result[0]["role"] == "system" + assert result[1]["role"] == "user" + + +# -- _build_prompt_moderation_payload ----------------------------------------- + + +class TestBuildPromptModerationPayload: + def test_payload_includes_tools_when_present(self): + inputs = { + "model": "gpt-4", + "structured_messages": [{"role": "user", "content": "hi"}], + "tools": [{"type": "function", "function": {"name": "fn"}}], + } + payload = RubrikLogger._build_prompt_moderation_payload(inputs, {}) + assert payload["tools"] == [{"type": "function", "function": {"name": "fn"}}] + + def test_payload_includes_user_when_present(self): + inputs = { + "structured_messages": [{"role": "user", "content": "hi"}], + } + request_data = {"user": "alice"} + payload = RubrikLogger._build_prompt_moderation_payload(inputs, request_data) + assert payload["user"] == "alice" + + def test_payload_uses_explicit_correlation_key(self): + inputs = {"structured_messages": [{"role": "user", "content": "hi"}]} + request_data = { + "correlation_key": "corr-123", + "litellm_call_id": "litellm-456", + } + payload = RubrikLogger._build_prompt_moderation_payload(inputs, request_data) + assert payload["correlation_key"] == "corr-123" + + def test_payload_falls_back_to_litellm_call_id(self): + """When correlation_key is absent, litellm_call_id is used.""" + inputs = {"structured_messages": [{"role": "user", "content": "hi"}]} + request_data = {"litellm_call_id": "litellm-789"} + payload = RubrikLogger._build_prompt_moderation_payload(inputs, request_data) + assert payload["correlation_key"] == "litellm-789" + + def test_payload_omits_optional_fields_when_absent(self): + inputs = {"structured_messages": [{"role": "user", "content": "hi"}]} + payload = RubrikLogger._build_prompt_moderation_payload(inputs, {}) + assert "tools" not in payload + assert "user" not in payload + assert "correlation_key" not in payload + + +# -- _extract_request_data tools preference ----------------------------------- + + +class TestExtractRequestDataToolsPreference: + def test_prefers_tools_from_request_data_over_optional_params(self): + """When 'tools' key exists in request_data, it wins over optional_params.""" + call_details = { + "messages": [{"role": "user", "content": "hi"}], + "model": "gpt-4", + "optional_params": { + "tools": [{"type": "function", "function": {"name": "from_optional"}}] + }, + } + request_data = { + "tools": [{"type": "function", "function": {"name": "from_request"}}] + } + result = RubrikLogger._extract_request_data(call_details, request_data) + assert result["tools"] == [ + {"type": "function", "function": {"name": "from_request"}} + ] + + def test_falls_back_to_optional_params_when_not_in_request_data(self): + call_details = { + "optional_params": { + "tools": [{"type": "function", "function": {"name": "from_optional"}}] + } + } + result = RubrikLogger._extract_request_data(call_details, {}) + assert result["tools"] == [ + {"type": "function", "function": {"name": "from_optional"}} + ] + + def test_explicit_empty_list_in_request_data_is_forwarded(self): + """An explicit empty tools list signals 'no tools' to the moderation service.""" + call_details = { + "optional_params": { + "tools": [{"type": "function", "function": {"name": "from_optional"}}] + } + } + request_data = {"tools": []} + result = RubrikLogger._extract_request_data(call_details, request_data) + assert result["tools"] == [] + + +# -- _extract_prompt_refusal -------------------------------------------------- + + +class TestExtractPromptRefusal: + def test_passthrough_response_returns_none(self): + """Empty dict (passthrough) → None.""" + assert RubrikLogger._extract_prompt_refusal({}) is None + + def test_no_choices_returns_none(self): + assert RubrikLogger._extract_prompt_refusal({"choices": []}) is None + + def test_block_response_returns_content(self): + service_response = { + "choices": [{"message": {"content": "Request blocked by Rubrik."}}] + } + result = RubrikLogger._extract_prompt_refusal(service_response) + assert result == "Request blocked by Rubrik." + + def test_empty_content_falls_back_to_default_message(self): + """When content is empty string or falsy, falls back to default refusal.""" + service_response = {"choices": [{"message": {"content": ""}}]} + result = RubrikLogger._extract_prompt_refusal(service_response) + assert result == "Request blocked by policy." + + def test_none_content_falls_back_to_default_message(self): + service_response = {"choices": [{"message": {"content": None}}]} + result = RubrikLogger._extract_prompt_refusal(service_response) + assert result == "Request blocked by policy." + + +# -- _prepend_system_prompt exception path ------------------------------------ + + +class TestPrependSystemPromptException: + def test_exception_during_unpack_is_caught_and_logged(self): + """When an exception is raised inside _prepend_system_prompt, it is swallowed.""" + + class ExplodingList(list): + def __iter__(self): + raise RuntimeError("iteration error!") + + payload = {"messages": ExplodingList()} + source = {"system": "You are an assistant."} + + # Must not raise + RubrikLogger._prepend_system_prompt(payload, source) + + +# -- _append_and_maybe_flush batch trigger ------------------------------------ + + +@pytest.mark.asyncio +class TestAppendAndMaybeFlush: + async def test_flush_triggered_when_queue_reaches_batch_size(self, handler): + """flush_queue is called when the queue length reaches batch_size.""" + handler.batch_size = 2 + handler.flush_queue = AsyncMock() + + await handler._append_and_maybe_flush({"msg": "a"}) + handler.flush_queue.assert_not_called() + + await handler._append_and_maybe_flush({"msg": "b"}) + handler.flush_queue.assert_called_once() + + async def test_no_flush_before_batch_size(self, handler): + handler.batch_size = 5 + handler.flush_queue = AsyncMock() + + for i in range(4): + await handler._append_and_maybe_flush({"msg": str(i)}) + + handler.flush_queue.assert_not_called() + + +# -- _enqueue_log_event exception handling ------------------------------------ + + +@pytest.mark.asyncio +class TestEnqueueLogEventExceptions: + async def test_exception_from_prepare_log_payload_is_caught(self, handler): + """Exceptions raised by _prepare_log_payload are caught and logged.""" + handler._prepare_log_payload = AsyncMock( + side_effect=RuntimeError("payload error") + ) + + # Must not raise + await handler._enqueue_log_event( + {"standard_logging_object": {"messages": [], "response": ""}}, "test" + ) + assert len(handler.log_queue) == 0 + + +# -- async_log_success_event skip when _rubrik_blocked ------------------------ + + +@pytest.mark.asyncio +class TestSuccessEventBlockedSkip: + async def test_skips_enqueue_when_rubrik_blocked_flag_set(self, handler): + """When kwargs['_rubrik_blocked'] is True, the event is not enqueued.""" + kwargs = { + "_rubrik_blocked": True, + "litellm_call_id": "blocked-call-123", + "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_post_call_failure_hook --------------------------------------------- + + +@pytest.mark.asyncio +class TestPostCallFailureHook: + async def test_non_modify_exception_returns_immediately(self, handler): + """Non-ModifyResponseException causes a no-op.""" + await handler.async_post_call_failure_hook( + request_data={"litellm_call_id": "test"}, + original_exception=ValueError("unrelated error"), + user_api_key_dict=None, + ) + assert len(handler.log_queue) == 0 + + async def test_modify_exception_without_stashed_logging_obj_emits_warning( + self, handler + ): + """ModifyResponseException with no _rubrik_logging_obj → warning, no enqueue.""" + request_data = {"litellm_call_id": "test-123", "model": "gpt-4"} + exc = ModifyResponseException( + message="blocked", + model="gpt-4", + request_data=request_data, + guardrail_name="rubrik", + ) + + await handler.async_post_call_failure_hook( + request_data=request_data, + original_exception=exc, + user_api_key_dict=None, + ) + assert len(handler.log_queue) == 0 + + async def test_modify_exception_with_valid_logging_obj_enqueues_payload( + self, handler + ): + """ModifyResponseException + stashed logging_obj → builds and enqueues.""" + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_call_id": "call-abc", + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "standard_logging_object": { + "id": "chatcmpl-original", + "model": "gpt-4", + "response": "original", + "messages": [{"role": "user", "content": "hi"}], + }, + "metadata": {}, + } + + request_data = {"_rubrik_logging_obj": logging_obj} + exc = ModifyResponseException( + message="blocked by policy", + model="gpt-4", + request_data=request_data, + guardrail_name="rubrik", + ) + + handler.batch_size = 10**6 # disable auto-flush + await handler.async_post_call_failure_hook( + request_data=request_data, + original_exception=exc, + user_api_key_dict=None, + ) + assert len(handler.log_queue) == 1 + assert "ModifyResponseException" in handler.log_queue[0]["response"] + + async def test_logging_obj_popped_from_request_data(self, handler): + """_rubrik_logging_obj must be popped from request_data so it is not + forwarded downstream.""" + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_call_id": "call-pop", + "model": "gpt-4", + "messages": [], + "standard_logging_object": { + "id": "chatcmpl-pop", + "model": "gpt-4", + "response": "text", + "messages": [], + }, + "metadata": {}, + } + + request_data = {"_rubrik_logging_obj": logging_obj} + exc = ModifyResponseException( + message="popped", + model="gpt-4", + request_data=request_data, + guardrail_name="rubrik", + ) + + handler.batch_size = 10**6 + await handler.async_post_call_failure_hook( + request_data=request_data, + original_exception=exc, + user_api_key_dict=None, + ) + assert "_rubrik_logging_obj" not in request_data + + async def test_build_and_enqueue_swallows_attribute_error_from_prepare_payload( + self, handler + ): + """When _prepare_block_failure_payload raises AttributeError/KeyError/TypeError, + the error is logged and the event is silently dropped (lines 806-812).""" + logging_obj = Mock() + # Make model_call_details.get() raise TypeError + logging_obj.model_call_details = None # .get() will raise AttributeError + + exc = ModifyResponseException( + message="blocked", + model="gpt-4", + request_data={}, + guardrail_name="rubrik", + ) + + # Must not raise + await handler._build_and_enqueue_block_event(logging_obj, exc, None) + assert len(handler.log_queue) == 0 + + async def test_build_and_enqueue_swallows_flush_exception(self, handler): + """When _append_and_maybe_flush raises, the error is logged (lines 816-817).""" + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_call_id": "call-flush-err", + "model": "gpt-4", + "messages": [], + "standard_logging_object": { + "id": "id-flush-err", + "model": "gpt-4", + "response": "text", + "messages": [], + }, + "metadata": {}, + } + + exc = ModifyResponseException( + message="blocked", + model="gpt-4", + request_data={}, + guardrail_name="rubrik", + ) + + handler._append_and_maybe_flush = AsyncMock( + side_effect=RuntimeError("flush failed") + ) + + # Must not raise + await handler._build_and_enqueue_block_event(logging_obj, exc, None) + + +# -- _prepare_block_failure_payload and _build_fallback_payload --------------- + + +class TestPrepareBlockFailurePayload: + def test_uses_standard_logging_object_when_present(self, handler): + """When standard_logging_object is on model_call_details, it is used as base.""" + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_call_id": "call-slo", + "model": "gpt-4", + "standard_logging_object": { + "id": "chatcmpl-original", + "model": "gpt-4", + "response": "original response", + "messages": [{"role": "user", "content": "hi"}], + }, + "metadata": {}, + } + exc = ModifyResponseException( + message="blocked", + model="gpt-4", + request_data={}, + guardrail_name="rubrik", + ) + + payload = handler._prepare_block_failure_payload(logging_obj, exc) + + assert "ModifyResponseException: blocked" in payload["response"] + assert payload["id"] == "call-slo" + + def test_uses_fallback_when_standard_logging_object_absent(self, handler): + """When standard_logging_object is absent, _build_fallback_payload is used.""" + from datetime import datetime + + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_call_id": "call-fallback", + "model": "claude-3", + "messages": [{"role": "user", "content": "question"}], + "optional_params": {"temperature": 0.5}, + "metadata": {"user_api_key_hash": "hash-abc"}, + "start_time": datetime(2024, 6, 1), + } + exc = ModifyResponseException( + message="prompt blocked", + model="claude-3", + request_data={}, + guardrail_name="rubrik", + ) + + payload = handler._prepare_block_failure_payload(logging_obj, exc) + + assert payload["id"] == "call-fallback" + assert payload["model"] == "claude-3" + assert payload["model_group"] == "claude-3" + assert "ModifyResponseException: prompt blocked" in payload["response"] + assert payload["metadata"]["user_api_key_hash"] == "hash-abc" + assert payload["status"] == "failure" + + def test_fallback_payload_without_start_time(self, handler): + """_build_fallback_payload handles missing start_time gracefully.""" + logging_obj = Mock() + logging_obj.model_call_details = { + "litellm_call_id": "call-notime", + "model": "gpt-4", + "messages": [], + "optional_params": {}, + "metadata": {}, + } + exc = ModifyResponseException( + message="blocked", + model="gpt-4", + request_data={}, + guardrail_name="rubrik", + ) + + payload = handler._prepare_block_failure_payload(logging_obj, exc) + assert payload["startTime"] is None + + +# -- async_send_batch empty queue and flush_queue edge cases ------------------ + + +@pytest.mark.asyncio +class TestQueueEdgeCases: + async def test_async_send_batch_returns_early_on_empty_queue(self, handler): + """async_send_batch is a no-op when the queue is empty.""" + handler.async_httpx_client = AsyncMock() + await handler.async_send_batch() + handler.async_httpx_client.post.assert_not_called() + + async def test_flush_queue_returns_early_when_flush_lock_is_none(self, handler): + """flush_queue is a no-op when flush_lock is None.""" + handler.flush_lock = None + handler.log_queue = [{"msg": "a"}] + handler.async_httpx_client = AsyncMock() + + await handler.flush_queue() + handler.async_httpx_client.post.assert_not_called() + + async def test_flush_queue_returns_early_when_queue_empty_inside_lock(self, handler): + """flush_queue acquires the lock then no-ops when the queue is empty.""" + handler.log_queue = [] + handler.async_httpx_client = AsyncMock() + + await handler.flush_queue() + handler.async_httpx_client.post.assert_not_called() + + +# -- _post_json non-dict response --------------------------------------------- + + +@pytest.mark.asyncio +class TestPostJson: + async def test_raises_type_error_for_list_response(self, handler): + """When the service returns a JSON array instead of a dict, TypeError is raised.""" + mock_client = AsyncMock() + mock_resp = Mock() + mock_resp.json.return_value = ["not", "a", "dict"] + mock_resp.raise_for_status = Mock() + mock_client.post = AsyncMock(return_value=mock_resp) + handler.moderation_client = mock_client + + with pytest.raises(TypeError, match="non-dict JSON"): + await handler._post_json( + handler.prompt_moderation_endpoint, {}, "Test service" + ) + + async def test_raises_type_error_for_string_response(self, handler): + """A bare string response also raises TypeError.""" + mock_client = AsyncMock() + mock_resp = Mock() + mock_resp.json.return_value = "blocked" + mock_resp.raise_for_status = Mock() + mock_client.post = AsyncMock(return_value=mock_resp) + handler.moderation_client = mock_client + + with pytest.raises(TypeError, match="non-dict JSON"): + await handler._post_json( + handler.response_moderation_endpoint, {}, "Test service" + )