jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

refactor: Update BaseInputMixin to include model_serializer and serialize_model method

Browse source 
This commit updates the `BaseInputMixin` class in the `input_mixin.py` file. It adds the `model_serializer` decorator and the `serialize_model` method to the mixin. The `model_serializer` decorator is used to specify the serialization behavior of the model, and the `serialize_model` method is responsible for serializing the model into a dictionary representation.

These additions enhance the flexibility and reusability of the `BaseInputMixin` class.
This commit is contained in:
ogabrielluiz 2024-06-12 16:11:39 -03:00
commit 01d396e037
1 changed files with 13 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

19
src/backend/base/langflow/inputs/input_mixin.py
View file

@ -1,7 +1,7 @@
from enum import Enum from enum import Enum
from typing import Annotated, Any, Optional from typing import Annotated, Any, Optional
from pydantic import BaseModel, ConfigDict, Field, PlainSerializer, field_validator from pydantic import BaseModel, ConfigDict, Field, PlainSerializer, field_validator, model_serializer
from langflow.field_typing.range_spec import RangeSpec from langflow.field_typing.range_spec import RangeSpec
@ -9,13 +9,13 @@ from langflow.field_typing.range_spec import RangeSpec
class FieldTypes(str, Enum): class FieldTypes(str, Enum):
TEXT = "str" TEXT = "str"
INTEGER = "int" INTEGER = "int"
PASSWORD = "SecretStr" PASSWORD = "str"
FLOAT = "float" FLOAT = "float"
BOOLEAN = "bool" BOOLEAN = "bool"
DICT = "dict" DICT = "dict"
NESTED_DICT = "NestedDict" NESTED_DICT = "NestedDict"
FILE = "file" FILE = "file"
PROMPT = "Prompt" PROMPT = "prompt"
SerializableFieldTypes = Annotated[FieldTypes, PlainSerializer(lambda v: v.value, return_type=str)] SerializableFieldTypes = Annotated[FieldTypes, PlainSerializer(lambda v: v.value, return_type=str)]
@ -69,7 +69,7 @@ class BaseInputMixin(BaseModel):
"""Specifies if the field should be displayed in title case. Defaults to True.""" """Specifies if the field should be displayed in title case. Defaults to True."""
def to_dict(self): def to_dict(self):
return self.model_dump(exclude_none=True) return self.model_dump(exclude_none=True, by_alias=True)
@field_validator("field_type", mode="before") @field_validator("field_type", mode="before")
@classmethod @classmethod
@ -78,15 +78,22 @@ class BaseInputMixin(BaseModel):
raise ValueError(f"field_type must be one of {FieldTypes}") raise ValueError(f"field_type must be one of {FieldTypes}")
return FieldTypes(v) return FieldTypes(v)
@model_serializer(mode="wrap")
def serialize_model(self, handler):
dump = handler(self)
if "field_type" in dump:
dump["type"] = dump.pop("field_type")
return dump
# Mixin for input fields that can be listable # Mixin for input fields that can be listable
class ListableInputMixin(BaseModel): class ListableInputMixin(BaseModel):
is_list: bool = Field(default=False) is_list: bool = Field(default=False, serialization_alias="list")
# Specific mixin for fields needing database interaction # Specific mixin for fields needing database interaction
class DatabaseLoadMixin(BaseModel): class DatabaseLoadMixin(BaseModel):
load_from_db: bool = Field(default=False) load_from_db: bool = Field(default=True)
# Specific mixin for fields needing file interaction # Specific mixin for fields needing file interaction