refactor: Update Template class fields and serialization

Update the Template class in base.py to include a serialization alias for the type_name field and use the InputTypes class for the fields list. This refactor improves the organization and maintainability of the codebase by providing clearer field definitions and serialization behavior.
This commit is contained in:
ogabrielluiz 2024-06-12 16:17:57 -03:00
commit cd5cba112e

View file

@ -1,15 +1,15 @@
from typing import Callable, Union from typing import Callable, Union
from pydantic import BaseModel, model_serializer from pydantic import BaseModel, Field, model_serializer
from langflow.inputs.input_mixin import BaseInputMixin from langflow.inputs.inputs import InputTypes
from langflow.template.field.base import Input from langflow.template.field.base import Input
from langflow.utils.constants import DIRECT_TYPES from langflow.utils.constants import DIRECT_TYPES
class Template(BaseModel): class Template(BaseModel):
type_name: str type_name: str = Field(serialization_alias="_type")
fields: list[Input | BaseInputMixin] fields: list[Input | InputTypes]
def process_fields( def process_fields(
self, self,
@ -30,7 +30,7 @@ class Template(BaseModel):
result = handler(self) result = handler(self)
for field in self.fields: for field in self.fields:
result[field.name] = field.model_dump(by_alias=True, exclude_none=True) result[field.name] = field.model_dump(by_alias=True, exclude_none=True)
result["_type"] = result.pop("type_name")
return result return result
# For backwards compatibility # For backwards compatibility