dify/api/dify_graph/runtime/variable_pool.py

from __future__ import annotations

import re
from collections import defaultdict
from collections.abc import Mapping, Sequence
from copy import deepcopy
from typing import Annotated, Any, Union, cast

from pydantic import BaseModel, Field

from dify_graph.constants import (
    CONVERSATION_VARIABLE_NODE_ID,
    ENVIRONMENT_VARIABLE_NODE_ID,
    RAG_PIPELINE_VARIABLE_NODE_ID,
    SYSTEM_VARIABLE_NODE_ID,
)
from dify_graph.file import File, FileAttribute, file_manager
from dify_graph.system_variable import SystemVariable
from dify_graph.variables import Segment, SegmentGroup, VariableBase
from dify_graph.variables.consts import SELECTORS_LENGTH
from dify_graph.variables.segments import FileSegment, ObjectSegment
from dify_graph.variables.variables import RAGPipelineVariableInput, Variable
from factories import variable_factory

VariableValue = Union[str, int, float, dict[str, object], list[object], File]

VARIABLE_PATTERN = re.compile(r"\{\{#([a-zA-Z0-9_]{1,50}(?:\.[a-zA-Z_][a-zA-Z0-9_]{0,29}){1,10})#\}\}")


class VariablePool(BaseModel):
    # Variable dictionary is a dictionary for looking up variables by their selector.
    # The first element of the selector is the node id, it's the first-level key in the dictionary.
    # Other elements of the selector are the keys in the second-level dictionary. To get the key, we hash the
    # elements of the selector except the first one.
    variable_dictionary: defaultdict[str, Annotated[dict[str, Variable], Field(default_factory=dict)]] = Field(
        description="Variables mapping",
        default=defaultdict(dict),
    )

    # The `user_inputs` is used only when constructing the inputs for the `StartNode`. It's not used elsewhere.
    user_inputs: Mapping[str, Any] = Field(
        description="User inputs",
        default_factory=dict,
    )
    system_variables: SystemVariable = Field(
        description="System variables",
        default_factory=SystemVariable.default,
    )
    environment_variables: Sequence[Variable] = Field(
        description="Environment variables.",
        default_factory=list[Variable],
    )
    conversation_variables: Sequence[Variable] = Field(
        description="Conversation variables.",
        default_factory=list[Variable],
    )
    rag_pipeline_variables: list[RAGPipelineVariableInput] = Field(
        description="RAG pipeline variables.",
        default_factory=list,
    )

    def model_post_init(self, context: Any, /):
        # Create a mapping from field names to SystemVariableKey enum values
        self._add_system_variables(self.system_variables)
        # Add environment variables to the variable pool
        for var in self.environment_variables:
            self.add((ENVIRONMENT_VARIABLE_NODE_ID, var.name), var)
        # Add conversation variables to the variable pool
        for var in self.conversation_variables:
            self.add((CONVERSATION_VARIABLE_NODE_ID, var.name), var)
        # Add rag pipeline variables to the variable pool
        if self.rag_pipeline_variables:
            rag_pipeline_variables_map: defaultdict[Any, dict[Any, Any]] = defaultdict(dict)
            for rag_var in self.rag_pipeline_variables:
                node_id = rag_var.variable.belong_to_node_id
                key = rag_var.variable.variable
                value = rag_var.value
                rag_pipeline_variables_map[node_id][key] = value
            for key, value in rag_pipeline_variables_map.items():
                self.add((RAG_PIPELINE_VARIABLE_NODE_ID, key), value)

    def add(self, selector: Sequence[str], value: Any, /):
        """
        Add a variable to the variable pool.

        This method accepts a selector path and a value, converting the value
        to a Variable object if necessary before storing it in the pool.

        Args:
            selector: A two-element sequence containing [node_id, variable_name].
                     The selector must have exactly 2 elements to be valid.
            value: The value to store. Can be a Variable, Segment, or any value
                  that can be converted to a Segment (str, int, float, dict, list, File).

        Raises:
            ValueError: If selector length is not exactly 2 elements.

        Note:
            While non-Segment values are currently accepted and automatically
            converted, it's recommended to pass Segment or Variable objects directly.
        """
        if len(selector) != SELECTORS_LENGTH:
            raise ValueError(
                f"Invalid selector: expected {SELECTORS_LENGTH} elements (node_id, variable_name), "
                f"got {len(selector)} elements"
            )

        if isinstance(value, VariableBase):
            variable = value
        elif isinstance(value, Segment):
            variable = variable_factory.segment_to_variable(segment=value, selector=selector)
        else:
            segment = variable_factory.build_segment(value)
            variable = variable_factory.segment_to_variable(segment=segment, selector=selector)

        node_id, name = self._selector_to_keys(selector)
        # Based on the definition of `Variable`,
        # `VariableBase` instances can be safely used as `Variable` since they are compatible.
        self.variable_dictionary[node_id][name] = cast(Variable, variable)

    @classmethod
    def _selector_to_keys(cls, selector: Sequence[str]) -> tuple[str, str]:
        return selector[0], selector[1]

    def _has(self, selector: Sequence[str]) -> bool:
        node_id, name = self._selector_to_keys(selector)
        if node_id not in self.variable_dictionary:
            return False
        if name not in self.variable_dictionary[node_id]:
            return False
        return True

    def get(self, selector: Sequence[str], /) -> Segment | None:
        """
        Retrieve a variable's value from the pool as a Segment.

        This method supports both simple selectors [node_id, variable_name] and
        extended selectors that include attribute access for FileSegment and
        ObjectSegment types.

        Args:
            selector: A sequence with at least 2 elements:
                     - [node_id, variable_name]: Returns the full segment
                     - [node_id, variable_name, attr, ...]: Returns a nested value
                       from FileSegment (e.g., 'url', 'name') or ObjectSegment

        Returns:
            The Segment associated with the selector, or None if not found.
            Returns None if selector has fewer than 2 elements.

        Raises:
            ValueError: If attempting to access an invalid FileAttribute.
        """
        if len(selector) < SELECTORS_LENGTH:
            return None

        node_id, name = self._selector_to_keys(selector)
        node_map = self.variable_dictionary.get(node_id)
        if node_map is None:
            return None

        segment: Segment | None = node_map.get(name)

        if segment is None:
            return None

        if len(selector) == 2:
            return segment

        if isinstance(segment, FileSegment):
            attr = selector[2]
            # Python support `attr in FileAttribute` after 3.12
            if attr not in {item.value for item in FileAttribute}:
                return None
            attr = FileAttribute(attr)
            attr_value = file_manager.get_attr(file=segment.value, attr=attr)
            return variable_factory.build_segment(attr_value)

        # Navigate through nested attributes
        result: Any = segment
        for attr in selector[2:]:
            result = self._extract_value(result)
            result = self._get_nested_attribute(result, attr)
            if result is None:
                return None

        # Return result as Segment
        return result if isinstance(result, Segment) else variable_factory.build_segment(result)

    def _extract_value(self, obj: Any):
        """Extract the actual value from an ObjectSegment."""
        return obj.value if isinstance(obj, ObjectSegment) else obj

    def _get_nested_attribute(self, obj: Mapping[str, Any], attr: str) -> Segment | None:
        """
        Get a nested attribute from a dictionary-like object.

        Args:
            obj: The dictionary-like object to search.
            attr: The key to look up.

        Returns:
            Segment | None:
                The corresponding Segment built from the attribute value if the key exists,
                otherwise None.
        """
        if not isinstance(obj, dict) or attr not in obj:
            return None
        return variable_factory.build_segment(obj.get(attr))

    def remove(self, selector: Sequence[str], /):
        """
        Remove variables from the variable pool based on the given selector.

        Args:
            selector (Sequence[str]): A sequence of strings representing the selector.

        Returns:
            None
        """
        if not selector:
            return
        if len(selector) == 1:
            self.variable_dictionary[selector[0]] = {}
            return
        key, hash_key = self._selector_to_keys(selector)
        self.variable_dictionary[key].pop(hash_key, None)

    def convert_template(self, template: str, /):
        parts = VARIABLE_PATTERN.split(template)
        segments: list[Segment] = []
        for part in filter(lambda x: x, parts):
            if "." in part and (variable := self.get(part.split("."))):
                segments.append(variable)
            else:
                segments.append(variable_factory.build_segment(part))
        return SegmentGroup(value=segments)

    def get_file(self, selector: Sequence[str], /) -> FileSegment | None:
        segment = self.get(selector)
        if isinstance(segment, FileSegment):
            return segment
        return None

    def get_by_prefix(self, prefix: str, /) -> Mapping[str, object]:
        """Return a copy of all variables stored under the given node prefix."""

        nodes = self.variable_dictionary.get(prefix)
        if not nodes:
            return {}

        result: dict[str, object] = {}
        for key, variable in nodes.items():
            value = variable.value
            result[key] = deepcopy(value)

        return result

    def _add_system_variables(self, system_variable: SystemVariable):
        sys_var_mapping = system_variable.to_dict()
        for key, value in sys_var_mapping.items():
            if value is None:
                continue
            selector = (SYSTEM_VARIABLE_NODE_ID, key)
            # If the system variable already exists, do not add it again.
            # This ensures that we can keep the id of the system variables intact.
            if self._has(selector):
                continue
            self.add(selector, value)

    @classmethod
    def empty(cls) -> VariablePool:
        """Create an empty variable pool."""
        return cls(system_variables=SystemVariable.default())