fix: improve reasoning traces validation for dialog rails

- Fix validation to properly handle both dictionary and Model object cases
- Add proper handling of cases where some dialog rail tasks have dedicated models and others fall back to main model
- Improve error messages to be more user-friendly:
  * Specify which model has the issue (main model or specific task model)
  * Reference config.yml in error messages
  * Provide clear YAML configuration instructions
  * Include specific task name when relevant
- Update test cases to match new error messages and validation logic
- Add proper handling of implicit dialog rails activation through user/bot messages and flows

This change ensures that reasoning traces are properly disabled when dialog rails are present,
regardless of whether they are explicitly configured or implicitly activated through user/bot
messages or flows. The improved error messages make it easier for users to understand and fix
configuration issues in their YAML files.

add more tests

Author: Pouyanpi

nemoguardrails/rails/llm/config.py

@@ -36,6 +36,7 @@
 from nemoguardrails.colang import parse_colang_file, parse_flow_elements
 from nemoguardrails.colang.v2_x.lang.utils import format_colang_parsing_error_message
 from nemoguardrails.colang.v2_x.runtime.errors import ColangParsingError
+from nemoguardrails.llm.types import Task
 
 log = logging.getLogger(__name__)
 
@@ -1140,37 +1141,66 @@ def check_reasoning_traces_with_dialog_rails(cls, values):
         rails = values.get("rails", {})
         dialog_rails = rails.get("dialog", {})
 
-        # check if any model has reasoning traces enabled
-        # TODO: we must check for models that are used in a specific dialog task
-        has_reasoning_traces = False
-        for model in models:
-            if isinstance(model, dict):
-                reasoning_config = model.get("reasoning_config", {})
-                if not reasoning_config.get("remove_thinking_traces", True):
-                    has_reasoning_traces = True
-                    break
-            elif hasattr(model, "reasoning_config"):
-                if not model.reasoning_config.remove_thinking_traces:
-                    has_reasoning_traces = True
-                    break
-
-        # check if dialog rails are present (explicitly or implicitly)
-        has_dialog_rails = bool(dialog_rails)
-
-        # check implicit dialog rails through user messages, bot messages, or flows
-        if not has_dialog_rails:
-            has_dialog_rails = (
-                bool(values.get("user_messages"))
-                or bool(values.get("bot_messages"))
-                or bool(values.get("flows"))
-            )
+        # dialog rail tasks that should not have reasoning traces
+        dialog_rail_tasks = [
+            Task.GENERATE_BOT_MESSAGE,
+            Task.GENERATE_USER_INTENT,
+            Task.GENERATE_NEXT_STEPS,
+            Task.GENERATE_INTENT_STEPS_MESSAGE,
+        ]
 
-        if has_reasoning_traces and has_dialog_rails:
-            raise ValueError(
-                "Reasoning traces cannot be enabled when dialog rails are present. "
-                "Please either disable reasoning traces or remove dialog rails."
+        # dialog rails are activated (explicitly or implicitly)
+        has_dialog_rails = (
+            bool(dialog_rails)
+            or bool(values.get("user_messages"))
+            or bool(values.get("bot_messages"))
+            or bool(values.get("flows"))
+        )
+
+        if has_dialog_rails:
+            # Get the main model if it exists
+            main_model = next(
+                (model for model in models if model.get("type") == "main"), None
             )
 
+            violations = []
+
+            for task in dialog_rail_tasks:
+                # Check if there's a dedicated model for this task
+                task_model = next(
+                    (model for model in models if model.get("type") == task.value), None
+                )
+
+                if task_model:
+                    # Handle both dictionary and Model object cases
+                    reasoning_config = (
+                        task_model.reasoning_config
+                        if hasattr(task_model, "reasoning_config")
+                        else task_model.get("reasoning_config", {})
+                    )
+                    if not reasoning_config.get("remove_thinking_traces", True):
+                        violations.append(
+                            f"Model '{task_model.get('type')}' has reasoning traces enabled in config.yml. "
+                            f"Reasoning traces must be disabled for dialog rail tasks. "
+                            f"Please update your config.yml to set 'remove_thinking_traces: true' under reasoning_config for this model."
+                        )
+                elif main_model:
+                    # Handle both dictionary and Model object cases
+                    reasoning_config = (
+                        main_model.reasoning_config
+                        if hasattr(main_model, "reasoning_config")
+                        else main_model.get("reasoning_config", {})
+                    )
+                    if not reasoning_config.get("remove_thinking_traces", True):
+                        violations.append(
+                            f"Main model has reasoning traces enabled in config.yml and is being used for dialog rail task '{task.value}'. "
+                            f"Reasoning traces must be disabled when dialog rails are present. "
+                            f"Please update your config.yml to set 'remove_thinking_traces: true' under reasoning_config for the main model."
+                        )
+
+            if violations:
+                raise ValueError("\n".join(violations))
+
         return values
 
     @root_validator(pre=True, allow_reuse=True)