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

rasa/tracing/config.py

@@ -30,8 +30,8 @@ from rasa.dialogue_understanding.commands import (
 )
 from rasa.dialogue_understanding.generator import (
     LLMCommandGenerator,
-    SingleStepLLMCommandGenerator,
     MultiStepLLMCommandGenerator,
+    SingleStepLLMCommandGenerator,
 )
 from rasa.dialogue_understanding.generator.flow_retrieval import FlowRetrieval
 from rasa.dialogue_understanding.generator.nlu_command_adapter import NLUCommandAdapter
@@ -40,11 +40,11 @@ from rasa.engine.training.graph_trainer import GraphTrainer
 from rasa.tracing.constants import (
     ENDPOINTS_ENDPOINT_KEY,
     ENDPOINTS_INSECURE_KEY,
+    ENDPOINTS_METRICS_KEY,
     ENDPOINTS_OTLP_BACKEND_TYPE,
+    ENDPOINTS_ROOT_CERTIFICATES_KEY,
     ENDPOINTS_TRACING_KEY,
     ENDPOINTS_TRACING_SERVICE_NAME_KEY,
-    ENDPOINTS_ROOT_CERTIFICATES_KEY,
-    ENDPOINTS_METRICS_KEY,
 )
 from rasa.tracing.instrumentation import instrumentation
 from rasa.tracing.metric_instrument_provider import MetricInstrumentProvider

rasa/tracing/instrumentation/attribute_extractors.py

@@ -21,6 +21,20 @@ from rasa.dialogue_understanding.commands import Command
 from rasa.dialogue_understanding.generator import LLMBasedCommandGenerator
 from rasa.dialogue_understanding.generator.constants import FLOW_RETRIEVAL_KEY
 from rasa.dialogue_understanding.stack.dialogue_stack import DialogueStack
