From 9bae1dff97bb517fd97e963e91eec3db70676115 Mon Sep 17 00:00:00 2001 From: Bowen Liang Date: Wed, 6 Nov 2024 23:55:36 +0800 Subject: [PATCH] fix UP042 by using StrEnum --- api/core/app/app_config/entities.py | 4 ++-- api/core/app/entities/queue_entities.py | 4 ++-- api/core/file/enums.py | 12 ++++++------ api/core/helper/code_executor/code_executor.py | 4 ++-- api/core/model_runtime/entities/message_entities.py | 4 ++-- api/core/model_runtime/entities/model_entities.py | 4 ++-- api/core/ops/entities/trace_entity.py | 4 ++-- .../langfuse_trace/entities/langfuse_trace_entity.py | 6 +++--- .../entities/langsmith_trace_entity.py | 4 ++-- api/core/prompt/simple_prompt_transform.py | 2 +- api/core/rag/datasource/keyword/keyword_type.py | 4 ++-- api/core/rag/datasource/vdb/vector_type.py | 4 ++-- api/core/rag/rerank/rerank_type.py | 4 ++-- api/core/tools/entities/tool_entities.py | 4 ++-- api/core/tools/tool/tool.py | 4 ++-- api/core/variables/types.py | 4 ++-- api/core/workflow/entities/node_entities.py | 4 ++-- api/core/workflow/enums.py | 4 ++-- api/core/workflow/nodes/enums.py | 4 ++-- api/core/workflow/nodes/iteration/entities.py | 4 ++-- .../workflow/nodes/variable_assigner/node_data.py | 4 ++-- api/extensions/storage/storage_type.py | 4 ++-- api/models/account.py | 6 +++--- api/models/dataset.py | 2 +- api/models/enums.py | 8 ++++---- api/models/model.py | 4 ++-- api/services/auth/auth_type.py | 4 ++-- api/services/recommend_app/recommend_app_type.py | 4 ++-- 28 files changed, 62 insertions(+), 62 deletions(-) diff --git a/api/core/app/app_config/entities.py b/api/core/app/app_config/entities.py index 9b72452d7a1a9e..15bd353484e372 100644 --- a/api/core/app/app_config/entities.py +++ b/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" diff --git a/api/core/app/entities/queue_entities.py b/api/core/app/entities/queue_entities.py index f1542ec5d8c578..319a67432f4641 100644 --- a/api/core/app/entities/queue_entities.py +++ b/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 """ diff --git a/api/core/file/enums.py b/api/core/file/enums.py index f4153f1676b620..06b99d3eb0b4b4 100644 --- a/api/core/file/enums.py +++ b/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" diff --git a/api/core/helper/code_executor/code_executor.py b/api/core/helper/code_executor/code_executor.py index 4932284540738a..3d2d47cbc4d654 100644 --- a/api/core/helper/code_executor/code_executor.py +++ b/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" diff --git a/api/core/model_runtime/entities/message_entities.py b/api/core/model_runtime/entities/message_entities.py index 3c244d368ef78b..ddd48e06485fdf 100644 --- a/api/core/model_runtime/entities/message_entities.py +++ b/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 @@ -93,7 +93,7 @@ class ImagePromptMessageContent(PromptMessageContent): Model class for image prompt message content. """ - class DETAIL(str, Enum): + class DETAIL(StrEnum): LOW = "low" HIGH = "high" diff --git a/api/core/model_runtime/entities/model_entities.py b/api/core/model_runtime/entities/model_entities.py index 52ea787c3ad572..1f0370912e2cc8 100644 --- a/api/core/model_runtime/entities/model_entities.py +++ b/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. """ diff --git a/api/core/ops/entities/trace_entity.py b/api/core/ops/entities/trace_entity.py index 256595286f324e..71ff03b6ef5160 100644 --- a/api/core/ops/entities/trace_entity.py +++ b/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 @@ -122,7 +122,7 @@ class TaskData(BaseModel): } -class TraceTaskName(str, Enum): +class TraceTaskName(StrEnum): CONVERSATION_TRACE = "conversation" WORKFLOW_TRACE = "workflow" MESSAGE_TRACE = "message" diff --git a/api/core/ops/langfuse_trace/entities/langfuse_trace_entity.py b/api/core/ops/langfuse_trace/entities/langfuse_trace_entity.py index 447b799f1f1187..f486da3a6d0c90 100644 --- a/api/core/ops/langfuse_trace/entities/langfuse_trace_entity.py +++ b/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" diff --git a/api/core/ops/langsmith_trace/entities/langsmith_trace_entity.py b/api/core/ops/langsmith_trace/entities/langsmith_trace_entity.py index 05c932fb99b424..348b7ba5012b6b 100644 --- a/api/core/ops/langsmith_trace/entities/langsmith_trace_entity.py +++ b/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" diff --git a/api/core/prompt/simple_prompt_transform.py b/api/core/prompt/simple_prompt_transform.py index 5a3481b96388e4..93dd92f188a9c6 100644 --- a/api/core/prompt/simple_prompt_transform.py +++ b/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" diff --git a/api/core/rag/datasource/keyword/keyword_type.py b/api/core/rag/datasource/keyword/keyword_type.py index d6deba3fb09fdf..d845c7111dd578 100644 --- a/api/core/rag/datasource/keyword/keyword_type.py +++ b/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" diff --git a/api/core/rag/datasource/vdb/vector_type.py b/api/core/rag/datasource/vdb/vector_type.py index 8e53e3ae8450d6..05183c03717fa3 100644 --- a/api/core/rag/datasource/vdb/vector_type.py +++ b/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" diff --git a/api/core/rag/rerank/rerank_type.py b/api/core/rag/rerank/rerank_type.py index d71eb2daa8f920..b2d13146540442 100644 --- a/api/core/rag/rerank/rerank_type.py +++ b/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" diff --git a/api/core/tools/entities/tool_entities.py b/api/core/tools/entities/tool_entities.py index d8637fd2cb0c52..4fc383f91baeba 100644 --- a/api/core/tools/entities/tool_entities.py +++ b/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" diff --git a/api/core/tools/tool/tool.py b/api/core/tools/tool/tool.py index 6cb6e18b6d4e84..f92ebfb39202b7 100644 --- a/api/core/tools/tool/tool.py +++ b/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" diff --git a/api/core/variables/types.py b/api/core/variables/types.py index 53c2e8a3aa6ddc..af6a2a29370846 100644 --- a/api/core/variables/types.py +++ b/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" diff --git a/api/core/workflow/entities/node_entities.py b/api/core/workflow/entities/node_entities.py index 7e10cddc712baa..f401ce084881ea 100644 --- a/api/core/workflow/entities/node_entities.py +++ b/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. """ diff --git a/api/core/workflow/enums.py b/api/core/workflow/enums.py index 213ed57f570968..9642efa1a5ca9e 100644 --- a/api/core/workflow/enums.py +++ b/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. """ diff --git a/api/core/workflow/nodes/enums.py b/api/core/workflow/nodes/enums.py index 208144655b5a59..9e9e52910ebcf8 100644 --- a/api/core/workflow/nodes/enums.py +++ b/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" diff --git a/api/core/workflow/nodes/iteration/entities.py b/api/core/workflow/nodes/iteration/entities.py index ebcb6f82fbc397..7a489dd725d653 100644 --- a/api/core/workflow/nodes/iteration/entities.py +++ b/api/core/workflow/nodes/iteration/entities.py @@ -1,4 +1,4 @@ -from enum import Enum +from enum import StrEnum from typing import Any, Optional from pydantic import Field @@ -6,7 +6,7 @@ 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" diff --git a/api/core/workflow/nodes/variable_assigner/node_data.py b/api/core/workflow/nodes/variable_assigner/node_data.py index 70ae29d45f47be..474ecefe76162a 100644 --- a/api/core/workflow/nodes/variable_assigner/node_data.py +++ b/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" diff --git a/api/extensions/storage/storage_type.py b/api/extensions/storage/storage_type.py index 415bf251f6e280..e7fa405afacdcc 100644 --- a/api/extensions/storage/storage_type.py +++ b/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" diff --git a/api/models/account.py b/api/models/account.py index 60b4f11aad2bdc..6684e8dd6ec65d 100644 --- a/api/models/account.py +++ b/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" diff --git a/api/models/dataset.py b/api/models/dataset.py index a1a626d7e46b2c..c4337ede564c19 100644 --- a/api/models/dataset.py +++ b/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" diff --git a/api/models/enums.py b/api/models/enums.py index a83d35e04245b7..7b9500ebe4d36a 100644 --- a/api/models/enums.py +++ b/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" diff --git a/api/models/model.py b/api/models/model.py index e909d53e3e29d3..0ba3d88bc95370 100644 --- a/api/models/model.py +++ b/api/models/model.py @@ -3,7 +3,7 @@ import uuid from collections.abc import Mapping from datetime import datetime -from enum import Enum +from enum import Enum, StrEnum from typing import Any, Literal, Optional import sqlalchemy as sa @@ -32,7 +32,7 @@ class DifySetup(db.Model): setup_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) -class AppMode(str, Enum): +class AppMode(StrEnum): COMPLETION = "completion" WORKFLOW = "workflow" CHAT = "chat" diff --git a/api/services/auth/auth_type.py b/api/services/auth/auth_type.py index 2d6e901447c369..2e1946841fbd15 100644 --- a/api/services/auth/auth_type.py +++ b/api/services/auth/auth_type.py @@ -1,6 +1,6 @@ -from enum import Enum +from enum import StrEnum -class AuthType(str, Enum): +class AuthType(StrEnum): FIRECRAWL = "firecrawl" JINA = "jinareader" diff --git a/api/services/recommend_app/recommend_app_type.py b/api/services/recommend_app/recommend_app_type.py index 7ea93b3f64b1d4..e60e435b3a02f8 100644 --- a/api/services/recommend_app/recommend_app_type.py +++ b/api/services/recommend_app/recommend_app_type.py @@ -1,7 +1,7 @@ -from enum import Enum +from enum import StrEnum -class RecommendAppType(str, Enum): +class RecommendAppType(StrEnum): REMOTE = "remote" BUILDIN = "builtin" DATABASE = "db"