From 1daf9151692496714318af7377209f2aefd617c2 Mon Sep 17 00:00:00 2001
From: Lucas Oliveira <62335616+lucaseduoli@users.noreply.github.com>
Date: Mon, 23 Dec 2024 15:54:40 -0300
Subject: [PATCH] fix: add data field to returned object on flowheader if flow
 is a component (#5373)

* Fixed FlowHeader model validator, adding data and removing it if flow is not a component

* [autofix.ci] apply automated fixes

* Fixed flowheader to use fieldvalidator instead of modelvalidator

* [autofix.ci] apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
---
 .../services/database/models/flow/model.py    | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/src/backend/base/langflow/services/database/models/flow/model.py b/src/backend/base/langflow/services/database/models/flow/model.py
index 132ea94de..7b154381c 100644
--- a/src/backend/base/langflow/services/database/models/flow/model.py
+++ b/src/backend/base/langflow/services/database/models/flow/model.py
@@ -9,7 +9,12 @@ import emoji
 from emoji import purely_emoji
 from fastapi import HTTPException, status
 from loguru import logger
-from pydantic import BaseModel, field_serializer, field_validator, model_validator
+from pydantic import (
+    BaseModel,
+    ValidationInfo,
+    field_serializer,
+    field_validator,
+)
 from sqlalchemy import Text, UniqueConstraint
 from sqlmodel import JSON, Column, Field, Relationship, SQLModel
 
@@ -203,16 +208,20 @@ class FlowHeader(BaseModel):
     id: UUID = Field(description="Unique identifier for the flow")
     name: str = Field(description="The name of the flow")
     folder_id: UUID | None = Field(
-        None, description="The ID of the folder containing the flow. None if not associated with a folder"
+        None,
+        description="The ID of the folder containing the flow. None if not associated with a folder",
     )
     is_component: bool | None = Field(None, description="Flag indicating whether the flow is a component")
     endpoint_name: str | None = Field(None, description="The name of the endpoint associated with this flow")
     description: str | None = Field(None, description="A description of the flow")
+    data: dict | None = Field(None, description="The data of the component, if is_component is True")
 
-    @model_validator(mode="before")
+    @field_validator("data", mode="before")
     @classmethod
-    def validate_flow_header(cls, data: dict):
-        return data
+    def validate_flow_header(cls, value: dict, info: ValidationInfo):
+        if not info.data["is_component"]:
+            return None
+        return value
 
 
 class FlowUpdate(SQLModel):