feat(agents-api): Allow execution_id filtering for child workflows

julep-ai · Feb 15, 2025 · 1dec2d1 · 1dec2d1
1 parent 54a4e75
commit 1dec2d1
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/agents-api/agents_api/workflows/task_execution/__init__.py b/agents-api/agents_api/workflows/task_execution/__init__.py
@@ -402,7 +402,7 @@ async def _handle_YieldStep(
         )
 
         result = await continue_as_child(
-            self.context,
+            self.context.execution_input,
             start=yield_next_target,
             current_input=output,
         )

diff --git a/agents-api/agents_api/workflows/task_execution/helpers.py b/agents-api/agents_api/workflows/task_execution/helpers.py
@@ -4,6 +4,7 @@
 
 from temporalio import workflow
 from temporalio.exceptions import ActivityError, ApplicationError, ChildWorkflowError
+from temporalio.common import SearchAttributeKey, SearchAttributePair, TypedSearchAttributes
 
 from ...common.retry_policies import DEFAULT_RETRY_POLICY
 
@@ -78,6 +79,10 @@ async def continue_as_child(
             info.workflow_type, *args, **kwargs
         )
 
+
+    execution_id = execution_input.execution.id
+    execution_id_key = SearchAttributeKey.for_keyword("CustomStringField")
+
     try:
         return await run(
             args=[
@@ -87,6 +92,9 @@ async def continue_as_child(
             ],
             retry_policy=DEFAULT_RETRY_POLICY,
             memo=workflow.memo() | user_state,
+            search_attributes=TypedSearchAttributes([
+            SearchAttributePair(execution_id_key, str(execution_id)),
+            ]),
         )
     except Exception as e:
         while isinstance(e, ChildWorkflowError) and getattr(e, "__cause__", None):