fix UP042 by using StrEnum

Author: bowenliang123

api/core/app/app_config/entities.py

@@ -1,5 +1,5 @@
 from collections.abc import Sequence
-from enum import Enum
+from enum import Enum, StrEnum
 from typing import Any, Optional
 
 from pydantic import BaseModel, Field, field_validator
@@ -88,7 +88,7 @@ def value_of(cls, value: str):
     advanced_completion_prompt_template: Optional[AdvancedCompletionPromptTemplateEntity] = None
 
 
-class VariableEntityType(str, Enum):
+class VariableEntityType(StrEnum):
     TEXT_INPUT = "text-input"
     SELECT = "select"
     PARAGRAPH = "paragraph"

api/core/app/entities/queue_entities.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from enum import Enum
+from enum import Enum, StrEnum
 from typing import Any, Optional
 
 from pydantic import BaseModel, field_validator
@@ -11,7 +11,7 @@
 from core.workflow.nodes.base import BaseNodeData
 
 
-class QueueEvent(str, Enum):
+class QueueEvent(StrEnum):
     """
     QueueEvent enum
     """

api/core/file/enums.py

@@ -1,7 +1,7 @@
-from enum import Enum
+from enum import StrEnum
 
 
-class FileType(str, Enum):
+class FileType(StrEnum):
     IMAGE = "image"
     DOCUMENT = "document"
     AUDIO = "audio"
@@ -16,7 +16,7 @@ def value_of(value):
         raise ValueError(f"No matching enum found for value '{value}'")
 
 
-class FileTransferMethod(str, Enum):
+class FileTransferMethod(StrEnum):
     REMOTE_URL = "remote_url"
     LOCAL_FILE = "local_file"
     TOOL_FILE = "tool_file"
@@ -29,7 +29,7 @@ def value_of(value):
         raise ValueError(f"No matching enum found for value '{value}'")
 
 
-class FileBelongsTo(str, Enum):
+class FileBelongsTo(StrEnum):
     USER = "user"
     ASSISTANT = "assistant"
 
@@ -41,7 +41,7 @@ def value_of(value):
         raise ValueError(f"No matching enum found for value '{value}'")
 
 
-class FileAttribute(str, Enum):
+class FileAttribute(StrEnum):
     TYPE = "type"
     SIZE = "size"
     NAME = "name"
@@ -51,5 +51,5 @@ class FileAttribute(str, Enum):
     EXTENSION = "extension"
 
 
-class ArrayFileAttribute(str, Enum):
+class ArrayFileAttribute(StrEnum):
     LENGTH = "length"

api/core/helper/code_executor/code_executor.py

@@ -1,5 +1,5 @@
 import logging
-from enum import Enum
+from enum import StrEnum
 from threading import Lock
 from typing import Optional
 
@@ -30,7 +30,7 @@ class Data(BaseModel):
     data: Data
 
 
-class CodeLanguage(str, Enum):
+class CodeLanguage(StrEnum):
     PYTHON3 = "python3"
     JINJA2 = "jinja2"
     JAVASCRIPT = "javascript"

api/core/model_runtime/entities/message_entities.py

@@ -1,5 +1,5 @@
 from abc import ABC
-from enum import Enum
+from enum import Enum, StrEnum
 from typing import Optional
 
 from pydantic import BaseModel, Field, field_validator
