rasa-pro 3.8.18__py3-none-any.whl → 3.9.14__py3-none-any.whl

rasa/tracing/instrumentation/metrics.py

@@ -5,11 +5,23 @@ import psutil
 from rasa.core.nlg.contextual_response_rephraser import ContextualResponseRephraser
 from rasa.core.policies.enterprise_search_policy import EnterpriseSearchPolicy
 from rasa.core.policies.intentless_policy import IntentlessPolicy
-from rasa.dialogue_understanding.generator import LLMCommandGenerator
+from rasa.dialogue_understanding.generator import (
+    LLMCommandGenerator,
+    SingleStepLLMCommandGenerator,
+    MultiStepLLMCommandGenerator,
+)
 from rasa.tracing.constants import (
     LLM_COMMAND_GENERATOR_CPU_USAGE_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_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,
@@ -17,14 +29,15 @@ from rasa.tracing.constants import (
     CONTEXTUAL_RESPONSE_REPHRASER_LLM_RESPONSE_DURATION_METRIC_NAME,
     RASA_CLIENT_REQUEST_BODY_SIZE_METRIC_NAME,
     RASA_CLIENT_REQUEST_DURATION_METRIC_NAME,
-    ENDPOINT_REQUEST_BODY_SIZE_IN_BYTES_ATTRIBUTE_NAME,
+    REQUEST_BODY_SIZE_IN_BYTES_ATTRIBUTE_NAME,
 )
 from rasa.tracing.metric_instrument_provider import MetricInstrumentProvider
 from rasa.utils.endpoints import EndpointConfig
 
 
-def record_llm_command_generator_cpu_usage(
+def record_llm_based_command_generator_cpu_usage(
     metric_instrument_provider: MetricInstrumentProvider,
+    metric_name: str,
 ) -> None:
     """Record CPU usage as a percentage.
 
@@ -32,11 +45,10 @@ def record_llm_command_generator_cpu_usage(
     This instrument is registered to the MetricInstrumentProvider internal singleton.
 
     :param metric_instrument_provider: The MetricInstrumentProvider instance
+    :param metric_name: The name of the metric instrument
     :return: None
     """
-    metric_instrument = metric_instrument_provider.get_instrument(
-        LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME
-    )
+    metric_instrument = metric_instrument_provider.get_instrument(metric_name)
     if not metric_instrument:
         return None
 
@@ -44,8 +56,9 @@ def record_llm_command_generator_cpu_usage(
     metric_instrument.record(amount=cpu_usage)
 
 
-def record_llm_command_generator_memory_usage(
+def record_llm_based_command_generator_memory_usage(
     metric_instrument_provider: MetricInstrumentProvider,
+    metric_name: str,
 ) -> None:
     """Record memory usage as a percentage.
 
@@ -53,11 +66,10 @@ def record_llm_command_generator_memory_usage(
     This instrument is registered to the MetricInstrumentProvider internal singleton.
 
     :param metric_instrument_provider: The MetricInstrumentProvider instance
+    :param metric_name: The name of the metric instrument
     :return: None
     """
-    metric_instrument = metric_instrument_provider.get_instrument(
-        LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME
-    )
+    metric_instrument = metric_instrument_provider.get_instrument(metric_name)
     if not metric_instrument:
         return None
 
@@ -65,9 +77,10 @@ def record_llm_command_generator_memory_usage(
     metric_instrument.record(amount=memory_usage)
 
 
-def record_llm_command_generator_prompt_token(
+def record_llm_based_command_generator_prompt_token(
     metric_instrument_provider: MetricInstrumentProvider,
     attributes: Dict[str, Any],
+    metric_name: str,
 ) -> None:
     """
     Record prompt token length.
@@ -77,14 +90,13 @@ def record_llm_command_generator_prompt_token(
 
     :param metric_instrument_provider: The MetricInstrumentProvider instance
     :param attributes: Extracted tracing attributes
+    :param metric_name: The name of the metric instrument
     :return: None
     """
     if PROMPT_TOKEN_LENGTH_ATTRIBUTE_NAME not in attributes:
         return None
 
-    metric_instrument = metric_instrument_provider.get_instrument(
-        LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME
-    )
+    metric_instrument = metric_instrument_provider.get_instrument(metric_name)
     if not metric_instrument:
         return None
 
@@ -115,9 +127,75 @@ def record_llm_command_generator_metrics(attributes: Dict[str, Any]) -> None:
     if not instrument_provider.instruments:
         return None
 
-    record_llm_command_generator_cpu_usage(instrument_provider)
-    record_llm_command_generator_memory_usage(instrument_provider)
-    record_llm_command_generator_prompt_token(instrument_provider, attributes)
+    record_llm_based_command_generator_cpu_usage(
+        instrument_provider, LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME
+    )
+    record_llm_based_command_generator_memory_usage(
+        instrument_provider, LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME
+    )
+    record_llm_based_command_generator_prompt_token(
+        instrument_provider,
+        attributes,
+        LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
+    )
+
+
+def record_single_step_llm_command_generator_metrics(
+    attributes: Dict[str, Any],
+) -> None:
+    """
+    Record measurements for SingleStepLLMCommandGenerator specific metrics.
+
+    The recording is done by the opentelemetry.metrics.Histogram instruments.
+    These instruments are registered to the MetricInstrumentProvider internal singleton.
+
+    :param attributes: Extracted tracing attributes
+    :return: None
+    """
+    instrument_provider = MetricInstrumentProvider()
+
+    if not instrument_provider.instruments:
+        return None
+
+    record_llm_based_command_generator_cpu_usage(
+        instrument_provider, SINGLE_STEP_LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME
+    )
+    record_llm_based_command_generator_memory_usage(
+        instrument_provider, SINGLE_STEP_LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME
+    )
+    record_llm_based_command_generator_prompt_token(
+        instrument_provider,
+        attributes,
+        SINGLE_STEP_LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
+    )
+
+
+def record_multi_step_llm_command_generator_metrics(attributes: Dict[str, Any]) -> None:
+    """
+    Record measurements for MultiStepLLMCommandGenerator specific metrics.
+
+    The recording is done by the opentelemetry.metrics.Histogram instruments.
+    These instruments are registered to the MetricInstrumentProvider internal singleton.
+
+    :param attributes: Extracted tracing attributes
+    :return: None
+    """
+    instrument_provider = MetricInstrumentProvider()
+
+    if not instrument_provider.instruments:
+        return None
+
+    record_llm_based_command_generator_cpu_usage(
+        instrument_provider, MULTI_STEP_LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME
+    )
+    record_llm_based_command_generator_memory_usage(
+        instrument_provider, MULTI_STEP_LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME
+    )
+    record_llm_based_command_generator_prompt_token(
+        instrument_provider,
+        attributes,
+        MULTI_STEP_LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
+    )
 
 
 def record_callable_duration_metrics(
@@ -126,6 +204,8 @@ def record_callable_duration_metrics(
     """
     Record duration of instrumented method calls invoked for the following components:
     - LLMCommandGenerator
+    - SingleStepLLMCommandGenerator
+    - MultiStepLLMCommandGenerator
     - EnterpriseSearchPolicy
     - IntentlessPolicy
     - ContextualResponseRephraser
@@ -145,11 +225,21 @@ def record_callable_duration_metrics(
     metric_instrument = None
     attributes = {}
 
-    if isinstance(self, LLMCommandGenerator):
+    if type(self) == LLMCommandGenerator:
         metric_instrument = instrument_provider.get_instrument(
             LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME
         )
 
+    if type(self) == SingleStepLLMCommandGenerator:
+        metric_instrument = instrument_provider.get_instrument(
+            SINGLE_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME
+        )
+
+    if type(self) == MultiStepLLMCommandGenerator:
+        metric_instrument = instrument_provider.get_instrument(
+            MULTI_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME
+        )
+
     if isinstance(self, EnterpriseSearchPolicy):
         metric_instrument = instrument_provider.get_instrument(
             ENTERPRISE_SEARCH_POLICY_LLM_RESPONSE_DURATION_METRIC_NAME
@@ -200,7 +290,5 @@ def record_request_size_in_bytes(attributes: Dict[str, Any]) -> None:
     if not metric_instrument:
         return None
 
-    request_body_size = attributes.pop(
-        ENDPOINT_REQUEST_BODY_SIZE_IN_BYTES_ATTRIBUTE_NAME, 0
-    )
+    request_body_size = attributes.pop(REQUEST_BODY_SIZE_IN_BYTES_ATTRIBUTE_NAME, 0)
     metric_instrument.record(amount=request_body_size, attributes=attributes)

rasa/tracing/metric_instrument_provider.py

@@ -9,7 +9,15 @@ from rasa.tracing.constants import (
     LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
     LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
     LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
-    LLM_COMMAND_GENERATOR_CPU_MEMORY_USAGE_UNIT_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_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,
@@ -34,6 +42,8 @@ class MetricInstrumentProvider(metaclass=Singleton):
 
         instruments = {
             **self._create_llm_command_generator_instruments(meter),
+            **self._create_single_step_llm_command_generator_instruments(meter),
+            **self._create_multi_step_llm_command_generator_instruments(meter),
             **self._create_llm_response_duration_instruments(meter),
             **self._create_client_request_instruments(meter),
         }
@@ -49,13 +59,13 @@ class MetricInstrumentProvider(metaclass=Singleton):
         llm_command_generator_cpu_usage = meter.create_histogram(
             name=LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME,
             description="CPU percentage for LLMCommandGenerator",
-            unit=LLM_COMMAND_GENERATOR_CPU_MEMORY_USAGE_UNIT_NAME,
+            unit=LLM_BASED_COMMAND_GENERATOR_CPU_MEMORY_USAGE_UNIT_NAME,
         )
 
         llm_command_generator_memory_usage = meter.create_histogram(
             name=LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
             description="RAM memory usage for LLMCommandGenerator",
-            unit=LLM_COMMAND_GENERATOR_CPU_MEMORY_USAGE_UNIT_NAME,
+            unit=LLM_BASED_COMMAND_GENERATOR_CPU_MEMORY_USAGE_UNIT_NAME,
         )
 
         llm_command_generator_prompt_token_usage = meter.create_histogram(
@@ -77,6 +87,76 @@ class MetricInstrumentProvider(metaclass=Singleton):
             LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME: llm_response_duration_llm_command_generator,  # noqa: E501
         }
 
+    @staticmethod
+    def _create_single_step_llm_command_generator_instruments(
+        meter: Meter,
+    ) -> Dict[str, Any]:
+        single_step_llm_command_generator_cpu_usage = meter.create_histogram(
+            name=SINGLE_STEP_LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME,
+            description="CPU percentage for SingleStepLLMCommandGenerator",
+            unit=LLM_BASED_COMMAND_GENERATOR_CPU_MEMORY_USAGE_UNIT_NAME,
+        )
+
+        single_step_llm_command_generator_memory_usage = meter.create_histogram(
+            name=SINGLE_STEP_LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
+            description="RAM memory usage for SingleStepLLMCommandGenerator",
+            unit=LLM_BASED_COMMAND_GENERATOR_CPU_MEMORY_USAGE_UNIT_NAME,
+        )
+
+        single_step_llm_command_generator_prompt_token_usage = meter.create_histogram(
+            name=SINGLE_STEP_LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
+            description="SingleStepLLMCommandGenerator prompt token length",
+            unit="1",
+        )
+
+        single_step_llm_response_duration_llm_command_generator = meter.create_histogram(  # noqa: E501
+            name=SINGLE_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
+            description="The duration of SingleStepLLMCommandGenerator's LLM call",
+            unit=DURATION_UNIT_NAME,
+        )
+
+        return {
+            SINGLE_STEP_LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME: single_step_llm_command_generator_cpu_usage,  # noqa: E501
+            SINGLE_STEP_LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME: single_step_llm_command_generator_memory_usage,  # noqa: E501
+            SINGLE_STEP_LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME: single_step_llm_command_generator_prompt_token_usage,  # noqa: E501
+            SINGLE_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME: single_step_llm_response_duration_llm_command_generator,  # noqa: E501
+        }
+
+    @staticmethod
+    def _create_multi_step_llm_command_generator_instruments(
+        meter: Meter,
+    ) -> Dict[str, Any]:
+        multi_step_llm_command_generator_cpu_usage = meter.create_histogram(
+            name=MULTI_STEP_LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME,
+            description="CPU percentage for MultiStepLLMCommandGenerator",
+            unit=LLM_BASED_COMMAND_GENERATOR_CPU_MEMORY_USAGE_UNIT_NAME,
+        )
+
+        multi_step_llm_command_generator_memory_usage = meter.create_histogram(
+            name=MULTI_STEP_LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME,
+            description="RAM memory usage for MultiStepLLMCommandGenerator",
+            unit=LLM_BASED_COMMAND_GENERATOR_CPU_MEMORY_USAGE_UNIT_NAME,
+        )
+
+        multi_step_llm_command_generator_prompt_token_usage = meter.create_histogram(
+            name=MULTI_STEP_LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME,
+            description="MultiStepLLMCommandGenerator prompt token length",
+            unit="1",
+        )
+
+        multi_step_llm_response_duration_llm_command_generator = meter.create_histogram(
+            name=MULTI_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME,
+            description="The duration of MultiStepLLMCommandGenerator's LLM call",
+            unit=DURATION_UNIT_NAME,
+        )
+
+        return {
+            MULTI_STEP_LLM_COMMAND_GENERATOR_CPU_USAGE_METRIC_NAME: multi_step_llm_command_generator_cpu_usage,  # noqa: E501
+            MULTI_STEP_LLM_COMMAND_GENERATOR_MEMORY_USAGE_METRIC_NAME: multi_step_llm_command_generator_memory_usage,  # noqa: E501
+            MULTI_STEP_LLM_COMMAND_GENERATOR_PROMPT_TOKEN_USAGE_METRIC_NAME: multi_step_llm_command_generator_prompt_token_usage,  # noqa: E501
+            MULTI_STEP_LLM_COMMAND_GENERATOR_LLM_RESPONSE_DURATION_METRIC_NAME: multi_step_llm_response_duration_llm_command_generator,  # noqa: E501
+        }
+
     @staticmethod
     def _create_llm_response_duration_instruments(meter: Meter) -> Dict[str, Any]:
         llm_response_duration_enterprise_search = meter.create_histogram(

rasa/utils/cli.py

@@ -20,7 +20,8 @@ def remove_argument_from_parser(
         for option in action.option_strings:
             if option == argument:
                 parser._handle_conflict_resolve(
-                    None, [(option, action)]  # type: ignore[arg-type]
+                    None,
+                    [(option, action)],
                 )
                 return
 

rasa/utils/common.py

@@ -520,7 +520,7 @@ class RepeatedLogFilter(logging.Filter):
 
 
 async def call_potential_coroutine(
-    coroutine_or_return_value: Union[Any, Coroutine]
+    coroutine_or_return_value: Union[Any, Coroutine],
 ) -> Any:
     """Awaits coroutine or returns value directly if it's not a coroutine.
 

rasa/utils/endpoints.py

@@ -17,7 +17,7 @@ structlogger = structlog.get_logger()
 def read_endpoint_config(
     filename: Text, endpoint_type: Text
 ) -> Optional["EndpointConfig"]:
-    """Read an endpoint configuration file from disk and extract one config."""  # noqa: E501
+    """Read an endpoint configuration file from disk and extract one config."""
     if not filename:
         return None
 
@@ -298,6 +298,5 @@ def int_arg(
     try:
         return int(str(arg))
     except (ValueError, TypeError):
-
         structlogger.warning("endpoint.int_arg.convert_failed", arg=arg, key=key)
         return default

rasa/utils/io.py

@@ -2,6 +2,7 @@ import asyncio
 import filecmp
 import logging
 import os
+import pickle
 import tempfile
 import warnings
 import re
@@ -97,6 +98,29 @@ def enable_async_loop_debugging(
     return event_loop
 
 
+def pickle_dump(filename: Union[Text, Path], obj: Any) -> None:
+    """Saves object to file.
+
+    Args:
+        filename: the filename to save the object to
+        obj: the object to store
+    """
+    with open(filename, "wb") as f:
+        pickle.dump(obj, f)
+
+
+def pickle_load(filename: Union[Text, Path]) -> Any:
+    """Loads an object from a file.
+
+    Args:
+        filename: the filename to load the object from
+
+    Returns: the loaded object
+    """
+    with open(filename, "rb") as f:
+        return pickle.load(f)
+
+
 def create_temporary_file(data: Any, suffix: Text = "", mode: Text = "w+") -> Text:
     """Creates a tempfile.NamedTemporaryFile object for data."""
     encoding = None if "b" in mode else rasa.shared.utils.io.DEFAULT_ENCODING
@@ -167,16 +191,58 @@ def create_validator(
     return FunctionValidator
 
 
+def json_unpickle(
+    file_name: Union[Text, Path], encode_non_string_keys: bool = False
+) -> Any:
+    """Unpickle an object from file using json.
+
+    Args:
+        file_name: the file to load the object from
+        encode_non_string_keys: If set to `True` then jsonpickle will encode non-string
+          dictionary keys instead of coercing them into strings via `repr()`.
+
+    Returns: the object
+    """
+    import jsonpickle.ext.numpy as jsonpickle_numpy
+    import jsonpickle
+
+    jsonpickle_numpy.register_handlers()
+
+    file_content = rasa.shared.utils.io.read_file(file_name)
+    return jsonpickle.loads(file_content, keys=encode_non_string_keys)
+
+
+def json_pickle(
+    file_name: Union[Text, Path], obj: Any, encode_non_string_keys: bool = False
+) -> None:
+    """Pickle an object to a file using json.
+
+    Args:
+        file_name: the file to store the object to
+        obj: the object to store
+        encode_non_string_keys: If set to `True` then jsonpickle will encode non-string
+          dictionary keys instead of coercing them into strings via `repr()`.
+    """
+    import jsonpickle.ext.numpy as jsonpickle_numpy
+    import jsonpickle
+
+    jsonpickle_numpy.register_handlers()
+
+    rasa.shared.utils.io.write_text_file(
+        jsonpickle.dumps(obj, keys=encode_non_string_keys), file_name
+    )
+
+
 def get_emoji_regex() -> Pattern:
     """Returns regex to identify emojis."""
     return re.compile(
         "["
-        "\U0001F600-\U0001F64F"  # emoticons
-        "\U0001F300-\U0001F5FF"  # symbols & pictographs
-        "\U0001F680-\U0001F6FF"  # transport & map symbols
-        "\U0001F1E0-\U0001F1FF"  # flags (iOS)
-        "\U00002702-\U000027B0"
-        "\U000024C2-\U0001F251"
+        "\U0001f600-\U0001f64f"  # emoticons
+        "\U0001f300-\U0001f5ff"  # symbols & pictographs
+        "\U0001f680-\U0001f6ff"  # transport & map symbols
+        "\U0001f1e0-\U0001f1ff"  # flags (iOS)
+        "\U00002702-\U000027b0"
+        "\U000024c2-\U0001f251"
         "\u200d"  # zero width joiner
         "\u200c"  # zero width non-joiner
         "]+",