Wire run_shield_moderation_v2 into Responses API endpoint

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.

Author: Jazzcort

src/app/endpoints/a2a.py

@@ -60,7 +60,7 @@
 from log import get_logger
 from models.api.requests import QueryRequest
 from models.config import Action
-from utils.agents.query import map_agent_inference_error
+from utils.agents.error_handler import map_agent_inference_error
 from utils.conversation_compaction import apply_compaction_blocking
 from utils.mcp_headers import McpHeaders, mcp_headers_dependency
 from utils.pydantic_ai_helpers import build_agent

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 = (

src/utils/agents/error_handler.py

@@ -0,0 +1,103 @@
+"""Error mapping for agent inference failures to structured API error responses."""
+
+from typing import TypeAlias
+
+from ogx_client import APIConnectionError, APIStatusError
+from pydantic_ai.exceptions import (
+    AgentRunError,
+    ContentFilterError,
+    IncompleteToolCall,
+    ModelAPIError,
+    ModelHTTPError,
+    UnexpectedModelBehavior,
+    UsageLimitExceeded,
+)
+
+from log import get_logger
+from models.api.responses.error import (
+    AbstractErrorResponse,
+    InternalServerErrorResponse,
+    PromptTooLongResponse,
+    QuotaExceededResponse,
+    ServiceUnavailableResponse,
+)
+from utils.query import (
+    handle_known_apistatus_errors,
+    is_context_length_error,
+)
+
+AgentInferenceError: TypeAlias = (
+    AgentRunError | APIStatusError | APIConnectionError | RuntimeError
+)
+
+logger = get_logger(__name__)
+
+
+def map_agent_inference_error(
+    exc: AgentInferenceError,
+    model_id: str,
+) -> AbstractErrorResponse:
+    """Map agent run failures from pydantic-ai or Llama Stack to an LCS error response.
+
+    Args:
+        exc: Agent, HTTP status, connection, or context-length runtime error.
+        model_id: Model identifier in provider/model format.
+
+    Returns:
+        Structured error response for HTTP or SSE error events.
+
+    Raises:
+        RuntimeError: Re-raised when ``exc`` is a non-agent ``RuntimeError`` that is
+            not a recognized context-length failure.
+    """
+    match exc:
+        case AgentRunError() as agent_exc:
+            return map_pydantic_agent_run_error(agent_exc, model_id)
+        case APIStatusError() as status_exc:
+            return handle_known_apistatus_errors(status_exc, model_id)
+        case APIConnectionError() as connection_exc:
+            return ServiceUnavailableResponse(
+                backend_name="OGX",
+                cause=str(connection_exc),
+            )
+        case RuntimeError() as runtime_exc if is_context_length_error(str(runtime_exc)):
+            return PromptTooLongResponse(model=model_id)
+        case _:
+            return InternalServerErrorResponse.generic()
+
+
+def map_pydantic_agent_run_error(
+    exc: AgentRunError, model_id: str
+) -> AbstractErrorResponse:
+    """Map pydantic-ai ``AgentRunError`` subclasses to LCS error responses.
+
+    Args:
+        exc: Agent exception to map.
+        model_id: Model identifier in provider/model format.
+
+    Returns:
+        Structured error response for HTTP or SSE error events.
+    """
+    match exc:
+        case ContentFilterError() as filter_exc:
+            return InternalServerErrorResponse.query_failed(str(filter_exc))
+        case IncompleteToolCall():
+            return PromptTooLongResponse(model=model_id)
+        case UnexpectedModelBehavior():
+            logger.error("Unexpected model behavior: %s", exc, exc_info=True)
+            return InternalServerErrorResponse.generic()
+        case UsageLimitExceeded():
+            return QuotaExceededResponse.model(model_id)
+        case ModelHTTPError() as http_exc if is_context_length_error(str(http_exc)):
+            return PromptTooLongResponse(model=model_id)
+        case ModelHTTPError(status_code=429):
+            return QuotaExceededResponse.model(model_id)
+        case ModelHTTPError():
+            return InternalServerErrorResponse.generic()
+        case ModelAPIError() as api_exc:
+            return ServiceUnavailableResponse(
+                backend_name="OGX",
+                cause=str(api_exc),
+            )
+        case _:
+            return InternalServerErrorResponse.query_failed(str(exc))

src/utils/agents/query.py

@@ -9,12 +9,6 @@
 from ogx_client import APIConnectionError, APIStatusError, AsyncOgxClient
 from pydantic_ai.exceptions import (
     AgentRunError,
-    ContentFilterError,
-    IncompleteToolCall,
-    ModelAPIError,
-    ModelHTTPError,
-    UnexpectedModelBehavior,
-    UsageLimitExceeded,
 )
 from pydantic_ai.messages import ModelRequest, ModelResponse, ToolReturnPart
 from pydantic_ai.run import AgentRunResult
@@ -27,14 +21,13 @@
     AbstractErrorResponse,
     InternalServerErrorResponse,
     PromptTooLongResponse,
-    QuotaExceededResponse,
-    ServiceUnavailableResponse,
 )
 from models.common.agents import AgentTurnAccumulator
 from models.common.moderation import ShieldModerationResult
 from models.common.responses.responses_api_params import ResponsesApiParams
 from models.common.responses.types import ResponseInput
 from models.common.turn_summary import TurnSummary
