rasa-pro 3.12.0.dev1__py3-none-any.whl

rasa/core/exporter.py

@@ -0,0 +1,284 @@
+import logging
+import uuid
+import datetime
+from typing import AsyncIterator, Text, Optional, List, Set, Dict, Any
+
+from tqdm import tqdm
+
+import rasa.shared.utils.cli
+import rasa.shared.utils.io
+from rasa.core.brokers.broker import EventBroker
+from rasa.core.brokers.pika import PikaEventBroker
+from rasa.core.constants import RASA_EXPORT_PROCESS_ID_HEADER_NAME
+from rasa.core.tracker_store import TrackerStore
+from rasa.shared.core.trackers import EventVerbosity
+from rasa.exceptions import (
+    NoEventsToMigrateError,
+    NoConversationsInTrackerStoreError,
+    PublishingError,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class Exporter:
+    """Manages the publishing of events in a tracker store to an event broker.
+
+    Attributes:
+        endpoints_path: Path to the endpoints file used to configure the event
+            broker and tracker store. If `None`, the default path ('endpoints.yml')
+            is used.
+        tracker_store: `TrackerStore` to export conversations from.
+        event_broker: `EventBroker` to export conversations to.
+        requested_conversation_ids: List of conversation IDs requested to be
+            processed.
+        minimum_timestamp: Minimum timestamp of events that are published.
+            If `None`, apply no such constraint.
+        maximum_timestamp: Maximum timestamp of events that are published.
+            If `None`, apply no such constraint.
+    """
+
+    def __init__(
+        self,
+        tracker_store: TrackerStore,
+        event_broker: EventBroker,
+        endpoints_path: Text,
+        requested_conversation_ids: Optional[Text] = None,
+        minimum_timestamp: Optional[float] = None,
+        maximum_timestamp: Optional[float] = None,
+        offset_timestamps_by_seconds: Optional[int] = None,
+    ) -> None:
+        self.endpoints_path = endpoints_path
+        self.tracker_store = tracker_store
+
+        self.event_broker = event_broker
+        self.requested_conversation_ids = requested_conversation_ids
+        self.minimum_timestamp = minimum_timestamp
+        self.maximum_timestamp = maximum_timestamp
+        self.offset_timestamps_by_seconds = offset_timestamps_by_seconds
+
+    async def publish_events(self) -> int:
+        """Publish events in a tracker store using an event broker.
+
+        Exits if the publishing of events is interrupted due to an error. In that case,
+        the CLI command to continue the export where it was interrupted is printed.
+
+        Returns:
+            The number of successfully published events.
+        """
+        self._print_offset_info()
+
+        published_events = 0
+        current_timestamp = None
+
+        headers = self._get_message_headers()
+
+        async for event in self._fetch_events_within_time_range():
+            # noinspection PyBroadException
+            try:
+                self._publish_with_message_headers(event, headers)
+                published_events += 1
+                current_timestamp = event["timestamp"]
+            except Exception as e:
+                logger.exception(e)
+                raise PublishingError(current_timestamp)
+
+        await self.event_broker.close()
+
+        return published_events
+
+    def _print_offset_info(self) -> None:
+        """Output information about the offset applied to event timestamps."""
+        if self.offset_timestamps_by_seconds is None:
+            return
+
+        delta = datetime.timedelta(seconds=abs(self.offset_timestamps_by_seconds))
+        operator = "-" if self.offset_timestamps_by_seconds > 0 else ""
+        rasa.shared.utils.cli.print_info(
+            f"All event timestamps will be offset by {operator}{delta}! ⏰"
+        )
+
+    def _get_message_headers(self) -> Optional[Dict[Text, Text]]:
+        """Generate a message header for publishing events to a `PikaEventBroker`.
+
+        Returns:
+            Message headers with a randomly generated uuid under the
+            `RASA_EXPORT_PROCESS_ID_HEADER_NAME` key if `self.event_broker` is a
+            `PikaEventBroker`, else `None`.
+
+        """
+        if isinstance(self.event_broker, PikaEventBroker):
+            return {RASA_EXPORT_PROCESS_ID_HEADER_NAME: uuid.uuid4().hex}
+
+        return None
+
+    def _publish_with_message_headers(
+        self, original_event: Dict[Text, Any], headers: Optional[Dict[Text, Text]]
+    ) -> None:
+        """Publish `event` to a message broker with `headers`.
+
+        Args:
+            original_event: Serialized event to be published.
+            headers: Message headers to be published if `self.event_broker` is a
+                `PikaEventBroker`.
+
+        """
+        if self.offset_timestamps_by_seconds is not None:
+            event = dict(original_event)
+            event["timestamp"] += self.offset_timestamps_by_seconds
+        else:
+            event = original_event
+
+        if isinstance(self.event_broker, PikaEventBroker):
+            self.event_broker.publish(event=event, headers=headers)
+        else:
+            self.event_broker.publish(event)
+
+    async def _get_conversation_ids_in_tracker(self) -> Set[Text]:
+        """Fetch conversation IDs in `self.tracker_store`.
+
+        Returns:
+            A set of conversation IDs in `self.tracker_store`.
+
+        Raises:
+            `NoConversationsInTrackerStoreError` if
+            `conversation_ids_in_tracker_store` is empty.
+
+        """
+        conversation_ids_in_tracker_store = set(await self.tracker_store.keys())
+
+        if conversation_ids_in_tracker_store:
+            return conversation_ids_in_tracker_store
+
+        raise NoConversationsInTrackerStoreError(
+            "Could not find any conversations in connected tracker store. "
+            "Please validate your `endpoints.yml` and make sure the defined "
+            "tracker store exists. Exiting."
+        )
+
+    def _validate_all_requested_ids_exist(
+        self, conversation_ids_in_tracker_store: Set[Text]
+    ) -> None:
+        """Warn user if `self.requested_conversation_ids` contains IDs not found in
+        `conversation_ids_in_tracker_store`.
+
+        Args:
+            conversation_ids_in_tracker_store: Set of conversation IDs contained in
+            the tracker store.
+
+        """
+        missing_ids_in_tracker_store = (
+            set(self.requested_conversation_ids) - conversation_ids_in_tracker_store
+        )
+        if missing_ids_in_tracker_store:
+            rasa.shared.utils.cli.print_warning(
+                f"Could not find the following requested "
+                f"conversation IDs in connected tracker store: "
+                f"{', '.join(sorted(missing_ids_in_tracker_store))}"
+            )
+
+    async def _get_conversation_ids_to_process(self) -> Set[Text]:
+        """Get conversation IDs that are good for processing.
+
+        Finds the intersection of events that are contained in the tracker store with
+        those events requested as a command-line argument.
+
+        Returns:
+            Conversation IDs that are both requested and contained in the tracker
+            store. If no conversation IDs are requested, all conversation IDs in the
+            tracker store are returned.
+
+        """
+        conversation_ids_in_tracker_store = (
+            await self._get_conversation_ids_in_tracker()
+        )
+
+        if not self.requested_conversation_ids:
+            return conversation_ids_in_tracker_store
+
+        self._validate_all_requested_ids_exist(conversation_ids_in_tracker_store)
+
+        conversation_ids_to_process = conversation_ids_in_tracker_store & set(
+            self.requested_conversation_ids
+        )
+
+        if not conversation_ids_to_process:
+            raise NoEventsToMigrateError(
+                "Could not find an overlap between the requested "
+                "conversation IDs and those found in the tracker store. Exiting."
+            )
+
+        return conversation_ids_to_process
+
+    async def _fetch_events_within_time_range(self) -> AsyncIterator[Dict[Text, Any]]:
+        """Fetch all events for `conversation_ids` within the supplied time range.
+
+        Returns:
+            Serialized events with added `sender_id` field.
+
+        """
+        conversation_ids_to_process = await self._get_conversation_ids_to_process()
+
+        rasa.shared.utils.cli.print_info(
+            f"Fetching events for {len(conversation_ids_to_process)} "
+            f"conversation IDs:"
+        )
+        for conversation_id in tqdm(conversation_ids_to_process, "conversation IDs"):
+            tracker = await self.tracker_store.retrieve_full_tracker(conversation_id)
+            if not tracker:
+                logger.info(
+                    f"Could not retrieve tracker for conversation ID "
+                    f"'{conversation_id}'. Skipping."
+                )
+                continue
+
+            _events = tracker.current_state(EventVerbosity.ALL)["events"]
+
+            if not _events:
+                logger.info(
+                    f"No events to migrate for conversation ID '{conversation_id}'."
+                )
+                continue
+
+            events = self._get_events_for_conversation_id(_events, conversation_id)
+
+            # the order of events was changed after ATO-2192
+            # more context: https://github.com/RasaHQ/rasa/pull/13019
+            # we should sort the events by timestamp to keep the order
+            events.sort(key=lambda x: x["timestamp"])
+
+            # the conversation IDs are needed in the event publishing
+            for event in events:
+                if (
+                    self.minimum_timestamp is not None
+                    and event["timestamp"] < self.minimum_timestamp
+                ):
+                    continue
+                if (
+                    self.maximum_timestamp is not None
+                    and event["timestamp"] >= self.maximum_timestamp
+                ):
+                    continue
+                yield event
+
+    @staticmethod
+    def _get_events_for_conversation_id(
+        events: List[Dict[Text, Any]], conversation_id: Text
+    ) -> List[Dict[Text, Any]]:
+        """Get serialised events with added `sender_id` key.
+
+        Args:
+            events: Events to modify.
+            conversation_id: Conversation ID to add to events.
+
+        Returns:
+            Events with added `sender_id` key.
+
+        """
+        events_with_conversation_id = []
+
+        for event in events:
+            event["sender_id"] = conversation_id
+            events_with_conversation_id.append(event)
+
+        return events_with_conversation_id

rasa/core/featurizers/precomputation.py

@@ -0,0 +1,410 @@
+from __future__ import annotations
+from typing import Optional, Text, Dict, List, Union, Iterable, Any
+from collections.abc import ValuesView, KeysView
+
+from rasa.engine.graph import GraphComponent
+from rasa.engine.storage.storage import ModelStorage
+from rasa.engine.storage.resource import Resource
+from rasa.engine.graph import ExecutionContext
+from rasa.shared.core.domain import Domain, SubState
+from rasa.shared.core.events import ActionExecuted, UserUttered, Event
+from rasa.shared.core.trackers import DialogueStateTracker
+from rasa.shared.core.training_data.structures import StoryGraph
+from rasa.shared.nlu.constants import ACTION_NAME, ACTION_TEXT, INTENT, TEXT
+from rasa.shared.nlu.training_data.message import Message
+from rasa.shared.nlu.training_data.training_data import TrainingData
+from rasa.shared.nlu.training_data.features import Features
+import rasa.shared.utils.io
+
+# TODO: make precomputations (MessageContainerForCoreFeaturization) cacheable
+
+
+class MessageContainerForCoreFeaturization:
+    """A key-value store for specific `Messages`.
+
+    This container can be only be used to store messages that contain exactly
+    one of the following attributes: `ACTION_NAME`, `ACTION_TEXT`, `TEXT`, or `INTENT`.
+    A combination of the key attribute and the corresponding value will be used as
+    key for the respective message.
+
+    Background/Motivation:
+    - Our policies only require these attributes to be tokenized and/or featurized
+      via NLU graph components, which is why we don't care about storing anything else.
+    - Our tokenizers and featurizers work independently for each attribute,
+      which is why we can separate them and ask for "exactly one" of the key
+      attributes.
+    - Our tokenizers add attributes (e.g. token sequences) and not just `Features`,
+      which is why we need messages and why we allow messages to contain more than
+      just the key attributes.
+    - Due to the way we use this datastructure, it won't contain all features that the
+      policies need (cf. `rasa.core.featurizers.SingleStateFeaturizer`) and sometimes
+      the messages will contain no features at all, which is the motivation for the
+      name of this class.
+    - Values for different attributes might coincide (e.g. 'greet' can appear as user
+      text as well as name of an intent), but attributes are not all tokenized and
+      featurized in the same way, which is why we use the combination of key attribute
+      and value to identify a message.
+
+    Usage:
+    - At the start of core's featurization pipeline, we use this container to
+      de-duplicate the given story data during training (e.g. "Hello" might appear very
+      often but it will end up in the training data only once) and to de-duplicate
+      the data given in the tracker (e.g. if a text appears repeatedly in the
+      dialogue, it will only be featurized once later).
+      See: `rasa.core.featurizers.precomputation.CoreFeaturizationInputConverter`.
+    - At the end of core's featurization pipeline, we wrap all resulting
+      (training data) messages into this container again.
+      See: `rasa.core.featurizers.precomputation.CoreFeaturizationCollector`.
+    """
+
+    KEY_ATTRIBUTES = [ACTION_NAME, ACTION_TEXT, TEXT, INTENT]  # noqa: RUF012
+
+    def __init__(self) -> None:
+        """Creates an empty container for precomputations."""
+        self._table: Dict[Text, Dict[Text, Message]] = {
+            key: {} for key in self.KEY_ATTRIBUTES
+        }
+        self._num_collisions_ignored = 0
+
+    def fingerprint(self) -> Text:
+        """Fingerprint the container.
+
+        Returns:
+            hex string as a fingerprint of the container.
+        """
+        message_fingerprints = [
+            message.fingerprint() for message in self.all_messages()
+        ]
+        return rasa.shared.utils.io.deep_container_fingerprint(message_fingerprints)
+
+    def __repr__(self) -> Text:
+        return f"{self.__class__.__name__}({self._table})"
+
+    def __len__(self) -> int:
+        return sum(
+            len(key_attribute_table) for key_attribute_table in self._table.values()
+        )
+
+    def messages(self, key_attribute: Optional[Text] = None) -> ValuesView:
+        """Returns a view of all messages."""
+        if key_attribute not in self._table:
+            raise ValueError(
+                f"Expected key attribute (i.e. one of {self.KEY_ATTRIBUTES}) "
+                f"but received {key_attribute}."
+            )
+        return self._table[key_attribute].values()
+
+    def all_messages(self) -> List[Message]:
+        """Returns a list containing all messages."""
+        return [
+            message
+            for key_attribute_table in self._table.values()
+            for message in key_attribute_table.values()
+        ]
+
+    def keys(self, key_attribute: Text) -> KeysView:
+        """Returns a view of the value keys for the given key attribute."""
+        if key_attribute not in self._table:
+            raise ValueError(
+                f"Expected key attribute (i.e. one of {self.KEY_ATTRIBUTES}) "
+                f"but received {key_attribute}."
+            )
+        return self._table[key_attribute].keys()
+
+    @property
+    def num_collisions_ignored(self) -> int:
+        """Returns the number of collisions that have been ignored."""
+        return self._num_collisions_ignored
+
+    def add(self, message_with_one_key_attribute: Message) -> None:
+        """Adds the given message if it is not already present.
+
+        Args:
+          message_with_one_key_attribute: The message we want to add to the lookup
+            table. It must have exactly one key attribute.
+
+        Raises:
+          `ValueError` if the given message does not contain exactly one key
+          attribute or if there is a collision with a message that has a different
+          hash value
+        """
+        # extract the key pair
+        attributes = message_with_one_key_attribute.data.keys()
+        key_attributes = set(attributes).intersection(self.KEY_ATTRIBUTES)
+        if not key_attributes or len(key_attributes) != 1:
+            raise ValueError(
+                f"Expected exactly one attribute out of "
+                f"{self.KEY_ATTRIBUTES} but received {len(attributes)} attributes "
+                f"({attributes})."
+            )
+        key_attribute = list(key_attributes)[0]  # noqa: RUF015
+        key_value = str(message_with_one_key_attribute.data[key_attribute])
+        # extract the message
+        existing_message = self._table[key_attribute].get(key_value)
+        if existing_message is not None:
+            if hash(existing_message) != hash(message_with_one_key_attribute):
+                raise ValueError(
+                    f"Expected added message to be consistent. "
+                    f"({key_attribute}, {key_value}) already maps "
+                    f"to {existing_message}, but we want to add "
+                    f"{message_with_one_key_attribute} now."
+                )
+            else:
+                self._num_collisions_ignored += 1
+        else:
+            self._table[key_attribute][key_value] = message_with_one_key_attribute
+
+    def add_all(self, messages_with_one_key_attribute: List[Message]) -> None:
+        """Adds the given messages.
+
+        Args:
+          messages_with_one_key_attribute: The messages that we want to add.
+            Each one must have exactly one key attribute.
+
+        Raises:
+          `ValueError` if we cannot create a key for the given message or if there is
+          a collisions with a message that has a different hash value
+        """
+        for message in messages_with_one_key_attribute:
+            self.add(message)
+
+    def collect_features(
+        self, sub_state: SubState, attributes: Optional[Iterable[Text]] = None
+    ) -> Dict[Text, List[Features]]:
+        """Collects features for all attributes in the given substate.
+
+        There might be be multiple messages in the container that contain features
+        relevant for the given substate, e.g. this is the case if `TEXT` and
+        `INTENT` are present in the given substate. All of those messages will be
+        collected and their features combined.
+
+        Args:
+          sub_state: substate for which we want to extract the relevent features
+          attributes: if not `None`, this specifies the list of the attributes of the
+            `Features` that we're interested in (i.e. all other `Features` contained
+            in the relevant messages will be ignored)
+
+        Returns:
+          a dictionary that maps all the (requested) attributes to a list of `Features`
+
+        Raises:
+          `ValueError`: if there exists some key pair (i.e. key attribute and
+            corresponding value) from the given substate cannot be found
+          `RuntimeError`: if features for the same attribute are found in two
+            different messages that are associated with the given substate
+        """
+        # If we specify a list of attributes, then we want a dict with one entry
+        # for each attribute back - even if the corresponding list of features is empty.
+        features: Dict[Text, List[Features]] = (
+            dict()
+            if attributes is None
+            else {attribute: [] for attribute in attributes}
+        )
+        # collect all relevant key attributes
+        key_attributes = set(sub_state.keys()).intersection(self.KEY_ATTRIBUTES)
+        for key_attribute in key_attributes:
+            key_value = str(sub_state[key_attribute])
+            message = self._table[key_attribute].get(key_value)
+            if not message:
+                raise ValueError(
+                    f"Unknown key ({key_attribute},{key_value}). Cannot retrieve "
+                    f"features for substate {sub_state}"
+                )
+            features_from_message = Features.groupby_attribute(
+                message.features, attributes=attributes
+            )
+            for feat_attribute, feat_value in features_from_message.items():
+                existing_values = features.get(feat_attribute)
+                # Note: the following if-s are needed because if we specify a list of
+                # attributes then `features_from_message` will contain one entry per
+                # attribute even if the corresponding feature list is empty.
+                if feat_value and existing_values:
+                    raise RuntimeError(
+                        f"Feature for attribute {feat_attribute} has already been "
+                        f"extracted from a different message stored under a key "
+                        f"in {key_attributes} "
+                        f"that is different from {key_attribute}. This means there's a "
+                        f"redundancy in the message container."
+                    )
+                if feat_value:
+                    features[feat_attribute] = feat_value
+        return features
+
+    def lookup_message(self, user_text: Text) -> Message:
+        """Returns a message that contains the given user text.
+
+        Args:
+          user_text: the text of a user utterance
+        Raises:
+          `ValueError` if there is no message associated with the given user text
+        """
+        message = self._table[TEXT].get(user_text)
+        if message is None:
+            raise ValueError(
+                f"Expected a message with key ({TEXT}, {user_text}) in lookup table."
+            )
+        return message
+
+    def derive_messages_from_domain_and_add(self, domain: Domain) -> None:
+        """Adds all lookup table entries that can be derived from the domain.
+
+        That is, all action names, action texts, and intents defined in the domain
+        will be turned into a (separate) messages and added to this lookup table.
+
+        Args:
+          domain: the domain from which we extract the substates
+        """
+        if (
+            domain.action_texts
+            and domain.action_names_or_texts[-len(domain.action_texts) :]
+            != domain.action_texts
+        ):
+            raise NotImplementedError(
+                "We assumed that domain's `action_names_or_texts` start with a list of "
+                "all action names, followed by the action texts. "
+                "Please update the code to grab the action_name and action_texts from "
+                "the domain correctly."
+            )
+        action_texts = domain.action_texts
+        action_names = domain.action_names_or_texts[
+            slice(0, -len(domain.action_texts) if domain.action_texts else None)
+        ]
+
+        for key_attribute, actions in [
+            (ACTION_NAME, action_names),
+            (ACTION_TEXT, action_texts),
+        ]:
+            for action in actions:
+                self.add(Message({key_attribute: action}))
+
+        for intent in domain.intent_properties.keys():
+            self.add(Message({INTENT: intent}))
+
+    def derive_messages_from_events_and_add(self, events: Iterable[Event]) -> None:
+        """Adds all relevant messages that can be derived from the given events.
+
+        That is, each action name, action text, user text and intent that can be
+        found in the given events will be turned into a (separate) message and added
+        to this container.
+
+        Args:
+          events: list of events to extract the substate from
+        """
+        for event in events:
+            key_value_list = []
+            if isinstance(event, UserUttered):
+                key_value_list = [(TEXT, event.text), (INTENT, event.intent_name)]
+            elif isinstance(event, ActionExecuted):
+                key_value_list = [
+                    (ACTION_TEXT, event.action_text),
+                    (ACTION_NAME, event.action_name),
+                ]
+            for key, value in key_value_list:
+                if value is not None:
+                    self.add(Message(data={key: value}))
+
+
+class CoreFeaturizationInputConverter(GraphComponent):
+    """Provides data for the featurization pipeline.
+
+    During training as well as during inference, the converter de-duplicates the given
+    data (i.e. story graph or list of messages) such that each text and intent from a
+    user message and each action name and action text appears exactly once.
+    """
+
+    @classmethod
+    def create(
+        cls,
+        config: Dict[Text, Any],
+        model_storage: ModelStorage,
+        resource: Resource,
+        execution_context: ExecutionContext,
+    ) -> CoreFeaturizationInputConverter:
+        """Creates a new instance (see parent class for full docstring)."""
+        return cls()
+
+    def convert_for_training(
+        self, domain: Domain, story_graph: StoryGraph
+    ) -> TrainingData:
+        """Creates de-duplicated training data.
+
+        Each possible user text and intent and each action name and action text
+        that can be found in the given domain and story graph appears exactly once
+        in the resulting training data. Moreover, each item is contained in a separate
+        messsage.
+
+        Args:
+           domain: the domain
+           story_graph: a story graph
+        Returns:
+           training data
+        """
+        container = MessageContainerForCoreFeaturization()
+
+        # collect all action and user (intent-only) substates known from domain
+        container.derive_messages_from_domain_and_add(domain=domain)
+
+        # collect all substates we see in the given data
+        all_events = (
+            event
+            for step in story_graph.story_steps
+            for event in step.events
+            if isinstance(event, UserUttered)
+            # because all action names and texts are known to the domain
+        )
+        container.derive_messages_from_events_and_add(events=all_events)
+
+        # Reminder: in case of complex recipes that train CountVectorizers, we'll have
+        # to make sure that there is at least one user substate with a TEXT to ensure
+        # `CountVectorizer` is trained...
+
+        return TrainingData(training_examples=container.all_messages())
+
+    def convert_for_inference(self, tracker: DialogueStateTracker) -> List[Message]:
+        """Creates a list of messages containing single user and action attributes.
+
+        Each possible user text and intent and each action name and action text
+        that can be found in the events of the given tracker will appear exactly once
+        in the resulting messages. Moreover, each item is contained in a separate
+        messsage.
+
+        Args:
+          tracker: a dialogue state tracker containing events
+        Returns:
+          a list of messages
+        """
+        # Note: `tracker.applied_events()` doesn't convert any events to a different
+        # type and hence just iterating over the events is quicker than "applying"
+        # events first and then iterating over results (again).
+        container = MessageContainerForCoreFeaturization()
+        container.derive_messages_from_events_and_add(tracker.events)
+        return container.all_messages()
+
+
+class CoreFeaturizationCollector(GraphComponent):
+    """Collects featurized messages for use by a policy."""
+
+    @classmethod
+    def create(
+        cls,
+        config: Dict[Text, Any],
+        model_storage: ModelStorage,
+        resource: Resource,
+        execution_context: ExecutionContext,
+    ) -> CoreFeaturizationCollector:
+        """Creates a new instance (see parent class for full docstring)."""
+        return cls()
+
+    def collect(
+        self, messages: Union[TrainingData, List[Message]]
+    ) -> MessageContainerForCoreFeaturization:
+        """Collects messages."""
+        if isinstance(messages, TrainingData):
+            messages = messages.training_examples
+        # Note that the input messages had been contained in a lookup table in
+        # `StoryToTrainingDataConverter. Hence, we don't need to worry about
+        # collisions here anymore.
+        container = MessageContainerForCoreFeaturization()
+        for message in messages:
+            container.add(message)
+        return container