fix(config): validate content safety and topic control configs at creation time (#1450)

Refactor internal error tests to work with new config-time validation.
Remove tests for missing prompts/models (now caught at config-time),
update single error message test to use runtime error injection.

Author: tgasser-nv

nemoguardrails/rails/llm/config.py

@@ -35,6 +35,7 @@
 
 from nemoguardrails import utils
 from nemoguardrails.colang import parse_colang_file, parse_flow_elements
+from nemoguardrails.colang.v1_0.runtime.flows import _normalize_flow_id
 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
@@ -1451,7 +1452,65 @@ class RailsConfig(BaseModel):
         description="Configuration for tracing.",
     )
 
-    @root_validator(pre=True, allow_reuse=True)
+    @root_validator(pre=True)
+    def check_model_exists_for_input_rails(cls, values):
+        """Make sure we have a model for each input rail where one is provided using $model=<model_type>"""
+        rails = values.get("rails", {})
+        input_flows = rails.get("input", {}).get("flows", [])
+
+        # If no flows have a model, early-out
+        input_flows_without_model = [
+            _get_flow_model(flow) is None for flow in input_flows
+        ]
+        if all(input_flows_without_model):
+            return values
+
+        models = values.get("models", []) or []
+        model_types = {
+            model.type if isinstance(model, Model) else model["type"]
+            for model in models
+        }
+
+        for flow in input_flows:
+            flow_model = _get_flow_model(flow)
+            if not flow_model:
+                continue
+            if flow_model not in model_types:
+                raise ValueError(
+                    f"No `{flow_model}` model provided for input flow `{_normalize_flow_id(flow)}`"
+                )
+        return values
+
+    @root_validator(pre=True)
+    def check_model_exists_for_output_rails(cls, values):
+        """Make sure we have a model for each output rail where one is provided using $model=<model_type>"""
+        rails = values.get("rails", {})
+        output_flows = rails.get("output", {}).get("flows", [])
+
+        # If no flows have a model, early-out
+        output_flows_without_model = [
+            _get_flow_model(flow) is None for flow in output_flows
+        ]
+        if all(output_flows_without_model):
+            return values
+
+        models = values.get("models", []) or []
+        model_types = {
+            model.type if isinstance(model, Model) else model["type"]
+            for model in models
+        }
+
+        for flow in output_flows:
+            flow_model = _get_flow_model(flow)
+            if not flow_model:
+                continue
+            if flow_model not in model_types:
+                raise ValueError(
+                    f"No `{flow_model}` model provided for output flow `{_normalize_flow_id(flow)}`"
+                )
+        return values
+
+    @root_validator(pre=True)
     def check_prompt_exist_for_self_check_rails(cls, values):
         rails = values.get("rails", {})
         prompts = values.get("prompts", []) or []
@@ -1477,6 +1536,16 @@ def check_prompt_exist_for_self_check_rails(cls, values):
                 "You must provide a `llama_guard_check_input` prompt template."
             )
 
+        # Only content-safety and topic-safety include a $model reference in the rail flow text
+        # Need to match rails with flow_id (excluding $model reference) and match prompts
+        # on the full flow_id (including $model reference)
+        _validate_rail_prompts(
+            enabled_input_rails, provided_task_prompts, "content safety check input"
+        )
+        _validate_rail_prompts(
+            enabled_input_rails, provided_task_prompts, "topic safety check input"
+        )
+
         # Output moderation prompt verification
         if (
             "self check output" in enabled_output_rails
@@ -1504,6 +1573,13 @@ def check_prompt_exist_for_self_check_rails(cls, values):
         ):
             raise ValueError("You must provide a `self_check_facts` prompt template.")
 
+        # Only content-safety and topic-safety include a $model reference in the rail flow text
+        # Need to match rails with flow_id (excluding $model reference) and match prompts
+        # on the full flow_id (including $model reference)
+        _validate_rail_prompts(
+            enabled_output_rails, provided_task_prompts, "content safety check output"
+        )
+
         return values
 
     @root_validator(pre=True, allow_reuse=True)
@@ -1833,3 +1909,28 @@ def _generate_rails_flows(flows):
         flow_definitions.insert(1, _LIBRARY_IMPORT + _NEWLINE * 2)
 
     return flow_definitions
