Refactor MessageModel to include a classmethod for creating instances from a Record
This commit is contained in:
parent
09061a41b9
commit
17ab6f4b31
1 changed files with 20 additions and 2 deletions
|
|
@ -1,9 +1,12 @@
|
||||||
import json
|
import json
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from typing import Any, Optional
|
from typing import TYPE_CHECKING, Any, Optional
|
||||||
|
|
||||||
from pydantic import BaseModel, Field, field_serializer, validator
|
from pydantic import BaseModel, Field, field_serializer, validator
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from langflow.schema import Record
|
||||||
|
|
||||||
|
|
||||||
class TransactionModel(BaseModel):
|
class TransactionModel(BaseModel):
|
||||||
id: Optional[int] = Field(default=None, alias="id")
|
id: Optional[int] = Field(default=None, alias="id")
|
||||||
|
|
@ -31,7 +34,7 @@ class TransactionModel(BaseModel):
|
||||||
class MessageModel(BaseModel):
|
class MessageModel(BaseModel):
|
||||||
id: Optional[int] = Field(default=None, alias="id")
|
id: Optional[int] = Field(default=None, alias="id")
|
||||||
timestamp: datetime = Field(default_factory=datetime.now)
|
timestamp: datetime = Field(default_factory=datetime.now)
|
||||||
sender_type: str
|
sender: str
|
||||||
sender_name: str
|
sender_name: str
|
||||||
session_id: str
|
session_id: str
|
||||||
message: str
|
message: str
|
||||||
|
|
@ -47,6 +50,21 @@ class MessageModel(BaseModel):
|
||||||
return json.loads(v)
|
return json.loads(v)
|
||||||
return 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):
|
class VertexBuildModel(BaseModel):
|
||||||
index: Optional[int] = Field(default=None, alias="index", exclude=True)
|
index: Optional[int] = Field(default=None, alias="index", exclude=True)
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue