scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-14 04:09:09 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 139

Merge remote-tracking branch 'origin/main' into litellm_test_d6496f

Browse source
This commit is contained in:
Yuneng Jiang 2026-10-10 02:45:53 -07:00
commit 700f615bc8
No known key found for this signature in database
37 changed files with 4993 additions and 385 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

72
enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py
View file

@ -12,7 +12,7 @@ Endpoints for /project operations
import json
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final
from typing import TYPE_CHECKING, Final, cast
from fastapi import APIRouter, Depends, HTTPException, Request
from pydantic import TypeAdapter
@ -31,6 +31,7 @@ from litellm.proxy.management_helpers.utils import (
management_endpoint_wrapper,
)
from litellm.proxy.utils import PrismaClient, handle_exception_on_proxy
from litellm.repositories.base_repository import record_to_dict
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
from litellm.repositories.prisma_protocols import TableActions
@ -57,6 +58,15 @@ def _project_table(prisma_client: PrismaClient) -> TableActions["prisma_models.L
return ProjectRepository(prisma_client).table
def _writer_project_table(
prisma_client: PrismaClient,
) -> TableActions["prisma_models.LiteLLM_ProjectTable"]:
return cast( # cast-ok: writer_db exposes generated Prisma tables through a dynamic wrapper
"TableActions[prisma_models.LiteLLM_ProjectTable]",
prisma_client.writer_db.litellm_projecttable,
)
def _verification_token_table(
prisma_client: PrismaClient,
) -> TableActions["prisma_models.LiteLLM_VerificationToken"]:
@ -824,6 +834,38 @@ async def update_project(
**({"max_budget": None} if "max_budget" in data.model_fields_set and data.max_budget is None else {}),
}
object_permission_data: Final = update_data.pop("object_permission", None)
object_permission_payload: Final = (
_OBJECT_PERMISSION_PAYLOAD.validate_python(object_permission_data) if object_permission_data else None
)
if data.team_id is not None:
current_project_record: Final = await _writer_project_table(prisma_client).find_unique(
where={"project_id": data.project_id}
)
current_project: Final = (
LiteLLM_ProjectTable.model_validate(record_to_dict(current_project_record))
if current_project_record is not None
else None
)
if current_project is not None and data.team_id != current_project.team_id:
mismatched_key_count: Final = await prisma_client.writer_db.litellm_verificationtoken.count(
where={
"project_id": data.project_id,
"OR": [{"team_id": {"not": data.team_id}}, {"team_id": None}],
}
)
if mismatched_key_count > 0:
raise ProxyException(
message=(
f"Project {data.project_id} has {mismatched_key_count} key(s) that do not belong to "
f"team {data.team_id}. Detach or delete them before moving the project."
),
type="bad_request",
code=400,
param="team_id",
)
if budget_updates and existing_project.budget_id:
# Update existing budget
await _budget_table(prisma_client).update(
@ -837,23 +879,17 @@ async def update_project(
for field in budget_updates.keys():
update_data.pop(field, None)
# Handle object permissions
if "object_permission" in update_data:
object_permission_data = update_data.pop("object_permission")
if object_permission_data:
object_permission_payload: Final = _OBJECT_PERMISSION_PAYLOAD.validate_python(object_permission_data)
if existing_project.object_permission_id:
# Update existing permission
await _object_permission_table(prisma_client).update(
where={"object_permission_id": existing_project.object_permission_id},
data=object_permission_payload,
)
else:
# Create new permission
created_permission: Final = await _object_permission_table(prisma_client).create(
data=object_permission_payload,
)
update_data["object_permission_id"] = created_permission.object_permission_id
if object_permission_payload is not None:
if existing_project.object_permission_id:
await _object_permission_table(prisma_client).update(
where={"object_permission_id": existing_project.object_permission_id},
data=object_permission_payload,
)
else:
created_permission: Final = await _object_permission_table(prisma_client).create(
data=object_permission_payload,
)
update_data["object_permission_id"] = created_permission.object_permission_id
# Handle metadata fields
for field in LiteLLM_ManagementEndpoint_MetadataFields:

6
litellm/integrations/custom_guardrail.py
View file

@ -223,6 +223,7 @@ class CustomGuardrail(CustomLogger):
records_own_guardrail_information: ClassVar[bool] = False
logging_only_scope: LoggingOnlyScope | None
logging_only_continue_on_input_failure: bool
stream_scope_default: GuardrailStreamScope = DEFAULT_GUARDRAIL_STREAM_SCOPE
stream_scope_by_hook: tuple[tuple[str, GuardrailStreamScope], ...] = ()
@ -305,6 +306,7 @@ class CustomGuardrail(CustomLogger):
stream_scope_arg: Final[object] = cast(object, kwargs.pop("stream_scope", None)) # cast-ok: config
self.apply_stream_scope(stream_scope_arg)
self.logging_only_scope = None
self.logging_only_continue_on_input_failure = False
if timeout is not None:
self.timeout = timeout
@ -1064,7 +1066,7 @@ class CustomGuardrail(CustomLogger):
except Exception as e:
if self.logging_only_scope == "output":
return None
if self.logging_only_scope == "both":
if self.logging_only_continue_on_input_failure:
verbose_logger.warning(
"Guardrail %s: logging_only request copy failed, skipping request scan: %s",
self.guardrail_name,
@ -1092,7 +1094,7 @@ class CustomGuardrail(CustomLogger):
"metadata": scratch_metadata,
}
if self.logging_only_scope != "output" and scratch_fields is not None:
if self.logging_only_scope == "both":
if self.logging_only_continue_on_input_failure:
try:
await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self)
except Exception as e: # noqa: BLE001 # one direction's scan failure must not drop the other direction's verdict

40
litellm/model_prices_and_context_window_backup.json
View file

@ -71786,5 +71786,45 @@
"supported_endpoints": [
"/v1/responses"
]
},
"azure_ai/Cohere-Embed-V5-Fast": {
"deprecation_date": "2027-01-23",
"input_cost_per_image_token": 4e-07,
"input_cost_per_token": 8e-08,
"litellm_provider": "azure_ai",
"max_input_tokens": 128000,
"max_tokens": 128000,
"mode": "embedding",
"output_cost_per_token": 0.0,
"output_vector_size": 2048.0,
"source": "https://prices.azure.com/api/retail/prices?$filter=productName%20eq%20'Cohere%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supported_endpoints": [
"/v1/embeddings"
],
"supported_modalities": [
"text",
"image"
],
"supports_embedding_image_input": true
},
"azure_ai/Cohere-Embed-V5-Pro": {
"deprecation_date": "2027-01-23",
"input_cost_per_image_token": 4e-07,
"input_cost_per_token": 1.2e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 128000,
"max_tokens": 128000,
"mode": "embedding",
"output_cost_per_token": 0.0,
"output_vector_size": 2048.0,
"source": "https://prices.azure.com/api/retail/prices?$filter=productName%20eq%20'Cohere%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supported_endpoints": [
"/v1/embeddings"
],
"supported_modalities": [
"text",
"image"
],
"supports_embedding_image_input": true
}
}

