Update tests

Author: christinaexyou

nemoguardrails/actions/llm/generation.py

@@ -887,9 +887,22 @@ async def generate_bot_message(
         # of course, it does not work when passed as context in `run_output_rails_in_streaming`
         # streaming_handler is set when stream_async method is used
 
+        has_streaming_handler = streaming_handler is not None
+        has_output_streaming = (
+            self.config.rails.output.streaming
+            and self.config.rails.output.streaming.enabled
+        )
+        log.info(
+            f"generate_bot_message: streaming_handler={has_streaming_handler}, output.streaming={has_output_streaming}"
+        )
         # if streaming_handler and len(self.config.rails.output.flows) > 0:
-        if streaming_handler and self.config.rails.output.streaming.enabled:
+        if streaming_handler and has_output_streaming:
+            log.info("Setting skip_output_rails = True")
             context_updates["skip_output_rails"] = True
+        else:
+            log.info(
+                f"NOT setting skip_output_rails: streaming_handler={has_streaming_handler}, output.streaming={has_output_streaming}"
+            )
 
         if bot_intent in self.config.bot_messages:
             # Choose a message randomly from self.config.bot_messages[bot_message]
@@ -970,7 +983,9 @@ async def generate_bot_message(
                                 new_event_dict("BotMessage", text=text)
                             )
 
-                            return ActionResult(events=output_events)
+                            return ActionResult(
+                                events=output_events, context_updates=context_updates
+                            )
                         else:
                             if streaming_handler:
                                 await streaming_handler.push_chunk(
@@ -987,7 +1002,9 @@ async def generate_bot_message(
                                 )
                             output_events.append(bot_message_event)
 
-                            return ActionResult(events=output_events)
+                            return ActionResult(
+                                events=output_events, context_updates=context_updates
+                            )
 
             # If we are in passthrough mode, we just use the input for prompting
             if self.config.passthrough:

nemoguardrails/colang/v1_0/runtime/flows.py

@@ -15,6 +15,7 @@
 
 """A simplified modeling of the CoFlows engine."""
 
+import logging
 import uuid
 from dataclasses import dataclass, field
 from enum import Enum
@@ -25,6 +26,8 @@
 from nemoguardrails.colang.v1_0.runtime.sliding import slide
 from nemoguardrails.utils import new_event_dict, new_uuid
 
+log = logging.getLogger(__name__)
+
 
 @dataclass
 class FlowConfig:
@@ -356,7 +359,15 @@ def compute_next_state(state: State, event: dict) -> State:
     if event["type"] == "ContextUpdate":
         # TODO: add support to also remove keys from the context.
         #  maybe with a special context key e.g. "__remove__": ["key1", "key2"]
+        if "skip_output_rails" in event["data"]:
+            log.info(
+                f"ContextUpdate setting skip_output_rails={event['data']['skip_output_rails']}"
+            )
         state.context.update(event["data"])
+        if "skip_output_rails" in state.context:
+            log.info(
+                f"After update, context skip_output_rails={state.context.get('skip_output_rails')}"
+            )
         state.context_updates = {}
         state.next_step = None
         return state
@@ -415,6 +426,12 @@ def compute_next_state(state: State, event: dict) -> State:
             _record_next_step(new_state, flow_state, flow_config, priority_modifier=0.9)
             continue
 
+        # Debug logging for BotMessage event and skip_output_rails
+        if event["type"] == "BotMessage":
+            log.info(
+                f"BotMessage event processing for flow '{flow_config.id}', skip_output_rails in context: {flow_state.context.get('skip_output_rails', 'NOT SET')}"
+            )
+
         # If we're at a branching point, we look at all individual heads.
         matching_head = None
 

nemoguardrails/colang/v1_0/runtime/runtime.py

@@ -523,6 +523,9 @@ async def _run_output_rails_in_parallel_streaming(
             flows_with_params: Dictionary mapping flow_id to {"action_name": str, "params": dict}
             events: The events list for context
         """
+        # Compute context from events so actions can access bot_message
+        context = compute_context(events)
+
         tasks = []
 
         async def run_single_rail(flow_id: str, action_info: dict) -> tuple:
@@ -532,8 +535,11 @@ async def run_single_rail(flow_id: str, action_info: dict) -> tuple:
                 action_name = action_info["action_name"]
                 params = action_info["params"]
 
+                # Merge context into params so actions have access to bot_message
+                params_with_context = {**params, "context": context}
+
                 result_tuple = await self.action_dispatcher.execute_action(
-                    action_name, params
+                    action_name, params_with_context
                 )
                 result, status = result_tuple
 
@@ -731,10 +737,19 @@ async def _process_start_action(self, events: List[dict]) -> List[dict]:
         return_events = []
         context_updates = {}
 
+        if action_name == "generate_bot_message":
+            log.info(
+                f"DEBUG: generate_bot_message returned, isinstance(ActionResult)={isinstance(result, ActionResult)}"
+            )
+
         if isinstance(result, ActionResult):
             return_value = result.return_value
             return_events = result.events
             context_updates.update(result.context_updates)
+            if action_name == "generate_bot_message":
+                log.info(
+                    f"generate_bot_message ActionResult: context_updates={context_updates}, skip_output_rails={'skip_output_rails' in context_updates}"
+                )
 
         # If we have an action result key, we also record the update.
         if action_result_key:

nemoguardrails/rails/llm/event_translator.py

@@ -1,3 +1,18 @@
+# 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.
+
 """Event translator for converting between messages and events."""
 
 import logging
@@ -24,7 +39,7 @@ def __init__(self, config: RailsConfig):
         self.events_history_cache: Dict[str, List[dict]] = {}
 
     def messages_to_events(
-        self, messages: List[dict], state: Optional[Any] = None
+        self, messages: List[dict[str, Any]], state: Optional[Any] = None
     ) -> List[dict]:
         """Convert messages to events.
 
@@ -156,7 +171,7 @@ def _messages_to_events_v1(self, messages: List[dict]) -> List[dict]:
         return events
 
     def _messages_to_events_v2(
-        self, messages: List[dict], state: Optional[Any]
+        self, messages: List[dict[str, Any]], state: Optional[Any]
     ) -> List[dict]:
         """Convert messages to events for Colang 2.x.