Add unit tests for the Record class

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

@@ -1,4 +1,5 @@
 import copy
+from typing import Optional
 
 from langchain_core.documents import Document
 from pydantic import BaseModel, model_validator
@@ -12,8 +13,9 @@ class Record(BaseModel):
         data (dict, optional): Additional data associated with the record.
     """
 
+    text_key: Optional[str] = "text"
     data: dict = {}
-    _default_value: str = ""
+    default_value: Optional[str] = ""
 
     @model_validator(mode="before")
     def validate_data(cls, values):
@@ -21,10 +23,22 @@ class Record(BaseModel):
             values["data"] = {}
         # Any other keyword should be added to the data dictionary
         for key in values:
-            if key not in values["data"] and key != "data":
+            if key not in values["data"] and key not in {"text_key", "data", "default_value"}:
                 values["data"][key] = values[key]
         return values
 
+    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)
+
     @classmethod
     def from_document(cls, document: Document) -> "Record":
         """
@@ -38,19 +52,27 @@ class Record(BaseModel):
         """
         data = document.metadata
         data["text"] = document.page_content
-        return cls(data=data)
+        return cls(data=data, text_key="text")
 
     def __add__(self, other: "Record") -> "Record":
         """
-        Concatenates the text of two records and combines their data.
-
-        Args:
-            other (Record): The other record to concatenate with.
-
-        Returns:
-            Record: The concatenated record.
+        Combines the data of two records 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, **other.data}
+        combined_data = self.data.copy()
+        for key, value in other.data.items():
+            # If the key exists in both records 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 Record(data=combined_data)
 
     def to_lc_document(self) -> Document:
@@ -60,17 +82,20 @@ class Record(BaseModel):
         Returns:
             Document: The converted Document.
         """
-        return Document(page_content=self.text, metadata=self.data)
+        text = self.data.pop(self.text_key, self.default_value)
+        return Document(page_content=text, metadata=self.data)
 
     def __getattr__(self, key):
         """
         Allows attribute-like access to the data dictionary.
         """
         try:
-            if key == "data" or key.startswith("_"):
+            if key.startswith("__"):
+                return self.__getattribute__(key)
+            if key in {"data", "text_key"} or key.startswith("_"):
                 return super().__getattr__(key)
 
-            return self.data.get(key, self._default_value)
+            return self.data.get(key, self.default_value)
         except KeyError:
             # Fallback to default behavior to raise AttributeError for undefined attributes
             raise AttributeError(f"'{type(self).__name__}' object has no attribute '{key}'")
@@ -80,7 +105,7 @@ class Record(BaseModel):
         Allows attribute-like setting of values in the data dictionary,
         while still allowing direct assignment to class attributes.
         """
-        if key == "data" or key.startswith("_"):
+        if key in {"data", "text_key"} or key.startswith("_"):
             super().__setattr__(key, value)
         else:
             self.data[key] = value
@@ -89,7 +114,7 @@ class Record(BaseModel):
         """
         Allows attribute-like deletion from the data dictionary.
         """
-        if key == "data" or key.startswith("_"):
+        if key in {"data", "text_key"} or key.startswith("_"):
             super().__delattr__(key)
         else:
             del self.data[key]
@@ -98,12 +123,8 @@ class Record(BaseModel):
         """
         Custom deepcopy implementation to handle copying of the Record object.
         """
-        cls = self.__class__
-        result = cls.__new__(cls)
-        memo[id(self)] = result
-        for k, v in self.__dict__.items():
-            setattr(result, k, copy.deepcopy(v, memo))
-        return result
+        # Create a new Record object with a deep copy of the data dictionary
+        return Record(data=copy.deepcopy(self.data, memo), text_key=self.text_key, default_value=self.default_value)
 
     def __str__(self) -> str:
         """
@@ -114,7 +135,8 @@ class Record(BaseModel):
         # build the string considering all keys in the data dictionary
         prefix = "Record("
         suffix = ")"
-        text = ", ".join([f"{k}={v}" for k, v in self.data.items()])
+        text = f"text_key={self.text_key}, "
+        text += ", ".join([f"{k}={v}" for k, v in self.data.items()])
         return prefix + text + suffix
 
     # check which attributes the Record has by checking the keys in the data dictionary