rasa-pro 3.11.0a4.dev3__py3-none-any.whl → 3.11.0rc1__py3-none-any.whl

rasa/core/featurizers/tracker_featurizers.py

@@ -1,11 +1,9 @@
 from __future__ import annotations
-from pathlib import Path
-from collections import defaultdict
-from abc import abstractmethod
-import jsonpickle
-import logging
 
-from tqdm import tqdm
+import logging
+from abc import abstractmethod
+from collections import defaultdict
+from pathlib import Path
 from typing import (
     Tuple,
     List,
@@ -18,25 +16,30 @@ from typing import (
     Set,
     DefaultDict,
     cast,
+    Type,
+    Callable,
+    ClassVar,
 )
+
 import numpy as np
+from tqdm import tqdm
 
-from rasa.core.featurizers.single_state_featurizer import SingleStateFeaturizer
-from rasa.core.featurizers.precomputation import MessageContainerForCoreFeaturization
-from rasa.core.exceptions import InvalidTrackerFeaturizerUsageError
 import rasa.shared.core.trackers
 import rasa.shared.utils.io
-from rasa.shared.nlu.constants import TEXT, INTENT, ENTITIES, ACTION_NAME
-from rasa.shared.nlu.training_data.features import Features
-from rasa.shared.core.trackers import DialogueStateTracker
-from rasa.shared.core.domain import State, Domain
-from rasa.shared.core.events import Event, ActionExecuted, UserUttered
+from rasa.core.exceptions import InvalidTrackerFeaturizerUsageError
+from rasa.core.featurizers.precomputation import MessageContainerForCoreFeaturization
+from rasa.core.featurizers.single_state_featurizer import SingleStateFeaturizer
 from rasa.shared.core.constants import (
     USER,
     ACTION_UNLIKELY_INTENT_NAME,
     PREVIOUS_ACTION,
 )
+from rasa.shared.core.domain import State, Domain
+from rasa.shared.core.events import Event, ActionExecuted, UserUttered
+from rasa.shared.core.trackers import DialogueStateTracker
 from rasa.shared.exceptions import RasaException
+from rasa.shared.nlu.constants import TEXT, INTENT, ENTITIES, ACTION_NAME
+from rasa.shared.nlu.training_data.features import Features
 from rasa.utils.tensorflow.constants import LABEL_PAD_ID
 from rasa.utils.tensorflow.model_data import ragged_array_to_ndarray
 
@@ -64,6 +67,10 @@ class InvalidStory(RasaException):
 class TrackerFeaturizer:
     """Base class for actual tracker featurizers."""
 
+    # Class registry to store all subclasses
+    _registry: ClassVar[Dict[str, Type["TrackerFeaturizer"]]] = {}
+    _featurizer_type: str = "TrackerFeaturizer"
+
     def __init__(
         self, state_featurizer: Optional[SingleStateFeaturizer] = None
     ) -> None:
@@ -74,6 +81,36 @@ class TrackerFeaturizer:
         """
         self.state_featurizer = state_featurizer
 
+    @classmethod
+    def register(cls, featurizer_type: str) -> Callable:
+        """Decorator to register featurizer subclasses."""
+
+        def wrapper(subclass: Type["TrackerFeaturizer"]) -> Type["TrackerFeaturizer"]:
+            cls._registry[featurizer_type] = subclass
+            # Store the type identifier in the class for serialization
+            subclass._featurizer_type = featurizer_type
+            return subclass
+
+        return wrapper
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "TrackerFeaturizer":
+        """Create featurizer instance from dictionary."""
+        featurizer_type = data.pop("type")
+
+        if featurizer_type not in cls._registry:
+            raise ValueError(f"Unknown featurizer type: {featurizer_type}")
+
+        # Get the correct subclass and instantiate it
+        subclass = cls._registry[featurizer_type]
+        return subclass.create_from_dict(data)
+
+    @classmethod
+    @abstractmethod
+    def create_from_dict(cls, data: Dict[str, Any]) -> "TrackerFeaturizer":
+        """Each subclass must implement its own creation from dict method."""
+        pass
+
     @staticmethod
     def _create_states(
         tracker: DialogueStateTracker,
@@ -465,9 +502,7 @@ class TrackerFeaturizer:
             self.state_featurizer.entity_tag_specs = []
 
         # noinspection PyTypeChecker
-        rasa.shared.utils.io.write_text_file(
-            str(jsonpickle.encode(self)), featurizer_file
-        )
+        rasa.shared.utils.io.dump_obj_as_json_to_file(featurizer_file, self.to_dict())
 
     @staticmethod
     def load(path: Union[Text, Path]) -> Optional[TrackerFeaturizer]:
@@ -481,7 +516,17 @@ class TrackerFeaturizer:
         """
         featurizer_file = Path(path) / FEATURIZER_FILE
         if featurizer_file.is_file():
-            return jsonpickle.decode(rasa.shared.utils.io.read_file(featurizer_file))
+            data = rasa.shared.utils.io.read_json_file(featurizer_file)
+
+            if "type" not in data:
+                logger.error(
+                    f"Couldn't load featurizer for policy. "
+                    f"File '{featurizer_file}' does not contain all "
+                    f"necessary information. 'type' is missing."
+                )
+                return None
+
+            return TrackerFeaturizer.from_dict(data)
 
         logger.error(
             f"Couldn't load featurizer for policy. "
@@ -508,7 +553,16 @@ class TrackerFeaturizer:
             )
         ]
 
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "type": self.__class__._featurizer_type,
+            "state_featurizer": (
+                self.state_featurizer.to_dict() if self.state_featurizer else None
+            ),
+        }
+
 
