rasa-pro 3.14.0.dev2__py3-none-any.whl → 3.14.0.dev4__py3-none-any.whl

rasa/core/channels/studio_chat.py

@@ -100,12 +100,12 @@ class StudioTrackerUpdatePlugin:
         """Remove tasks that have already completed."""
         self.tasks = [task for task in self.tasks if not task.done()]
 
-    @hookimpl  # type: ignore[misc]
+    @hookimpl
     def after_new_user_message(self, tracker: "DialogueStateTracker") -> None:
         """Triggers a tracker update notification after a new user message."""
         self.handle_tracker_update(tracker)
 
-    @hookimpl  # type: ignore[misc]
+    @hookimpl
     def after_action_executed(self, tracker: "DialogueStateTracker") -> None:
         """Triggers a tracker update notification after an action is executed."""
         self.handle_tracker_update(tracker)
@@ -125,7 +125,7 @@ class StudioTrackerUpdatePlugin:
         self.tasks.append(task)
         self._cleanup_tasks()
 
-    @hookimpl  # type: ignore[misc]
+    @hookimpl
     def after_server_stop(self) -> None:
         """Cancels all remaining tasks when the server stops."""
         self._cancel_tasks()
@@ -446,7 +446,7 @@ class StudioChatInput(SocketIOInput, VoiceInputChannel):
         if sid in self.active_connections:
             del self.active_connections[sid]
 
-    @hookimpl  # type: ignore[misc]
+    @hookimpl
     def after_server_stop(self) -> None:
         """Cleanup background tasks and active connections when the server stops."""
         structlogger.info("studio_chat.after_server_stop.cleanup")
@@ -531,7 +531,7 @@ class StudioVoiceOutputChannel(VoiceOutputChannel):
 
     def create_marker_message(self, recipient_id: str) -> Tuple[str, str]:
         message_id = uuid.uuid4().hex
