From 8e83b1695f67889c9f6b07c2bd12352e3c681ca0 Mon Sep 17 00:00:00 2001
From: Diwank Tomer <diwank@julep.ai>
Date: Tue, 28 May 2024 23:22:04 -0700
Subject: [PATCH] feat(agents-api): Add cozo migrations for tasks schema

Signed-off-by: Diwank Tomer <diwank@julep.ai>
---
 .../migrate_1716939839_task_relations.py      | 80 +++++++++++++++++++
 1 file changed, 80 insertions(+)
 create mode 100644 agents-api/migrations/migrate_1716939839_task_relations.py

diff --git a/agents-api/migrations/migrate_1716939839_task_relations.py b/agents-api/migrations/migrate_1716939839_task_relations.py
new file mode 100644
index 000000000..6ea3fab6e
--- /dev/null
+++ b/agents-api/migrations/migrate_1716939839_task_relations.py
@@ -0,0 +1,80 @@
+# /usr/bin/env python3
+
+MIGRATION_ID = "task_relations"
+CREATED_AT = 1716939839.690704
+
+
+def run(client, queries):
+    joiner = "}\n\n{"
+
+    query = joiner.join(queries)
+    query = f"{{\n{query}\n}}"
+    client.run(query)
+
+
+create_task_relation_query = dict(
+    up="""
+    :create tasks {
+        agent_id: Uuid,
+        task_id: Uuid,
+        updated_at_ms: Validity default [floor(now() * 1000), true],
+        =>
+        name: String,
+        description: String? default null,
+        input_schema: Json,
+        tools_available: [Uuid] default [],
+        workflows: [Json],
+        created_at: Float default now(),
+    }
+    """,
+    down="::remove tasks",
+)
+
+create_execution_relation_query = dict(
+    up="""
+    :create executions {
+        task_id: Uuid,
+        execution_id: Uuid,
+        =>
+        status: String default 'queued',
+        # one of: "queued", "starting", "running", "waiting-for-input", "success", "failed"
+
+        arguments: Json,
+        created_at: Float default now(),
+        updated_at: Float default now(),
+    }
+    """,
+    down="::remove executions",
+)
+
+create_transition_relation_query = dict(
+    up="""
+    :create transitions {
+        execution_id: Uuid,
+        transition_id: Uuid,
+        =>
+        type: String,
+        # one of: "finished", "waiting", "error", "step"
+
+        from: (String, Int),
+        to: (String, Int)?,
+        output: Json,
+        created_at: Float default now(),
+    }
+    """,
+    down="::remove transitions",
+)
+
+queries = [
+    create_task_relation_query,
+    create_execution_relation_query,
+    create_transition_relation_query,
+]
+
+
+def up(client):
+    run(client, [q["up"] for q in queries])
+
+
+def down(client):
+    run(client, [q["down"] for q in reversed(queries)])