+
+
+MODEL_PREFIX = "$model="
+
+
+def _get_flow_model(flow_text) -> Optional[str]:
+    """Helper to return a model name from a flow definition"""
+    if MODEL_PREFIX not in flow_text:
+        return None
+    return flow_text.split(MODEL_PREFIX)[-1].strip()
+
+
+def _validate_rail_prompts(
+    rails: list[str], prompts: list[Any], validation_rail: str
+) -> None:
+    for rail in rails:
+        flow_id = _normalize_flow_id(rail)
+        flow_model = _get_flow_model(rail)
+        if flow_id == validation_rail:
+            prompt_flow_id = flow_id.replace(" ", "_")
+            expected_prompt = f"{prompt_flow_id} $model={flow_model}"
+            if expected_prompt not in prompts:
+                raise ValueError(
+                    f"You must provide a `{expected_prompt}` prompt template."
+                )

tests/test_internal_error_parallel_rails.py

@@ -78,48 +78,6 @@ async def test_internal_error_stops_execution():
     not _has_langchain_openai or not _has_openai_key,
     reason="langchain-openai not available",
 )
-@pytest.mark.asyncio
-async def test_content_safety_missing_prompt():
-    config_data = {
-        "instructions": [
-            {"type": "general", "content": "You are a helpful assistant."}
-        ],
-        "models": [
-            {"type": "main", "engine": "openai", "model": "gpt-3.5-turbo"},
-            {"type": "content_safety", "engine": "openai", "model": "gpt-3.5-turbo"},
-        ],
-        "rails": {
-            "input": {
-                "flows": [
-                    "content safety check input $model=content_safety",
-                    "self check input",
-                ],
-                "parallel": True,
-            }
-        },
-    }
-
-    config = RailsConfig.from_content(
-        config=config_data,
-        yaml_content="prompts:\n  - task: self_check_input\n    content: 'Is the user input safe? Answer Yes or No.'",
-    )
-
-    chat = TestChat(config, llm_completions=["Safe response"])
-    chat >> "test message"
-
-    result = await chat.app.generate_async(messages=chat.history, options=OPTIONS)
-
-    assert result is not None
-    assert "internal error" in result.response[0]["content"].lower()
-
-    stop_events = [
-        event
-        for event in result.log.internal_events
-        if event.get("type") == "BotIntent" and event.get("intent") == "stop"
-    ]
-    assert len(stop_events) > 0
-
-
 @pytest.mark.asyncio
 async def test_no_app_llm_request_on_internal_error():
     """Test that App LLM request is not sent when internal error occurs."""
@@ -164,48 +122,6 @@ async def test_no_app_llm_request_on_internal_error():
             ), "Expected BotIntent stop event after internal error"
 
 
-@pytest.mark.asyncio
-async def test_content_safety_missing_model():
-    """Test content safety with missing model configuration."""
-    config_data = {
-        "instructions": [
-            {"type": "general", "content": "You are a helpful assistant."}
-        ],
-        "models": [
-            {"type": "main", "engine": "openai", "model": "gpt-3.5-turbo"}
-            # missing content_safety model
-        ],
-        "rails": {
-            "input": {
-                "flows": ["content safety check input $model=content_safety"],
-                "parallel": True,
-            }
-        },
-    }
-
-    config = RailsConfig.from_content(
-        config=config_data,
-        yaml_content="prompts:\n  - task: content_safety_check_input $model=content_safety\n    content: 'Check if this is safe: {{ user_input }}'",
-    )
-
-    chat = TestChat(config, llm_completions=["Response"])
-    chat >> "test message"
-
-    result = await chat.app.generate_async(messages=chat.history, options=OPTIONS)
-
-    # should get internal error due to missing model
-    assert result is not None
-    assert "internal error" in result.response[0]["content"].lower()
-
-    # verify stop event was generated
-    stop_events = [
-        event
-        for event in result.log.internal_events
-        if event.get("type") == "BotIntent" and event.get("intent") == "stop"
-    ]
-    assert len(stop_events) > 0
-
-
 @pytest.mark.asyncio
 async def test_parallel_rails_partial_failure():
     """Test that partial failure in parallel rails is handled properly."""
@@ -343,21 +259,28 @@ async def test_action_execution_returns_failed():
         ), "Expected BotIntent stop event after action failure"
 
 
+@pytest.mark.skipif(
+    not _has_langchain_openai or not _has_openai_key,
+    reason="langchain-openai not available",
+)
 @pytest.mark.asyncio
 async def test_single_error_message_not_multiple():
     """Test that we get exactly one error message, not multiple for each failed rail.
 
     Before the fix, if we had multiple rails failing, we'd get multiple error messages.
     This test verifies we only get one error message even with multiple parallel rails.
+    Now with config-time validation, we provide valid config and trigger runtime failures.
     """
     config_data = {
-        "models": [{"type": "main", "engine": "openai", "model": "gpt-3.5-turbo"}],
+        "models": [
+            {"type": "main", "engine": "openai", "model": "gpt-3.5-turbo"},
+            {"type": "content_safety", "engine": "openai", "model": "gpt-3.5-turbo"},
+        ],
         "rails": {
             "input": {
                 "flows": [
                     "self check input",
                     "content safety check input $model=content_safety",
-                    "llama guard check input $model=llama_guard",
                 ],
                 "parallel": True,
             }
@@ -366,44 +289,56 @@ async def test_single_error_message_not_multiple():
             {
                 "task": "self_check_input",
                 "content": "Is the user input safe? Answer Yes or No.",
-            }
-            # missing prompts and models for content_safety and llama_guard
+            },
+            {
+                "task": "content_safety_check_input $model=content_safety",
+                "content": "Check content safety: {{ user_input }}",
+            },
         ],
     }
 
     config = RailsConfig.from_content(config=config_data)
 
-    chat = TestChat(config, llm_completions=["Test response"])
-    chat >> "test message"
+    with patch(
+        "nemoguardrails.llm.taskmanager.LLMTaskManager.render_task_prompt"
+    ) as mock_render:
+        mock_render.side_effect = Exception("Runtime error in multiple rails")
 
-    result = await chat.app.generate_async(messages=chat.history, options=OPTIONS)
+        chat = TestChat(config, llm_completions=["Test response"])
+        chat >> "test message"
 
-    # should get exactly one response, not multiple
-    assert result is not None
-    assert len(result.response) == 1, f"Expected 1 response, got {len(result.response)}"
-
-    # that single response should be an internal error
-    assert "internal error" in result.response[0]["content"].lower()
-
-    # count how many times "internal error" appears in the response
-    error_count = result.response[0]["content"].lower().count("internal error")
-    assert error_count == 1, f"Expected 1 'internal error' message, found {error_count}"
-
-    # verify stop event was generated
-    stop_events = [
-        event
-        for event in result.log.internal_events
-        if event.get("type") == "BotIntent" and event.get("intent") == "stop"
-    ]
-    assert len(stop_events) >= 1, "Expected at least one BotIntent stop event"
-
-    # verify we don't have multiple StartUtteranceBotAction events with error messages
-    error_utterances = [
-        event
-        for event in result.log.internal_events
-        if event.get("type") == "StartUtteranceBotAction"
-        and "internal error" in event.get("script", "").lower()
-    ]
-    assert (
-        len(error_utterances) == 1
-    ), f"Expected 1 error utterance, found {len(error_utterances)}"
+        result = await chat.app.generate_async(messages=chat.history, options=OPTIONS)
+
+        # should get exactly one response, not multiple
+        assert result is not None
+        assert (
+            len(result.response) == 1
+        ), f"Expected 1 response, got {len(result.response)}"
+
+        # that single response should be an internal error
+        assert "internal error" in result.response[0]["content"].lower()
+
+        # count how many times "internal error" appears in the response
+        error_count = result.response[0]["content"].lower().count("internal error")
+        assert (
+            error_count == 1
+        ), f"Expected 1 'internal error' message, found {error_count}"
+
+        # verify stop event was generated
+        stop_events = [
+            event
+            for event in result.log.internal_events
+            if event.get("type") == "BotIntent" and event.get("intent") == "stop"
+        ]
+        assert len(stop_events) >= 1, "Expected at least one BotIntent stop event"
+
+        # verify we don't have multiple StartUtteranceBotAction events with error messages
+        error_utterances = [
+            event
+            for event in result.log.internal_events
+            if event.get("type") == "StartUtteranceBotAction"
+            and "internal error" in event.get("script", "").lower()
+        ]
+        assert (
+            len(error_utterances) == 1
+        ), f"Expected 1 error utterance, found {len(error_utterances)}"