feat(agents-api): Add migrations for updated typespec models

Signed-off-by: Diwank Tomer <[email protected]>

agents-api/migrations/migrate_1721576813_extended_tool_relations.py

@@ -0,0 +1,90 @@
+#/usr/bin/env python3
+
+MIGRATION_ID = "extended_tool_relations"
+CREATED_AT = 1721576813.383905
+
+
+drop_agent_functions_hnsw_index = dict(
+    up="""
+    ::hnsw drop agent_functions:embedding_space
+    """,
+    down="""
+    ::hnsw create agent_functions:embedding_space {
+        fields: [embedding],
+        filter: !is_null(embedding),
+        dim: 768,
+        distance: Cosine,
+        m: 64,
+        ef_construction: 256,
+        extend_candidates: false,
+        keep_pruned_connections: false,
+    }
+    """,
+)
+
+create_tools_relation = dict(
+    up="""
+    ?[agent_id, tool_id, tool_type, name, spec, updated_at, created_at] := *agent_functions{
+        agent_id, tool_id, name, description, parameters, updated_at, created_at
+    }, tool_type = "function",
+    spec = {"description": description, "parameters": parameters}
+
+    :create tools {
+        agent_id: Uuid,
+        tool_id: Uuid,
+        =>
+        tool_type: String,
+        name: String,
+        spec: Json,
+
+        updated_at: Float default now(),
+        created_at: Float default now(),
+    }
+    """,
+    down="""
+    ::remove tools
+    """,
+)
+
+drop_agent_functions_table = dict(
+    up="""
+    ::remove agent_functions
+    """,
+    down="""
+    :create agent_functions {
+        agent_id: Uuid,
+        tool_id: Uuid,
+        =>
+        name: String,
+        description: String,
+        parameters: Json,
+        embed_instruction: String default 'Transform this tool description for retrieval: ',
+        embedding: <F32; 768>? default null,
+        updated_at: Float default now(),
+        created_at: Float default now(),
+    }
+    """,
+)
+
+
+queries_to_run = [
+    drop_agent_functions_hnsw_index,
+    create_tools_relation,
+    drop_agent_functions_table,
+]
+
+
+def run(client, *queries):
+    joiner = "}\n\n{"
+
+    query = joiner.join(queries)
+    query = f"{{\n{query}\n}}"
+    client.run(query)
+
+
+def up(client):
+    run(client, *[q["up"] for q in queries_to_run])
+
+
+def down(client):
+    run(client, *[q["down"] for q in reversed(queries_to_run)])

agents-api/migrations/migrate_1721609661_task_tool_ref_by_name.py

@@ -0,0 +1,109 @@
+#/usr/bin/env python3
+
+MIGRATION_ID = "task_tool_ref_by_name"
+CREATED_AT = 1721609661.768934
+
+
+
+# - add metadata
+# - add inherit_tools bool
+# - rename tools_available to tools
+update_tasks_relation = dict(
+    up="""
+    ?[
+        agent_id,
+        task_id,
+        updated_at_ms,
+        name,
+        description,
+        input_schema,
+        tools,
+        inherit_tools,
+        workflows,
+        created_at,
+        metadata,
+    ] := *tasks {
+        agent_id,
+        task_id,
+        updated_at_ms,
+        name,
+        description,
+        input_schema,
+        tools_available: tools,
+        workflows,
+        created_at,
+    }, metadata = {},
+    inherit_tools = true
+
+    :replace 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: [Uuid] default [],
+        inherit_tools: Bool default true,
+        workflows: [Json],
+        created_at: Float default now(),
+        metadata: Json default {},
+    }
+    """,
+    down="""
+    ?[
+        agent_id,
+        task_id,
+        updated_at_ms,
+        name,
+        description,
+        input_schema,
+        tools_available,
+        workflows,
+        created_at,
+    ] := *tasks {
+        agent_id,
+        task_id,
+        updated_at_ms,
+        name,
+        description,
+        input_schema,
+        tools: tools_available,
+        workflows,
+        created_at,
+    }
+
+    :replace 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(),
+    }
+    """,
+)
+
+queries_to_run = [
+    update_tasks_relation,
+]
+
+
+def run(client, *queries):
+    joiner = "}\n\n{"
+
+    query = joiner.join(queries)
+    query = f"{{\n{query}\n}}"
+    client.run(query)
+
+
+def up(client):
+    run(client, *[q["up"] for q in queries_to_run])
+
+
+def down(client):
+    run(client, *[q["down"] for q in reversed(queries_to_run)])

agents-api/migrations/migrate_1721609675_multi_agent_multi_user_session.py

@@ -0,0 +1,79 @@
+#/usr/bin/env python3
+
+MIGRATION_ID = "multi_agent_multi_user_session"
+CREATED_AT = 1721609675.213755
+
+add_multiple_participants_in_session = dict(
+    up="""
+    ?[session_id, participant_id, participant_type] :=
+        *session_lookup {
+            agent_id: participant_id,
+            user_id: null,
+            session_id,
+        }, participant_type = 'agent'
+
+    ?[session_id, participant_id, participant_type] :=
+        *session_lookup {
+            agent_id,
+            user_id: participant_id,
+            session_id,
+        }, participant_type = 'user',
+        participant_id != null
+
+    :replace session_lookup  {
+        session_id: Uuid,
+        participant_type: String,
+        participant_id: Uuid,
+    }
+    """,
+    down="""
+    users[user_id, session_id] :=
+        *session_lookup {
+            session_id,
+            participant_type: "user",
+            participant_id: user_id,
+        }
+
+    agents[agent_id, session_id] :=
+        *session_lookup {
+            session_id,
+            participant_type: "agent",
+            participant_id: agent_id,
+        }
+
+    ?[agent_id, user_id, session_id] :=
+        agents[agent_id, session_id],
+        users[user_id, session_id]
+
+    ?[agent_id, user_id, session_id] :=
+        agents[agent_id, session_id],
+        not users[_, session_id],
+        user_id = null
+
+    :replace session_lookup  {
+        agent_id: Uuid,
+        user_id: Uuid? default null,
+        session_id: Uuid,
+    }
+    """,
+)
+
+queries_to_run = [
+    add_multiple_participants_in_session,
+]
+
+
+def run(client, *queries):
+    joiner = "}\n\n{"
+
+    query = joiner.join(queries)
+    query = f"{{\n{query}\n}}"
+    client.run(query)
+
+
+def up(client):
+    run(client, *[q["up"] for q in queries_to_run])
+
+
+def down(client):
+    run(client, *[q["down"] for q in reversed(queries_to_run)])