rasa-pro 3.11.0rc1__py3-none-any.whl → 3.11.0rc2__py3-none-any.whl

rasa/core/policies/intentless_policy.py

@@ -18,10 +18,6 @@ from rasa.core.constants import (
     UTTER_SOURCE_METADATA_KEY,
 )
 from rasa.core.policies.policy import Policy, PolicyPrediction, SupportedData
-from rasa.dialogue_understanding.generator.constants import (
-    TRAINED_MODEL_NAME_CONFIG_KEY,
-    TRAINED_EMBEDDINGS_CONFIG_KEY,
-)
 from rasa.dialogue_understanding.patterns.chitchat import FLOW_PATTERN_CHITCHAT
 from rasa.dialogue_understanding.stack.frames import (
     ChitChatStackFrame,
@@ -64,6 +60,10 @@ from rasa.shared.providers.embedding._langchain_embedding_client_adapter import
     _LangchainEmbeddingClientAdapter,
 )
 from rasa.shared.providers.llm.llm_client import LLMClient
+from rasa.shared.utils.health_check.embeddings_health_check_mixin import (
+    EmbeddingsHealthCheckMixin,
+)
+from rasa.shared.utils.health_check.llm_health_check_mixin import LLMHealthCheckMixin
 from rasa.shared.utils.io import deep_container_fingerprint
 from rasa.shared.utils.llm import (
     AI,
@@ -79,12 +79,6 @@ from rasa.shared.utils.llm import (
     tracker_as_readable_transcript,
     resolve_model_client_config,
 )
-from rasa.shared.utils.health_check import (
-    perform_training_time_llm_health_check,
-    perform_training_time_embeddings_health_check,
-    perform_inference_time_llm_health_check,
-    perform_inference_time_embeddings_health_check,
-)
 from rasa.utils.log_utils import log_llm
 from rasa.utils.ml_utils import (
     extract_ai_response_examples,
@@ -383,7 +377,7 @@ def conversation_as_prompt(conversation: Conversation) -> str:
 @DefaultV1Recipe.register(
     DefaultV1Recipe.ComponentType.POLICY_WITH_END_TO_END_SUPPORT, is_trainable=True
 )
-class IntentlessPolicy(Policy):
+class IntentlessPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Policy):
     """Policy which uses a language model to generate the next action.
 
     The policy uses the OpenAI API to generate the next action based on the
@@ -516,10 +510,8 @@ class IntentlessPolicy(Policy):
             A policy must return its resource locator so that potential children nodes
             can load the policy from the resource.
         """
-        (
-            self.config[TRAINED_MODEL_NAME_CONFIG_KEY],
-            self.config[TRAINED_EMBEDDINGS_CONFIG_KEY],
-        ) = self._perform_training_time_health_checks()
+        # Perform health checks of both LLM and embeddings client configs
+        self._perform_health_checks(self.config, "intentless_policy.train")
 
         responses = filter_responses(responses, forms, flows or FlowsList([]))
         telemetry.track_intentless_policy_train()
@@ -952,10 +944,13 @@ class IntentlessPolicy(Policy):
         **kwargs: Any,
     ) -> "IntentlessPolicy":
         """Loads a trained policy (see parent class for full docstring)."""
+
+        # Perform health checks of both LLM and embeddings client configs
+        cls._perform_health_checks(config, "intentless_policy.load")
+
         responses_docsearch = None
         samples_docsearch = None
         prompt_template = None
