rasa-pro 3.10.15__py3-none-any.whl → 3.11.0__py3-none-any.whl

rasa/tracing/instrumentation/attribute_extractors.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Text, Tuple, Union
 
 import tiktoken
 from numpy import ndarray
+from rasa.dialogue_understanding.generator.constants import FLOW_RETRIEVAL_KEY
 from rasa_sdk.grpc_py import action_webhook_pb2
 
 from rasa.core.actions.action import DirectCustomActionExecutor
@@ -19,6 +20,7 @@ from rasa.core.processor import MessageProcessor
 from rasa.core.tracker_store import TrackerStore
 from rasa.dialogue_understanding.commands import Command
 from rasa.dialogue_understanding.stack.dialogue_stack import DialogueStack
+from rasa.dialogue_understanding.generator import LLMBasedCommandGenerator
 from rasa.engine.graph import ExecutionContext, GraphModelConfiguration, GraphNode
 from rasa.engine.training.graph_trainer import GraphTrainer
 from rasa.shared.constants import (
@@ -27,6 +29,8 @@ from rasa.shared.constants import (
     PROVIDER_CONFIG_KEY,
     TIMEOUT_CONFIG_KEY,
     DEPLOYMENT_CONFIG_KEY,
+    MODEL_GROUP_ID_CONFIG_KEY,
+    LLM_CONFIG_KEY,
 )
 from rasa.shared.core.constants import REQUESTED_SLOT
 from rasa.shared.core.domain import Domain
@@ -50,10 +54,7 @@ if TYPE_CHECKING:
     from rasa.core.policies.enterprise_search_policy import EnterpriseSearchPolicy
     from rasa.core.policies.intentless_policy import IntentlessPolicy
     from rasa.core.policies.policy import PolicyPrediction
-    from rasa.dialogue_understanding.generator import (
-        CommandGenerator,
-        LLMBasedCommandGenerator,
-    )
+    from rasa.dialogue_understanding.generator import CommandGenerator
 
 # This file contains all attribute extractors for tracing instrumentation.
 # These are functions that are applied to the arguments of the wrapped function to be
@@ -300,22 +301,49 @@ def extract_attrs_for_command(
     }
 
 
-def extract_llm_config(self: Any, default_llm_config: Dict[str, Any]) -> Dict[str, Any]:
+def extract_llm_config(
+    self: Any,
+    default_llm_config: Dict[str, Any],
+    default_embeddings_config: Dict[str, Any],
+) -> Dict[str, Any]:
     if isinstance(self, ContextualResponseRephraser):
-        config = self.nlg_endpoint.kwargs
+        # ContextualResponseRephraser is not a graph component, so it's
+        # not having a full config.
+        config = {"llm": self.llm_config}
     else:
         config = self.config
 
     llm_property = combine_custom_and_default_config(
-        config.get("llm"), default_llm_config
+        config.get(LLM_CONFIG_KEY), default_llm_config
     )
 
+    if isinstance(self, LLMBasedCommandGenerator):
+        flow_retrieval_config = config.get(FLOW_RETRIEVAL_KEY, {}) or {}
+        embeddings_property = combine_custom_and_default_config(
+            flow_retrieval_config.get(EMBEDDINGS_CONFIG_KEY),
+            default_embeddings_config,
+        )
+    else:
+        embeddings_property = combine_custom_and_default_config(
+            config.get(EMBEDDINGS_CONFIG_KEY), default_embeddings_config
+        )
+
     attributes = {
         "class_name": self.__class__.__name__,
+        # llm client attributes
         "llm_model": str(llm_property.get(MODEL_CONFIG_KEY)),
         "llm_type": str(llm_property.get(PROVIDER_CONFIG_KEY)),
-        "embeddings": json.dumps(config.get(EMBEDDINGS_CONFIG_KEY, {})),
+        "llm_model_group_id": str(llm_property.get(MODEL_GROUP_ID_CONFIG_KEY)),
         "llm_temperature": str(llm_property.get("temperature")),
+        "llm_request_timeout": str(llm_property.get(TIMEOUT_CONFIG_KEY)),
+        # embedding client attributes
+        "embeddings_model": str(embeddings_property.get(MODEL_CONFIG_KEY)),
+        "embeddings_type": str(embeddings_property.get(PROVIDER_CONFIG_KEY)),
+        "embeddings_model_group_id": str(
+            embeddings_property.get(MODEL_GROUP_ID_CONFIG_KEY)
+        ),
+        # TODO: Keeping this to avoid potential breaking changes
+        "embeddings": json.dumps(embeddings_property, sort_keys=True),
         "request_timeout": str(llm_property.get(TIMEOUT_CONFIG_KEY)),
     }
 
@@ -329,11 +357,16 @@ def extract_attrs_for_llm_based_command_generator(
     self: "LLMBasedCommandGenerator",
     prompt: str,
 ) -> Dict[str, Any]:
-    from rasa.dialogue_understanding.generator.constants import (
-        DEFAULT_LLM_CONFIG,
+    from rasa.dialogue_understanding.generator.constants import DEFAULT_LLM_CONFIG
+    from rasa.dialogue_understanding.generator.flow_retrieval import (
+        DEFAULT_EMBEDDINGS_CONFIG,
     )
 
-    attributes = extract_llm_config(self, default_llm_config=DEFAULT_LLM_CONFIG)
+    attributes = extract_llm_config(
+        self,
+        default_llm_config=DEFAULT_LLM_CONFIG,
+        default_embeddings_config=DEFAULT_EMBEDDINGS_CONFIG,
+    )
 
     return extend_attributes_with_prompt_tokens_length(self, attributes, prompt)
 
@@ -344,7 +377,12 @@ def extract_attrs_for_contextual_response_rephraser(
 ) -> Dict[str, Any]:
     from rasa.core.nlg.contextual_response_rephraser import DEFAULT_LLM_CONFIG
 
-    attributes = extract_llm_config(self, default_llm_config=DEFAULT_LLM_CONFIG)
+    attributes = extract_llm_config(
+        self,
+        default_llm_config=DEFAULT_LLM_CONFIG,
+        # rephraser is not using embeddings
+        default_embeddings_config={},
+    )
 
     return extend_attributes_with_prompt_tokens_length(self, attributes, prompt)
 
@@ -355,7 +393,12 @@ def extract_attrs_for_create_history(
 ) -> Dict[str, Any]:
     from rasa.core.nlg.contextual_response_rephraser import DEFAULT_LLM_CONFIG
 
-    return extract_llm_config(self, default_llm_config=DEFAULT_LLM_CONFIG)
+    return extract_llm_config(
+        self,
+        default_llm_config=DEFAULT_LLM_CONFIG,
+        # rephraser is not using embeddings
+        default_embeddings_config={},
+    )
 
 
 def extract_attrs_for_generate(
@@ -371,6 +414,26 @@ def extract_attrs_for_generate(
     }
 
 
+def extract_attrs_for_performing_health_check(
+    custom_config: Optional[Dict[str, Any]],
+    default_config: Dict[str, Any],
+    log_source_method: str,
+    log_source_component: str,
+) -> Dict[str, Any]:
+    from rasa.shared.utils.health_check.health_check import is_api_health_check_enabled
+
+    attrs = {
+        "api_health_check_enabled": is_api_health_check_enabled(),
+        "health_check_trigger_component": log_source_component,
+        "health_check_trigger_method": log_source_method,
+    }
+    if is_api_health_check_enabled():
+        attrs["config"] = json.dumps(
+            combine_custom_and_default_config(custom_config, default_config)
+        )
+    return attrs
+
+
 def extract_attrs_for_execute_commands(
     tracker: DialogueStateTracker,
     all_flows: FlowsList,
@@ -580,9 +643,16 @@ def extract_attrs_for_intentless_policy_find_closest_response(
 def extract_attrs_for_intentless_policy_generate_llm_answer(
     self: "IntentlessPolicy", llm: "BaseLLM", prompt: str
 ) -> Dict[str, Any]:
-    from rasa.core.policies.intentless_policy import DEFAULT_LLM_CONFIG
+    from rasa.core.policies.intentless_policy import (
+        DEFAULT_LLM_CONFIG,
+        DEFAULT_EMBEDDINGS_CONFIG,
+    )
 
-    attributes = extract_llm_config(self, default_llm_config=DEFAULT_LLM_CONFIG)
+    attributes = extract_llm_config(
+        self,
+        default_llm_config=DEFAULT_LLM_CONFIG,
+        default_embeddings_config=DEFAULT_EMBEDDINGS_CONFIG,
+    )
 
     return extend_attributes_with_prompt_tokens_length(self, attributes, prompt)
 
@@ -590,9 +660,16 @@ def extract_attrs_for_intentless_policy_generate_llm_answer(
 def extract_attrs_for_enterprise_search_generate_llm_answer(
     self: "EnterpriseSearchPolicy", llm: "BaseLLM", prompt: str
 ) -> Dict[str, Any]:
-    from rasa.core.policies.enterprise_search_policy import DEFAULT_LLM_CONFIG
+    from rasa.core.policies.enterprise_search_policy import (
+        DEFAULT_LLM_CONFIG,
+        DEFAULT_EMBEDDINGS_CONFIG,
+    )
 
-    attributes = extract_llm_config(self, default_llm_config=DEFAULT_LLM_CONFIG)
+    attributes = extract_llm_config(
+        self,
+        default_llm_config=DEFAULT_LLM_CONFIG,
+        default_embeddings_config=DEFAULT_EMBEDDINGS_CONFIG,
+    )
 
     return extend_attributes_with_prompt_tokens_length(self, attributes, prompt)
 
@@ -664,7 +741,7 @@ def extract_attrs_for_custom_action_executor_run(
 
     attrs: Dict[str, Any] = {
         "class_name": self.__class__.__name__,
-        "action_name": self.action_name,
+        "action_name": self.action_name if hasattr(self, "action_name") else "None",
         "sender_id": tracker.sender_id,
         "url": str(url),
         "actions_module": str(actions_module),

rasa/tracing/instrumentation/instrumentation.py

@@ -45,6 +45,7 @@ from rasa.dialogue_understanding.generator import (
     MultiStepLLMCommandGenerator,
     SingleStepLLMCommandGenerator,
 )
+from rasa.dialogue_understanding.generator.flow_retrieval import FlowRetrieval
 from rasa.dialogue_understanding.generator.nlu_command_adapter import NLUCommandAdapter
 from rasa.engine.graph import GraphNode
 from rasa.engine.training.graph_trainer import GraphTrainer
@@ -283,6 +284,7 @@ SingleStepLLMCommandGeneratorType = TypeVar(
 MultiStepLLMCommandGeneratorType = TypeVar(
     "MultiStepLLMCommandGeneratorType", bound=MultiStepLLMCommandGenerator
 )
+FlowRetrievalType = TypeVar("FlowRetrievalType", bound=FlowRetrieval)
 CommandType = TypeVar("CommandType", bound=Command)
 PolicyType = TypeVar("PolicyType", bound=Policy)
 InformationRetrievalType = TypeVar(
@@ -317,6 +319,7 @@ def instrument(
     custom_action_executor_subclasses: Optional[
         List[Type[CustomActionExecutor]]
     ] = None,
+    flow_retrieval_class: Optional[Type[FlowRetrievalType]] = None,
 ) -> None:
     """Substitute methods to be traced by their traced counterparts.
 
@@ -445,6 +448,12 @@ def instrument(
             "_check_commands_against_startable_flows",
             attribute_extractors.extract_attrs_for_check_commands_against_startable_flows,
         )
+        _instrument_perform_health_check_method_for_component(
+            tracer_provider.get_tracer(llm_command_generator_class.__module__),
+            llm_command_generator_class,
+            "perform_llm_health_check",
+            attribute_extractors.extract_attrs_for_performing_health_check,
+        )
         mark_class_as_instrumented(llm_command_generator_class)
 
     if (
@@ -468,6 +477,14 @@ def instrument(
             "_check_commands_against_startable_flows",
             attribute_extractors.extract_attrs_for_check_commands_against_startable_flows,
         )
+        _instrument_perform_health_check_method_for_component(
+            tracer_provider.get_tracer(
+                single_step_llm_command_generator_class.__module__
+            ),
+            single_step_llm_command_generator_class,
+            "perform_llm_health_check",
+            attribute_extractors.extract_attrs_for_performing_health_check,
+        )
         mark_class_as_instrumented(single_step_llm_command_generator_class)
 
     if multi_step_llm_command_generator_class is not None and not class_is_instrumented(
@@ -488,8 +505,36 @@ def instrument(
             ),
             multi_step_llm_command_generator_class,
         )
+        _instrument_perform_health_check_method_for_component(
+            tracer_provider.get_tracer(
+                multi_step_llm_command_generator_class.__module__
+            ),
+            multi_step_llm_command_generator_class,
+            "perform_llm_health_check",
+            attribute_extractors.extract_attrs_for_performing_health_check,
+        )
         mark_class_as_instrumented(multi_step_llm_command_generator_class)
 
+    if (
+        any(
+            llm_based_command_generator_class is not None
+            for llm_based_command_generator_class in (
+                llm_command_generator_class,
+                single_step_llm_command_generator_class,
+                multi_step_llm_command_generator_class,
+            )
+        )
+        and flow_retrieval_class is not None
+        and not class_is_instrumented(flow_retrieval_class)
+    ):
+        _instrument_perform_health_check_method_for_component(
+            tracer_provider.get_tracer(flow_retrieval_class.__module__),
+            flow_retrieval_class,
+            "perform_embeddings_health_check",
+            attribute_extractors.extract_attrs_for_performing_health_check,
+        )
+        mark_class_as_instrumented(flow_retrieval_class)
+
     if command_subclasses:
         for command_subclass in command_subclasses:
             if command_subclass is not None and not class_is_instrumented(
@@ -524,6 +569,12 @@ def instrument(
             "generate",
             attribute_extractors.extract_attrs_for_generate,
         )
+        _instrument_perform_health_check_method_for_component(
+            tracer_provider.get_tracer(contextual_response_rephraser_class.__module__),
+            contextual_response_rephraser_class,
+            "perform_llm_health_check",
+            attribute_extractors.extract_attrs_for_performing_health_check,
+        )
         mark_class_as_instrumented(contextual_response_rephraser_class)
 
     if not module_is_instrumented(COMMAND_PROCESSOR_MODULE_NAME):
@@ -755,6 +806,18 @@ def _instrument_enterprise_search_policy(
         "_generate_llm_answer",
         attribute_extractors.extract_attrs_for_enterprise_search_generate_llm_answer,
     )
+    _instrument_perform_health_check_method_for_component(
+        tracer_provider.get_tracer(policy_class.__module__),
+        policy_class,
+        "perform_embeddings_health_check",
+        attribute_extractors.extract_attrs_for_performing_health_check,
+    )
+    _instrument_perform_health_check_method_for_component(
+        tracer_provider.get_tracer(policy_class.__module__),
+        policy_class,
+        "perform_llm_health_check",
+        attribute_extractors.extract_attrs_for_performing_health_check,
+    )
 
 
 def _instrument_intentless_policy(
@@ -787,6 +850,18 @@ def _instrument_intentless_policy(
         "_generate_llm_answer",
         attribute_extractors.extract_attrs_for_intentless_policy_generate_llm_answer,
     )
+    _instrument_perform_health_check_method_for_component(
+        tracer_provider.get_tracer(policy_class.__module__),
+        policy_class,
+        "perform_embeddings_health_check",
+        attribute_extractors.extract_attrs_for_performing_health_check,
+    )
+    _instrument_perform_health_check_method_for_component(
+        tracer_provider.get_tracer(policy_class.__module__),
+        policy_class,
+        "perform_llm_health_check",
+        attribute_extractors.extract_attrs_for_performing_health_check,
+    )
 
 
 def _instrument_processor(
@@ -1139,6 +1214,52 @@ def _instrument_grpc_custom_action_executor(
     logger.debug(f"Instrumented '{grpc_custom_action_executor_class.__name__}.run.")
 
 
+def _instrument_perform_health_check_method_for_component(
+    tracer: Tracer,
+    instrumented_class: Type,
+    method_name: Text,
+    attr_extractor: Optional[Callable] = None,
+    return_value_attr_extractor: Optional[Callable] = None,
+) -> None:
+    def tracing_perform_health_check_for_component(
+        fn: Callable[..., S],
+    ) -> Callable[..., S]:
+        @functools.wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> S:
+            # Check the first argument to adjust for self/cls depending on how
+            # the static method from LLMHealthCheckMixin / EmbeddingsLLMHealthCheckMixin
+            # is called.
+            if args and isinstance(
+                args[0], (instrumented_class, type(instrumented_class))
+            ):
+                # The first argument is self/cls; align args to match the signature
+                args = args[1:]
+
+            span_name = f"{instrumented_class.__name__}.{fn.__name__}"
+            extracted_attrs = attr_extractor(*args, **kwargs) if attr_extractor else {}
+
+            with tracer.start_as_current_span(span_name) as span:
+                result = fn(*args, **kwargs)
+
+                # Extract attributes from the return value, if an extractor is provided
+                return_value_attributes = (
+                    return_value_attr_extractor(result, *args, **kwargs)
+                    if return_value_attr_extractor
+                    else {}
+                )
+
+                span.set_attributes({**extracted_attrs, **return_value_attributes})
+                return result
+
+        return wrapper
+
+    method_to_trace = getattr(instrumented_class, method_name)
+    traced_method = tracing_perform_health_check_for_component(method_to_trace)
+    setattr(instrumented_class, method_name, traced_method)
+
+    logger.debug(f"Instrumented '{instrumented_class.__name__}.{method_name}'.")
+
+
 def _mangled_instrumented_boolean_attribute_name(instrumented_class: Type) -> Text:
     # see https://peps.python.org/pep-0008/#method-names-and-instance-variables
     # and https://stackoverflow.com/a/50401073

rasa/utils/common.py

@@ -90,6 +90,11 @@ EXPECTED_WARNINGS: List[Tuple[Type[Warning], str]] = [
     # Ignore Keras DeprecationWarning since it requires that we
     # upgrade tensorflow-macos to 2.13.0 version.
     (DeprecationWarning, "invalid escape sequence*"),
+    # Ignore importlib open_text and read_text warnings for now
+    (
+        DeprecationWarning,
+        "https://importlib-resources.readthedocs.io/en/latest/using.html#migrating-from-legacy",
+    ),
 ]
 
 PYTHON_LOGGING_SCHEMA_DOCS = (

rasa/utils/endpoints.py

@@ -1,7 +1,7 @@
 import os
 import ssl
 from types import ModuleType
-from typing import Any, Optional, Text, Dict, Union
+from typing import Any, List, Optional, Text, Dict, Union
 
 import aiohttp
 import structlog
@@ -41,6 +41,32 @@ def read_endpoint_config(
         return None
 
 
+def read_property_config_from_endpoints_file(
+    filename: str, property_name: str
+) -> Optional[Union[Dict[str, Any], List]]:
+    """Read a property from an endpoint configuration file."""
+    if not filename:
+        return None
+
+    try:
+        content = read_config_file(filename)
+
+        if content.get(property_name) is None:
+            return None
+
+        return content[property_name]
+    except FileNotFoundError:
+        structlogger.error(
+            "endpoint.read.failed_no_such_file",
+            filename=os.path.abspath(filename),
+            event_info=(
+                "Failed to read endpoint configuration file - "
+                "the file was not found."
+            ),
+        )
+        return None
+
+
 def concat_url(base: Text, subpath: Optional[Text]) -> Text:
     """Append a subpath to a base url.
 

rasa/utils/io.py

@@ -18,7 +18,6 @@ from typing import (
     Type,
     Callable,
     TYPE_CHECKING,
-    Pattern,
 )
 
 from ruamel import yaml
@@ -167,21 +166,14 @@ def create_validator(
     return FunctionValidator
 
 
-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"
-        "\u200d"  # zero width joiner
-        "\u200c"  # zero width non-joiner
-        "]+",
-        flags=re.UNICODE,
-    )
+def remove_emojis(s: str) -> str:
+    import demoji
+
+    replaced = demoji.replace(s)
+    if replaced == s:
+        return s
+    # remove duplicate or trailing whitespaces if emojis were removed
+    return re.sub(r" +", " ", replaced).strip()
 
 
 def are_directories_equal(dir1: Path, dir2: Path) -> bool:

rasa/utils/log_utils.py

@@ -78,6 +78,7 @@ def _anonymizer(
 
 def configure_structlog(
     log_level: Optional[int] = None,
+    include_time: bool = False,
 ) -> None:
     """Configure logging of the server."""
     if log_level is None:  # Log level NOTSET is 0 so we use `is None` here
@@ -114,6 +115,9 @@ def configure_structlog(
         SentryProcessor(event_level=logging.FATAL),
     ]
 
+    if include_time:
+        shared_processors.append(structlog.processors.TimeStamper(fmt="iso"))
+
     if not FORCE_JSON_LOGGING and sys.stderr.isatty():
         # Pretty printing when we run in a terminal session.
         # Automatically prints pretty tracebacks when "rich" is installed
@@ -143,11 +147,14 @@ def configure_structlog(
         # logger.
         cache_logger_on_first_use=True,
     )
+    # doing logger creation inline, to prevent usage of unconfigured logger
+    structlog.get_logger().debug("structlog.configured")
 
 
 def log_llm(logger: Any, log_module: str, log_event: str, **kwargs: Any) -> None:
-    """Logs LLM-specific events depending on a flag passed through an environment
-    variable. If the module's flag is set to INFO (e.g.
+    """Logs LLM-specific events depending on a flag passed through an env var.
+
+    If the module's flag is set to INFO (e.g.
     LOG_PROMPT_LLM_COMMAND_GENERATOR=INFO), its prompt is logged at INFO level,
     overriding the general log level setting.
 

rasa/utils/sanic_error_handler.py

@@ -0,0 +1,32 @@
+from sanic import Sanic
+from sanic.handlers import ErrorHandler
+from sanic.request import Request
+from sanic.exceptions import ServerError
+
+
+# TODO: remove custom handler when upgrading to sanic >= 24
+#       the underlying issue https://github.com/sanic-org/sanic/issues/2572
+#       has been fixed in sanic 24
+class IgnoreWSServerErrorHandler(ErrorHandler):
+    @staticmethod
+    def log(request: Request, exception: Exception) -> None:
+        try:
+            if (
+                request.url.startswith("ws")
+                and isinstance(exception, ServerError)
+                and exception.args
+                and (
+                    exception.args[0]
+                    == "Invalid response type None (need HTTPResponse)"
+                )
+            ):
+                # in case we are in a websocket connection, we don't want to log the
+                # the error, as this is a bug in sanic
+                return
+        except Exception:
+            pass
+        ErrorHandler.log(request, exception)  # type: ignore
+
+
+def register_custom_sanic_error_handler(app: Sanic) -> None:
+    app.error_handler = IgnoreWSServerErrorHandler()