diff --git a/.circleci/config.yml b/.circleci/config.yml index c3a34a97b3b..790cd6c7010 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -88,6 +88,36 @@ commands: rm -f /tmp/uv-install.sh echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$BASH_ENV" export PATH="$HOME/.local/bin:$PATH" + install_rust: + description: "Install pinned rustup (1.28.2) and Rust toolchain (1.97.1) with checksum verification. Adds ~/.cargo/bin to PATH. Run this before any `uv sync` or `uv build` of the workspace: the root package builds litellm-rust through maturin, and on an image without cargo maturin fetches an unpinned rustup and a floating toolchain by itself." + steps: + - run: + name: Install Rust (rustup 1.28.2, toolchain 1.97.1) + command: | + case "$(uname -m)" in + x86_64) + RUSTUP_TRIPLE=x86_64-unknown-linux-gnu + RUSTUP_SHA256=20a06e644b0d9bd2fbdbfd52d42540bdde820ea7df86e92e533c073da0cdd43c + ;; + aarch64) + RUSTUP_TRIPLE=aarch64-unknown-linux-gnu + RUSTUP_SHA256=e3853c5a252fca15252d07cb23a1bdd9377a8c6f3efa01531109281ae47f841c + ;; + *) + echo "install_rust: unsupported architecture $(uname -m)" >&2 + exit 1 + ;; + esac + curl -sSLf -o /tmp/rustup-init \ + "https://static.rust-lang.org/rustup/archive/1.28.2/${RUSTUP_TRIPLE}/rustup-init" + echo "${RUSTUP_SHA256} /tmp/rustup-init" | sha256sum -c - + chmod +x /tmp/rustup-init + /tmp/rustup-init -y --no-modify-path --profile minimal --default-toolchain 1.97.1 + rm -f /tmp/rustup-init + echo 'export PATH="$HOME/.cargo/bin:$PATH"' >> "$BASH_ENV" + export PATH="$HOME/.cargo/bin:$PATH" + rustc --version + cargo --version start_postgres: description: "Start a postgres-db container on port 5432 and wait until it accepts connections." parameters: @@ -163,6 +193,26 @@ commands: done echo "fake OpenAI endpoint did not become ready" >&2 exit 1 + start_cost_center_service: + description: "Start the stand-in cost center validation service (tests/store_model_in_db_tests/cost_center_service.py) on host port 9414 and wait until healthy. The proxy's team-metadata validator (team_metadata_validator_e2e.py, impl 'http') reaches it via TEAM_METADATA_VALIDATION_SERVICE_URL=http://host.docker.internal:9414/validate. Run after uv deps are synced." + steps: + - run: + name: Start cost center validation service + background: true + command: | + uv run --no-sync python tests/store_model_in_db_tests/cost_center_service.py --host 0.0.0.0 --port 9414 + - run: + name: Wait for cost center validation service + command: | + for i in $(seq 1 30); do + if curl -sf http://localhost:9414/health >/dev/null 2>&1; then + echo "cost center validation service is up" + exit 0 + fi + sleep 1 + done + echo "cost center validation service did not become ready" >&2 + exit 1 setup_litellm_enterprise_pip: steps: - run: @@ -178,6 +228,7 @@ commands: - checkout - setup_google_dns - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} @@ -292,6 +343,7 @@ jobs: - checkout - setup_google_dns - install_uv + - install_rust - run: name: Build the wheel environment: @@ -324,6 +376,7 @@ jobs: keys: - v1-uv-cache-{{ checksum "uv.lock" }} - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -397,6 +450,7 @@ jobs: keys: - v1-uv-cache-{{ checksum "uv.lock" }} - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -471,6 +525,7 @@ jobs: keys: - v1-uv-cache-{{ checksum "uv.lock" }} - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -522,6 +577,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -588,6 +644,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -628,6 +685,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -669,6 +727,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -702,6 +761,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} @@ -752,6 +812,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} @@ -803,6 +864,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -836,6 +898,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} @@ -882,6 +945,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -928,6 +992,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -970,6 +1035,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1016,6 +1082,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1063,6 +1130,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} @@ -1103,6 +1171,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1148,6 +1217,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1192,6 +1262,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1224,6 +1295,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1267,6 +1339,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1311,6 +1384,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1355,6 +1429,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1386,6 +1461,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1432,6 +1508,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1477,6 +1554,7 @@ jobs: keys: - v1-uv-cache-{{ checksum "uv.lock" }} - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1527,6 +1605,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1551,6 +1630,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1577,6 +1657,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1678,6 +1759,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1773,6 +1855,7 @@ jobs: at: ~/project - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1861,6 +1944,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1944,6 +2028,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -2076,6 +2161,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -2162,6 +2248,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -2258,12 +2345,14 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | uv sync --frozen --all-groups --all-extras --python 3.12 - start_postgres - start_fake_openai_endpoint + - start_cost_center_service - attach_workspace: at: ~/project - run: @@ -2283,11 +2372,13 @@ jobs: -e STORE_MODEL_IN_DB="True" \ -e LITELLM_MASTER_KEY="sk-1234" \ -e FAKE_OPENAI_API_BASE=http://host.docker.internal:8190 \ + -e TEAM_METADATA_VALIDATION_SERVICE_URL=http://host.docker.internal:9414/validate \ -e LITELLM_LICENSE=$LITELLM_LICENSE \ -e LITELLM_LOG=ERROR \ --add-host host.docker.internal:host-gateway \ --name my-app \ -v $(pwd)/litellm/proxy/example_config_yaml/store_model_db_config.yaml:/app/config.yaml \ + -v $(pwd)/litellm/proxy/example_config_yaml/team_metadata_validator_e2e.py:/app/team_metadata_validator_e2e.py \ litellm-docker-database:ci \ --config /app/config.yaml \ --port 4000 @@ -2333,6 +2424,7 @@ jobs: - setup_google_dns # Remove Docker CLI installation since it's already available in machine executor - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -2414,6 +2506,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -2553,6 +2646,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -2743,6 +2837,7 @@ jobs: category: client - setup_google_dns - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} @@ -2885,6 +2980,7 @@ jobs: category: client - setup_google_dns - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} 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/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 07f04136bd0..848af54cbcf 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1975,6 +1975,7 @@ "prompt_cache_min_tokens": 2048 }, "anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -2011,6 +2012,7 @@ "prompt_cache_min_tokens": 1024 }, "global.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -2047,6 +2049,7 @@ "prompt_cache_min_tokens": 1024 }, "us.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -2083,6 +2086,7 @@ "prompt_cache_min_tokens": 1024 }, "eu.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -2119,6 +2123,7 @@ "prompt_cache_min_tokens": 1024 }, "au.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -2155,6 +2160,7 @@ "prompt_cache_min_tokens": 1024 }, "jp.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3ccf3ea9952..268b08d0f2e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -543,6 +543,7 @@ class LiteLLMRoutes(enum.Enum): "/v2/team/list", "/organization/list", "/team/available", + "/team/metadata_schema", "/user/info", "/v2/user/info", "/model/info", @@ -609,6 +610,7 @@ class LiteLLMRoutes(enum.Enum): "/team/block", "/team/unblock", "/team/available", + "/team/metadata_schema", "/team/permissions_list", "/team/permissions_update", "/team/permissions_bulk_update", diff --git a/litellm/proxy/example_config_yaml/custom_team_metadata_validate.py b/litellm/proxy/example_config_yaml/custom_team_metadata_validate.py new file mode 100644 index 00000000000..1ca21bc3104 --- /dev/null +++ b/litellm/proxy/example_config_yaml/custom_team_metadata_validate.py @@ -0,0 +1,40 @@ +"""Example validator for `general_settings.custom_team_metadata_validate`. + +Wire it up in the proxy config: + +```yaml +general_settings: + custom_team_metadata_validate: custom_team_metadata_validate.validate_team_metadata + team_metadata_validation_timeout: 5 + team_metadata_validation_error_message: "Validation service unavailable, contact your admin." +``` + +Return `valid=False` with an `error_message` to reject the write with that +message (HTTP 400). Raise any exception (for example, when the upstream +validation service is unreachable) to fail closed with the generic +`team_metadata_validation_error_message` (HTTP 503). +""" + +from litellm.proxy.management_helpers.team_metadata_validation import ( + TeamMetadataValidationPayload, + TeamMetadataValidationResult, +) + +VALID_COST_CENTERS = frozenset({"CC-1001", "CC-1002", "CC-2001"}) + + +async def validate_team_metadata( + payload: TeamMetadataValidationPayload, +) -> TeamMetadataValidationResult: + cost_center = payload.metadata.get("cost_center") + if cost_center is None: + return TeamMetadataValidationResult( + valid=False, + error_message="Team metadata must include a cost_center. Contact the FinOps team.", + ) + if cost_center not in VALID_COST_CENTERS: + return TeamMetadataValidationResult( + valid=False, + error_message=f"Cost center {cost_center} is not recognized. Contact the FinOps team.", + ) + return TeamMetadataValidationResult(valid=True) diff --git a/litellm/proxy/example_config_yaml/store_model_db_config.yaml b/litellm/proxy/example_config_yaml/store_model_db_config.yaml index 5b77a53b4b7..3f19ab39a02 100644 --- a/litellm/proxy/example_config_yaml/store_model_db_config.yaml +++ b/litellm/proxy/example_config_yaml/store_model_db_config.yaml @@ -7,4 +7,7 @@ model_list: general_settings: store_model_in_db: true + custom_team_metadata_validate: team_metadata_validator_e2e.validate_team_metadata + team_metadata_validation_timeout: 5 + team_metadata_validation_error_message: "Cost center validation is unavailable right now; the team was not saved. Contact FinOps." diff --git a/litellm/proxy/example_config_yaml/team_metadata_validator_e2e.py b/litellm/proxy/example_config_yaml/team_metadata_validator_e2e.py new file mode 100644 index 00000000000..23390ffec1c --- /dev/null +++ b/litellm/proxy/example_config_yaml/team_metadata_validator_e2e.py @@ -0,0 +1,107 @@ +"""Dispatching team metadata validator for the store_model_in_db e2e suite. + +The suite runs one proxy with one config, so a single registered validator +dispatches to one of three independent implementations chosen per request via +the `_e2e_validator_impl` metadata key: + +- `allowlist`: requires `cost_center` and checks it against a static set +- `http`: POSTs the metadata to the cost center service at + `TEAM_METADATA_VALIDATION_SERVICE_URL`; transport errors raise (fail closed) +- `http_down`: like `http` but targets a closed port, proving the 503 path +- `immutable`: requires `cost_center` and forbids changing it once set + +A request whose metadata carries no `_e2e_validator_impl` key is accepted +untouched, so the rest of the suite's team operations are unaffected. An +unknown impl value raises, which the proxy converts to the fail-closed 503. +""" + +import os + +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.proxy.management_helpers.team_metadata_validation import ( + TeamMetadataValidationPayload, + TeamMetadataValidationResult, +) +from litellm.types.llms.custom_http import httpxSpecialProvider + +ALLOWED_COST_CENTERS = frozenset({"CC-1001", "CC-1002"}) +CLOSED_PORT_URL = "http://127.0.0.1:9/validate" + + +async def _validate_allowlist(payload: TeamMetadataValidationPayload) -> TeamMetadataValidationResult: + cost_center = payload.metadata.get("cost_center") + if cost_center is None: + return TeamMetadataValidationResult( + valid=False, + error_message="cost_center is required in team metadata. Contact the FinOps team.", + ) + if cost_center not in ALLOWED_COST_CENTERS: + return TeamMetadataValidationResult( + valid=False, + error_message=f"Cost center {cost_center} is not recognized. Contact the FinOps team.", + ) + return TeamMetadataValidationResult(valid=True) + + +async def _validate_via_http(payload: TeamMetadataValidationPayload, service_url: str) -> TeamMetadataValidationResult: + client = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + response = await client.post( + service_url, + json={ # mutable-ok: httpx serializes the request body from a plain dict + "operation": payload.operation, + "metadata": payload.metadata, + }, + timeout=2.0, + ) + body = response.json() + if body.get("ok") is True: + return TeamMetadataValidationResult(valid=True) + return TeamMetadataValidationResult( + valid=False, + error_message=body.get("reason", "Rejected by the cost center service."), + ) + + +class _ImmutableCostCenterValidator: + def __init__(self, immutable_key: str = "cost_center") -> None: + self.immutable_key = immutable_key + + async def __call__(self, payload: TeamMetadataValidationPayload) -> TeamMetadataValidationResult: + current = payload.metadata.get(self.immutable_key) + if current is None: + return TeamMetadataValidationResult( + valid=False, + error_message=f"{self.immutable_key} is required in team metadata. Contact the FinOps team.", + ) + if payload.operation == "update" and payload.existing_metadata is not None: + prior = payload.existing_metadata.get(self.immutable_key) + if prior is not None and prior != current: + return TeamMetadataValidationResult( + valid=False, + error_message=( + f"{self.immutable_key} is immutable once set " + f"(stored: {prior}, requested: {current}). Contact the FinOps team." + ), + ) + return TeamMetadataValidationResult(valid=True) + + +_IMMUTABLE_VALIDATOR = _ImmutableCostCenterValidator() + + +async def validate_team_metadata( + payload: TeamMetadataValidationPayload, +) -> TeamMetadataValidationResult: + impl = payload.metadata.get("_e2e_validator_impl") + if impl is None: + return TeamMetadataValidationResult(valid=True) + if impl == "allowlist": + return await _validate_allowlist(payload) + if impl == "http": + service_url = os.environ.get("TEAM_METADATA_VALIDATION_SERVICE_URL", "http://localhost:9414/validate") + return await _validate_via_http(payload, service_url) + if impl == "http_down": + return await _validate_via_http(payload, CLOSED_PORT_URL) + if impl == "immutable": + return await _IMMUTABLE_VALIDATOR(payload) + raise ValueError(f"unknown _e2e_validator_impl: {impl}") 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/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 50109b02189..475c5813cdd 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -61,6 +61,7 @@ from litellm.repositories.team_repository import TeamRepository from litellm.router import Router from litellm.router_utils.auto_router_model_naming import ( STRATEGY_ROUTER_PARAM_FIELDS, + validate_complexity_router_config_write, validate_strategy_router_model_write, ) from litellm.types.proxy.management_endpoints.model_management_endpoints import ( @@ -117,11 +118,20 @@ def _strategy_router_write_violation( An auto-router deployment's ``litellm_params.model`` (``auto_router/...``) is the discriminator the router loads it by; a write that mangles it makes the router drop the deployment silently under ``ignore_invalid_deployments``. - Only writes that supply ``litellm_params.model`` are judged, against the - merged (stored + incoming) params, so partial patches and restores of an - already-corrupted row stay legal. Returns the violation, or None. + Only writes that supply ``litellm_params.model`` are judged on the naming + contract, against the merged (stored + incoming) params, so partial patches + and restores of an already-corrupted row stay legal. A config is judged only + when the write carries one, for the same reason: a rename must not be held + hostage by a stored config it does not touch. Returns the violation, or None. """ - if incoming_params is None or incoming_params.model is None: + if incoming_params is None: + return None + config_violation = validate_complexity_router_config_write( + complexity_router_config=incoming_params.complexity_router_config + ) + if config_violation is not None: + return config_violation + if incoming_params.model is None: return None present_fields = frozenset( field diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 4dd86e5769d..bdd7e18a850 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -113,6 +113,10 @@ from litellm.proxy.management_helpers.object_permission_utils import ( from litellm.proxy.management_helpers.team_member_permission_checks import ( TeamMemberPermissionChecks, ) +from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_SCHEMA_REGISTRY, + validate_team_metadata_if_configured, +) from litellm.proxy.management_helpers.utils import ( add_new_member, management_endpoint_wrapper, @@ -147,6 +151,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( TeamListResponse, TeamMemberAddResult, TeamMemberInfoResponse, + TeamMetadataSchemaResponse, UpdateTeamMemberPermissionsRequest, ) @@ -1287,6 +1292,18 @@ async def new_team( _check_passthrough_routes_caller_permission(data, user_api_key_dict, entity="team") + if isinstance(data.metadata, dict): + TeamMemberBudgetHandler.strip_system_managed_metadata_keys(data.metadata) + + await validate_team_metadata_if_configured( + operation="create", + metadata=data.metadata, + existing_metadata=None, + team_id=data.team_id, + team_alias=data.team_alias, + user_api_key_dict=user_api_key_dict, + ) + ## ADD TO MODEL TABLE _model_id = None if data.model_aliases is not None and isinstance(data.model_aliases, dict): @@ -1301,9 +1318,6 @@ async def new_team( _model_id = model_dict.id - ## Create Team Member Budget Table - if isinstance(data.metadata, dict): - TeamMemberBudgetHandler.strip_system_managed_metadata_keys(data.metadata) data_json = data.json() ## Handle Object Permission - MCP, Vector Stores etc. @@ -1965,6 +1979,25 @@ async def update_team( if isinstance(updated_kv.get("metadata"), dict): TeamMemberBudgetHandler.strip_system_managed_metadata_keys(updated_kv["metadata"]) + if "metadata" in updated_kv: + stored_metadata = ( + { # mutable-ok: the validator payload's isinstance guard requires a plain dict + key: value + for key, value in existing_team_row.metadata.items() + if key not in TeamMemberBudgetHandler.SYSTEM_MANAGED_METADATA_KEYS + } + if isinstance(existing_team_row.metadata, dict) + else None + ) + await validate_team_metadata_if_configured( + operation="update", + metadata=updated_kv.get("metadata"), + existing_metadata=stored_metadata, + team_id=data.team_id, + team_alias=data.team_alias if data.team_alias is not None else existing_team_row.team_alias, + user_api_key_dict=user_api_key_dict, + ) + # Check budget_duration and budget_reset_at _set_budget_reset_at(data, updated_kv) @@ -4179,6 +4212,24 @@ async def unblock_team( return record +@router.get( + "/team/metadata_schema", + tags=["team management"], # mutable-ok: fastapi's decorator signature types tags as a list + dependencies=(Depends(user_api_key_auth),), + response_model=TeamMetadataSchemaResponse, +) +async def get_team_metadata_schema(): + """ + Get the team metadata fields declared in ``general_settings.team_metadata_schema``. + + The UI uses this to prepopulate the team metadata form with the declared + keys. Returns an empty ``fields`` list when no schema is configured. This + schema is advisory; server-side enforcement stays with + ``custom_team_metadata_validate``. + """ + return TeamMetadataSchemaResponse(fields=TEAM_METADATA_SCHEMA_REGISTRY.get()) + + @router.get("/team/available") async def list_available_teams( http_request: Request, diff --git a/litellm/proxy/management_helpers/team_metadata_validation.py b/litellm/proxy/management_helpers/team_metadata_validation.py new file mode 100644 index 00000000000..2b8ba05c602 --- /dev/null +++ b/litellm/proxy/management_helpers/team_metadata_validation.py @@ -0,0 +1,193 @@ +"""Custom validation of team metadata on team create/update. + +Operators point `general_settings.custom_team_metadata_validate` at an async +Python function (loaded via `get_instance_fn`, like `custom_key_generate`). +The function receives a `TeamMetadataValidationPayload` and returns a +`TeamMetadataValidationResult`. The proxy awaits it before committing a team +write and fails closed: a rejected value surfaces the function's own message +(HTTP 400), while any raised exception or timeout blocks the write with a +generic message (HTTP 503). +""" + +import asyncio +import inspect +from collections.abc import Awaitable, Mapping +from types import MappingProxyType +from typing import Literal, Protocol + +from fastapi import HTTPException, status +from pydantic import BaseModel, JsonValue, TypeAdapter + +from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth +from litellm.types.proxy.management_endpoints.team_endpoints import ( + TeamMetadataFieldSchema, +) + +DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS = 5.0 +DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE = ( + "Team metadata validation is currently unavailable, so the team was not saved. Contact your proxy admin." +) +DEFAULT_TEAM_METADATA_VALIDATION_REJECTED_MESSAGE = "Team metadata failed validation." + + +class TeamMetadataRequester(BaseModel): + user_id: str | None = None + user_email: str | None = None + user_role: str | None = None + + +class TeamMetadataValidationPayload(BaseModel): + operation: Literal["create", "update"] + metadata: Mapping[str, JsonValue] + existing_metadata: Mapping[str, JsonValue] | None = None + team_id: str | None = None + team_alias: str | None = None + requester: TeamMetadataRequester + + +class TeamMetadataValidationResult(BaseModel): + valid: bool + error_message: str | None = None + + +_EMPTY_METADATA: Mapping[str, JsonValue] = MappingProxyType({}) + + +class TeamMetadataValidator(Protocol): + def __call__(self, payload: TeamMetadataValidationPayload, /) -> Awaitable[TeamMetadataValidationResult]: ... + + +class TeamMetadataValidatorRegistry: + def __init__(self) -> None: + self._validator: TeamMetadataValidator | None = None + + def set(self, validator: TeamMetadataValidator | None) -> None: + self._validator = validator + + def get(self) -> TeamMetadataValidator | None: + return self._validator + + +TEAM_METADATA_VALIDATOR_REGISTRY = TeamMetadataValidatorRegistry() + +_TEAM_METADATA_SCHEMA_ADAPTER: TypeAdapter[tuple[TeamMetadataFieldSchema, ...]] = TypeAdapter( + tuple[TeamMetadataFieldSchema, ...] +) + + +def parse_team_metadata_schema(raw_schema: object) -> tuple[TeamMetadataFieldSchema, ...]: + """Parse ``general_settings.team_metadata_schema``; raises on a malformed schema so config load fails fast.""" + if raw_schema is None: + return () + fields = _TEAM_METADATA_SCHEMA_ADAPTER.validate_python(raw_schema) + keys = tuple(field.key for field in fields) + duplicate_keys = sorted(frozenset(key for key in keys if keys.count(key) > 1)) + if duplicate_keys: + raise ValueError(f"team_metadata_schema contains duplicate keys: {', '.join(duplicate_keys)}") + return fields + + +class TeamMetadataSchemaRegistry: + def __init__(self) -> None: + self._fields: tuple[TeamMetadataFieldSchema, ...] = () + + def set(self, fields: tuple[TeamMetadataFieldSchema, ...]) -> None: + self._fields = fields + + def get(self) -> tuple[TeamMetadataFieldSchema, ...]: + return self._fields + + +TEAM_METADATA_SCHEMA_REGISTRY = TeamMetadataSchemaRegistry() + + +async def run_team_metadata_validation( + validator: TeamMetadataValidator, + payload: TeamMetadataValidationPayload, + premium_user: bool, + timeout_seconds: float, + unavailable_message: str, +) -> None: + if premium_user is not True: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ # mutable-ok: HTTPException.detail has no immutable form + "error": f"custom_team_metadata_validate is an Enterprise feature. {CommonProxyErrors.not_premium_user.value}" + }, + ) + if not ( + inspect.iscoroutinefunction(validator) or inspect.iscoroutinefunction(getattr(validator, "__call__", None)) + ): + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={ # mutable-ok: HTTPException.detail has no immutable form + "error": "custom_team_metadata_validate must be an async function" + }, + ) + + try: + raw_result = await asyncio.wait_for(validator(payload), timeout=timeout_seconds) + result = TeamMetadataValidationResult.model_validate(raw_result) + except Exception: # noqa: BLE001 # fail closed: any validator failure must block the team write + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail={"error": unavailable_message}, # mutable-ok: HTTPException.detail has no immutable form + ) + + if not result.valid: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ # mutable-ok: HTTPException.detail has no immutable form + "error": result.error_message or DEFAULT_TEAM_METADATA_VALIDATION_REJECTED_MESSAGE + }, + ) + + +def _read_timeout_seconds(general_settings: Mapping[str, object]) -> float: + raw_timeout = general_settings.get("team_metadata_validation_timeout") + if isinstance(raw_timeout, (int, float)) and not isinstance(raw_timeout, bool) and raw_timeout > 0: + return float(raw_timeout) + return DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS + + +def _read_unavailable_message(general_settings: Mapping[str, object]) -> str: + raw_message = general_settings.get("team_metadata_validation_error_message") + if isinstance(raw_message, str) and raw_message.strip(): + return raw_message + return DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE + + +async def validate_team_metadata_if_configured( + operation: Literal["create", "update"], + metadata: Mapping[str, JsonValue] | None, + existing_metadata: Mapping[str, JsonValue] | None, + team_id: str | None, + team_alias: str | None, + user_api_key_dict: UserAPIKeyAuth, + registry: TeamMetadataValidatorRegistry = TEAM_METADATA_VALIDATOR_REGISTRY, +) -> None: + from litellm.proxy.proxy_server import general_settings, premium_user + + validator = registry.get() + if validator is None: + return + + payload = TeamMetadataValidationPayload( + operation=operation, + metadata=metadata if isinstance(metadata, dict) else _EMPTY_METADATA, + existing_metadata=existing_metadata if isinstance(existing_metadata, dict) else None, + team_id=team_id, + team_alias=team_alias, + requester=TeamMetadataRequester( + user_id=user_api_key_dict.user_id, + user_email=user_api_key_dict.user_email, + user_role=user_api_key_dict.user_role.value if user_api_key_dict.user_role is not None else None, + ), + ) + await run_team_metadata_validation( + validator=validator, + payload=payload, + premium_user=premium_user, + timeout_seconds=_read_timeout_seconds(general_settings), + unavailable_message=_read_unavailable_message(general_settings), + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ac45898ce0b..2ff39164d80 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -461,6 +461,11 @@ from litellm.proxy.management_helpers.audit_logs import ( create_audit_log_for_update, create_object_audit_log, ) +from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_SCHEMA_REGISTRY, + TEAM_METADATA_VALIDATOR_REGISTRY, + parse_team_metadata_schema, +) from litellm.proxy.memory.memory_endpoints import router as memory_router from litellm.proxy.middleware.billable_request_metrics_middleware import ( BillableRequestMetricsMiddleware, @@ -783,6 +788,8 @@ def cleanup_router_config_variables(): user_custom_auth_path = None user_custom_key_generate = None user_custom_key_update = None + TEAM_METADATA_VALIDATOR_REGISTRY.set(None) + TEAM_METADATA_SCHEMA_REGISTRY.set(()) user_custom_sso = None user_custom_ui_sso_sign_in_handler = None use_background_health_checks = None @@ -3499,6 +3506,7 @@ _DB_OVERLAY_REMOTE_MODULE_STR_FIELDS: dict[str, tuple[str, ...]] = { "custom_auth", "custom_key_generate", "custom_key_update", + "custom_team_metadata_validate", "custom_sso", "custom_ui_sso_sign_in_handler", ), @@ -4829,6 +4837,14 @@ class ProxyConfig: if custom_key_update is not None: user_custom_key_update = get_instance_fn(value=custom_key_update, config_file_path=config_file_path) + custom_team_metadata_validate = general_settings.get("custom_team_metadata_validate", None) + TEAM_METADATA_VALIDATOR_REGISTRY.set( + get_instance_fn(value=custom_team_metadata_validate, config_file_path=config_file_path) + if custom_team_metadata_validate is not None + else None + ) + TEAM_METADATA_SCHEMA_REGISTRY.set(parse_team_metadata_schema(general_settings.get("team_metadata_schema"))) + custom_sso = general_settings.get("custom_sso", None) if custom_sso is not None: user_custom_sso = get_instance_fn(value=custom_sso, config_file_path=config_file_path) diff --git a/litellm/router_utils/auto_router_model_naming.py b/litellm/router_utils/auto_router_model_naming.py index b89e1c154f0..51a3f127b29 100644 --- a/litellm/router_utils/auto_router_model_naming.py +++ b/litellm/router_utils/auto_router_model_naming.py @@ -62,12 +62,42 @@ def classify_strategy_router_model(model: str) -> StrategyRouterKind | None: return "semantic" +def validate_complexity_router_config_write(complexity_router_config: Mapping[str, object] | None) -> str | None: + """Reject a complexity config the router would refuse to build a deployment from. + + Parsed with the router's own ``ComplexityRouterConfig`` rather than a copy of + its rules, so the boundary rejects exactly what the load would. Plugins are + resolved from dotted paths only on the config.yaml path, so a written config + reaches this function in the same shape the load hands to the same model. + Judged on the config alone: a patch may write one without naming a model, and + the stored model is encrypted at rest, so it cannot be classified here. + """ + from pydantic import ValidationError + + from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig + + if complexity_router_config is None: + return None + try: + _ = ComplexityRouterConfig.model_validate(complexity_router_config) + except ValidationError as exc: + first = exc.errors()[0] + location = ".".join(str(part) for part in first.get("loc", ())) or "complexity_router_config" + return ( + f"complexity_router_config is invalid at {location}: {first.get('msg', 'invalid value')}. " + "The router would drop this deployment at load time, so the write is rejected instead." + ) + return None + + def validate_strategy_router_model_write(model: str, present_fields: frozenset[str]) -> str | None: """Check that writing ``model`` leaves a deployment the router can load. ``present_fields`` is the set of strategy-router param fields that are non-None on the deployment after the write (stored fields merged with the incoming ones). Returns a human-readable violation, or None when coherent. + A config's contents are ``validate_complexity_router_config_write``'s to + judge, since a write may carry one without naming a model at all. """ kind = classify_strategy_router_model(model) if kind is None: diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index d9f8229dbed..21de28a9b29 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -301,7 +301,7 @@ class BedrockToolSpec(dict): "name": name, "description": description, } - if supports_strict_tools and strict is not None: + if supports_strict_tools and strict: tool_spec["strict"] = strict super().__init__(toolSpec=tool_spec) diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index 2d9387da956..f3d0529a7ea 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -1,13 +1,12 @@ from typing import Any, Dict, List, Literal, Optional, Union -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict, Field from litellm.proxy._types import ( KeyManagementRoutes, LiteLLM_DeletedTeamTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, - LiteLLM_UserTable, Member, ) @@ -125,3 +124,22 @@ class TeamMemberInfoResponse(LiteLLM_TeamMembership): role: Optional[str] = None user_email: Optional[str] = None team_alias: Optional[str] = None + + +class TeamMetadataFieldSchema(BaseModel): + """One declared team metadata field from ``general_settings.team_metadata_schema``. + + Advisory only: the UI uses it to prepopulate the team metadata form. + Enforcement stays with ``custom_team_metadata_validate``. + """ + + model_config = ConfigDict(extra="forbid") + + key: str = Field(min_length=1) + label: Optional[str] = None + + +class TeamMetadataSchemaResponse(BaseModel): + """Response for GET /team/metadata_schema; ``fields`` is empty when no schema is configured.""" + + fields: tuple[TeamMetadataFieldSchema, ...] diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 346f613ea3e..56e0391f419 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1975,6 +1975,7 @@ "prompt_cache_min_tokens": 2048 }, "anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -2011,6 +2012,7 @@ "prompt_cache_min_tokens": 1024 }, "global.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_1hr": 4e-06, "cache_read_input_token_cost": 2e-07, @@ -2047,6 +2049,7 @@ "prompt_cache_min_tokens": 1024 }, "us.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -2083,6 +2086,7 @@ "prompt_cache_min_tokens": 1024 }, "eu.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -2119,6 +2123,7 @@ "prompt_cache_min_tokens": 1024 }, "au.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, @@ -2155,6 +2160,7 @@ "prompt_cache_min_tokens": 1024 }, "jp.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_1hr": 4.4e-06, "cache_read_input_token_cost": 2.2e-07, diff --git a/tests/store_model_in_db_tests/cost_center_service.py b/tests/store_model_in_db_tests/cost_center_service.py new file mode 100644 index 00000000000..c2820dfbe7e --- /dev/null +++ b/tests/store_model_in_db_tests/cost_center_service.py @@ -0,0 +1,54 @@ +"""Stand-in cost center validation service for the team metadata e2e tests. + +Accepts POST /validate with {"operation": ..., "metadata": {...}} and answers +{"ok": true} or {"ok": false, "reason": ...} based on a static allowlist. +GET /health answers 200 for the CI wait loop. +""" + +import argparse +import json +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +ALLOWED_COST_CENTERS = {"CC-1001", "CC-1002"} + + +class CostCenterHandler(BaseHTTPRequestHandler): + def _respond(self, status: int, body: dict) -> None: + payload = json.dumps(body).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def do_GET(self): + if self.path == "/health": + self._respond(200, {"status": "healthy"}) + return + self._respond(404, {"error": "not found"}) + + def do_POST(self): + if self.path != "/validate": + self._respond(404, {"error": "not found"}) + return + length = int(self.headers.get("Content-Length", 0)) + body = json.loads(self.rfile.read(length) or b"{}") + cost_center = (body.get("metadata") or {}).get("cost_center") + if cost_center is None: + self._respond(200, {"ok": False, "reason": "cost_center missing per cost center service"}) + elif cost_center not in ALLOWED_COST_CENTERS: + self._respond(200, {"ok": False, "reason": f"cost center {cost_center} rejected by cost center service"}) + else: + self._respond(200, {"ok": True}) + + def log_message(self, format, *args): + pass + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--host", default="0.0.0.0") + parser.add_argument("--port", type=int, default=9414) + args = parser.parse_args() + print(f"cost center service listening on {args.host}:{args.port}") + ThreadingHTTPServer((args.host, args.port), CostCenterHandler).serve_forever() diff --git a/tests/store_model_in_db_tests/test_team_metadata_validation_e2e.py b/tests/store_model_in_db_tests/test_team_metadata_validation_e2e.py new file mode 100644 index 00000000000..050b59ef426 --- /dev/null +++ b/tests/store_model_in_db_tests/test_team_metadata_validation_e2e.py @@ -0,0 +1,171 @@ +"""E2E matrix for custom team metadata validation against the DB-backed proxy. + +The proxy (see store_model_db_config.yaml) registers +team_metadata_validator_e2e.validate_team_metadata, which dispatches per +request to one of three independent implementations via the +`_e2e_validator_impl` metadata key: a static allowlist function, an +HTTP-backed function calling the cost center service started by CI, and an +immutability-enforcing class instance. Metadata without the dispatch key is +accepted untouched, so the rest of this suite is unaffected. +""" + +import os +import uuid + +import httpx +import pytest + +PROXY_BASE_URL = os.getenv("PROXY_BASE_URL", "http://localhost:4000") +MASTER_KEY = os.getenv("LITELLM_MASTER_KEY", "sk-1234") +HEADERS = {"Authorization": f"Bearer {MASTER_KEY}", "Content-Type": "application/json"} + +UNAVAILABLE_MESSAGE = "Cost center validation is unavailable right now; the team was not saved. Contact FinOps." + +IMPLS = ["allowlist", "http", "immutable"] + +REQUIRED_MESSAGES = { + "allowlist": "cost_center is required in team metadata", + "http": "cost_center missing per cost center service", + "immutable": "cost_center is required in team metadata", +} +UNKNOWN_MESSAGES = { + "allowlist": "is not recognized", + "http": "rejected by cost center service", +} + + +def _meta(impl, **fields): + return {"_e2e_validator_impl": impl, **fields} + + +def _create_team(metadata, team_id=None): + body = {"team_alias": f"meta-validate-{uuid.uuid4().hex[:8]}"} + if team_id is not None: + body["team_id"] = team_id + if metadata is not None: + body["metadata"] = metadata + return httpx.post(f"{PROXY_BASE_URL}/team/new", headers=HEADERS, json=body, timeout=30) + + +def _patch_team(team_id, body): + return httpx.patch(f"{PROXY_BASE_URL}/team/{team_id}", headers=HEADERS, json=body, timeout=30) + + +def _post_update(team_id, body): + return httpx.post(f"{PROXY_BASE_URL}/team/update", headers=HEADERS, json={"team_id": team_id, **body}, timeout=30) + + +def _team_info(team_id): + return httpx.get(f"{PROXY_BASE_URL}/team/info", headers=HEADERS, params={"team_id": team_id}, timeout=30) + + +def _delete_team(team_id): + httpx.post(f"{PROXY_BASE_URL}/team/delete", headers=HEADERS, json={"team_ids": [team_id]}, timeout=30) + + +@pytest.fixture +def team_with_cost_center(request): + impl = request.param + team_id = f"meta-validate-{impl}-{uuid.uuid4().hex[:8]}" + response = _create_team(metadata=_meta(impl, cost_center="CC-1001"), team_id=team_id) + assert response.status_code == 200, response.text + yield impl, team_id + _delete_team(team_id) + + +@pytest.mark.parametrize("impl", IMPLS) +def test_create_with_valid_cost_center_succeeds(impl): + response = _create_team(metadata=_meta(impl, cost_center="CC-1001")) + assert response.status_code == 200, response.text + team_id = response.json()["team_id"] + try: + assert response.json()["metadata"]["cost_center"] == "CC-1001" + finally: + _delete_team(team_id) + + +@pytest.mark.parametrize("impl", IMPLS) +def test_create_without_cost_center_is_rejected(impl): + team_id = f"meta-validate-reject-{impl}-{uuid.uuid4().hex[:8]}" + response = _create_team(metadata=_meta(impl), team_id=team_id) + assert response.status_code == 400, response.text + assert REQUIRED_MESSAGES[impl] in response.text + info = _team_info(team_id) + assert info.status_code == 404, "rejected create must not leave a team row behind" + + +@pytest.mark.parametrize("impl", IMPLS) +def test_create_with_unknown_cost_center(impl): + response = _create_team(metadata=_meta(impl, cost_center="CC-9999")) + if impl == "immutable": + assert response.status_code == 200, response.text + _delete_team(response.json()["team_id"]) + return + assert response.status_code == 400, response.text + assert UNKNOWN_MESSAGES[impl] in response.text + + +@pytest.mark.parametrize("team_with_cost_center", IMPLS, indirect=True) +def test_patch_changing_cost_center(team_with_cost_center): + impl, team_id = team_with_cost_center + response = _patch_team(team_id, {"metadata": {"cost_center": "CC-1002"}}) + if impl == "immutable": + assert response.status_code == 400, response.text + assert "immutable once set" in response.text + info = _team_info(team_id).json()["team_info"]["metadata"] + assert info["cost_center"] == "CC-1001", "blocked update must leave stored metadata intact" + return + assert response.status_code == 200, response.text + assert response.json()["metadata"]["cost_center"] == "CC-1002" + + +@pytest.mark.parametrize("team_with_cost_center", IMPLS, indirect=True) +def test_patch_unrelated_key_validates_merged_result(team_with_cost_center): + impl, team_id = team_with_cost_center + response = _patch_team(team_id, {"metadata": {"team_notes": "hello"}}) + assert response.status_code == 200, response.text + merged = response.json()["metadata"] + assert merged["cost_center"] == "CC-1001" + assert merged["team_notes"] == "hello" + + +@pytest.mark.parametrize("team_with_cost_center", IMPLS, indirect=True) +def test_patch_null_deleting_cost_center_is_rejected(team_with_cost_center): + impl, team_id = team_with_cost_center + response = _patch_team(team_id, {"metadata": {"cost_center": None}}) + assert response.status_code == 400, response.text + assert REQUIRED_MESSAGES[impl] in response.text + info = _team_info(team_id).json()["team_info"]["metadata"] + assert info["cost_center"] == "CC-1001" + + +@pytest.mark.parametrize("team_with_cost_center", IMPLS, indirect=True) +def test_post_update_dropping_cost_center_is_rejected(team_with_cost_center): + impl, team_id = team_with_cost_center + response = _post_update(team_id, {"metadata": _meta(impl, team_notes="only-notes")}) + assert response.status_code == 400, response.text + assert REQUIRED_MESSAGES[impl] in response.text + + +@pytest.mark.parametrize("team_with_cost_center", IMPLS, indirect=True) +def test_update_without_metadata_skips_validation(team_with_cost_center): + impl, team_id = team_with_cost_center + response = _post_update(team_id, {"tpm_limit": 55}) + assert response.status_code == 200, response.text + + +def test_http_service_outage_fails_closed_with_configured_message(): + response = _create_team(metadata=_meta("http_down", cost_center="CC-1001")) + assert response.status_code == 503, response.text + assert UNAVAILABLE_MESSAGE in response.text + + +def test_metadata_without_dispatch_key_is_untouched(): + response = _create_team(metadata={"any_key": "any_value"}) + assert response.status_code == 200, response.text + team_id = response.json()["team_id"] + try: + update = _post_update(team_id, {"metadata": {"any_key": "changed"}}) + assert update.status_code == 200, update.text + finally: + _delete_team(team_id) 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" + ) diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py index 791982fc3dc..370ec4b6f60 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_bedrock_converse_strict_tools_opus_47_48.py @@ -4,6 +4,11 @@ Bedrock Converse routes Claude Opus 4.7/4.8 and Claude Sonnet 4 through an Anthropic-compatible validator that rejects ``toolSpec.strict`` even though Anthropic's native API accepts ``strict`` as a top-level tool field. See BerriAI/litellm#31582. + +That per-model gate only covers models whose cost-map entry carries the flag, so a +``strict: false`` that litellm itself synthesized still broke unflagged models. Since +``strict: false`` is the Chat Completions default, it is now dropped for every model +rather than forwarded as a no-op the provider can reject. See BerriAI/litellm#33193. """ import pytest @@ -31,6 +36,16 @@ _STRICT_TOOL = [ } ] +_NON_STRICT_TOOL = [ + { + "type": "function", + "function": { + **_STRICT_TOOL[0]["function"], + "strict": False, + }, + } +] + @pytest.mark.parametrize( "model_id", @@ -48,12 +63,17 @@ _STRICT_TOOL = [ "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0", "bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0", "bedrock/apac.anthropic.claude-sonnet-4-20250514-v1:0", + # Sonnet 5 rejects it too, verified live against Bedrock in us-east-1 + "anthropic.claude-sonnet-5", + "bedrock/us.anthropic.claude-sonnet-5", + "bedrock/eu.anthropic.claude-sonnet-5", + "bedrock/jp.anthropic.claude-sonnet-5", ], ) def test_bedrock_tools_pt_strict_dropped_for_strict_unsupported_models( model_id: str, ) -> None: - """Opus 4.7/4.8 and Sonnet 4 reject toolSpec.strict and additionalProperties.""" + """Opus 4.7/4.8, Sonnet 4 and Sonnet 5 reject toolSpec.strict and additionalProperties.""" result = _bedrock_tools_pt(_STRICT_TOOL, model=model_id) tool_spec = result[0]["toolSpec"] assert ( @@ -81,6 +101,55 @@ def test_bedrock_tools_pt_strict_kept_for_other_anthropic(model_id: str) -> None ), f"strict missing for {model_id}: {result[0]['toolSpec']}" +@pytest.mark.parametrize( + "model_id", + [ + "bedrock/us.anthropic.claude-sonnet-5", + "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "anthropic.claude-sonnet-4-5-20250929-v1:0", + "bedrock/us.anthropic.claude-sonnet-4-6", + "bedrock/us.anthropic.claude-opus-4-8", + ], +) +def test_bedrock_tools_pt_falsy_strict_always_dropped(model_id: str) -> None: + """``strict: false`` is the Chat Completions default, so forwarding it says nothing + the provider does not already assume. Bedrock Converse rejects the key's presence + for a growing set of Claude models, so it is dropped for every model, including the + ones whose cost-map entry still allows ``strict: true`` through.""" + result = _bedrock_tools_pt(_NON_STRICT_TOOL, model=model_id) + tool_spec = result[0]["toolSpec"] + assert ( + "strict" not in tool_spec + ), f"no-op strict: false leaked into toolSpec for {model_id}: {tool_spec}" + + +def test_responses_bridge_function_tool_does_not_reach_bedrock_with_strict() -> None: + """The Responses-to-Chat-Completions bridge stamps ``strict: false`` onto every + function tool even when the caller never sent one, which is how Codex CLI requests + acquired the key. Assert the fabricated value does not survive to toolSpec.""" + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + responses_tool = { + "type": "function", + "name": "get_weather", + "description": "Get the weather for a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + } + chat_tools, _ = ( + LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + [responses_tool] + ) + ) + result = _bedrock_tools_pt(chat_tools, model="bedrock/us.anthropic.claude-sonnet-5") + assert "strict" not in result[0]["toolSpec"] + + @pytest.mark.parametrize( "model_id", [ @@ -129,6 +198,15 @@ def test_bedrock_converse_supports_strict_tools_helper() -> None: ) is False ) + assert bedrock_converse_supports_strict_tools("anthropic.claude-sonnet-5") is False + assert ( + bedrock_converse_supports_strict_tools("bedrock/us.anthropic.claude-sonnet-5") + is False + ) + assert ( + bedrock_converse_supports_strict_tools("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0") + is True + ) @pytest.mark.parametrize( @@ -143,6 +221,12 @@ def test_bedrock_converse_supports_strict_tools_helper() -> None: "us.anthropic.claude-sonnet-4-20250514-v1:0", "eu.anthropic.claude-sonnet-4-20250514-v1:0", "apac.anthropic.claude-sonnet-4-20250514-v1:0", + "anthropic.claude-sonnet-5", + "global.anthropic.claude-sonnet-5", + "us.anthropic.claude-sonnet-5", + "eu.anthropic.claude-sonnet-5", + "au.anthropic.claude-sonnet-5", + "jp.anthropic.claude-sonnet-5", ], ) def test_strict_tools_flag_set_in_model_cost_map(cost_map_key: str) -> None: diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 8dbf38555b3..95405a3b016 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -3477,6 +3477,140 @@ class TestStrategyRouterWriteValidation: is None ) + def test_create_with_empty_keyword_rule_rejected(self): + """LIT-5133: the router refuses to build a rule with no keyword, but only at load time. + Without this the row is written, dropped on reload, and the caller gets a 500 plus a + deployment that can never come back.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _strategy_router_write_violation, + ) + + violation = _strategy_router_write_violation( + incoming_params=LiteLLM_Params( + model="auto_router/complexity_router", + complexity_router_default_model="gpt-4o-mini", + complexity_router_config={ + "tiers": {"SIMPLE": ["gpt-4o-mini"]}, + "keyword_tier_rules": [{"keywords": [], "tier": "COMPLEX"}], + }, + ), + existing_params=None, + ) + assert violation is not None + assert "complexity_router_config is invalid" in violation + assert "keyword_tier_rules" in violation + + def test_patch_that_only_renames_does_not_judge_the_stored_config(self): + """Only a config the write actually carries is judged. A row stored before this validation + existed is already unloadable, and holding its rename hostage would break the restore path + this function documents; the repair is a write that supplies a good config.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _strategy_router_write_violation, + ) + from litellm.types.router import updateLiteLLMParams + + stored_bad = LiteLLM_Params( + model="auto_router/complexity_router", + complexity_router_config={ + "tiers": {"SIMPLE": ["gpt-4o-mini"]}, + "keyword_tier_rules": [{"keywords": [" "], "tier": "COMPLEX"}], + }, + ) + assert ( + _strategy_router_write_violation( + incoming_params=updateLiteLLMParams(model="auto_router/complexity_router"), + existing_params=stored_bad, + ) + is None + ) + + def test_incoming_config_replaces_stored_rather_than_merging(self): + """The field is written wholesale, so a good incoming config must clear a bad stored one.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _strategy_router_write_violation, + ) + from litellm.types.router import updateLiteLLMParams + + stored_bad = LiteLLM_Params( + model="auto_router/complexity_router", + complexity_router_config={ + "tiers": {"SIMPLE": ["gpt-4o-mini"]}, + "keyword_tier_rules": [{"keywords": [], "tier": "COMPLEX"}], + }, + ) + assert ( + _strategy_router_write_violation( + incoming_params=updateLiteLLMParams( + model="auto_router/complexity_router", + complexity_router_config={ + "tiers": {"SIMPLE": ["gpt-4o-mini"]}, + "keyword_tier_rules": [{"keywords": ["invoice"], "tier": "COMPLEX"}], + }, + ), + existing_params=stored_bad, + ) + is None + ) + + def test_config_only_patch_is_judged_without_a_model_in_the_payload(self): + """A patch may carry a config and no model, which is what a caller updating only the + routing rules sends. That path skipped the naming contract, so it has to be judged on the + config alone against the stored model, or it overwrites a working router with one that + cannot load and takes it out of service.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _strategy_router_write_violation, + ) + from litellm.types.router import updateLiteLLMParams + + violation = _strategy_router_write_violation( + incoming_params=updateLiteLLMParams( + complexity_router_config={ + "tiers": {"SIMPLE": ["gpt-4o-mini"]}, + "keyword_tier_rules": [{"keywords": [], "tier": "COMPLEX"}], + } + ), + existing_params=self._stored_complexity_params(), + ) + assert violation is not None + assert "complexity_router_config is invalid" in violation + + def test_config_only_patch_with_a_loadable_config_is_allowed(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _strategy_router_write_violation, + ) + from litellm.types.router import updateLiteLLMParams + + assert ( + _strategy_router_write_violation( + incoming_params=updateLiteLLMParams( + complexity_router_config={ + "tiers": {"SIMPLE": ["gpt-4o-mini"]}, + "keyword_tier_rules": [{"keywords": ["invoice"], "tier": "COMPLEX"}], + } + ), + existing_params=self._stored_complexity_params(), + ) + is None + ) + + def test_config_only_patch_is_judged_on_the_config_alone(self): + """The stored model is encrypted at rest, so a patch that names no model cannot be + classified from the row. An unloadable config is rejected on its own merits instead, + which is also the only reading that closes the path regardless of what is stored.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _strategy_router_write_violation, + ) + from litellm.types.router import updateLiteLLMParams + + violation = _strategy_router_write_violation( + incoming_params=updateLiteLLMParams( + complexity_router_config={"keyword_tier_rules": [{"keywords": [], "tier": "COMPLEX"}]} + ), + existing_params=LiteLLM_Params(model="c2VjcmV0-encrypted-at-rest"), + ) + assert violation is not None + assert "complexity_router_config is invalid" in violation + def test_create_semantic_router_missing_embedding_rejected(self): from litellm.proxy.management_endpoints.model_management_endpoints import ( _strategy_router_write_violation, diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 0de2c1ac71d..fc018a5333e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -9863,11 +9863,14 @@ async def _drive_team_write( raw_body=None, user=None, find_returns_none=False, + mock_sink=None, ): """Drive POST ``update_team`` or PATCH ``patch_team`` against a mocked team. Returns ``(endpoint_result, update_mock)``; propagates whatever the endpoint raises. Inspect ``update_mock.call_args.kwargs["data"]`` for the DB write. + Pass a dict as ``mock_sink`` to receive the update mock even when the + endpoint raises. """ from unittest.mock import AsyncMock, MagicMock, Mock from unittest.mock import patch as _patch @@ -9913,6 +9916,8 @@ async def _drive_team_write( return_value=LiteLLM_TeamTable(team_id=_PATCH_TEAM_ID, team_alias="t") ) pc.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + if mock_sink is not None: + mock_sink["update"] = pc.db.litellm_teamtable.update req = Mock(spec=Request) if kind == "post": @@ -10175,6 +10180,249 @@ async def test_patch_returns_full_team_object_not_wrapper(): assert result.team_id == _PATCH_TEAM_ID +# --------------------------------------------------------------------------- +# custom_team_metadata_validate wiring: the configured validator must gate +# every team write path (POST /team/new, POST /team/update, PATCH /team/{id}) +# and must see the metadata that will actually be written. +# --------------------------------------------------------------------------- + +from contextlib import contextmanager + +from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_VALIDATOR_REGISTRY, + TeamMetadataValidationResult, +) + + +@contextmanager +def _configured_team_metadata_validator(validator): + TEAM_METADATA_VALIDATOR_REGISTRY.set(validator) + try: + with ( + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + yield + finally: + TEAM_METADATA_VALIDATOR_REGISTRY.set(None) + + +def _recording_validator(recorded, valid=True, error_message=None): + async def validator(payload): + recorded.append(payload) + return TeamMetadataValidationResult(valid=valid, error_message=error_message) + + return validator + + +@pytest.mark.asyncio +async def test_update_validator_sees_replacement_on_post_and_merged_on_patch(): + """POST hands the validator the wholesale replacement; PATCH hands it the + RFC 7386 merged result including preserved keys.""" + existing = {"cost_center": "OLD", "keep": 1} + body = {"metadata": {"cost_center": "NEW"}} + + recorded_post = [] + with _configured_team_metadata_validator(_recording_validator(recorded_post)): + await _drive_team_write("post", existing_metadata=existing, payload=body) + + recorded_patch = [] + with _configured_team_metadata_validator(_recording_validator(recorded_patch)): + await _drive_team_write("patch", existing_metadata=existing, payload=body) + + assert len(recorded_post) == 1 + assert recorded_post[0].operation == "update" + assert recorded_post[0].metadata == {"cost_center": "NEW"} + assert recorded_post[0].existing_metadata == existing + + assert len(recorded_patch) == 1 + assert recorded_patch[0].operation == "update" + assert recorded_patch[0].metadata == {"cost_center": "NEW", "keep": 1} + assert recorded_patch[0].existing_metadata == existing + + +@pytest.mark.asyncio +async def test_patch_null_delete_removes_key_from_validated_metadata(): + """Deleting a key via PATCH null must be visible to the validator as the + key's absence in the resulting metadata, so a required key cannot be + silently dropped.""" + recorded = [] + with _configured_team_metadata_validator(_recording_validator(recorded)): + await _drive_team_write( + "patch", + existing_metadata={"cost_center": "OLD", "keep": 1}, + payload={"metadata": {"cost_center": None}}, + ) + + assert len(recorded) == 1 + assert recorded[0].metadata == {"keep": 1} + assert recorded[0].existing_metadata == {"cost_center": "OLD", "keep": 1} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["post", "patch"]) +async def test_update_without_metadata_skips_validator(kind): + recorded = [] + with _configured_team_metadata_validator(_recording_validator(recorded)): + await _drive_team_write(kind, existing_metadata={"k": "v"}, payload={"tpm_limit": 5}) + + assert recorded == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["post", "patch"]) +async def test_update_validator_rejection_blocks_db_write(kind): + recorded = [] + validator = _recording_validator(recorded, valid=False, error_message="cost center rejected, contact FinOps") + sink = {} + + with _configured_team_metadata_validator(validator): + with pytest.raises(ProxyException) as exc_info: + await _drive_team_write( + kind, + existing_metadata={"cost_center": "OLD"}, + payload={"metadata": {"cost_center": "BAD"}}, + mock_sink=sink, + ) + + assert str(exc_info.value.code) == "400" + assert "cost center rejected, contact FinOps" in str(exc_info.value.message) + assert len(recorded) == 1 + sink["update"].assert_not_awaited() + + +@pytest.mark.asyncio +async def test_new_team_validator_runs_without_metadata_and_rejection_blocks_create(): + """Create always validates, even when the request carries no metadata, so a + required-key policy can reject a team created without one.""" + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + recorded = [] + validator = _recording_validator(recorded, valid=False, error_message="cost_center is required") + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server._license_check") as mock_license, + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + _configured_team_metadata_validator(validator), + ): + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_prisma.db.litellm_teamtable.create = AsyncMock() + mock_license.is_team_count_over_limit.return_value = False + + with pytest.raises(ProxyException) as exc_info: + await new_team( + data=NewTeamRequest(team_alias="no-metadata-team"), + http_request=MagicMock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1"), + ) + + assert str(exc_info.value.code) == "400" + assert "cost_center is required" in str(exc_info.value.message) + assert len(recorded) == 1 + assert recorded[0].operation == "create" + assert recorded[0].metadata == {} + assert recorded[0].existing_metadata is None + mock_prisma.db.litellm_teamtable.create.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_new_team_validator_accept_proceeds_to_create(mock_db_client, mock_admin_auth): + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + mock_db_client.jsonify_team_object = lambda db_data: db_data + mock_db_client.get_data = AsyncMock(return_value=None) + mock_db_client.update_data = AsyncMock(return_value=MagicMock()) + mock_db_client.db = MagicMock() + + team_create_result = MagicMock(team_id="team-accept-1") + team_create_result.model_dump.return_value = {"team_id": "team-accept-1"} + mock_db_client.db.litellm_teamtable = MagicMock() + mock_db_client.db.litellm_teamtable.create = AsyncMock(return_value=team_create_result) + mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_create_result) + mock_db_client.db.litellm_usertable = MagicMock() + mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) + + recorded = [] + with _configured_team_metadata_validator(_recording_validator(recorded)): + await new_team( + data=NewTeamRequest(team_alias="accepted-team", metadata={"cost_center": "CC-1001"}), + http_request=MagicMock(spec=Request), + user_api_key_dict=mock_admin_auth, + ) + + assert len(recorded) == 1 + assert recorded[0].operation == "create" + assert recorded[0].metadata == {"cost_center": "CC-1001"} + mock_db_client.db.litellm_teamtable.create.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_new_team_rejection_precedes_model_alias_write(): + """A rejected create must not leave an orphaned LiteLLM_ModelTable row: + validation runs before the model_aliases insert.""" + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + validator = _recording_validator([], valid=False, error_message="cost_center is required") + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server._license_check") as mock_license, + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + _configured_team_metadata_validator(validator), + ): + mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_prisma.db.litellm_teamtable.create = AsyncMock() + mock_prisma.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model-1")) + mock_license.is_team_count_over_limit.return_value = False + + with pytest.raises(ProxyException): + await new_team( + data=NewTeamRequest( + team_alias="alias-orphan-check", + model_aliases={"alias-model": "gpt-4o"}, + ), + http_request=MagicMock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1"), + ) + + mock_prisma.db.litellm_modeltable.create.assert_not_awaited() + mock_prisma.db.litellm_teamtable.create.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["post", "patch"]) +async def test_update_existing_metadata_excludes_system_managed_keys(kind): + """The validator's existing_metadata must be symmetric with metadata: + server-owned keys (team_member_budget_id) are stripped from both, so a + key-preservation validator never sees them 'disappear'.""" + recorded = [] + stored = {"cost_center": "CC-1001", "team_member_budget_id": "budget-abc"} + + with _configured_team_metadata_validator(_recording_validator(recorded)): + await _drive_team_write( + kind, + existing_metadata=dict(stored), + payload={"metadata": {"notes": "x"}}, + ) + + assert len(recorded) == 1 + assert recorded[0].existing_metadata == {"cost_center": "CC-1001"} + assert "team_member_budget_id" not in recorded[0].metadata + + # --------------------------------------------------------------------------- # PATCH body is validated through PatchTeamRequest before it is handed to # update_team. The write below must stay byte-identical to what the untyped @@ -10340,6 +10588,76 @@ async def test_list_available_teams_filters_joined_and_validates_rows(monkeypatc assert find_many_kwargs["where"] == {"team_id": {"in": ["team-open"]}} +@pytest.mark.asyncio +async def test_get_team_metadata_schema_returns_configured_fields(): + from litellm.proxy.management_endpoints.team_endpoints import get_team_metadata_schema + from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_SCHEMA_REGISTRY, + parse_team_metadata_schema, + ) + + TEAM_METADATA_SCHEMA_REGISTRY.set( + parse_team_metadata_schema( + [ + {"key": "cost_center", "label": "Cost Center"}, + {"key": "app_name", "label": "Application Name"}, + ] + ) + ) + try: + result = await get_team_metadata_schema() + finally: + TEAM_METADATA_SCHEMA_REGISTRY.set(()) + + assert [field.key for field in result.fields] == ["cost_center", "app_name"] + assert result.fields[0].label == "Cost Center" + assert result.fields[1].label == "Application Name" + + +@pytest.mark.asyncio +async def test_get_team_metadata_schema_empty_when_unconfigured(): + from litellm.proxy.management_endpoints.team_endpoints import get_team_metadata_schema + from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_SCHEMA_REGISTRY, + ) + + TEAM_METADATA_SCHEMA_REGISTRY.set(()) + result = await get_team_metadata_schema() + + assert result.fields == () + + +def test_get_team_metadata_schema_route_requires_auth(): + from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_SCHEMA_REGISTRY, + parse_team_metadata_schema, + ) + + with patch("litellm.proxy.proxy_server.master_key", "sk-1234"): + response = client.get("/team/metadata_schema") + assert response.status_code == 401 + + TEAM_METADATA_SCHEMA_REGISTRY.set(parse_team_metadata_schema([{"key": "cost_center", "label": "Cost Center"}])) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1" + ) + try: + authed = client.get("/team/metadata_schema") + finally: + app.dependency_overrides.pop(user_api_key_auth, None) + TEAM_METADATA_SCHEMA_REGISTRY.set(()) + + assert authed.status_code == 200 + assert authed.json() == {"fields": [{"key": "cost_center", "label": "Cost Center"}]} + + +def test_team_metadata_schema_route_is_readable_by_non_admins(): + from litellm.proxy._types import LiteLLMRoutes + + assert "/team/metadata_schema" in LiteLLMRoutes.info_routes.value + assert "/team/metadata_schema" in LiteLLMRoutes.management_routes.value + + def _provisioning_caller(role: LitellmUserRoles) -> UserAPIKeyAuth: return UserAPIKeyAuth(user_id="caller-1", user_role=role) diff --git a/tests/test_litellm/proxy/management_helpers/team_metadata_validator_impls.py b/tests/test_litellm/proxy/management_helpers/team_metadata_validator_impls.py new file mode 100644 index 00000000000..39ebeb97ec1 --- /dev/null +++ b/tests/test_litellm/proxy/management_helpers/team_metadata_validator_impls.py @@ -0,0 +1,92 @@ +"""Three independent `custom_team_metadata_validate` implementations. + +Used by the matrix tests in `test_team_metadata_validation.py` and loadable +directly from a proxy config via `get_instance_fn` for live verification: + +- `validate_allowlist`: plain async function; requires `cost_center` and + checks it against a static allowlist. +- `validate_via_http`: async function that POSTs the metadata to an external + validation service (`TEAM_METADATA_VALIDATION_SERVICE_URL`); any transport + error or non-2xx response raises, exercising the fail-closed path. +- `IMMUTABLE_COST_CENTER_VALIDATOR`: class instance with an async + `__call__`; requires `cost_center` and forbids changing it once set, using + `existing_metadata` and `operation`. +""" + +import os + +import httpx + +from litellm.proxy.management_helpers.team_metadata_validation import ( + TeamMetadataValidationPayload, + TeamMetadataValidationResult, +) + +ALLOWED_COST_CENTERS = frozenset({"CC-1001", "CC-1002"}) +DEFAULT_SERVICE_URL = "http://localhost:9414/validate" + + +async def validate_allowlist( + payload: TeamMetadataValidationPayload, +) -> TeamMetadataValidationResult: + cost_center = payload.metadata.get("cost_center") + if cost_center is None: + return TeamMetadataValidationResult( + valid=False, + error_message="cost_center is required in team metadata. Contact the FinOps team.", + ) + if cost_center not in ALLOWED_COST_CENTERS: + return TeamMetadataValidationResult( + valid=False, + error_message=f"Cost center {cost_center} is not recognized. Contact the FinOps team.", + ) + return TeamMetadataValidationResult(valid=True) + + +async def validate_via_http( + payload: TeamMetadataValidationPayload, +) -> TeamMetadataValidationResult: + service_url = os.environ.get("TEAM_METADATA_VALIDATION_SERVICE_URL", DEFAULT_SERVICE_URL) + async with httpx.AsyncClient(timeout=2.0) as client: + response = await client.post( + service_url, + json={"operation": payload.operation, "metadata": payload.metadata}, + ) + response.raise_for_status() + body = response.json() + if body.get("ok") is True: + return TeamMetadataValidationResult(valid=True) + return TeamMetadataValidationResult( + valid=False, + error_message=body.get("reason", "Rejected by the cost center service."), + ) + + +class ImmutableCostCenterValidator: + def __init__(self, immutable_key: str = "cost_center") -> None: + self.immutable_key = immutable_key + + async def __call__( + self, + payload: TeamMetadataValidationPayload, + ) -> TeamMetadataValidationResult: + current = payload.metadata.get(self.immutable_key) + if current is None: + return TeamMetadataValidationResult( + valid=False, + error_message=f"{self.immutable_key} is required in team metadata. Contact the FinOps team.", + ) + if payload.operation == "update" and payload.existing_metadata is not None: + prior = payload.existing_metadata.get(self.immutable_key) + if prior is not None and prior != current: + return TeamMetadataValidationResult( + valid=False, + error_message=( + f"{self.immutable_key} is immutable once set " + f"(stored: {prior}, requested: {current}). Contact the FinOps team." + ), + ) + return TeamMetadataValidationResult(valid=True) + + +IMMUTABLE_COST_CENTER_VALIDATOR = ImmutableCostCenterValidator() diff --git a/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py b/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py new file mode 100644 index 00000000000..d22c3db0f7e --- /dev/null +++ b/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py @@ -0,0 +1,704 @@ +import asyncio +import os +import sys +from unittest.mock import patch + +import pytest +from fastapi import HTTPException + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.management_helpers.team_metadata_validation import ( + DEFAULT_TEAM_METADATA_VALIDATION_REJECTED_MESSAGE, + DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS, + DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE, + TeamMetadataRequester, + TeamMetadataValidationPayload, + TeamMetadataValidationResult, + TeamMetadataValidatorRegistry, + _read_timeout_seconds, + _read_unavailable_message, + run_team_metadata_validation, + validate_team_metadata_if_configured, +) + + +def _registry_with(validator): + registry = TeamMetadataValidatorRegistry() + registry.set(validator) + return registry + +UNAVAILABLE_MESSAGE = "validation system down, contact ops" + + +def _payload(**overrides): + values = { + "operation": "create", + "metadata": {"cost_center": "CC-1001"}, + "existing_metadata": None, + "team_id": "team-1", + "team_alias": "alias-1", + "requester": TeamMetadataRequester(user_id="u1"), + } + values.update(overrides) + return TeamMetadataValidationPayload(**values) + + +async def _run(validator, payload=None, premium_user=True, timeout_seconds=1.0): + await run_team_metadata_validation( + validator=validator, + payload=payload or _payload(), + premium_user=premium_user, + timeout_seconds=timeout_seconds, + unavailable_message=UNAVAILABLE_MESSAGE, + ) + + +@pytest.mark.asyncio +async def test_valid_result_passes(): + async def validator(payload): + return TeamMetadataValidationResult(valid=True) + + await _run(validator) + + +@pytest.mark.asyncio +async def test_rejection_raises_400_with_validator_message(): + async def validator(payload): + return TeamMetadataValidationResult(valid=False, error_message="cost center rejected, contact FinOps") + + with pytest.raises(HTTPException) as exc_info: + await _run(validator) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == {"error": "cost center rejected, contact FinOps"} + + +@pytest.mark.asyncio +async def test_rejection_without_message_uses_default(): + async def validator(payload): + return TeamMetadataValidationResult(valid=False) + + with pytest.raises(HTTPException) as exc_info: + await _run(validator) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == {"error": DEFAULT_TEAM_METADATA_VALIDATION_REJECTED_MESSAGE} + + +@pytest.mark.asyncio +async def test_dict_shaped_return_is_accepted(): + async def validator(payload): + return {"valid": False, "error_message": "rejected via dict"} + + with pytest.raises(HTTPException) as exc_info: + await _run(validator) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == {"error": "rejected via dict"} + + +@pytest.mark.asyncio +async def test_validator_exception_fails_closed_with_generic_message(): + async def validator(payload): + raise RuntimeError("internal validation service is down") + + with pytest.raises(HTTPException) as exc_info: + await _run(validator) + assert exc_info.value.status_code == 503 + assert exc_info.value.detail == {"error": UNAVAILABLE_MESSAGE} + + +@pytest.mark.asyncio +async def test_validator_timeout_fails_closed(): + async def validator(payload): + await asyncio.sleep(1.0) + return TeamMetadataValidationResult(valid=True) + + with pytest.raises(HTTPException) as exc_info: + await _run(validator, timeout_seconds=0.01) + assert exc_info.value.status_code == 503 + assert exc_info.value.detail == {"error": UNAVAILABLE_MESSAGE} + + +@pytest.mark.asyncio +async def test_malformed_return_shape_fails_closed(): + async def validator(payload): + return "not-a-validation-result" + + with pytest.raises(HTTPException) as exc_info: + await _run(validator) + assert exc_info.value.status_code == 503 + assert exc_info.value.detail == {"error": UNAVAILABLE_MESSAGE} + + +@pytest.mark.asyncio +async def test_non_premium_user_is_rejected(): + async def validator(payload): + return TeamMetadataValidationResult(valid=True) + + with pytest.raises(HTTPException) as exc_info: + await _run(validator, premium_user=False) + assert exc_info.value.status_code == 400 + assert CommonProxyErrors.not_premium_user.value in exc_info.value.detail["error"] + + +@pytest.mark.asyncio +async def test_non_coroutine_validator_is_rejected(): + def validator(payload): + return TeamMetadataValidationResult(valid=True) + + with pytest.raises(HTTPException) as exc_info: + await _run(validator) + assert exc_info.value.status_code == 500 + assert exc_info.value.detail == {"error": "custom_team_metadata_validate must be an async function"} + + +@pytest.mark.parametrize( + "general_settings, expected", + [ + ({}, DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS), + ({"team_metadata_validation_timeout": 2}, 2.0), + ({"team_metadata_validation_timeout": 0.5}, 0.5), + ({"team_metadata_validation_timeout": True}, DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS), + ({"team_metadata_validation_timeout": -1}, DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS), + ({"team_metadata_validation_timeout": 0}, DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS), + ({"team_metadata_validation_timeout": "3"}, DEFAULT_TEAM_METADATA_VALIDATION_TIMEOUT_SECONDS), + ], +) +def test_read_timeout_seconds(general_settings, expected): + assert _read_timeout_seconds(general_settings) == expected + + +@pytest.mark.parametrize( + "general_settings, expected", + [ + ({}, DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE), + ( + {"team_metadata_validation_error_message": "call the help desk"}, + "call the help desk", + ), + ({"team_metadata_validation_error_message": " "}, DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE), + ({"team_metadata_validation_error_message": None}, DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE), + ], +) +def test_read_unavailable_message(general_settings, expected): + assert _read_unavailable_message(general_settings) == expected + + +@pytest.mark.asyncio +async def test_adapter_is_noop_when_unconfigured(): + calls = [] + + async def validator(payload): + calls.append(payload) + return TeamMetadataValidationResult(valid=True) + + await validate_team_metadata_if_configured( + operation="create", + metadata={"cost_center": "CC-1001"}, + existing_metadata=None, + team_id="team-1", + team_alias="alias-1", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="u1"), + registry=TeamMetadataValidatorRegistry(), + ) + assert calls == [] + + +@pytest.mark.asyncio +async def test_adapter_builds_payload_and_reads_settings(): + recorded = [] + + async def validator(payload): + recorded.append(payload) + return TeamMetadataValidationResult(valid=True) + + with ( + patch("litellm.proxy.proxy_server.premium_user", True), + patch( + "litellm.proxy.proxy_server.general_settings", + {"team_metadata_validation_timeout": 3, "team_metadata_validation_error_message": "ops msg"}, + ), + ): + await validate_team_metadata_if_configured( + operation="update", + metadata={"cost_center": "CC-2001"}, + existing_metadata={"cost_center": "CC-1001", "keep": 1}, + team_id="team-9", + team_alias="alias-9", + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user-9", + user_email="user-9@example.com", + ), + registry=_registry_with(validator), + ) + + assert len(recorded) == 1 + payload = recorded[0] + assert payload.operation == "update" + assert payload.metadata == {"cost_center": "CC-2001"} + assert payload.existing_metadata == {"cost_center": "CC-1001", "keep": 1} + assert payload.team_id == "team-9" + assert payload.team_alias == "alias-9" + assert payload.requester.user_id == "user-9" + assert payload.requester.user_email == "user-9@example.com" + assert payload.requester.user_role == LitellmUserRoles.INTERNAL_USER.value + + +@pytest.mark.asyncio +async def test_adapter_normalizes_non_dict_metadata_to_empty_dict(): + recorded = [] + + async def validator(payload): + recorded.append(payload) + return TeamMetadataValidationResult(valid=True) + + with ( + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + await validate_team_metadata_if_configured( + operation="create", + metadata=None, + existing_metadata=None, + team_id="team-1", + team_alias=None, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="u1"), + registry=_registry_with(validator), + ) + + assert len(recorded) == 1 + assert recorded[0].metadata == {} + assert recorded[0].existing_metadata is None + + +# --------------------------------------------------------------------------- +# Validator implementation matrix: three independent implementations +# (allowlist function, HTTP-service-backed function, immutability class +# instance) driven through the real team write endpoints. +# --------------------------------------------------------------------------- + +import json as _json +import socket +import threading +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from unittest.mock import AsyncMock, MagicMock, Mock + +import team_metadata_validator_impls as impls + +from litellm.proxy._types import ProxyException +from litellm.proxy.management_helpers.team_metadata_validation import ( + TEAM_METADATA_VALIDATOR_REGISTRY, +) + + +class _CostCenterServiceHandler(BaseHTTPRequestHandler): + def do_POST(self): + length = int(self.headers.get("Content-Length", 0)) + body = _json.loads(self.rfile.read(length) or b"{}") + cost_center = (body.get("metadata") or {}).get("cost_center") + if cost_center is None: + resp = {"ok": False, "reason": "cost_center missing per cost center service"} + elif cost_center not in impls.ALLOWED_COST_CENTERS: + resp = {"ok": False, "reason": f"cost center {cost_center} rejected by cost center service"} + else: + resp = {"ok": True} + payload = _json.dumps(resp).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, format, *args): + pass + + +@pytest.fixture(scope="module") +def cost_center_service_url(): + server = ThreadingHTTPServer(("127.0.0.1", 0), _CostCenterServiceHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}/validate" + finally: + server.shutdown() + thread.join(timeout=5) + + +def _closed_port_url(): + probe = socket.socket() + probe.bind(("127.0.0.1", 0)) + port = probe.getsockname()[1] + probe.close() + return f"http://127.0.0.1:{port}/validate" + + +@contextmanager +def _configured(validator): + TEAM_METADATA_VALIDATOR_REGISTRY.set(validator) + try: + with ( + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + yield + finally: + TEAM_METADATA_VALIDATOR_REGISTRY.set(None) + + +async def _drive_create(metadata, mock_sink=None): + from fastapi import Request + + from litellm.proxy._types import LiteLLM_TeamTable, NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as pc, + patch("litellm.proxy.proxy_server._license_check") as lic, + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + ): + team_row = MagicMock(team_id="matrix-team-1") + team_row.model_dump.return_value = {"team_id": "matrix-team-1"} + pc.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + pc.get_data = AsyncMock(return_value=None) + pc.update_data = AsyncMock(return_value=MagicMock()) + pc.db.litellm_teamtable.create = AsyncMock(return_value=team_row) + pc.db.litellm_teamtable.count = AsyncMock(return_value=0) + pc.db.litellm_teamtable.update = AsyncMock(return_value=team_row) + pc.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) + pc.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model-1")) + lic.is_team_count_over_limit.return_value = False + if mock_sink is not None: + mock_sink["team_create"] = pc.db.litellm_teamtable.create + mock_sink["model_create"] = pc.db.litellm_modeltable.create + + request_kwargs = {"team_alias": "matrix-team"} + if metadata is not None: + request_kwargs["metadata"] = metadata + return await new_team( + data=NewTeamRequest(**request_kwargs), + http_request=MagicMock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + +async def _drive_update(kind, existing_metadata, payload): + from fastapi import Request + + from litellm.proxy._types import LiteLLM_TeamTable, PatchTeamRequest, UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import patch_team, update_team + + team_id = "matrix-team-upd" + existing = LiteLLM_TeamTable( + team_id=team_id, + team_alias="matrix", + metadata=existing_metadata, + organization_id=None, + ) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1") + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as pc, + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.management_endpoints.team_endpoints._refresh_cached_team", + new=AsyncMock(), + ), + ): + pc.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing) + pc.db.litellm_teamtable.update = AsyncMock( + return_value=LiteLLM_TeamTable(team_id=team_id, team_alias="matrix") + ) + pc.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + + req = Mock(spec=Request) + if kind == "post": + return await update_team( + data=UpdateTeamRequest(team_id=team_id, **payload), + http_request=req, + user_api_key_dict=auth, + litellm_changed_by=None, + ) + return await patch_team( + team_id=team_id, + data=PatchTeamRequest.model_validate(dict(payload)), + http_request=req, + user_api_key_dict=auth, + litellm_changed_by=None, + ) + + +OK = ("ok", None) + +_MATRIX_IMPLS = { + "allowlist": lambda: impls.validate_allowlist, + "http": lambda: impls.validate_via_http, + "immutable_class": lambda: impls.IMMUTABLE_COST_CENTER_VALIDATOR, +} + +# (scenario, kind, existing_metadata, request payload, {impl: expected}) +# expected is ("ok", None) or ("reject", ) +_MATRIX_SCENARIOS = [ + ( + "create-valid-cost-center", + "create", + None, + {"cost_center": "CC-1001"}, + {"allowlist": OK, "http": OK, "immutable_class": OK}, + ), + ( + "create-missing-cost-center", + "create", + None, + None, + { + "allowlist": ("reject", "cost_center is required"), + "http": ("reject", "cost_center missing per cost center service"), + "immutable_class": ("reject", "cost_center is required"), + }, + ), + ( + "create-unknown-cost-center", + "create", + None, + {"cost_center": "CC-9999"}, + { + "allowlist": ("reject", "is not recognized"), + "http": ("reject", "rejected by cost center service"), + "immutable_class": OK, + }, + ), + ( + "patch-change-cost-center", + "patch", + {"cost_center": "CC-1001"}, + {"metadata": {"cost_center": "CC-1002"}}, + { + "allowlist": OK, + "http": OK, + "immutable_class": ("reject", "immutable once set"), + }, + ), + ( + "patch-unrelated-key-preserves-cost-center", + "patch", + {"cost_center": "CC-1001"}, + {"metadata": {"notes": "hello"}}, + {"allowlist": OK, "http": OK, "immutable_class": OK}, + ), + ( + "patch-null-deletes-cost-center", + "patch", + {"cost_center": "CC-1001"}, + {"metadata": {"cost_center": None}}, + { + "allowlist": ("reject", "cost_center is required"), + "http": ("reject", "cost_center missing per cost center service"), + "immutable_class": ("reject", "cost_center is required"), + }, + ), + ( + "post-replace-drops-cost-center", + "post", + {"cost_center": "CC-1001"}, + {"metadata": {"notes": "only-notes"}}, + { + "allowlist": ("reject", "cost_center is required"), + "http": ("reject", "cost_center missing per cost center service"), + "immutable_class": ("reject", "cost_center is required"), + }, + ), + ( + "update-without-metadata-skips-validation", + "post", + {"cost_center": "CC-1001"}, + {"tpm_limit": 5}, + {"allowlist": OK, "http": OK, "immutable_class": OK}, + ), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("impl_name", sorted(_MATRIX_IMPLS)) +@pytest.mark.parametrize( + "scenario, kind, existing_metadata, request_payload, expectations", + _MATRIX_SCENARIOS, + ids=[row[0] for row in _MATRIX_SCENARIOS], +) +async def test_validator_implementation_matrix( + monkeypatch, + cost_center_service_url, + impl_name, + scenario, + kind, + existing_metadata, + request_payload, + expectations, +): + monkeypatch.setenv("TEAM_METADATA_VALIDATION_SERVICE_URL", cost_center_service_url) + validator = _MATRIX_IMPLS[impl_name]() + outcome, message_part = expectations[impl_name] + + async def drive(): + if kind == "create": + return await _drive_create(metadata=request_payload) + return await _drive_update(kind, existing_metadata, request_payload) + + with _configured(validator): + if outcome == "ok": + await drive() + else: + with pytest.raises(ProxyException) as exc_info: + await drive() + assert str(exc_info.value.code) == "400", f"{scenario} x {impl_name}" + assert message_part in str(exc_info.value.message), f"{scenario} x {impl_name}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "kind, existing_metadata, request_payload", + [ + ("create", None, {"cost_center": "CC-1001"}), + ("patch", {"cost_center": "CC-1001"}, {"metadata": {"notes": "x"}}), + ], +) +async def test_http_validator_service_outage_fails_closed(monkeypatch, kind, existing_metadata, request_payload): + monkeypatch.setenv("TEAM_METADATA_VALIDATION_SERVICE_URL", _closed_port_url()) + + with _configured(impls.validate_via_http): + with pytest.raises(ProxyException) as exc_info: + if kind == "create": + await _drive_create(metadata=request_payload) + else: + await _drive_update(kind, existing_metadata, request_payload) + + assert str(exc_info.value.code) == "503" + assert DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_class_instance_with_async_call_is_accepted(): + await _run(impls.ImmutableCostCenterValidator(), payload=_payload(metadata={"cost_center": "CC-1"})) + + +@pytest.mark.asyncio +async def test_class_instance_with_sync_call_is_rejected(): + class SyncValidator: + def __call__(self, payload): + return TeamMetadataValidationResult(valid=True) + + with pytest.raises(HTTPException) as exc_info: + await _run(SyncValidator()) + assert exc_info.value.status_code == 500 + + +from litellm.proxy.management_helpers.team_metadata_validation import ( + TeamMetadataSchemaRegistry, + parse_team_metadata_schema, +) + + +def test_parse_schema_none_returns_empty(): + assert parse_team_metadata_schema(None) == () + + +def test_parse_schema_round_trips_fields_in_order(): + raw = [ + {"key": "cost_center", "label": "Cost Center"}, + {"key": "app_name"}, + ] + + fields = parse_team_metadata_schema(raw) + + assert [field.key for field in fields] == ["cost_center", "app_name"] + assert fields[0].label == "Cost Center" + assert fields[1].label is None + + +@pytest.mark.parametrize( + "raw", + [ + "cost_center", + {"key": "cost_center"}, + [{"label": "missing key"}], + [{"key": ""}], + [{"key": "cost_center", "required": True}], + [{"key": "cost_center", "description": "Cost center code"}], + [{"key": "cost_center", "allowed_values": ["CC-1001"]}], + ], +) +def test_parse_schema_malformed_raises(raw): + with pytest.raises(Exception): + parse_team_metadata_schema(raw) + + +def test_parse_schema_duplicate_keys_raise(): + with pytest.raises(ValueError, match="duplicate"): + parse_team_metadata_schema([{"key": "cost_center"}, {"key": "app_name"}, {"key": "cost_center"}]) + + +def test_schema_registry_defaults_empty_and_round_trips(): + registry = TeamMetadataSchemaRegistry() + assert registry.get() == () + + fields = parse_team_metadata_schema([{"key": "cost_center"}]) + registry.set(fields) + assert registry.get() == fields + + registry.set(()) + assert registry.get() == () + + +@pytest.mark.asyncio +async def test_non_callable_validator_is_rejected_with_clean_500(): + class NotCallable: + pass + + with pytest.raises(HTTPException) as exc_info: + await _run(NotCallable()) + assert exc_info.value.status_code == 500 + assert exc_info.value.detail == {"error": "custom_team_metadata_validate must be an async function"} + + +def test_parse_schema_duplicate_error_lists_offending_keys(): + with pytest.raises(ValueError) as exc_info: + parse_team_metadata_schema( + [{"key": "cost_center"}, {"key": "app_name"}, {"key": "cost_center"}, {"key": "app_name"}] + ) + assert str(exc_info.value) == "team_metadata_schema contains duplicate keys: app_name, cost_center" + + +@pytest.mark.asyncio +async def test_adapter_applies_configured_timeout_to_slow_validator(): + async def slow_validator(payload): + await asyncio.sleep(0.2) + return TeamMetadataValidationResult(valid=True) + + registry = TeamMetadataValidatorRegistry() + registry.set(slow_validator) + + with ( + patch("litellm.proxy.proxy_server.premium_user", True), + patch( + "litellm.proxy.proxy_server.general_settings", + {"team_metadata_validation_timeout": 0.01}, + ), + ): + with pytest.raises(HTTPException) as exc_info: + await validate_team_metadata_if_configured( + operation="create", + metadata={"cost_center": "CC-1001"}, + existing_metadata=None, + team_id="team-1", + team_alias="alias-1", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="u1"), + registry=registry, + ) + + assert exc_info.value.status_code == 503 diff --git a/tests/test_litellm/router_utils/test_auto_router_model_naming.py b/tests/test_litellm/router_utils/test_auto_router_model_naming.py index ca290caac0b..73c9742876c 100644 --- a/tests/test_litellm/router_utils/test_auto_router_model_naming.py +++ b/tests/test_litellm/router_utils/test_auto_router_model_naming.py @@ -2,6 +2,7 @@ import pytest from litellm.router_utils.auto_router_model_naming import ( classify_strategy_router_model, + validate_complexity_router_config_write, validate_strategy_router_model_write, ) @@ -75,3 +76,74 @@ def test_validate_rejects_incoherent_writes(model, present_fields, expected_frag ) def test_validate_accepts_coherent_writes(model, present_fields): assert validate_strategy_router_model_write(model=model, present_fields=present_fields) is None + + +VALID_TIERS = { + "SIMPLE": ["gpt-4o-mini"], + "MEDIUM": ["gpt-4o-mini"], + "COMPLEX": ["gpt-4o"], + "REASONING": ["gpt-4o"], +} + + +@pytest.mark.parametrize( + "keyword_tier_rules,expected_fragment", + [ + ([{"keywords": [], "tier": "COMPLEX"}], "at least 1 item"), + ([{"keywords": [" "], "tier": "COMPLEX"}], "non-empty keyword"), + ( + [{"keywords": ["invoice"], "tier": "MEDIUM"}, {"keywords": [], "tier": "COMPLEX"}], + "at least 1 item", + ), + ], +) +def test_validate_rejects_unloadable_complexity_config(keyword_tier_rules, expected_fragment): + """A rule with no keyword makes ComplexityRouterConfig unbuildable, so the row must never be + written: without this the deployment is persisted, dropped at load, and the caller gets a 500.""" + violation = validate_complexity_router_config_write( + complexity_router_config={ + "tiers": VALID_TIERS, + "classifier_type": "heuristic", + "keyword_tier_rules": keyword_tier_rules, + } + ) + assert violation is not None + assert "complexity_router_config is invalid" in violation + assert expected_fragment in violation + + +@pytest.mark.parametrize( + "complexity_router_config", + [ + {"tiers": VALID_TIERS, "classifier_type": "heuristic"}, + { + "tiers": VALID_TIERS, + "classifier_type": "heuristic", + "keyword_tier_rules": [{"keywords": ["invoice", "refund"], "tier": "MEDIUM"}], + }, + # extra="allow" on the model, so an unrecognised key is not this gate's business + {"tiers": VALID_TIERS, "classifier_type": "heuristic", "some_future_key": "value"}, + ], +) +def test_validate_accepts_loadable_complexity_config(complexity_router_config): + assert validate_complexity_router_config_write(complexity_router_config=complexity_router_config) is None + + +def test_naming_check_ignores_the_config_entirely(): + """The naming contract and the config's contents are separate questions with separate owners; + a write may carry a config without naming a model, so neither can stand in for the other.""" + violation = validate_strategy_router_model_write( + model="auto_router/complexity_router", present_fields=frozenset() + ) + assert violation is not None + assert "requires" in violation + + +def test_config_check_ignores_the_model_entirely(): + assert validate_complexity_router_config_write(complexity_router_config=None) is None + assert ( + validate_complexity_router_config_write( + complexity_router_config={"tiers": VALID_TIERS, "keyword_tier_rules": [{"keywords": [], "tier": "COMPLEX"}]} + ) + is not None + ) diff --git a/tests/test_litellm/test_circleci_rust_toolchain.py b/tests/test_litellm/test_circleci_rust_toolchain.py new file mode 100644 index 00000000000..c35ced51e16 --- /dev/null +++ b/tests/test_litellm/test_circleci_rust_toolchain.py @@ -0,0 +1,148 @@ +"""Static guardrails for how CircleCI provisions Rust. + +The root package builds `litellm-rust` through maturin, so any job that runs +`uv sync` or `uv build` compiles the bridge. The `cimg/python` images ship no +Rust toolchain, and when cargo is missing maturin's `puccinialin` helper +quietly provisions one itself: it fetches `rustup-init` from the unversioned +`https://static.rust-lang.org/rustup/dist//` path with no checksum and +installs a floating `stable` toolchain. uv suppresses build-backend output on a +successful sync, so that happens with nothing in the job log to show for it, +and the compiler a job builds with changes whenever upstream publishes. + +Two invariants are pinned here: + + 1. No step list (job or reusable command) reaches a `uv sync` / `uv build` + without a Rust toolchain already provisioned ahead of it. That is the + `install_rust` command on Linux and an inline pinned rustup install in the + Windows job, so the check accepts either. A new job that syncs without one + falls back to the unpinned path, which is exactly the regression a static + check catches at PR time and a green CI run does not. + 2. `install_rust` itself pins what it downloads: an explicit rustup version in + the URL, a verified SHA-256, and an exact toolchain version rather than a + channel name. + +The Windows job predates `install_rust` and provisions its toolchain inline, so +invariant 2 is scoped to `install_rust`; invariant 1 covers both. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import pytest +import yaml + +REPO_ROOT = Path(__file__).resolve().parents[2] +CONFIG = REPO_ROOT / ".circleci" / "config.yml" + +BUILDS_WORKSPACE = re.compile(r"\buv\s+(?:sync|build)\b") +RUSTUP_ARCHIVE_URL = re.compile(r"https://static\.rust-lang\.org/rustup/archive/\d+\.\d+\.\d+/") +EXACT_TOOLCHAIN = re.compile(r"--default-toolchain\s+\"?\d+\.\d+\.\d+\"?") + + +def _config() -> dict[str, object]: + return yaml.safe_load(CONFIG.read_text()) + + +def _step_text(step: object) -> str: + """Flatten one step into the shell text it runs, or '' for a command reference.""" + if isinstance(step, dict): + run = step.get("run") + if isinstance(run, str): + return run + if isinstance(run, dict): + command = run.get("command") + return command if isinstance(command, str) else "" + return "" + + +def _without_comments(text: str) -> str: + return "\n".join(line for line in text.splitlines() if not line.lstrip().startswith("#")) + + +def _provisions_rust(step: object) -> bool: + if step == "install_rust": + return True + text = _step_text(step) + return "rustup-init" in text and ("sha256sum" in text or "SHA256" in text) + + +def _step_lists() -> dict[str, list[object]]: + config = _config() + lists: dict[str, list[object]] = {} + for kind in ("jobs", "commands"): + section = config.get(kind) + if not isinstance(section, dict): + continue + for name, body in section.items(): + steps = body.get("steps") if isinstance(body, dict) else None + if isinstance(steps, list): + lists[f"{kind[:-1]} {name}"] = steps + return lists + + +def _first_unprovisioned_build(steps: list[object]) -> str | None: + """Return the shell text of the first workspace build reached without Rust, if any.""" + rust_ready = False + for step in steps: + if _provisions_rust(step): + rust_ready = True + text = _step_text(step) + if BUILDS_WORKSPACE.search(_without_comments(text)) and not rust_ready: + return text + return None + + +def test_step_lists_exist() -> None: + lists = _step_lists() + assert "command install_rust" in lists + building = { + name + for name, steps in lists.items() + if any(BUILDS_WORKSPACE.search(_without_comments(_step_text(s))) for s in steps) + } + assert len(building) > 10, f"expected many workspace-building step lists, found {sorted(building)}" + + +def test_no_workspace_build_without_a_provisioned_rust_toolchain() -> None: + offenders = { + name: build for name, steps in _step_lists().items() if (build := _first_unprovisioned_build(steps)) is not None + } + assert not offenders, ( + "these CircleCI step lists run `uv sync`/`uv build` with no Rust toolchain provisioned first, " + "so maturin will download an unpinned rustup and a floating toolchain instead: " + f"{ {name: build.strip().splitlines()[0] for name, build in offenders.items()} }" + ) + + +@pytest.fixture(name="install_rust_command") +def _install_rust_command() -> str: + steps = _step_lists()["command install_rust"] + return "\n".join(_step_text(step) for step in steps) + + +def test_install_rust_pins_the_rustup_version_in_the_url(install_rust_command: str) -> None: + assert RUSTUP_ARCHIVE_URL.search(install_rust_command), ( + "install_rust must download rustup-init from a version-pinned /rustup/archive// URL; " + "the /rustup/dist/ path always serves whatever rustup is current" + ) + assert "/rustup/dist/" not in install_rust_command + + +def test_install_rust_verifies_the_installer_checksum(install_rust_command: str) -> None: + assert "sha256sum -c" in install_rust_command + assert re.search(r"RUSTUP_SHA256=[0-9a-f]{64}\b", install_rust_command), ( + "install_rust must compare the downloaded installer against a hardcoded SHA-256 " + "taken from rust-lang's published .sha256 sidecar" + ) + checksum_index = install_rust_command.index("sha256sum -c") + execute_index = install_rust_command.index("/tmp/rustup-init -y") + assert checksum_index < execute_index, "the checksum must be verified before the installer is executed" + + +def test_install_rust_pins_an_exact_toolchain_version(install_rust_command: str) -> None: + assert EXACT_TOOLCHAIN.search(install_rust_command), ( + "install_rust must pin an exact toolchain version (e.g. 1.97.1); a channel name like " + "stable/beta/nightly makes the compiler drift with whatever upstream published that day" + ) diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 107f66b8f1a..50b652ff57e 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2862,6 +2862,16 @@ "count": 2 } }, + "src/components/common_components/MetadataKeyValueFields.test.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/components/common_components/MetadataKeyValueFields.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, "src/components/common_components/ModelAliasManager.tsx": { "no-restricted-imports": { "count": 1 @@ -4142,11 +4152,6 @@ "count": 1 } }, - "src/components/view_logs/LogDetailsDrawer/prettyMessagesUtils.ts": { - "no-nested-ternary": { - "count": 1 - } - }, "src/components/view_logs/LogDetailsDrawer/useKeyboardNavigation.ts": { "react-hooks/immutability": { "count": 2 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.test.tsx new file mode 100644 index 00000000000..603cddb7b89 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.test.tsx @@ -0,0 +1,142 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderWithProviders } from "@/../tests/test-utils"; +import { screen, fireEvent } from "@testing-library/react"; +import { TeamGuardrailsTab } from "./TeamGuardrailsTab"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +vi.mock("@/components/networking", () => ({ + listGuardrailSubmissions: vi.fn(), + approveGuardrailSubmission: vi.fn(), + rejectGuardrailSubmission: vi.fn(), + updateGuardrailCall: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/guardrails/useRegisterGuardrail", () => ({ + useRegisterGuardrail: () => ({ + mutateAsync: vi.fn(), + isPending: false, + }), +})); + +vi.mock("@/components/common_components/team_dropdown", () => ({ + default: () => null, +})); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: vi.fn(), +})); + +import { listGuardrailSubmissions } from "@/components/networking"; + +const pendingSubmission = { + guardrail_id: "guard-1", + guardrail_name: "test-pending-guardrail", + status: "pending_review", + team_id: "team-1", + team_guardrail: true, + litellm_params: { + guardrail: "generic_guardrail_api", + mode: "pre_call", + api_base: "https://example.com/guard", + headers: { "X-API-Key": "secret" }, + extra_headers: ["x-request-id"], + }, + guardrail_info: {}, + submitted_at: "2026-05-09T00:00:00Z", +}; + +const baseAuth = { + token: "test-token", + accessToken: "test-token", + userId: "user-1", + userEmail: "user@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, +}; + +describe("TeamGuardrailsTab — approve/reject role gate", () => { + const mockUseAuthorized = vi.mocked(useAuthorized); + + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(listGuardrailSubmissions).mockResolvedValue({ + submissions: [pendingSubmission], + summary: { total: 1, pending_review: 1, active: 0, rejected: 0 }, + }); + }); + + it("hides Approve and Reject buttons for an internal user on a pending submission", async () => { + mockUseAuthorized.mockReturnValue({ ...baseAuth, userRole: "Internal User" }); + renderWithProviders(); + + await screen.findByText("test-pending-guardrail"); + + expect(screen.queryByRole("button", { name: /approve/i })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /reject/i })).not.toBeInTheDocument(); + }); + + it("hides Approve and Reject buttons for an Admin Viewer, whom the backend rejects with 403", async () => { + mockUseAuthorized.mockReturnValue({ ...baseAuth, userRole: "Admin Viewer" }); + renderWithProviders(); + + await screen.findByText("test-pending-guardrail"); + + expect(screen.queryByRole("button", { name: /approve/i })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /reject/i })).not.toBeInTheDocument(); + }); + + it("shows Approve and Reject buttons for an admin on a pending submission", async () => { + mockUseAuthorized.mockReturnValue({ ...baseAuth, userRole: "Admin" }); + renderWithProviders(); + + await screen.findByText("test-pending-guardrail"); + + expect(screen.getByRole("button", { name: /approve/i })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /reject/i })).toBeInTheDocument(); + }); + + it("hides Approve and Reject buttons when userRole is undefined (defaults to non-admin)", async () => { + mockUseAuthorized.mockReturnValue({ ...baseAuth, userRole: undefined }); + renderWithProviders(); + + await screen.findByText("test-pending-guardrail"); + + expect(screen.queryByRole("button", { name: /approve/i })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /reject/i })).not.toBeInTheDocument(); + }); + + it("disables all admin-only write controls for a non-admin, including the detail panel", async () => { + mockUseAuthorized.mockReturnValue({ ...baseAuth, userRole: "Internal User" }); + renderWithProviders(); + + await screen.findByText("test-pending-guardrail"); + expect(screen.getByRole("switch")).toBeDisabled(); + + fireEvent.click(screen.getByRole("button", { name: "Review" })); + await screen.findByText("Forward LiteLLM API Key"); + + expect(screen.queryByRole("button", { name: /approve/i })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /reject/i })).not.toBeInTheDocument(); + screen.getAllByRole("switch").forEach((toggle) => expect(toggle).toBeDisabled()); + expect(screen.queryByRole("button", { name: "Add" })).not.toBeInTheDocument(); + expect(screen.queryByLabelText(/^Remove/)).not.toBeInTheDocument(); + expect(screen.queryByPlaceholderText("e.g. x-request-id")).not.toBeInTheDocument(); + }); + + it("keeps all write controls enabled for an admin in the detail panel", async () => { + mockUseAuthorized.mockReturnValue({ ...baseAuth, userRole: "Admin" }); + renderWithProviders(); + + await screen.findByText("test-pending-guardrail"); + + fireEvent.click(screen.getByRole("button", { name: "Review" })); + await screen.findByText("Forward LiteLLM API Key"); + + expect(screen.getAllByRole("button", { name: /approve/i }).length).toBeGreaterThanOrEqual(2); + screen.getAllByRole("switch").forEach((toggle) => expect(toggle).toBeEnabled()); + expect(screen.getAllByRole("button", { name: "Add" })).toHaveLength(2); + expect(screen.getByLabelText("Remove X-API-Key")).toBeInTheDocument(); + expect(screen.getByLabelText("Remove x-request-id")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx index 4217a765732..496e1129371 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx @@ -27,6 +27,8 @@ import { import NotificationsManager from "@/components/molecules/notifications_manager"; import TeamDropdown from "@/components/common_components/team_dropdown"; import { useRegisterGuardrail } from "@/app/(dashboard)/hooks/guardrails/useRegisterGuardrail"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { isProxyAdminRole } from "@/utils/roles"; type GuardrailStatus = "active" | "pending" | "rejected"; @@ -188,16 +190,25 @@ function StatCard({ label, value, color }: { label: string; value: number; color ); } -function Toggle({ enabled, onToggle }: { enabled: boolean; onToggle: () => void }) { +function Toggle({ + enabled, + onToggle, + disabled = false, +}: { + enabled: boolean; + onToggle: () => void; + disabled?: boolean; +}) { return ( - {g.status === "pending" && ( + {isAdmin && g.status === "pending" && ( <> + {isAdmin && ( + + )} ))} )} -
- setNewStaticHeaderKey(e.target.value)} - placeholder="Header name (e.g. X-API-Key)" - className="flex-1 min-w-0 text-xs font-mono border border-gray-200 rounded-sm px-2 py-1.5 text-gray-700 placeholder-gray-400 focus:outline-hidden focus:ring-1 focus:ring-blue-500" - onKeyDown={(e) => { - if (e.key === "Enter") { - e.preventDefault(); + {isAdmin && ( +
+ setNewStaticHeaderKey(e.target.value)} + placeholder="Header name (e.g. X-API-Key)" + className="flex-1 min-w-0 text-xs font-mono border border-gray-200 rounded-sm px-2 py-1.5 text-gray-700 placeholder-gray-400 focus:outline-hidden focus:ring-1 focus:ring-blue-500" + onKeyDown={(e) => { + if (e.key === "Enter") { + e.preventDefault(); + const key = newStaticHeaderKey.trim(); + const value = newStaticHeaderValue.trim(); + if (key && !g.customHeaders.some((h) => h.key.toLowerCase() === key.toLowerCase())) { + onUpdateCustomHeaders([...g.customHeaders, { key, value }]); + setNewStaticHeaderKey(""); + setNewStaticHeaderValue(""); + } + } + }} + /> + setNewStaticHeaderValue(e.target.value)} + placeholder="Value" + className="flex-1 min-w-0 text-xs font-mono border border-gray-200 rounded-sm px-2 py-1.5 text-gray-700 placeholder-gray-400 focus:outline-hidden focus:ring-1 focus:ring-blue-500" + onKeyDown={(e) => { + if (e.key === "Enter") { + e.preventDefault(); + const key = newStaticHeaderKey.trim(); + const value = newStaticHeaderValue.trim(); + if (key && !g.customHeaders.some((h) => h.key.toLowerCase() === key.toLowerCase())) { + onUpdateCustomHeaders([...g.customHeaders, { key, value }]); + setNewStaticHeaderKey(""); + setNewStaticHeaderValue(""); + } + } + }} + /> + -
+ }} + className="text-xs font-medium text-blue-600 hover:text-blue-700 border border-blue-200 bg-blue-50 hover:bg-blue-100 px-2 py-1.5 rounded-sm transition-colors shrink-0" + > + Add + +
+ )}
@@ -546,50 +565,54 @@ function DetailPanel({ className="flex items-center justify-between gap-2 text-xs font-mono bg-gray-50 border border-gray-200 rounded-sm px-2 py-1.5" > {name} - + {isAdmin && ( + + )} ))} )} -
- setNewExtraHeader(e.target.value)} - placeholder="e.g. x-request-id" - className="flex-1 min-w-0 text-xs font-mono border border-gray-200 rounded-sm px-2 py-1.5 text-gray-700 placeholder-gray-400 focus:outline-hidden focus:ring-1 focus:ring-blue-500" - onKeyDown={(e) => { - if (e.key === "Enter") { - e.preventDefault(); + {isAdmin && ( +
+ setNewExtraHeader(e.target.value)} + placeholder="e.g. x-request-id" + className="flex-1 min-w-0 text-xs font-mono border border-gray-200 rounded-sm px-2 py-1.5 text-gray-700 placeholder-gray-400 focus:outline-hidden focus:ring-1 focus:ring-blue-500" + onKeyDown={(e) => { + if (e.key === "Enter") { + e.preventDefault(); + const name = newExtraHeader.trim().toLowerCase(); + if (name && !g.extraHeaders.map((h) => h.toLowerCase()).includes(name)) { + onUpdateExtraHeaders([...g.extraHeaders, name]); + setNewExtraHeader(""); + } + } + }} + /> + -
+ }} + className="text-xs font-medium text-blue-600 hover:text-blue-700 border border-blue-200 bg-blue-50 hover:bg-blue-100 px-2 py-1.5 rounded-sm transition-colors" + > + Add + +
+ )}
- {g.status === "pending" && ( + {isAdmin && g.status === "pending" && (
diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index 1f89403a358..50689f5cf09 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -1,4 +1,4 @@ -import { renderWithProviders, screen, waitFor, testQueryClient } from "../../../tests/test-utils"; +import { renderWithProviders, screen, waitFor, testQueryClient, within } from "../../../tests/test-utils"; import { act, fireEvent } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { vi } from "vitest"; @@ -123,14 +123,22 @@ describe("AddAutoRouterTab", () => { mockHandleAddAutoRouterSubmit.mockResolvedValue(undefined); }); - it("flags every mandatory field when Add Auto Router is clicked with nothing filled", async () => { + // Nothing is filled in, so there is nothing to submit. The button reports that itself instead of + // accepting a click and answering with a toast. + it("offers no submit at all until every tier has a model", async () => { + renderWithProviders(); + + expect(screen.getByRole("button", { name: /add auto router/i })).toBeDisabled(); + }); + + it("still flags the router name once the config no longer blocks the submit", async () => { const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); renderWithProviders(); await user.click(screen.getByRole("button", { name: /add auto router/i })); expect(await screen.findByText("Auto router name is required")).toBeInTheDocument(); - expect(screen.getAllByText("This tier is required")).toHaveLength(4); expect(NotificationManager.fromBackend).toHaveBeenCalledWith("Please enter an Auto Router Name"); }); @@ -517,6 +525,85 @@ describe("AddAutoRouterTab", () => { expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0]).toMatchObject({ team_id: "team-1" }); }); + // LIT-5133: "Add keyword rule" seeds a row with no keywords, and the semantic toggle that used + // to be the only thing checking them is off by default. The row was dropped on the way to the + // payload, so the create succeeded and the caller's rule was gone with nothing said about it. + it("takes the submit away while a keyword rule is left empty", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "keyword-router"); + await user.click(screen.getByText("Advanced: Keyword/Semantic Matching")); + await user.click(screen.getByRole("button", { name: /add keyword rule/i })); + + expect(screen.getByRole("button", { name: /add auto router/i })).toBeDisabled(); + // The row says so on its own; there is no failed submit left to surface it. + expect(await screen.findByText("At least one keyword is required")).toBeInTheDocument(); + expect(handleAddAutoRouterSubmit).not.toHaveBeenCalled(); + }); + + it("gives the submit back once that keyword rule is filled", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "keyword-router"); + await user.click(screen.getByText("Advanced: Keyword/Semantic Matching")); + await user.click(screen.getByRole("button", { name: /add keyword rule/i })); + expect(screen.getByRole("button", { name: /add auto router/i })).toBeDisabled(); + + await user.type( + within(screen.getByText("Keywords 1").closest("div") as HTMLElement).getByRole("combobox"), + "invoice{enter}", + ); + + expect(screen.getByRole("button", { name: /add auto router/i })).toBeEnabled(); + expect(screen.queryByText("At least one keyword is required")).not.toBeInTheDocument(); + }); + + it("marks only the offending keyword row, leaving a filled one alone", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "keyword-router"); + await user.click(screen.getByText("Advanced: Keyword/Semantic Matching")); + await user.click(screen.getByRole("button", { name: /add keyword rule/i })); + await user.type( + within(screen.getByText("Keywords 1").closest("div") as HTMLElement).getByRole("combobox"), + "invoice{enter}", + ); + await user.click(screen.getByRole("button", { name: /add keyword rule/i })); + + expect(await screen.findAllByText("At least one keyword is required")).toHaveLength(1); + expect(screen.getByRole("button", { name: /add auto router/i })).toBeDisabled(); + }); + + it("creates the router once that keyword rule is filled in", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + openTemplateDropdown(); + fireEvent.click(optionByLabel("Custom Configuration")!); + await user.type(screen.getByPlaceholderText(/smart_router/i), "keyword-router"); + await user.click(screen.getByText("Advanced: Keyword/Semantic Matching")); + await user.click(screen.getByRole("button", { name: /add keyword rule/i })); + const keywordsField = screen.getByText("Keywords 1").closest("div") as HTMLElement; + await user.type(within(keywordsField).getByRole("combobox"), "invoice{enter}"); + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0]).toMatchObject({ + complexity_router_config: { keyword_tier_rules: [{ keywords: ["invoice"], tier: "COMPLEX" }] }, + }); + }); + it("blocks the submit when a team admin has not picked a team", async () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 86f1e7b5512..905365ba8f9 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -20,6 +20,7 @@ import { DEFAULT_ESCALATION_KEYWORDS } from "./EscalationKeywords"; import { DEFAULT_MATCH_THRESHOLD } from "./SemanticKeywordMatching"; import { buildComplexityRouterConfig, + getKeywordTierRulesError, getMissingTiersError, getSemanticConfigError, } from "./build_complexity_router_config"; @@ -214,6 +215,11 @@ const AddAutoRouterTab: React.FC = ({ return getMissingModelsInPreset(preset, freshSet).length === 0; }; + // Why the submit is unavailable, or null when it is available. The button reads this to disable + // itself and to say what is missing, so the two can never give different answers. + const submitBlockedReason = + getMissingTiersError(complexityRouterConfig.tiers) ?? getKeywordTierRulesError(keywordTierRules); + const submitRecommendedRouter = async (name: string) => { if (!selectedPreset) { setShowValidationErrors(true); @@ -257,6 +263,13 @@ const AddAutoRouterTab: React.FC = ({ return; } + const keywordRulesError = getKeywordTierRulesError(keywordTierRules); + if (keywordRulesError) { + setShowValidationErrors(true); + NotificationManager.fromBackend(keywordRulesError); + return; + } + const semanticError = getSemanticConfigError({ semanticMatchingEnabled, embeddingModel, keywordTierRules }); if (semanticError) { setShowValidationErrors(true); @@ -502,15 +515,18 @@ const AddAutoRouterTab: React.FC = ({ Test Connection } - + + +
diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index 9d784b57903..4cbe54ad4a6 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -1,5 +1,6 @@ import { buildComplexityRouterConfig, + getKeywordTierRulesError, getMissingTiersError, getSemanticConfigError, BuildComplexityRouterConfigParams, @@ -183,7 +184,7 @@ describe("buildComplexityRouterConfig", () => { expect(config.keyword_tier_rules).toBeUndefined(); }); - it("trims keywords and drops rules left empty, so unfilled rows never 400 the backend", () => { + it("trims keywords but keeps rules left empty, so a dropped row can never pass for a saved one", () => { const params: BuildComplexityRouterConfigParams = { ...baseParams, keywordTierRules: [ @@ -193,17 +194,13 @@ describe("buildComplexityRouterConfig", () => { ], }; const config = buildComplexityRouterConfig(params); - // r1 keeps only its real keyword (trimmed); r2 and r3 are dropped entirely. - expect(config.keyword_tier_rules).toEqual([{ keywords: ["deploy to k8s"], tier: "REASONING" }]); - }); - - it("omits keyword_tier_rules entirely when every rule is empty", () => { - const params: BuildComplexityRouterConfigParams = { - ...baseParams, - keywordTierRules: [{ id: "r1", keywords: ["", " "], tier: "COMPLEX" }], - }; - const config = buildComplexityRouterConfig(params); - expect(config.keyword_tier_rules).toBeUndefined(); + // getKeywordTierRulesError blocks this submit; r2 and r3 survive here so the backend rejects + // them loudly rather than the caller's rows vanishing on a successful save. + expect(config.keyword_tier_rules).toEqual([ + { keywords: ["deploy to k8s"], tier: "REASONING" }, + { keywords: [], tier: "COMPLEX" }, + { keywords: [], tier: "SIMPLE" }, + ]); }); it("omits adaptive fields when adaptive is disabled even if weights linger in state", () => { @@ -318,17 +315,6 @@ describe("getSemanticConfigError", () => { ).toMatch(/keyword tier rule/i); }); - it("errors when a rule has no non-empty keywords", () => { - const emptyRule = { id: "r2", keywords: ["", " "], tier: "SIMPLE" as const }; - expect( - getSemanticConfigError({ - semanticMatchingEnabled: true, - embeddingModel: "voyage-3-5", - keywordTierRules: [emptyRule], - }), - ).toMatch(/at least one keyword/i); - }); - it("returns null when enabled with both an embedding model and rules", () => { expect( getSemanticConfigError({ semanticMatchingEnabled: true, embeddingModel: "voyage-3-5", keywordTierRules: [rule] }), @@ -336,6 +322,52 @@ describe("getSemanticConfigError", () => { }); }); +describe("getKeywordTierRulesError", () => { + it("returns null when every rule carries a keyword", () => { + expect( + getKeywordTierRulesError([ + { id: "r1", keywords: ["invoice"], tier: "MEDIUM" }, + { id: "r2", keywords: ["deploy to k8s"], tier: "REASONING" }, + ]), + ).toBeNull(); + }); + + it("returns null when there are no rules at all, since the section is optional", () => { + expect(getKeywordTierRulesError([])).toBeNull(); + }); + + // The whole point of the ticket: the semantic toggle is off by default, and an unfilled row + // used to be discarded silently on an otherwise successful create. + it("rejects a row left empty while semantic matching is off", () => { + expect(getKeywordTierRulesError([{ id: "r1", keywords: [], tier: "COMPLEX" }])).toBe( + "Add at least one keyword to keyword rule(s): 1", + ); + }); + + it.each([ + ["whitespace only", [" "]], + ["blank strings, as an unfilled row between filled ones leaves behind", ["", " ", ""]], + ])("treats %s as empty rather than as a keyword", (_label, keywords) => { + expect(getKeywordTierRulesError([{ id: "r1", keywords, tier: "SIMPLE" }])).toMatch(/keyword rule\(s\): 1/); + }); + + // Row numbers have to survive rules that are fine, or the message points at the wrong input. + it("names each offending row by its position among all rules", () => { + expect( + getKeywordTierRulesError([ + { id: "r1", keywords: ["invoice"], tier: "MEDIUM" }, + { id: "r2", keywords: [], tier: "COMPLEX" }, + { id: "r3", keywords: ["billing"], tier: "SIMPLE" }, + { id: "r4", keywords: [" "], tier: "REASONING" }, + ]), + ).toBe("Add at least one keyword to keyword rule(s): 2, 4"); + }); + + it("keeps a keyword whose surrounding whitespace is the only thing trimmed", () => { + expect(getKeywordTierRulesError([{ id: "r1", keywords: [" invoice "], tier: "MEDIUM" }])).toBeNull(); + }); +}); + describe("buildComplexityRouterConfig assistant turns", () => { const llmParams: BuildComplexityRouterConfigParams = { ...baseParams, diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index cd6c697b377..dcec58479a6 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -1,5 +1,5 @@ import { KeywordTierRule } from "./KeywordTierRules"; -import { serializeKeywordTierRules } from "./complexity_router_keywords"; +import { emptyKeywordTierRuleIndexes, serializeKeywordTierRules } from "./complexity_router_keywords"; import { AdaptiveEligible, AdaptiveRouterWeights, @@ -58,6 +58,12 @@ export const getMissingTiersError = (tiers: ComplexityTiers): string | null => { return `Select a model for the following tier(s): ${missing.join(", ")}`; }; +export const getKeywordTierRulesError = (keywordTierRules: KeywordTierRule[]): string | null => { + const emptyRows = emptyKeywordTierRuleIndexes(keywordTierRules); + if (emptyRows.length === 0) return null; + return `Add at least one keyword to keyword rule(s): ${emptyRows.map((index) => index + 1).join(", ")}`; +}; + export const getSemanticConfigError = ({ semanticMatchingEnabled, embeddingModel, @@ -68,8 +74,6 @@ export const getSemanticConfigError = ({ if (!semanticMatchingEnabled) return null; if (!embeddingModel) return "Select an embedding model to use semantic keyword matching"; if (keywordTierRules.length === 0) return "Add at least one keyword tier rule to use semantic keyword matching"; - if (keywordTierRules.some((rule) => !rule.keywords.some((keyword) => keyword.trim()))) - return "Every keyword tier rule needs at least one keyword"; return null; }; @@ -94,7 +98,6 @@ export const buildComplexityRouterConfig = ({ returnRawModelName, }: BuildComplexityRouterConfigParams): ComplexityRouterConfigPayload => { const cleanedEscalationKeywords = escalationKeywords.map((keyword) => keyword.trim()).filter(Boolean); - // Trim keywords and drop empty ones; drop any rule left with no keywords. Clicking const cleanedKeywordTierRules = serializeKeywordTierRules(keywordTierRules); return { diff --git a/ui/litellm-dashboard/src/components/add_model/complexity_router_keywords.ts b/ui/litellm-dashboard/src/components/add_model/complexity_router_keywords.ts index 9cfdaed4e23..6fe93cddae3 100644 --- a/ui/litellm-dashboard/src/components/add_model/complexity_router_keywords.ts +++ b/ui/litellm-dashboard/src/components/add_model/complexity_router_keywords.ts @@ -19,13 +19,19 @@ const asKeywords = (value: unknown): string[] => : []; /** - * Drop the React-only id, trim keywords, and discard rules left empty. "Add keyword rule" - * seeds a row with no keywords, and the backend validator rejects those with a 400. + * Drop the React-only id and trim keywords, leaving one entry per rule. A rule left empty stays + * empty rather than disappearing, so getKeywordTierRulesError can name the row it came from. */ export const serializeKeywordTierRules = (rules: KeywordTierRule[]): StoredKeywordTierRule[] => - rules - .map((rule) => ({ keywords: asKeywords(rule.keywords).filter(Boolean), tier: rule.tier })) - .filter((rule) => rule.keywords.length > 0); + rules.map((rule) => ({ keywords: asKeywords(rule.keywords).filter(Boolean), tier: rule.tier })); + +/** + * Positions of the rules left without a keyword, as indexes into the caller's own array. The + * submit-time message and the inline error on the row both read this, so the row the message + * names is always the row that lights up. + */ +export const emptyKeywordTierRuleIndexes = (rules: KeywordTierRule[]): number[] => + serializeKeywordTierRules(rules).flatMap((rule, index) => (rule.keywords.length === 0 ? [index] : [])); export const hydrateKeywordTierRules = (value: unknown): KeywordTierRule[] => { if (!Array.isArray(value)) return []; diff --git a/ui/litellm-dashboard/src/components/common_components/MetadataKeyValueFields.test.tsx b/ui/litellm-dashboard/src/components/common_components/MetadataKeyValueFields.test.tsx new file mode 100644 index 00000000000..365cc8bbff5 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/MetadataKeyValueFields.test.tsx @@ -0,0 +1,282 @@ +import { render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { Form } from "antd"; +import React from "react"; +import { describe, expect, it, vi } from "vitest"; +import { TeamMetadataField } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema"; +import MetadataKeyValueFields, { + MetadataPair, + metadataObjectToPairs, + metadataPairsToObject, +} from "./MetadataKeyValueFields"; + +describe("metadataObjectToPairs", () => { + it("returns an empty list for null or undefined metadata", () => { + expect(metadataObjectToPairs(null)).toEqual([]); + expect(metadataObjectToPairs(undefined)).toEqual([]); + }); + + it("keeps plain string values as-is", () => { + expect(metadataObjectToPairs({ department: "research" })).toEqual([{ key: "department", value: "research" }]); + }); + + it("serializes non-string values as JSON", () => { + expect( + metadataObjectToPairs({ + tier: 3, + beta: true, + config: { region: "us" }, + tags: ["a", "b"], + empty: null, + }), + ).toEqual([ + { key: "tier", value: "3" }, + { key: "beta", value: "true" }, + { key: "config", value: '{"region":"us"}' }, + { key: "tags", value: '["a","b"]' }, + { key: "empty", value: "null" }, + ]); + }); + + it("quotes string values that would otherwise parse as JSON, so types round-trip", () => { + expect(metadataObjectToPairs({ code: "42", flag: "true" })).toEqual([ + { key: "code", value: '"42"' }, + { key: "flag", value: '"true"' }, + ]); + }); + + it("filters out excluded keys", () => { + expect( + metadataObjectToPairs({ department: "research", logging: [{ callback_name: "langfuse" }] }, new Set(["logging"])), + ).toEqual([{ key: "department", value: "research" }]); + }); +}); + +describe("metadataPairsToObject", () => { + it("returns an empty object for undefined pairs", () => { + expect(metadataPairsToObject(undefined)).toEqual({}); + }); + + it("keeps plain text values as strings", () => { + expect(metadataPairsToObject([{ key: "department", value: "research" }])).toEqual({ department: "research" }); + }); + + it("parses JSON values into their typed form", () => { + expect( + metadataPairsToObject([ + { key: "tier", value: "3" }, + { key: "beta", value: "true" }, + { key: "config", value: '{"region":"us"}' }, + { key: "code", value: '"42"' }, + ]), + ).toEqual({ tier: 3, beta: true, config: { region: "us" }, code: "42" }); + }); + + it("skips rows without a key and defaults a missing value to an empty string", () => { + expect(metadataPairsToObject([{ key: "", value: "orphan" }, undefined, { key: "kept" }])).toEqual({ kept: "" }); + }); + + it("round-trips a mixed-type metadata object losslessly", () => { + const metadata = { + department: "research", + code: "42", + tier: 3, + beta: true, + config: { region: "us", replicas: 2 }, + }; + expect(metadataPairsToObject(metadataObjectToPairs(metadata))).toEqual(metadata); + }); +}); + +interface HarnessProps { + onFinish: (values: { metadata?: MetadataPair[] }) => void; + initialMetadata?: MetadataPair[]; + schemaFields?: TeamMetadataField[]; + schemaLoading?: boolean; +} + +const Harness: React.FC = ({ onFinish, initialMetadata, schemaFields, schemaLoading }) => { + const [form] = Form.useForm(); + return ( +
+ + + + ); +}; + +describe("MetadataKeyValueFields", () => { + it("renders one row per existing pair", () => { + render( + , + ); + + const keyInputs = screen.getAllByPlaceholderText("Key"); + const valueInputs = screen.getAllByPlaceholderText("Value"); + expect(keyInputs.map((input) => (input as HTMLInputElement).value)).toEqual(["department", "tier"]); + expect(valueInputs.map((input) => (input as HTMLInputElement).value)).toEqual(["research", "3"]); + }); + + it("adds a row and submits the entered pair", async () => { + const user = userEvent.setup(); + const onFinish = vi.fn(); + render(); + + await user.click(screen.getByRole("button", { name: /add key-value pair/i })); + await user.type(screen.getByPlaceholderText("Key"), "cost_center"); + await user.type(screen.getByPlaceholderText("Value"), "eng-1"); + await user.click(screen.getByRole("button", { name: "Save" })); + + await waitFor(() => { + expect(onFinish).toHaveBeenCalledWith({ metadata: [{ key: "cost_center", value: "eng-1" }] }); + }); + }); + + it("removes a row when its remove icon is clicked", async () => { + const user = userEvent.setup(); + const onFinish = vi.fn(); + render( + , + ); + + await user.click(screen.getAllByLabelText("Remove key-value pair")[0]); + await user.click(screen.getByRole("button", { name: "Save" })); + + await waitFor(() => { + expect(onFinish).toHaveBeenCalledWith({ metadata: [{ key: "tier", value: "3" }] }); + }); + }); + + it("blocks submission on duplicate keys", async () => { + const user = userEvent.setup(); + const onFinish = vi.fn(); + render( + , + ); + + await user.click(screen.getByRole("button", { name: "Save" })); + + await waitFor(() => { + expect(screen.getAllByText("Duplicate key").length).toBeGreaterThan(0); + }); + expect(onFinish).not.toHaveBeenCalled(); + }); + + it("blocks submission when a row is missing its key", async () => { + const user = userEvent.setup(); + const onFinish = vi.fn(); + render(); + + await user.click(screen.getByRole("button", { name: /add key-value pair/i })); + await user.type(screen.getByPlaceholderText("Value"), "orphan"); + await user.click(screen.getByRole("button", { name: "Save" })); + + await waitFor(() => { + expect(screen.getByText("Missing key")).toBeInTheDocument(); + }); + expect(onFinish).not.toHaveBeenCalled(); + }); +}); + +describe("MetadataKeyValueFields with a declared schema", () => { + const schema: TeamMetadataField[] = [ + { key: "cost_center", label: "Cost Center" }, + { key: "app_name", label: "Application Name" }, + ]; + + it("should prepopulate one ordinary editable pair row per declared key", async () => { + render(); + + await waitFor(() => { + expect(screen.getAllByPlaceholderText("Key").map((input) => (input as HTMLInputElement).value)).toEqual([ + "cost_center", + "app_name", + ]); + }); + screen.getAllByPlaceholderText("Key").forEach((input) => expect(input).toBeEnabled()); + expect(screen.getAllByLabelText("Remove key-value pair")).toHaveLength(2); + }); + + it("should submit a prepopulated key with its typed value", async () => { + const user = userEvent.setup(); + const onFinish = vi.fn(); + render(); + + await user.type(await screen.findByPlaceholderText("Value"), "CC-1001"); + await user.click(screen.getByRole("button", { name: "Save" })); + + await waitFor(() => { + expect(onFinish).toHaveBeenCalledWith({ metadata: [{ key: "cost_center", value: "CC-1001" }] }); + }); + }); + + it("should not add a second row for keys already present in the form", async () => { + render( + , + ); + + await waitFor(() => { + expect(screen.getAllByPlaceholderText("Key").map((input) => (input as HTMLInputElement).value)).toEqual([ + "cost_center", + "app_name", + ]); + }); + expect(screen.getAllByPlaceholderText("Value").map((input) => (input as HTMLInputElement).value)).toEqual([ + "CC-1001", + "", + ]); + }); + + it("should let the user remove a prepopulated row", async () => { + const user = userEvent.setup(); + render(); + + await screen.findAllByPlaceholderText("Key"); + await user.click(screen.getAllByLabelText("Remove key-value pair")[0]); + + await waitFor(() => { + expect(screen.getAllByPlaceholderText("Key").map((input) => (input as HTMLInputElement).value)).toEqual([ + "app_name", + ]); + }); + }); + + it("should show a skeleton instead of the editor while the schema is loading", () => { + render(); + + expect(screen.getByTestId("metadata-schema-skeleton")).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /add key-value pair/i })).not.toBeInTheDocument(); + }); + + it("should seed rows when the schema arrives after an initial loading state", async () => { + const onFinish = vi.fn(); + const { rerender } = render(); + + rerender(); + + await waitFor(() => { + expect(screen.getAllByPlaceholderText("Key").map((input) => (input as HTMLInputElement).value)).toEqual([ + "cost_center", + "app_name", + ]); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/MetadataKeyValueFields.tsx b/ui/litellm-dashboard/src/components/common_components/MetadataKeyValueFields.tsx new file mode 100644 index 00000000000..da085f95ad8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/MetadataKeyValueFields.tsx @@ -0,0 +1,135 @@ +import { MinusCircleOutlined, PlusOutlined } from "@ant-design/icons"; +import { Button, Form, FormInstance, Input, Skeleton, Space } from "antd"; +import React, { useEffect, useRef } from "react"; + +import { TeamMetadataField } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema"; + +export interface MetadataPair { + key: string; + value: string; +} + +function formatMetadataValue(value: unknown): string { + if (typeof value !== "string") { + return JSON.stringify(value) ?? ""; + } + try { + JSON.parse(value); + return JSON.stringify(value); + } catch { + return value; + } +} + +function parseMetadataValue(raw: string): unknown { + try { + return JSON.parse(raw); + } catch { + return raw; + } +} + +export function metadataObjectToPairs( + metadata: Record | null | undefined, + excludedKeys: ReadonlySet = new Set(), +): MetadataPair[] { + return Object.entries(metadata ?? {}) + .filter(([key]) => !excludedKeys.has(key)) + .map(([key, value]) => ({ key, value: formatMetadataValue(value) })); +} + +export function metadataPairsToObject( + pairs: readonly (Partial | undefined)[] | undefined, +): Record { + return Object.fromEntries( + (pairs ?? []) + .filter((pair): pair is Partial & { key: string } => Boolean(pair?.key)) + .map((pair) => [pair.key, parseMetadataValue(pair.value ?? "")]), + ); +} + +interface MetadataKeyValueFieldsProps { + form: FormInstance; + name?: string; + schemaFields?: readonly TeamMetadataField[]; + schemaLoading?: boolean; +} + +const MetadataKeyValueFields: React.FC = ({ + form, + name = "metadata", + schemaFields = [], + schemaLoading = false, +}) => { + const seededRef = useRef(false); + + useEffect(() => { + if (seededRef.current || schemaLoading || schemaFields.length === 0) return; + seededRef.current = true; + const pairs: (Partial | undefined)[] = form.getFieldValue(name) ?? []; + if (!Array.isArray(pairs)) return; + const existingKeys = new Set(pairs.map((pair) => pair?.key).filter(Boolean)); + const seeded = schemaFields + .filter((field) => !existingKeys.has(field.key)) + .map((field) => ({ key: field.key, value: "" })); + if (seeded.length > 0) { + form.setFieldValue(name, [...pairs, ...seeded]); + } + }, [form, name, schemaFields, schemaLoading]); + + if (schemaLoading) { + return ( +
+ +
+ ); + } + + return ( + + {(fields, { add, remove }) => ( + <> + {fields.map(({ key, name: fieldName, ...restField }) => ( + + { + if (!value) return Promise.resolve(); + const all: (Partial | undefined)[] = form.getFieldValue(name) ?? []; + const dupes = all.filter((entry) => entry?.key === value); + if (dupes.length > 1) { + return Promise.reject(new Error("Duplicate key")); + } + return Promise.resolve(); + }, + }, + ]} + > + + + + + + remove(fieldName)} + style={{ color: "#ef4444" }} + /> + + ))} + + + + + )} + + ); +}; + +export default MetadataKeyValueFields; diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index 971c833a0de..818dcd1f648 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -54,13 +54,16 @@ describe("buildUpdatedComplexityRouterConfig keyword matching", () => { expect(result.keyword_tier_rules).toEqual([{ keywords: ["chargeback"], tier: "COMPLEX" }]); }); - it("drops a rule left empty rather than shipping one the backend 400s on", () => { + // getKeywordTierRulesError blocks this save, so the builder never runs on a real edit. Keeping + // the rule here means that if a caller ever reaches it anyway, the stored rules are replaced by + // something the backend rejects out loud rather than by silence that reads as a clean save. + it("keeps a rule left empty rather than quietly dropping the caller's row", () => { const result = buildUpdatedComplexityRouterConfig(STORED, FORM_VALUE, undefined, { ...hydratedState, keywordTierRules: [{ id: "new-1", keywords: [" "], tier: "SIMPLE" }], }); - expect(result.keyword_tier_rules).toBeUndefined(); + expect(result.keyword_tier_rules).toEqual([{ keywords: [], tier: "SIMPLE" }]); }); it("removes the semantic trio when the toggle is turned off", () => { diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx index c0806befa52..3976e4c1381 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx @@ -118,6 +118,80 @@ describe("EditAutoRouterModal keyword matching", () => { await waitFor(() => expect(NotificationsManager.fromBackend).toHaveBeenCalled()); expect(modelPatchUpdateCall).not.toHaveBeenCalled(); }); + + // LIT-5133, edit side. Semantic matching is off here on purpose: it used to be the only thing + // that checked a rule for keywords, so with it on this save was already blocked and the test + // would pass without the fix. Off, the unfilled row was dropped and the save reported success. + it("blocks a save that adds a keyword rule and leaves it empty", async () => { + const user = userEvent.setup(); + renderWithProviders( + , + ); + + await screen.findByText(/Escalation Keywords/i); + fireEvent.click(screen.getByText("Advanced: Keyword/Semantic Matching")); + await user.click(screen.getByRole("button", { name: /add keyword rule/i })); + + // The modal renders the same controls as the create form, so it owes the same treatment: + // the row says what is missing and the save is not offered while it is. + expect(await screen.findByText("At least one keyword is required")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /save changes/i })).toBeDisabled(); + expect(modelPatchUpdateCall).not.toHaveBeenCalled(); + }); + + it("gives the save back once the added keyword rule is filled", async () => { + const user = userEvent.setup(); + renderWithProviders( + , + ); + + await screen.findByText(/Escalation Keywords/i); + fireEvent.click(screen.getByText("Advanced: Keyword/Semantic Matching")); + await user.click(screen.getByRole("button", { name: /add keyword rule/i })); + expect(screen.getByRole("button", { name: /save changes/i })).toBeDisabled(); + + await user.type( + within(screen.getByText("Keywords 2").closest("div") as HTMLElement).getByRole("combobox"), + "chargeback{enter}", + ); + + expect(screen.getByRole("button", { name: /save changes/i })).toBeEnabled(); + expect(screen.queryByText("At least one keyword is required")).not.toBeInTheDocument(); + }); }); describe("EditAutoRouterModal classifier context window", () => { diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index a70fc31d6fe..2c58cd70cb9 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -1,12 +1,12 @@ import React, { useEffect, useState } from "react"; -import { Modal, Form, Button, Select as AntdSelect } from "antd"; +import { Modal, Form, Button, Select as AntdSelect, Tooltip } from "antd"; import { Text, TextInput } from "@tremor/react"; import { modelAvailableCall, modelPatchUpdateCall } from "../networking"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; import RouterConfigBuilder from "../add_model/RouterConfigBuilder"; import { normalizeTierModels } from "../add_model/complexity_router_tiers"; import { isComplexityRouter } from "../add_model/auto_router_strategies"; -import { getSemanticConfigError } from "../add_model/build_complexity_router_config"; +import { getKeywordTierRulesError, getSemanticConfigError } from "../add_model/build_complexity_router_config"; import { KeywordTierRule } from "../add_model/KeywordTierRules"; import { DEFAULT_MATCH_THRESHOLD } from "../add_model/SemanticKeywordMatching"; import { hydrateKeywordTierRules, serializeKeywordTierRules } from "../add_model/complexity_router_keywords"; @@ -118,8 +118,8 @@ export const buildUpdatedComplexityRouterConfig = ( }), ...(value.return_raw_model_name && { return_raw_model_name: true }), ...(keywordMatching && { - // Mirrors buildComplexityRouterConfig: rules only when non-empty (the backend rejects - // an empty rule with a 400), escalation keywords always, semantic trio only when on. + // Mirrors buildComplexityRouterConfig: the key only when there is a rule to write, + // escalation keywords always, semantic trio only when on. ...(storedKeywordRules.length > 0 && { keyword_tier_rules: storedKeywordRules }), escalation_keywords: keywordMatching.escalationKeywords.map((k) => k.trim()).filter(Boolean), ...(keywordMatching.semanticMatchingEnabled && { @@ -145,6 +145,7 @@ const EditAutoRouterModal: React.FC = ({ const [modelInfo, setModelInfo] = useState([]); const [showCustomDefaultModel, setShowCustomDefaultModel] = useState(false); const [showCustomEmbeddingModel, setShowCustomEmbeddingModel] = useState(false); + const [showValidationErrors, setShowValidationErrors] = useState(false); const [routerConfig, setRouterConfig] = useState(null); const [customTechnicalKeywords, setCustomTechnicalKeywords] = useState([]); const [keywordTierRules, setKeywordTierRules] = useState([]); @@ -158,6 +159,15 @@ const EditAutoRouterModal: React.FC = ({ }); const isComplexityRouterModel = isComplexityRouter(modelData?.litellm_params); + // Mirrors the create form: the button says why it is unavailable and disables on the same + // answer. Tiers use this modal's own rule, which allows a partly filled router, so an edit that + // is legal today stays legal. + const submitBlockedReason = !isComplexityRouterModel + ? null + : (Object.values(complexityRouterConfig.tiers).every((models) => models.length === 0) + ? "Please select at least one model for a complexity tier" + : null) ?? getKeywordTierRulesError(keywordTierRules); + useEffect(() => { if (isVisible && modelData) { initializeForm(); @@ -295,24 +305,29 @@ const EditAutoRouterModal: React.FC = ({ if (isComplexityRouterModel) { const { tiers, classifier_type, classifier_llm_config } = complexityRouterConfig; if (Object.values(tiers).every((models) => models.length === 0)) { + setShowValidationErrors(true); NotificationsManager.fromBackend("Please select at least one model for a complexity tier"); return; } if (classifier_type === "llm" && !classifier_llm_config?.model) { + setShowValidationErrors(true); NotificationsManager.fromBackend("Please select a classifier model, or switch back to Heuristic"); return; } - // Same guard the create form applies (add_auto_router_tab.tsx). The backend rejects - // semantic_keyword_matching without an embedding model or keyword rules - // (complexity_router/config.py), so without this a save fails as a raw 400 instead of - // an inline message. + // Same guards the create form applies (add_auto_router_tab.tsx). The backend rejects a + // keyword rule with no keyword, and semantic_keyword_matching without an embedding model + // or keyword rules (complexity_router/config.py), so without these a save fails as a raw + // 400 instead of an inline message. + const keywordRulesError = getKeywordTierRulesError(keywordTierRules); + if (keywordRulesError) { + setShowValidationErrors(true); + NotificationsManager.fromBackend(keywordRulesError); + return; + } - // Same guard the create form applies (add_auto_router_tab.tsx). The backend rejects - // semantic_keyword_matching without an embedding model or keyword rules - // (complexity_router/config.py), so without this a save fails as a raw 400 instead of - // an inline message. const semanticError = getSemanticConfigError({ semanticMatchingEnabled, embeddingModel, keywordTierRules }); if (semanticError) { + setShowValidationErrors(true); NotificationsManager.fromBackend(semanticError); return; } @@ -410,9 +425,11 @@ const EditAutoRouterModal: React.FC = ({ , - , + + + , ]} width={1000} destroyOnHidden @@ -436,6 +453,7 @@ const EditAutoRouterModal: React.FC = ({ /* Complexity Router Configuration */
{ diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 03cf0e9583c..5a2d33ee4bd 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -38,7 +38,7 @@ import type { CoordinationRedisTestResponse, } from "@/app/(dashboard)/caching/_components/coordination_redis_settings/types"; import { MCP_TOOLS_PREVIEW_FORBIDDEN_MESSAGE } from "./mcp_tools/constants"; -import { createApiClient, deriveErrorMessage } from "@/lib/http/client"; +import { createApiClient, deriveErrorMessage, unwrapProxyErrorMessage } from "@/lib/http/client"; import { resolveApiBase } from "@/lib/http/resolveApiBase"; import { registerAuthHeaderNameGetter, @@ -2643,7 +2643,7 @@ export const teamUpdateCall = async ( const errorData = await response.text(); handleError(errorData); console.error("Error response from the server:", errorData); - NotificationsManager.fromBackend("Failed to update team settings: " + errorData); + NotificationsManager.fromBackend("Failed to update team settings: " + unwrapProxyErrorMessage(errorData)); throw new Error(errorData); } const data = (await response.json()) as { data: Team; team_id: string }; diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index 712cff80649..513719a2ad9 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -1,3 +1,4 @@ +import { useTeamMetadataSchema } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema"; import * as networking from "@/components/networking"; import { screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; @@ -26,6 +27,10 @@ vi.mock("@/components/utils/dataUtils", () => ({ formatNumberWithCommas: vi.fn((value: number) => value.toLocaleString()), })); +vi.mock("@/app/(dashboard)/hooks/teams/useTeamMetadataSchema", () => ({ + useTeamMetadataSchema: vi.fn(() => ({ data: [], isLoading: false })), +})); + vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ useAllProxyModels: vi.fn(), })); @@ -220,6 +225,7 @@ describe("TeamInfoView", () => { isFetching: false, refetch: vi.fn(), } as any); + vi.mocked(useTeamMetadataSchema).mockReturnValue({ data: [], isLoading: false } as any); vi.mocked(networking.getGuardrailsList).mockResolvedValue({ guardrails: [] }); vi.mocked(networking.getPoliciesList).mockResolvedValue({ policies: [] }); @@ -893,6 +899,137 @@ describe("TeamInfoView", () => { }); }); + describe("metadata key-value editing", () => { + const openSettingsEditor = async (user: ReturnType) => { + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + await user.click(screen.getByRole("tab", { name: "Settings" })); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + }); + + await user.click(screen.getByRole("button", { name: /edit settings/i })); + + await waitFor(() => { + expect(screen.getByLabelText("Team Name")).toBeInTheDocument(); + }); + }; + + it("prefills pairs from team metadata, hides UI-managed keys, and round-trips typed values on save", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + metadata: { + department: "research", + tier: 3, + beta: true, + config: { region: "us" }, + logging: [{ callback_name: "langfuse", callback_type: "success", callback_vars: {} }], + guardrails: ["g1"], + disable_global_guardrails: false, + model_tpm_limit: { "gpt-4": 100 }, + }, + models: ["gpt-4"], + }), + ); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + await openSettingsEditor(user); + + const keyValues = screen.getAllByPlaceholderText("Key").map((input) => (input as HTMLInputElement).value); + expect(keyValues).toEqual(["department", "tier", "beta", "config"]); + const valueValues = screen.getAllByPlaceholderText("Value").map((input) => (input as HTMLInputElement).value); + expect(valueValues).toEqual(["research", "3", "true", '{"region":"us"}']); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalled(); + }); + + const updateArg = vi.mocked(networking.teamUpdateCall).mock.calls[0][1]; + expect(updateArg.metadata).toMatchObject({ + department: "research", + tier: 3, + beta: true, + config: { region: "us" }, + logging: [{ callback_name: "langfuse", callback_type: "success", callback_vars: {} }], + }); + expect(updateArg.metadata).not.toHaveProperty("model_tpm_limit"); + expect(updateArg.model_tpm_limit).toEqual({ "gpt-4": 100 }); + }); + + it("includes a newly added pair in the team update", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: ["gpt-4"] })); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + await openSettingsEditor(user); + + await user.click(screen.getByRole("button", { name: /add key-value pair/i })); + await user.type(screen.getByPlaceholderText("Key"), "cost_center"); + await user.type(screen.getByPlaceholderText("Value"), "eng-1"); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalled(); + }); + + expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].metadata).toMatchObject({ cost_center: "eng-1" }); + }); + + it("should keep declared keys as ordinary prefilled rows and submit the edited value", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(useTeamMetadataSchema).mockReturnValue({ + data: [ + { key: "cost_center", label: "Cost Center" }, + { key: "app_name", label: "Application Name" }, + ], + isLoading: false, + } as any); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + metadata: { cost_center: "CC-OLD", department: "research" }, + models: ["gpt-4"], + }), + ); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + await openSettingsEditor(user); + + await waitFor(() => { + expect(screen.getAllByPlaceholderText("Key").map((input) => (input as HTMLInputElement).value)).toEqual([ + "cost_center", + "department", + "app_name", + ]); + }); + expect(screen.getAllByPlaceholderText("Value")[0]).toHaveValue("CC-OLD"); + + await user.clear(screen.getAllByPlaceholderText("Value")[0]); + await user.type(screen.getAllByPlaceholderText("Value")[0], "CC-NEW"); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalled(); + }); + + expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].metadata).toMatchObject({ + cost_center: "CC-NEW", + department: "research", + app_name: "", + }); + }); + }); + describe("model aliases", () => { const openSettingsEditor = async (user: ReturnType) => { await waitFor(() => { diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 34570043f52..bbe5dc05a88 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -35,6 +35,11 @@ import { CheckIcon, CopyIcon } from "lucide-react"; import React, { useEffect, useMemo, useState } from "react"; import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils"; import AccessGroupSelector from "../common_components/AccessGroupSelector"; +import MetadataKeyValueFields, { + metadataObjectToPairs, + metadataPairsToObject, +} from "../common_components/MetadataKeyValueFields"; +import { useTeamMetadataSchema } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema"; import ModelAliasManager from "../common_components/ModelAliasManager"; import AgentSelector from "../agent_management/AgentSelector"; import DeleteResourceModal from "../common_components/DeleteResourceModal"; @@ -66,6 +71,18 @@ import { import TeamMembersComponent from "./TeamMemberTab"; import { TeamVirtualKeysTable } from "./TeamVirtualKeysTable"; +const UI_MANAGED_METADATA_KEYS: ReadonlySet = new Set([ + "logging", + "secret_manager_settings", + "soft_budget_alerting_emails", + "model_tpm_limit", + "model_rpm_limit", + "allowed_passthrough_routes", + "guardrails", + "opted_out_global_guardrails", + "disable_global_guardrails", +]); + export interface TeamMembership { user_id: string; team_id: string; @@ -203,6 +220,7 @@ const TeamInfoView: React.FC = ({ const [organization, setOrganization] = useState(null); const { userRole, userId } = useAuthorized(); const { data: userOrganizations = [] } = useOrganizations(); + const { data: teamMetadataSchemaFields = [], isLoading: isTeamMetadataSchemaLoading } = useTeamMetadataSchema(); const queryClient = useQueryClient(); // Check if user is org admin for this team's organization @@ -461,16 +479,7 @@ const TeamInfoView: React.FC = ({ if (!accessToken) return; setIsTeamSaving(true); - let parsedMetadata = {}; - try { - const rawMetadata = values.metadata ? JSON.parse(values.metadata) : {}; - // Exclude soft_budget_alerting_emails from parsed metadata since it's handled separately - const { soft_budget_alerting_emails, ...rest } = rawMetadata; - parsedMetadata = rest; - } catch (e) { - NotificationsManager.fromBackend("Invalid JSON in metadata field"); - return; - } + const parsedMetadata = metadataPairsToObject(values.metadata); let secretManagerSettings: Record | undefined; if (typeof values.secret_manager_settings === "string") { @@ -980,21 +989,7 @@ const TeamInfoView: React.FC = ({ soft_budget_alerting_emails: Array.isArray(info.metadata?.soft_budget_alerting_emails) ? info.metadata.soft_budget_alerting_emails.join(", ") : "", - metadata: info.metadata - ? JSON.stringify( - (({ - logging, - secret_manager_settings, - soft_budget_alerting_emails, - model_tpm_limit, - model_rpm_limit, - allowed_passthrough_routes, - ...rest - }) => rest)(info.metadata), - null, - 2, - ) - : "", + metadata: metadataObjectToPairs(info.metadata, UI_MANAGED_METADATA_KEYS), logging_settings: info.metadata?.logging || [], secret_manager_settings: info.metadata?.secret_manager_settings ? JSON.stringify(info.metadata.secret_manager_settings, null, 2) @@ -1170,6 +1165,17 @@ const TeamInfoView: React.FC = ({ + + + + = ({ /> - - - -