feat(tasks): Update execution transition relation (#397)

* wip Signed-off-by: Diwank Tomer <[email protected]> * wip Signed-off-by: Diwank Tomer <[email protected]> * wip Signed-off-by: Julep Developers <[email protected]> * feat(tasks): Update execution transition relation Signed-off-by: Diwank Tomer <[email protected]> --------- Signed-off-by: Diwank Tomer <[email protected]> Signed-off-by: Julep Developers <[email protected]> Co-authored-by: Diwank Tomer <[email protected]> Co-authored-by: Julep Developers <[email protected]>
julep-ai · Jun 22, 2024 · a6b17e9 · a6b17e9
1 parent c72546e
commit a6b17e9
Show file tree

Hide file tree

Showing 63 changed files with 274 additions and 750 deletions.
diff --git a/agents-api/agents_api/autogen/openapi_model.py b/agents-api/agents_api/autogen/openapi_model.py
@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  openapi.yaml
-#   timestamp: 2024-06-21T03:13:24+00:00
+#   timestamp: 2024-06-21T21:41:04+00:00
 
 from __future__ import annotations
 
@@ -841,7 +841,7 @@ class ImageUrl(BaseModel):
     """
     URL or base64 data url (e.g. `data:image/jpeg;base64,<the base64 encoded image>`)
     """
-    detail: Detail | None = "auto"
+    detail: Detail | None = "auto"  # pytype: disable=annotation-type-mismatch
     """
     image detail to feed into the model can be low | high | auto
     """
@@ -1160,17 +1160,23 @@ class Execution(BaseModel):
 class ExecutionTransition(BaseModel):
     id: UUID
     execution_id: UUID
-    created_at: AwareDatetime
-    outputs: Dict[str, Any]
+    type: Annotated[str, Field(pattern="^(finish|wait|error|step)$")]
     """
-    Outputs from an Execution Transition
+    Execution Status
     """
     from_: Annotated[List[str | int], Field(alias="from")]
     to: List[str | int]
-    type: Annotated[str, Field(pattern="^(finish|wait|error|step)$")]
+    task_token: str | None = None
+    outputs: Dict[str, Any]
     """
-    Execution Status
+    Outputs from an Execution Transition
+    """
+    metadata: Dict[str, Any] | None = None
     """
+    (Optional) metadata
+    """
+    created_at: AwareDatetime
+    updated_at: AwareDatetime | None = None
 
 
 class PromptWorkflowStep(BaseModel):

diff --git a/agents-api/agents_api/common/utils/cozo.py b/agents-api/agents_api/common/utils/cozo.py
@@ -19,6 +19,6 @@
     Client._process_mutate_data(_fake_client, data)
 )
 
-uuid_int_list_to_uuid4 = lambda data: UUID(
-    bytes=b"".join([i.to_bytes(1, "big") for i in data])
-)
+
+def uuid_int_list_to_uuid4(data):
+    return UUID(bytes=b"".join([i.to_bytes(1, "big") for i in data]))
diff --git a/agents-api/agents_api/models/execution/create_execution_transition.py b/agents-api/agents_api/models/execution/create_execution_transition.py
@@ -16,6 +16,8 @@ def create_execution_transition_query(
     from_: tuple[str, int],
     to: tuple[str, int] | None,
     outputs: Dict[str, Any] | None,
+    task_token: str | None = None,
+    metadata: Dict[str, Any] = {},
 ) -> tuple[str, dict]:
     # TODO: Check for agent in developer ID; Assert whether dev can access agent and by relation the task
     # TODO: Check for task and execution
@@ -25,17 +27,19 @@ def create_execution_transition_query(
 
     query = """
 {
-    ?[execution_id, transition_id, type, from, to, output] <- [[
+    ?[execution_id, transition_id, type, from, to, output, task_token, metadata] <- [[
         to_uuid($execution_id),
         to_uuid($transition_id),
         $type,
         $from,
         $to,
         $output,
+        $task_token,
+        $metadata,
     ]]
 
     :insert transitions {
-        execution_id, transition_id, type, from, to, output
+        execution_id, transition_id, type, from, to, output, task_token, metadata,
     }
 }
 """
@@ -48,5 +52,7 @@ def create_execution_transition_query(
             "from": from_,
             "to": to,
             "output": outputs,
+            "task_token": task_token,
+            "metadata": metadata,
         },
     )
diff --git a/agents-api/agents_api/models/execution/get_execution_input.py b/agents-api/agents_api/models/execution/get_execution_input.py
@@ -1,9 +1,9 @@
 from uuid import UUID
 
-from ..utils import cozo_query
-
 from beartype import beartype
 
+from ..utils import cozo_query
+
 
 @cozo_query
 @beartype

diff --git a/agents-api/agents_api/models/execution/get_execution_transition.py b/agents-api/agents_api/models/execution/get_execution_transition.py
@@ -19,7 +19,11 @@ def get_execution_transition_query(
         type,
         from,
         to,
-        output
+        output,
+        metadata,
+        task_token,
+        created_at,
+        updated_at,
     }
 }
 """

diff --git a/agents-api/migrations/migrate_1716939839_task_relations.py b/agents-api/migrations/migrate_1716939839_task_relations.py
@@ -60,6 +60,11 @@ def run(client, queries):
         from: (String, Int),
         to: (String, Int)?,
         output: Json,
+
+        task_token: String? default null,
+
+        # should store: an Activity Id, a Workflow Id, and optionally a Run Id.
+        metadata: Json default {},
         created_at: Float default now(),
         updated_at: Float default now(),
     }

diff --git a/agents-api/poetry.lock b/agents-api/poetry.lock
diff --git a/mock_openapi.yaml b/mock_openapi.yaml
@@ -1374,6 +1374,7 @@ components:
         metadata:
           type: object
           properties: {}
+          additionalProperties: true
           description: (Optional) metadata
       required:
         - id
@@ -1410,6 +1411,7 @@ components:
         metadata:
           type: object
           properties: {}
+          additionalProperties: true
           description: Optional metadata
         instructions:
           oneOf:
@@ -1505,6 +1507,7 @@ components:
         metadata:
           type: object
           properties: {}
+          additionalProperties: true
           description: Optional metadata
         render_templates:
           type: boolean
@@ -1539,6 +1542,7 @@ components:
         metadata:
           type: object
           properties: {}
+          additionalProperties: true
           description: (Optional) metadata
       description: A valid request payload for creating a user
     CreateSessionRequest:
@@ -1559,6 +1563,7 @@ components:
         metadata:
           type: object
           properties: {}
+          additionalProperties: true
           description: Optional metadata
         render_templates:
           type: boolean
@@ -1608,6 +1613,7 @@ components:
         metadata:
           type: object
           properties: {}
+          additionalProperties: true
           description: (Optional) metadata
         instructions:
           oneOf:
@@ -1630,6 +1636,7 @@ components:
         metadata:
           type: object
           properties: {}
+          additionalProperties: true
           description: Optional metadata
         token_budget:
           type: integer
@@ -1659,6 +1666,7 @@ components:
         metadata:
           type: object
           properties: {}
+          additionalProperties: true
           description: Optional metadata
         instructions:
           oneOf:
@@ -1685,6 +1693,7 @@ components:
         metadata:
           type: object
           properties: {}
+          additionalProperties: true
           description: Optional metadata
       description: A valid request payload for updating a user
       required:
@@ -2322,6 +2331,7 @@ components:
         metadata:
           type: object
           properties: {}
+          additionalProperties: true
           description: optional metadata
       required:
         - title
@@ -2347,6 +2357,7 @@ components:
         metadata:
           type: object
           properties: {}
+          additionalProperties: true
           description: Optional metadata
       required:
         - title
@@ -2512,6 +2523,7 @@ components:
         metadata:
           type: object
           properties: {}
+          additionalProperties: true
           description: Optional metadata
       description: A request for patching a user
     PatchAgentRequest:
@@ -2533,6 +2545,7 @@ components:
         metadata:
           type: object
           properties: {}
+          additionalProperties: true
           description: Optional metadata
         instructions:
           oneOf:
@@ -2553,6 +2566,7 @@ components:
         metadata:
           type: object
           properties: {}
+          additionalProperties: true
           description: Optional metadata
         token_budget:
           type: integer
@@ -2793,13 +2807,8 @@ components:
         execution_id:
           type: string
           format: uuid
-        created_at:
-          type: string
-          format: date-time
-        outputs:
-          type: object
-          additionalProperties: true
-          description: Outputs from an Execution Transition
+        type:
+          $ref: '#/components/schemas/TransitionType'
         from:
           type: array
           items:
@@ -2815,8 +2824,25 @@ components:
               description: Canonical name of the workflow
             - type: integer
               description: Step index
-        type:
-          $ref: '#/components/schemas/TransitionType'
+        task_token:
+          type: string
+          default: null
+          nullable: true
+        outputs:
+          type: object
+          additionalProperties: true
+          description: Outputs from an Execution Transition
+        metadata:
+          type: object
+          properties: {}
+          additionalProperties: true
+          description: (Optional) metadata
+        created_at:
+          type: string
+          format: date-time
+        updated_at:
+          type: string
+          format: date-time
       required:
         - id
         - execution_id