rasa-pro 3.10.16__py3-none-any.whl → 3.11.0__py3-none-any.whl

rasa/core/policies/enterprise_search_policy.py

@@ -2,43 +2,41 @@ import importlib.resources
 import json
 import re
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Text
-
 import dotenv
 import structlog
 from jinja2 import Template
 from pydantic import ValidationError
 
-from rasa.shared.providers.embedding._langchain_embedding_client_adapter import (
-    _LangchainEmbeddingClientAdapter,
-)
-from rasa.shared.providers.llm.llm_client import LLMClient
-
 import rasa.shared.utils.io
-from rasa.telemetry import (
-    track_enterprise_search_policy_predict,
-    track_enterprise_search_policy_train_completed,
-    track_enterprise_search_policy_train_started,
-)
-from rasa.shared.exceptions import RasaException
 from rasa.core.constants import (
     POLICY_MAX_HISTORY,
     POLICY_PRIORITY,
     SEARCH_POLICY_PRIORITY,
     UTTER_SOURCE_METADATA_KEY,
 )
+from rasa.core.information_retrieval import (
+    InformationRetrieval,
+    SearchResult,
+    InformationRetrievalException,
+    create_from_endpoint_config,
+)
+from rasa.core.information_retrieval.faiss import FAISS_Store
 from rasa.core.policies.policy import Policy, PolicyPrediction
 from rasa.core.utils import AvailableEndpoints
