hotfix for hotfix: Add custom attributes to the workflow

Signed-off-by: Diwank Singh Tomer <[email protected]>

agents-api/agents_api/clients/temporal.py

@@ -2,6 +2,11 @@
 from uuid import UUID
 
 from temporalio.client import Client, TLSConfig
+from temporalio.common import (
+    SearchAttributeKey,
+    SearchAttributePair,
+    TypedSearchAttributes,
+)
 
 from ..autogen.openapi_model import TransitionTarget
 from ..common.protocol.tasks import ExecutionInput
@@ -48,6 +53,7 @@ async def run_task_execution_workflow(
     from ..workflows.task_execution import TaskExecutionWorkflow
 
     client = client or (await get_client())
+    execution_id_key = SearchAttributeKey.for_keyword("CustomStringField")
 
     return await client.start_workflow(
         TaskExecutionWorkflow.run,
@@ -56,7 +62,13 @@ async def run_task_execution_workflow(
         id=str(job_id),
         run_timeout=timedelta(days=31),
         retry_policy=DEFAULT_RETRY_POLICY,
-        # TODO: Should add search_attributes for queryability
+        search_attributes=TypedSearchAttributes(
+            [
+                SearchAttributePair(
+                    execution_id_key, str(execution_input.execution.id)
+                ),
+            ]
+        ),
     )
 
 

agents-api/agents_api/workflows/task_execution/__init__.py

@@ -126,14 +126,6 @@ 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}]"