From 8f4515e6555e4bee58ce4021196d2a65e4f63e5a Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Wed, 5 Mar 2025 10:49:23 -0300 Subject: [PATCH] fix: update data validation in Data model (#6789) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: enhance data validation in Data model * fix: Add warning for invalid data format in Data schema Modify data validation to log a warning when data is not a dictionary, preparing for a future error in version 1.3.0 * ✅ (test_data_class.py): add assertion to check if data attribute is an empty dictionary after calling get_text() method * 🐛 (data.py): fix issue where values["data"] could be None, causing potential errors in the code. Now ensuring "data" key exists and is initialized as an empty dictionary if it is not present. --------- Co-authored-by: Cristhian Zanforlin Lousa --- src/backend/base/langflow/schema/data.py | 12 ++++++++++-- src/backend/tests/unit/test_data_class.py | 1 + 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/src/backend/base/langflow/schema/data.py b/src/backend/base/langflow/schema/data.py index 1ea72ac15..dea76b4f7 100644 --- a/src/backend/base/langflow/schema/data.py +++ b/src/backend/base/langflow/schema/data.py @@ -8,7 +8,7 @@ from uuid import UUID from langchain_core.documents import Document from langchain_core.messages import AIMessage, BaseMessage, HumanMessage from loguru import logger -from pydantic import BaseModel, model_serializer, model_validator +from pydantic import BaseModel, ConfigDict, model_serializer, model_validator from langflow.utils.constants import MESSAGE_SENDER_AI, MESSAGE_SENDER_USER from langflow.utils.image import create_data_url @@ -21,6 +21,8 @@ class Data(BaseModel): data (dict, optional): Additional data associated with the record. """ + model_config = ConfigDict(validate_assignment=True) + text_key: str = "text" data: dict = {} default_value: str | None = "" @@ -31,8 +33,14 @@ class Data(BaseModel): if not isinstance(values, dict): msg = "Data must be a dictionary" raise ValueError(msg) # noqa: TRY004 - if not values.get("data"): + if "data" not in values or values["data"] is None: values["data"] = {} + if not isinstance(values["data"], dict): + msg = ( + f"Invalid data format: expected dictionary but got {type(values).__name__}." + " This will raise an error in version langflow==1.3.0." + ) + logger.warning(msg) # Any other keyword should be added to the data dictionary for key in values: if key not in values["data"] and key not in {"text_key", "data", "default_value"}: diff --git a/src/backend/tests/unit/test_data_class.py b/src/backend/tests/unit/test_data_class.py index 9af22f7dd..03fc80cff 100644 --- a/src/backend/tests/unit/test_data_class.py +++ b/src/backend/tests/unit/test_data_class.py @@ -138,3 +138,4 @@ def test_get_text_with_none_data(): schema = Data(data=data, text_key="text", default_value="default") result = schema.get_text() assert result == "default" + assert schema.data == {}