refactor: Update TransactionModel and TransactionModelResponse fields

This commit is contained in:
ogabrielluiz 2024-06-07 15:57:08 -03:00
commit 0db23d2756
2 changed files with 12 additions and 32 deletions

View file

@ -11,10 +11,9 @@ if TYPE_CHECKING:
class TransactionModel(BaseModel):
index: Optional[int] = Field(default=None)
timestamp: Optional[datetime] = Field(default_factory=datetime.now, alias="timestamp")
flow_id: str
source: str
target: str
target_args: dict
vertex_id: str
inputs: dict
outputs: dict
status: str
error: Optional[str] = None
@ -23,13 +22,13 @@ class TransactionModel(BaseModel):
populate_by_name = True
# validate target_args in case it is a JSON
@field_validator("target_args", mode="before")
@field_validator("outputs", "inputs", mode="before")
def validate_target_args(cls, v):
if isinstance(v, str):
return json.loads(v)
return v
@field_serializer("target_args")
@field_serializer("outputs", "inputs")
def serialize_target_args(v):
if isinstance(v, dict):
return json.dumps(v)
@ -39,10 +38,9 @@ class TransactionModel(BaseModel):
class TransactionModelResponse(BaseModel):
index: Optional[int] = Field(default=None)
timestamp: Optional[datetime] = Field(default_factory=datetime.now, alias="timestamp")
flow_id: str
source: str
target: str
target_args: dict
vertex_id: str
inputs: dict
outputs: dict
status: str
error: Optional[str] = None
@ -51,7 +49,7 @@ class TransactionModelResponse(BaseModel):
populate_by_name = True
# validate target_args in case it is a JSON
@field_validator("target_args", mode="before")
@field_validator("outputs", "inputs", mode="before")
def validate_target_args(cls, v):
if isinstance(v, str):
return json.loads(v)
@ -76,7 +74,6 @@ class MessageModel(BaseModel):
session_id: str
message: str
files: list[str] = []
artifacts: dict
class Config:
from_attributes = True
@ -88,12 +85,6 @@ class MessageModel(BaseModel):
return json.loads(v)
return v
@field_validator("artifacts", mode="before")
def validate_target_args(cls, v):
if isinstance(v, str):
return json.loads(v)
return v
@classmethod
def from_record(cls, record: "Record", flow_id: Optional[str] = None):
# first check if the record has all the required fields
@ -114,12 +105,6 @@ class MessageModel(BaseModel):
class MessageModelResponse(MessageModel):
index: Optional[int] = Field(default=None)
@field_validator("artifacts", mode="before")
def serialize_artifacts(v):
if isinstance(v, str):
return json.loads(v)
return v
@field_validator("index", mode="before")
def validate_id(cls, v):
if isinstance(v, float):

View file

@ -119,21 +119,16 @@ async def log_message(
sender_name: str,
message: str,
session_id: str,
artifacts: Optional[dict] = None,
files: Optional[list] = None,
flow_id: Optional[str] = None,
):
try:
from langflow.graph.vertex.base import Vertex
if isinstance(session_id, Vertex):
session_id = await session_id.build() # type: ignore
monitor_service = get_monitor_service()
row = {
"sender": sender,
"sender_name": sender_name,
"message": message,
"artifacts": artifacts or {},
"files": files or [],
"session_id": session_id,
"timestamp": monitor_service.get_timestamp(),
"flow_id": flow_id,
@ -190,7 +185,7 @@ def log_transaction(vertex: "Vertex", status, error=None):
data = {
"vertex_id": vertex.id,
"inputs": clean_params,
"output": str(vertex.result),
"outputs": vertex.result.model_dump_json(),
"timestamp": monitor_service.get_timestamp(),
"status": status,
"error": error,