Explorar o código

fix: replace Enum with StrEnum

Yeuoly hai 4 meses
pai
achega
a61da6cf95

+ 2 - 2
api/core/agent/entities.py

@@ -1,4 +1,4 @@
-from enum import Enum
+from enum import StrEnum
 from typing import Any, Optional, Union
 
 from pydantic import BaseModel
@@ -69,7 +69,7 @@ class AgentEntity(BaseModel):
     Agent Entity.
     """
 
-    class Strategy(Enum):
+    class Strategy(StrEnum):
         """
         Agent Strategy.
         """

+ 7 - 6
api/core/entities/parameter_entities.py

@@ -1,7 +1,7 @@
-from enum import Enum
+from enum import StrEnum
 
 
-class CommonParameterType(Enum):
+class CommonParameterType(StrEnum):
     SECRET_INPUT = "secret-input"
     TEXT_INPUT = "text-input"
     SELECT = "select"
@@ -12,19 +12,20 @@ class CommonParameterType(Enum):
     SYSTEM_FILES = "system-files"
     BOOLEAN = "boolean"
     APP_SELECTOR = "app-selector"
-    # TOOL_SELECTOR = "tool-selector"
     MODEL_SELECTOR = "model-selector"
     TOOLS_SELECTOR = "array[tools]"
 
+    # TOOL_SELECTOR = "tool-selector"
+
 
-class AppSelectorScope(Enum):
+class AppSelectorScope(StrEnum):
     ALL = "all"
     CHAT = "chat"
     WORKFLOW = "workflow"
     COMPLETION = "completion"
 
 
-class ModelSelectorScope(Enum):
+class ModelSelectorScope(StrEnum):
     LLM = "llm"
     TEXT_EMBEDDING = "text-embedding"
     RERANK = "rerank"
@@ -34,7 +35,7 @@ class ModelSelectorScope(Enum):
     VISION = "vision"
 
 
-class ToolSelectorScope(Enum):
+class ToolSelectorScope(StrEnum):
     ALL = "all"
     CUSTOM = "custom"
     BUILTIN = "builtin"

+ 3 - 3
api/core/plugin/entities/parameters.py

@@ -59,7 +59,7 @@ class PluginParameter(BaseModel):
         return v
 
 
-def as_normal_type(typ: enum.Enum):
+def as_normal_type(typ: enum.StrEnum):
     if typ.value in {
         PluginParameterType.SECRET_INPUT,
         PluginParameterType.SELECT,
@@ -68,7 +68,7 @@ def as_normal_type(typ: enum.Enum):
     return typ.value
 
 
-def cast_parameter_value(typ: enum.Enum, value: Any, /):
+def cast_parameter_value(typ: enum.StrEnum, value: Any, /):
     try:
         match typ.value:
             case PluginParameterType.STRING | PluginParameterType.SECRET_INPUT | PluginParameterType.SELECT:
@@ -127,7 +127,7 @@ def cast_parameter_value(typ: enum.Enum, value: Any, /):
         raise ValueError(f"The tool parameter value {value} is not in correct type of {as_normal_type(typ)}.")
 
 
-def init_frontend_parameter(rule: PluginParameter, type: enum.Enum, value: Any):
+def init_frontend_parameter(rule: PluginParameter, type: enum.StrEnum, value: Any):
     """
     init frontend parameter by rule
     """

+ 3 - 4
api/core/plugin/entities/plugin_daemon.py

@@ -1,6 +1,5 @@
-import enum
 from datetime import datetime
-from enum import Enum
+from enum import StrEnum
 from typing import Generic, Optional, TypeVar
 
 from pydantic import BaseModel, ConfigDict, Field
@@ -31,7 +30,7 @@ class InstallPluginMessage(BaseModel):
     Message for installing a plugin.
     """
 
-    class Event(Enum):
+    class Event(StrEnum):
         Info = "info"
         Done = "done"
         Error = "error"
@@ -127,7 +126,7 @@ class PluginDaemonInnerError(Exception):
         self.message = message
 
 
-class PluginInstallTaskStatus(enum.StrEnum):
+class PluginInstallTaskStatus(StrEnum):
     Pending = "pending"
     Running = "running"
     Success = "success"