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

rasa/studio/upload.py

@@ -15,11 +15,18 @@ from rasa.shared.constants import (
     DEFAULT_DOMAIN_PATHS,
     DEFAULT_CONFIG_PATH,
 )
-from rasa.shared.core.flows.yaml_flows_io import YamlFlowsWriter
+from rasa.shared.core.domain import Domain
+from rasa.shared.core.flows.yaml_flows_io import YAMLFlowsReader, YamlFlowsWriter
 from rasa.shared.exceptions import RasaException
 from rasa.shared.importers.importer import TrainingDataImporter, FlowSyncImporter
-from rasa.shared.nlu.training_data.formats.rasa_yaml import RasaYAMLWriter
-from rasa.shared.utils.yaml import dump_obj_as_yaml_to_string, read_yaml_file
+from rasa.shared.nlu.training_data.formats.rasa_yaml import (
+    RasaYAMLReader,
+    RasaYAMLWriter,
+)
+from rasa.shared.utils.yaml import (
+    dump_obj_as_yaml_to_string,
+    read_yaml_file,
+)
 from rasa.studio import results_logger
 from rasa.studio.auth import KeycloakTokenReader
 from rasa.studio.config import StudioConfig
@@ -27,6 +34,27 @@ from rasa.studio.results_logger import StudioResult, with_studio_error_handler
 
 structlogger = structlog.get_logger()
 
