data-designer 0.1.0__py3-none-any.whl

data_designer/engine/processing/ginja/environment.py

@@ -0,0 +1,461 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Callable
+from functools import partial, wraps
+import re
+from typing import Any
+
+from jinja2 import meta
+from jinja2 import nodes as j_nodes
+from jinja2.exceptions import SecurityError, TemplateSyntaxError
+from jinja2.nodes import Template
+from jinja2.sandbox import ImmutableSandboxedEnvironment
+from jsonpath_rust_bindings import Finder
+
+from data_designer.engine.processing.ginja.ast import (
+    ast_count_name_references,
+    ast_descendant_count,
+    ast_max_depth,
+)
+from data_designer.engine.processing.ginja.exceptions import (
+    UserTemplateError,
+    UserTemplateUnsupportedFiltersError,
+    maybe_handle_missing_filter_exception,
+)
+from data_designer.engine.processing.ginja.record import sanitize_record
+
+MAX_RENDERED_LEN = 512_000
+MAX_AST_NODE_COUNT = 600
+MAX_AST_DEPTH = 10
+ALLOWED_JINJA_FILTERS = [
+    ## Jinja2 Builtin Filters
+    "abs",
+    "capitalize",
+    "escape",
+    "first",
+    "float",
+    "forceescape",
+    "int",
+    "items",
+    "last",
+    "length",
+    "list",
+    "lower",
+    "max",
+    "min",
+    "random",
+    "replace",
+    "reverse",
+    "round",
+    "sort",
+    "string",
+    "title",
+    "trim",
+    "truncate",
+    "unique",
+    "urlencode",
+    ## Custom Filters
+    "jsonpath",
+]
+
+USER_PROMPT_TEMPLATE_ERROR_MESSAGE = """\
+User provided prompt generation template is invalid.\
+"""
+UNSUPPORTED_AST_NODES = [
+    j_nodes.Import,  # No {% include ... %}
+    j_nodes.Macro,  # No {% macro ... %}
+    j_nodes.Assign,  # No {% set ... %}
+    j_nodes.Extends,  # No {% extends ... %}
+    j_nodes.Block,  # No {% block ... %}
+]
+
+
+def jsonpath_jinja_filter(data: dict, expression: str) -> list[Any]:
+    """Defines JSONPath-based operations on variables.
+
+    Args:
+        data (dict): data object to filter.
+        expression (str): a valid JSONPath string.
+
+    Returns:
+        list[Any]: A list of JSONPath match values.
+    """
+    if not isinstance(data, dict):
+        raise ValueError("Cannot perform JSONPath filter on non-structured data.")
+
+    return [result.data for result in Finder(data).find(expression)]
+
+
+def is_jinja_template(user_template: str) -> bool:
+    """Determine if a prompt template is a Jinja2 template from heuristics.
+
+    This function is intended to help migration from format strings->Jinja.
+    If we only support Jinja2, then this function is not needed.
+
+    Args:
+        user_template (str): A user-provided template string to test.
+
+    Returns:
+        True if the heuristic believes it is a Jinja2 template.
+    """
+    jinja_pattern_pairs = [("{{", "}}"), ("{%", "%}"), ("{#", "#}")]
+    for open_pattern, close_pattern in jinja_pattern_pairs:
+        if open_pattern in user_template and close_pattern in user_template:
+            return True
+
+    return False
+
+
+class UserTemplateSandboxEnvironment(ImmutableSandboxedEnvironment):
+    """Defines a robust environment for rendering Gretel's Jinja2 subset.
+
+    The use of Jinja2 sandboxing is critical. We are taking Jinja2
+    templates from users -- we need to take steps to ensure that users
+    are not able to break containment or exfiltrate server-side secrets.
+
+    This Environment definition attempts to lock down as much as we can
+    for a pure python implementation by extending restrictions past
+    that of the `ImmutableSandboxedEnviornment.` While that environment
+    provides a base layer of protections, including:
+
+        - No references to private attributes
+        - Restrictions on loop iterations (OverflowError)
+
+    We enforce further precautions:
+
+        - Forced auto-escaping templates (preventing some injection attacks).
+        - Prevents access to the template's `self` attribute.
+        - Prevents reference to variables except for a provided white-list.
+        - Removes support for: include, extend, macro, set, block, and nested loops
+        - Errors on too-long rendered templates (e.g. >128k chars).
+        - Remove all default Jinja filter operations except for JSONPath (negotiable).
+        - Uses AST static analysis to threshold the complexity of allowed templates.
+
+    """
+
+    max_rendered_len: int
+    max_ast_node_count: int
+    max_ast_depth: int
+    allowed_references: list[str]
+
+    def __init__(
+        self,
+        allowed_references: list[str] | None = None,
+        max_rendered_len: int = MAX_RENDERED_LEN,
+        max_ast_node_count: int = MAX_AST_NODE_COUNT,
+        max_ast_depth: int = MAX_AST_DEPTH,
+        **kwargs,
+    ):
+        """Args:
+        max_rendered_len (int): The maximum allowable character count for
+            rendered templates.
+
+        allowed_references (optional, list[str]): If set, indicates which variables
+            are allowed to be referenced by the Jinja2 template. If not specified,
+            defaults to [], which indicates that the Jinja2 template is not
+            allowed to refer to _any_ variables outside of itself.
+
+        max_ast_node_count (optional, int): Parameter for static analysis of
+            Jinja2 template complexity -- counts the number of distinct nodes
+            in the parsed Jinja2 AST. A large number of nodes indicates many
+            distinct operations within the provided user template, which can
+            cause long compute times, or may be malicious in nature. If not
+            specified, defaults to MAX_AST_NODE_COUNT set by this module.
+
+        max_ast_depth (optional, int): Parameter for static analysis of
+            Jinja2 template complexity -- measures the maximum depth of the
+            parsed Jinja2 AST. A high depth indicates a high degree of nesting
+            within the user template. This may can cause long compute times,
+            or may be malicious in nature. If not specified, defaults to
+            MAX_AST_DEPTH set by this module.
+
+        **kwargs: Additional kwargs passed to ImmutableSandboxedEnvironment.
+        """
+        super().__init__(autoescape=False, **kwargs)
+        self.max_rendered_len = max_rendered_len
+        self.max_ast_node_count = max_ast_node_count
+        self.max_ast_depth = max_ast_depth
+        self.allowed_references = allowed_references if allowed_references else []
+
+        ## Add on our supported filters
+        self.filters["jsonpath"] = jsonpath_jinja_filter
+
+        ## Cut out all but approved Jinja filters
+        self.filters = {k: v for k, v in self.filters.items() if k in ALLOWED_JINJA_FILTERS}
+
+    def _assert_template_has_valid_references(self, ast: Template) -> None:
+        """Assert that all named variable references are allowed.
+
+        Checks against the environment's allowed reference list created
+        at initialization.
+        """
+        template_vars = meta.find_undeclared_variables(ast)
+        unallowed_vars = set(template_vars) - set(self.allowed_references)
+        if len(unallowed_vars) > 0:
+            raise UserTemplateError(f"Unknown variable references in Jinja template: {unallowed_vars}")
+
+    def _assert_template_has_valid_ast_nodes(self, ast: Template) -> None:
+        """Assert that un-allowed operations aren't in the template."""
+        black_list_node_count = sum(ast_descendant_count(ast, node_type) for node_type in UNSUPPORTED_AST_NODES)
+
+        if black_list_node_count != 0:
+            raise UserTemplateError("Non-permitted operations in Jinja template.")
+
+    def _assert_template_has_no_recursive_for(self, ast: Template) -> None:
+        """Assert that the template does not use {% for ... recursive %}"""
+        if any(node.recursive for node in ast.find_all(j_nodes.For)):
+            raise UserTemplateError("Non-permitted operations in Jinja template.")
+
+    def _assert_template_has_no_nested_for(self, ast: Template) -> None:
+        """Assert that the template does not contain nested loops.
+
+        This assertion is made to ensure that templates cannot combinatorially
+        explode. High-range values are controlled by the `MAX_RANGE` setting
+        on `SandboxedEnvironment`.
+        """
+        # Check each For node in the AST to see if it has For descendants
+        for node in ast.find_all(j_nodes.For):
+            if ast_descendant_count(node, only_type=j_nodes.For):
+                raise UserTemplateError("Non-permitted operations in Jinja template (nested-for).")
+
+    def _assert_template_ast_complexity(self, ast: Template) -> None:
+        """Assert that the AST tree parsed from the template is not overly complex.
+
+        Complexity is measured by the depth of the tree (measure of nesting),
+        as well as the number of nodes it contains (how many distinct operations).
+        If either is over a fixed limit specified at initialization, the assert fails.
+        """
+        node_count = ast_descendant_count(ast)
+        max_depth = ast_max_depth(ast)
+
+        if node_count > self.max_ast_node_count or max_depth > self.max_ast_depth:
+            raise UserTemplateError("Jinja template too complex, simplify your template.")
+
+    def _assert_template_has_no_self_reference(self, ast: Template) -> None:
+        """Assert that the template cannot refer to its own settings.
+
+        Templates may attempt to use {{ self }} references to gain
+        access to properties of the template object itself. This
+        is disallowed.
+        """
+        if ast_count_name_references(ast, "self") != 0:
+            raise UserTemplateError("Non-permitted operations in Jinja template.")
+
+    def validate_template(self, user_template: str) -> None:
+        """Template validations are run against the template object itself.
+        First-layer injection attacks are (on the parse operation) are
+        prevented by using `autoescape=True` on environment creation.
+
+        Afterwards, we can analyze the AST of the parsed template to detect
+        and mitigate a wide range of attacks.
+
+        Args:
+            user_template (str): A submitted user Jinja2 template.
+
+        Raises:
+            TemplateSyntaxError: If the provided template is malformed or
+                not parseable as a Jinja2 template.
+            UserTemplateError: If any of the assertions fail.
+        """
+        try:
+            ast = self.parse(user_template)
+            self._assert_template_has_valid_ast_nodes(ast)
+            self._assert_template_has_no_recursive_for(ast)
+            self._assert_template_has_no_nested_for(ast)
+            self._assert_template_ast_complexity(ast)
+            self._assert_template_has_no_self_reference(ast)
+            self._assert_template_has_valid_references(ast)
+        except Exception as exception:
+            maybe_handle_missing_filter_exception(exception, available_jinja_filters=list(self.filters.keys()))
+            raise exception
+
+    def _assert_rendered_text_length(self, rendered_text: str) -> None:
+        """Check against the length of the rendered string."""
+        rendered_len = len(rendered_text)
+        if rendered_len > self.max_rendered_len:
+            raise UserTemplateError(f"Rendered Jinja template too large ({rendered_len} > {self.max_rendered_len}).")
+
+    def _assert_rendered_text_has_no_builtin_descriptions(self, rendered_text: str) -> None:
+        """Check to make sure that the outputs aren't descriptions of methods.
+
+        In the event that the user types the name of a __builtin__
+        object method, but doesn't call it, we don't want to report
+        information about the system's memory contents.
+
+        Further, if the user made a mistake, we'd rather error out
+        rather than continue task processing, for instance.
+        """
+        patterns = [
+            r"<built-in method (.*?) of (.*?) object at 0x(.*?)>",
+            r"<function (.*?) at (.*?)>",
+        ]
+        for pattern in patterns:
+            matches = re.search(pattern, rendered_text)
+            if bool(matches):
+                raise UserTemplateError("User template has uncalled __builtin__ method.")
+
+    def _assert_rendered_text_not_empty(self, rendered_text: str) -> None:
+        """Check to make sure the resulting text isn't an empty string"""
+        if len(rendered_text) == 0:
+            raise UserTemplateError("User template renders to empty text.")
+
+    def validate_rendered_text(self, rendered_text: str) -> None:
+        """Raises UserTemplateError on invalid renders.
+
+        This is used as a post-processing step for capturing and
+        acting on strings before they go out the door.
+        """
+        self._assert_rendered_text_not_empty(rendered_text)
+        self._assert_rendered_text_length(rendered_text)
+        self._assert_rendered_text_has_no_builtin_descriptions(rendered_text)
+
+    def safe_render(self, user_template: str, record: dict, skip_template_validation: bool = False) -> str:
+        """Attempt to safely render a user's template.
+
+        Args:
+            user_template (str): The user submitted Jinja2 template string.
+            record (dict): a record of fields which are able to be referenced by the template.
+            skip_template_validation (optional, bool): If true, then AST checks against the
+                template itself will not be performed. WARNING: this should ONLY be set to true
+                if the template has already been validated.
+
+        Raises:
+            UserTemplateError: If the template cannot be rendered because the
+                user template does not conform to Gretel's Jinja2 subset,
+                is too long, or contains some attempted malicious payload.
+                If skip_template_validation is False, this error may also indicate
+                that the template itself has failed static analysis. See the error
+                message for more details.
+
+            RecordContentsError: If there is a system-internal error with
+                the supplied record data. This error is raised to prevent Jinja2
+                processing of potentially insecure data objects.
+        """
+        if not skip_template_validation:
+            self.validate_template(user_template)
+
+        record = sanitize_record(record)
+
+        try:
+            template = self.from_string(user_template)
+            rendered_text = template.render(record)
+        except SecurityError:
+            raise UserTemplateError("Non-permitted operations in Jinja template.")
+        except OverflowError:
+            raise UserTemplateError("Template too large.")
+        except Exception as exception:
+            maybe_handle_missing_filter_exception(exception, available_jinja_filters=list(self.filters.keys()))
+            raise exception
+
+        self.validate_rendered_text(rendered_text)
+
+        return rendered_text
+
+    def get_references(self, user_template: str) -> set[str]:
+        """Get all referenced variables from the provided template.
+
+        Args:
+            user_template (str): A user-provided Jinja template.
+
+        Returns:
+            set[str]: A set of all variable names referenced in
+                the supplied Jinja template. If no variables are
+                referenced, then this will be an empty list.
+        """
+        ast = self.parse(user_template)
+        return meta.find_undeclared_variables(ast)
+
+
+def sanitize_user_exceptions(func):
+    """Sanitize returned user-space exceptions."""
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except UserTemplateUnsupportedFiltersError as exception:
+            ## Informative messaging is already handled in this
+            ## specific case.
+            raise exception
+        except (UserTemplateError, TemplateSyntaxError):
+            ## All other details are wrapped in a generic error message
+            raise UserTemplateError(USER_PROMPT_TEMPLATE_ERROR_MESSAGE)
+
+    return wrapper
+
+
+class WithJinja2UserTemplateRendering:
+    """Mixin class to support user-supplied Jinja2 rendering.
+
+    Provides `self.render_template(record: dict)` to the receiving
+    class, which can be used to safely render user-provided Jinja2
+    templates using `UserTemplateSandboxedEnvironment`.
+
+    This mixin also provides error message sanitization for exceptions
+    raised by the rendering environment.
+
+    Usage:
+
+        class Foo(WithJinja2UserTemplateRendering):
+            def my_func(self, user_template: str, records: list[dict]):
+
+                ## Call once per template -- must be prepared before
+                ## being able to call self.render_template
+                self.prepare_jinja2_template_renderer(user_template)
+
+                ## Can call many times after
+                for record in records:
+                    self.render_template(record)
+    """
+
+    _template_render_fn: Callable
+
+    @sanitize_user_exceptions
+    def prepare_jinja2_template_renderer(self, prompt_template: str, dataset_variables: list[str]) -> None:
+        """Build Jinja2 template render function."""
+        jinja_render_env = UserTemplateSandboxEnvironment(allowed_references=dataset_variables)
+        jinja_render_env.validate_template(prompt_template)
+        self._template_render_fn = partial(
+            jinja_render_env.safe_render,
+            prompt_template,
+            skip_template_validation=True,
+        )
+
+    @sanitize_user_exceptions
+    def render_template(self, record: dict) -> str:
+        return self._template_render_fn(record)
+
+    @sanitize_user_exceptions
+    def prepare_jinja2_multi_template_renderer(
+        self,
+        template_name: str,
+        prompt_template: str,
+        dataset_variables: list[str],
+    ) -> None:
+        if not self._template_prepared_in_multi_template_renderer(template_name):
+            self._create_render_func_registry()
+            jinja_render_env = UserTemplateSandboxEnvironment(allowed_references=dataset_variables)
+            jinja_render_env.validate_template(prompt_template)
+            self._render_func_registry[template_name] = partial(
+                jinja_render_env.safe_render,
+                prompt_template,
+                skip_template_validation=True,
+            )
+
+    @sanitize_user_exceptions
+    def render_multi_template(self, template_name: str, record: dict) -> str:
+        if not hasattr(self, "_render_func_registry"):
+            raise UserTemplateError("Multi-template renderer not prepared.")
+        if template_name not in self._render_func_registry:
+            raise UserTemplateError(f"Template {template_name} not prepared.")
+        return self._render_func_registry[template_name](record)
+
+    def _template_prepared_in_multi_template_renderer(self, template_name: str) -> bool:
+        if not hasattr(self, "_render_func_registry"):
+            return False
+        return template_name in self._render_func_registry
+
+    def _create_render_func_registry(self) -> None:
+        if not hasattr(self, "_render_func_registry"):
+            self._render_func_registry = {}

data_designer/engine/processing/ginja/exceptions.py

@@ -0,0 +1,54 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+
+from jinja2 import TemplateAssertionError
+
+
+class UserTemplateError(Exception):
+    """Exception for user-induced template flaws, intentional or not."""
+
+
+class UserTemplateUnsupportedFiltersError(UserTemplateError):
+    """Specific exception for the case of unsupported filters."""
+
+
+class RecordContentsError(Exception):
+    """Exception for cases involving the record providing template context."""
+
+
+def maybe_handle_missing_filter_exception(exception: BaseException, available_jinja_filters: list[str]) -> None:
+    """Interpret and handle the possible case of a missing filter exception.
+
+    If this wasn't a missing filter exception, then this function will do
+    nothing.
+
+    Args:
+        exception (BaseException): The caught exception.
+        available_jinja_filters (list[str]): The list of Jinja filters that
+            are known to be available within the environment.
+
+    Raises:
+        UserTemplateUnsupportedFiltersError: If the exception was specifically for an unknown
+            or unsupported Jinja2 filter.
+    """
+    if not isinstance(exception, TemplateAssertionError):
+        return
+
+    exc_message = exception.message or ""
+
+    ## The missing filter message has the format:
+    ## "No filter named '____'"
+    match = re.search(r"No filter named '([^']+)'", exc_message)
+    if not match:
+        return
+    else:
+        missing_filter_name = match.group(1)
+        available_filter_str = ", ".join(available_jinja_filters)
+        raise UserTemplateUnsupportedFiltersError(
+            (
+                f"The Jinja2 filter `{{{{ ... | {missing_filter_name} }}}}` "
+                f"is not a permitted operation. Available filters: {available_filter_str}"
+            )
+        ) from exception

