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

@@ -6,14 +6,34 @@
 
 @cozo_query
 def create_execution_query(
+    agent_id: UUID,
     task_id: UUID,
     execution_id: UUID,
+    developer_id: UUID,
     status: Literal[
         "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/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

@@ -5,15 +5,25 @@
 
 @cozo_query
 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,30 @@
+from uuid import UUID
+
+from ..utils import cozo_query
+
+
+@cozo_query
+def list_execution_transition_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)})

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

@@ -0,0 +1,42 @@
+from uuid import UUID
+
+from ..utils import cozo_query
+from typing import Literal, Dict, Any
+
+
+@cozo_query
+def update_execution_status_query(
+    task_id: UUID,
+    execution_id: UUID,
+    status: Literal[
+        "queued", "starting", "running", "waiting-for-input", "success", "failed"
+    ],
+    arguments: Dict[str, Any] = {},
+) -> tuple[str, dict]:
+    query = """
+{
+    ?[execution_id, task_id, status, updated_at] <- [[
+        to_uuid($execution_id),
+        to_uuid($task_id),
+        $status,
+        now()
+    ]]
+
+    :update executions {
+        execution_id,
+        task_id,
+        status,
+        updated_at
+    }
+}
+
+"""
+    return (
+        query,
+        {
+            "task_id": str(task_id),
+            "execution_id": str(execution_id),
+            "status": status,
+            "arguments": arguments,
+        },
+    )

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

@@ -0,0 +1,47 @@
+from uuid import UUID
+
+from agents_api.common.utils.datetime import utcnow
+
+from ..utils import cozo_query
+from ...common.utils.cozo import cozo_process_mutate_data
+
+
+@cozo_query
+def update_execution_transition_query(
+    execution_id: UUID, transition_id: UUID, **update_data
+) -> tuple[str, dict]:
+
+    transition_update_cols, transition_update_vals = cozo_process_mutate_data(
+        {
+            **{k: v for k, v, in update_data.items() if v is not None},
+            "execution_id": str(execution_id),
+            "transition_id": str(transition_id),
+            "updated_at": utcnow().timestamp(),
+        }
+    )
+    query = f"""
+    {{
+        input[{transition_update_cols}] <- $transition_update_vals
+
+        ?[{transition_update_cols}] := input[{transition_update_cols}],
+        *transitions {{
+            execution_id: to_uuid($execution_id),
+            transition_id: to_uuid($transition_id),
+        }}
+
+        :update transitions {{
+            {transition_update_cols}
+        }}
+        :returning
+    }}
+
+"""
+
+    return (
+        query,
+        {
+            "transitioon_update_vals": transition_update_vals,
+            "execution_id": str(execution_id),
+            "transition_id": str(transition_id),
+        },
+    )

agents-api/agents_api/models/task/create_task.py

@@ -4,21 +4,52 @@
 """
 
 from uuid import UUID
-from typing import List, Optional, Dict, Any
-
+from typing import List, Dict, Any
+from beartype import beartype
 
 from ..utils import cozo_query
 
 
 @cozo_query
+@beartype
 def create_task_query(
+    agent_id: UUID,
     task_id: UUID,
     developer_id: UUID,
-    agent_id: UUID,
     name: str,
     description: str,
     input_schema: Dict[str, any],
     tools_available: List[UUID] = [],
     workflows: List[Dict[str, Any]] = [],
 ) -> tuple[str, dict]:
-    pass
+    # TODO: Check for agent in developer ID; Assert whether dev can access agent and by relation the task
+    query = """
+{
+    ?[agent_id, task_id, name, description, input_schema, tools_available, workflows] <- [[
+        to_uuid($agent_id), to_uuid($task_id), $name, $description, $input_schema, [to_uuid($tools_available)], $workflows
+        ]]
+
+    :insert tasks {
+        agent_id,
+        task_id,
+        name,
+        description,
+        input_schema,
+        tools_available,
+        workflows
+    }
+}
+"""
+
+    return (
+        query,
+        {
+            "agent_id": str(agent_id),
+            "task_id": str(task_id),
+            "name": name,
+            "description": description,
+            "input_schema": input_schema,
+            "tools_available": tools_available,
+            "workflows": workflows,
+        },
+    )

agents-api/agents_api/models/task/get_task.py

@@ -1,9 +1,37 @@
 from uuid import UUID
 
+from beartype import beartype
+
 from ..utils import cozo_query
 
 
 @cozo_query
-def get_task_query(developer_id: UUID, task_id: UUID) -> tuple[str, dict]:
-    query = """"""
-    return (query, {"developer_id": str(developer_id), "task_id": str(task_id)})
+@beartype
+def get_task_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 = """
+    ?[
+        name,
+        description,
+        input_schema,
+        tools_available,
+        workflows,
+        created_at,
+        updated_at,
+    ] := *tasks {
+            agent_id: to_uuid($agent_id),
+            task_id: to_uuid($task_id),
+            updated_at_ms,
+            name,
+            description,
+            input_schema,
+            tools_available,
+            workflows,
+            created_at,
+            @ 'NOW'
+        }, updated_at = to_int(updated_at_ms) / 1000
+    """
+
+    return (query, {"agent_id": str(agent_id), "task_id": str(task_id)})