From dc1c092fe7c416fba8b36915fe7c8d9d0dc6a6de Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Wed, 6 Mar 2024 10:55:07 -0300
Subject: [PATCH] Add icon and icon_bg_color fields to Flow model

---
 ...fd30_add_icon_and_icon_bg_color_to_flow.py | 56 +++++++++++++++++++
 .../services/database/models/flow/model.py    | 38 +++++++++++++
 2 files changed, 94 insertions(+)
 create mode 100644 src/backend/langflow/alembic/versions/63b9c451fd30_add_icon_and_icon_bg_color_to_flow.py

diff --git a/src/backend/langflow/alembic/versions/63b9c451fd30_add_icon_and_icon_bg_color_to_flow.py b/src/backend/langflow/alembic/versions/63b9c451fd30_add_icon_and_icon_bg_color_to_flow.py
new file mode 100644
index 000000000..3deb66346
--- /dev/null
+++ b/src/backend/langflow/alembic/versions/63b9c451fd30_add_icon_and_icon_bg_color_to_flow.py
@@ -0,0 +1,56 @@
+"""Add icon and icon_bg_color to Flow
+
+Revision ID: 63b9c451fd30
+Revises: bc2f01c40e4a
+Create Date: 2024-03-06 10:53:47.148658
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+from sqlalchemy.engine.reflection import Inspector
+
+# revision identifiers, used by Alembic.
+revision: str = "63b9c451fd30"
+down_revision: Union[str, None] = "bc2f01c40e4a"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    inspector = Inspector.from_engine(conn)  # type: ignore
+    table_names = inspector.get_table_names()
+    column_names = [column["name"] for column in inspector.get_columns("flow")]
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("flow", schema=None) as batch_op:
+        if "icon" not in column_names:
+            batch_op.add_column(
+                sa.Column("icon", sqlmodel.sql.sqltypes.AutoString(), nullable=True)
+            )
+        if "icon_bg_color" not in column_names:
+            batch_op.add_column(
+                sa.Column(
+                    "icon_bg_color", sqlmodel.sql.sqltypes.AutoString(), nullable=True
+                )
+            )
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    conn = op.get_bind()
+    inspector = Inspector.from_engine(conn)  # type: ignore
+    table_names = inspector.get_table_names()
+    column_names = [column["name"] for column in inspector.get_columns("flow")]
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("flow", schema=None) as batch_op:
+        if "icon" in column_names:
+            batch_op.drop_column("icon")
+        if "icon_bg_color" in column_names:
+            batch_op.drop_column("icon_bg_color")
+
+    # ### end Alembic commands ###
diff --git a/src/backend/langflow/services/database/models/flow/model.py b/src/backend/langflow/services/database/models/flow/model.py
index c7ee08141..b76fd11cf 100644
--- a/src/backend/langflow/services/database/models/flow/model.py
+++ b/src/backend/langflow/services/database/models/flow/model.py
@@ -1,5 +1,6 @@
 # Path: src/backend/langflow/database/models/flow.py
 
+import re
 from datetime import datetime
 from typing import TYPE_CHECKING, Dict, Optional
 from uuid import UUID, uuid4
@@ -7,6 +8,7 @@ from uuid import UUID, uuid4
 from pydantic import field_serializer, field_validator
 from sqlmodel import JSON, Column, Field, Relationship, SQLModel
 
+from langflow.interface.custom.attributes import validate_icon
 from langflow.schema.schema import Record
 
 if TYPE_CHECKING:
@@ -16,6 +18,8 @@ if TYPE_CHECKING:
 class FlowBase(SQLModel):
     name: str = Field(index=True)
     description: Optional[str] = Field(index=True, nullable=True, default=None)
+    icon: Optional[str] = Field(default=None, nullable=True)
+    icon_bg_color: Optional[str] = Field(default=None, nullable=True)
     data: Optional[Dict] = Field(default=None, nullable=True)
     is_component: Optional[bool] = Field(default=False, nullable=True)
     updated_at: Optional[datetime] = Field(
@@ -23,6 +27,40 @@ class FlowBase(SQLModel):
     )
     folder: Optional[str] = Field(default=None, nullable=True)
 
+    @field_validator("icon_bg_color")
+    def validate_icon_bg_color(cls, v):
+        if v is not None and not isinstance(v, str):
+            raise ValueError("Icon background color must be a string")
+        # validate that is is a hex color
+        if v and not v.startswith("#"):
+            raise ValueError("Icon background color must start with #")
+
+        # validate that it is a valid hex color
+        if v and len(v) != 7:
+            raise ValueError("Icon background color must be 7 characters long")
+        return v
+
+    @field_validator("icon")
+    def validate_icon_atr(cls, v):
+        #   const emojiRegex = /\p{Emoji}/u;
+        # const isEmoji = emojiRegex.test(data?.node?.icon!);
+        # emoji pattern in Python
+        pattern = r"\p{Emoji}"
+        emoji = validate_icon(v)
+        is_emoji = re.search(pattern, emoji)
+        if is_emoji:
+            # this is indeed an emoji
+            return emoji
+        # otherwise it should be a valid lucide icon
+        if v is not None and not isinstance(v, str):
+            raise ValueError("Icon must be a string")
+        # is should be lowercase and contain only letters and hyphens
+        if v and not v.islower():
+            raise ValueError("Icon must be lowercase")
+        if v and not v.replace("-", "").isalpha():
+            raise ValueError("Icon must contain only letters and hyphens")
+        return v
+
     @field_validator("data")
     def validate_json(v):
         if not v: