Underground-Digital/Workflow-Engine

1import re2from collections.abc import Mapping, Sequence3from typing import Any4 5from core.workflow.entities.variable_entities import VariableSelector6 7REGEX = re.compile(r"\{\{(#[a-zA-Z0-9_]{1,50}(\.[a-zA-Z_][a-zA-Z0-9_]{0,29}){1,10}#)\}\}")8 9SELECTOR_PATTERN = re.compile(r"\{\{(#[a-zA-Z0-9_]{1,50}(?:\.[a-zA-Z_][a-zA-Z0-9_]{0,29}){1,10}#)\}\}")10 11 12def extract_selectors_from_template(template: str, /) -> Sequence[VariableSelector]:13    parts = SELECTOR_PATTERN.split(template)14    selectors = []15    for part in filter(lambda x: x, parts):16        if "." in part and part[0] == "#" and part[-1] == "#":17            selectors.append(VariableSelector(variable=f"{part}", value_selector=part[1:-1].split(".")))18    return selectors19 20 21class VariableTemplateParser:22    """23    !NOTE: Consider to use the new `segments` module instead of this class.24 25    A class for parsing and manipulating template variables in a string.26 27    Rules:28 29    1. Template variables must be enclosed in `{{}}`.30    2. The template variable Key can only be: #node_id.var1.var2#.31    3. The template variable Key cannot contain new lines or spaces, and must comply with rule 2.32 33    Example usage:34 35    template = "Hello, {{#node_id.query.name#}}! Your age is {{#node_id.query.age#}}."36    parser = VariableTemplateParser(template)37 38    # Extract template variable keys39    variable_keys = parser.extract()40    print(variable_keys)41    # Output: ['#node_id.query.name#', '#node_id.query.age#']42 43    # Extract variable selectors44    variable_selectors = parser.extract_variable_selectors()45    print(variable_selectors)46    # Output: [VariableSelector(variable='#node_id.query.name#', value_selector=['node_id', 'query', 'name']),47    #          VariableSelector(variable='#node_id.query.age#', value_selector=['node_id', 'query', 'age'])]48 49    # Format the template string50    inputs = {'#node_id.query.name#': 'John', '#node_id.query.age#': 25}}51    formatted_string = parser.format(inputs)52    print(formatted_string)53    # Output: "Hello, John! Your age is 25."54    """55 56    def __init__(self, template: str):57        self.template = template58        self.variable_keys = self.extract()59 60    def extract(self) -> list:61        """62        Extracts all the template variable keys from the template string.63 64        Returns:65            A list of template variable keys.66        """67        # Regular expression to match the template rules68        matches = re.findall(REGEX, self.template)69 70        first_group_matches = [match[0] for match in matches]71 72        return list(set(first_group_matches))73 74    def extract_variable_selectors(self) -> list[VariableSelector]:75        """76        Extracts the variable selectors from the template variable keys.77 78        Returns:79            A list of VariableSelector objects representing the variable selectors.80        """81        variable_selectors = []82        for variable_key in self.variable_keys:83            remove_hash = variable_key.replace("#", "")84            split_result = remove_hash.split(".")85            if len(split_result) < 2:86                continue87 88            variable_selectors.append(VariableSelector(variable=variable_key, value_selector=split_result))89 90        return variable_selectors91 92    def format(self, inputs: Mapping[str, Any]) -> str:93        """94        Formats the template string by replacing the template variables with their corresponding values.95 96        Args:97            inputs: A dictionary containing the values for the template variables.98            remove_template_variables: A boolean indicating whether to remove the template variables from the output.99 100        Returns:101            The formatted string with template variables replaced by their values.102        """103 104        def replacer(match):105            key = match.group(1)106            value = inputs.get(key, match.group(0))  # return original matched string if key not found107 108            if value is None:109                value = ""110            # convert the value to string111            if isinstance(value, list | dict | bool | int | float):112                value = str(value)113 114            # remove template variables if required115            return VariableTemplateParser.remove_template_variables(value)116 117        prompt = re.sub(REGEX, replacer, self.template)118        return re.sub(r"<\|.*?\|>", "", prompt)119 120    @classmethod121    def remove_template_variables(cls, text: str):122        """123        Removes the template variables from the given text.124 125        Args:126            text: The text from which to remove the template variables.127 128        Returns:129            The text with template variables removed.130        """131        return re.sub(REGEX, r"{\1}", text)132