From 348079ec38f9c2471affb27d817df65e1dee1f64 Mon Sep 17 00:00:00 2001 From: Diwank Singh Tomer Date: Thu, 3 Oct 2024 15:27:38 -0400 Subject: [PATCH 1/3] feat(agents-api): Naive speed improvement to cozo queries by adding limit=1 to all relevant queries Signed-off-by: Diwank Singh Tomer --- ...-and-format.yml => lint-agents-api-pr.yml} | 18 ------ .github/workflows/test-agents-api-pr.yml | 48 ++++++++++++++++ .github/workflows/typecheck-agents-api-pr.yml | 56 +++++++++++++++++++ .../agents_api/models/agent/get_agent.py | 2 + .../models/chat/get_cached_response.py | 1 + .../models/chat/prepare_chat_context.py | 2 + .../models/developer/get_developer.py | 2 + .../agents_api/models/docs/embed_snippets.py | 1 + agents-api/agents_api/models/docs/get_doc.py | 2 + .../agents_api/models/entry/get_history.py | 4 +- .../execution/create_execution_transition.py | 2 + .../execution/get_execution_transition.py | 4 +- .../execution/get_paused_execution_token.py | 1 + .../execution/get_temporal_workflow_data.py | 2 + .../models/execution/lookup_temporal_data.py | 2 + .../execution/prepare_execution_input.py | 2 + .../models/execution/update_execution.py | 2 + .../agents_api/models/session/get_session.py | 5 +- .../models/session/prepare_session_data.py | 2 + .../agents_api/models/tools/create_tools.py | 1 + .../agents_api/models/tools/get_tool.py | 2 + .../tools/get_tool_args_from_metadata.py | 4 ++ agents-api/agents_api/models/user/get_user.py | 5 +- agents-api/agents_api/models/utils.py | 4 ++ 24 files changed, 152 insertions(+), 22 deletions(-) rename .github/workflows/{lint-and-format.yml => lint-agents-api-pr.yml} (74%) create mode 100644 .github/workflows/test-agents-api-pr.yml create mode 100644 .github/workflows/typecheck-agents-api-pr.yml diff --git a/.github/workflows/lint-and-format.yml b/.github/workflows/lint-agents-api-pr.yml similarity index 74% rename from .github/workflows/lint-and-format.yml rename to .github/workflows/lint-agents-api-pr.yml index 65e4983c4..90fdd0a4b 100644 --- a/.github/workflows/lint-and-format.yml +++ b/.github/workflows/lint-agents-api-pr.yml @@ -33,35 +33,17 @@ jobs: restore-keys: | ${{ runner.os }}-agents-api-poetry- - - name: Cache pytype - uses: actions/cache@v4 - with: - path: agents-api/.pytype - key: ${{ runner.os }}-agents-api-pytype-${{ hashFiles('agents-api/**/*.py') }} - restore-keys: | - ${{ runner.os }}-agents-api-pytype- - - name: Install dependencies run: | cd agents-api poetry install - - name: Typecheck - run: | - cd agents-api - poetry run poe typecheck - - name: Lint and format run: | cd agents-api poetry run poe format poetry run poe lint - - name: Run tests - run: | - cd agents-api - poetry run poe test - - uses: stefanzweifel/git-auto-commit-action@v4 with: commit_message: "refactor: Lint agents-api (CI)" diff --git a/.github/workflows/test-agents-api-pr.yml b/.github/workflows/test-agents-api-pr.yml new file mode 100644 index 000000000..654711dad --- /dev/null +++ b/.github/workflows/test-agents-api-pr.yml @@ -0,0 +1,48 @@ +name: Test agents-api +run-name: ${{ github.actor }} is testing the code + +# TODO: Fix CI github actions +# SCRUM-26 + +on: [pull_request] + +jobs: + Test: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install and configure Poetry + uses: snok/install-poetry@v1 + + - name: Configure Poetry to use .venv + run: | + cd agents-api + poetry config virtualenvs.in-project true + + - name: Cache Poetry virtualenv + uses: actions/cache@v4 + with: + path: agents-api/.venv + key: ${{ runner.os }}-agents-api-poetry-${{ hashFiles('agents-api/poetry.lock') }} + restore-keys: | + ${{ runner.os }}-agents-api-poetry- + + - name: Install dependencies + run: | + cd agents-api + poetry install + + - name: Run tests + run: | + cd agents-api + poetry run poe test + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} diff --git a/.github/workflows/typecheck-agents-api-pr.yml b/.github/workflows/typecheck-agents-api-pr.yml new file mode 100644 index 000000000..985740b87 --- /dev/null +++ b/.github/workflows/typecheck-agents-api-pr.yml @@ -0,0 +1,56 @@ +name: Typecheck agents-api +run-name: ${{ github.actor }} is typechecking the code + +# TODO: Fix CI github actions +# SCRUM-26 + +on: [pull_request] + +jobs: + Typecheck: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install and configure Poetry + uses: snok/install-poetry@v1 + + - name: Configure Poetry to use .venv + run: | + cd agents-api + poetry config virtualenvs.in-project true + + - name: Cache Poetry virtualenv + uses: actions/cache@v4 + with: + path: agents-api/.venv + key: ${{ runner.os }}-agents-api-poetry-${{ hashFiles('agents-api/poetry.lock') }} + restore-keys: | + ${{ runner.os }}-agents-api-poetry- + + - name: Cache pytype + uses: actions/cache@v4 + with: + path: agents-api/.pytype + key: ${{ runner.os }}-agents-api-pytype-${{ hashFiles('agents-api/**/*.py') }} + restore-keys: | + ${{ runner.os }}-agents-api-pytype- + + - name: Install dependencies + run: | + cd agents-api + poetry install + + - name: Typecheck + run: | + cd agents-api + poetry run poe typecheck + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} diff --git a/agents-api/agents_api/models/agent/get_agent.py b/agents-api/agents_api/models/agent/get_agent.py index 956fa46a5..f8127751a 100644 --- a/agents-api/agents_api/models/agent/get_agent.py +++ b/agents-api/agents_api/models/agent/get_agent.py @@ -101,6 +101,8 @@ def get_agent(*, developer_id: UUID, agent_id: UUID) -> tuple[list[str], dict]: "min_p": min_p, "preset": preset, } + + :limit 1 """ queries = [ diff --git a/agents-api/agents_api/models/chat/get_cached_response.py b/agents-api/agents_api/models/chat/get_cached_response.py index 5440a6703..368c88567 100644 --- a/agents-api/agents_api/models/chat/get_cached_response.py +++ b/agents-api/agents_api/models/chat/get_cached_response.py @@ -9,6 +9,7 @@ def get_cached_response(key: str) -> tuple[str, dict]: query = """ input[key] <- [[$key]] ?[key, value] := input[key], *session_cache{key, value} + :limit 1 """ return (query, {"key": key}) diff --git a/agents-api/agents_api/models/chat/prepare_chat_context.py b/agents-api/agents_api/models/chat/prepare_chat_context.py index 4d521acb0..9be2d64aa 100644 --- a/agents-api/agents_api/models/chat/prepare_chat_context.py +++ b/agents-api/agents_api/models/chat/prepare_chat_context.py @@ -119,6 +119,8 @@ def prepare_chat_context( ?[{', '.join(session_data_fields)}, toolsets] := *_session_data_json {{ {', '.join(session_data_fields)} }}, *_toolsets_json {{ toolsets }} + + :limit 1 """ queries = [ diff --git a/agents-api/agents_api/models/developer/get_developer.py b/agents-api/agents_api/models/developer/get_developer.py index 31ade5334..0ae5421aa 100644 --- a/agents-api/agents_api/models/developer/get_developer.py +++ b/agents-api/agents_api/models/developer/get_developer.py @@ -67,6 +67,8 @@ def get_developer( created_at, updated_at, } + + :limit 1 """ return (query, {"developer_id": developer_id}) diff --git a/agents-api/agents_api/models/docs/embed_snippets.py b/agents-api/agents_api/models/docs/embed_snippets.py index e810d0379..993e76396 100644 --- a/agents-api/agents_api/models/docs/embed_snippets.py +++ b/agents-api/agents_api/models/docs/embed_snippets.py @@ -81,6 +81,7 @@ def embed_snippets( }}, index > {max(snippet_indices)} + :limit 1 :assert none """ diff --git a/agents-api/agents_api/models/docs/get_doc.py b/agents-api/agents_api/models/docs/get_doc.py index 4b3996c27..e81084985 100644 --- a/agents-api/agents_api/models/docs/get_doc.py +++ b/agents-api/agents_api/models/docs/get_doc.py @@ -87,6 +87,8 @@ def get_doc( metadata, }, snippets[snippet_data] + + :limit 1 """ queries = [ diff --git a/agents-api/agents_api/models/entry/get_history.py b/agents-api/agents_api/models/entry/get_history.py index 8918d357f..bd3c0397a 100644 --- a/agents-api/agents_api/models/entry/get_history.py +++ b/agents-api/agents_api/models/entry/get_history.py @@ -131,7 +131,9 @@ def get_history( session_relations[relations], session_id = to_uuid($session_id), created_at = now() - """ + + :limit 1 + """ queries = [ verify_developer_id_query(developer_id), 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 42c6f1c22..f40395126 100644 --- a/agents-api/agents_api/models/execution/create_execution_transition.py +++ b/agents-api/agents_api/models/execution/create_execution_transition.py @@ -141,6 +141,8 @@ def create_execution_transition( found = length(prev_transitions), valid = if($next_type == "init", found == 0, found > 0), assert(valid, "Invalid transition"), + + :limit 1 """ # Prepare the insert query diff --git a/agents-api/agents_api/models/execution/get_execution_transition.py b/agents-api/agents_api/models/execution/get_execution_transition.py index d3ebf3783..e2b38789a 100644 --- a/agents-api/agents_api/models/execution/get_execution_transition.py +++ b/agents-api/agents_api/models/execution/get_execution_transition.py @@ -65,7 +65,9 @@ def get_execution_transition( is_null(next_tuple), null, {"workflow": next_tuple->0, "step": next_tuple->1}, - ), + ) + + :limit 1 """ get_query += filter 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 b4c9f9081..44eb8a4da 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 @@ -45,6 +45,7 @@ def get_paused_execution_token( execution_id = to_uuid($execution_id), status = "awaiting_input" + :limit 1 :assert some """ diff --git a/agents-api/agents_api/models/execution/get_temporal_workflow_data.py b/agents-api/agents_api/models/execution/get_temporal_workflow_data.py index bb0a462ef..8b1bf4604 100644 --- a/agents-api/agents_api/models/execution/get_temporal_workflow_data.py +++ b/agents-api/agents_api/models/execution/get_temporal_workflow_data.py @@ -45,6 +45,8 @@ def get_temporal_workflow_data( result_run_id, first_execution_run_id, } + + :limit 1 """ return ( diff --git a/agents-api/agents_api/models/execution/lookup_temporal_data.py b/agents-api/agents_api/models/execution/lookup_temporal_data.py index e3c369b94..35f09129b 100644 --- a/agents-api/agents_api/models/execution/lookup_temporal_data.py +++ b/agents-api/agents_api/models/execution/lookup_temporal_data.py @@ -43,6 +43,8 @@ def lookup_temporal_data( *temporal_executions_lookup { id, execution_id, run_id, first_execution_run_id, result_run_id } + + :limit 1 """ queries = [ diff --git a/agents-api/agents_api/models/execution/prepare_execution_input.py b/agents-api/agents_api/models/execution/prepare_execution_input.py index b763a7508..5f30e7f83 100644 --- a/agents-api/agents_api/models/execution/prepare_execution_input.py +++ b/agents-api/agents_api/models/execution/prepare_execution_input.py @@ -187,6 +187,8 @@ def prepare_execution_input( user = null, session = null, arguments = execution->"input" + + :limit 1 """ queries = [ diff --git a/agents-api/agents_api/models/execution/update_execution.py b/agents-api/agents_api/models/execution/update_execution.py index c2f40c7ff..4d0367ae4 100644 --- a/agents-api/agents_api/models/execution/update_execution.py +++ b/agents-api/agents_api/models/execution/update_execution.py @@ -82,6 +82,8 @@ def update_execution( ?[num] := valid_status[num], assert(num > 0, 'Invalid status') + + :limit 1 """ update_query = f""" diff --git a/agents-api/agents_api/models/session/get_session.py b/agents-api/agents_api/models/session/get_session.py index 2d8956eb7..45d971c3f 100644 --- a/agents-api/agents_api/models/session/get_session.py +++ b/agents-api/agents_api/models/session/get_session.py @@ -95,7 +95,10 @@ def get_session( token_budget, context_overflow, @ "END" - }, updated_at = to_int(validity) + }, + updated_at = to_int(validity) + + :limit 1 """ queries = [ diff --git a/agents-api/agents_api/models/session/prepare_session_data.py b/agents-api/agents_api/models/session/prepare_session_data.py index e8cfb7fc7..bbbd9c4cd 100644 --- a/agents-api/agents_api/models/session/prepare_session_data.py +++ b/agents-api/agents_api/models/session/prepare_session_data.py @@ -209,6 +209,8 @@ def prepare_session_data( session_data[session], user_data[users], agent_data[agents] + + :limit 1 """ queries = [ diff --git a/agents-api/agents_api/models/tools/create_tools.py b/agents-api/agents_api/models/tools/create_tools.py index dd8397797..87ea3a9cf 100644 --- a/agents-api/agents_api/models/tools/create_tools.py +++ b/agents-api/agents_api/models/tools/create_tools.py @@ -80,6 +80,7 @@ def create_tools( name, } + :limit 1 :assert none """ diff --git a/agents-api/agents_api/models/tools/get_tool.py b/agents-api/agents_api/models/tools/get_tool.py index 5ea009064..465fd2efe 100644 --- a/agents-api/agents_api/models/tools/get_tool.py +++ b/agents-api/agents_api/models/tools/get_tool.py @@ -68,6 +68,8 @@ def get_tool( updated_at, created_at, } + + :limit 1 """ queries = [ diff --git a/agents-api/agents_api/models/tools/get_tool_args_from_metadata.py b/agents-api/agents_api/models/tools/get_tool_args_from_metadata.py index ade296a29..2cdb92cb9 100644 --- a/agents-api/agents_api/models/tools/get_tool_args_from_metadata.py +++ b/agents-api/agents_api/models/tools/get_tool_args_from_metadata.py @@ -46,6 +46,8 @@ def tool_args_for_task( # Right values overwrite left values # See: https://docs.cozodb.org/en/latest/functions.html#Func.Vector.concat values = concat(agent_{arg_type}, task_{arg_type}), + + :limit 1 """ queries = [ @@ -88,6 +90,8 @@ def tool_args_for_session( # Right values overwrite left values # See: https://docs.cozodb.org/en/latest/functions.html#Func.Vector.concat values = concat(agent_{arg_type}, session_{arg_type}), + + :limit 1 """ queries = [ diff --git a/agents-api/agents_api/models/user/get_user.py b/agents-api/agents_api/models/user/get_user.py index cc7c6f970..2b4f59c83 100644 --- a/agents-api/agents_api/models/user/get_user.py +++ b/agents-api/agents_api/models/user/get_user.py @@ -67,7 +67,10 @@ def get_user( created_at, updated_at, metadata, - }""" + } + + :limit 1 + """ queries = [ verify_developer_id_query(developer_id), diff --git a/agents-api/agents_api/models/utils.py b/agents-api/agents_api/models/utils.py index c97d92451..c163642c0 100644 --- a/agents-api/agents_api/models/utils.py +++ b/agents-api/agents_api/models/utils.py @@ -127,6 +127,8 @@ def verify_developer_id_query(developer_id: UUID | str) -> str: matched[num], exists = num > 0, assert(exists, "Developer does not exist") + + :limit 1 """ @@ -162,6 +164,8 @@ def verify_developer_owns_resource_query( found[num], exists = num > 0, assert(exists, "Developer does not own resource {resource} with {resource_id_key} {resource_id_value}") + + :limit 1 """ rule = rule_head + rule_body + assertion From 248288f3721ab88127b0db1137c5fd579224efb5 Mon Sep 17 00:00:00 2001 From: Diwank Singh Tomer Date: Thu, 3 Oct 2024 20:04:29 -0400 Subject: [PATCH 2/3] fix(.github): Remove hash for pytype cache Signed-off-by: Diwank Singh Tomer --- .github/workflows/typecheck-agents-api-pr.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/typecheck-agents-api-pr.yml b/.github/workflows/typecheck-agents-api-pr.yml index 985740b87..884ef3311 100644 --- a/.github/workflows/typecheck-agents-api-pr.yml +++ b/.github/workflows/typecheck-agents-api-pr.yml @@ -37,7 +37,7 @@ jobs: uses: actions/cache@v4 with: path: agents-api/.pytype - key: ${{ runner.os }}-agents-api-pytype-${{ hashFiles('agents-api/**/*.py') }} + key: ${{ runner.os }}-agents-api-pytype- restore-keys: | ${{ runner.os }}-agents-api-pytype- From 4cd6bb0379e5afe3fb9941a338a1d7ab89f8ca0d Mon Sep 17 00:00:00 2001 From: Diwank Singh Tomer Date: Thu, 3 Oct 2024 20:07:55 -0400 Subject: [PATCH 3/3] fix(.github): Add base_ref as key for pytype cache and prevent tests from running after first fail Signed-off-by: Diwank Singh Tomer --- .github/workflows/test-agents-api-pr.yml | 2 +- .github/workflows/typecheck-agents-api-pr.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test-agents-api-pr.yml b/.github/workflows/test-agents-api-pr.yml index 654711dad..b0f815634 100644 --- a/.github/workflows/test-agents-api-pr.yml +++ b/.github/workflows/test-agents-api-pr.yml @@ -41,7 +41,7 @@ jobs: - name: Run tests run: | cd agents-api - poetry run poe test + poetry run poe test --fail-limit 1 concurrency: group: ${{ github.workflow }}-${{ github.ref }} diff --git a/.github/workflows/typecheck-agents-api-pr.yml b/.github/workflows/typecheck-agents-api-pr.yml index 884ef3311..513390883 100644 --- a/.github/workflows/typecheck-agents-api-pr.yml +++ b/.github/workflows/typecheck-agents-api-pr.yml @@ -37,7 +37,7 @@ jobs: uses: actions/cache@v4 with: path: agents-api/.pytype - key: ${{ runner.os }}-agents-api-pytype- + key: ${{ runner.os }}-agents-api-pytype-${{ github.base_ref }} restore-keys: | ${{ runner.os }}-agents-api-pytype-