refactor(llm): remove LLMs isolation for actions (#1408)

* refactor(llm): remove isolated LLMs for actions

Removes the logic for creating isolated LLM instances for actions that
require an 'llm' parameter in `LLMRails`. This includes deleting the
`_create_isolated_llms_for_actions`, `_detect_llm_requiring_actions`,
`_get_action_function`, and `_create_action_llm_copy` methods, as well
as removing their invocation from the class initialization.

Also deletes all related tests, including unit, integration, and e2e
tests for LLM isolation and model_kwargs handling. This simplifies the
codebase and test suite by eliminating support for per-action LLM
isolation.

* Revert tests/utils.py to state before #1336

Author: Pouyanpi

nemoguardrails/rails/llm/llmrails.py

@@ -286,8 +286,6 @@ def __init__(
         # We also register the kb as a parameter that can be passed to actions.
         self.runtime.register_action_param("kb", self.kb)
 
-        # detect actions that need isolated LLM instances and create them
-        self._create_isolated_llms_for_actions()
         # Reference to the general ExplainInfo object.
         self.explain_info = None
 
@@ -511,147 +509,6 @@ def _init_llms(self):
 
         self.runtime.register_action_param("llms", llms)
 
-    def _create_isolated_llms_for_actions(self):
-        """Create isolated LLM copies for all actions that accept 'llm' parameter."""
-        if not self.llm:
-            log.debug("No main LLM available for creating isolated copies")
-            return
-
-        try:
-            actions_needing_llms = self._detect_llm_requiring_actions()
-            log.info(
-                "%d actions requiring isolated LLMs: %s",
-                len(actions_needing_llms),
-                list(actions_needing_llms),
-            )
-
-            created_count = 0
-
-            configured_actions_names = []
-            try:
-                if self.config.flows:
-                    get_action_details = partial(
-                        get_action_details_from_flow_id, flows=self.config.flows
-                    )
-                    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)
-                else:
-                    # for configurations without flow definitions, use all actions that need LLMs
-                    log.info(
-                        "No flow definitions found, creating isolated LLMs for all actions requiring them"
-                    )
-                    configured_actions_names = list(actions_needing_llms)
-            except Exception as e:
-                # if flow matching fails, fall back to all actions that need LLMs
-                log.info(
-                    "Flow matching failed (%s), creating isolated LLMs for all actions requiring them",
-                    e,
-                )
-                configured_actions_names = list(actions_needing_llms)
-
-            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:
-                        self.runtime.register_action_param(
-                            f"{action_name}_llm", isolated_llm
-                        )
-                        created_count += 1
-                        log.debug("Created isolated LLM for action: %s", action_name)
-                else:
-                    log.debug(
-                        "Action %s already has dedicated LLM, skipping isolation",
-                        action_name,
-                    )
-
-            log.info("Created %d isolated LLM instances for actions", created_count)
-
-        except Exception as e:
-            log.warning("Failed to create isolated LLMs for actions: %s", e)
-
-    def _detect_llm_requiring_actions(self):
-        """Auto-detect actions that have 'llm' parameter."""
-        import inspect
-
-        actions_needing_llms = set()
-
-        if (
-            not hasattr(self.runtime, "action_dispatcher")
-            or not self.runtime.action_dispatcher
-        ):
-            log.debug("Action dispatcher not available")
-            return actions_needing_llms
-
-        for (
-            action_name,
-            action_info,
-        ) in self.runtime.action_dispatcher.registered_actions.items():
-            action_func = self._get_action_function(action_info)
-            if not action_func:
-                continue
-
-            try:
-                sig = inspect.signature(action_func)
-                if "llm" in sig.parameters:
-                    actions_needing_llms.add(action_name)
-                    log.debug("Action %s has 'llm' parameter", action_name)
-
-            except Exception as e:
-                log.debug("Could not inspect action %s: %s", action_name, e)
-
-        return actions_needing_llms
-
-    def _get_action_function(self, action_info):
-        """Extract the actual function from action info."""
-        return action_info if callable(action_info) else None
-
-    def _create_action_llm_copy(
-        self, main_llm: Union[BaseLLM, BaseChatModel], action_name: str
-    ) -> Optional[Union[BaseLLM, BaseChatModel]]:
-        """Create an isolated copy of main LLM for a specific action."""
-        import copy
-
-        try:
-            # shallow copy to preserve HTTP clients, credentials, etc.
-            # but create new instance to avoid shared state
-            isolated_llm = copy.copy(main_llm)
-
-            # isolate model_kwargs to prevent shared mutable state
-            if (
-                hasattr(isolated_llm, "model_kwargs")
-                and isolated_llm.model_kwargs is not None
-            ):
-                isolated_llm.model_kwargs = isolated_llm.model_kwargs.copy()
-
-            log.debug(
-                "Successfully created isolated LLM copy for action: %s", action_name
-            )
-            return isolated_llm
-
-        except Exception as e:
-            error_msg = (
-                "Failed to create isolated LLM instance for action '%s'. "
-                "This is required to prevent parameter contamination between different actions. "
-                "\n\nPossible solutions:"
-                "\n1. If using a custom LLM class, ensure it supports copy.copy() operation"
-                "\n2. Check that your LLM configuration doesn't contain non-copyable objects"
-                "\n3. Consider using a dedicated LLM configuration for action '%s'"
-                "\n\nOriginal error: %s"
-                "\n\nTo use a dedicated LLM for this action, add to your config:"
-                "\nmodels:"
-                "\n  - type: %s"
-                "\n    engine: <your_engine>"
-                "\n    model: <your_model>"
-            ) % (action_name, action_name, e, action_name)
-            log.error(error_msg)
-            raise RuntimeError(error_msg)
-
     def _get_embeddings_search_provider_instance(
         self, esp_config: Optional[EmbeddingSearchProvider] = None
     ) -> EmbeddingsIndex:

tests/runnable_rails/test_runnable_rails.py

@@ -633,11 +633,11 @@ def test_chat_prompt_template_with_runnable_rails_fixed():
 
     llm = FakeLLM(
         responses=[
-            "Hi!",
+            "no",
             "express greeting",
+            "no",
             "Welcome to our clinic! I'm so glad you're here.",
-            "Additional response",
-            "Another response",
+            "no",
         ]
     )