From b362f4abade79056b3d6566c057aa4266c14e317 Mon Sep 17 00:00:00 2001 From: Rohit Behera <126186063+r0h1tb@users.noreply.github.com> Date: Sun, 2 Aug 2026 00:29:02 +0530 Subject: [PATCH] feat(cli): real progress bars for indexing (#5) Uses the existing rich dependency rather than adding tqdm. Parsing previously used console.status(), an indeterminate spinner. It now reports a bar with count, elapsed and ETA alongside the current file name. The bigger gap was embeddings. build_embeddings() ran behind a single 'Building embeddings...' spinner with no output until it finished. On a first run that also downloads the model this is many minutes of apparent hang -- I had to query Qdrant directly to confirm the process was alive. It now takes an optional progress_callback(done, total), invoked after each batch, and the CLI renders it as a bar. Default is None, so the API is unchanged for existing callers. The callback also fires once with (0, 0) when there is nothing embeddable, so callers always get a terminal update instead of a bar that never resolves. Bars are transient, so they clear on completion and leave the existing summary lines as the only residue. Verified by re-indexing a 105-file Spring Boot project end to end. Suite: 3 failed, 178 passed (baseline on main: 3 failed, 174 passed) -- same three pre-existing failures, fixed separately in #51. --- ast_rag/cli.py | 46 +++++++++++++++-- ast_rag/services/embedding_manager.py | 14 ++++- tests/test_indexing_progress.py | 74 +++++++++++++++++++++++++++ 3 files changed, 128 insertions(+), 6 deletions(-) create mode 100644 tests/test_indexing_progress.py diff --git a/ast_rag/cli.py b/ast_rag/cli.py index e22741b..7a516c0 100644 --- a/ast_rag/cli.py +++ b/ast_rag/cli.py @@ -29,6 +29,15 @@ import typer from rich.console import Console +from rich.progress import ( + BarColumn, + MofNCompleteColumn, + Progress, + SpinnerColumn, + TextColumn, + TimeElapsedColumn, + TimeRemainingColumn, +) from ast_rag.models import ProjectConfig from ast_rag.services.parsing.parser_manager import ParserManager, walk_source_files @@ -52,6 +61,26 @@ console = Console() +def _index_progress() -> Progress: + """Progress bar used by the indexing phases. + + Indexing a large repository can run for many minutes -- embedding in + particular -- so the phases report count, elapsed and ETA rather than an + indeterminate spinner. `transient` keeps the finished bar from cluttering + the summary output. + """ + return Progress( + SpinnerColumn(), + TextColumn("[progress.description]{task.description}"), + BarColumn(), + MofNCompleteColumn(), + TimeElapsedColumn(), + TimeRemainingColumn(), + console=console, + transient=True, + ) + + def _get_humanize_callback() -> callable: def callback(value: bool) -> bool: return value @@ -151,9 +180,11 @@ def init( all_blocks = [] all_block_edges = [] - with console.status(f"Parsing {len(files)} files...") as status: - for i, (fp, lang) in enumerate(files): - status.update(f"Parsing [{i + 1}/{len(files)}] {os.path.relpath(fp, root)}") + with _index_progress() as progress: + task = progress.add_task("Parsing", total=len(files)) + for fp, lang in files: + progress.update(task, description=f"Parsing {os.path.relpath(fp, root)}") + progress.advance(task) tree = pm.parse_file(fp) if tree is None: continue @@ -198,8 +229,13 @@ def init( # 4. Build embeddings embed = EmbeddingManager(cfg.qdrant, cfg.embedding, neo4j_driver=driver) - with console.status("Building embeddings..."): - count = embed.build_embeddings(all_nodes) + with _index_progress() as progress: + task = progress.add_task("Building embeddings", total=None) + + def _on_batch(done: int, total: int) -> None: + progress.update(task, completed=done, total=total or None) + + count = embed.build_embeddings(all_nodes, progress_callback=_on_batch) console.print(f"[green]Indexed {count} node embeddings.[/green]") console.rule("[bold green]Done[/bold green]") diff --git a/ast_rag/services/embedding_manager.py b/ast_rag/services/embedding_manager.py index c888c38..6e213ae 100644 --- a/ast_rag/services/embedding_manager.py +++ b/ast_rag/services/embedding_manager.py @@ -21,7 +21,7 @@ import logging import uuid -from typing import Optional +from typing import Callable, Optional import numpy as np from neo4j import Driver @@ -279,13 +279,23 @@ def build_embeddings( self, nodes: list[ASTNode], batch_size: int = 64, + progress_callback: Optional[Callable[[int, int], None]] = None, ) -> int: """Build embeddings for all embeddable nodes (bulk). + Args: + nodes: Nodes to consider; non-embeddable kinds are skipped. + batch_size: Nodes encoded per batch. + progress_callback: Optional ``fn(done, total)`` invoked after each + batch. Embedding a large repository is the slowest phase of + indexing, so callers need a way to report progress. + Returns the number of nodes indexed. """ to_embed = [n for n in nodes if n.kind in EMBEDDABLE_KINDS] if not to_embed: + if progress_callback: + progress_callback(0, 0) return 0 logger.info("Building embeddings for %d nodes...", len(to_embed)) @@ -308,6 +318,8 @@ def build_embeddings( ] client.upsert(collection_name=name, points=points, wait=True) count += len(batch) + if progress_callback: + progress_callback(count, len(to_embed)) logger.debug( "Embedded batch %d/%d", i // batch_size + 1, diff --git a/tests/test_indexing_progress.py b/tests/test_indexing_progress.py new file mode 100644 index 0000000..338f51b --- /dev/null +++ b/tests/test_indexing_progress.py @@ -0,0 +1,74 @@ +"""Progress reporting for long indexing phases (issue #5).""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import numpy as np + +from ast_rag.cli import _index_progress +from ast_rag.models import ASTNode, NodeKind, Language +from ast_rag.services.embedding_manager import EmbeddingManager + + +def _node(i: int) -> ASTNode: + return ASTNode( + id=f"{i:024x}", + kind=NodeKind.FUNCTION, + name=f"fn_{i}", + qualified_name=f"mod.fn_{i}", + file_path="mod.py", + start_line=i, + end_line=i + 1, + start_byte=i * 10, + end_byte=i * 10 + 9, + lang=Language.PYTHON, + ) + + +def _manager() -> EmbeddingManager: + qcfg, ecfg = MagicMock(), MagicMock() + qcfg.collection_name = "test" + ecfg.hybrid_search = False # otherwise __init__ compares MagicMock weights + em = EmbeddingManager(qcfg, ecfg) + em._get_client = MagicMock(return_value=MagicMock()) + em._encode = lambda texts: np.zeros((len(texts), 4), dtype=np.float32) + return em + + +def test_progress_callback_reports_monotonic_completion(): + nodes = [_node(i) for i in range(10)] + seen: list[tuple[int, int]] = [] + + with patch("ast_rag.services.embedding_manager._node_to_payload", return_value={}): + count = _manager().build_embeddings( + nodes, batch_size=4, progress_callback=lambda d, t: seen.append((d, t)) + ) + + assert count == 10 + assert seen, "progress_callback was never invoked" + assert [d for d, _ in seen] == sorted(d for d, _ in seen), "completion went backwards" + assert seen[-1] == (10, 10), f"final callback should report completion, got {seen[-1]}" + assert all(t == 10 for _, t in seen), "total changed mid-run" + + +def test_progress_callback_is_optional(): + """Omitting the callback must preserve the original behaviour.""" + nodes = [_node(i) for i in range(3)] + with patch("ast_rag.services.embedding_manager._node_to_payload", return_value={}): + assert _manager().build_embeddings(nodes, batch_size=2) == 3 + + +def test_progress_callback_fires_on_empty_input(): + seen: list[tuple[int, int]] = [] + assert _manager().build_embeddings([], progress_callback=lambda d, t: seen.append((d, t))) == 0 + assert seen == [(0, 0)], "callers need a terminal update even when there is nothing to embed" + + +def test_index_progress_reports_counts_and_eta(): + progress = _index_progress() + columns = {type(c).__name__ for c in progress.columns} + assert "BarColumn" in columns + assert "MofNCompleteColumn" in columns, "issue #5 asks for a file count" + assert "TimeRemainingColumn" in columns + assert progress.live.transient