From f09267cc07a55f86a4b4c62ef4a77241298e7e69 Mon Sep 17 00:00:00 2001
From: Diwank Singh Tomer <diwank.singh@gmail.com>
Date: Wed, 16 Oct 2024 15:03:57 -0400
Subject: [PATCH] hotfix: Add custom attributes to the workflow

Signed-off-by: Diwank Singh Tomer <diwank.singh@gmail.com>
---
 .../agents_api/workflows/task_execution/__init__.py       | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/agents-api/agents_api/workflows/task_execution/__init__.py b/agents-api/agents_api/workflows/task_execution/__init__.py
index 155b49397..90a4e0bad 100644
--- a/agents-api/agents_api/workflows/task_execution/__init__.py
+++ b/agents-api/agents_api/workflows/task_execution/__init__.py
@@ -126,6 +126,14 @@ async def run(
         start: TransitionTarget = TransitionTarget(workflow="main", step=0),
         previous_inputs: list[Any] = [],
     ) -> Any:
+        # Add metadata to the workflow run
+        workflow.upsert_search_attributes(
+            {
+                "task_id": execution_input.task.id,
+                "execution_id": execution_input.execution.id,
+            }
+        )
+
         workflow.logger.info(
             f"TaskExecutionWorkflow for task {execution_input.task.id}"
             f" [LOC {start.workflow}.{start.step}]"