Merge pull request #9043 from RasaHQ/misc/mypy-index

fix `mypy` index issues

Author: wochinge

rasa/cli/x.py

@@ -2,7 +2,7 @@
 import asyncio
 import importlib.util
 import logging
-from multiprocessing import get_context
+from multiprocessing import get_context, Process
 import os
 import signal
 import sys
@@ -196,9 +196,10 @@ def _is_correct_event_broker(event_broker: EndpointConfig) -> bool:
     )
 
 
-def start_rasa_for_local_rasa_x(args: argparse.Namespace, rasa_x_token: Text) -> None:
+def start_rasa_for_local_rasa_x(
+    args: argparse.Namespace, rasa_x_token: Text
+) -> Process:
     """Starts the Rasa X API with Rasa as a background process."""
-
     credentials_path, endpoints_path = _get_credentials_and_endpoints_paths(args)
     endpoints = AvailableEndpoints.read_endpoints(endpoints_path)
 

rasa/core/actions/action.py

@@ -58,6 +58,7 @@
     from rasa.shared.core.trackers import DialogueStateTracker
     from rasa.core.nlg import NaturalLanguageGenerator
     from rasa.core.channels.channel import OutputChannel
+    from rasa.shared.core.events import IntentPrediction
 
 logger = logging.getLogger(__name__)
 
@@ -851,9 +852,17 @@ async def run(
         domain: "Domain",
     ) -> List[Event]:
         """Runs action. Please see parent class for the full docstring."""
-        intent_to_affirm = tracker.latest_message.intent.get(INTENT_NAME_KEY)
+        latest_message = tracker.latest_message
+        if latest_message is None:
+            raise TypeError(
+                "Cannot find last user message for detecting fallback affirmation."
+            )
+
+        intent_to_affirm = latest_message.intent.get(INTENT_NAME_KEY)
 
-        intent_ranking = tracker.latest_message.parse_data.get(INTENT_RANKING_KEY, [])
+        intent_ranking: List["IntentPrediction"] = latest_message.parse_data.get(
+            INTENT_RANKING_KEY
+        ) or []
         if (
             intent_to_affirm == DEFAULT_NLU_FALLBACK_INTENT_NAME
             and len(intent_ranking) > 1

rasa/core/actions/two_stage_fallback.py

@@ -24,6 +24,7 @@
     ACTION_DEFAULT_ASK_REPHRASE_NAME,
     ACTION_TWO_STAGE_FALLBACK_NAME,
 )
+from rasa.shared.nlu.constants import INTENT, PREDICTED_CONFIDENCE_KEY
 from rasa.utils.endpoints import EndpointConfig
 
 
@@ -124,7 +125,7 @@ def _last_intent_name(tracker: DialogueStateTracker) -> Optional[Text]:
     if not last_message:
         return None
 
-    return last_message.intent.get("name")
+    return last_message.intent_name
 
 
 def _two_fallbacks_in_a_row(tracker: DialogueStateTracker) -> bool:
@@ -179,6 +180,6 @@ def _message_clarification(tracker: DialogueStateTracker) -> List[Event]:
         )
 
     clarification = copy.deepcopy(latest_message)
-    clarification.parse_data["intent"]["confidence"] = 1.0
+    clarification.parse_data[INTENT][PREDICTED_CONFIDENCE_KEY] = 1.0
     clarification.timestamp = time.time()
     return [ActionExecuted(ACTION_LISTEN_NAME), clarification]

rasa/core/channels/console.py

@@ -8,7 +8,7 @@
 import questionary
 from aiohttp import ClientTimeout
 from prompt_toolkit.styles import Style
-from typing import Any
+from typing import Any, Generator
 from typing import Text, Optional, Dict, List
 
 import rasa.shared.utils.cli
