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

rasa/studio/upload.py

@@ -1,9 +1,11 @@
 import argparse
 import base64
-import logging
-from typing import Dict, Iterable, List, Set, Text, Tuple, Union
+import sys
+from typing import Dict, Iterable, List, Set, Text, Tuple, Union, Any
 
+import questionary
 import requests
+import structlog
 
 import rasa.cli.telemetry
 import rasa.cli.utils
@@ -11,16 +13,19 @@ import rasa.shared.utils.cli
 import rasa.shared.utils.io
 from rasa.shared.constants import (
     DEFAULT_DOMAIN_PATHS,
+    DEFAULT_CONFIG_PATH,
 )
 from rasa.shared.core.flows.yaml_flows_io import 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
+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
+from rasa.studio.results_logger import StudioResult, with_studio_error_handler
 
-logger = logging.getLogger(__name__)
+structlogger = structlog.get_logger()
 
 
 def _get_selected_entities_and_intents(
@@ -32,37 +37,59 @@ def _get_selected_entities_and_intents(
 
     if entities is None or len(entities) == 0:
         entities = entities_from_files
-        logger.info("No entities specified. Using all entities from files.")
+        structlogger.info(
+            "rasa.studio.upload.entities_empty",
+            event_info="No entities specified. Using all entities from files.",
+        )
 
     intents = args.intents
 
     if intents is None or len(intents) == 0:
         intents = intents_from_files
-        logger.info("No intents specified. Using all intents from files.")
+        structlogger.info(
+            "rasa.studio.upload.intents_empty",
+            event_info="No intents specified. Using all intents from files.",
+        )
 
     return list(entities), list(intents)
 
 
 def handle_upload(args: argparse.Namespace) -> None:
     """Uploads primitives to rasa studio."""
-    assistant_name = args.assistant_name[0]
     endpoint = StudioConfig.read_config().studio_url
     if not endpoint:
         rasa.shared.utils.cli.print_error_and_exit(
             "No GraphQL endpoint found in config. Please run `rasa studio config`."
         )
     else:
-        logger.info("Loading data...")
+        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
         )
 
+        args.config = rasa.cli.utils.get_validated_path(
+            args.config, "config", DEFAULT_CONFIG_PATH
+        )
+
         # check safely if args.calm is set and not fail if not
         if hasattr(args, "calm") and args.calm:
-            upload_calm_assistant(args, assistant_name, endpoint)
+            upload_calm_assistant(args, endpoint)
         else:
-            upload_nlu_assistant(args, assistant_name, endpoint)
+            upload_nlu_assistant(args, endpoint)
+
+
+config_keys = [
+    "recipe",
+    "language",
+    "pipeline",
+    "llm",
+    "policies",
+    "model_name",
+    "assistant_id",
+]
 
 
 def extract_values(data: Dict, keys: List[Text]) -> Dict:
@@ -70,9 +97,36 @@ def extract_values(data: Dict, keys: List[Text]) -> Dict:
     return {key: data.get(key) for key in keys if data.get(key)}
 
 
-def upload_calm_assistant(
-    args: argparse.Namespace, assistant_name: str, endpoint: str
-) -> None:
+def _get_assistant_name(config: Dict[Text, Any]) -> str:
+    config_assistant_id = config.get("assistant_id", "")
+    assistant_name = questionary.text(
+        "Please provide the assistant name", default=config_assistant_id
+    ).ask()
+    if not assistant_name:
+        structlogger.error(
+            "rasa.studio.upload.assistant_name_empty",
+            event_info="Assistant name cannot be empty.",
+        )
+        sys.exit(1)
+
+    # if assistant_name exists and different from config assistant_id,
+    # notify user and upload with new assistant_name
+    same = assistant_name == config_assistant_id
+    if not same and config_assistant_id != "":
+        structlogger.info(
+            "rasa.studio.upload.assistant_name_mismatch",
+            event_info=(
+                f"Assistant name '{assistant_name}' is different"
+                f" from the one in the config file: '{config_assistant_id}'."
+            ),
+        )
+
+    structlogger.info(f"Uploading assistant with the name '{assistant_name}'.")
+    return assistant_name
+
+
+@with_studio_error_handler
+def upload_calm_assistant(args: argparse.Namespace, endpoint: str) -> StudioResult:
     """Uploads the CALM assistant data to Rasa Studio.
 
     Args:
@@ -82,101 +136,92 @@ def upload_calm_assistant(
             - flows: The path to the flows
             - endpoints: The path to the endpoints
             - config: The path to the config
-        assistant_name: The name of the assistant
         endpoint: The studio endpoint
     Returns:
         None
     """
-    logger.info("Parsing CALM assistant data...")
+    structlogger.info(
+        "rasa.studio.upload.loading_data", event_info="Parsing CALM assistant data..."
+    )
 
-    try:
-        importer = TrainingDataImporter.load_from_dict(
-            domain_path=args.domain,
-            config_path=args.config,
-        )
+    importer = TrainingDataImporter.load_from_dict(
+        domain_path=args.domain,
+        config_path=args.config,
+    )
 
-        # Prepare config and domain
-        config_from_files = importer.get_config()
-        domain_from_files = importer.get_domain().as_dict()
-
-        # Extract domain and config values
-        domain_keys = [
-            "version",
-            "actions",
-            "responses",
-            "slots",
-            "intents",
-            "entities",
-            "forms",
-            "session_config",
-        ]
-        config_keys = [
-            "recipe",
-            "language",
-            "pipeline",
-            "llm",
-            "policies",
-            "model_name",
-        ]
-
-        domain = extract_values(domain_from_files, domain_keys)
-        config = extract_values(config_from_files, config_keys)
-
-        training_data_paths = args.data
-
-        if isinstance(training_data_paths, list):
-            training_data_paths.append(args.flows)
-        elif isinstance(training_data_paths, str):
+    # Prepare config and domain
+    config = importer.get_config()
+    domain_from_files = importer.get_domain().as_dict()
+    endpoints_from_files = read_yaml_file(args.endpoints)
+    config_from_files = read_yaml_file(args.config)
+
+    # 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")
 
-        # Prepare flows
-        flow_importer = FlowSyncImporter.load_from_dict(
-            training_data_paths=training_data_paths
-        )
-
-        user_flows = flow_importer.get_flows().user_flows
-        flows = list(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
-        )
-        nlu_data = nlu_importer.get_nlu_data()
+    # Prepare flows
+    flow_importer = FlowSyncImporter.load_from_dict(
+        training_data_paths=training_data_paths
+    )
 
-        intents_from_files = nlu_data.intents
+    user_flows = flow_importer.get_flows().user_flows
+    flows = list(user_flows)
 
-        nlu_examples = nlu_data.filter_training_examples(
-            lambda ex: ex.get("intent") in intents_from_files
-        )
+    # 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
+    )
+    nlu_data = nlu_importer.get_nlu_data()
 
-        nlu_examples_yaml = RasaYAMLWriter().dumps(nlu_examples)
+    intents_from_files = nlu_data.intents
 
-        # Build GraphQL request
-        graphql_req = build_import_request(
-            assistant_name,
-            flows_yaml=YamlFlowsWriter().dumps(flows),
-            domain_yaml=dump_obj_as_yaml_to_string(domain),
-            config_yaml=dump_obj_as_yaml_to_string(config),
-            nlu_yaml=nlu_examples_yaml,
-        )
+    nlu_examples = nlu_data.filter_training_examples(
+        lambda ex: ex.get("intent") in intents_from_files
+    )
 
-        logger.info("Uploading to Rasa Studio...")
-        response, status = make_request(endpoint, graphql_req)
+    nlu_examples_yaml = RasaYAMLWriter().dumps(nlu_examples)
 
-        if status:
-            rasa.shared.utils.cli.print_success(response)
-        else:
-            logger.error(f"Failed to upload to Rasa Studio: {response}")
-            rasa.shared.utils.cli.print_error(response)
+    # Build GraphQL request
+    graphql_req = build_import_request(
+        assistant_name,
+        flows_yaml=YamlFlowsWriter().dumps(flows),
+        domain_yaml=dump_obj_as_yaml_to_string(domain),
+        config_yaml=dump_obj_as_yaml_to_string(config_from_files),
+        endpoints=dump_obj_as_yaml_to_string(endpoints_from_files),
+        nlu_yaml=nlu_examples_yaml,
+    )
 
-    except Exception as e:
-        logger.error(f"An error occurred while uploading the CALM assistant: {e}")
+    structlogger.info("Uploading to Rasa Studio...")
+    return make_request(endpoint, graphql_req)
 
 
-def upload_nlu_assistant(
-    args: argparse.Namespace, assistant_name: str, endpoint: str
-) -> None:
+@with_studio_error_handler
+def upload_nlu_assistant(args: argparse.Namespace, endpoint: str) -> StudioResult:
     """Uploads the classic (dm1) assistant data to Rasa Studio.
 
     Args:
@@ -185,14 +230,13 @@ def upload_nlu_assistant(
             - domain: The path to the domain
             - intents: The intents to upload
             - entities: The entities to upload
-        assistant_name: The name of the assistant
         endpoint: The studio endpoint
     Returns:
         None
     """
