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

rasa/dialogue_understanding/generator/flow_retrieval.py

@@ -0,0 +1,466 @@
+"""The module is primarily centered around the `FlowRetrieval` class which handles the
+initialization, configuration validation, vector store management, and flow retrieval
+logic. It integrates components for managing embeddings, vector stores, and
+flow-specific templates, facilitating semantic search functionalities.
+Key Features:
+- Configurable embedding strategies for dialogue components.
+- Seamless interaction with model storage and resource management.
+- Supports dynamic loading and persistence of vector stores.
+- Enables population of vector stores based on specified dialogue flows and domain
+information.
+- Implements flow retrieval mechanisms including semantic search based on dialogue
+context.
+Usage:
+Interaction with this class typically involves creating an instance with a
+configuration dict, model storage, and a resource reference, then using methods
+like `populate`, `persist`, or dynamic retrieval methods to manage or utilize
+flows within a conversational context.
+"""
+
+import importlib
+from typing import Dict, Text, Any, List, Optional
+
+import structlog
+from jinja2 import Template
+from langchain.docstore.document import Document
+from langchain.schema.embeddings import Embeddings
+from langchain_community.vectorstores.faiss import FAISS
+from langchain_community.vectorstores.utils import DistanceStrategy
+
+from rasa.engine.storage.resource import Resource
+from rasa.engine.storage.storage import ModelStorage
+import rasa.shared.utils.io
+from rasa.shared.constants import (
+    EMBEDDINGS_CONFIG_KEY,
+    PROVIDER_CONFIG_KEY,
+    OPENAI_PROVIDER,
+)
+from rasa.shared.core.domain import Domain
+from rasa.shared.core.flows import FlowsList
+from rasa.shared.core.trackers import DialogueStateTracker
+from rasa.shared.exceptions import ProviderClientAPIException
+from rasa.shared.nlu.constants import TEXT, FLOWS_FROM_SEMANTIC_SEARCH
+from rasa.shared.nlu.training_data.message import Message
+from rasa.shared.providers.embedding._langchain_embedding_client_adapter import (
+    _LangchainEmbeddingClientAdapter,
+)
+from rasa.shared.utils.health_check.embeddings_health_check_mixin import (
+    EmbeddingsHealthCheckMixin,
+)
+from rasa.shared.utils.llm import (
+    tracker_as_readable_transcript,
+    embedder_factory,
+    DEFAULT_OPENAI_EMBEDDING_MODEL_NAME,
+    USER,
+    get_prompt_template,
+    allowed_values_for_slot,
+    resolve_model_client_config,
+)
+
+DEFAULT_FLOW_DOCUMENT_TEMPLATE = importlib.resources.read_text(
+    "rasa.dialogue_understanding.generator", "flow_document_template.jinja2"
+)
+
+FLOW_RETRIEVAL_CONFIG_FILE_NAME = "flow_retrieval_config.json"
+
+DEFAULT_EMBEDDINGS_CONFIG = {
+    PROVIDER_CONFIG_KEY: OPENAI_PROVIDER,
+    "model": DEFAULT_OPENAI_EMBEDDING_MODEL_NAME,
+}
+
+MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY = "num_flows"
+SHOULD_EMBED_SLOTS_KEY = "should_embed_slots"
+TURNS_TO_EMBED_KEY = "turns_to_embed"
+
+DEFAULT_MAX_FLOWS_FROM_SEMANTIC_SEARCH = 20
+DEFAULT_TURNS_TO_EMBED = 1
+DEFAULT_SHOULD_EMBED_SLOTS = True
+
+
+structlogger = structlog.get_logger()
+
+
+class FlowRetrieval(EmbeddingsHealthCheckMixin):
+    @classmethod
+    def get_default_config(cls) -> Dict[str, Any]:
+        """The default config for the flow retrieval."""
+        return {
+            EMBEDDINGS_CONFIG_KEY: DEFAULT_EMBEDDINGS_CONFIG,
+            MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY: DEFAULT_MAX_FLOWS_FROM_SEMANTIC_SEARCH,
+            TURNS_TO_EMBED_KEY: DEFAULT_TURNS_TO_EMBED,
+            SHOULD_EMBED_SLOTS_KEY: DEFAULT_SHOULD_EMBED_SLOTS,
+        }
+
+    def __init__(
+        self,
+        config: Dict[str, Any],
+        model_storage: ModelStorage,
+        resource: Resource,
+    ):
+        config = {**self.get_default_config(), **config}
+        self.config = self.validate_config(config)
+        self.config[EMBEDDINGS_CONFIG_KEY] = resolve_model_client_config(
+            self.config.get(EMBEDDINGS_CONFIG_KEY), FlowRetrieval.__name__
+        )
+        self.vector_store: Optional[FAISS] = None
+        self.flow_document_template = get_prompt_template(
+            None, DEFAULT_FLOW_DOCUMENT_TEMPLATE
+        )
+        self._model_storage = model_storage
+        self._resource = resource
+
+    @classmethod
+    def validate_config(cls, config: Dict[Text, Any]) -> Dict[Text, Any]:
+        if config[MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY] < 0:
+            config[MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY] = (
+                DEFAULT_MAX_FLOWS_FROM_SEMANTIC_SEARCH
+            )
+            structlogger.error(
+                f"flow_retrieval.validate_config.{MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY}.set_as_negative",
+                event_info=(
+                    f"The `{MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY}` is a "
+                    f"negative value. Setting it to the default value "
+                    f"({DEFAULT_MAX_FLOWS_FROM_SEMANTIC_SEARCH})."
+                ),
+                old_value=config[MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY],
+                new_value=DEFAULT_MAX_FLOWS_FROM_SEMANTIC_SEARCH,
+            )
+
+        if config[TURNS_TO_EMBED_KEY] < 1:
+            config[TURNS_TO_EMBED_KEY] = DEFAULT_TURNS_TO_EMBED
+            structlogger.error(
+                f"flow_retrieval.validate_config.{TURNS_TO_EMBED_KEY}.less_than_one",
+                event_info=(
+                    f"The `{TURNS_TO_EMBED_KEY}` is less than 1."
+                    f"Setting it to the default value ({DEFAULT_TURNS_TO_EMBED})."
+                    f"Only the latest user utterance will be used."
+                ),
+                old_value=config[TURNS_TO_EMBED_KEY],
+                new_value=DEFAULT_TURNS_TO_EMBED,
+            )
+
+        return config
+
+    @classmethod
+    def load(
+        cls,
+        config: Dict[Text, Any],
+        model_storage: ModelStorage,
+        resource: Resource,
+        **kwargs: Any,
+    ) -> "FlowRetrieval":
+        """Load flow retrieval with previously populated FAISS vector store."""
+
+        # Perform health check on resolved embedding client config
+        embeddings_config = resolve_model_client_config(
+            config.get(EMBEDDINGS_CONFIG_KEY, {})
+        )
+        cls.perform_embeddings_health_check(
+            embeddings_config,
+            DEFAULT_EMBEDDINGS_CONFIG,
+            "flow_retrieval.load",
+            FlowRetrieval.__name__,
+        )
+
+        # initialize base flow retrieval
+        flow_retrieval = FlowRetrieval(config, model_storage, resource)
+        # load vector store
+        vector_store = cls._load_vector_store(
+            flow_retrieval.config, model_storage, resource
+        )
+        flow_retrieval.vector_store = vector_store
+
+        return flow_retrieval
+
+    @classmethod
+    def _load_vector_store(
+        cls, config: Dict[Text, Any], model_storage: ModelStorage, resource: Resource
+    ) -> Optional[FAISS]:
+        """Loads a FAISS vector store from a specified local path."""
+        embeddings = cls._create_embedder(config)
+        try:
+            with model_storage.read_from(resource) as model_path:
+                return FAISS.load_local(
+                    folder_path=model_path,
+                    embeddings=embeddings,
+                    distance_strategy=DistanceStrategy.MAX_INNER_PRODUCT,
+                    allow_dangerous_deserialization=True,
+                )
+        except Exception as e:
+            structlogger.warning(
+                "flow_retrieval.load_vector_store.failed",
+                error=e,
+                resource=resource.name,
+            )
+            return None
+
+    @classmethod
+    def _create_embedder(cls, config: Dict[Text, Any]) -> Embeddings:
+        """Creates an embedder.
+
+        Returns:
+            The embedder.
+        """
+        # Copy the config so original config is not modified
+        config = config.copy()
+        # Resolve config and instantiate the embedding client
+        config[EMBEDDINGS_CONFIG_KEY] = resolve_model_client_config(
+            config.get(EMBEDDINGS_CONFIG_KEY), FlowRetrieval.__name__
+        )
+        client = embedder_factory(
+            config.get(EMBEDDINGS_CONFIG_KEY), DEFAULT_EMBEDDINGS_CONFIG
+        )
+        # Wrap the embedding client in the adapter
+        return _LangchainEmbeddingClientAdapter(client)
+
+    def persist(self) -> None:
+        self._persist_vector_store()
+        self._persist_config()
+
+    def _persist_vector_store(self) -> None:
+        """Persists the FAISS vector store."""
+        if self.vector_store is not None:
+            with self._model_storage.write_to(self._resource) as model_path:
+                self.vector_store.save_local(model_path)
+        else:
+            structlogger.warning(
+                "flow_retrieval.persist_vector_store.not_initialized",
+                event_info="Vector store is None, not persisted.",
+            )
+
+    def _persist_config(self) -> None:
+        with self._model_storage.write_to(self._resource) as path:
+            rasa.shared.utils.io.dump_obj_as_json_to_file(
+                path / FLOW_RETRIEVAL_CONFIG_FILE_NAME, self.config
+            )
+
+    def populate(self, flows: FlowsList, domain: Domain) -> None:
+        """Populates the vector store with embeddings generated from
+        documents based on the flow descriptions, and flow slots
+        descriptions.
+
+        Args:
+            flows: List of flows to populate the vector store with.
+            domain: The domain containing relevant slot information.
+        """
+        # Perform health check before populating the vector store with flows
+        self.perform_embeddings_health_check(
+            self.config.get(EMBEDDINGS_CONFIG_KEY),
+            DEFAULT_EMBEDDINGS_CONFIG,
+            "flow_retrieval.train",
+            FlowRetrieval.__name__,
+        )
+
+        flows_to_embedd = flows.exclude_link_only_flows()
+        embeddings = self._create_embedder(self.config)
+        documents = self._generate_flow_documents(flows_to_embedd, domain)
+        try:
+            self.vector_store = FAISS.from_documents(
+                documents=documents,
+                embedding=embeddings,
+                distance_strategy=DistanceStrategy.MAX_INNER_PRODUCT,
+            )
+        except Exception as e:
+            error_type = e.__class__.__name__
+            structlogger.error(
+                "flow_retrieval.populate_vector_store.not_populated",
+                event_info=(
+                    "Failed to populate the FAISS store with the provided flows."
+                ),
+                error_type=error_type,
+                error=e,
+            )
+            raise
+
+    def _generate_flow_documents(
+        self, flows: FlowsList, domain: Domain
+    ) -> List[Document]:
+        """Converts each flow in the provided list into an embeddable document. These
+        documents include the following information for each flow: flow's name,
+        flow's description, and associated slots and their descriptions and allowed
+        values.
+
+        Args:
+            flows: List of flows.
+            domain: The domain containing relevant slot information.
+
+        Returns:
+            List of documents, each representing the converted information of a single
+            flow.
+        """
+        slots = {slot.name: slot for slot in domain.slots}
+        flow_docs = []
+
+        for flow in flows:
+            flow_info: Dict[Text, Any] = {
+                "flow": {
+                    "name": flow.name,
+                    "description": flow.description,
+                    "slots": [],
+                }
+            }
+            if self.config[SHOULD_EMBED_SLOTS_KEY]:
+                flow_info["flow"]["slots"] = [
+                    {
+                        "name": q.collect,
+                        "description": q.description,
+                        "allowed_values": allowed_values_for_slot(slots[q.collect]),
+                    }
+                    for q in flow.get_collect_steps()
+                ]
+
+            flow_docs.append(
+                Document(
+                    page_content=Template(self.flow_document_template)
+                    .render(flow_info)
+                    .strip(),
+                    metadata={"flow_id": flow.id},
+                )
+            )
+
+        return flow_docs
+
+    async def filter_flows(
+        self, tracker: DialogueStateTracker, message: Message, flows: FlowsList
+    ) -> FlowsList:
+        """Filters the given flows.
+        The filtered flows with the following rules:
+        - Exclude flows that reachable only via link
+        - Include flows set to be always included in the prompt
+        - Include flows started during the conversation,
+        - (Optionally) The rest of flows are filter so only the top 'k' with
+          the highest similarity to the ongoing conversation are included.
+          Set 'max_flows_from_semantic_search' to negative number to disable
+          the option.
+
+        Args:
+            tracker: the tracker
+            message: the user message
+            flows: available flows
+
+        Returns:
+            List of flows that includes always-included flows, previously started flows
+            and top `k` relevant flows for the current conversation.
+        """
+        # apply basic filtering
+        flows = flows.exclude_link_only_flows()
+        always_included_flows = flows.get_flows_always_included_in_prompt()
+        previously_started_flows = tracker.get_previously_started_flows(flows)
+        # apply semantic search filtering
+        most_similar_flows = await self.find_most_similar_flows(tracker, message, flows)
+        return FlowsList.from_multiple_flows_lists(
+            always_included_flows,
+            previously_started_flows,
+            most_similar_flows,
+        )
+
+    async def find_most_similar_flows(
+        self, tracker: DialogueStateTracker, message: Message, flows: FlowsList
+    ) -> FlowsList:
+        """Filters the given flows so only the top 'k' most similar
+        flows to the current conversation are left.
+        retrieved from the vector store.
+
+        Args:
+            tracker: the tracker
+            message: the user message
+
+        Returns:
+            The most similar flows to the current conversation.
+        """
+        query = self._prepare_query(tracker, message)
+        documents_with_scores = await self._query_vector_store(query)
+        # filter out None i.e. more flows were embedded during training than are
+        # available during prediction
+        most_similar_flows_with_scores = [
+            (flow, score)
+            for doc, score in documents_with_scores
+            if (flow := flows.flow_by_id(doc.metadata["flow_id"])) is not None
+        ]
+        # sort by score decreasing
+        most_similar_flows_with_scores.sort(key=lambda x: x[1], reverse=True)
+        # add the relevant flows to the message for evaluation purposes
+        message.set(
+            prop=FLOWS_FROM_SEMANTIC_SEARCH,
+            info=[
+                (flow.id, float(score))
+                for flow, score in most_similar_flows_with_scores
+            ],
+            add_to_output=True,
+        )
+        return FlowsList([f for f, _ in most_similar_flows_with_scores])
+
+    def _prepare_query(self, tracker: DialogueStateTracker, message: Message) -> Text:
+        """Prepares the query for vector store. The query is composed
+        of the conversation within the specified number of turns.
+
+        Args:
+            tracker: the tracker
+            message: the user message
+
+        Returns:
+            The query for vector store.
+        """
+        if int(self.config[TURNS_TO_EMBED_KEY]) > 1:
+            current_conversation = tracker_as_readable_transcript(
+                tracker,
+                human_prefix=USER,
+                max_turns=int(self.config[TURNS_TO_EMBED_KEY]),
+            )
+            current_conversation += f"\n{USER}: {message.data[TEXT]}"
+            return current_conversation
+
+        return f"{message.data[TEXT]}"
+
+    async def _query_vector_store(self, query: Text) -> List:
+        """Compares the query with all flows using a vector store
+        and returns the top k relevant flows for the current conversation.
+
+        Args:
+            query: query
+
+        Returns:
+            The top k documents with similarity scores.
+        """
+        if self.vector_store is None:
+            structlogger.error(
+                "flow_retrieval.query_vector_store.vector_store_not_configured",
+                event_info="Vector store is not configured",
+            )
+            return []
+        try:
+            documents_with_scores = (
+                await self.vector_store.asimilarity_search_with_score(
+                    query, k=int(self.config[MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY])
+                )
+            )
+            structlogger.debug(
+                "flow_retrieval.query_vector_store.fetched",
+                event_info=(
+                    f"Fetched the top {self.config[MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY]}"
+                    f"similar flows from the vector store"
+                ),
+                query=query,
+                top_k=self.config[MAX_FLOWS_FROM_SEMANTIC_SEARCH_KEY],
+                results=[
+                    {
+                        "flow_id": document.metadata["flow_id"],
+                        "score": score,
+                        "content": document.page_content,
+                    }
+                    for document, score in documents_with_scores
+                ],
+            )
+            return documents_with_scores
+        except Exception as e:
+            error_type = e.__class__.__name__
+            structlogger.error(
+                "flow_retrieval.query_vector_store.error",
+                event_info="Cannot fetch flows from vector store",
+                error_type=error_type,
+                error=e,
+                query=query,
+            )
+            raise ProviderClientAPIException(
+                message="Cannot fetch flows from vector store", original_exception=e
+            )