diff --git a/nemoguardrails/library/self_check/input_check/actions.py b/nemoguardrails/library/self_check/input_check/actions.py index bd7401d3d0..65d0fbcc8d 100644 --- a/nemoguardrails/library/self_check/input_check/actions.py +++ b/nemoguardrails/library/self_check/input_check/actions.py @@ -14,39 +14,24 @@ # limitations under the License. import logging -from typing import Dict, Optional +from typing import Dict, List, Optional from nemoguardrails import RailsConfig from nemoguardrails.actions.actions import ActionResult, action -from nemoguardrails.actions.llm.utils import llm_call, warn_if_truncated -from nemoguardrails.context import llm_call_info_var +from nemoguardrails.library.self_check.utils import ( + SELF_CHECK_INPUT_DEFAULT_TASK, + SELF_CHECK_INPUT_FLOW, + SELF_CHECK_INPUT_TASK_PARAM, + resolve_self_check_task, + run_self_check_task, +) from nemoguardrails.llm.taskmanager import LLMTaskManager -from nemoguardrails.logging.explain import LLMCallInfo from nemoguardrails.types import LLMModel from nemoguardrails.utils import new_event_dict log = logging.getLogger(__name__) -DEFAULT_TASK = "self_check_input" - - -def _get_llm( - llms: Dict[str, LLMModel], task: str, default_task: str, default_llm: Optional[LLMModel] -) -> Optional[LLMModel]: - if task in llms: - return llms[task] - elif default_task in llms: - log.debug(f"No model found with type={task}, falling back to default {default_task} model") - return llms[default_task] - elif default_llm is not None: - log.debug(f"No model found with type={task} or type={default_task}, falling back to main model") - return default_llm - else: - error_msg = ( - f"No matching model for task={task} found. " - f"Please configure a model with type={task}, type={default_task}, or type=main" - ) - raise ValueError(error_msg) +DEFAULT_TASK = SELF_CHECK_INPUT_DEFAULT_TASK @action(is_system_action=True) @@ -54,9 +39,10 @@ async def self_check_input( llms: Dict[str, LLMModel], llm_task_manager: LLMTaskManager, context: Optional[dict] = None, + events: Optional[List[dict]] = None, llm: Optional[LLMModel] = None, config: Optional[RailsConfig] = None, - task: str = DEFAULT_TASK, + task: Optional[str] = None, **kwargs, ): """Checks the input from the user. @@ -68,55 +54,38 @@ async def self_check_input( True if the input should be allowed, False otherwise. """ - _MAX_TOKENS = 1024 + context = context or {} user_input = context.get("user_message") - # guard against an unset $task variable - if task.startswith("$") and task.endswith("_task"): - task = DEFAULT_TASK - - llm = _get_llm(llms, task, default_task=DEFAULT_TASK, default_llm=llm) + task = resolve_self_check_task( + task, + context, + events, + triggered_rail_key="triggered_input_rail", + start_rail_event_type="StartInputRail", + flow_id=SELF_CHECK_INPUT_FLOW, + task_param=SELF_CHECK_INPUT_TASK_PARAM, + default_task=DEFAULT_TASK, + ) if user_input: - prompt = llm_task_manager.render_task_prompt( + is_safe, response = await run_self_check_task( task=task, - context={ + prompt_context={ "user_input": user_input, }, + llms=llms, + default_task=DEFAULT_TASK, + main_llm=llm, + llm_task_manager=llm_task_manager, + lowest_temperature=config.lowest_temperature, ) - stop = llm_task_manager.get_stop_tokens(task=task) - max_tokens = llm_task_manager.get_max_tokens(task=task) - max_tokens = max_tokens or _MAX_TOKENS - - # Initialize the LLMCallInfo object - llm_call_info_var.set(LLMCallInfo(task=task)) - - llm_response = await llm_call( - llm, - prompt, - stop=stop, - llm_params={ - "temperature": config.lowest_temperature, - "max_tokens": max_tokens, - }, - ) - warn_if_truncated(llm_response, task) - response = llm_response.content if task == DEFAULT_TASK: log.info(f"Input self-checking result is: `{response}`.") else: log.info(f"Input self-checking result for task={task} is: `{response}`.") - # for sake of backward compatibility - # if the output_parser is not registered we will use the default one - if llm_task_manager.has_output_parser(task): - result = llm_task_manager.parse_task_output(task, output=response) - else: - result = llm_task_manager.parse_task_output(task, output=response, forced_output_parser="is_content_safe") - - is_safe = result[0] - if not is_safe: return ActionResult( return_value=False, diff --git a/nemoguardrails/library/self_check/input_check/flows.co b/nemoguardrails/library/self_check/input_check/flows.co index 1ee83f0b9b..dd6a3defdf 100644 --- a/nemoguardrails/library/self_check/input_check/flows.co +++ b/nemoguardrails/library/self_check/input_check/flows.co @@ -1,5 +1,5 @@ -flow self check input $input_task="self_check_input" - $allowed = await SelfCheckInputAction(task=$input_task) +flow self check input $task="self_check_input" + $allowed = await SelfCheckInputAction(task=$task) if not $allowed if $system.config.enable_rails_exceptions diff --git a/nemoguardrails/library/self_check/input_check/flows.v1.co b/nemoguardrails/library/self_check/input_check/flows.v1.co index 5ab8a4f227..dbd04ea1e7 100644 --- a/nemoguardrails/library/self_check/input_check/flows.v1.co +++ b/nemoguardrails/library/self_check/input_check/flows.v1.co @@ -2,10 +2,7 @@ define bot refuse to respond "I'm sorry, I can't respond to that." define flow self check input - if not $input_task - $input_task = "self_check_input" - - $allowed = execute self_check_input(task=$input_task) + $allowed = execute self_check_input if not $allowed if $config.enable_rails_exceptions diff --git a/nemoguardrails/library/self_check/output_check/actions.py b/nemoguardrails/library/self_check/output_check/actions.py index 16ed64f946..d56785960f 100644 --- a/nemoguardrails/library/self_check/output_check/actions.py +++ b/nemoguardrails/library/self_check/output_check/actions.py @@ -14,38 +14,23 @@ # limitations under the License. import logging -from typing import Dict, Optional +from typing import Dict, List, Optional from nemoguardrails import RailsConfig from nemoguardrails.actions import action -from nemoguardrails.actions.llm.utils import llm_call, warn_if_truncated -from nemoguardrails.context import llm_call_info_var +from nemoguardrails.library.self_check.utils import ( + SELF_CHECK_OUTPUT_DEFAULT_TASK, + SELF_CHECK_OUTPUT_FLOW, + SELF_CHECK_OUTPUT_TASK_PARAM, + resolve_self_check_task, + run_self_check_task, +) from nemoguardrails.llm.taskmanager import LLMTaskManager -from nemoguardrails.logging.explain import LLMCallInfo from nemoguardrails.types import LLMModel log = logging.getLogger(__name__) -DEFAULT_TASK = "self_check_output" - - -def _get_llm( - llms: Dict[str, LLMModel], task: str, default_task: str, default_llm: Optional[LLMModel] -) -> Optional[LLMModel]: - if task in llms: - return llms[task] - elif default_task in llms: - log.debug(f"No model found with type={task}, falling back to default {default_task} model") - return llms[default_task] - elif default_llm is not None: - log.debug(f"No model found with type={task} or type={default_task}, falling back to main model") - return default_llm - else: - error_msg = ( - f"No matching model for task={task} found. " - f"Please configure a model with type={task}, type={default_task}, or type=main" - ) - raise ValueError(error_msg) +DEFAULT_TASK = SELF_CHECK_OUTPUT_DEFAULT_TASK @action(is_system_action=True, output_mapping=lambda value: not value) @@ -53,9 +38,10 @@ async def self_check_output( llms: Dict[str, LLMModel], llm_task_manager: LLMTaskManager, context: Optional[dict] = None, + events: Optional[List[dict]] = None, llm: Optional[LLMModel] = None, config: Optional[RailsConfig] = None, - task: str = DEFAULT_TASK, + task: Optional[str] = None, **kwargs, ): """Checks if the output from the bot. @@ -70,57 +56,40 @@ async def self_check_output( True if the output should be allowed, False otherwise. """ - _MAX_TOKENS = 1024 + context = context or {} bot_response = context.get("bot_message") user_input = context.get("user_message") bot_thinking = context.get("bot_thinking") - # guard against an unset $task variable - if task.startswith("$") and task.endswith("_task"): - task = DEFAULT_TASK - - llm = _get_llm(llms, task, default_task=DEFAULT_TASK, default_llm=llm) + task = resolve_self_check_task( + task, + context, + events, + triggered_rail_key="triggered_output_rail", + start_rail_event_type="StartOutputRail", + flow_id=SELF_CHECK_OUTPUT_FLOW, + task_param=SELF_CHECK_OUTPUT_TASK_PARAM, + default_task=DEFAULT_TASK, + ) if bot_response: - prompt = llm_task_manager.render_task_prompt( + is_safe, response = await run_self_check_task( task=task, - context={ + prompt_context={ "user_input": user_input, "bot_response": bot_response, "bot_thinking": bot_thinking, }, + llms=llms, + default_task=DEFAULT_TASK, + main_llm=llm, + llm_task_manager=llm_task_manager, + lowest_temperature=config.lowest_temperature, ) - stop = llm_task_manager.get_stop_tokens(task=task) - max_tokens = llm_task_manager.get_max_tokens(task=task) - max_tokens = max_tokens or _MAX_TOKENS - - # Initialize the LLMCallInfo object - llm_call_info_var.set(LLMCallInfo(task=task)) - - llm_response = await llm_call( - llm, - prompt, - stop=stop, - llm_params={ - "temperature": config.lowest_temperature, - "max_tokens": max_tokens, - }, - ) - warn_if_truncated(llm_response, task) - response = llm_response.content if task == DEFAULT_TASK: log.info(f"Output self-checking result is: `{response}`.") else: log.info(f"Output self-checking result for task={task} is: `{response}`.") - # for sake of backward compatibility - # if the output_parser is not registered we will use the default one - if llm_task_manager.has_output_parser(task): - result = llm_task_manager.parse_task_output(task, output=response) - else: - result = llm_task_manager.parse_task_output(task, output=response, forced_output_parser="is_content_safe") - - is_safe = result[0] - return is_safe diff --git a/nemoguardrails/library/self_check/output_check/flows.co b/nemoguardrails/library/self_check/output_check/flows.co index 64759f6bcc..8b8bafb294 100644 --- a/nemoguardrails/library/self_check/output_check/flows.co +++ b/nemoguardrails/library/self_check/output_check/flows.co @@ -1,5 +1,5 @@ -flow self check output $output_task="self_check_output" - $allowed = await SelfCheckOutputAction(task=$output_task) +flow self check output $task="self_check_output" + $allowed = await SelfCheckOutputAction(task=$task) if not $allowed if $system.config.enable_rails_exceptions diff --git a/nemoguardrails/library/self_check/output_check/flows.v1.co b/nemoguardrails/library/self_check/output_check/flows.v1.co index a40d7e41bc..3d39a98ae6 100644 --- a/nemoguardrails/library/self_check/output_check/flows.v1.co +++ b/nemoguardrails/library/self_check/output_check/flows.v1.co @@ -2,10 +2,7 @@ define bot refuse to respond "I'm sorry, I can't respond to that." define flow self check output - if not $output_task - $output_task = "self_check_output" - - $allowed = execute self_check_output(task=$output_task) + $allowed = execute self_check_output if not $allowed if $config.enable_rails_exceptions diff --git a/nemoguardrails/library/self_check/utils.py b/nemoguardrails/library/self_check/utils.py new file mode 100644 index 0000000000..1708b7ea44 --- /dev/null +++ b/nemoguardrails/library/self_check/utils.py @@ -0,0 +1,147 @@ +# 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. + +import logging +from typing import Any, Dict, List, Optional + +from nemoguardrails.actions.llm.utils import llm_call, warn_if_truncated +from nemoguardrails.colang.v1_0.runtime.flows import _get_flow_params, _normalize_flow_id +from nemoguardrails.context import llm_call_info_var +from nemoguardrails.logging.explain import LLMCallInfo +from nemoguardrails.types import LLMModel + +log = logging.getLogger(__name__) + +SELF_CHECK_INPUT_FLOW = "self check input" +SELF_CHECK_OUTPUT_FLOW = "self check output" +SELF_CHECK_INPUT_TASK_PARAM = "task" +SELF_CHECK_OUTPUT_TASK_PARAM = "task" +SELF_CHECK_INPUT_DEFAULT_TASK = "self_check_input" +SELF_CHECK_OUTPUT_DEFAULT_TASK = "self_check_output" + + +def get_self_check_task_from_rail(flow: Any, flow_id: str, task_param: str, default_task: str) -> Optional[str]: + if not isinstance(flow, str) or _normalize_flow_id(flow) != flow_id: + return None + + return _get_flow_params(flow).get(task_param) or default_task + + +def resolve_self_check_task( + task: Optional[str], + context: Optional[dict], + events: Optional[List[dict]], + triggered_rail_key: str, + start_rail_event_type: str, + flow_id: str, + task_param: str, + default_task: str, +) -> str: + if task and not task.startswith("$"): + return task + + context = context or {} + context_task = get_self_check_task_from_rail( + context.get(triggered_rail_key), + flow_id=flow_id, + task_param=task_param, + default_task=default_task, + ) + if context_task: + return context_task + + for event in reversed(events or []): + if event.get("type") == "start_flow" and event.get("flow_id") == flow_id: + event_params = event.get("params") or {} + return event_params.get(task_param) or default_task + + if event.get("type") == start_rail_event_type: + event_task = get_self_check_task_from_rail( + event.get("flow_id"), + flow_id=flow_id, + task_param=task_param, + default_task=default_task, + ) + if event_task: + return event_task + + return default_task + + +def get_self_check_llm( + llms: Dict[str, LLMModel], + task: str, + default_task: str, + main_llm: Optional[LLMModel], +) -> LLMModel: + if task in llms: + return llms[task] + + if default_task in llms: + log.debug("No model found with type=%s, falling back to default %s model", task, default_task) + return llms[default_task] + + if main_llm is not None: + log.debug("No model found with type=%s or type=%s, falling back to main model", task, default_task) + return main_llm + + raise ValueError( + f"No matching model for task={task} found. " + f"Please configure a model with type={task}, type={default_task}, or type=main" + ) + + +def parse_self_check_output(llm_task_manager: Any, task: str, response: str): + if llm_task_manager.has_output_parser(task): + return llm_task_manager.parse_task_output(task, output=response) + + return llm_task_manager.parse_task_output(task, output=response, forced_output_parser="is_content_safe") + + +async def run_self_check_task( + task: str, + prompt_context: Dict[str, Any], + llms: Dict[str, LLMModel], + default_task: str, + main_llm: Optional[LLMModel], + llm_task_manager: Any, + lowest_temperature: float, + max_tokens: int = 1024, +) -> tuple[bool, str]: + llm = get_self_check_llm(llms, task, default_task=default_task, main_llm=main_llm) + + prompt = llm_task_manager.render_task_prompt( + task=task, + context=prompt_context, + ) + stop = llm_task_manager.get_stop_tokens(task=task) + task_max_tokens = llm_task_manager.get_max_tokens(task=task) or max_tokens + + llm_call_info_var.set(LLMCallInfo(task=task)) + + llm_response = await llm_call( + llm, + prompt, + stop=stop, + llm_params={ + "temperature": lowest_temperature, + "max_tokens": task_max_tokens, + }, + ) + warn_if_truncated(llm_response, task) + response = llm_response.content + + result = parse_self_check_output(llm_task_manager, task, response) + return bool(result[0]), response diff --git a/nemoguardrails/rails/llm/llmrails.py b/nemoguardrails/rails/llm/llmrails.py index 63c0fa27b7..fbc6791356 100644 --- a/nemoguardrails/rails/llm/llmrails.py +++ b/nemoguardrails/rails/llm/llmrails.py @@ -53,7 +53,7 @@ from nemoguardrails.actions.v2_x.generation import LLMGenerationActionsV2dotx from nemoguardrails.base_guardrails import BaseGuardrails from nemoguardrails.colang import parse_colang_file -from nemoguardrails.colang.v1_0.runtime.flows import _normalize_flow_id, compute_context +from nemoguardrails.colang.v1_0.runtime.flows import _get_flow_params, _normalize_flow_id, compute_context from nemoguardrails.colang.v1_0.runtime.runtime import Runtime, RuntimeV1_0 from nemoguardrails.colang.v2_x.runtime.flows import Action, State from nemoguardrails.colang.v2_x.runtime.runtime import RuntimeV2_x @@ -1847,10 +1847,11 @@ def _prepare_params( bot_response_chunk: str, prompt: Optional[str] = None, messages: Optional[List[dict]] = None, - action_params: Dict[str, Any] = {}, + action_params: Optional[Dict[str, Any]] = None, ): context_message = _get_last_context_message(messages) user_message = prompt or _get_latest_user_message(messages) + flow_params = _get_flow_params(flow_id) context = { "user_message": user_message, @@ -1859,17 +1860,20 @@ def _prepare_params( if context_message: context.update(context_message["content"]) + context["triggered_output_rail"] = flow_id model_name = flow_id.split("$")[-1].split("=")[-1].strip('"') - # we pass action params that are defined in the flow - # caveate, e.g. prmpt_security uses bot_response=$bot_message - # to resolve replace placeholders in action_params - for key, value in action_params.items(): - if value == "$bot_message": - action_params[key] = bot_response_chunk - elif value == "$user_message": - action_params[key] = user_message + context_params = { + "bot_message": bot_response_chunk, + "user_message": user_message, + **flow_params, + } + resolved_action_params = {} + for key, value in (action_params or {}).items(): + if isinstance(value, str) and value.startswith("$"): + value = context_params.get(value[1:], value) + resolved_action_params[key] = value return { # TODO:: are there other context variables that need to be passed? @@ -1881,7 +1885,7 @@ def _prepare_params( "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), - **action_params, + **resolved_action_params, } buffer_strategy = get_buffer_strategy(output_rails_streaming_config) diff --git a/tests/integrations/langchain/test_streaming.py b/tests/integrations/langchain/test_streaming.py index b1802edee9..4c58895d47 100644 --- a/tests/integrations/langchain/test_streaming.py +++ b/tests/integrations/langchain/test_streaming.py @@ -23,6 +23,7 @@ from nemoguardrails.actions import action from nemoguardrails.exceptions import StreamingNotSupportedError from nemoguardrails.streaming import StreamingHandler +from nemoguardrails.testing.fake_model import FakeLLMModel from tests.utils import TestChat @@ -381,6 +382,75 @@ async def test_sequential_streaming_output_rails_allowed( await asyncio.gather(*asyncio.all_tasks() - {asyncio.current_task()}) +@pytest.mark.asyncio +async def test_streaming_output_rails_preserve_custom_task(): + config = RailsConfig.from_content( + config={ + "models": [], + "rails": { + "output": { + "flows": ["self check output $task=check_data_leakage"], + "streaming": { + "enabled": True, + "chunk_size": 4, + "context_size": 2, + "stream_first": False, + }, + } + }, + "prompts": [ + { + "task": "check_data_leakage", + "content": """ + Bot response: {{ bot_response }} + Answer No. + """, + }, + { + "task": "self_check_output", + "content": """ + Bot response: {{ bot_response }} + Answer Yes. + """, + }, + ], + }, + colang_content=""" + define user express greeting + "hi" + + define flow + user express greeting + bot tell joke + """, + ) + custom_llm = FakeLLMModel(responses=["No", "No"]) + default_llm = FakeLLMModel(responses=["Yes"]) + + chat = TestChat( + config, + llm_completions=[ + ' express greeting\nbot express greeting\n "Hi, how are you doing?"', + ' "This response should stream safely."', + ], + streaming=True, + ) + chat.app.runtime.registered_action_params["llms"]["check_data_leakage"] = custom_llm + chat.app.runtime.registered_action_params["llms"]["self_check_output"] = default_llm + + chunks = [] + async for chunk in chat.app.stream_async( + messages=[{"role": "user", "content": "Hi!"}], + ): + chunks.append(chunk) + + assert "".join(chunks) == "This response should stream safely." + assert custom_llm.inference_count > 0 + assert default_llm.inference_count == 0 + + await asyncio.gather(*asyncio.all_tasks() - {asyncio.current_task()}) + + @pytest.mark.asyncio async def test_streaming_output_rails_blocked(output_rails_streaming_config): """This test checks if the streaming output rails block the completions when a BLOCK keyword is present. diff --git a/tests/test_multiple_self_check_rails.py b/tests/test_multiple_self_check_rails.py index f93438e2d4..cfa67ec51a 100644 --- a/tests/test_multiple_self_check_rails.py +++ b/tests/test_multiple_self_check_rails.py @@ -18,11 +18,19 @@ import pytest from nemoguardrails import RailsConfig +from nemoguardrails.library.self_check.utils import ( + SELF_CHECK_INPUT_DEFAULT_TASK, + SELF_CHECK_INPUT_FLOW, + SELF_CHECK_INPUT_TASK_PARAM, + get_self_check_llm, + get_self_check_task_from_rail, + resolve_self_check_task, + run_self_check_task, +) +from nemoguardrails.llm.taskmanager import LLMTaskManager from nemoguardrails.testing.fake_model import FakeLLMModel from tests.utils import TestChat -# --- Multiple input rails --- - multi_input_config = RailsConfig.from_content( """ define user express greeting @@ -41,8 +49,8 @@ rails: input: flows: - - self check input $input_task=check_harmful - - self check input $input_task=check_off_topic + - self check input $task=check_harmful + - self check input $task=check_off_topic prompts: - task: check_harmful content: | @@ -65,8 +73,8 @@ def test_multiple_input_rails_both_pass(): chat = TestChat( multi_input_config, llm_completions=[ - "No", # check_harmful passes - "No", # check_off_topic passes + "No", + "No", " express greeting", ' "Hey!"', ], @@ -83,7 +91,7 @@ def test_multiple_input_rails_first_blocks(): chat = TestChat( multi_input_config, llm_completions=[ - "Yes", # check_harmful blocks + "Yes", ], ) @@ -99,8 +107,8 @@ def test_multiple_input_rails_second_blocks(): chat = TestChat( multi_input_config, llm_completions=[ - "No", # check_harmful passes - "Yes", # check_off_topic blocks + "No", + "Yes", ], ) @@ -111,8 +119,6 @@ def test_multiple_input_rails_second_blocks(): assert new_message["content"]["type"] == "InputRailException" -# --- Multiple output rails --- - multi_output_config = RailsConfig.from_content( """ define user ask question @@ -127,8 +133,8 @@ def test_multiple_input_rails_second_blocks(): rails: output: flows: - - self check output $output_task=check_inappropriate - - self check output $output_task=check_data_leakage + - self check output $task=check_inappropriate + - self check output $task=check_data_leakage prompts: - task: check_inappropriate content: | @@ -153,8 +159,8 @@ def test_multiple_output_rails_both_pass(): llm_completions=[ " ask question", " Here is the answer.", - "No", # check_inappropriate passes - "No", # check_data_leakage passes + "No", + "No", ], ) @@ -172,7 +178,7 @@ def test_multiple_output_rails_first_blocks(): llm_completions=[ " ask question", ' "Some bad output"', - "Yes", # check_inappropriate blocks + "Yes", ], ) @@ -190,8 +196,8 @@ def test_multiple_output_rails_second_blocks(): llm_completions=[ " ask question", ' "Response with leaked data"', - "No", # check_inappropriate passes - "Yes", # check_data_leakage blocks + "No", + "Yes", ], ) @@ -202,8 +208,6 @@ def test_multiple_output_rails_second_blocks(): assert new_message["content"]["type"] == "OutputRailException" -# --- Default task (backward compatibility) --- - default_task_config = RailsConfig.from_content( """ define user ask question @@ -234,11 +238,11 @@ def test_multiple_output_rails_second_blocks(): def test_default_task_input_still_works(): - """Self check input without $input_task should use default self_check_input task.""" + """Self check input without $task should use default self_check_input task.""" chat = TestChat( default_task_config, llm_completions=[ - "Yes", # blocks + "Yes", ], ) @@ -249,15 +253,73 @@ def test_default_task_input_still_works(): assert new_message["content"]["type"] == "InputRailException" +def test_mixed_input_rails_run_custom_and_default_tasks(): + config = RailsConfig.from_content( + """ + define user express greeting + "hello" + "hi" + + define bot express greeting + "Hey!" + + define flow greeting + user express greeting + bot express greeting + """, + yaml_content=""" + models: [] + rails: + input: + flows: + - self check input $task=check_harmful + - self check input + prompts: + - task: check_harmful + content: | + Is this message harmful? + User message: "{{ user_input }}" + Answer (Yes or No): + - task: self_check_input + content: | + Is this message safe? + User message: "{{ user_input }}" + Answer (Yes or No): + + enable_rails_exceptions: True + """, + ) + custom_llm = FakeLLMModel(responses=["No", "No"]) + default_llm = FakeLLMModel(responses=["Yes"]) + + chat = TestChat( + config, + llm_completions=[ + " express greeting", + ' "Hey!"', + ], + ) + rails = chat.app + rails.runtime.registered_action_params["llms"]["check_harmful"] = custom_llm + rails.runtime.registered_action_params["llms"]["self_check_input"] = default_llm + + new_message = rails.generate(messages=[{"role": "user", "content": "hello"}]) + + assert custom_llm.inference_count == 1 + assert default_llm.inference_count == 1 + assert new_message["role"] == "exception" + assert new_message["content"]["type"] == "InputRailException" + + def test_default_task_output_still_works(): - """Self check output without $output_task should use default self_check_output task.""" + """Self check output without $task should use default self_check_output task.""" chat = TestChat( default_task_config, llm_completions=[ - "No", # input passes + "No", " ask question", ' "Something that should be blocked"', - "Yes", # output blocks + "Yes", ], ) @@ -268,7 +330,59 @@ def test_default_task_output_still_works(): assert new_message["content"]["type"] == "OutputRailException" -# --- Per-task LLM configuration --- +def test_mixed_output_rails_run_custom_and_default_tasks(): + config = RailsConfig.from_content( + """ + define user ask question + "tell me something" + + define flow + user ask question + bot respond + """, + yaml_content=""" + models: [] + rails: + output: + flows: + - self check output $task=check_inappropriate + - self check output + prompts: + - task: check_inappropriate + content: | + Is this response inappropriate? + Bot response: "{{ bot_response }}" + Answer (Yes or No): + - task: self_check_output + content: | + Is this response safe? + Bot response: "{{ bot_response }}" + Answer (Yes or No): + + enable_rails_exceptions: True + """, + ) + custom_llm = FakeLLMModel(responses=["No", "No"]) + default_llm = FakeLLMModel(responses=["Yes"]) + + chat = TestChat( + config, + llm_completions=[ + " ask question", + ' "Response that the default rail blocks"', + ], + ) + rails = chat.app + rails.runtime.registered_action_params["llms"]["check_inappropriate"] = custom_llm + rails.runtime.registered_action_params["llms"]["self_check_output"] = default_llm + + new_message = rails.generate(messages=[{"role": "user", "content": "tell me something"}]) + + assert custom_llm.inference_count == 1 + assert default_llm.inference_count == 1 + assert new_message["role"] == "exception" + assert new_message["content"]["type"] == "OutputRailException" + per_task_input_config = RailsConfig.from_content( """ @@ -288,8 +402,8 @@ def test_default_task_output_still_works(): rails: input: flows: - - self check input $input_task=check_harmful - - self check input $input_task=check_off_topic + - self check input $task=check_harmful + - self check input $task=check_off_topic prompts: - task: check_harmful content: | @@ -367,8 +481,8 @@ def test_per_task_llm_input_first_blocks_skips_second(): rails: output: flows: - - self check output $output_task=check_inappropriate - - self check output $output_task=check_data_leakage + - self check output $task=check_inappropriate + - self check output $task=check_data_leakage prompts: - task: check_inappropriate content: | @@ -436,13 +550,124 @@ def test_per_task_llm_output_first_blocks_skips_second(): assert data_leakage_llm.inference_count == 0 +def test_parallel_input_rail_uses_custom_task(): + config = RailsConfig.from_content( + """ + define user express greeting + "hello" + "hi" + + define bot express greeting + "Hey!" + + define flow greeting + user express greeting + bot express greeting + """, + yaml_content=""" + models: [] + rails: + input: + parallel: true + flows: + - self check input $task=check_harmful + prompts: + - task: check_harmful + content: | + Is this message harmful? + User message: "{{ user_input }}" + Answer (Yes or No): + - task: self_check_input + content: | + Is this message safe? + User message: "{{ user_input }}" + Answer (Yes or No): + + enable_rails_exceptions: True + """, + ) + custom_llm = FakeLLMModel(responses=["No"]) + default_llm = FakeLLMModel(responses=["Yes"]) + + chat = TestChat( + config, + llm_completions=[ + " express greeting", + ' "Hey!"', + ], + ) + rails = chat.app + rails.runtime.registered_action_params["llms"]["check_harmful"] = custom_llm + rails.runtime.registered_action_params["llms"]["self_check_input"] = default_llm + + new_message = rails.generate(messages=[{"role": "user", "content": "hello"}]) + + assert new_message["role"] == "assistant" + assert custom_llm.inference_count == 1 + assert default_llm.inference_count == 0 + + +def test_parallel_output_rail_uses_custom_task(): + config = RailsConfig.from_content( + """ + define user ask question + "tell me something" + + define flow + user ask question + bot respond + """, + yaml_content=""" + models: [] + rails: + output: + parallel: true + flows: + - self check output $task=check_inappropriate + prompts: + - task: check_inappropriate + content: | + Is this response inappropriate? + Bot response: "{{ bot_response }}" + Answer (Yes or No): + - task: self_check_output + content: | + Is this response safe? + Bot response: "{{ bot_response }}" + Answer (Yes or No): + + enable_rails_exceptions: True + """, + ) + custom_llm = FakeLLMModel(responses=["No"]) + default_llm = FakeLLMModel(responses=["Yes"]) + + chat = TestChat( + config, + llm_completions=[ + " ask question", + " Here is the answer.", + ], + ) + rails = chat.app + rails.runtime.registered_action_params["llms"]["check_inappropriate"] = custom_llm + rails.runtime.registered_action_params["llms"]["self_check_output"] = default_llm + + new_message = rails.generate(messages=[{"role": "user", "content": "tell me something"}]) + + assert new_message["role"] == "assistant" + assert new_message["content"] == "Here is the answer." + assert custom_llm.inference_count == 1 + assert default_llm.inference_count == 0 + + def test_per_task_llm_falls_back_to_main_when_not_configured(): """When no task-specific LLM is in the llms dict, it should fall back to the main LLM.""" chat = TestChat( per_task_input_config, llm_completions=[ - "No", # check_harmful via main LLM - "No", # check_off_topic via main LLM + "No", + "No", " express greeting", ], ) @@ -454,9 +679,6 @@ def test_per_task_llm_falls_back_to_main_when_not_configured(): assert chat.llm.inference_count == 3 -# --- Model fallback chain --- - - def test_input_fallback_to_default_task_model(): """When a custom task has no model, fall back to the self_check_input model.""" default_llm = FakeLLMModel(responses=["No", "No"]) @@ -550,18 +772,130 @@ def test_output_fallback_chain_prefers_task_over_default(): assert chat.llm.inference_count == 2 -def test_input_no_model_raises_error(): - """When no model is available at any fallback level, _get_llm raises ValueError.""" - from nemoguardrails.library.self_check.input_check.actions import _get_llm +def _resolve_input_task(task=None, context=None, events=None): + return resolve_self_check_task( + task, + context, + events, + triggered_rail_key="triggered_input_rail", + start_rail_event_type="StartInputRail", + flow_id=SELF_CHECK_INPUT_FLOW, + task_param=SELF_CHECK_INPUT_TASK_PARAM, + default_task=SELF_CHECK_INPUT_DEFAULT_TASK, + ) + +@pytest.mark.asyncio +async def test_run_self_check_task_uses_concrete_task_without_runtime_context(): + task_llm = FakeLLMModel(responses=["No"]) + default_llm = FakeLLMModel(responses=["Yes"]) + + is_safe, response = await run_self_check_task( + task="check_harmful", + prompt_context={"user_input": "hello"}, + llms={ + "check_harmful": task_llm, + "self_check_input": default_llm, + }, + default_task=SELF_CHECK_INPUT_DEFAULT_TASK, + main_llm=None, + llm_task_manager=LLMTaskManager(per_task_input_config), + lowest_temperature=per_task_input_config.lowest_temperature, + ) + + assert is_safe + assert response == "No" + assert task_llm.inference_count == 1 + assert default_llm.inference_count == 0 + + +def test_get_self_check_task_from_rail_resolves_custom_and_default_tasks(): + assert ( + get_self_check_task_from_rail( + "self check input $task=check_harmful", + flow_id=SELF_CHECK_INPUT_FLOW, + task_param=SELF_CHECK_INPUT_TASK_PARAM, + default_task=SELF_CHECK_INPUT_DEFAULT_TASK, + ) + == "check_harmful" + ) + + assert ( + get_self_check_task_from_rail( + "self check input", + flow_id=SELF_CHECK_INPUT_FLOW, + task_param=SELF_CHECK_INPUT_TASK_PARAM, + default_task=SELF_CHECK_INPUT_DEFAULT_TASK, + ) + == SELF_CHECK_INPUT_DEFAULT_TASK + ) + + assert ( + get_self_check_task_from_rail( + "self check output $task=check_inappropriate", + flow_id=SELF_CHECK_INPUT_FLOW, + task_param=SELF_CHECK_INPUT_TASK_PARAM, + default_task=SELF_CHECK_INPUT_DEFAULT_TASK, + ) + is None + ) + + +def test_resolve_self_check_task_prefers_explicit_task(): + task = _resolve_input_task( + task="check_harmful", + context={"triggered_input_rail": "self check input $task=check_off_topic"}, + ) + + assert task == "check_harmful" + + +def test_resolve_self_check_task_uses_triggered_rail_context(): + task = _resolve_input_task( + task="$task", + context={"triggered_input_rail": "self check input $task=check_harmful"}, + ) + + assert task == "check_harmful" + + +def test_resolve_self_check_task_uses_latest_start_rail_event(): + task = _resolve_input_task( + task="$task", + events=[ + {"type": "StartInputRail", "flow_id": "self check input $task=check_off_topic"}, + {"type": "SomeOtherEvent", "flow_id": "self check input $task=ignored"}, + {"type": "StartInputRail", "flow_id": "self check input $task=check_harmful"}, + ], + ) + + assert task == "check_harmful" + + +def test_resolve_self_check_task_uses_start_flow_params(): + task = _resolve_input_task( + events=[ + {"type": "start_flow", "flow_id": SELF_CHECK_INPUT_FLOW, "params": {"task": "check_harmful"}}, + ], + ) + + assert task == "check_harmful" + + +def test_resolve_self_check_task_defaults_unresolved_placeholders(): + assert _resolve_input_task(context={"triggered_input_rail": "self check input"}) == SELF_CHECK_INPUT_DEFAULT_TASK + assert _resolve_input_task(task="$task") == SELF_CHECK_INPUT_DEFAULT_TASK + assert _resolve_input_task() == SELF_CHECK_INPUT_DEFAULT_TASK + + +def test_input_no_model_raises_error(): + """When no model is available at any fallback level, get_self_check_llm raises ValueError.""" with pytest.raises(ValueError, match="No matching model"): - _get_llm({}, "check_harmful", "self_check_input", default_llm=None) + get_self_check_llm({}, "check_harmful", "self_check_input", main_llm=None) def test_get_llm_fallback_chain(): - """_get_llm resolves models in order: task -> default_task -> main -> ValueError.""" - from nemoguardrails.library.self_check.input_check.actions import _get_llm - + """get_self_check_llm resolves models in order: task -> default_task -> main -> ValueError.""" task_llm = FakeLLMModel(responses=[]) default_llm = FakeLLMModel(responses=[]) main_llm = FakeLLMModel(responses=[]) @@ -572,17 +906,17 @@ def test_get_llm_fallback_chain(): } # Level 1: exact task match - assert _get_llm(all_llms, "check_harmful", "self_check_input", default_llm=main_llm) is task_llm + assert get_self_check_llm(all_llms, "check_harmful", "self_check_input", main_llm=main_llm) is task_llm # Level 2: fall back to default task - assert _get_llm(all_llms, "check_off_topic", "self_check_input", default_llm=main_llm) is default_llm + assert get_self_check_llm(all_llms, "check_off_topic", "self_check_input", main_llm=main_llm) is default_llm # Level 3: fall back to default_llm (the llm action param) - assert _get_llm({}, "check_harmful", "self_check_input", default_llm=main_llm) is main_llm + assert get_self_check_llm({}, "check_harmful", "self_check_input", main_llm=main_llm) is main_llm # Level 4: no model raises ValueError with pytest.raises(ValueError, match="No matching model"): - _get_llm({}, "check_harmful", "self_check_input", default_llm=None) + get_self_check_llm({}, "check_harmful", "self_check_input", main_llm=None) with pytest.raises(ValueError, match="No matching model"): - _get_llm({}, "check_inappropriate", "self_check_output", default_llm=None) + get_self_check_llm({}, "check_inappropriate", "self_check_output", main_llm=None)