feat(agents-api): cozo queries for tasks

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

@@ -0,0 +1,25 @@
+from uuid import UUID
+
+from ..utils import cozo_query
+from typing import Literal, Dict, Any
+
+
+@cozo_query
+def create_execution_query(
+    task_id: UUID,
+    execution_id: UUID,
+    status: Literal[
+        "queued", "starting", "running", "waiting-for-input", "success", "failed"
+    ] = "queued",
+    arguments: Dict[str, Any] = {},
+) -> tuple[str, dict]:
+    query = """"""
+    return (
+        query,
+        {
+            "task_id": str(task_id),
+            "execution_id": str(execution_id),
+            "status": status,
+            "arguments": arguments,
+        },
+    )

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

@@ -0,0 +1,19 @@
+from uuid import UUID
+
+from ..utils import cozo_query
+
+
+@cozo_query
+def get_execution_transition_query(
+    execution_id: UUID, transition_id: UUID, developer_id: UUID
+) -> tuple[str, dict]:
+
+    query = """"""
+    return (
+        query,
+        {
+            "execution_id": str(execution_id),
+            "transition_id": str(transition_id),
+            "developer_id": str(developer_id),
+        },
+    )

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

@@ -0,0 +1,25 @@
+from uuid import UUID
+
+from ..utils import cozo_query
+from typing import Literal, Dict, Any
+
+
+@cozo_query
+def create_execution_query(
+    task_id: UUID,
+    execution_id: UUID,
+    status: Literal[
+        "queued", "starting", "running", "waiting-for-input", "success", "failed"
+    ] = "queued",
+    arguments: Dict[str, Any] = {},
+) -> tuple[str, dict]:
+    query = """"""
+    return (
+        query,
+        {
+            "task_id": str(task_id),
+            "execution_id": str(execution_id),
+            "status": status,
+            "arguments": arguments,
+        },
+    )

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

@@ -0,0 +1,25 @@
+from uuid import UUID
+
+from ..utils import cozo_query
+from typing import Literal, Dict, Any
+
+
+@cozo_query
+def create_execution_query(
+    task_id: UUID,
+    execution_id: UUID,
+    status: Literal[
+        "queued", "starting", "running", "waiting-for-input", "success", "failed"
+    ] = "queued",
+    arguments: Dict[str, Any] = {},
+) -> tuple[str, dict]:
+    query = """"""
+    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,19 @@
+from uuid import UUID
+
+from ..utils import cozo_query
+
+
+@cozo_query
+def get_execution_transition_query(
+    execution_id: UUID, transition_id: UUID, developer_id: UUID
+) -> tuple[str, dict]:
+
+    query = """"""
+    return (
+        query,
+        {
+            "execution_id": str(execution_id),
+            "transition_id": str(transition_id),
+            "developer_id": str(developer_id),
+        },
+    )

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

@@ -4,6 +4,28 @@
 
 
 @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)})
+def get_task_query(agent_id: UUID, task_id: UUID) -> tuple[str, dict]:
+    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)})

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

@@ -7,7 +7,7 @@
 
 @cozo_query
 def list_tasks_query(
-    developer_id: UUID,
+    agent_id: UUID,
     limit: int = 100,
     offset: int = 0,
     # metadata_filter: dict[str, Any] = {},
@@ -21,4 +21,32 @@ def list_tasks_query(
     Returns:
         pd.DataFrame: A DataFrame containing the queried task data.
     """
-    pass
+
+    query = """
+    ?[
+        task_id,
+        name,
+        description,
+        input_schema,
+        tools_available,
+        workflows,
+        created_at,
+        updated_at,
+    ] := *tasks {
+            agent_id: to_uuid($agent_id),
+            task_id,
+            updated_at_ms,
+            name,
+            description,
+            input_schema,
+            tools_available,
+            workflows,
+            created_at,
+            @ 'NOW'
+        }, updated_at = to_int(updated_at_ms) / 1000
+
+      :limit $limit
+      :offset $offset
+    """
+
+    return (query, {"agent_id": str(agent_id), "limit": limit, "offset": offset})

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

@@ -0,0 +1,24 @@
+"""
+This module contains the functionality for creating a new Task in the 'cozodb` database.
+It constructs and executes a datalog query to insert Task data.
+"""
+
+from uuid import UUID
+from typing import List, Optional, Dict, Any
+
+
+from ..utils import cozo_query
+
+
+@cozo_query
+def create_task_query(
+    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

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

@@ -58,6 +58,9 @@ async def create_task(
     x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
 ) -> ResourceCreatedResponse:
     task_id = uuid4()
+
+    # TODO: Do thorough validation of the task spec
+
     resp: pd.DataFrame = create_task_query(
         task_id=task_id,
         developer_id=x_developer_id,
@@ -73,6 +76,9 @@ async def create_task_execution(
     request: CreateExecution,
     x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
 ) -> ResourceCreatedResponse:
+    # TODO: Do thorough validation of the input against task input schema
+    # DO NOT let the user specify the status
+
     resp = create_execution_query()
     return ResourceCreatedResponse(
         id=resp["execution_id"][0], created_at=resp["created_at"][0]
@@ -98,3 +104,7 @@ async def get_execution_transition(
     resp = get_execution_transition_query(execution_id, transition_id, x_developer_id)
 
     pass
+
+# @router.get("/executions/{execution_id}")  # -> get the execution object
+# @router.put("/executions/{execution_id}/transitions/{transition_id}")  # -> update a waiting transition
+# @router.get("/executions/{execution_id}/transitions")  # -> get all transitions for an execution