From 603ffd1f1b7b8a1915bf5d13a1a354ac8a18bcb8 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Wed, 7 Feb 2024 22:33:07 -0300
Subject: [PATCH] Add Record class and conversion functions

---
 src/backend/langflow/schema.py | 52 ++++++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)
 create mode 100644 src/backend/langflow/schema.py

diff --git a/src/backend/langflow/schema.py b/src/backend/langflow/schema.py
new file mode 100644
index 000000000..de1be7767
--- /dev/null
+++ b/src/backend/langflow/schema.py
@@ -0,0 +1,52 @@
+from typing import Optional
+
+from langchain_core.documents import Document
+from pydantic import BaseModel
+
+
+class Record(BaseModel):
+    """
+    Represents a record with text and optional data.
+
+    Attributes:
+        text (str): The text of the record.
+        data (dict, optional): Additional data associated with the record.
+    """
+
+    text: str
+    data: Optional[dict] = None
+
+    @classmethod
+    def from_document(cls, document: Document) -> "Record":
+        """
+        Converts a Document to a Record.
+
+        Args:
+            document (Document): The Document to convert.
+
+        Returns:
+            Record: The converted Record.
+        """
+        return cls(text=document.page_content, data=document.metadata)
+
+    def to_lc_document(self) -> Document:
+        """
+        Converts the Record to a Document.
+
+        Returns:
+            Document: The converted Document.
+        """
+        return Document(page_content=self.text, metadata=self.data)
+
+
+def docs_to_records(documents: list[Document]) -> list[Record]:
+    """
+    Converts a list of Documents to a list of Records.
+
+    Args:
+        documents (list[Document]): The list of Documents to convert.
+
+    Returns:
+        list[Record]: The converted list of Records.
+    """
+    return [Record.from_document(document) for document in documents]