From b6546e456e79202064f5256b62773ab9b54ee963 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
Date: Thu, 3 Oct 2024 15:20:38 -0300
Subject: [PATCH] feat: add 'tags' column to 'flow' table and update models
 (#3986)

* Add 'tags' column to 'flow' table and update models

- Added migration script to include 'tags' column in 'flow' table.
- Updated `Flow` model to include `tags` field.
- Introduced `Tags` enum in `schema.py` for predefined tag values.

* Update `tags` column to use JSON type in Flow model

* Add conditional checks for 'tags' column in Alembic migration script

* Make 'tags' field nullable in Flow model

* Add default value for 'tags' field in Flow model

* Update default values for 'tags' field in Flow model
---
 .../d2d475a1f7c0_add_tags_column_to_flow.py   | 41 +++++++++++++++++++
 .../services/database/models/flow/model.py    |  2 +
 .../services/database/models/flow/schema.py   |  6 +++
 3 files changed, 49 insertions(+)
 create mode 100644 src/backend/base/langflow/alembic/versions/d2d475a1f7c0_add_tags_column_to_flow.py
 create mode 100644 src/backend/base/langflow/services/database/models/flow/schema.py

diff --git a/src/backend/base/langflow/alembic/versions/d2d475a1f7c0_add_tags_column_to_flow.py b/src/backend/base/langflow/alembic/versions/d2d475a1f7c0_add_tags_column_to_flow.py
new file mode 100644
index 000000000..d314930ff
--- /dev/null
+++ b/src/backend/base/langflow/alembic/versions/d2d475a1f7c0_add_tags_column_to_flow.py
@@ -0,0 +1,41 @@
+"""add tags column to flow
+
+Revision ID: d2d475a1f7c0
+Revises: d3dbf656a499
+Create Date: 2024-10-03 13:33:59.517261
+
+"""
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+from sqlalchemy.engine.reflection import Inspector
+
+from langflow.utils import migration
+
+# revision identifiers, used by Alembic.
+revision: str = 'd2d475a1f7c0'
+down_revision: Union[str, None] = 'd3dbf656a499'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('flow', schema=None) as batch_op:
+        if not migration.column_exists(table_name='flow', column_name='tags', conn=conn):
+            batch_op.add_column(sa.Column('tags', sa.JSON(), nullable=True))
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    conn = op.get_bind()
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('flow', schema=None) as batch_op:
+        if migration.column_exists(table_name='flow', column_name='tags', conn=conn):
+            batch_op.drop_column('tags')
+
+    # ### end Alembic commands ###
diff --git a/src/backend/base/langflow/services/database/models/flow/model.py b/src/backend/base/langflow/services/database/models/flow/model.py
index 1898c94f3..6d50c64fd 100644
--- a/src/backend/base/langflow/services/database/models/flow/model.py
+++ b/src/backend/base/langflow/services/database/models/flow/model.py
@@ -34,6 +34,7 @@ class FlowBase(SQLModel):
     updated_at: datetime | None = Field(default_factory=lambda: datetime.now(timezone.utc), nullable=True)
     webhook: bool | None = Field(default=False, nullable=True, description="Can be used on the webhook endpoint")
     endpoint_name: str | None = Field(default=None, nullable=True, index=True)
+    tags: list[str] | None = None
 
     @field_validator("endpoint_name")
     @classmethod
@@ -152,6 +153,7 @@ class Flow(FlowBase, table=True):  # type: ignore
     data: dict | None = Field(default=None, sa_column=Column(JSON))
     user_id: UUID | None = Field(index=True, foreign_key="user.id", nullable=True)
     user: "User" = Relationship(back_populates="flows")
+    tags: list[str] | None = Field(sa_column=Column(JSON), default=[])
     folder_id: UUID | None = Field(default=None, foreign_key="folder.id", nullable=True, index=True)
     folder: Optional["Folder"] = Relationship(back_populates="flows")
     messages: list["MessageTable"] = Relationship(back_populates="flow")
diff --git a/src/backend/base/langflow/services/database/models/flow/schema.py b/src/backend/base/langflow/services/database/models/flow/schema.py
new file mode 100644
index 000000000..aa030caa1
--- /dev/null
+++ b/src/backend/base/langflow/services/database/models/flow/schema.py
@@ -0,0 +1,6 @@
+from enum import Enum
+
+
+class Tags(str, Enum):
+    CHATBOTS = "chatbots"
+    AGENTS = "agents"