feat(base.py): add methods to get, update, and upsert fields in the Template class for easier field manipulation

🔧 chore(base.py): refactor Template class to improve code readability and maintainability
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-12-11 12:30:03 -03:00
commit 17c6ed1db3

View file

@ -1,8 +1,9 @@
from typing import Callable, Union
from pydantic import BaseModel, model_serializer
from langflow.template.field.base import TemplateField
from langflow.utils.constants import DIRECT_TYPES
from pydantic import BaseModel, model_serializer
class Template(BaseModel):
@ -39,3 +40,25 @@ class Template(BaseModel):
def add_field(self, field: TemplateField) -> None:
self.fields.append(field)
def get_field(self, field_name: str) -> TemplateField:
"""Returns the field with the given name."""
field = next((field for field in self.fields if field.name == field_name), None)
if field is None:
raise ValueError(f"Field {field_name} not found in template {self.type_name}")
return field
def update_field(self, field_name: str, field: TemplateField) -> None:
"""Updates the field with the given name."""
for idx, template_field in enumerate(self.fields):
if template_field.name == field_name:
self.fields[idx] = field
return
raise ValueError(f"Field {field_name} not found in template {self.type_name}")
def upsert_field(self, field_name: str, field: TemplateField) -> None:
"""Updates the field with the given name or adds it if it doesn't exist."""
try:
self.update_field(field_name, field)
except ValueError:
self.add_field(field)