@@ -86,7 +86,7 @@ class ImagePromptMessageContent(PromptMessageContent):
     Model class for image prompt message content.
     """
 
-    class DETAIL(str, Enum):
+    class DETAIL(StrEnum):
         LOW = "low"
         HIGH = "high"
 

api/core/model_runtime/entities/model_entities.py

@@ -1,5 +1,5 @@
 from decimal import Decimal
-from enum import Enum
+from enum import Enum, StrEnum
 from typing import Any, Optional
 
 from pydantic import BaseModel, ConfigDict
@@ -89,7 +89,7 @@ class ModelFeature(Enum):
     STREAM_TOOL_CALL = "stream-tool-call"
 
 
-class DefaultParameterName(str, Enum):
+class DefaultParameterName(StrEnum):
     """
     Enum class for parameter template variable.
     """

api/core/ops/entities/trace_entity.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from enum import Enum
+from enum import StrEnum
 from typing import Any, Optional, Union
 
 from pydantic import BaseModel, ConfigDict, field_validator
@@ -111,7 +111,7 @@ class GenerateNameTraceInfo(BaseTraceInfo):
 }
 
 
-class TraceTaskName(str, Enum):
+class TraceTaskName(StrEnum):
     CONVERSATION_TRACE = "conversation"
     WORKFLOW_TRACE = "workflow"
     MESSAGE_TRACE = "message"

api/core/ops/langfuse_trace/entities/langfuse_trace_entity.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from enum import Enum
+from enum import StrEnum
 from typing import Any, Optional, Union
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator
@@ -39,7 +39,7 @@ def validate_input_output(v, field_name):
     return v
 
 
-class LevelEnum(str, Enum):
+class LevelEnum(StrEnum):
     DEBUG = "DEBUG"
     WARNING = "WARNING"
     ERROR = "ERROR"
@@ -178,7 +178,7 @@ def ensure_dict(cls, v, info: ValidationInfo):
         return validate_input_output(v, field_name)
 
 
-class UnitEnum(str, Enum):
+class UnitEnum(StrEnum):
     CHARACTERS = "CHARACTERS"
     TOKENS = "TOKENS"
     SECONDS = "SECONDS"

api/core/ops/langsmith_trace/entities/langsmith_trace_entity.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from enum import Enum
+from enum import StrEnum
 from typing import Any, Optional, Union
 
 from pydantic import BaseModel, Field, field_validator
@@ -8,7 +8,7 @@
 from core.ops.utils import replace_text_with_content
 
 
-class LangSmithRunType(str, Enum):
+class LangSmithRunType(StrEnum):
     tool = "tool"
     chain = "chain"
     llm = "llm"

api/core/prompt/simple_prompt_transform.py

@@ -23,7 +23,7 @@
     from core.file.models import File
 
 
-class ModelMode(str, enum.Enum):
+class ModelMode(enum.StrEnum):
     COMPLETION = "completion"
     CHAT = "chat"
 

api/core/rag/datasource/keyword/keyword_type.py

@@ -1,5 +1,5 @@
-from enum import Enum
+from enum import StrEnum
 
 
-class KeyWordType(str, Enum):
+class KeyWordType(StrEnum):
     JIEBA = "jieba"

api/core/rag/datasource/vdb/vector_type.py

@@ -1,7 +1,7 @@
-from enum import Enum
+from enum import StrEnum
 
 
-class VectorType(str, Enum):
+class VectorType(StrEnum):
     ANALYTICDB = "analyticdb"
     CHROMA = "chroma"
     MILVUS = "milvus"

api/core/rag/rerank/rerank_type.py

@@ -1,6 +1,6 @@
-from enum import Enum
+from enum import StrEnum
 
 
-class RerankMode(str, Enum):
+class RerankMode(StrEnum):
     RERANKING_MODEL = "reranking_model"
     WEIGHTED_SCORE = "weighted_score"

api/core/tools/entities/tool_entities.py

@@ -1,4 +1,4 @@
-from enum import Enum
+from enum import Enum, StrEnum
 from typing import Any, Optional, Union, cast
 
 from pydantic import BaseModel, Field, field_validator
@@ -137,7 +137,7 @@ def transform_id_to_str(cls, value) -> str:
 
 
 class ToolParameter(BaseModel):
-    class ToolParameterType(str, Enum):
+    class ToolParameterType(StrEnum):
         STRING = "string"
         NUMBER = "number"
         BOOLEAN = "boolean"

api/core/tools/tool/tool.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Mapping
 from copy import deepcopy
-from enum import Enum
+from enum import Enum, StrEnum
 from typing import TYPE_CHECKING, Any, Optional, Union
 
 from pydantic import BaseModel, ConfigDict, field_validator
@@ -62,7 +62,7 @@ def __init__(self, **data: Any):
     def __init__(self, **data: Any):
         super().__init__(**data)
 
-    class VariableKey(str, Enum):
+    class VariableKey(StrEnum):
         IMAGE = "image"
         DOCUMENT = "document"
         VIDEO = "video"

api/core/variables/types.py

@@ -1,7 +1,7 @@
-from enum import Enum
+from enum import StrEnum
 
 
-class SegmentType(str, Enum):
+class SegmentType(StrEnum):
     NONE = "none"
     NUMBER = "number"
     STRING = "string"

api/core/workflow/entities/node_entities.py

@@ -1,5 +1,5 @@
 from collections.abc import Mapping
-from enum import Enum
+from enum import StrEnum
 from typing import Any, Optional
 
 from pydantic import BaseModel
@@ -8,7 +8,7 @@
 from models.workflow import WorkflowNodeExecutionStatus
 
 
-class NodeRunMetadataKey(str, Enum):
+class NodeRunMetadataKey(StrEnum):
     """
     Node Run Metadata Key.
     """

api/core/workflow/enums.py

@@ -1,7 +1,7 @@
-from enum import Enum
+from enum import StrEnum
 
 
-class SystemVariableKey(str, Enum):
+class SystemVariableKey(StrEnum):
     """
     System Variables.
     """

api/core/workflow/nodes/enums.py

@@ -1,7 +1,7 @@
-from enum import Enum
+from enum import StrEnum
 
 
-class NodeType(str, Enum):
+class NodeType(StrEnum):
     START = "start"
     END = "end"
     ANSWER = "answer"

api/core/workflow/nodes/iteration/entities.py

@@ -1,12 +1,12 @@
-from enum import Enum
+from enum import StrEnum
 from typing import Any, Optional
 
 from pydantic import Field
 
 from core.workflow.nodes.base import BaseIterationNodeData, BaseIterationState, BaseNodeData
 
 
-class ErrorHandleMode(str, Enum):
+class ErrorHandleMode(StrEnum):
     TERMINATED = "terminated"
     CONTINUE_ON_ERROR = "continue-on-error"
     REMOVE_ABNORMAL_OUTPUT = "remove-abnormal-output"

api/core/workflow/nodes/variable_assigner/node_data.py

@@ -1,11 +1,11 @@
 from collections.abc import Sequence
-from enum import Enum
+from enum import StrEnum
 from typing import Optional
 
 from core.workflow.nodes.base import BaseNodeData
 
 
-class WriteMode(str, Enum):
+class WriteMode(StrEnum):
     OVER_WRITE = "over-write"
     APPEND = "append"
     CLEAR = "clear"

api/extensions/storage/storage_type.py

@@ -1,7 +1,7 @@
-from enum import Enum
+from enum import StrEnum
 
 
-class StorageType(str, Enum):
+class StorageType(StrEnum):
     ALIYUN_OSS = "aliyun-oss"
     AZURE_BLOB = "azure-blob"
     BAIDU_OBS = "baidu-obs"

api/models/account.py

@@ -8,7 +8,7 @@
 from .types import StringUUID
 
 
-class AccountStatus(str, enum.Enum):
+class AccountStatus(enum.StrEnum):
     PENDING = "pending"
     UNINITIALIZED = "uninitialized"
     ACTIVE = "active"
@@ -121,12 +121,12 @@ def is_dataset_operator(self):
         return self._current_tenant.current_role == TenantAccountRole.DATASET_OPERATOR
 
 
-class TenantStatus(str, enum.Enum):
+class TenantStatus(enum.StrEnum):
     NORMAL = "normal"
     ARCHIVE = "archive"
 
 
-class TenantAccountRole(str, enum.Enum):
+class TenantAccountRole(enum.StrEnum):
     OWNER = "owner"
     ADMIN = "admin"
     EDITOR = "editor"

api/models/dataset.py

@@ -23,7 +23,7 @@
 from .types import StringUUID
 
 
-class DatasetPermissionEnum(str, enum.Enum):
+class DatasetPermissionEnum(enum.StrEnum):
     ONLY_ME = "only_me"
     ALL_TEAM = "all_team_members"
     PARTIAL_TEAM = "partial_members"

api/models/enums.py

@@ -1,16 +1,16 @@
-from enum import Enum
+from enum import StrEnum
 
 
-class CreatedByRole(str, Enum):
+class CreatedByRole(StrEnum):
     ACCOUNT = "account"
     END_USER = "end_user"
 
 
-class UserFrom(str, Enum):
+class UserFrom(StrEnum):
     ACCOUNT = "account"
     END_USER = "end-user"
 
 
-class WorkflowRunTriggeredFrom(str, Enum):
+class WorkflowRunTriggeredFrom(StrEnum):
     DEBUGGING = "debugging"
     APP_RUN = "app-run"