feat(llmrails): isolate LLMs only for configured actions (#1342)

Update isolated LLM creation to only target actions defined in rails config
flows. This ensures that LLMs are not unnecessarily created for actions not
present in the configuration. Adds comprehensive tests to verify correct
behavior, including handling of empty configs and skipping already-registered
LLMs.

Author: Pouyanpi

nemoguardrails/rails/llm/llmrails.py

@@ -523,7 +523,21 @@ def _create_isolated_llms_for_actions(self):
             )
 
             created_count = 0
-            for action_name in actions_needing_llms:
+            # Get the actions from flows defined in rails config
+            get_action_details = partial(
+                get_action_details_from_flow_id, flows=self.config.flows
+            )
+            configured_actions_names = []
+            for flow_id in self.config.rails.input.flows:
+                action_name, _ = get_action_details(flow_id)
+                configured_actions_names.append(action_name)
+            for flow_id in self.config.rails.output.flows:
+                action_name, _ = get_action_details(flow_id)
+                configured_actions_names.append(action_name)
+
+            for action_name in configured_actions_names:
+                if action_name not in actions_needing_llms:
+                    continue
                 if f"{action_name}_llm" not in self.runtime.registered_action_params:
                     isolated_llm = self._create_action_llm_copy(self.llm, action_name)
                     if isolated_llm:

tests/test_llm_isolation.py

@@ -255,48 +255,89 @@ def test_create_isolated_llms_for_actions_integration(self, rails_with_mock_llm)
         """Test the full isolated LLM creation process."""
         rails = rails_with_mock_llm
 
+        # Mock rails configuration with flows
+        rails.config.rails = Mock()
+        rails.config.rails.input = Mock()
+        rails.config.rails.output = Mock()
+        rails.config.rails.input.flows = ["input_flow_1", "input_flow_2"]
+        rails.config.rails.output.flows = ["output_flow_1"]
+
         rails.runtime = Mock()
         rails.runtime.action_dispatcher = MockActionDispatcher()
         rails.runtime.registered_action_params = {}
         rails.runtime.register_action_param = Mock()
 