-from rasa.dialogue_understanding.patterns.internal_error import (
-    InternalErrorPatternFlowStackFrame,
+from rasa.dialogue_understanding.generator.constants import (
+    LLM_CONFIG_KEY,
 )
 from rasa.dialogue_understanding.patterns.cannot_handle import (
     CannotHandlePatternFlowStackFrame,
 )
-from rasa.dialogue_understanding.stack.frames import PatternFlowStackFrame
+from rasa.dialogue_understanding.patterns.internal_error import (
+    InternalErrorPatternFlowStackFrame,
+)
 from rasa.dialogue_understanding.stack.frames import (
     DialogueStackFrame,
     SearchStackFrame,
 )
+from rasa.dialogue_understanding.stack.frames import PatternFlowStackFrame
 from rasa.engine.graph import ExecutionContext
 from rasa.engine.recipes.default_recipe import DefaultV1Recipe
 from rasa.engine.storage.resource import Resource
@@ -47,13 +45,13 @@ from rasa.graph_components.providers.forms_provider import Forms
 from rasa.graph_components.providers.responses_provider import Responses
 from rasa.shared.constants import (
     EMBEDDINGS_CONFIG_KEY,
-    LLM_CONFIG_KEY,
     MODEL_CONFIG_KEY,
-    MODEL_NAME_CONFIG_KEY,
     PROMPT_CONFIG_KEY,
     PROVIDER_CONFIG_KEY,
     OPENAI_PROVIDER,
     TIMEOUT_CONFIG_KEY,
+    MODEL_NAME_CONFIG_KEY,
+    MODEL_GROUP_ID_CONFIG_KEY,
 )
 from rasa.shared.core.constants import (
     ACTION_CANCEL_FLOW,
@@ -64,8 +62,17 @@ from rasa.shared.core.domain import Domain
 from rasa.shared.core.events import Event, UserUttered, BotUttered
 from rasa.shared.core.generator import TrackerWithCachedStates
 from rasa.shared.core.trackers import DialogueStateTracker, EventVerbosity
+from rasa.shared.exceptions import RasaException, FileIOException
 from rasa.shared.nlu.training_data.training_data import TrainingData
+from rasa.shared.providers.embedding._langchain_embedding_client_adapter import (
+    _LangchainEmbeddingClientAdapter,
+)
+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,
@@ -75,15 +82,12 @@ from rasa.shared.utils.llm import (
     llm_factory,
     sanitize_message_for_prompt,
     tracker_as_readable_transcript,
-    try_instantiate_llm_client,
-    try_instantiate_embedder,
+    resolve_model_client_config,
 )
-from rasa.core.information_retrieval.faiss import FAISS_Store
-from rasa.core.information_retrieval import (
-    InformationRetrieval,
-    SearchResult,
-    InformationRetrievalException,
-    create_from_endpoint_config,
+from rasa.telemetry import (
+    track_enterprise_search_policy_predict,
+    track_enterprise_search_policy_train_completed,
+    track_enterprise_search_policy_train_started,
 )
 
 if TYPE_CHECKING:
@@ -128,6 +132,7 @@ DEFAULT_EMBEDDINGS_CONFIG = {
 }
 
 ENTERPRISE_SEARCH_PROMPT_FILE_NAME = "enterprise_search_policy_prompt.jinja2"
+ENTERPRISE_SEARCH_CONFIG_FILE_NAME = "config.json"
 
 SEARCH_RESULTS_METADATA_KEY = "search_results"
 SEARCH_QUERY_METADATA_KEY = "search_query"
@@ -152,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
@@ -198,24 +203,35 @@ class EnterpriseSearchPolicy(Policy):
         """Constructs a new Policy object."""
         super().__init__(config, model_storage, resource, execution_context, featurizer)
 
+        # Resolve LLM config
+        self.config[LLM_CONFIG_KEY] = resolve_model_client_config(
+            self.config.get(LLM_CONFIG_KEY), EnterpriseSearchPolicy.__name__
+        )
+        # Resolve embeddings config
+        self.config[EMBEDDINGS_CONFIG_KEY] = resolve_model_client_config(
+            self.config.get(EMBEDDINGS_CONFIG_KEY), EnterpriseSearchPolicy.__name__
+        )
+
         # Vector store object and configuration
         self.vector_store = vector_store
-        self.vector_store_config = config.get(
+        self.vector_store_config = self.config.get(
             VECTOR_STORE_PROPERTY, DEFAULT_VECTOR_STORE
         )
+
         # Embeddings configuration for encoding the search query
-        self.embeddings_config = self.config.get(
-            EMBEDDINGS_CONFIG_KEY, DEFAULT_EMBEDDINGS_CONFIG
+        self.embeddings_config = (
+            self.config[EMBEDDINGS_CONFIG_KEY] or DEFAULT_EMBEDDINGS_CONFIG
         )
+
+        # LLM Configuration for response generation
+        self.llm_config = self.config[LLM_CONFIG_KEY] or DEFAULT_LLM_CONFIG
+
         # Maximum number of turns to include in the prompt
         self.max_history = self.config.get(POLICY_MAX_HISTORY)
 
         # Maximum number of messages to include in the search query
         self.max_messages_in_query = self.config.get(MAX_MESSAGES_IN_QUERY_KEY, 2)
 
-        # LLM Configuration for response generation
-        self.llm_config = self.config.get(LLM_CONFIG_KEY, DEFAULT_LLM_CONFIG)
-
         # boolean to enable/disable tracing of prompt tokens
         self.trace_prompt_tokens = self.config.get(TRACE_TOKENS_PROPERTY, False)
 
@@ -244,9 +260,16 @@ class EnterpriseSearchPolicy(Policy):
         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), EnterpriseSearchPolicy.__name__
+        )
         client = embedder_factory(
             config.get(EMBEDDINGS_CONFIG_KEY), DEFAULT_EMBEDDINGS_CONFIG
         )
+        # Wrap the embedding client in the adapter
         return _LangchainEmbeddingClientAdapter(client)
 
     def train(  # type: ignore[override]
@@ -273,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
@@ -292,14 +318,6 @@ class EnterpriseSearchPolicy(Policy):
                 f"required environment variables. Error: {e}"
             )
 
-        # validate llm configuration
-        try_instantiate_llm_client(
-            self.config.get(LLM_CONFIG_KEY),
-            DEFAULT_LLM_CONFIG,
-            "enterprise_search_policy.train",
-            "EnterpriseSearchPolicy",
-        )
-
         if store_type == DEFAULT_VECTOR_STORE_TYPE:
             logger.info("enterprise_search_policy.train.faiss")
             with self._model_storage.write_to(self._resource) as path:
@@ -318,9 +336,13 @@ class EnterpriseSearchPolicy(Policy):
             embeddings_type=self.embeddings_config.get(PROVIDER_CONFIG_KEY),
             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_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_ID_CONFIG_KEY),
             citation_enabled=self.citation_enabled,
         )
         self.persist()
@@ -332,6 +354,9 @@ class EnterpriseSearchPolicy(Policy):
             rasa.shared.utils.io.write_text_file(
                 self.prompt_template, path / ENTERPRISE_SEARCH_PROMPT_FILE_NAME
             )
+            rasa.shared.utils.io.dump_obj_as_json_to_file(
+                path / ENTERPRISE_SEARCH_CONFIG_FILE_NAME, self.config
+            )
 
     def _prepare_slots_for_template(
         self, tracker: DialogueStateTracker
@@ -512,9 +537,13 @@ class EnterpriseSearchPolicy(Policy):
             embeddings_type=self.embeddings_config.get(PROVIDER_CONFIG_KEY),
             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_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_ID_CONFIG_KEY),
             citation_enabled=self.citation_enabled,
         )
         return self._create_prediction(
@@ -662,25 +691,28 @@ class EnterpriseSearchPolicy(Policy):
         execution_context: ExecutionContext,
         **kwargs: Any,
     ) -> "EnterpriseSearchPolicy":
