rasa-pro 3.11.0rc1__py3-none-any.whl → 3.11.0rc3__py3-none-any.whl

rasa/shared/utils/llm.py

@@ -690,14 +690,16 @@ def resolve_model_client_config(
 ) -> Optional[Dict[str, Any]]:
     """Resolve the model group in the model config.
 
-    If the config is pointing to a model group, the corresponding model group
+    1. If the config is pointing to a model group, the corresponding model group
     of the endpoints.yml is returned.
-    If the config is using the old syntax, e.g. defining the llm
+    2. If the config is using the old syntax, e.g. defining the llm
     directly in config.yml, the config is returned as is.
+    3. If the config is already resolved, return it as is.
 
     Args:
         model_config: The model config to be resolved.
         component_name: The name of the component.
+        component_name: The method of the component.
 
     Returns:
         The resolved llm config.
@@ -718,6 +720,7 @@ def resolve_model_client_config(
     if model_config is None:
         return None
 
+    # Config is already resolved or defines a client without model groups
     if MODEL_GROUP_CONFIG_KEY not in model_config:
         return model_config
 

rasa/shared/utils/yaml.py

@@ -1,4 +1,5 @@
 import datetime
+import io
 import logging
 import os
 import re
@@ -8,19 +9,13 @@ from dataclasses import field
 from functools import lru_cache
 from io import StringIO
 from pathlib import Path
-from typing import Dict, List, Optional, Any, Callable, Tuple, Union
+from typing import Any, List, Optional, Tuple, Dict, Callable, Union
 
 import jsonschema
 from importlib_resources import files
 from packaging import version
 from pykwalify.core import Core
 from pykwalify.errors import SchemaError
-from ruamel import yaml as yaml
-from ruamel.yaml import RoundTripRepresenter, YAMLError
-from ruamel.yaml.comments import CommentedSeq, CommentedMap
-from ruamel.yaml.constructor import DuplicateKeyError, BaseConstructor, ScalarNode
-from ruamel.yaml.loader import SafeLoader
-
 from rasa.shared.constants import (
     ASSERTIONS_SCHEMA_EXTENSIONS_FILE,
     ASSERTIONS_SCHEMA_FILE,
@@ -51,6 +46,11 @@ from rasa.shared.utils.io import (
     raise_warning,
     read_json_file,
 )
+from ruamel import yaml as yaml
+from ruamel.yaml import YAML, RoundTripRepresenter, YAMLError
+from ruamel.yaml.comments import CommentedSeq, CommentedMap
+from ruamel.yaml.constructor import DuplicateKeyError, BaseConstructor, ScalarNode
+from ruamel.yaml.loader import SafeLoader
 
 logger = logging.getLogger(__name__)
 
@@ -64,8 +64,17 @@ SENSITIVE_DATA = [API_KEY]
 
 @dataclass
 class PathWithError:
+    """Represents a validation error at a specific location in the YAML content.
+
+    Attributes:
+        message (str): A description of the validation error.
+        path (List[str]): Path to the node where the error occurred.
+        key (Optional[str]): The specific key associated with the error, if any.
+    """
+
     message: str
     path: List[str] = field(default_factory=list)
+    key: Optional[str] = None
 
 
 def fix_yaml_loader() -> None:
@@ -146,21 +155,72 @@ class YamlValidationException(YamlException, ValueError):
         if self.validation_errors:
             unique_errors = {}
             for error in self.validation_errors:
-                line_number = self._line_number_for_path(self.content, error.path)
+                line_number = self._line_number_for_path(
+                    self.content, error.path, error.key
+                )
 
                 if line_number and self.filename:
-                    error_representation = f"  in {self.filename}:{line_number}:\n"
+                    error_location = f"  in {self.filename}:{line_number}:\n"
                 elif line_number:
-                    error_representation = f"  in Line {line_number}:\n"
+                    error_location = f"  in Line {line_number}:\n"
                 else:
-                    error_representation = ""
+                    error_location = ""
 
-                error_representation += f"      {error.message}"
-                unique_errors[error.message] = error_representation
+                code_snippet = self._get_code_snippet(line_number)
+                error_message = f"{error_location}\n{code_snippet}{error.message}\n"
+                unique_errors[error.message] = error_message
             error_msg = "\n".join(unique_errors.values())
             msg += f":\n{error_msg}"
         return msg
 
+    def _get_code_snippet(
+        self,
+        error_line: Optional[int],
+        context_lines: int = 2,
+    ) -> str:
+        """Extract code snippet from the YAML lines around the error.
+
+        Args:
+            error_line: Line number where the error occurred (1-based).
+            context_lines: Number of context lines before and after the error line.
+                Default is 2, balancing context and readability. Adjust as needed.
+
+        Returns:
+            A string containing the code snippet with the error highlighted.
+        """
+        yaml_lines = self._get_serialized_yaml_lines()
+        if not yaml_lines or error_line is None:
+            return ""
+
+        start = max(error_line - context_lines - 1, 0)
+        end = min(error_line + context_lines, len(yaml_lines))
+        snippet_lines = yaml_lines[start:end]
+        snippet = ""
+        for idx, line_content in enumerate(snippet_lines, start=start + 1):
+            prefix = ">>> " if idx == error_line else "    "
+            line_number_str = str(idx)
+            snippet += f"{prefix}{line_number_str} | {line_content}\n"
+        return snippet
+
+    def _get_serialized_yaml_lines(self) -> List[str]:
+        """Serialize the content back to YAML and return the lines."""
+        yaml_lines = []
+        try:
+            yaml = YAML()
+            yaml.default_flow_style = False
+            # Set width to 1000, so we don't break the lines of the original YAML file
+            yaml.width = 1000  # type: ignore[assignment]
+            yaml.indent(mapping=2, sequence=4, offset=2)
+            stream = io.StringIO()
+            yaml.dump(self.content, stream)
+            serialized_yaml = stream.getvalue()
+            yaml_lines = serialized_yaml.splitlines()
+            return yaml_lines
+        except Exception as exc:
+            logger.debug(f"Error serializing YAML content: {exc}")
+
+        return yaml_lines
+
     def _calculate_number_of_lines(
         self,
         current: Union[CommentedSeq, CommentedMap],
@@ -228,7 +288,9 @@ class YamlValidationException(YamlException, ValueError):
         # Return the calculated child offset and True indicating a line number was found
         return child_offset, True
 
-    def _line_number_for_path(self, current: Any, path: List[str]) -> Optional[int]:
+    def _line_number_for_path(
+        self, current: Any, path: List[str], key: Optional[str] = None
+    ) -> Optional[int]:
         """Get line number for a yaml path in the current content.
 
         Implemented using recursion: algorithm goes down the path navigating to the
