fix(runtime): set stop flag for exception-based rails in parallel mode (#1487)

Author: Pouyanpi

nemoguardrails/colang/v1_0/runtime/runtime.py

@@ -313,7 +313,8 @@ async def task_call_helper(flow_uid, post_event, func, *args, **kwargs):
             result = await func(*args, **kwargs)
 
             has_stop = any(
-                event["type"] == "BotIntent" and event["intent"] == "stop"
+                (event["type"] == "BotIntent" and event["intent"] == "stop")
+                or event["type"].endswith("Exception")
                 for event in result
             )
 
@@ -377,7 +378,8 @@ async def task_call_helper(flow_uid, post_event, func, *args, **kwargs):
 
                     # Check if this rail requested to stop
                     has_stop = any(
-                        event["type"] == "BotIntent" and event["intent"] == "stop"
+                        (event["type"] == "BotIntent" and event["intent"] == "stop")
+                        or event["type"].endswith("Exception")
                         for event in result
                     )
 
@@ -402,6 +404,7 @@ async def task_call_helper(flow_uid, post_event, func, *args, **kwargs):
                                     if v == pending_task:
                                         del unique_flow_ids[k]
                                         break
+                        # Remove the stopped flow from unique_flow_ids so it's not in finished_task_results
                         del unique_flow_ids[flow_id]
                         break
                     else:
@@ -446,15 +449,14 @@ async def task_call_helper(flow_uid, post_event, func, *args, **kwargs):
 
             def filter_and_append(logs, target_log):
                 for plog in logs:
-                    # Filter out "Listen" and "start_flow" events from task processing log
                     if plog["type"] == "event" and (
-                        plog["data"]["type"] == "Listen"
-                        or plog["data"]["type"] == "start_flow"
+                        plog["data"]["type"] == "start_flow"
                     ):
                         continue
                     target_log.append(plog)
 
-            filter_and_append(stopped_task_processing_logs, processing_log)
+            # Only append finished rails logs. Stopped rail logs should not be appended
+            # again since they're already in the processing log from when they started
             filter_and_append(finished_task_processing_logs, processing_log)
 
         # We pack all events into a single event to add it to the event history.
@@ -463,6 +465,7 @@ def filter_and_append(logs, target_log):
             data={"events": finished_task_results},
         )
 
+        # Return stopped_task_results separately so the caller knows to stop processing
         return ActionResult(
             events=[history_events] + stopped_task_results,
             context_updates=context_updates,

nemoguardrails/logging/processing_log.py

@@ -191,6 +191,15 @@ def compute_generation_log(processing_log: List[dict]) -> GenerationLog:
             elif event_type == "OutputRailsFinished":
                 output_rails_finished_at = event["timestamp"]
 
+            elif event_type.endswith("Exception"):
+                if activated_rail is not None and activated_rail.type in [
+                    "input",
+                    "output",
+                ]:
+                    activated_rail.stop = True
+                    if "stop" not in activated_rail.decisions:
+                        activated_rail.decisions.append("stop")
+
         elif event["type"] == "llm_call_info":
             if executed_action is not None:
                 executed_action.llm_calls.append(event["data"])
@@ -210,7 +219,8 @@ def compute_generation_log(processing_log: List[dict]) -> GenerationLog:
 
         if activated_rail.type in ["input", "output"]:
             activated_rail.stop = True
-            activated_rail.decisions.append("stop")
+            if "stop" not in activated_rail.decisions:
+                activated_rail.decisions.append("stop")
 
     # If we have input rails, we also record the general stats
     if input_rails_started_at:

tests/test_configs/parallel_rails_with_exceptions/actions.py

@@ -0,0 +1,69 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemoguardrails.actions import action
+
+
+@action(is_system_action=True)
+async def check_safety_action(context: dict):
+    user_message = context.get("user_message", "")
+
+    unsafe_terms = ["unsafe", "dangerous", "harmful", "kill", "violence"]
+    is_safe = not any(term in user_message.lower() for term in unsafe_terms)
+
+    return is_safe
+
+
+@action(is_system_action=True)
+async def check_topic_action(context: dict):
+    user_message = context.get("user_message", "")
+
+    off_topic_terms = ["offtopic", "irrelevant", "unrelated", "stupid", "idiot"]
+    is_on_topic = not any(term in user_message.lower() for term in off_topic_terms)
+
+    return is_on_topic
+
+
+@action(is_system_action=True)
+async def check_with_context_update(context: dict):
+    user_message = context.get("user_message", "")
+
+    violation_count = context.get("violation_count", 0)
+    context["violation_count"] = violation_count + 1
+
+    blocked_terms = ["blocked", "forbidden"]
+    is_allowed = not any(term in user_message.lower() for term in blocked_terms)
+
+    return is_allowed
+
+
+@action(is_system_action=True)
+async def check_output_safety_action(context: dict):
+    bot_message = context.get("bot_message", "")
+
+    unsafe_terms = ["harmful", "dangerous", "unsafe", "violence"]
+    is_safe = not any(term in bot_message.lower() for term in unsafe_terms)
+
+    return is_safe
+
+
+@action(is_system_action=True)
+async def check_output_length_action(context: dict):
+    bot_message = context.get("bot_message", "")
+
+    max_length = 500
+    is_valid = len(bot_message) <= max_length
+
+    return is_valid

tests/test_configs/parallel_rails_with_exceptions/config.yml

@@ -0,0 +1,19 @@
+models:
+  - type: main
+    engine: openai
+    model: gpt-3.5-turbo-instruct
+
+rails:
+  input:
+    flows:
+      - check safety with exception
+      - check topic with exception
+    parallel: True
+
+  output:
+    flows:
+      - check output safety with exception
+      - check output length with exception
+    parallel: True
+
+enable_rails_exceptions: True

tests/test_configs/parallel_rails_with_exceptions/rails.co

@@ -0,0 +1,34 @@
+define flow check safety with exception
+  $is_safe = execute check_safety_action
+
+  if not $is_safe
+    create event SafetyCheckException(message="Input blocked by safety check")
+    stop
+
+define flow check topic with exception
+  $is_on_topic = execute check_topic_action
+
+  if not $is_on_topic
+    create event TopicCheckException(message="Input blocked by topic check")
+    stop
+
+define flow check with context update
+  $is_allowed = execute check_with_context_update
+
+  if not $is_allowed
+    create event ContextUpdateException(message="Input blocked with context update")
+    stop
+
+define flow check output safety with exception
+  $is_safe = execute check_output_safety_action
+
+  if not $is_safe
+    create event OutputSafetyException(message="Output blocked by safety check")
+    stop
+
+define flow check output length with exception
+  $is_valid_length = execute check_output_length_action
+
+  if not $is_valid_length
+    create event OutputLengthException(message="Output blocked due to length")
+    stop