rasa-pro 3.11.0__py3-none-any.whl → 3.11.0a1__py3-none-any.whl

rasa/core/utils.py

@@ -1,9 +1,8 @@
-import structlog
 import logging
 import os
 from pathlib import Path
 from socket import SOCK_DGRAM, SOCK_STREAM
-from typing import Any, Dict, Optional, List, Set, TYPE_CHECKING, Text, Tuple, Union
+from typing import Any, Dict, Optional, Set, TYPE_CHECKING, Text, Tuple, Union
 
 import numpy as np
 from sanic import Sanic
@@ -20,18 +19,14 @@ from rasa.core.constants import (
 from rasa.core.lock_store import LockStore, RedisLockStore, InMemoryLockStore
 from rasa.shared.constants import DEFAULT_ENDPOINTS_PATH, TCP_PROTOCOL
 from rasa.shared.core.trackers import DialogueStateTracker
-from rasa.utils.endpoints import (
-    EndpointConfig,
-    read_endpoint_config,
-    read_property_config_from_endpoints_file,
-)
+from rasa.utils.endpoints import EndpointConfig, read_endpoint_config
 from rasa.utils.io import write_yaml
 
 if TYPE_CHECKING:
     from rasa.core.nlg import NaturalLanguageGenerator
     from rasa.shared.core.domain import Domain
 
-structlogger = structlog.get_logger()
+logger = logging.getLogger(__name__)
 
 
 def configure_file_logging(
@@ -129,17 +124,15 @@ def list_routes(app: Sanic) -> Dict[Text, Text]:
         for arg in route._params:
             options[arg] = f"[{arg}]"
 
-        name = route.name.replace("rasa_server.", "")
-        methods = ",".join(route.methods)
+        handlers = [(next(iter(route.methods)), route.name.replace("rasa_server.", ""))]
 
-        full_endpoint = "/" + "/".join(endpoint)
-        line = unquote(f"{full_endpoint:50s} {methods:30s} {name}")
-        output[name] = line
+        for method, name in handlers:
+            full_endpoint = "/" + "/".join(endpoint)
+            line = unquote(f"{full_endpoint:50s} {method:30s} {name}")
+            output[name] = line
 
     url_table = "\n".join(output[url] for url in sorted(output))
-    structlogger.debug(
-        "server.routes", event_info=f"Available web server routes: \n{url_table}"
-    )
+    logger.debug(f"Available web server routes: \n{url_table}")
 
     return output
 
@@ -178,8 +171,6 @@ def is_limit_reached(num_messages: int, limit: Optional[int]) -> bool:
 class AvailableEndpoints:
     """Collection of configured endpoints."""
 
-    _instance = None
-
     @classmethod
     def read_endpoints(cls, endpoint_file: Text) -> "AvailableEndpoints":
         """Read the different endpoints from a yaml file."""
@@ -193,9 +184,6 @@ class AvailableEndpoints:
         lock_store = read_endpoint_config(endpoint_file, endpoint_type="lock_store")
         event_broker = read_endpoint_config(endpoint_file, endpoint_type="event_broker")
         vector_store = read_endpoint_config(endpoint_file, endpoint_type="vector_store")
-        model_groups = read_property_config_from_endpoints_file(
-            endpoint_file, property_name="model_groups"
-        )
 
         return cls(
             nlg,
@@ -206,7 +194,6 @@ class AvailableEndpoints:
             lock_store,
             event_broker,
             vector_store,
-            model_groups,
         )
 
     def __init__(
@@ -219,7 +206,6 @@ class AvailableEndpoints:
         lock_store: Optional[EndpointConfig] = None,
         event_broker: Optional[EndpointConfig] = None,
         vector_store: Optional[EndpointConfig] = None,
-        model_groups: Optional[List[Dict[str, Any]]] = None,
     ) -> None:
         """Create an `AvailableEndpoints` object."""
         self.model = model
@@ -230,15 +216,6 @@ class AvailableEndpoints:
         self.lock_store = lock_store
         self.event_broker = event_broker
         self.vector_store = vector_store
-        self.model_groups = model_groups
-
-    @classmethod
-    def get_instance(cls, endpoint_file: Optional[Text] = None) -> "AvailableEndpoints":
-        """Get the singleton instance of AvailableEndpoints."""
-        # Ensure that the instance is initialized only once.
-        if cls._instance is None:
-            cls._instance = cls.read_endpoints(endpoint_file)
-        return cls._instance
 
 
 def read_endpoints_from_path(
@@ -257,7 +234,7 @@ def read_endpoints_from_path(
     endpoints_config_path = cli_utils.get_validated_path(
         endpoints_path, "endpoints", DEFAULT_ENDPOINTS_PATH, True
     )
-    return AvailableEndpoints.get_instance(endpoints_config_path)
+    return AvailableEndpoints.read_endpoints(endpoints_config_path)
 
 
 def _lock_store_is_multi_worker_compatible(
@@ -286,22 +263,17 @@ def number_of_sanic_workers(lock_store: Union[EndpointConfig, LockStore, None])
     """
 
     def _log_and_get_default_number_of_workers() -> int:
-        structlogger.debug(
-            "server.worker.set_count",
-            number_of_workers=DEFAULT_SANIC_WORKERS,
-            event_info=f"Using the default number of Sanic workers "
-            f"({DEFAULT_SANIC_WORKERS}).",
+        logger.debug(
+            f"Using the default number of Sanic workers ({DEFAULT_SANIC_WORKERS})."
         )
         return DEFAULT_SANIC_WORKERS
 
     try:
         env_value = int(os.environ.get(ENV_SANIC_WORKERS, DEFAULT_SANIC_WORKERS))
     except ValueError:
-        structlogger.error(
-            "server.worker.set_count.error",
-            number_of_workers=os.environ[ENV_SANIC_WORKERS],
-            event_info=f"Cannot convert environment variable `{ENV_SANIC_WORKERS}` "
-            f"to int ('{os.environ[ENV_SANIC_WORKERS]}').",
+        logger.error(
+            f"Cannot convert environment variable `{ENV_SANIC_WORKERS}` "
+            f"to int ('{os.environ[ENV_SANIC_WORKERS]}')."
         )
         return _log_and_get_default_number_of_workers()
 
@@ -309,28 +281,20 @@ def number_of_sanic_workers(lock_store: Union[EndpointConfig, LockStore, None])
         return _log_and_get_default_number_of_workers()
 
     if env_value < 1:
-        structlogger.warning(
-            "server.worker.set_count.error_less_than_one",
-            number_of_workers=env_value,
-            event_info=f"Cannot set number of Sanic workers to the desired value "
-            f"({env_value}). The number of workers must be at least 1.",
+        logger.debug(
+            f"Cannot set number of Sanic workers to the desired value "
+            f"({env_value}). The number of workers must be at least 1."
         )
         return _log_and_get_default_number_of_workers()
 
     if _lock_store_is_multi_worker_compatible(lock_store):
-        structlogger.debug(
-            "server.worker.set_count.success",
-            event_info=f"Using {env_value} Sanic workers.",
-            num_workers=env_value,
-        )
+        logger.debug(f"Using {env_value} Sanic workers.")
         return env_value
 
-    structlogger.warning(
-        "server.worker.set_count.error_no_lock_store",
-        event_info=f"Unable to assign desired number of Sanic workers ({env_value}) as "
+    logger.debug(
+        f"Unable to assign desired number of Sanic workers ({env_value}) as "
         f"no `RedisLockStore` or custom `LockStore` endpoint "
-        f"configuration has been found.",
-        num_workers=env_value,
+        f"configuration has been found."
     )
     return _log_and_get_default_number_of_workers()
 

rasa/dialogue_understanding/coexistence/llm_based_router.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import importlib
+import os
 from typing import Any, Dict, List, Optional
 
 import structlog
@@ -15,14 +16,13 @@ from rasa.dialogue_understanding.coexistence.constants import (
 )
 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,
-)
+from rasa.dialogue_understanding.generator.constants import LLM_CONFIG_KEY
 from rasa.engine.graph import ExecutionContext, GraphComponent
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.storage.resource import Resource
 from rasa.engine.storage.storage import ModelStorage
 from rasa.shared.constants import (
+    LLM_API_HEALTH_CHECK_ENV_VAR,
     ROUTE_TO_CALM_SLOT,
     PROMPT_CONFIG_KEY,
     PROVIDER_CONFIG_KEY,
@@ -35,13 +35,12 @@ 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,
     get_prompt_template,
+    llm_api_health_check,
     llm_factory,
-    resolve_model_client_config,
+    try_instantiate_llm_client,
 )
 from rasa.utils.log_utils import log_llm
 
@@ -49,7 +48,6 @@ LLM_BASED_ROUTER_PROMPT_FILE_NAME = "llm_based_router_prompt.jinja2"
 DEFAULT_COMMAND_PROMPT_TEMPLATE = importlib.resources.read_text(
     "rasa.dialogue_understanding.coexistence", "router_template.jinja2"
 )
-LLM_BASED_ROUTER_CONFIG_FILE_NAME = "config.json"
 
 # Token ids for gpt 3.5 and gpt 4 corresponding to space + capitalized Letter
 A_TO_C_TOKEN_IDS_CHATGPT = [
@@ -76,7 +74,7 @@ structlogger = structlog.get_logger()
     ],
     is_trainable=True,
 )
-class LLMBasedRouter(LLMHealthCheckMixin, GraphComponent):
+class LLMBasedRouter(GraphComponent):
     @staticmethod
     def get_default_config() -> Dict[str, Any]:
         """The component's default config (see parent class for full docstring)."""
@@ -98,9 +96,6 @@ class LLMBasedRouter(LLMHealthCheckMixin, GraphComponent):
         prompt_template: Optional[str] = None,
     ) -> None:
         self.config = {**self.get_default_config(), **config}
-        self.config[LLM_CONFIG_KEY] = resolve_model_client_config(
-            self.config.get(LLM_CONFIG_KEY), LLMBasedRouter.__name__
-        )
 
         self.prompt_template = (
             prompt_template
@@ -134,18 +129,20 @@ class LLMBasedRouter(LLMHealthCheckMixin, GraphComponent):
             rasa.shared.utils.io.write_text_file(
                 self.prompt_template, path / LLM_BASED_ROUTER_PROMPT_FILE_NAME
             )
-            rasa.shared.utils.io.dump_obj_as_json_to_file(
-                path / LLM_BASED_ROUTER_CONFIG_FILE_NAME, self.config
-            )
 
     def train(self, training_data: TrainingData) -> Resource:
         """Train the intent classifier on a data set."""
-        self.perform_llm_health_check(
+        # Validate llm configuration
+        llm_client = try_instantiate_llm_client(
             self.config.get(LLM_CONFIG_KEY),
             DEFAULT_LLM_CONFIG,
             "llm_based_router.train",
             LLMBasedRouter.__name__,
         )
+        if os.getenv(LLM_API_HEALTH_CHECK_ENV_VAR, "true").lower() == "true":
+            llm_api_health_check(
+                llm_client, "llm_based_router.train", LLMBasedRouter.__name__
+            )
 
         self.persist()
         return self._resource
@@ -160,16 +157,6 @@ class LLMBasedRouter(LLMHealthCheckMixin, 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
         try:
             with model_storage.read_from(resource) as path:
@@ -311,17 +298,3 @@ class LLMBasedRouter(LLMHealthCheckMixin, GraphComponent):
             # we have to catch all exceptions here
             structlogger.error("llm_based_router.llm.error", error=e)
             return None
-
-    @classmethod
-    def fingerprint_addon(cls, config: Dict[str, Any]) -> Optional[str]:
-        """Add a fingerprint of llm based router for the graph."""
-        prompt_template = get_prompt_template(
-            config.get(PROMPT_CONFIG_KEY),
-            DEFAULT_COMMAND_PROMPT_TEMPLATE,
-        )
-
-        llm_config = resolve_model_client_config(
-            config.get(LLM_CONFIG_KEY), LLMBasedRouter.__name__
-        )
-
-        return deep_container_fingerprint([prompt_template, llm_config])

rasa/dialogue_understanding/commands/__init__.py

@@ -33,9 +33,6 @@ from rasa.dialogue_understanding.commands.session_start_command import (
     SessionStartCommand,
 )
 from rasa.dialogue_understanding.commands.session_end_command import SessionEndCommand
-from rasa.dialogue_understanding.commands.repeat_bot_messages_command import (
-    RepeatBotMessagesCommand,
-)
 
 __all__ = [
     "Command",
@@ -56,6 +53,5 @@ __all__ = [
     "ChangeFlowCommand",
     "SessionStartCommand",
     "SessionEndCommand",
-    "RepeatBotMessagesCommand",
     "RestartCommand",
 ]

rasa/dialogue_understanding/commands/change_flow_command.py

@@ -36,9 +36,3 @@ class ChangeFlowCommand(Command):
         # the change flow command is not actually pushing anything to the tracker,
         # but it is predicted by the MultiStepLLMCommandGenerator and used internally
         return []
-
-    def __eq__(self, other: Any) -> bool:
-        return isinstance(other, ChangeFlowCommand)
-
-    def __hash__(self) -> int:
-        return hash(self.command())

rasa/dialogue_understanding/commands/utils.py

@@ -11,7 +11,6 @@ from rasa.dialogue_understanding.commands import (
     SkipQuestionCommand,
     RestartCommand,
 )
-from rasa.dialogue_understanding.commands.user_silence_command import UserSilenceCommand
 from rasa.dialogue_understanding.patterns.cancel import CancelPatternFlowStackFrame
 from rasa.dialogue_understanding.patterns.cannot_handle import (
     CannotHandlePatternFlowStackFrame,
@@ -28,13 +27,9 @@ from rasa.dialogue_understanding.patterns.session_start import (
 from rasa.dialogue_understanding.patterns.skip_question import (
     SkipQuestionPatternFlowStackFrame,
 )
-from rasa.dialogue_understanding.patterns.user_silence import (
-    UserSilencePatternFlowStackFrame,
-)
 
 triggerable_pattern_to_command_class: Dict[str, Type[Command]] = {
     SessionStartPatternFlowStackFrame.flow_id: SessionStartCommand,
-    UserSilencePatternFlowStackFrame.flow_id: UserSilenceCommand,
     CancelPatternFlowStackFrame.flow_id: CancelFlowCommand,
     ChitchatPatternFlowStackFrame.flow_id: ChitChatAnswerCommand,
     HumanHandoffPatternFlowStackFrame.flow_id: HumanHandoffCommand,

rasa/dialogue_understanding/generator/constants.py

@@ -23,5 +23,3 @@ USER_INPUT_CONFIG_KEY = "user_input"
 FLOW_RETRIEVAL_KEY = "flow_retrieval"
 FLOW_RETRIEVAL_ACTIVE_KEY = "active"
 FLOW_RETRIEVAL_EMBEDDINGS_CONFIG_KEY = "embeddings"
-
-FLOW_RETRIEVAL_FLOW_THRESHOLD = 20

rasa/dialogue_understanding/generator/flow_retrieval.py

@@ -1,4 +1,5 @@
-"""The module is primarily centered around the `FlowRetrieval` class which handles the
+"""
+The module is primarily centered around the `FlowRetrieval` class which handles the
 initialization, configuration validation, vector store management, and flow retrieval
 logic. It integrates components for managing embeddings, vector stores, and
 flow-specific templates, facilitating semantic search functionalities.
@@ -26,10 +27,8 @@ from langchain.docstore.document import Document
 from langchain.schema.embeddings import Embeddings
 from langchain_community.vectorstores.faiss import FAISS
 from langchain_community.vectorstores.utils import DistanceStrategy
-
 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,
@@ -38,15 +37,12 @@ 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
 from rasa.shared.nlu.constants import TEXT, FLOWS_FROM_SEMANTIC_SEARCH
 from rasa.shared.nlu.training_data.message import Message
+from rasa.shared.exceptions import ProviderClientAPIException
 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,
@@ -54,15 +50,12 @@ from rasa.shared.utils.llm import (
     USER,
     get_prompt_template,
     allowed_values_for_slot,
-    resolve_model_client_config,
 )
 
 DEFAULT_FLOW_DOCUMENT_TEMPLATE = importlib.resources.read_text(
     "rasa.dialogue_understanding.generator", "flow_document_template.jinja2"
 )
 
-FLOW_RETRIEVAL_CONFIG_FILE_NAME = "flow_retrieval_config.json"
-
 DEFAULT_EMBEDDINGS_CONFIG = {
     PROVIDER_CONFIG_KEY: OPENAI_PROVIDER,
     "model": DEFAULT_OPENAI_EMBEDDING_MODEL_NAME,
@@ -80,7 +73,7 @@ DEFAULT_SHOULD_EMBED_SLOTS = True
 structlogger = structlog.get_logger()
 
 
-class FlowRetrieval(EmbeddingsHealthCheckMixin):
+class FlowRetrieval:
     @classmethod
     def get_default_config(cls) -> Dict[str, Any]:
         """The default config for the flow retrieval."""
@@ -99,9 +92,6 @@ class FlowRetrieval(EmbeddingsHealthCheckMixin):
     ):
         config = {**self.get_default_config(), **config}
         self.config = self.validate_config(config)
-        self.config[EMBEDDINGS_CONFIG_KEY] = resolve_model_client_config(
-            self.config.get(EMBEDDINGS_CONFIG_KEY), FlowRetrieval.__name__
-        )
         self.vector_store: Optional[FAISS] = None
         self.flow_document_template = get_prompt_template(
             None, DEFAULT_FLOW_DOCUMENT_TEMPLATE
@@ -150,18 +140,6 @@ class FlowRetrieval(EmbeddingsHealthCheckMixin):
         **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
@@ -169,7 +147,6 @@ class FlowRetrieval(EmbeddingsHealthCheckMixin):
             flow_retrieval.config, model_storage, resource
         )
         flow_retrieval.vector_store = vector_store
-
         return flow_retrieval
 
     @classmethod
@@ -201,21 +178,13 @@ class FlowRetrieval(EmbeddingsHealthCheckMixin):
         Returns:
             The embedder.
         """
-        # Copy the config so original config is not modified
-        config = config.copy()
-        # Resolve config and instantiate the embedding client
-        config[EMBEDDINGS_CONFIG_KEY] = resolve_model_client_config(
-            config.get(EMBEDDINGS_CONFIG_KEY), FlowRetrieval.__name__
-        )
         client = embedder_factory(
             config.get(EMBEDDINGS_CONFIG_KEY), DEFAULT_EMBEDDINGS_CONFIG
         )
-        # Wrap the embedding client in the adapter
         return _LangchainEmbeddingClientAdapter(client)
 
     def persist(self) -> None:
         self._persist_vector_store()
-        self._persist_config()
 
     def _persist_vector_store(self) -> None:
         """Persists the FAISS vector store."""
@@ -228,12 +197,6 @@ class FlowRetrieval(EmbeddingsHealthCheckMixin):
                 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
@@ -243,14 +206,6 @@ class FlowRetrieval(EmbeddingsHealthCheckMixin):
             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

@@ -2,6 +2,7 @@ from abc import ABC, abstractmethod
 from functools import lru_cache
 from typing import Dict, Any, List, Optional, Tuple, Union, Text
 
+import os
 import structlog
 from jinja2 import Template
 
@@ -16,13 +17,13 @@ from rasa.dialogue_understanding.generator.constants import (
     LLM_CONFIG_KEY,
     FLOW_RETRIEVAL_KEY,
     FLOW_RETRIEVAL_ACTIVE_KEY,
-    FLOW_RETRIEVAL_FLOW_THRESHOLD,
 )
 from rasa.dialogue_understanding.generator.flow_retrieval import FlowRetrieval
 from rasa.engine.graph import GraphComponent, ExecutionContext
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.storage.resource import Resource
 from rasa.engine.storage.storage import ModelStorage
+from rasa.shared.constants import LLM_API_HEALTH_CHECK_ENV_VAR
 from rasa.shared.core.domain import Domain
 from rasa.shared.core.flows import FlowStep, Flow, FlowsList
 from rasa.shared.core.flows.steps.collect import CollectInformationFlowStep
@@ -32,11 +33,11 @@ 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_api_health_check,
     llm_factory,
-    resolve_model_client_config,
+    try_instantiate_llm_client,
 )
 from rasa.utils.log_utils import log_llm
 
@@ -49,9 +50,7 @@ structlogger = structlog.get_logger()
     ],
     is_trainable=True,
 )
-class LLMBasedCommandGenerator(
-    LLMHealthCheckMixin, GraphComponent, CommandGenerator, ABC
-):
+class LLMBasedCommandGenerator(GraphComponent, CommandGenerator, ABC):
     """An abstract class defining interface and common functionality
     of an LLM-based command generators.
     """
@@ -65,9 +64,6 @@ class LLMBasedCommandGenerator(
     ) -> None:
         super().__init__(config)
         self.config = {**self.get_default_config(), **config}
-        self.config[LLM_CONFIG_KEY] = resolve_model_client_config(
-            self.config.get(LLM_CONFIG_KEY), LLMBasedCommandGenerator.__name__
-        )
         self._model_storage = model_storage
         self._resource = resource
         self.flow_retrieval: Optional[FlowRetrieval]
@@ -77,9 +73,17 @@ class LLMBasedCommandGenerator(
                 self.config[FLOW_RETRIEVAL_KEY], model_storage, resource
             )
             structlogger.info("llm_based_command_generator.flow_retrieval.enabled")
-            self.config[FLOW_RETRIEVAL_KEY] = self.flow_retrieval.config
         else:
             self.flow_retrieval = None
+            structlogger.warn(
+                "llm_based_command_generator.flow_retrieval.disabled",
+                event_info=(
+                    "Disabling flow retrieval can cause issues when there are a "
+                    "large number of flows to be included in the prompt. For more"
+                    "information see:\n"
+                    "https://rasa.com/docs/rasa-pro/concepts/dialogue-understanding#how-the-llmcommandgenerator-works"
+                ),
+            )
 
     ### Abstract methods
     @staticmethod
@@ -167,32 +171,18 @@ class LLMBasedCommandGenerator(
         """Train the llm based command generator. Stores all flows into a vector
         store.
         """
-        self.perform_llm_health_check(
+        # Validate llm configuration
+        llm_client = try_instantiate_llm_client(
             self.config.get(LLM_CONFIG_KEY),
             DEFAULT_LLM_CONFIG,
             "llm_based_command_generator.train",
             LLMBasedCommandGenerator.__name__,
         )
-
-        if (
-            self.flow_retrieval is None
-            and len(flows.user_flows) > FLOW_RETRIEVAL_FLOW_THRESHOLD
-        ):
-            structlogger.warn(
-                "llm_based_command_generator.flow_retrieval.disabled",
-                event_info=(
-                    f"You have {len(flows.user_flows)} user flows but flow "
-                    f"retrieval is disabled. "
-                    f"It is recommended to enable flow retrieval if the "
-                    f"total number of user flows exceed "
-                    f"{FLOW_RETRIEVAL_FLOW_THRESHOLD}. "
-                    f"Keeping it disabled can result in deterioration of "
-                    f"command generator's functional "
-                    f"performance and higher costs because of increased "
-                    f"number of tokens in the prompt. For more"
-                    "information see:\n"
-                    "https://rasa.com/docs/rasa-pro/concepts/dialogue-understanding#how-the-llmcommandgenerator-works"
-                ),
+        if os.getenv(LLM_API_HEALTH_CHECK_ENV_VAR, "true").lower() == "true":
+            llm_api_health_check(
+                llm_client,
+                "llm_based_command_generator.train",
+                LLMBasedCommandGenerator.__name__,
             )
 
         # flow retrieval is populated with only user-defined flows
@@ -202,11 +192,10 @@ class LLMBasedCommandGenerator(
         except Exception as e:
             structlogger.error(
                 "llm_based_command_generator.train.failed",
-                event_info="Flow retrieval store is inaccessible.",
+                event_info=("Flow retrieval store isinaccessible."),
                 error=e,
             )
             raise
-
         self.persist()
         return self._resource
 
@@ -244,10 +233,7 @@ class LLMBasedCommandGenerator(
 
     @classmethod
     def load_flow_retrival(
-        cls,
-        config: Dict[str, Any],
-        model_storage: ModelStorage,
-        resource: Resource,
+        cls, config: Dict[Text, Any], model_storage: ModelStorage, resource: Resource
     ) -> Optional[FlowRetrieval]:
         """Load the FlowRetrieval component if it is enabled in the configuration."""
         enable_flow_retrieval = config.get(FLOW_RETRIEVAL_KEY, {}).get(