From d76af08784e3a494c4eb15cd5c917f0523de3213 Mon Sep 17 00:00:00 2001
From: Novice Lee <novicelee@NoviPro.local>
Date: Thu, 9 Jan 2025 16:50:05 +0800
Subject: [PATCH] feat: add agent log icon

---
 api/core/app/apps/workflow_app_runner.py    |  1 +
 api/core/workflow/nodes/agent/agent_node.py | 14 ++++++++-
 api/core/workflow/nodes/tool/tool_node.py   | 32 +++++++++++++++++++++
 3 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/api/core/app/apps/workflow_app_runner.py b/api/core/app/apps/workflow_app_runner.py
index 8dde1acfa4..2d0bc95524 100644
--- a/api/core/app/apps/workflow_app_runner.py
+++ b/api/core/app/apps/workflow_app_runner.py
@@ -385,6 +385,7 @@ class WorkflowBasedAppRunner(AppRunner):
                     error=event.error,
                     status=event.status,
                     data=event.data,
+                    metadata=event.metadata,
                 )
             )
         elif isinstance(event, ParallelBranchRunStartedEvent):
diff --git a/api/core/workflow/nodes/agent/agent_node.py b/api/core/workflow/nodes/agent/agent_node.py
index dabaacb31e..8c29ce5176 100644
--- a/api/core/workflow/nodes/agent/agent_node.py
+++ b/api/core/workflow/nodes/agent/agent_node.py
@@ -7,6 +7,7 @@ from core.agent.plugin_entities import AgentStrategyParameter
 from core.model_manager import ModelManager
 from core.model_runtime.entities.model_entities import ModelType
 from core.plugin.manager.exc import PluginDaemonClientSideError
+from core.plugin.manager.plugin import PluginInstallationManager
 from core.tools.entities.tool_entities import ToolProviderType
 from core.tools.tool_manager import ToolManager
 from core.workflow.entities.node_entities import NodeRunResult
@@ -89,9 +90,20 @@ class AgentNode(ToolNode):
 
         try:
             # convert tool messages
+            manager = PluginInstallationManager()
+            plugins = manager.list_plugins(self.tenant_id)
+            current_plugin = next(
+                plugin
+                for plugin in plugins
+                if f"{plugin.plugin_id}/{plugin.name}"
+                == cast(AgentNodeData, self.node_data).agent_strategy_provider_name
+            )
             yield from self._transform_message(
                 message_stream,
-                {"provider": (cast(AgentNodeData, self.node_data)).agent_strategy_provider_name},
+                {
+                    "icon": current_plugin.declaration.icon,
+                    "agent_strategy": cast(AgentNodeData, self.node_data).agent_strategy_name,
+                },
                 parameters_for_log,
             )
         except PluginDaemonClientSideError as e:
diff --git a/api/core/workflow/nodes/tool/tool_node.py b/api/core/workflow/nodes/tool/tool_node.py
index 4536b8ffb0..8518db5d80 100644
--- a/api/core/workflow/nodes/tool/tool_node.py
+++ b/api/core/workflow/nodes/tool/tool_node.py
@@ -7,6 +7,7 @@ from sqlalchemy.orm import Session
 from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler
 from core.file import File, FileTransferMethod
 from core.plugin.manager.exc import PluginDaemonClientSideError
+from core.plugin.manager.plugin import PluginInstallationManager
 from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter
 from core.tools.tool_engine import ToolEngine
 from core.tools.utils.message_transformer import ToolFileMessageTransformer
@@ -24,6 +25,7 @@ from extensions.ext_database import db
 from factories import file_factory
 from models import ToolFile
 from models.workflow import WorkflowNodeExecutionStatus
+from services.tools.builtin_tools_manage_service import BuiltinToolManageService
 
 from .entities import ToolNodeData
 from .exc import (
@@ -297,6 +299,36 @@ class ToolNode(BaseNode[ToolNodeData]):
                 files.append(message.meta["file"])
             elif message.type == ToolInvokeMessage.MessageType.LOG:
                 assert isinstance(message.message, ToolInvokeMessage.LogMessage)
+                if message.message.metadata:
+                    icon = tool_info.get("icon", "")
+                    dict_metadata = dict(message.message.metadata)
+                    if dict_metadata.get("provider"):
+                        manager = PluginInstallationManager()
+                        plugins = manager.list_plugins(self.tenant_id)
+                        try:
+                            current_plugin = next(
+                                plugin
+                                for plugin in plugins
+                                if f"{plugin.plugin_id}/{plugin.name}" == dict_metadata["provider"]
+                            )
+                            icon = current_plugin.declaration.icon
+                        except StopIteration:
+                            pass
+                        try:
+                            builtin_tool = next(
+                                provider
+                                for provider in BuiltinToolManageService.list_builtin_tools(
+                                    self.user_id,
+                                    self.tenant_id,
+                                )
+                                if provider.name == dict_metadata["provider"]
+                            )
+                            icon = builtin_tool.icon
+                        except StopIteration:
+                            pass
+
+                        dict_metadata["icon"] = icon
+                        message.message.metadata = dict_metadata
                 agent_log = AgentLogEvent(
                     id=message.message.id,
                     node_execution_id=self.id,