dify/api/core/workflow/system_variable.py

from __future__ import annotations

from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Any
from uuid import uuid4

from pydantic import AliasChoices, BaseModel, ConfigDict, Field, model_validator

from core.file.models import File
from core.workflow.enums import SystemVariableKey


class SystemVariable(BaseModel):
    """A model for managing system variables.

    Fields with a value of `None` are treated as absent and will not be included
    in the variable pool.
    """

    model_config = ConfigDict(
        extra="forbid",
        serialize_by_alias=True,
        validate_by_alias=True,
    )

    user_id: str | None = None

    # Ideally, `app_id` and `workflow_id` should be required and not `None`.
    # However, there are scenarios in the codebase where these fields are not set.
    # To maintain compatibility, they are marked as optional here.
    app_id: str | None = None
    workflow_id: str | None = None

    timestamp: int | None = None

    files: Sequence[File] = Field(default_factory=list)

    # NOTE: The `workflow_execution_id` field was previously named `workflow_run_id`.
    # To maintain compatibility with existing workflows, it must be serialized
    # as `workflow_run_id` in dictionaries or JSON objects, and also referenced
    # as `workflow_run_id` in the variable pool.
    workflow_execution_id: str | None = Field(
        validation_alias=AliasChoices("workflow_execution_id", "workflow_run_id"),
        serialization_alias="workflow_run_id",
        default=None,
    )
    # Chatflow related fields.
    query: str | None = None
    conversation_id: str | None = None
    dialogue_count: int | None = None
    document_id: str | None = None
    original_document_id: str | None = None
    dataset_id: str | None = None
    batch: str | None = None
    datasource_type: str | None = None
    datasource_info: Mapping[str, Any] | None = None
    invoke_from: str | None = None

    @model_validator(mode="before")
    @classmethod
    def validate_json_fields(cls, data):
        if isinstance(data, dict):
            # For JSON validation, only allow workflow_run_id
            if "workflow_execution_id" in data and "workflow_run_id" not in data:
                # This is likely from direct instantiation, allow it
                return data
            elif "workflow_execution_id" in data and "workflow_run_id" in data:
                # Both present, remove workflow_execution_id
                data = data.copy()
                data.pop("workflow_execution_id")
                return data
        return data

    @classmethod
    def default(cls) -> SystemVariable:
        return cls(workflow_execution_id=str(uuid4()))

    def to_dict(self) -> dict[SystemVariableKey, Any]:
        # NOTE: This method is provided for compatibility with legacy code.
        # New code should use the `SystemVariable` object directly instead of converting
        # it to a dictionary, as this conversion results in the loss of type information
        # for each key, making static analysis more difficult.

        d: dict[SystemVariableKey, Any] = {
            SystemVariableKey.FILES: self.files,
        }
        if self.user_id is not None:
            d[SystemVariableKey.USER_ID] = self.user_id
        if self.app_id is not None:
            d[SystemVariableKey.APP_ID] = self.app_id
        if self.workflow_id is not None:
            d[SystemVariableKey.WORKFLOW_ID] = self.workflow_id
        if self.workflow_execution_id is not None:
            d[SystemVariableKey.WORKFLOW_EXECUTION_ID] = self.workflow_execution_id
        if self.query is not None:
            d[SystemVariableKey.QUERY] = self.query
        if self.conversation_id is not None:
            d[SystemVariableKey.CONVERSATION_ID] = self.conversation_id
        if self.dialogue_count is not None:
            d[SystemVariableKey.DIALOGUE_COUNT] = self.dialogue_count
        if self.document_id is not None:
            d[SystemVariableKey.DOCUMENT_ID] = self.document_id
        if self.original_document_id is not None:
            d[SystemVariableKey.ORIGINAL_DOCUMENT_ID] = self.original_document_id
        if self.dataset_id is not None:
            d[SystemVariableKey.DATASET_ID] = self.dataset_id
        if self.batch is not None:
            d[SystemVariableKey.BATCH] = self.batch
        if self.datasource_type is not None:
            d[SystemVariableKey.DATASOURCE_TYPE] = self.datasource_type
        if self.datasource_info is not None:
            d[SystemVariableKey.DATASOURCE_INFO] = self.datasource_info
        if self.invoke_from is not None:
            d[SystemVariableKey.INVOKE_FROM] = self.invoke_from
        if self.timestamp is not None:
            d[SystemVariableKey.TIMESTAMP] = self.timestamp
        return d

    def as_view(self) -> SystemVariableReadOnlyView:
        return SystemVariableReadOnlyView(self)


class SystemVariableReadOnlyView:
    """
    A read-only view of a SystemVariable that implements the ReadOnlySystemVariable protocol.

    This class wraps a SystemVariable instance and provides read-only access to all its fields.
    It always reads the latest data from the wrapped instance and prevents any write operations.
    """

    def __init__(self, system_variable: SystemVariable) -> None:
        """
        Initialize the read-only view with a SystemVariable instance.

        Args:
            system_variable: The SystemVariable instance to wrap
        """
        self._system_variable = system_variable

    @property
    def user_id(self) -> str | None:
        return self._system_variable.user_id

    @property
    def app_id(self) -> str | None:
        return self._system_variable.app_id

    @property
    def workflow_id(self) -> str | None:
        return self._system_variable.workflow_id

    @property
    def workflow_execution_id(self) -> str | None:
        return self._system_variable.workflow_execution_id

    @property
    def query(self) -> str | None:
        return self._system_variable.query

    @property
    def conversation_id(self) -> str | None:
        return self._system_variable.conversation_id

    @property
    def dialogue_count(self) -> int | None:
        return self._system_variable.dialogue_count

    @property
    def document_id(self) -> str | None:
        return self._system_variable.document_id

    @property
    def original_document_id(self) -> str | None:
        return self._system_variable.original_document_id

    @property
    def dataset_id(self) -> str | None:
        return self._system_variable.dataset_id

    @property
    def batch(self) -> str | None:
        return self._system_variable.batch

    @property
    def datasource_type(self) -> str | None:
        return self._system_variable.datasource_type

    @property
    def invoke_from(self) -> str | None:
        return self._system_variable.invoke_from

    @property
    def files(self) -> Sequence[File]:
        """
        Get a copy of the files from the wrapped SystemVariable.

        Returns:
            A defensive copy of the files sequence to prevent modification
        """
        return tuple(self._system_variable.files)  # Convert to immutable tuple

    @property
    def datasource_info(self) -> Mapping[str, Any] | None:
        """
        Get a copy of the datasource info from the wrapped SystemVariable.

        Returns:
            A view of the datasource info mapping to prevent modification
        """
        if self._system_variable.datasource_info is None:
            return None
        return MappingProxyType(self._system_variable.datasource_info)

    def __repr__(self) -> str:
        """Return a string representation of the read-only view."""
        return f"SystemVariableReadOnlyView(system_variable={self._system_variable!r})"