Merge pull request #1078 from julep-ai/f/workflow-refactoring

Task execution workflow refactoring

Author: whiterabbit1983

agents-api/agents_api/activities/excecute_api_call.py renamed to agents-api/agents_api/activities/execute_api_call.py

@@ -120,7 +120,7 @@ async def run_task_execution_workflow(
 
     return await client.start_workflow(
         TaskExecutionWorkflow.run,
-        args=[execution_input, start, previous_inputs],
+        args=[execution_input, start, previous_inputs[-1]],
         task_queue=temporal_task_queue,
         id=str(job_id),
         run_timeout=timedelta(days=31),

agents-api/agents_api/clients/temporal.py

@@ -0,0 +1,103 @@
+from typing import Any
+from uuid import UUID
+
+from beartype import beartype
+from pydantic import BaseModel
+
+from ...autogen.openapi_model import (
+    Agent,
+    CreateTaskRequest,
+    CreateToolRequest,
+    Execution,
+    PartialTaskSpecDef,
+    PatchTaskRequest,
+    Session,
+    Task,
+    TaskSpec,
+    TaskSpecDef,
+    TaskToolDef,
+    Tool,
+    UpdateTaskRequest,
+    User,
+    Workflow,
+)
+
+
+class ExecutionInput(BaseModel):
+    developer_id: UUID
+    execution: Execution | None = None
+    task: TaskSpecDef | None = None
+    agent: Agent
+    agent_tools: list[Tool | CreateToolRequest]
+    arguments: dict[str, Any]
+
+    # Not used at the moment
+    user: User | None = None
+    session: Session | None = None
+
+
+@beartype
+def task_to_spec(
+    task: Task | CreateTaskRequest | UpdateTaskRequest | PatchTaskRequest, **model_opts
+) -> TaskSpecDef | PartialTaskSpecDef:
+    task_data = task.model_dump(
+        **model_opts, exclude={"version", "developer_id", "task_id", "id", "agent_id"}
+    )
+
+    if "tools" in task_data:
+        del task_data["tools"]
+
+    tools = []
+    for tool in task.tools:
+        tool_spec = getattr(tool, tool.type)
+
+        tool_obj = dict(
+            type=tool.type,
+            spec=tool_spec.model_dump(),
+            **tool.model_dump(exclude={"type"}),
+        )
+        tools.append(TaskToolDef(**tool_obj))
+
+    workflows = [Workflow(name="main", steps=task_data.pop("main"))]
+
+    for key, steps in list(task_data.items()):
+        if key not in TaskSpec.model_fields:
+            workflows.append(Workflow(name=key, steps=steps))
+            del task_data[key]
+
+    cls = PartialTaskSpecDef if isinstance(task, PatchTaskRequest) else TaskSpecDef
+
+    return cls(
+        workflows=workflows,
+        tools=tools,
+        **task_data,
+    )
+
+
+def spec_to_task_data(spec: dict) -> dict:
+    task_id = spec.pop("task_id", None)
+
+    workflows = spec.pop("workflows")
+    workflows_dict = {workflow["name"]: workflow["steps"] for workflow in workflows}
+
+    tools = spec.pop("tools", []) or []
+    tools = [{tool["type"]: tool.pop("spec"), **tool} for tool in tools if tool]
+
+    return {
+        "id": task_id,
+        "tools": tools,
+        **spec,
+        **workflows_dict,
+    }
+
+
+def spec_to_task(**spec) -> Task | CreateTaskRequest:
+    if not spec.get("id"):
+        spec["id"] = spec.pop("task_id", None)
+
+    if not spec.get("updated_at"):
+        [updated_at_ms, _] = spec.pop("updated_at_ms", None)
+        spec["updated_at"] = updated_at_ms and (updated_at_ms / 1000)
+
+    cls = Task if spec["id"] else CreateTaskRequest
+    return cls(**spec_to_task_data(spec))

agents-api/agents_api/common/protocol/models.py

@@ -1,7 +1,5 @@
 from typing import Annotated, Any, Literal
-from uuid import UUID
 
-from beartype import beartype
 from temporalio import workflow
 from temporalio.exceptions import ApplicationError
 
@@ -10,29 +8,20 @@
     from pydantic_partial import create_partial_model
 
     from ...autogen.openapi_model import (
-        Agent,
-        CreateTaskRequest,
         CreateToolRequest,
         CreateTransitionRequest,
-        Execution,
         ExecutionStatus,
-        PartialTaskSpecDef,
-        PatchTaskRequest,
-        Session,
-        Task,
-        TaskSpec,
-        TaskSpecDef,
-        TaskToolDef,
         Tool,
         ToolRef,
         TransitionTarget,
         TransitionType,
-        UpdateTaskRequest,
-        User,
         Workflow,
         WorkflowStep,
     )
 
+from ...queries.executions import list_execution_transitions
+from .models import ExecutionInput
+
 # TODO: Maybe we should use a library for this
 
 # State Machine
@@ -145,23 +134,10 @@ class PartialTransition(create_partial_model(CreateTransitionRequest)):
     user_state: dict[str, Any] = Field(default_factory=dict)
 
 
-class ExecutionInput(BaseModel):
-    developer_id: UUID
-    execution: Execution | None = None
-    task: TaskSpecDef | None = None
-    agent: Agent
-    agent_tools: list[Tool | CreateToolRequest]
-    arguments: dict[str, Any]
-
-    # Not used at the moment
-    user: User | None = None
-    session: Session | None = None
-
-
 class StepContext(BaseModel):
     execution_input: ExecutionInput
-    inputs: list[Any]
     cursor: TransitionTarget
+    current_input: Any
 
     @computed_field
     @property
@@ -197,16 +173,6 @@ def tools(self) -> list[Tool | CreateToolRequest]:
 
         return filtered_tools + task_tools
 
-    @computed_field
-    @property
-    def outputs(self) -> list[dict[str, Any]]:  # included in dump
-        return self.inputs[1:]
-
-    @computed_field
-    @property
-    def current_input(self) -> dict[str, Any]:  # included in dump
-        return self.inputs[-1]
-
     @computed_field
     @property
     def current_workflow(self) -> Annotated[Workflow, Field(exclude=True)]:
@@ -239,13 +205,22 @@ def model_dump(self, *args, **kwargs) -> dict[str, Any]:
 
         return dump | execution_input
 
+    async def get_inputs(self) -> list[Any]:
+        transitions = await list_execution_transitions(
+            execution_id=self.execution_input.execution.id,
+            limit=1000,
+            direction="asc",
+        )
+        return [t.output for t in transitions]
+
     async def prepare_for_step(self, *args, **kwargs) -> dict[str, Any]:
         current_input = self.current_input
-        inputs = self.inputs
+        inputs = await self.get_inputs()
 
         # Merge execution inputs into the dump dict
         dump = self.model_dump(*args, **kwargs)
         dump["inputs"] = inputs
+        dump["outputs"] = inputs[1:]
         prepared = dump | {"_": current_input}
 
         for i, input in enumerate(inputs):
@@ -260,70 +235,3 @@ class StepOutcome(BaseModel):
     error: str | None = None
     output: Any = None
     transition_to: tuple[TransitionType, TransitionTarget] | None = None
-
-
-@beartype
-def task_to_spec(
-    task: Task | CreateTaskRequest | UpdateTaskRequest | PatchTaskRequest, **model_opts
-) -> TaskSpecDef | PartialTaskSpecDef:
-    task_data = task.model_dump(
-        **model_opts, exclude={"version", "developer_id", "task_id", "id", "agent_id"}
-    )
-
-    if "tools" in task_data:
-        del task_data["tools"]
-
-    tools = []
-    for tool in task.tools:
-        tool_spec = getattr(tool, tool.type)
-
-        tool_obj = dict(
-            type=tool.type,
-            spec=tool_spec.model_dump(),
-            **tool.model_dump(exclude={"type"}),
-        )
-        tools.append(TaskToolDef(**tool_obj))
-
-    workflows = [Workflow(name="main", steps=task_data.pop("main"))]
-
-    for key, steps in list(task_data.items()):
-        if key not in TaskSpec.model_fields:
-            workflows.append(Workflow(name=key, steps=steps))
-            del task_data[key]
-
-    cls = PartialTaskSpecDef if isinstance(task, PatchTaskRequest) else TaskSpecDef
-
-    return cls(
-        workflows=workflows,
-        tools=tools,
-        **task_data,
-    )
-
-
-def spec_to_task_data(spec: dict) -> dict:
-    task_id = spec.pop("task_id", None)
-
-    workflows = spec.pop("workflows")
-    workflows_dict = {workflow["name"]: workflow["steps"] for workflow in workflows}
-
-    tools = spec.pop("tools", []) or []
-    tools = [{tool["type"]: tool.pop("spec"), **tool} for tool in tools if tool]
-
-    return {
-        "id": task_id,
-        "tools": tools,
-        **spec,
-        **workflows_dict,
-    }
-
-
-def spec_to_task(**spec) -> Task | CreateTaskRequest:
-    if not spec.get("id"):
-        spec["id"] = spec.pop("task_id", None)
-
-    if not spec.get("updated_at"):
-        [updated_at_ms, _] = spec.pop("updated_at_ms", None)
-        spec["updated_at"] = updated_at_ms and (updated_at_ms / 1000)
-
-    cls = Task if spec["id"] else CreateTaskRequest
-    return cls(**spec_to_task_data(spec))

agents-api/agents_api/common/protocol/tasks.py

@@ -2,7 +2,7 @@
 
 from beartype import beartype
 
-from ...common.protocol.tasks import ExecutionInput
+from ...common.protocol.models import ExecutionInput
 from ...common.utils.db_exceptions import common_db_exceptions
 from ..utils import pg_query, rewrap_exceptions, wrap_in_class
 

agents-api/agents_api/queries/executions/prepare_execution_input.py

@@ -5,7 +5,7 @@
 from uuid_extensions import uuid7
 
 from ...autogen.openapi_model import CreateOrUpdateTaskRequest, ResourceUpdatedResponse
-from ...common.protocol.tasks import task_to_spec
+from ...common.protocol.models import task_to_spec
 from ...common.utils.db_exceptions import common_db_exceptions
 from ...metrics.counters import increase_counter
 from ..utils import generate_canonical_name, pg_query, rewrap_exceptions, wrap_in_class

agents-api/agents_api/queries/tasks/create_or_update_task.py

@@ -5,7 +5,7 @@
 from uuid_extensions import uuid7
 
 from ...autogen.openapi_model import CreateTaskRequest, ResourceCreatedResponse
-from ...common.protocol.tasks import task_to_spec
+from ...common.protocol.models import task_to_spec
 from ...common.utils.db_exceptions import common_db_exceptions
 from ...metrics.counters import increase_counter
 from ..utils import (

agents-api/agents_api/queries/tasks/create_task.py

@@ -3,7 +3,7 @@
 
 from beartype import beartype
 
-from ...common.protocol.tasks import spec_to_task
+from ...common.protocol.models import spec_to_task
 from ...common.utils.db_exceptions import common_db_exceptions
 from ..utils import pg_query, rewrap_exceptions, wrap_in_class
 

agents-api/agents_api/queries/tasks/get_task.py

@@ -4,7 +4,7 @@
 from beartype import beartype
 from fastapi import HTTPException
 
-from ...common.protocol.tasks import spec_to_task
+from ...common.protocol.models import spec_to_task
 from ...common.utils.db_exceptions import common_db_exceptions
 from ..utils import pg_query, rewrap_exceptions, wrap_in_class
 

agents-api/agents_api/queries/tasks/list_tasks.py

@@ -4,7 +4,7 @@
 from beartype import beartype
 
 from ...autogen.openapi_model import PatchTaskRequest, ResourceUpdatedResponse
-from ...common.protocol.tasks import task_to_spec
+from ...common.protocol.models import task_to_spec
 from ...common.utils.datetime import utcnow
 from ...common.utils.db_exceptions import common_db_exceptions
 from ...metrics.counters import increase_counter

agents-api/agents_api/queries/tasks/patch_task.py

@@ -4,7 +4,7 @@
 from beartype import beartype
 
 from ...autogen.openapi_model import ResourceUpdatedResponse, UpdateTaskRequest
-from ...common.protocol.tasks import task_to_spec
+from ...common.protocol.models import task_to_spec
 from ...common.utils.datetime import utcnow
 from ...common.utils.db_exceptions import common_db_exceptions
 from ...metrics.counters import increase_counter

agents-api/agents_api/queries/tasks/update_task.py

@@ -19,7 +19,7 @@
 )
 from ...clients.temporal import run_task_execution_workflow
 from ...common.protocol.developers import Developer
-from ...common.protocol.tasks import task_to_spec
+from ...common.protocol.models import task_to_spec
 from ...dependencies.developer_id import get_developer_id
 from ...env import max_free_executions
 from ...queries.developers.get_developer import get_developer

agents-api/agents_api/routers/tasks/create_task_execution.py

@@ -21,7 +21,7 @@ def create_worker(client: Client) -> Any:
 
     from ..activities import task_steps
     from ..activities.demo import demo_activity
-    from ..activities.excecute_api_call import execute_api_call
+    from ..activities.execute_api_call import execute_api_call
     from ..activities.execute_integration import execute_integration
     from ..activities.execute_system import execute_system
     from ..activities.sync_items_remote import load_inputs_remote, save_inputs_remote