Wire run_shield_moderation_v2 into Responses API endpoint

Jazzcort · Jazzcort · commit 0237198c5d57 · 2026-07-24T16:08:49.000-04:00
Add run_shield_moderation_v2 and build_shield to utils/shields.py to
run shield moderation through AbstractSafetyCapability instances instead
of the Llama Stack client. Update the responses endpoint to call the new
function with shield configs directly. Include unit tests covering pass,
block, filtering, and error handling paths.
diff --git a/src/app/endpoints/responses.py b/src/app/endpoints/responses.py
@@ -105,7 +105,7 @@
     select_model_for_responses,
 )
 from utils.rh_identity import get_rh_identity_context
-from utils.shields import run_shield_moderation
+from utils.shields import run_shield_moderation_v2
 from utils.suid import (
     normalize_conversation_id,
 )
@@ -424,11 +424,11 @@ async def responses_endpoint_handler(
     attachments_text = extract_attachments_text(original_request.input)
 
     endpoint_path = ENDPOINT_PATH_RESPONSES
-    moderation_result = await run_shield_moderation(
-        client,
+
+    moderation_result = await run_shield_moderation_v2(
         input_text + "\n\n" + attachments_text,
-        endpoint_path,
-        original_request.shield_ids,
+        configuration.configuration.shields,
+        responses_request.shield_ids,
     )
 
     filter_server_tools = (
diff --git a/src/utils/shields.py b/src/utils/shields.py
@@ -10,17 +10,31 @@
 from ogx_client import (
     APIStatusError as LLSApiStatusError,
 )
+from pydantic_ai import ModelAPIError, ModelHTTPError, UnexpectedModelBehavior
 
 from configuration import AppConfig
+from log import get_logger
 from models.api.requests import QueryRequest
 from models.api.responses.error import (
     InternalServerErrorResponse,
     NotFoundResponse,
     ServiceUnavailableResponse,
     UnprocessableEntityResponse,
 )
-from models.common import ShieldModerationPassed, ShieldModerationResult
-from models.config import ShieldConfiguration
+from models.common.moderation import (
+    ShieldModerationPassed,
+    ShieldModerationResult,
+)
+from models.config import QuestionValidityConfig, RedactionConfig, ShieldConfiguration
+from pydantic_ai_lightspeed.capabilities.base import AbstractSafetyCapability
+from pydantic_ai_lightspeed.capabilities.question_validity._capability import (
+    QuestionValidity,
+)
+from pydantic_ai_lightspeed.capabilities.redaction._capability import (
+    PiiRedactionCapability,
+)
+
+logger = get_logger(__name__)
 
 
 def validate_shield_ids_override(
@@ -59,6 +73,74 @@ def validate_shield_ids_override(
         raise HTTPException(**response.model_dump())
 
 
+async def run_shield_moderation_v2(
+    input_text: str,
+    shield_configs: list[ShieldConfiguration],
+    selected_shield_ids: Optional[list[str]] = None,
+) -> ShieldModerationResult:
+    """Run v2 shield moderation on input text.
+
+    Iterates through configured shields and runs moderation checks.
+
+    Parameters:
+        input_text: The text to moderate.
+        shield_configs: List of shield configurations to evaluate.
+        selected_shield_ids: Optional list of shield names to filter by.
+
+    Returns:
+        Result indicating if content was blocked or passed.
+
+    Raises:
+        HTTPException: 503 if the shield model is unreachable or returns an
+            HTTP error, 500 if the model returns an unexpected response.
+    """
+    selected_shield_configs = (
+        [c for c in shield_configs if c.name in selected_shield_ids]
+        if selected_shield_ids is not None
+        else shield_configs
+    )
+
+    try:
+        for shield_config in selected_shield_configs:
+            shield = build_shield(shield_config)
+            shield_result = await shield.run(input_text)
+
+            if shield_result.decision == "blocked":
+                return shield_result
+    except (ModelHTTPError, ModelAPIError) as e:
+        logger.error("Shield moderation model request failed: %s", e)
+        error_response = ServiceUnavailableResponse(
+            backend_name="shield model",
+            cause=str(e),
+        )
+        raise HTTPException(**error_response.model_dump()) from e
+    except UnexpectedModelBehavior as e:
+        logger.error("Shield moderation received unexpected model response: %s", e)
+        error_response = InternalServerErrorResponse(
+            response="Shield moderation failed",
+            cause=str(e),
+        )
+        raise HTTPException(**error_response.model_dump()) from e
+
+    return ShieldModerationPassed()
+
+
+def build_shield(shield_config: ShieldConfiguration) -> AbstractSafetyCapability:
+    """Build a safety capability instance from a shield configuration.
+
+    Parameters:
+        shield_config: The shield configuration to build from.
+
+    Returns:
+        The constructed safety capability.
+    """
+    match shield_config.config:
+        case QuestionValidityConfig():
+            return QuestionValidity(shield_config.config)
+        case RedactionConfig():
+            return PiiRedactionCapability(shield_config.config)
+
+
 async def run_shield_moderation(
     _client: AsyncOgxClient,
     _input_text: str,
diff --git a/tests/unit/utils/test_shields.py b/tests/unit/utils/test_shields.py
@@ -2,17 +2,24 @@
 
 import pytest
 from fastapi import HTTPException, status
+from pydantic_ai import ModelAPIError, ModelHTTPError, UnexpectedModelBehavior
 from pytest_mock import MockerFixture
 
-from models.config import QuestionValidityConfig, QuestionValidityShieldConfiguration
+from models.common.moderation import ShieldModerationBlocked, ShieldModerationPassed
+from models.config import (
+    QuestionValidityConfig,
+    QuestionValidityShieldConfiguration,
+    ShieldConfiguration,
+)
 from utils.shields import (
     append_turn_to_conversation,
     get_shields_for_request,
+    run_shield_moderation_v2,
     validate_shield_ids_override,
 )
 
 
-def _shield(name: str) -> QuestionValidityShieldConfiguration:
+def _shield_config(name: str) -> QuestionValidityShieldConfiguration:
     """Build a minimal question-validity shield configuration for tests."""
     return QuestionValidityShieldConfiguration(
         name=name,
@@ -133,30 +140,174 @@ def test_raises_422_when_empty_list_shield_ids_and_override_disabled(
         assert exc_info.value.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
 
 
+class TestRunShieldModerationV2:
+    """Tests for run_shield_moderation_v2 function."""
+
+    @pytest.mark.asyncio
+    async def test_returns_passed_when_no_shields(self) -> None:
+        """Return ShieldModerationPassed when shield list is empty."""
+        result = await run_shield_moderation_v2("test input", [])
+        assert isinstance(result, ShieldModerationPassed)
+
+    @pytest.mark.asyncio
+    async def test_returns_passed_when_all_shields_pass(
+        self, mocker: MockerFixture
+    ) -> None:
+        """Return ShieldModerationPassed when every shield passes."""
+        mock_shield = mocker.Mock()
+        mock_shield.run = mocker.AsyncMock(return_value=ShieldModerationPassed())
+        mocker.patch("utils.shields.build_shield", return_value=mock_shield)
+
+        shields: list[ShieldConfiguration] = [
+            _shield_config("s1"),
+            _shield_config("s2"),
+        ]
+        result = await run_shield_moderation_v2("test input", shields)
+
+        assert isinstance(result, ShieldModerationPassed)
+        assert mock_shield.run.call_count == 2
+
+    @pytest.mark.asyncio
+    async def test_returns_blocked_on_first_block(self, mocker: MockerFixture) -> None:
+        """Return blocked result from first shield that blocks."""
+        blocked = ShieldModerationBlocked(message="rejected", moderation_id="modr-123")
+        mock_shield = mocker.Mock()
+        mock_shield.run = mocker.AsyncMock(return_value=blocked)
+        mocker.patch("utils.shields.build_shield", return_value=mock_shield)
+
+        shields: list[ShieldConfiguration] = [
+            _shield_config("s1"),
+            _shield_config("s2"),
+        ]
+        result = await run_shield_moderation_v2("test input", shields)
+
+        assert isinstance(result, ShieldModerationBlocked)
+        assert result.message == "rejected"
+        mock_shield.run.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_filters_by_selected_shield_ids(self, mocker: MockerFixture) -> None:
+        """Only run shields matching the selected IDs."""
+        mock_shield = mocker.Mock()
+        mock_shield.run = mocker.AsyncMock(return_value=ShieldModerationPassed())
+        mocker.patch("utils.shields.build_shield", return_value=mock_shield)
+
+        shields: list[ShieldConfiguration] = [
+            _shield_config("s1"),
+            _shield_config("s2"),
+            _shield_config("s3"),
+        ]
+        result = await run_shield_moderation_v2(
+            "test input", shields, selected_shield_ids=["s2"]
+        )
+
+        assert isinstance(result, ShieldModerationPassed)
+        mock_shield.run.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_shields_stops_on_first_block(self, mocker: MockerFixture) -> None:
+        """Stop at the first blocking shield."""
+        blocked = ShieldModerationBlocked(message="rejected", moderation_id="modr-789")
+        mock_qv_shield = mocker.Mock()
+        mock_qv_shield.run = mocker.AsyncMock(return_value=blocked)
+
+        mock_redact_shield = mocker.Mock()
+        mock_redact_shield.run = mocker.AsyncMock(return_value=ShieldModerationPassed())
+
+        mocker.patch(
+            "utils.shields.build_shield",
+            side_effect=[mock_qv_shield, mock_redact_shield],
+        )
+
+        shields: list[ShieldConfiguration] = [
+            _shield_config("s-1"),
+            _shield_config("s-2"),
+        ]
+        result = await run_shield_moderation_v2("test input", shields)
+
+        assert isinstance(result, ShieldModerationBlocked)
+        mock_qv_shield.run.assert_called_once()
+        mock_redact_shield.run.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_raises_503_on_model_http_error(self, mocker: MockerFixture) -> None:
+        """Raise HTTP 503 when the shield model returns an HTTP error."""
+        mock_shield = mocker.Mock()
+        mock_shield.run = mocker.AsyncMock(
+            side_effect=ModelHTTPError(
+                status_code=500, model_name="test-model", body=None
+            )
+        )
+        mocker.patch("utils.shields.build_shield", return_value=mock_shield)
+
+        with pytest.raises(HTTPException) as exc_info:
+            await run_shield_moderation_v2("test input", [_shield_config("s1")])
+
+        assert exc_info.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE
+
+    @pytest.mark.asyncio
+    async def test_raises_503_on_model_api_error(self, mocker: MockerFixture) -> None:
+        """Raise HTTP 503 when the shield model is unreachable."""
+        mock_shield = mocker.Mock()
+        mock_shield.run = mocker.AsyncMock(
+            side_effect=ModelAPIError(
+                model_name="test-model", message="Connection refused"
+            )
+        )
+        mocker.patch("utils.shields.build_shield", return_value=mock_shield)
+
+        with pytest.raises(HTTPException) as exc_info:
+            await run_shield_moderation_v2("test input", [_shield_config("s1")])
+
+        assert exc_info.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE
+
+    @pytest.mark.asyncio
+    async def test_raises_500_on_unexpected_model_behavior(
+        self, mocker: MockerFixture
+    ) -> None:
+        """Raise HTTP 500 when the shield model returns an unexpected response."""
+        mock_shield = mocker.Mock()
+        mock_shield.run = mocker.AsyncMock(
+            side_effect=UnexpectedModelBehavior("bad response format")
+        )
+        mocker.patch("utils.shields.build_shield", return_value=mock_shield)
+
+        with pytest.raises(HTTPException) as exc_info:
+            await run_shield_moderation_v2("test input", [_shield_config("s1")])
+
+        assert exc_info.value.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR
+
+
 class TestGetShieldsForRequest:
     """Tests for get_shields_for_request function."""
 
     def test_returns_all_shields_when_shield_ids_none(self) -> None:
         """Return all configured shields when shield_ids is None."""
-        shields = [_shield("shield-1"), _shield("shield-2")]
+        shields = [
+            _shield_config("shield-1"),
+            _shield_config("shield-2"),
+        ]
 
         result = get_shields_for_request(shields, shield_ids=None)
 
         assert result == shields
 
     def test_returns_empty_list_when_shield_ids_empty(self) -> None:
         """Return no shields when an empty shield_ids list is provided."""
-        shields = [_shield("shield-1"), _shield("shield-2")]
+        shields = [
+            _shield_config("shield-1"),
+            _shield_config("shield-2"),
+        ]
 
         result = get_shields_for_request(shields, shield_ids=[])
 
         assert result == []
 
     def test_filters_to_requested_shields_when_all_exist(self) -> None:
         """Return only shields whose names appear in shield_ids."""
-        shield1 = _shield("shield-1")
-        shield2 = _shield("shield-2")
-        shield3 = _shield("shield-3")
+        shield1 = _shield_config("shield-1")
+        shield2 = _shield_config("shield-2")
+        shield3 = _shield_config("shield-3")
 
         result = get_shields_for_request(
             [shield1, shield2, shield3], shield_ids=["shield-1", "shield-3"]
@@ -168,7 +319,8 @@ def test_raises_404_when_requested_shield_not_configured(self) -> None:
         """Raise 404 when a requested shield name is not configured."""
         with pytest.raises(HTTPException) as exc_info:
             get_shields_for_request(
-                [_shield("shield-1")], shield_ids=["shield-1", "missing-shield"]
+                [_shield_config("shield-1")],
+                shield_ids=["shield-1", "missing-shield"],
             )
 
         assert exc_info.value.status_code == status.HTTP_404_NOT_FOUND