+@TrackerFeaturizer.register("FullDialogueTrackerFeaturizer")
 class FullDialogueTrackerFeaturizer(TrackerFeaturizer):
     """Creates full dialogue training data for time distributed architectures.
 
@@ -646,7 +700,20 @@ class FullDialogueTrackerFeaturizer(TrackerFeaturizer):
 
         return trackers_as_states
 
+    def to_dict(self) -> Dict[str, Any]:
+        return super().to_dict()
 
+    @classmethod
+    def create_from_dict(cls, data: Dict[str, Any]) -> "FullDialogueTrackerFeaturizer":
+        state_featurizer = SingleStateFeaturizer.create_from_dict(
+            data["state_featurizer"]
+        )
+        return cls(
+            state_featurizer,
+        )
+
+
+@TrackerFeaturizer.register("MaxHistoryTrackerFeaturizer")
 class MaxHistoryTrackerFeaturizer(TrackerFeaturizer):
     """Truncates the tracker history into `max_history` long sequences.
 
@@ -884,7 +951,25 @@ class MaxHistoryTrackerFeaturizer(TrackerFeaturizer):
 
         return trackers_as_states
 
+    def to_dict(self) -> Dict[str, Any]:
+        data = super().to_dict()
+        data.update(
+            {
+                "remove_duplicates": self.remove_duplicates,
+                "max_history": self.max_history,
+            }
+        )
+        return data
+
+    @classmethod
+    def create_from_dict(cls, data: Dict[str, Any]) -> "MaxHistoryTrackerFeaturizer":
+        state_featurizer = SingleStateFeaturizer.create_from_dict(
+            data["state_featurizer"]
+        )
+        return cls(state_featurizer, data["max_history"], data["remove_duplicates"])
 
