From f36a7038f632613f986bc739a1c353ff4b5031ea Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Mon, 8 Jun 2026 21:55:31 +0200 Subject: [PATCH] refactor(llmrails): extract streaming workflow --- nemoguardrails/rails/llm/llmrails.py | 374 +------------- .../rails/llm/streaming/__init__.py | 34 ++ .../rails/llm/streaming/generation_stream.py | 162 ++++++ .../llm/streaming/streaming_output_rails.py | 378 ++++++++++++++ tests/rails/llm/test_generation_stream.py | 364 +++++++++++++ .../rails/llm/test_streaming_output_rails.py | 480 ++++++++++++++++++ 6 files changed, 1444 insertions(+), 348 deletions(-) create mode 100644 nemoguardrails/rails/llm/streaming/__init__.py create mode 100644 nemoguardrails/rails/llm/streaming/generation_stream.py create mode 100644 nemoguardrails/rails/llm/streaming/streaming_output_rails.py create mode 100644 tests/rails/llm/test_generation_stream.py create mode 100644 tests/rails/llm/test_streaming_output_rails.py diff --git a/nemoguardrails/rails/llm/llmrails.py b/nemoguardrails/rails/llm/llmrails.py index 3f9101c61b..d1f60a3a2a 100644 --- a/nemoguardrails/rails/llm/llmrails.py +++ b/nemoguardrails/rails/llm/llmrails.py @@ -15,16 +15,12 @@ """LLM Rails entry point.""" -import asyncio -import json import logging import warnings -from functools import partial from typing import ( Any, AsyncIterator, Callable, - Dict, List, Literal, Optional, @@ -36,19 +32,16 @@ from typing_extensions import Self -from nemoguardrails.actions.output_mapping import is_output_blocked from nemoguardrails.base_guardrails import BaseGuardrails from nemoguardrails.colang.v1_0.runtime.runtime import Runtime from nemoguardrails.colang.v2_x.runtime.flows import State from nemoguardrails.embeddings.index import EmbeddingsIndex from nemoguardrails.embeddings.providers import register_embedding_provider from nemoguardrails.embeddings.providers.base import EmbeddingModel -from nemoguardrails.exceptions import StreamingNotSupportedError from nemoguardrails.llm.models.initializer import init_llm_model from nemoguardrails.logging.explain import ExplainInfo from nemoguardrails.logging.verbose import set_verbose from nemoguardrails.patch_asyncio import check_sync_call_from_async_loop -from nemoguardrails.rails.llm.buffer import get_buffer_strategy from nemoguardrails.rails.llm.config import ( OutputRailsStreamingConfig, RailsConfig, @@ -90,15 +83,14 @@ sync_update_llm_bindings, ) from nemoguardrails.rails.llm.startup.tracing import create_startup_tracing_adapters -from nemoguardrails.rails.llm.utils import ( - get_action_details_from_flow_id, +from nemoguardrails.rails.llm.streaming.generation_stream import ( + generation_token_stream, + validate_streaming_with_output_rails, ) -from nemoguardrails.streaming import END_OF_STREAM, StreamingHandler +from nemoguardrails.rails.llm.streaming.streaming_output_rails import run_output_rails_in_streaming +from nemoguardrails.streaming import StreamingHandler from nemoguardrails.types import LLMModel -from nemoguardrails.utils import ( - extract_error_json, - get_or_create_event_loop, -) +from nemoguardrails.utils import get_or_create_event_loop log = logging.getLogger(__name__) @@ -465,17 +457,6 @@ async def generate_async( finally: await request_context.close() - def _validate_streaming_with_output_rails(self) -> None: - if len(self.config.rails.output.flows) > 0 and ( - not self.config.rails.output.streaming or not self.config.rails.output.streaming.enabled - ): - raise StreamingNotSupportedError( - "stream_async() cannot be used when output rails are configured but " - "rails.output.streaming.enabled is False. Either set " - "rails.output.streaming.enabled to True in your configuration, or use " - "generate_async() instead of stream_async()." - ) - @overload def stream_async( self, @@ -511,89 +492,18 @@ def stream_async( include_generation_metadata: Optional[bool] = None, ) -> AsyncIterator[Union[str, dict]]: """Simplified interface for getting directly the streamed tokens from the LLM.""" - - if include_generation_metadata is not None: - warnings.warn( - "include_generation_metadata is deprecated, use include_metadata instead. " - "It will be removed in version 0.22.0.", - DeprecationWarning, - stacklevel=2, - ) - include_metadata = include_generation_metadata - - self._validate_streaming_with_output_rails() - self._validate_public_state(state) - # if an external generator is provided, use it directly - if generator: - if self.config.rails.output.streaming and self.config.rails.output.streaming.enabled: - return self._run_output_rails_in_streaming( - streaming_handler=generator, - output_rails_streaming_config=self.config.rails.output.streaming, - messages=messages, - prompt=prompt, - ) - else: - return generator - - self._explain_info = self._ensure_explain_info() - - streaming_handler = StreamingHandler(include_metadata=include_metadata) - - # Create a properly managed task with exception handling - async def _generation_task(): - try: - await self.generate_async( - prompt=prompt, - messages=messages, - streaming_handler=streaming_handler, - options=options, - state=state, - ) - except Exception as e: - # If an exception occurs during generation, push it to the streaming handler as a json string - # This ensures the streaming pipeline is properly terminated - log.error(f"Error in generation task: {e}", exc_info=True) - error_message = str(e) - error_dict = extract_error_json(error_message) - error_payload = json.dumps(error_dict) - await streaming_handler.push_chunk(error_payload) - await streaming_handler.push_chunk(END_OF_STREAM) # type: ignore - - task = asyncio.create_task(_generation_task()) - - # Store task reference to prevent garbage collection and ensure proper cleanup - if not hasattr(self, "_active_tasks"): - self._active_tasks = set() - self._active_tasks.add(task) - - # Clean up task when it's done - def task_done_callback(task): - self._active_tasks.discard(task) - - task.add_done_callback(task_done_callback) - - # when we have output rails we wrap the streaming handler - # if len(self.config.rails.output.flows) > 0: - # - if self.config.rails.output.streaming and self.config.rails.output.streaming.enabled: - base_iterator = self._run_output_rails_in_streaming( - streaming_handler=streaming_handler, - output_rails_streaming_config=self.config.rails.output.streaming, - messages=messages, - prompt=prompt, - ) - else: - base_iterator = streaming_handler - - async def wrapped_iterator(): - try: - async for chunk in base_iterator: - if chunk is not None: - yield chunk - finally: - await task - - return wrapped_iterator() + validate_streaming_with_output_rails(self.config) + validate_public_state(self.config, state) + return generation_token_stream( + self, + prompt=prompt, + messages=messages, + options=options, + state=state, + include_metadata=include_metadata, + generator=generator, + include_generation_metadata=include_generation_metadata, + ) def generate( self, @@ -891,247 +801,15 @@ async def _run_output_rails_in_streaming( 2. Runs sequential (parallel for colang 2.0 in future) flows for each chunk. 3. Yields the chunk if not blocked, or STOP if blocked. """ - - def _get_last_context_message( - messages: Optional[List[dict]] = None, - ) -> dict: - if messages is None: - return {} - - for message in reversed(messages): - if message.get("role") == "context": - return message - return {} - - def _get_latest_user_message( - messages: Optional[List[dict]] = None, - ) -> str: - if messages is None: - return "" - for message in reversed(messages): - if message.get("role") == "user": - return message.get("content", "") - return "" - - def _prepare_context_for_parallel_rails( - chunk_str: str, - prompt: Optional[str] = None, - messages: Optional[List[dict]] = None, - ) -> dict: - """Prepare context for parallel rails execution.""" - context_message = _get_last_context_message(messages) - user_message = prompt or _get_latest_user_message(messages) - - context = { - "user_message": user_message, - "bot_message": chunk_str, - } - - if context_message: - context.update(context_message["content"]) - - return context - - def _create_events_for_chunk(chunk_str: str, context: dict) -> List[dict]: - """Create events for running output rails on a chunk.""" - return [ - {"type": "ContextUpdate", "data": context}, - {"type": "BotMessage", "text": chunk_str}, - ] - - def _prepare_params( - flow_id: str, - action_name: str, - bot_response_chunk: str, - prompt: Optional[str] = None, - messages: Optional[List[dict]] = None, - action_params: Dict[str, Any] = {}, + async for chunk in run_output_rails_in_streaming( + self, + streaming_handler=streaming_handler, + output_rails_streaming_config=output_rails_streaming_config, + prompt=prompt, + messages=messages, + stream_first=stream_first, ): - context_message = _get_last_context_message(messages) - user_message = prompt or _get_latest_user_message(messages) - - context = { - "user_message": user_message, - "bot_message": bot_response_chunk, - } - - if context_message: - context.update(context_message["content"]) - - model_name = flow_id.split("$")[-1].split("=")[-1].strip('"') - - # Resolve $bot_message / $user_message into a new dict. action_params - # is the shared flow config (reused across chunks and requests) and - # must not be mutated in place. - resolved_params = dict(action_params or {}) - for key, value in resolved_params.items(): - if value == "$bot_message": - resolved_params[key] = bot_response_chunk - elif value == "$user_message": - resolved_params[key] = user_message - - return { - # TODO:: are there other context variables that need to be passed? - # passing events to compute context was not successful - # context var failed due to different context - "context": context, - "llm_task_manager": self.runtime.llm_task_manager, - "config": self.config, - "model_name": model_name, - "llms": self.runtime.registered_action_params.get("llms", {}), - "llm": self.runtime.registered_action_params.get(f"{action_name}_llm", self.llm), - **resolved_params, - } - - buffer_strategy = get_buffer_strategy(output_rails_streaming_config) - output_rails_flows_id = self.config.rails.output.flows - stream_first = stream_first or output_rails_streaming_config.stream_first - get_action_details = partial(get_action_details_from_flow_id, flows=self.config.flows) - - parallel_mode = getattr(self.config.rails.output, "parallel", False) - - async for chunk_batch in buffer_strategy(streaming_handler): - user_output_chunks = chunk_batch.user_output_chunks - # format processing_context for output rails processing (needs full context) - bot_response_chunk = buffer_strategy.format_chunks(chunk_batch.processing_context) - - # check if user_output_chunks is a list of individual chunks - # or if it's a JSON string, by convention this means an error occurred and the error dict is stored as a JSON - if not isinstance(user_output_chunks, list): - try: - json.loads(user_output_chunks) - yield user_output_chunks - return - except (json.JSONDecodeError, TypeError): - # if it's not JSON, treat it as empty list - user_output_chunks = [] - - if stream_first: - # yield the individual chunks directly from the buffer strategy - for chunk in user_output_chunks: - yield chunk - - if parallel_mode: - try: - context = _prepare_context_for_parallel_rails(bot_response_chunk, prompt, messages) - events = _create_events_for_chunk(bot_response_chunk, context) - - flows_with_params = {} - for flow_id in output_rails_flows_id: - action_name, action_params = get_action_details(flow_id) - params = _prepare_params( - flow_id=flow_id, - action_name=action_name, - bot_response_chunk=bot_response_chunk, - prompt=prompt, - messages=messages, - action_params=action_params, - ) - flows_with_params[flow_id] = { - "action_name": action_name, - "params": params, - } - - result_tuple = await self.runtime.action_dispatcher.execute_action( - "run_output_rails_in_parallel_streaming", - { - "flows_with_params": flows_with_params, - "events": events, - }, - ) - - # ActionDispatcher.execute_action always returns (result, status) - result, status = result_tuple - - if status != "success": - log.error(f"Parallel rails execution failed with status: {status}") - # continue processing the chunk even if rails fail - pass - else: - # if there are any stop events, content was blocked or internal error occurred - result_events = getattr(result, "events", None) - if result_events: - # extract the flow info from the first stop event - stop_event = result_events[0] - blocked_flow = stop_event.get("flow_id", "output rails") - error_type = stop_event.get("error_type") - - if error_type == "internal_error": - error_message = stop_event.get("error_message", "Unknown error") - reason = f"Internal error in {blocked_flow} rail: {error_message}" - error_code = "rail_execution_failure" - error_type = "internal_error" - else: - reason = f"Blocked by {blocked_flow} rails." - error_code = "content_blocked" - error_type = "guardrails_violation" - - error_data = { - "error": { - "message": reason, - "type": error_type, - "param": blocked_flow, - "code": error_code, - } - } - yield json.dumps(error_data) - return - - except Exception as e: - log.error(f"Error in parallel rail execution: {e}") - # don't block the stream for rail execution errors - # continue processing the chunk - pass - - # update explain info for parallel mode - self._explain_info = self._ensure_explain_info() - - else: - for flow_id in output_rails_flows_id: - action_name, action_params = get_action_details(flow_id) - - params = _prepare_params( - flow_id=flow_id, - action_name=action_name, - bot_response_chunk=bot_response_chunk, - prompt=prompt, - messages=messages, - action_params=action_params, - ) - - result = await self.runtime.action_dispatcher.execute_action(action_name, params) - self._explain_info = self._ensure_explain_info() - - action_func = self.runtime.action_dispatcher.get_action(action_name) - - # Use the mapping to decide if the result indicates blocked content. - if is_output_blocked(result, action_func): - reason = f"Blocked by {flow_id} rails." - - # return the error as a plain JSON string (not in SSE format) - # NOTE: When integrating with the OpenAI Python client, the server code should: - # 1. detect this JSON error object in the stream - # 2. terminate the stream - # 3. format the error following OpenAI's SSE format - # the OpenAI client will then properly raise an APIError with this error message - - error_data = { - "error": { - "message": reason, - "type": "guardrails_violation", - "param": flow_id, - "code": "content_blocked", - } - } - - # return as plain JSON: the server should detect this JSON and convert it to an HTTP error - yield json.dumps(error_data) - return - - if not stream_first: - # yield the individual chunks directly from the buffer strategy - for chunk in user_output_chunks: - yield chunk + yield chunk def _determine_rails_from_messages(messages: List[dict]) -> Optional[dict]: diff --git a/nemoguardrails/rails/llm/streaming/__init__.py b/nemoguardrails/rails/llm/streaming/__init__.py new file mode 100644 index 0000000000..412cc58a1f --- /dev/null +++ b/nemoguardrails/rails/llm/streaming/__init__.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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.rails.llm.streaming.generation_stream import ( + GenerationStreamRails, + generation_token_stream, +) +from nemoguardrails.rails.llm.streaming.streaming_output_rails import ( + StreamingOutputActionDispatcher, + StreamingOutputRails, + StreamingOutputRuntime, + run_output_rails_in_streaming, +) + +__all__ = [ + "GenerationStreamRails", + "StreamingOutputActionDispatcher", + "StreamingOutputRails", + "StreamingOutputRuntime", + "generation_token_stream", + "run_output_rails_in_streaming", +] diff --git a/nemoguardrails/rails/llm/streaming/generation_stream.py b/nemoguardrails/rails/llm/streaming/generation_stream.py new file mode 100644 index 0000000000..a560691fa9 --- /dev/null +++ b/nemoguardrails/rails/llm/streaming/generation_stream.py @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 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. + +"""Generation token stream lifecycle helpers.""" + +import asyncio +import json +import logging +import warnings +from typing import Any, AsyncIterator, Optional, Protocol, Union, cast + +from nemoguardrails.exceptions import StreamingNotSupportedError +from nemoguardrails.rails.llm.options import GenerationOptions +from nemoguardrails.rails.llm.streaming.streaming_output_rails import ( + StreamingOutputRails, + run_output_rails_in_streaming, +) +from nemoguardrails.streaming import END_OF_STREAM, StreamingHandler +from nemoguardrails.utils import extract_error_json + +log = logging.getLogger(__name__) + +__all__ = [ + "GenerationStreamRails", + "generation_token_stream", + "validate_streaming_with_output_rails", +] + + +class GenerationStreamRails(StreamingOutputRails, Protocol): + async def generate_async( + self, + *, + prompt: Optional[str] = None, + messages: Optional[list[dict]] = None, + options: Optional[Union[dict, GenerationOptions]] = None, + state: Optional[Any] = None, + streaming_handler: Optional[StreamingHandler] = None, + ) -> object: ... + + +def validate_streaming_with_output_rails(config: Any) -> None: + if len(config.rails.output.flows) > 0 and ( + not config.rails.output.streaming or not config.rails.output.streaming.enabled + ): + raise StreamingNotSupportedError( + "stream_async() cannot be used when output rails are configured but " + "rails.output.streaming.enabled is False. Either set " + "rails.output.streaming.enabled to True in your configuration, or use " + "generate_async() instead of stream_async()." + ) + + +def generation_token_stream( + rails: GenerationStreamRails, + *, + prompt: Optional[str] = None, + messages: Optional[list[dict]] = None, + options: Optional[Union[dict, GenerationOptions]] = None, + state: Optional[Any] = None, + include_metadata: Optional[bool] = False, + generator: Optional[AsyncIterator[str]] = None, + include_generation_metadata: Optional[bool] = None, +) -> AsyncIterator[Union[str, dict]]: + """Return the token stream for a generation request.""" + if include_generation_metadata is not None: + warnings.warn( + "include_generation_metadata is deprecated, use include_metadata instead. " + "It will be removed in version 0.22.0.", + DeprecationWarning, + stacklevel=2, + ) + include_metadata = include_generation_metadata + + validate_streaming_with_output_rails(rails.config) + + if generator: + if rails.config.rails.output.streaming and rails.config.rails.output.streaming.enabled: + return run_output_rails_in_streaming( + rails, + streaming_handler=generator, + output_rails_streaming_config=rails.config.rails.output.streaming, + messages=messages, + prompt=prompt, + ) + return generator + + rails._explain_info = rails._ensure_explain_info() + + streaming_handler = StreamingHandler(include_metadata=include_metadata) + + async def _generation_task(): + try: + await rails.generate_async( + prompt=prompt, + messages=messages, + streaming_handler=streaming_handler, + options=options, + state=state, + ) + except Exception as e: + # If an exception occurs during generation, push it to the streaming + # handler as a json string. This ensures the streaming pipeline is + # properly terminated. + log.error(f"Error in generation task: {e}", exc_info=True) + error_message = str(e) + error_dict = extract_error_json(error_message) + error_payload = json.dumps(error_dict) + await streaming_handler.push_chunk(error_payload) + await streaming_handler.push_chunk(END_OF_STREAM) # type: ignore + + task = asyncio.create_task(_generation_task()) + _track_generation_task(rails, task) + + output_rail_streaming_enabled = bool( + rails.config.rails.output.streaming and rails.config.rails.output.streaming.enabled + ) + if output_rail_streaming_enabled: + base_iterator = run_output_rails_in_streaming( + rails, + streaming_handler=streaming_handler, + output_rails_streaming_config=rails.config.rails.output.streaming, + messages=messages, + prompt=prompt, + ) + else: + base_iterator = streaming_handler + + async def wrapped_iterator(): + try: + async for chunk in base_iterator: + if chunk is not None: + yield chunk + finally: + await task + + return wrapped_iterator() + + +def _track_generation_task(rails: GenerationStreamRails, task: asyncio.Task) -> None: + """Track background stream tasks so they are not garbage collected.""" + task_holder = cast(Any, rails) + if not hasattr(task_holder, "_active_tasks"): + task_holder._active_tasks = set() + task_holder._active_tasks.add(task) + + def task_done_callback(task): + task_holder._active_tasks.discard(task) + + task.add_done_callback(task_done_callback) diff --git a/nemoguardrails/rails/llm/streaming/streaming_output_rails.py b/nemoguardrails/rails/llm/streaming/streaming_output_rails.py new file mode 100644 index 0000000000..0afa0181b3 --- /dev/null +++ b/nemoguardrails/rails/llm/streaming/streaming_output_rails.py @@ -0,0 +1,378 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 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. + +"""Buffered output rail streaming helpers.""" + +import json +import logging +from functools import partial +from typing import Any, AsyncIterator, Callable, Dict, List, Optional, Protocol + +from nemoguardrails.actions.output_mapping import is_output_blocked +from nemoguardrails.rails.llm.buffer import get_buffer_strategy +from nemoguardrails.rails.llm.config import OutputRailsStreamingConfig +from nemoguardrails.rails.llm.utils import get_action_details_from_flow_id + +log = logging.getLogger(__name__) + +__all__ = [ + "StreamingOutputActionDispatcher", + "StreamingOutputRails", + "StreamingOutputRuntime", + "run_output_rails_in_streaming", +] + + +class StreamingOutputActionDispatcher(Protocol): + async def execute_action(self, action_name: str, params: Dict[str, Any]) -> Any: ... + + def get_action(self, name: str, /) -> Any: ... + + +class StreamingOutputRuntime(Protocol): + @property + def action_dispatcher(self) -> StreamingOutputActionDispatcher: ... + + @property + def llm_task_manager(self) -> Any: ... + + @property + def registered_action_params(self) -> Dict[str, Any]: ... + + +class StreamingOutputRails(Protocol): + _explain_info: Any + + @property + def config(self) -> Any: ... + + @property + def runtime(self) -> StreamingOutputRuntime: ... + + @property + def llm(self) -> Any: ... + + def _ensure_explain_info(self) -> Any: ... + + +async def run_output_rails_in_streaming( + rails: StreamingOutputRails, + streaming_handler: AsyncIterator[str], + output_rails_streaming_config: OutputRailsStreamingConfig, + prompt: Optional[str] = None, + messages: Optional[List[dict]] = None, + stream_first: Optional[bool] = None, +) -> AsyncIterator[str]: + """ + 1. Buffers tokens from 'streaming_handler' via BufferStrategy. + 2. Runs sequential (parallel for colang 2.0 in future) flows for each chunk. + 3. Yields the chunk if not blocked, or STOP if blocked. + """ + buffer_strategy = get_buffer_strategy(output_rails_streaming_config) + output_rails_flows_id = rails.config.rails.output.flows + stream_first = stream_first or output_rails_streaming_config.stream_first + get_action_details = partial(get_action_details_from_flow_id, flows=rails.config.flows) + + parallel_mode = getattr(rails.config.rails.output, "parallel", False) + + async for chunk_batch in buffer_strategy(streaming_handler): + user_output_chunks = chunk_batch.user_output_chunks + # format processing_context for output rails processing (needs full context) + bot_response_chunk = buffer_strategy.format_chunks(chunk_batch.processing_context) + + # check if user_output_chunks is a list of individual chunks + # or if it's a JSON string, by convention this means an error occurred and the error dict is stored as a JSON + if not isinstance(user_output_chunks, list): + try: + json.loads(user_output_chunks) + yield user_output_chunks + return + except (json.JSONDecodeError, TypeError): + # if it's not JSON, treat it as empty list + user_output_chunks = [] + + if stream_first: + # yield the individual chunks directly from the buffer strategy + for chunk in user_output_chunks: + yield chunk + + if parallel_mode: + async for chunk in _run_parallel_output_rails( + rails=rails, + output_rails_flows_id=output_rails_flows_id, + get_action_details=get_action_details, + bot_response_chunk=bot_response_chunk, + prompt=prompt, + messages=messages, + ): + yield chunk + return + else: + async for chunk in _run_sequential_output_rails( + rails=rails, + output_rails_flows_id=output_rails_flows_id, + get_action_details=get_action_details, + bot_response_chunk=bot_response_chunk, + prompt=prompt, + messages=messages, + ): + yield chunk + return + + if not stream_first: + # yield the individual chunks directly from the buffer strategy + for chunk in user_output_chunks: + yield chunk + + +async def _run_parallel_output_rails( + *, + rails: StreamingOutputRails, + output_rails_flows_id: List[str], + get_action_details: Callable[[str], tuple[str, Dict[str, Any]]], + bot_response_chunk: str, + prompt: Optional[str], + messages: Optional[List[dict]], +) -> AsyncIterator[str]: + try: + context = _prepare_context_for_parallel_rails(bot_response_chunk, prompt, messages) + events = _create_events_for_chunk(bot_response_chunk, context) + + flows_with_params = {} + for flow_id in output_rails_flows_id: + action_name, action_params = get_action_details(flow_id) + params = _prepare_params( + rails=rails, + flow_id=flow_id, + action_name=action_name, + bot_response_chunk=bot_response_chunk, + prompt=prompt, + messages=messages, + action_params=action_params, + ) + flows_with_params[flow_id] = { + "action_name": action_name, + "params": params, + } + + result_tuple = await rails.runtime.action_dispatcher.execute_action( + "run_output_rails_in_parallel_streaming", + { + "flows_with_params": flows_with_params, + "events": events, + }, + ) + + # ActionDispatcher.execute_action always returns (result, status) + result, status = result_tuple + + if status != "success": + log.error(f"Parallel rails execution failed with status: {status}") + else: + # if there are any stop events, content was blocked or internal error occurred + result_events = getattr(result, "events", None) + if result_events: + yield _parallel_stop_error(result_events[0]) + return + + except Exception as e: + log.error(f"Error in parallel rail execution: {e}") + + # update explain info for parallel mode + rails._explain_info = rails._ensure_explain_info() + + +async def _run_sequential_output_rails( + *, + rails: StreamingOutputRails, + output_rails_flows_id: List[str], + get_action_details: Callable[[str], tuple[str, Dict[str, Any]]], + bot_response_chunk: str, + prompt: Optional[str], + messages: Optional[List[dict]], +) -> AsyncIterator[str]: + for flow_id in output_rails_flows_id: + action_name, action_params = get_action_details(flow_id) + + params = _prepare_params( + rails=rails, + flow_id=flow_id, + action_name=action_name, + bot_response_chunk=bot_response_chunk, + prompt=prompt, + messages=messages, + action_params=action_params, + ) + + action_result = await rails.runtime.action_dispatcher.execute_action(action_name, params) + rails._explain_info = rails._ensure_explain_info() + + # Use the mapping to decide if the result indicates blocked content. + action_func = rails.runtime.action_dispatcher.get_action(action_name) + if is_output_blocked(action_result, action_func): + yield _blocked_output_error(flow_id) + return + + +def _get_last_context_message( + messages: Optional[List[dict]] = None, +) -> dict: + if messages is None: + return {} + + for message in reversed(messages): + if message.get("role") == "context": + return message + return {} + + +def _get_latest_user_message( + messages: Optional[List[dict]] = None, +) -> str: + if messages is None: + return "" + for message in reversed(messages): + if message.get("role") == "user": + return message.get("content", "") + return "" + + +def _prepare_context_for_parallel_rails( + chunk_str: str, + prompt: Optional[str] = None, + messages: Optional[List[dict]] = None, +) -> dict: + """Prepare context for parallel rails execution.""" + context_message = _get_last_context_message(messages) + user_message = prompt or _get_latest_user_message(messages) + + context = { + "user_message": user_message, + "bot_message": chunk_str, + } + + if context_message: + context.update(context_message["content"]) + + return context + + +def _create_events_for_chunk(chunk_str: str, context: dict) -> List[dict]: + """Create events for running output rails on a chunk.""" + return [ + {"type": "ContextUpdate", "data": context}, + {"type": "BotMessage", "text": chunk_str}, + ] + + +def _prepare_params( + *, + rails: StreamingOutputRails, + flow_id: str, + action_name: str, + bot_response_chunk: str, + prompt: Optional[str] = None, + messages: Optional[List[dict]] = None, + action_params: Dict[str, Any], +): + context_message = _get_last_context_message(messages) + user_message = prompt or _get_latest_user_message(messages) + + context = { + "user_message": user_message, + "bot_message": bot_response_chunk, + } + + if context_message: + context.update(context_message["content"]) + + model_name = flow_id.split("$")[-1].split("=")[-1].strip('"') + + resolved_params = dict(action_params or {}) + for key, value in resolved_params.items(): + if value == "$bot_message": + resolved_params[key] = bot_response_chunk + elif value == "$user_message": + resolved_params[key] = user_message + + return { + # TODO:: are there other context variables that need to be passed? + # passing events to compute context was not successful + # context var failed due to different context + "context": context, + "llm_task_manager": rails.runtime.llm_task_manager, + "config": rails.config, + "model_name": model_name, + "llms": rails.runtime.registered_action_params.get("llms", {}), + "llm": rails.runtime.registered_action_params.get(f"{action_name}_llm", rails.llm), + **resolved_params, + } + + +def _blocked_output_error(flow_id: str) -> str: + reason = f"Blocked by {flow_id} rails." + + # return the error as a plain JSON string (not in SSE format) + # NOTE: When integrating with the OpenAI Python client, the server code should: + # 1. detect this JSON error object in the stream + # 2. terminate the stream + # 3. format the error following OpenAI's SSE format + # the OpenAI client will then properly raise an APIError with this error message + error_data = { + "error": { + "message": reason, + "type": "guardrails_violation", + "param": flow_id, + "code": "content_blocked", + } + } + + # return as plain JSON: the server should detect this JSON and convert it to an HTTP error + return json.dumps(error_data) + + +def _internal_rail_error(flow_id: str, error_message: str) -> str: + error_data = { + "error": { + "message": f"Internal error in {flow_id} rail: {error_message}", + "type": "internal_error", + "param": flow_id, + "code": "rail_execution_failure", + } + } + return json.dumps(error_data) + + +def _parallel_stop_error(stop_event: dict) -> str: + blocked_flow = stop_event.get("flow_id", "output rails") + error_type = stop_event.get("error_type") + + if error_type == "internal_error": + error_message = stop_event.get("error_message", "Unknown error") + return _internal_rail_error(blocked_flow, error_message) + else: + reason = f"Blocked by {blocked_flow} rails." + error_code = "content_blocked" + error_type = "guardrails_violation" + + error_data = { + "error": { + "message": reason, + "type": error_type, + "param": blocked_flow, + "code": error_code, + } + } + return json.dumps(error_data) diff --git a/tests/rails/llm/test_generation_stream.py b/tests/rails/llm/test_generation_stream.py new file mode 100644 index 0000000000..a08b433c91 --- /dev/null +++ b/tests/rails/llm/test_generation_stream.py @@ -0,0 +1,364 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 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. + +import asyncio +import json +from types import SimpleNamespace +from typing import Any, AsyncIterator, Optional, cast + +import pytest + +from nemoguardrails.exceptions import StreamingNotSupportedError +from nemoguardrails.rails.llm.config import RailsConfig +from nemoguardrails.rails.llm.llmrails import LLMRails +from nemoguardrails.rails.llm.streaming import generation_stream +from nemoguardrails.rails.llm.streaming.generation_stream import ( + generation_token_stream, + validate_streaming_with_output_rails, +) +from nemoguardrails.streaming import END_OF_STREAM + + +class FakeStreamingHandler: + instances = [] + + def __init__(self, include_metadata: Optional[bool] = False): + self.include_metadata = include_metadata + self.queue = asyncio.Queue() + self.pushed_chunks = [] + self.__class__.instances.append(self) + + async def push_chunk(self, chunk): + self.pushed_chunks.append(chunk) + await self.queue.put(chunk) + + def __aiter__(self): + return self + + async def __anext__(self): + chunk = await self.queue.get() + if chunk == END_OF_STREAM: + raise StopAsyncIteration + return chunk + + +class FakeRuntime: + def __init__(self): + self.action_dispatcher: Any = None + self.llm_task_manager: Any = None + self.registered_action_params: dict[str, Any] = {} + + +class FakeRails: + def __init__(self, *, streaming=None, output_flows=None, generate_error: bool = False): + self.config = SimpleNamespace( + rails=SimpleNamespace( + output=SimpleNamespace( + flows=output_flows or [], + streaming=streaming, + ) + ) + ) + self.generate_error = generate_error + self.generate_calls = [] + self.output_rail_calls = [] + self._explain_info = None + self.runtime = FakeRuntime() + self.llm = None + + def _ensure_explain_info(self): + return {"explain": True} + + async def generate_async( + self, + *, + prompt=None, + messages=None, + streaming_handler=None, + options=None, + state=None, + ): + self.generate_calls.append( + { + "prompt": prompt, + "messages": messages, + "streaming_handler": streaming_handler, + "options": options, + "state": state, + } + ) + if self.generate_error: + raise RuntimeError('Error code: 500 - {"error": {"message": "boom"}}') + assert streaming_handler is not None + await streaming_handler.push_chunk(None) + await streaming_handler.push_chunk("hello") + await streaming_handler.push_chunk(END_OF_STREAM) + + +class SlowFakeRails(FakeRails): + def __init__(self): + super().__init__() + self.cancelled = asyncio.Event() + self.release = asyncio.Event() + + async def generate_async( + self, + *, + prompt=None, + messages=None, + streaming_handler=None, + options=None, + state=None, + ): + self.generate_calls.append( + { + "prompt": prompt, + "messages": messages, + "streaming_handler": streaming_handler, + "options": options, + "state": state, + } + ) + assert streaming_handler is not None + await streaming_handler.push_chunk("hello") + try: + await self.release.wait() + except asyncio.CancelledError: + self.cancelled.set() + raise + await streaming_handler.push_chunk(END_OF_STREAM) + + +class BlockingOutputRailsFakeRails(SlowFakeRails): + def __init__(self): + super().__init__() + self.config.rails.output.streaming = SimpleNamespace(enabled=True) + + +async def _collect(iterator): + return [chunk async for chunk in iterator] + + +async def _tokens(*chunks) -> AsyncIterator[str]: + for chunk in chunks: + yield chunk + + +@pytest.fixture(autouse=True) +def fake_streaming_handler_and_output_rails(monkeypatch): + FakeStreamingHandler.instances = [] + monkeypatch.setattr(generation_stream, "StreamingHandler", FakeStreamingHandler) + + def fake_run_output_rails_in_streaming( + rails, + streaming_handler, + output_rails_streaming_config, + prompt=None, + messages=None, + stream_first=None, + ): + rails.output_rail_calls.append( + { + "streaming_handler": streaming_handler, + "output_rails_streaming_config": output_rails_streaming_config, + "messages": messages, + "prompt": prompt, + "stream_first": stream_first, + } + ) + + async def _wrapped(): + async for chunk in streaming_handler: + if chunk is None: + continue + if isinstance(rails, BlockingOutputRailsFakeRails): + yield json.dumps( + { + "error": { + "message": "Blocked by self check output rails.", + "type": "guardrails_violation", + "param": "self check output", + "code": "content_blocked", + } + } + ) + return + yield f"wrapped:{chunk}" + + return _wrapped() + + monkeypatch.setattr( + generation_stream, + "run_output_rails_in_streaming", + fake_run_output_rails_in_streaming, + ) + + +@pytest.mark.asyncio +async def test_generation_token_stream_runs_generation_task_and_tracks_it(): + rails = FakeRails() + + stream = generation_token_stream( + rails, + messages=[{"role": "user", "content": "Hi"}], + include_metadata=True, + ) + chunks = await _collect(stream) + await asyncio.sleep(0) + + assert chunks == ["hello"] + assert rails._explain_info == {"explain": True} + assert FakeStreamingHandler.instances[0].include_metadata is True + assert rails.generate_calls[0]["messages"] == [{"role": "user", "content": "Hi"}] + assert rails.generate_calls[0]["streaming_handler"] is FakeStreamingHandler.instances[0] + assert cast(set, getattr(rails, "_active_tasks")) == set() + + +@pytest.mark.asyncio +async def test_generation_token_stream_honors_deprecated_metadata_flag(): + rails = FakeRails() + + with pytest.warns(DeprecationWarning, match="include_generation_metadata is deprecated"): + stream = generation_token_stream(rails, include_generation_metadata=True) + + await _collect(stream) + + assert FakeStreamingHandler.instances[0].include_metadata is True + + +@pytest.mark.asyncio +async def test_generation_token_stream_returns_external_generator_without_output_rails(): + rails = FakeRails() + generator = _tokens("a", "b") + + stream = generation_token_stream(rails, generator=generator) + + assert stream is generator + assert await _collect(stream) == ["a", "b"] + + +@pytest.mark.asyncio +async def test_generation_token_stream_wraps_external_generator_when_output_rail_streaming_enabled(): + streaming_config = SimpleNamespace(enabled=True) + rails = FakeRails(streaming=streaming_config) + generator = _tokens("a") + + stream = generation_token_stream( + rails, + generator=generator, + messages=[{"role": "user", "content": "Hi"}], + prompt=None, + ) + + assert await _collect(stream) == ["wrapped:a"] + assert rails.output_rail_calls == [ + { + "streaming_handler": generator, + "output_rails_streaming_config": streaming_config, + "messages": [{"role": "user", "content": "Hi"}], + "prompt": None, + "stream_first": None, + } + ] + + +@pytest.mark.asyncio +async def test_generation_token_stream_wraps_internal_handler_when_output_rail_streaming_enabled(): + rails = FakeRails(streaming=SimpleNamespace(enabled=True)) + + stream = generation_token_stream(rails, prompt="Hi") + + assert await _collect(stream) == ["wrapped:hello"] + assert rails.output_rail_calls[0]["streaming_handler"] is FakeStreamingHandler.instances[0] + assert rails.generate_calls[0]["prompt"] == "Hi" + + +def test_validate_streaming_with_output_rails_rejects_non_streaming_output_rails(): + config = SimpleNamespace( + rails=SimpleNamespace( + output=SimpleNamespace( + flows=["self check output"], + streaming=SimpleNamespace(enabled=False), + ) + ) + ) + + with pytest.raises(StreamingNotSupportedError, match="stream_async\\(\\) cannot be used"): + validate_streaming_with_output_rails(config) + + +def test_stream_async_validates_streaming_support_before_state(): + config = RailsConfig(models=[], colang_version="1.0") + config.rails.output.flows = ["self check output"] + rails = LLMRails(config) + + with pytest.raises(StreamingNotSupportedError, match="stream_async\\(\\) cannot be used"): + rails.stream_async( + messages=[{"role": "user", "content": "Hi"}], + state={"foo": 1}, + ) + + +@pytest.mark.asyncio +async def test_generation_token_stream_pushes_error_payload_when_generation_fails(): + rails = FakeRails(generate_error=True) + + chunks = await _collect(generation_token_stream(rails, messages=[])) + + assert len(chunks) == 1 + assert json.loads(chunks[0]) == {"error": {"message": "boom"}} + assert FakeStreamingHandler.instances[0].pushed_chunks == [ + chunks[0], + END_OF_STREAM, + ] + + +@pytest.mark.asyncio +async def test_generation_token_stream_waits_for_generation_when_consumer_closes_early(): + rails = SlowFakeRails() + stream = generation_token_stream(rails, messages=[{"role": "user", "content": "Hi"}]) + + assert await stream.__anext__() == "hello" + close_task = asyncio.create_task(cast(Any, stream).aclose()) + await asyncio.sleep(0) + + assert not close_task.done() + assert not rails.cancelled.is_set() + + rails.release.set() + await asyncio.wait_for(close_task, timeout=1) + + assert not rails.cancelled.is_set() + assert cast(set, getattr(rails, "_active_tasks")) == set() + + +@pytest.mark.asyncio +async def test_generation_token_stream_waits_for_generation_after_output_rail_error(): + rails = BlockingOutputRailsFakeRails() + stream = generation_token_stream(rails, messages=[{"role": "user", "content": "Hi"}]) + + collect_task = asyncio.create_task(_collect(stream)) + await asyncio.sleep(0) + + assert not collect_task.done() + assert not rails.cancelled.is_set() + + rails.release.set() + chunks = await asyncio.wait_for(collect_task, timeout=1) + + assert len(chunks) == 1 + assert json.loads(chunks[0])["error"]["code"] == "content_blocked" + assert not rails.cancelled.is_set() + assert cast(set, getattr(rails, "_active_tasks")) == set() diff --git a/tests/rails/llm/test_streaming_output_rails.py b/tests/rails/llm/test_streaming_output_rails.py new file mode 100644 index 0000000000..1b42ec3d49 --- /dev/null +++ b/tests/rails/llm/test_streaming_output_rails.py @@ -0,0 +1,480 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023-2026 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. + +import asyncio +import json +from types import SimpleNamespace +from typing import Any, AsyncIterator, cast + +import pytest + +from nemoguardrails.rails.llm.buffer import ChunkBatch +from nemoguardrails.rails.llm.config import OutputRailsStreamingConfig +from nemoguardrails.rails.llm.streaming import streaming_output_rails +from nemoguardrails.rails.llm.streaming.streaming_output_rails import ( + run_output_rails_in_streaming, +) + + +class FakeBufferStrategy: + def __init__(self, batches: list[ChunkBatch]): + self.batches = batches + + async def __call__(self, streaming_handler: AsyncIterator[str]) -> AsyncIterator[ChunkBatch]: + del streaming_handler + for batch in self.batches: + yield batch + + def format_chunks(self, chunks: list[str]) -> str: + return "".join(chunks) + + +class FakeActionDispatcher: + def __init__( + self, + *, + result: Any = True, + status: str = "success", + parallel_result: Any = None, + parallel_status: str = "success", + parallel_exception: Exception | None = None, + ): + self.result = result + self.status = status + self.parallel_result = parallel_result + self.parallel_status = parallel_status + self.parallel_exception = parallel_exception + self.calls = [] + + async def execute_action(self, action_name: str, params: dict[str, Any]): + self.calls.append((action_name, params)) + + if action_name == "run_output_rails_in_parallel_streaming": + if self.parallel_exception: + raise self.parallel_exception + return self.parallel_result, self.parallel_status + + return self.result, self.status + + def get_action(self, action_name: str): + del action_name + return SimpleNamespace(action_meta={}) + + +class FakeRuntime: + action_dispatcher: FakeActionDispatcher + llm_task_manager: Any + registered_action_params: dict[str, Any] + + def __init__(self, dispatcher: FakeActionDispatcher): + self.action_dispatcher = dispatcher + self.llm_task_manager = "task-manager" + self.registered_action_params = { + "llms": {"main": "registered-llm"}, + "self_check_output_llm": "rail-llm", + } + + +class FakeRails: + def __init__(self, dispatcher: FakeActionDispatcher, *, parallel: bool = False): + self.config = SimpleNamespace( + flows=[], + rails=SimpleNamespace( + output=SimpleNamespace( + flows=["self check output"], + parallel=parallel, + ) + ), + ) + self.runtime = FakeRuntime(dispatcher) + self.llm = "main-llm" + self._explain_info = None + self.ensure_explain_info_calls = 0 + + def _ensure_explain_info(self): + self.ensure_explain_info_calls += 1 + return {"ensured": self.ensure_explain_info_calls} + + +async def _empty_stream() -> AsyncIterator[str]: + if False: + yield "" + + +async def _tracked_stream(closed: Any, *chunks: str) -> AsyncIterator[str]: + try: + for chunk in chunks: + yield chunk + await asyncio.sleep(0) + finally: + closed.set() + + +async def _collect(iterator: AsyncIterator[str]) -> list[str]: + return [chunk async for chunk in iterator] + + +def _patch_buffer_strategy(monkeypatch, batches: list[ChunkBatch]) -> FakeBufferStrategy: + strategy = FakeBufferStrategy(batches) + monkeypatch.setattr(streaming_output_rails, "get_buffer_strategy", lambda config: strategy) + return strategy + + +def _patch_action_details(monkeypatch): + def fake_get_action_details_from_flow_id(flow_id: str, flows: list): + del flow_id, flows + return ( + "self_check_output", + { + "bot_response": "$bot_message", + "user_input": "$user_message", + "literal": "value", + }, + ) + + monkeypatch.setattr( + streaming_output_rails, + "get_action_details_from_flow_id", + fake_get_action_details_from_flow_id, + ) + + +def _patch_shared_action_details(monkeypatch, action_params: dict): + def fake_get_action_details_from_flow_id(flow_id: str, flows: list): + del flow_id, flows + return "self_check_output", action_params + + monkeypatch.setattr( + streaming_output_rails, + "get_action_details_from_flow_id", + fake_get_action_details_from_flow_id, + ) + + +@pytest.mark.asyncio +async def test_output_rails_stream_runs_sequential_rails_before_yielding(monkeypatch): + _patch_buffer_strategy( + monkeypatch, + [ChunkBatch(processing_context=["He", "llo"], user_output_chunks=["He", "llo"])], + ) + _patch_action_details(monkeypatch) + dispatcher = FakeActionDispatcher(result=True) + rails = FakeRails(dispatcher) + + chunks = await _collect( + run_output_rails_in_streaming( + rails, + _empty_stream(), + OutputRailsStreamingConfig(stream_first=False), + messages=[ + {"role": "context", "content": {"account_type": "enterprise"}}, + {"role": "user", "content": "Hi"}, + ], + ) + ) + + assert chunks == ["He", "llo"] + assert dispatcher.calls[0][0] == "self_check_output" + params = dispatcher.calls[0][1] + assert params["bot_response"] == "Hello" + assert params["user_input"] == "Hi" + assert params["literal"] == "value" + assert params["context"] == { + "account_type": "enterprise", + "bot_message": "Hello", + "user_message": "Hi", + } + assert params["llm_task_manager"] == "task-manager" + assert params["llms"] == {"main": "registered-llm"} + assert params["llm"] == "rail-llm" + assert rails._explain_info == {"ensured": 1} + + +@pytest.mark.asyncio +async def test_output_rails_stream_does_not_mutate_action_param_templates(monkeypatch): + _patch_buffer_strategy( + monkeypatch, + [ + ChunkBatch(processing_context=["He"], user_output_chunks=["He"]), + ChunkBatch(processing_context=["llo"], user_output_chunks=["llo"]), + ], + ) + shared_action_params = { + "bot_response": "$bot_message", + "user_input": "$user_message", + } + _patch_shared_action_details(monkeypatch, shared_action_params) + dispatcher = FakeActionDispatcher(result=True) + rails = FakeRails(dispatcher) + + await _collect( + run_output_rails_in_streaming( + rails, + _empty_stream(), + OutputRailsStreamingConfig(stream_first=False), + messages=[{"role": "user", "content": "Hi"}], + ) + ) + + assert shared_action_params == { + "bot_response": "$bot_message", + "user_input": "$user_message", + } + assert dispatcher.calls[0][1]["bot_response"] == "He" + assert dispatcher.calls[1][1]["bot_response"] == "llo" + assert dispatcher.calls[0][1]["user_input"] == "Hi" + assert dispatcher.calls[1][1]["user_input"] == "Hi" + + +@pytest.mark.asyncio +async def test_output_rails_stream_passes_through_json_error_chunks(monkeypatch): + error_chunk = '{"error": {"message": "upstream failure"}}' + _patch_buffer_strategy( + monkeypatch, + [ + ChunkBatch( + processing_context=[], + user_output_chunks=cast(Any, error_chunk), + ) + ], + ) + _patch_action_details(monkeypatch) + dispatcher = FakeActionDispatcher() + rails = FakeRails(dispatcher) + + chunks = await _collect( + run_output_rails_in_streaming( + rails, + _empty_stream(), + OutputRailsStreamingConfig(stream_first=False), + ) + ) + + assert chunks == [error_chunk] + assert dispatcher.calls == [] + assert rails._explain_info is None + + +@pytest.mark.asyncio +async def test_output_rails_stream_yields_block_error_for_sequential_rails(monkeypatch): + _patch_buffer_strategy( + monkeypatch, + [ChunkBatch(processing_context=["blocked"], user_output_chunks=["blocked"])], + ) + _patch_action_details(monkeypatch) + dispatcher = FakeActionDispatcher(result=False) + rails = FakeRails(dispatcher) + + chunks = await _collect( + run_output_rails_in_streaming( + rails, + _empty_stream(), + OutputRailsStreamingConfig(stream_first=False), + ) + ) + + assert len(chunks) == 1 + assert json.loads(chunks[0]) == { + "error": { + "message": "Blocked by self check output rails.", + "type": "guardrails_violation", + "param": "self check output", + "code": "content_blocked", + } + } + assert rails._explain_info == {"ensured": 1} + + +@pytest.mark.asyncio +async def test_output_rails_stream_uses_config_stream_first_when_explicit_false(monkeypatch): + _patch_buffer_strategy( + monkeypatch, + [ChunkBatch(processing_context=["blocked"], user_output_chunks=["blocked"])], + ) + _patch_action_details(monkeypatch) + dispatcher = FakeActionDispatcher(result=False) + rails = FakeRails(dispatcher) + + chunks = await _collect( + run_output_rails_in_streaming( + rails, + _empty_stream(), + OutputRailsStreamingConfig(stream_first=True), + stream_first=False, + ) + ) + + assert len(chunks) == 2 + assert chunks[0] == "blocked" + assert json.loads(chunks[1])["error"]["code"] == "content_blocked" + + +@pytest.mark.asyncio +async def test_output_rails_stream_closes_external_generator_after_block(monkeypatch): + _patch_action_details(monkeypatch) + dispatcher = FakeActionDispatcher(result=False) + rails = FakeRails(dispatcher) + closed = asyncio.Event() + + chunks = await _collect( + run_output_rails_in_streaming( + rails, + _tracked_stream(closed, "blocked", "unused"), + OutputRailsStreamingConfig(stream_first=False, chunk_size=1, context_size=0), + ) + ) + await asyncio.wait_for(closed.wait(), timeout=1) + + assert len(chunks) == 1 + assert json.loads(chunks[0])["error"]["code"] == "content_blocked" + + +@pytest.mark.asyncio +async def test_output_rails_stream_closes_external_generator_when_consumer_closes(monkeypatch): + _patch_action_details(monkeypatch) + dispatcher = FakeActionDispatcher(result=True) + rails = FakeRails(dispatcher) + closed = asyncio.Event() + stream = run_output_rails_in_streaming( + rails, + _tracked_stream(closed, "first", "unused"), + OutputRailsStreamingConfig(stream_first=False, chunk_size=1, context_size=0), + ) + + assert await stream.__anext__() == "first" + await asyncio.wait_for(cast(Any, stream).aclose(), timeout=1) + await asyncio.wait_for(closed.wait(), timeout=1) + + +@pytest.mark.asyncio +async def test_output_rails_stream_continues_for_sequential_action_failure(monkeypatch): + _patch_buffer_strategy( + monkeypatch, + [ChunkBatch(processing_context=["unsafe"], user_output_chunks=["unsafe"])], + ) + _patch_action_details(monkeypatch) + dispatcher = FakeActionDispatcher(result=None, status="failed") + rails = FakeRails(dispatcher) + + chunks = await _collect( + run_output_rails_in_streaming( + rails, + _empty_stream(), + OutputRailsStreamingConfig(stream_first=False), + ) + ) + + assert chunks == ["unsafe"] + assert rails._explain_info == {"ensured": 1} + + +@pytest.mark.asyncio +async def test_output_rails_stream_yields_parallel_stop_event_error(monkeypatch): + _patch_buffer_strategy( + monkeypatch, + [ChunkBatch(processing_context=["He", "llo"], user_output_chunks=["He", "llo"])], + ) + _patch_action_details(monkeypatch) + dispatcher = FakeActionDispatcher( + parallel_result=SimpleNamespace( + events=[ + { + "flow_id": "self check output", + "error_type": "internal_error", + "error_message": "action failed", + } + ] + ) + ) + rails = FakeRails(dispatcher, parallel=True) + + chunks = await _collect( + run_output_rails_in_streaming( + rails, + _empty_stream(), + OutputRailsStreamingConfig(stream_first=False), + prompt="Hi", + messages=[{"role": "context", "content": {"account_type": "enterprise"}}], + ) + ) + + assert len(chunks) == 1 + assert json.loads(chunks[0]) == { + "error": { + "message": "Internal error in self check output rail: action failed", + "type": "internal_error", + "param": "self check output", + "code": "rail_execution_failure", + } + } + action_name, params = dispatcher.calls[0] + assert action_name == "run_output_rails_in_parallel_streaming" + assert params["events"] == [ + { + "type": "ContextUpdate", + "data": { + "account_type": "enterprise", + "bot_message": "Hello", + "user_message": "Hi", + }, + }, + {"type": "BotMessage", "text": "Hello"}, + ] + assert params["flows_with_params"]["self check output"]["action_name"] == "self_check_output" + assert params["flows_with_params"]["self check output"]["params"]["bot_response"] == "Hello" + + +@pytest.mark.asyncio +async def test_output_rails_stream_continues_for_parallel_action_failure(monkeypatch): + _patch_buffer_strategy( + monkeypatch, + [ChunkBatch(processing_context=["unsafe"], user_output_chunks=["unsafe"])], + ) + _patch_action_details(monkeypatch) + dispatcher = FakeActionDispatcher(parallel_status="failed") + rails = FakeRails(dispatcher, parallel=True) + + chunks = await _collect( + run_output_rails_in_streaming( + rails, + _empty_stream(), + OutputRailsStreamingConfig(stream_first=False), + ) + ) + + assert chunks == ["unsafe"] + assert rails._explain_info == {"ensured": 1} + + +@pytest.mark.asyncio +async def test_output_rails_stream_continues_for_parallel_exception(monkeypatch): + _patch_buffer_strategy( + monkeypatch, + [ChunkBatch(processing_context=["unsafe"], user_output_chunks=["unsafe"])], + ) + _patch_action_details(monkeypatch) + dispatcher = FakeActionDispatcher(parallel_exception=RuntimeError("parallel boom")) + rails = FakeRails(dispatcher, parallel=True) + + chunks = await _collect( + run_output_rails_in_streaming( + rails, + _empty_stream(), + OutputRailsStreamingConfig(stream_first=False), + ) + ) + + assert chunks == ["unsafe"] + assert rails._explain_info == {"ensured": 1}