Refactor of shield moderation

Author: asimurka

pyproject.toml

@@ -28,9 +28,9 @@ dependencies = [
     # Used by authentication/k8s integration
     "kubernetes>=30.1.0",
     # Used to call Llama Stack APIs
-    "llama-stack==0.4.3",
-    "llama-stack-client==0.4.3",
-    "llama-stack-api==0.4.4",
+    "llama-stack==0.5.2",
+    "llama-stack-client==0.5.2",
+    "llama-stack-api==0.5.2",
     # Used by Logger
     "rich>=14.0.0",
     # Used by JWK token auth handler

src/app/endpoints/query.py

@@ -36,6 +36,7 @@
     UnauthorizedResponse,
     UnprocessableEntityResponse,
 )
+from utils.conversations import append_turn_items_to_conversation
 from utils.endpoints import (
     check_configuration_loaded,
     validate_and_retrieve_conversation,
@@ -59,14 +60,11 @@
     get_topic_summary,
     prepare_responses_params,
 )
-from utils.shields import (
-    append_turn_to_conversation,
-    run_shield_moderation,
-    validate_shield_ids_override,
-)
+from utils.shields import run_shield_moderation, validate_shield_ids_override
 from utils.suid import normalize_conversation_id
 from utils.types import (
     ResponsesApiParams,
+    ShieldModerationResult,
     TurnSummary,
 )
 from utils.vector_search import build_rag_context