data_designer/engine/processing/ginja/record.py

@@ -0,0 +1,30 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+
+from data_designer.config.utils.io_helpers import serialize_data
+from data_designer.engine.processing.ginja.exceptions import RecordContentsError
+
+
+def sanitize_record(record: dict) -> dict:
+    """Sanitize a record into basic types.
+
+    To prevent any unexpected attributes from being callable from
+    the template, we apply a serdes step to ensure that the record
+    used as context for the rendering step consists of basic
+    python types (e.g. those that can be represented via JSON).
+
+    Args:
+        record (dict): A dictionary object which can be serialized.
+
+    Raises:
+        RecordContentsError if the record contents are not able
+            to be represented with JSON.
+    """
+    try:
+        ser = serialize_data(record)
+    except (TypeError, ValueError) as e:
+        raise RecordContentsError("Unexpected unserializable content found in record.") from e
+
+    return json.loads(ser)

data_designer/engine/processing/gsonschema/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

data_designer/engine/processing/gsonschema/exceptions.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from jsonschema import ValidationError
+
+
+class JSONSchemaValidationError(ValidationError):
+    """Alias of ValidationError to ease imports."""

data_designer/engine/processing/gsonschema/schema_transformers.py

@@ -0,0 +1,81 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from copy import deepcopy
+from typing import Any
+
+from data_designer.engine.processing.gsonschema.types import JSONSchemaT
+
+
+def _is_bare_dictionary_schema(schema_part: Any) -> bool:
+    """Classify bare dictionary schemas
+
+    A bare dictionary schema is one which looks like the following:
+
+        { "title": ... , "type": "object" }
+        { "type": "object" }
+
+    These schemas do not specify any "properties", just their existence.
+    """
+    if not isinstance(schema_part, dict):
+        return False
+
+    if schema_part.get("type") != "object":
+        return False
+
+    if ("title" in schema_part and len(schema_part) == 2) or (len(schema_part) == 1):
+        return True
+
+    return False
+
+
+def forbid_additional_properties(schema: JSONSchemaT) -> JSONSchemaT:
+    """Transform the provided schema into one which forbids additional properties.
+
+    Args:
+        schema (JSONSchemaT): A JSONSchema to transform.
+
+    Returns:
+        JSONSchemaT: A new JSONSchema matching the provided one, but
+            with `additionalProperties: False` set everywhere.
+    """
+    new_schema = deepcopy(schema)
+
+    def _enforce(schema_part: Any) -> None:
+        if isinstance(schema_part, dict):
+            if schema_part.get("type") == "object" or "properties" in schema_part:
+                ## We need to handle the special case that the schema specifies just a bare
+                ## dictionary. In those cases, the implication is that _all_ dictionaries
+                ## are valid, so we should not forbid extra properties in that case.
+                allow_additional_properties = _is_bare_dictionary_schema(schema_part)
+                schema_part["additionalProperties"] = allow_additional_properties
+
+            # Traverse into nested schemas.
+            for key, value in schema_part.items():
+                if key in ("properties", "patternProperties"):
+                    if isinstance(value, dict):
+                        for sub_schema in value.values():
+                            _enforce(sub_schema)
+                elif key == "items":
+                    if isinstance(value, dict):
+                        _enforce(value)
+                    elif isinstance(value, list):
+                        for item in value:
+                            _enforce(item)
+                elif key in ("allOf", "anyOf", "oneOf"):
+                    if isinstance(value, list):
+                        for item in value:
+                            _enforce(item)
+                elif key in ("not", "if", "then", "else"):
+                    if isinstance(value, dict):
+                        _enforce(value)
+                elif key == "$defs":
+                    for sub_schema in value.values():
+                        _enforce(sub_schema)
+
+        elif isinstance(schema_part, list):
+            for item in schema_part:
+                _enforce(item)
+
+    _enforce(new_schema)
+    return new_schema

data_designer/engine/processing/gsonschema/types.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any, TypeVar
+
+T_primitive = TypeVar("T_primitive", str, int, float, bool)
+DataObjectT = dict | list | str | int | float | bool
+JSONSchemaT = dict[str, Any]