From 4863e1775ba4630df7d6596b164b509d836357a2 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 04:33:16 +0000 Subject: [PATCH 1/9] feat(guardrails): add TypeSafe Jev relevance-based compaction guardrail Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- .../guardrails/auto_router_compression.py | 2 +- .../guardrail_hooks/typesafe/__init__.py | 80 ++++ .../guardrail_hooks/typesafe/typesafe.py | 390 ++++++++++++++++++ litellm/types/guardrails.py | 7 +- .../guardrails/guardrail_hooks/typesafe.py | 58 +++ .../guardrail_hooks/test_typesafe.py | 274 ++++++++++++ .../add_model/buildAutoRouterCompression.ts | 2 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 9 files changed, 812 insertions(+), 5 deletions(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/typesafe.py create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index b244678e201..a215cc573d7 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -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" diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index c37b9fff1f0..335419c6372 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -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 diff --git a/litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py new file mode 100644 index 00000000000..c347c863a1b --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py @@ -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, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py b/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py new file mode 100644 index 00000000000..ef192030e95 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py @@ -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 "" + 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 diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 6346e13f3ba..400cadd69e7 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -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, diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/typesafe.py b/litellm/types/proxy/guardrails/guardrail_hooks/typesafe.py new file mode 100644 index 00000000000..4e742bfc7be --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/typesafe.py @@ -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" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py new file mode 100644 index 00000000000..293e4f8c344 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py @@ -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, 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 diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts index c3503f78afb..292b497f6df 100644 --- a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts @@ -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()); diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index fd882937e79..16d61bd9b0f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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} */ From dffb6a38d95f5e2fa2c259f06f50243c46ca2ab2 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 04:35:58 +0000 Subject: [PATCH 2/9] refactor(guardrails): tighten typesafe guardrail typing and error handling Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/typesafe/__init__.py | 42 +++++-------- .../guardrail_hooks/typesafe/typesafe.py | 63 +++++++------------ .../guardrail_hooks/test_typesafe.py | 3 +- 3 files changed, 40 insertions(+), 68 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py index c347c863a1b..40f040eb676 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py @@ -1,12 +1,17 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Final, cast +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 @@ -24,40 +29,27 @@ def _coerce_event_hook( 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 _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 = getattr(litellm_params, "optional_params", None) + 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_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"), + 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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py b/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py index ef192030e95..360219e0ad5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py +++ b/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py @@ -3,13 +3,7 @@ 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. +tool results Jev judges no longer relevant. """ from __future__ import annotations @@ -24,7 +18,6 @@ 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 ( @@ -56,8 +49,6 @@ 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]" @@ -72,9 +63,11 @@ def _is_object_list(value: object) -> TypeGuard[list[object]]: # guard-ok: isin return isinstance(value, list) -def _safe_response_text(response: object, limit: int = 500) -> str: +def _safe_response_text(response: HttpxResponse | None, limit: int = 500) -> str: + if response is None: + return "" try: - text: Final = getattr(response, "text", "") + text: Final = response.text except httpx.DecodingError: return "" return (text or "")[:limit] @@ -120,13 +113,7 @@ def _tool_call_entries(assistant_message: Mapping[str, object]) -> list[dict[str 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. - """ + """``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 @@ -180,8 +167,7 @@ class TypeSafeGuardrail(CustomGuardrail): ) 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.""" + """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", @@ -190,16 +176,10 @@ class TypeSafeGuardrail(CustomGuardrail): ) return verbose_proxy_logger.error("TypeSafe: %s. detail=%s", error, log_detail) - raise HTTPException(status_code=500, detail={"error": error}) + raise HTTPException(status_code=502, 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. - """ + """Completed tool exchanges eligible for evaluation: unprotected, and long enough to be worth a call.""" protected: Final = _protected_indices(messages) candidates: Final[list[tuple[int, ...]]] = [] for group in group_tool_exchanges(messages): @@ -245,8 +225,7 @@ class TypeSafeGuardrail(CustomGuardrail): 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.""" + """Returns the response, or None when the service failed and fail_open applies.""" payload: Final[dict[str, object]] = { "model": self.jev_model, "state": state, @@ -267,18 +246,18 @@ class TypeSafeGuardrail(CustomGuardrail): ) 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)}) + detail: Final[dict[str, object]] = ( + { + "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)} + ) + self._handle_failure("TypeSafe evaluation service request failed", detail) return None if not 200 <= raw_response.status_code < 300: self._handle_failure( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py index 293e4f8c344..8b732091d7c 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py @@ -227,8 +227,9 @@ async def test_fail_closed_raises_http_exception(): 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): + 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 From ca8c9062d80f1ee5f2e34765d0ba3dbc4fd7af72 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 04:48:11 +0000 Subject: [PATCH 3/9] fix(guardrails): satisfy strict lint gates in typesafe guardrail Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/typesafe/typesafe.py | 44 ++++++++++++------- 1 file changed, 28 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py b/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py index 360219e0ad5..9fca6eaed20 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py +++ b/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py @@ -11,7 +11,7 @@ from __future__ import annotations import asyncio import time from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Annotated, Final, Literal, TypeGuard, cast +from typing import TYPE_CHECKING, Annotated, Final, Literal import httpx from fastapi import HTTPException @@ -55,12 +55,22 @@ DROPPED_RESULT_TEXT: Final = ( ) -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) +_STR_OBJECT_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) -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 _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: @@ -99,15 +109,16 @@ def _question_instructions(question_id: str) -> str: 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): + tool_calls: Final = _as_object_list(assistant_message.get("tool_calls")) + if tool_calls is None: return [] entries: Final[list[dict[str, object]]] = [] for tool_call in tool_calls: - if not _is_str_object_dict(tool_call): + parsed_call = _as_str_object_dict(tool_call) + if parsed_call is None: continue - function = tool_call.get("function") - fn = function if _is_str_object_dict(function) else tool_call + function = _as_str_object_dict(parsed_call.get("function")) + fn = function if function is not None else parsed_call entries.append({"name": fn.get("name"), "arguments": fn.get("arguments")}) return entries @@ -137,7 +148,7 @@ class TypeSafeGuardrail(CustomGuardrail): 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") @@ -266,7 +277,7 @@ class TypeSafeGuardrail(CustomGuardrail): ) return None try: - body: Final = cast(object, raw_response.json()) + 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", @@ -293,12 +304,13 @@ class TypeSafeGuardrail(CustomGuardrail): if input_type != "request": return inputs - structured_messages: Final = inputs.get("structured_messages") - if not _is_object_list(structured_messages) or not structured_messages: + structured_messages: Final = _as_object_list(inputs.get("structured_messages")) + if 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): + parsed_messages: Final = [_as_str_object_dict(m) for m in structured_messages] + if any(m is None for m in parsed_messages): return inputs + messages: Final = [m for m in parsed_messages if m is not None] candidates: Final = self._candidate_exchanges(messages) if not candidates: From fd4476b1308feba1da8e86cd162401ce1bb2da4c Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 04:49:50 +0000 Subject: [PATCH 4/9] fix(guardrails): bound typesafe tuning params, preserve result tail, log fail-open status Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/typesafe/typesafe.py | 26 +++++++++++++++++- .../guardrails/guardrail_hooks/typesafe.py | 7 ++++- .../guardrail_hooks/test_typesafe.py | 27 ++++++++++++------- 3 files changed, 49 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py b/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py index 9fca6eaed20..b3cf5bcf800 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py +++ b/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py @@ -53,6 +53,7 @@ _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]) @@ -99,6 +100,17 @@ class _JevSystemOneResponse(BaseModel): _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 " @@ -231,7 +243,7 @@ class TypeSafeGuardrail(CustomGuardrail): ) tool_exchanges[f"e{ordinal}"] = { "tool_calls": _tool_call_entries(messages[group[0]]), - "result": result_text[: self.max_result_chars_in_state], + "result": _truncate_for_state(result_text, self.max_result_chars_in_state), } return {"task": task, "system": system, "tool_exchanges": tool_exchanges} @@ -324,6 +336,18 @@ class TypeSafeGuardrail(CustomGuardrail): 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={ + "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( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/typesafe.py b/litellm/types/proxy/guardrails/guardrail_hooks/typesafe.py index 4e742bfc7be..59482d2e190 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/typesafe.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/typesafe.py @@ -10,6 +10,8 @@ class TypeSafeGuardrailOptionalParams(BaseModel): 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." @@ -17,14 +19,17 @@ class TypeSafeGuardrailOptionalParams(BaseModel): ) 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. Defaults to 4000." + "Tool result text is truncated to this many characters when sent to the Jev evaluator, " + "keeping the head and tail. Defaults to 4000." ), ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py index 8b732091d7c..e2e5fc2fefc 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py @@ -40,7 +40,7 @@ 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]: +def _exchange(call_id: str, tool_text: str, name: str = "web_search") -> list[dict[str, object]]: return [ { "role": "assistant", @@ -57,7 +57,7 @@ def _exchange(call_id: str, tool_text: str, name: str = "web_search") -> list[di ] -def _messages(*, tail: list | None = None) -> list[dict]: +def _messages(*, tail: list[dict[str, object]] | None = None) -> list[dict[str, object]]: base = [ {"role": "system", "content": SYSTEM_TEXT}, {"role": "user", "content": USER_TEXT}, @@ -65,16 +65,21 @@ def _messages(*, tail: list | None = None) -> list[dict]: return base + (tail or []) -def _make_guardrail(handler: MagicMock | None = None, **kwargs) -> TypeSafeGuardrail: - defaults = dict( +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, ) - defaults.update(kwargs) - return TypeSafeGuardrail(**defaults) def _make_handler(answers: dict[str, float], status: int = 200) -> MagicMock: @@ -91,11 +96,13 @@ def _make_handler(answers: dict[str, float], status: int = 200) -> MagicMock: return handler -def _inputs(messages: list) -> GenericGuardrailAPIInputs: +def _inputs(messages: list[dict[str, object]]) -> GenericGuardrailAPIInputs: return GenericGuardrailAPIInputs(structured_messages=messages) -async def _apply(guardrail: TypeSafeGuardrail, messages: list, input_type: str = "request"): +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={}, @@ -188,7 +195,9 @@ async def test_request_body_shape_and_truncation(): 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 len(exchange["result"]) == 50 + assert exchange["result"].startswith(TOOL_OUTPUT_LONG[:10]) + assert exchange["result"].endswith(TOOL_OUTPUT_LONG[-11:]) assert exchange["tool_calls"] == [{"name": "web_search", "arguments": '{"query": "ev"}'}] From 91d4c579e43bd5429741254ea0ff94956edc96fc Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 05:06:48 +0000 Subject: [PATCH 5/9] refactor(guardrails): freeze or suppress mutable constructions in typesafe guardrail Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/typesafe/__init__.py | 15 +- .../guardrail_hooks/typesafe/typesafe.py | 148 ++++++++++-------- .../guardrail_hooks/test_typesafe.py | 2 +- 3 files changed, 91 insertions(+), 74 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py index 40f040eb676..c1b05597a6c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py @@ -1,5 +1,6 @@ from __future__ import annotations +from types import MappingProxyType from typing import TYPE_CHECKING, Final from pydantic import BaseModel @@ -25,7 +26,9 @@ def _coerce_event_hook( if isinstance(mode, Mode): return mode if isinstance(mode, list): - return [GuardrailEventHooks(item) for item in mode] + return [ + GuardrailEventHooks(item) for item in mode + ] # mutable-ok: CustomGuardrail event_hook contract wants a list return GuardrailEventHooks(mode) @@ -63,10 +66,8 @@ def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> return _callback -guardrail_initializer_registry: Final = { - SupportedGuardrailIntegrations.TYPESAFE.value: initialize_guardrail, -} +guardrail_initializer_registry: Final = MappingProxyType( + {SupportedGuardrailIntegrations.TYPESAFE.value: initialize_guardrail} +) -guardrail_class_registry: Final = { - SupportedGuardrailIntegrations.TYPESAFE.value: TypeSafeGuardrail, -} +guardrail_class_registry: Final = MappingProxyType({SupportedGuardrailIntegrations.TYPESAFE.value: TypeSafeGuardrail}) diff --git a/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py b/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py index b3cf5bcf800..384bfd2c54b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py +++ b/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py @@ -11,6 +11,7 @@ from __future__ import annotations import asyncio import time from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Literal import httpx @@ -120,19 +121,20 @@ def _question_instructions(question_id: str) -> str: ) -def _tool_call_entries(assistant_message: Mapping[str, object]) -> list[dict[str, object]]: +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 [] - entries: Final[list[dict[str, object]]] = [] - for tool_call in tool_calls: - parsed_call = _as_str_object_dict(tool_call) - if parsed_call is None: - continue - function = _as_str_object_dict(parsed_call.get("function")) - fn = function if function is not None else parsed_call - entries.append({"name": fn.get("name"), "arguments": fn.get("arguments")}) - return entries + 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]: @@ -199,30 +201,32 @@ class TypeSafeGuardrail(CustomGuardrail): ) return verbose_proxy_logger.error("TypeSafe: %s. detail=%s", error, log_detail) - raise HTTPException(status_code=502, detail={"error": error}) + raise HTTPException(status_code=502, detail={"error": error}) # mutable-ok: FastAPI wants a dict detail - def _candidate_exchanges(self, messages: list[dict[str, object]]) -> list[tuple[int, ...]]: + 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[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) + 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:] - def _build_state(self, messages: list[dict[str, object]], candidates: list[tuple[int, ...]]) -> dict[str, object]: + @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")) @@ -234,26 +238,29 @@ class TypeSafeGuardrail(CustomGuardrail): 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_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(result_text, self.max_result_chars_in_state), + "result": _truncate_for_state( + self._exchange_tool_text(messages, group), self.max_result_chars_in_state + ), } - return {"task": task, "system": system, "tool_exchanges": tool_exchanges} + 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: list[str]) -> _JevSystemOneResponse | None: + 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]] = { + payload: Final[dict[str, object]] = { # mutable-ok: serialized to JSON by httpx "model": self.jev_model, "state": state, - "questions": { - question_id: {"type": "noul", "instructions": _question_instructions(question_id)} + "questions": { # mutable-ok: serialized to JSON + question_id: { + "type": "noul", + "instructions": _question_instructions(question_id), + } # mutable-ok: serialized to JSON for question_id in question_ids }, } @@ -261,7 +268,7 @@ class TypeSafeGuardrail(CustomGuardrail): 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={ + headers={ # mutable-ok: httpx header contract is a dict "Authorization": f"Bearer {self.typesafe_api_key}", "Content-Type": "application/json", }, @@ -271,21 +278,24 @@ class TypeSafeGuardrail(CustomGuardrail): 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)} + 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", - {"status_code": raw_response.status_code, "body": _safe_response_text(raw_response)}, + { + "status_code": raw_response.status_code, + "body": _safe_response_text(raw_response), + }, # mutable-ok: log detail record ) return None try: @@ -293,7 +303,7 @@ class TypeSafeGuardrail(CustomGuardrail): except (ValueError, httpx.DecodingError, RecursionError): self._handle_failure( "TypeSafe evaluation service returned an unreadable response", - {"body": _safe_response_text(raw_response)}, + {"body": _safe_response_text(raw_response)}, # mutable-ok: log detail record ) return None try: @@ -301,7 +311,7 @@ class TypeSafeGuardrail(CustomGuardrail): except ValidationError: self._handle_failure( "TypeSafe evaluation service returned unexpected response shape", - {"body": _safe_response_text(raw_response)}, + {"body": _safe_response_text(raw_response)}, # mutable-ok: log detail record ) return None @@ -319,17 +329,17 @@ class TypeSafeGuardrail(CustomGuardrail): structured_messages: Final = _as_object_list(inputs.get("structured_messages")) if not structured_messages: return inputs - parsed_messages: Final = [_as_str_object_dict(m) for m in structured_messages] + 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 = [m for m in parsed_messages if m is not None] + 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 = [f"e{ordinal}" for ordinal in range(len(candidates))] + 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() @@ -337,10 +347,12 @@ class TypeSafeGuardrail(CustomGuardrail): 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={ - "error": "TypeSafe evaluation unavailable; request forwarded uncompacted", - "model": self.jev_model, - }, + guardrail_json_response=MappingProxyType( + { + "error": "TypeSafe evaluation unavailable; request forwarded uncompacted", + "model": self.jev_model, + } + ), request_data=request_data, guardrail_status="guardrail_failed_to_respond", guardrail_provider="typesafe", @@ -365,8 +377,10 @@ class TypeSafeGuardrail(CustomGuardrail): 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 + compacted_messages: Final = [ # mutable-ok: structured_messages contract is a list of dicts + {**message, "content": DROPPED_RESULT_TEXT} + if index in dropped_tool_indices + else message # mutable-ok: JSON message row for index, message in enumerate(messages) ] chars_removed: Final = sum( @@ -381,12 +395,14 @@ class TypeSafeGuardrail(CustomGuardrail): 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, - }, + guardrail_json_response=MappingProxyType( + { + "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", @@ -394,7 +410,7 @@ class TypeSafeGuardrail(CustomGuardrail): end_time=end_time, duration=end_time - start_time, ) - return {**inputs, "structured_messages": compacted_messages} # pyright: ignore[reportReturnType] # plain dicts satisfy AllMessageValues at runtime + 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: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py index e2e5fc2fefc..92840c489c3 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py @@ -198,7 +198,7 @@ async def test_request_body_shape_and_truncation(): assert len(exchange["result"]) == 50 assert exchange["result"].startswith(TOOL_OUTPUT_LONG[:10]) assert exchange["result"].endswith(TOOL_OUTPUT_LONG[-11:]) - assert exchange["tool_calls"] == [{"name": "web_search", "arguments": '{"query": "ev"}'}] + assert list(exchange["tool_calls"]) == [{"name": "web_search", "arguments": '{"query": "ev"}'}] @pytest.mark.asyncio From 351afc85196a26b6180e82183c2794b5be3b8958 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 05:10:48 +0000 Subject: [PATCH 6/9] test(guardrails): cover typesafe failure paths and edge shapes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/test_typesafe.py | 116 +++++++++++++++++- 1 file changed, 115 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py index 92840c489c3..a936d725bd3 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py @@ -16,7 +16,7 @@ Tests cover: - response input_type passthrough and initialize_guardrail wiring """ -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, PropertyMock import pytest from fastapi import HTTPException @@ -282,3 +282,117 @@ def test_initialize_guardrail_applies_optional_params_and_registry_keys(): 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 From 87e1c6b3ba19b184735601323eef6d3002f5a0c2 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 05:25:20 +0000 Subject: [PATCH 7/9] fix(guardrails): pin mutable-ok suppressions to constructed literals Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrails/guardrail_hooks/typesafe/__init__.py | 4 ++-- .../guardrails/guardrail_hooks/typesafe/typesafe.py | 12 ++++++------ 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py index c1b05597a6c..abfc60c669a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py @@ -26,9 +26,9 @@ def _coerce_event_hook( if isinstance(mode, Mode): return mode if isinstance(mode, list): - return [ + return [ # mutable-ok: CustomGuardrail event_hook contract wants a list GuardrailEventHooks(item) for item in mode - ] # mutable-ok: CustomGuardrail event_hook contract wants a list + ] return GuardrailEventHooks(mode) diff --git a/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py b/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py index 384bfd2c54b..cb4f21b3261 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py +++ b/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py @@ -257,10 +257,10 @@ class TypeSafeGuardrail(CustomGuardrail): "model": self.jev_model, "state": state, "questions": { # mutable-ok: serialized to JSON - question_id: { + question_id: { # mutable-ok: serialized to JSON "type": "noul", "instructions": _question_instructions(question_id), - } # mutable-ok: serialized to JSON + } for question_id in question_ids }, } @@ -292,10 +292,10 @@ class TypeSafeGuardrail(CustomGuardrail): 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), - }, # mutable-ok: log detail record + }, ) return None try: @@ -378,9 +378,9 @@ class TypeSafeGuardrail(CustomGuardrail): return inputs compacted_messages: Final = [ # mutable-ok: structured_messages contract is a list of dicts - {**message, "content": DROPPED_RESULT_TEXT} + {**message, "content": DROPPED_RESULT_TEXT} # mutable-ok: JSON message row if index in dropped_tool_indices - else message # mutable-ok: JSON message row + else message for index, message in enumerate(messages) ] chars_removed: Final = sum( From 22e6947fef68c2e3af137cdca78b3d8811fbb8a4 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 05:44:25 +0000 Subject: [PATCH 8/9] fix(guardrails): keep typesafe guardrail log payloads JSON-serializable Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/typesafe/typesafe.py | 25 ++++++++----------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py b/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py index cb4f21b3261..9df5c204a77 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py +++ b/litellm/proxy/guardrails/guardrail_hooks/typesafe/typesafe.py @@ -11,7 +11,6 @@ from __future__ import annotations import asyncio import time from collections.abc import Mapping, Sequence -from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Literal import httpx @@ -347,12 +346,10 @@ class TypeSafeGuardrail(CustomGuardrail): 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=MappingProxyType( - { - "error": "TypeSafe evaluation unavailable; request forwarded uncompacted", - "model": self.jev_model, - } - ), + 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", @@ -395,14 +392,12 @@ class TypeSafeGuardrail(CustomGuardrail): chars_removed, ) self.add_standard_logging_guardrail_information_to_request_data( # pyright: ignore[reportUnknownMemberType] # untyped base helper - guardrail_json_response=MappingProxyType( - { - "exchanges_evaluated": len(candidates), - "exchanges_dropped": exchanges_dropped, - "chars_removed": chars_removed, - "model": self.jev_model, - } - ), + 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", From 0765f6d571b3a27de6bc9a7456da977bf62705e4 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 06:05:43 +0000 Subject: [PATCH 9/9] fix(guardrails): keep typesafe registries as dicts so guardrail discovery finds them Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/typesafe/__init__.py | 11 +++---- .../guardrail_hooks/test_typesafe.py | 29 +++++++++++++------ 2 files changed, 26 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py index abfc60c669a..dcea75d3a98 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/typesafe/__init__.py @@ -1,6 +1,5 @@ from __future__ import annotations -from types import MappingProxyType from typing import TYPE_CHECKING, Final from pydantic import BaseModel @@ -66,8 +65,10 @@ def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> return _callback -guardrail_initializer_registry: Final = MappingProxyType( - {SupportedGuardrailIntegrations.TYPESAFE.value: initialize_guardrail} -) +guardrail_initializer_registry: Final = { # mutable-ok: guardrail_registry discovery checks isinstance(registry, dict) + SupportedGuardrailIntegrations.TYPESAFE.value: initialize_guardrail, +} -guardrail_class_registry: Final = MappingProxyType({SupportedGuardrailIntegrations.TYPESAFE.value: TypeSafeGuardrail}) +guardrail_class_registry: Final = { # mutable-ok: guardrail_registry discovery checks isinstance(registry, dict) + SupportedGuardrailIntegrations.TYPESAFE.value: TypeSafeGuardrail, +} diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py index a936d725bd3..2d1db07a1a0 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_typesafe.py @@ -36,7 +36,7 @@ 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_LONG = "Result: EV range comparison. " * 40 TOOL_OUTPUT_SHORT = "short" @@ -141,8 +141,6 @@ async def test_low_noul_exchange_blanked_high_kept_and_input_not_mutated(): 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) @@ -303,9 +301,15 @@ def test_get_config_model_and_ui_name(): 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 + 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 + 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(): @@ -322,7 +326,9 @@ def test_odd_tool_call_shapes_yield_no_entries(): 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"}])) + 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] @@ -363,9 +369,7 @@ async def test_http_status_error_includes_status_and_undecodable_body(): 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) - ) + 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) @@ -396,3 +400,10 @@ def test_optional_params_defaults_and_event_hook_coercion(): 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