Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions litellm/litellm_core_utils/litellm_logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -5865,6 +5865,7 @@ def _get_status_fields(
# Mapping for legacy guardrail status values to new GuardrailStatus values
GUARDRAIL_STATUS_MAP: Final[dict[str, GuardrailStatus]] = {
"success": "success",
"guardrail_flagged": "guardrail_flagged",
"blocked": "guardrail_intervened", # legacy
"guardrail_intervened": "guardrail_intervened", # direct
"failure": "guardrail_failed_to_respond", # legacy
Expand All @@ -5886,6 +5887,7 @@ def _get_status_fields(
GUARDRAIL_STATUS_SEVERITY: Final[tuple[GuardrailStatus, ...]] = (
"not_run",
"success",
"guardrail_flagged",
"guardrail_failed_to_respond",
"guardrail_intervened",
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ async def apply_guardrail(inputs, request_data, input_type):

import asyncio
import threading
import time
from collections.abc import Callable, Mapping
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast

Expand Down Expand Up @@ -93,6 +94,7 @@ class CustomCodeGuardrail(CustomGuardrail):
that returns one of:
- allow() - let the request/response through
- block(reason) - reject with a message
- flag(reason) - let it through but log a non-blocking violation
- modify(texts=...) - transform the content

Example:
Expand Down Expand Up @@ -227,6 +229,7 @@ async def apply_guardrail(
raise CustomCodeExecutionError(f"Custom code guardrail not compiled: {self._compile_error}")
raise CustomCodeExecutionError("Custom code guardrail not compiled")

start_time: Final = time.time()
try:
# Prepare inputs dict for the function

Expand All @@ -245,6 +248,7 @@ async def apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type=input_type,
start_time=start_time,
)

except HTTPException:
Expand Down Expand Up @@ -290,6 +294,7 @@ def _process_result(
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
start_time: float,
) -> GenericGuardrailAPIInputs:
"""
Process the result from the custom code function.
Expand All @@ -299,6 +304,7 @@ def _process_result(
inputs: The original inputs
request_data: The request data
input_type: "request" or "response"
start_time: Unix timestamp of when the guardrail started running, used for the flagged log entry

Returns:
GenericGuardrailAPIInputs - possibly modified
Expand Down Expand Up @@ -348,6 +354,27 @@ def _process_result(
},
)

elif action == "flag":
flag_reason: Final = result.get("reason", "Flagged by custom code guardrail")
verbose_proxy_logger.info(
"Custom code guardrail '%s': Flagging %s - %s", self.guardrail_name, input_type, flag_reason
)
end_time: Final = time.time()
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response={ # mutable-ok: logging helper requires a dict
"action": "flag",
"reason": flag_reason,
"input_type": input_type,
"metadata": result.get("metadata") or {}, # mutable-ok: logging helper requires a dict
},
request_data=request_data,
guardrail_status="guardrail_flagged",
start_time=start_time,
end_time=end_time,
duration=end_time - start_time,
)
return inputs

elif action == "modify":
verbose_proxy_logger.debug("Custom code guardrail '%s': Modifying %s", self.guardrail_name, input_type)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
import json
import re
from collections.abc import Mapping, Sequence
from typing import Any, Final
from typing import Any, Final, Literal
from urllib.parse import urlparse

import httpx
Expand Down Expand Up @@ -51,6 +51,31 @@ def block(reason: str, detection_info: Mapping[str, object] | None = None) -> di
return result


class FlagResult(TypedDict):
action: ReadOnly[Literal["flag"]]
reason: ReadOnly[str]
metadata: ReadOnly[Mapping[str, object]]


def flag(reason: str, metadata: Mapping[str, object] | None = None) -> FlagResult:
"""
Let the request/response proceed unchanged but record a non-blocking violation.

Args:
reason: Human-readable reason for flagging
metadata: Optional structured metadata stored alongside the reason

Returns:
Dict indicating the request should be flagged but allowed
"""
result: Final[FlagResult] = {
"action": "flag",
"reason": reason,
"metadata": metadata if metadata is not None else {},
}
return result


def modify(
texts: Sequence[str] | None = None,
images: Sequence[object] | None = None,
Expand Down Expand Up @@ -787,6 +812,7 @@ def get_custom_code_primitives() -> dict[str, object]:
# Result types
"allow": allow,
"block": block,
"flag": flag,
"modify": modify,
# Regex
"regex_match": regex_match,
Expand Down
18 changes: 8 additions & 10 deletions litellm/proxy/guardrails/usage_endpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.guardrails.usage_tracking import guardrail_status_to_action
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import (
DailyGuardrailMetricsRepository,
Expand All @@ -41,6 +42,7 @@
router: Final = APIRouter()

_EMPTY_UNITS: Final[Mapping[str, int]] = MappingProxyType({})
_ACTION_SEVERITY: Final[Mapping[str, int]] = MappingProxyType({"passed": 0, "flagged": 1, "blocked": 2})

_USAGE_MAX_RANGE_DAYS: Final = 366

Expand Down Expand Up @@ -680,21 +682,17 @@ def _usage_log_entry_from_row(
except Exception:
meta = {}
guardrail_info_list: Final[Sequence[_GuardrailRunInfo]] = (meta or {}).get("guardrail_information") or []
entry_for_guardrail: _GuardrailRunInfo | None = None
for gi in guardrail_info_list:
if (gi.get("guardrail_id") or gi.get("guardrail_name")) == r.guardrail_id:
entry_for_guardrail = gi
break
entry_for_guardrail: Final[_GuardrailRunInfo | None] = max(
(gi for gi in guardrail_info_list if (gi.get("guardrail_id") or gi.get("guardrail_name")) == r.guardrail_id),
key=lambda gi: _ACTION_SEVERITY[guardrail_status_to_action(gi.get("guardrail_status"))],
default=None,
)
action_val = "passed"
score_val = None
latency_val = None
reason_val = None
if entry_for_guardrail:
st: Final = (entry_for_guardrail.get("guardrail_status") or "").lower()
if "intervened" in st or "block" in st:
action_val = "blocked"
elif "fail" in st or "error" in st:
action_val = "flagged"
action_val = guardrail_status_to_action(entry_for_guardrail.get("guardrail_status"))
duration: Final = entry_for_guardrail.get("duration")
if duration is not None:
latency_val = round(float(duration) * 1000, 0)
Expand Down
6 changes: 3 additions & 3 deletions litellm/proxy/guardrails/usage_tracking.py
Original file line number Diff line number Diff line change
Expand Up @@ -160,14 +160,14 @@ async def _upsert_rows_with_retry(
return await _upsert_rows_with_retry(retryable, upsert_row, label, sleep, retries_left - 1)


def _guardrail_status_to_action(status: str | None) -> str:
def guardrail_status_to_action(status: str | None) -> str:
"""Map StandardLogging guardrail_status to blocked/passed/flagged."""
if not status:
return "passed"
s: Final = (status or "").lower()
if "intervened" in s or "block" in s:
return "blocked"
if "fail" in s or "error" in s:
if "flagged" in s or "fail" in s or "error" in s:
return "flagged"
return "passed"

Expand Down Expand Up @@ -318,7 +318,7 @@ async def process_spend_logs_guardrail_usage(
continue
key = _MetricsKey(guardrail_id, date_key)
daily_guardrail[key]["requests_evaluated"] += 1
action = _guardrail_status_to_action(entry.get("guardrail_status"))
action = guardrail_status_to_action(entry.get("guardrail_status"))
if action == "passed":
daily_guardrail[key]["passed_count"] += 1
elif action == "blocked":
Expand Down
5 changes: 4 additions & 1 deletion litellm/types/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -3074,7 +3074,9 @@ class GuardrailMode(TypedDict, total=False):
default: str | list[str] | None


GuardrailStatus = Literal["success", "guardrail_intervened", "guardrail_failed_to_respond", "not_run"]
GuardrailStatus = Literal[
"success", "guardrail_flagged", "guardrail_intervened", "guardrail_failed_to_respond", "not_run"
]


class StandardLoggingGuardrailInformation(TypedDict, total=False):
Expand Down Expand Up @@ -3264,6 +3266,7 @@ class StandardLoggingPayloadStatusFields(TypedDict, total=False):
"""
Status of guardrail execution:
- 'success': Guardrail ran and allowed content through
- 'guardrail_flagged': Guardrail allowed content through but recorded a non-blocking violation
- 'guardrail_intervened': Guardrail blocked or modified content
- 'guardrail_failed_to_respond': Guardrail had technical failure
- 'not_run': No guardrail was run
Expand Down
18 changes: 17 additions & 1 deletion tests/test_litellm/litellm_core_utils/test_litellm_logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,10 @@
from litellm.constants import SENTRY_DENYLIST, SENTRY_PII_DENYLIST
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
from litellm.litellm_core_utils.litellm_logging import set_callbacks
from litellm.litellm_core_utils.litellm_logging import (
_get_status_fields,
set_callbacks,
)
from litellm.types.utils import ModelResponse, TextCompletionResponse


Expand Down Expand Up @@ -6277,3 +6280,16 @@ def test_passthrough_embeddings_result_swapped_for_callbacks():

assert isinstance(swapped_result, EmbeddingResponse)
assert swapped_result.data[0]["embedding"] == [0.1, 0.2, 0.3]


def test_get_status_fields_ranks_guardrail_flagged_between_success_and_intervened():
"""LIT-6894: a non-blocking flagged verdict must outrank success in the
request-level guardrail_status but never mask an intervention."""
flagged = {"guardrail_status": "guardrail_flagged"}

assert _get_status_fields(
"success", [{"guardrail_status": "success"}, flagged], None
)["guardrail_status"] == "guardrail_flagged"
assert _get_status_fields(
"success", [flagged, {"guardrail_status": "guardrail_intervened"}], None
)["guardrail_status"] == "guardrail_intervened"
65 changes: 65 additions & 0 deletions tests/test_litellm/proxy/guardrails/test_custom_code_security.py
Original file line number Diff line number Diff line change
Expand Up @@ -197,6 +197,71 @@ async def test_custom_code_post_call_block_raises_http_400():
}


FLAG_CODE = (
"def apply_guardrail(inputs, request_data, input_type):\n"
' return flag("audit hit", metadata={"category": "topic"})\n'
)


@pytest.mark.asyncio
@pytest.mark.parametrize("input_type", ["request", "response"])
async def test_custom_code_flag_passes_content_through_and_records_flagged_entry(input_type):
"""LIT-6894: flag() must not raise, must return the content unchanged and must log
exactly one guardrail_flagged entry (the decorator must not add a second "success")."""
guardrail = CustomCodeGuardrail(custom_code=FLAG_CODE, guardrail_name="t", event_hook=["pre_call", "post_call"])
request_data = {"model": "test-model", "litellm_metadata": {}}

result = await guardrail.apply_guardrail(
inputs={"texts": ["hello"]},
request_data=request_data,
input_type=input_type,
)

assert result == {"texts": ["hello"]}
entries = request_data["litellm_metadata"]["standard_logging_guardrail_information"]
assert len(entries) == 1
entry = entries[0]
assert entry["guardrail_status"] == "guardrail_flagged"
assert entry["guardrail_name"] == "t"
assert entry["guardrail_mode"] == ["pre_call", "post_call"]
assert entry["guardrail_response"] == {
"action": "flag",
"reason": "audit hit",
"input_type": input_type,
"metadata": {"category": "topic"},
}
assert entry["duration"] is not None and entry["duration"] >= 0


@pytest.mark.asyncio
async def test_custom_code_flag_default_reason_and_empty_metadata():
code = "def apply_guardrail(inputs, request_data, input_type):\n return flag('just a note')\n"
guardrail = _compile(code)
request_data = {"model": "m", "litellm_metadata": {}}

await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data=request_data, input_type="request")

entry = request_data["litellm_metadata"]["standard_logging_guardrail_information"][0]
assert entry["guardrail_response"] == {
"action": "flag",
"reason": "just a note",
"input_type": "request",
"metadata": {},
}


@pytest.mark.asyncio
async def test_custom_code_allow_still_records_success_not_flagged():
code = "def apply_guardrail(inputs, request_data, input_type):\n return allow()\n"
guardrail = _compile(code)
request_data = {"model": "m", "litellm_metadata": {}}

await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data=request_data, input_type="request")

entries = request_data["litellm_metadata"]["standard_logging_guardrail_information"]
assert [e["guardrail_status"] for e in entries] == ["success"]


def test_typical_sync_guardrail_still_works():
code = (
"def apply_guardrail(inputs, request_data, input_type):\n"
Expand Down
Loading
Loading