rasa-pro 3.14.0a15__py3-none-any.whl → 3.14.0a16__py3-none-any.whl

rasa/core/channels/inspector/src/components/LatencyDisplay.tsx

@@ -2,40 +2,47 @@ import {
   Box,
   Flex,
   FlexProps,
-  Text,
-  useColorModeValue,
-  Tooltip,
   Table,
   Tbody,
-  Tr,
   Td,
+  Text,
+  Tooltip,
+  Tr,
+  useColorModeValue,
 } from '@chakra-ui/react'
 import { useOurTheme } from '../theme'
-import { LatencyData } from '../types'
+import {
+  isRasaLatency,
+  isVoiceLatency,
+  RasaLatency,
+  VoiceLatency,
+} from '../types'
 
-interface Props extends FlexProps {
-  latency: LatencyData
+interface MinimalDisplayProps extends FlexProps {
+  latency: RasaLatency
 }
 
 /**
  * Simple latency display for text-only conversations.
  * Shows a single response time value.
  */
-const MinimalDisplay = ({ latency, sx, ...props }: Props) => {
+const MinimalDisplay = ({ latency, sx, ...props }: MinimalDisplayProps) => {
   const containerSx = {
     ...sx,
     display: 'flex',
     alignItems: 'center',
   }
-  
+
   const getLatencyColor = (latency: number) => {
     if (latency < 1500) return 'green.500'
     if (latency < 2500) return 'orange.500'
     return 'red.500'
   }
-  
-  const value = Math.round(latency.rasa_processing_latency_ms || 0);
-  const color = getLatencyColor(value);
+
+  const value = Math.round(
+    (latency.command_processor || 0) + (latency.prediction_loop || 0),
+  )
+  const color = getLatencyColor(value)
 
   return (
     <Flex sx={containerSx} {...props}>
@@ -49,11 +56,15 @@ const MinimalDisplay = ({ latency, sx, ...props }: Props) => {
   )
 }
 
+interface WaterfallDisplayProps extends FlexProps {
+  latency: VoiceLatency
+}
+
 /**
  * Detailed latency waterfall chart for voice conversations.
  * Displays processing times for ASR, Rasa, and TTS components.
  */
-const WaterfallDisplay = ({ latency, sx, ...props }: Props) => {
+const WaterfallDisplay = ({ latency, sx, ...props }: WaterfallDisplayProps) => {
   const { rasaSpace } = useOurTheme()
 
   const containerSx = {
@@ -61,13 +72,13 @@ const WaterfallDisplay = ({ latency, sx, ...props }: Props) => {
     flexDirection: 'column',
     gap: rasaSpace[1],
   }
-  
+
   const headerSx = {
     fontSize: 'sm',
     fontWeight: 'bold',
     color: useColorModeValue('gray.700', 'gray.300'),
   }
-  
+
   const waterfallBarSx = {
     height: '24px',
     borderRadius: '4px',
@@ -75,10 +86,10 @@ const WaterfallDisplay = ({ latency, sx, ...props }: Props) => {
     border: '1px solid',
     borderColor: useColorModeValue('gray.200', 'gray.600'),
   }
-  
+
   const legendTableSx = {
     size: 'sm',
-    mt: rasaSpace[0.5]
+    mt: rasaSpace[0.5],
   }
 
   const getLatencyColor = (type: string) => {
@@ -95,8 +106,10 @@ const WaterfallDisplay = ({ latency, sx, ...props }: Props) => {
     const descriptions: { [key: string]: string } = {
       asr: 'Time from the first Partial Transcript event to the Final Transcript event from Speech Recognition. It also includes the time taken by the user to speak.',
       rasa: 'Time taken by Rasa to process the text from the Final Transcript event from Speech Recognition until a text response is generated.',
-      tts_first: 'Time between the request sent to Text-to-Speech processing and the first byte of audio received by Rasa.',
-      tts_complete: 'Time taken by Text-to-Speech to complete audio generation. It depends on the length of the text and could overlap with the Bot speaking time.'
+      tts_first:
+        'Time between the request sent to Text-to-Speech processing and the first byte of audio received by Rasa.',
+      tts_complete:
+        'Time taken by Text-to-Speech to complete audio generation. It depends on the length of the text and could overlap with the Bot speaking time.',
     }
     return descriptions[type] || ''
   }
@@ -151,10 +164,10 @@ const WaterfallDisplay = ({ latency, sx, ...props }: Props) => {
   )
 
   // Calculate total latency for title (Rasa + TTS First Byte)
-  const totalDisplayLatency = 
-    (latency.rasa_processing_latency_ms || 0) + 
-    (latency.tts_first_byte_latency_ms || 0);
-    
+  const totalDisplayLatency =
+    (latency.rasa_processing_latency_ms || 0) +
+    (latency.tts_first_byte_latency_ms || 0)
+
   return (
     <Flex sx={containerSx} {...props}>
       <Text sx={headerSx}>
@@ -241,28 +254,43 @@ const WaterfallDisplay = ({ latency, sx, ...props }: Props) => {
   )
 }
 
+interface LatencyDisplayProps extends FlexProps {
+  voiceLatency?: any
+  rasaLatency?: any
+}
+
 /**
  * Displays processing latency information for the conversation.
- * Shows either a detailed waterfall chart for voice conversations or 
+ * Shows either a detailed waterfall chart for voice conversations or
  * a simpler display for text-only conversations.
  */
 export const LatencyDisplay = ({
   sx,
-  latency,
+  voiceLatency,
+  rasaLatency,
   ...props
-}: Props) => {
-  if (!latency) {
-    console.warn('Latency data is not available')
-    return null
+}: LatencyDisplayProps) => {
+  // We give priority to voice latency as most comprehensive.
+  if (isVoiceLatency(voiceLatency)) {
+    return <WaterfallDisplay latency={voiceLatency} sx={sx} {...props} />
   }
 
-  // Show waterfall if voice metrics are available, otherwise show minimal display
-  const isVoiceMetricsAvailable =
-    latency.asr_latency_ms && latency.tts_complete_latency_ms
+  // If voice latency is not available, we use rasa latency.
+  if (isRasaLatency(rasaLatency)) {
+    return <MinimalDisplay latency={rasaLatency} sx={sx} {...props} />
+  }
 
-  if (isVoiceMetricsAvailable) {
-    return <WaterfallDisplay latency={latency} sx={sx} {...props} />
+  // If rasa latency is not available, we are either waiting for first data from voice stream or the conversation is not started yet.
+  if (!rasaLatency) {
+    return <Text>Rasa latency data is not available yet</Text>
   }
 
-  return <MinimalDisplay latency={latency} sx={sx} {...props} />
+  // If rasa latency is available but did not pass type guard, we are in an unexpected state
+  // and we should log an error but gracefully fallback to showing no data.
+  console.warn(
+    `Latency data has unexpected format. Raw data of rasaLatency: ${rasaLatency},
+    raw data of voice latency: ${voiceLatency}`,
+  )
+
+  return <Text>Latency data is not available yet</Text>
 }

rasa/core/channels/inspector/src/types.ts

@@ -17,7 +17,8 @@ export interface Event {
   timestamp: string
   update?: string
   parse_data?: { commands: Command[] }
-  metadata?: { utter_action?: string }
+  metadata?: { utter_action?: string; execution_times?: RasaLatency }
+  name?: string
 }
 
 export interface Command {
@@ -42,11 +43,16 @@ export interface Stack {
   ended: boolean
 }
 
-export interface LatencyData {
-  rasa_processing_latency_ms?: number
-  asr_latency_ms?: number
-  tts_first_byte_latency_ms?: number
-  tts_complete_latency_ms?: number
+export interface RasaLatency {
+  command_processor: number
+  prediction_loop: number
+}
+
+export interface VoiceLatency {
+  rasa_processing_latency_ms: number
+  asr_latency_ms: number
+  tts_first_byte_latency_ms: number
+  tts_complete_latency_ms: number
 }
 
 export interface Tracker {
@@ -54,7 +60,6 @@ export interface Tracker {
   slots: { [key: string]: unknown }
   events: Event[]
   stack: Stack[]
-  latency?: LatencyData
 }
 
 export interface Flow {
@@ -96,3 +101,23 @@ interface Step {
   utter: string
   set_slots?: unknown
 }
+
+export function isRasaLatency(obj: any): obj is RasaLatency {
+  return (
+    obj !== null &&
+    typeof obj === 'object' &&
+    typeof obj.command_processor === 'number' &&
+    typeof obj.prediction_loop === 'number'
+  )
+}
+
+export function isVoiceLatency(obj: any): obj is VoiceLatency {
+  return (
+    obj !== null &&
+    typeof obj === 'object' &&
+    typeof obj.rasa_processing_latency_ms === 'number' &&
+    typeof obj.asr_latency_ms === 'number' &&
+    typeof obj.tts_first_byte_latency_ms === 'number' &&
+    typeof obj.tts_complete_latency_ms === 'number'
+  )
+}

rasa/core/channels/studio_chat.py

@@ -4,7 +4,6 @@ import asyncio
 import audioop
 import base64
 import json
-import time
 import uuid
 from functools import partial
 from typing import (
@@ -53,9 +52,7 @@ if TYPE_CHECKING:
 structlogger = structlog.get_logger()
 
 
-def tracker_as_dump(
-    tracker: "DialogueStateTracker", latency: Optional[float] = None
-) -> Dict[str, Any]:
+def tracker_as_dump(tracker: "DialogueStateTracker") -> Dict[str, Any]:
     """Create a dump of the tracker state."""
     from rasa.shared.core.trackers import get_trackers_for_conversation_sessions
 
@@ -66,13 +63,7 @@ def tracker_as_dump(
     else:
         last_tracker = multiple_tracker_sessions[-1]
 
-    # TODO: this is a bug: the bridge converts this back to json, but it
-    # should be json in the first place
     state = last_tracker.current_state(EventVerbosity.AFTER_RESTART)
-
-    if latency is not None:
-        state["latency"] = {"rasa_processing_latency_ms": latency}
-
     return state
 
 
@@ -227,32 +218,16 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
     def _register_tracker_update_hook(self) -> None:
         plugin_manager().register(StudioTrackerUpdatePlugin(self))
 
-    async def on_tracker_updated(
-        self, tracker: "DialogueStateTracker", latency: Optional[float] = None
-    ) -> None:
+    async def on_tracker_updated(self, tracker: "DialogueStateTracker") -> None:
         """Triggers a tracker update notification after a change to the tracker."""
-        await self.publish_tracker_update(
-            tracker.sender_id, tracker_as_dump(tracker, latency)
-        )
+        await self.publish_tracker_update(tracker.sender_id, tracker_as_dump(tracker))
 
-    async def publish_tracker_update(self, sender_id: str, tracker_dump: str) -> None:
+    async def publish_tracker_update(
+        self, sender_id: str, tracker_dump: Dict[str, Any]
+    ) -> None:
         """Publishes a tracker update notification to the websocket."""
         await self.emit("tracker", tracker_dump, room=sender_id)
 
-    def _record_turn_start_time(self, sender_id: Text) -> None:
-        """Records the start time of a new turn."""
-        self._turn_start_times[sender_id] = time.time()
-
-    def _get_latency(self, sender_id: Text) -> Optional[float]:
-        """Returns the latency of the current turn in milliseconds."""
-        if sender_id not in self._turn_start_times:
-            return None
-
-        latency = (time.time() - self._turn_start_times[sender_id]) * 1000
-        # The turn is over, so we can remove the start time
-        del self._turn_start_times[sender_id]
-        return latency
-
     async def on_message_proxy(
         self,
         on_new_message: Callable[[UserMessage], Awaitable[Any]],
@@ -262,7 +237,6 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
 
         Triggers a tracker update notification after processing the message.
         """
-        self._record_turn_start_time(message.sender_id)
         try:
             await on_new_message(message)
         except Exception as e:
@@ -288,8 +262,7 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
             structlogger.error("studio_chat.on_message_proxy.tracker_not_found")
             return
 
-        latency = self._get_latency(message.sender_id)
-        await self.on_tracker_updated(tracker, latency)
+        await self.on_tracker_updated(tracker)
 
     async def emit_error(self, message: str, room: str, e: Exception) -> None:
         await self.emit(
@@ -389,10 +362,10 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
                 call_state.is_bot_speaking = True
         return ContinueConversationAction()
 
-    def _create_output_channel(
+    def create_output_channel(
         self, voice_websocket: "Websocket", tts_engine: TTSEngine
     ) -> VoiceOutputChannel:
-        """Create a voice output channel."""
+        """Create a voice output channel. This is used by VoiceInputChannel."""
         return StudioVoiceOutputChannel(
             voice_websocket,
             tts_engine,
@@ -469,8 +442,9 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
     def blueprint(
         self, on_new_message: Callable[[UserMessage], Awaitable[Any]]
     ) -> SocketBlueprint:
-        proxied_on_message = partial(self.on_message_proxy, on_new_message)
-        socket_blueprint = super().blueprint(proxied_on_message)
+        socket_blueprint = super().blueprint(
+            partial(self.on_message_proxy, on_new_message)
+        )
 
         if not self.sio_server:
             structlogger.error("studio_chat.blueprint.sio_not_initialized")
@@ -506,7 +480,7 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
 
             # start a voice session if requested
             if data and data.get("is_voice", False):
-                self._start_voice_session(data["session_id"], sid, proxied_on_message)
+                self._start_voice_session(data["session_id"], sid, on_new_message)
 
         @self.sio_server.on(self.user_message_evt, namespace=self.namespace)
         async def handle_message(sid: Text, data: Dict) -> None:
@@ -521,7 +495,7 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
 
             try:
                 # Handle text messages
-                await self.handle_user_message(sid, data, proxied_on_message)
+                await self.handle_user_message(sid, data, on_new_message)
             except Exception as e:
                 structlogger.exception(
                     "studio_chat.sio.handle_message.error",

rasa/core/constants.py

@@ -112,3 +112,9 @@ ACTIVE_FLOW_METADATA_KEY = "active_flow"
 STEP_ID_METADATA_KEY = "step_id"
 KEY_IS_CALM_SYSTEM = "is_calm_system"
 KEY_IS_COEXISTENCE_ASSISTANT = "is_coexistence_assistant"
+
+IAM_CLOUD_PROVIDER_ENV_VAR_NAME = "IAM_CLOUD_PROVIDER"
+SQL_TRACKER_STORE_SSL_MODE_ENV_VAR_NAME = "SQL_TRACKER_STORE_SSL_MODE"
+SQL_TRACKER_STORE_SSL_ROOT_CERTIFICATE_ENV_VAR_NAME = (
+    "SQL_TRACKER_STORE_SSL_ROOT_CERTIFICATE"
+)

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/processor.py

@@ -72,6 +72,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,
@@ -207,6 +208,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:
@@ -1154,6 +1156,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):
@@ -1403,6 +1406,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,
@@ -1448,6 +1479,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