hotfix: Add custom attributes to the workflow

Signed-off-by: Diwank Singh Tomer <[email protected]>
julep-ai · Oct 16, 2024 · f09267c · f09267c
1 parent fffa9b4
commit f09267c
Showing 1 changed file with 8 additions and 0 deletions.
diff --git 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}]"