feat(agents-api): cozo queries for tasks

agents-api/agents_api/autogen/openapi_model.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  openapi.yaml
-#   timestamp: 2024-05-30T00:31:23+00:00
+#   timestamp: 2024-05-30T20:05:30+00:00
 
 from __future__ import annotations
 
@@ -805,7 +805,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
     """

agents-api/agents_api/models/execution/create_execution.py

@@ -1,19 +1,42 @@
+from typing import Literal, Dict, Any
 from uuid import UUID
 
+from beartype import beartype
+
 from ..utils import cozo_query
-from typing import Literal, Dict, Any
 
 
 @cozo_query
+@beartype
 def create_execution_query(
+    developer_id: UUID,
+    agent_id: UUID,
     task_id: UUID,
     execution_id: UUID,
     status: Literal[
-        "queued", "starting", "running", "waiting-for-input", "success", "failed"
+        "queued", "starting", "running", "waiting_for_input", "success", "failed"
     ] = "queued",
     arguments: Dict[str, Any] = {},
 ) -> tuple[str, dict]:
-    query = """"""
+    # TODO: Check for agent in developer ID; Assert whether dev can access agent and by relation the task
+
+    query = """
+{
+    ?[task_id, execution_id, status, arguments] <- [[
+        to_uuid($task_id),
+        to_uuid($execution_id),
+        $status,
+        $arguments
+    ]]
+
+    :insert executions {
+        task_id,
+        execution_id,
+        status,
+        arguments
+    }
+}
+"""
     return (
         query,
         {

agents-api/agents_api/models/execution/create_execution_transition.py

@@ -0,0 +1,49 @@
+from typing import Literal, Dict, Any
+from uuid import UUID
+
+from beartype import beartype
+
+from ..utils import cozo_query
+
+
+@cozo_query
+@beartype
+def create_execution_transition_query(
+    developer_id: UUID,
+    execution_id: UUID,
+    transition_id: UUID,
+    type_: Literal["finished", "waiting", "error", "step"],
+    from_: tuple[str, int],
+    to: tuple[str, int] | None,
+    output: 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
+
+    query = """
+{
+    ?[execution_id, transition_id, type, from, to, output] <- [[
+        to_uuid($execution_id),
+        to_uuid($transition_id),
+        $type,
+        $from,
+        $to,
+        $output,
+    ]]
+
+    :insert transitions {
+        execution_id, transition_id, type, from, to, output
+    }
+}
+"""
+    return (
+        query,
+        {
+            "execution_id": str(execution_id),
+            "transition_id": str(transition_id),
+            "type": type_,
+            "from": from_,
+            "to": to,
+            "output": output,
+        },
+    )

agents-api/agents_api/models/execution/get_execution.py

@@ -0,0 +1,32 @@
+from uuid import UUID
+
+from ..utils import cozo_query
+
+from beartype import beartype
+
+
+@cozo_query
+@beartype
+def get_execution_query(
+    task_id: UUID,
+    execution_id: UUID,
+) -> tuple[str, dict]:
+    query = """
+{
+    ?[status, arguments, created_at, updated_at] := *executions {
+        task_id: to_uuid($task_id),
+        execution_id: to_uuid($execution_id),
+        status,
+        arguments,
+        created_at,
+        updated_at,
+    }
+}
+"""
+    return (
+        query,
+        {
+            "task_id": str(task_id),
+            "execution_id": str(execution_id),
+        },
+    )

agents-api/agents_api/models/execution/get_execution_status.py

@@ -1,12 +1,21 @@
-from typing import Literal
 from uuid import UUID
+from beartype import beartype
 
 from ..utils import cozo_query
 
 
 @cozo_query
-def get_execution_status_query(task_id: UUID, developer_id: UUID) -> tuple[str, dict]:
+@beartype
+def get_execution_status_query(task_id: UUID, execution_id: UUID) -> tuple[str, dict]:
     task_id = str(task_id)
-    developer_id = str(developer_id)
-    query = """"""
-    return (query, {"task_id": task_id, "developer_id": developer_id})
+    execution_id = str(execution_id)
+    query = """
+{
+    ?[status] := *executions {
+        task_id: to_uuid($task_id),
+        execution_id: to_uuid($execution_id),
+        status,
+    }
+}
+"""
+    return (query, {"task_id": task_id, "execution_id": execution_id})

agents-api/agents_api/models/execution/get_execution_transition.py

@@ -1,19 +1,32 @@
 from uuid import UUID
 
+from beartype import beartype
+
 from ..utils import cozo_query
 
 
 @cozo_query
+@beartype
 def get_execution_transition_query(
-    execution_id: UUID, transition_id: UUID, developer_id: UUID
+    execution_id: UUID, transition_id: UUID
 ) -> tuple[str, dict]:
 
-    query = """"""
+    query = """
+{
+    ?[type, from, to, output] := *transitions {
+        execution_id: to_uuid($execution_id),
+        transition_id: to_uuid($transition_id),
+        type,
+        from,
+        to,
+        output
+    }
+}
+"""
     return (
         query,
         {
             "execution_id": str(execution_id),
             "transition_id": str(transition_id),
-            "developer_id": str(developer_id),
         },
     )

agents-api/agents_api/models/execution/list_execution_transitions.py

@@ -0,0 +1,33 @@
+from uuid import UUID
+
+from beartype import beartype
+
+from ..utils import cozo_query
+
+
+@cozo_query
+@beartype
+def list_execution_transitions_query(execution_id: UUID) -> tuple[str, dict]:
+
+    query = """
+{
+    ?[transition_id, type, from, to, output, updated_at, created_at] := *transitions {
+        execution_id: to_uuid($execution_id),
+        transition_id,
+        type,
+        from,
+        to,
+        output,
+        updated_at,
+        created_at,
+    }
+    :limit 100
+    :offset 0
+}
+"""
+    return (
+        query,
+        {
+            "execution_id": str(execution_id),
+        },
+    )

agents-api/agents_api/models/execution/list_executions.py

@@ -0,0 +1,35 @@
+from uuid import UUID
+
+from beartype import beartype
+
+from ..utils import cozo_query
+
+
+@cozo_query
+@beartype
+def list_task_executions_query(
+    agent_id: UUID, task_id: UUID, developer_id: UUID
+) -> tuple[str, dict]:
+    # TODO: Check for agent in developer ID; Assert whether dev can access agent and by relation the task
+    query = """
+{
+    ?[
+        execution_id,
+        status,
+        arguments,
+        created_at,
+        updated_at,
+    ] := *executions {
+            task_id: to_uuid($task_id),
+            execution_id,
+            status,
+            arguments,
+            created_at,
+            updated_at,
+        }
+        :limit 10
+        :offset 0
+
+}
+"""
+    return (query, {"task_id": str(task_id)})