From d2fdd77aee4f6814bebf6ec7313051230cd067b9 Mon Sep 17 00:00:00 2001
From: Novice <novice12185727@gmail.com>
Date: Tue, 15 Apr 2025 13:53:21 +0800
Subject: [PATCH] chore(agent_node): refactor the code supported old version
 SDK

---
 api/core/workflow/nodes/agent/agent_node.py | 10 ++++++----
 api/core/workflow/nodes/agent/entities.py   | 15 +++++++++++++++
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/api/core/workflow/nodes/agent/agent_node.py b/api/core/workflow/nodes/agent/agent_node.py
index 942e6412d0..7262d6d25b 100644
--- a/api/core/workflow/nodes/agent/agent_node.py
+++ b/api/core/workflow/nodes/agent/agent_node.py
@@ -6,7 +6,7 @@ from core.agent.entities import AgentToolEntity
 from core.agent.plugin_entities import AgentStrategyParameter
 from core.memory.token_buffer_memory import TokenBufferMemory
 from core.model_manager import ModelInstance, ModelManager
-from core.model_runtime.entities.model_entities import AIModelEntity, ModelFeature, ModelType
+from core.model_runtime.entities.model_entities import AIModelEntity, ModelType
 from core.plugin.manager.exc import PluginDaemonClientSideError
 from core.plugin.manager.plugin import PluginInstallationManager
 from core.provider_manager import ProviderManager
@@ -16,7 +16,7 @@ from core.variables.segments import StringSegment
 from core.workflow.entities.node_entities import NodeRunResult
 from core.workflow.entities.variable_pool import VariablePool
 from core.workflow.enums import SystemVariableKey
-from core.workflow.nodes.agent.entities import AgentNodeData, ParamsAutoGenerated
+from core.workflow.nodes.agent.entities import AgentNodeData, AgentOldVersionModelFeatures, ParamsAutoGenerated
 from core.workflow.nodes.base.entities import BaseNodeData
 from core.workflow.nodes.enums import NodeType
 from core.workflow.nodes.event.event import RunCompletedEvent
@@ -253,8 +253,10 @@ class AgentNode(ToolNode):
                     value["history_prompt_messages"] = history_prompt_messages
                     if model_schema:
                         # remove structured output feature to support old version agent plugin
-                        if model_schema.features and ModelFeature.STRUCTURED_OUTPUT in model_schema.features:
-                            model_schema.features.remove(ModelFeature.STRUCTURED_OUTPUT)
+                        if model_schema.features:
+                            for feature in model_schema.features:
+                                if feature.value not in AgentOldVersionModelFeatures:
+                                    model_schema.features.remove(feature)
                         value["entity"] = model_schema.model_dump(mode="json")
                     else:
                         value["entity"] = None
diff --git a/api/core/workflow/nodes/agent/entities.py b/api/core/workflow/nodes/agent/entities.py
index 87cc7e9824..77e94375bf 100644
--- a/api/core/workflow/nodes/agent/entities.py
+++ b/api/core/workflow/nodes/agent/entities.py
@@ -24,3 +24,18 @@ class AgentNodeData(BaseNodeData):
 class ParamsAutoGenerated(Enum):
     CLOSE = 0
     OPEN = 1
+
+
+class AgentOldVersionModelFeatures(Enum):
+    """
+    Enum class for old SDK version llm feature.
+    """
+
+    TOOL_CALL = "tool-call"
+    MULTI_TOOL_CALL = "multi-tool-call"
+    AGENT_THOUGHT = "agent-thought"
+    VISION = "vision"
+    STREAM_TOOL_CALL = "stream-tool-call"
+    DOCUMENT = "document"
+    VIDEO = "video"
+    AUDIO = "audio"