From 7cd1c321f208b6354b18fa2aeb0a40386d39a523 Mon Sep 17 00:00:00 2001
From: Diwank Tomer <diwank@julep.ai>
Date: Wed, 29 May 2024 19:31:06 -0700
Subject: [PATCH] wip(agents-api): cozo queries for tasks

Signed-off-by: Diwank Tomer <diwank@julep.ai>
---
 .../agents_api/models/execution/__init__.py   |  0
 agents-api/agents_api/models/task/__init__.py |  0
 agents-api/agents_api/models/task/get_task.py | 28 ++++++++++++++--
 .../agents_api/models/task/list_tasks.py      | 32 +++++++++++++++++--
 4 files changed, 55 insertions(+), 5 deletions(-)
 create mode 100644 agents-api/agents_api/models/execution/__init__.py
 create mode 100644 agents-api/agents_api/models/task/__init__.py

diff --git a/agents-api/agents_api/models/execution/__init__.py b/agents-api/agents_api/models/execution/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/agents-api/agents_api/models/task/__init__.py b/agents-api/agents_api/models/task/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/agents-api/agents_api/models/task/get_task.py b/agents-api/agents_api/models/task/get_task.py
index be9509353..f89da19ad 100644
--- a/agents-api/agents_api/models/task/get_task.py
+++ b/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)})
diff --git a/agents-api/agents_api/models/task/list_tasks.py b/agents-api/agents_api/models/task/list_tasks.py
index 33fbf3ae0..64f3d8010 100644
--- a/agents-api/agents_api/models/task/list_tasks.py
+++ b/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})