diff --git a/agents-api/agents_api/clients/temporal.py b/agents-api/agents_api/clients/temporal.py index 4bb25cbc9..deb4809f1 100644 --- a/agents-api/agents_api/clients/temporal.py +++ b/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) + ), + ] + ), ) diff --git a/agents-api/agents_api/workflows/task_execution/__init__.py b/agents-api/agents_api/workflows/task_execution/__init__.py index 90a4e0bad..155b49397 100644 --- a/agents-api/agents_api/workflows/task_execution/__init__.py +++ b/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}]"