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

rasa/core/nlg/contextual_response_rephraser.py

@@ -2,7 +2,6 @@ from typing import Any, Dict, Optional, Text
 
 import structlog
 from jinja2 import Template
-
 from rasa import telemetry
 from rasa.core.nlg.response import TemplatedNaturalLanguageGenerator
 from rasa.core.nlg.summarize import summarize_conversation
@@ -14,11 +13,12 @@ from rasa.shared.constants import (
     PROVIDER_CONFIG_KEY,
     OPENAI_PROVIDER,
     TIMEOUT_CONFIG_KEY,
-    MODEL_GROUP_CONFIG_KEY,
+    MODEL_GROUP_ID_CONFIG_KEY,
 )
 from rasa.shared.core.domain import KEY_RESPONSES_TEXT, Domain
 from rasa.shared.core.events import BotUttered, UserUttered
 from rasa.shared.core.trackers import DialogueStateTracker
+from rasa.shared.utils.health_check.llm_health_check_mixin import LLMHealthCheckMixin
 from rasa.shared.utils.llm import (
     DEFAULT_OPENAI_GENERATE_MODEL_NAME,
     DEFAULT_OPENAI_MAX_GENERATED_TOKENS,
@@ -28,7 +28,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
 from rasa.shared.utils.llm import (
     tracker_as_readable_transcript,
 )
@@ -44,6 +43,8 @@ RESPONSE_REPHRASING_TEMPLATE_KEY = "rephrase_prompt"
 RESPONSE_SUMMARISE_CONVERSATION_KEY = "summarize_conversation"
 
 DEFAULT_REPHRASE_ALL = False
+DEFAULT_SUMMARIZE_HISTORY = True
+DEFAULT_MAX_HISTORICAL_TURNS = 5
 
 DEFAULT_LLM_CONFIG = {
     PROVIDER_CONFIG_KEY: OPENAI_PROVIDER,
@@ -68,7 +69,9 @@ Suggested AI Response: {{suggested_response}}
 Rephrased AI Response:"""
 
 
-class ContextualResponseRephraser(TemplatedNaturalLanguageGenerator):
+class ContextualResponseRephraser(
+    LLMHealthCheckMixin, TemplatedNaturalLanguageGenerator
+):
     """Generates responses based on modified templates.
 
     The templates are filled with the entities and slots that are available in the
@@ -102,13 +105,19 @@ class ContextualResponseRephraser(TemplatedNaturalLanguageGenerator):
         self.trace_prompt_tokens = self.nlg_endpoint.kwargs.get(
             "trace_prompt_tokens", False
         )
+        self.summarize_history = self.nlg_endpoint.kwargs.get(
+            "summarize_history", DEFAULT_SUMMARIZE_HISTORY
+        )
+        self.max_historical_turns = self.nlg_endpoint.kwargs.get(
+            "max_historical_turns", DEFAULT_MAX_HISTORICAL_TURNS
+        )
 
         self.llm_config = resolve_model_client_config(
             self.nlg_endpoint.kwargs.get(LLM_CONFIG_KEY),
             ContextualResponseRephraser.__name__,
         )
 
-        perform_training_time_llm_health_check(
+        self.perform_llm_health_check(
             self.llm_config,
             DEFAULT_LLM_CONFIG,
             "contextual_response_rephraser.init",
@@ -213,18 +222,17 @@ class ContextualResponseRephraser(TemplatedNaturalLanguageGenerator):
         prompt_template_text = self._template_for_response_rephrasing(response)
 
         # Retrieve inputs for the dynamic prompt
-        transcript = tracker_as_readable_transcript(tracker, max_turns=5)
         latest_message = self._last_message_if_human(tracker)
         current_input = f"{USER}: {latest_message}" if latest_message else ""
 
         # Only summarise conversation history if flagged
-        summarize_conversation_flag = response.get("metadata", {}).get(
-            RESPONSE_SUMMARISE_CONVERSATION_KEY, False
-        )
-        if summarize_conversation_flag:
+        if self.summarize_history:
             history = await self._create_history(tracker)
         else:
-            history = transcript
+            # make sure the transcript/history contains the last user utterance
+            max_turns = max(self.max_historical_turns, 1)
+            history = tracker_as_readable_transcript(tracker, max_turns=max_turns)
+            # the history already contains the current input
             current_input = ""
 
         prompt = Template(prompt_template_text).render(
@@ -245,7 +253,7 @@ class ContextualResponseRephraser(TemplatedNaturalLanguageGenerator):
             llm_type=self.llm_property(PROVIDER_CONFIG_KEY),
             llm_model=self.llm_property(MODEL_CONFIG_KEY)
             or self.llm_property(MODEL_NAME_CONFIG_KEY),
-            llm_model_group_id=self.llm_property(MODEL_GROUP_CONFIG_KEY),
+            llm_model_group_id=self.llm_property(MODEL_GROUP_ID_CONFIG_KEY),
         )
         if not (updated_text := await self._generate_llm_response(prompt)):
             # If the LLM fails to generate a response, we

rasa/core/policies/enterprise_search_policy.py

@@ -1,7 +1,7 @@
 import importlib.resources
 import json
 import re
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Text, Tuple
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Text
 import dotenv
 import structlog
 from jinja2 import Template
@@ -25,8 +25,6 @@ from rasa.core.policies.policy import Policy, PolicyPrediction
 from rasa.core.utils import AvailableEndpoints
 from rasa.dialogue_understanding.generator.constants import (
     LLM_CONFIG_KEY,
-    TRAINED_MODEL_NAME_CONFIG_KEY,
-    TRAINED_EMBEDDINGS_CONFIG_KEY,
 )
 from rasa.dialogue_understanding.patterns.cannot_handle import (
     CannotHandlePatternFlowStackFrame,
@@ -53,7 +51,7 @@ from rasa.shared.constants import (
     OPENAI_PROVIDER,
     TIMEOUT_CONFIG_KEY,
     MODEL_NAME_CONFIG_KEY,
-    MODEL_GROUP_CONFIG_KEY,
+    MODEL_GROUP_ID_CONFIG_KEY,
 )
 from rasa.shared.core.constants import (
     ACTION_CANCEL_FLOW,
@@ -71,6 +69,10 @@ from rasa.shared.providers.embedding._langchain_embedding_client_adapter import
 )
 from rasa.shared.providers.llm.llm_client import LLMClient
 from rasa.shared.utils.cli import print_error_and_exit
+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 (
     DEFAULT_OPENAI_CHAT_MODEL_NAME,
@@ -82,12 +84,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.telemetry import (
     track_enterprise_search_policy_predict,
     track_enterprise_search_policy_train_completed,
@@ -161,7 +157,7 @@ class VectorStoreConfigurationError(RasaException):
 @DefaultV1Recipe.register(
     DefaultV1Recipe.ComponentType.POLICY_WITH_END_TO_END_SUPPORT, is_trainable=True
 )
-class EnterpriseSearchPolicy(Policy):
+class EnterpriseSearchPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Policy):
     """Policy which uses a vector store and LLMs to respond to user messages.
 
     The policy uses a vector store and LLMs to respond to user messages. The
@@ -300,6 +296,9 @@ class EnterpriseSearchPolicy(Policy):
             A policy must return its resource locator so that potential children nodes
             can load the policy from the resource.
         """
+        # Perform health checks for both LLM and embeddings client configs
+        self._perform_health_checks(self.config, "enterprise_search_policy.train")
+
         store_type = self.vector_store_config.get(VECTOR_STORE_TYPE_PROPERTY)
 
         # telemetry call to track training start
@@ -319,11 +318,6 @@ class EnterpriseSearchPolicy(Policy):
                 f"required environment variables. Error: {e}"
             )
 
-        (
-            self.config[TRAINED_MODEL_NAME_CONFIG_KEY],
-            self.config[TRAINED_EMBEDDINGS_CONFIG_KEY],
-        ) = self._perform_training_time_health_checks()
-
         if store_type == DEFAULT_VECTOR_STORE_TYPE:
             logger.info("enterprise_search_policy.train.faiss")
             with self._model_storage.write_to(self._resource) as path:
@@ -343,12 +337,12 @@ class EnterpriseSearchPolicy(Policy):
             embeddings_model=self.embeddings_config.get(MODEL_CONFIG_KEY)
             or self.embeddings_config.get(MODEL_NAME_CONFIG_KEY),
             embeddings_model_group_id=self.embeddings_config.get(
-                MODEL_GROUP_CONFIG_KEY
+                MODEL_GROUP_ID_CONFIG_KEY
             ),
             llm_type=self.llm_config.get(PROVIDER_CONFIG_KEY),
             llm_model=self.llm_config.get(MODEL_CONFIG_KEY)
             or self.llm_config.get(MODEL_NAME_CONFIG_KEY),
-            llm_model_group_id=self.llm_config.get(MODEL_GROUP_CONFIG_KEY),
+            llm_model_group_id=self.llm_config.get(MODEL_GROUP_ID_CONFIG_KEY),
             citation_enabled=self.citation_enabled,
         )
         self.persist()
@@ -544,12 +538,12 @@ class EnterpriseSearchPolicy(Policy):
             embeddings_model=self.embeddings_config.get(MODEL_CONFIG_KEY)
             or self.embeddings_config.get(MODEL_NAME_CONFIG_KEY),
             embeddings_model_group_id=self.embeddings_config.get(
-                MODEL_GROUP_CONFIG_KEY
+                MODEL_GROUP_ID_CONFIG_KEY
             ),
             llm_type=self.llm_config.get(PROVIDER_CONFIG_KEY),
             llm_model=self.llm_config.get(MODEL_CONFIG_KEY)
             or self.llm_config.get(MODEL_NAME_CONFIG_KEY),
-            llm_model_group_id=self.llm_config.get(MODEL_GROUP_CONFIG_KEY),
+            llm_model_group_id=self.llm_config.get(MODEL_GROUP_ID_CONFIG_KEY),
             citation_enabled=self.citation_enabled,
         )
         return self._create_prediction(
@@ -698,16 +692,16 @@ class EnterpriseSearchPolicy(Policy):
         **kwargs: Any,
     ) -> "EnterpriseSearchPolicy":
         """Loads a trained policy (see parent class for full docstring)."""
+
+        # Perform health checks for both LLM and embeddings client configs
+        cls._perform_health_checks(config, "enterprise_search_policy.load")
+
         prompt_template = None
-        persisted_config = None
         try:
             with model_storage.read_from(resource) as path:
                 prompt_template = rasa.shared.utils.io.read_file(
                     path / ENTERPRISE_SEARCH_PROMPT_FILE_NAME
                 )
-                persisted_config = rasa.shared.utils.io.read_json_file(
-                    path / ENTERPRISE_SEARCH_CONFIG_FILE_NAME
-                )
         except (FileNotFoundError, FileIOException) as e:
             logger.warning(
                 "enterprise_search_policy.load.failed", error=e, resource=resource.name
@@ -737,7 +731,7 @@ class EnterpriseSearchPolicy(Policy):
                 embeddings=embeddings,
             )  # type: ignore
 
-        policy = cls(
+        return cls(
             config,
             model_storage,
             resource,
@@ -746,14 +740,6 @@ class EnterpriseSearchPolicy(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 _get_local_knowledge_data(cls, config: Dict[str, Any]) -> Optional[List[str]]:
         """This is required only for local knowledge base types.
@@ -894,52 +880,26 @@ class EnterpriseSearchPolicy(Policy):
 
         return joined_answer + joined_sources
 
-    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,
-            "enterprise_search_policy.train",
-            EnterpriseSearchPolicy.__name__,
-        )
-        train_embedding_name = perform_training_time_embeddings_health_check(
-            self.config.get(EMBEDDINGS_CONFIG_KEY),
-            DEFAULT_EMBEDDINGS_CONFIG,
-            "enterprise_search_policy.train",
-            EnterpriseSearchPolicy.__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,
-            "enterprise_search_policy.load",
+            log_source_method,
             EnterpriseSearchPolicy.__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,
-            "enterprise_search_policy.load",
+            log_source_method,
             EnterpriseSearchPolicy.__name__,
         )

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,
@@ -43,7 +39,7 @@ from rasa.shared.constants import (
     PROVIDER_CONFIG_KEY,
     OPENAI_PROVIDER,
     TIMEOUT_CONFIG_KEY,
-    MODEL_GROUP_CONFIG_KEY,
+    MODEL_GROUP_ID_CONFIG_KEY,
 )
 from rasa.shared.core.constants import ACTION_LISTEN_NAME
 from rasa.shared.core.constants import ACTION_TRIGGER_CHITCHAT
@@ -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()
@@ -566,11 +558,13 @@ class IntentlessPolicy(Policy):
             embeddings_type=self.embeddings_property(PROVIDER_CONFIG_KEY),
             embeddings_model=self.embeddings_property(MODEL_CONFIG_KEY)
             or self.embeddings_property(MODEL_NAME_CONFIG_KEY),
-            embeddings_model_group_id=self.embeddings_property(MODEL_GROUP_CONFIG_KEY),
+            embeddings_model_group_id=self.embeddings_property(
+                MODEL_GROUP_ID_CONFIG_KEY
+            ),
             llm_type=self.llm_property(PROVIDER_CONFIG_KEY),
             llm_model=self.llm_property(MODEL_CONFIG_KEY)
             or self.llm_property(MODEL_NAME_CONFIG_KEY),
-            llm_model_group_id=self.llm_property(MODEL_GROUP_CONFIG_KEY),
+            llm_model_group_id=self.llm_property(MODEL_GROUP_ID_CONFIG_KEY),
         )
 
         self.persist()
@@ -650,11 +644,13 @@ class IntentlessPolicy(Policy):
             embeddings_type=self.embeddings_property(PROVIDER_CONFIG_KEY),
             embeddings_model=self.embeddings_property(MODEL_CONFIG_KEY)
             or self.embeddings_property(MODEL_NAME_CONFIG_KEY),
-            embeddings_model_group_id=self.embeddings_property(MODEL_GROUP_CONFIG_KEY),
+            embeddings_model_group_id=self.embeddings_property(
+                MODEL_GROUP_ID_CONFIG_KEY
+            ),
             llm_type=self.llm_property(PROVIDER_CONFIG_KEY),
             llm_model=self.llm_property(MODEL_CONFIG_KEY)
             or self.llm_property(MODEL_NAME_CONFIG_KEY),
-            llm_model_group_id=self.llm_property(MODEL_GROUP_CONFIG_KEY),
+            llm_model_group_id=self.llm_property(MODEL_GROUP_ID_CONFIG_KEY),
             score=score,
         )
 
@@ -952,10 +948,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 +972,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 +987,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 +1006,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"