From 93c36991282c21d8510a5d96125777882e175b9d Mon Sep 17 00:00:00 2001
From: Yeuoly <admin@srmxy.cn>
Date: Sun, 15 Dec 2024 18:12:29 +0800
Subject: [PATCH] feat: add label to agent log

---
 api/core/app/apps/workflow_app_runner.py            | 1 +
 api/core/app/entities/queue_entities.py             | 1 +
 api/core/app/entities/task_entities.py              | 1 +
 api/core/app/task_pipeline/workflow_cycle_manage.py | 2 ++
 api/core/tools/entities/tool_entities.py            | 1 +
 api/core/workflow/graph_engine/entities/event.py    | 1 +
 api/core/workflow/nodes/tool/tool_node.py           | 1 +
 7 files changed, 8 insertions(+)

diff --git a/api/core/app/apps/workflow_app_runner.py b/api/core/app/apps/workflow_app_runner.py
index 0eafdbc198..45a5e05dca 100644
--- a/api/core/app/apps/workflow_app_runner.py
+++ b/api/core/app/apps/workflow_app_runner.py
@@ -301,6 +301,7 @@ class WorkflowBasedAppRunner(AppRunner):
             self._publish_event(
                 QueueAgentLogEvent(
                     id=event.id,
+                    label=event.label,
                     node_execution_id=event.node_execution_id,
                     parent_id=event.parent_id,
                     error=event.error,
diff --git a/api/core/app/entities/queue_entities.py b/api/core/app/entities/queue_entities.py
index 406c1b02d6..6ed68574a1 100644
--- a/api/core/app/entities/queue_entities.py
+++ b/api/core/app/entities/queue_entities.py
@@ -308,6 +308,7 @@ class QueueAgentLogEvent(AppQueueEvent):
 
     event: QueueEvent = QueueEvent.AGENT_LOG
     id: str
+    label: str
     node_execution_id: str
     parent_id: str | None
     error: str | None
diff --git a/api/core/app/entities/task_entities.py b/api/core/app/entities/task_entities.py
index ceef3c9353..9e506f6dda 100644
--- a/api/core/app/entities/task_entities.py
+++ b/api/core/app/entities/task_entities.py
@@ -640,6 +640,7 @@ class AgentLogStreamResponse(StreamResponse):
 
         node_execution_id: str
         id: str
+        label: str
         parent_id: str | None
         error: str | None
         status: str
diff --git a/api/core/app/task_pipeline/workflow_cycle_manage.py b/api/core/app/task_pipeline/workflow_cycle_manage.py
index 3e3a5166aa..279ea17af2 100644
--- a/api/core/app/task_pipeline/workflow_cycle_manage.py
+++ b/api/core/app/task_pipeline/workflow_cycle_manage.py
@@ -796,6 +796,7 @@ class WorkflowCycleManage:
                     error=event.error,
                     status=event.status,
                     data=event.data,
+                    label=event.label,
                 )
             )
 
@@ -807,6 +808,7 @@ class WorkflowCycleManage:
                 node_execution_id=node_execution_id,
                 id=event.id,
                 parent_id=event.parent_id,
+                label=event.label,
                 error=event.error,
                 status=event.status,
                 data=event.data,
diff --git a/api/core/tools/entities/tool_entities.py b/api/core/tools/entities/tool_entities.py
index d53c70c4c9..050a43bad0 100644
--- a/api/core/tools/entities/tool_entities.py
+++ b/api/core/tools/entities/tool_entities.py
@@ -160,6 +160,7 @@ class ToolInvokeMessage(BaseModel):
             SUCCESS = "success"
 
         id: str
+        label: str = Field(..., description="The label of the log")
         parent_id: Optional[str] = Field(default=None, description="Leave empty for root log")
         error: Optional[str] = Field(default=None, description="The error message")
         status: LogStatus = Field(..., description="The status of the log")
diff --git a/api/core/workflow/graph_engine/entities/event.py b/api/core/workflow/graph_engine/entities/event.py
index ff4ac6878d..2211b024df 100644
--- a/api/core/workflow/graph_engine/entities/event.py
+++ b/api/core/workflow/graph_engine/entities/event.py
@@ -181,6 +181,7 @@ class BaseAgentEvent(GraphEngineEvent):
 
 class AgentLogEvent(BaseAgentEvent):
     id: str = Field(..., description="id")
+    label: str = Field(..., description="label")
     node_execution_id: str = Field(..., description="node execution id")
     parent_id: str | None = Field(..., description="parent id")
     error: str | None = Field(..., description="error")
diff --git a/api/core/workflow/nodes/tool/tool_node.py b/api/core/workflow/nodes/tool/tool_node.py
index 5c256cbac8..388b25d57b 100644
--- a/api/core/workflow/nodes/tool/tool_node.py
+++ b/api/core/workflow/nodes/tool/tool_node.py
@@ -287,6 +287,7 @@ class ToolNode(BaseNode[ToolNodeData]):
                     error=message.message.error,
                     status=message.message.status.value,
                     data=message.message.data,
+                    label=message.message.label,
                 )
 
         yield RunCompletedEvent(