@@ -237,6 +299,7 @@ class YamlValidationException(YamlException, ValueError):
         Args:
             current: current content
             path: path to traverse within the content
+            key: the key associated with the error, if any
 
         Returns:
             the line number of the path in the content.
@@ -247,6 +310,10 @@ class YamlValidationException(YamlException, ValueError):
         this_line = current.lc.line + 1 if hasattr(current, "lc") else None
 
         if not path:
+            if key and hasattr(current, "lc"):
+                if hasattr(current.lc, "data") and key in current.lc.data:
+                    key_line_no = current.lc.data[key][0] + 1
+                    return key_line_no
             return this_line
 
         head, tail = path[0], path[1:]
@@ -256,7 +323,7 @@ class YamlValidationException(YamlException, ValueError):
 
         if head:
             if isinstance(current, dict) and head in current:
-                line = self._line_number_for_path(current[head], tail)
+                line = self._line_number_for_path(current[head], tail, key)
                 if line is None:
                     line_offset, found_lc = self._calculate_number_of_lines(
                         current, head
@@ -266,10 +333,13 @@ class YamlValidationException(YamlException, ValueError):
                     return this_line + line_offset
                 return line
             elif isinstance(current, list) and head.isdigit():
-                return self._line_number_for_path(current[int(head)], tail) or this_line
+                return (
+                    self._line_number_for_path(current[int(head)], tail, key)
+                    or this_line
+                )
             else:
                 return this_line
-        return self._line_number_for_path(current, tail) or this_line
+        return self._line_number_for_path(current, tail, key) or this_line
 
 
 def read_schema_file(
@@ -331,13 +401,26 @@ def validate_yaml_content_using_schema(
     try:
         core.validate(raise_exception=True)
     except SchemaError:
+        # PyKwalify propagates each validation error up the data hierarchy, resulting
+        # in multiple redundant errors for a single issue. To present a clear message
+        # about the root cause, we use only the first error.
+        error = core.errors[0]
+
+        # Increment numeric indices by 1 to convert from 0-based to 1-based indexing
+        error_message = re.sub(
+            r"(/)(\d+)", lambda m: f"/{int(m.group(2)) + 1}", str(error)
+        )
+
         raise YamlValidationException(
             "Please make sure the file is correct and all "
             "mandatory parameters are specified. Here are the errors "
             "found during validation",
             [
-                PathWithError(message=str(e), path=e.path.split("/"))
-                for e in core.errors
+                PathWithError(
+                    message=error_message,
+                    path=error.path.removeprefix("/").split("/"),
+                    key=getattr(error, "key", None),
+                )
             ],
             content=yaml_content,
         )
@@ -424,46 +507,6 @@ def validate_raw_yaml_using_schema_file_with_responses(
     )
 
 
-def process_content(content: str) -> str:
-    """Process the content to handle both Windows paths and emojis.
-    Windows paths are processed by escaping backslashes but emojis are left untouched.
-
-    Args:
-        content: yaml content to be processed
-    """
-    # Detect common Windows path patterns: e.g., C:\ or \\
-    UNESCAPED_WINDOWS_PATH_PATTERN = re.compile(
-        r"(?<!\w)[a-zA-Z]:(\\[a-zA-Z0-9_ -]+)*(\\)?(?!\\n)"
-    )
-    ESCAPED_WINDOWS_PATH_PATTERN = re.compile(
-        r"(?<!\w)[a-zA-Z]:(\\\\[a-zA-Z0-9_ -]+)+\\\\?(?!\\n)"
-    )
-
-    # Function to escape backslashes in Windows paths but leave other content as is
-    def escape_windows_paths(match: re.Match) -> str:
-        path = str(match.group(0))
-        return path.replace("\\", "\\\\")  # Escape backslashes only in Windows paths
-
-    def unescape_windows_paths(match: re.Match) -> str:
-        path = str(match.group(0))
-        return path.replace("\\\\", "\\")
-
-    # First, process Windows paths by escaping backslashes
-    content = re.sub(UNESCAPED_WINDOWS_PATH_PATTERN, escape_windows_paths, content)
-
-    # Ensure proper handling of emojis by decoding Unicode sequences
-    content = (
-        content.encode("utf-8")
-        .decode("raw_unicode_escape")
-        .encode("utf-16", "surrogatepass")
-        .decode("utf-16")
-    )
-
-    content = re.sub(ESCAPED_WINDOWS_PATH_PATTERN, unescape_windows_paths, content)
-
-    return content
-
-
 def read_yaml(
     content: str,
     reader_type: Union[str, List[str]] = "safe",
@@ -479,9 +522,6 @@ def read_yaml(
     Raises:
         ruamel.yaml.parser.ParserError: If there was an error when parsing the YAML.
     """
-    if _is_ascii(content):
-        content = process_content(content)
-
     custom_constructor = kwargs.get("custom_constructor", None)
 
     # Create YAML parser with custom constructor

rasa/studio/auth.py

@@ -23,12 +23,10 @@ from rasa.studio.results_logger import with_studio_error_handler, StudioResult
 class StudioAuth:
     """Handles the authentication with the Rasa Studio authentication server."""
 
-    def __init__(
-        self,
-        studio_config: StudioConfig,
-        verify: bool = True,
-    ) -> None:
+    def __init__(self, studio_config: StudioConfig) -> None:
         self.config = studio_config
+        verify = not studio_config.disable_verify
+
         self.keycloak_openid = KeycloakOpenID(
             server_url=studio_config.authentication_server_url,
             client_id=studio_config.client_id,

rasa/studio/config.py

@@ -2,13 +2,14 @@ from __future__ import annotations
 
 import os
 from dataclasses import dataclass
-from typing import Dict, Optional, Text
+from typing import Any, Dict, Optional, Text
 
 from rasa.utils.common import read_global_config_value, write_global_config_value
 
 from rasa.studio.constants import (
     RASA_STUDIO_AUTH_SERVER_URL_ENV,
     RASA_STUDIO_CLI_CLIENT_ID_KEY_ENV,
+    RASA_STUDIO_CLI_DISABLE_VERIFY_KEY_ENV,
     RASA_STUDIO_CLI_REALM_NAME_KEY_ENV,
     RASA_STUDIO_CLI_STUDIO_URL_ENV,
     STUDIO_CONFIG_KEY,
@@ -19,6 +20,7 @@ STUDIO_URL_KEY = "studio_url"
 CLIENT_ID_KEY = "client_id"
 REALM_NAME_KEY = "realm_name"
 CLIENT_SECRET_KEY = "client_secret"
+DISABLE_VERIFY = "disable_verify"
 
 
 @dataclass
@@ -27,13 +29,15 @@ class StudioConfig:
     studio_url: Optional[Text]
     client_id: Optional[Text]
     realm_name: Optional[Text]
+    disable_verify: bool = False
 
-    def to_dict(self) -> Dict[Text, Optional[Text]]:
+    def to_dict(self) -> Dict[Text, Optional[Any]]:
         return {
             AUTH_SERVER_URL_KEY: self.authentication_server_url,
             STUDIO_URL_KEY: self.studio_url,
             CLIENT_ID_KEY: self.client_id,
             REALM_NAME_KEY: self.realm_name,
+            DISABLE_VERIFY: self.disable_verify,
         }
 
     @classmethod
@@ -43,6 +47,7 @@ class StudioConfig:
             studio_url=data[STUDIO_URL_KEY],
             client_id=data[CLIENT_ID_KEY],
             realm_name=data[REALM_NAME_KEY],
+            disable_verify=data.get(DISABLE_VERIFY, False),
         )
 
     def write_config(self) -> None:
@@ -73,7 +78,7 @@ class StudioConfig:
         config = read_global_config_value(STUDIO_CONFIG_KEY, unavailable_ok=True)
 
         if config is None:
-            return StudioConfig(None, None, None, None)
+            return StudioConfig(None, None, None, None, False)
 
         if not isinstance(config, dict):
             raise ValueError(
@@ -83,7 +88,7 @@ class StudioConfig:
             )
 
         for key in config:
-            if not isinstance(config[key], str):
+            if not isinstance(config[key], str) and key != DISABLE_VERIFY:
                 raise ValueError(
                     "Invalid config file format. "
                     f"Key '{key}' is not a text value."
@@ -102,6 +107,9 @@ class StudioConfig:
             studio_url=StudioConfig._read_env_value(RASA_STUDIO_CLI_STUDIO_URL_ENV),
             client_id=StudioConfig._read_env_value(RASA_STUDIO_CLI_CLIENT_ID_KEY_ENV),
             realm_name=StudioConfig._read_env_value(RASA_STUDIO_CLI_REALM_NAME_KEY_ENV),
+            disable_verify=bool(
+                os.getenv(RASA_STUDIO_CLI_DISABLE_VERIFY_KEY_ENV, False)
+            ),
         )
 
     @staticmethod
@@ -124,4 +132,5 @@ class StudioConfig:
             studio_url=self.studio_url or other.studio_url,
             client_id=self.client_id or other.client_id,
             realm_name=self.realm_name or other.realm_name,
+            disable_verify=self.disable_verify or other.disable_verify,
         )

rasa/studio/constants.py

@@ -10,6 +10,7 @@ RASA_STUDIO_AUTH_SERVER_URL_ENV = "RASA_STUDIO_AUTH_SERVER_URL"
 RASA_STUDIO_CLI_STUDIO_URL_ENV = "RASA_STUDIO_CLI_STUDIO_URL"
 RASA_STUDIO_CLI_REALM_NAME_KEY_ENV = "RASA_STUDIO_CLI_REALM_NAME_KEY"
 RASA_STUDIO_CLI_CLIENT_ID_KEY_ENV = "RASA_STUDIO_CLI_CLIENT_ID_KEY"
+RASA_STUDIO_CLI_DISABLE_VERIFY_KEY_ENV = "RASA_STUDIO_CLI_DISABLE_VERIFY_KEY"
 
 STUDIO_NLU_FILENAME = "studio_nlu.yml"
 STUDIO_DOMAIN_FILENAME = "studio_domain.yml"

rasa/studio/data_handler.py

@@ -76,7 +76,9 @@ class StudioDataHandler:
 
         return request
 
-    def _make_request(self, GQL_req: Dict[Any, Any]) -> Dict[Any, Any]:
+    def _make_request(
+        self, GQL_req: Dict[Any, Any], verify: bool = True
+    ) -> Dict[Any, Any]:
         token = KeycloakTokenReader().get_token()
         if token.is_expired():
             token = self.refresh_token(token)
@@ -93,6 +95,7 @@ class StudioDataHandler:
                 "Authorization": f"{token.token_type} {token.access_token}",
                 "Content-Type": "application/json",
             },
+            verify=verify,
         )
 
         if res.status_code != 200:
@@ -128,7 +131,9 @@ class StudioDataHandler:
             The data from Rasa Studio.
         """
         GQL_req = self._build_request()
-        response = self._make_request(GQL_req)
+        verify = not self.studio_config.disable_verify
+
+        response = self._make_request(GQL_req, verify=verify)
         self._extract_data(response)
 
     def request_data(
@@ -145,7 +150,9 @@ class StudioDataHandler:
             The data from Rasa Studio.
         """
         GQL_req = self._build_request(intent_names, entity_names)
-        response = self._make_request(GQL_req)
+        verify = not self.studio_config.disable_verify
+
+        response = self._make_request(GQL_req, verify=verify)
         self._extract_data(response)
 
     def get_config(self) -> Optional[str]:

rasa/studio/upload.py

@@ -56,14 +56,17 @@ def _get_selected_entities_and_intents(
 
 def handle_upload(args: argparse.Namespace) -> None:
     """Uploads primitives to rasa studio."""
-    endpoint = StudioConfig.read_config().studio_url
+    studio_config = StudioConfig.read_config()
+    endpoint = studio_config.studio_url
+    verify = not studio_config.disable_verify
+
     if not endpoint:
         rasa.shared.utils.cli.print_error_and_exit(
             "No GraphQL endpoint found in config. Please run `rasa studio config`."
         )
         return
 
-    if not is_auth_working(endpoint):
+    if not is_auth_working(endpoint, verify):
         rasa.shared.utils.cli.print_error_and_exit(
             "Authentication is invalid or expired. Please run `rasa studio login`."
         )
@@ -81,9 +84,9 @@ def handle_upload(args: argparse.Namespace) -> None:
 
     # check safely if args.calm is set and not fail if not
     if hasattr(args, "calm") and args.calm:
-        upload_calm_assistant(args, endpoint)
+        upload_calm_assistant(args, endpoint, verify=verify)
     else:
-        upload_nlu_assistant(args, endpoint)
+        upload_nlu_assistant(args, endpoint, verify=verify)
 
 
 config_keys = [
@@ -135,7 +138,9 @@ def _get_assistant_name(config: Dict[Text, Any]) -> str:
 
 
 @with_studio_error_handler
-def upload_calm_assistant(args: argparse.Namespace, endpoint: str) -> StudioResult:
+def upload_calm_assistant(
+    args: argparse.Namespace, endpoint: str, verify: bool = True
+) -> StudioResult:
     """Uploads the CALM assistant data to Rasa Studio.
 
     Args:
@@ -227,11 +232,13 @@ def upload_calm_assistant(args: argparse.Namespace, endpoint: str) -> StudioResu
     structlogger.info(
         "rasa.studio.upload.calm", event_info="Uploading to Rasa Studio..."
     )
-    return make_request(endpoint, graphql_req)
+    return make_request(endpoint, graphql_req, verify)
 
 
 @with_studio_error_handler
-def upload_nlu_assistant(args: argparse.Namespace, endpoint: str) -> StudioResult:
+def upload_nlu_assistant(
+    args: argparse.Namespace, endpoint: str, verify: bool = True
+) -> StudioResult:
     """Uploads the classic (dm1) assistant data to Rasa Studio.
 
     Args:
@@ -241,6 +248,7 @@ def upload_nlu_assistant(args: argparse.Namespace, endpoint: str) -> StudioResul
             - intents: The intents to upload
             - entities: The entities to upload
         endpoint: The studio endpoint
+        verify: Whether to verify SSL
     Returns:
         None
     """
@@ -286,10 +294,10 @@ def upload_nlu_assistant(args: argparse.Namespace, endpoint: str) -> StudioResul
     structlogger.info(
         "rasa.studio.upload.nlu", event_info="Uploading to Rasa Studio..."
     )
-    return make_request(endpoint, graphql_req)
+    return make_request(endpoint, graphql_req, verify)
 
 
-def is_auth_working(endpoint: str) -> bool:
+def is_auth_working(endpoint: str, verify: bool = True) -> bool:
     """Send a test request to Studio to check if auth is working."""
     result = make_request(
         endpoint,
@@ -306,16 +314,18 @@ def is_auth_working(endpoint: str) -> bool:
             ),
             "variables": {},
         },
+        verify,
     )
     return result.was_successful
 
 
-def make_request(endpoint: str, graphql_req: Dict) -> StudioResult:
+def make_request(endpoint: str, graphql_req: Dict, verify: bool = True) -> StudioResult:
     """Makes a request to the studio endpoint to upload data.
 
     Args:
         endpoint: The studio endpoint
         graphql_req: The graphql request
+        verify: Whether to verify SSL
     """
     token = KeycloakTokenReader().get_token()
     res = requests.post(
@@ -325,6 +335,7 @@ def make_request(endpoint: str, graphql_req: Dict) -> StudioResult:
             "Authorization": f"{token.token_type} {token.access_token}",
             "Content-Type": "application/json",
         },
+        verify=verify,
     )
 
     if results_logger.response_has_errors(res.json()):

rasa/telemetry.py

@@ -35,9 +35,9 @@ from rasa.constants import (
 from rasa.shared.constants import (
     PROMPT_CONFIG_KEY,
     PROMPT_TEMPLATE_CONFIG_KEY,
-    MODEL_GROUP_CONFIG_KEY,
     LLM_API_HEALTH_CHECK_ENV_VAR,
     LLM_API_HEALTH_CHECK_DEFAULT_VALUE,
+    MODEL_GROUP_CONFIG_KEY,
 )
 from rasa.engine.storage.local_model_storage import LocalModelStorage
 from rasa.shared.constants import DOCS_URL_TELEMETRY, UTTER_ASK_PREFIX
@@ -112,6 +112,7 @@ TELEMETRY_INTERACTIVE_LEARNING_STARTED_EVENT = "Interactive Learning Started"
 TELEMETRY_SERVER_STARTED_EVENT = "Server Started"
 TELEMETRY_PROJECT_CREATED_EVENT = "Project Created"
 TELEMETRY_SHELL_STARTED_EVENT = "Shell Started"
+TELEMETRY_INSPECT_STARTED_EVENT = "Inspect Started"
 TELEMETRY_VISUALIZATION_STARTED_EVENT = "Story Visualization Started"
 TELEMETRY_TEST_CORE_EVENT = "Model Core Tested"
 TELEMETRY_TEST_NLU_EVENT = "Model NLU Tested"
@@ -1132,6 +1133,7 @@ def _get_llm_command_generator_config(config: Dict[str, Any]) -> Optional[Dict]:
     def extract_llm_command_generator_llm_client_settings(component: Dict) -> Dict:
         """Extracts settings related to LLM command generator."""
         llm_config = component.get(LLM_CONFIG_KEY, {})
+        # Config at this stage is not yet resolved, so read from `model_group`
         llm_model_group_id = llm_config.get(MODEL_GROUP_CONFIG_KEY)
         llm_model_name = llm_config.get(MODEL_CONFIG_KEY) or llm_config.get(
             MODEL_NAME_CONFIG_KEY
@@ -1173,6 +1175,7 @@ def _get_llm_command_generator_config(config: Dict[str, Any]) -> Optional[Dict]:
             if flow_retrieval_enabled
             else None
         )
+        # Config at this stage is not yet resolved, so read from `model_group`
         flow_retrieval_embedding_model_group_id = embeddings_config.get(
             MODEL_GROUP_CONFIG_KEY
         )
@@ -1378,6 +1381,17 @@ def track_shell_started(model_type: Text) -> None:
     _track(TELEMETRY_SHELL_STARTED_EVENT, {"type": model_type})
 
 
+@ensure_telemetry_enabled
+def track_inspect_started(model_type: Text) -> None:
+    """Track when a user starts a bot using rasa inspect.
+
+    Args:
+        channel: Channel name `socketio` (used for chat assistants)
+         or `browser_audio` (used for voice).
+    """
+    _track(TELEMETRY_INSPECT_STARTED_EVENT, {"type": model_type})
+
+
 @ensure_telemetry_enabled
 def track_visualization() -> None:
     """Track when a user runs the visualization."""

rasa/tracing/config.py

@@ -33,6 +33,7 @@ from rasa.dialogue_understanding.generator import (
     SingleStepLLMCommandGenerator,
     MultiStepLLMCommandGenerator,
 )
+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
@@ -111,6 +112,7 @@ def configure_tracing(tracer_provider: Optional[TracerProvider]) -> None:
         single_step_llm_command_generator_class=SingleStepLLMCommandGenerator,
         multi_step_llm_command_generator_class=MultiStepLLMCommandGenerator,
         custom_action_executor_subclasses=custom_action_executor_subclasses,
+        flow_retrieval_class=FlowRetrieval,
     )
 
 
@@ -129,7 +131,7 @@ def get_tracer_provider(endpoints_file: Text) -> Optional[TracerProvider]:
 
     if not cfg:
         logger.info(
-            f"No endpoint for tracing type available in {endpoints_file},"
+            f"No endpoint for tracing type available in {endpoints_file}, "
             f"tracing will not be configured."
         )
         return None

rasa/tracing/instrumentation/attribute_extractors.py

@@ -414,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,