-        marker_data = {"marker": message_id}
+        marker_data: Dict[str, Any] = {"marker": message_id}
 
         # Include comprehensive latency information if available
         latency_data = {
@@ -546,7 +546,7 @@ class StudioVoiceOutputChannel(VoiceOutputChannel):
 
         # Add latency data to marker if any metrics are available
         if latency_data:
-            marker_data["latency"] = latency_data  # type: ignore[assignment]
+            marker_data["latency"] = latency_data
 
         return json.dumps(marker_data), message_id
 

rasa/core/channels/voice_stream/genesys.py

@@ -274,7 +274,7 @@ class GenesysInputChannel(VoiceInputChannel):
 
     def handle_ping(self, ws: Websocket, message: dict) -> None:
         """Handle ping message from Genesys."""
-        response = {
+        response: Dict[str, Any] = {
             "version": "2",
             "type": "pong",
             "seq": self._get_next_sequence(),

rasa/core/policies/flow_policy.py

@@ -138,7 +138,7 @@ class FlowPolicy(Policy):
         # create executor and predict next action
         try:
             prediction = await flow_executor.advance_flows(
-                tracker, domain.action_names_or_texts, flows
+                tracker, domain.action_names_or_texts, flows, domain.slots
             )
             return self._create_prediction_result(
                 prediction.action_name,
@@ -165,7 +165,7 @@ class FlowPolicy(Policy):
             events = tracker.create_stack_updated_events(updated_stack)
             tracker.update_with_events(events)
             prediction = await flow_executor.advance_flows(
-                tracker, domain.action_names_or_texts, flows
+                tracker, domain.action_names_or_texts, flows, domain.slots
             )
             collected_events = events + (prediction.events or [])
             return self._create_prediction_result(

rasa/core/policies/flows/flow_executor.py

@@ -10,12 +10,15 @@ from structlog.contextvars import bound_contextvars
 
 from rasa.agents.agent_manager import AgentManager
 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,
+    MAX_AGENT_RETRY_DELAY_SECONDS,
 )
 from rasa.agents.core.types import AgentStatus, ProtocolType
 from rasa.agents.schemas import AgentInput, AgentOutput
+from rasa.agents.schemas.agent_input import AgentInputSlot
 from rasa.core.available_agents import AvailableAgents
 from rasa.core.available_endpoints import AvailableEndpoints
 from rasa.core.constants import ACTIVE_FLOW_METADATA_KEY, STEP_ID_METADATA_KEY
@@ -103,7 +106,7 @@ from rasa.shared.core.flows.steps import (
     SetSlotsFlowStep,
 )
 from rasa.shared.core.flows.steps.constants import START_STEP
-from rasa.shared.core.slots import Slot, SlotRejection
+from rasa.shared.core.slots import CategoricalSlot, Slot, SlotRejection
 from rasa.shared.core.trackers import DialogueStateTracker
 from rasa.shared.utils.llm import tracker_as_readable_transcript
 
@@ -111,7 +114,6 @@ structlogger = structlog.get_logger()
 
 MAX_NUMBER_OF_STEPS = 250
 
-MAX_AGENT_RETRY_DELAY_SECONDS = 5
 MAX_AGENT_RETRIES = 3
 
 # Slots that should not be forwarded to sub-agents via AgentInput
@@ -387,7 +389,10 @@ def reset_scoped_slots(
 
 
 async def advance_flows(
-    tracker: DialogueStateTracker, available_actions: List[str], flows: FlowsList
+    tracker: DialogueStateTracker,
+    available_actions: List[str],
+    flows: FlowsList,
+    slots: List[Slot],
 ) -> FlowActionPrediction:
     """Advance the current flows until the next action.
 
@@ -395,6 +400,7 @@ async def advance_flows(
         tracker: The tracker to get the next action for.
         available_actions: The actions that are available in the domain.
         flows: All flows.
+        slots: The slots that are available in the domain.
 
     Returns:
     The predicted action and the events to run.
@@ -404,13 +410,16 @@ async def advance_flows(
         # if there are no flows, there is nothing to do
         return FlowActionPrediction(None, 0.0)
 
-    return await advance_flows_until_next_action(tracker, available_actions, flows)
+    return await advance_flows_until_next_action(
+        tracker, available_actions, flows, slots
+    )
 
 
 async def advance_flows_until_next_action(
     tracker: DialogueStateTracker,
     available_actions: List[str],
     flows: FlowsList,
+    slots: List[Slot],
 ) -> FlowActionPrediction:
     """Advance the flow and select the next action to execute.
 
@@ -476,6 +485,7 @@ async def advance_flows_until_next_action(
                     available_actions,
                     flows,
                     previous_step_id,
+                    slots,
                 )
                 new_events = step_result.events
                 if (
@@ -586,6 +596,7 @@ async def run_step(
     available_actions: List[str],
     flows: FlowsList,
     previous_step_id: str,
+    slots: List[Slot],
 ) -> FlowStepResult:
     """Run a single step of a flow.
 
@@ -604,6 +615,7 @@ async def run_step(
         available_actions: The actions that are available in the domain.
         flows: All flows.
         previous_step_id: The ID of the previous step.
+        slots: The slots that are available in the domain.
 
     Returns:
     A result of running the step describing where to transition to.
@@ -643,7 +655,7 @@ async def run_step(
         return _run_link_step(initial_events, stack, step)
 
     elif isinstance(step, CallFlowStep):
-        return await _run_call_step(initial_events, stack, step, tracker)
+        return await _run_call_step(initial_events, stack, step, tracker, slots)
 
     elif isinstance(step, SetSlotsFlowStep):
         return _run_set_slot_step(initial_events, step)
@@ -717,12 +729,13 @@ async def _run_call_step(
     stack: DialogueStack,
     step: CallFlowStep,
     tracker: DialogueStateTracker,
+    slots: List[Slot],
 ) -> FlowStepResult:
     structlogger.debug("flow.step.run.call")
     if step.is_calling_mcp_tool():
         return await call_mcp_tool(initial_events, stack, step, tracker)
     elif step.is_calling_agent():
-        return await run_agent(initial_events, stack, step, tracker)
+        return await run_agent(initial_events, stack, step, tracker, slots)
     else:
         stack.push(
             UserFlowStackFrame(
@@ -885,6 +898,7 @@ async def run_agent(
     stack: DialogueStack,
     step: CallFlowStep,
     tracker: DialogueStateTracker,
+    slots: List[Slot],
 ) -> FlowStepResult:
     """Run an agent call step."""
     structlogger.debug(
@@ -927,6 +941,9 @@ async def run_agent(
         if agent_stack_frame and agent_stack_frame.metadata
         else {}
     )
+    _update_agent_input_metadata_with_events(
+        agent_input_metadata, step.call, step.flow_id, tracker
+    )
     if step.exit_if:
         # TODO: this is a temporary fix to reset the slots covered by the exit_if
         if (
@@ -942,7 +959,7 @@ async def run_agent(
         user_message=tracker.latest_message.text or ""
         if tracker.latest_message
         else "",
-        slots=_filter_slots_for_agent(tracker.current_slot_values()),
+        slots=_prepare_slots_for_agent(tracker.current_slot_values(), slots),
         conversation_history=tracker_as_readable_transcript(tracker),
         events=tracker.current_state().get("events") or [],
         metadata=agent_input_metadata,
@@ -983,6 +1000,7 @@ async def run_agent(
             output.response_message or ""
         )
         output.metadata[AGENT_METADATA_TOOL_RESULTS_KEY] = output.tool_results or []
+        _update_agent_events(final_events, output.metadata)
 
         top_stack_frame = stack.top()
         # update the agent stack frame if it is already on the stack
@@ -1009,6 +1027,8 @@ async def run_agent(
         )
         return PauseFlowReturnPrediction(action_prediction)
     elif output.status == AgentStatus.COMPLETED:
+        output.metadata = output.metadata or {}
+        _update_agent_events(final_events, output.metadata)
         structlogger.debug(
             "flow.step.run_agent.completed",
             agent_name=step.call,
@@ -1027,6 +1047,8 @@ async def run_agent(
         else:
             return ContinueFlowWithNextStep(events=final_events)
     elif output.status == AgentStatus.FATAL_ERROR:
+        output.metadata = output.metadata or {}
+        _update_agent_events(final_events, output.metadata)
         # the agent failed, trigger pattern_internal_error
         structlogger.error(
             "flow.step.run_agent.fatal_error",
@@ -1044,6 +1066,8 @@ async def run_agent(
         stack.push(InternalErrorPatternFlowStackFrame())
         return ContinueFlowWithNextStep(events=final_events)
     else:
+        output.metadata = output.metadata or {}
+        _update_agent_events(final_events, output.metadata)
         structlogger.error(
             "flow.step.run_agent.unknown_status",
             agent_name=step.call,
@@ -1156,18 +1180,73 @@ async def _call_agent_with_retry(
     )
 
 
-def _filter_slots_for_agent(slots: Dict[str, Any]) -> Dict[str, Any]:
-    """Filter out slots that should not be forwarded to agents.
+def _prepare_slots_for_agent(
+    slot_values: Dict[str, Any], slot_definitions: List[Slot]
+) -> List[AgentInputSlot]:
+    """Prepare the slots for the agent.
+
+    Filter out slots that should not be forwarded to agents.
+    Add the slot type and allowed values to the slot dictionary.
 
     Args:
-        slots: The full slot dictionary from the tracker.
+        slot_values: The full slot dictionary from the tracker.
+        slot_definitions: The slot definitions from the domain.
 
     Returns:
-        A copy of the slot dictionary excluding slots listed in
-        `SLOTS_EXCLUDED_FOR_AGENT`.
+        A list of slots containing the name, current value, type, and allowed values.
     """
-    return {
-        key: value
-        for key, value in slots.items()
-        if key not in SLOTS_EXCLUDED_FOR_AGENT
-    }
+
+    def _get_slot_definition(slot_name: str) -> Optional[Slot]:
+        for slot in slot_definitions:
+            if slot.name == slot_name:
+                return slot
+        return None
+
+    filtered_slots: List[AgentInputSlot] = []
+    for key, value in slot_values.items():
+        if key in SLOTS_EXCLUDED_FOR_AGENT:
+            continue
+        slot_definition = _get_slot_definition(key)
+        if slot_definition:
+            filtered_slots.append(
+                AgentInputSlot(
+                    name=key,
+                    value=value,
+                    type=slot_definition.type_name if slot_definition else "any",
+                    allowed_values=slot_definition.values
+                    if isinstance(slot_definition, CategoricalSlot)
+                    else None,
+                )
+            )
+
+    return filtered_slots
+
+
+def _update_agent_events(events: List[Event], metadata: Dict[str, Any]) -> None:
+    """Update the agent events based on the agent output metadata if needed."""
+    if A2A_AGENT_CONTEXT_ID_KEY in metadata:
+        # If the context ID is present, we need to store it in the AgentStarted
+        # event, so that it can be re-used later in case the agent is restarted.
+        for event in events:
+            if isinstance(event, AgentStarted):
+                event.context_id = metadata[A2A_AGENT_CONTEXT_ID_KEY]
+
+
+def _update_agent_input_metadata_with_events(
+    metadata: Dict[str, Any], agent_id: str, flow_id: str, tracker: DialogueStateTracker
+) -> None:
+    """Update the agent input metadata with the events."""
+    agent_started_events = [
+        event
+        for event in tracker.events
+        if type(event) == AgentStarted
+        and event.agent_id == agent_id
+        and event.flow_id == flow_id
+    ]
+    if agent_started_events:
+        # If we have context ID from the previous agent run, we want to
+        # include it in the metadata so that the agent can continue the same
+        # context.
+        agent_started_event = agent_started_events[-1]
+        if agent_started_event.context_id:
+            metadata[A2A_AGENT_CONTEXT_ID_KEY] = agent_started_event.context_id

rasa/core/policies/flows/mcp_tool_executor.py

@@ -1,4 +1,6 @@
 import json
+import operator
+from functools import reduce
 from typing import Any, Dict, List, Optional
 
 import structlog
@@ -17,10 +19,12 @@ from rasa.shared.core.events import Event, SlotSet
 from rasa.shared.core.flows.steps import CallFlowStep
 from rasa.shared.core.trackers import DialogueStateTracker
 from rasa.shared.utils.mcp.server_connection import MCPServerConnection
+from rasa.utils.common import ensure_jsonified_iterable
 
 structlogger = structlog.get_logger()
 
-SEPARATOR = "\n###\n"
+CONFIG_RESULT_KEY = "result_key"
+CONFIG_SLOT = "slot"
 
 
 async def call_mcp_tool(
@@ -101,7 +105,7 @@ async def _execute_mcp_tool_call(
         result = await mcp_server.call_tool(step.call, arguments)
 
         # Handle tool execution result
-        if result.isError:
+        if result is None or result.isError:
             return _handle_mcp_tool_error(
                 stack,
                 initial_events,
@@ -121,12 +125,20 @@ async def _execute_mcp_tool_call(
                 tool_name=step.call,
                 mcp_server=step.mcp_server,
                 result_content=result.content,
+                result_structured_content=result.structuredContent,
             )
 
             # Process successful result
-            set_slot_event = _process_tool_result(result, step.mapping["output"])
-            if set_slot_event:
-                initial_events.append(set_slot_event)
+            if set_slot_event := _process_tool_result(result, step.mapping["output"]):
+                initial_events.extend(set_slot_event)
+            else:
+                return _handle_mcp_tool_error(
+                    stack,
+                    initial_events,
+                    f"Failed to process tool result for '{step.call}'.",
+                    tool_name=step.call,
+                    mcp_server=step.mcp_server,
+                )
 
         return ContinueFlowWithNextStep(events=initial_events)
 
@@ -203,16 +215,41 @@ def _prepare_tool_arguments(
     return arguments
 
 
+def _jsonify_slot_value(value: Any) -> str | int | float | bool | None:
+    """Prepare value for SlotSet: iterables -> JSON string, primitives -> as-is"""
+    if isinstance(value, (list, dict)) and len(value):
+        return json.dumps(ensure_jsonified_iterable(value))
+    return value
+
+
 def _process_tool_result(
     result: CallToolResult,
-    output_mapping: str,
-) -> Optional[SlotSet]:
+    output_mapping: List[Dict[str, str]],
+) -> Optional[List[SlotSet]]:
     """Create a SetSlot event for the tool result."""
     try:
-        content_as_string = [
-            json.dumps(content_part.model_dump()) for content_part in result.content
-        ]
-        return SlotSet(key=output_mapping, value=SEPARATOR.join(content_as_string))
+        _result_as_dict = {"result": result.model_dump()}
+        slots = []
+        for mapping in output_mapping:
+            try:
+                # Use reduce to navigate through nested keys in the result
+                slot_value = reduce(
+                    operator.getitem,
+                    mapping[CONFIG_RESULT_KEY].split("."),
+                    _result_as_dict,
+                )
+                slots.append(
+                    SlotSet(mapping[CONFIG_SLOT], _jsonify_slot_value(slot_value))
+                )
+            except (KeyError, TypeError):
+                structlogger.error(
+                    "call_mcp_tool.result_key_not_found_in_tool_result",
+                    slot=mapping[CONFIG_SLOT],
+                    result_key=mapping[CONFIG_RESULT_KEY],
+                    result=_result_as_dict,
+                )
+                return None
+        return slots
     except Exception as e:
         structlogger.error(
             "call_mcp_tool.result_processing_failed",

rasa/core/policies/intentless_policy.py

@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Text, Tuple
 
 import structlog
 import tiktoken
-from deprecated import deprecated  # type: ignore[import]
+from deprecated import deprecated  # type: ignore[import-untyped]
 from jinja2 import Template
 from langchain.docstore.document import Document
 from langchain.schema.embeddings import Embeddings

rasa/core/policies/unexpected_intent_policy.py

@@ -624,6 +624,7 @@ class UnexpecTEDIntentPolicy(TEDPolicy):
         query_intent = (
             last_user_uttered_event.intent_name
             if last_user_uttered_event is not None
+            and isinstance(last_user_uttered_event, UserUttered)
             else ""
         )
         is_unlikely_intent = self._check_unlikely_intent(

rasa/core/processor.py

@@ -730,7 +730,7 @@ class MessageProcessor:
         if not self.domain or self.domain.is_empty():
             return
 
-        intent = parse_data["intent"][INTENT_NAME_KEY]
+        intent = parse_data[INTENT][INTENT_NAME_KEY]
         if intent and intent not in self.domain.intents:
             rasa.shared.utils.io.raise_warning(
                 f"Parsed an intent '{intent}' "
@@ -739,7 +739,7 @@ class MessageProcessor:
                 docs=DOCS_URL_DOMAINS,
             )
 
-        entities = parse_data["entities"] or []
+        entities = parse_data[ENTITIES] or []
         for element in entities:
             entity = element["entity"]
             if entity and entity not in self.domain.entities:
@@ -823,9 +823,9 @@ class MessageProcessor:
         self._update_full_retrieval_intent(parse_data)
         structlogger.debug(
             "processor.message.parse",
-            parse_data_text=copy.deepcopy(parse_data["text"]),
-            parse_data_intent=parse_data["intent"],
-            parse_data_entities=copy.deepcopy(parse_data["entities"]),
+            parse_data_text=copy.deepcopy(parse_data[TEXT]),
+            parse_data_intent=parse_data[INTENT],
+            parse_data_entities=copy.deepcopy(parse_data[ENTITIES]),
         )
 
         self._check_for_unseen_features(parse_data)
@@ -974,7 +974,7 @@ class MessageProcessor:
                     f"invalid intent: {parse_data[INTENT]['name']}. "
                     f"Returning CannotHandleCommand() as a fallback."
                 ),
-                invalid_intent=parse_data[INTENT]["name"],
+                invalid_intent=parse_data[INTENT][INTENT_NAME_KEY],
             )
             commands.append(
                 CannotHandleCommand(RASA_PATTERN_CANNOT_HANDLE_INVALID_INTENT)
@@ -984,7 +984,7 @@ class MessageProcessor:
 
     def _contains_undefined_intent(self, message: Message) -> bool:
         """Checks if the message contains an undefined intent."""
-        intent_name = message.get(INTENT, {}).get("name")
+        intent_name = message.get(INTENT, {}).get(INTENT_NAME_KEY)
         return intent_name is not None and intent_name not in self.domain.intents
 
     async def _parse_message_with_graph(
@@ -1034,8 +1034,8 @@ class MessageProcessor:
         tracker.update(
             UserUttered(
                 message.text,
-                parse_data["intent"],
-                parse_data["entities"],
+                parse_data[INTENT],
+                parse_data[ENTITIES],
                 parse_data,
                 input_channel=message.input_channel,
                 message_id=message.message_id,
@@ -1044,7 +1044,7 @@ class MessageProcessor:
             self.domain,
         )
 
-        if parse_data["entities"]:
+        if parse_data[ENTITIES]:
             self._log_slots(tracker)
 
         plugin_manager().hook.after_new_user_message(tracker=tracker)
@@ -1472,11 +1472,9 @@ class MessageProcessor:
             # tracker has never expired if sessions are disabled
             return False
 
-        user_uttered_event: Optional[UserUttered] = tracker.get_last_event_for(
-            UserUttered
-        )
+        user_uttered_event = tracker.get_last_event_for(UserUttered)
 
-        if not user_uttered_event:
+        if not user_uttered_event or not isinstance(user_uttered_event, UserUttered):
             # there is no user event so far so the session should not be considered
             # expired
             return False

rasa/core/tracker_stores/tracker_store.py

@@ -542,7 +542,7 @@ class FailSafeTrackerStore(TrackerStore):
         return self._tracker_store.domain
 
     @domain.setter
-    def domain(self, domain: Domain) -> None:
+    def domain(self, domain: Optional[Domain]) -> None:
         self._tracker_store.domain = domain
 
         if self._fallback_tracker_store:
@@ -805,9 +805,7 @@ class AwaitableTrackerStore(TrackerStore):
     async def retrieve(self, sender_id: Text) -> Optional[DialogueStateTracker]:
         """Wrapper to call `retrieve` method of primary tracker store."""
         result = self._tracker_store.retrieve(sender_id)
-        return (
-            await result if isawaitable(result) else result  # type: ignore[return-value, misc]
-        )
+        return await result if isawaitable(result) else result
 
     async def keys(self) -> Iterable[Text]:
         """Wrapper to call `keys` method of primary tracker store."""
@@ -834,6 +832,4 @@ class AwaitableTrackerStore(TrackerStore):
     ) -> Optional[DialogueStateTracker]:
         """Wrapper to call `retrieve_full_tracker` method of primary tracker store."""
         result = self._tracker_store.retrieve_full_tracker(conversation_id)
-        return (
-            await result if isawaitable(result) else result  # type: ignore[return-value, misc]
-        )
+        return await result if isawaitable(result) else result

rasa/core/train.py

@@ -46,7 +46,7 @@ async def train_comparison_models(
                     output=str(Path(output_path, f"run_{r +1}")),
                     fixed_model_name=config_name + PERCENTAGE_KEY + str(percentage),
                     additional_arguments={
-                        **additional_arguments,
+                        **(additional_arguments or {}),
                         "exclusion_percentage": percentage,
                     },
                 )

rasa/core/training/interactive.py

@@ -91,7 +91,7 @@ from rasa.shared.core.training_data.visualization import (
 )
 from rasa.shared.exceptions import InvalidConfigException
 from rasa.shared.importers.rasa import TrainingDataImporter
-from rasa.shared.nlu.constants import INTENT_NAME_KEY, TEXT
+from rasa.shared.nlu.constants import ENTITIES, INTENT, INTENT_NAME_KEY, TEXT
 
 # noinspection PyProtectedMember
 from rasa.shared.nlu.training_data import loading
@@ -789,7 +789,7 @@ def _collect_messages(events: List[Dict[Text, Any]]) -> List[Message]:
             data = event.get("parse_data", {})
             rasa_nlu_training_data_utils.remove_untrainable_entities_from(data)
             msg = Message.build(
-                data["text"], data["intent"][INTENT_NAME_KEY], data["entities"]
+                data[TEXT], data[INTENT][INTENT_NAME_KEY], data[ENTITIES]
             )
             messages.append(msg)
         elif event.get("event") == UserUtteranceReverted.type_name and messages:
@@ -901,13 +901,13 @@ def _get_nlu_target_format(export_path: Text) -> Text:
 
 def _entities_from_messages(messages: List[Message]) -> List[Text]:
     """Return all entities that occur in at least one of the messages."""
-    return list({e["entity"] for m in messages for e in m.data.get("entities", [])})
+    return list({e["entity"] for m in messages for e in m.data.get(ENTITIES, [])})
 
 
 def _intents_from_messages(messages: List[Message]) -> Set[Text]:
     """Return all intents that occur in at least one of the messages."""
     # set of distinct intents
-    distinct_intents = {m.data["intent"] for m in messages if "intent" in m.data}
+    distinct_intents = {m.data[INTENT] for m in messages if INTENT in m.data}
 
     return distinct_intents
 
@@ -1191,11 +1191,11 @@ def _as_md_message(parse_data: Dict[Text, Any]) -> Text:
     """Display the parse data of a message in markdown format."""
     from rasa.shared.nlu.training_data.formats.readerwriter import TrainingDataWriter
 
-    if parse_data.get("text", "").startswith(INTENT_MESSAGE_PREFIX):
-        return parse_data["text"]
+    if parse_data.get(TEXT, "").startswith(INTENT_MESSAGE_PREFIX):
+        return parse_data[TEXT]
 
-    if not parse_data.get("entities"):
-        parse_data["entities"] = []
+    if not parse_data.get(ENTITIES):
+        parse_data[ENTITIES] = []
 
     return TrainingDataWriter.generate_message(parse_data)
 
@@ -1207,7 +1207,7 @@ def _validate_user_regex(latest_message: Dict[Text, Any], intents: List[Text]) -
     `/greet`. Return `True` if the intent is a known one.
     """
     parse_data = latest_message.get("parse_data", {})
-    intent = parse_data.get("intent", {}).get(INTENT_NAME_KEY)
+    intent = parse_data.get(INTENT, {}).get(INTENT_NAME_KEY)
 
     if intent in intents:
         return True
@@ -1224,8 +1224,8 @@ async def _validate_user_text(
     """
     parse_data = latest_message.get("parse_data", {})
     text = _as_md_message(parse_data)
-    intent = parse_data.get("intent", {}).get(INTENT_NAME_KEY)
-    entities = parse_data.get("entities", [])
+    intent = parse_data.get(INTENT, {}).get(INTENT_NAME_KEY)
+    entities = parse_data.get(ENTITIES, [])
     if entities:
         message = (
             f"Is the intent '{intent}' correct for '{text}' and are "
@@ -1276,9 +1276,9 @@ async def _validate_nlu(
 
         entities = await _correct_entities(latest_message, endpoint, conversation_id)
         corrected_nlu = {
-            "intent": corrected_intent,
-            "entities": entities,
-            "text": latest_message.get("text"),
+            INTENT: corrected_intent,
+            ENTITIES: entities,
+            TEXT: latest_message.get("text"),
         }
 
         await _correct_wrong_nlu(corrected_nlu, events, endpoint, conversation_id)
@@ -1315,9 +1315,9 @@ def _merge_annotated_and_original_entities(
     # overwrite entities which have already been
     # annotated in the original annotation to preserve
     # additional entity parser information
-    entities = parse_annotated.get("entities", [])[:]
+    entities = parse_annotated.get(ENTITIES, [])[:]
     for i, entity in enumerate(entities):
-        for original_entity in parse_original.get("entities", []):
+        for original_entity in parse_original.get(ENTITIES, []):
             if _is_same_entity_annotation(entity, original_entity):
                 entities[i] = original_entity
                 break