diff --git a/src/backend/langflow/template/template/base.py b/src/backend/langflow/template/template/base.py index bb5d35375..d7632e239 100644 --- a/src/backend/langflow/template/template/base.py +++ b/src/backend/langflow/template/template/base.py @@ -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)