@@ -121,9 +121,9 @@ async def send_message_receive_block(
             return await resp.json()
 
 
-async def send_message_receive_stream(
+async def _send_message_receive_stream(
     server_url: Text, auth_token: Text, sender_id: Text, message: Text
-) -> None:
+) -> Generator[Dict[Text, Any], None, None]:
     payload = {"sender": sender_id, "message": message}
 
     url = f"{server_url}/webhooks/rest/webhook?stream=true&token={auth_token}"
@@ -175,7 +175,7 @@ async def record_messages(
             break
 
         if use_response_stream:
-            bot_responses = send_message_receive_stream(
+            bot_responses = _send_message_receive_stream(
                 server_url, auth_token, sender_id, text
             )
             previous_response = None

rasa/core/policies/ensemble.py

@@ -12,7 +12,7 @@
 import rasa.core
 import rasa.core.training.training
 from rasa.core.constants import FALLBACK_POLICY_PRIORITY
-from rasa.shared.exceptions import RasaException
+from rasa.shared.exceptions import RasaException, InvalidConfigException
 import rasa.shared.utils.common
 import rasa.shared.utils.io
 import rasa.utils.io
@@ -636,7 +636,13 @@ def _pick_best_policy(
             if form_confidence > best_confidence:
                 best_policy_name = form_policy_name
 
-        best_prediction = predictions[best_policy_name]
+        best_prediction = predictions.get(best_policy_name)
+
+        if not best_prediction:
+            raise InvalidConfigException(
+                f"No prediction for policy '{best_policy_name}' found. Please check "
+                f"your model configuration."
+            )
 
         policy_events += best_prediction.optional_events
 

rasa/core/policies/rule_policy.py

@@ -145,7 +145,7 @@ def __init__(
         self._restrict_rules = restrict_rules
         self._check_for_contradictions = check_for_contradictions
 
-        self._rules_sources = None
+        self._rules_sources = defaultdict(list)
 
         # max history is set to `None` in order to capture any lengths of rule stories
         super().__init__(

rasa/core/policies/ted_policy.py

@@ -684,7 +684,7 @@ def predict_action_probabilities(
             tracker, domain, interpreter
         )
         model_data = self._create_model_data(tracker_state_features)
-        outputs = self.model.run_inference(model_data)
+        outputs: Dict[Text, np.ndarray] = self.model.run_inference(model_data)
 
         # take the last prediction in the sequence
         similarities = outputs["similarities"][:, -1, :]

rasa/core/policies/two_stage_fallback.py

@@ -22,7 +22,7 @@
     ACTION_DEFAULT_ASK_REPHRASE_NAME,
 )
 from rasa.shared.core.domain import InvalidDomain, Domain
-from rasa.shared.nlu.constants import ACTION_NAME, INTENT_NAME_KEY
+from rasa.shared.nlu.constants import ACTION_NAME, INTENT_NAME_KEY, INTENT
 
 if TYPE_CHECKING:
     from rasa.core.policies.ensemble import PolicyEnsemble
@@ -124,7 +124,7 @@ def predict_action_probabilities(
     ) -> PolicyPrediction:
         """Predicts the next action if NLU confidence is low."""
         nlu_data = tracker.latest_message.parse_data
-        last_intent_name = nlu_data["intent"].get(INTENT_NAME_KEY, None)
+        last_intent_name = nlu_data[INTENT].get(INTENT_NAME_KEY, None)
         should_nlu_fallback = self.should_nlu_fallback(
             nlu_data, tracker.latest_action.get(ACTION_NAME)
         )

rasa/core/test.py

@@ -18,11 +18,9 @@
 )
 from rasa.shared.core.domain import Domain
 from rasa.nlu.constants import (
-    ENTITY_ATTRIBUTE_TEXT,
     RESPONSE_SELECTOR_DEFAULT_INTENT,
     RESPONSE_SELECTOR_RETRIEVAL_INTENTS,
     TOKENS_NAMES,
-    ENTITY_ATTRIBUTE_CONFIDENCE,
 )
 from rasa.shared.nlu.constants import (
     INTENT,
@@ -38,8 +36,7 @@
     RESPONSE_SELECTOR,
     FULL_RETRIEVAL_INTENT_NAME_KEY,
     TEXT,
-    ENTITY_ATTRIBUTE_GROUP,
-    ENTITY_ATTRIBUTE_ROLE,
+    ENTITY_ATTRIBUTE_TEXT,
 )
 from rasa.constants import RESULTS_FILE, PERCENTAGE_KEY
 from rasa.shared.core.events import (
@@ -56,24 +53,7 @@
     from rasa.core.agent import Agent
     from rasa.core.processor import MessageProcessor
     from rasa.shared.core.generator import TrainingDataGenerator
-
-    from typing_extensions import TypedDict
-
-    EntityPrediction = TypedDict(
-        "EntityPrediction",
-        {
-            ENTITY_ATTRIBUTE_TEXT: Text,
-            ENTITY_ATTRIBUTE_START: Optional[float],
-            ENTITY_ATTRIBUTE_END: Optional[float],
-            ENTITY_ATTRIBUTE_VALUE: Text,
-            ENTITY_ATTRIBUTE_CONFIDENCE: float,
-            ENTITY_ATTRIBUTE_TYPE: Text,
-            ENTITY_ATTRIBUTE_GROUP: Optional[Text],
-            ENTITY_ATTRIBUTE_ROLE: Optional[Text],
-            "additional_info": Any,
-        },
-        total=False,
-    )
+    from rasa.shared.core.events import EntityPrediction
 
 CONFUSION_MATRIX_STORIES_FILE = "story_confusion_matrix.png"
 REPORT_STORIES_FILE = "story_report.json"
@@ -911,7 +891,7 @@ async def test(
             num_failed = len(story_evaluation.failed_stories)
             num_correct = len(story_evaluation.successful_stories)
             num_convs = num_failed + num_correct
-            if num_convs:
+            if num_convs and isinstance(report, Dict):
                 conv_accuracy = num_correct / num_convs
                 report["conversation_accuracy"] = {
                     "accuracy": conv_accuracy,

rasa/core/training/interactive.py

@@ -961,7 +961,7 @@ async def _predict_till_next_listen(
         predictions = result.get("scores")
         probabilities = [prediction["score"] for prediction in predictions]
         pred_out = int(np.argmax(probabilities))
-        action_name = predictions[pred_out].get("action")
+        action_name = predictions.get(pred_out, {}).get("action")
         policy = result.get("policy")
         confidence = result.get("confidence")
 

rasa/nlu/classifiers/diet_classifier.py

@@ -341,7 +341,7 @@ def __init__(
         self._check_config_parameters()
 
         # transform numbers to labels
-        self.index_label_id_mapping = index_label_id_mapping
+        self.index_label_id_mapping = index_label_id_mapping or {}
 
         self._entity_tag_specs = entity_tag_specs
 
@@ -647,7 +647,8 @@ def _create_label_data(
         return label_data
 
     def _use_default_label_features(self, label_ids: np.ndarray) -> List[FeatureArray]:
-        all_label_features = self._label_data.get(LABEL, SENTENCE)[0]
+        feature_arrays: List[FeatureArray] = self._label_data.get(LABEL, SENTENCE)
+        all_label_features = feature_arrays[0]
         return [
             FeatureArray(
                 np.array([all_label_features[label_id] for label_id in label_ids]),

rasa/nlu/constants.py

@@ -1,12 +1,10 @@
 import rasa.shared.nlu.constants
-
+from rasa.shared.nlu.constants import ENTITY_ATTRIBUTE_CONFIDENCE
 
 BILOU_ENTITIES = "bilou_entities"
 BILOU_ENTITIES_ROLE = "bilou_entities_role"
 BILOU_ENTITIES_GROUP = "bilou_entities_group"
 
-ENTITY_ATTRIBUTE_TEXT = "text"
-ENTITY_ATTRIBUTE_CONFIDENCE = "confidence"
 ENTITY_ATTRIBUTE_CONFIDENCE_TYPE = (
     f"{ENTITY_ATTRIBUTE_CONFIDENCE}_{rasa.shared.nlu.constants.ENTITY_ATTRIBUTE_TYPE}"
 )

rasa/nlu/featurizers/sparse_featurizer/count_vectors_featurizer.py

@@ -169,7 +169,7 @@ def _check_attribute_vocabulary(self, attribute: Text) -> bool:
         """Checks if trained vocabulary exists in attribute's count vectorizer."""
         try:
             return hasattr(self.vectorizers[attribute], "vocabulary_")
-        except (AttributeError, TypeError):
+        except (AttributeError, KeyError):
             return False
 
     def _get_attribute_vocabulary(self, attribute: Text) -> Optional[Dict[Text, int]]:
@@ -240,7 +240,7 @@ def __init__(
         self._attributes = self._attributes_for(self.analyzer)
 
         # declare class instance for CountVectorizer
-        self.vectorizers = vectorizers
+        self.vectorizers = vectorizers or {}
 
         self.finetune_mode = finetune_mode
 

rasa/nlu/featurizers/sparse_featurizer/lexical_syntactic_featurizer.py

@@ -63,7 +63,7 @@ def required_components(cls) -> List[Type[Component]]:
         "pos": lambda token: token.data.get(POS_TAG_KEY)
         if POS_TAG_KEY in token.data
         else None,
-        "pos2": lambda token: token.data.get(POS_TAG_KEY)[:2]
+        "pos2": lambda token: token.data.get(POS_TAG_KEY, [])[:2]
         if "pos" in token.data
         else None,
         "upper": lambda token: token.text.isupper(),

rasa/nlu/model.py

@@ -358,6 +358,7 @@ def _update_metadata_epochs(
         new_config: Optional[Dict] = None,
         finetuning_epoch_fraction: float = 1.0,
     ) -> Metadata:
+        new_config = new_config or {}
         for old_component_config, new_component_config in zip(
             model_metadata.metadata["pipeline"], new_config["pipeline"]
         ):

rasa/nlu/test.py

@@ -35,7 +35,6 @@
     ENTITY_ATTRIBUTE_CONFIDENCE_TYPE,
     ENTITY_ATTRIBUTE_CONFIDENCE_ROLE,
     ENTITY_ATTRIBUTE_CONFIDENCE_GROUP,
-    ENTITY_ATTRIBUTE_TEXT,
 )
 from rasa.shared.nlu.constants import (
     TEXT,
@@ -50,6 +49,7 @@
     NO_ENTITY_TAG,
     INTENT_NAME_KEY,
     PREDICTED_CONFIDENCE_KEY,
+    ENTITY_ATTRIBUTE_TEXT,
 )
 from rasa.model import get_model
 from rasa.nlu.components import ComponentBuilder

rasa/shared/core/events.py

@@ -49,11 +49,49 @@
     INTENT_NAME_KEY,
     ENTITY_ATTRIBUTE_ROLE,
     ENTITY_ATTRIBUTE_GROUP,
+    PREDICTED_CONFIDENCE_KEY,
+    INTENT_RANKING_KEY,
+    ENTITY_ATTRIBUTE_TEXT,
+    ENTITY_ATTRIBUTE_START,
+    ENTITY_ATTRIBUTE_CONFIDENCE,
+    ENTITY_ATTRIBUTE_END,
 )
 
 if TYPE_CHECKING:
+    from typing_extensions import TypedDict
+
     from rasa.shared.core.trackers import DialogueStateTracker
 
+    EntityPrediction = TypedDict(
+        "EntityPrediction",
+        {
+            ENTITY_ATTRIBUTE_TEXT: Text,
+            ENTITY_ATTRIBUTE_START: Optional[float],
+            ENTITY_ATTRIBUTE_END: Optional[float],
+            ENTITY_ATTRIBUTE_VALUE: Text,
+            ENTITY_ATTRIBUTE_CONFIDENCE: float,
+            ENTITY_ATTRIBUTE_TYPE: Text,
+            ENTITY_ATTRIBUTE_GROUP: Optional[Text],
+            ENTITY_ATTRIBUTE_ROLE: Optional[Text],
+            "additional_info": Any,
+        },
+        total=False,
+    )
+
+    IntentPrediction = TypedDict(
+        "IntentPrediction", {INTENT_NAME_KEY: Text, PREDICTED_CONFIDENCE_KEY: float,},
+    )
+    NLUPredictionData = TypedDict(
+        "NLUPredictionData",
+        {
+            INTENT: IntentPrediction,
+            INTENT_RANKING_KEY: List[IntentPrediction],
+            ENTITIES: List[EntityPrediction],
+            "message_id": Optional[Text],
+            "metadata": Dict,
+        },
+        total=False,
+    )
 logger = logging.getLogger(__name__)
 
 
@@ -369,7 +407,7 @@ def __init__(
         text: Optional[Text] = None,
         intent: Optional[Dict] = None,
         entities: Optional[List[Dict]] = None,
-        parse_data: Optional[Dict[Text, Any]] = None,
+        parse_data: Optional["NLUPredictionData"] = None,
         timestamp: Optional[float] = None,
         input_channel: Optional[Text] = None,
         message_id: Optional[Text] = None,
@@ -410,7 +448,7 @@ def __init__(
             # happens during training
             self.use_text_for_featurization = False
 
-        self.parse_data = {
+        self.parse_data: "NLUPredictionData" = {
             INTENT: self.intent,
             # Copy entities so that changes to `self.entities` don't affect
             # `self.parse_data` and hence don't get persisted
@@ -426,7 +464,7 @@ def __init__(
     @staticmethod
     def _from_parse_data(
         text: Text,
-        parse_data: Dict[Text, Any],
+        parse_data: "NLUPredictionData",
         timestamp: Optional[float] = None,
         input_channel: Optional[Text] = None,
         message_id: Optional[Text] = None,