src/backend/base/langflow/schema/data.py

import copy
import json
from datetime import datetime
from decimal import Decimal
from typing import cast
from uuid import UUID

from langchain_core.documents import Document
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
from loguru import logger
from pydantic import BaseModel, model_serializer, model_validator

from langflow.utils.constants import MESSAGE_SENDER_AI, MESSAGE_SENDER_USER
from langflow.utils.image import create_data_url


class Data(BaseModel):
    """Represents a record with text and optional data.

    Attributes:
        data (dict, optional): Additional data associated with the record.
    """

    text_key: str = "text"
    data: dict = {}
    default_value: str | None = ""

    @model_validator(mode="before")
    @classmethod
    def validate_data(cls, values):
        if not isinstance(values, dict):
            msg = "Data must be a dictionary"
            raise ValueError(msg)  # noqa: TRY004
        if not values.get("data"):
            values["data"] = {}
        # Any other keyword should be added to the data dictionary
        for key in values:
            if key not in values["data"] and key not in {"text_key", "data", "default_value"}:
                values["data"][key] = values[key]
        return values

    @model_serializer(mode="plain", when_used="json")
    def serialize_model(self):
        return {k: v.to_json() if hasattr(v, "to_json") else v for k, v in self.data.items()}

    def get_text(self):
        """Retrieves the text value from the data dictionary.

        If the text key is present in the data dictionary, the corresponding value is returned.
        Otherwise, the default value is returned.

        Returns:
            The text value from the data dictionary or the default value.
        """
        return self.data.get(self.text_key, self.default_value)

    def set_text(self, text: str | None) -> str:
        r"""Sets the text value in the data dictionary.

        The object's `text` value is set to `text parameter as given, with the following modifications:

         - `text` value of `None` is converted to an empty string.
         - `text` value is converted to `str` type.

        Args:
            text (str): The text to be set in the data dictionary.

        Returns:
            str: The text value that was set in the data dictionary.
        """
        new_text = "" if text is None else str(text)
        self.data[self.text_key] = new_text
        return new_text

    @classmethod
    def from_document(cls, document: Document) -> "Data":
        """Converts a Document to a Data.

        Args:
            document (Document): The Document to convert.

        Returns:
            Data: The converted Data.
        """
        data = document.metadata
        data["text"] = document.page_content
        return cls(data=data, text_key="text")

    @classmethod
    def from_lc_message(cls, message: BaseMessage) -> "Data":
        """Converts a BaseMessage to a Data.

        Args:
            message (BaseMessage): The BaseMessage to convert.

        Returns:
            Data: The converted Data.
        """
        data: dict = {"text": message.content}
        data["metadata"] = cast("dict", message.to_json())
        return cls(data=data, text_key="text")

    def __add__(self, other: "Data") -> "Data":
        """Combines the data of two data by attempting to add values for overlapping keys.

        Combines the data of two data by attempting to add values for overlapping keys
        for all types that support the addition operation. Falls back to the value from 'other'
        record when addition is not supported.
        """
        combined_data = self.data.copy()
        for key, value in other.data.items():
            # If the key exists in both data and both values support the addition operation
            if key in combined_data:
                try:
                    combined_data[key] += value
                except TypeError:
                    # Fallback: Use the value from 'other' record if addition is not supported
                    combined_data[key] = value
            else:
                # If the key is not in the first record, simply add it
                combined_data[key] = value

        return Data(data=combined_data)

    def to_lc_document(self) -> Document:
        """Converts the Data to a Document.

        Returns:
            Document: The converted Document.
        """
        data_copy = self.data.copy()
        text = data_copy.pop(self.text_key, self.default_value)
        if isinstance(text, str):
            return Document(page_content=text, metadata=data_copy)
        return Document(page_content=str(text), metadata=data_copy)

    def to_lc_message(
        self,
    ) -> BaseMessage:
        """Converts the Data to a BaseMessage.

        Returns:
            BaseMessage: The converted BaseMessage.
        """
        # The idea of this function is to be a helper to convert a Data to a BaseMessage
        # It will use the "sender" key to determine if the message is Human or AI
        # If the key is not present, it will default to AI
        # But first we check if all required keys are present in the data dictionary
        # they are: "text", "sender"
        if not all(key in self.data for key in ["text", "sender"]):
            msg = f"Missing required keys ('text', 'sender') in Data: {self.data}"
            raise ValueError(msg)
        sender = self.data.get("sender", MESSAGE_SENDER_AI)
        text = self.data.get("text", "")
        files = self.data.get("files", [])
        if sender == MESSAGE_SENDER_USER:
            if files:
                contents = [{"type": "text", "text": text}]
                for file_path in files:
                    image_url = create_data_url(file_path)
                    contents.append({"type": "image_url", "image_url": {"url": image_url}})
                human_message = HumanMessage(content=contents)
            else:
                human_message = HumanMessage(
                    content=[{"type": "text", "text": text}],
                )

            return human_message

        return AIMessage(content=text)

    def __getattr__(self, key):
        """Allows attribute-like access to the data dictionary."""
        try:
            if key.startswith("__"):
                return self.__getattribute__(key)
            if key in {"data", "text_key"} or key.startswith("_"):
                return super().__getattr__(key)
            return self.data[key]
        except KeyError as e:
            # Fallback to default behavior to raise AttributeError for undefined attributes
            msg = f"'{type(self).__name__}' object has no attribute '{key}'"
            raise AttributeError(msg) from e

    def __setattr__(self, key, value) -> None:
        """Set attribute-like values in the data dictionary.

        Allows attribute-like setting of values in the data dictionary.
        while still allowing direct assignment to class attributes.
        """
        if key in {"data", "text_key"} or key.startswith("_"):
            super().__setattr__(key, value)
        elif key in self.model_fields:
            self.data[key] = value
            super().__setattr__(key, value)
        else:
            self.data[key] = value

    def __delattr__(self, key) -> None:
        """Allows attribute-like deletion from the data dictionary."""
        if key in {"data", "text_key"} or key.startswith("_"):
            super().__delattr__(key)
        else:
            del self.data[key]

    def __deepcopy__(self, memo):
        """Custom deepcopy implementation to handle copying of the Data object."""
        # Create a new Data object with a deep copy of the data dictionary
        return Data(data=copy.deepcopy(self.data, memo), text_key=self.text_key, default_value=self.default_value)

    # check which attributes the Data has by checking the keys in the data dictionary
    def __dir__(self):
        return super().__dir__() + list(self.data.keys())

    def __str__(self) -> str:
        # return a JSON string representation of the Data atributes
        try:
            data = {k: v.to_json() if hasattr(v, "to_json") else v for k, v in self.data.items()}
            return serialize_data(data)  # use the custom serializer
        except Exception:  # noqa: BLE001
            logger.opt(exception=True).debug("Error converting Data to JSON")
            return str(self.data)

    def __contains__(self, key) -> bool:
        return key in self.data

    def __eq__(self, /, other):
        return isinstance(other, Data) and self.data == other.data


def custom_serializer(obj):
    if isinstance(obj, datetime):
        return obj.astimezone().isoformat()
    if isinstance(obj, Decimal):
        return float(obj)
    if isinstance(obj, UUID):
        return str(obj)
    if isinstance(obj, BaseModel):
        return obj.model_dump()
    # Add more custom serialization rules as needed
    msg = f"Type {type(obj)} not serializable"
    raise TypeError(msg)


def serialize_data(data):
    return json.dumps(data, indent=4, default=custom_serializer)