feat(library): namespace custom task prompts in self-check

nemoguardrails/library/self_check/input_check/actions.py

@@ -21,7 +21,7 @@
 from nemoguardrails.library.self_check.utils import (
     SELF_CHECK_INPUT_DEFAULT_TASK,
     SELF_CHECK_INPUT_FLOW,
-    SELF_CHECK_INPUT_TASK_PARAM,
+    SELF_CHECK_INPUT_VARIANT_PARAM,
     resolve_self_check_task,
     run_self_check_task,
 )
@@ -64,7 +64,7 @@ async def self_check_input(
         triggered_rail_key="triggered_input_rail",
         start_rail_event_type="StartInputRail",
         flow_id=SELF_CHECK_INPUT_FLOW,
-        task_param=SELF_CHECK_INPUT_TASK_PARAM,
+        variant_param=SELF_CHECK_INPUT_VARIANT_PARAM,
         default_task=DEFAULT_TASK,
     )
 

nemoguardrails/library/self_check/input_check/flows.co

@@ -1,5 +1,5 @@
-flow self check input $task="self_check_input"
-  $allowed = await SelfCheckInputAction(task=$task)
+flow self check input $variant="self_check_input"
+  $allowed = await SelfCheckInputAction(task=$variant)
 
   if not $allowed
     if $system.config.enable_rails_exceptions

nemoguardrails/library/self_check/output_check/actions.py

@@ -21,7 +21,7 @@
 from nemoguardrails.library.self_check.utils import (
     SELF_CHECK_OUTPUT_DEFAULT_TASK,
     SELF_CHECK_OUTPUT_FLOW,
-    SELF_CHECK_OUTPUT_TASK_PARAM,
+    SELF_CHECK_OUTPUT_VARIANT_PARAM,
     resolve_self_check_task,
     run_self_check_task,
 )
@@ -68,7 +68,7 @@ async def self_check_output(
         triggered_rail_key="triggered_output_rail",
         start_rail_event_type="StartOutputRail",
         flow_id=SELF_CHECK_OUTPUT_FLOW,
-        task_param=SELF_CHECK_OUTPUT_TASK_PARAM,
+        variant_param=SELF_CHECK_OUTPUT_VARIANT_PARAM,
         default_task=DEFAULT_TASK,
     )
 

nemoguardrails/library/self_check/output_check/flows.co

@@ -1,5 +1,5 @@
-flow self check output $task="self_check_output"
-  $allowed = await SelfCheckOutputAction(task=$task)
+flow self check output $variant="self_check_output"
+  $allowed = await SelfCheckOutputAction(task=$variant)
 
   if not $allowed
     if $system.config.enable_rails_exceptions

nemoguardrails/library/self_check/utils.py

@@ -14,6 +14,8 @@
 # limitations under the License.
 
 import logging
+import warnings
+from collections.abc import Collection
 from typing import Any, Dict, List, Optional
 
 from nemoguardrails.actions.llm.utils import llm_call, warn_if_truncated
@@ -26,17 +28,41 @@
 
 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_VARIANT_PARAM = "variant"
+SELF_CHECK_INPUT_VARIANT_PARAM = SELF_CHECK_VARIANT_PARAM
+SELF_CHECK_OUTPUT_VARIANT_PARAM = SELF_CHECK_VARIANT_PARAM
 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]:
+def get_self_check_task_from_rail(flow: Any, flow_id: str, variant_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
+    return _get_flow_params(flow).get(variant_param) or default_task
+
+
+def get_self_check_prompt_task(
+    task: str, default_task: str, available_prompt_tasks: Optional[Collection[str]] = None
+) -> str:
+    """Resolve the prompt task name while supporting deprecated bare custom tasks."""
+    if task == default_task:
+        return task
+
+    prompt_task = f"{default_task} ${SELF_CHECK_VARIANT_PARAM}={task}"
+    if (
+        available_prompt_tasks is not None
+        and prompt_task not in available_prompt_tasks
+        and task in available_prompt_tasks
+    ):
+        warnings.warn(
+            f"The `{task}` self-check prompt task is deprecated; rename it to `{prompt_task}`.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return task
+
+    return prompt_task
 
 
 def resolve_self_check_task(
@@ -46,7 +72,7 @@ def resolve_self_check_task(
     triggered_rail_key: str,
     start_rail_event_type: str,
     flow_id: str,
-    task_param: str,
+    variant_param: str,
     default_task: str,
 ) -> str:
     if task and not task.startswith("$"):
@@ -56,7 +82,7 @@ def resolve_self_check_task(
     context_task = get_self_check_task_from_rail(
         context.get(triggered_rail_key),
         flow_id=flow_id,
-        task_param=task_param,
+        variant_param=variant_param,
         default_task=default_task,
     )
     if context_task:
@@ -65,13 +91,13 @@ def resolve_self_check_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
+            return event_params.get(variant_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,
+                variant_param=variant_param,
                 default_task=default_task,
             )
             if event_task:
@@ -121,15 +147,17 @@ async def run_self_check_task(
     max_tokens: int = 1024,
 ) -> tuple[bool, str]:
     llm = get_self_check_llm(llms, task, default_task=default_task, main_llm=main_llm)
+    available_prompt_tasks = {prompt.task for prompt in llm_task_manager.config.prompts or []}
+    prompt_task = get_self_check_prompt_task(task, default_task, available_prompt_tasks)
 
     prompt = llm_task_manager.render_task_prompt(
-        task=task,
+        task=prompt_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
+    stop = llm_task_manager.get_stop_tokens(task=prompt_task)
+    task_max_tokens = llm_task_manager.get_max_tokens(task=prompt_task) or max_tokens
 
-    llm_call_info_var.set(LLMCallInfo(task=task))
+    llm_call_info_var.set(LLMCallInfo(task=prompt_task))
 
     llm_response = await llm_call(
         llm,
@@ -140,8 +168,8 @@ async def run_self_check_task(
             "max_tokens": task_max_tokens,
         },
     )
-    warn_if_truncated(llm_response, task)
+    warn_if_truncated(llm_response, prompt_task)
     response = llm_response.content
 
-    result = parse_self_check_output(llm_task_manager, task, response)
+    result = parse_self_check_output(llm_task_manager, prompt_task, response)
     return bool(result[0]), response

nemoguardrails/rails/llm/config.py

@@ -43,6 +43,7 @@
     InvalidModelConfigurationError,
     InvalidRailsConfigurationError,
 )
+from nemoguardrails.library.self_check.utils import get_self_check_prompt_task
 
 log = logging.getLogger(__name__)
 
@@ -1960,8 +1961,8 @@ def check_prompt_exist_for_self_check_rails(cls, values):
         provided_task_prompts = [prompt.task if hasattr(prompt, "task") else prompt.get("task") for prompt in prompts]
 
         # Input moderation prompt verification
-        _validate_self_check_prompts(
-            enabled_input_rails, provided_task_prompts, "self check input", "self_check_input", "task"
+        _validate_self_check_rail_prompts(
+            enabled_input_rails, provided_task_prompts, "self check input", "self_check_input"
         )
         if "llama guard check input" in enabled_input_rails and "llama_guard_check_input" not in provided_task_prompts:
             raise InvalidRailsConfigurationError(
@@ -1975,8 +1976,8 @@ def check_prompt_exist_for_self_check_rails(cls, values):
         _validate_rail_prompts(enabled_input_rails, provided_task_prompts, "topic safety check input")
 
         # Output moderation prompt verification
-        _validate_self_check_prompts(
-            enabled_output_rails, provided_task_prompts, "self check output", "self_check_output", "task"
+        _validate_self_check_rail_prompts(
+            enabled_output_rails, provided_task_prompts, "self check output", "self_check_output"
         )
         if (
             "llama guard check output" in enabled_output_rails
@@ -2363,22 +2364,23 @@ def _get_flow_model(flow_text) -> Optional[str]:
     return flow_text.split(MODEL_PREFIX)[-1].strip()
 
 
-def _validate_self_check_prompts(
-    rails: list[str],
-    prompts: list[Any],
-    flow_name: str,
-    default_task: str,
-    task_param: str,
+def _validate_self_check_rail_prompts(
+    rails: list[str], prompts: list[Any], validation_rail: str, default_task: str
 ) -> None:
+    """Ensure every enabled self-check rail has a prompt for its resolved task.
+
+    A self-check rail may select a custom task via `$variant=...`; the task falls
+    back to `default_task` when omitted. Custom prompts are scoped under the
+    default task name.
+    """
     for rail in rails:
-        normalized = _normalize_flow_id(rail)
-        if normalized != flow_name:
+        if _normalize_flow_id(rail) != validation_rail:
             continue
-        custom_task = _get_flow_params(rail).get(task_param)
-        expected = custom_task if custom_task else default_task
-        if expected not in prompts:
+        task = _get_flow_params(rail).get("variant") or default_task
+        prompt_task = get_self_check_prompt_task(task, default_task, prompts)
+        if prompt_task not in prompts:
             raise InvalidRailsConfigurationError(
-                f"Missing a `{expected}` prompt template, which is required for the `{rail}` rail."
+                f"Missing a `{prompt_task}` prompt template, which is required for the `{rail}` rail."
             )