+from rasa.dialogue_understanding_test.du_test_result import (
+    KEY_TEST_CASES_ACCURACY,
+    KEY_USER_UTTERANCES_ACCURACY,
+    OUTPUT_COMMAND_METRICS,
+    OUTPUT_NAMES_OF_FAILED_TESTS,
+    OUTPUT_NAMES_OF_PASSED_TESTS,
+    OUTPUT_NUMBER_OF_FAILED_TESTS,
+    OUTPUT_NUMBER_OF_FAILED_USER_UTTERANCES,
+    OUTPUT_NUMBER_OF_PASSED_TESTS,
+    OUTPUT_NUMBER_OF_PASSED_USER_UTTERANCES,
+    OUTPUT_TEST_CASES_ACCURACY,
+    OUTPUT_USER_UTTERANCES_ACCURACY,
+    DialogueUnderstandingTestSuiteResult,
+)
 from rasa.engine.graph import ExecutionContext, GraphModelConfiguration, GraphNode
 from rasa.engine.training.graph_trainer import GraphTrainer
 from rasa.shared.constants import (
@@ -571,6 +585,44 @@ def extract_attrs_for_advance_flows(
     }
 
 
+def extract_attrs_for_du_print_test_results(
+    test_suite_result: DialogueUnderstandingTestSuiteResult,
+    output_prompt: bool,
+) -> Dict[str, Any]:
+    """Extract the attributes for
+    `rasa.dialogue_understanding_test.io.print_test_results` function.
+    """
+    from rasa.tracing.instrumentation.instrumentation import (
+        DIALOG_UNDERSTANDING_TEST_IO_MODULE_NAME,
+    )
+
+    attributes_dict = {
+        "module_name": DIALOG_UNDERSTANDING_TEST_IO_MODULE_NAME,
+        OUTPUT_TEST_CASES_ACCURACY: test_suite_result.accuracy[KEY_TEST_CASES_ACCURACY],
+        OUTPUT_USER_UTTERANCES_ACCURACY: test_suite_result.accuracy[
+            KEY_USER_UTTERANCES_ACCURACY
+        ],
+        OUTPUT_NUMBER_OF_PASSED_TESTS: test_suite_result.number_of_passed_tests,
+        OUTPUT_NUMBER_OF_FAILED_TESTS: test_suite_result.number_of_failed_tests,
+        OUTPUT_NUMBER_OF_PASSED_USER_UTTERANCES: test_suite_result.number_of_passed_user_utterances,  # noqa: E501
+        OUTPUT_NUMBER_OF_FAILED_USER_UTTERANCES: test_suite_result.number_of_failed_user_utterances,  # noqa: E501
+        OUTPUT_NAMES_OF_PASSED_TESTS: json.dumps(
+            test_suite_result.names_of_passed_tests
+        ),
+        OUTPUT_NAMES_OF_FAILED_TESTS: json.dumps(
+            test_suite_result.names_of_failed_tests
+        ),
+    }
+    if test_suite_result.command_metrics:
+        attributes_dict[OUTPUT_COMMAND_METRICS] = json.dumps(
+            {
+                key: value.as_dict()
+                for key, value in test_suite_result.command_metrics.items()
+            }
+        )
+    return attributes_dict
+
+
 def extract_attrs_for_run_step(
     step: FlowStep,
     flow: Flow,
@@ -578,6 +630,7 @@ def extract_attrs_for_run_step(
     tracker: DialogueStateTracker,
     available_actions: List[str],
     flows: FlowsList,
+    previous_step_id: Text,
 ) -> Dict[str, Any]:
     current_context = extract_current_context_attribute(stack)
 
@@ -586,6 +639,7 @@ def extract_attrs_for_run_step(
         "step_description": step.description if step.description else "None",
         "current_flow_id": flow.id,
         "current_context": json.dumps(current_context),
+        "previous_step_id": previous_step_id,
     }
 
 

rasa/tracing/instrumentation/instrumentation.py

@@ -24,7 +24,7 @@ from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.trace import SpanKind, Tracer
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 
-from rasa.core.actions.action import Action, RemoteAction, CustomActionExecutor
+from rasa.core.actions.action import Action, CustomActionExecutor, RemoteAction
 from rasa.core.actions.custom_action_executor import RetryCustomActionExecutor
 from rasa.core.actions.grpc_custom_action_executor import GRPCCustomActionExecutor
 from rasa.core.agent import Agent
@@ -63,13 +63,13 @@ from rasa.tracing.instrumentation.intentless_policy_instrumentation import (
     _instrument_select_response_examples,
 )
 from rasa.tracing.instrumentation.metrics import (
+    record_callable_duration_metrics,
     record_llm_command_generator_metrics,
-    record_single_step_llm_command_generator_metrics,
     record_multi_step_llm_command_generator_metrics,
-    record_callable_duration_metrics,
     record_request_size_in_bytes,
+    record_single_step_llm_command_generator_metrics,
 )
-from rasa.utils.endpoints import concat_url, EndpointConfig
+from rasa.utils.endpoints import EndpointConfig, concat_url
 
 # The `TypeVar` representing the return type for a function to be wrapped.
 S = TypeVar("S")
@@ -84,6 +84,7 @@ COMMAND_PROCESSOR_MODULE_NAME = (
     "rasa.dialogue_understanding.processor.command_processor"
 )
 FLOW_EXECUTOR_MODULE_NAME = "rasa.core.policies.flows.flow_executor"
+DIALOG_UNDERSTANDING_TEST_IO_MODULE_NAME = "rasa.dialogue_understanding_test.io"
 
 
 def _check_extractor_argument_list(
@@ -151,7 +152,11 @@ def traceable(
         )
 
         module_name = attrs.pop("module_name", "")
-        if module_name in ["command_processor", FLOW_EXECUTOR_MODULE_NAME]:
+        if module_name in [
+            "command_processor",
+            FLOW_EXECUTOR_MODULE_NAME,
+            DIALOG_UNDERSTANDING_TEST_IO_MODULE_NAME,
+        ]:
             span_name = f"{module_name}.{fn.__name__}"
         else:
             span_name = f"{self.__class__.__name__}.{fn.__name__}"
@@ -583,6 +588,9 @@ def instrument(
     if not module_is_instrumented(FLOW_EXECUTOR_MODULE_NAME):
         _instrument_flow_executor_module(tracer_provider)
 
+    if not module_is_instrumented(DIALOG_UNDERSTANDING_TEST_IO_MODULE_NAME):
+        _instrument_dialog_understanding_test_io_module(tracer_provider)
+
     if policy_subclasses:
         for policy_subclass in policy_subclasses:
             if policy_subclass is not None and not class_is_instrumented(
@@ -974,6 +982,18 @@ def _instrument_flow_executor_module(tracer_provider: TracerProvider) -> None:
     mark_module_as_instrumented(FLOW_EXECUTOR_MODULE_NAME)
 
 
+def _instrument_dialog_understanding_test_io_module(
+    tracer_provider: TracerProvider,
+) -> None:
+    _instrument_function(
+        tracer_provider.get_tracer(DIALOG_UNDERSTANDING_TEST_IO_MODULE_NAME),
+        DIALOG_UNDERSTANDING_TEST_IO_MODULE_NAME,
+        "print_test_results",
+        attribute_extractors.extract_attrs_for_du_print_test_results,
+    )
+    mark_module_as_instrumented(DIALOG_UNDERSTANDING_TEST_IO_MODULE_NAME)
+
+
 def _instrument_advance_flows_until_next_action(
     tracer: Tracer,
     module_name: str,

rasa/tracing/instrumentation/intentless_policy_instrumentation.py

@@ -1,7 +1,7 @@
 import functools
 import json
 import logging
-from typing import Callable, List, Optional, TYPE_CHECKING, Type
+from typing import TYPE_CHECKING, Callable, List, Optional, Type
 
 from opentelemetry.trace import Tracer
 

rasa/tracing/instrumentation/metrics.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any
+from typing import Any, Dict
 
 import psutil
 
@@ -7,29 +7,29 @@ from rasa.core.policies.enterprise_search_policy import EnterpriseSearchPolicy
 from rasa.core.policies.intentless_policy import IntentlessPolicy
 from rasa.dialogue_understanding.generator import (
     LLMCommandGenerator,
-    SingleStepLLMCommandGenerator,
     MultiStepLLMCommandGenerator,
+    SingleStepLLMCommandGenerator,
 )
 from rasa.tracing.constants import (
+    CONTEXTUAL_RESPONSE_REPHRASER_LLM_RESPONSE_DURATION_METRIC_NAME,
+    ENTERPRISE_SEARCH_POLICY_LLM_RESPONSE_DURATION_METRIC_NAME,
+    INTENTLESS_POLICY_LLM_RESPONSE_DURATION_METRIC_NAME,
     LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME,
+    LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
     LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
     LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
-    SINGLE_STEP_LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME,
-    SINGLE_STEP_LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
-    SINGLE_STEP_LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
-    SINGLE_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
     MULTI_STEP_LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME,
+    MULTI_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
     MULTI_STEP_LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
     MULTI_STEP_LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
-    MULTI_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
     PROMPT_TOKEN_LENGTH_ATTRIBUTE_NAME,
-    LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
-    ENTERPRISE_SEARCH_POLICY_LLM_RESPONSE_DURATION_METRIC_NAME,
-    INTENTLESS_POLICY_LLM_RESPONSE_DURATION_METRIC_NAME,
-    CONTEXTUAL_RESPONSE_REPHRASER_LLM_RESPONSE_DURATION_METRIC_NAME,
     RASA_CLIENT_REQUEST_BODY_SIZE_METRIC_NAME,
     RASA_CLIENT_REQUEST_DURATION_METRIC_NAME,
     REQUEST_BODY_SIZE_IN_BYTES_ATTRIBUTE_NAME,
+    SINGLE_STEP_LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME,
+    SINGLE_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
+    SINGLE_STEP_LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
+    SINGLE_STEP_LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
 )
 from rasa.tracing.metric_instrument_provider import MetricInstrumentProvider
 from rasa.utils.endpoints import EndpointConfig

rasa/tracing/metric_instrument_provider.py

@@ -1,30 +1,30 @@
-from typing import ClassVar, Dict, Any
+from typing import Any, ClassVar, Dict
 
 from opentelemetry.metrics import get_meter_provider
 from opentelemetry.sdk.metrics import Meter
 
-from rasa.utils.singleton import Singleton
 from rasa.tracing.constants import (
+    CONTEXTUAL_RESPONSE_REPHRASER_LLM_RESPONSE_DURATION_METRIC_NAME,
+    DURATION_UNIT_NAME,
+    ENTERPRISE_SEARCH_POLICY_LLM_RESPONSE_DURATION_METRIC_NAME,
+    INTENTLESS_POLICY_LLM_RESPONSE_DURATION_METRIC_NAME,
+    LLM_BASED_COMMAND_GENERATOR_CPU_MEMORY_USAGE_UNIT_NAME,
     LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME,
+    LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
     LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
     LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
-    LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
-    LLM_BASED_COMMAND_GENERATOR_CPU_MEMORY_USAGE_UNIT_NAME,
-    SINGLE_STEP_LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME,
-    SINGLE_STEP_LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
-    SINGLE_STEP_LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
-    SINGLE_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
     MULTI_STEP_LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME,
+    MULTI_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
     MULTI_STEP_LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
     MULTI_STEP_LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
-    MULTI_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
-    DURATION_UNIT_NAME,
-    ENTERPRISE_SEARCH_POLICY_LLM_RESPONSE_DURATION_METRIC_NAME,
-    INTENTLESS_POLICY_LLM_RESPONSE_DURATION_METRIC_NAME,
-    CONTEXTUAL_RESPONSE_REPHRASER_LLM_RESPONSE_DURATION_METRIC_NAME,
-    RASA_CLIENT_REQUEST_DURATION_METRIC_NAME,
     RASA_CLIENT_REQUEST_BODY_SIZE_METRIC_NAME,
+    RASA_CLIENT_REQUEST_DURATION_METRIC_NAME,
+    SINGLE_STEP_LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME,
+    SINGLE_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
+    SINGLE_STEP_LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
+    SINGLE_STEP_LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
 )
+from rasa.utils.singleton import Singleton
 
 
 class MetricInstrumentProvider(metaclass=Singleton):

rasa/utils/common.py

@@ -8,38 +8,39 @@ import shutil
 import tempfile
 import warnings
 from pathlib import Path
+from socket import SOCK_DGRAM, SOCK_STREAM
 from types import TracebackType
 from typing import (
     Any,
+    ContextManager,
     Coroutine,
     Dict,
     List,
     Optional,
+    Set,
     Text,
+    Tuple,
     Type,
     TypeVar,
     Union,
-    ContextManager,
-    Set,
-    Tuple,
 )
 
-from socket import SOCK_DGRAM, SOCK_STREAM
 import numpy as np
+
+import rasa.shared.utils.io
 import rasa.utils.io
 from rasa.anonymization import ENV_LOG_LEVEL_FAKER, ENV_LOG_LEVEL_PRESIDIO
 from rasa.constants import (
     DEFAULT_LOG_LEVEL_LIBRARIES,
+    ENV_LOG_LEVEL_KAFKA,
     ENV_LOG_LEVEL_LIBRARIES,
     ENV_LOG_LEVEL_MATPLOTLIB,
     ENV_LOG_LEVEL_MLFLOW,
     ENV_LOG_LEVEL_RABBITMQ,
-    ENV_LOG_LEVEL_KAFKA,
 )
 from rasa.shared.constants import DEFAULT_LOG_LEVEL, ENV_LOG_LEVEL, TCP_PROTOCOL
 from rasa.shared.exceptions import RasaException
-import rasa.shared.utils.io
-from rasa.shared.utils.yaml import read_yaml_file, write_yaml, read_config_file
+from rasa.shared.utils.yaml import read_config_file, read_yaml_file, write_yaml
 
 logger = logging.getLogger(__name__)
 
@@ -146,7 +147,9 @@ def read_global_config(path: Text) -> Dict[Text, Any]:
     """
     # noinspection PyBroadException
     try:
-        return read_config_file(path)
+        # set expand_env_vars to False because this is a global config file
+        # used for telemetry, in which case we don't need to expand env vars
+        return read_config_file(path, expand_env_vars=False)
     except Exception:
         # if things go south we pretend there is no config
         return {}
@@ -299,7 +302,7 @@ def update_sanic_log_level(
     syslog_protocol: Optional[Text] = None,
 ) -> None:
     """Set the log level to 'LOG_LEVEL_LIBRARIES' environment variable ."""
-    from sanic.log import logger, error_logger, access_logger
+    from sanic.log import access_logger, error_logger, logger
 
     log_level = os.environ.get(ENV_LOG_LEVEL_LIBRARIES, DEFAULT_LOG_LEVEL_LIBRARIES)
 

rasa/utils/endpoints.py

@@ -1,7 +1,7 @@
 import os
 import ssl
 from types import ModuleType
-from typing import Any, List, Optional, Text, Dict, Union
+from typing import Any, Dict, List, Optional, Text, Union
 
 import aiohttp
 import structlog

rasa/utils/io.py

@@ -2,22 +2,22 @@ import asyncio
 import filecmp
 import logging
 import os
+import re
 import tempfile
 import warnings
-import re
 from asyncio import AbstractEventLoop
 from collections import OrderedDict
 from io import StringIO
 from pathlib import Path
 from typing import (
-    Optional,
-    Text,
+    TYPE_CHECKING,
     Any,
-    Union,
+    Callable,
     List,
+    Optional,
+    Text,
     Type,
-    Callable,
-    TYPE_CHECKING,
+    Union,
 )
 
 from ruamel import yaml
@@ -153,8 +153,8 @@ def create_validator(
     """Helper method to create `Validator` classes from callable functions. Should be
     removed when questionary supports `Validator` objects.
     """
-    from prompt_toolkit.validation import Validator, ValidationError
     from prompt_toolkit.document import Document
+    from prompt_toolkit.validation import ValidationError, Validator
 
     class FunctionValidator(Validator):
         @staticmethod

rasa/utils/licensing.py

@@ -1,4 +1,3 @@
-from asyncio import AbstractEventLoop
 import hashlib
 import os
 import random
@@ -6,19 +5,19 @@ import re
 import time
 import typing
 import uuid
+from asyncio import AbstractEventLoop
 from datetime import datetime, timezone
 from typing import Any, Callable, Dict, Optional, Set, Text, TypeVar
 
 import jwt
+import structlog
 from dotenv import dotenv_values
 from sanic import Sanic
-import structlog
-from rasa import telemetry
 
+from rasa import telemetry
 from rasa.core import jobs
 from rasa.shared.utils.cli import print_error_and_exit
 
-
 if typing.TYPE_CHECKING:
     from rasa.core.tracker_store import TrackerStore
 

rasa/utils/log_utils.py

@@ -1,22 +1,23 @@
 from __future__ import annotations
-import os
+
 import logging
+import os
 import sys
 from typing import Any, Dict, Optional
 
 import structlog
-from structlog_sentry import SentryProcessor
 from structlog.dev import ConsoleRenderer
 from structlog.typing import EventDict, WrappedLogger
+from structlog_sentry import SentryProcessor
+
+from rasa.plugin import plugin_manager
 from rasa.shared.constants import (
-    ENV_LOG_LEVEL,
     DEFAULT_LOG_LEVEL,
+    DEFAULT_LOG_LEVEL_LLM,
+    ENV_LOG_LEVEL,
     ENV_LOG_LEVEL_LLM,
     ENV_LOG_LEVEL_LLM_MODULE_NAMES,
-    DEFAULT_LOG_LEVEL_LLM,
 )
-from rasa.plugin import plugin_manager
-
 
 FORCE_JSON_LOGGING = os.environ.get("FORCE_JSON_LOGGING")
 

rasa/utils/ml_utils.py

@@ -5,6 +5,7 @@ import numpy as np
 import structlog
 from langchain.schema.embeddings import Embeddings
 from langchain_community.vectorstores import FAISS
+
 from rasa.shared.constants import REQUIRED_SLOTS_KEY
 from rasa.shared.core.domain import KEY_RESPONSES_TEXT, Domain
 from rasa.shared.utils.llm import AI

rasa/utils/plotting.py

@@ -1,11 +1,11 @@
-import logging
 import itertools
+import logging
 import os
 from functools import wraps
+from typing import Any, Callable, List, Optional, Text, Tuple, TypeVar, Union
 
-import numpy as np
-from typing import Any, Callable, List, Optional, Text, TypeVar, Union, Tuple
 import matplotlib
+import numpy as np
 from matplotlib.ticker import FormatStrFormatter
 
 import rasa.shared.utils.io

rasa/utils/sanic_error_handler.py

@@ -1,7 +1,7 @@
 from sanic import Sanic
+from sanic.exceptions import ServerError
 from sanic.handlers import ErrorHandler
 from sanic.request import Request
-from sanic.exceptions import ServerError
 
 
 # TODO: remove custom handler when upgrading to sanic >= 24

rasa/utils/tensorflow/callback.py

@@ -1,7 +1,7 @@
+import logging
 from pathlib import Path
-from typing import Dict, Text, Any, Optional
+from typing import Any, Dict, Optional, Text
 
-import logging
 import tensorflow as tf
 from tqdm import tqdm
 

rasa/utils/tensorflow/crf.py

@@ -1,8 +1,8 @@
+from typing import Any, List, Optional, Tuple, Union
+
 import tensorflow as tf
 from tensorflow import TensorShape
 from tensorflow.types.experimental import TensorLike
-from typing import Tuple, Any, List, Union, Optional
-
 
 # original code taken from
 # https://github.com/tensorflow/addons/blob/b8cab7fd61af4f697a1cdae4f51c37c346b9c6f0/tensorflow_addons/text/crf.py

rasa/utils/tensorflow/data_generator.py

@@ -1,13 +1,13 @@
+import logging
 import math
-from typing import List, Union, Text, Optional, Any, Tuple, Dict, cast
+from typing import Any, Dict, List, Optional, Text, Tuple, Union, cast
 
-import logging
-import scipy.sparse
 import numpy as np
+import scipy.sparse
 from tensorflow.keras.utils import Sequence
 
-from rasa.utils.tensorflow.constants import SEQUENCE, BALANCED
-from rasa.utils.tensorflow.model_data import RasaModelData, Data, FeatureArray
+from rasa.utils.tensorflow.constants import BALANCED, SEQUENCE
+from rasa.utils.tensorflow.model_data import Data, FeatureArray, RasaModelData
 
 logger = logging.getLogger(__name__)
 

rasa/utils/tensorflow/environment.py

@@ -1,16 +1,16 @@
 import logging
 import os
-from typing import Text, Dict
 import typing
+from typing import Dict, Text
 
 import rasa.shared.utils.io
 from rasa.constants import (
-    ENV_GPU_CONFIG,
     ENV_CPU_INTER_OP_CONFIG,
     ENV_CPU_INTRA_OP_CONFIG,
+    ENV_GPU_CONFIG,
 )
-from rasa.utils.tensorflow.constants import TF_DETERMINISTIC_OPS
 from rasa.shared.utils import io as shared_io_utils
+from rasa.utils.tensorflow.constants import TF_DETERMINISTIC_OPS
 
 if typing.TYPE_CHECKING:
     from tensorflow import config as tf_config

rasa/utils/tensorflow/feature_array.py

@@ -1,9 +1,8 @@
-from typing import Dict, Any, List, Tuple, Optional, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
-from safetensors.numpy import save_file
 import numpy as np
-from safetensors.numpy import load_file
 import scipy.sparse
+from safetensors.numpy import load_file, save_file
 
 import rasa.shared.utils.io
 

rasa/utils/tensorflow/layers.py

@@ -1,29 +1,35 @@
 import logging
-from typing import List, Optional, Text, Tuple, Callable, Union, Any
+from typing import Any, Callable, List, Optional, Text, Tuple, Union
+
 import tensorflow as tf
+import tensorflow.keras.backend as K
 
 # TODO: The following is not (yet) available via tf.keras
 from keras.src.utils.control_flow_util import smart_cond
-import tensorflow.keras.backend as K
 
 import rasa.utils.tensorflow.crf
+import rasa.utils.tensorflow.layers_utils as layers_utils
+from rasa.core.constants import DIALOGUE
+from rasa.shared.nlu.constants import (
+    ACTION_NAME,
+    ACTION_TEXT,
+    FEATURE_TYPE_SENTENCE,
+    FEATURE_TYPE_SEQUENCE,
+    INTENT,
+    TEXT,
+)
 from rasa.utils.tensorflow.constants import (
-    SOFTMAX,
-    MARGIN,
     COSINE,
-    INNER,
     CROSS_ENTROPY,
+    INNER,
     LABEL,
     LABEL_PAD_ID,
+    MARGIN,
+    SOFTMAX,
 )
-from rasa.core.constants import DIALOGUE
-from rasa.shared.nlu.constants import FEATURE_TYPE_SENTENCE, FEATURE_TYPE_SEQUENCE
-from rasa.shared.nlu.constants import TEXT, INTENT, ACTION_NAME, ACTION_TEXT
-
-from rasa.utils.tensorflow.metrics import F1Score
-from rasa.utils.tensorflow.exceptions import TFLayerConfigException
-import rasa.utils.tensorflow.layers_utils as layers_utils
 from rasa.utils.tensorflow.crf import crf_log_likelihood
+from rasa.utils.tensorflow.exceptions import TFLayerConfigException
+from rasa.utils.tensorflow.metrics import F1Score
 
 logger = logging.getLogger(__name__)
 

rasa/utils/tensorflow/layers_utils.py

@@ -1,6 +1,7 @@
+from typing import Optional, Union
+
 import tensorflow as tf
 from tensorflow import Tensor
-from typing import Union, Optional
 
 
 def random_indices(

rasa/utils/tensorflow/metrics.py

@@ -1,8 +1,8 @@
+from typing import Any, Dict, Optional
+
 import tensorflow as tf
 from tensorflow.keras import backend as K
 from tensorflow.types.experimental import TensorLike
-from typing import Any, Dict, Optional
-
 
 # original code taken from
 # https://github.com/tensorflow/addons/blob/f30df4322b5580b3e5946530a60f7126035dd73b/tensorflow_addons/metrics/f_scores.py

rasa/utils/tensorflow/model_data.py

@@ -1,20 +1,20 @@
 import logging
+from collections import OrderedDict, defaultdict
 from typing import (
-    Optional,
+    Any,
     DefaultDict,
     Dict,
+    ItemsView,
     Iterable,
-    Text,
     List,
+    NamedTuple,
+    Optional,
+    Text,
     Tuple,
-    Any,
     Union,
-    NamedTuple,
-    ItemsView,
-    overload,
     cast,
+    overload,
 )
-from collections import defaultdict, OrderedDict
 
 import numpy as np
 import scipy.sparse