feat(agents-api): Naive speed improvement to cozo queries by adding limit=1 to all relevant queries

.github/workflows/lint-and-format.yml → .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)"

.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' }}

.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' }}

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 = [

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})

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 = [

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})

agents-api/agents_api/models/docs/embed_snippets.py

@@ -81,6 +81,7 @@ def embed_snippets(
             }},
             index > {max(snippet_indices)}
 
+        :limit 1
         :assert none
     """
 

agents-api/agents_api/models/docs/get_doc.py

@@ -87,6 +87,8 @@ def get_doc(
                 metadata,
             },
             snippets[snippet_data]
+
+        :limit 1
     """
 
     queries = [

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),

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

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

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
     """
 

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 (

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 = [

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 = [

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"""

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 = [

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 = [

agents-api/agents_api/models/tools/create_tools.py

@@ -80,6 +80,7 @@ def create_tools(
                 name,
             }
 
+        :limit 1
         :assert none
     """
 

agents-api/agents_api/models/tools/get_tool.py

@@ -68,6 +68,8 @@ def get_tool(
                 updated_at,
                 created_at,
             }
+
+        :limit 1
     """
 
     queries = [

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 = [

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),

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