Add shields support to the responses API implementation

It includes both streaming and not streaming support, by
leveraging the refusal field on the response

Author: luis5tb

src/app/endpoints/query_v2.py

@@ -317,7 +317,7 @@ async def retrieve_response(  # pylint: disable=too-many-locals,too-many-branche
     given query, handling shield configuration, tool usage, and
     attachment validation.
 
-    This function configures system prompts and toolgroups
+    This function configures system prompts, shields, and toolgroups
     (including RAG and MCP integration) as needed based on
     the query request and system configuration. It
     validates attachments, manages conversation and session
@@ -337,8 +337,12 @@ async def retrieve_response(  # pylint: disable=too-many-locals,too-many-branche
         and the conversation ID, the list of parsed referenced documents,
         and token usage information.
     """
-    # TODO(ltomasbo): implement shields support once available in Responses API
-    logger.info("Shields are not yet supported in Responses API. Disabling safety")
+    # 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)
 
     # use system prompt from request or default one
     system_prompt = get_system_prompt(query_request, configuration)
@@ -376,6 +380,10 @@ async def retrieve_response(  # pylint: disable=too-many-locals,too-many-branche
     if query_request.conversation_id:
         create_kwargs["previous_response_id"] = query_request.conversation_id
 
+    # Add shields to extra_body if available
+    if available_shields:
+        create_kwargs["extra_body"] = {"guardrails": available_shields}
+
     response = await client.responses.create(**create_kwargs)
     response = cast(OpenAIResponseObject, response)
 
@@ -401,6 +409,15 @@ 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)
+
     logger.info(
         "Response processing complete - Tool calls: %d, Response length: %d chars",
         len(tool_calls),

src/app/endpoints/streaming_query_v2.py

@@ -34,6 +34,7 @@
 from authorization.middleware import authorize
 from configuration import configuration
 from constants import MEDIA_TYPE_JSON
+import metrics
 from models.cache_entry import CacheEntry
 from models.config import Action
 from models.context import ResponseGeneratorContext
@@ -247,6 +248,18 @@ async def response_generator(  # pylint: disable=too-many-branches,too-many-stat
             elif event_type == "response.completed":
                 # Capture the response object for token usage extraction
                 latest_response_object = getattr(chunk, "response", None)
+
+                # 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)
+
                 if not emitted_turn_complete:
                     final_message = summary.llm_response or "".join(text_parts)
                     yield format_stream_data(
@@ -394,11 +407,11 @@ async def retrieve_response(
     Asynchronously retrieves a streaming response and conversation
     ID from the Llama Stack agent for a given user query.
 
-    This function configures input/output shields, system prompt,
-    and tool usage based on the request and environment. It
-    prepares the agent with appropriate headers and toolgroups,
-    validates attachments if present, and initiates a streaming
-    turn with the user's query and any provided documents.
+    This function configures shields, system prompt, and tool usage
+    based on the request and environment. It prepares the agent with
+    appropriate headers and toolgroups, validates attachments if
+    present, and initiates a streaming turn with the user's query
+    and any provided documents.
 
     Parameters:
         model_id (str): Identifier of the model to use for the query.
@@ -411,7 +424,12 @@ async def retrieve_response(
         tuple: A tuple containing the streaming response object
         and the conversation ID.
     """
-    logger.info("Shields are not yet supported in Responses API.")
+    # 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)
 
     # use system prompt from request or default one
     system_prompt = get_system_prompt(query_request, configuration)
@@ -448,6 +466,10 @@ async def retrieve_response(
     if query_request.conversation_id:
         create_params["previous_response_id"] = query_request.conversation_id
 
+    # Add shields to extra_body if available
+    if available_shields:
+        create_params["extra_body"] = {"guardrails": available_shields}
+
     response = await client.responses.create(**create_params)
     response_stream = cast(AsyncIterator[OpenAIResponseObjectStream], response)
 

tests/unit/app/endpoints/test_query_v2.py

@@ -109,6 +109,8 @@ async def test_retrieve_response_no_tools_bypasses_tools(mocker):
     mock_vector_stores = mocker.Mock()
     mock_vector_stores.data = []
     mock_client.vector_stores.list = mocker.AsyncMock(return_value=mock_vector_stores)
+    # Mock shields.list
+    mock_client.shields.list = mocker.AsyncMock(return_value=[])
 
     # Ensure system prompt resolution does not require real config
     mocker.patch("app.endpoints.query_v2.get_system_prompt", return_value="PROMPT")
@@ -143,6 +145,8 @@ async def test_retrieve_response_builds_rag_and_mcp_tools(mocker):
     mock_vector_stores = mocker.Mock()
     mock_vector_stores.data = [mocker.Mock(id="dbA")]
     mock_client.vector_stores.list = mocker.AsyncMock(return_value=mock_vector_stores)
+    # Mock shields.list
+    mock_client.shields.list = mocker.AsyncMock(return_value=[])
 
     mocker.patch("app.endpoints.query_v2.get_system_prompt", return_value="PROMPT")
     mock_cfg = mocker.Mock()
@@ -207,6 +211,8 @@ async def test_retrieve_response_parses_output_and_tool_calls(mocker):
     mock_vector_stores = mocker.Mock()
     mock_vector_stores.data = []
     mock_client.vector_stores.list = mocker.AsyncMock(return_value=mock_vector_stores)
+    # Mock shields.list
+    mock_client.shields.list = mocker.AsyncMock(return_value=[])
 
     mocker.patch("app.endpoints.query_v2.get_system_prompt", return_value="PROMPT")
     mocker.patch("app.endpoints.query_v2.configuration", mocker.Mock(mcp_servers=[]))
@@ -252,6 +258,8 @@ async def test_retrieve_response_with_usage_info(mocker):
     mock_vector_stores = mocker.Mock()
     mock_vector_stores.data = []
     mock_client.vector_stores.list = mocker.AsyncMock(return_value=mock_vector_stores)
+    # Mock shields.list
+    mock_client.shields.list = mocker.AsyncMock(return_value=[])
 
     mocker.patch("app.endpoints.query_v2.get_system_prompt", return_value="PROMPT")
     mocker.patch("app.endpoints.query_v2.configuration", mocker.Mock(mcp_servers=[]))
@@ -289,6 +297,8 @@ async def test_retrieve_response_with_usage_dict(mocker):
     mock_vector_stores = mocker.Mock()
     mock_vector_stores.data = []
     mock_client.vector_stores.list = mocker.AsyncMock(return_value=mock_vector_stores)
+    # Mock shields.list
+    mock_client.shields.list = mocker.AsyncMock(return_value=[])
 
     mocker.patch("app.endpoints.query_v2.get_system_prompt", return_value="PROMPT")
     mocker.patch("app.endpoints.query_v2.configuration", mocker.Mock(mcp_servers=[]))
@@ -326,6 +336,8 @@ async def test_retrieve_response_with_empty_usage_dict(mocker):
     mock_vector_stores = mocker.Mock()
     mock_vector_stores.data = []
     mock_client.vector_stores.list = mocker.AsyncMock(return_value=mock_vector_stores)
+    # Mock shields.list
+    mock_client.shields.list = mocker.AsyncMock(return_value=[])
 
     mocker.patch("app.endpoints.query_v2.get_system_prompt", return_value="PROMPT")
     mocker.patch("app.endpoints.query_v2.configuration", mocker.Mock(mcp_servers=[]))
@@ -354,6 +366,8 @@ async def test_retrieve_response_validates_attachments(mocker):
     mock_vector_stores = mocker.Mock()
     mock_vector_stores.data = []
     mock_client.vector_stores.list = mocker.AsyncMock(return_value=mock_vector_stores)
+    # Mock shields.list
+    mock_client.shields.list = mocker.AsyncMock(return_value=[])
 
     mocker.patch("app.endpoints.query_v2.get_system_prompt", return_value="PROMPT")
     mocker.patch("app.endpoints.query_v2.configuration", mocker.Mock(mcp_servers=[]))
@@ -459,3 +473,177 @@ def _raise(*_args, **_kwargs):
     assert exc.value.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR
     assert "Unable to connect to Llama Stack" in str(exc.value.detail)
     fail_metric.inc.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_retrieve_response_with_shields_available(mocker):
+    """Test that shields are listed and passed to responses API when available."""
+    mock_client = mocker.Mock()
+
+    # Mock shields.list to return available shields
+    shield1 = mocker.Mock()
+    shield1.identifier = "shield-1"
+    shield2 = mocker.Mock()
+    shield2.identifier = "shield-2"
+    mock_client.shields.list = mocker.AsyncMock(return_value=[shield1, shield2])
+
+    output_item = mocker.Mock()
+    output_item.type = "message"
+    output_item.role = "assistant"
+    output_item.content = "Safe response"
+
+    response_obj = mocker.Mock()
+    response_obj.id = "resp-shields"
+    response_obj.output = [output_item]
+    response_obj.usage = None
+
+    mock_client.responses.create = mocker.AsyncMock(return_value=response_obj)
+    mock_vector_stores = mocker.Mock()
+    mock_vector_stores.data = []
+    mock_client.vector_stores.list = mocker.AsyncMock(return_value=mock_vector_stores)
+
+    mocker.patch("app.endpoints.query_v2.get_system_prompt", return_value="PROMPT")
+    mocker.patch("app.endpoints.query_v2.configuration", mocker.Mock(mcp_servers=[]))
+
+    qr = QueryRequest(query="hello")
+    summary, conv_id, _referenced_docs, _token_usage = await retrieve_response(
+        mock_client, "model-shields", qr, token="tkn", provider_id="test-provider"
+    )
+
+    assert conv_id == "resp-shields"
+    assert summary.llm_response == "Safe response"
+
+    # Verify that shields were passed in extra_body
+    kwargs = mock_client.responses.create.call_args.kwargs
+    assert "extra_body" in kwargs
+    assert "guardrails" in kwargs["extra_body"]
+    assert kwargs["extra_body"]["guardrails"] == ["shield-1", "shield-2"]
+
+
+@pytest.mark.asyncio
+async def test_retrieve_response_with_no_shields_available(mocker):
+    """Test that no extra_body is added when no shields are available."""
+    mock_client = mocker.Mock()
+
+    # Mock shields.list to return no shields
+    mock_client.shields.list = mocker.AsyncMock(return_value=[])
+
+    output_item = mocker.Mock()
+    output_item.type = "message"
+    output_item.role = "assistant"
+    output_item.content = "Response without shields"
+
+    response_obj = mocker.Mock()
+    response_obj.id = "resp-no-shields"
+    response_obj.output = [output_item]
+    response_obj.usage = None
+
+    mock_client.responses.create = mocker.AsyncMock(return_value=response_obj)
+    mock_vector_stores = mocker.Mock()
+    mock_vector_stores.data = []
+    mock_client.vector_stores.list = mocker.AsyncMock(return_value=mock_vector_stores)
+
+    mocker.patch("app.endpoints.query_v2.get_system_prompt", return_value="PROMPT")
+    mocker.patch("app.endpoints.query_v2.configuration", mocker.Mock(mcp_servers=[]))
+
+    qr = QueryRequest(query="hello")
+    summary, conv_id, _referenced_docs, _token_usage = await retrieve_response(
+        mock_client, "model-no-shields", qr, token="tkn", provider_id="test-provider"
+    )
+
+    assert conv_id == "resp-no-shields"
+    assert summary.llm_response == "Response without shields"
+
+    # Verify that no extra_body was added
+    kwargs = mock_client.responses.create.call_args.kwargs
+    assert "extra_body" not in kwargs
+
+
+@pytest.mark.asyncio
+async def test_retrieve_response_detects_shield_violation(mocker):
+    """Test that shield violations are detected and metrics are incremented."""
+    mock_client = mocker.Mock()
+
+    # Mock shields.list to return available shields
+    shield1 = mocker.Mock()
+    shield1.identifier = "safety-shield"
+    mock_client.shields.list = mocker.AsyncMock(return_value=[shield1])
+
+    # Create output with shield violation (refusal)
+    output_item = mocker.Mock()
+    output_item.type = "message"
+    output_item.role = "assistant"
+    output_item.content = "I cannot help with that request"
+    output_item.refusal = "Content violates safety policy"
+
+    response_obj = mocker.Mock()
+    response_obj.id = "resp-violation"
+    response_obj.output = [output_item]
+    response_obj.usage = None
+
+    mock_client.responses.create = mocker.AsyncMock(return_value=response_obj)
+    mock_vector_stores = mocker.Mock()
+    mock_vector_stores.data = []
+    mock_client.vector_stores.list = mocker.AsyncMock(return_value=mock_vector_stores)
+
+    mocker.patch("app.endpoints.query_v2.get_system_prompt", return_value="PROMPT")
+    mocker.patch("app.endpoints.query_v2.configuration", mocker.Mock(mcp_servers=[]))
+
+    # Mock the validation error metric
+    validation_metric = mocker.patch("metrics.llm_calls_validation_errors_total")
+
+    qr = QueryRequest(query="dangerous query")
+    summary, conv_id, _referenced_docs, _token_usage = await retrieve_response(
+        mock_client, "model-violation", qr, token="tkn", provider_id="test-provider"
+    )
+
+    assert conv_id == "resp-violation"
+    assert summary.llm_response == "I cannot help with that request"
+
+    # Verify that the validation error metric was incremented
+    validation_metric.inc.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_retrieve_response_no_violation_with_shields(mocker):
+    """Test that no metric is incremented when there's no shield violation."""
+    mock_client = mocker.Mock()
+
+    # Mock shields.list to return available shields
+    shield1 = mocker.Mock()
+    shield1.identifier = "safety-shield"
+    mock_client.shields.list = mocker.AsyncMock(return_value=[shield1])
+
+    # Create output without shield violation
+    output_item = mocker.Mock()
+    output_item.type = "message"
+    output_item.role = "assistant"
+    output_item.content = "Safe response"
+    output_item.refusal = None  # No violation
+
+    response_obj = mocker.Mock()
+    response_obj.id = "resp-safe"
+    response_obj.output = [output_item]
+    response_obj.usage = None
+
+    mock_client.responses.create = mocker.AsyncMock(return_value=response_obj)
+    mock_vector_stores = mocker.Mock()
+    mock_vector_stores.data = []
+    mock_client.vector_stores.list = mocker.AsyncMock(return_value=mock_vector_stores)
+
+    mocker.patch("app.endpoints.query_v2.get_system_prompt", return_value="PROMPT")
+    mocker.patch("app.endpoints.query_v2.configuration", mocker.Mock(mcp_servers=[]))
+
+    # Mock the validation error metric
+    validation_metric = mocker.patch("metrics.llm_calls_validation_errors_total")
+
+    qr = QueryRequest(query="safe query")
+    summary, conv_id, _referenced_docs, _token_usage = await retrieve_response(
+        mock_client, "model-safe", qr, token="tkn", provider_id="test-provider"
+    )
+
+    assert conv_id == "resp-safe"
+    assert summary.llm_response == "Safe response"
+
+    # Verify that the validation error metric was NOT incremented
+    validation_metric.inc.assert_not_called()