+
+@TrackerFeaturizer.register("IntentMaxHistoryTrackerFeaturizer")
 class IntentMaxHistoryTrackerFeaturizer(MaxHistoryTrackerFeaturizer):
     """Truncates the tracker history into `max_history` long sequences.
 
@@ -1159,6 +1244,18 @@ class IntentMaxHistoryTrackerFeaturizer(MaxHistoryTrackerFeaturizer):
 
         return trackers_as_states
 
+    def to_dict(self) -> Dict[str, Any]:
+        return super().to_dict()
+
+    @classmethod
+    def create_from_dict(
+        cls, data: Dict[str, Any]
+    ) -> "IntentMaxHistoryTrackerFeaturizer":
+        state_featurizer = SingleStateFeaturizer.create_from_dict(
+            data["state_featurizer"]
+        )
+        return cls(state_featurizer, data["max_history"], data["remove_duplicates"])
+
 
 def _is_prev_action_unlikely_intent_in_state(state: State) -> bool:
     prev_action_name = state.get(PREVIOUS_ACTION, {}).get(ACTION_NAME)

rasa/core/nlg/contextual_response_rephraser.py

@@ -1,13 +1,12 @@
 from typing import Any, Dict, Optional, Text
 
-import os
 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
 from rasa.shared.constants import (
-    LLM_API_HEALTH_CHECK_ENV_VAR,
     LLM_CONFIG_KEY,
     MODEL_CONFIG_KEY,
     MODEL_NAME_CONFIG_KEY,
@@ -15,6 +14,7 @@ from rasa.shared.constants import (
     PROVIDER_CONFIG_KEY,
     OPENAI_PROVIDER,
     TIMEOUT_CONFIG_KEY,
+    MODEL_GROUP_CONFIG_KEY,
 )
 from rasa.shared.core.domain import KEY_RESPONSES_TEXT, Domain
 from rasa.shared.core.events import BotUttered, UserUttered
@@ -25,17 +25,14 @@ from rasa.shared.utils.llm import (
     USER,
     combine_custom_and_default_config,
     get_prompt_template,
-    llm_api_health_check,
     llm_factory,
-    try_instantiate_llm_client,
+    resolve_model_client_config,
 )
-from rasa.utils.endpoints import EndpointConfig
+from rasa.shared.utils.health_check import perform_training_time_llm_health_check
 from rasa.shared.utils.llm import (
     tracker_as_readable_transcript,
 )
-
-from rasa.core.nlg.summarize import summarize_conversation
-
+from rasa.utils.endpoints import EndpointConfig
 from rasa.utils.log_utils import log_llm
 
 structlogger = structlog.get_logger()
@@ -105,18 +102,18 @@ class ContextualResponseRephraser(TemplatedNaturalLanguageGenerator):
         self.trace_prompt_tokens = self.nlg_endpoint.kwargs.get(
             "trace_prompt_tokens", False
         )
-        llm_client = try_instantiate_llm_client(
+
+        self.llm_config = resolve_model_client_config(
             self.nlg_endpoint.kwargs.get(LLM_CONFIG_KEY),
+            ContextualResponseRephraser.__name__,
+        )
+
+        perform_training_time_llm_health_check(
+            self.llm_config,
             DEFAULT_LLM_CONFIG,
             "contextual_response_rephraser.init",
             ContextualResponseRephraser.__name__,
         )
-        if os.getenv(LLM_API_HEALTH_CHECK_ENV_VAR, "true").lower() == "true":
-            llm_api_health_check(
-                llm_client,
-                "contextual_response_rephraser.init",
-                ContextualResponseRephraser.__name__,
-            )
 
     def _last_message_if_human(self, tracker: DialogueStateTracker) -> Optional[str]:
         """Returns the latest message from the tracker.
@@ -145,9 +142,7 @@ class ContextualResponseRephraser(TemplatedNaturalLanguageGenerator):
         Returns:
             generated text
         """
-        llm = llm_factory(
-            self.nlg_endpoint.kwargs.get(LLM_CONFIG_KEY), DEFAULT_LLM_CONFIG
-        )
+        llm = llm_factory(self.llm_config, DEFAULT_LLM_CONFIG)
 
         try:
             llm_response = await llm.acompletion(prompt)
@@ -161,7 +156,7 @@ class ContextualResponseRephraser(TemplatedNaturalLanguageGenerator):
     def llm_property(self, prop: str) -> Optional[str]:
         """Returns a property of the LLM provider."""
         return combine_custom_and_default_config(
-            self.nlg_endpoint.kwargs.get(LLM_CONFIG_KEY), DEFAULT_LLM_CONFIG
+            self.llm_config, DEFAULT_LLM_CONFIG
         ).get(prop)
 
     def custom_prompt_template(self, prompt_template: str) -> Optional[str]:
@@ -194,9 +189,7 @@ class ContextualResponseRephraser(TemplatedNaturalLanguageGenerator):
         Returns:
         The history for the prompt.
         """
-        llm = llm_factory(
-            self.nlg_endpoint.kwargs.get(LLM_CONFIG_KEY), DEFAULT_LLM_CONFIG
-        )
+        llm = llm_factory(self.llm_config, DEFAULT_LLM_CONFIG)
         return await summarize_conversation(tracker, llm, max_turns=5)
 
     async def rephrase(
@@ -252,6 +245,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),
         )
         if not (updated_text := await self._generate_llm_response(prompt)):
             # If the LLM fails to generate a response, we

