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>
This commit is contained in:
yassin 2026-09-18 06:05:43 +00:00
parent 22e6947fef
commit 0765f6d571
2 changed files with 26 additions and 14 deletions

View file

@ -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,
}

View file

@ -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