From c3bc32ce5dcbe7674193f8908015d18a90d954da Mon Sep 17 00:00:00 2001 From: Diwank Tomer Date: Sun, 21 Jul 2024 21:16:10 -0400 Subject: [PATCH] feat(agents-api): Add migrations for updated typespec models Signed-off-by: Diwank Tomer --- ...rate_1721576813_extended_tool_relations.py | 90 +++++++++++++++ ...igrate_1721609661_task_tool_ref_by_name.py | 109 ++++++++++++++++++ ...21609675_multi_agent_multi_user_session.py | 79 +++++++++++++ 3 files changed, 278 insertions(+) create mode 100644 agents-api/migrations/migrate_1721576813_extended_tool_relations.py create mode 100644 agents-api/migrations/migrate_1721609661_task_tool_ref_by_name.py create mode 100644 agents-api/migrations/migrate_1721609675_multi_agent_multi_user_session.py diff --git a/agents-api/migrations/migrate_1721576813_extended_tool_relations.py b/agents-api/migrations/migrate_1721576813_extended_tool_relations.py new file mode 100644 index 000000000..1798b598c --- /dev/null +++ b/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: ? 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)]) diff --git a/agents-api/migrations/migrate_1721609661_task_tool_ref_by_name.py b/agents-api/migrations/migrate_1721609661_task_tool_ref_by_name.py new file mode 100644 index 000000000..466ad511c --- /dev/null +++ b/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)]) diff --git a/agents-api/migrations/migrate_1721609675_multi_agent_multi_user_session.py b/agents-api/migrations/migrate_1721609675_multi_agent_multi_user_session.py new file mode 100644 index 000000000..e6a02b45c --- /dev/null +++ b/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)])