-    logger.info("Found DM1 assistant data, parsing...")
+    structlogger.info("Found DM1 assistant data, parsing...")
     importer = TrainingDataImporter.load_from_dict(
-        domain_path=args.domain, training_data_paths=args.data
+        domain_path=args.domain, training_data_paths=args.data, config_path=args.config
     )
 
     intents_from_files = importer.get_nlu_data().intents
@@ -202,7 +246,12 @@ def upload_nlu_assistant(
         args, intents_from_files, entities_from_files
     )
 
-    logger.info("Validating data...")
+    config_from_files = importer.get_config()
+    config = extract_values(config_from_files, config_keys)
+
+    assistant_name = _get_assistant_name(config)
+
+    structlogger.info("Validating data...")
     _check_for_missing_primitives(
         intents, entities, intents_from_files, entities_from_files
     )
@@ -219,15 +268,11 @@ def upload_nlu_assistant(
 
     graphql_req = build_request(assistant_name, nlu_examples_yaml, domain_yaml)
 
-    logger.info("Uploading to Rasa Studio...")
-    response, status = make_request(endpoint, graphql_req)
-    if status:
-        rasa.shared.utils.cli.print_success(response)
-    else:
-        rasa.shared.utils.cli.print_error(response)
+    structlogger.info("Uploading to Rasa Studio...")
+    return make_request(endpoint, graphql_req)
 
 
-def make_request(endpoint: str, graphql_req: Dict) -> Tuple[str, bool]:
+def make_request(endpoint: str, graphql_req: Dict) -> StudioResult:
     """Makes a request to the studio endpoint to upload data.
 
     Args:
@@ -235,7 +280,6 @@ def make_request(endpoint: str, graphql_req: Dict) -> Tuple[str, bool]:
         graphql_req: The graphql request
     """
     token = KeycloakTokenReader().get_token()
-
     res = requests.post(
         endpoint,
         json=graphql_req,
@@ -245,30 +289,9 @@ def make_request(endpoint: str, graphql_req: Dict) -> Tuple[str, bool]:
         },
     )
 
-    response = res.json()
-
-    if _response_has_errors(response):
-        print_errors_from_response(response)
-
-    if res.status_code != 200:
-        return f"Upload failed with status code {res.status_code}", False
-    elif _response_has_errors(response):
-        return "Error while uploading data!", False
-
-    return "Upload successful!", True
-
-
-def _response_has_errors(response: Dict) -> bool:
-    return (
-        "errors" in response
-        and isinstance(response["errors"], list)
-        and len(response["errors"]) > 0
-    )
-
-
-def print_errors_from_response(response: Dict) -> None:
-    for error in response["errors"]:
-        logger.error(error["message"])
+    if results_logger.response_has_errors(res.json()):
+        return StudioResult.error(res.json())
+    return StudioResult.success("Upload successful")
 
 
 def _add_missing_entities(
@@ -278,7 +301,7 @@ def _add_missing_entities(
     all_entities.extend(entities)
     for entity in entities_from_intents:
         if entity not in entities:
-            logger.warning(
+            structlogger.warning(
                 f"Adding entity '{entity}' to upload since it is used in an intent."
             )
             all_entities.append(entity)
@@ -290,6 +313,7 @@ def build_import_request(
     flows_yaml: str,
     domain_yaml: str,
     config_yaml: str,
+    endpoints: str,
     nlu_yaml: str = "",
 ) -> Dict:
     # b64encode expects bytes and returns bytes so we need to decode to string
@@ -297,6 +321,7 @@ def build_import_request(
     base64_flows = base64.b64encode(flows_yaml.encode("utf-8")).decode("utf-8")
     base64_config = base64.b64encode(config_yaml.encode("utf-8")).decode("utf-8")
     base64_nlu = base64.b64encode(nlu_yaml.encode("utf-8")).decode("utf-8")
+    base64_endpoints = base64.b64encode(endpoints.encode("utf-8")).decode("utf-8")
 
     graphql_req = {
         "query": (
@@ -310,6 +335,7 @@ def build_import_request(
                 "flows": base64_flows,
                 "nlu": base64_nlu,
                 "config": base64_config,
+                "endpoints": base64_endpoints,
             }
         },
     }
@@ -347,16 +373,15 @@ def _filter_domain(
     entities: List[Union[str, Dict]], intents: List[str], domain_from_files: Dict
 ) -> Dict:
     """Filters the domain to only include the selected entities and intents."""
-    domain = {
+    selected_entities = _remove_not_selected_entities(
+        entities, domain_from_files.get("entities", [])
+    )
+    return {
         "version": domain_from_files["version"],
         "intents": intents,
-        "entities": _remove_not_selected_entities(
-            entities, domain_from_files["entities"]
-        ),
+        "entities": selected_entities,
     }
 
-    return domain
-
 
 def _check_for_missing_primitives(
     intents: Iterable[str],