rasa-pro 3.12.21__py3-none-any.whl → 3.13.0__py3-none-any.whl

rasa/core/tracker_stores/mongo_tracker_store.py

@@ -0,0 +1,223 @@
+from __future__ import annotations
+
+import itertools
+from datetime import datetime
+from typing import Any, Dict, Iterable, Iterator, List, Optional, Text
+
+import structlog
+from pymongo.synchronous.collection import Collection
+
+from rasa.core.brokers.broker import EventBroker
+from rasa.core.tracker_stores.tracker_store import SerializedTrackerAsText, TrackerStore
+from rasa.shared.core.domain import Domain
+from rasa.shared.core.events import SessionStarted
+from rasa.shared.core.trackers import DialogueStateTracker, EventVerbosity
+
+structlogger = structlog.get_logger(__name__)
+
+
+class MongoTrackerStore(TrackerStore, SerializedTrackerAsText):
+    """Stores conversation history in Mongo.
+
+    Property methods:
+        conversations: returns the current conversation
+    """
+
+    def __init__(
+        self,
+        domain: Domain,
+        host: Optional[Text] = "mongodb://localhost:27017",
+        db: Optional[Text] = "rasa",
+        username: Optional[Text] = None,
+        password: Optional[Text] = None,
+        auth_source: Optional[Text] = "admin",
+        collection: Text = "conversations",
+        event_broker: Optional[EventBroker] = None,
+        **kwargs: Dict[Text, Any],
+    ) -> None:
+        from pymongo import MongoClient
+        from pymongo.database import Database
+
+        self.client: MongoClient = MongoClient(
+            host,
+            username=username,
+            password=password,
+            authSource=auth_source,
+            # delay connect until process forking is done
+            connect=False,
+        )
+
+        self.db = Database(self.client, db)
+        self.collection = collection
+        super().__init__(domain, event_broker, **kwargs)
+
+        self._ensure_indices()
+
+    @property
+    def conversations(self) -> Collection:
+        """Returns the current conversation."""
+        return self.db[self.collection]
+
+    def _ensure_indices(self) -> None:
+        """Create an index on the sender_id."""
+        self.conversations.create_index("sender_id")
+
+    @staticmethod
+    def _current_tracker_state_without_events(tracker: DialogueStateTracker) -> Dict:
+        # get current tracker state and remove `events` key from state
+        # since events are pushed separately in the `update_one()` operation
+        state = tracker.current_state(EventVerbosity.ALL)
+        state.pop("events", None)
+
+        return state
+
+    async def delete(self, sender_id: Text) -> None:
+        """Delete tracker for the given sender_id.
+
+        Args:
+            sender_id: Sender id of the tracker to be deleted.
+        """
+        if not await self.exists(sender_id):
+            structlogger.info(
+                "mongo_tracker_store.delete.no_tracker_for_sender_id",
+                event_info=f"Could not find tracker for conversation ID '{sender_id}'.",
+            )
+            return None
+        self.conversations.delete_one({"sender_id": sender_id})
+
+        structlogger.info(
+            "mongo_tracker_store.delete.deleted_tracker",
+            sender_id=sender_id,
+        )
+
+    async def save(self, tracker: DialogueStateTracker) -> None:
+        """Saves the current conversation state."""
+        await self.stream_events(tracker)
+
+        additional_events = self._additional_events(tracker)
+
+        self.conversations.update_one(
+            {"sender_id": tracker.sender_id},
+            {
+                "$set": self._current_tracker_state_without_events(tracker),
+                "$push": {
+                    "events": {"$each": [e.as_dict() for e in additional_events]}
+                },
+            },
+            upsert=True,
+        )
+
+    def _additional_events(self, tracker: DialogueStateTracker) -> Iterator:
+        """Return events from the tracker which aren't currently stored.
+
+        Args:
+            tracker: Tracker to inspect.
+
+        Returns:
+            List of serialised events that aren't currently stored.
+
+        """
+        stored = self.conversations.find_one({"sender_id": tracker.sender_id}) or {}
+        all_events = self._events_from_serialized_tracker(stored)
+
+        number_events_since_last_session = len(
+            self._events_since_last_session_start(all_events)
+        )
+
+        return itertools.islice(
+            tracker.events, number_events_since_last_session, len(tracker.events)
+        )
+
+    @staticmethod
+    def _events_from_serialized_tracker(serialised: Dict) -> List[Dict]:
+        return serialised.get("events", [])
+
+    @staticmethod
+    def _events_since_last_session_start(events: List[Dict]) -> List[Dict]:
+        """Retrieve events since and including the latest `SessionStart` event.
+
+        Args:
+            events: All events for a conversation ID.
+
+        Returns:
+            List of serialised events since and including the latest `SessionStarted`
+            event. Returns all events if no such event is found.
+
+        """
+        events_after_session_start = []
+        for event in reversed(events):
+            events_after_session_start.append(event)
+            if event["event"] == SessionStarted.type_name:
+                break
+
+        return list(reversed(events_after_session_start))
+
+    async def _retrieve(
+        self, sender_id: Text, fetch_events_from_all_sessions: bool
+    ) -> Optional[List[Dict[Text, Any]]]:
+        stored = self.conversations.find_one({"sender_id": sender_id})
+
+        # look for conversations which have used an `int` sender_id in the past
+        # and update them.
+        if not stored and sender_id.isdigit():
+            from pymongo import ReturnDocument
+
+            stored = self.conversations.find_one_and_update(
+                {"sender_id": int(sender_id)},
+                {"$set": {"sender_id": str(sender_id)}},
+                return_document=ReturnDocument.AFTER,
+            )
+
+        if not stored:
+            return None
+
+        events = self._events_from_serialized_tracker(stored)
+
+        if not fetch_events_from_all_sessions:
+            events = self._events_since_last_session_start(events)
+
+        return events
+
+    async def retrieve(self, sender_id: Text) -> Optional[DialogueStateTracker]:
+        """Retrieves tracker for the latest conversation session."""
+        events = await self._retrieve(sender_id, fetch_events_from_all_sessions=False)
+
+        if not events:
+            return None
+
+        return DialogueStateTracker.from_dict(sender_id, events, self.domain.slots)
+
+    async def retrieve_full_tracker(
+        self, conversation_id: Text
+    ) -> Optional[DialogueStateTracker]:
+        """Fetching all tracker events across conversation sessions."""
+        events = await self._retrieve(
+            conversation_id, fetch_events_from_all_sessions=True
+        )
+
+        if not events:
+            return None
+
+        return DialogueStateTracker.from_dict(
+            conversation_id, events, self.domain.slots
+        )
+
+    async def keys(self) -> Iterable[Text]:
+        """Returns sender_ids of the Mongo Tracker Store."""
+        return [c["sender_id"] for c in self.conversations.find()]
+
+    async def update(self, tracker: DialogueStateTracker) -> None:
+        """Overwrites the tracker for the given sender_id."""
+        self.conversations.replace_one(
+            {"sender_id": tracker.sender_id},
+            tracker.current_state(EventVerbosity.ALL),
+            upsert=True,
+        )
+
+        first_event_timestamp = str(datetime.fromtimestamp(tracker.events[0].timestamp))
+
+        structlogger.info(
+            "redis_tracker_store.update.updated_tracker",
+            sender_id=tracker.sender_id,
+            first_event_timestamp=first_event_timestamp,
+        )

