mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
feat(guardrails): add TypeSafe Jev relevance-based compaction guardrail
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin
parent
db37977307
commit
4863e1775b
9 changed files with 812 additions and 5 deletions
|
|
@ -10006,7 +10006,7 @@
|
|||
},
|
||||
"unreachable_fallback": {
|
||||
"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": [
|
||||
"fail_closed",
|
||||
"fail_open"
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ if TYPE_CHECKING:
|
|||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.router import Router
|
||||
|
||||
COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr"})
|
||||
COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr", "typesafe"})
|
||||
_NO_COMPRESSION: Final = "none"
|
||||
|
||||
# A ContextVar, not metadata: metadata reaches spend logs the caller can read, and a
|
||||
|
|
|
|||
|
|
@ -0,0 +1,80 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Final, cast
|
||||
|
||||
from litellm.types.guardrails import (
|
||||
GuardrailEventHooks,
|
||||
Mode,
|
||||
SupportedGuardrailIntegrations,
|
||||
)
|
||||
|
||||
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 [GuardrailEventHooks(item) for item in mode]
|
||||
return GuardrailEventHooks(mode)
|
||||
|
||||
|
||||
def _get_optional_value(litellm_params: LitellmParams, optional_params: object | None, attribute_name: str) -> object:
|
||||
if optional_params is not None:
|
||||
value: Final = getattr(optional_params, attribute_name, None)
|
||||
if value is not None:
|
||||
return cast(object, value)
|
||||
return cast(object, getattr(litellm_params, attribute_name, None))
|
||||
|
||||
|
||||
def _optional_float(litellm_params: LitellmParams, optional_params: object | None, attribute_name: str) -> float | None:
|
||||
value: Final = _get_optional_value(litellm_params, optional_params, attribute_name)
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
||||
return None
|
||||
return float(value)
|
||||
|
||||
|
||||
def _optional_int(litellm_params: LitellmParams, optional_params: object | None, attribute_name: str) -> int | None:
|
||||
value: Final = _get_optional_value(litellm_params, optional_params, attribute_name)
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
return None
|
||||
return value
|
||||
|
||||
|
||||
def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> TypeSafeGuardrail:
|
||||
import litellm
|
||||
|
||||
optional_params: Final = getattr(litellm_params, "optional_params", None)
|
||||
|
||||
_callback: Final = TypeSafeGuardrail(
|
||||
api_base=litellm_params.api_base,
|
||||
api_key=litellm_params.api_key,
|
||||
model=litellm_params.model,
|
||||
relevance_threshold=_optional_float(litellm_params, optional_params, "relevance_threshold"),
|
||||
min_chars_to_evaluate=_optional_int(litellm_params, optional_params, "min_chars_to_evaluate"),
|
||||
max_result_chars_in_state=_optional_int(litellm_params, 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 = {
|
||||
SupportedGuardrailIntegrations.TYPESAFE.value: initialize_guardrail,
|
||||
}
|
||||
|
||||
guardrail_class_registry: Final = {
|
||||
SupportedGuardrailIntegrations.TYPESAFE.value: TypeSafeGuardrail,
|
||||
}
|
||||
390
litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py
Normal file
390
litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py
Normal file
|
|
@ -0,0 +1,390 @@
|
|||
"""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. The assistant tool-call rows stay
|
||||
intact, so the conversation remains well-formed while the dead context stops
|
||||
consuming input tokens.
|
||||
|
||||
Exchanges follow litellm's own compression protection policy: system rows, the
|
||||
last user row, and the last assistant row (which, expanded over its tool
|
||||
exchange, covers the most recent exchange) are never evaluated or rewritten.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Annotated, Final, Literal, TypeGuard, cast
|
||||
|
||||
import httpx
|
||||
from fastapi import HTTPException
|
||||
from httpx import Response as HttpxResponse
|
||||
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
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
|
||||
# The shared GuardrailCallback client carries no per-call bound; an on-request
|
||||
# guardrail must not hold the caller's request for the client's pooled timeout.
|
||||
_JEV_TIMEOUT_SECONDS: Final = 30.0
|
||||
DROPPED_RESULT_TEXT: Final = (
|
||||
"[Tool result removed by TypeSafe compaction: judged no longer relevant to the current task]"
|
||||
)
|
||||
|
||||
|
||||
def _is_str_object_dict(value: object) -> TypeGuard[dict[str, object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip
|
||||
return isinstance(value, dict)
|
||||
|
||||
|
||||
def _is_object_list(value: object) -> TypeGuard[list[object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip
|
||||
return isinstance(value, list)
|
||||
|
||||
|
||||
def _safe_response_text(response: object, limit: int = 500) -> str:
|
||||
try:
|
||||
text: Final = getattr(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 _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_entries(assistant_message: Mapping[str, object]) -> list[dict[str, object]]:
|
||||
tool_calls: Final = assistant_message.get("tool_calls")
|
||||
if not _is_object_list(tool_calls):
|
||||
return []
|
||||
entries: Final[list[dict[str, object]]] = []
|
||||
for tool_call in tool_calls:
|
||||
if not _is_str_object_dict(tool_call):
|
||||
continue
|
||||
function = tool_call.get("function")
|
||||
fn = function if _is_str_object_dict(function) else tool_call
|
||||
entries.append({"name": fn.get("name"), "arguments": fn.get("arguments")})
|
||||
return entries
|
||||
|
||||
|
||||
def _protected_indices(messages: Sequence[Mapping[str, object]]) -> frozenset[int]:
|
||||
"""Rows typesafe must not rewrite, expanded over whole tool exchanges.
|
||||
|
||||
``get_protected_indices`` covers system rows, the last user row, the last
|
||||
assistant row, and cache_control prefixes. Expanding over exchanges keeps an
|
||||
exchange atomic: the last assistant row protects its own tool results too,
|
||||
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,
|
||||
):
|
||||
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 the caller forwards uncompacted; fail_closed raises.
|
||||
Upstream bodies go to server logs only; the raised HTTPException is generic."""
|
||||
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=500, detail={"error": error})
|
||||
|
||||
def _candidate_exchanges(self, messages: list[dict[str, object]]) -> list[tuple[int, ...]]:
|
||||
"""Message-index groups eligible for relevance evaluation, oldest first.
|
||||
|
||||
A candidate is a completed tool exchange: an assistant row that made
|
||||
tool calls plus at least one ``tool``/``function`` row answering it,
|
||||
with no member protected, and enough combined tool-result text to be
|
||||
worth an evaluation call.
|
||||
"""
|
||||
protected: Final = _protected_indices(messages)
|
||||
candidates: Final[list[tuple[int, ...]]] = []
|
||||
for group in group_tool_exchanges(messages):
|
||||
if len(group) < 2:
|
||||
continue
|
||||
if messages[group[0]].get("role") != "assistant":
|
||||
continue
|
||||
if any(member in protected for member in group):
|
||||
continue
|
||||
tool_text = "".join(
|
||||
content_to_text(messages[index].get("content"))
|
||||
for index in group[1:]
|
||||
if messages[index].get("role") in ("tool", "function")
|
||||
)
|
||||
if not tool_text or len(tool_text) < self.min_chars_to_evaluate:
|
||||
continue
|
||||
candidates.append(group)
|
||||
return candidates[-_MAX_EXCHANGES_EVALUATED:]
|
||||
|
||||
def _build_state(self, messages: list[dict[str, object]], candidates: list[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[dict[str, object]] = {}
|
||||
for ordinal, group in enumerate(candidates):
|
||||
result_text = "".join(
|
||||
content_to_text(messages[index].get("content"))
|
||||
for index in group[1:]
|
||||
if messages[index].get("role") in ("tool", "function")
|
||||
)
|
||||
tool_exchanges[f"e{ordinal}"] = {
|
||||
"tool_calls": _tool_call_entries(messages[group[0]]),
|
||||
"result": result_text[: self.max_result_chars_in_state],
|
||||
}
|
||||
return {"task": task, "system": system, "tool_exchanges": tool_exchanges}
|
||||
|
||||
async def _call_systemone(self, state: dict[str, object], question_ids: list[str]) -> _JevSystemOneResponse | None:
|
||||
"""Evaluate each exchange. Returns the response, or None when the service
|
||||
failed and fail_open applies."""
|
||||
payload: Final[dict[str, object]] = {
|
||||
"model": self.jev_model,
|
||||
"state": state,
|
||||
"questions": {
|
||||
question_id: {"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={
|
||||
"Authorization": f"Bearer {self.typesafe_api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
timeout=_JEV_TIMEOUT_SECONDS,
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except httpx.HTTPStatusError as e:
|
||||
resp: Final = getattr(e, "response", None)
|
||||
self._handle_failure(
|
||||
"TypeSafe evaluation service returned an error",
|
||||
{"status_code": getattr(resp, "status_code", None), "body": _safe_response_text(resp)},
|
||||
)
|
||||
return None
|
||||
except (httpx.RequestError, litellm.Timeout) as e:
|
||||
self._handle_failure("TypeSafe evaluation service request failed", {"detail": str(e)})
|
||||
return None
|
||||
except Exception as e:
|
||||
self._handle_failure("TypeSafe evaluation service request failed", {"detail": str(e)})
|
||||
return None
|
||||
if not 200 <= raw_response.status_code < 300:
|
||||
self._handle_failure(
|
||||
"TypeSafe evaluation service returned an error",
|
||||
{"status_code": raw_response.status_code, "body": _safe_response_text(raw_response)},
|
||||
)
|
||||
return None
|
||||
try:
|
||||
body: Final = cast(object, raw_response.json())
|
||||
except (ValueError, httpx.DecodingError, RecursionError):
|
||||
self._handle_failure(
|
||||
"TypeSafe evaluation service returned an unreadable response",
|
||||
{"body": _safe_response_text(raw_response)},
|
||||
)
|
||||
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)},
|
||||
)
|
||||
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 = inputs.get("structured_messages")
|
||||
if not _is_object_list(structured_messages) or not structured_messages:
|
||||
return inputs
|
||||
messages: Final = [m for m in structured_messages if _is_str_object_dict(m)]
|
||||
if len(messages) != len(structured_messages):
|
||||
return inputs
|
||||
|
||||
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 = [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:
|
||||
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 = [
|
||||
{**message, "content": DROPPED_RESULT_TEXT} 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={
|
||||
"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] # 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 (
|
||||
ToolPermissionGuardrailConfigModel,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.typesafe import (
|
||||
TypeSafeGuardrailConfigModel,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.vigil_guard import (
|
||||
VigilGuardGuardrailConfigModel,
|
||||
)
|
||||
|
|
@ -138,6 +141,7 @@ class SupportedGuardrailIntegrations(Enum):
|
|||
SINGULR = "singulr"
|
||||
HEADROOM = "headroom"
|
||||
COMPRESR = "compresr"
|
||||
TYPESAFE = "typesafe"
|
||||
STRAIKER = "straiker"
|
||||
ALICE = "alice"
|
||||
AGENT_365 = "agent_365"
|
||||
|
|
@ -1055,7 +1059,7 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up
|
|||
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'. "
|
||||
"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."
|
||||
),
|
||||
)
|
||||
|
|
@ -1171,6 +1175,7 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o
|
|||
LakeraV2GuardrailConfigModel,
|
||||
HeadroomGuardrailConfigModel,
|
||||
CompresrGuardrailConfigModel,
|
||||
TypeSafeGuardrailConfigModel,
|
||||
RepelloAIGuardrailConfigModel,
|
||||
LassoGuardrailConfigModel,
|
||||
DeepKeepGuardrailConfigModel,
|
||||
|
|
|
|||
58
litellm/types/proxy/guardrails/guardrail_hooks/typesafe.py
Normal file
58
litellm/types/proxy/guardrails/guardrail_hooks/typesafe.py
Normal file
|
|
@ -0,0 +1,58 @@
|
|||
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,
|
||||
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,
|
||||
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,
|
||||
description=(
|
||||
"Tool result text is truncated to this many characters when sent to the Jev evaluator. 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,274 @@
|
|||
"""
|
||||
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
|
||||
|
||||
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 # > 200 chars
|
||||
TOOL_OUTPUT_SHORT = "short"
|
||||
|
||||
|
||||
def _exchange(call_id: str, tool_text: str, name: str = "web_search") -> list[dict]:
|
||||
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 | None = None) -> list[dict]:
|
||||
base = [
|
||||
{"role": "system", "content": SYSTEM_TEXT},
|
||||
{"role": "user", "content": USER_TEXT},
|
||||
]
|
||||
return base + (tail or [])
|
||||
|
||||
|
||||
def _make_guardrail(handler: MagicMock | None = None, **kwargs) -> TypeSafeGuardrail:
|
||||
defaults = dict(
|
||||
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}),
|
||||
)
|
||||
defaults.update(kwargs)
|
||||
return TypeSafeGuardrail(**defaults)
|
||||
|
||||
|
||||
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) -> GenericGuardrailAPIInputs:
|
||||
return GenericGuardrailAPIInputs(structured_messages=messages)
|
||||
|
||||
|
||||
async def _apply(guardrail: TypeSafeGuardrail, messages: list, input_type: str = "request"):
|
||||
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)
|
||||
# Ends on a tool result: the last assistant row is protected, so the whole
|
||||
# last exchange is out of scope even though its text is long.
|
||||
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 exchange["result"] == TOOL_OUTPUT_LONG[:50]
|
||||
assert 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):
|
||||
await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request", logging_obj=None)
|
||||
|
||||
|
||||
@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
|
||||
|
|
@ -14,7 +14,7 @@ export const NO_COMPRESSION = "none";
|
|||
|
||||
/** Guardrail providers that compress prompts, mirroring COMPRESSION_GUARDRAIL_PROVIDERS in
|
||||
* 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 =>
|
||||
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
|
|
@ -24272,7 +24272,7 @@ export interface components {
|
|||
timeout?: number | null;
|
||||
/**
|
||||
* 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
|
||||
* @enum {string}
|
||||
*/
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue