wip(agents-api): cozo queries for tasks

Signed-off-by: Diwank Tomer <[email protected]>
julep-ai · May 30, 2024 · 7cd1c32 · 7cd1c32
1 parent 170c301
commit 7cd1c32
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 5 deletions.
diff --git a/agents-api/agents_api/models/execution/__init__.py b/agents-api/agents_api/models/execution/__init__.py
diff --git a/agents-api/agents_api/models/task/__init__.py b/agents-api/agents_api/models/task/__init__.py
diff --git 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
@@ -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})