diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index dd809e4882b..fdefbf904dd 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -1173,7 +1173,8 @@ def _build_field_dict( # Check for custom UI type override field_json_schema_extra = getattr(field, "json_schema_extra", {}) if field_json_schema_extra and "ui_type" in field_json_schema_extra: - field_type = field_json_schema_extra["ui_type"].value + ui_type = field_json_schema_extra["ui_type"] + field_type = ui_type.value if hasattr(ui_type, "value") else ui_type elif field_json_schema_extra and "type" in field_json_schema_extra: field_type = field_json_schema_extra["type"] diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 0ac3637b380..59bca884db7 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -1103,4 +1103,47 @@ async def test_get_guardrail_info_endpoint_db_guardrail(mocker): assert isinstance(result, GuardrailInfoResponse) assert result.guardrail_id == "test-db-guardrail" assert result.guardrail_name == "Test DB Guardrail" - assert result.guardrail_definition_location == "db" \ No newline at end of file + assert result.guardrail_definition_location == "db" + + +class TestBuildFieldDict: + """Test _build_field_dict handles both enum and string ui_type values.""" + + def test_build_field_dict_with_string_ui_type(self): + """Test that _build_field_dict works when ui_type is a plain string (e.g. BlockCodeExecutionGuardrailConfigModel).""" + from unittest.mock import MagicMock + + from litellm.proxy.guardrails.guardrail_endpoints import _build_field_dict + + field = MagicMock() + field.json_schema_extra = {"ui_type": "multiselect", "options": ["python", "javascript"]} + + result = _build_field_dict( + field=field, + field_annotation=str, + description="Test field", + required=False, + ) + + assert result["type"] == "multiselect" + assert result["description"] == "Test field" + + def test_build_field_dict_with_enum_ui_type(self): + """Test that _build_field_dict works when ui_type is a GuardrailParamUITypes enum.""" + from unittest.mock import MagicMock + + from litellm.proxy.guardrails.guardrail_endpoints import _build_field_dict + from litellm.types.guardrails import GuardrailParamUITypes + + field = MagicMock() + field.json_schema_extra = {"ui_type": GuardrailParamUITypes.BOOL} + + result = _build_field_dict( + field=field, + field_annotation=bool, + description="Test bool field", + required=True, + ) + + assert result["type"] == "bool" + assert result["required"] is True \ No newline at end of file