rasa-pro 3.14.0.dev7__py3-none-any.whl → 3.14.0.dev8__py3-none-any.whl

rasa/core/iam_credentials_providers/aws_iam_credentials_providers.py

@@ -0,0 +1,66 @@
+from typing import Optional
+
+import boto3
+import structlog
+from botocore.exceptions import BotoCoreError
+
+from rasa.core.iam_credentials_providers.credentials_provider_protocol import (
+    IAMCredentialsProvider,
+    IAMCredentialsProviderInput,
+    SupportedServiceType,
+    TemporaryCredentials,
+)
+
+structlogger = structlog.get_logger(__name__)
+
+
+class AWSRDSIAMCredentialsProvider(IAMCredentialsProvider):
+    """Generates temporary credentials for AWS RDS using IAM roles."""
+
+    def __init__(self, username: str, host: str, port: int):
+        """Initializes the provider."""
+        self.username = username
+        self.host = host
+        self.port = port
+
+    def get_credentials(self) -> TemporaryCredentials:
+        """Generates temporary credentials for AWS RDS."""
+        structlogger.debug(
+            "rasa.core.aws_rds_iam_credentials_provider.get_credentials",
+            event_info="IAM authentication for AWS RDS enabled. "
+            "Generating temporary auth token...",
+        )
+
+        try:
+            client = boto3.client("rds")
+            auth_token = client.generate_db_auth_token(
+                DBHostname=self.host,
+                Port=self.port,
+                DBUsername=self.username,
+            )
+            structlogger.info(
+                "rasa.core.aws_rds_iam_credentials_provider.generated_credentials",
+                event_info="Successfully generated temporary auth token for AWS RDS.",
+            )
+            return TemporaryCredentials(auth_token=auth_token)
+        except (BotoCoreError, ValueError) as exc:
+            structlogger.error(
+                "rasa.core.aws_rds_iam_credentials_provider.error_generating_credentials",
+                event_info="Failed to generate temporary auth token for AWS RDS.",
+                error=str(exc),
+            )
+            return TemporaryCredentials(auth_token=None)
+
+
+def create_aws_iam_credentials_provider(
+    provider_input: "IAMCredentialsProviderInput",
+) -> Optional["IAMCredentialsProvider"]:
+    """Factory function to create an AWS IAM credentials provider."""
+    if provider_input.service_name == SupportedServiceType.TRACKER_STORE:
+        return AWSRDSIAMCredentialsProvider(
+            username=provider_input.username,
+            host=provider_input.host,
+            port=provider_input.port,
+        )
+
+    return None

rasa/core/iam_credentials_providers/credentials_provider_protocol.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import os
+from enum import Enum
+from typing import Optional, Protocol, runtime_checkable
+
+import structlog
+from pydantic import BaseModel
+
+from rasa.core.constants import IAM_CLOUD_PROVIDER_ENV_VAR_NAME
+
+structlogger = structlog.get_logger(__name__)
+
+
+class TemporaryCredentials(BaseModel):
+    """Dataclass storing temporary credentials."""
+
+    auth_token: Optional[str] = None
+    expiration: Optional[int] = None
+    username: Optional[str] = None
+    presigned_url: Optional[str] = None
+
+
+@runtime_checkable
+class IAMCredentialsProvider(Protocol):
+    """Interface for generating temporary credentials using IAM roles."""
+
+    def get_credentials(self) -> TemporaryCredentials:
+        """Generates temporary credentials using IAM roles."""
+        ...
+
+
+class IAMCredentialsProviderType(Enum):
+    """Enum for supported IAM credentials provider types."""
+
+    AWS = "aws"
+
+
+class SupportedServiceType(Enum):
+    """Enum for supported services using IAM credentials providers."""
+
+    TRACKER_STORE = "tracker_store"
+    EVENT_BROKER = "event_broker"
+    LOCK_STORE = "lock_store"
+
+
+class IAMCredentialsProviderInput(BaseModel):
+    """Input data for creating an IAM credentials provider."""
+
+    service_name: SupportedServiceType
+    username: Optional[str] = None
+    host: Optional[str] = None
+    port: Optional[int] = None
+
+
+def create_iam_credentials_provider(
+    provider_input: IAMCredentialsProviderInput,
+) -> Optional[IAMCredentialsProvider]:
+    """Factory function to create an IAM credentials provider.
+
+    Args:
+        provider_input: Input data for creating an IAM credentials provider.
+
+    Returns:
+        An instance of the specified IAM credentials provider or
+        None if the type is unsupported.
+    """
+    iam_cloud_provider = os.getenv(IAM_CLOUD_PROVIDER_ENV_VAR_NAME)
+
+    if iam_cloud_provider is None:
+        return None
+
+    try:
+        provider_type = IAMCredentialsProviderType(iam_cloud_provider.lower())
+    except ValueError:
+        structlogger.warning(
+            "rasa.core.iam_credentials_provider.create_iam_credentials_provider.unsupported_provider",
+            event_info=f"Unsupported IAM cloud provider: {iam_cloud_provider}",
+        )
+        return None
+
+    if provider_type == IAMCredentialsProviderType.AWS:
+        from rasa.core.iam_credentials_providers.aws_iam_credentials_providers import (
+            create_aws_iam_credentials_provider,
+        )
+
+        return create_aws_iam_credentials_provider(provider_input)
+
+    return None

