Shields impl using Moderations API

Author: asimurka

docker-compose-library.yaml

@@ -18,7 +18,7 @@ services:
       - TAVILY_SEARCH_API_KEY=${TAVILY_SEARCH_API_KEY:-}
       # OpenAI
       - OPENAI_API_KEY=${OPENAI_API_KEY}
-      - E2E_OPENAI_MODEL=${E2E_OPENAI_MODEL:-gpt-4-turbo}
+      - E2E_OPENAI_MODEL=${E2E_OPENAI_MODEL:-gpt-4o-mini}
       # Azure
       - AZURE_API_KEY=${AZURE_API_KEY:-}
       # RHAIIS

docker-compose.yaml

@@ -16,7 +16,7 @@ services:
       - TAVILY_SEARCH_API_KEY=${TAVILY_SEARCH_API_KEY:-}
       # OpenAI
       - OPENAI_API_KEY=${OPENAI_API_KEY}
-      - E2E_OPENAI_MODEL=${E2E_OPENAI_MODEL}
+      - E2E_OPENAI_MODEL=${E2E_OPENAI_MODEL:-gpt-4o-mini}
       # Azure
       - AZURE_API_KEY=${AZURE_API_KEY}
       # RHAIIS

run.yaml

@@ -131,8 +131,15 @@ storage:
       namespace: prompts
       backend: kv_default
 registered_resources:
-  models: []
-  shields: []
+  models:
+  - model_id: gpt-4o-mini
+    provider_id: openai
+    model_type: llm
+    provider_model_id: gpt-4o-mini
+  shields:
+  - shield_id: llama-guard
+    provider_id: llama-guard
+    provider_shield_id: openai/gpt-4o-mini
   vector_dbs: []
   datasets: []
   scoring_fns: []

src/app/endpoints/query.py

@@ -8,11 +8,11 @@
 from typing import Annotated, Any, Optional, cast
 
 from fastapi import APIRouter, Depends, HTTPException, Request
-from litellm.exceptions import RateLimitError
 from llama_stack_client import (
     APIConnectionError,
     APIStatusError,
-    AsyncLlamaStackClient,  # type: ignore
+    AsyncLlamaStackClient,
+    RateLimitError,  # type: ignore
 )
 from llama_stack_client.types import Shield, UserMessage  # type: ignore
 from llama_stack_client.types.alpha.agents.turn import Turn

src/app/endpoints/query_v2.py

@@ -10,7 +10,7 @@
 from llama_stack.apis.agents.openai_responses import (
     OpenAIResponseObject,
 )
-from llama_stack_client import AsyncLlamaStackClient  # type: ignore
+from llama_stack_client import AsyncLlamaStackClient
 
 import metrics
 from app.endpoints.query import (
@@ -42,7 +42,10 @@
 )
 from utils.mcp_headers import mcp_headers_dependency
 from utils.responses import extract_text_from_response_output_item
-from utils.shields import detect_shield_violations, get_available_shields
+from utils.shields import (
+    append_turn_to_conversation,
+    run_shield_moderation,
+)
 from utils.suid import normalize_conversation_id, to_llama_stack_conversation_id
 from utils.token_counter import TokenCounter
 from utils.types import RAGChunk, ToolCallSummary, ToolResultSummary, TurnSummary