rasa/core/tracker_stores/redis_tracker_store.py

@@ -0,0 +1,252 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any, Dict, Iterable, Optional, Text
+
+import structlog
+
+import rasa.shared
+from rasa.core.brokers.broker import EventBroker
+from rasa.core.tracker_stores.tracker_store import SerializedTrackerAsText, TrackerStore
+from rasa.shared.core.domain import Domain
+from rasa.shared.core.trackers import DialogueStateTracker
+
+structlogger = structlog.get_logger(__name__)
+
+# default value for key prefix in RedisTrackerStore
+DEFAULT_REDIS_TRACKER_STORE_KEY_PREFIX = "tracker:"
+
+
+class RedisTrackerStore(TrackerStore, SerializedTrackerAsText):
+    """Stores conversation history in Redis."""
+
+    def __init__(
+        self,
+        domain: Domain,
+        host: Text = "localhost",
+        port: int = 6379,
+        db: int = 0,
+        username: Optional[Text] = None,
+        password: Optional[Text] = None,
+        event_broker: Optional[EventBroker] = None,
+        record_exp: Optional[float] = None,
+        key_prefix: Optional[Text] = None,
+        use_ssl: bool = False,
+        ssl_keyfile: Optional[Text] = None,
+        ssl_certfile: Optional[Text] = None,
+        ssl_ca_certs: Optional[Text] = None,
+        **kwargs: Dict[Text, Any],
+    ) -> None:
+        """Initializes the tracker store."""
+        import redis
+
+        self.red = redis.StrictRedis(
+            host=host,
+            port=port,
+            db=db,
+            username=username,
+            password=password,
+            ssl=use_ssl,
+            ssl_keyfile=ssl_keyfile,
+            ssl_certfile=ssl_certfile,
+            ssl_ca_certs=ssl_ca_certs,
+            decode_responses=True,
+        )
+        self.record_exp = record_exp
+
+        self.key_prefix = DEFAULT_REDIS_TRACKER_STORE_KEY_PREFIX
+        if key_prefix:
+            structlogger.debug(
+                "redis_tracker_store.init.custom_key_prefix",
+                event_info=f"Setting non-default redis key prefix: '{key_prefix}'.",
+            )
+            self._set_key_prefix(key_prefix)
+
+        super().__init__(domain, event_broker, **kwargs)
+
+    def _set_key_prefix(self, key_prefix: Text) -> None:
+        if isinstance(key_prefix, str) and key_prefix.isalnum():
+            self.key_prefix = key_prefix + ":" + DEFAULT_REDIS_TRACKER_STORE_KEY_PREFIX
+        else:
+            structlogger.warning(
+                "redis_tracker_store.init.invalid_key_prefix",
+                event_info=(
+                    f"Omitting provided non-alphanumeric "
+                    f"redis key prefix: '{key_prefix}'. "
+                    f"Using default '{self.key_prefix}' instead."
+                ),
+            )
+
+    def _get_key_prefix(self) -> Text:
+        return self.key_prefix
+
+    async def save(
+        self, tracker: DialogueStateTracker, timeout: Optional[float] = None
+    ) -> None:
+        """Saves the current conversation state."""
+        await self.stream_events(tracker)
+
+        if not timeout and self.record_exp:
+            timeout = self.record_exp
+
+        # if the sender_id starts with the key prefix, we remove it
+        # this is used to avoid storing the prefix twice
+        sender_id = tracker.sender_id
+        if sender_id.startswith(self.key_prefix):
+            sender_id = sender_id[len(self.key_prefix) :]
+
+        stored = self.red.get(self.key_prefix + sender_id)
+
+        if stored is not None:
+            prior_tracker = self.deserialise_tracker(sender_id, stored)
+
+            tracker = self._merge_trackers(prior_tracker, tracker)
+
+        serialised_tracker = self.serialise_tracker(tracker)
+        self.red.set(self.key_prefix + sender_id, serialised_tracker, ex=timeout)
+
+    async def delete(self, sender_id: Text) -> None:
+        """Delete tracker for the given sender_id.
+
+        Args:
+            sender_id: Sender id of the tracker to be deleted.
+        """
+        if not await self.exists(sender_id):
+            structlogger.info(
+                "redis_tracker_store.delete.no_tracker_for_sender_id",
+                event_info=f"Could not find tracker for conversation ID '{sender_id}'.",
+            )
+            return None
+
+        if sender_id.startswith(self.key_prefix):
+            sender_id = sender_id[len(self.key_prefix) :]
+
+        self.red.delete(self.key_prefix + sender_id)
+        structlogger.info(
+            "redis_tracker_store.delete.deleted_tracker",
+            sender_id=sender_id,
+        )
+
+    async def retrieve(self, sender_id: Text) -> Optional[DialogueStateTracker]:
+        """Retrieves tracker for the latest conversation session.
+
+        The Redis key is formed by appending a prefix to sender_id.
+
+        Args:
+            sender_id: Conversation ID to fetch the tracker for.
+
+        Returns:
+            Tracker containing events from the latest conversation sessions.
+        """
+        return await self._retrieve(sender_id, fetch_all_sessions=False)
+
+    async def retrieve_full_tracker(
+        self, sender_id: Text
+    ) -> Optional[DialogueStateTracker]:
+        """Retrieves tracker for all conversation sessions.
+
+        The Redis key is formed by appending a prefix to sender_id.
+
+        Args:
+            sender_id: Conversation ID to fetch the tracker for.
+
+        Returns:
+            Tracker containing events from all conversation sessions.
+        """
+        return await self._retrieve(sender_id, fetch_all_sessions=True)
+
+    async def _retrieve(
+        self, sender_id: Text, fetch_all_sessions: bool
+    ) -> Optional[DialogueStateTracker]:
+        """Returns tracker matching sender_id.
+
+        Args:
+            sender_id: Conversation ID to fetch the tracker for.
+            fetch_all_sessions: Whether to fetch all sessions or only the last one.
+        """
+        if sender_id.startswith(self.key_prefix):
+            sender_id = sender_id[len(self.key_prefix) :]
+
+        stored = self.red.get(self.key_prefix + sender_id)
+        if stored is None:
+            structlogger.debug(
+                "redis_tracker_store.retrieve.no_tracker_for_sender_id",
+                event_info=f"Could not find tracker for conversation ID '{sender_id}'.",
+            )
+            return None
+
+        tracker = self.deserialise_tracker(sender_id, stored)
+        if fetch_all_sessions:
+            return tracker
+
+        # only return the last session
+        multiple_tracker_sessions = (
+            rasa.shared.core.trackers.get_trackers_for_conversation_sessions(tracker)
+        )
+
+        if len(multiple_tracker_sessions) <= 1:
+            return tracker
+
+        return multiple_tracker_sessions[-1]
+
+    async def keys(self) -> Iterable[Text]:
+        """Returns keys of the Redis Tracker Store."""
+        return self.red.keys(self.key_prefix + "*")
+
+    @staticmethod
+    def _merge_trackers(
+        prior_tracker: DialogueStateTracker, tracker: DialogueStateTracker
+    ) -> DialogueStateTracker:
+        """Merges two trackers.
+
+        Args:
+            prior_tracker: Tracker containing events from the previous conversation
+                sessions.
+            tracker: Tracker containing events from the current conversation session.
+        """
+        if not prior_tracker.events:
+            return tracker
+
+        last_event_timestamp = prior_tracker.events[-1].timestamp
+        past_tracker = tracker.travel_back_in_time(target_time=last_event_timestamp)
+
+        if past_tracker.events == prior_tracker.events:
+            return tracker
+
+        merged = tracker.init_copy()
+        merged.update_with_events(list(prior_tracker.events), override_timestamp=False)
+
+        for new_event in tracker.events:
+            # Event subclasses implement `__eq__` method that make it difficult
+            # to compare events. We use `as_dict` to compare events.
+            if all(
+                [
+                    new_event.as_dict() != existing_event.as_dict()
+                    for existing_event in merged.events
+                ]
+            ):
+                merged.update(new_event)
+
+        return merged
+
+    async def update(self, tracker: DialogueStateTracker) -> None:
+        """Overwrites the tracker for the given sender_id."""
+        serialised_tracker = self.serialise_tracker(tracker)
+
+        # if the sender_id starts with the key prefix, we remove it
+        # this is used to avoid storing the prefix twice
+        sender_id = tracker.sender_id
+        if sender_id.startswith(self.key_prefix):
+            sender_id = sender_id[len(self.key_prefix) :]
+
+        self.red.set(
+            self.key_prefix + sender_id, serialised_tracker, ex=self.record_exp
+        )
+
+        first_event_timestamp = str(datetime.fromtimestamp(tracker.events[0].timestamp))
+
+        structlogger.info(
+            "redis_tracker_store.update.updated_tracker",
+            sender_id=tracker.sender_id,
+            first_event_timestamp=first_event_timestamp,
+        )