-        persisted_config = None
         try:
             with model_storage.read_from(resource) as path:
                 responses_docsearch = load_faiss_vector_store(
@@ -973,15 +968,12 @@ class IntentlessPolicy(Policy):
                 prompt_template = rasa.shared.utils.io.read_file(
                     path / INTENTLESS_PROMPT_TEMPLATE_FILE_NAME
                 )
-                persisted_config = rasa.shared.utils.io.read_json_file(
-                    path / INTENTLESS_CONFIG_FILE_NAME
-                )
         except (ValueError, FileNotFoundError, FileIOException) as e:
             structlogger.warning(
                 "intentless_policy.load.failed", error=e, resource_name=resource.name
             )
 
-        policy = cls(
+        return cls(
             config,
             model_storage,
             resource,
@@ -991,14 +983,6 @@ class IntentlessPolicy(Policy):
             prompt_template=prompt_template,
         )
 
-        cls._perform_inference_time_health_checks(
-            persisted_config,
-            policy.config.get(LLM_CONFIG_KEY),
-            policy.config.get(EMBEDDINGS_CONFIG_KEY),
-        )
-
-        return policy
-
     @classmethod
     def fingerprint_addon(cls, config: Dict[str, Any]) -> Optional[str]:
         """Add a fingerprint of intentless policy for the graph."""
@@ -1018,52 +1002,26 @@ class IntentlessPolicy(Policy):
             [prompt_template, llm_config, embedding_config]
         )
 
-    def _perform_training_time_health_checks(
-        self,
-    ) -> Tuple[Optional[str], Optional[str]]:
-        train_model_name = perform_training_time_llm_health_check(
-            self.config.get(LLM_CONFIG_KEY),
-            DEFAULT_LLM_CONFIG,
-            "intentless_policy.train",
-            IntentlessPolicy.__name__,
-        )
-        train_embedding_name = perform_training_time_embeddings_health_check(
-            self.config.get(EMBEDDINGS_CONFIG_KEY),
-            DEFAULT_EMBEDDINGS_CONFIG,
-            "intentless_policy.train",
-            IntentlessPolicy.__name__,
-        )
-        return train_model_name, train_embedding_name
-
     @classmethod
-    def _perform_inference_time_health_checks(
-        cls,
-        persisted_config: Optional[Dict[str, Any]],
-        resolved_llm_config: Optional[Dict[str, Any]],
-        resolved_embeddings_config: Optional[Dict[str, Any]],
+    def _perform_health_checks(
+        cls, config: Dict[Text, Any], log_source_method: str
     ) -> None:
-        train_model_name = (
-            persisted_config.get(TRAINED_MODEL_NAME_CONFIG_KEY, None)
-            if persisted_config
-            else None
-        )
-        perform_inference_time_llm_health_check(
-            resolved_llm_config,
+        # Perform health check of the LLM client config
+        llm_config = resolve_model_client_config(config.get(LLM_CONFIG_KEY, {}))
+        cls.perform_llm_health_check(
+            llm_config,
             DEFAULT_LLM_CONFIG,
-            train_model_name,
-            "intentless_policy.load",
+            log_source_method,
             IntentlessPolicy.__name__,
         )
 
-        train_embeddings_name = (
-            persisted_config.get(TRAINED_EMBEDDINGS_CONFIG_KEY, None)
-            if persisted_config
-            else None
+        # Perform health check of the embeddings client config
+        embeddings_config = resolve_model_client_config(
+            config.get(EMBEDDINGS_CONFIG_KEY, {})
         )
-        perform_inference_time_embeddings_health_check(
-            resolved_embeddings_config,
+        cls.perform_embeddings_health_check(
+            embeddings_config,
             DEFAULT_EMBEDDINGS_CONFIG,
-            train_embeddings_name,
-            "intentless_policy.load",
+            log_source_method,
             IntentlessPolicy.__name__,
         )

rasa/dialogue_understanding/coexistence/llm_based_router.py

@@ -17,7 +17,6 @@ from rasa.dialogue_understanding.commands import Command, SetSlotCommand
 from rasa.dialogue_understanding.commands.noop_command import NoopCommand
 from rasa.dialogue_understanding.generator.constants import (
     LLM_CONFIG_KEY,
-    TRAINED_MODEL_NAME_CONFIG_KEY,
 )
 from rasa.engine.graph import ExecutionContext, GraphComponent
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
@@ -36,6 +35,7 @@ from rasa.shared.exceptions import InvalidConfigException, FileIOException
 from rasa.shared.nlu.constants import COMMANDS, TEXT
 from rasa.shared.nlu.training_data.message import Message
 from rasa.shared.nlu.training_data.training_data import TrainingData
+from rasa.shared.utils.health_check.llm_health_check_mixin import LLMHealthCheckMixin
 from rasa.shared.utils.io import deep_container_fingerprint
 from rasa.shared.utils.llm import (
     DEFAULT_OPENAI_CHAT_MODEL_NAME,
@@ -43,10 +43,6 @@ from rasa.shared.utils.llm import (
     llm_factory,
     resolve_model_client_config,
 )
-from rasa.shared.utils.health_check import (
-    perform_training_time_llm_health_check,
-    perform_inference_time_llm_health_check,
-)
 from rasa.utils.log_utils import log_llm
 
 LLM_BASED_ROUTER_PROMPT_FILE_NAME = "llm_based_router_prompt.jinja2"
@@ -80,7 +76,7 @@ structlogger = structlog.get_logger()
     ],
     is_trainable=True,
 )
-class LLMBasedRouter(GraphComponent):
+class LLMBasedRouter(LLMHealthCheckMixin, GraphComponent):
     @staticmethod
     def get_default_config() -> Dict[str, Any]:
         """The component's default config (see parent class for full docstring)."""
@@ -144,13 +140,11 @@ class LLMBasedRouter(GraphComponent):
 
     def train(self, training_data: TrainingData) -> Resource:
         """Train the intent classifier on a data set."""
-        self.config[TRAINED_MODEL_NAME_CONFIG_KEY] = (
-            perform_training_time_llm_health_check(
-                self.config.get(LLM_CONFIG_KEY),
-                DEFAULT_LLM_CONFIG,
-                "llm_based_router.train",
-                LLMBasedRouter.__name__,
-            )
+        self.perform_llm_health_check(
+            self.config.get(LLM_CONFIG_KEY),
+            DEFAULT_LLM_CONFIG,
+            "llm_based_router.train",
+            LLMBasedRouter.__name__,
         )
 
         self.persist()
@@ -166,37 +160,28 @@ class LLMBasedRouter(GraphComponent):
         **kwargs: Any,
     ) -> "LLMBasedRouter":
         """Loads trained component (see parent class for full docstring)."""
+
+        # Perform health check on the resolved LLM client config
+        llm_config = resolve_model_client_config(config.get(LLM_CONFIG_KEY, {}))
+        cls.perform_llm_health_check(
+            llm_config,
+            DEFAULT_LLM_CONFIG,
+            "llm_based_router.load",
+            LLMBasedRouter.__name__,
+        )
+
         prompt_template = None
-        persisted_config = None
         try:
             with model_storage.read_from(resource) as path:
                 prompt_template = rasa.shared.utils.io.read_file(
                     path / LLM_BASED_ROUTER_PROMPT_FILE_NAME
                 )
-                persisted_config = rasa.shared.utils.io.read_json_file(
-                    path / LLM_BASED_ROUTER_CONFIG_FILE_NAME
-                )
         except (FileNotFoundError, FileIOException) as e:
             structlogger.warning(
                 "llm_based_router.load.failed", error=e, resource=resource.name
             )
 
-        router = cls(config, model_storage, resource, prompt_template=prompt_template)
-
-        train_model_name = (
-            persisted_config.get(TRAINED_MODEL_NAME_CONFIG_KEY, None)
-            if persisted_config
-            else None
-        )
-        perform_inference_time_llm_health_check(
-            router.config.get(LLM_CONFIG_KEY),
-            DEFAULT_LLM_CONFIG,
-            train_model_name,
-            "llm_based_router.load",
-            LLMBasedRouter.__name__,
-        )
-
-        return router
+        return cls(config, model_storage, resource, prompt_template=prompt_template)
 
     @classmethod
     def create(

rasa/dialogue_understanding/generator/constants.py

@@ -18,8 +18,6 @@ DEFAULT_LLM_CONFIG = {
 }
 
 LLM_CONFIG_KEY = "llm"
-TRAINED_MODEL_NAME_CONFIG_KEY = "trained_llm_model_name"
-TRAINED_EMBEDDINGS_CONFIG_KEY = "trained_embeddings_model_name"
 USER_INPUT_CONFIG_KEY = "user_input"
 
 FLOW_RETRIEVAL_KEY = "flow_retrieval"

rasa/dialogue_understanding/generator/flow_retrieval.py

@@ -27,12 +27,9 @@ from langchain.schema.embeddings import Embeddings
 from langchain_community.vectorstores.faiss import FAISS
 from langchain_community.vectorstores.utils import DistanceStrategy
 
-from rasa.dialogue_understanding.generator.constants import (
-    TRAINED_EMBEDDINGS_CONFIG_KEY,
-)
 from rasa.engine.storage.resource import Resource
 from rasa.engine.storage.storage import ModelStorage
-
+import rasa.shared.utils.io
 from rasa.shared.constants import (
     EMBEDDINGS_CONFIG_KEY,
     PROVIDER_CONFIG_KEY,
@@ -41,12 +38,15 @@ from rasa.shared.constants import (
 from rasa.shared.core.domain import Domain
 from rasa.shared.core.flows import FlowsList
 from rasa.shared.core.trackers import DialogueStateTracker
-from rasa.shared.exceptions import ProviderClientAPIException, FileIOException
+from rasa.shared.exceptions import ProviderClientAPIException
 from rasa.shared.nlu.constants import TEXT, FLOWS_FROM_SEMANTIC_SEARCH
 from rasa.shared.nlu.training_data.message import Message
 from rasa.shared.providers.embedding._langchain_embedding_client_adapter import (
     _LangchainEmbeddingClientAdapter,
 )
+from rasa.shared.utils.health_check.embeddings_health_check_mixin import (
+    EmbeddingsHealthCheckMixin,
+)
 from rasa.shared.utils.llm import (
     tracker_as_readable_transcript,
     embedder_factory,
@@ -56,11 +56,6 @@ from rasa.shared.utils.llm import (
     allowed_values_for_slot,
     resolve_model_client_config,
 )
-from rasa.shared.utils.health_check import (
-    perform_training_time_embeddings_health_check,
-    perform_inference_time_embeddings_health_check,
-)
-from rasa.shared.utils.io import dump_obj_as_json_to_file, read_json_file
 
 DEFAULT_FLOW_DOCUMENT_TEMPLATE = importlib.resources.read_text(
     "rasa.dialogue_understanding.generator", "flow_document_template.jinja2"
@@ -85,7 +80,7 @@ DEFAULT_SHOULD_EMBED_SLOTS = True
 structlogger = structlog.get_logger()
 
 
-class FlowRetrieval:
+class FlowRetrieval(EmbeddingsHealthCheckMixin):
     @classmethod
     def get_default_config(cls) -> Dict[str, Any]:
         """The default config for the flow retrieval."""
@@ -94,7 +89,6 @@ class FlowRetrieval:
             MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY: DEFAULT_MAX_FLOWS_FROM_SEMANTIC_SEARCH,
             TURNS_TO_EMBED_KEY: DEFAULT_TURNS_TO_EMBED,
             SHOULD_EMBED_SLOTS_KEY: DEFAULT_SHOULD_EMBED_SLOTS,
-            TRAINED_EMBEDDINGS_CONFIG_KEY: None,
         }
 
     def __init__(
@@ -147,16 +141,6 @@ class FlowRetrieval:
 
         return config
 
-    def train(self) -> None:
-        self.config[TRAINED_EMBEDDINGS_CONFIG_KEY] = (
-            perform_training_time_embeddings_health_check(
-                self.config.get(EMBEDDINGS_CONFIG_KEY),
-                DEFAULT_EMBEDDINGS_CONFIG,
-                "flow_retrieval.train",
-                FlowRetrieval.__name__,
-            )
-        )
-
     @classmethod
     def load(
         cls,
@@ -166,6 +150,18 @@ class FlowRetrieval:
         **kwargs: Any,
     ) -> "FlowRetrieval":
         """Load flow retrieval with previously populated FAISS vector store."""
+
+        # Perform health check on resolved embedding client config
+        embeddings_config = resolve_model_client_config(
+            config.get(EMBEDDINGS_CONFIG_KEY, {})
+        )
+        cls.perform_embeddings_health_check(
+            embeddings_config,
+            DEFAULT_EMBEDDINGS_CONFIG,
+            "flow_retrieval.load",
+            FlowRetrieval.__name__,
+        )
+
         # initialize base flow retrieval
         flow_retrieval = FlowRetrieval(config, model_storage, resource)
         # load vector store
@@ -174,30 +170,6 @@ class FlowRetrieval:
         )
         flow_retrieval.vector_store = vector_store
 
-        persisted_config = None
-        try:
-            with model_storage.read_from(resource) as path:
-                persisted_config = read_json_file(
-                    path / FLOW_RETRIEVAL_CONFIG_FILE_NAME
-                )
-        except (FileNotFoundError, FileIOException) as e:
-            structlogger.warning(
-                "flow_retrieval.load.failed", error=e, resource=resource.name
-            )
-
-        train_embeddings_name = (
-            persisted_config.get(TRAINED_EMBEDDINGS_CONFIG_KEY, None)
-            if persisted_config
-            else None
-        )
-        perform_inference_time_embeddings_health_check(
-            flow_retrieval.config.get(EMBEDDINGS_CONFIG_KEY),
-            DEFAULT_EMBEDDINGS_CONFIG,
-            train_embeddings_name,
-            "flow_retrieval.load",
-            FlowRetrieval.__name__,
-        )
-
         return flow_retrieval
 
     @classmethod
@@ -243,10 +215,7 @@ class FlowRetrieval:
 
     def persist(self) -> None:
         self._persist_vector_store()
-        with self._model_storage.write_to(self._resource) as path:
-            dump_obj_as_json_to_file(
-                path / FLOW_RETRIEVAL_CONFIG_FILE_NAME, self.config
-            )
+        self._persist_config()
 
     def _persist_vector_store(self) -> None:
         """Persists the FAISS vector store."""
@@ -259,6 +228,12 @@ class FlowRetrieval:
                 event_info="Vector store is None, not persisted.",
             )
 
+    def _persist_config(self) -> None:
+        with self._model_storage.write_to(self._resource) as path:
+            rasa.shared.utils.io.dump_obj_as_json_to_file(
+                path / FLOW_RETRIEVAL_CONFIG_FILE_NAME, self.config
+            )
+
     def populate(self, flows: FlowsList, domain: Domain) -> None:
         """Populates the vector store with embeddings generated from
         documents based on the flow descriptions, and flow slots
@@ -268,6 +243,14 @@ class FlowRetrieval:
             flows: List of flows to populate the vector store with.
             domain: The domain containing relevant slot information.
         """
+        # Perform health check before populating the vector store with flows
+        self.perform_embeddings_health_check(
+            self.config.get(EMBEDDINGS_CONFIG_KEY),
+            DEFAULT_EMBEDDINGS_CONFIG,
+            "flow_retrieval.train",
+            FlowRetrieval.__name__,
+        )
+
         flows_to_embedd = flows.exclude_link_only_flows()
         embeddings = self._create_embedder(self.config)
         documents = self._generate_flow_documents(flows_to_embedd, domain)

rasa/dialogue_understanding/generator/llm_based_command_generator.py

@@ -17,7 +17,6 @@ from rasa.dialogue_understanding.generator.constants import (
     FLOW_RETRIEVAL_KEY,
     FLOW_RETRIEVAL_ACTIVE_KEY,
     FLOW_RETRIEVAL_FLOW_THRESHOLD,
-    TRAINED_MODEL_NAME_CONFIG_KEY,
 )
 from rasa.dialogue_understanding.generator.flow_retrieval import FlowRetrieval
 from rasa.engine.graph import GraphComponent, ExecutionContext
@@ -33,27 +32,26 @@ from rasa.shared.exceptions import ProviderClientAPIException
 from rasa.shared.nlu.constants import FLOWS_IN_PROMPT
 from rasa.shared.nlu.training_data.message import Message
 from rasa.shared.nlu.training_data.training_data import TrainingData
+from rasa.shared.utils.health_check.llm_health_check_mixin import LLMHealthCheckMixin
 from rasa.shared.utils.llm import (
     allowed_values_for_slot,
     llm_factory,
     resolve_model_client_config,
 )
-from rasa.shared.utils.health_check import perform_training_time_llm_health_check
 from rasa.utils.log_utils import log_llm
 
 structlogger = structlog.get_logger()
 
 
-LLM_BASED_COMMAND_GENERATOR_CONFIG_FILE = "config.json"
-
-
 @DefaultV1Recipe.register(
     [
         DefaultV1Recipe.ComponentType.COMMAND_GENERATOR,
     ],
     is_trainable=True,
 )
-class LLMBasedCommandGenerator(GraphComponent, CommandGenerator, ABC):
+class LLMBasedCommandGenerator(
+    LLMHealthCheckMixin, GraphComponent, CommandGenerator, ABC
+):
     """An abstract class defining interface and common functionality
     of an LLM-based command generators.
     """
@@ -106,11 +104,7 @@ class LLMBasedCommandGenerator(GraphComponent, CommandGenerator, ABC):
     @abstractmethod
     def persist(self) -> None:
         """Persist the component to disk for future loading."""
-        # persist the config to store the resolved llm and embedding config
-        with self._model_storage.write_to(self._resource) as path:
-            rasa.shared.utils.io.dump_obj_as_json_to_file(
-                path / LLM_BASED_COMMAND_GENERATOR_CONFIG_FILE, self.config
-            )
+        pass
 
     @abstractmethod
     async def predict_commands(
@@ -173,13 +167,11 @@ class LLMBasedCommandGenerator(GraphComponent, CommandGenerator, ABC):
         """Train the llm based command generator. Stores all flows into a vector
         store.
         """
-        self.config[TRAINED_MODEL_NAME_CONFIG_KEY] = (
-            perform_training_time_llm_health_check(
-                self.config.get(LLM_CONFIG_KEY),
-                DEFAULT_LLM_CONFIG,
-                "llm_based_command_generator.train",
-                LLMBasedCommandGenerator.__name__,
-            )
+        self.perform_llm_health_check(
+            self.config.get(LLM_CONFIG_KEY),
+            DEFAULT_LLM_CONFIG,
+            "llm_based_command_generator.train",
+            LLMBasedCommandGenerator.__name__,
         )
 
         if (
@@ -210,12 +202,11 @@ class LLMBasedCommandGenerator(GraphComponent, CommandGenerator, ABC):
         except Exception as e:
             structlogger.error(
                 "llm_based_command_generator.train.failed",
-                event_info="Flow retrieval store isinaccessible.",
+                event_info="Flow retrieval store is inaccessible.",
                 error=e,
             )
             raise
-        if self.flow_retrieval is not None:
-            self.flow_retrieval.train()
+
         self.persist()
         return self._resource
 
@@ -251,25 +242,6 @@ class LLMBasedCommandGenerator(GraphComponent, CommandGenerator, ABC):
             )
         return None
 
-    @classmethod
-    def load_config_from_model_storage(
-        cls,
-        model_storage: ModelStorage,
-        resource: Resource,
-    ) -> Optional[Text]:
-        try:
-            with model_storage.read_from(resource) as path:
-                return rasa.shared.utils.io.read_json_file(
-                    path / LLM_BASED_COMMAND_GENERATOR_CONFIG_FILE
-                )
-        except (FileNotFoundError, FileIOException) as e:
-            structlogger.warning(
-                "llm_based_command_generator.load_config.failed",
-                error=e,
-                resource=resource.name,
-            )
-        return None
-
     @classmethod
     def load_flow_retrival(
         cls,

rasa/dialogue_understanding/generator/multi_step/multi_step_llm_command_generator.py

@@ -24,7 +24,6 @@ from rasa.dialogue_understanding.generator.constants import (
     LLM_CONFIG_KEY,
     USER_INPUT_CONFIG_KEY,
     FLOW_RETRIEVAL_KEY,
-    TRAINED_MODEL_NAME_CONFIG_KEY,
     DEFAULT_LLM_CONFIG,
 )
 from rasa.dialogue_understanding.generator.flow_retrieval import FlowRetrieval
@@ -60,7 +59,6 @@ from rasa.shared.utils.llm import (
     allowed_values_for_slot,
     resolve_model_client_config,
 )
-from rasa.shared.utils.health_check import perform_inference_time_llm_health_check
 
 # multistep template keys
 HANDLE_FLOWS_KEY = "handle_flows"
@@ -77,6 +75,7 @@ DEFAULT_HANDLE_FLOWS_TEMPLATE = importlib.resources.read_text(
 DEFAULT_FILL_SLOTS_TEMPLATE = importlib.resources.read_text(
     "rasa.dialogue_understanding.generator.multi_step", "fill_slots_prompt.jinja2"
 ).strip()
+MULTI_STEP_LLM_COMMAND_GENERATOR_CONFIG_FILE = "config.json"
 
 # dictionary of template names and associated file names and default values
 PROMPT_TEMPLATES = {
@@ -145,15 +144,18 @@ class MultiStepLLMCommandGenerator(LLMBasedCommandGenerator):
         **kwargs: Any,
     ) -> "MultiStepLLMCommandGenerator":
         """Loads trained component (see parent class for full docstring)."""
-        prompts = cls._load_prompt_templates(model_storage, resource)
 
-        persisted_config = cls.load_config_from_model_storage(model_storage, resource)
-        train_model_name = (
-            persisted_config.get(TRAINED_MODEL_NAME_CONFIG_KEY, None)
-            if persisted_config
-            else None
+        # Perform health check of the LLM client config
+        llm_config = resolve_model_client_config(config.get(LLM_CONFIG_KEY, {}))
+        cls.perform_llm_health_check(
+            llm_config,
+            DEFAULT_LLM_CONFIG,
+            "multi_step_llm_command_generator.load",
+            MultiStepLLMCommandGenerator.__name__,
         )
 
+        prompts = cls._load_prompt_templates(model_storage, resource)
+
         # init base command generator
         command_generator = cls(config, model_storage, resource, prompts)
         # load flow retrieval if enabled
@@ -162,23 +164,12 @@ class MultiStepLLMCommandGenerator(LLMBasedCommandGenerator):
                 command_generator.config, model_storage, resource
             )
 
-        perform_inference_time_llm_health_check(
-            command_generator.config.get(LLM_CONFIG_KEY),
-            DEFAULT_LLM_CONFIG,
-            train_model_name,
-            "multi_step_llm_command_generator.load",
-            MultiStepLLMCommandGenerator.__name__,
-        )
-
         return command_generator
 
     def persist(self) -> None:
         """Persist this component to disk for future loading."""
-        super().persist()
-
-        # persist prompt template
         self._persist_prompt_templates()
-        # persist flow retrieval
+        self._persist_config()
         if self.flow_retrieval is not None:
             self.flow_retrieval.persist()
 
@@ -411,6 +402,13 @@ class MultiStepLLMCommandGenerator(LLMBasedCommandGenerator):
                 file_path = path / file_name
                 rasa.shared.utils.io.write_text_file(template, file_path)
 
+    def _persist_config(self) -> None:
+        """Persist config as a source of truth for resolved clients."""
+        with self._model_storage.write_to(self._resource) as path:
+            rasa.shared.utils.io.dump_obj_as_json_to_file(
+                path / MULTI_STEP_LLM_COMMAND_GENERATOR_CONFIG_FILE, self.config
+            )
+
     async def _predict_commands_with_multi_step(
         self,
         message: Message,

rasa/dialogue_understanding/generator/nlu_command_adapter.py

@@ -19,6 +19,7 @@ from rasa.engine.storage.storage import ModelStorage
 from rasa.shared.constants import ROUTE_TO_CALM_SLOT
 from rasa.shared.core.domain import Domain
 from rasa.shared.core.flows.flows_list import FlowsList
+from rasa.shared.core.flows.steps import CollectInformationFlowStep
 from rasa.shared.core.slot_mappings import (
     SlotFillingManager,
     extract_slot_value,
@@ -217,7 +218,24 @@ def _issue_set_slot_commands(
     commands: List[Command] = []
     domain = domain if domain else Domain.empty()
     slot_filling_manager = SlotFillingManager(domain, tracker, message)
-    available_slot_names = flows.available_slot_names()
+
+    # only use slots that don't have ask_before_filling set to True
+    available_slot_names = flows.available_slot_names(ask_before_filling=False)
+
+    # check if the current step is a CollectInformationFlowStep
+    # in case it has ask_before_filling set to True, we need to add the
+    # slot to the available_slot_names
+    if tracker.active_flow:
+        flow = flows.flow_by_id(tracker.active_flow)
+        step_id = tracker.current_step_id
+        if flow is not None:
+            current_step = flow.step_by_id(step_id)
+            if (
+                current_step
+                and isinstance(current_step, CollectInformationFlowStep)
+                and current_step.ask_before_filling
+            ):
+                available_slot_names.add(current_step.collect)
 
     for _, slot in tracker.slots.items():
         # if a slot is not collected in available flows,