From 17ab6f4b317c9113687798812e6f3c71c8df875a Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Wed, 21 Feb 2024 16:01:29 -0300
Subject: [PATCH] Refactor MessageModel to include a classmethod for creating
 instances from a Record

---
 .../langflow/services/monitor/schema.py       | 22 +++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/src/backend/langflow/services/monitor/schema.py b/src/backend/langflow/services/monitor/schema.py
index 0cca430f6..6fbb842ec 100644
--- a/src/backend/langflow/services/monitor/schema.py
+++ b/src/backend/langflow/services/monitor/schema.py
@@ -1,9 +1,12 @@
 import json
 from datetime import datetime
-from typing import Any, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 from pydantic import BaseModel, Field, field_serializer, validator
 
+if TYPE_CHECKING:
+    from langflow.schema import Record
+
 
 class TransactionModel(BaseModel):
     id: Optional[int] = Field(default=None, alias="id")
@@ -31,7 +34,7 @@ class TransactionModel(BaseModel):
 class MessageModel(BaseModel):
     id: Optional[int] = Field(default=None, alias="id")
     timestamp: datetime = Field(default_factory=datetime.now)
-    sender_type: str
+    sender: str
     sender_name: str
     session_id: str
     message: str
@@ -47,6 +50,21 @@ class MessageModel(BaseModel):
             return json.loads(v)
         return v
 
+    @classmethod
+    def from_record(cls, record: "Record"):
+        # first check if the record has all the required fields
+        if "sender" not in record.data and "sender_name" not in record.data:
+            raise ValueError(
+                "The record does not have the required fields 'sender' and 'sender_name' in the data."
+            )
+        return cls(
+            sender=record.data["sender"],
+            sender_name=record.data["sender_name"],
+            message=record.text,
+            session_id=record.data.get("session_id", ""),
+            artifacts=record.data.get("artifacts", {}),
+        )
+
 
 class VertexBuildModel(BaseModel):
     index: Optional[int] = Field(default=None, alias="index", exclude=True)