@@ -322,9 +325,6 @@ 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.
     """
-    # List available shields for Responses API
-    available_shields = await get_available_shields(client)
-
     # use system prompt from request or default one
     system_prompt = get_system_prompt(query_request, configuration)
     logger.debug("Using system prompt: %s", system_prompt)
@@ -370,6 +370,26 @@ async def retrieve_response(  # pylint: disable=too-many-locals,too-many-branche
             conversation_id,
         )
 
+    # Run shield moderation before calling LLM
+    moderation_result = await run_shield_moderation(client, input_text)
+    if moderation_result.blocked:
+        violation_message = moderation_result.message or ""
+        await append_turn_to_conversation(
+            client, llama_stack_conv_id, input_text, violation_message
+        )
+        summary = TurnSummary(
+            llm_response=violation_message,
+            tool_calls=[],
+            tool_results=[],
+            rag_chunks=[],
+        )
+        return (
+            summary,
+            normalize_conversation_id(conversation_id),
+            [],
+            TokenCounter(),
+        )
+
     # Create OpenAI response using responses API
     create_kwargs: dict[str, Any] = {
         "input": input_text,
@@ -381,10 +401,6 @@ async def retrieve_response(  # pylint: disable=too-many-locals,too-many-branche
         "conversation": llama_stack_conv_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)
     logger.info("Response: %s", response)
@@ -410,9 +426,6 @@ async def retrieve_response(  # pylint: disable=too-many-locals,too-many-branche
         if tool_result:
             tool_results.append(tool_result)
 
-    # 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",
         len(tool_calls),

src/app/endpoints/streaming_query.py

@@ -11,10 +11,10 @@
 
 from fastapi import APIRouter, Depends, Request
 from fastapi.responses import StreamingResponse
-from litellm.exceptions import RateLimitError
 from llama_stack_client import (
     APIConnectionError,
-    AsyncLlamaStackClient,  # type: ignore
+    AsyncLlamaStackClient,
+    RateLimitError,  # type: ignore
 )
 from llama_stack_client.types import UserMessage  # type: ignore
 from llama_stack_client.types.alpha.agents.agent_turn_response_stream_chunk import (

src/app/endpoints/streaming_query_v2.py

@@ -6,8 +6,14 @@
 from fastapi import APIRouter, Depends, Request
 from fastapi.responses import StreamingResponse
 from llama_stack.apis.agents.openai_responses import (
+    OpenAIResponseContentPartOutputText,
+    OpenAIResponseMessage,
     OpenAIResponseObject,
     OpenAIResponseObjectStream,
+    OpenAIResponseObjectStreamResponseCompleted,
+    OpenAIResponseObjectStreamResponseContentPartAdded,
+    OpenAIResponseObjectStreamResponseOutputTextDelta,
+    OpenAIResponseOutputMessageContentOutputText,
 )
 from llama_stack_client import AsyncLlamaStackClient
 
@@ -53,7 +59,10 @@
 from utils.quota import consume_tokens, get_available_quotas
 from utils.suid import normalize_conversation_id, to_llama_stack_conversation_id
 from utils.mcp_headers import mcp_headers_dependency
-from utils.shields import detect_shield_violations, get_available_shields
+from utils.shields import (
+    append_turn_to_conversation,
+    run_shield_moderation,
+)
 from utils.token_counter import TokenCounter
 from utils.transcripts import store_transcript
 from utils.types import ToolCallSummary, TurnSummary
@@ -234,12 +243,6 @@ async def response_generator(  # pylint: disable=too-many-branches,too-many-stat
                 # 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:
-                    output = getattr(latest_response_object, "output", None)
-                    if output is not None:
-                        detect_shield_violations(output)
-
                 if not emitted_turn_complete:
                     final_message = summary.llm_response or "".join(text_parts)
                     if not final_message:
@@ -394,9 +397,6 @@ async def retrieve_response(  # pylint: disable=too-many-locals
         tuple: A tuple containing the streaming response object
         and the conversation ID.
     """
-    # List available shields for Responses API
-    available_shields = await get_available_shields(client)
-
     # use system prompt from request or default one
     system_prompt = get_system_prompt(query_request, configuration)
     logger.debug("Using system prompt: %s", system_prompt)
@@ -441,6 +441,18 @@ async def retrieve_response(  # pylint: disable=too-many-locals
             conversation_id,
         )
 
