From 43b4821a55368da01645ef92ec67448dbcd25129 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Tue, 9 Jun 2026 17:45:47 +0200 Subject: [PATCH 1/2] fix(streaming): isolate policy followups Keep streaming behavior changes separate from the LLMRails decomposition stack. --- .../rails/llm/streaming/generation_stream.py | 40 ++++++++++++++++- .../llm/streaming/streaming_output_rails.py | 35 ++++++++++++--- tests/rails/llm/test_generation_stream.py | 24 +++------- .../rails/llm/test_streaming_output_rails.py | 45 ++++++++++++++----- 4 files changed, 108 insertions(+), 36 deletions(-) diff --git a/nemoguardrails/rails/llm/streaming/generation_stream.py b/nemoguardrails/rails/llm/streaming/generation_stream.py index 091539caef..0ba6a65a10 100644 --- a/nemoguardrails/rails/llm/streaming/generation_stream.py +++ b/nemoguardrails/rails/llm/streaming/generation_stream.py @@ -127,16 +127,36 @@ async def _generation_task(): base_iterator = streaming_handler async def wrapped_iterator(): + stream_completed = False + terminal_output_rail_error = False try: async for chunk in base_iterator: if chunk is not None: + if output_rail_streaming_enabled and _is_stream_error_chunk(chunk): + terminal_output_rail_error = True yield chunk + stream_completed = True finally: - await task + await _finish_generation_task( + task, + stream_completed=stream_completed and not terminal_output_rail_error, + ) return wrapped_iterator() +def _is_stream_error_chunk(chunk: Union[str, dict]) -> bool: + if not isinstance(chunk, str): + return False + + try: + chunk_data = json.loads(chunk) + except json.JSONDecodeError: + return False + + return isinstance(chunk_data, dict) and "error" in chunk_data + + def _track_generation_task(rails: GenerationStreamSurface, task: asyncio.Task) -> None: """Track background stream tasks so they are not garbage collected.""" task_holder = cast(Any, rails) @@ -148,3 +168,21 @@ def task_done_callback(task): task_holder._active_tasks.discard(task) task.add_done_callback(task_done_callback) + + +async def _finish_generation_task( + task: asyncio.Task, + *, + stream_completed: bool, +) -> None: + """Finalize a background generation task for a stream consumer.""" + cancelled_by_stream_close = False + if not stream_completed and not task.done(): + task.cancel() + cancelled_by_stream_close = True + + try: + await task + except asyncio.CancelledError: + if not cancelled_by_stream_close: + raise diff --git a/nemoguardrails/rails/llm/streaming/streaming_output_rails.py b/nemoguardrails/rails/llm/streaming/streaming_output_rails.py index 7c59e261f3..afebf2896c 100644 --- a/nemoguardrails/rails/llm/streaming/streaming_output_rails.py +++ b/nemoguardrails/rails/llm/streaming/streaming_output_rails.py @@ -49,7 +49,8 @@ async def run_output_rails_in_streaming( """ 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 + if stream_first is None: + stream_first = 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) @@ -146,7 +147,10 @@ async def _run_parallel_output_rails( result, status = result_tuple if status != "success": - log.error(f"Parallel rails execution failed with status: {status}") + error_message = f"Parallel rails execution failed with status: {status}" + log.error(error_message) + yield _internal_rail_error("output rails", error_message) + return else: # if there are any stop events, content was blocked or internal error occurred result_events = getattr(result, "events", None) @@ -155,7 +159,10 @@ async def _run_parallel_output_rails( return except Exception as e: - log.error(f"Error in parallel rail execution: {e}") + error_message = f"Error in parallel rail execution: {e}" + log.error(error_message) + yield _internal_rail_error("output rails", error_message) + return # update explain info for parallel mode rails._explain_info = rails._ensure_explain_info() @@ -186,9 +193,16 @@ async def _run_sequential_output_rails( 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): + result, status = _sequential_action_result(action_result) + if status != "success": + error_message = f"Action {action_name} failed with status: {status}" + log.error(error_message) + yield _internal_rail_error(flow_id, error_message) + return + + # Use the mapping to decide if the result indicates blocked content. + if is_output_blocked(result, action_func): yield _blocked_output_error(flow_id) return @@ -322,6 +336,17 @@ def _internal_rail_error(flow_id: str, error_message: str) -> str: return json.dumps(error_data) +def _sequential_action_result(action_result: Any) -> tuple[Any, str]: + if ( + isinstance(action_result, tuple) + and len(action_result) == 2 + and isinstance(action_result[1], str) + and action_result[1] in {"success", "failed"} + ): + return action_result + return action_result, "success" + + def _parallel_stop_error(stop_event: dict) -> str: blocked_flow = stop_event.get("flow_id", "output rails") error_type = stop_event.get("error_type") diff --git a/tests/rails/llm/test_generation_stream.py b/tests/rails/llm/test_generation_stream.py index a08b433c91..fba9833bc2 100644 --- a/tests/rails/llm/test_generation_stream.py +++ b/tests/rails/llm/test_generation_stream.py @@ -326,39 +326,27 @@ async def test_generation_token_stream_pushes_error_payload_when_generation_fail @pytest.mark.asyncio -async def test_generation_token_stream_waits_for_generation_when_consumer_closes_early(): +async def test_generation_token_stream_cancels_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.wait_for(cast(Any, stream).aclose(), timeout=1) + await asyncio.wait_for(rails.cancelled.wait(), timeout=1) 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(): +async def test_generation_token_stream_cancels_generation_after_output_rail_error(): rails = BlockingOutputRailsFakeRails() stream = generation_token_stream(rails, messages=[{"role": "user", "content": "Hi"}]) - collect_task = asyncio.create_task(_collect(stream)) + chunks = await asyncio.wait_for(_collect(stream), timeout=1) + await asyncio.wait_for(rails.cancelled.wait(), timeout=1) 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 index 1b42ec3d49..3d54907ab3 100644 --- a/tests/rails/llm/test_streaming_output_rails.py +++ b/tests/rails/llm/test_streaming_output_rails.py @@ -298,7 +298,7 @@ async def test_output_rails_stream_yields_block_error_for_sequential_rails(monke @pytest.mark.asyncio -async def test_output_rails_stream_uses_config_stream_first_when_explicit_false(monkeypatch): +async def test_output_rails_stream_honors_explicit_stream_first_false(monkeypatch): _patch_buffer_strategy( monkeypatch, [ChunkBatch(processing_context=["blocked"], user_output_chunks=["blocked"])], @@ -316,9 +316,8 @@ async def test_output_rails_stream_uses_config_stream_first_when_explicit_false( ) ) - assert len(chunks) == 2 - assert chunks[0] == "blocked" - assert json.loads(chunks[1])["error"]["code"] == "content_blocked" + assert len(chunks) == 1 + assert json.loads(chunks[0])["error"]["code"] == "content_blocked" @pytest.mark.asyncio @@ -359,7 +358,7 @@ async def test_output_rails_stream_closes_external_generator_when_consumer_close @pytest.mark.asyncio -async def test_output_rails_stream_continues_for_sequential_action_failure(monkeypatch): +async def test_output_rails_stream_yields_internal_error_for_sequential_action_failure(monkeypatch): _patch_buffer_strategy( monkeypatch, [ChunkBatch(processing_context=["unsafe"], user_output_chunks=["unsafe"])], @@ -376,7 +375,15 @@ async def test_output_rails_stream_continues_for_sequential_action_failure(monke ) ) - assert chunks == ["unsafe"] + assert len(chunks) == 1 + assert json.loads(chunks[0]) == { + "error": { + "message": "Internal error in self check output rail: Action self_check_output failed with status: failed", + "type": "internal_error", + "param": "self check output", + "code": "rail_execution_failure", + } + } assert rails._explain_info == {"ensured": 1} @@ -437,7 +444,7 @@ async def test_output_rails_stream_yields_parallel_stop_event_error(monkeypatch) @pytest.mark.asyncio -async def test_output_rails_stream_continues_for_parallel_action_failure(monkeypatch): +async def test_output_rails_stream_yields_internal_error_for_parallel_action_failure(monkeypatch): _patch_buffer_strategy( monkeypatch, [ChunkBatch(processing_context=["unsafe"], user_output_chunks=["unsafe"])], @@ -454,12 +461,19 @@ async def test_output_rails_stream_continues_for_parallel_action_failure(monkeyp ) ) - assert chunks == ["unsafe"] - assert rails._explain_info == {"ensured": 1} + assert len(chunks) == 1 + assert json.loads(chunks[0]) == { + "error": { + "message": "Internal error in output rails rail: Parallel rails execution failed with status: failed", + "type": "internal_error", + "param": "output rails", + "code": "rail_execution_failure", + } + } @pytest.mark.asyncio -async def test_output_rails_stream_continues_for_parallel_exception(monkeypatch): +async def test_output_rails_stream_yields_internal_error_for_parallel_exception(monkeypatch): _patch_buffer_strategy( monkeypatch, [ChunkBatch(processing_context=["unsafe"], user_output_chunks=["unsafe"])], @@ -476,5 +490,12 @@ async def test_output_rails_stream_continues_for_parallel_exception(monkeypatch) ) ) - assert chunks == ["unsafe"] - assert rails._explain_info == {"ensured": 1} + assert len(chunks) == 1 + assert json.loads(chunks[0]) == { + "error": { + "message": "Internal error in output rails rail: Error in parallel rail execution: parallel boom", + "type": "internal_error", + "param": "output rails", + "code": "rail_execution_failure", + } + } From 28a45403570ec6b3420c453a246c96b134bc11b3 Mon Sep 17 00:00:00 2001 From: Pouyanpi <13303554+Pouyanpi@users.noreply.github.com> Date: Wed, 10 Jun 2026 11:24:34 +0200 Subject: [PATCH 2/2] fix(startup): avoid duplicate colang flow loading --- .../rails/llm/startup/colang_flows.py | 19 +++++++++++-- tests/rails/llm/test_config_preparation.py | 27 +++++++++++++++++-- 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/nemoguardrails/rails/llm/startup/colang_flows.py b/nemoguardrails/rails/llm/startup/colang_flows.py index c0afbe34c0..d6fa258dbd 100644 --- a/nemoguardrails/rails/llm/startup/colang_flows.py +++ b/nemoguardrails/rails/llm/startup/colang_flows.py @@ -67,6 +67,21 @@ def _iter_colang_resources(resource: Traversable) -> Iterator[Traversable]: yield from _iter_colang_resources(child) +def _flow_config_key(flow_config: dict) -> str | None: + return flow_config.get("id") or flow_config.get("name") + + +def _extend_missing_flow_configs(config: RailsConfig, flow_configs: list[dict]) -> None: + existing_flow_keys = {_flow_config_key(flow_config) for flow_config in config.flows} + for flow_config in flow_configs: + flow_key = _flow_config_key(flow_config) + if flow_key is not None and flow_key in existing_flow_keys: + continue + + config.flows.append(flow_config) + existing_flow_keys.add(flow_key) + + def load_default_colang_1_flows(config: RailsConfig) -> None: """Load the built-in Colang 1.0 LLM flows into the rails config.""" if config.colang_version != "1.0": @@ -79,7 +94,7 @@ def load_default_colang_1_flows(config: RailsConfig) -> None: for flow_config in default_flows: flow_config["is_system_flow"] = True - config.flows.extend(default_flows) + _extend_missing_flow_configs(config, default_flows) def load_guardrails_library_flows_and_bot_messages(config: RailsConfig) -> None: @@ -101,7 +116,7 @@ def load_guardrails_library_flows_and_bot_messages(config: RailsConfig) -> None: for flow_config in content["flows"]: flow_config["is_system_flow"] = True - config.flows.extend(content["flows"]) + _extend_missing_flow_configs(config, content["flows"]) for message_id, utterances in content.get("bot_messages", {}).items(): if message_id not in config.bot_messages: diff --git a/tests/rails/llm/test_config_preparation.py b/tests/rails/llm/test_config_preparation.py index dfe7a991f9..f69bb13288 100644 --- a/tests/rails/llm/test_config_preparation.py +++ b/tests/rails/llm/test_config_preparation.py @@ -26,7 +26,7 @@ def test_config_preparation_startup_module_exports_public_helpers(): ] -def test_prepare_llmrails_config_loads_colang_flows_each_time_in_place(): +def test_prepare_llmrails_config_deduplicates_loaded_colang_flows(): config = RailsConfig(models=[]) first = prepare_llmrails_config( @@ -41,7 +41,30 @@ def test_prepare_llmrails_config_loads_colang_flows_each_time_in_place(): assert first is config assert second is config assert first_flow_count > 0 - assert len(config.flows) > first_flow_count + assert len(config.flows) == first_flow_count + + +def test_prepare_llmrails_config_can_reload_flows_after_caller_clears_them(): + config = RailsConfig(models=[]) + + prepare_llmrails_config(config=config) + flow_count = len(config.flows) + config.flows = [] + + prepare_llmrails_config(config=config) + + assert len(config.flows) == flow_count + + +def test_prepare_llmrails_config_copy_does_not_duplicate_loaded_flows(): + config = RailsConfig(models=[]) + + prepare_llmrails_config(config=config) + flow_count = len(config.flows) + prepared = prepare_llmrails_config(config=config, in_place=False) + + assert prepared is not config + assert len(prepared.flows) == flow_count def test_prepare_llmrails_config_marks_rail_flows_each_time():