rasa-pro 3.11.0rc3__py3-none-any.whl → 3.11.2__py3-none-any.whl

rasa/shared/utils/health_check/health_check.py

@@ -1,4 +1,5 @@
 import os
+import sys
 from typing import Optional, Dict, Any
 
 from rasa.shared.constants import (
@@ -9,7 +10,6 @@ from rasa.shared.constants import (
 from rasa.shared.exceptions import ProviderClientValidationError
 from rasa.shared.providers.embedding.embedding_client import EmbeddingClient
 from rasa.shared.providers.llm.llm_client import LLMClient
-from rasa.shared.utils.cli import print_error_and_exit
 from rasa.shared.utils.llm import llm_factory, structlogger, embedder_factory
 
 
@@ -25,15 +25,15 @@ def try_instantiate_llm_client(
     except (ProviderClientValidationError, ValueError) as e:
         structlogger.error(
             f"{log_source_function}.llm_instantiation_failed",
-            message="Unable to instantiate LLM client.",
+            event_info=(
+                f"Unable to create the LLM client for component - "
+                f"{log_source_component}. "
+                f"Please make sure you specified the required environment variables "
+                f"and configuration keys. "
+            ),
             error=e,
         )
-        print_error_and_exit(
-            f"Unable to create the LLM client for component - {log_source_component}. "
-            f"Please make sure you specified the required environment variables "
-            f"and configuration keys. "
-            f"Error: {e}"
-        )
+        sys.exit(1)
 
 
 def try_instantiate_embedder(
@@ -48,14 +48,14 @@ def try_instantiate_embedder(
     except (ProviderClientValidationError, ValueError) as e:
         structlogger.error(
             f"{log_source_function}.embedder_instantiation_failed",
-            message="Unable to instantiate Embedding client.",
+            event_info=(
+                f"Unable to create the Embedding client for component - "
+                f"{log_source_component}. Please make sure you specified the required "
+                f"environment variables and configuration keys."
+            ),
             error=e,
         )
-        print_error_and_exit(
-            f"Unable to create the Embedding client for component - "
-            f"{log_source_component}. Please make sure you specified the required "
-            f"environment variables and configuration keys. Error: {e}"
-        )
+        sys.exit(1)
 
 
 def perform_llm_health_check(
@@ -202,13 +202,14 @@ def send_test_llm_api_request(
     except Exception as e:
         structlogger.error(
             f"{log_source_function}.send_test_llm_api_request_failed",
-            event_info="Test call to the LLM API failed.",
+            event_info=(
+                f"Test call to the LLM API failed for component - "
+                f"{log_source_component}.",
+            ),
+            config=llm_client.config,
             error=e,
         )
-        print_error_and_exit(
-            f"Test call to the LLM API failed for component - {log_source_component}. "
-            f"Error: {e}"
-        )
+        sys.exit(1)
 
 
 def send_test_embeddings_api_request(
@@ -232,13 +233,14 @@ def send_test_embeddings_api_request(
     except Exception as e:
         structlogger.error(
             f"{log_source_function}.send_test_llm_api_request_failed",
-            event_info="Test call to the Embeddings API failed.",
+            event_info=(
+                f"Test call to the Embeddings API failed for component - "
+                f"{log_source_component}."
+            ),
+            config=embedder.config,
             error=e,
         )
-        print_error_and_exit(
-            f"Test call to the Embeddings API failed for component - "
-            f"{log_source_component}. Error: {e}"
-        )
+        sys.exit(1)
 
 
 def is_api_health_check_enabled() -> bool:

rasa/shared/utils/yaml.py

@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 import datetime
 import io
 import logging
@@ -9,7 +10,7 @@ from dataclasses import field
 from functools import lru_cache
 from io import StringIO
 from pathlib import Path
-from typing import Any, List, Optional, Tuple, Dict, Callable, Union
+from typing import Any, Generator, List, Optional, Tuple, Dict, Callable, Union
 
 import jsonschema
 from importlib_resources import files
@@ -26,7 +27,7 @@ from rasa.shared.constants import (
     LATEST_TRAINING_DATA_FORMAT_VERSION,
     SCHEMA_EXTENSIONS_FILE,
     RESPONSES_SCHEMA_FILE,
-    API_KEY,
+    SENSITIVE_DATA,
 )
 from rasa.shared.exceptions import (
     YamlException,
@@ -59,7 +60,6 @@ YAML_VERSION = (1, 2)
 READ_YAML_FILE_CACHE_MAXSIZE = os.environ.get(
     READ_YAML_FILE_CACHE_MAXSIZE_ENV_VAR, DEFAULT_READ_YAML_FILE_CACHE_MAXSIZE
 )
-SENSITIVE_DATA = [API_KEY]
 
 
 @dataclass
@@ -88,13 +88,19 @@ def fix_yaml_loader() -> None:
     yaml.Loader.add_constructor("tag:yaml.org,2002:str", construct_yaml_str)
     yaml.SafeLoader.add_constructor("tag:yaml.org,2002:str", construct_yaml_str)
 
+    _add_env_var_resolver()
 
-def replace_environment_variables() -> None:
-    """Enable yaml loader to process the environment variables in the yaml."""
+
+def _add_env_var_resolver() -> None:
+    """Enable yaml loader to detect the environment variables in the yaml."""
     # eg. ${USER_NAME}, ${PASSWORD}
     env_var_pattern = re.compile(r"^(.*)\$\{(.*)\}(.*)$")
     yaml.Resolver.add_implicit_resolver("!env_var", env_var_pattern, None)
 
+
+def _add_yaml_constructor_to_replace_environment_variables() -> None:
+    """Enable yaml loader to replace the environment variables in the yaml."""
+
     def env_var_constructor(loader: BaseConstructor, node: ScalarNode) -> str:
         """Process environment variables found in the YAML."""
         value = loader.construct_scalar(node)
@@ -121,7 +127,6 @@ def replace_environment_variables() -> None:
 
 
 fix_yaml_loader()
-replace_environment_variables()
 
 
 class YamlValidationException(YamlException, ValueError):
@@ -343,7 +348,7 @@ class YamlValidationException(YamlException, ValueError):
 
 
 def read_schema_file(
-    schema_file: str, package_name: str = PACKAGE_NAME
+    schema_file: str, package_name: str = PACKAGE_NAME, expand_env_vars: bool = True
 ) -> Union[List[Any], Dict[str, Any]]:
     """Read a schema file from the package.
 
@@ -351,12 +356,13 @@ def read_schema_file(
         schema_file: The schema file to read.
         package_name: the name of the package the schema is located in. defaults
             to `rasa`.
+        expand_env_vars: Whether to expand environment variables in the file.
 
     Returns:
         The schema as a dictionary.
     """
     schema_path = str(files(package_name).joinpath(schema_file))
-    return read_yaml_file(schema_path)
+    return read_yaml_file(schema_path, expand_env_vars=expand_env_vars)
 
 
 def parse_raw_yaml(raw_yaml_content: str) -> Dict[str, Any]:
@@ -430,6 +436,7 @@ def validate_raw_yaml_using_schema(
     raw_yaml_content: str,
     schema_content: Dict[str, Any],
     schema_extensions: Optional[List[str]] = None,
+    expand_env_vars: bool = True,
 ) -> None:
     """Validate raw yaml content using a schema.
 
@@ -439,6 +446,7 @@ def validate_raw_yaml_using_schema(
         raw_yaml_content: the raw YAML content to be validated (usually a string)
         schema_content: the schema for the yaml_file_content
         schema_extensions: pykwalify schema extension files
+        expand_env_vars: Whether to expand environment variables.
     """
     try:
         # we need "rt" since
@@ -446,7 +454,11 @@ def validate_raw_yaml_using_schema(
         # will include e.g. at which line an object was parsed. this is very
         # helpful when we validate files later on and want to point the user to the
         # right line
-        yaml_data = read_yaml(raw_yaml_content, reader_type=["safe", "rt"])
+        yaml_data = read_yaml(
+            raw_yaml_content,
+            reader_type=["safe", "rt"],
+            expand_env_vars=expand_env_vars,
+        )
     except (YAMLError, DuplicateKeyError) as e:
         raise YamlSyntaxException(underlying_yaml_exception=e)
 
@@ -454,7 +466,10 @@ def validate_raw_yaml_using_schema(
 
 
 def validate_raw_yaml_using_schema_file(
-    raw_yaml_content: str, schema_path: str, package_name: str = PACKAGE_NAME
+    raw_yaml_content: str,
+    schema_path: str,
+    package_name: str = PACKAGE_NAME,
+    expand_env_vars: bool = True,
 ) -> None:
     """Validate raw yaml content using a schema from file.
 
@@ -463,15 +478,21 @@ def validate_raw_yaml_using_schema_file(
         schema_path: the schema used for validation
         package_name: the name of the package the schema is located in. defaults
             to `rasa`.
+        expand_env_vars: Whether to expand environment variables in the file.
     """
-    schema_content = read_schema_file(schema_path, package_name)
-    validate_raw_yaml_using_schema(raw_yaml_content, schema_content)
+    schema_content = read_schema_file(
+        schema_path, package_name, expand_env_vars=expand_env_vars
+    )
+    validate_raw_yaml_using_schema(
+        raw_yaml_content, schema_content, expand_env_vars=expand_env_vars
+    )
 
 
 def validate_raw_yaml_content_using_schema_with_responses(
     raw_yaml_content: str,
     schema_content: Union[List[Any], Dict[str, Any]],
     package_name: str = PACKAGE_NAME,
+    expand_env_vars: bool = True,
 ) -> None:
     """Validate raw yaml content using a schema with responses sub-schema.
 
@@ -480,18 +501,26 @@ def validate_raw_yaml_content_using_schema_with_responses(
         schema_content: the content of the YAML schema
         package_name: the name of the package the schema is located in. defaults
         to `rasa`.
+        expand_env_vars: Whether to expand environment variables in the file.
     """
     # bot responses are part of the schema extension
     # it will be included if the schema explicitly references it with include: responses
-    bot_responses_schema_content = read_schema_file(RESPONSES_SCHEMA_FILE, package_name)
+    bot_responses_schema_content = read_schema_file(
+        RESPONSES_SCHEMA_FILE, package_name, expand_env_vars=expand_env_vars
+    )
     schema_content = dict(schema_content, **bot_responses_schema_content)
     schema_extensions = [str(files(package_name).joinpath(SCHEMA_EXTENSIONS_FILE))]
 
-    validate_raw_yaml_using_schema(raw_yaml_content, schema_content, schema_extensions)
+    validate_raw_yaml_using_schema(
+        raw_yaml_content, schema_content, schema_extensions, expand_env_vars
+    )
 
 
 def validate_raw_yaml_using_schema_file_with_responses(
-    raw_yaml_content: str, schema_path: str, package_name: str = PACKAGE_NAME
+    raw_yaml_content: str,
+    schema_path: str,
+    package_name: str = PACKAGE_NAME,
+    expand_env_vars: bool = True,
 ) -> None:
     """Validate domain yaml content using a schema from file with responses sub-schema.
 
@@ -500,13 +529,32 @@ def validate_raw_yaml_using_schema_file_with_responses(
         schema_path: the schema of the yaml file
         package_name: the name of the package the schema is located in. defaults
             to `rasa`.
+        expand_env_vars: Whether to expand environment variables in the file.
     """
-    schema_content = read_schema_file(schema_path, package_name)
+    schema_content = read_schema_file(schema_path, package_name, expand_env_vars)
     validate_raw_yaml_content_using_schema_with_responses(
-        raw_yaml_content, schema_content, package_name
+        raw_yaml_content, schema_content, package_name, expand_env_vars
     )
 
 
+@contextmanager
+def environment_variables_replaced(
+    yaml_parser: yaml.YAML,
+) -> Generator[None, None, None]:
+    """Replace environment variables during yaml loading.
+
+    Resets the environment variable constructor after the context manager exits.
+    """
+    try:
+        _add_yaml_constructor_to_replace_environment_variables()
+        yield
+    finally:
+        # replace env var constructor with one that does not expand env vars
+        yaml_parser.constructor.add_constructor(
+            "!env_var", lambda loader, node: loader.construct_scalar(node)
+        )
+
+
 def read_yaml(
     content: str,
     reader_type: Union[str, List[str]] = "safe",
@@ -523,12 +571,17 @@ def read_yaml(
         ruamel.yaml.parser.ParserError: If there was an error when parsing the YAML.
     """
     custom_constructor = kwargs.get("custom_constructor", None)
+    expand_env_vars = kwargs.get("expand_env_vars", True)
 
     # Create YAML parser with custom constructor
     yaml_parser, reset_constructors = create_yaml_parser(
         reader_type, custom_constructor
     )
-    yaml_content = yaml_parser.load(content) or {}
+    if expand_env_vars:
+        with environment_variables_replaced(yaml_parser):
+            yaml_content = yaml_parser.load(content) or {}
+    else:
+        yaml_content = yaml_parser.load(content) or {}
 
     # Reset to default constructors
     reset_constructors()
@@ -544,9 +597,9 @@ def create_yaml_parser(
 
     Args:
         reader_type (str): The type of the reader
-        (e.g., 'safe', 'rt', 'unsafe').
+            (e.g., 'safe', 'rt', 'unsafe').
         custom_constructor (Optional[Callable]):
-        A custom constructor function for YAML parsing.
+            A custom constructor function for YAML parsing.
 
     Returns:
         Tuple[yaml.YAML, Callable[[], None]]: A tuple containing
@@ -617,7 +670,9 @@ def _is_ascii(text: str) -> bool:
 
 @lru_cache(maxsize=READ_YAML_FILE_CACHE_MAXSIZE)
 def read_yaml_file(
-    filename: Union[str, Path], reader_type: Union[str, Tuple[str]] = "safe"
+    filename: Union[str, Path],
+    reader_type: Union[str, Tuple[str]] = "safe",
+    expand_env_vars: bool = True,
 ) -> Union[List[Any], Dict[str, Any]]:
     """Parses a yaml file.
 
@@ -626,6 +681,7 @@ def read_yaml_file(
     Args:
         filename: The path to the file which should be read.
         reader_type: Reader type to use. By default "safe" will be used.
+        expand_env_vars: Whether to expand environment variables in the file.
 
     Returns:
         Parsed content of the file.
@@ -634,7 +690,11 @@ def read_yaml_file(
         fixed_reader_type = (
             list(reader_type) if isinstance(reader_type, tuple) else reader_type
         )
-        return read_yaml(read_file(filename, DEFAULT_ENCODING), fixed_reader_type)
+        return read_yaml(
+            read_file(filename, DEFAULT_ENCODING),
+            fixed_reader_type,
+            expand_env_vars=expand_env_vars,
+        )
     except (YAMLError, DuplicateKeyError) as e:
         raise YamlSyntaxException(filename, e)
 
@@ -657,11 +717,14 @@ def read_config_file(
     return read_validated_yaml(filename, CONFIG_SCHEMA_FILE, reader_type)
 
 
-def read_model_configuration(filename: Union[Path, str]) -> Dict[str, Any]:
+def read_model_configuration(
+    filename: Union[Path, str], expand_env_vars: bool = True
+) -> Dict[str, Any]:
     """Parses a model configuration file.
 
     Args:
         filename: The path to the file which should be read.
+        expand_env_vars: Whether to expand environment variables in the file.
 
     Raises:
         YamlValidationException: In case the model configuration doesn't match the
@@ -670,24 +733,34 @@ def read_model_configuration(filename: Union[Path, str]) -> Dict[str, Any]:
     Returns:
         Parsed config file.
     """
-    return read_validated_yaml(filename, MODEL_CONFIG_SCHEMA_FILE)
+    return read_validated_yaml(
+        filename, MODEL_CONFIG_SCHEMA_FILE, expand_env_vars=expand_env_vars
+    )
 
 
 def dump_obj_as_yaml_to_string(
-    obj: Any, should_preserve_key_order: bool = False
+    obj: Any,
+    should_preserve_key_order: bool = False,
+    transform: Optional[Callable] = None,
 ) -> str:
     """Writes data (python dict) to a yaml string.
 
     Args:
         obj: The object to dump. Has to be serializable.
         should_preserve_key_order: Whether to force preserve key order in `data`.
+        transform: A function to transform the data before writing it to the file.
 
     Returns:
         The object converted to a YAML string.
     """
     buffer = StringIO()
 
-    write_yaml(obj, buffer, should_preserve_key_order=should_preserve_key_order)
+    write_yaml(
+        obj,
+        buffer,
+        should_preserve_key_order=should_preserve_key_order,
+        transform=transform,
+    )
 
     return buffer.getvalue()
 
@@ -708,6 +781,7 @@ def write_yaml(
     data: Any,
     target: Union[str, Path, StringIO],
     should_preserve_key_order: bool = False,
+    transform: Optional[Callable[[Any], Any]] = None,
 ) -> None:
     """Writes a yaml to the file or to the stream.
 
@@ -715,6 +789,7 @@ def write_yaml(
         data: The data to write.
         target: The path to the file which should be written or a stream object
         should_preserve_key_order: Whether to force preserve key order in `data`.
+        transform: A function to transform the data before writing it to the file.
     """
     _enable_ordered_dict_yaml_dumping()
 
@@ -732,11 +807,11 @@ def write_yaml(
     )
 
     if isinstance(target, StringIO):
-        dumper.dump(data, target)
+        dumper.dump(data, target, transform=transform)
         return
 
     with Path(target).open("w", encoding=DEFAULT_ENCODING) as outfile:
-        dumper.dump(data, outfile)
+        dumper.dump(data, outfile, transform=transform)
 
 
 def is_key_in_yaml(file_path: Union[str, Path], *keys: str) -> bool:
@@ -768,6 +843,7 @@ def read_validated_yaml(
     filename: Union[str, Path],
     schema: str,
     reader_type: Union[str, List[str]] = "safe",
+    expand_env_vars: bool = True,
 ) -> Any:
     """Validates YAML file content and returns parsed content.
 
@@ -776,6 +852,7 @@ def read_validated_yaml(
         schema: The path to the schema file which should be used for validating the
             file content.
         reader_type: Reader type to use. By default, "safe" will be used.
+        expand_env_vars: Whether to expand environment variables in the file.
 
     Returns:
         The parsed file content.
@@ -786,8 +863,10 @@ def read_validated_yaml(
     """
     content = read_file(filename)
 
-    validate_raw_yaml_using_schema_file(content, schema)
-    return read_yaml(content, reader_type)
+    validate_raw_yaml_using_schema_file(
+        content, schema, expand_env_vars=expand_env_vars
+    )
+    return read_yaml(content, reader_type, expand_env_vars=expand_env_vars)
 
 
 def validate_training_data(json_data: Dict[str, Any], schema: Dict[str, Any]) -> None:
@@ -906,6 +985,7 @@ def validate_yaml_with_jsonschema(
     humanize_error: Callable[
         [jsonschema.ValidationError], str
     ] = default_error_humanizer,
+    expand_env_vars: bool = True,
 ) -> None:
     """Validate data format.
 
@@ -916,6 +996,7 @@ def validate_yaml_with_jsonschema(
             to `rasa`.
         humanize_error: a function to convert a jsonschema.ValidationError into a
             human-readable error message. Defaults to `default_error_humanizer`.
+        expand_env_vars: Whether to expand environment variables in the file.
 
     Raises:
         YamlSyntaxException: if the yaml file is not valid.
@@ -933,7 +1014,11 @@ def validate_yaml_with_jsonschema(
         # will include e.g. at which line an object was parsed. this is very
         # helpful when we validate files later on and want to point the user to the
         # right line
-        source_data = read_yaml(yaml_file_content, reader_type=["safe", "rt"])
+        source_data = read_yaml(
+            yaml_file_content,
+            reader_type=["safe", "rt"],
+            expand_env_vars=expand_env_vars,
+        )
     except (YAMLError, DuplicateKeyError) as e:
         raise YamlSyntaxException(underlying_yaml_exception=e)
 

rasa/studio/data_handler.py

@@ -360,7 +360,9 @@ def import_data_from_studio(
     studio_domain.persist(domain_file)
 
     data_from_studio = TrainingDataImporter.load_from_dict(
-        domain_path=str(domain_file), training_data_paths=data_paths
+        domain_path=str(domain_file),
+        training_data_paths=data_paths,
+        expand_env_vars=False,
     )
 
     return data_from_studio, data_original