From 3181d8d0a0f7adc78e28cb56c57cdad089e6d5aa Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Mon, 9 Dec 2024 18:11:45 +0300 Subject: [PATCH 01/14] feat: Add indices migration --- .../migrate_1733755642_transition_indices.py | 68 +++++++++++++++++++ 1 file changed, 68 insertions(+) create mode 100644 agents-api/migrations/migrate_1733755642_transition_indices.py diff --git a/agents-api/migrations/migrate_1733755642_transition_indices.py b/agents-api/migrations/migrate_1733755642_transition_indices.py new file mode 100644 index 000000000..797d91644 --- /dev/null +++ b/agents-api/migrations/migrate_1733755642_transition_indices.py @@ -0,0 +1,68 @@ +#/usr/bin/env python3 + +MIGRATION_ID = "transition_indices" +CREATED_AT = 1733755642.881131 + + +create_transition_indices = dict( + up=""" + ::index create executions:execution_id_status_idx { status, execution_id } + ::index create executions:execution_id_input_idx { input, execution_id } + ::index create executions:execution_id_output_idx { output, execution_id } + ::index create executions:execution_id_error_idx { error, execution_id } + ::index create executions:execution_id_session_id_idx { session_id, execution_id } + ::index create executions:execution_id_metadata_idx { metadata, execution_id } + + ::index create agents:developer_id_name_idx { name, developer_id } + ::index create agents:developer_id_about_idx { about, developer_id } + ::index create agents:developer_id_instructions_idx { instructions, developer_id } + ::index create agents:developer_id_model_idx { model, developer_id } + ::index create agents:developer_id_metadata_idx { metadata, developer_id } + + ::index create developers:developer_id_email_idx { email, developer_id } + ::index create developers:developer_id_active_idx { active, developer_id } + ::index create developers:developer_id_tags_idx { tags, developer_id } + ::index create developers:developer_id_settings_idx { settings, developer_id } + + ::index create transitions:transition_id_type_idx { type, transition_id } + ::index create transitions:transition_id_current_idx { current, transition_id } + ::index create transitions:transition_id_next_idx { next, transition_id } + ::index create transitions:transition_id_output_idx { output, transition_id } + ::index create transitions:transition_id_task_token_idx { task_token, transition_id } + ::index create transitions:transition_id_metadata_idx { metadata, transition_id } + """, + down=""" + ::index drop executions:execution_id_status_id + ::index drop executions:execution_id_input_idx + ::index drop executions:execution_id_output_idx + ::index drop executions:execution_id_error_idx + ::index drop executions:execution_id_session_id_idx + ::index drop executions:execution_id_metadata_idx + + ::index drop agents:developer_id_name_idx + ::index drop agents:developer_id_about_idx + ::index drop agents:developer_id_instructions_idx + ::index drop agents:developer_id_model_idx + ::index drop agents:developer_id_metadata_idx + + ::index drop developers:developer_id_email_idx + ::index drop developers:developer_id_active_idx + ::index drop developers:developer_id_tags_idx + ::index drop developers:developer_id_settings_idx + + ::index drop transitions:transition_id_type_idx + ::index drop transitions:transition_id_current_idx + ::index drop transitions:transition_id_next_idx + ::index drop transitions:transition_id_output_idx + ::index drop transitions:transition_id_task_token_idx + ::index drop transitions:transition_id_metadata_idx + """, +) + + +def up(client): + client.run(create_transition_indices["up"]) + + +def down(client): + client.run(create_transition_indices["down"]) From 6d78d276323892d726c730131ef5ca64ca4ba396 Mon Sep 17 00:00:00 2001 From: whiterabbit1983 Date: Mon, 9 Dec 2024 16:18:10 +0000 Subject: [PATCH 02/14] refactor: Lint agents-api (CI) --- agents-api/migrations/migrate_1733755642_transition_indices.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agents-api/migrations/migrate_1733755642_transition_indices.py b/agents-api/migrations/migrate_1733755642_transition_indices.py index 797d91644..56c0dca55 100644 --- a/agents-api/migrations/migrate_1733755642_transition_indices.py +++ b/agents-api/migrations/migrate_1733755642_transition_indices.py @@ -1,4 +1,4 @@ -#/usr/bin/env python3 +# /usr/bin/env python3 MIGRATION_ID = "transition_indices" CREATED_AT = 1733755642.881131 From 3bd06684a09edc5435ea8597c8efca15b0a6eb6e Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Tue, 10 Dec 2024 09:58:43 +0300 Subject: [PATCH 03/14] feat: Reduce the number of indices and apply them to queries --- .../execution/create_execution_transition.py | 6 ++- .../models/execution/update_execution.py | 2 +- .../migrate_1733755642_transition_indices.py | 50 ++----------------- 3 files changed, 9 insertions(+), 49 deletions(-) diff --git a/agents-api/agents_api/models/execution/create_execution_transition.py b/agents-api/agents_api/models/execution/create_execution_transition.py index d885379fb..98a497632 100644 --- a/agents-api/agents_api/models/execution/create_execution_transition.py +++ b/agents-api/agents_api/models/execution/create_execution_transition.py @@ -140,7 +140,7 @@ def create_execution_transition( ] last_transition_type[min_cost(type_created_at)] := - *transitions {{ + *transitions:execution_id_type_idx {{ execution_id: to_uuid("{str(execution_id)}"), type, created_at, @@ -214,6 +214,10 @@ def create_execution_transition( insert_query, ] + with open("out.txt", "w") as f: + print(queries, file=f) + raise Exception + return ( queries, { diff --git a/agents-api/agents_api/models/execution/update_execution.py b/agents-api/agents_api/models/execution/update_execution.py index 35deab259..8135e8cae 100644 --- a/agents-api/agents_api/models/execution/update_execution.py +++ b/agents-api/agents_api/models/execution/update_execution.py @@ -79,7 +79,7 @@ def update_execution( validate_status_query = """ valid_status[count(status)] := - *executions { + *executions:execution_id_status_idx { status, execution_id: to_uuid($execution_id), }, diff --git a/agents-api/migrations/migrate_1733755642_transition_indices.py b/agents-api/migrations/migrate_1733755642_transition_indices.py index 56c0dca55..9ab69aff5 100644 --- a/agents-api/migrations/migrate_1733755642_transition_indices.py +++ b/agents-api/migrations/migrate_1733755642_transition_indices.py @@ -7,55 +7,11 @@ create_transition_indices = dict( up=""" ::index create executions:execution_id_status_idx { status, execution_id } - ::index create executions:execution_id_input_idx { input, execution_id } - ::index create executions:execution_id_output_idx { output, execution_id } - ::index create executions:execution_id_error_idx { error, execution_id } - ::index create executions:execution_id_session_id_idx { session_id, execution_id } - ::index create executions:execution_id_metadata_idx { metadata, execution_id } - - ::index create agents:developer_id_name_idx { name, developer_id } - ::index create agents:developer_id_about_idx { about, developer_id } - ::index create agents:developer_id_instructions_idx { instructions, developer_id } - ::index create agents:developer_id_model_idx { model, developer_id } - ::index create agents:developer_id_metadata_idx { metadata, developer_id } - - ::index create developers:developer_id_email_idx { email, developer_id } - ::index create developers:developer_id_active_idx { active, developer_id } - ::index create developers:developer_id_tags_idx { tags, developer_id } - ::index create developers:developer_id_settings_idx { settings, developer_id } - - ::index create transitions:transition_id_type_idx { type, transition_id } - ::index create transitions:transition_id_current_idx { current, transition_id } - ::index create transitions:transition_id_next_idx { next, transition_id } - ::index create transitions:transition_id_output_idx { output, transition_id } - ::index create transitions:transition_id_task_token_idx { task_token, transition_id } - ::index create transitions:transition_id_metadata_idx { metadata, transition_id } + ::index create transitions:execution_id_type_idx { type, transition_id } """, down=""" - ::index drop executions:execution_id_status_id - ::index drop executions:execution_id_input_idx - ::index drop executions:execution_id_output_idx - ::index drop executions:execution_id_error_idx - ::index drop executions:execution_id_session_id_idx - ::index drop executions:execution_id_metadata_idx - - ::index drop agents:developer_id_name_idx - ::index drop agents:developer_id_about_idx - ::index drop agents:developer_id_instructions_idx - ::index drop agents:developer_id_model_idx - ::index drop agents:developer_id_metadata_idx - - ::index drop developers:developer_id_email_idx - ::index drop developers:developer_id_active_idx - ::index drop developers:developer_id_tags_idx - ::index drop developers:developer_id_settings_idx - - ::index drop transitions:transition_id_type_idx - ::index drop transitions:transition_id_current_idx - ::index drop transitions:transition_id_next_idx - ::index drop transitions:transition_id_output_idx - ::index drop transitions:transition_id_task_token_idx - ::index drop transitions:transition_id_metadata_idx + ::index drop executions:execution_id_status_idx + ::index drop transitions:execution_id_type_idx """, ) From 74595f10cba5631cd6f3f6f00100f474a2aa7a37 Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Tue, 10 Dec 2024 10:05:31 +0300 Subject: [PATCH 04/14] feat: Add metadata index --- agents-api/migrations/migrate_1733755642_transition_indices.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/agents-api/migrations/migrate_1733755642_transition_indices.py b/agents-api/migrations/migrate_1733755642_transition_indices.py index 9ab69aff5..8d15ad5d0 100644 --- a/agents-api/migrations/migrate_1733755642_transition_indices.py +++ b/agents-api/migrations/migrate_1733755642_transition_indices.py @@ -8,10 +8,12 @@ up=""" ::index create executions:execution_id_status_idx { status, execution_id } ::index create transitions:execution_id_type_idx { type, transition_id } + ::index create executions:execution_id_metadata_idx { metadata, execution_id } """, down=""" ::index drop executions:execution_id_status_idx ::index drop transitions:execution_id_type_idx + ::index drop executions:execution_id_metadata_idx """, ) From d8df03879798fd7689f241876e87039c793420ac Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Tue, 10 Dec 2024 10:11:26 +0300 Subject: [PATCH 05/14] fix: Remove debugging --- .../models/execution/create_execution_transition.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/agents-api/agents_api/models/execution/create_execution_transition.py b/agents-api/agents_api/models/execution/create_execution_transition.py index 98a497632..86a13312a 100644 --- a/agents-api/agents_api/models/execution/create_execution_transition.py +++ b/agents-api/agents_api/models/execution/create_execution_transition.py @@ -214,10 +214,6 @@ def create_execution_transition( insert_query, ] - with open("out.txt", "w") as f: - print(queries, file=f) - raise Exception - return ( queries, { From e8f1d87297414c54e39fcd2cfb1780c7cd415742 Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Tue, 10 Dec 2024 10:11:37 +0300 Subject: [PATCH 06/14] fix: Fix index field --- agents-api/migrations/migrate_1733755642_transition_indices.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agents-api/migrations/migrate_1733755642_transition_indices.py b/agents-api/migrations/migrate_1733755642_transition_indices.py index 8d15ad5d0..dfa004cdc 100644 --- a/agents-api/migrations/migrate_1733755642_transition_indices.py +++ b/agents-api/migrations/migrate_1733755642_transition_indices.py @@ -7,7 +7,7 @@ create_transition_indices = dict( up=""" ::index create executions:execution_id_status_idx { status, execution_id } - ::index create transitions:execution_id_type_idx { type, transition_id } + ::index create transitions:execution_id_type_idx { type, execution_id } ::index create executions:execution_id_metadata_idx { metadata, execution_id } """, down=""" From 098b44828bd166ebebcc59d3dfcd689dfc6ac52d Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Tue, 10 Dec 2024 11:39:03 +0300 Subject: [PATCH 07/14] feat: Add more indices --- .../migrate_1733755642_transition_indices.py | 22 ++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/agents-api/migrations/migrate_1733755642_transition_indices.py b/agents-api/migrations/migrate_1733755642_transition_indices.py index dfa004cdc..085c62f69 100644 --- a/agents-api/migrations/migrate_1733755642_transition_indices.py +++ b/agents-api/migrations/migrate_1733755642_transition_indices.py @@ -6,14 +6,26 @@ create_transition_indices = dict( up=""" - ::index create executions:execution_id_status_idx { status, execution_id } - ::index create transitions:execution_id_type_idx { type, execution_id } - ::index create executions:execution_id_metadata_idx { metadata, execution_id } + ::index create executions:execution_id_status_idx { execution_id, status } + ::index create executions:execution_id_task_id_idx { execution_id, task_id } + ::index create tasks:task_id_agent_id_idx { task_id, agent_id } + ::index create agents:agent_id_developer_id_idx { agent_id, developer_id } + ::index create sessions:session_id_developer_id_idx { session_id, developer_id } + ::index create docs:owner_id_metadata_doc_id_idx { owner_id, metadata, doc_id } + ::index create agents:developer_id_metadata_agent_id_idx { developer_id, metadata, agent_id } + ::index create users:developer_id_metadata_user_id_idx { developer_id, metadata, user_id } + ::index create transitions:execution_id_type_created_at_idx { execution_id, type, created_at } """, down=""" ::index drop executions:execution_id_status_idx - ::index drop transitions:execution_id_type_idx - ::index drop executions:execution_id_metadata_idx + ::index drop executions:execution_id_task_id_idx + ::index drop tasks:task_id_agent_id_idx + ::index drop agents:agent_id_developer_id_idx + ::index drop sessions:session_id_developer_id_idx + ::index drop docs:owner_id_metadata_doc_id_idx + ::index drop agents:developer_id_metadata_agent_id_idx + ::index drop users:developer_id_metadata_user_id_idx + ::index drop transitions:execution_id_type_created_at_idx """, ) From 2406d7af0cc1811865374eb0dad6771a37b7e3ee Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Tue, 10 Dec 2024 11:39:24 +0300 Subject: [PATCH 08/14] chore: Refactor create execution transition queries --- .../execution/create_execution_transition.py | 279 +++++------------- 1 file changed, 73 insertions(+), 206 deletions(-) diff --git a/agents-api/agents_api/models/execution/create_execution_transition.py b/agents-api/agents_api/models/execution/create_execution_transition.py index 86a13312a..cd2e2a8e1 100644 --- a/agents-api/agents_api/models/execution/create_execution_transition.py +++ b/agents-api/agents_api/models/execution/create_execution_transition.py @@ -25,62 +25,8 @@ from .update_execution import update_execution -def validate_transition_targets(data: CreateTransitionRequest) -> None: - # Make sure the current/next targets are valid - match data.type: - case "finish_branch": - pass # TODO: Implement - case "finish" | "error" | "cancelled": - pass - - ### FIXME: HACK: Fix this and uncomment - - ### assert ( - ### data.next is None - ### ), "Next target must be None for finish/finish_branch/error/cancelled" - - case "init_branch" | "init": - assert ( - data.next and data.current.step == data.next.step == 0 - ), "Next target must be same as current for init_branch/init and step 0" - - case "wait": - assert data.next is None, "Next target must be None for wait" - - case "resume" | "step": - assert data.next is not None, "Next target must be provided for resume/step" - - if data.next.workflow == data.current.workflow: - assert ( - data.next.step > data.current.step - ), "Next step must be greater than current" - - case _: - raise ValueError(f"Invalid transition type: {data.type}") - - -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - Transition, - transform=lambda d: { - **d, - "id": d["transition_id"], - "current": {"workflow": d["current"][0], "step": d["current"][1]}, - "next": d["next"] and {"workflow": d["next"][0], "step": d["next"][1]}, - }, - one=True, - _kind="inserted", -) -@cozo_query -@increase_counter("create_execution_transition") @beartype -def create_execution_transition( +def _create_execution_transition( *, developer_id: UUID, execution_id: UUID, @@ -140,7 +86,7 @@ def create_execution_transition( ] last_transition_type[min_cost(type_created_at)] := - *transitions:execution_id_type_idx {{ + *transitions:execution_id_type_created_at_idx {{ execution_id: to_uuid("{str(execution_id)}"), type, created_at, @@ -225,167 +171,88 @@ def create_execution_transition( ) -@rewrap_exceptions( - { - QueryException: partialclass(HTTPException, status_code=400), - ValidationError: partialclass(HTTPException, status_code=400), - TypeError: partialclass(HTTPException, status_code=400), - } -) -@wrap_in_class( - Transition, - transform=lambda d: { - **d, - "id": d["transition_id"], - "current": {"workflow": d["current"][0], "step": d["current"][1]}, - "next": d["next"] and {"workflow": d["next"][0], "step": d["next"][1]}, - }, - one=True, - _kind="inserted", -) -@cozo_query_async -@increase_counter("create_execution_transition_async") -@beartype -async def create_execution_transition_async( - *, - developer_id: UUID, - execution_id: UUID, - data: CreateTransitionRequest, - # Only one of these needed - transition_id: UUID | None = None, - task_token: str | None = None, - # Only required for updating the execution status as well - update_execution_status: bool = False, - task_id: UUID | None = None, -) -> tuple[list[str | None], dict]: - transition_id = transition_id or uuid4() - data.metadata = data.metadata or {} - data.execution_id = execution_id - - # Dump to json - if isinstance(data.output, list): - data.output = [ - item.model_dump(mode="json") if hasattr(item, "model_dump") else item - for item in data.output - ] - - elif hasattr(data.output, "model_dump"): - data.output = data.output.model_dump(mode="json") - - # TODO: This is a hack to make sure the transition is valid - # (parallel transitions are whack, we should do something better) - is_parallel = data.current.workflow.startswith("PAR:") - - # Prepare the transition data - transition_data = data.model_dump(exclude_unset=True, exclude={"id"}) - - # Parse the current and next targets - validate_transition_targets(data) - current_target = transition_data.pop("current") - next_target = transition_data.pop("next") - - transition_data["current"] = (current_target["workflow"], current_target["step"]) - transition_data["next"] = next_target and ( - next_target["workflow"], - next_target["step"], - ) - - columns, transition_values = cozo_process_mutate_data( - { - **transition_data, - "task_token": str(task_token), # Converting to str for JSON serialisation - "transition_id": str(transition_id), - "execution_id": str(execution_id), - } - ) - - # Make sure the transition is valid - check_last_transition_query = f""" - valid_transition[start, end] <- [ - {", ".join(f'["{start}", "{end}"]' for start, ends in valid_transitions.items() for end in ends)} - ] +def validate_transition_targets(data: CreateTransitionRequest) -> None: + # Make sure the current/next targets are valid + match data.type: + case "finish_branch": + pass # TODO: Implement + case "finish" | "error" | "cancelled": + pass - last_transition_type[min_cost(type_created_at)] := - *transitions {{ - execution_id: to_uuid("{str(execution_id)}"), - type, - created_at, - }}, - type_created_at = [type, -created_at] + ### FIXME: HACK: Fix this and uncomment - matched[collect(last_type)] := - last_transition_type[data], - last_type_data = first(data), - last_type = if(is_null(last_type_data), "init", last_type_data), - valid_transition[last_type, $next_type] + ### assert ( + ### data.next is None + ### ), "Next target must be None for finish/finish_branch/error/cancelled" - ?[valid] := - matched[prev_transitions], - found = length(prev_transitions), - valid = if($next_type == "init", found == 0, found > 0), - assert(valid, "Invalid transition"), + case "init_branch" | "init": + assert ( + data.next and data.current.step == data.next.step == 0 + ), "Next target must be same as current for init_branch/init and step 0" - :limit 1 - """ + case "wait": + assert data.next is None, "Next target must be None for wait" - # Prepare the insert query - insert_query = f""" - ?[{columns}] <- $transition_values + case "resume" | "step": + assert data.next is not None, "Next target must be provided for resume/step" - :insert transitions {{ - {columns} - }} - - :returning - """ + if data.next.workflow == data.current.workflow: + assert ( + data.next.step > data.current.step + ), "Next step must be greater than current" - validate_status_query, update_execution_query, update_execution_params = ( - "", - "", - {}, - ) + case _: + raise ValueError(f"Invalid transition type: {data.type}") - if update_execution_status: - assert ( - task_id is not None - ), "task_id is required for updating the execution status" - # Prepare the execution update query - [*_, validate_status_query, update_execution_query], update_execution_params = ( - update_execution.__wrapped__( - developer_id=developer_id, - task_id=task_id, - execution_id=execution_id, - data=UpdateExecutionRequest( - status=transition_to_execution_status[data.type] - ), - output=data.output if data.type != "error" else None, - error=str(data.output) - if data.type == "error" and data.output - else None, +create_execution_transition = rewrap_exceptions( + { + QueryException: partialclass(HTTPException, status_code=400), + ValidationError: partialclass(HTTPException, status_code=400), + TypeError: partialclass(HTTPException, status_code=400), + } +)( + wrap_in_class( + Transition, + transform=lambda d: { + **d, + "id": d["transition_id"], + "current": {"workflow": d["current"][0], "step": d["current"][1]}, + "next": d["next"] and {"workflow": d["next"][0], "step": d["next"][1]}, + }, + one=True, + _kind="inserted", + )( + cozo_query( + increase_counter("create_execution_transition")( + _create_execution_transition ) ) + ) +) - queries = [ - verify_developer_id_query(developer_id), - verify_developer_owns_resource_query( - developer_id, - "executions", - execution_id=execution_id, - parents=[("agents", "agent_id"), ("tasks", "task_id")], - ), - validate_status_query if not is_parallel else None, - update_execution_query if not is_parallel else None, - check_last_transition_query if not is_parallel else None, - insert_query, - ] - - return ( - queries, - { - "transition_values": transition_values, - "next_type": data.type, - "valid_transitions": valid_transitions, - **update_execution_params, +create_execution_transition_async = rewrap_exceptions( + { + QueryException: partialclass(HTTPException, status_code=400), + ValidationError: partialclass(HTTPException, status_code=400), + TypeError: partialclass(HTTPException, status_code=400), + } +)( + wrap_in_class( + Transition, + transform=lambda d: { + **d, + "id": d["transition_id"], + "current": {"workflow": d["current"][0], "step": d["current"][1]}, + "next": d["next"] and {"workflow": d["next"][0], "step": d["next"][1]}, }, + one=True, + _kind="inserted", + )( + cozo_query_async( + increase_counter("create_execution_transition")( + _create_execution_transition + ) + ) ) +) From 7d1d09eec9273aa196e41960896b3007a0719776 Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Tue, 10 Dec 2024 13:42:12 +0300 Subject: [PATCH 09/14] chore: Add one more index --- agents-api/migrations/migrate_1733755642_transition_indices.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/agents-api/migrations/migrate_1733755642_transition_indices.py b/agents-api/migrations/migrate_1733755642_transition_indices.py index 085c62f69..eedebc565 100644 --- a/agents-api/migrations/migrate_1733755642_transition_indices.py +++ b/agents-api/migrations/migrate_1733755642_transition_indices.py @@ -8,6 +8,7 @@ up=""" ::index create executions:execution_id_status_idx { execution_id, status } ::index create executions:execution_id_task_id_idx { execution_id, task_id } + ::index create executions:task_id_execution_id_idx { task_id, execution_id } ::index create tasks:task_id_agent_id_idx { task_id, agent_id } ::index create agents:agent_id_developer_id_idx { agent_id, developer_id } ::index create sessions:session_id_developer_id_idx { session_id, developer_id } @@ -19,6 +20,7 @@ down=""" ::index drop executions:execution_id_status_idx ::index drop executions:execution_id_task_id_idx + ::index drop executions:task_id_execution_id_idx ::index drop tasks:task_id_agent_id_idx ::index drop agents:agent_id_developer_id_idx ::index drop sessions:session_id_developer_id_idx From f10f1ba773af3e7022711c19d99816992d775977 Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Tue, 10 Dec 2024 13:42:50 +0300 Subject: [PATCH 10/14] feat: Make queries use indices --- agents-api/agents_api/models/agent/get_agent.py | 2 +- agents-api/agents_api/models/agent/list_agents.py | 2 +- agents-api/agents_api/models/docs/list_docs.py | 2 +- .../agents_api/models/docs/search_docs_by_embedding.py | 6 +++--- agents-api/agents_api/models/docs/search_docs_by_text.py | 4 ++-- agents-api/agents_api/models/execution/count_executions.py | 2 +- agents-api/agents_api/models/execution/get_execution.py | 2 +- .../models/execution/get_paused_execution_token.py | 4 ++-- .../models/execution/list_execution_transitions.py | 2 +- agents-api/agents_api/models/task/get_task.py | 2 +- agents-api/agents_api/models/user/get_user.py | 2 +- agents-api/agents_api/models/user/list_users.py | 2 +- agents-api/agents_api/models/user/patch_user.py | 2 +- 13 files changed, 17 insertions(+), 17 deletions(-) diff --git a/agents-api/agents_api/models/agent/get_agent.py b/agents-api/agents_api/models/agent/get_agent.py index 008e39454..0c14b45f0 100644 --- a/agents-api/agents_api/models/agent/get_agent.py +++ b/agents-api/agents_api/models/agent/get_agent.py @@ -70,7 +70,7 @@ def get_agent(*, developer_id: UUID, agent_id: UUID) -> tuple[list[str], dict]: default_settings, instructions, ] := input[id], - *agents { + *agents:developer_id_metadata_agent_id_idx { developer_id: to_uuid($developer_id), agent_id: id, model, diff --git a/agents-api/agents_api/models/agent/list_agents.py b/agents-api/agents_api/models/agent/list_agents.py index 882b6c8c6..f7915a84e 100644 --- a/agents-api/agents_api/models/agent/list_agents.py +++ b/agents-api/agents_api/models/agent/list_agents.py @@ -76,7 +76,7 @@ def list_agents( default_settings, instructions, ] := input[developer_id], - *agents {{ + *agents:developer_id_metadata_agent_id_idx {{ developer_id, agent_id: id, model, diff --git a/agents-api/agents_api/models/docs/list_docs.py b/agents-api/agents_api/models/docs/list_docs.py index dd389d58c..7dfd31e37 100644 --- a/agents-api/agents_api/models/docs/list_docs.py +++ b/agents-api/agents_api/models/docs/list_docs.py @@ -106,7 +106,7 @@ def list_docs( ] := owner_type = $owner_type, owner_id = to_uuid($owner_id), - *docs {{ + *docs:owner_id_metadata_doc_id_idx {{ owner_type, owner_id, doc_id: id, diff --git a/agents-api/agents_api/models/docs/search_docs_by_embedding.py b/agents-api/agents_api/models/docs/search_docs_by_embedding.py index 992e12f9d..d21c8bc0c 100644 --- a/agents-api/agents_api/models/docs/search_docs_by_embedding.py +++ b/agents-api/agents_api/models/docs/search_docs_by_embedding.py @@ -91,7 +91,7 @@ def search_docs_by_embedding( snippet_counter[count(item)] := owners[owner_type, owner_id_str], owner_id = to_uuid(owner_id_str), - *docs {{ + *docs:owner_id_metadata_doc_id_idx {{ owner_type, owner_id, doc_id: item, @@ -142,7 +142,7 @@ def search_docs_by_embedding( input[owner_type, owner_id, query], # Restrict the search to all documents that match the owner - *docs {{ + *docs:owner_id_metadata_doc_id_idx {{ owner_type, owner_id, doc_id, @@ -203,7 +203,7 @@ def search_docs_by_embedding( input[owner_type, owner_id, query], # Restrict the search to all documents that match the owner - *docs {{ + *docs:owner_id_metadata_doc_id_idx {{ owner_type, owner_id, doc_id, diff --git a/agents-api/agents_api/models/docs/search_docs_by_text.py b/agents-api/agents_api/models/docs/search_docs_by_text.py index ac1a9f54f..6d3d93c98 100644 --- a/agents-api/agents_api/models/docs/search_docs_by_text.py +++ b/agents-api/agents_api/models/docs/search_docs_by_text.py @@ -90,7 +90,7 @@ def search_docs_by_text( candidate[doc_id] := input[owner_type, owner_id], - *docs {{ + *docs:owner_id_metadata_doc_id_idx {{ owner_type, owner_id, doc_id, @@ -145,7 +145,7 @@ def search_docs_by_text( metadata, ] := candidate[doc_id], - *docs {{ + *docs:owner_id_metadata_doc_id_idx {{ owner_type, owner_id, doc_id, diff --git a/agents-api/agents_api/models/execution/count_executions.py b/agents-api/agents_api/models/execution/count_executions.py index 7f10e5bfa..d130f0359 100644 --- a/agents-api/agents_api/models/execution/count_executions.py +++ b/agents-api/agents_api/models/execution/count_executions.py @@ -39,7 +39,7 @@ def count_executions( counter[count(id)] := input[task_id], - *executions { + *executions:task_id_execution_id_idx { task_id, execution_id: id, } diff --git a/agents-api/agents_api/models/execution/get_execution.py b/agents-api/agents_api/models/execution/get_execution.py index db0279b1f..c4e30ab64 100644 --- a/agents-api/agents_api/models/execution/get_execution.py +++ b/agents-api/agents_api/models/execution/get_execution.py @@ -52,7 +52,7 @@ def get_execution( ?[id, task_id, status, input, output, error, session_id, metadata, created_at, updated_at] := input[execution_id], - *executions { + *executions:execution_id_status_idx { task_id, execution_id, status, diff --git a/agents-api/agents_api/models/execution/get_paused_execution_token.py b/agents-api/agents_api/models/execution/get_paused_execution_token.py index 44eb8a4da..526a6012c 100644 --- a/agents-api/agents_api/models/execution/get_paused_execution_token.py +++ b/agents-api/agents_api/models/execution/get_paused_execution_token.py @@ -38,7 +38,7 @@ def get_paused_execution_token( check_status_query = """ ?[execution_id, status] := - *executions { + *executions:execution_id_status_idx { execution_id, status, }, @@ -55,7 +55,7 @@ def get_paused_execution_token( *executions { execution_id, }, - *transitions { + *transitions:execution_id_type_created_at_idx { execution_id, created_at, task_token, diff --git a/agents-api/agents_api/models/execution/list_execution_transitions.py b/agents-api/agents_api/models/execution/list_execution_transitions.py index 8931676f6..d30e8595e 100644 --- a/agents-api/agents_api/models/execution/list_execution_transitions.py +++ b/agents-api/agents_api/models/execution/list_execution_transitions.py @@ -35,7 +35,7 @@ def list_execution_transitions( query = f""" ?[id, execution_id, type, current, next, output, metadata, updated_at, created_at] := - *transitions {{ + *transitions:execution_id_type_created_at_idx {{ execution_id, transition_id: id, type, diff --git a/agents-api/agents_api/models/task/get_task.py b/agents-api/agents_api/models/task/get_task.py index 460fdc38b..ab6edb2c3 100644 --- a/agents-api/agents_api/models/task/get_task.py +++ b/agents-api/agents_api/models/task/get_task.py @@ -63,7 +63,7 @@ def get_task( metadata, ] := input[task_id], - *tasks { + *tasks:task_id_agent_id_idx { agent_id, task_id, updated_at_ms, diff --git a/agents-api/agents_api/models/user/get_user.py b/agents-api/agents_api/models/user/get_user.py index 69b3da883..89f49dae1 100644 --- a/agents-api/agents_api/models/user/get_user.py +++ b/agents-api/agents_api/models/user/get_user.py @@ -85,7 +85,7 @@ def get_user( updated_at, metadata, ] := input[developer_id, id], - *users { + *users:developer_id_metadata_user_id_idx { user_id: id, developer_id, name, diff --git a/agents-api/agents_api/models/user/list_users.py b/agents-api/agents_api/models/user/list_users.py index f1e06adf4..cc857b1a1 100644 --- a/agents-api/agents_api/models/user/list_users.py +++ b/agents-api/agents_api/models/user/list_users.py @@ -88,7 +88,7 @@ def list_users( metadata, ] := input[developer_id], - *users {{ + *users:developer_id_metadata_user_id_idx {{ user_id: id, developer_id, name, diff --git a/agents-api/agents_api/models/user/patch_user.py b/agents-api/agents_api/models/user/patch_user.py index e091edc63..bd3fc0246 100644 --- a/agents-api/agents_api/models/user/patch_user.py +++ b/agents-api/agents_api/models/user/patch_user.py @@ -91,7 +91,7 @@ def patch_user( ?[{user_update_cols}, metadata] := input[{user_update_cols}], - *users {{ + *users:developer_id_metadata_user_id_idx {{ developer_id: to_uuid($developer_id), user_id: to_uuid($user_id), metadata: md, From 629cbf83f64ff74fcce03bba715580713dc3b7df Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Tue, 10 Dec 2024 14:38:22 +0300 Subject: [PATCH 11/14] fix: Add task_id constraint --- agents-api/agents_api/models/execution/update_execution.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/agents-api/agents_api/models/execution/update_execution.py b/agents-api/agents_api/models/execution/update_execution.py index 8135e8cae..b485e34f8 100644 --- a/agents-api/agents_api/models/execution/update_execution.py +++ b/agents-api/agents_api/models/execution/update_execution.py @@ -82,6 +82,7 @@ def update_execution( *executions:execution_id_status_idx { status, execution_id: to_uuid($execution_id), + task_id: to_uuid($task_id), }, status in $valid_previous_statuses @@ -124,5 +125,6 @@ def update_execution( "values": values, "valid_previous_statuses": valid_previous_statuses, "execution_id": str(execution_id), + "task_id": task_id, }, ) From 2dc7b458f9745a1a5bccce353f584a15aff89ee1 Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Tue, 10 Dec 2024 15:23:34 +0300 Subject: [PATCH 12/14] fix: Fix timeseries name --- .../agents_api/models/execution/create_execution_transition.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agents-api/agents_api/models/execution/create_execution_transition.py b/agents-api/agents_api/models/execution/create_execution_transition.py index cd2e2a8e1..59a63ed09 100644 --- a/agents-api/agents_api/models/execution/create_execution_transition.py +++ b/agents-api/agents_api/models/execution/create_execution_transition.py @@ -250,7 +250,7 @@ def validate_transition_targets(data: CreateTransitionRequest) -> None: _kind="inserted", )( cozo_query_async( - increase_counter("create_execution_transition")( + increase_counter("create_execution_transition_async")( _create_execution_transition ) ) From 2faf2248dcbee7b41585c4d81a1d051a80ef4e3b Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Tue, 10 Dec 2024 15:43:43 +0300 Subject: [PATCH 13/14] fix: Fix migration queries --- .../migrate_1733755642_transition_indices.py | 54 ++++++++++--------- 1 file changed, 28 insertions(+), 26 deletions(-) diff --git a/agents-api/migrations/migrate_1733755642_transition_indices.py b/agents-api/migrations/migrate_1733755642_transition_indices.py index eedebc565..1b33f4646 100644 --- a/agents-api/migrations/migrate_1733755642_transition_indices.py +++ b/agents-api/migrations/migrate_1733755642_transition_indices.py @@ -5,36 +5,38 @@ create_transition_indices = dict( - up=""" - ::index create executions:execution_id_status_idx { execution_id, status } - ::index create executions:execution_id_task_id_idx { execution_id, task_id } - ::index create executions:task_id_execution_id_idx { task_id, execution_id } - ::index create tasks:task_id_agent_id_idx { task_id, agent_id } - ::index create agents:agent_id_developer_id_idx { agent_id, developer_id } - ::index create sessions:session_id_developer_id_idx { session_id, developer_id } - ::index create docs:owner_id_metadata_doc_id_idx { owner_id, metadata, doc_id } - ::index create agents:developer_id_metadata_agent_id_idx { developer_id, metadata, agent_id } - ::index create users:developer_id_metadata_user_id_idx { developer_id, metadata, user_id } - ::index create transitions:execution_id_type_created_at_idx { execution_id, type, created_at } - """, - down=""" - ::index drop executions:execution_id_status_idx - ::index drop executions:execution_id_task_id_idx - ::index drop executions:task_id_execution_id_idx - ::index drop tasks:task_id_agent_id_idx - ::index drop agents:agent_id_developer_id_idx - ::index drop sessions:session_id_developer_id_idx - ::index drop docs:owner_id_metadata_doc_id_idx - ::index drop agents:developer_id_metadata_agent_id_idx - ::index drop users:developer_id_metadata_user_id_idx - ::index drop transitions:execution_id_type_created_at_idx - """, + up=[ + "::index create executions:execution_id_status_idx { execution_id, status }", + "::index create executions:execution_id_task_id_idx { execution_id, task_id }", + "::index create executions:task_id_execution_id_idx { task_id, execution_id }", + "::index create tasks:task_id_agent_id_idx { task_id, agent_id }", + "::index create agents:agent_id_developer_id_idx { agent_id, developer_id }", + "::index create sessions:session_id_developer_id_idx { session_id, developer_id }", + "::index create docs:owner_id_metadata_doc_id_idx { owner_id, metadata, doc_id }", + "::index create agents:developer_id_metadata_agent_id_idx { developer_id, metadata, agent_id }", + "::index create users:developer_id_metadata_user_id_idx { developer_id, metadata, user_id }", + "::index create transitions:execution_id_type_created_at_idx { execution_id, type, created_at }", + ], + down=[ + "::index drop executions:execution_id_status_idx", + "::index drop executions:execution_id_task_id_idx", + "::index drop executions:task_id_execution_id_idx", + "::index drop tasks:task_id_agent_id_idx", + "::index drop agents:agent_id_developer_id_idx", + "::index drop sessions:session_id_developer_id_idx", + "::index drop docs:owner_id_metadata_doc_id_idx", + "::index drop agents:developer_id_metadata_agent_id_idx", + "::index drop users:developer_id_metadata_user_id_idx", + "::index drop transitions:execution_id_type_created_at_idx", + ], ) def up(client): - client.run(create_transition_indices["up"]) + for q in create_transition_indices["up"]: + client.run(q) def down(client): - client.run(create_transition_indices["down"]) + for q in create_transition_indices["down"]: + client.run(q) From 2f88b8c23b389d1a7485bbc3872b9aae243944ca Mon Sep 17 00:00:00 2001 From: Dmitry Paramonov Date: Tue, 10 Dec 2024 15:44:03 +0300 Subject: [PATCH 14/14] fix: Remove indices as they cause no field error --- agents-api/agents_api/models/agent/get_agent.py | 2 +- agents-api/agents_api/models/agent/list_agents.py | 2 +- agents-api/agents_api/models/docs/list_docs.py | 2 +- agents-api/agents_api/models/docs/search_docs_by_embedding.py | 4 ++-- agents-api/agents_api/models/docs/search_docs_by_text.py | 2 +- 5 files changed, 6 insertions(+), 6 deletions(-) diff --git a/agents-api/agents_api/models/agent/get_agent.py b/agents-api/agents_api/models/agent/get_agent.py index 0c14b45f0..008e39454 100644 --- a/agents-api/agents_api/models/agent/get_agent.py +++ b/agents-api/agents_api/models/agent/get_agent.py @@ -70,7 +70,7 @@ def get_agent(*, developer_id: UUID, agent_id: UUID) -> tuple[list[str], dict]: default_settings, instructions, ] := input[id], - *agents:developer_id_metadata_agent_id_idx { + *agents { developer_id: to_uuid($developer_id), agent_id: id, model, diff --git a/agents-api/agents_api/models/agent/list_agents.py b/agents-api/agents_api/models/agent/list_agents.py index f7915a84e..882b6c8c6 100644 --- a/agents-api/agents_api/models/agent/list_agents.py +++ b/agents-api/agents_api/models/agent/list_agents.py @@ -76,7 +76,7 @@ def list_agents( default_settings, instructions, ] := input[developer_id], - *agents:developer_id_metadata_agent_id_idx {{ + *agents {{ developer_id, agent_id: id, model, diff --git a/agents-api/agents_api/models/docs/list_docs.py b/agents-api/agents_api/models/docs/list_docs.py index 7dfd31e37..dd389d58c 100644 --- a/agents-api/agents_api/models/docs/list_docs.py +++ b/agents-api/agents_api/models/docs/list_docs.py @@ -106,7 +106,7 @@ def list_docs( ] := owner_type = $owner_type, owner_id = to_uuid($owner_id), - *docs:owner_id_metadata_doc_id_idx {{ + *docs {{ owner_type, owner_id, doc_id: id, diff --git a/agents-api/agents_api/models/docs/search_docs_by_embedding.py b/agents-api/agents_api/models/docs/search_docs_by_embedding.py index d21c8bc0c..49ec069df 100644 --- a/agents-api/agents_api/models/docs/search_docs_by_embedding.py +++ b/agents-api/agents_api/models/docs/search_docs_by_embedding.py @@ -142,7 +142,7 @@ def search_docs_by_embedding( input[owner_type, owner_id, query], # Restrict the search to all documents that match the owner - *docs:owner_id_metadata_doc_id_idx {{ + *docs {{ owner_type, owner_id, doc_id, @@ -203,7 +203,7 @@ def search_docs_by_embedding( input[owner_type, owner_id, query], # Restrict the search to all documents that match the owner - *docs:owner_id_metadata_doc_id_idx {{ + *docs {{ owner_type, owner_id, doc_id, diff --git a/agents-api/agents_api/models/docs/search_docs_by_text.py b/agents-api/agents_api/models/docs/search_docs_by_text.py index 6d3d93c98..ce5319673 100644 --- a/agents-api/agents_api/models/docs/search_docs_by_text.py +++ b/agents-api/agents_api/models/docs/search_docs_by_text.py @@ -145,7 +145,7 @@ def search_docs_by_text( metadata, ] := candidate[doc_id], - *docs:owner_id_metadata_doc_id_idx {{ + *docs {{ owner_type, owner_id, doc_id,