@@ -158,14 +156,21 @@ async def query_endpoint_handler(
 
     client = AsyncLlamaStackClientHolder().get_client()
 
-    # Build RAG context from Inline RAG sources
-    inline_rag_context = await build_rag_context(
-        client, query_request.query, query_request.vector_store_ids, query_request.solr
-    )
-
     # Moderation input is the raw user content (query + attachments) without injected RAG
     # context, to avoid false positives from retrieved document content.
     moderation_input = prepare_input(query_request)
+    moderation_result = await run_shield_moderation(
+        client, moderation_input, query_request.shield_ids
+    )
+
+    # Build RAG context from Inline RAG sources
+    inline_rag_context = await build_rag_context(
+        client,
+        moderation_result.decision,
+        query_request.query,
+        query_request.vector_store_ids,
+        query_request.solr,
+    )
 
     # Prepare API request parameters
     responses_params = await prepare_responses_params(
@@ -177,7 +182,7 @@ async def query_endpoint_handler(
         stream=False,
         store=True,
         request_headers=request.headers,
-        inline_rag_context=inline_rag_context.context_text or None,
+        inline_rag_context=inline_rag_context.context_text,
     )
 
     # Handle Azure token refresh if needed
@@ -189,32 +194,22 @@ async def query_endpoint_handler(
     ):
         client = await update_azure_token(client)
 
-    # Build index identification mapping for RAG source resolution
-    vector_store_ids = extract_vector_store_ids_from_tools(responses_params.tools)
-    rag_id_mapping = configuration.rag_id_mapping
-
     # Retrieve response using Responses API
-    turn_summary = await retrieve_response(
-        client,
-        responses_params,
-        query_request.shield_ids,
-        vector_store_ids,
-        rag_id_mapping,
-        moderation_input=moderation_input,
-    )
-
-    # Combine inline RAG results (BYOK + Solr) with tool-based RAG results for the transcript
-    rag_chunks = inline_rag_context.rag_chunks
-    tool_rag_chunks = turn_summary.rag_chunks or []
-    logger.info("RAG as a tool retrieved %d chunks", len(tool_rag_chunks))
-    turn_summary.rag_chunks = rag_chunks + tool_rag_chunks
-
-    # Add tool-based RAG documents and chunks
-    rag_documents = inline_rag_context.referenced_documents
-    tool_rag_documents = turn_summary.referenced_documents or []
-    turn_summary.referenced_documents = deduplicate_referenced_documents(
-        rag_documents + tool_rag_documents
-    )
+    turn_summary = await retrieve_response(client, responses_params, moderation_result)
+
+    if moderation_result.decision == "passed":
+        # Combine inline RAG results (BYOK + Solr) with tool-based RAG results for the transcript
+        rag_chunks = inline_rag_context.rag_chunks
+        tool_rag_chunks = turn_summary.rag_chunks
+        logger.info("RAG as a tool retrieved %d chunks", len(tool_rag_chunks))
+        turn_summary.rag_chunks = rag_chunks + tool_rag_chunks
+
+        # Add tool-based RAG documents and chunks
+        rag_documents = inline_rag_context.referenced_documents
+        tool_rag_documents = turn_summary.referenced_documents
+        turn_summary.referenced_documents = deduplicate_referenced_documents(
+            rag_documents + tool_rag_documents
+        )
 
     # Get topic summary for new conversation
     if not user_conversation and query_request.generate_topic_summary:
@@ -274,10 +269,7 @@ async def query_endpoint_handler(
 async def retrieve_response(  # pylint: disable=too-many-locals
     client: AsyncLlamaStackClient,
     responses_params: ResponsesApiParams,
-    shield_ids: Optional[list[str]] = None,
-    vector_store_ids: Optional[list[str]] = None,
-    rag_id_mapping: Optional[dict[str, str]] = None,
-    moderation_input: Optional[str] = None,
+    moderation_result: ShieldModerationResult,
 ) -> TurnSummary:
     """
     Retrieve response from LLMs and agents.
@@ -288,33 +280,23 @@ async def retrieve_response(  # pylint: disable=too-many-locals
     Parameters:
         client: The AsyncLlamaStackClient to use for the request.
         responses_params: The Responses API parameters.
-        shield_ids: Optional list of shield IDs for moderation.
-        vector_store_ids: Vector store IDs used in the query for source resolution.
-        rag_id_mapping: Mapping from vector_db_id to user-facing rag_id.
-        moderation_input: Text to moderate. Should be the raw user content (query +
-            attachments) without injected RAG context to avoid false positives.
-            Falls back to responses_params.input if not provided.
+        moderation_result: The moderation result.
 
     Returns:
         TurnSummary: Summary of the LLM response content
     """
     response: Optional[OpenAIResponseObject] = None
-    try:
-        moderation_result = await run_shield_moderation(
+    if moderation_result.decision == "blocked":
+        await append_turn_items_to_conversation(
             client,
-            moderation_input or cast(str, responses_params.input),
-            shield_ids,
+            responses_params.conversation,
+            responses_params.input,
+            [moderation_result.refusal_response],
+        )
+        return TurnSummary(
+            id=moderation_result.moderation_id, llm_response=moderation_result.message
         )
-        if moderation_result.decision == "blocked":
-            # Handle shield moderation blocking
-            violation_message = moderation_result.message
-            await append_turn_to_conversation(
-                client,
-                responses_params.conversation,
-                cast(str, responses_params.input),
-                violation_message,
-            )
-            return TurnSummary(llm_response=violation_message)
+    try:
         response = await client.responses.create(
             **responses_params.model_dump(exclude_none=True)
         )
@@ -335,6 +317,8 @@ async def retrieve_response(  # pylint: disable=too-many-locals
         error_response = handle_known_apistatus_errors(e, responses_params.model)
         raise HTTPException(**error_response.model_dump()) from e
 
+    vector_store_ids = extract_vector_store_ids_from_tools(responses_params.tools)
+    rag_id_mapping = configuration.rag_id_mapping
     return build_turn_summary(
         response, responses_params.model, vector_store_ids, rag_id_mapping
     )

src/app/endpoints/streaming_query.py

@@ -57,6 +57,7 @@
     UnauthorizedResponse,
     UnprocessableEntityResponse,
 )
+from utils.conversations import append_turn_items_to_conversation
 from utils.endpoints import (
     check_configuration_loaded,
     validate_and_retrieve_conversation,
@@ -189,10 +190,22 @@ async def streaming_query_endpoint_handler(  # pylint: disable=too-many-locals
 
     client = AsyncLlamaStackClientHolder().get_client()
 
+    # Moderation input is the raw user content (query + attachments) without injected RAG
+    # context, to avoid false positives from retrieved document content.
+    moderation_input = prepare_input(query_request)
+    moderation_result = await run_shield_moderation(
+        client, moderation_input, query_request.shield_ids
+    )
+
     # Build RAG context from Inline RAG sources
     inline_rag_context = await build_rag_context(
-        client, query_request.query, query_request.vector_store_ids, query_request.solr
+        client,
+        moderation_result.decision,
+        query_request.query,
+        query_request.vector_store_ids,
+        query_request.solr,
     )
+
     # Prepare API request parameters
     responses_params = await prepare_responses_params(
         client=client,
@@ -203,7 +216,7 @@ async def streaming_query_endpoint_handler(  # pylint: disable=too-many-locals
         stream=True,
         store=True,
         request_headers=request.headers,
-        inline_rag_context=inline_rag_context.context_text or None,
+        inline_rag_context=inline_rag_context.context_text,
     )
 
     # Handle Azure token refresh if needed
@@ -227,6 +240,7 @@ async def streaming_query_endpoint_handler(  # pylint: disable=too-many-locals
         query_request=query_request,
         started_at=started_at,
         client=client,
+        moderation_result=moderation_result,
         vector_store_ids=extract_vector_store_ids_from_tools(responses_params.tools),
         rag_id_mapping=configuration.rag_id_mapping,
     )
@@ -240,9 +254,15 @@ async def streaming_query_endpoint_handler(  # pylint: disable=too-many-locals
     generator, turn_summary = await retrieve_response_generator(
         responses_params=responses_params,
         context=context,
-        inline_rag_documents=inline_rag_context.referenced_documents,
+        inline_rag_docs=inline_rag_context.referenced_documents,
     )
 
+    # Combine inline RAG results (BYOK + Solr) with tool-based results
+    if context.moderation_result.decision == "passed":
+        turn_summary.referenced_documents = deduplicate_referenced_documents(
+            inline_rag_context.referenced_documents + turn_summary.referenced_documents
+        )
+
     response_media_type = (
         MEDIA_TYPE_TEXT
         if query_request.media_type == MEDIA_TYPE_TEXT
@@ -263,7 +283,7 @@ async def streaming_query_endpoint_handler(  # pylint: disable=too-many-locals
 async def retrieve_response_generator(
     responses_params: ResponsesApiParams,
     context: ResponseGeneratorContext,
-    inline_rag_documents: list[ReferencedDocument],
+    inline_rag_docs: list[ReferencedDocument],
 ) -> tuple[AsyncIterator[str], TurnSummary]:
     """
     Retrieve the appropriate response generator.
@@ -275,40 +295,44 @@ async def retrieve_response_generator(
     Args:
         responses_params: The Responses API parameters
         context: The response generator context
-        inline_rag_documents: Referenced documents from inline  RAG (BYOK + Solr)
-
+        inline_rag_docs: Inline RAG (BYOK + Solr) documents
     Returns:
         tuple[AsyncIterator[str], TurnSummary]: The response generator and turn summary
 
     """
     turn_summary = TurnSummary()
     try:
-        moderation_result = await run_shield_moderation(
-            context.client,
-            prepare_input(context.query_request),
-            context.query_request.shield_ids,
-        )
-        if moderation_result.decision == "blocked":
-            turn_summary.llm_response = moderation_result.message
-            await append_turn_to_conversation(
+        if context.moderation_result.decision == "blocked":
+            turn_summary.llm_response = context.moderation_result.message
+            turn_summary.id = context.moderation_result.moderation_id
+            await append_turn_items_to_conversation(
                 context.client,
                 responses_params.conversation,
-                cast(str, responses_params.input),
-                moderation_result.message,
+                responses_params.input,
+                [context.moderation_result.refusal_response],
             )
             media_type = context.query_request.media_type or MEDIA_TYPE_JSON
             return (
-                shield_violation_generator(moderation_result.message, media_type),
+                shield_violation_generator(
+                    context.moderation_result.message,
+                    media_type,
+                ),
                 turn_summary,
             )
         # Retrieve response stream (may raise exceptions)
         response = await context.client.responses.create(
             **responses_params.model_dump(exclude_none=True)
         )
         # Store pre-RAG documents for later merging with tool-based RAG
-        turn_summary.inline_rag_documents = inline_rag_documents
-        return response_generator(response, context, turn_summary), turn_summary
-
+        return (
+            response_generator(
+                response,
+                context,
+                turn_summary,
+                inline_rag_docs,
+            ),
+            turn_summary,
+        )
     # Handle know LLS client errors only at stream creation time and shield execution
     except RuntimeError as e:  # library mode wraps 413 into runtime error
         if "context_length" in str(e).lower():
@@ -559,6 +583,7 @@ async def response_generator(  # pylint: disable=too-many-branches,too-many-stat
     turn_response: AsyncIterator[OpenAIResponseObjectStream],
     context: ResponseGeneratorContext,
     turn_summary: TurnSummary,
+    inline_rag_docs: list[ReferencedDocument],
 ) -> AsyncIterator[str]:
     """Generate SSE formatted streaming response.
 
@@ -570,7 +595,7 @@ async def response_generator(  # pylint: disable=too-many-branches,too-many-stat
         turn_response: The streaming response from Llama Stack
         context: The response generator context
         turn_summary: TurnSummary to populate during streaming
-
+        inline_rag_docs: Inline RAG (BYOK + Solr) documents
     Yields:
         SSE-formatted strings for tokens, tool calls, tool results,
         turn completion, and error events.
@@ -741,15 +766,15 @@ async def response_generator(  # pylint: disable=too-many-branches,too-many-stat
     turn_summary.token_usage = extract_token_usage(
         latest_response_object.usage, context.model_id
     )
-    tool_based_documents = parse_referenced_documents(
+    # Parse tool-based referenced documents from the final response object
+    tool_rag_docs = parse_referenced_documents(
         latest_response_object,
         vector_store_ids=context.vector_store_ids,
         rag_id_mapping=context.rag_id_mapping,
     )
-
-    # Merge pre-RAG documents with tool-based documents and deduplicate
+    # Combine inline RAG results (BYOK + Solr) with tool-based results
     turn_summary.referenced_documents = deduplicate_referenced_documents(
-        turn_summary.inline_rag_documents + tool_based_documents
+        inline_rag_docs + tool_rag_docs
     )
 
 

src/constants.py

@@ -2,7 +2,7 @@
 
 # Minimal and maximal supported Llama Stack version
 MINIMAL_SUPPORTED_LLAMA_STACK_VERSION = "0.2.17"
-MAXIMAL_SUPPORTED_LLAMA_STACK_VERSION = "0.4.3"
+MAXIMAL_SUPPORTED_LLAMA_STACK_VERSION = "0.5.2"
 
 UNABLE_TO_PROCESS_RESPONSE = "Unable to process this request"
 
@@ -214,3 +214,5 @@
 # Environment variable to force StreamHandler instead of RichHandler
 # Set to any non-empty value to disable RichHandler
 LIGHTSPEED_STACK_DISABLE_RICH_HANDLER_ENV_VAR = "LIGHTSPEED_STACK_DISABLE_RICH_HANDLER"
+
+DEFAULT_VIOLATION_MESSAGE = "I cannot process this request due to policy restrictions."

src/models/context.py

@@ -4,6 +4,7 @@
 from llama_stack_client import AsyncLlamaStackClient
 
 from models.requests import QueryRequest
+from utils.types import ShieldModerationResult
 
 
 @dataclass
@@ -23,6 +24,7 @@ class ResponseGeneratorContext:  # pylint: disable=too-many-instance-attributes
         query_request: The query request object
         started_at: Timestamp when the request started (ISO 8601 format)
         client: The Llama Stack client for API interactions
+        moderation_result: The moderation result
         vector_store_ids: Vector store IDs used in the query for source resolution.
         rag_id_mapping: Mapping from vector_db_id to user-facing rag_id.
     """
@@ -42,6 +44,7 @@ class ResponseGeneratorContext:  # pylint: disable=too-many-instance-attributes
 
     # Dependencies & State
     client: AsyncLlamaStackClient
+    moderation_result: ShieldModerationResult
 
     # RAG index identification
     vector_store_ids: list[str] = field(default_factory=list)