fix(self-check): preserve custom task routing

nemoguardrails/library/self_check/input_check/actions.py

@@ -14,49 +14,35 @@
 # 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)
 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,

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

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

nemoguardrails/library/self_check/output_check/actions.py

@@ -14,48 +14,34 @@
 # 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)
 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

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

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