Skip to content

Commit

Permalink
fix UP042 by using StrEnum
Browse files Browse the repository at this point in the history
  • Loading branch information
bowenliang123 committed Nov 8, 2024
1 parent d491e87 commit 9bae1df
Show file tree
Hide file tree
Showing 28 changed files with 62 additions and 62 deletions.
4 changes: 2 additions & 2 deletions api/core/app/app_config/entities.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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"
Expand Down
4 changes: 2 additions & 2 deletions api/core/app/entities/queue_entities.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -11,7 +11,7 @@
from core.workflow.nodes.base import BaseNodeData


class QueueEvent(str, Enum):
class QueueEvent(StrEnum):
"""
QueueEvent enum
"""
Expand Down
12 changes: 6 additions & 6 deletions api/core/file/enums.py
Original file line number Diff line number Diff line change
@@ -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"
Expand All @@ -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"
Expand All @@ -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"

Expand All @@ -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"
Expand All @@ -51,5 +51,5 @@ class FileAttribute(str, Enum):
EXTENSION = "extension"


class ArrayFileAttribute(str, Enum):
class ArrayFileAttribute(StrEnum):
LENGTH = "length"
4 changes: 2 additions & 2 deletions api/core/helper/code_executor/code_executor.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import logging
from enum import Enum
from enum import StrEnum
from threading import Lock
from typing import Optional

Expand Down Expand Up @@ -30,7 +30,7 @@ class Data(BaseModel):
data: Data


class CodeLanguage(str, Enum):
class CodeLanguage(StrEnum):
PYTHON3 = "python3"
JINJA2 = "jinja2"
JAVASCRIPT = "javascript"
Expand Down
4 changes: 2 additions & 2 deletions api/core/model_runtime/entities/message_entities.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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"

Expand Down
4 changes: 2 additions & 2 deletions api/core/model_runtime/entities/model_entities.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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.
"""
Expand Down
4 changes: 2 additions & 2 deletions api/core/ops/entities/trace_entity.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -122,7 +122,7 @@ class TaskData(BaseModel):
}


class TraceTaskName(str, Enum):
class TraceTaskName(StrEnum):
CONVERSATION_TRACE = "conversation"
WORKFLOW_TRACE = "workflow"
MESSAGE_TRACE = "message"
Expand Down
6 changes: 3 additions & 3 deletions api/core/ops/langfuse_trace/entities/langfuse_trace_entity.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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"
Expand Down
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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"
Expand Down
2 changes: 1 addition & 1 deletion api/core/prompt/simple_prompt_transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
from core.file.models import File


class ModelMode(str, enum.Enum):
class ModelMode(enum.StrEnum):
COMPLETION = "completion"
CHAT = "chat"

Expand Down
4 changes: 2 additions & 2 deletions api/core/rag/datasource/keyword/keyword_type.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from enum import Enum
from enum import StrEnum


class KeyWordType(str, Enum):
class KeyWordType(StrEnum):
JIEBA = "jieba"
4 changes: 2 additions & 2 deletions api/core/rag/datasource/vdb/vector_type.py
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
4 changes: 2 additions & 2 deletions api/core/rag/rerank/rerank_type.py
Original file line number Diff line number Diff line change
@@ -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"
4 changes: 2 additions & 2 deletions api/core/tools/entities/tool_entities.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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"
Expand Down
4 changes: 2 additions & 2 deletions api/core/tools/tool/tool.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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"
Expand Down
4 changes: 2 additions & 2 deletions api/core/variables/types.py
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
4 changes: 2 additions & 2 deletions api/core/workflow/entities/node_entities.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -8,7 +8,7 @@
from models.workflow import WorkflowNodeExecutionStatus


class NodeRunMetadataKey(str, Enum):
class NodeRunMetadataKey(StrEnum):
"""
Node Run Metadata Key.
"""
Expand Down
4 changes: 2 additions & 2 deletions api/core/workflow/enums.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from enum import Enum
from enum import StrEnum


class SystemVariableKey(str, Enum):
class SystemVariableKey(StrEnum):
"""
System Variables.
"""
Expand Down
4 changes: 2 additions & 2 deletions api/core/workflow/nodes/enums.py
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
4 changes: 2 additions & 2 deletions api/core/workflow/nodes/iteration/entities.py
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
4 changes: 2 additions & 2 deletions api/core/workflow/nodes/variable_assigner/node_data.py
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
4 changes: 2 additions & 2 deletions api/extensions/storage/storage_type.py
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
6 changes: 3 additions & 3 deletions api/models/account.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
from .types import StringUUID


class AccountStatus(str, enum.Enum):
class AccountStatus(enum.StrEnum):
PENDING = "pending"
UNINITIALIZED = "uninitialized"
ACTIVE = "active"
Expand Down Expand Up @@ -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"
Expand Down
2 changes: 1 addition & 1 deletion api/models/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
8 changes: 4 additions & 4 deletions api/models/enums.py
Original file line number Diff line number Diff line change
@@ -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"
Loading

0 comments on commit 9bae1df

Please sign in to comment.