rasa/core/policies/enterprise_search_policy.py

@@ -783,7 +783,7 @@ class EnterpriseSearchPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Po
         if not os.path.exists(docs_folder) or not os.path.isdir(docs_folder):
             error_message = (
                 f"Document source directory does not exist or is not a "
-                f"directory: '{docs_folder}'. "
+                f"directory: '{os.path.abspath(docs_folder)}'. "
                 "Please specify a valid path to the documents source directory in the "
                 "vector_store configuration."
             )
@@ -1132,8 +1132,7 @@ class EnterpriseSearchPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Po
         embeddings_config: Dict[Text, Any],
         log_source_method: str,
     ) -> None:
-        """
-        Perform the health checks using resolved LLM and embeddings configurations.
+        """Perform the health checks using resolved LLM and embeddings configurations.
         Resolved means the configuration is either:
         - A reference to a model group that has already been expanded into
           its corresponding configuration using the information from
@@ -1162,8 +1161,7 @@ class EnterpriseSearchPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Po
 
     @classmethod
     def get_system_default_prompt_based_on_config(cls, config: Dict[str, Any]) -> str:
-        """
-        Resolves the default prompt template for Enterprise Search Policy based on
+        """Resolves the default prompt template for Enterprise Search Policy based on
         the component's configuration.
 
         - The old prompt is selected when both citation and relevancy check are either
@@ -1194,8 +1192,7 @@ class EnterpriseSearchPolicy(LLMHealthCheckMixin, EmbeddingsHealthCheckMixin, Po
         relevancy_check_enabled: bool,
         citation_enabled: bool,
     ) -> str:
-        """
-        Returns the appropriate default prompt template based on the feature flags.
+        """Returns the appropriate default prompt template based on the feature flags.
 
         The selection follows this priority:
         1. If relevancy check is enabled, return the prompt that includes both relevancy

rasa/core/policies/flows/flow_executor.py

@@ -13,7 +13,7 @@ from rasa.agents.constants import (
     A2A_AGENT_CONTEXT_ID_KEY,
     AGENT_METADATA_AGENT_RESPONSE_KEY,
     AGENT_METADATA_EXIT_IF_KEY,
-    AGENT_METADATA_TOOL_RESULTS_KEY,
+    AGENT_METADATA_STRUCTURED_RESULTS_KEY,
     MAX_AGENT_RETRY_DELAY_SECONDS,
 )
 from rasa.agents.core.types import AgentStatus, ProtocolType
@@ -176,8 +176,8 @@ def select_next_step_id(
     # if the current step is a call step to an agent, and we already have an
     # AgentStackFrame on top of the stack, we need to return the current
     # step id again in order to loop back to the agent.
-    agent_stack_frame = tracker.stack.top()
-    if agent_stack_frame and isinstance(agent_stack_frame, AgentStackFrame):
+    top_stack_frame = tracker.stack.top()
+    if top_stack_frame and isinstance(top_stack_frame, AgentStackFrame):
         return current.id
 
     next_step = current.next
@@ -810,7 +810,14 @@ def _run_action_step(
             # do not log about non-existing validation actions of collect steps
             utter_action_name = render_template_variables("{{context.utter}}", context)
             if utter_action_name not in available_actions:
-                structlogger.warning("flow.step.run.action.unknown", action=action_name)
+                structlogger.warning(
+                    "flow.step.run.action.unknown",
+                    action=action_name,
+                    event_info=(
+                        f"The action '{action_name}' is not defined in the domain but "
+                        f"getting triggered by the flow '{step.flow_id}'."
+                    ),
+                )
         return ContinueFlowWithNextStep(events=initial_events)
 
 
@@ -1020,7 +1027,9 @@ async def run_agent(
         output.metadata[AGENT_METADATA_AGENT_RESPONSE_KEY] = (
             output.response_message or ""
         )
-        output.metadata[AGENT_METADATA_TOOL_RESULTS_KEY] = output.tool_results or []
+        output.metadata[AGENT_METADATA_STRUCTURED_RESULTS_KEY] = (
+            output.structured_results or []
+        )
         _update_agent_events(final_events, output.metadata)
 
         top_stack_frame = stack.top()

rasa/core/policies/ted_policy.py

@@ -943,14 +943,16 @@ class TEDPolicy(Policy):
 
         with self._model_storage.write_to(self._resource) as model_path:
             model_filename = self._metadata_filename()
-            tf_model_file = model_path / f"{model_filename}.tf_model"
+            tf_model_file = model_path / f"{model_filename}.weights.h5"
 
             rasa.shared.utils.io.create_directory_for_file(tf_model_file)
 
             self.featurizer.persist(model_path)
 
             if self.config[CHECKPOINT_MODEL] and self.tmp_checkpoint_dir:
-                self.model.load_weights(self.tmp_checkpoint_dir / "checkpoint.tf_model")
+                self.model.load_weights(
+                    self.tmp_checkpoint_dir / "checkpoint.weights.h5"
+                )
                 # Save an empty file to flag that this model has been
                 # produced using checkpointing
                 checkpoint_marker = model_path / f"{model_filename}.from_checkpoint.pkl"
@@ -1009,7 +1011,7 @@ class TEDPolicy(Policy):
         Args:
             model_path: Path where model is to be persisted.
         """
-        tf_model_file = model_path / f"{cls._metadata_filename()}.tf_model"
+        tf_model_file = model_path / f"{cls._metadata_filename()}.weights.h5"
 
         # load data example
         loaded_data = deserialize_nested_feature_arrays(
@@ -1463,7 +1465,7 @@ class TED(TransformerRasaModel):
 
         dialogue_transformed, attention_weights = self._tf_layers[
             f"transformer.{DIALOGUE}"
-        ](dialogue_in, 1 - mask, self._training)
+        ](dialogue_in, 1 - mask, training=self._training)
         dialogue_transformed = tf.nn.gelu(dialogue_transformed)
 
         if self.max_history_featurizer_is_used:
@@ -1708,7 +1710,7 @@ class TED(TransformerRasaModel):
 
         if attribute in SENTENCE_FEATURES_TO_ENCODE + LABEL_FEATURES_TO_ENCODE:
             attribute_features = self._tf_layers[f"encoding_layer.{attribute}"](
-                attribute_features, self._training
+                attribute_features, training=self._training
             )
 
         # attribute features have shape

rasa/core/processor.py

@@ -73,6 +73,7 @@ from rasa.shared.core.constants import (
     ACTION_CORRECT_FLOW_SLOT,
     ACTION_EXTRACT_SLOTS,
     ACTION_LISTEN_NAME,
+    ACTION_METADATA_EXECUTION_TIME,
     ACTION_SESSION_START_NAME,
     FOLLOWUP_ACTION,
     SESSION_START_METADATA_SLOT,
@@ -208,6 +209,7 @@ class MessageProcessor:
     ) -> Optional[List[Dict[Text, Any]]]:
         """Handle a single message with this processor."""
         # preprocess message if necessary
+        self.time_turn_start = time.time()
         tracker = await self.log_message(message, should_save_tracker=False)
 
         if self.model_metadata.training_type == TrainingType.NLU:
@@ -1155,6 +1157,7 @@ class MessageProcessor:
         should_predict_another_action = True
 
         tracker = await self.run_command_processor(tracker)
+        self.time_command_processor = time.time()
 
         # action loop. predicts actions until we hit action listen
         while should_predict_another_action and self._should_handle_message(tracker):
@@ -1408,6 +1411,34 @@ class MessageProcessor:
         plugin_manager().hook.after_action_executed(tracker=tracker)
         return self.should_predict_another_action(action.name())
 
+    def _add_metadata_if_action_listen(
+        self, action: Action, prediction: PolicyPrediction
+    ) -> None:
+        """Adds execution times to the ActionExecuted event metadata."""
+        if not hasattr(self, "time_turn_start"):
+            return
+
+        if not hasattr(self, "time_command_processor"):
+            return
+
+        if not action.name() == ACTION_LISTEN_NAME:
+            return
+
+        if prediction.action_metadata is None:
+            prediction.action_metadata = {}
+
+        # calculate execution times
+        execution_time_prediction_loop = (
+            time.time() - self.time_command_processor
+        ) * 1000
+        execution_time_command_processor = (
+            self.time_command_processor - self.time_turn_start
+        ) * 1000
+        prediction.action_metadata[ACTION_METADATA_EXECUTION_TIME] = {
+            "command_processor": execution_time_command_processor,
+            "prediction_loop": execution_time_prediction_loop,
+        }
+
     def _log_action_and_events_on_tracker(
         self,
         tracker: DialogueStateTracker,
@@ -1453,6 +1484,7 @@ class MessageProcessor:
         tracker.update_with_events(prediction.events)
 
         # log the action and its produced events
+        self._add_metadata_if_action_listen(action, prediction)
         tracker.update(
             action.event_for_successful_execution(
                 prediction, was_successful, error_message