Multi-Language Parser Architecture with Tree-sitter for AST Extraction in code-graph-rag

The code-graph-rag project implements a language-agnostic AST extraction pipeline using Tree-sitter, built around a three-layer architecture: a centralized language registry (TreeSitterModule), a dynamic parser loader with grammar caching, and a unified node protocol that abstracts Tree-sitter specifics.

The vitali87/code-graph-rag repository solves the problem of parsing dozens of programming languages without duplicating analysis logic. At its core lies a multi-language parser architecture with Tree-sitter for AST extraction that decouples grammar specifics from graph-building algorithms. This design lets contributors add new languages by simply registering a grammar and implementing a thin adapter—no changes to the core graph engine required.

---

Core Components of the Parser Architecture

The architecture separates concerns into three well-defined layers. Understanding each layer is essential for extending the system or debugging parser behavior.

Tree-sitter Module Registry

Language support is declared centrally in codebase_rag/constants/languages.py. The TreeSitterModule enum acts as the single source of truth for which languages the system understands and which grammars map to each.

# From codebase_rag/constants/languages.py (simplified)

class TreeSitterModule(Enum):
    PYTHON = "python"
    JAVASCRIPT = "javascript"
    TYPESCRIPT = "typescript"
    RUST = "rust"
    GO = "go"
    # ... additional languages

Each variant carries metadata linking it to a compiled grammar file (.so or .dll) stored under the grammars/ directory. This eliminates scattered string constants and enables IDE autocompletion when selecting parsers.

Parser Loader with Grammar Caching

The codebase_rag/parser_loader.py module handles runtime grammar selection and instantiation. When load_parser() receives a file path, it:

1. Detects the language via extension or shebang analysis
2. Maps the language to its TreeSitterModule entry
3. Loads the compiled grammar (with LRU caching to avoid repeated disk I/O)
4. Returns a TreeSitterLanguage instance ready for parsing

from codebase_rag.parser_loader import load_parser

# Auto-detects .py extension → TreeSitterModule.PYTHON

parser = load_parser("src/models.py")

Grammar caching matters for performance: re-parsing thousands of files in a large repository would be prohibitively slow without reusing compiled Tree-sitter parsers. The implementation at line 104 of parser_loader.py maintains this cache keyed by grammar path.

Unified AST Interface via Protocol

Raw Tree-sitter nodes expose language-specific type names and inconsistent method signatures. The TreeSitterNodeProtocol in codebase_rag/types_defs.py (line 191) solves this by defining a language-neutral contract:

# Conceptual protocol definition from codebase_rag/types_defs.py

class TreeSitterNodeProtocol:
    def type(self) -> str: ...
    def text(self) -> str: ...
    def children(self) -> List["TreeSitterNodeProtocol"]: ...
    def start_point(self) -> Tuple[int, int]: ...
    def end_point(self) -> Tuple[int, int]: ...

Helper utilities in codebase_rag/parsers/utils.py wrap raw Tree-sitter nodes to satisfy this protocol. Line 23 and surrounding code handle edge cases like multi-byte character encoding and malformed AST regions.

---

End-to-End AST Extraction Pipeline

The complete flow from file discovery to queryable graph proceeds through four stages:

Stage 1: Repository Scanning and Language Detection

codebase_rag/graph_loader.GraphLoader orchestrates the process. It walks the repository tree, classifies each file by language, and routes it to the appropriate parser.

Stage 2: Parser Instantiation

For each file, the loader delegates to parser_loader.py, which selects and caches the correct Tree-sitter grammar. This stage is transparent to calling code.

Stage 3: AST Parsing and Normalization

The parser generates a raw Tree-sitter tree. Utility functions in parsers/utils.py recursively traverse this tree, wrapping each node in the protocol interface. The result is a language-agnostic AST that downstream components can traverse without importing Tree-sitter types.

Stage 4: Language-Specific Analysis

Each language package (e.g., parsers/python/, parsers/rust/) implements analysis passes using only the protocol API. For example, parsers/python/type_inference.py builds Python-specific type information while remaining decoupled from the core graph engine. This separation means:

• Adding Python features requires changes only in parsers/python/
• Adding a new language requires no changes to existing language packages
• Cross-language queries work uniformly through the protocol abstraction

---

Practical Code Examples

Load and Traverse a Single File

from codebase_rag.parser_loader import load_parser
from pathlib import Path