+CONFIG_KEYS = [
+    "recipe",
+    "language",
+    "pipeline",
+    "llm",
+    "policies",
+    "model_name",
+    "assistant_id",
+]
+
+DOMAIN_KEYS = [
+    "version",
+    "actions",
+    "responses",
+    "slots",
+    "intents",
+    "entities",
+    "forms",
+    "session_config",
+]
+
 
 def _get_selected_entities_and_intents(
     args: argparse.Namespace,
@@ -54,6 +82,44 @@ def _get_selected_entities_and_intents(
     return list(entities), list(intents)
 
 
+def run_validation(args: argparse.Namespace) -> None:
+    """Run the validation before uploading to Studio.
+
+    This is to avoid uploading invalid assistant data
+    that would raise errors during Rasa Pro training in Studio.
+
+    The validation checks that were selected to be run before uploading
+    maintain parity with the features that are supported in Studio.
+    """
+    from rasa.validator import Validator
+
+    training_data_importer = TrainingDataImporter.load_from_dict(
+        domain_path=args.domain,
+        training_data_paths=args.data,
+        config_path=args.config,
+        expand_env_vars=False,
+    )
+
+    structlogger.info(
+        "rasa.studio.upload.validating_data",
+        event_info="Validating domain and training data...",
+    )
+
+    validator = Validator.from_importer(training_data_importer)
+
+    if not validator.verify_studio_supported_validations():
+        structlogger.error(
+            "rasa.studio.upload.validate_files.project_validation_error",
+            event_info="Project validation completed with errors.",
+        )
+        sys.exit(1)
+
+    structlogger.info(
+        "rasa.studio.upload.validate_files.success",
+        event_info="Project validation completed successfully.",
+    )
+
+
 def handle_upload(args: argparse.Namespace) -> None:
     """Uploads primitives to rasa studio."""
     studio_config = StudioConfig.read_config()
@@ -64,35 +130,33 @@ def handle_upload(args: argparse.Namespace) -> None:
         rasa.shared.utils.cli.print_error_and_exit(
             "No GraphQL endpoint found in config. Please run `rasa studio config`."
         )
-    else:
-        structlogger.info(
-            "rasa.studio.upload.loading_data", event_info="Loading data..."
-        )
+        return
 
-        args.domain = rasa.cli.utils.get_validated_path(
-            args.domain, "domain", DEFAULT_DOMAIN_PATHS
+    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`."
         )
+        return
 
-        args.config = rasa.cli.utils.get_validated_path(
-            args.config, "config", DEFAULT_CONFIG_PATH
-        )
+    structlogger.info("rasa.studio.upload.loading_data", event_info="Loading data...")
+
+    args.domain = rasa.cli.utils.get_validated_path(
+        args.domain, "domain", DEFAULT_DOMAIN_PATHS
+    )
 
-        # check safely if args.calm is set and not fail if not
-        if hasattr(args, "calm") and args.calm:
-            upload_calm_assistant(args, endpoint, verify=verify)
-        else:
-            upload_nlu_assistant(args, endpoint, verify=verify)
+    args.config = rasa.cli.utils.get_validated_path(
+        args.config, "config", DEFAULT_CONFIG_PATH
+    )
 
+    Domain.expand_env_vars = False
+    RasaYAMLReader.expand_env_vars = False
+    YAMLFlowsReader.expand_env_vars = False
 
-config_keys = [
-    "recipe",
-    "language",
-    "pipeline",
-    "llm",
-    "policies",
-    "model_name",
-    "assistant_id",
-]
+    # check safely if args.calm is set and not fail if not
+    if hasattr(args, "calm") and args.calm:
+        upload_calm_assistant(args, endpoint, verify=verify)
+    else:
+        upload_nlu_assistant(args, endpoint, verify=verify)
 
 
 def extract_values(data: Dict, keys: List[Text]) -> Dict:
@@ -124,7 +188,11 @@ def _get_assistant_name(config: Dict[Text, Any]) -> str:
             ),
         )
 
-    structlogger.info(f"Uploading assistant with the name '{assistant_name}'.")
+    structlogger.info(
+        "rasa.studio.upload.name_selected",
+        event_info=f"Uploading assistant with the name '{assistant_name}'.",
+        assistant_name=assistant_name,
+    )
     return assistant_name
 
 
@@ -132,7 +200,7 @@ def _get_assistant_name(config: Dict[Text, Any]) -> str:
 def upload_calm_assistant(
     args: argparse.Namespace, endpoint: str, verify: bool = True
 ) -> StudioResult:
-    """Uploads the CALM assistant data to Rasa Studio.
+    """Validates and uploads the CALM assistant data to Rasa Studio.
 
     Args:
         args: The command line arguments
@@ -145,6 +213,8 @@ def upload_calm_assistant(
     Returns:
         None
     """
+    run_validation(args)
+
     structlogger.info(
         "rasa.studio.upload.loading_data", event_info="Parsing CALM assistant data..."
     )
@@ -152,64 +222,37 @@ def upload_calm_assistant(
     importer = TrainingDataImporter.load_from_dict(
         domain_path=args.domain,
         config_path=args.config,
+        expand_env_vars=False,
     )
 
     # Prepare config and domain
     config = importer.get_config()
     domain_from_files = importer.get_user_domain().as_dict()
-    endpoints_from_files = read_yaml_file(args.endpoints)
-    config_from_files = read_yaml_file(args.config)
+    endpoints_from_files = read_yaml_file(args.endpoints, expand_env_vars=False)
+    config_from_files = read_yaml_file(args.config, expand_env_vars=False)
 
     # Extract domain and config values
-    domain_keys = [
-        "version",
-        "actions",
-        "responses",
-        "slots",
-        "intents",
-        "entities",
-        "forms",
-        "session_config",
-    ]
-
-    domain = extract_values(domain_from_files, domain_keys)
-
-    assistant_name = _get_assistant_name(config)
-
-    training_data_paths = args.data
-
-    if isinstance(training_data_paths, list):
-        training_data_paths.append(args.flows)
-    elif isinstance(training_data_paths, str):
-        if isinstance(args.flows, list):
-            training_data_paths = [training_data_paths] + args.flows
-        elif isinstance(args.flows, str):
-            training_data_paths = [training_data_paths, args.flows]
-        else:
-            raise RasaException("Invalid flows path")
+    domain = extract_values(domain_from_files, DOMAIN_KEYS)
 
     # Prepare flows
     flow_importer = FlowSyncImporter.load_from_dict(
-        training_data_paths=training_data_paths
+        training_data_paths=args.data, expand_env_vars=False
     )
-
     flows = list(flow_importer.get_user_flows())
 
     # We instantiate the TrainingDataImporter again on purpose to avoid
     # adding patterns to domain's actions. More info https://t.ly/W8uuc
     nlu_importer = TrainingDataImporter.load_from_dict(
-        domain_path=args.domain, training_data_paths=args.data
+        training_data_paths=args.data, expand_env_vars=False
     )
     nlu_data = nlu_importer.get_nlu_data()
-
-    intents_from_files = nlu_data.intents
-
     nlu_examples = nlu_data.filter_training_examples(
-        lambda ex: ex.get("intent") in intents_from_files
+        lambda ex: ex.get("intent") in nlu_data.intents
     )
-
     nlu_examples_yaml = RasaYAMLWriter().dumps(nlu_examples)
 
+    assistant_name = _get_assistant_name(config)
+
     # Build GraphQL request
     graphql_req = build_import_request(
         assistant_name,
@@ -220,7 +263,9 @@ def upload_calm_assistant(
         nlu_yaml=nlu_examples_yaml,
     )
 
-    structlogger.info("Uploading to Rasa Studio...")
+    structlogger.info(
+        "rasa.studio.upload.calm", event_info="Uploading to Rasa Studio..."
+    )
     return make_request(endpoint, graphql_req, verify)
 
 
@@ -237,27 +282,37 @@ def upload_nlu_assistant(
             - intents: The intents to upload
             - entities: The entities to upload
         endpoint: The studio endpoint
+        verify: Whether to verify SSL
     Returns:
         None
     """
-    structlogger.info("Found DM1 assistant data, parsing...")
+    structlogger.info(
+        "rasa.studio.upload.nlu_data_read",
+        event_info="Found DM1 assistant data, parsing...",
+    )
     importer = TrainingDataImporter.load_from_dict(
-        domain_path=args.domain, training_data_paths=args.data, config_path=args.config
+        domain_path=args.domain,
+        training_data_paths=args.data,
+        config_path=args.config,
+        expand_env_vars=False,
     )
 
     intents_from_files = importer.get_nlu_data().intents
-    entities_from_files = importer.get_domain().entities
 
+    domain_from_files = importer.get_domain()
+    entities_from_files = domain_from_files.entities
     entities, intents = _get_selected_entities_and_intents(
         args, intents_from_files, entities_from_files
     )
 
     config_from_files = importer.get_config()
-    config = extract_values(config_from_files, config_keys)
+    config = extract_values(config_from_files, CONFIG_KEYS)
 
     assistant_name = _get_assistant_name(config)
 
-    structlogger.info("Validating data...")
+    structlogger.info(
+        "rasa.studio.upload.nlu_data_validate", event_info="Validating data..."
+    )
     _check_for_missing_primitives(
         intents, entities, intents_from_files, entities_from_files
     )
@@ -269,15 +324,39 @@ def upload_nlu_assistant(
     all_entities = _add_missing_entities(nlu_examples.entities, entities)
     nlu_examples_yaml = RasaYAMLWriter().dumps(nlu_examples)
 
-    domain = _filter_domain(all_entities, intents, importer.get_domain().as_dict())
+    domain = _filter_domain(all_entities, intents, domain_from_files.as_dict())
     domain_yaml = dump_obj_as_yaml_to_string(domain)
 
     graphql_req = build_request(assistant_name, nlu_examples_yaml, domain_yaml)
 
-    structlogger.info("Uploading to Rasa Studio...")
+    structlogger.info(
+        "rasa.studio.upload.nlu", event_info="Uploading to Rasa Studio..."
+    )
     return make_request(endpoint, graphql_req, verify)
 
 
+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,
+        {
+            "operationName": "LicenseDetails",
+            "query": (
+                "query LicenseDetails {\n"
+                "  licenseDetails {\n"
+                "    valid\n"
+                "    scopes\n"
+                "    __typename\n"
+                "  }\n"
+                "}"
+            ),
+            "variables": {},
+        },
+        verify,
+    )
+    return result.was_successful
+
+
 def make_request(endpoint: str, graphql_req: Dict, verify: bool = True) -> StudioResult:
     """Makes a request to the studio endpoint to upload data.
 
@@ -310,7 +389,12 @@ def _add_missing_entities(
     for entity in entities_from_intents:
         if entity not in entities:
             structlogger.warning(
-                f"Adding entity '{entity}' to upload since it is used in an intent."
+                "rasa.studio.upload.adding_missing_entity",
+                event_info=(
+                    f"Adding entity '{entity}' to upload "
+                    "since it is used in an intent."
+                ),
+                entity=entity,
             )
             all_entities.append(entity)
     return all_entities
@@ -378,7 +462,9 @@ def build_request(
 
 
 def _filter_domain(
-    entities: List[Union[str, Dict]], intents: List[str], domain_from_files: Dict
+    entities: List[Union[str, Dict]],
+    intents: List[str],
+    domain_from_files: Dict[str, Any],
 ) -> Dict:
     """Filters the domain to only include the selected entities and intents."""
     selected_entities = _remove_not_selected_entities(

rasa/telemetry.py

@@ -32,7 +32,13 @@ from rasa.constants import (
     CONFIG_TELEMETRY_ENABLED,
     CONFIG_TELEMETRY_ID,
 )
-from rasa.shared.constants import PROMPT_CONFIG_KEY, PROMPT_TEMPLATE_CONFIG_KEY
+from rasa.shared.constants import (
+    PROMPT_CONFIG_KEY,
+    PROMPT_TEMPLATE_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
 from rasa.shared.core.flows import Flow
@@ -106,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"
@@ -157,6 +164,7 @@ NUM_LINK_STEPS = "num_link_steps"
 NUM_CALL_STEPS = "num_call_steps"
 NUM_SHARED_SLOTS_BETWEEN_FLOWS = "num_shared_slots_between_flows"
 LLM_COMMAND_GENERATOR_MODEL_NAME = "llm_command_generator_model_name"
+LLM_COMMAND_GENERATOR_MODEL_GROUP_ID = "llm_command_generator_model_group_id"
 LLM_COMMAND_GENERATOR_CUSTOM_PROMPT_USED = "llm_command_generator_custom_prompt_used"
 MULTI_STEP_LLM_COMMAND_GENERATOR_HANDLE_FLOWS_PROMPT_USED = (
     "multi_step_llm_command_generator_custom_handle_flows_prompt_used"
@@ -166,6 +174,7 @@ MULTI_STEP_LLM_COMMAND_GENERATOR_FILL_SLOTS_PROMPT_USED = (
 )
 FLOW_RETRIEVAL_ENABLED = "flow_retrieval_enabled"
 FLOW_RETRIEVAL_EMBEDDING_MODEL_NAME = "flow_retrieval_embedding_model_name"
+FLOW_RETRIEVAL_EMBEDDING_MODEL_GROUP_ID = "flow_retrieval_embedding_model_group_id"
 TRACING_BACKEND = "tracing_backend"
 METRICS_BACKEND = "metrics_backend"
 VERSION = "version"
@@ -960,6 +969,13 @@ def track_model_training(
         "policies": config.get("policies"),
         "train_schema": config.get("train_schema"),
         "predict_schema": config.get("predict_schema"),
+        "model_groups": rasa.core.utils.AvailableEndpoints.get_instance().model_groups,
+        "api_health_check_enabled": (
+            os.getenv(
+                LLM_API_HEALTH_CHECK_ENV_VAR, LLM_API_HEALTH_CHECK_DEFAULT_VALUE
+            ).lower()
+            == "true"
+        ),
         "num_intent_examples": len(nlu_data.intent_examples),
         "num_entity_examples": len(nlu_data.entity_examples),
         "num_actions": len(domain.action_names_or_texts),
@@ -1114,46 +1130,78 @@ def _get_llm_command_generator_config(config: Dict[str, Any]) -> Optional[Dict]:
                 return component
         return None
 
-    def extract_settings(component: Dict) -> Dict:
-        """Extracts the settings from the command generator component."""
+    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, {})
-        llm_model_name = (
-            llm_config.get(MODEL_CONFIG_KEY)
-            or llm_config.get(MODEL_NAME_CONFIG_KEY)
-            or DEFAULT_LLM_CONFIG[MODEL_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
         )
+        if llm_model_group_id is None and llm_model_name is None:
+            llm_model_name = DEFAULT_LLM_CONFIG[MODEL_CONFIG_KEY]
+
+        custom_prompt_used = (
+            PROMPT_CONFIG_KEY in component or PROMPT_TEMPLATE_CONFIG_KEY in component
+        )
+        return {
+            LLM_COMMAND_GENERATOR_MODEL_NAME: llm_model_name,
+            LLM_COMMAND_GENERATOR_MODEL_GROUP_ID: llm_model_group_id,
+            LLM_COMMAND_GENERATOR_CUSTOM_PROMPT_USED: custom_prompt_used,
+        }
+
+    def extract_multistep_command_generator_prompt_settings(component: Dict) -> Dict:
+        """Extracts settings related to multistep command generator."""
+        prompt_templates = component.get("prompt_templates", {})
+        handle_flows_prompt_used = HANDLE_FLOWS_KEY in prompt_templates
+        fill_slots_prompt_used = FILL_SLOTS_KEY in prompt_templates
+        return {
+            MULTI_STEP_LLM_COMMAND_GENERATOR_HANDLE_FLOWS_PROMPT_USED: handle_flows_prompt_used,  # noqa: E501
+            MULTI_STEP_LLM_COMMAND_GENERATOR_FILL_SLOTS_PROMPT_USED: fill_slots_prompt_used,  # noqa: E501
+        }
+
+    def extract_flow_retrieval_settings(component: Dict) -> Dict:
+        """Extracts settings related to flow retrieval."""
         flow_retrieval_config = component.get(FLOW_RETRIEVAL_KEY, {})
         flow_retrieval_enabled = flow_retrieval_config.get("active", True)
-        flow_retrieval_embeddings_config = flow_retrieval_config.get(
+        embeddings_config = flow_retrieval_config.get(
             EMBEDDINGS_CONFIG_KEY, DEFAULT_EMBEDDINGS_CONFIG
         )
         flow_retrieval_embedding_model_name = (
             (
-                flow_retrieval_embeddings_config.get(MODEL_NAME_CONFIG_KEY)
-                or flow_retrieval_embeddings_config.get(MODEL_CONFIG_KEY)
+                embeddings_config.get(MODEL_NAME_CONFIG_KEY)
+                or embeddings_config.get(MODEL_CONFIG_KEY)
             )
             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
+        )
         return {
-            LLM_COMMAND_GENERATOR_MODEL_NAME: llm_model_name,
-            LLM_COMMAND_GENERATOR_CUSTOM_PROMPT_USED: PROMPT_CONFIG_KEY in component
-            or PROMPT_TEMPLATE_CONFIG_KEY in component,
-            MULTI_STEP_LLM_COMMAND_GENERATOR_HANDLE_FLOWS_PROMPT_USED: HANDLE_FLOWS_KEY
-            in component.get("prompt_templates", {}),
-            MULTI_STEP_LLM_COMMAND_GENERATOR_FILL_SLOTS_PROMPT_USED: FILL_SLOTS_KEY
-            in component.get("prompt_templates", {}),
             FLOW_RETRIEVAL_ENABLED: flow_retrieval_enabled,
             FLOW_RETRIEVAL_EMBEDDING_MODEL_NAME: flow_retrieval_embedding_model_name,
+            FLOW_RETRIEVAL_EMBEDDING_MODEL_GROUP_ID: flow_retrieval_embedding_model_group_id,  # noqa: E501
         }
 
+    def extract_settings(component: Dict) -> Dict:
+        """Extracts the settings from the command generator component."""
+        settings = {}
+        settings.update(extract_llm_command_generator_llm_client_settings(component))
+        settings.update(extract_multistep_command_generator_prompt_settings(component))
+        settings.update(extract_flow_retrieval_settings(component))
+        return settings
+
     command_generator_config = {
         LLM_COMMAND_GENERATOR_MODEL_NAME: None,
+        LLM_COMMAND_GENERATOR_MODEL_GROUP_ID: None,
         LLM_COMMAND_GENERATOR_CUSTOM_PROMPT_USED: None,
         MULTI_STEP_LLM_COMMAND_GENERATOR_HANDLE_FLOWS_PROMPT_USED: None,
         MULTI_STEP_LLM_COMMAND_GENERATOR_FILL_SLOTS_PROMPT_USED: None,
         FLOW_RETRIEVAL_ENABLED: None,
         FLOW_RETRIEVAL_EMBEDDING_MODEL_NAME: None,
+        FLOW_RETRIEVAL_EMBEDDING_MODEL_GROUP_ID: None,
     }
 
     pipeline = config.get("pipeline", [])
@@ -1333,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."""
@@ -1553,6 +1612,7 @@ def track_response_rephrase(
     custom_prompt_template: Optional[str],
     llm_type: Optional[str],
     llm_model: Optional[str],
+    llm_model_group_id: Optional[str],
 ) -> None:
     """Track when a user rephrases a response."""
     _track(
@@ -1562,6 +1622,7 @@ def track_response_rephrase(
             "custom_prompt_template": custom_prompt_template,
             "llm_type": llm_type,
             "llm_model": llm_model,
+            "llm_model_group_id": llm_model_group_id,
         },
     )
 
@@ -1576,8 +1637,10 @@ def track_intentless_policy_train() -> None:
 def track_intentless_policy_train_completed(
     embeddings_type: Optional[str],
     embeddings_model: Optional[str],
+    embeddings_model_group_id: Optional[str],
     llm_type: Optional[str],
     llm_model: Optional[str],
+    llm_model_group_id: Optional[str],
 ) -> None:
     """Track when a user trains a policy."""
     _track(
@@ -1585,8 +1648,10 @@ def track_intentless_policy_train_completed(
         {
             "embeddings_type": embeddings_type,
             "embeddings_model": embeddings_model,
+            "embeddings_model_group_id": embeddings_model_group_id,
             "llm_type": llm_type,
             "llm_model": llm_model,
+            "llm_model_group_id": llm_model_group_id,
         },
     )
 
@@ -1595,8 +1660,10 @@ def track_intentless_policy_train_completed(
 def track_intentless_policy_predict(
     embeddings_type: Optional[str],
     embeddings_model: Optional[str],
+    embeddings_model_group_id: Optional[str],
     llm_type: Optional[str],
     llm_model: Optional[str],
+    llm_model_group_id: Optional[str],
     score: float,
 ) -> None:
     """Track when a user trains a policy."""
@@ -1605,8 +1672,10 @@ def track_intentless_policy_predict(
         {
             "embeddings_type": embeddings_type,
             "embeddings_model": embeddings_model,
+            "embeddings_model_group_id": embeddings_model_group_id,
             "llm_type": llm_type,
             "llm_model": llm_model,
+            "llm_model_group_id": llm_model_group_id,
             "score": score,
         },
     )
@@ -1696,8 +1765,10 @@ def track_enterprise_search_policy_train_completed(
     vector_store_type: Optional[str],
     embeddings_type: Optional[str],
     embeddings_model: Optional[str],
+    embeddings_model_group_id: Optional[str],
     llm_type: Optional[str],
     llm_model: Optional[str],
+    llm_model_group_id: Optional[str],
     citation_enabled: Optional[bool],
 ) -> None:
     """Track when a user completes training Enterprise Search policy."""
@@ -1707,8 +1778,10 @@ def track_enterprise_search_policy_train_completed(
             "vector_store_type": vector_store_type,
             "embeddings_type": embeddings_type,
             "embeddings_model": embeddings_model,
+            "embeddings_model_group_id": embeddings_model_group_id,
             "llm_type": llm_type,
             "llm_model": llm_model,
+            "llm_model_group_id": llm_model_group_id,
             "citation_enabled": citation_enabled,
         },
     )
@@ -1719,8 +1792,10 @@ def track_enterprise_search_policy_predict(
     vector_store_type: Optional[str],
     embeddings_type: Optional[str],
     embeddings_model: Optional[str],
+    embeddings_model_group_id: Optional[str],
     llm_type: Optional[str],
     llm_model: Optional[str],
+    llm_model_group_id: Optional[str],
     citation_enabled: Optional[bool],
 ) -> None:
     """Track when a user predicts the next action using Enterprise Search policy."""
@@ -1730,8 +1805,10 @@ def track_enterprise_search_policy_predict(
             "vector_store_type": vector_store_type,
             "embeddings_type": embeddings_type,
             "embeddings_model": embeddings_model,
+            "embeddings_model_group_id": embeddings_model_group_id,
             "llm_type": llm_type,
             "llm_model": llm_model,
+            "llm_model_group_id": llm_model_group_id,
             "citation_enabled": citation_enabled,
         },
     )

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