rasa/core/persistor.py

@@ -30,7 +30,7 @@ from rasa.shared.utils.io import raise_warning
 
 if TYPE_CHECKING:
     from azure.storage.blob import ContainerClient
-    import botocore
+    from botocore.exceptions import ClientError
 
 structlogger = structlog.get_logger()
 
@@ -233,13 +233,13 @@ class AWSPersistor(Persistor):
     def _ensure_bucket_exists(
         self, bucket_name: Text, region_name: Optional[Text] = None
     ) -> None:
-        import botocore
+        from botocore import exceptions
 
         # noinspection PyUnresolvedReferences
         try:
             self.s3.meta.client.head_bucket(Bucket=bucket_name)
-        except botocore.exceptions.ClientError as e:
-            if self.error_code(e) == HTTP_STATUS_FORBIDDEN:
+        except exceptions.ClientError as exc:
+            if self._error_code(exc) == HTTP_STATUS_FORBIDDEN:
                 log = (
                     f"Access to the specified bucket '{bucket_name}' is forbidden. "
                     "Please make sure you have the necessary "
@@ -251,7 +251,7 @@ class AWSPersistor(Persistor):
                     event_info=log,
                 )
                 raise RasaException(log)
-            elif self.error_code(e) == HTTP_STATUS_NOT_FOUND:
+            elif self._error_code(exc) == HTTP_STATUS_NOT_FOUND:
                 log = (
                     f"The specified bucket '{bucket_name}' does not exist. "
                     "Please make sure to create the bucket first."
@@ -264,7 +264,7 @@ class AWSPersistor(Persistor):
                 raise RasaException(log)
 
     @staticmethod
-    def error_code(e: botocore.exceptions.ClientError) -> int:
+    def _error_code(e: "ClientError") -> int:
         return int(e.response["Error"]["Code"])
 
     def _persist_tar(self, file_key: Text, tar_path: Text) -> None:
@@ -274,26 +274,48 @@ class AWSPersistor(Persistor):
 
     def _retrieve_tar(self, model_path: Text) -> None:
         """Downloads a model that has previously been persisted to s3."""
-        import botocore
+        from botocore import exceptions
 
         target_filename = os.path.basename(model_path)
-        try:
-            with open(target_filename, "wb") as f:
-                self.bucket.download_fileobj(model_path, f)
-        except botocore.exceptions.ClientError as e:
-            if self.error_code(e) == HTTP_STATUS_NOT_FOUND:
-                log = (
-                    f"Model '{target_filename}' not found in the specified bucket "
-                    f"'{self.bucket_name}'. Please make sure the model exists "
-                    f"in the bucket."
-                )
-                structlogger.error(
-                    "gcp_persistor.retrieve_tar.model_not_found",
-                    bucket_name=self.bucket_name,
-                    target_filename=target_filename,
-                    event_info=log,
-                )
-                raise ModelNotFound() from e
+        bucket_objects = list(self.bucket.objects.all())
+
+        model_found = False
+
+        log = (
+            f"Model '{target_filename}' not found in the specified bucket "
+            f"'{self.bucket_name}'. Please make sure the model exists "
+            f"in the bucket."
+        )
+
+        for obj in bucket_objects:
+            if model_path not in obj.key:
+                continue
+            structlogger.debug(
+                "aws_persistor.retrieve_tar.object_found", object_key=obj.key
+            )
+
+            try:
+                with open(target_filename, "wb") as f:
+                    self.bucket.download_fileobj(obj.key, f)
+                    model_found = True
+                    break
+            except exceptions.ClientError as exc:
+                if self._error_code(exc) == HTTP_STATUS_NOT_FOUND:
+                    structlogger.error(
+                        "aws_persistor.retrieve_tar.model_not_found",
+                        bucket_name=self.bucket_name,
+                        target_filename=target_filename,
+                        event_info=log,
+                    )
+                    raise ModelNotFound() from exc
+        if not model_found:
+            structlogger.error(
+                "aws_persistor.retrieve_tar.model_not_found",
+                bucket_name=self.bucket_name,
+                target_filename=target_filename,
+                event_info=log,
+            )
+            raise ModelNotFound()
 
 
 class GCSPersistor(Persistor):
@@ -322,7 +344,7 @@ class GCSPersistor(Persistor):
 
         try:
             self.storage_client.get_bucket(bucket_name)
-        except auth_exceptions.GoogleAuthError as e:
+        except auth_exceptions.GoogleAuthError as exc:
             log = (
                 f"An error occurred while authenticating with Google Cloud "
                 f"Storage. Please make sure you have the necessary credentials "
@@ -333,8 +355,8 @@ class GCSPersistor(Persistor):
                 bucket_name=bucket_name,
                 event_info=log,
             )
-            raise RasaException(log) from e
-        except exceptions.NotFound as e:
+            raise RasaException(log) from exc
+        except exceptions.NotFound as exc:
             log = (
                 f"The specified Google Cloud Storage bucket '{bucket_name}' "
                 f"does not exist. Please make sure to create the bucket first or "
@@ -345,20 +367,20 @@ class GCSPersistor(Persistor):
                 bucket_name=bucket_name,
                 event_info=log,
             )
-            raise RasaException(log) from e
-        except exceptions.Forbidden as e:
+            raise RasaException(log) from exc
+        except exceptions.Forbidden as exc:
             log = (
                 f"Access to the specified Google Cloud storage bucket '{bucket_name}' "
                 f"is forbidden. Please make sure you have the necessary "
-                f"permission to access the bucket. "
+                f"permissions to access the bucket. "
             )
             structlogger.error(
                 "gcp_persistor.ensure_bucket_exists.bucket_access_forbidden",
                 bucket_name=bucket_name,
                 event_info=log,
             )
-            raise RasaException(log) from e
-        except ValueError as e:
+            raise RasaException(log) from exc
+        except ValueError as exc:
             # bucket_name is None
             log = (
                 "The specified Google Cloud Storage bucket name is None. Please "
@@ -368,7 +390,7 @@ class GCSPersistor(Persistor):
                 "gcp_persistor.ensure_bucket_exists.bucket_name_none",
                 event_info=log,
             )
-            raise RasaException(log) from e
+            raise RasaException(log) from exc
 
     def _persist_tar(self, file_key: Text, tar_path: Text) -> None:
         """Uploads a model persisted in the `target_dir` to GCS."""
@@ -382,7 +404,7 @@ class GCSPersistor(Persistor):
         blob = self.bucket.blob(target_filename)
         try:
             blob.download_to_filename(target_filename)
-        except exceptions.NotFound as e:
+        except exceptions.NotFound as exc:
             log = (
                 f"Model '{target_filename}' not found in the specified bucket "
                 f"'{self.bucket_name}'. Please make sure the model exists "
@@ -394,7 +416,7 @@ class GCSPersistor(Persistor):
                 target_filename=target_filename,
                 event_info=log,
             )
-            raise ModelNotFound() from e
+            raise ModelNotFound() from exc
 
 
 class AzurePersistor(Persistor):
@@ -440,8 +462,33 @@ class AzurePersistor(Persistor):
 
     def _retrieve_tar(self, target_filename: Text) -> None:
         """Downloads a model that has previously been persisted to Azure."""
-        blob_client = self._container_client().get_blob_client(target_filename)
+        try:
+            blob_list = self._container_client().list_blobs()
+
+            for blob in blob_list:
+                if target_filename not in blob.name:
+                    continue
+
+                structlogger.debug(
+                    "azure_persistor.retrieve_tar.blob_found", blob_name=blob.name
+                )
 
-        with open(target_filename, "wb") as blob:
-            download_stream = blob_client.download_blob()
-            blob.write(download_stream.readall())
+                with open(target_filename, "wb") as model_file:
+                    blob_client = self._container_client().get_blob_client(blob.name)
+                    download_stream = blob_client.download_blob()
+                    model_file.write(download_stream.readall())
+        except Exception as exc:
+            log = (
+                f"An exception occurred while trying to download "
+                f"the model '{target_filename}' in the specified container "
+                f"'{self.container_name}'. Please make sure the model exists "
+                f"in the container."
+            )
+            structlogger.error(
+                "azure_persistor.retrieve_tar.model_download_error",
+                container_name=self.container_name,
+                target_filename=target_filename,
+                event_info=log,
+                exception=exc,
+            )
+            raise ModelNotFound() from exc