# 1. Load parser for detected language

parser = load_parser("example.py")  # Returns TreeSitterLanguage

# 2. Read and parse source

source = Path("example.py").read_text(encoding="utf-8")
root = parser.parse(source)  # Returns TreeSitterNodeProtocol

# 3. Language-neutral tree traversal

def walk(node, depth=0):
    indent = "  " * depth
    start = node.start_point()
    end = node.end_point()
    print(f"{indent}{node.type()} [{start}–{end}]")
    for child in node.children():
        walk(child, depth + 1)

walk(root)

This example demonstrates the complete pipeline: automatic language detection, grammar loading, parsing, and protocol-wrapped traversal. The walk() function works identically for .py, .rs, .js, or any supported extension.

Build a Cross-Language Code Graph

from codebase_rag.graph_loader import GraphLoader

loader = GraphLoader(root_path="my_monorepo/")
graph = loader.build_graph()  # Parses all files, normalizes AST nodes

# Query functions across all languages

functions = [
    node for node in graph.nodes
    if node.kind == "function_definition"  # Unified kind name

]

print(f"Found {len(functions)} functions across {len(set(f.language for f in functions))} languages")

The GraphLoader internally orchestrates parser selection per file. The resulting graph uses normalized kind values—function_definition covers Python def, JavaScript function, Rust fn, and equivalent constructs—enabling truly language-agnostic analysis.

---

Adding a New Language: Step-by-Step

The architecture minimizes friction for language extensions. Follow this pattern to add support for a new language like Swift:

1. Register the language: Add SWIFT = "swift" to TreeSitterModule in constants/languages.py

2. Provide the grammar: Compile the Tree-sitter Swift grammar and place the shared object at grammars/swift.so

3. Implement adapter utilities (if needed): Create parsers/swift/ with analysis passes that consume TreeSitterNodeProtocol

No changes to graph_loader.py, parser_loader.py, or existing language packages are required.

---

Key Files and Responsibilities

File	Purpose
codebase_rag/constants/languages.py	TreeSitterModule enum—master language registry
codebase_rag/parser_loader.py	Dynamic grammar loading and caching
codebase_rag/types_defs.py	TreeSitterNodeProtocol definition
codebase_rag/parsers/utils.py	Protocol wrapping and traversal utilities
codebase_rag/graph_loader.py	Repository scanning and orchestration
codebase_rag/parsers/<language>/	Language-specific analysis modules

---

Summary

• Centralized registry: TreeSitterModule in languages.py declares all supported grammars
• Efficient loading: parser_loader.py caches compiled grammars to avoid redundant disk I/O
• Language abstraction: TreeSitterNodeProtocol in types_defs.py decouples analysis logic from Tree-sitter internals
• Extensible design: New languages require only registry entry + grammar file + optional adapter
• Unified queries: Cross-language analysis works through normalized node kinds and protocol methods

---

Frequently Asked Questions

How does code-graph-rag detect which Tree-sitter grammar to use for a file?

GraphLoader examines file extensions and shebang lines, then maps the detected language to its TreeSitterModule entry. parser_loader.py uses this enum value to locate and load the corresponding compiled grammar. The mapping is explicit in constants/languages.py, avoiding brittle string matching scattered through the codebase.

What performance optimizations exist for parsing large repositories?

Grammar caching in parser_loader.py prevents reloading the same .so files repeatedly across thousands of source files. Parsed ASTs are processed streaming where possible, and the protocol wrapper in parsers/utils.py avoids expensive deep copies of node data. For incremental builds, future enhancements could cache serialized ASTs keyed by file hash.

Can I use the Tree-sitter protocol abstraction outside of graph construction?

Yes. The TreeSitterNodeProtocol and utility functions in parsers/utils.py are designed as standalone components. Import load_parser and the protocol types directly for custom analysis tools, linting, or code transformation pipelines that need multi-language parsing without the full graph overhead.

Why use Tree-sitter instead of language-specific parsers like libclang or typescript?

Tree-sitter provides consistent APIs across all languages, incremental parsing for editor-like use cases, and maintained open-source grammars for 50+ languages. While language-specific tools may offer richer semantics (e.g., full type resolution), Tree-sitter hits the sweet spot for structural analysis—precise enough for call graphs and symbol extraction, fast enough for interactive tools, and uniform enough to scale across a polyglot codebase.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}