+from utils.agents.error_handler import map_agent_inference_error
 from utils.agents.tool_processor import (
     process_function_tool_call,
     process_function_tool_result,
@@ -45,8 +38,6 @@
 from utils.pydantic_ai_helpers import build_agent
 from utils.query import (
     extract_provider_and_model_from_model_id,
-    handle_known_apistatus_errors,
-    is_context_length_error,
 )
 from utils.responses import extract_vector_store_ids_from_tools
 from utils.token_counter import TokenCounter
@@ -68,76 +59,6 @@ class AgentFinishReason(str, Enum):
     ERROR = "error"
 
 
-def map_agent_inference_error(
-    exc: AgentInferenceError,
-    model_id: str,
-) -> AbstractErrorResponse:
-    """Map agent run failures from pydantic-ai or Llama Stack to an LCS error response.
-
-    Args:
-        exc: Agent, HTTP status, connection, or context-length runtime error.
-        model_id: Model identifier in provider/model format.
-
-    Returns:
-        Structured error response for HTTP or SSE error events.
-
-    Raises:
-        RuntimeError: Re-raised when ``exc`` is a non-agent ``RuntimeError`` that is
-            not a recognized context-length failure.
-    """
-    match exc:
-        case AgentRunError() as agent_exc:
-            return map_pydantic_agent_run_error(agent_exc, model_id)
-        case APIStatusError() as status_exc:
-            return handle_known_apistatus_errors(status_exc, model_id)
-        case APIConnectionError() as connection_exc:
-            return ServiceUnavailableResponse(
-                backend_name="OGX",
-                cause=str(connection_exc),
-            )
-        case RuntimeError() as runtime_exc if is_context_length_error(str(runtime_exc)):
-            return PromptTooLongResponse(model=model_id)
-        case _:
-            return InternalServerErrorResponse.generic()
-
-
-def map_pydantic_agent_run_error(
-    exc: AgentRunError, model_id: str
-) -> AbstractErrorResponse:
-    """Map pydantic-ai ``AgentRunError`` subclasses to LCS error responses.
-
-    Args:
-        exc: Agent exception to map.
-        model_id: Model identifier in provider/model format.
-
-    Returns:
-        Structured error response for HTTP or SSE error events.
-    """
-    match exc:
-        case ContentFilterError() as filter_exc:
-            return InternalServerErrorResponse.query_failed(str(filter_exc))
-        case IncompleteToolCall():
-            return PromptTooLongResponse(model=model_id)
-        case UnexpectedModelBehavior():
-            logger.error("Unexpected model behavior: %s", exc, exc_info=True)
-            return InternalServerErrorResponse.generic()
-        case UsageLimitExceeded():
-            return QuotaExceededResponse.model(model_id)
-        case ModelHTTPError() as http_exc if is_context_length_error(str(http_exc)):
-            return PromptTooLongResponse(model=model_id)
-        case ModelHTTPError(status_code=429):
-            return QuotaExceededResponse.model(model_id)
-        case ModelHTTPError():
-            return InternalServerErrorResponse.generic()
-        case ModelAPIError() as api_exc:
-            return ServiceUnavailableResponse(
-                backend_name="OGX",
-                cause=str(api_exc),
-            )
-        case _:
-            return InternalServerErrorResponse.query_failed(str(exc))
-
-
 def get_agent_finish_reason(response: ModelResponse) -> AgentFinishReason:
     """Get the finish reason from a completed agent model response.
 

src/utils/agents/streaming.py

@@ -43,12 +43,12 @@
 from models.common.responses.contexts import ResponseGeneratorContext
 from models.common.responses.responses_api_params import ResponsesApiParams
 from models.common.turn_summary import TurnSummary
+from utils.agents.error_handler import map_agent_inference_error
 from utils.agents.query import (
     AgentFinishReason,
     extract_agent_token_usage,
     get_agent_finish_reason,
     get_finish_reason_error,
-    map_agent_inference_error,
 )
 from utils.agents.tool_processor import (
     process_function_tool_call,

src/utils/shields.py

@@ -3,24 +3,33 @@
 from typing import Optional
 
 from fastapi import HTTPException
-from ogx_client import (
-    APIConnectionError,
-    AsyncOgxClient,
-)
-from ogx_client import (
-    APIStatusError as LLSApiStatusError,
-)
+from ogx_client import APIConnectionError, APIStatusError, AsyncOgxClient
+from pydantic_ai.exceptions import AgentRunError
 
 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,
+)
+from utils.agents.error_handler import map_agent_inference_error
+
+logger = get_logger(__name__)
 
 
 def validate_shield_ids_override(
@@ -59,6 +68,63 @@ 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.
+    """
+    selected_shield_configs = get_shields_for_request(
+        shield_configs, selected_shield_ids
+    )
+
+    for shield_config in selected_shield_configs:
+        shield = build_shield(shield_config)
+
+        try:
+            shield_result = await shield.run(input_text)
+        # APIConnectionError and APIStatusError from ogx should not be raised from model_request,
+        # because they will be caught inside AsyncOpenAI and transferred into openai's
+        # APIConnectionError. The openai's exceptions will further transferred into ModelHTTPError
+        # or ModelAPIError by _map_api_errors in OpenAIResponseModel.
+        except (AgentRunError, RuntimeError) as exc:
+            model_id = getattr(shield_config.config, "model_id", "unknown-shield-model")
+            response = map_agent_inference_error(exc, model_id)
+            raise HTTPException(**response.model_dump()) from exc
+
+        if shield_result.decision == "blocked":
+            return shield_result
+
+    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,
@@ -124,7 +190,7 @@ async def append_turn_to_conversation(
             cause=str(e),
         )
         raise HTTPException(**error_response.model_dump()) from e
-    except LLSApiStatusError as e:
+    except APIStatusError as e:
         error_response = InternalServerErrorResponse.generic()
         raise HTTPException(**error_response.model_dump()) from e
 

tests/integration/endpoints/test_responses_integration.py

@@ -173,7 +173,7 @@ def _configure_shield_blocked(
         moderation_id=moderation_id,
     )
     mocker.patch(
-        "app.endpoints.responses.run_shield_moderation",
+        "app.endpoints.responses.run_shield_moderation_v2",
         return_value=blocked,
     )