-        try_instantiate_llm_client(
-            config.get(LLM_CONFIG_KEY),
-            DEFAULT_LLM_CONFIG,
-            "enterprise_search_policy.load",
-            EnterpriseSearchPolicy.__name__,
-        )
-        try_instantiate_embedder(
-            config.get(EMBEDDINGS_CONFIG_KEY),
-            DEFAULT_EMBEDDINGS_CONFIG,
-            "enterprise_search_policy.load",
-            EnterpriseSearchPolicy.__name__,
-        )
         """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
+        try:
+            with model_storage.read_from(resource) as path:
+                prompt_template = rasa.shared.utils.io.read_file(
+                    path / ENTERPRISE_SEARCH_PROMPT_FILE_NAME
+                )
+        except (FileNotFoundError, FileIOException) as e:
+            logger.warning(
+                "enterprise_search_policy.load.failed", error=e, resource=resource.name
+            )
+
         store_type = config.get(VECTOR_STORE_PROPERTY, {}).get(
             VECTOR_STORE_TYPE_PROPERTY
         )
 
         embeddings = cls._create_plain_embedder(config)
+
         logger.info("enterprise_search_policy.load", config=config)
         if store_type == DEFAULT_VECTOR_STORE_TYPE:
             # if a vector store is not specified,
@@ -698,16 +730,7 @@ class EnterpriseSearchPolicy(Policy):
                 config_type=store_type,
                 embeddings=embeddings,
             )  # type: ignore
-        try:
-            with model_storage.read_from(resource) as path:
-                prompt_template = rasa.shared.utils.io.read_file(
-                    path / ENTERPRISE_SEARCH_PROMPT_FILE_NAME
-                )
 
-        except (FileNotFoundError, FileNotFoundError) as e:
-            logger.warning(
-                "enterprise_search_policy.load.failed", error=e, resource=resource.name
-            )
         return cls(
             config,
             model_storage,
@@ -748,14 +771,23 @@ class EnterpriseSearchPolicy(Policy):
 
     @classmethod
     def fingerprint_addon(cls, config: Dict[str, Any]) -> Optional[str]:
-        """Add a fingerprint of the knowledge base and prompt template for the graph."""
+        """Add a fingerprint of enterprise search policy for the graph."""
         local_knowledge_data = cls._get_local_knowledge_data(config)
 
         prompt_template = get_prompt_template(
             config.get(PROMPT_CONFIG_KEY),
             DEFAULT_ENTERPRISE_SEARCH_PROMPT_TEMPLATE,
         )
-        return deep_container_fingerprint([prompt_template, local_knowledge_data])
+
+        llm_config = resolve_model_client_config(
+            config.get(LLM_CONFIG_KEY), EnterpriseSearchPolicy.__name__
+        )
+        embedding_config = resolve_model_client_config(
+            config.get(EMBEDDINGS_CONFIG_KEY), EnterpriseSearchPolicy.__name__
+        )
+        return deep_container_fingerprint(
+            [prompt_template, local_knowledge_data, llm_config, embedding_config]
+        )
 
     @staticmethod
     def post_process_citations(llm_answer: str) -> str:
@@ -847,3 +879,27 @@ class EnterpriseSearchPolicy(Policy):
         joined_sources = "\n".join(new_sources)
 
         return joined_answer + joined_sources
+
+    @classmethod
+    def _perform_health_checks(
+        cls, config: Dict[Text, Any], log_source_method: str
+    ) -> None:
+        # 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,
+            log_source_method,
+            EnterpriseSearchPolicy.__name__,
+        )
+
+        # Perform health check of the embeddings client config
+        embeddings_config = resolve_model_client_config(
+            config.get(EMBEDDINGS_CONFIG_KEY, {})
+        )
+        cls.perform_embeddings_health_check(
+            embeddings_config,
+            DEFAULT_EMBEDDINGS_CONFIG,
+            log_source_method,
+            EnterpriseSearchPolicy.__name__,
+        )

rasa/core/policies/flows/flow_executor.py

@@ -330,24 +330,27 @@ def reset_scoped_slots(
     events: List[Event] = []
 
     not_resettable_slot_names = set()
+    flow_persistable_slots = current_flow.persisted_slots
 
     for step in current_flow.steps_with_calls_resolved:
         if isinstance(step, CollectInformationFlowStep):
             # reset all slots scoped to the flow
-            if step.reset_after_flow_ends:
-                _reset_slot(step.collect, tracker)
+            slot_name = step.collect
+            if step.reset_after_flow_ends and slot_name not in flow_persistable_slots:
+                _reset_slot(slot_name, tracker)
             else:
-                not_resettable_slot_names.add(step.collect)
+                not_resettable_slot_names.add(slot_name)
 
     # slots set by the set slots step should be reset after the flow ends
     # unless they are also used in a collect step where `reset_after_flow_ends`
-    # is set to `False`
+    # is set to `False` or set in the `persisted_slots` list.
     resettable_set_slots = [
         slot["key"]
         for step in current_flow.steps_with_calls_resolved
         if isinstance(step, SetSlotsFlowStep)
         for slot in step.slots
         if slot["key"] not in not_resettable_slot_names
+        and slot["key"] not in flow_persistable_slots
     ]
 
     for name in resettable_set_slots:
@@ -484,8 +487,7 @@ def validate_collect_step(
     step: CollectInformationFlowStep,
     stack: DialogueStack,
     available_actions: List[str],
-    slots: Dict[str, Slot],
-    flow_name: str,
+    slots: Dict[Text, Slot],
 ) -> bool:
     """Validate that a collect step can be executed.
 