-        rails._create_isolated_llms_for_actions()
-
-        expected_calls = [
+        # Mock get_action_details_from_flow_id to return actions that need LLMs
+        def mock_get_action_details(flow_id, flows):
+            mapping = {
+                "input_flow_1": ("action_with_llm", {}),
+                "input_flow_2": ("generate_user_intent", {}),
+                "output_flow_1": ("self_check_output", {}),
+            }
+            return mapping.get(flow_id, ("unknown_action", {}))
+
+        with patch(
+            "nemoguardrails.rails.llm.llmrails.get_action_details_from_flow_id",
+            side_effect=mock_get_action_details,
+        ):
+            rails._create_isolated_llms_for_actions()
+
+        expected_llm_params = [
             "action_with_llm_llm",
             "generate_user_intent_llm",
             "self_check_output_llm",
         ]
 
-        actual_calls = [
+        registered_llm_params = [
             call[0][0] for call in rails.runtime.register_action_param.call_args_list
         ]
 
-        for expected_call in expected_calls:
-            assert expected_call in actual_calls
+        for expected_param in expected_llm_params:
+            assert expected_param in registered_llm_params
 
     def test_create_isolated_llms_skips_existing_specialized_llms(
         self, rails_with_mock_llm
     ):
         """Test that existing specialized LLMs are not overridden."""
         rails = rails_with_mock_llm
 
+        # Mock rails configuration with flows
+        rails.config.rails = Mock()
+        rails.config.rails.input = Mock()
+        rails.config.rails.output = Mock()
+        rails.config.rails.input.flows = ["input_flow_1", "input_flow_2"]
+        rails.config.rails.output.flows = ["output_flow_1"]
+
         rails.runtime = Mock()
         rails.runtime.action_dispatcher = MockActionDispatcher()
         rails.runtime.registered_action_params = {"self_check_output_llm": Mock()}
         rails.runtime.register_action_param = Mock()
 
-        rails._create_isolated_llms_for_actions()
-
-        # verify self_check_output_llm was NOT re-registered
-        actual_calls = [
+        # Mock get_action_details_from_flow_id to return actions that need LLMs
+        def mock_get_action_details(flow_id, flows):
+            mapping = {
+                "input_flow_1": ("action_with_llm", {}),
+                "input_flow_2": ("generate_user_intent", {}),
+                "output_flow_1": (
+                    "self_check_output",
+                    {},
+                ),  # This one already has an LLM
+            }
+            return mapping.get(flow_id, ("unknown_action", {}))
+
+        with patch(
+            "nemoguardrails.rails.llm.llmrails.get_action_details_from_flow_id",
+            side_effect=mock_get_action_details,
+        ):
+            rails._create_isolated_llms_for_actions()
+
+        registered_llm_params = [
             call[0][0] for call in rails.runtime.register_action_param.call_args_list
         ]
-        assert "self_check_output_llm" not in actual_calls
 
-        # but other actions should still get isolated LLMs
-        assert "action_with_llm_llm" in actual_calls
-        assert "generate_user_intent_llm" in actual_calls
+        assert "self_check_output_llm" not in registered_llm_params
+        assert "action_with_llm_llm" in registered_llm_params
+        assert "generate_user_intent_llm" in registered_llm_params
 
     def test_create_isolated_llms_handles_no_main_llm(self, mock_config):
         """Test graceful handling when no main LLM is available."""
@@ -411,3 +452,87 @@ def test_action_detection_parametrized(
             assert action_name in actions_needing_llms
         else:
             assert action_name not in actions_needing_llms
+
+    def test_create_isolated_llms_for_configured_actions_only(
+        self, rails_with_mock_llm
+    ):
+        """Test that isolated LLMs are created only for actions configured in rails flows."""
+        rails = rails_with_mock_llm
+
+        rails.config.rails = Mock()
+        rails.config.rails.input = Mock()
+        rails.config.rails.output = Mock()
+        rails.config.rails.input.flows = [
+            "input_flow_1",
+            "input_flow_2",
+            "input_flow_3",
+        ]
+        rails.config.rails.output.flows = ["output_flow_1", "output_flow_2"]
+
+        rails.runtime = Mock()
+        rails.runtime.action_dispatcher = MockActionDispatcher()
+        rails.runtime.registered_action_params = {}
+        rails.runtime.register_action_param = Mock()
+
+        def mock_get_action_details(flow_id, flows):
+            mapping = {
+                "input_flow_1": ("action_with_llm", {}),
+                "input_flow_2": ("action_without_llm", {}),
+                "input_flow_3": ("self_check_output", {}),
+                "output_flow_1": ("generate_user_intent", {}),
+                "output_flow_2": ("non_configured_action", {}),
+            }
+            return mapping.get(flow_id, ("unknown_action", {}))
+
+        with patch(
+            "nemoguardrails.rails.llm.llmrails.get_action_details_from_flow_id",
+            side_effect=mock_get_action_details,
+        ):
+            rails._create_isolated_llms_for_actions()
+
+        registered_llm_params = [
+            call[0][0] for call in rails.runtime.register_action_param.call_args_list
+        ]
+
+        expected_isolated_llm_params = [
+            "action_with_llm_llm",
+            "generate_user_intent_llm",
+            "self_check_output_llm",
+        ]
+
+        for expected_param in expected_isolated_llm_params:
+            assert (
+                expected_param in registered_llm_params
+            ), f"Expected {expected_param} to be registered as action param"
+
+        assert "action_without_llm_llm" not in registered_llm_params
+        assert "non_configured_action_llm" not in registered_llm_params
+
+        assert len(registered_llm_params) == 3, (
+            f"Should only create isolated LLMs for actions from config flows that need LLMs. "
+            f"Got {registered_llm_params}"
+        )
+
+    def test_create_isolated_llms_handles_empty_rails_config(self, rails_with_mock_llm):
+        """Test that the method handles empty rails configuration gracefully."""
+        rails = rails_with_mock_llm
+
+        rails.config.rails = Mock()
+        rails.config.rails.input = Mock()
+        rails.config.rails.output = Mock()
+        rails.config.rails.input.flows = []
+        rails.config.rails.output.flows = []
+
+        rails.runtime = Mock()
+        rails.runtime.action_dispatcher = MockActionDispatcher()
+        rails.runtime.registered_action_params = {}
+        rails.runtime.register_action_param = Mock()
+
+        with patch(
+            "nemoguardrails.rails.llm.llmrails.get_action_details_from_flow_id"
+        ) as mock_get_action:
+            rails._create_isolated_llms_for_actions()
+
+        mock_get_action.assert_not_called()
+
+        rails.runtime.register_action_param.assert_not_called()

tests/test_llm_isolation_e2e.py

@@ -25,7 +25,7 @@
 from nemoguardrails import LLMRails
 from nemoguardrails.rails.llm.config import RailsConfig
 
-TEST_LIVE_MODE = os.environ.get("TEST_LIVE_MODE")
+LIVE_TEST_MODE = os.environ.get("TEST_LIVE_MODE")
 
 
 @pytest.fixture
@@ -75,7 +75,7 @@ def test_config_path(test_config_content):
 
 
 @pytest.mark.skipif(
-    not TEST_LIVE_MODE,
+    not LIVE_TEST_MODE,
     reason="This test requires TEST_LIVE_MODE environment variable to be set for live testing",
 )
 class TestLLMIsolationE2E:
@@ -383,7 +383,7 @@ async def test_parameter_isolation_multiple_iterations(
 
 
 @pytest.mark.skipif(
-    not TEST_LIVE_MODE,
+    not LIVE_TEST_MODE,
     reason="This test requires TEST_LIVE_MODE environment variable to be set for live testing",
 )
 class TestLLMIsolationErrorHandling:
@@ -474,5 +474,135 @@ async def run_parameter_contamination_test():
         )
 
 
+@pytest.mark.skipif(
+    not LIVE_TEST_MODE,
+    reason="This test requires TEST_LIVE_MODE environment variable to be set for live testing",
+)
+class TestLLMIsolationConfiguredActionsOnly:
+    """Test that isolated LLMs are created only for actions configured in rails flows."""
+
+    @staticmethod
+    def _create_rails_with_config(config_content: str) -> LLMRails:
+        """Helper to create LLMRails instance from config content."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            config_path = Path(temp_dir) / "config.yml"
+            config_path.write_text(config_content)
+            config = RailsConfig.from_path(str(temp_dir))
+            return LLMRails(config, verbose=False)
+
+    @staticmethod
+    def _get_isolated_llm_params(
+        rails: LLMRails, exclude_specialized: bool = False
+    ) -> list:
+        """Helper to get isolated LLM parameters from rails instance."""
+        registered_params = rails.runtime.registered_action_params
+        isolated_llm_params = [
+            key
+            for key in registered_params.keys()
+            if key.endswith("_llm") and key != "llm" and key != "llms"
+        ]
+
+        if exclude_specialized:
+            specialized_llms = ["content_safety_llm", "topic_safety_llm"]
+            isolated_llm_params = [
+                param for param in isolated_llm_params if param not in specialized_llms
+            ]
+
+        return isolated_llm_params
+
+    def test_only_configured__rail_actions_get_isolated_llms(self):
+        """Test that only actions from output rails flows get isolated LLMs."""
+        config_content = """
+        models:
+          - type: main
+            engine: openai
+            model: gpt-4o-mini
+
+        rails:
+          output:
+            flows:
+              - self check output
+              - self check input
+
+        prompts:
+          - task: self_check_output
+            content: |
+              Check if output is safe.
+              Output: {{ bot_message }}
+              Safe? (Yes/No):
+          - task: self_check_input
+            content: |
+              Check if input is safe.
+              Input: {{ user_input }}
+              Safe? (Yes/No):
+        """
+
+        rails = self._create_rails_with_config(config_content)
+        isolated_llm_params = self._get_isolated_llm_params(rails)
+
+        assert "self_check_output_llm" in isolated_llm_params
+        assert "self_check_input_llm" in isolated_llm_params
+        assert "self_check_facts_llm" not in isolated_llm_params
+
+    def test_no_isolated_llms_when_no_rails_configured(self):
+        """Test that no isolated LLMs are created when no rails are configured."""
+        config_content = """
+        models:
+          - type: main
+            engine: openai
+            model: gpt-4o-mini
+        """
+
+        rails = self._create_rails_with_config(config_content)
+        isolated_llm_params = self._get_isolated_llm_params(
+            rails, exclude_specialized=True
+        )
+
+        assert (
+            len(isolated_llm_params) == 0
+        ), f"Unexpected isolated LLMs created: {isolated_llm_params}"
+
+    def test_empty_rails_flows_creates_no_isolated_llms(self):
+        """Test that empty rails flows list creates no isolated LLMs."""
+        config_content = """
+        models:
+          - type: main
+            engine: openai
+            model: gpt-4o-mini
+
+        rails:
+          input:
+            flows: []
+          output:
+            flows: []
+        """
+
+        rails = self._create_rails_with_config(config_content)
+        isolated_llm_params = self._get_isolated_llm_params(
+            rails, exclude_specialized=True
+        )
+
+        assert (
+            len(isolated_llm_params) == 0
+        ), f"Unexpected isolated LLMs created: {isolated_llm_params}"
+
+    def test_non_llm_requiring_actions_dont_get_isolated_llms(self):
+        """Test that even valid flows don't get isolated LLMs if actions don't require LLMs."""
+        config_content = """
+        models:
+          - type: main
+            engine: openai
+            model: gpt-4o-mini
+        """
+
+        rails = self._create_rails_with_config(config_content)
+
+        # retrieve_relevant_chunks action exists but doesn't require LLM
+        # so it should never get an isolated LLM even if it were configured
+        assert (
+            "retrieve_relevant_chunks_llm" not in rails.runtime.registered_action_params
+        )
+
+
 if __name__ == "__main__":
     asyncio.run(run_parameter_contamination_test())