mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-28 01:32:17 +00:00
Merge pull request #41757 from BerriAI/litellm_typesafe_compaction_guardrail
feat(guardrails): add TypeSafe Jev relevance-based compaction guardrail
This commit is contained in:
commit
2edda5aec3
9 changed files with 972 additions and 5 deletions
|
|
@ -10030,7 +10030,7 @@
|
||||||
},
|
},
|
||||||
"unreachable_fallback": {
|
"unreachable_fallback": {
|
||||||
"default": "fail_closed",
|
"default": "fail_closed",
|
||||||
"description": "Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed.",
|
"description": "Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', 'compresr', and 'typesafe'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed.",
|
||||||
"enum": [
|
"enum": [
|
||||||
"fail_closed",
|
"fail_closed",
|
||||||
"fail_open"
|
"fail_open"
|
||||||
|
|
|
||||||
|
|
@ -21,7 +21,7 @@ if TYPE_CHECKING:
|
||||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||||
from litellm.router import Router
|
from litellm.router import Router
|
||||||
|
|
||||||
COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr"})
|
COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr", "typesafe"})
|
||||||
_NO_COMPRESSION: Final = "none"
|
_NO_COMPRESSION: Final = "none"
|
||||||
|
|
||||||
# A ContextVar, not metadata: metadata reaches spend logs the caller can read, and a
|
# A ContextVar, not metadata: metadata reaches spend logs the caller can read, and a
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,74 @@
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING, Final
|
||||||
|
|
||||||
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
from litellm.types.guardrails import (
|
||||||
|
GuardrailEventHooks,
|
||||||
|
Mode,
|
||||||
|
SupportedGuardrailIntegrations,
|
||||||
|
)
|
||||||
|
from litellm.types.proxy.guardrails.guardrail_hooks.typesafe import (
|
||||||
|
TypeSafeGuardrailOptionalParams,
|
||||||
|
)
|
||||||
|
|
||||||
|
from .typesafe import TypeSafeGuardrail
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from litellm.types.guardrails import Guardrail, LitellmParams
|
||||||
|
|
||||||
|
|
||||||
|
def _coerce_event_hook(
|
||||||
|
mode: str | list[str] | Mode,
|
||||||
|
) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode:
|
||||||
|
if isinstance(mode, Mode):
|
||||||
|
return mode
|
||||||
|
if isinstance(mode, list):
|
||||||
|
return [ # mutable-ok: CustomGuardrail event_hook contract wants a list
|
||||||
|
GuardrailEventHooks(item) for item in mode
|
||||||
|
]
|
||||||
|
return GuardrailEventHooks(mode)
|
||||||
|
|
||||||
|
|
||||||
|
def _optional_params(litellm_params: LitellmParams) -> TypeSafeGuardrailOptionalParams:
|
||||||
|
value: Final = litellm_params.optional_params
|
||||||
|
if isinstance(value, TypeSafeGuardrailOptionalParams):
|
||||||
|
return value
|
||||||
|
if isinstance(value, BaseModel):
|
||||||
|
return TypeSafeGuardrailOptionalParams.model_validate(value.model_dump())
|
||||||
|
return TypeSafeGuardrailOptionalParams()
|
||||||
|
|
||||||
|
|
||||||
|
def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> TypeSafeGuardrail:
|
||||||
|
import litellm
|
||||||
|
|
||||||
|
optional_params: Final = _optional_params(litellm_params)
|
||||||
|
|
||||||
|
_callback: Final = TypeSafeGuardrail(
|
||||||
|
api_base=litellm_params.api_base,
|
||||||
|
api_key=litellm_params.api_key,
|
||||||
|
model=litellm_params.model,
|
||||||
|
relevance_threshold=optional_params.relevance_threshold,
|
||||||
|
min_chars_to_evaluate=optional_params.min_chars_to_evaluate,
|
||||||
|
max_result_chars_in_state=optional_params.max_result_chars_in_state,
|
||||||
|
guardrail_name=guardrail["guardrail_name"],
|
||||||
|
event_hook=_coerce_event_hook(litellm_params.mode),
|
||||||
|
default_on=litellm_params.default_on or False,
|
||||||
|
unreachable_fallback=(
|
||||||
|
litellm_params.unreachable_fallback if "unreachable_fallback" in litellm_params.model_fields_set else None
|
||||||
|
),
|
||||||
|
)
|
||||||
|
litellm.logging_callback_manager.add_litellm_callback( # pyright: ignore[reportUnknownMemberType] # callback manager is untyped
|
||||||
|
_callback
|
||||||
|
)
|
||||||
|
return _callback
|
||||||
|
|
||||||
|
|
||||||
|
guardrail_initializer_registry: Final = { # mutable-ok: guardrail_registry discovery checks isinstance(registry, dict)
|
||||||
|
SupportedGuardrailIntegrations.TYPESAFE.value: initialize_guardrail,
|
||||||
|
}
|
||||||
|
|
||||||
|
guardrail_class_registry: Final = { # mutable-ok: guardrail_registry discovery checks isinstance(registry, dict)
|
||||||
|
SupportedGuardrailIntegrations.TYPESAFE.value: TypeSafeGuardrail,
|
||||||
|
}
|
||||||
416
litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py
Normal file
416
litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py
Normal file
|
|
@ -0,0 +1,416 @@
|
||||||
|
"""TypeSafe (Jev) relevance-based compaction guardrail.
|
||||||
|
|
||||||
|
Instead of summarizing tool output, the guardrail asks TypeSafe's Jev model
|
||||||
|
one yes/no question per completed tool exchange ("is this result still needed
|
||||||
|
for the current task?") over ``POST {api_base}/v1/systemone`` and blanks the
|
||||||
|
tool results Jev judges no longer relevant.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import time
|
||||||
|
from collections.abc import Mapping, Sequence
|
||||||
|
from typing import TYPE_CHECKING, Annotated, Final, Literal
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
from fastapi import HTTPException
|
||||||
|
from httpx import Response as HttpxResponse
|
||||||
|
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
|
||||||
|
|
||||||
|
from litellm._logging import verbose_proxy_logger
|
||||||
|
from litellm.compression.compress import get_protected_indices
|
||||||
|
from litellm.integrations.custom_guardrail import (
|
||||||
|
CustomGuardrail,
|
||||||
|
log_guardrail_information, # pyright: ignore[reportUnknownVariableType] # decorator is untyped in custom_guardrail
|
||||||
|
)
|
||||||
|
from litellm.litellm_core_utils.prompt_templates.factory import group_tool_exchanges
|
||||||
|
from litellm.llms.custom_httpx.http_handler import (
|
||||||
|
AsyncHTTPHandler,
|
||||||
|
get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # helper is untyped in http_handler
|
||||||
|
httpxSpecialProvider,
|
||||||
|
)
|
||||||
|
from litellm.proxy.guardrails.guardrail_hooks.content_text import content_to_text
|
||||||
|
from litellm.secret_managers.main import get_secret_str
|
||||||
|
from litellm.types.guardrails import GuardrailEventHooks, Mode
|
||||||
|
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from litellm.litellm_core_utils.litellm_logging import (
|
||||||
|
Logging as LiteLLMLoggingObj,
|
||||||
|
)
|
||||||
|
from litellm.types.proxy.guardrails.guardrail_hooks.typesafe import (
|
||||||
|
TypeSafeGuardrailConfigModel,
|
||||||
|
)
|
||||||
|
|
||||||
|
DEFAULT_API_BASE: Final = "https://api.typesafe.ai"
|
||||||
|
DEFAULT_MODEL: Final = "jev-latest"
|
||||||
|
DEFAULT_RELEVANCE_THRESHOLD: Final = 0.2
|
||||||
|
DEFAULT_MIN_CHARS_TO_EVALUATE: Final = 200
|
||||||
|
DEFAULT_MAX_RESULT_CHARS_IN_STATE: Final = 4000
|
||||||
|
_MAX_EXCHANGES_EVALUATED: Final = 200
|
||||||
|
_JEV_TIMEOUT_SECONDS: Final = 30.0
|
||||||
|
DROPPED_RESULT_TEXT: Final = (
|
||||||
|
"[Tool result removed by TypeSafe compaction: judged no longer relevant to the current task]"
|
||||||
|
)
|
||||||
|
_ELISION_MARKER: Final = "\n... [middle truncated] ...\n"
|
||||||
|
|
||||||
|
|
||||||
|
_STR_OBJECT_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||||
|
_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
|
||||||
|
|
||||||
|
|
||||||
|
def _as_str_object_dict(value: object) -> dict[str, object] | None:
|
||||||
|
try:
|
||||||
|
return _STR_OBJECT_DICT_ADAPTER.validate_python(value)
|
||||||
|
except ValidationError:
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def _as_object_list(value: object) -> list[object] | None:
|
||||||
|
try:
|
||||||
|
return _OBJECT_LIST_ADAPTER.validate_python(value)
|
||||||
|
except ValidationError:
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def _safe_response_text(response: HttpxResponse | None, limit: int = 500) -> str:
|
||||||
|
if response is None:
|
||||||
|
return ""
|
||||||
|
try:
|
||||||
|
text: Final = response.text
|
||||||
|
except httpx.DecodingError:
|
||||||
|
return "<undecodable response body>"
|
||||||
|
return (text or "")[:limit]
|
||||||
|
|
||||||
|
|
||||||
|
class _JevNoulAnswer(BaseModel):
|
||||||
|
model_config = ConfigDict(frozen=True, allow_inf_nan=False)
|
||||||
|
|
||||||
|
type: Literal["noul"]
|
||||||
|
noul: Annotated[float, Field(ge=0.0, le=1.0)]
|
||||||
|
|
||||||
|
|
||||||
|
class _JevSystemOneResponse(BaseModel):
|
||||||
|
model_config = ConfigDict(frozen=True)
|
||||||
|
|
||||||
|
answers: Mapping[str, _JevNoulAnswer]
|
||||||
|
|
||||||
|
|
||||||
|
_JEV_RESPONSE_ADAPTER: Final = TypeAdapter(_JevSystemOneResponse)
|
||||||
|
|
||||||
|
|
||||||
|
def _truncate_for_state(text: str, max_chars: int) -> str:
|
||||||
|
"""Keeps the head and tail within ``max_chars`` so Jev sees both ends of a long result."""
|
||||||
|
if len(text) <= max_chars:
|
||||||
|
return text
|
||||||
|
if max_chars <= len(_ELISION_MARKER):
|
||||||
|
return text[:max_chars]
|
||||||
|
budget: Final = max_chars - len(_ELISION_MARKER)
|
||||||
|
head: Final = budget // 2
|
||||||
|
return text[:head] + _ELISION_MARKER + text[len(text) - (budget - head) :]
|
||||||
|
|
||||||
|
|
||||||
|
def _question_instructions(question_id: str) -> str:
|
||||||
|
return (
|
||||||
|
f"Is tool exchange `{question_id}` in `tool_exchanges` still needed by the assistant to "
|
||||||
|
"complete `task`? Answer yes if its result contains information the assistant has not yet "
|
||||||
|
"fully used or will need again; answer no if it is off-topic, superseded, or already "
|
||||||
|
"incorporated into later messages."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _tool_call_entry(tool_call: object) -> dict[str, object] | None:
|
||||||
|
parsed_call = _as_str_object_dict(tool_call)
|
||||||
|
if parsed_call is None:
|
||||||
|
return None
|
||||||
|
function = _as_str_object_dict(parsed_call.get("function"))
|
||||||
|
fn = function if function is not None else parsed_call
|
||||||
|
return {"name": fn.get("name"), "arguments": fn.get("arguments")} # mutable-ok: serialized to JSON
|
||||||
|
|
||||||
|
|
||||||
|
def _tool_call_entries(assistant_message: Mapping[str, object]) -> tuple[dict[str, object], ...]:
|
||||||
|
tool_calls: Final = _as_object_list(assistant_message.get("tool_calls"))
|
||||||
|
if tool_calls is None:
|
||||||
|
return ()
|
||||||
|
return tuple(entry for tool_call in tool_calls if (entry := _tool_call_entry(tool_call)) is not None)
|
||||||
|
|
||||||
|
|
||||||
|
def _protected_indices(messages: Sequence[Mapping[str, object]]) -> frozenset[int]:
|
||||||
|
"""``get_protected_indices`` expanded over whole tool exchanges, so the most recent exchange is never evaluated."""
|
||||||
|
protected: Final = frozenset(get_protected_indices(messages))
|
||||||
|
return protected | frozenset(
|
||||||
|
index
|
||||||
|
for group in group_tool_exchanges(messages)
|
||||||
|
if any(member in protected for member in group)
|
||||||
|
for index in group
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class TypeSafeGuardrail(CustomGuardrail):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
api_base: str | None = None,
|
||||||
|
api_key: str | None = None,
|
||||||
|
model: str | None = None,
|
||||||
|
relevance_threshold: float | None = None,
|
||||||
|
min_chars_to_evaluate: int | None = None,
|
||||||
|
max_result_chars_in_state: int | None = None,
|
||||||
|
unreachable_fallback: str | None = None,
|
||||||
|
guardrail_name: str | None = None,
|
||||||
|
event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None = None,
|
||||||
|
default_on: bool = False,
|
||||||
|
async_handler: AsyncHTTPHandler | None = None,
|
||||||
|
) -> None:
|
||||||
|
raw_api_base: Final = (api_base or get_secret_str("TYPESAFE_API_BASE") or DEFAULT_API_BASE).rstrip("/")
|
||||||
|
self.typesafe_api_base = raw_api_base
|
||||||
|
self.typesafe_api_key = api_key or get_secret_str("TYPESAFE_API_KEY")
|
||||||
|
if not self.typesafe_api_key:
|
||||||
|
raise ValueError(
|
||||||
|
"TypeSafe guardrail requires an API key. Set `api_key` in the "
|
||||||
|
"guardrail config or the TYPESAFE_API_KEY env var."
|
||||||
|
)
|
||||||
|
self.jev_model = model or DEFAULT_MODEL
|
||||||
|
self.relevance_threshold = DEFAULT_RELEVANCE_THRESHOLD if relevance_threshold is None else relevance_threshold
|
||||||
|
self.min_chars_to_evaluate = (
|
||||||
|
DEFAULT_MIN_CHARS_TO_EVALUATE if min_chars_to_evaluate is None else min_chars_to_evaluate
|
||||||
|
)
|
||||||
|
self.max_result_chars_in_state = (
|
||||||
|
DEFAULT_MAX_RESULT_CHARS_IN_STATE if max_result_chars_in_state is None else max_result_chars_in_state
|
||||||
|
)
|
||||||
|
self.unreachable_fallback: Literal["fail_closed", "fail_open"] = (
|
||||||
|
"fail_closed" if unreachable_fallback == "fail_closed" else "fail_open"
|
||||||
|
)
|
||||||
|
self.async_handler: AsyncHTTPHandler = async_handler or get_async_httpx_client(
|
||||||
|
llm_provider=httpxSpecialProvider.GuardrailCallback,
|
||||||
|
)
|
||||||
|
super().__init__( # pyright: ignore[reportUnknownMemberType] # CustomGuardrail.__init__ is untyped
|
||||||
|
guardrail_name=guardrail_name,
|
||||||
|
event_hook=event_hook,
|
||||||
|
default_on=default_on,
|
||||||
|
)
|
||||||
|
|
||||||
|
def _handle_failure(self, error: str, log_detail: dict[str, object]) -> None:
|
||||||
|
"""fail_open logs and returns; fail_closed raises a generic 502 (upstream bodies stay in server logs)."""
|
||||||
|
if self.unreachable_fallback == "fail_open":
|
||||||
|
verbose_proxy_logger.warning(
|
||||||
|
"TypeSafe: %s; fail_open configured, forwarding request uncompacted. detail=%s",
|
||||||
|
error,
|
||||||
|
log_detail,
|
||||||
|
)
|
||||||
|
return
|
||||||
|
verbose_proxy_logger.error("TypeSafe: %s. detail=%s", error, log_detail)
|
||||||
|
raise HTTPException(status_code=502, detail={"error": error}) # mutable-ok: FastAPI wants a dict detail
|
||||||
|
|
||||||
|
def _candidate_exchanges(self, messages: Sequence[dict[str, object]]) -> tuple[tuple[int, ...], ...]:
|
||||||
|
"""Completed tool exchanges eligible for evaluation: unprotected, and long enough to be worth a call."""
|
||||||
|
protected: Final = _protected_indices(messages)
|
||||||
|
candidates: Final = tuple(
|
||||||
|
group
|
||||||
|
for group in group_tool_exchanges(messages)
|
||||||
|
if len(group) >= 2
|
||||||
|
and messages[group[0]].get("role") == "assistant"
|
||||||
|
and not any(member in protected for member in group)
|
||||||
|
and len(self._exchange_tool_text(messages, group)) >= self.min_chars_to_evaluate
|
||||||
|
)
|
||||||
|
return candidates[-_MAX_EXCHANGES_EVALUATED:]
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _exchange_tool_text(messages: Sequence[dict[str, object]], group: tuple[int, ...]) -> str:
|
||||||
|
return "".join(
|
||||||
|
content_to_text(messages[index].get("content"))
|
||||||
|
for index in group[1:]
|
||||||
|
if messages[index].get("role") in ("tool", "function")
|
||||||
|
)
|
||||||
|
|
||||||
|
def _build_state(
|
||||||
|
self, messages: Sequence[dict[str, object]], candidates: tuple[tuple[int, ...], ...]
|
||||||
|
) -> dict[str, object]:
|
||||||
|
task: Final = next(
|
||||||
|
(
|
||||||
|
content_to_text(messages[index].get("content"))
|
||||||
|
for index in range(len(messages) - 1, -1, -1)
|
||||||
|
if messages[index].get("role") == "user"
|
||||||
|
),
|
||||||
|
"",
|
||||||
|
)
|
||||||
|
system: Final = "\n\n".join(
|
||||||
|
content_to_text(message.get("content")) for message in messages if message.get("role") == "system"
|
||||||
|
)
|
||||||
|
tool_exchanges: Final = { # mutable-ok: accumulated once, serialized to JSON
|
||||||
|
f"e{ordinal}": { # mutable-ok: serialized to JSON
|
||||||
|
"tool_calls": _tool_call_entries(messages[group[0]]),
|
||||||
|
"result": _truncate_for_state(
|
||||||
|
self._exchange_tool_text(messages, group), self.max_result_chars_in_state
|
||||||
|
),
|
||||||
|
}
|
||||||
|
for ordinal, group in enumerate(candidates)
|
||||||
|
}
|
||||||
|
return {"task": task, "system": system, "tool_exchanges": tool_exchanges} # mutable-ok: serialized to JSON
|
||||||
|
|
||||||
|
async def _call_systemone(
|
||||||
|
self, state: dict[str, object], question_ids: Sequence[str]
|
||||||
|
) -> _JevSystemOneResponse | None:
|
||||||
|
"""Returns the response, or None when the service failed and fail_open applies."""
|
||||||
|
payload: Final[dict[str, object]] = { # mutable-ok: serialized to JSON by httpx
|
||||||
|
"model": self.jev_model,
|
||||||
|
"state": state,
|
||||||
|
"questions": { # mutable-ok: serialized to JSON
|
||||||
|
question_id: { # mutable-ok: serialized to JSON
|
||||||
|
"type": "noul",
|
||||||
|
"instructions": _question_instructions(question_id),
|
||||||
|
}
|
||||||
|
for question_id in question_ids
|
||||||
|
},
|
||||||
|
}
|
||||||
|
try:
|
||||||
|
raw_response: HttpxResponse = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post is untyped
|
||||||
|
url=f"{self.typesafe_api_base}/v1/systemone",
|
||||||
|
json=payload,
|
||||||
|
headers={ # mutable-ok: httpx header contract is a dict
|
||||||
|
"Authorization": f"Bearer {self.typesafe_api_key}",
|
||||||
|
"Content-Type": "application/json",
|
||||||
|
},
|
||||||
|
timeout=_JEV_TIMEOUT_SECONDS,
|
||||||
|
)
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
raise
|
||||||
|
except Exception as e:
|
||||||
|
detail: Final[dict[str, object]] = (
|
||||||
|
{ # mutable-ok: log detail record
|
||||||
|
"error_type": type(e).__name__,
|
||||||
|
"detail": str(e),
|
||||||
|
"status_code": e.response.status_code,
|
||||||
|
"body": _safe_response_text(e.response),
|
||||||
|
}
|
||||||
|
if isinstance(e, httpx.HTTPStatusError)
|
||||||
|
else {"error_type": type(e).__name__, "detail": str(e)} # mutable-ok: log detail record
|
||||||
|
)
|
||||||
|
self._handle_failure("TypeSafe evaluation service request failed", detail)
|
||||||
|
return None
|
||||||
|
if not 200 <= raw_response.status_code < 300:
|
||||||
|
self._handle_failure(
|
||||||
|
"TypeSafe evaluation service returned an error",
|
||||||
|
{ # mutable-ok: log detail record
|
||||||
|
"status_code": raw_response.status_code,
|
||||||
|
"body": _safe_response_text(raw_response),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
return None
|
||||||
|
try:
|
||||||
|
body: Final[object] = raw_response.json() # pyright: ignore[reportAny] # httpx Response.json() is untyped
|
||||||
|
except (ValueError, httpx.DecodingError, RecursionError):
|
||||||
|
self._handle_failure(
|
||||||
|
"TypeSafe evaluation service returned an unreadable response",
|
||||||
|
{"body": _safe_response_text(raw_response)}, # mutable-ok: log detail record
|
||||||
|
)
|
||||||
|
return None
|
||||||
|
try:
|
||||||
|
return _JEV_RESPONSE_ADAPTER.validate_python(body)
|
||||||
|
except ValidationError:
|
||||||
|
self._handle_failure(
|
||||||
|
"TypeSafe evaluation service returned unexpected response shape",
|
||||||
|
{"body": _safe_response_text(raw_response)}, # mutable-ok: log detail record
|
||||||
|
)
|
||||||
|
return None
|
||||||
|
|
||||||
|
@log_guardrail_information
|
||||||
|
async def apply_guardrail(
|
||||||
|
self,
|
||||||
|
inputs: GenericGuardrailAPIInputs,
|
||||||
|
request_data: dict[str, object],
|
||||||
|
input_type: Literal["request", "response"],
|
||||||
|
logging_obj: LiteLLMLoggingObj | None = None,
|
||||||
|
) -> GenericGuardrailAPIInputs:
|
||||||
|
if input_type != "request":
|
||||||
|
return inputs
|
||||||
|
|
||||||
|
structured_messages: Final = _as_object_list(inputs.get("structured_messages"))
|
||||||
|
if not structured_messages:
|
||||||
|
return inputs
|
||||||
|
parsed_messages: Final = tuple(_as_str_object_dict(m) for m in structured_messages)
|
||||||
|
if any(m is None for m in parsed_messages):
|
||||||
|
return inputs
|
||||||
|
messages: Final = tuple(m for m in parsed_messages if m is not None)
|
||||||
|
|
||||||
|
candidates: Final = self._candidate_exchanges(messages)
|
||||||
|
if not candidates:
|
||||||
|
verbose_proxy_logger.debug("TypeSafe: no completed tool exchanges eligible for evaluation")
|
||||||
|
return inputs
|
||||||
|
|
||||||
|
question_ids: Final = tuple(f"e{ordinal}" for ordinal in range(len(candidates)))
|
||||||
|
state: Final = self._build_state(messages, candidates)
|
||||||
|
|
||||||
|
start_time: Final = time.monotonic()
|
||||||
|
response: Final = await self._call_systemone(state, question_ids)
|
||||||
|
end_time: Final = time.monotonic()
|
||||||
|
if response is None:
|
||||||
|
self.add_standard_logging_guardrail_information_to_request_data( # pyright: ignore[reportUnknownMemberType] # untyped base helper
|
||||||
|
guardrail_json_response={ # mutable-ok: must stay JSON-serializable for shared logging
|
||||||
|
"error": "TypeSafe evaluation unavailable; request forwarded uncompacted",
|
||||||
|
"model": self.jev_model,
|
||||||
|
},
|
||||||
|
request_data=request_data,
|
||||||
|
guardrail_status="guardrail_failed_to_respond",
|
||||||
|
guardrail_provider="typesafe",
|
||||||
|
start_time=start_time,
|
||||||
|
end_time=end_time,
|
||||||
|
duration=end_time - start_time,
|
||||||
|
)
|
||||||
|
return inputs
|
||||||
|
|
||||||
|
dropped_ordinals: Final = frozenset(
|
||||||
|
ordinal
|
||||||
|
for ordinal in range(len(candidates))
|
||||||
|
if (answer := response.answers.get(f"e{ordinal}")) is not None and answer.noul < self.relevance_threshold
|
||||||
|
)
|
||||||
|
dropped_tool_indices: Final[frozenset[int]] = frozenset(
|
||||||
|
index
|
||||||
|
for ordinal in dropped_ordinals
|
||||||
|
for index in candidates[ordinal][1:]
|
||||||
|
if messages[index].get("role") in ("tool", "function")
|
||||||
|
)
|
||||||
|
if not dropped_tool_indices:
|
||||||
|
verbose_proxy_logger.debug("TypeSafe: all evaluated exchanges still relevant; request unchanged")
|
||||||
|
return inputs
|
||||||
|
|
||||||
|
compacted_messages: Final = [ # mutable-ok: structured_messages contract is a list of dicts
|
||||||
|
{**message, "content": DROPPED_RESULT_TEXT} # mutable-ok: JSON message row
|
||||||
|
if index in dropped_tool_indices
|
||||||
|
else message
|
||||||
|
for index, message in enumerate(messages)
|
||||||
|
]
|
||||||
|
chars_removed: Final = sum(
|
||||||
|
len(content_to_text(messages[index].get("content"))) - len(DROPPED_RESULT_TEXT)
|
||||||
|
for index in dropped_tool_indices
|
||||||
|
)
|
||||||
|
exchanges_dropped: Final = len(dropped_ordinals)
|
||||||
|
verbose_proxy_logger.info(
|
||||||
|
"TypeSafe: evaluated %s tool exchange(s), dropped %s, ~%s chars removed",
|
||||||
|
len(candidates),
|
||||||
|
exchanges_dropped,
|
||||||
|
chars_removed,
|
||||||
|
)
|
||||||
|
self.add_standard_logging_guardrail_information_to_request_data( # pyright: ignore[reportUnknownMemberType] # untyped base helper
|
||||||
|
guardrail_json_response={ # mutable-ok: must stay JSON-serializable for shared logging
|
||||||
|
"exchanges_evaluated": len(candidates),
|
||||||
|
"exchanges_dropped": exchanges_dropped,
|
||||||
|
"chars_removed": chars_removed,
|
||||||
|
"model": self.jev_model,
|
||||||
|
},
|
||||||
|
request_data=request_data,
|
||||||
|
guardrail_status="success",
|
||||||
|
guardrail_provider="typesafe",
|
||||||
|
start_time=start_time,
|
||||||
|
end_time=end_time,
|
||||||
|
duration=end_time - start_time,
|
||||||
|
)
|
||||||
|
return {**inputs, "structured_messages": compacted_messages} # pyright: ignore[reportReturnType] # mutable-ok: inputs protocol is a plain dict # plain dicts satisfy AllMessageValues at runtime
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def get_config_model() -> type[TypeSafeGuardrailConfigModel] | None:
|
||||||
|
from litellm.types.proxy.guardrails.guardrail_hooks.typesafe import (
|
||||||
|
TypeSafeGuardrailConfigModel,
|
||||||
|
)
|
||||||
|
|
||||||
|
return TypeSafeGuardrailConfigModel
|
||||||
|
|
@ -62,6 +62,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.singulr import (
|
||||||
from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (
|
from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (
|
||||||
ToolPermissionGuardrailConfigModel,
|
ToolPermissionGuardrailConfigModel,
|
||||||
)
|
)
|
||||||
|
from litellm.types.proxy.guardrails.guardrail_hooks.typesafe import (
|
||||||
|
TypeSafeGuardrailConfigModel,
|
||||||
|
)
|
||||||
from litellm.types.proxy.guardrails.guardrail_hooks.vigil_guard import (
|
from litellm.types.proxy.guardrails.guardrail_hooks.vigil_guard import (
|
||||||
VigilGuardGuardrailConfigModel,
|
VigilGuardGuardrailConfigModel,
|
||||||
)
|
)
|
||||||
|
|
@ -138,6 +141,7 @@ class SupportedGuardrailIntegrations(Enum):
|
||||||
SINGULR = "singulr"
|
SINGULR = "singulr"
|
||||||
HEADROOM = "headroom"
|
HEADROOM = "headroom"
|
||||||
COMPRESR = "compresr"
|
COMPRESR = "compresr"
|
||||||
|
TYPESAFE = "typesafe"
|
||||||
STRAIKER = "straiker"
|
STRAIKER = "straiker"
|
||||||
ALICE = "alice"
|
ALICE = "alice"
|
||||||
AGENT_365 = "agent_365"
|
AGENT_365 = "agent_365"
|
||||||
|
|
@ -1055,7 +1059,7 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up
|
||||||
default="fail_closed",
|
default="fail_closed",
|
||||||
description=(
|
description=(
|
||||||
"Behavior when a guardrail endpoint is unreachable due to network errors. "
|
"Behavior when a guardrail endpoint is unreachable due to network errors. "
|
||||||
"Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. "
|
"Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', 'compresr', and 'typesafe'. "
|
||||||
"'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed."
|
"'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed."
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
@ -1171,6 +1175,7 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o
|
||||||
LakeraV2GuardrailConfigModel,
|
LakeraV2GuardrailConfigModel,
|
||||||
HeadroomGuardrailConfigModel,
|
HeadroomGuardrailConfigModel,
|
||||||
CompresrGuardrailConfigModel,
|
CompresrGuardrailConfigModel,
|
||||||
|
TypeSafeGuardrailConfigModel,
|
||||||
RepelloAIGuardrailConfigModel,
|
RepelloAIGuardrailConfigModel,
|
||||||
LassoGuardrailConfigModel,
|
LassoGuardrailConfigModel,
|
||||||
DeepKeepGuardrailConfigModel,
|
DeepKeepGuardrailConfigModel,
|
||||||
|
|
|
||||||
63
litellm/types/proxy/guardrails/guardrail_hooks/typesafe.py
Normal file
63
litellm/types/proxy/guardrails/guardrail_hooks/typesafe.py
Normal file
|
|
@ -0,0 +1,63 @@
|
||||||
|
from typing import Literal
|
||||||
|
|
||||||
|
from pydantic import BaseModel, Field
|
||||||
|
|
||||||
|
from .base import GuardrailConfigModel
|
||||||
|
|
||||||
|
|
||||||
|
class TypeSafeGuardrailOptionalParams(BaseModel):
|
||||||
|
"""Optional tuning knobs for the TypeSafe (Jev) compaction guardrail."""
|
||||||
|
|
||||||
|
relevance_threshold: float | None = Field(
|
||||||
|
default=None,
|
||||||
|
ge=0.0,
|
||||||
|
le=1.0,
|
||||||
|
description=(
|
||||||
|
"Relevance cutoff in [0, 1]. A completed tool exchange is dropped when Jev "
|
||||||
|
"scores the probability that it is still needed below this value. Defaults to 0.2."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
min_chars_to_evaluate: int | None = Field(
|
||||||
|
default=None,
|
||||||
|
ge=0,
|
||||||
|
description=(
|
||||||
|
"Skip tool exchanges whose combined tool-result text is shorter than this many characters. Defaults to 200."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
max_result_chars_in_state: int | None = Field(
|
||||||
|
default=None,
|
||||||
|
ge=1,
|
||||||
|
description=(
|
||||||
|
"Tool result text is truncated to this many characters when sent to the Jev evaluator, "
|
||||||
|
"keeping the head and tail. Defaults to 4000."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class TypeSafeGuardrailConfigModel(GuardrailConfigModel[TypeSafeGuardrailOptionalParams]):
|
||||||
|
api_key: str | None = Field(
|
||||||
|
default=None,
|
||||||
|
description="TypeSafe API key, sent as a Bearer token. Falls back to the TYPESAFE_API_KEY env var.",
|
||||||
|
)
|
||||||
|
api_base: str | None = Field(
|
||||||
|
default=None,
|
||||||
|
description=(
|
||||||
|
"Base URL of the TypeSafe API. Falls back to the TYPESAFE_API_BASE env var, then https://api.typesafe.ai."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
model: str | None = Field(
|
||||||
|
default=None,
|
||||||
|
description="TypeSafe evaluation model (not the LLM). Defaults to 'jev-latest'.",
|
||||||
|
)
|
||||||
|
unreachable_fallback: Literal["fail_closed", "fail_open"] = Field(
|
||||||
|
default="fail_open",
|
||||||
|
description=(
|
||||||
|
"Behavior when the TypeSafe evaluation service is unreachable or errors. "
|
||||||
|
"'fail_open' (default) forwards the request uncompacted. 'fail_closed' "
|
||||||
|
"raises an error instead."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def ui_friendly_name() -> str:
|
||||||
|
return "TypeSafe (Jev) Compaction"
|
||||||
|
|
@ -0,0 +1,409 @@
|
||||||
|
"""
|
||||||
|
Unit tests for the TypeSafe (Jev) compaction guardrail.
|
||||||
|
|
||||||
|
Tests cover:
|
||||||
|
- exchanges scored below relevance_threshold have their tool rows blanked while
|
||||||
|
assistant tool-call rows and kept exchanges pass through verbatim, without
|
||||||
|
mutating the caller's message list
|
||||||
|
- protected rows (system, last user, and the last tool exchange via the
|
||||||
|
last-assistant rule) are never sent to Jev even when long
|
||||||
|
- exchanges under min_chars_to_evaluate are skipped
|
||||||
|
- request shape: POST {api_base}/v1/systemone with Bearer auth, one noul
|
||||||
|
question per candidate keyed e<i>, task = last user text, results truncated
|
||||||
|
to max_result_chars_in_state
|
||||||
|
- identity return when there are no candidates or nothing is dropped
|
||||||
|
- fail_open forwards uncompacted on service failure; fail_closed raises
|
||||||
|
- response input_type passthrough and initialize_guardrail wiring
|
||||||
|
"""
|
||||||
|
|
||||||
|
from unittest.mock import AsyncMock, MagicMock, PropertyMock
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from fastapi import HTTPException
|
||||||
|
|
||||||
|
from litellm.proxy.guardrails.guardrail_hooks.typesafe import (
|
||||||
|
TypeSafeGuardrail,
|
||||||
|
guardrail_class_registry,
|
||||||
|
guardrail_initializer_registry,
|
||||||
|
initialize_guardrail,
|
||||||
|
)
|
||||||
|
from litellm.proxy.guardrails.guardrail_hooks.typesafe.typesafe import DROPPED_RESULT_TEXT
|
||||||
|
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||||
|
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||||
|
|
||||||
|
FAKE_API_BASE = "https://typesafe.example.com"
|
||||||
|
FAKE_API_KEY = "ts_test-key"
|
||||||
|
|
||||||
|
SYSTEM_TEXT = "You are a research assistant."
|
||||||
|
USER_TEXT = "Which 2026 EV has the longest range?"
|
||||||
|
TOOL_OUTPUT_LONG = "Result: EV range comparison. " * 40
|
||||||
|
TOOL_OUTPUT_SHORT = "short"
|
||||||
|
|
||||||
|
|
||||||
|
def _exchange(call_id: str, tool_text: str, name: str = "web_search") -> list[dict[str, object]]:
|
||||||
|
return [
|
||||||
|
{
|
||||||
|
"role": "assistant",
|
||||||
|
"content": None,
|
||||||
|
"tool_calls": [
|
||||||
|
{
|
||||||
|
"id": call_id,
|
||||||
|
"type": "function",
|
||||||
|
"function": {"name": name, "arguments": '{"query": "ev"}'},
|
||||||
|
}
|
||||||
|
],
|
||||||
|
},
|
||||||
|
{"role": "tool", "tool_call_id": call_id, "name": name, "content": tool_text},
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def _messages(*, tail: list[dict[str, object]] | None = None) -> list[dict[str, object]]:
|
||||||
|
base = [
|
||||||
|
{"role": "system", "content": SYSTEM_TEXT},
|
||||||
|
{"role": "user", "content": USER_TEXT},
|
||||||
|
]
|
||||||
|
return base + (tail or [])
|
||||||
|
|
||||||
|
|
||||||
|
def _make_guardrail(
|
||||||
|
handler: MagicMock | None = None,
|
||||||
|
*,
|
||||||
|
max_result_chars_in_state: int | None = None,
|
||||||
|
unreachable_fallback: str | None = None,
|
||||||
|
) -> TypeSafeGuardrail:
|
||||||
|
return TypeSafeGuardrail(
|
||||||
|
api_base=FAKE_API_BASE,
|
||||||
|
api_key=FAKE_API_KEY,
|
||||||
|
guardrail_name="typesafe",
|
||||||
|
default_on=True,
|
||||||
|
async_handler=handler or _make_handler({"e0": 0.9}),
|
||||||
|
max_result_chars_in_state=max_result_chars_in_state,
|
||||||
|
unreachable_fallback=unreachable_fallback,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _make_handler(answers: dict[str, float], status: int = 200) -> MagicMock:
|
||||||
|
response = MagicMock()
|
||||||
|
response.status_code = status
|
||||||
|
response.json.return_value = {
|
||||||
|
"model": "jev-1.13.0",
|
||||||
|
"answers": {qid: {"type": "noul", "noul": score} for qid, score in answers.items()},
|
||||||
|
"usage": {"input_tokens": 10, "output_tokens": 1},
|
||||||
|
}
|
||||||
|
response.text = ""
|
||||||
|
handler = MagicMock()
|
||||||
|
handler.post = AsyncMock(return_value=response)
|
||||||
|
return handler
|
||||||
|
|
||||||
|
|
||||||
|
def _inputs(messages: list[dict[str, object]]) -> GenericGuardrailAPIInputs:
|
||||||
|
return GenericGuardrailAPIInputs(structured_messages=messages)
|
||||||
|
|
||||||
|
|
||||||
|
async def _apply(
|
||||||
|
guardrail: TypeSafeGuardrail, messages: list[dict[str, object]], input_type: str = "request"
|
||||||
|
) -> GenericGuardrailAPIInputs:
|
||||||
|
return await guardrail.apply_guardrail(
|
||||||
|
inputs=_inputs(messages),
|
||||||
|
request_data={},
|
||||||
|
input_type=input_type, # pyright: ignore[reportArgumentType] # test uses the same literal domain
|
||||||
|
logging_obj=None,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_low_noul_exchange_blanked_high_kept_and_input_not_mutated():
|
||||||
|
handler = _make_handler({"e0": 0.1, "e1": 0.95})
|
||||||
|
guardrail = _make_guardrail(handler)
|
||||||
|
messages = _messages(
|
||||||
|
tail=[
|
||||||
|
*_exchange("call_1", TOOL_OUTPUT_LONG),
|
||||||
|
*_exchange("call_2", TOOL_OUTPUT_LONG),
|
||||||
|
{"role": "assistant", "content": "still thinking"},
|
||||||
|
]
|
||||||
|
)
|
||||||
|
snapshot = [dict(m) for m in messages]
|
||||||
|
|
||||||
|
result = await _apply(guardrail, messages)
|
||||||
|
out = result["structured_messages"]
|
||||||
|
|
||||||
|
assert out[3]["content"] == DROPPED_RESULT_TEXT
|
||||||
|
assert out[3]["tool_call_id"] == "call_1"
|
||||||
|
assert out[3]["role"] == "tool"
|
||||||
|
assert out[5]["content"] == TOOL_OUTPUT_LONG
|
||||||
|
assert out[2] == messages[2]
|
||||||
|
assert out[4] == messages[4]
|
||||||
|
assert out[6]["content"] == "still thinking"
|
||||||
|
assert messages == snapshot
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_last_exchange_and_protected_rows_never_evaluated():
|
||||||
|
handler = _make_handler({"e0": 0.05})
|
||||||
|
guardrail = _make_guardrail(handler)
|
||||||
|
messages = _messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), *_exchange("call_2", TOOL_OUTPUT_LONG)])
|
||||||
|
|
||||||
|
result = await _apply(guardrail, messages)
|
||||||
|
|
||||||
|
payload = handler.post.call_args.kwargs["json"]
|
||||||
|
assert list(payload["questions"]) == ["e0"]
|
||||||
|
assert list(payload["state"]["tool_exchanges"]) == ["e0"]
|
||||||
|
assert payload["state"]["task"] == USER_TEXT
|
||||||
|
assert payload["state"]["system"] == SYSTEM_TEXT
|
||||||
|
out = result["structured_messages"]
|
||||||
|
assert out[3]["content"] == DROPPED_RESULT_TEXT
|
||||||
|
assert out[5]["content"] == TOOL_OUTPUT_LONG
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_short_exchange_not_sent():
|
||||||
|
handler = _make_handler({"e0": 0.9})
|
||||||
|
guardrail = _make_guardrail(handler)
|
||||||
|
messages = _messages(
|
||||||
|
tail=[
|
||||||
|
*_exchange("call_1", TOOL_OUTPUT_SHORT),
|
||||||
|
*_exchange("call_2", TOOL_OUTPUT_LONG),
|
||||||
|
{"role": "assistant", "content": "done"},
|
||||||
|
]
|
||||||
|
)
|
||||||
|
result = await _apply(guardrail, messages)
|
||||||
|
payload = handler.post.call_args.kwargs["json"]
|
||||||
|
assert list(payload["questions"]) == ["e0"]
|
||||||
|
exchange = payload["state"]["tool_exchanges"]["e0"]
|
||||||
|
assert exchange["result"] == TOOL_OUTPUT_LONG
|
||||||
|
assert result is not None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_request_body_shape_and_truncation():
|
||||||
|
handler = _make_handler({"e0": 0.9})
|
||||||
|
guardrail = _make_guardrail(handler, max_result_chars_in_state=50)
|
||||||
|
messages = _messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "done"}])
|
||||||
|
await _apply(guardrail, messages)
|
||||||
|
|
||||||
|
kwargs = handler.post.call_args.kwargs
|
||||||
|
assert kwargs["url"].endswith("/v1/systemone")
|
||||||
|
assert kwargs["url"].startswith(FAKE_API_BASE)
|
||||||
|
assert kwargs["headers"]["Authorization"] == f"Bearer {FAKE_API_KEY}"
|
||||||
|
assert kwargs["headers"]["Content-Type"] == "application/json"
|
||||||
|
payload = kwargs["json"]
|
||||||
|
assert payload["model"] == "jev-latest"
|
||||||
|
assert list(payload["questions"]) == ["e0"]
|
||||||
|
assert payload["questions"]["e0"]["type"] == "noul"
|
||||||
|
assert "e0" in payload["questions"]["e0"]["instructions"]
|
||||||
|
assert payload["state"]["task"] == USER_TEXT
|
||||||
|
exchange = payload["state"]["tool_exchanges"]["e0"]
|
||||||
|
assert len(exchange["result"]) == 50
|
||||||
|
assert exchange["result"].startswith(TOOL_OUTPUT_LONG[:10])
|
||||||
|
assert exchange["result"].endswith(TOOL_OUTPUT_LONG[-11:])
|
||||||
|
assert list(exchange["tool_calls"]) == [{"name": "web_search", "arguments": '{"query": "ev"}'}]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_no_candidates_returns_identity_and_skips_http():
|
||||||
|
handler = _make_handler({})
|
||||||
|
guardrail = _make_guardrail(handler)
|
||||||
|
inputs = _inputs(_messages(tail=[{"role": "assistant", "content": "plain answer"}]))
|
||||||
|
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
|
||||||
|
assert result is inputs
|
||||||
|
handler.post.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_all_above_threshold_returns_identity():
|
||||||
|
handler = _make_handler({"e0": 0.9})
|
||||||
|
guardrail = _make_guardrail(handler)
|
||||||
|
inputs = _inputs(_messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "x"}]))
|
||||||
|
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
|
||||||
|
assert result is inputs
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_fail_open_returns_inputs_on_exception():
|
||||||
|
handler = MagicMock()
|
||||||
|
handler.post = AsyncMock(side_effect=Exception("connection refused"))
|
||||||
|
guardrail = _make_guardrail(handler, unreachable_fallback="fail_open")
|
||||||
|
inputs = _inputs(_messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "x"}]))
|
||||||
|
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
|
||||||
|
assert result is inputs
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_fail_closed_raises_http_exception():
|
||||||
|
handler = MagicMock()
|
||||||
|
handler.post = AsyncMock(side_effect=Exception("connection refused"))
|
||||||
|
guardrail = _make_guardrail(handler, unreachable_fallback="fail_closed")
|
||||||
|
inputs = _inputs(_messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "x"}]))
|
||||||
|
with pytest.raises(HTTPException) as exc_info:
|
||||||
|
await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
|
||||||
|
assert exc_info.value.status_code == 502
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_fail_open_on_non_2xx():
|
||||||
|
handler = _make_handler({"e0": 0.9}, status=500)
|
||||||
|
guardrail = _make_guardrail(handler)
|
||||||
|
inputs = _inputs(_messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "x"}]))
|
||||||
|
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
|
||||||
|
assert result is inputs
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_response_input_type_passthrough():
|
||||||
|
handler = _make_handler({"e0": 0.05})
|
||||||
|
guardrail = _make_guardrail(handler)
|
||||||
|
inputs = _inputs(_messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG)]))
|
||||||
|
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="response", logging_obj=None)
|
||||||
|
assert result is inputs
|
||||||
|
handler.post.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
def test_initialize_guardrail_applies_optional_params_and_registry_keys():
|
||||||
|
from litellm.types.guardrails import LitellmParams
|
||||||
|
|
||||||
|
litellm_params = LitellmParams(
|
||||||
|
guardrail="typesafe",
|
||||||
|
mode="pre_call",
|
||||||
|
api_key=FAKE_API_KEY,
|
||||||
|
api_base=FAKE_API_BASE,
|
||||||
|
optional_params={
|
||||||
|
"relevance_threshold": 0.5,
|
||||||
|
"min_chars_to_evaluate": 10,
|
||||||
|
"max_result_chars_in_state": 100,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
callback = initialize_guardrail(litellm_params, {"guardrail_name": "jev-compaction"})
|
||||||
|
assert isinstance(callback, TypeSafeGuardrail)
|
||||||
|
assert callback.relevance_threshold == 0.5
|
||||||
|
assert callback.min_chars_to_evaluate == 10
|
||||||
|
assert callback.max_result_chars_in_state == 100
|
||||||
|
assert callback.unreachable_fallback == "fail_open"
|
||||||
|
assert guardrail_initializer_registry[SupportedGuardrailIntegrations.TYPESAFE.value] is initialize_guardrail
|
||||||
|
assert guardrail_class_registry[SupportedGuardrailIntegrations.TYPESAFE.value] is TypeSafeGuardrail
|
||||||
|
|
||||||
|
|
||||||
|
def test_missing_api_key_raises(monkeypatch):
|
||||||
|
monkeypatch.delenv("TYPESAFE_API_KEY", raising=False)
|
||||||
|
with pytest.raises(ValueError, match="requires an API key"):
|
||||||
|
TypeSafeGuardrail(api_key=None)
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_config_model_and_ui_name():
|
||||||
|
from litellm.types.proxy.guardrails.guardrail_hooks.typesafe import (
|
||||||
|
TypeSafeGuardrailConfigModel,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert TypeSafeGuardrail.get_config_model() is TypeSafeGuardrailConfigModel
|
||||||
|
assert TypeSafeGuardrailConfigModel.ui_friendly_name() == "TypeSafe (Jev) Compaction"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_non_list_and_non_dict_messages_return_identity():
|
||||||
|
guardrail = _make_guardrail()
|
||||||
|
not_a_list = GenericGuardrailAPIInputs(structured_messages={"role": "user"})
|
||||||
|
assert (
|
||||||
|
await guardrail.apply_guardrail(inputs=not_a_list, request_data={}, input_type="request", logging_obj=None)
|
||||||
|
is not_a_list
|
||||||
|
)
|
||||||
|
with_bad_row = _inputs(_messages(tail=[["not", "a", "dict"]]))
|
||||||
|
assert (
|
||||||
|
await guardrail.apply_guardrail(inputs=with_bad_row, request_data={}, input_type="request", logging_obj=None)
|
||||||
|
is with_bad_row
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_odd_tool_call_shapes_yield_no_entries():
|
||||||
|
from litellm.proxy.guardrails.guardrail_hooks.typesafe.typesafe import _tool_call_entries
|
||||||
|
|
||||||
|
assert _tool_call_entries({"tool_calls": "not-a-list"}) == ()
|
||||||
|
assert _tool_call_entries({"tool_calls": None}) == ()
|
||||||
|
assert list(_tool_call_entries({"tool_calls": [42]})) == []
|
||||||
|
entries = _tool_call_entries({"tool_calls": [{"function": {"name": "web_search", "arguments": "{}"}}]})
|
||||||
|
assert list(entries) == [{"name": "web_search", "arguments": "{}"}]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_short_max_chars_uses_prefix_slice():
|
||||||
|
handler = _make_handler({"e0": 0.9})
|
||||||
|
guardrail = _make_guardrail(handler, max_result_chars_in_state=5)
|
||||||
|
await _apply(
|
||||||
|
guardrail, _messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "x"}])
|
||||||
|
)
|
||||||
|
result = handler.post.call_args.kwargs["json"]["state"]["tool_exchanges"]["e0"]["result"]
|
||||||
|
assert result == TOOL_OUTPUT_LONG[:5]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_unreadable_json_body_fails_open():
|
||||||
|
handler = MagicMock()
|
||||||
|
response = MagicMock()
|
||||||
|
response.status_code = 200
|
||||||
|
response.text = "not json"
|
||||||
|
response.json.side_effect = ValueError("no json")
|
||||||
|
handler.post = AsyncMock(return_value=response)
|
||||||
|
guardrail = _make_guardrail(handler)
|
||||||
|
inputs = _inputs(_messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "x"}]))
|
||||||
|
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
|
||||||
|
assert result is inputs
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_malformed_answers_shape_fails_open():
|
||||||
|
handler = MagicMock()
|
||||||
|
response = MagicMock()
|
||||||
|
response.status_code = 200
|
||||||
|
response.text = '{"answers": "oops"}'
|
||||||
|
response.json.return_value = {"answers": "oops"}
|
||||||
|
handler.post = AsyncMock(return_value=response)
|
||||||
|
guardrail = _make_guardrail(handler)
|
||||||
|
inputs = _inputs(_messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "x"}]))
|
||||||
|
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
|
||||||
|
assert result is inputs
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_http_status_error_includes_status_and_undecodable_body():
|
||||||
|
import httpx
|
||||||
|
|
||||||
|
response = MagicMock()
|
||||||
|
response.status_code = 503
|
||||||
|
type(response).text = PropertyMock(side_effect=httpx.DecodingError("bad codec"))
|
||||||
|
handler = MagicMock()
|
||||||
|
handler.post = AsyncMock(side_effect=httpx.HTTPStatusError("unavailable", request=MagicMock(), response=response))
|
||||||
|
guardrail = _make_guardrail(handler)
|
||||||
|
inputs = _inputs(_messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "x"}]))
|
||||||
|
result = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
|
||||||
|
assert result is inputs
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_cancelled_jev_call_propagates():
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
handler = MagicMock()
|
||||||
|
handler.post = AsyncMock(side_effect=asyncio.CancelledError())
|
||||||
|
guardrail = _make_guardrail(handler)
|
||||||
|
inputs = _inputs(_messages(tail=[*_exchange("call_1", TOOL_OUTPUT_LONG), {"role": "assistant", "content": "x"}]))
|
||||||
|
with pytest.raises(asyncio.CancelledError):
|
||||||
|
await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
|
||||||
|
|
||||||
|
|
||||||
|
def test_optional_params_defaults_and_event_hook_coercion():
|
||||||
|
from litellm.proxy.guardrails.guardrail_hooks.typesafe import _coerce_event_hook, _optional_params
|
||||||
|
from litellm.types.guardrails import GuardrailEventHooks, LitellmParams
|
||||||
|
|
||||||
|
assert _coerce_event_hook("pre_call") is GuardrailEventHooks.pre_call
|
||||||
|
assert _coerce_event_hook(["pre_call", "post_call"]) == [
|
||||||
|
GuardrailEventHooks.pre_call,
|
||||||
|
GuardrailEventHooks.post_call,
|
||||||
|
]
|
||||||
|
litellm_params = LitellmParams(guardrail="typesafe", mode="pre_call", api_key=FAKE_API_KEY)
|
||||||
|
params = _optional_params(litellm_params)
|
||||||
|
assert params.relevance_threshold is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_typesafe_initializer_discoverable_via_hook_registries():
|
||||||
|
from litellm.proxy.guardrails.guardrail_registry import get_guardrail_initializer_from_hooks
|
||||||
|
|
||||||
|
initializers = get_guardrail_initializer_from_hooks()
|
||||||
|
assert initializers["typesafe"] is initialize_guardrail
|
||||||
|
|
@ -14,7 +14,7 @@ export const NO_COMPRESSION = "none";
|
||||||
|
|
||||||
/** Guardrail providers that compress prompts, mirroring COMPRESSION_GUARDRAIL_PROVIDERS in
|
/** Guardrail providers that compress prompts, mirroring COMPRESSION_GUARDRAIL_PROVIDERS in
|
||||||
* litellm/proxy/guardrails/auto_router_compression.py. Both are selectable per hop. */
|
* litellm/proxy/guardrails/auto_router_compression.py. Both are selectable per hop. */
|
||||||
export const COMPRESSION_GUARDRAIL_PROVIDERS: readonly string[] = ["headroom", "compresr"];
|
export const COMPRESSION_GUARDRAIL_PROVIDERS: readonly string[] = ["headroom", "compresr", "typesafe"];
|
||||||
|
|
||||||
export const isCompressionGuardrailProvider = (provider: unknown): boolean =>
|
export const isCompressionGuardrailProvider = (provider: unknown): boolean =>
|
||||||
typeof provider === "string" && COMPRESSION_GUARDRAIL_PROVIDERS.includes(provider.toLowerCase());
|
typeof provider === "string" && COMPRESSION_GUARDRAIL_PROVIDERS.includes(provider.toLowerCase());
|
||||||
|
|
|
||||||
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -24352,7 +24352,7 @@ export interface components {
|
||||||
timeout?: number | null;
|
timeout?: number | null;
|
||||||
/**
|
/**
|
||||||
* Unreachable Fallback
|
* Unreachable Fallback
|
||||||
* @description Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed.
|
* @description Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', 'compresr', and 'typesafe'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed.
|
||||||
* @default fail_closed
|
* @default fail_closed
|
||||||
* @enum {string}
|
* @enum {string}
|
||||||
*/
|
*/
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue