From dffb6a38d95f5e2fa2c259f06f50243c46ca2ab2 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 04:35:58 +0000 Subject: [PATCH] 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