From 1a2e63073e07fc58c4ca6d103dde3d532c68e727 Mon Sep 17 00:00:00 2001 From: Rohit Behera <126186063+r0h1tb@users.noreply.github.com> Date: Sat, 8 Aug 2026 15:50:46 +0530 Subject: [PATCH 1/4] fix(config): make a fresh clone run without editing anything. Fixes #62 The committed ast_rag_config.json pointed at four private LAN addresses (Neo4j, Qdrant, the summarizer LLM and a remote embedding server). A fresh clone could not reach any of them, and because that file shadows the built-in defaults it broke the one path that already worked: the defaults in dto/config.py have always been localhost. - untrack ast_rag_config.json and gitignore it, so the defaults apply and an existing local copy survives a pull - add ast_rag_config.example.json pointing at the documented Docker services - add docker-compose.yml so `docker compose up -d` starts exactly what the defaults expect. The docker/*.sh scripts invoke podman, which the README never mentions installing - give Neo4jConfig a connection_timeout and pass it to the driver; the driver default left the CLI apparently hung, with `ast-rag stats` against an unreachable host taking 34s before failing README: install now uses compose and needs no config file, and the language table lists Go, which shipped in #17/#55 but was never added there. Fixes #66 --- .gitignore | 1 + README.md | 31 ++++++++---- ast_rag/dto/config.py | 4 ++ ast_rag/repositories/neo4j_helpers.py | 3 ++ ...config.json => ast_rag_config.example.json | 9 ++-- docker-compose.yml | 48 +++++++++++++++++++ 6 files changed, 82 insertions(+), 14 deletions(-) rename ast_rag_config.json => ast_rag_config.example.json (82%) create mode 100644 docker-compose.yml diff --git a/.gitignore b/.gitignore index bf23a8f..4046dc6 100644 --- a/.gitignore +++ b/.gitignore @@ -58,3 +58,4 @@ config.local.json # Archive (not for publication) docs/archive/ +ast_rag_config.json diff --git a/README.md b/README.md index be8332f..0a9b3ca 100644 --- a/README.md +++ b/README.md @@ -26,26 +26,39 @@ Context-aware code intelligence system for AI agents and developers. Provides se | **Python** | `.py` | ⭐⭐ Good | Classes, functions, imports, type hints | | **TypeScript** | `.ts` | ⭐⭐ Good | Classes, interfaces, functions, imports | | **TSX / JSX** | `.tsx` `.jsx` | ⭐⭐ Good | Classes, interfaces, functions, imports, JSX elements | +| **Go** | `.go` | ⭐⭐ Good | Functions, methods, structs, interfaces, imports | Files with other extensions are skipped during indexing with a warning listing the supported languages. ## 📦 Installation ```bash -# 1. Clone repository -git clone && cd raged - -# 2. Create virtual environment +# 1. Clone and install +git clone https://github.com/lexasub/raged && cd raged python -m venv venv && source venv/bin/activate - -# 3. Install dependencies pip install -e . -# 4. Start Neo4j and Qdrant (Docker) -docker run -d --name neo4j -p 7687:7687 -p 7474:7474 neo4j:latest -docker run -d --name qdrant -p 6333:6333 qdrant/qdrant:latest +# 2. Start Neo4j and Qdrant +docker compose up -d + +# 3. Index and query — no config file needed +ast-rag index-folder ./ast_rag +ast-rag query "batch upsert nodes" +``` + +The defaults point at the services `docker compose up -d` starts, so this works +on a fresh clone with no configuration. To point elsewhere, set environment +variables rather than editing a tracked file: + +```bash +export AST_RAG_NEO4J_URI=bolt://myhost:7687 +export AST_RAG_NEO4J_PASSWORD=secret +export AST_RAG_QDRANT_URL=http://myhost:6333 ``` +Or copy `ast_rag_config.example.json` to `ast_rag_config.json` and edit it +(it is gitignored). + ## ⚙️ Configuration Create `ast_rag_config.json` in project root: diff --git a/ast_rag/dto/config.py b/ast_rag/dto/config.py index 78c3c68..ff0dd39 100644 --- a/ast_rag/dto/config.py +++ b/ast_rag/dto/config.py @@ -26,6 +26,10 @@ class Neo4jConfig(BaseModel): password: str = "password" database: str = "neo4j" project_id: str = "default" + # Kept short deliberately: the driver's default leaves the CLI apparently + # hung for tens of seconds when the service is simply not running, which is + # the usual state on a fresh clone. + connection_timeout: float = 10.0 class QdrantConfig(BaseModel): diff --git a/ast_rag/repositories/neo4j_helpers.py b/ast_rag/repositories/neo4j_helpers.py index 22a8757..a294236 100644 --- a/ast_rag/repositories/neo4j_helpers.py +++ b/ast_rag/repositories/neo4j_helpers.py @@ -45,6 +45,7 @@ def create_driver(config: Neo4jConfig, create_if_not_exists: bool = True) -> Dri config.uri, auth=(config.user, config.password), database="neo4j", + connection_timeout=config.connection_timeout, ) if create_if_not_exists and config.database != "neo4j": @@ -66,12 +67,14 @@ def create_driver(config: Neo4jConfig, create_if_not_exists: bool = True) -> Dri config.uri, auth=(config.user, config.password), database=config.database, + connection_timeout=config.connection_timeout, ) except Exception: return GraphDatabase.driver( config.uri, auth=(config.user, config.password), database="neo4j", + connection_timeout=config.connection_timeout, ) diff --git a/ast_rag_config.json b/ast_rag_config.example.json similarity index 82% rename from ast_rag_config.json rename to ast_rag_config.example.json index e97b4ae..038f45b 100644 --- a/ast_rag_config.json +++ b/ast_rag_config.example.json @@ -1,8 +1,8 @@ { "summarizer": { - "enabled": true, + "enabled": false, "llm": { - "base_url": "http://192.168.2.109:1113/v1", + "base_url": "http://localhost:11434/v1", "model": "qwen2.5-coder:14b", "api_key": "ollama", "timeout": 120, @@ -31,20 +31,19 @@ "db_path": ".ast_rag_parse_cache.sqlite" }, "neo4j": { - "uri": "bolt://192.168.2.109:7687", + "uri": "bolt://localhost:7687", "user": "neo4j", "password": "password", "connection_timeout": 30, "max_connection_pool_size": 50 }, "qdrant": { - "url": "http://192.168.2.109:6333", + "url": "http://localhost:6333", "collection_name": "ast_rag_nodes", "timeout": 60 }, "embedding": { "model_name": "bge-m3", - "remote_url": "http://192.168.2.6:1113/v1/embeddings", "dimension": 1024, "remote_batch_size": 32, "timeout": 120 diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..eb1f70c --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,48 @@ +# Services AST-RAG needs, matching the defaults in ast_rag/dto/config.py. +# +# docker compose up -d +# +# Nothing else to configure: with these running, `ast-rag` works with no +# config file at all. Override per environment with AST_RAG_NEO4J_URI / +# AST_RAG_NEO4J_PASSWORD / AST_RAG_QDRANT_URL rather than editing a file. + +services: + neo4j: + image: neo4j:5.18-community + container_name: ast_rag_neo4j + restart: unless-stopped + ports: + - "7474:7474" # HTTP browser + - "7687:7687" # Bolt + environment: + NEO4J_AUTH: neo4j/password + NEO4J_PLUGINS: '["apoc"]' + NEO4J_server_memory_heap_initial__size: 512m + NEO4J_server_memory_heap_max__size: 2G + NEO4J_server_memory_pagecache_size: 512m + volumes: + - neo4j_data:/data + - neo4j_logs:/logs + healthcheck: + # cypher-shell exits non-zero until Bolt is actually serving, which is + # later than the container being "up". + test: ["CMD-SHELL", "cypher-shell -u neo4j -p password 'RETURN 1' || exit 1"] + interval: 10s + timeout: 10s + retries: 12 + start_period: 30s + + qdrant: + image: qdrant/qdrant:latest + container_name: ast_rag_qdrant + restart: unless-stopped + ports: + - "6333:6333" # REST + - "6334:6334" # gRPC + volumes: + - qdrant_data:/qdrant/storage + +volumes: + neo4j_data: + neo4j_logs: + qdrant_data: From 1573d6277023f303f490d8472b6c6ec112946b53 Mon Sep 17 00:00:00 2001 From: Rohit Behera <126186063+r0h1tb@users.noreply.github.com> Date: Sat, 8 Aug 2026 15:50:59 +0530 Subject: [PATCH 2/4] fix(schema): repair index DDL so indexes are actually created Every index creation failed at indexing time while `index-folder` reported "Errors: 0", so the graph ran with no indexes at all. Two causes: STANDARD_INDEXES is unpacked as (label, property, name) and fed to create_index, but it also held one fulltext entry shaped (name, [labels], [properties]). Unpacked as a B-tree index that produced CREATE INDEX ['name', 'qualified_name'] IF NOT EXISTS FOR (n:ast_symbol_fulltext) ON (n.['Function', 'Class', 'Method']) which Neo4j rejects. The entry also duplicated an explicit create_fulltext_index call a few lines below. Fulltext definitions now live in their own STANDARD_FULLTEXT_INDEXES list, and the dead loop that tried to filter them back out of STANDARD_INDEXES is gone. create_fulltext_index emitted invalid Cypher on three counts: CREATE FULLTEXT INDEX IF NOT EXISTS ast_symbol_fulltext FOR ([Function:Class:Method]) ON EACH [name, qualified_name] Neo4j 5 wants the name before IF NOT EXISTS -- the same ordering rule as CREATE CONSTRAINT, which test_schema_cypher already covers at the sibling call site -- labels alternated with | on a bound variable, and qualified property references. Analyzer options also belong under indexConfig. Verified against Neo4j 5.18: SHOW INDEXES now lists ast_symbol_fulltext (FULLTEXT) plus the six RANGE indexes; indexing logs no schema errors. --- ast_rag/repositories/schema_manager.py | 63 ++++++++++++-------------- tests/test_schema_cypher.py | 48 ++++++++++++++++++++ 2 files changed, 78 insertions(+), 33 deletions(-) diff --git a/ast_rag/repositories/schema_manager.py b/ast_rag/repositories/schema_manager.py index a33d690..ae3c2cd 100644 --- a/ast_rag/repositories/schema_manager.py +++ b/ast_rag/repositories/schema_manager.py @@ -64,7 +64,16 @@ class SchemaManager: ("Function", "name", "function_name_idx"), ("Function", "signature", "function_signature_idx"), ("Class", "name", "class_name_idx"), - # Full-text indexes + ] + + # Full-text indexes are (name, labels, properties) and are created through + # create_fulltext_index. They are kept out of STANDARD_INDEXES, whose + # entries are (label, property, name) and are fed to create_index: mixing + # the two shapes in one list meant the fulltext entry was unpacked as a + # B-tree index, producing + # CREATE INDEX ['name','qualified_name'] ... FOR (n:ast_symbol_fulltext) + # which Neo4j rejects, so no index was ever created. + STANDARD_FULLTEXT_INDEXES = [ ("ast_symbol_fulltext", ["Function", "Class", "Method"], ["name", "qualified_name"]), ] @@ -275,18 +284,23 @@ def create_fulltext_index( if analyzer is None: analyzer = "standard" - label_filters = ":".join(labels) if labels else "" - properties_str = ", ".join(properties) + # Neo4j 5 syntax: the name comes *before* IF NOT EXISTS, labels are + # alternated with | on a bound variable, and ON EACH takes qualified + # property references. The previous form emitted + # CREATE FULLTEXT INDEX IF NOT EXISTS name FOR ([A:B]) ON EACH [p] + # which fails to parse on every count, so the index never existed. + label_filters = "|".join(labels) if labels else "" + properties_str = ", ".join(f"n.{p}" for p in properties) query_parts = [ "CREATE FULLTEXT INDEX", - "IF NOT EXISTS" if if_not_exists else "", f"{index_name}", + "IF NOT EXISTS" if if_not_exists else "", "FOR", - f"([{label_filters}])", + f"(n:{label_filters})", "ON EACH", f"[{properties_str}]", - f"OPTIONS {{ analyzer: '{analyzer}' }}", + f"OPTIONS {{ indexConfig: {{ `fulltext.analyzer`: '{analyzer}' }} }}", ] query = " ".join(part for part in query_parts if part) @@ -437,33 +451,16 @@ def create_standard_indexes(self) -> dict[str, Any]: stats["errors"].append(f"{index_name}: {exc}") logger.error("Failed to create standard index %s: %s", index_name, exc) - # Create fulltext indexes - for index_name, labels, properties in [ - idx for idx in self.STANDARD_INDEXES if isinstance(idx[1], list) - ]: - # Actually, STANDARD_INDEXES doesn't have fulltext in the same format - # Let's handle the fulltext index separately - pass - - # Handle the fulltext index from STANDARD_INDEXES - # We know the fulltext index is defined as: - # ("ast_symbol_fulltext", ["Function", "Class", "Method"], ["name", "qualified_name"]) - # But our STANDARD_INDEXES structure is different for fulltext - # Let's just create it directly - try: - if self.create_fulltext_index( - "ast_symbol_fulltext", - ["Function", "Class", "Method"], - ["name", "qualified_name"], - if_not_exists=True, - ): - stats["created"] += 1 - else: - stats["skipped"] += 1 - except Exception as exc: - stats["failed"] += 1 - stats["errors"].append(f"ast_symbol_fulltext: {exc}") - logger.error("Failed to create standard fulltext index: %s", exc) + for index_name, labels, properties in self.STANDARD_FULLTEXT_INDEXES: + try: + if self.create_fulltext_index(index_name, labels, properties, if_not_exists=True): + stats["created"] += 1 + else: + stats["skipped"] += 1 + except Exception as exc: + stats["failed"] += 1 + stats["errors"].append(f"{index_name}: {exc}") + logger.error("Failed to create standard fulltext index %s: %s", index_name, exc) logger.info( "Standard indexes: %d created, %d skipped, %d failed", diff --git a/tests/test_schema_cypher.py b/tests/test_schema_cypher.py index 733c3fc..e332a73 100644 --- a/tests/test_schema_cypher.py +++ b/tests/test_schema_cypher.py @@ -72,3 +72,51 @@ def test_constraint_name_not_immediately_after_if_not_exists(): query = _captured_query(fn, **kwargs) assert "IF NOT EXISTS c_name" not in query assert "IF NOT EXISTS i_name" not in query + + +def test_fulltext_index_name_precedes_if_not_exists(): + """Same ordering rule as CREATE CONSTRAINT, at the sibling call site. + + The fulltext builder emitted ``CREATE FULLTEXT INDEX IF NOT EXISTS ``, + which Neo4j 5 rejects, so the symbol fulltext index was never created while + indexing reported success. + """ + query = _captured_query( + "create_fulltext_index", + index_name="ast_symbol_fulltext", + labels=["Function", "Class", "Method"], + properties=["name", "qualified_name"], + ) + assert re.search(r"CREATE FULLTEXT INDEX\s+ast_symbol_fulltext\s+IF NOT EXISTS", query), query + + +def test_fulltext_index_uses_label_alternation_and_qualified_properties(): + """``FOR ([A:B])`` and bare ``ON EACH [p]`` are both invalid Cypher.""" + query = _captured_query( + "create_fulltext_index", + index_name="ast_symbol_fulltext", + labels=["Function", "Class", "Method"], + properties=["name", "qualified_name"], + ) + assert "FOR (n:Function|Class|Method)" in query, query + assert "ON EACH [n.name, n.qualified_name]" in query, query + assert "([" not in query, f"label filter still bracketed: {query}" + + +def test_standard_indexes_are_all_btree_shaped(): + """STANDARD_INDEXES is unpacked as (label, property, name) and fed to + create_index. A fulltext entry (name, [labels], [properties]) in that list + produced CREATE INDEX ['name','qualified_name'] FOR (n:ast_symbol_fulltext). + """ + from ast_rag.repositories.schema_manager import SchemaManager + + for entry in SchemaManager.STANDARD_INDEXES: + label, property_name, index_name = entry + assert isinstance(label, str), entry + assert isinstance(property_name, str), entry + assert isinstance(index_name, str), entry + + for index_name, labels, properties in SchemaManager.STANDARD_FULLTEXT_INDEXES: + assert isinstance(index_name, str) + assert isinstance(labels, list) and labels + assert isinstance(properties, list) and properties From 0206a994d7c65e2ad9499b8d64eecb6a8f6b1bf1 Mon Sep 17 00:00:00 2001 From: Rohit Behera <126186063+r0h1tb@users.noreply.github.com> Date: Sat, 8 Aug 2026 15:51:13 +0530 Subject: [PATCH 3/4] fix(index): resolve cross-file references in index-folder `init` indexes in two phases: parse every file, collect a project-wide name -> id map, then resolve edges against it. `index-folder` did not -- it parsed and extracted edges together, per file, in the worker process, so the resolver only ever saw the current file's symbols. The comment in `init` states the consequence exactly: "when that map only holds the current file's nodes, any reference to a symbol defined elsewhere is silently dropped". Indexing this repo with `index-folder` produced 330 CALLS edges and zero crossing a file boundary, so `callers`, `refs` and `call-graph` returned nothing for any symbol used from another module -- the common case. AGENTS.md points agents at `index-folder`. index-folder now runs the same two phases. Phase 1 extracts nodes and builds the symbol map; phase 2 re-parses and resolves edges against it. Trees are not picklable across processes, so phase 2 re-parses rather than carrying them over; the map is published through a ProcessPoolExecutor initializer so it is pickled once per worker instead of once per file. Indexing ./ast_rag before and after: edges 1,488 -> 2,223 cross-file CALLS 0 -> 376 `ast-rag callers create_driver` returned "No callers found" before and now lists its callers across mcp/server.py and services/watcher_service.py. Cost is a second parse pass: 6s -> 11s for 62 files. Also drops a hardcoded /home/su/src/local/raged fallback from the worker sys.path setup. --- ast_rag/cli.py | 304 ++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 235 insertions(+), 69 deletions(-) diff --git a/ast_rag/cli.py b/ast_rag/cli.py index 0346221..cabd8a6 100644 --- a/ast_rag/cli.py +++ b/ast_rag/cli.py @@ -24,10 +24,12 @@ import json import logging import os +import time from pathlib import Path from typing import Optional import typer +from neo4j.exceptions import AuthError from rich.console import Console from rich.progress import ( BarColumn, @@ -102,22 +104,93 @@ def callback(value: bool) -> bool: def _load_config(config_path: Optional[str] = None) -> ProjectConfig: - """Load project config from JSON file or return defaults.""" + """Load project config from JSON file or return defaults. + + Resolution order, lowest priority first: built-in defaults (which point at + localhost, so a fresh clone works against the documented Docker setup with + no config file at all), then ``ast_rag_config.json`` in the CWD, then an + explicit ``--config``, then ``AST_RAG_*`` environment variables. + + The env layer is last so that a container or CI job can point the same + checkout at different services without editing a tracked file. + """ if config_path and Path(config_path).exists(): - return ProjectConfig.model_validate_json(Path(config_path).read_text()) - # Check for ast_rag_config.json in CWD - default = Path("ast_rag_config.json") - if default.exists(): - return ProjectConfig.model_validate_json(default.read_text()) - return ProjectConfig() + cfg = ProjectConfig.model_validate_json(Path(config_path).read_text()) + else: + default = Path("ast_rag_config.json") + if default.exists(): + cfg = ProjectConfig.model_validate_json(default.read_text()) + else: + cfg = ProjectConfig() + return _apply_env_overrides(cfg) -def _build_api(cfg: ProjectConfig) -> ASTRagAPI: + +def _apply_env_overrides(cfg: ProjectConfig) -> ProjectConfig: + """Overlay ``AST_RAG_*`` environment variables onto a loaded config.""" + env_map = { + "AST_RAG_NEO4J_URI": (cfg.neo4j, "uri"), + "AST_RAG_NEO4J_USER": (cfg.neo4j, "user"), + "AST_RAG_NEO4J_PASSWORD": (cfg.neo4j, "password"), + "AST_RAG_NEO4J_DATABASE": (cfg.neo4j, "database"), + "AST_RAG_QDRANT_URL": (cfg.qdrant, "url"), + "AST_RAG_QDRANT_COLLECTION": (cfg.qdrant, "collection_name"), + } + for var, (section, field) in env_map.items(): + value = os.environ.get(var) + if value: + setattr(section, field, value) + return cfg + + +def _connect(cfg: ProjectConfig): + """Create a Neo4j driver and fail fast with a usable message if it is down. + + Every command needs a driver, so the reachability check lives here rather + than being repeated at each call site. + """ driver = create_driver(cfg.neo4j) + _verify_neo4j(driver, cfg) + return driver + + +def _build_api(cfg: ProjectConfig) -> ASTRagAPI: + driver = _connect(cfg) embed = EmbeddingManager(cfg.qdrant, cfg.embedding, neo4j_driver=driver) return ASTRagAPI(driver, embed) +def _verify_neo4j(driver, cfg: ProjectConfig) -> None: + """Fail fast with a usable message when Neo4j is not reachable. + + Without this the first query blocks for the driver's default connection + timeout and then raises ServiceUnavailable, which Typer renders as a full + traceback. For anyone who has not started the services yet -- the common + case on a fresh clone -- that is a wall of Bolt internals rather than + "start Neo4j". + """ + try: + driver.verify_connectivity() + except AuthError: + console.print( + f"[red]Neo4j rejected the credentials for[/red] {cfg.neo4j.uri} " + f"(user: {cfg.neo4j.user}).\n" + "Set AST_RAG_NEO4J_USER / AST_RAG_NEO4J_PASSWORD, or edit " + "ast_rag_config.json." + ) + raise typer.Exit(code=1) + except Exception as exc: # ServiceUnavailable and friends + console.print( + f"[red]Cannot reach Neo4j at[/red] {cfg.neo4j.uri}\n\n" + "Start the services:\n" + " [cyan]docker compose up -d[/cyan]\n\n" + "Or point AST-RAG somewhere else:\n" + " [cyan]export AST_RAG_NEO4J_URI=bolt://host:7687[/cyan]\n\n" + f"[dim]{type(exc).__name__}: {str(exc).splitlines()[0]}[/dim]" + ) + raise typer.Exit(code=1) + + # --------------------------------------------------------------------------- # init command # --------------------------------------------------------------------------- @@ -154,7 +227,7 @@ def init( console.rule(f"[bold blue]AST-RAG init[/bold blue]: {root}") # 1. Apply schema - driver = create_driver(cfg.neo4j) + driver = _connect(cfg) with console.status("Applying Neo4j schema..."): apply_schema(driver) @@ -351,7 +424,7 @@ def update( logging.basicConfig(level=logging.WARNING) cfg = _load_config(config) - driver = create_driver(cfg.neo4j) + driver = _connect(cfg) embed = EmbeddingManager(cfg.qdrant, cfg.embedding, neo4j_driver=driver) console.rule(f"[bold blue]AST-RAG update[/bold blue] {from_commit[:8]}..{to_commit[:8]}") @@ -502,7 +575,7 @@ def refs( ) -> None: """Find all references/usages of a symbol.""" cfg = _load_config(config) - driver = create_driver(cfg.neo4j) + driver = _connect(cfg) embed = EmbeddingManager(cfg.qdrant, cfg.embedding, neo4j_driver=driver) api = ASTRagAPI(driver, embed) @@ -697,7 +770,7 @@ def workspace( logging.basicConfig(level=logging.WARNING) cfg = _load_config(config) - driver = create_driver(cfg.neo4j) + driver = _connect(cfg) root = os.path.abspath(path) console.rule(f"[bold blue]AST-RAG workspace[/bold blue]: {root}") @@ -774,7 +847,6 @@ def evaluate( from pathlib import Path from ast_rag.api import ASTRagAPI - from ast_rag.repositories import create_driver from ast_rag.services.embedding_manager import EmbeddingManager # Load configuration @@ -784,7 +856,7 @@ def evaluate( console.rule("[bold blue]AST-RAG QUALITY EVALUATION[/bold blue]") console.print("[yellow]Initializing Neo4j and EmbeddingManager...[/yellow]") - driver = create_driver(cfg.neo4j) + driver = _connect(cfg) embed = EmbeddingManager(cfg.qdrant, cfg.embedding, neo4j_driver=driver) api = ASTRagAPI(driver, embed) @@ -1071,16 +1143,44 @@ def signature_search( # --------------------------------------------------------------------------- -def _parse_file_for_multiprocessing(args): - """Parse file for multiprocessing - must be at module level for pickle.""" - import sys +def _print_batch_progress(done, total, n_nodes, n_edges, batch_start, start_time): + """One progress line per batch, shared by both indexing phases.""" + elapsed = time.time() - start_time + rate = done / elapsed if elapsed > 0 else 0 + console.print( + f"[cyan][{done:>6}/{total}][/cyan] " + f"+{n_nodes:>4} nodes, +{n_edges:>4} edges | " + f"{rate:>5.1f} files/s | " + f"Batch: {time.time() - batch_start:.2f}s" + ) + + +# Set once per worker process by _init_symbol_worker so the project-wide symbol +# map is pickled per worker rather than per file. +_WORKER_SYMBOLS: dict[str, str] = {} + + +def _init_symbol_worker(symbols: dict[str, str]) -> None: + """ProcessPoolExecutor initializer: publish the symbol map to this worker.""" + global _WORKER_SYMBOLS + _WORKER_SYMBOLS = symbols + + +def _worker_sys_path() -> None: + """Make ast_rag importable in a subprocess started with spawn.""" import os + import sys - # Add project root to path for subprocess - project_root = os.environ.get("AST_RAG_PROJECT_ROOT", "/home/su/src/local/raged") - if project_root not in sys.path: + project_root = os.environ.get("AST_RAG_PROJECT_ROOT") + if project_root and project_root not in sys.path: sys.path.insert(0, project_root) + +def _parse_nodes_for_multiprocessing(args): + """Phase 1: extract nodes only. Must be module level for pickle.""" + import os + + _worker_sys_path() file_path, lang, source = args commit = os.environ.get("AST_RAG_COMMIT", "INIT") project_id = os.environ.get("AST_RAG_PROJECT_ID", "default") @@ -1090,10 +1190,37 @@ def _parse_file_for_multiprocessing(args): pm = ParserManager(project_id=project_id) tree = pm.parse_file(file_path, source=source) if tree is None: - return (file_path, [], []) + return (file_path, [], None) nodes = pm.extract_nodes(tree, file_path, lang, source, commit) - edges = pm.extract_edges(tree, nodes, file_path, lang, source, commit) - return (file_path, nodes, edges) + return (file_path, nodes, None) + except Exception as e: + return (file_path, [], str(e)) + + +def _parse_edges_for_multiprocessing(args): + """Phase 2: extract edges, resolving against the project-wide symbol map. + + The tree is re-parsed rather than carried over from phase 1 because + tree-sitter trees are not picklable across processes. + """ + import os + + _worker_sys_path() + file_path, lang, source = args + commit = os.environ.get("AST_RAG_COMMIT", "INIT") + project_id = os.environ.get("AST_RAG_PROJECT_ID", "default") + try: + from ast_rag.services.parsing.parser_manager import ParserManager + + pm = ParserManager(project_id=project_id) + tree = pm.parse_file(file_path, source=source) + if tree is None: + return (file_path, [], None) + nodes = pm.extract_nodes(tree, file_path, lang, source, commit) + edges = pm.extract_edges( + tree, nodes, file_path, lang, source, commit, global_symbols=_WORKER_SYMBOLS + ) + return (file_path, edges, None) except Exception as e: return (file_path, [], str(e)) @@ -1132,7 +1259,7 @@ def index_folder( os.environ["AST_RAG_PROJECT_ROOT"] = str(Path(__file__).parent.parent.parent) from ast_rag.services.parsing.parser_manager import EXT_TO_LANG - from ast_rag.repositories import create_driver, apply_schema + from ast_rag.repositories import apply_schema from ast_rag.services.graph_updater_service import ( _nodes_to_batch_by_label, batch_upsert_nodes, @@ -1156,7 +1283,7 @@ def index_folder( # Connect to Neo4j console.print("[yellow]Connecting to Neo4j...[/yellow]") - driver = create_driver(cfg.neo4j) + driver = _connect(cfg) if not no_schema: console.print("[yellow]Applying schema...[/yellow]") apply_schema(driver) @@ -1206,69 +1333,108 @@ def index_folder( errors = 0 start_time = time.time() - for i in range(0, len(all_files), batch_size): - batch = all_files[i : i + batch_size] - batch_start = time.time() + # Two-phase index, matching `init`. Edge resolution matches references + # against a name -> id map; when that map holds only the current file's + # nodes, every reference to a symbol defined elsewhere is silently dropped. + # Indexing this repo that way produced 330 CALLS edges and *zero* of them + # crossing a file boundary. Collecting all symbols first is what lets + # cross-file references link at all. See #56. + sources: dict[str, bytes] = {} + for fp, lang in all_files: + try: + with open(fp, "rb") as f: + sources[fp] = f.read() + except Exception: + errors += 1 - # Read files - files_with_source = [] - for fp, lang in batch: - try: - with open(fp, "rb") as f: - source = f.read() - files_with_source.append((fp, lang, source)) - except Exception: - errors += 1 + indexable = [(fp, lang) for fp, lang in all_files if fp in sources] + global_symbols: dict[str, str] = {} - # Parse in parallel - parsed = [] - with ProcessPoolExecutor(max_workers=workers) as executor: - futures = [ - executor.submit(_parse_file_for_multiprocessing, args) for args in files_with_source - ] - for future in as_completed(futures): - parsed.append(future.result()) + # ---- Phase 1: nodes, and the project-wide symbol map ------------------- + console.print("[yellow]Phase 1/2: extracting symbols...[/yellow]") + for i in range(0, len(indexable), batch_size): + batch = indexable[i : i + batch_size] + batch_start = time.time() + args_list = [(fp, lang, sources[fp]) for fp, lang in batch] - # Collect nodes and edges batch_nodes = [] - batch_edges = [] - for fp, nodes, edges_or_error in parsed: - if isinstance(edges_or_error, str): - errors += 1 - else: - batch_nodes.extend(nodes) - batch_edges.extend(edges_or_error) + with ProcessPoolExecutor(max_workers=workers) as executor: + for future in as_completed( + executor.submit(_parse_nodes_for_multiprocessing, a) for a in args_list + ): + fp, nodes, error = future.result() + if error: + errors += 1 + else: + batch_nodes.extend(nodes) + + # First definition of a name wins; files are walked in a stable order + # so the choice is deterministic across runs. + for node in batch_nodes: + global_symbols.setdefault(node.name, node.id) - # Insert to Neo4j if batch_nodes: try: with driver.session() as session: by_label = _nodes_to_batch_by_label(batch_nodes) for label, props_list in by_label.items(): batch_upsert_nodes(session, {label: props_list}) + total_nodes += len(batch_nodes) + except Exception as e: + console.print(f"[red]Neo4j error: {e}[/red]") + errors += 1 - all_edge_dicts = [e.to_neo4j_props() for e in batch_edges] - batch_upsert_edges(session, all_edge_dicts) + _print_batch_progress( + min(i + batch_size, len(indexable)), + len(indexable), + len(batch_nodes), + 0, + batch_start, + start_time, + ) - total_nodes += len(batch_nodes) + # ---- Phase 2: edges, resolved against every symbol --------------------- + console.print( + f"[yellow]Phase 2/2: resolving edges against {len(global_symbols)} symbols...[/yellow]" + ) + for i in range(0, len(indexable), batch_size): + batch = indexable[i : i + batch_size] + batch_start = time.time() + args_list = [(fp, lang, sources[fp]) for fp, lang in batch] + + batch_edges = [] + with ProcessPoolExecutor( + max_workers=workers, initializer=_init_symbol_worker, initargs=(global_symbols,) + ) as executor: + for future in as_completed( + executor.submit(_parse_edges_for_multiprocessing, a) for a in args_list + ): + fp, edges, error = future.result() + if error: + errors += 1 + else: + batch_edges.extend(edges) + + if batch_edges: + try: + with driver.session() as session: + batch_upsert_edges(session, [e.to_neo4j_props() for e in batch_edges]) total_edges += len(batch_edges) except Exception as e: console.print(f"[red]Neo4j error: {e}[/red]") errors += 1 - # Progress - elapsed = time.time() - start_time - files_done = min(i + batch_size, len(all_files)) - files_per_sec = files_done / elapsed if elapsed > 0 else 0 - - console.print( - f"[cyan][{files_done:>6}/{len(all_files)}][/cyan] " - f"+{len(batch_nodes):>4} nodes, +{len(batch_edges):>4} edges | " - f"{files_per_sec:>5.1f} files/s | " - f"Batch: {time.time() - batch_start:.2f}s" + _print_batch_progress( + min(i + batch_size, len(indexable)), + len(indexable), + 0, + len(batch_edges), + batch_start, + start_time, ) total_time = time.time() - start_time + files_done = len(indexable) console.rule("[bold green]FOLDER COMPLETE[/bold green]") console.print(f"[bold]Time:[/bold] {total_time / 60:.1f} minutes") @@ -1756,7 +1922,7 @@ def analyze_stacktrace( cfg = _load_config(config) try: - driver = create_driver(cfg.neo4j) + driver = _connect(cfg) embed = EmbeddingManager(cfg.qdrant, cfg.embedding, neo4j_driver=driver) from ast_rag.stack_trace import StackTraceService @@ -1939,7 +2105,7 @@ def stats( logging.basicConfig(level=logging.WARNING) cfg = _load_config(config) - driver = create_driver(cfg.neo4j) + driver = _connect(cfg) try: data = _collect_index_stats(driver) finally: From acc447acc1a3535173cd286d1c8f3884c7ebcb6c Mon Sep 17 00:00:00 2001 From: Rohit Behera <126186063+r0h1tb@users.noreply.github.com> Date: Sun, 9 Aug 2026 19:19:02 +0530 Subject: [PATCH 4/4] fix(index): extract code blocks in index-folder `init` extracts blocks for Python and Rust files and stores them with their CONTAINS_BLOCK edges. `index-folder` never did, so an index built that way contained no Block nodes at all and both commands that read them returned nothing: $ ast-rag blocks _verify_neo4j No blocks found. $ ast-rag lambdas No lambdas found. AGENTS.md points agents at `index-folder`, so this was the common path. Phase 2 already re-parses each file and holds its nodes, which is everything extract_blocks needs, so blocks are collected there rather than in a third pass. Indexing ./ast_rag now yields 1,047 blocks (620 if, 180 for, 149 with, 82 try, 6 lambda, 4 while) where it previously yielded zero, and both commands return real results. Same class as the cross-file symbol table in the previous commit: work that `init` does and `index-folder` silently skipped. --- ast_rag/cli.py | 43 ++++++++++++++++++++++++++++++++++--------- 1 file changed, 34 insertions(+), 9 deletions(-) diff --git a/ast_rag/cli.py b/ast_rag/cli.py index cabd8a6..9c5362e 100644 --- a/ast_rag/cli.py +++ b/ast_rag/cli.py @@ -1198,10 +1198,12 @@ def _parse_nodes_for_multiprocessing(args): def _parse_edges_for_multiprocessing(args): - """Phase 2: extract edges, resolving against the project-wide symbol map. + """Phase 2: extract edges and blocks, resolving against the symbol map. The tree is re-parsed rather than carried over from phase 1 because - tree-sitter trees are not picklable across processes. + tree-sitter trees are not picklable across processes. Blocks are extracted + here too rather than in a third pass, since this phase already holds both + the tree and the file's nodes, which is everything extract_blocks needs. """ import os @@ -1215,14 +1217,18 @@ def _parse_edges_for_multiprocessing(args): pm = ParserManager(project_id=project_id) tree = pm.parse_file(file_path, source=source) if tree is None: - return (file_path, [], None) + return (file_path, [], [], [], None) nodes = pm.extract_nodes(tree, file_path, lang, source, commit) edges = pm.extract_edges( tree, nodes, file_path, lang, source, commit, global_symbols=_WORKER_SYMBOLS ) - return (file_path, edges, None) + blocks: list = [] + block_edges: list = [] + if lang in ("python", "rust"): + blocks, block_edges = pm.extract_blocks(tree, nodes, file_path, lang, source, commit) + return (file_path, edges, blocks, block_edges, None) except Exception as e: - return (file_path, [], str(e)) + return (file_path, [], [], [], str(e)) @app.command() @@ -1264,10 +1270,15 @@ def index_folder( _nodes_to_batch_by_label, batch_upsert_nodes, batch_upsert_edges, + batch_upsert_blocks, + batch_upsert_block_edges, ) cfg = _load_config(config) folder_path = Path(path).resolve() + # Workers read this from the environment; keep the parent in step so + # blocks are stamped with the same commit as the nodes and edges. + commit = os.environ.get("AST_RAG_COMMIT", "INIT") if not folder_path.exists(): console.print(f"[red]Error: Folder not found: {folder_path}[/red]") @@ -1330,6 +1341,7 @@ def index_folder( # Index in batches total_nodes = 0 total_edges = 0 + total_blocks = 0 errors = 0 start_time = time.time() @@ -1403,23 +1415,35 @@ def index_folder( args_list = [(fp, lang, sources[fp]) for fp, lang in batch] batch_edges = [] + batch_blocks = [] + batch_block_edges = [] with ProcessPoolExecutor( max_workers=workers, initializer=_init_symbol_worker, initargs=(global_symbols,) ) as executor: for future in as_completed( executor.submit(_parse_edges_for_multiprocessing, a) for a in args_list ): - fp, edges, error = future.result() + fp, edges, blocks, block_edges, error = future.result() if error: errors += 1 else: batch_edges.extend(edges) + batch_blocks.extend(blocks) + batch_block_edges.extend(block_edges) - if batch_edges: + if batch_edges or batch_blocks: try: with driver.session() as session: - batch_upsert_edges(session, [e.to_neo4j_props() for e in batch_edges]) - total_edges += len(batch_edges) + if batch_edges: + batch_upsert_edges(session, [e.to_neo4j_props() for e in batch_edges]) + total_edges += len(batch_edges) + # `init` stores blocks; index-folder never did, so + # `ast-rag blocks` and `ast-rag lambdas` returned nothing + # for any index built this way. + if batch_blocks: + batch_upsert_blocks(session, batch_blocks, commit) + batch_upsert_block_edges(session, batch_block_edges) + total_blocks += len(batch_blocks) except Exception as e: console.print(f"[red]Neo4j error: {e}[/red]") errors += 1 @@ -1442,6 +1466,7 @@ def index_folder( console.print(f"[bold]Speed:[/bold] {files_done / total_time:.1f} files/s") console.print(f"[bold]Nodes:[/bold] {total_nodes:,}") console.print(f"[bold]Edges:[/bold] {total_edges:,}") + console.print(f"[bold]Blocks:[/bold] {total_blocks:,}") console.print(f"[bold]Errors:[/bold] {errors}") driver.close()