Skip to content

perf(cli): stop importing sentence_transformers for every command - #60

Merged
r0h1tb merged 1 commit into
mainfrom
perf/lazy-sentence-transformers
Aug 2, 2026
Merged

perf(cli): stop importing sentence_transformers for every command#60
r0h1tb merged 1 commit into
mainfrom
perf/lazy-sentence-transformers

Conversation

@r0h1tb

@r0h1tb r0h1tb commented Aug 2, 2026

Copy link
Copy Markdown
Collaborator

Found while running a full UAT sweep of the CLI against a live index — every command had a hard ~2.9s floor, including pure graph lookups.

Stacked on #50 and #56. Review those first; this branch contains them.

Cause

import ast_rag.cli costs 3.65s, and -X importtime attributes 3.05s of it to sentence_transformers (and transitively torch/transformers), reached via services/__init__embedding_manager:

    3651 ms   ast_rag.cli
    3508 ms     ast_rag.services.parsing.parser_manager
    3462 ms           ast_rag.services.embedding_manager
    3054 ms             sentence_transformers
    1447 ms                   transformers.configuration_utils

Only semantic search needs the model. goto, callers, refs, sig, blocks, cache-stats are pure Neo4j lookups and never touch it — but they all paid for it.

Fix

Move the import into _get_model(), where it's actually used, with TYPE_CHECKING for the annotations. Same shape as the watchdog fix in #50.

import ast_rag.cli   3.65s -> 0.61s

Measured end to end

15 CLI scenarios against a live index (raged indexing itself — 965 nodes, 3,782 edges):

before after
median command latency 3.068s 0.558s
fastest command 2.864s 0.487s
goto 3.371s 0.655s
callers 2.887s 0.487s

query still pays the model load on first use, as it must — it's unchanged at ~12s cold, and it still returns its full result set, so the lazy path is exercised and works.

For an interactive tool, and especially for an agent calling goto/callers in a loop, a 5.5x cut on the common path is the difference between usable and not.

Verification

pytest tests/    3 failed, 191 passed, 1 xfailed
ruff check       All checks passed!

Baseline on main is 3 failed, 174 passed — same three pre-existing failures (fixed in #51). No test changes here; the existing suite covers that embeddings still work, and the UAT sweep confirms query behaves identically.

@r0h1tb

r0h1tb commented Aug 2, 2026

Copy link
Copy Markdown
Collaborator Author

Rebased onto the rebased #56 (this PR is stacked on it, so #56 needs to merge first — until then the diff here shows #56's commits too).

Resolved one conflict in embedding_manager.py: #56 added Callable to the typing import while this branch added TYPE_CHECKING. Both are used, so the import is now from typing import TYPE_CHECKING, Callable, Optional.

Measured effect of this PR — importing ast_rag.services on each branch:

branch pulls in sentence_transformers? import time
main yes 5.80s
#56 yes 4.78s
this branch no 0.48s

That's ~12x on every CLI invocation, including --help, because services/__init__.py imported embedding_manager eagerly and that imports sentence_transformers at module scope.

Incidental confirmation that this is worth having: while setting up to verify these PRs I couldn't run the parser tests without installing torch + sentence-transformers, purely because of this import chain. On this branch they run without the ML stack present.

Verified against main (6f71a9b): ruff check and ruff format --check clean; pytest tests/ gives 216 passed, 1 skipped, 1 xfailed vs. main's 209/1/1 — the +7 are #56's tests carried in the stack, nothing pre-existing changed state.

Importing ast_rag.cli took 3.65s, of which 3.05s was sentence_transformers
(and transitively torch and transformers), reached via
services/__init__ -> embedding_manager.

Only semantic search needs it. Every other command -- goto, callers, refs,
sig, blocks, cache-stats -- paid the full cost and never used the model.

Moves the import into _get_model(), where it is actually needed, with
TYPE_CHECKING for the annotations.

  import ast_rag.cli   3.65s -> 0.61s

Measured end to end over 15 CLI scenarios against a live index:

  median command latency   3.068s -> 0.558s
  fastest command          2.864s -> 0.487s

Semantic search is unaffected; the model loads on first use and 'query' still
returns its full result set.
@r0h1tb
r0h1tb force-pushed the perf/lazy-sentence-transformers branch from d726ec0 to 0237383 Compare August 2, 2026 12:13
@r0h1tb
r0h1tb merged commit ff8cf65 into main Aug 2, 2026
1 check passed
@r0h1tb
r0h1tb deleted the perf/lazy-sentence-transformers branch August 2, 2026 12:16
@github-project-automation github-project-automation Bot moved this from Backlog to Done in raged kanban Aug 2, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Status: Done

Development

Successfully merging this pull request may close these issues.

1 participant