34
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -11620,13 +11620,24 @@
"description": "Google Cloud location/region (e.g., us-central1)",
"title": "Location"
},
"logging_only_continue_on_input_failure": {
"anyOf": [
{
"type": "boolean"
},
{
"type": "null"
}
],
"description": "when True, a flagged or raising logging_only request scan is logged and the response is still scanned, so both verdicts land. Only applies to mode logging_only and is ignored when logging_only_scope is 'input' or 'output'.",
"title": "Logging Only Continue On Input Failure"
},
"logging_only_scope": {
"anyOf": [
{
"enum": [
"input",
"output",
"both"
"output"
],
"type": "string"
},
@ -11634,7 +11645,7 @@
"type": "null"
}
],
"description": "which direction a logging_only scan observes: 'input' (request), 'output' (response), or 'both' (default). Only applies to mode logging_only; pre_call/post_call on the same guardrail keep blocking.",
"description": "which direction a logging_only scan observes: 'input' (request) or 'output' (response); unset scans both directions. Only applies to mode logging_only; pre_call/post_call on the same guardrail keep blocking.",
"title": "Logging Only Scope"
},
"mask_request_content": {
@ -14027,13 +14038,24 @@
"description": "Google Cloud location/region (e.g., us-central1)",
"title": "Location"
},
"logging_only_continue_on_input_failure": {
"anyOf": [
{
"type": "boolean"
},
{
"type": "null"
}
],
"description": "when True, a flagged or raising logging_only request scan is logged and the response is still scanned, so both verdicts land. Only applies to mode logging_only and is ignored when logging_only_scope is 'input' or 'output'.",
"title": "Logging Only Continue On Input Failure"
},
"logging_only_scope": {
"anyOf": [
{
"enum": [
"input",
"output",
"both"
"output"
],
"type": "string"
},
@ -14041,7 +14063,7 @@
"type": "null"
}
],
"description": "which direction a logging_only scan observes: 'input' (request), 'output' (response), or 'both' (default). Only applies to mode logging_only; pre_call/post_call on the same guardrail keep blocking.",
"description": "which direction a logging_only scan observes: 'input' (request) or 'output' (response); unset scans both directions. Only applies to mode logging_only; pre_call/post_call on the same guardrail keep blocking.",
"title": "Logging Only Scope"
},
"mask": {

27
litellm/proxy/auth/auth_utils.py
View file

@ -5,6 +5,7 @@ import sys
from collections.abc import Collection, Iterator, Mapping
from functools import lru_cache
from logging import Logger
from types import MappingProxyType
from typing import Any, Final, Protocol
from fastapi import HTTPException, Request, status
@ -37,7 +38,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
)
from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, Deployment, server_owned_wif_fields_present
from litellm.types.router import reject_server_owned_wif_params as _reject_server_owned_wif_params
from litellm.types.utils import CustomPricingLiteLLMParams
from litellm.types.utils import CustomPricingLiteLLMParams, oauth_token_exchange_litellm_params
def is_invalid_virtual_key_error(exception: BaseException | None) -> bool:
@ -279,6 +280,7 @@ _DEPLOYMENT_MANAGEMENT_ROUTES: Final[frozenset[str]] = frozenset(
("/model/new", "/model/update", "/model/delete", "/health/test_connection")
)
_DEPLOYMENT_ID_UPDATE_ROUTE: Final = re.compile(r"^/model/[^/]+/update$")
_MCP_SERVER_MANAGEMENT_ROUTE: Final = re.compile(r"^/v1/mcp/server(?:/.*)?$")
def route_manages_deployments(route: str | None) -> bool:
@ -297,6 +299,10 @@ def route_manages_deployments(route: str | None) -> bool:
)
def route_manages_mcp_servers(route: str | None) -> bool:
return route is not None and _MCP_SERVER_MANAGEMENT_ROUTE.fullmatch(route) is not None
_NESTED_CONFIG_KEYS: Final[tuple[str, ...]] = ("litellm_embedding_config", "extra_body")
# Metadata containers that carry per-request configuration consumed by the
@ -456,19 +462,25 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = (
def _check_banned_params(
body: dict,
body: Mapping[str, object],
general_settings: dict,
llm_router: Router | None,
model: str,
*,
manages_deployments: bool = False,
mcp_server_fields_allowed: bool = False,
) -> None:
"""Raise ``ValueError`` if ``body`` carries a banned param without admin opt-in.
Shared between the root-level check and the nested-config check so a
new banned param only needs to be added in one place.
"""
reject_server_owned_wif_params(body)
wif_body: Final[Mapping[str, object]] = (
MappingProxyType({key: value for key, value in body.items() if key not in oauth_token_exchange_litellm_params})
if mcp_server_fields_allowed
else body
)
reject_server_owned_wif_params(wif_body)
if not manages_deployments:
reject_federated_credential_reference(body)
for param in _BANNED_REQUEST_BODY_PARAMS:
@ -613,7 +625,14 @@ def is_request_body_safe(
if "model_list" in request_body:
raise ValueError("Rejected Request: model_list is not allowed in the request body.")
manages_deployments: Final = route_manages_deployments(route)
_check_banned_params(request_body, general_settings, llm_router, model, manages_deployments=manages_deployments)
_check_banned_params(
request_body,
general_settings,
llm_router,
model,
manages_deployments=manages_deployments,
mcp_server_fields_allowed=route_manages_mcp_servers(route),
)
for nested_key in _NESTED_CONFIG_KEYS:
nested = _coerce_metadata_to_dict(request_body.get(nested_key))
if nested is not None:

49
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -143,14 +143,17 @@ def _get_guardrails_list_response(
unmasked_length=4,
number_of_asterisks=4,
)
guardrail_configs.append(
GuardrailInfoResponse(
guardrail_id=guardrail.get("guardrail_id"),
guardrail_name=guardrail.get("guardrail_name"),
litellm_params=with_tolerated_stream_scope(masked_params),
guardrail_info=guardrail.get("guardrail_info"),
)
guardrail_info_response = GuardrailInfoResponse(
guardrail_id=guardrail.get("guardrail_id"),
guardrail_name=guardrail.get("guardrail_name"),
guardrail_info=guardrail.get("guardrail_info"),
)
guardrail_info_response.litellm_params = parse_tolerant_litellm_params(
with_tolerated_stream_scope(masked_params),
guardrail_info_response.guardrail_name,
params_model=BaseLitellmParams,
)
guardrail_configs.append(guardrail_info_response)
return ListGuardrailsResponse(guardrails=guardrail_configs)
@ -1282,14 +1285,30 @@ async def patch_guardrail(
status_code=422,
detail=f"Invalid guardrail configuration, update rejected: {validation_error}",
) from validation_error
clear_stored_scope: Final = (
"logging_only_scope" not in requested_litellm_params
and parsed_litellm_params.logging_only_scope is not None
and GuardrailEventHooks.logging_only.value not in configured_event_hooks(parsed_litellm_params.mode)
mode_lacks_logging_only: Final = GuardrailEventHooks.logging_only.value not in configured_event_hooks(
parsed_litellm_params.mode
)
merged_overrides: Final = MappingProxyType(
{
**(
{"logging_only_scope": None}
if "logging_only_scope" not in requested_litellm_params
and parsed_litellm_params.logging_only_scope is not None
and mode_lacks_logging_only
else {}
),
**(
{"logging_only_continue_on_input_failure": None}
if "logging_only_continue_on_input_failure" not in requested_litellm_params
and parsed_litellm_params.logging_only_continue_on_input_failure
and mode_lacks_logging_only
else {}
),
}
)
litellm_params: Final = (
LitellmParams(**MappingProxyType({**merged_litellm_params, "logging_only_scope": None}))
if clear_stored_scope
LitellmParams(**MappingProxyType({**merged_litellm_params, **merged_overrides}))
if merged_overrides
else parsed_litellm_params
)
@ -1320,7 +1339,9 @@ async def patch_guardrail(
try:
IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(
guardrail=guardrail,
reject_invalid_logging_only_scope="logging_only_scope" in requested_litellm_params,
reject_invalid_logging_only_scope=bool(
{"logging_only_scope", "logging_only_continue_on_input_failure"} & set(requested_litellm_params)
),
)
verbose_proxy_logger.info(
"Immediate sync: Successfully updated guardrail '%s' (ID: %s)", guardrail_name, guardrail_id

69
litellm/proxy/guardrails/guardrail_registry.py
View file

@ -587,7 +587,7 @@ def _logging_only_scope_error(
f"Guardrail {guardrail_name}: logging_only_scope is set, but mode does not include logging_only, "
"so it would never apply. Add logging_only to mode or remove logging_only_scope."
)
if logging_only_scope in ("input", "output") and not custom_guardrail_callback.supports_logging_only_scope():
if logging_only_scope is not None and not custom_guardrail_callback.supports_logging_only_scope():
return (
f"Guardrail {guardrail_name}: logging_only_scope={logging_only_scope!r} is not supported by this "
"guardrail, whose logging_only hook scans on its own. Remove logging_only_scope."
@ -595,6 +595,18 @@ def _logging_only_scope_error(
return None
def _logging_only_continue_error(guardrail_name: str, litellm_params: LitellmParams) -> str | None:
if litellm_params.logging_only_continue_on_input_failure is True and (
GuardrailEventHooks.logging_only.value not in configured_event_hooks(litellm_params.mode)
):
return (
f"Guardrail {guardrail_name}: logging_only_continue_on_input_failure is set, but mode does not "
"include logging_only, so it would never apply. Add logging_only to mode or remove "
"logging_only_continue_on_input_failure."
)
return None
def _configure_callback_scoping(
custom_guardrail_callback: CustomGuardrail,
guardrail_name: str,
@ -602,20 +614,31 @@ def _configure_callback_scoping(
*,
reject_invalid_logging_only_scope: bool = False,
) -> None:
logging_only_scope: Final = litellm_params.logging_only_scope
logging_only_scope_error: Final = _logging_only_scope_error(
custom_guardrail_callback, guardrail_name, litellm_params
)
logging_only_continue_error: Final = _logging_only_continue_error(guardrail_name, litellm_params)
if reject_invalid_logging_only_scope:
for error in (logging_only_scope_error, logging_only_continue_error):
if error is not None:
raise ValueError(error)
if logging_only_scope_error is not None:
if reject_invalid_logging_only_scope:
raise ValueError(logging_only_scope_error)
verbose_proxy_logger.error(
"%s Ignoring logging_only_scope; the guardrail keeps its configured mode.",
logging_only_scope_error.replace("\r", "").replace("\n", ""),
)
custom_guardrail_callback.logging_only_scope = None
else:
custom_guardrail_callback.logging_only_scope = logging_only_scope
if logging_only_continue_error is not None:
verbose_proxy_logger.error(
"%s Ignoring logging_only_continue_on_input_failure; the guardrail keeps its configured mode.",
logging_only_continue_error.replace("\r", "").replace("\n", ""),
)
logging_only_scope: Final = None if logging_only_scope_error is not None else litellm_params.logging_only_scope
custom_guardrail_callback.logging_only_scope = logging_only_scope
custom_guardrail_callback.logging_only_continue_on_input_failure = (
logging_only_continue_error is None
and logging_only_scope is None
and litellm_params.logging_only_continue_on_input_failure is True
)
for scoping_param in (
"skip_system_message_in_guardrail",
"skip_tool_message_in_guardrail",
@ -651,15 +674,24 @@ def parse_tolerant_litellm_params(
try:
return params_model(**litellm_params_data)
except ValidationError as validation_error:
if any(tuple(error["loc"]) != ("logging_only_scope",) for error in validation_error.errors()):
tolerated_fields: Final = ("logging_only_scope", "logging_only_continue_on_input_failure")
error_locs: Final = {tuple(error["loc"]) for error in validation_error.errors()}
tolerated_locs: Final = frozenset((field,) for field in tolerated_fields)
if any(loc not in tolerated_locs for loc in error_locs):
raise
verbose_proxy_logger.error(
"Guardrail %s: logging_only_scope=%r is not one of 'input', 'output' or 'both'. "
"Ignoring logging_only_scope; the guardrail keeps its configured mode.",
guardrail_name.replace("\r", "").replace("\n", ""),
str(litellm_params_data.get("logging_only_scope")).replace("\r", "").replace("\n", "")[:100],
sanitized: Final = MappingProxyType(
{**litellm_params_data, **{field: None for field in tolerated_fields if (field,) in error_locs}}
)
return params_model(**MappingProxyType({**litellm_params_data, "logging_only_scope": None}))
for field in tolerated_fields:
if (field,) in error_locs:
verbose_proxy_logger.error(
"Guardrail %s: %s=%r is not a valid value. Ignoring %s; the guardrail keeps its configured mode.",
guardrail_name.replace("\r", "").replace("\n", ""),
field,
str(litellm_params_data.get(field)).replace("\r", "").replace("\n", "")[:100],
field,
)
return params_model(**sanitized)
class InMemoryGuardrailHandler:
@ -829,9 +861,12 @@ class InMemoryGuardrailHandler:
litellm_params: Final = LitellmParams(**params) if isinstance(params, dict) else params
guardrail_name: Final = guardrail.get("guardrail_name", "Unknown")
for custom_guardrail_callback in self._tracked_callbacks(guardrail_id):
scope_error = _logging_only_scope_error(custom_guardrail_callback, guardrail_name, litellm_params)
if scope_error is not None:
raise ValueError(scope_error)
for error in (
_logging_only_scope_error(custom_guardrail_callback, guardrail_name, litellm_params),
_logging_only_continue_error(guardrail_name, litellm_params),
):
if error is not None:
raise ValueError(error)
def initialize_custom_guardrail(
self,

296
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -43,6 +43,7 @@ from litellm.constants import (
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.models.credentials import CredentialItem
from litellm.models.project import LiteLLM_ProjectTable
from litellm.proxy._experimental.mcp_server.db import (
rotate_mcp_server_credentials_master_key,
rotate_mcp_user_credentials_master_key,
@ -122,10 +123,11 @@ from litellm.proxy.management_helpers.access_group_key_sync import (
)
from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: F401 # legacy module exports
ObjectPermissionUpsert,
_set_object_permission, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
attach_object_permission_to_dict,
handle_update_object_permission_common,
invalidate_cached_object_permissions,
prepare_object_permission_upsert,
set_object_permission,
validate_key_mcp_servers_against_team,
validate_key_search_tools_against_team,
@ -150,11 +152,12 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports
hash_token_if_needed,
is_valid_api_key,
)
from litellm.repositories.base_repository import BaseRepository
from litellm.repositories.base_repository import BaseRepository, record_to_dict
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.config_repository import ConfigParam, ConfigRepository
from litellm.repositories.credentials_repository import CredentialsRepository
from litellm.repositories.model_repository import ModelRepository
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import (
DeletedVerificationTokenRepository,
@ -279,6 +282,19 @@ def _prisma_table(
)
def _writer_project_table(
prisma_client: PrismaClient,
) -> "TableActions[prisma_models.LiteLLM_ProjectTable]":
return cast( # cast-ok: writer_db exposes generated Prisma tables through a dynamic wrapper
"TableActions[prisma_models.LiteLLM_ProjectTable]",
prisma_client.writer_db.litellm_projecttable,
)
class KeyProjectBindingError(HTTPException):
pass
def _deleted_verification_token_table(
prisma_client: PrismaClient,
) -> "TableActions[prisma_models.LiteLLM_DeletedVerificationToken]":
@ -1361,7 +1377,8 @@ async def _common_key_generation_helper(
apply_enterprise_key_management_params,
)
data = apply_enterprise_key_management_params(data, team_table)
enterprise_data: Final[object] = apply_enterprise_key_management_params(data, team_table)
data = GenerateKeyRequest.model_validate(enterprise_data)
except Exception as e:
verbose_proxy_logger.debug(
"litellm.proxy.proxy_server.generate_key_fn(): Enterprise key management params not applied - %s", e
@ -1396,6 +1413,10 @@ async def _common_key_generation_helper(
)
_budget_id = getattr(_budget, "budget_id", None)
created_budget_id: Final[str | None] = (
_budget_id if prisma_client is not None and data.soft_budget is not None else None
)
# ADD METADATA FIELDS
# Set Management Endpoint Metadata Fields
for field in LiteLLM_ManagementEndpoint_MetadataFields_Premium:
@ -1501,10 +1522,19 @@ async def _common_key_generation_helper(
for _op_field, _op_default_value in _default_object_permission.items():
_caller_object_permission.setdefault(_op_field, _op_default_value)
should_create_object_permission: Final = prisma_client is not None and isinstance(
data_json.get("object_permission"), dict
)
data_json = await set_object_permission( # rebind-ok: pre-existing rebinding on a rename-only line
data_json=data_json,
prisma_client=prisma_client,
)
written_object_permission_id: Final = data_json.get("object_permission_id")
created_object_permission_id: Final = (
written_object_permission_id
if should_create_object_permission and isinstance(written_object_permission_id, str)
else None
)
_validate_key_alias_format(key_alias=data_json.get("key_alias", None))
@ -1572,7 +1602,19 @@ async def _common_key_generation_helper(
prisma_client=prisma_client,
)
response = await generate_key_helper_fn(request_type="key", **data_json, table_name="key", llm_router=llm_router)
try:
response = await generate_key_helper_fn(
request_type="key", **data_json, table_name="key", llm_router=llm_router
)
except KeyProjectBindingError:
if prisma_client is not None:
if created_object_permission_id is not None:
await ObjectPermissionRepository(prisma_client).table.delete(
where={"object_permission_id": created_object_permission_id}
)
if created_budget_id is not None:
await BudgetRepository(prisma_client).table.delete(where={"budget_id": created_budget_id})
raise
response["soft_budget"] = data.soft_budget # include the user-input soft budget in the response
@ -1834,6 +1876,85 @@ async def _check_project_key_limits(
)
async def _check_key_project_team(
project_id: str,
key_team_id: str | None,
prisma_client: PrismaClient,
) -> str | None:
project_record: Final = await _writer_project_table(prisma_client).find_unique(where={"project_id": project_id})
if project_record is None:
raise KeyProjectBindingError(
status_code=404,
detail={"error": f"Project not found, project_id={project_id}"},
)
project_obj: Final = LiteLLM_ProjectTable.model_validate(record_to_dict(project_record))
if project_obj.team_id is None or project_obj.team_id == key_team_id:
return project_obj.team_id
raise KeyProjectBindingError(
status_code=400,
detail={
"error": (
f"Project {project_id} belongs to team {project_obj.team_id}, but the key belongs to "
f"{key_team_id if key_team_id is not None else 'no team'}. "
"A key can only be attached to a project owned by its own team."
)
},
)
async def _check_key_project_team_on_mutation(
data: UpdateKeyRequest | RegenerateKeyRequest,
existing_key_row: LiteLLM_VerificationToken,
prisma_client: PrismaClient,
) -> str | None:
fields_set: Final = data.model_fields_set
team_changed: Final = "team_id" in fields_set and data.team_id != existing_key_row.team_id
project_changed: Final = "project_id" in fields_set and data.project_id != existing_key_row.project_id
if not team_changed and not project_changed:
return None
project_id: Final = data.project_id if "project_id" in fields_set else existing_key_row.project_id
if project_id is None:
return None
team_id: Final = data.team_id if "team_id" in fields_set else existing_key_row.team_id
return await _check_key_project_team(
project_id=project_id,
key_team_id=team_id,
prisma_client=prisma_client,
)
async def _check_caller_in_project_team(
project_team_id: str,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
) -> None:
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
return
team_table: Final = await get_team_object(
team_id=project_team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_dict.parent_otel_span,
check_db_only=True,
)
if get_caller_team_role(team_table=team_table, user_api_key_dict=user_api_key_dict) is None:
raise HTTPException(
status_code=403,
detail={
"error": (
f"Only a proxy admin or a member of team {project_team_id} can attach a key to a project "
"owned by that team."
)
},
)
def check_org_key_model_specific_limits(
keys: Sequence[LiteLLM_VerificationToken],
org_table: LiteLLM_OrganizationTable,
@ -2488,6 +2609,18 @@ async def _update_key_row_with_soft_budget(
return result
async def _update_key_row(
prisma_client: PrismaClient,
key: str,
update_values: Mapping[str, object],
) -> _KeyUpdateResult | None:
key_update_data: Final = MappingProxyType({**update_values, "token": key})
response: Final = await prisma_client.update_data(token=key, data=key_update_data)
if response is None:
return None
return cast("_KeyUpdateResult", response) # cast-ok: key update_data returns token and data
async def prepare_key_update_data(
data: UpdateKeyRequest | RegenerateKeyRequest,
existing_key_row: LiteLLM_VerificationToken,
@ -2594,27 +2727,42 @@ async def prepare_key_update_data(
return non_default_values
async def _handle_update_object_permission(
data_json: dict,
async def _prepare_key_update_object_permission(
object_permission_data: object,
existing_key_row: LiteLLM_VerificationToken,
prisma_client: PrismaClient,
) -> dict:
"""Persist the requested object permission row and swap it for its id, only after the key policy allowed the write."""
if "object_permission" not in data_json:
return data_json
) -> ObjectPermissionUpsert | None:
if object_permission_data is None:
return None
object_permission_id: Final = await handle_update_object_permission_common(
data_json=data_json,
parsed_object_permission: Final[object] = (
json.loads(object_permission_data) if isinstance(object_permission_data, str) else object_permission_data
)
permission_data: Final[dict[str, object]] = (
TypeAdapter(dict[str, object]).validate_python(parsed_object_permission)
if isinstance(parsed_object_permission, dict)
else {}
)
return await prepare_object_permission_upsert(
new_object_permission=permission_data,
existing_object_permission_id=existing_key_row.object_permission_id,
prisma_client=prisma_client,
)
# Add the object_permission_id to data_json if one was created/updated
if object_permission_id is not None:
data_json["object_permission_id"] = object_permission_id
verbose_proxy_logger.debug("updated object_permission_id: %s", object_permission_id)
return data_json
async def _write_prepared_key_update_object_permission(
data_json: Mapping[str, object],
upsert: ObjectPermissionUpsert | None,
prisma_client: PrismaClient,
) -> Mapping[str, object]:
if upsert is None:
return data_json
await ObjectPermissionRepository(prisma_client).table.upsert(
where={"object_permission_id": upsert.object_permission_id},
data={"create": upsert.record, "update": upsert.record},
)
return MappingProxyType({**data_json, "object_permission_id": upsert.object_permission_id})
def is_different_team(data: UpdateKeyRequest, existing_key_row: LiteLLM_VerificationToken) -> bool:
@ -2726,6 +2874,23 @@ def _resolve_token_to_update(data: UpdateKeyRequest, existing_key_row: LiteLLM_V
return existing_key_row.token
async def _check_single_key_update_team_permissions(
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient | None,
existing_key_row: LiteLLM_VerificationToken,
user_api_key_cache: UserApiKeyCache,
) -> None:
if prisma_client is None:
return
await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
user_api_key_dict=user_api_key_dict,
route=KeyManagementRoutes.KEY_UPDATE,
prisma_client=prisma_client,
existing_key_row=existing_key_row,
user_api_key_cache=user_api_key_cache,
)
async def _process_single_key_update(
update_key_request: UpdateKeyRequest,
user_api_key_dict: UserAPIKeyAuth,
@ -2789,15 +2954,12 @@ async def _process_single_key_update(
entity="key",
)
# Check team member permissions
if prisma_client is not None:
await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
user_api_key_dict=user_api_key_dict,
route=KeyManagementRoutes.KEY_UPDATE,
prisma_client=prisma_client,
existing_key_row=existing_key_row,
user_api_key_cache=user_api_key_cache,
)
await _check_single_key_update_team_permissions(
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
existing_key_row=existing_key_row,
user_api_key_cache=user_api_key_cache,
)
check_denied_passthrough_routes_caller_permission(
update_key_request,
user_api_key_dict,
@ -2887,11 +3049,24 @@ async def _process_single_key_update(
detail={"error": "Database not connected"},
)
update_values: Final = await _handle_update_object_permission(
data_json=non_default_values,
object_permission_upsert: Final = await _prepare_key_update_object_permission(
object_permission_data=non_default_values.get("object_permission"),
existing_key_row=existing_key_row,
prisma_client=prisma_client,
)
key_update_values: Final = MappingProxyType(
{field: value for field, value in non_default_values.items() if field != "object_permission"}
)
await _check_key_project_team_on_mutation(
data=key_request,
existing_key_row=existing_key_row,
prisma_client=prisma_client,
)
update_values: Final = await _write_prepared_key_update_object_permission(
data_json=key_update_values,
upsert=object_permission_upsert,
prisma_client=prisma_client,
)
_data: Final = {**update_values, "token": key_request.key}
response: Final[Mapping[str, object] | None] = cast( # cast-ok: every update_data branch returns a str-keyed dict
"Mapping[str, object] | None",
@ -2902,7 +3077,7 @@ async def _process_single_key_update(
await invalidate_cached_object_permissions(
object_permission_ids=(
existing_key_row.object_permission_id,
non_default_values.get("object_permission_id"),
update_values.get("object_permission_id"),
),
user_api_key_cache=user_api_key_cache,
)
@ -3563,11 +3738,24 @@ async def update_key_fn(
if prisma_client is None:
raise Exception("Not connected to DB!")
update_values: Final = await _handle_update_object_permission(
data_json=non_default_values,
object_permission_upsert: Final = await _prepare_key_update_object_permission(
object_permission_data=non_default_values.get("object_permission"),
existing_key_row=existing_key_row,
prisma_client=prisma_client,
)
key_update_values: Final = MappingProxyType(
{field: value for field, value in non_default_values.items() if field != "object_permission"}
)
await _check_key_project_team_on_mutation(
data=data,
existing_key_row=existing_key_row,
prisma_client=prisma_client,
)
update_values: Final = await _write_prepared_key_update_object_permission(
data_json=key_update_values,
upsert=object_permission_upsert,
prisma_client=prisma_client,
)
changed_by: Final = user_api_key_dict.user_id or litellm_proxy_admin_name
response: Final = (
await _update_key_row_with_soft_budget(
@ -3579,7 +3767,11 @@ async def update_key_fn(
changed_by=changed_by,
)
if "soft_budget" in data.model_fields_set
else await prisma_client.update_data(token=key, data=MappingProxyType({**update_values, "token": key}))
else await _update_key_row(
prisma_client=prisma_client,
key=key,
update_values=update_values,
)
)
# Delete - key from cache, since it's been updated!
@ -3587,7 +3779,7 @@ async def update_key_fn(
await invalidate_cached_object_permissions(
object_permission_ids=(
existing_key_row.object_permission_id,
non_default_values.get("object_permission_id"),
update_values.get("object_permission_id"),
),
user_api_key_cache=user_api_key_cache,
)
@ -4846,6 +5038,13 @@ async def generate_key_helper_fn(
# the LiteLLM_VerificationToken table will increase in size if we don't do this check
return user_data
if project_id is not None:
await _check_key_project_team(
project_id=project_id,
key_team_id=team_id,
prisma_client=prisma_client,
)
## CREATE KEY
verbose_proxy_logger.debug(
"prisma_client: Creating Key= %s",
@ -5655,7 +5854,6 @@ async def _execute_virtual_key_regeneration(
new_token: Final = await get_new_token(data=data)
new_token_hash: Final = hash_token(new_token)
new_token_key_name: Final = abbreviate_api_key(api_key=new_token)
update_data = {"token": new_token_hash, "key_name": new_token_key_name}
non_default_values = {}
if data is not None:
@ -5681,12 +5879,34 @@ async def _execute_virtual_key_regeneration(
request=data if data is not None else RegenerateKeyRequest(),
),
)
update_values: Final = await _handle_update_object_permission(
data_json=non_default_values,
object_permission_upsert: Final = await _prepare_key_update_object_permission(
object_permission_data=non_default_values.get("object_permission"),
existing_key_row=key_in_db,
prisma_client=prisma_client,
)
update_data.update(update_values)
key_update_values: Final = MappingProxyType(
{field: value for field, value in non_default_values.items() if field != "object_permission"}
)
if data is not None:
project_team_id: Final = await _check_key_project_team_on_mutation(
data=data,
existing_key_row=key_in_db,
prisma_client=prisma_client,
)
if project_team_id is not None:
await _check_caller_in_project_team(
project_team_id=project_team_id,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
update_values: Final = await _write_prepared_key_update_object_permission(
data_json=key_update_values,
upsert=object_permission_upsert,
prisma_client=prisma_client,
)
update_data: Final = MappingProxyType({"token": new_token_hash, "key_name": new_token_key_name, **update_values})
jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data)
# Snapshot before the token update: the FK cascade rewrites mapping rows to the new hash,
@ -5724,7 +5944,7 @@ async def _execute_virtual_key_regeneration(
await invalidate_cached_object_permissions(
object_permission_ids=(
key_in_db.object_permission_id,
non_default_values.get("object_permission_id"),
update_values.get("object_permission_id"),
),
user_api_key_cache=user_api_key_cache,
)

2
litellm/proxy/spend_tracking/spend_management_endpoints.py
View file

@ -3767,7 +3767,7 @@ async def global_spend_refresh():
},
)
await new_client.db.connect()
await _query_raw(new_client, sql_query)
await new_client.db.execute_raw(sql_query)
verbose_proxy_logger.info("MonthlyGlobalSpend view refreshed")
return {
"message": "MonthlyGlobalSpend view refreshed",

16
litellm/types/guardrails.py
View file

@ -989,7 +989,7 @@ def runtime_stream_scope(
return DEFAULT_GUARDRAIL_STREAM_SCOPE, MappingProxyType(coerced)
LoggingOnlyScope = Literal["input", "output", "both"]
LoggingOnlyScope = Literal["input", "output"]
class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch update guardrails
@ -1250,8 +1250,18 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up
logging_only_scope: LoggingOnlyScope | None = Field(
default=None,
description=(
"which direction a logging_only scan observes: 'input' (request), 'output' (response), or 'both' "
"(default). Only applies to mode logging_only; pre_call/post_call on the same guardrail keep blocking."
"which direction a logging_only scan observes: 'input' (request) or 'output' (response); "
"unset scans both directions. Only applies to mode logging_only; pre_call/post_call on the "
"same guardrail keep blocking."
),
)
logging_only_continue_on_input_failure: bool | None = Field(
default=None,
description=(
"when True, a flagged or raising logging_only request scan is logged and the response is "
"still scanned, so both verdicts land. Only applies to mode logging_only and is ignored "
"when logging_only_scope is 'input' or 'output'."
),
)

40
model_prices_and_context_window.json
View file

@ -71786,5 +71786,45 @@
"supported_endpoints": [
"/v1/responses"
]
},
"azure_ai/Cohere-Embed-V5-Fast": {
"deprecation_date": "2027-01-23",
"input_cost_per_image_token": 4e-07,
"input_cost_per_token": 8e-08,
"litellm_provider": "azure_ai",
"max_input_tokens": 128000,
"max_tokens": 128000,
"mode": "embedding",
"output_cost_per_token": 0.0,
"output_vector_size": 2048.0,
"source": "https://prices.azure.com/api/retail/prices?$filter=productName%20eq%20'Cohere%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supported_endpoints": [
"/v1/embeddings"
],
"supported_modalities": [
"text",
"image"
],
"supports_embedding_image_input": true
},
"azure_ai/Cohere-Embed-V5-Pro": {
"deprecation_date": "2027-01-23",
"input_cost_per_image_token": 4e-07,
"input_cost_per_token": 1.2e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 128000,
"max_tokens": 128000,
"mode": "embedding",
"output_cost_per_token": 0.0,
"output_vector_size": 2048.0,
"source": "https://prices.azure.com/api/retail/prices?$filter=productName%20eq%20'Cohere%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supported_endpoints": [
"/v1/embeddings"
],
"supported_modalities": [
"text",
"image"
],
"supports_embedding_image_input": true
}
}

182
tests/integration/database/test_global_spend_refresh_writer.py Normal file
View file

@ -0,0 +1,182 @@
import os
import uuid
from collections.abc import Generator
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Final
from urllib.parse import urlsplit, urlunsplit
import psycopg
import pytest
from psycopg import sql
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows, scratch_database
from integration._support.process import owned_proxy
pytestmark: Final = pytest.mark.timeout(240)
READER_PASSWORD: Final = "integration-reader-password"
SPEND: Final = 1.25
REFRESHED: Final = {"message": "MonthlyGlobalSpend view refreshed", "status": "success"}
FAILED: Final = {"message": "Failed to refresh materialized view", "status": "failure"}
@dataclass(frozen=True, slots=True)
class SpendDatabase:
candidate: Gateway
writer_url: str
read_url: str
@contextmanager
def readonly_role() -> Generator[str]:
role: Final = f"integration_reader_{uuid.uuid4().hex}"
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as admin:
admin.execute(
sql.SQL("CREATE ROLE {} LOGIN PASSWORD {} NOSUPERUSER NOINHERIT").format(
sql.Identifier(role), sql.Literal(READER_PASSWORD)
)
)
try:
admin.execute(sql.SQL("ALTER ROLE {} SET default_transaction_read_only = on").format(sql.Identifier(role)))
yield role
finally:
admin.execute(sql.SQL("DROP ROLE {}").format(sql.Identifier(role)))
@contextmanager
def spend_database(gateway: Gateway, tmp_path: Path, *, with_reader: bool) -> Generator[SpendDatabase]:
"""A proxy on its own database with its spend views created, reading through a real read-only role when
`with_reader`, and with no read replica at all otherwise"""
with readonly_role() as role, scratch_database() as writer_url:
parsed: Final = urlsplit(writer_url)
host: Final = parsed.netloc.rpartition("@")[2]
reader_url: Final = urlunsplit(parsed._replace(netloc=f"{role}:{READER_PASSWORD}@{host}"))
with psycopg.connect(writer_url, autocommit=True) as owner:
owner.execute(sql.SQL("GRANT USAGE ON SCHEMA public TO {}").format(sql.Identifier(role)))
owner.execute(
sql.SQL("ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT ON TABLES TO {}").format(
sql.Identifier(role)
)
)
with owned_proxy(
gateway,
tmp_path,
{"DATABASE_URL": writer_url, **({"DATABASE_URL_READ_REPLICA": reader_url} if with_reader else {})},
remove_environment=("DATABASE_URL_READ_REPLICA",),
) as candidate:
eventually(
lambda: read_rows(
"SELECT relname FROM pg_class WHERE relname=%s AND relkind='v'",
("MonthlyGlobalSpend",),
database_url=writer_url,
),
bool,
seconds=60,
)
connected: Final = bool(read_rows("SELECT pid FROM pg_stat_activity WHERE usename=%s", (role,)))
assert connected == with_reader, f"read-only reader connected={connected}, expected {with_reader}"
yield SpendDatabase(candidate, writer_url, reader_url if with_reader else writer_url)
def materialize_stock_view(writer_url: str) -> None:
with psycopg.connect(writer_url, autocommit=True) as owner:
owner.execute("""
DO $$
DECLARE definition text := pg_get_viewdef('"MonthlyGlobalSpend"'::regclass);
BEGIN
EXECUTE 'DROP VIEW "MonthlyGlobalSpend"';
EXECUTE format('CREATE MATERIALIZED VIEW "MonthlyGlobalSpend" AS %s WITH NO DATA', rtrim(definition, '; '));
END $$
""")
def insert_spend(writer_url: str) -> None:
with psycopg.connect(writer_url, autocommit=True) as owner:
owner.execute(
'INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, "startTime", "endTime", spend) '
"VALUES (%s, 'acompletion', now(), now(), %s)",
(f"integration-{uuid.uuid4().hex}", SPEND),
)
def test_global_spend_refresh_runs_on_the_writer_with_a_real_readonly_reader(gateway: Gateway, tmp_path: Path) -> None:
with spend_database(gateway, tmp_path, with_reader=True) as database:
materialize_stock_view(database.writer_url)
insert_spend(database.writer_url)
with psycopg.connect(database.read_url) as reader:
with pytest.raises(psycopg.errors.ReadOnlySqlTransaction):
reader.execute('REFRESH MATERIALIZED VIEW "MonthlyGlobalSpend"')
response: Final = database.candidate.request("POST", "/global/spend/refresh")
assert response.status_code == 200 and response.json() == REFRESHED, response.text
assert read_rows('SELECT spend FROM "MonthlyGlobalSpend"', (), database_url=database.read_url) == [
{"spend": SPEND}
]
def test_concurrent_global_spend_refreshes_all_run_on_the_writer(gateway: Gateway, tmp_path: Path) -> None:
with spend_database(gateway, tmp_path, with_reader=True) as database:
materialize_stock_view(database.writer_url)
insert_spend(database.writer_url)
with ThreadPoolExecutor(max_workers=5) as executor:
pending: Final = [
executor.submit(database.candidate.request, "POST", "/global/spend/refresh") for _ in range(5)
]
responses: Final = [future.result() for future in pending]
assert [(response.status_code, response.json()) for response in responses] == [(200, REFRESHED)] * 5
assert read_rows('SELECT spend FROM "MonthlyGlobalSpend"', (), database_url=database.read_url) == [
{"spend": SPEND}
]
def test_global_spend_refresh_without_a_read_replica_still_refreshes(gateway: Gateway, tmp_path: Path) -> None:
with spend_database(gateway, tmp_path, with_reader=False) as database:
materialize_stock_view(database.writer_url)
insert_spend(database.writer_url)
response: Final = database.candidate.request("POST", "/global/spend/refresh")
assert response.status_code == 200 and response.json() == REFRESHED, response.text
assert read_rows('SELECT spend FROM "MonthlyGlobalSpend"', (), database_url=database.writer_url) == [
{"spend": SPEND}
]
def test_global_spend_refresh_leaves_the_stock_plain_view_alone(gateway: Gateway, tmp_path: Path) -> None:
with spend_database(gateway, tmp_path, with_reader=True) as database:
response: Final = database.candidate.request("POST", "/global/spend/refresh")
assert response.status_code == 200 and response.json() is None, response.text
assert read_rows(
"SELECT relkind FROM pg_class WHERE relname=%s", ("MonthlyGlobalSpend",), database_url=database.writer_url
) == [{"relkind": "v"}]
def test_global_spend_refresh_reports_failure_when_the_writer_refresh_errors(gateway: Gateway, tmp_path: Path) -> None:
with spend_database(gateway, tmp_path, with_reader=True) as database:
with psycopg.connect(database.writer_url, autocommit=True) as owner:
owner.execute("CREATE SEQUENCE refresh_attempts")
owner.execute('DROP VIEW "MonthlyGlobalSpend"')
owner.execute(
'CREATE MATERIALIZED VIEW "MonthlyGlobalSpend" AS '
"SELECT nextval('refresh_attempts') / 0 AS spend WITH NO DATA"
)
response: Final = database.candidate.request("POST", "/global/spend/refresh")
assert response.status_code == 200 and response.json() == FAILED, response.text
assert read_rows("SELECT is_called FROM refresh_attempts", (), database_url=database.writer_url) == [
{"is_called": True}
], "the refresh never ran on the writer"
assert read_rows(
"SELECT relispopulated FROM pg_class WHERE relname=%s",
("MonthlyGlobalSpend",),
database_url=database.writer_url,
) == [{"relispopulated": False}]

1605
tests/integration/management/test_project_lifecycle.py
View file

File diff suppressed because it is too large Load diff

20
tests/integration/observability/_logging_only_scope_support.py
View file

@ -240,7 +240,7 @@ def wire_server(respond: Callable[[Request], Reply], port: int = 0, *, policy_ed
def _directions_for_scope(base_default: tuple[Direction, ...], scope: str | None) -> tuple[Direction, ...]:
if scope is None or scope == "both":
if scope is None:
return base_default
selected_direction: Final = "request" if scope == "input" else "response"
return tuple(direction for direction in base_default if direction == selected_direction)
@ -629,6 +629,7 @@ def _configuration(
scope: str | None,
*,
include_scope: bool = True,
continue_on_input_failure: bool | None = None,
default_on: bool = True,
mode: str | list[str] = "logging_only",
cache: bool = False,
@ -648,6 +649,11 @@ def _configuration(
"api_key": "synthetic-guardrail-key",
"extra_headers": ["x-litellm-call-id"],
**({"logging_only_scope": scope} if include_scope else {}),
**(
{"logging_only_continue_on_input_failure": continue_on_input_failure}
if continue_on_input_failure is not None
else {}
),
}
config["guardrails"] = [{"guardrail_name": identity, "litellm_params": params}]
scope_name: Final = scope if scope is not None else "unset"
@ -760,6 +766,7 @@ def _insert_database_guardrail(
*,
mode: str = "pre_call",
default_on: bool = True,
continue_on_input_failure: bool | None = None,
) -> None:
params: Final = {
"guardrail": "generic_guardrail_api",
@ -769,6 +776,11 @@ def _insert_database_guardrail(
"api_key": "synthetic-guardrail-key",
"extra_headers": ["x-litellm-call-id"],
**({"logging_only_scope": scope} if scope is not None else {}),
**(
{"logging_only_continue_on_input_failure": continue_on_input_failure}
if continue_on_input_failure is not None
else {}
),
}
write_rows(
'INSERT INTO "LiteLLM_GuardrailsTable" '
@ -805,6 +817,7 @@ def _post_guardrail_body(
api_base: str,
scope: JsonValue,
include_scope: bool = True,
continue_on_input_failure: bool | None = None,
) -> dict[str, JsonValue]:
params: Final = {
"guardrail": provider,
@ -821,6 +834,11 @@ def _post_guardrail_body(
),
**({"extra_headers": ["x-litellm-call-id"]} if provider == "generic_guardrail_api" else {}),
**({"logging_only_scope": scope} if include_scope else {}),
**(
{"logging_only_continue_on_input_failure": continue_on_input_failure}
if continue_on_input_failure is not None
else {}
),
}
return {
"guardrail": {

13
tests/integration/observability/test_guardrail_effects.py
View file

@ -1807,11 +1807,16 @@ def test_responses_pre_call_denial_stream_survives_worker_kill(gateway: Gateway,
assert response.status_code == 200, response.text
assert response.headers["content-type"].startswith("text/event-stream"), response.text
@pytest.mark.parametrize(
("logging_only_scope", "scanned_directions"),
(("input", ("request",)), ("output", ("response",)), ("both", ("request", "response"))),
("logging_only_settings", "scanned_directions"),
(
({"logging_only_scope": "input"}, ("request",)),
({"logging_only_scope": "output"}, ("response",)),
({"logging_only_continue_on_input_failure": True}, ("request", "response")),
),
ids=("input", "output", "continue_on_input_failure"),
)
def test_logging_only_scope_observes_only_the_configured_direction_without_blocking(
gateway: Gateway, tmp_path: Path, logging_only_scope: str, scanned_directions: tuple[str, ...]
gateway: Gateway, tmp_path: Path, logging_only_settings: dict[str, object], scanned_directions: tuple[str, ...]
) -> None:
identity: Final = "guardrail" + uuid.uuid4().hex
prompt: Final = "synthetic observed prompt " + identity
@ -1848,7 +1853,7 @@ def test_logging_only_scope_observes_only_the_configured_direction_without_block
"litellm_params": {
"guardrail": "generic_guardrail_api",
"mode": "logging_only",
"logging_only_scope": logging_only_scope,
**logging_only_settings,
"default_on": True,
"api_base": policy.url,
"api_key": "synthetic-guardrail-key",

172
tests/integration/observability/test_logging_only_scope_config.py
View file

@ -879,14 +879,14 @@ def test_G6_database_guardrail_polling_normalizes_invalid_scope_without_reinitia
id="H4-post-rejects-presidio-input-scope",
),
pytest.param(
"H4",
"presidio",
"H2",
"generic_guardrail_api",
"logging_only",
"both",
200,
"both",
"",
id="H4-post-accepts-presidio-both-scope",
422,
None,
"logging_only_scope",
id="H2-post-rejects-removed-both-scope",
),
),
)
@ -1589,3 +1589,163 @@ def test_H13_unauthenticated_management_and_chat_requests_do_not_scan(gateway: G
assert _drain_upstream(gateway.upstream_url) == ()
finally:
delete_scenario(upstream_handle)
def test_K1_post_rejects_continue_flag_outside_logging_only(gateway: Gateway, tmp_path: Path) -> None:
identity: Final = f"logging-scope-k1-{uuid.uuid4().hex}"
config: Final = _empty_proxy_configuration(tmp_path, identity)
try:
with owned_proxy(gateway, tmp_path, {}, config=config, workers=1) as candidate:
before_list: Final = candidate.get("/v2/guardrails/list")["guardrails"]
response: Final = candidate.request(
"POST",
"/guardrails",
_post_guardrail_body(
identity,
"generic_guardrail_api",
"pre_call",
"http://127.0.0.1:9",
None,
include_scope=False,
continue_on_input_failure=True,
),
)
assert response.status_code == 400, response.text
assert "logging_only_continue_on_input_failure" in response.text, response.text
assert _management_guardrail_rows(identity) == ()
after_list: Final = candidate.get("/v2/guardrails/list")["guardrails"]
assert after_list == before_list, (before_list, after_list)
assert identity not in {object_value(item)["guardrail_name"] for item in after_list}, after_list
finally:
_delete_database_guardrail(identity)
def test_K3_patch_flag_rejection_preserves_raw_stored_row(gateway: Gateway, tmp_path: Path) -> None:
identity: Final = f"logging-scope-k3-{uuid.uuid4().hex}"
config: Final = _empty_proxy_configuration(tmp_path, identity)
try:
with owned_proxy(gateway, tmp_path, {}, config=config, workers=1) as candidate:
created: Final = _create_guardrail(
candidate,
identity,
{
"guardrail": "generic_guardrail_api",
"mode": "pre_call",
"default_on": True,
"api_base": "http://127.0.0.1:9",
"api_key": "synthetic-guardrail-key",
},
)
guardrail_id: Final = string_value(created["guardrail_id"])
before_raw: Final = read_rows(
'SELECT litellm_params::text AS raw_params FROM "LiteLLM_GuardrailsTable" WHERE guardrail_name=%s',
(identity,),
)
assert len(before_raw) == 1, before_raw
rejected: Final = candidate.request(
"PATCH",
f"/guardrails/{guardrail_id}",
{"litellm_params": {"logging_only_continue_on_input_failure": True}},
)
assert rejected.status_code == 422, rejected.text
assert "logging_only_continue_on_input_failure" in rejected.text, rejected.text
after_raw: Final = read_rows(
'SELECT litellm_params::text AS raw_params FROM "LiteLLM_GuardrailsTable" WHERE guardrail_name=%s',
(identity,),
)
assert len(after_raw) == 1, after_raw
# Rollback re-encrypts api_key, so ciphertext bytes differ even though the row is unchanged.
before_params: Final = JSON_OBJECT.validate_json(before_raw[0]["raw_params"])
after_params: Final = JSON_OBJECT.validate_json(after_raw[0]["raw_params"])
assert set(before_params) == set(after_params), (before_params, after_params)
assert {
key: before_params[key] for key in before_params if key != "api_key"
} == {key: after_params[key] for key in after_params if key != "api_key"}, (before_params, after_params)
assert str(after_params["api_key"]).startswith("litellm_enc::"), after_params
finally:
_delete_database_guardrail(identity)
def test_K4_patch_dropping_logging_only_mode_clears_stored_flag(gateway: Gateway, tmp_path: Path) -> None:
identity: Final = f"logging-scope-k4-{uuid.uuid4().hex}"
config: Final = _empty_proxy_configuration(tmp_path, identity)
try:
with owned_proxy(gateway, tmp_path, {}, config=config, workers=1) as candidate:
created: Final = _create_guardrail(
candidate,
identity,
{
"guardrail": "generic_guardrail_api",
"mode": ["pre_call", "logging_only"],
"default_on": True,
"api_base": "http://127.0.0.1:9",
"api_key": "synthetic-guardrail-key",
"logging_only_continue_on_input_failure": True,
},
)
guardrail_id: Final = string_value(created["guardrail_id"])
before: Final = _management_guardrail_rows(identity)
assert len(before) == 1, before
assert object_value(before[0]["litellm_params"]).get("logging_only_continue_on_input_failure") is True, (
before
)
patched: Final = candidate.request(
"PATCH",
f"/guardrails/{guardrail_id}",
{"litellm_params": {"mode": ["pre_call"]}},
)
assert patched.status_code == 200, patched.text
persisted: Final = _management_guardrail_rows(identity)
assert len(persisted) == 1, persisted
params: Final = object_value(persisted[0]["litellm_params"])
assert params.get("logging_only_continue_on_input_failure") is None, persisted
assert params["mode"] == ["pre_call"], persisted
finally:
_delete_database_guardrail(identity)
def test_K5_yaml_flag_on_blocking_guardrail_boots_and_keeps_blocking(gateway: Gateway, tmp_path: Path) -> None:
identity: Final = f"logging-scope-k5-{uuid.uuid4().hex}"
prompt: Final = f"synthetic yaml flag marker {identity}"
reply: Final = f"synthetic yaml flag response {identity}"
scenario_id: Final = f"phase12-k5-{uuid.uuid4().hex}"
upstream_handle: Final = register_scenario(scenario_id, _provider_response("chat", scenario_id, reply, False))
def policy(request: Request) -> Reply:
assert request.target == "/beta/litellm_basic_guardrail_api", request.target
return Reply(body=b'{"action":"BLOCKED","blocked_reason":"synthetic yaml flag denial"}')
try:
with gateway.scenario() as scenario:
model: Final = scenario.model(
model="openai/gpt-4o-mini",
api_base=f"{upstream_handle.api_base()}/v1",
api_key="synthetic-provider-key",
)
with wire_server(policy) as guardrail:
config: Final = _configuration(
tmp_path,
identity,
guardrail.url,
None,
include_scope=False,
continue_on_input_failure=True,
mode="pre_call",
)
with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=1) as owned:
candidate: Final = owned.gateway
guarded: Final = candidate.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": prompt}]},
headers={"x-litellm-call-id": f"{scenario_id}-guarded"},
)
assert guarded.status_code == 400 and "synthetic yaml flag denial" in guarded.text, guarded.text
assert _drain_upstream(gateway.upstream_url) == ()
payloads: Final = tuple(JSON_OBJECT.validate_json(call.body) for call in guardrail.drain())
assert tuple(_direction(payload) for payload in payloads) == ("request",), payloads
log_text: Final = owned.log.read_text()
assert "Ignoring logging_only_continue_on_input_failure" in log_text, log_text
finally:
_delete_database_guardrail(identity)
delete_scenario(upstream_handle)

210
tests/integration/observability/test_logging_only_scope_runtime.py
View file

@ -64,7 +64,6 @@ from tests.integration.cost_calculation.cost_tracking_case import JsonResponse
(
pytest.param("A1", "chat", False, "openai_sync", "input", True, (), id="A1-chat-input"),
pytest.param("A2", "chat", False, "openai_sync", "output", True, (), id="A2-chat-output"),
pytest.param("A3", "chat", False, "openai_sync", "both", True, (), id="A3-chat-both"),
pytest.param("A4", "chat", False, "httpx", None, False, (), id="A4-chat-missing-scope"),
pytest.param("A5", "chat", False, "httpx", None, True, (), id="A5-chat-null-scope"),
pytest.param("A6", "chat", True, "openai_async", "input", True, (), id="A6-chat-stream-async-input"),
@ -97,16 +96,6 @@ from tests.integration.cost_calculation.cost_tracking_case import JsonResponse
pytest.param("A15", "responses", True, "openai_sync", "input", True, (), id="A15-responses-stream-input"),
pytest.param("B1", "chat", False, "openai_sync", "input", True, ("request",), id="B1-logging-block-input"),
pytest.param("B2", "chat", False, "openai_sync", "output", True, ("response",), id="B2-logging-block-output"),
pytest.param(
"B3",
"chat",
False,
"openai_sync",
"both",
True,
("request",),
id="B3-logging-block-both",
),
),
)
def test_runtime_directional_scope_matches_client_call_and_spend_log(
@ -911,10 +900,9 @@ def test_logging_only_scope_respects_guardrail_selection_level(
delete_scenario(upstream_handle)
def test_X1_missing_null_and_both_scope_have_identical_scans(gateway: Gateway, tmp_path: Path) -> None:
def test_X1_missing_and_null_scope_have_identical_scans(gateway: Gateway, tmp_path: Path) -> None:
identity: Final = f"logging-scope-x1-{uuid.uuid4().hex}"
variants: Final = (
("both", "both", True),
("missing", None, False),
("null", None, True),
)
@ -1577,3 +1565,199 @@ def test_S4_logging_only_input_scope_scans_every_multipart_text_part(gateway: Ga
) == tuple((identity, "logging_only", "guardrail_intervened") for _ in expected_directions), entries
finally:
delete_scenario(upstream_handle)
@pytest.mark.parametrize(
("row_id", "scope", "include_scope", "continue_flag", "expected_directions"),
(
pytest.param("J1", None, False, None, ("request",), id="J1-unset-flag-stops-after-flagged-input"),
pytest.param(
"J2",
None,
False,
True,
("request", "response"),
id="J2-flag-continues-past-flagged-input",
),
pytest.param("J5", "input", True, True, ("request",), id="J5-input-scope-ignores-flag"),
pytest.param("J6", "output", True, True, ("response",), id="J6-output-scope-ignores-flag"),
),
)
def test_continue_on_input_failure_observes_flagged_input_and_response(
gateway: Gateway,
tmp_path: Path,
row_id: str,
scope: str | None,
include_scope: bool,
continue_flag: bool | None,
expected_directions: tuple[Direction, ...],
) -> None:
identity: Final = f"logging-scope-{row_id.lower()}-{uuid.uuid4().hex}"
prompt: Final = f"synthetic flagged request {identity}"
reply: Final = f"synthetic flagged response {identity}"
scenario_id: Final = f"phase12-{row_id.lower()}-{uuid.uuid4().hex}"
guarded_call_id: Final = f"{scenario_id}-guarded"
upstream_handle: Final = register_scenario(
scenario_id, _provider_response("chat", scenario_id, reply, False)
)
def policy(_request: Request) -> Reply:
return Reply(body=b'{"action":"BLOCKED","blocked_reason":"synthetic flagged denial"}')
try:
with gateway.scenario() as scenario:
model: Final = scenario.model(
model="openai/gpt-4o-mini",
api_base=f"{upstream_handle.api_base()}/v1",
api_key="synthetic-provider-key",
)
with wire_server(policy) as guardrail:
config: Final = _configuration(
tmp_path,
identity,
guardrail.url,
scope,
include_scope=include_scope,
continue_on_input_failure=continue_flag,
)
with owned_proxy(gateway, tmp_path, {}, config=config, workers=1) as candidate:
baseline: Final = _call_client(
"openai_sync", "chat", gateway, model, prompt, False, f"{scenario_id}-baseline"
)
assert baseline.status == 200 and baseline.text == reply, baseline
assert len(_drain_upstream(gateway.upstream_url)) == 1
result: Final = _call_client(
"openai_sync", "chat", candidate, model, prompt, False, guarded_call_id
)
assert (result.status, _response_body_without_ids(result.body)) == (
baseline.status,
_response_body_without_ids(baseline.body),
), (result, baseline)
assert result.text == reply, result
upstream: Final = _drain_upstream(gateway.upstream_url)
assert len(upstream) == 1 and prompt in json.dumps(upstream[0]["body"]), upstream
eventually(
lambda: guardrail.received.qsize(),
lambda count: count >= len(expected_directions),
seconds=20,
)
calls: Final = guardrail.drain()
assert len(calls) == len(expected_directions), calls
payloads: Final = tuple(JSON_OBJECT.validate_json(call.body) for call in calls)
observed_directions: Final = tuple(_direction(payload) for payload in payloads)
assert tuple(sorted(observed_directions)) == tuple(sorted(expected_directions)), payloads
assert all(_policy_call_id_matches(payload, guarded_call_id) for payload in payloads), payloads
assert tuple(payload["texts"] for payload in payloads) == tuple(
[prompt] if direction == "request" else [reply] for direction in observed_directions
), payloads
rows: Final = _spend_rows(model, 2)
guarded_rows: Final = tuple(row for row in rows if _guardrail_entries(row))
assert len(guarded_rows) == 1, rows
_assert_response_id("chat", str(guarded_rows[0]["request_id"]), result.response_id, scenario_id)
entries: Final = _guardrail_entries(guarded_rows[0])
assert tuple(
(entry["guardrail_name"], entry["guardrail_mode"], entry["guardrail_status"])
for entry in entries
) == tuple(
(identity, "logging_only", "guardrail_intervened") for _ in expected_directions
), entries
finally:
delete_scenario(upstream_handle)
@pytest.mark.parametrize(
("row_id", "continue_flag", "expected_directions", "expected_statuses"),
(
pytest.param(
"J7",
None,
("request",),
("guardrail_failed_to_respond",),
id="J7-raising-input-aborts-response-scan",
),
pytest.param(
"J8",
True,
("request", "response"),
("guardrail_failed_to_respond", "success"),
id="J8-flag-continues-past-raising-input",
),
),
)
def test_continue_on_input_failure_recovers_from_raising_input_scan(
gateway: Gateway,
tmp_path: Path,
row_id: str,
continue_flag: bool | None,
expected_directions: tuple[Direction, ...],
expected_statuses: tuple[str, ...],
) -> None:
identity: Final = f"logging-scope-{row_id.lower()}-{uuid.uuid4().hex}"
prompt: Final = f"synthetic outage request {identity}"
reply: Final = f"synthetic outage response {identity}"
scenario_id: Final = f"phase12-{row_id.lower()}-{uuid.uuid4().hex}"
guarded_call_id: Final = f"{scenario_id}-guarded"
upstream_handle: Final = register_scenario(
scenario_id, _provider_response("chat", scenario_id, reply, False)
)
def policy(request: Request) -> Reply:
payload: Final = JSON_OBJECT.validate_json(request.body)
if _direction(payload) == "request":
return Reply(status=500, body=b"policy edge outage")
return Reply(body=b'{"action":"NONE"}')
try:
with gateway.scenario() as scenario:
model: Final = scenario.model(
model="openai/gpt-4o-mini",
api_base=f"{upstream_handle.api_base()}/v1",
api_key="synthetic-provider-key",
)
with wire_server(policy) as guardrail:
config: Final = _configuration(
tmp_path,
identity,
guardrail.url,
None,
include_scope=False,
continue_on_input_failure=continue_flag,
)
with owned_proxy(gateway, tmp_path, {}, config=config, workers=1) as candidate:
baseline: Final = _call_client(
"openai_sync", "chat", gateway, model, prompt, False, f"{scenario_id}-baseline"
)
assert baseline.status == 200 and baseline.text == reply, baseline
assert len(_drain_upstream(gateway.upstream_url)) == 1
result: Final = _call_client(
"openai_sync", "chat", candidate, model, prompt, False, guarded_call_id
)
assert (result.status, _response_body_without_ids(result.body)) == (
baseline.status,
_response_body_without_ids(baseline.body),
), (result, baseline)
assert result.text == reply, result
upstream: Final = _drain_upstream(gateway.upstream_url)
assert len(upstream) == 1 and prompt in json.dumps(upstream[0]["body"]), upstream
eventually(
lambda: guardrail.received.qsize(),
lambda count: count >= len(expected_directions),
seconds=20,
)
calls: Final = guardrail.drain()
assert len(calls) == len(expected_directions), calls
payloads: Final = tuple(JSON_OBJECT.validate_json(call.body) for call in calls)
assert tuple(_direction(payload) for payload in payloads) == expected_directions, payloads
assert all(_policy_call_id_matches(payload, guarded_call_id) for payload in payloads), payloads
rows: Final = _spend_rows(model, 2)
guarded_rows: Final = tuple(row for row in rows if _guardrail_entries(row))
assert len(guarded_rows) == 1, rows
_assert_response_id("chat", str(guarded_rows[0]["request_id"]), result.response_id, scenario_id)
entries: Final = _guardrail_entries(guarded_rows[0])
assert tuple(entry["guardrail_status"] for entry in entries) == expected_statuses, entries
assert all(
entry["guardrail_name"] == identity and entry["guardrail_mode"] == "logging_only"
for entry in entries
), entries
finally:
delete_scenario(upstream_handle)

211
tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py
View file

@ -1,3 +1,4 @@
import asyncio
import os
import traceback
from collections.abc import Mapping
@ -10,6 +11,9 @@ from unittest import mock
from dotenv import load_dotenv
from fastapi import HTTPException, Request
from litellm.constants import PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS
from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker
load_dotenv()
import time
@ -36,6 +40,7 @@ from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.caching.caching import DualCache
from litellm.proxy._types import (
LiteLLM_ProjectTable,
NewProjectRequest,
UpdateProjectRequest,
DeleteProjectRequest,
@ -1241,6 +1246,11 @@ def _project_update_mocks(monkeypatch, stored_metadata: dict) -> mock.MagicMock:
mock_prisma.jsonify_object = lambda data: data
mock_prisma.db.litellm_projecttable.find_unique = mock.AsyncMock(return_value=existing_row)
mock_prisma.db.litellm_projecttable.update = mock.AsyncMock(return_value=mock.MagicMock())
mock_prisma.writer_db = mock.MagicMock()
mock_prisma.writer_db.litellm_projecttable.find_unique = mock.AsyncMock(
return_value={"project_id": "project-update-test", "team_id": None}
)
mock_prisma.writer_db.litellm_verificationtoken.count = mock.AsyncMock(return_value=0)
monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True)
monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma)
@ -1264,6 +1274,207 @@ def _written_project_data(mock_prisma: mock.MagicMock) -> dict:
return mock_prisma.db.litellm_projecttable.update.await_args.kwargs["data"]
@pytest.mark.asyncio
async def test_update_project_object_permission_validation_precedes_budget_write(
monkeypatch: pytest.MonkeyPatch,
) -> None:
project_id: Final = "project-object-permission-validation"
mock_prisma: Final = _project_update_mocks(monkeypatch, {})
mock_prisma.db.litellm_projecttable.find_unique.return_value.budget_id = "budget-project"
budget_table: Final = mock.MagicMock()
budget_table.update = mock.AsyncMock()
mock_prisma.db.litellm_budgettable = budget_table
def jsonify_object_permission_as_string(payload: dict[str, object]) -> dict[str, object]:
return {**payload, "object_permission": '{"vector_stores": ["replacement-store"]}'}
mock_prisma.jsonify_object = jsonify_object_permission_as_string
with pytest.raises(ProxyException) as error:
await _run_project_update(
project_id,
max_budget=50,
object_permission={"vector_stores": ["replacement-store"]},
)
assert error.value.code == "500"
assert "Input should be a valid dictionary" in error.value.message
assert "input_type=str" in error.value.message
budget_table.update.assert_not_awaited()
mock_prisma.db.litellm_projecttable.update.assert_not_awaited()
@pytest.mark.asyncio
async def test_update_project_rejects_move_when_attached_teamless_key_exists(
monkeypatch: pytest.MonkeyPatch,
) -> None:
project_id: Final = "project-teamless-key"
destination_team_id: Final = "team-b"
mock_prisma: Final = _project_update_mocks(monkeypatch, {})
mock_prisma.db.litellm_projecttable.find_unique.return_value.team_id = "team-a"
mock_prisma.db.litellm_projecttable.find_unique.return_value.budget_id = "budget-project"
budget_table: Final = mock.MagicMock()
budget_table.update = mock.AsyncMock()
mock_prisma.db.litellm_budgettable = budget_table
mock_prisma.db.litellm_teamtable.find_unique = mock.AsyncMock(
return_value=LiteLLM_TeamTable(team_id=destination_team_id)
)
mock_prisma.db.litellm_verificationtoken.count = mock.AsyncMock(return_value=0)
mock_prisma.writer_db = mock.MagicMock()
mock_prisma.writer_db.litellm_projecttable.find_unique = mock.AsyncMock(
return_value=LiteLLM_ProjectTable(project_id=project_id, team_id="team-a")
)
async def count_teamless_keys(*, where: Mapping[str, object]) -> int:
conditions: Final = where.get("OR")
return int(isinstance(conditions, list) and {"team_id": None} in conditions)
mock_prisma.writer_db.litellm_verificationtoken.count = mock.AsyncMock(side_effect=count_teamless_keys)
with pytest.raises(ProxyException) as error:
await _run_project_update(
project_id,
team_id=destination_team_id,
max_budget=50,
)
expected_detail: Final = {
"error": (
f"Project {project_id} has 1 key(s) that do not belong to team {destination_team_id}. "
"Detach or delete them before moving the project."
)
}
assert error.value.code == "400"
assert expected_detail["error"] in error.value.message
mock_prisma.writer_db.litellm_verificationtoken.count.assert_awaited_once()
mock_prisma.db.litellm_verificationtoken.count.assert_not_awaited()
budget_table.update.assert_not_awaited()
mock_prisma.db.litellm_projecttable.update.assert_not_awaited()
@pytest.mark.asyncio
async def test_update_project_rejects_move_when_writer_team_differs_from_stale_reader(
monkeypatch: pytest.MonkeyPatch,
) -> None:
project_id: Final = "project-replica-lag"
destination_team_id: Final = "team-a"
mock_prisma: Final = _project_update_mocks(monkeypatch, {})
mock_prisma.db.litellm_projecttable.find_unique.return_value.team_id = destination_team_id
mock_prisma.db.litellm_teamtable.find_unique = mock.AsyncMock(
return_value=LiteLLM_TeamTable(team_id=destination_team_id)
)
mock_prisma.db.litellm_verificationtoken.count = mock.AsyncMock(return_value=0)
mock_prisma.writer_db.litellm_projecttable.find_unique = mock.AsyncMock(
return_value=LiteLLM_ProjectTable(project_id=project_id, team_id="team-b")
)
mock_prisma.writer_db.litellm_verificationtoken.count = mock.AsyncMock(return_value=1)
with pytest.raises(ProxyException) as error:
await _run_project_update(project_id, team_id=destination_team_id)
assert error.value.code == "400"
assert (
f"Project {project_id} has 1 key(s) that do not belong to team {destination_team_id}. "
"Detach or delete them before moving the project."
) in error.value.message
mock_prisma.writer_db.litellm_projecttable.find_unique.assert_awaited_once()
mock_prisma.writer_db.litellm_verificationtoken.count.assert_awaited_once()
mock_prisma.db.litellm_verificationtoken.count.assert_not_awaited()
mock_prisma.db.litellm_projecttable.update.assert_not_awaited()
@pytest.mark.asyncio
async def test_update_project_slow_writer_ownership_reads_do_not_stall_db_tracker(
monkeypatch: pytest.MonkeyPatch,
) -> None:
project_id: Final = "project-slow-writer"
source_team_id: Final = "team-source"
destination_team_id: Final = "team-destination"
mock_prisma: Final = _project_update_mocks(monkeypatch, {})
mock_prisma.db.litellm_projecttable.find_unique.return_value.team_id = source_team_id
mock_prisma.db.litellm_teamtable.find_unique = mock.AsyncMock(
return_value=LiteLLM_TeamTable(team_id=destination_team_id, models=[])
)
async def slow_project_lookup(*, where: Mapping[str, object]) -> LiteLLM_ProjectTable:
assert where == {"project_id": project_id}
await asyncio.sleep(0)
await asyncio.sleep(0)
return LiteLLM_ProjectTable(project_id=project_id, team_id=source_team_id)
async def slow_key_count(*, where: Mapping[str, object]) -> int:
assert where == {
"project_id": project_id,
"OR": [{"team_id": {"not": destination_team_id}}, {"team_id": None}],
}
await asyncio.sleep(0)
await asyncio.sleep(0)
return 0
mock_prisma.writer_db.litellm_projecttable.find_unique = mock.AsyncMock(side_effect=slow_project_lookup)
mock_prisma.writer_db.litellm_verificationtoken.count = mock.AsyncMock(side_effect=slow_key_count)
monkeypatch.setattr("litellm.proxy.db.db_lookup_gate.PROXY_DB_LOOKUP_DEADLINE_SECONDS", 0.0)
db_lookup_stall_tracker.clear()
try:
await _run_project_update(project_id, team_id=destination_team_id)
assert mock_prisma.writer_db.litellm_projecttable.find_unique.await_count == 1
assert mock_prisma.writer_db.litellm_verificationtoken.count.await_count == 1
mock_prisma.db.litellm_projecttable.update.assert_awaited_once()
assert db_lookup_stall_tracker.stalled_within(PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS) is False
finally:
db_lookup_stall_tracker.clear()
@pytest.mark.asyncio
async def test_update_project_team_limit_error_precedes_mismatched_key_guard(
monkeypatch: pytest.MonkeyPatch,
) -> None:
project_id: Final = "project-team-limit-precedence"
source_team_id: Final = "team-source"
destination_team_id: Final = "team-destination"
mock_prisma: Final = _project_update_mocks(monkeypatch, {})
mock_prisma.db.litellm_projecttable.find_unique.return_value.team_id = source_team_id
mock_prisma.db.litellm_teamtable.find_unique = mock.AsyncMock(
return_value=LiteLLM_TeamTable(
team_id=destination_team_id,
models=["allowed-model"],
)
)
mock_prisma.writer_db.litellm_projecttable.find_unique = mock.AsyncMock(
return_value=LiteLLM_ProjectTable(project_id=project_id, team_id=source_team_id)
)
mock_prisma.writer_db.litellm_verificationtoken.count = mock.AsyncMock(return_value=1)
with pytest.raises(ProxyException, match="not in team's allowed models") as error:
await _run_project_update(project_id, team_id=destination_team_id, models=["disallowed-model"])
assert error.value.code == "400"
mock_prisma.writer_db.litellm_verificationtoken.count.assert_not_awaited()
mock_prisma.db.litellm_projecttable.update.assert_not_awaited()
@pytest.mark.asyncio
async def test_update_project_allows_move_when_no_attached_keys_exist(monkeypatch: pytest.MonkeyPatch) -> None:
project_id: Final = "project-without-keys"
destination_team_id: Final = "team-b"
mock_prisma: Final = _project_update_mocks(monkeypatch, {})
mock_prisma.db.litellm_projecttable.find_unique.return_value.team_id = "team-a"
mock_prisma.db.litellm_teamtable.find_unique = mock.AsyncMock(
return_value=LiteLLM_TeamTable(team_id=destination_team_id)
)
mock_prisma.db.litellm_verificationtoken.count = mock.AsyncMock(return_value=0)
mock_prisma.writer_db.litellm_projecttable.find_unique = mock.AsyncMock(
return_value=LiteLLM_ProjectTable(project_id=project_id, team_id="team-a")
)
await _run_project_update(project_id, team_id=destination_team_id)
mock_prisma.writer_db.litellm_verificationtoken.count.assert_awaited_once()
mock_prisma.db.litellm_verificationtoken.count.assert_not_awaited()
mock_prisma.db.litellm_projecttable.update.assert_awaited_once()
@pytest.mark.asyncio
async def test_update_project_clears_model_itpm_limit_sent_as_an_empty_map(monkeypatch):
"""

62
tests/unit/integrations/test_custom_guardrail.py
View file

@ -2812,22 +2812,24 @@ class TestLoggingOnlyApplyGuardrail:
assert kwargs["standard_logging_object"] == {"guardrail_information": None}
@pytest.mark.parametrize(
"scope,expected_calls",
"scope,continue_on_input_failure,expected_calls",
(
(None, [("request", ["hello there"]), ("response", ["general kenobi"])]),
("both", [("request", ["hello there"]), ("response", ["general kenobi"])]),
("input", [("request", ["hello there"])]),
("output", [("response", ["general kenobi"])]),
(None, False, [("request", ["hello there"]), ("response", ["general kenobi"])]),
(None, True, [("request", ["hello there"]), ("response", ["general kenobi"])]),
("input", False, [("request", ["hello there"])]),
("output", False, [("response", ["general kenobi"])]),
),
)
@pytest.mark.asyncio
async def test_logging_only_scope_scans_configured_directions(
self,
scope: LoggingOnlyScope | None,
continue_on_input_failure: bool,
expected_calls: list[tuple[str, list[str]]],
) -> None:
guardrail: Final = _ApplyOnlyObserver()
guardrail.logging_only_scope = scope
guardrail.logging_only_continue_on_input_failure = continue_on_input_failure
kwargs, response = _logged_call([{"role": "user", "content": "hello there"}])
out_kwargs, _ = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value)
@ -2881,21 +2883,37 @@ class TestLoggingOnlyApplyGuardrail:
assert [entry["guardrail_name"] for entry in entries] == ["apply-only-observer"]
assert [entry["guardrail_status"] for entry in entries] == ["success"]
@pytest.mark.parametrize(
"scope,continue_on_input_failure,expected_calls,expected_statuses",
(
(None, True, [("response", ["general kenobi"])], ["success"]),
("input", True, [], []),
("output", True, [("response", ["general kenobi"])], ["success"]),
(None, False, [], []),
),
)
@pytest.mark.asyncio
async def test_request_copy_failure_does_not_drop_the_response_scan_for_explicit_both_scope(self):
async def test_request_copy_failure_drops_the_response_scan_only_without_continue(
self,
scope: LoggingOnlyScope | None,
continue_on_input_failure: bool,
expected_calls: list[tuple[str, list[str]]],
expected_statuses: list[str],
):
import threading
guardrail: Final = _ApplyOnlyObserver()
guardrail.logging_only_scope = "both"
guardrail.logging_only_scope = scope
guardrail.logging_only_continue_on_input_failure = continue_on_input_failure
call: Final = _logged_call([{"role": "user", "content": "hello there", "lock": threading.Lock()}])
kwargs: Final = call[0]
response: Final = call[1]
out_kwargs, _ = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value)
assert guardrail.calls == [("response", ["general kenobi"])]
entries: Final = out_kwargs["standard_logging_object"]["guardrail_information"]
assert [entry["guardrail_status"] for entry in entries] == ["success"]
assert guardrail.calls == expected_calls
entries: Final = out_kwargs["standard_logging_object"].get("guardrail_information") or []
assert [entry["guardrail_status"] for entry in entries] == expected_statuses
@pytest.mark.asyncio
async def test_block_verdict_is_recorded_without_raising(self):
@ -2927,8 +2945,23 @@ class TestLoggingOnlyApplyGuardrail:
entries = out_kwargs["standard_logging_object"]["guardrail_information"]
assert [e["guardrail_status"] for e in entries] == ["guardrail_failed_to_respond"]
@pytest.mark.parametrize(
"scope,continue_on_input_failure,expected_calls,expected_statuses",
(
(None, True, [("request", ["hello there"]), ("response", ["general kenobi"])], ["guardrail_failed_to_respond", "success"]),
("output", True, [("response", ["general kenobi"])], ["success"]),
("input", True, [("request", ["hello there"])], ["guardrail_failed_to_respond"]),
(None, False, [("request", ["hello there"])], ["guardrail_failed_to_respond"]),
),
)
@pytest.mark.asyncio
async def test_input_scan_error_does_not_drop_the_response_scan_for_explicit_both_scope(self):
async def test_input_scan_error_drops_the_response_scan_only_without_continue(
self,
scope: LoggingOnlyScope | None,
continue_on_input_failure: bool,
expected_calls: list[tuple[str, list[str]]],
expected_statuses: list[str],
):
class _FailingBothObserver(_ApplyOnlyObserver):
@log_guardrail_information
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
@ -2938,14 +2971,15 @@ class TestLoggingOnlyApplyGuardrail:
return GenericGuardrailAPIInputs(texts=[])
guardrail = _FailingBothObserver()
guardrail.logging_only_scope = "both"
guardrail.logging_only_scope = scope
guardrail.logging_only_continue_on_input_failure = continue_on_input_failure
kwargs, response = _logged_call([{"role": "user", "content": "hello there"}])
out_kwargs, _ = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value)
assert guardrail.calls == [("request", ["hello there"]), ("response", ["general kenobi"])]
assert guardrail.calls == expected_calls
entries = out_kwargs["standard_logging_object"]["guardrail_information"]
assert [e["guardrail_status"] for e in entries] == ["guardrail_failed_to_respond", "success"]
assert [e["guardrail_status"] for e in entries] == expected_statuses
@pytest.mark.asyncio
async def test_call_type_without_translation_is_skipped(self):

54
tests/unit/proxy/auth/test_auth_utils.py
View file

@ -64,6 +64,60 @@ def test_token_exchange_settings_in_request_body_are_rejected(field: str) -> Non
assert field in str(error.value)
@pytest.mark.parametrize("field", oauth_token_exchange_litellm_params)
@pytest.mark.parametrize("route", ("/v1/mcp/server", "/v1/mcp/server/abc123"))
def test_mcp_server_management_body_may_set_token_exchange_settings(field: str, route: str) -> None:
assert (
is_request_body_safe(
request_body={"server_id": "abc123", "auth_type": "oauth2_token_exchange", field: "configured-value"},
general_settings={},
llm_router=None,
model="",
route=route,
)
is True
)
@pytest.mark.parametrize("field", oauth_token_exchange_litellm_params)
@pytest.mark.parametrize("route", ("/chat/completions", "/v1/mcp/serverfoo"))
def test_token_exchange_settings_remain_banned_outside_mcp_server_management(field: str, route: str) -> None:
with pytest.raises(ValueError, match=field):
is_request_body_safe(
request_body={"server_id": "abc123", "auth_type": "oauth2_token_exchange", field: "configured-value"},
general_settings={},
llm_router=None,
model="",
route=route,
)
def test_mcp_server_management_body_still_rejects_anthropic_wif_settings() -> None:
with pytest.raises(ValueError, match="anthropic_federation_rule_id"):
is_request_body_safe(
request_body={"server_id": "abc123", "anthropic_federation_rule_id": "fdrl_admin"},
general_settings={},
llm_router=None,
model="",
route="/v1/mcp/server",
)
def test_mcp_server_management_body_still_rejects_nested_token_exchange_settings() -> None:
with pytest.raises(ValueError, match="token_exchange_endpoint"):
is_request_body_safe(
request_body={
"server_id": "abc123",
"auth_type": "oauth2_token_exchange",
"metadata": {"token_exchange_endpoint": "https://identity.example.com/token"},
},
general_settings={},
llm_router=None,
model="",
route="/v1/mcp/server",
)
def test_model_opt_in_cannot_allow_a_token_exchange_endpoint_in_a_request_body() -> None:
router: Final = Router(
model_list=[

157
tests/unit/proxy/guardrails/test_guardrail_endpoints.py
View file

@ -848,6 +848,40 @@ def test_get_guardrails_list_response_tolerates_invalid_config_stream_scope():
assert response.guardrails[1].litellm_params.stream_scope == "streaming"
def test_get_guardrails_list_response_tolerates_invalid_config_logging_only_scope():
from litellm.proxy.guardrails.guardrail_endpoints import (
_get_guardrails_list_response,
)
response = _get_guardrails_list_response(
[
{
"guardrail_id": "removed-both-scope",
"guardrail_name": "removed-both-scope",
"litellm_params": {
"guardrail": "generic_guardrail_api",
"mode": "logging_only",
"logging_only_scope": "both",
},
},
{
"guardrail_id": "input-scope",
"guardrail_name": "input-scope",
"litellm_params": {
"guardrail": "generic_guardrail_api",
"mode": "logging_only",
"logging_only_scope": "input",
},
},
]
)
assert response.guardrails[0].litellm_params is not None
assert response.guardrails[0].litellm_params.logging_only_scope is None
assert response.guardrails[1].litellm_params is not None
assert response.guardrails[1].litellm_params.logging_only_scope == "input"
def test_get_provider_specific_params():
"""Test getting provider-specific parameters"""
from litellm.proxy.guardrails.guardrail_endpoints import _get_fields_from_model
@ -1727,6 +1761,129 @@ async def test_patch_guardrail_rejects_explicit_unsupported_scope_and_rolls_back
handler.delete_in_memory_guardrail(stored_guardrail["guardrail_id"])
@pytest.mark.asyncio
async def test_patch_guardrail_rejects_invalid_continue_flag_with_422(mocker, monkeypatch, mock_guardrail_registry):
handler, stored_guardrail = _setup_patch_scope_guardrail(
mocker,
monkeypatch,
mock_guardrail_registry,
_PatchScopeSupportedGuardrail,
"patch_scope_supported_test",
{"mode": "pre_call", "default_on": True},
)
request = PatchGuardrailRequest(
litellm_params=BaseLitellmParams(mode="pre_call", logging_only_continue_on_input_failure=True)
)
try:
with pytest.raises(HTTPException) as exc_info:
await patch_guardrail(
stored_guardrail["guardrail_id"],
request,
user_api_key_dict=MOCK_ADMIN_USER,
)
assert exc_info.value.status_code == 422
assert "update rejected" in str(exc_info.value.detail)
assert "logging_only_continue_on_input_failure is set, but mode does not include logging_only" in str(
exc_info.value.detail
)
assert mock_guardrail_registry.update_guardrail_in_db.call_count == 2
restored_guardrail = mock_guardrail_registry.update_guardrail_in_db.call_args_list[-1].kwargs["guardrail"]
assert restored_guardrail["litellm_params"] == stored_guardrail["litellm_params"]
callback = handler.guardrail_id_to_custom_guardrail[stored_guardrail["guardrail_id"]]
assert callback.logging_only_continue_on_input_failure is False
finally:
handler.delete_in_memory_guardrail(stored_guardrail["guardrail_id"])
@pytest.mark.asyncio
async def test_patch_guardrail_clears_continue_flag_when_logging_only_mode_is_removed(
mocker, monkeypatch, mock_guardrail_registry
):
handler, stored_guardrail = _setup_patch_scope_guardrail(
mocker,
monkeypatch,
mock_guardrail_registry,
_PatchScopeSupportedGuardrail,
"patch_scope_supported_test",
{
"mode": ["pre_call", "logging_only"],
"logging_only_continue_on_input_failure": True,
"default_on": True,
},
)
request = PatchGuardrailRequest(litellm_params=BaseLitellmParams(mode=["pre_call"]))
try:
result = await patch_guardrail(
stored_guardrail["guardrail_id"],
request,
user_api_key_dict=MOCK_ADMIN_USER,
)
assert result["guardrail_id"] == stored_guardrail["guardrail_id"]
persisted_guardrail = mock_guardrail_registry.update_guardrail_in_db.call_args.kwargs["guardrail"]
assert persisted_guardrail["litellm_params"].logging_only_continue_on_input_failure is None
callback = handler.guardrail_id_to_custom_guardrail[stored_guardrail["guardrail_id"]]
assert callback.logging_only_continue_on_input_failure is False
finally:
handler.delete_in_memory_guardrail(stored_guardrail["guardrail_id"])
@pytest.mark.asyncio
async def test_add_guardrail_rejects_continue_flag_outside_logging_only_with_400(
mocker, monkeypatch, mock_guardrail_registry
):
import litellm
from litellm.proxy.guardrails import guardrail_registry as registry_module
guardrail_type: Final = "patch_scope_supported_test"
litellm_params: Final = {
"guardrail": guardrail_type,
"mode": "pre_call",
"logging_only_continue_on_input_failure": True,
}
mock_guardrail_registry.add_guardrail_to_db = AsyncMock(
return_value={
"guardrail_id": "flag-400-test",
"guardrail_name": "flag-400-test",
"litellm_params": litellm_params,
"guardrail_info": {},
}
)
monkeypatch.setitem(
registry_module.guardrail_initializer_registry,
guardrail_type,
_patch_scope_initializer(_PatchScopeSupportedGuardrail),
)
handler: Final = InMemoryGuardrailHandler()
mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.Mock())
mocker.patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", handler)
delete_row = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints._delete_guardrail_row", AsyncMock())
with pytest.raises(HTTPException) as exc_info:
await create_guardrail(
CreateGuardrailRequest(
guardrail={
"guardrail_name": "flag-400-test",
"litellm_params": litellm_params,
"guardrail_info": {},
}
),
user_api_key_dict=MOCK_ADMIN_USER,
)
assert exc_info.value.status_code == 400
assert "logging_only_continue_on_input_failure is set, but mode does not include logging_only" in str(
exc_info.value.detail
)
assert "flag-400-test" not in handler.guardrail_id_to_custom_guardrail
assert all(getattr(callback, "guardrail_name", None) != "flag-400-test" for callback in litellm.callbacks)
delete_row.assert_awaited_once()
@pytest.mark.asyncio
async def test_patch_guardrail_rejected_update_restores_invalid_stored_scope_verbatim(
mocker, monkeypatch, mock_guardrail_registry

94
tests/unit/proxy/guardrails/test_guardrail_registry.py
View file

@ -1273,6 +1273,7 @@ class TestLoggingOnlyScopeValidation:
callback_type: type[CustomGuardrail] = _LoggingOnlyScopeSupportedGuardrail,
reject_invalid_logging_only_scope: bool = False,
assert_registered: bool = False,
continue_on_input_failure: bool | None = None,
) -> CustomGuardrail:
import litellm
from litellm.proxy.guardrails import guardrail_registry as registry_module
@ -1306,6 +1307,7 @@ class TestLoggingOnlyScopeValidation:
"guardrail": guardrail_type,
"mode": mode,
"logging_only_scope": scope,
"logging_only_continue_on_input_failure": continue_on_input_failure,
},
},
reject_invalid_logging_only_scope=reject_invalid_logging_only_scope,
@ -1377,14 +1379,16 @@ class TestLoggingOnlyScopeValidation:
"guardrail, whose logging_only hook scans on its own. Remove logging_only_scope."
)
def test_both_scope_accepted_when_guardrail_owns_logging_hook(self) -> None:
def test_continue_flag_accepted_when_guardrail_owns_logging_hook(self) -> None:
callback: Final = self._initialize(
mode="logging_only",
scope="both",
scope=None,
continue_on_input_failure=True,
callback_type=_LoggingOnlyScopeUnsupportedGuardrail,
)
assert callback.logging_only_scope == "both"
assert callback.logging_only_scope is None
assert callback.logging_only_continue_on_input_failure is True
def test_output_scope_accepted_for_native_lifecycle_guardrail(self) -> None:
callback: Final = self._initialize(
@ -1395,9 +1399,10 @@ class TestLoggingOnlyScopeValidation:
assert callback.logging_only_scope == "output"
def test_invalid_scope_fails_litellm_params_validation(self) -> None:
@pytest.mark.parametrize("scope", ("request", "both"))
def test_invalid_scope_fails_litellm_params_validation(self, scope: str) -> None:
with pytest.raises(ValidationError):
LitellmParams(guardrail="test", mode="logging_only", logging_only_scope="request")
LitellmParams(guardrail="test", mode="logging_only", logging_only_scope=scope)
def test_invalid_scope_literal_keeps_content_filter_registered_and_blocking(self) -> None:
import litellm
@ -1438,6 +1443,85 @@ class TestLoggingOnlyScopeValidation:
assert all(callback_list == snapshot for callback_list, snapshot in zip(callback_lists, callback_snapshots))
@pytest.mark.parametrize(
("scope", "continue_on_input_failure", "expected_direction", "expected_continue"),
(
(None, None, None, False),
(None, True, None, True),
(None, False, None, False),
("input", True, "input", False),
("output", True, "output", False),
("input", None, "input", False),
),
)
def test_logging_only_settings_table(
self,
scope: LoggingOnlyScope | None,
continue_on_input_failure: bool | None,
expected_direction: str | None,
expected_continue: bool,
) -> None:
callback: Final = self._initialize(
mode="logging_only",
scope=scope,
continue_on_input_failure=continue_on_input_failure,
)
assert callback.logging_only_scope == expected_direction
assert callback.logging_only_continue_on_input_failure is expected_continue
def test_continue_flag_without_logging_only_mode_is_ignored_at_load(self) -> None:
callback: Final = self._initialize(
mode="pre_call", continue_on_input_failure=True, scope=None, assert_registered=True
)
assert callback.logging_only_scope is None
assert callback.logging_only_continue_on_input_failure is False
assert callback.should_run_guardrail(data={}, event_type=GuardrailEventHooks.pre_call) is True
def test_continue_flag_without_logging_only_mode_is_rejected_for_api_writes(self) -> None:
with pytest.raises(ValueError, match="logging_only_continue_on_input_failure is set") as exc_info:
self._initialize(
mode="pre_call",
scope=None,
continue_on_input_failure=True,
reject_invalid_logging_only_scope=True,
)
assert str(exc_info.value) == (
"Guardrail logging-only-scope-guardrail: logging_only_continue_on_input_failure is set, but mode "
"does not include logging_only, so it would never apply. Add logging_only to mode or remove "
"logging_only_continue_on_input_failure."
)
def test_continue_flag_is_accepted_when_guardrail_owns_logging_hook(self) -> None:
callback: Final = self._initialize(
mode="logging_only",
scope=None,
continue_on_input_failure=True,
callback_type=_LoggingOnlyScopeUnsupportedGuardrail,
)
assert callback.logging_only_scope is None
assert callback.logging_only_continue_on_input_failure is True
def test_invalid_continue_flag_is_tolerated_by_parse_tolerant(self) -> None:
params: Final = {
"guardrail": "litellm_content_filter",
"mode": "logging_only",
"logging_only_continue_on_input_failure": "banana",
}
parsed: Final = parse_tolerant_litellm_params(params, "invalid-flag-test")
assert parsed.logging_only_continue_on_input_failure is None
def test_invalid_continue_flag_raises_without_tolerance(self) -> None:
with pytest.raises(ValidationError):
LitellmParams(
guardrail="test", mode="logging_only", logging_only_continue_on_input_failure="banana"
)
def test_invalid_scope_literal_does_not_tolerate_other_litellm_params_errors(self) -> None:
with pytest.raises(ValidationError):
parse_tolerant_litellm_params(

1463
tests/unit/proxy/management_endpoints/test_key_management_endpoints.py
View file

File diff suppressed because it is too large Load diff

60
tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py
View file

@ -9,6 +9,7 @@ import sqlite3
from datetime import timezone
from unittest.mock import AsyncMock, MagicMock, patch
import prisma
import pytest
from fastapi import HTTPException
from fastapi.testclient import TestClient
@ -277,6 +278,7 @@ from litellm.proxy._types import (
from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
from litellm.proxy.management.teams import authz as team_access
from litellm.proxy.proxy_server import app
from litellm.proxy.utils import PrismaClient
from litellm.proxy.spend_tracking import spend_management_endpoints
from litellm.router import Router
from litellm.types.utils import BudgetConfig
@ -7999,3 +8001,61 @@ async def test_management_team_lookup_without_memberships_keeps_own_user_scope()
scope = await resolve_owned_read_scope(auth.user_id, lookup)
assert scope == OwnedRows("caller")
assert team_reads == []
class _PrismaJournal:
def __init__(self) -> None:
self.statements: tuple[tuple[str, str], ...] = ()
def __call__(self, **kwargs: object) -> "_ReadOnlyReplicaPrisma":
return _ReadOnlyReplicaPrisma(self, "reader" if "datasource" in kwargs else "writer")
def record(self, role: str, query: str) -> None:
self.statements = (*self.statements, (role, query.strip()))
class _ReadOnlyReplicaPrisma:
def __init__(self, journal: _PrismaJournal, role: str) -> None:
self.journal = journal
self.role = role
def is_connected(self) -> bool:
return False
async def connect(self, timeout: object = None) -> None:
return None
def _run(self, query: str) -> None:
self.journal.record(self.role, query)
if self.role == "reader" and not query.lstrip().upper().startswith("SELECT"):
raise Exception(f"cannot execute {query.split()[0]} in a read-only transaction")
async def query_raw(self, query: str, *args: object) -> list[dict[str, str]]:
self._run(query)
return [{"relname": "MonthlyGlobalSpend", "relkind": "m"}]
async def execute_raw(self, query: str, *args: object) -> int:
self._run(query)
return 0
@pytest.mark.asyncio
async def test_global_spend_refresh_runs_the_refresh_on_the_writer_when_a_reader_is_configured(
monkeypatch: pytest.MonkeyPatch,
) -> None:
journal: Final = _PrismaJournal()
monkeypatch.setattr(prisma, "Prisma", journal)
monkeypatch.setenv("DATABASE_URL", "postgresql://writer@localhost/litellm")
monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://reader@localhost/litellm")
monkeypatch.setattr(
ps,
"prisma_client",
PrismaClient(database_url="postgresql://writer@localhost/litellm", proxy_logging_obj=ps.proxy_logging_obj),
)
response: Final = await spend_management_endpoints.global_spend_refresh()
assert response == {"message": "MonthlyGlobalSpend view refreshed", "status": "success"}, journal.statements
assert [statement for statement in journal.statements if statement[1].startswith("REFRESH")] == [
("writer", 'REFRESH MATERIALIZED VIEW "MonthlyGlobalSpend";')
]

58
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailFormField.tsx
View file

@ -2,9 +2,16 @@
import { CircleHelp } from "lucide-react";
import React, { useEffect, useId } from "react";
import { useController, type Control, type ControllerRenderProps, type RegisterOptions } from "react-hook-form";
import {
useController,
useWatch,
type Control,
type ControllerRenderProps,
type RegisterOptions,
} from "react-hook-form";
import { Field, FieldDescription, FieldError, FieldLabel } from "@/components/ui/field";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { Switch } from "@/components/ui/switch";
import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
import {
getLoggingOnlyScopeOptions,
@ -22,6 +29,7 @@ export interface GuardrailCriterion {
export interface GuardrailFormValues extends Record<string, unknown> {
criteria?: GuardrailCriterion[];
logging_only_scope_choice?: LoggingOnlyScopeChoice;
logging_only_continue_on_input_failure?: boolean;
}
export type GuardrailFormControl = Control<GuardrailFormValues>;
export type GuardrailFieldRules = Pick<RegisterOptions<GuardrailFormValues, string>, "validate">;
@ -172,20 +180,44 @@ export const LoggingOnlyScopeField: React.FC<{
mode: unknown;
directionalScopeSupported: boolean;
}> = ({ control, mode, directionalScopeSupported }) => {
const scopeChoice = useWatch({ control, name: "logging_only_scope_choice" });
const continueDisabled = scopeChoice === "input" || scopeChoice === "output";
if (!modeIncludesLoggingOnly(mode)) return null;
return (
<GuardrailField
control={control}
name="logging_only_scope_choice"
label={labelWithHint(
"Logging only scope",
"Which direction a logging_only scan observes. Observe-only scans never block; pre_call and post_call on this guardrail still block.",
)}
>
{(fieldControl) => (
<LoggingOnlyScopeSelect control={fieldControl} directionalScopeSupported={directionalScopeSupported} />
)}
</GuardrailField>
<>
<GuardrailField
control={control}
name="logging_only_scope_choice"
label={labelWithHint(
"Logging only scope",
"Which direction a logging_only scan observes. Observe-only scans never block; pre_call and post_call on this guardrail still block.",
)}
>
{(fieldControl) => (
<LoggingOnlyScopeSelect control={fieldControl} directionalScopeSupported={directionalScopeSupported} />
)}
</GuardrailField>
<GuardrailField
control={control}
name="logging_only_continue_on_input_failure"
label={labelWithHint(
"Continue observing the response after a flagged request",
"When on, a flagged or failing request scan is logged and the response is still scanned, so both verdicts land.",
)}
>
{(fieldControl) => (
<Switch
id={fieldControl.id}
checked={!continueDisabled && fieldControl.value === true}
disabled={continueDisabled}
onCheckedChange={fieldControl.onChange}
aria-invalid={fieldControl["aria-invalid"]}
aria-describedby={fieldControl["aria-describedby"]}
aria-label="Continue observing the response after a flagged request"
/>
)}
</GuardrailField>
</>
);
};

12
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailModeDisplay.tsx
View file

@ -1,12 +1,18 @@
import React from "react";
import { Badge } from "@/components/ui/badge";
import { Card } from "@/components/ui/card";
import { formatGuardrailMode, formatLoggingOnlyScope, modeIncludesLoggingOnly } from "./guardrail_info_helpers";
import {
formatGuardrailMode,
formatLoggingOnlyScope,
loggingOnlyContinueFromParams,
modeIncludesLoggingOnly,
} from "./guardrail_info_helpers";
type GuardrailModeParams = {
mode?: unknown;
default_on?: boolean;
logging_only_scope?: string | null;
logging_only_continue_on_input_failure?: boolean | null;
};
export const GuardrailModeCard: React.FC<{ litellmParams: GuardrailModeParams }> = ({ litellmParams }) => (
@ -22,6 +28,8 @@ export const GuardrailModeCard: React.FC<{ litellmParams: GuardrailModeParams }>
<div className="mt-4">
<p>Logging only scope</p>
<h3 className="text-lg font-medium">{formatLoggingOnlyScope(litellmParams.logging_only_scope)}</h3>
<p className="mt-2">Continue after a flagged request</p>
<h3 className="text-lg font-medium">{loggingOnlyContinueFromParams(litellmParams) ? "On" : "Off"}</h3>
</div>
)}
</Card>
@ -37,6 +45,8 @@ export const GuardrailModeRows: React.FC<{ litellmParams: GuardrailModeParams }>
<div>
<p className="font-medium">Logging only scope</p>
<div>{formatLoggingOnlyScope(litellmParams.logging_only_scope)}</div>
<p className="font-medium">Continue after a flagged request</p>
<div>{loggingOnlyContinueFromParams(litellmParams) ? "On" : "Off"}</div>
</div>
)}
</>

6
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailReadOnlyDetails.tsx
View file

@ -19,7 +19,11 @@ export const GuardrailReadOnlyDetails = ({
guardrailId: string;
guardrailName: string;
displayName: string;
litellmParams: { mode?: unknown; logging_only_scope?: string | null };
litellmParams: {
mode?: unknown;
logging_only_scope?: string | null;
logging_only_continue_on_input_failure?: boolean | null;
};
streamScope: unknown;
defaultOn: boolean | undefined;
piiEntityCount: number;

22
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.characterization.test.tsx
View file

@ -145,6 +145,26 @@ describe("AddGuardrailForm create payload characterization", () => {
expect(payload()?.litellm_params.logging_only_scope).toBe("output");
});
it("disables and unchecks the continue toggle under a directional scope", async () => {
vi.mocked(networking.getGuardrailUISettings).mockResolvedValue({
...uiSettings,
supported_modes: ["pre_call", "logging_only"],
});
const user = userEvent.setup({ delay: null });
renderForm();
await pickProvider(user, "Bedrock Guardrail");
await user.click(screen.getByLabelText("Mode"));
await user.click((await screen.findAllByText("logging_only")).at(-1) as HTMLElement);
await chooseSelectOption(user, await screen.findByLabelText("Logging only scope"), "Input only (request)");
const toggle = await screen.findByRole("switch", {
name: "Continue observing the response after a flagged request",
});
expect(toggle).toHaveAttribute("aria-disabled", "true");
expect(toggle).not.toBeChecked();
});
it("hides directional scope choices for providers that do not support them", async () => {
vi.mocked(networking.getGuardrailUISettings).mockResolvedValue({
...uiSettings,
@ -163,7 +183,7 @@ describe("AddGuardrailForm create payload characterization", () => {
await user.click((await screen.findAllByText("logging_only")).at(-1) as HTMLElement);
await user.click(await screen.findByLabelText("Logging only scope"));
expect(await screen.findByRole("option", { name: "Both (request and response)" })).toBeInTheDocument();
expect(await screen.findByRole("option", { name: "Default (request and response)" })).toBeInTheDocument();
expect(screen.queryByRole("option", { name: "Input only (request)" })).not.toBeInTheDocument();
expect(screen.queryByRole("option", { name: "Output only (response)" })).not.toBeInTheDocument();
});

12
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx
View file

@ -17,6 +17,7 @@ import {
getGuardrailProviders,
getSupportedModesForProvider,
guardrail_provider_map,
effectiveLoggingOnlyContinue,
modeIncludesLoggingOnly,
populateGuardrailProviderMap,
populateGuardrailProviders,
@ -173,6 +174,7 @@ const INITIAL_VALUES: GuardrailFormValues = {
mode: "pre_call",
default_on: false,
logging_only_scope_choice: "default",
logging_only_continue_on_input_failure: false,
skip_system_message_choice: "inherit",
skip_tool_message_choice: "inherit",
stream_scope_by_mode: {},
@ -294,6 +296,7 @@ const AddGuardrailForm: React.FC<AddGuardrailFormProps> = ({ visible, onClose, a
mode: preset.mode,
default_on: preset.defaultOn,
logging_only_scope_choice: "default",
logging_only_continue_on_input_failure: false,
skip_system_message_choice: "inherit",
skip_tool_message_choice: "inherit",
};
@ -494,6 +497,15 @@ const AddGuardrailForm: React.FC<AddGuardrailFormProps> = ({ visible, onClose, a
if (modeIncludesLoggingOnly(values.mode) && loggingOnlyScope !== null) {
guardrailData.litellm_params.logging_only_scope = loggingOnlyScope;
}
if (
modeIncludesLoggingOnly(values.mode) &&
effectiveLoggingOnlyContinue(
values.logging_only_scope_choice as LoggingOnlyScopeChoice | undefined,
values.logging_only_continue_on_input_failure,
)
) {
guardrailData.litellm_params.logging_only_continue_on_input_failure = true;
}
// For Presidio PII, add the entity and action configurations
if (providerKey === "PresidioPII" && selectedEntities.length > 0) {

23
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeLoggingOnlyScope.tsx
View file

@ -1,6 +1,7 @@
import React from "react";
import {
choiceToLoggingOnlyScope,
effectiveLoggingOnlyContinue,
getLoggingOnlyScopeOptions,
getLoggingOnlyScopeUpdate,
type LoggingOnlyScope,
@ -18,7 +19,7 @@ export const CustomCodeLoggingOnlyScopeSelect: React.FC<{
value: LoggingOnlyScopeChoice;
onChange: (choice: LoggingOnlyScopeChoice) => void;
}> = ({ value, onChange }) => (
<div className="w-[200px]">
<div className="w-[240px]">
<label className="mb-1 block text-xs font-medium text-muted-foreground">Logging only scope</label>
<Select
items={LOGGING_ONLY_SCOPE_ITEMS}
@ -44,15 +45,25 @@ export const CustomCodeLoggingOnlyScopeSelect: React.FC<{
export const getCustomCodeLoggingOnlyScopeCreate = (
mode: string[],
choice: LoggingOnlyScopeChoice,
): { logging_only_scope?: LoggingOnlyScope } => {
continueOnInputFailure: boolean,
): { logging_only_scope?: LoggingOnlyScope; logging_only_continue_on_input_failure?: boolean } => {
if (!mode.includes("logging_only")) return {};
const scope = choiceToLoggingOnlyScope(choice);
return scope === null ? {} : { logging_only_scope: scope };
return {
...(scope === null ? {} : { logging_only_scope: scope }),
...(effectiveLoggingOnlyContinue(choice, continueOnInputFailure)
? { logging_only_continue_on_input_failure: true }
: {}),
};
};
export const getCustomCodeLoggingOnlyScopeUpdate = (
mode: string[],
litellmParams: { logging_only_scope?: string | null } | null | undefined,
litellmParams:
| { logging_only_scope?: string | null; logging_only_continue_on_input_failure?: boolean | null }
| null
| undefined,
choice: LoggingOnlyScopeChoice,
): { logging_only_scope?: LoggingOnlyScope | null } =>
mode.includes("logging_only") ? getLoggingOnlyScopeUpdate(litellmParams, choice) : {};
continueOnInputFailure: boolean,
): { logging_only_scope?: LoggingOnlyScope | null; logging_only_continue_on_input_failure?: boolean } =>
mode.includes("logging_only") ? getLoggingOnlyScopeUpdate(litellmParams, choice, continueOnInputFailure) : {};

51
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.test.tsx
View file

@ -208,7 +208,7 @@ describe("CustomCodeModal", () => {
await waitFor(() => expect(mockUpdate).toHaveBeenCalledTimes(1));
expect(mockUpdate.mock.calls[0][2]).toMatchObject({
litellm_params: { logging_only_scope: "output" },
litellm_params: { logging_only_scope: "output", logging_only_continue_on_input_failure: false },
});
});
@ -232,6 +232,7 @@ describe("CustomCodeModal", () => {
await waitFor(() => expect(mockUpdate).toHaveBeenCalledTimes(1));
expect(mockUpdate.mock.calls[0][2]).toHaveProperty("litellm_params.logging_only_scope", null);
expect(mockUpdate.mock.calls[0][2]).toHaveProperty("litellm_params.logging_only_continue_on_input_failure", false);
});
it("should omit an unchanged logging-only scope from the edit payload", async () => {
@ -287,6 +288,54 @@ describe("CustomCodeModal", () => {
expect(mockCreate.mock.calls[0][1]).toMatchObject({
litellm_params: { logging_only_scope: "input" },
});
expect(mockCreate.mock.calls[0][1]).not.toHaveProperty("litellm_params.logging_only_continue_on_input_failure");
});
it("disables and unchecks the continue toggle under a directional scope", async () => {
const user = userEvent.setup();
renderModal();
await user.click(screen.getAllByRole("combobox")[0]);
await user.keyboard("logging_only");
await user.click(await screen.findByRole("option", { name: "logging_only" }));
await user.click(await screen.findByRole("combobox", { name: "Logging only scope" }));
await user.click(await screen.findByRole("option", { name: "Input only (request)" }));
const toggle = await screen.findByRole("switch", {
name: "Continue observing the response after a flagged request",
});
expect(toggle).toHaveAttribute("aria-disabled", "true");
expect(toggle).not.toBeChecked();
});
it("should load and turn off a stored continue flag in edit mode", async () => {
const user = userEvent.setup();
renderModal({
editData: {
guardrail_id: "g-1",
guardrail_name: "existing-guardrail",
litellm_params: {
mode: ["logging_only"],
logging_only_continue_on_input_failure: true,
custom_code: "def apply_guardrail(): pass",
},
},
});
const scopeSelect = await screen.findByRole("combobox", { name: "Logging only scope" });
expect(scopeSelect).toHaveTextContent("Default (request and response)");
const toggle = await screen.findByRole("switch", {
name: "Continue observing the response after a flagged request",
});
expect(toggle).toBeChecked();
await user.click(toggle);
await user.click(screen.getByRole("button", { name: /update guardrail/i }));
await waitFor(() => expect(mockUpdate).toHaveBeenCalledTimes(1));
expect(mockUpdate.mock.calls[0][2]).toMatchObject({
litellm_params: { logging_only_scope: null, logging_only_continue_on_input_failure: false },
});
});
it("should omit the default logging-only scope when creating a guardrail", async () => {

31
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx
View file

@ -2,7 +2,7 @@ import React, { useState, useRef, useEffect } from "react";
import { CheckCircle2, ChevronRight, Code, ExternalLink, PlayCircle, Save, Users, XCircle } from "lucide-react";
import { createGuardrailCall, updateGuardrailCall, testCustomCodeGuardrail } from "@/components/networking";
import { toast } from "@/lib/toast";
import { loggingOnlyScopeToChoice } from "../guardrail_info_helpers";
import { loggingOnlyContinueFromParams, loggingOnlyScopeToChoice } from "../guardrail_info_helpers";
import type { LoggingOnlyScope, LoggingOnlyScopeChoice } from "../guardrail_info_helpers";
import {
CustomCodeLoggingOnlyScopeSelect,
@ -84,6 +84,7 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
const [mode, setMode] = useState<string[]>(["pre_call"]);
const [streamScopeByMode, setStreamScopeByMode] = useState<Record<string, GuardrailStreamScope>>({});
const [loggingOnlyScopeChoice, setLoggingOnlyScopeChoice] = useState<LoggingOnlyScopeChoice>("default");
const [loggingOnlyContinue, setLoggingOnlyContinue] = useState(false);
const [defaultOn, setDefaultOn] = useState(false);
const [selectedTemplate, setSelectedTemplate] = useState<string>("empty");
const [code, setCode] = useState(CODE_TEMPLATES.empty.code);
@ -218,6 +219,7 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
),
);
setLoggingOnlyScopeChoice(loggingOnlyScopeToChoice(editData.litellm_params?.logging_only_scope));
setLoggingOnlyContinue(loggingOnlyContinueFromParams(editData.litellm_params));
setDefaultOn(editData.litellm_params?.default_on || false);
setCode(editData.litellm_params?.custom_code || CODE_TEMPLATES.empty.code);
setSelectedTemplate(""); // No template selected in edit mode
@ -227,6 +229,7 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
setMode(["pre_call"]);
setStreamScopeByMode({});
setLoggingOnlyScopeChoice("default");
setLoggingOnlyContinue(false);
setDefaultOn(false);
setSelectedTemplate("empty");
setCode(CODE_TEMPLATES.empty.code);
@ -284,7 +287,12 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
const updateData: any = {
litellm_params: {
custom_code: code,
...getCustomCodeLoggingOnlyScopeUpdate(mode, editData.litellm_params, loggingOnlyScopeChoice),
...getCustomCodeLoggingOnlyScopeUpdate(
mode,
editData.litellm_params,
loggingOnlyScopeChoice,
loggingOnlyContinue,
),
},
};
@ -323,7 +331,7 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
default_on: defaultOn,
custom_code: code,
...(streamScope !== undefined ? { stream_scope: streamScope } : {}),
...getCustomCodeLoggingOnlyScopeCreate(mode, loggingOnlyScopeChoice),
...getCustomCodeLoggingOnlyScopeCreate(mode, loggingOnlyScopeChoice, loggingOnlyContinue),
},
guardrail_info: {},
};
@ -428,8 +436,8 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
</DialogHeader>
{/* Top Controls */}
<div className="flex items-center gap-4 border-b border-border py-4">
<div className="max-w-[200px] flex-1">
<div className="flex flex-wrap items-center gap-4 border-b border-border py-4">
<div className="min-w-[160px] max-w-[200px] flex-1">
<label className="mb-1 block text-xs font-medium text-muted-foreground">Guardrail Name</label>
<Input
value={guardrailName}
@ -474,6 +482,19 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
{mode.includes("logging_only") && (
<CustomCodeLoggingOnlyScopeSelect value={loggingOnlyScopeChoice} onChange={setLoggingOnlyScopeChoice} />
)}
{mode.includes("logging_only") && (
<div className="flex shrink-0 items-center gap-2 pt-5">
<span className="text-sm text-muted-foreground">
Continue observing the response after a flagged request
</span>
<Switch
checked={loggingOnlyScopeChoice === "default" && loggingOnlyContinue}
disabled={loggingOnlyScopeChoice !== "default"}
onCheckedChange={setLoggingOnlyContinue}
aria-label="Continue observing the response after a flagged request"
/>
</div>
)}
<div className="w-[180px]">
<label className="mb-1 block text-xs font-medium text-muted-foreground">Template</label>
<Select

68
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.characterization.test.tsx
View file

@ -182,7 +182,9 @@ describe("GuardrailInfoView update payload characterization", () => {
await saveChanges(user);
await waitFor(() => expect(networking.updateGuardrailCall).toHaveBeenCalledTimes(1));
expect(lastPayload()).toEqual({ litellm_params: { logging_only_scope: "output" } });
expect(lastPayload()).toEqual({
litellm_params: { logging_only_scope: "output", logging_only_continue_on_input_failure: false },
});
});
it("clears the logging-only scope when the edit choice returns to default", async () => {
@ -201,7 +203,9 @@ describe("GuardrailInfoView update payload characterization", () => {
await saveChanges(user);
await waitFor(() => expect(networking.updateGuardrailCall).toHaveBeenCalledTimes(1));
expect(lastPayload()).toEqual({ litellm_params: { logging_only_scope: null } });
expect(lastPayload()).toEqual({
litellm_params: { logging_only_scope: null, logging_only_continue_on_input_failure: false },
});
});
it("clears a stored directional scope for a provider that does not support it", async () => {
@ -218,7 +222,65 @@ describe("GuardrailInfoView update payload characterization", () => {
await saveChanges(user);
await waitFor(() => expect(networking.updateGuardrailCall).toHaveBeenCalledTimes(1));
expect(lastPayload()).toEqual({ litellm_params: { logging_only_scope: null } });
expect(lastPayload()).toEqual({
litellm_params: { logging_only_scope: null, logging_only_continue_on_input_failure: false },
});
});
it("shows a stored continue flag as Default with the toggle on and saves no change", async () => {
vi.mocked(networking.getGuardrailInfo).mockResolvedValue(
guardrail({ guardrailIdentifier: "gr-abc", mode: "logging_only", logging_only_continue_on_input_failure: true }),
);
const user = userEvent.setup({ delay: null });
renderView();
expect(await screen.findAllByText("Default (request and response)")).not.toHaveLength(0);
await openEditor(user);
const toggle = await screen.findByRole("switch", {
name: "Continue observing the response after a flagged request",
});
expect(toggle).toBeChecked();
await saveChanges(user);
expect(networking.updateGuardrailCall).not.toHaveBeenCalled();
});
it("sends both scope keys when a stored continue flag is switched off", async () => {
vi.mocked(networking.getGuardrailInfo).mockResolvedValue(
guardrail({ guardrailIdentifier: "gr-abc", mode: "logging_only", logging_only_continue_on_input_failure: true }),
);
const user = userEvent.setup({ delay: null });
renderView();
await openEditor(user);
await user.click(
await screen.findByRole("switch", { name: "Continue observing the response after a flagged request" }),
);
await saveChanges(user);
await waitFor(() => expect(networking.updateGuardrailCall).toHaveBeenCalledTimes(1));
expect(lastPayload()).toEqual({
litellm_params: { logging_only_scope: null, logging_only_continue_on_input_failure: false },
});
});
it("sends the continue flag when the toggle is switched on", async () => {
vi.mocked(networking.getGuardrailInfo).mockResolvedValue(
guardrail({ guardrailIdentifier: "gr-abc", mode: "logging_only" }),
);
const user = userEvent.setup({ delay: null });
renderView();
await openEditor(user);
await user.click(
await screen.findByRole("switch", { name: "Continue observing the response after a flagged request" }),
);
await saveChanges(user);
await waitFor(() => expect(networking.updateGuardrailCall).toHaveBeenCalledTimes(1));
expect(lastPayload()).toEqual({
litellm_params: { logging_only_scope: null, logging_only_continue_on_input_failure: true },
});
});
it("parses the guardrail information textarea into an object", async () => {

17
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx
View file

@ -39,6 +39,7 @@ import {
getLoggingOnlyScopeUpdate,
getGuardrailLogoAndName,
guardrail_provider_map,
loggingOnlyContinueFromParams,
loggingOnlyScopeToChoice,
skipSystemMessageToChoice,
skipToolMessageToChoice,
@ -173,8 +174,8 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
// Only if there are entities configured
if (Object.keys(piiConfig).length > 0) {
const entities: string[] = [];
const actions: { [key: string]: string } = {};
const entities: string[] = [],
actions: { [key: string]: string } = {};
Object.entries(piiConfig).forEach(([entity, action]: [string, any]) => {
entities.push(entity);
@ -226,14 +227,13 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
fetchGuardrailUISettings();
}, [guardrailId, accessToken]);
// Reset form when guardrail data or provider params change. Only the names this form actually
// binds are seeded: an unbound key would otherwise be submitted as if the user had set it.
useEffect(() => {
if (!guardrailData) return;
form.setValue("guardrail_name", guardrailData.guardrail_name);
form.setValue("default_on", guardrailData.litellm_params?.default_on);
const storedLoggingOnlyScope = guardrailData.litellm_params?.logging_only_scope;
form.setValue("logging_only_scope_choice", loggingOnlyScopeToChoice(storedLoggingOnlyScope));
const storedParams = guardrailData.litellm_params;
form.setValue("logging_only_scope_choice", loggingOnlyScopeToChoice(storedParams?.logging_only_scope));
form.setValue("logging_only_continue_on_input_failure", loggingOnlyContinueFromParams(storedParams));
form.setValue(
"skip_system_message_choice",
skipSystemMessageToChoice(guardrailData.litellm_params?.skip_system_message_in_guardrail),
@ -302,12 +302,11 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
try {
if (!accessToken) return;
// Prepare update data object - only include changed fields
const { logging_only_scope_choice: scopeChoice, logging_only_continue_on_input_failure: continueFlag } = values;
const updateData: any = {
litellm_params: getLoggingOnlyScopeUpdate(guardrailData.litellm_params, values.logging_only_scope_choice),
litellm_params: getLoggingOnlyScopeUpdate(guardrailData.litellm_params, scopeChoice, continueFlag),
};
// Only include guardrail_name if it has changed
if (values.guardrail_name !== guardrailData.guardrail_name) {
updateData.guardrail_name = values.guardrail_name;
}

49
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx
View file

@ -22,6 +22,7 @@ import {
streamScopePayload,
loggingOnlyScopeToChoice,
choiceToLoggingOnlyScope,
loggingOnlyContinueFromParams,
getLoggingOnlyScopeUpdate,
getLoggingOnlyScopeOptions,
formatLoggingOnlyScope,
@ -257,7 +258,6 @@ describe("guardrail_info_helpers", () => {
it("normalizes directional choices only when the provider does not support them", () => {
expect(normalizeLoggingOnlyScopeChoice("input", false)).toBe("default");
expect(normalizeLoggingOnlyScopeChoice("output", false)).toBe("default");
expect(normalizeLoggingOnlyScopeChoice("both", false)).toBe("both");
expect(normalizeLoggingOnlyScopeChoice("default", false)).toBe("default");
expect(normalizeLoggingOnlyScopeChoice("input", true)).toBe("input");
expect(normalizeLoggingOnlyScopeChoice("output", true)).toBe("output");
@ -266,7 +266,6 @@ describe("guardrail_info_helpers", () => {
it("maps API scope values to choices and back", () => {
expect(loggingOnlyScopeToChoice("input")).toBe("input");
expect(loggingOnlyScopeToChoice("output")).toBe("output");
expect(loggingOnlyScopeToChoice("both")).toBe("both");
expect(loggingOnlyScopeToChoice(undefined)).toBe("default");
expect(loggingOnlyScopeToChoice(null)).toBe("default");
expect(loggingOnlyScopeToChoice("invalid")).toBe("default");
@ -275,21 +274,48 @@ describe("guardrail_info_helpers", () => {
expect(choiceToLoggingOnlyScope(undefined)).toBeNull();
expect(choiceToLoggingOnlyScope("input")).toBe("input");
expect(choiceToLoggingOnlyScope("output")).toBe("output");
expect(choiceToLoggingOnlyScope("both")).toBe("both");
expect(getLoggingOnlyScopeUpdate({ logging_only_scope: "input" }, "input")).toEqual({});
expect(getLoggingOnlyScopeUpdate({ logging_only_scope: "input" }, "output")).toEqual({
expect(getLoggingOnlyScopeUpdate({ logging_only_scope: "input" }, "input", false)).toEqual({});
expect(getLoggingOnlyScopeUpdate({ logging_only_scope: "input" }, "output", false)).toEqual({
logging_only_scope: "output",
logging_only_continue_on_input_failure: false,
});
expect(getLoggingOnlyScopeUpdate({ logging_only_scope: "input" }, "default")).toEqual({
expect(getLoggingOnlyScopeUpdate({ logging_only_scope: "input" }, "default", false)).toEqual({
logging_only_scope: null,
logging_only_continue_on_input_failure: false,
});
expect(getLoggingOnlyScopeUpdate(undefined, "default", true)).toEqual({
logging_only_scope: null,
logging_only_continue_on_input_failure: true,
});
expect(
getLoggingOnlyScopeUpdate(
{ logging_only_scope: "input", logging_only_continue_on_input_failure: true },
"input",
false,
),
).toEqual({});
expect(getLoggingOnlyScopeUpdate({ logging_only_continue_on_input_failure: true }, "default", true)).toEqual({});
expect(getLoggingOnlyScopeUpdate({ logging_only_continue_on_input_failure: true }, "default", false)).toEqual({
logging_only_scope: null,
logging_only_continue_on_input_failure: false,
});
});
it("reads the continue toggle from the stored flag", () => {
expect(loggingOnlyContinueFromParams({ logging_only_continue_on_input_failure: true })).toBe(true);
expect(loggingOnlyContinueFromParams({ logging_only_scope: "input" })).toBe(false);
expect(
loggingOnlyContinueFromParams({ logging_only_scope: "input", logging_only_continue_on_input_failure: true }),
).toBe(false);
expect(loggingOnlyContinueFromParams({ logging_only_continue_on_input_failure: false })).toBe(false);
expect(loggingOnlyContinueFromParams(undefined)).toBe(false);
expect(loggingOnlyContinueFromParams(null)).toBe(false);
});
it("formats every scope and falls back to default for missing or unknown values", () => {
expect(formatLoggingOnlyScope("input")).toBe("Input only (request)");
expect(formatLoggingOnlyScope("output")).toBe("Output only (response)");
expect(formatLoggingOnlyScope("both")).toBe("Both (request and response)");
expect(formatLoggingOnlyScope(undefined)).toBe("Default (request and response)");
expect(formatLoggingOnlyScope(null)).toBe("Default (request and response)");
expect(formatLoggingOnlyScope("invalid")).toBe("Default (request and response)");
@ -308,13 +334,8 @@ describe("guardrail_info_helpers", () => {
});
it("filters directional options for unsupported providers and keeps all options otherwise", () => {
expect(getLoggingOnlyScopeOptions(false).map((option) => option.value)).toEqual(["default", "both"]);
expect(getLoggingOnlyScopeOptions(true).map((option) => option.value)).toEqual([
"default",
"input",
"output",
"both",
]);
expect(getLoggingOnlyScopeOptions(false).map((option) => option.value)).toEqual(["default"]);
expect(getLoggingOnlyScopeOptions(true).map((option) => option.value)).toEqual(["default", "input", "output"]);
expect(
supportsDirectionalLoggingOnlyScope(

53
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx
View file

@ -116,41 +116,70 @@ export const toModeArray = (raw: unknown): string[] => {
return [];
};
export type LoggingOnlyScope = "input" | "output" | "both";
export type LoggingOnlyScope = "input" | "output";
export type LoggingOnlyScopeChoice = "default" | LoggingOnlyScope;
export type LoggingOnlyScopeOption = { label: string; value: LoggingOnlyScopeChoice };
export const normalizeLoggingOnlyScopeChoice = (
choice: LoggingOnlyScopeChoice,
directionalScopeSupported: boolean,
): LoggingOnlyScopeChoice =>
directionalScopeSupported || choice === "default" || choice === "both" ? choice : "default";
): LoggingOnlyScopeChoice => (directionalScopeSupported || choice === "default" ? choice : "default");
const LOGGING_ONLY_SCOPE_OPTIONS: LoggingOnlyScopeOption[] = [
{ label: "Default (request and response)", value: "default" },
{ label: "Input only (request)", value: "input" },
{ label: "Output only (response)", value: "output" },
{ label: "Both (request and response)", value: "both" },
];
export const loggingOnlyScopeToChoice = (v: string | null | undefined): LoggingOnlyScopeChoice =>
v === "input" || v === "output" || v === "both" ? v : "default";
v === "input" || v === "output" ? v : "default";
export const choiceToLoggingOnlyScope = (choice: LoggingOnlyScopeChoice | undefined): LoggingOnlyScope | null =>
choice === "input" || choice === "output" || choice === "both" ? choice : null;
choice === "input" || choice === "output" ? choice : null;
export const loggingOnlyContinueFromParams = (
litellmParams:
| {
logging_only_scope?: string | null;
logging_only_continue_on_input_failure?: boolean | null;
}
| null
| undefined,
): boolean =>
litellmParams?.logging_only_continue_on_input_failure === true &&
(litellmParams?.logging_only_scope === null || litellmParams?.logging_only_scope === undefined);
export const effectiveLoggingOnlyContinue = (
choice: LoggingOnlyScopeChoice | undefined,
toggle: boolean | undefined,
): boolean => choice === "default" && toggle === true;
export const getLoggingOnlyScopeUpdate = (
litellmParams: { logging_only_scope?: string | null } | null | undefined,
litellmParams:
| {
logging_only_scope?: string | null;
logging_only_continue_on_input_failure?: boolean | null;
}
| null
| undefined,
choice: LoggingOnlyScopeChoice | undefined,
): { logging_only_scope?: LoggingOnlyScope | null } => {
if (choice === undefined || choice === loggingOnlyScopeToChoice(litellmParams?.logging_only_scope)) return {};
return { logging_only_scope: choiceToLoggingOnlyScope(choice) };
continueToggle: boolean | undefined,
): { logging_only_scope?: LoggingOnlyScope | null; logging_only_continue_on_input_failure?: boolean } => {
if (
choice === undefined ||
(choice === loggingOnlyScopeToChoice(litellmParams?.logging_only_scope) &&
effectiveLoggingOnlyContinue(choice, continueToggle) === loggingOnlyContinueFromParams(litellmParams))
)
return {};
return {
logging_only_scope: choiceToLoggingOnlyScope(choice),
logging_only_continue_on_input_failure: effectiveLoggingOnlyContinue(choice, continueToggle),
};
};
export const formatLoggingOnlyScope = (v: string | null | undefined): string => {
if (v === "input") return "Input only (request)";
if (v === "output") return "Output only (response)";
if (v === "both") return "Both (request and response)";
return "Default (request and response)";
};
@ -169,7 +198,7 @@ export const modeIncludesLoggingOnly = (raw: unknown): boolean => {
export const getLoggingOnlyScopeOptions = (directionalScopeSupported: boolean): LoggingOnlyScopeOption[] =>
directionalScopeSupported
? LOGGING_ONLY_SCOPE_OPTIONS
: LOGGING_ONLY_SCOPE_OPTIONS.filter((option) => option.value === "default" || option.value === "both");
: LOGGING_ONLY_SCOPE_OPTIONS.filter((option) => option.value === "default");
export const supportsDirectionalLoggingOnlyScope = (
settings: { providers_without_directional_logging_only_scope?: string[] } | null,

22
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -26688,10 +26688,15 @@ export interface components {
*/
location?: string | null;
/**
* Logging Only Scope
* @description which direction a logging_only scan observes: 'input' (request), 'output' (response), or 'both' (default). Only applies to mode logging_only; pre_call/post_call on the same guardrail keep blocking.
* Logging Only Continue On Input Failure
* @description when True, a flagged or raising logging_only request scan is logged and the response is still scanned, so both verdicts land. Only applies to mode logging_only and is ignored when logging_only_scope is 'input' or 'output'.
*/
logging_only_scope?: ("input" | "output" | "both") | null;
logging_only_continue_on_input_failure?: boolean | null;
/**
* Logging Only Scope
* @description which direction a logging_only scan observes: 'input' (request) or 'output' (response); unset scans both directions. Only applies to mode logging_only; pre_call/post_call on the same guardrail keep blocking.
*/
logging_only_scope?: ("input" | "output") | null;
/**
* Mask Request Content
* @description Will mask request content if guardrail makes any changes
@ -36013,10 +36018,15 @@ export interface components {
*/
location?: string | null;
/**
* Logging Only Scope
* @description which direction a logging_only scan observes: 'input' (request), 'output' (response), or 'both' (default). Only applies to mode logging_only; pre_call/post_call on the same guardrail keep blocking.
* Logging Only Continue On Input Failure
* @description when True, a flagged or raising logging_only request scan is logged and the response is still scanned, so both verdicts land. Only applies to mode logging_only and is ignored when logging_only_scope is 'input' or 'output'.
*/
logging_only_scope?: ("input" | "output" | "both") | null;
logging_only_continue_on_input_failure?: boolean | null;
/**
* Logging Only Scope
* @description which direction a logging_only scan observes: 'input' (request) or 'output' (response); unset scans both directions. Only applies to mode logging_only; pre_call/post_call on the same guardrail keep blocking.
*/
logging_only_scope?: ("input" | "output") | null;
/**
* Mask
* @description Enable content masking using Lasso classifix API