@@ -508,12 +510,12 @@ def validate_collect_step(
         slot_name=step.collect,
     )
 
-    cancel_flow_and_push_internal_error(stack, flow_name)
+    cancel_flow_and_push_internal_error(stack)
 
     return False
 
 
-def cancel_flow_and_push_internal_error(stack: DialogueStack, flow_name: str) -> None:
+def cancel_flow_and_push_internal_error(stack: DialogueStack) -> None:
     """Cancel the top user flow and push the internal error pattern."""
     top_frame = stack.top()
 
@@ -525,7 +527,7 @@ def cancel_flow_and_push_internal_error(stack: DialogueStack, flow_name: str) ->
         canceled_frames = CancelFlowCommand.select_canceled_frames(stack)
         stack.push(
             CancelPatternFlowStackFrame(
-                canceled_name=flow_name,
+                canceled_name=top_frame.flow_id,
                 canceled_frames=canceled_frames,
             )
         )
@@ -537,7 +539,6 @@ def validate_custom_slot_mappings(
     stack: DialogueStack,
     tracker: DialogueStateTracker,
     available_actions: List[str],
-    flow_name: str,
 ) -> bool:
     """Validate a slot with custom mappings.
 
@@ -558,7 +559,7 @@ def validate_custom_slot_mappings(
                     action=step.collect_action,
                     collect=step.collect,
                 )
-                cancel_flow_and_push_internal_error(stack, flow_name)
+                cancel_flow_and_push_internal_error(stack)
                 return False
 
     return True
@@ -598,12 +599,7 @@ def run_step(
 
     if isinstance(step, CollectInformationFlowStep):
         return _run_collect_information_step(
-            available_actions,
-            initial_events,
-            stack,
-            step,
-            tracker,
-            flow.readable_name(),
+            available_actions, initial_events, stack, step, tracker
         )
 
     elif isinstance(step, ActionFlowStep):
@@ -723,18 +719,15 @@ def _run_collect_information_step(
     stack: DialogueStack,
     step: CollectInformationFlowStep,
     tracker: DialogueStateTracker,
-    flow_name: str,
 ) -> FlowStepResult:
-    is_step_valid = validate_collect_step(
-        step, stack, available_actions, tracker.slots, flow_name
-    )
+    is_step_valid = validate_collect_step(step, stack, available_actions, tracker.slots)
 
     if not is_step_valid:
         # if we return any other FlowStepResult, the assistant will stay silent
         # instead of triggering the internal error pattern
         return ContinueFlowWithNextStep(events=initial_events)
     is_mapping_valid = validate_custom_slot_mappings(
-        step, stack, tracker, available_actions, flow_name
+        step, stack, tracker, available_actions
     )
 
     if not is_mapping_valid: