Responses API: avoid shields code duplication

Author: luis5tb

src/app/endpoints/query_v2.py

@@ -34,6 +34,7 @@
     get_topic_summary_system_prompt,
 )
 from utils.mcp_headers import mcp_headers_dependency
+from utils.shields import detect_shield_violations, get_available_shields
 from utils.token_counter import TokenCounter
 from utils.types import TurnSummary, ToolCallSummary
 
@@ -343,11 +344,7 @@ async def retrieve_response(  # pylint: disable=too-many-locals,too-many-branche
         and token usage information.
     """
     # List available shields for Responses API
-    available_shields = [shield.identifier for shield in await client.shields.list()]
-    if not available_shields:
-        logger.info("No available shields. Disabling safety")
-    else:
-        logger.info("Available shields: %s", available_shields)
+    available_shields = await get_available_shields(client)
 
     # use system prompt from request or default one
     system_prompt = get_system_prompt(query_request, configuration)
@@ -414,14 +411,8 @@ async def retrieve_response(  # pylint: disable=too-many-locals,too-many-branche
         if tool_summary:
             tool_calls.append(tool_summary)
 
-        # Check for shield violations
-        item_type = getattr(output_item, "type", None)
-        if item_type == "message":
-            refusal = getattr(output_item, "refusal", None)
-            if refusal:
-                # Metric for LLM validation errors (shield violations)
-                metrics.llm_calls_validation_errors_total.inc()
-                logger.warning("Shield violation detected: %s", refusal)
+    # Check for shield violations across all output items
+    detect_shield_violations(response.output)
 
     logger.info(
         "Response processing complete - Tool calls: %d, Response length: %d chars",

src/app/endpoints/streaming_query_v2.py

@@ -32,7 +32,6 @@
 from authorization.middleware import authorize
 from configuration import configuration
 from constants import MEDIA_TYPE_JSON
-import metrics
 from models.config import Action
 from models.context import ResponseGeneratorContext
 from models.requests import QueryRequest
@@ -42,6 +41,7 @@
     get_system_prompt,
 )
 from utils.mcp_headers import mcp_headers_dependency
+from utils.shields import detect_shield_violations, get_available_shields
 from utils.token_counter import TokenCounter
 from utils.transcripts import store_transcript
 from utils.types import TurnSummary, ToolCallSummary
@@ -247,14 +247,9 @@ async def response_generator(  # pylint: disable=too-many-branches,too-many-stat
 
                 # Check for shield violations in the completed response
                 if latest_response_object:
-                    for output_item in getattr(latest_response_object, "output", []):
-                        item_type = getattr(output_item, "type", None)
-                        if item_type == "message":
-                            refusal = getattr(output_item, "refusal", None)
-                            if refusal:
-                                # Metric for LLM validation errors (shield violations)
-                                metrics.llm_calls_validation_errors_total.inc()
-                                logger.warning("Shield violation detected: %s", refusal)
+                    detect_shield_violations(
+                        getattr(latest_response_object, "output", [])
+                    )
 
                 if not emitted_turn_complete:
                     final_message = summary.llm_response or "".join(text_parts)
@@ -379,11 +374,7 @@ async def retrieve_response(
         and the conversation ID.
     """
     # List available shields for Responses API
-    available_shields = [shield.identifier for shield in await client.shields.list()]
-    if not available_shields:
-        logger.info("No available shields. Disabling safety")
-    else:
-        logger.info("Available shields: %s", available_shields)
+    available_shields = await get_available_shields(client)
 
     # use system prompt from request or default one
     system_prompt = get_system_prompt(query_request, configuration)

src/utils/shields.py

@@ -0,0 +1,54 @@
+"""Utility functions for working with Llama Stack shields."""
+
+import logging
+from typing import Any
+
+from llama_stack_client import AsyncLlamaStackClient
+
+import metrics
+
+logger = logging.getLogger(__name__)
+
+
+async def get_available_shields(client: AsyncLlamaStackClient) -> list[str]:
+    """
+    Discover and return available shield identifiers.
+
+    Args:
+        client: The Llama Stack client to query for available shields.
+
+    Returns:
+        List of shield identifiers that are available.
+    """
+    available_shields = [shield.identifier for shield in await client.shields.list()]
+    if not available_shields:
+        logger.info("No available shields. Disabling safety")
+    else:
+        logger.info("Available shields: %s", available_shields)
+    return available_shields
+
+
+def detect_shield_violations(output_items: list[Any]) -> bool:
+    """
+    Check output items for shield violations and update metrics.
+
+    Iterates through output items looking for message items with refusal
+    attributes. If a refusal is found, increments the validation error
+    metric and logs a warning.
+
+    Args:
+        output_items: List of output items from the LLM response to check.
+
+    Returns:
+        True if a shield violation was detected, False otherwise.
+    """
+    for output_item in output_items:
+        item_type = getattr(output_item, "type", None)
+        if item_type == "message":
+            refusal = getattr(output_item, "refusal", None)
+            if refusal:
+                # Metric for LLM validation errors (shield violations)
+                metrics.llm_calls_validation_errors_total.inc()
+                logger.warning("Shield violation detected: %s", refusal)
+                return True
+    return False