+    # Run shield moderation before calling LLM
+    moderation_result = await run_shield_moderation(client, input_text)
+    if moderation_result.blocked:
+        violation_message = moderation_result.message or ""
+        await append_turn_to_conversation(
+            client, llama_stack_conv_id, input_text, violation_message
+        )
+        return (
+            create_violation_stream(violation_message, moderation_result.shield_model),
+            normalize_conversation_id(conversation_id),
+        )
+
     create_params: dict[str, Any] = {
         "input": input_text,
         "model": model_id,
@@ -451,14 +463,55 @@ async def retrieve_response(  # pylint: disable=too-many-locals
         "conversation": llama_stack_conv_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)
-    # async for chunk in response_stream:
-    #     logger.error("Chunk: %s", chunk.model_dump_json())
-    # Return the normalized conversation_id
-    # The response_generator will emit it in the start event
+
     return response_stream, normalize_conversation_id(conversation_id)
+
+
+async def create_violation_stream(
+    message: str,
+    shield_model: str | None = None,
+) -> AsyncIterator[OpenAIResponseObjectStream]:
+    """Create a minimal response stream for shield violations."""
+    response_id = "resp_shield_violation"
+    item_id = "msg_shield_violation"
+
+    # Content part added (triggers empty initial token)
+    yield OpenAIResponseObjectStreamResponseContentPartAdded(
+        content_index=0,
+        response_id=response_id,
+        item_id=item_id,
+        output_index=0,
+        part=OpenAIResponseContentPartOutputText(text=""),
+        sequence_number=0,
+    )
+
+    # Text delta
+    yield OpenAIResponseObjectStreamResponseOutputTextDelta(
+        content_index=0,
+        delta=message,
+        item_id=item_id,
+        output_index=0,
+        sequence_number=1,
+    )
+
+    # Completed response
+    yield OpenAIResponseObjectStreamResponseCompleted(
+        response=OpenAIResponseObject(
+            id=response_id,
+            created_at=0,
+            model=shield_model or "shield",
+            output=[
+                OpenAIResponseMessage(
+                    id=item_id,
+                    content=[
+                        OpenAIResponseOutputMessageContentOutputText(text=message)
+                    ],
+                    role="assistant",
+                    status="completed",
+                )
+            ],
+            status="completed",
+        )
+    )

src/utils/shields.py

@@ -1,14 +1,20 @@
 """Utility functions for working with Llama Stack shields."""
 
 import logging
-from typing import Any
+from typing import Any, cast
 
-from llama_stack_client import AsyncLlamaStackClient
+from fastapi import HTTPException
+from llama_stack_client import AsyncLlamaStackClient, BadRequestError
+from llama_stack_client.types import CreateResponse
 
 import metrics
+from models.responses import NotFoundResponse
+from utils.types import ShieldModerationResult
 
 logger = logging.getLogger(__name__)
 
+DEFAULT_VIOLATION_MESSAGE = "I cannot process this request due to policy restrictions."
+
 
 async def get_available_shields(client: AsyncLlamaStackClient) -> list[str]:
     """
@@ -52,3 +58,100 @@ def detect_shield_violations(output_items: list[Any]) -> bool:
                 logger.warning("Shield violation detected: %s", refusal)
                 return True
     return False
+
+
+async def run_shield_moderation(
+    client: AsyncLlamaStackClient,
+    input_text: str,
+) -> ShieldModerationResult:
+    """
+    Run shield moderation on input text.
+
+    Iterates through all configured shields and runs moderation checks.
+    Raises HTTPException if shield model is not found.
+
+    Parameters:
+        client: The Llama Stack client.
+        input_text: The text to moderate.
+
+    Returns:
+        ShieldModerationResult: Result indicating if content was blocked and the message.
+
+    Raises:
+        HTTPException: If shield's provider_resource_id is not configured or model not found.
+    """
+    available_models = {model.identifier for model in await client.models.list()}
+
+    for shield in await client.shields.list():
+        if (
+            not shield.provider_resource_id
+            or shield.provider_resource_id not in available_models
+        ):
+            response = NotFoundResponse(
+                resource="Shield model", resource_id=shield.provider_resource_id or ""
+            )
+            raise HTTPException(**response.model_dump())
+
+        try:
+            moderation = await client.moderations.create(
+                input=input_text, model=shield.provider_resource_id
+            )
+            moderation_result = cast(CreateResponse, moderation)
+
+            if moderation_result.results and moderation_result.results[0].flagged:
+                result = moderation_result.results[0]
+                metrics.llm_calls_validation_errors_total.inc()
+                logger.warning(
+                    "Shield '%s' flagged content: categories=%s",
+                    shield.identifier,
+                    result.categories,
+                )
+                violation_message = result.user_message or DEFAULT_VIOLATION_MESSAGE
+                return ShieldModerationResult(
+                    blocked=True,
+                    message=violation_message,
+                    shield_model=shield.provider_resource_id,
+                )
+
+        # Known Llama Stack bug: BadRequestError is raised when violation is present
+        # in the shield LLM response but has wrong format that cannot be parsed.
+        except BadRequestError:
+            logger.warning(
+                "Shield '%s' returned BadRequestError, treating as blocked",
+                shield.identifier,
+            )
+            metrics.llm_calls_validation_errors_total.inc()
+            return ShieldModerationResult(
+                blocked=True,
+                message=DEFAULT_VIOLATION_MESSAGE,
+                shield_model=shield.provider_resource_id,
+            )
+
+    return ShieldModerationResult(blocked=False)
+
+
+async def append_turn_to_conversation(
+    client: AsyncLlamaStackClient,
+    conversation_id: str,
+    user_message: str,
+    assistant_message: str,
+) -> None:
+    """
+    Append a user/assistant turn to a conversation after shield violation.
+
+    Used to record the conversation turn when a shield blocks the request,
+    storing both the user's original message and the violation response.
+
+    Parameters:
+        client: The Llama Stack client.
+        conversation_id: The Llama Stack conversation ID.
+        user_message: The user's input message.
+        assistant_message: The shield violation response message.
+    """
+    await client.conversations.items.create(
+        conversation_id,
+        items=[
+            {"type": "message", "role": "user", "content": user_message},
+            {"type": "message", "role": "assistant", "content": assistant_message},
+        ],
+    )