An accurate Retrieval-Augmented Generation (RAG) system that analyzes multi-language codebases using Tree-sitter, builds comprehensive knowledge graphs, and enables natural language querying of codebase structure and relationships.

• 🌍 Multi-Language Support: Supports Python, JavaScript, TypeScript, Rust, and Go codebases
• 🌳 Tree-sitter Parsing: Uses Tree-sitter for robust, language-agnostic AST parsing
• 📊 Knowledge Graph Storage: Uses Memgraph to store codebase structure as an interconnected graph
• 🗣️ Natural Language Querying: Ask questions about your codebase in plain English
• 🤖 AI-Powered Cypher Generation: Leverages Google Gemini to translate natural language to Cypher queries
• 📝 Code Snippet Retrieval: Retrieves actual source code snippets for found functions/methods
• 🔗 Dependency Analysis: Parses pyproject.toml to understand external dependencies
• 🎯 Nested Function Support: Handles complex nested functions and class hierarchies
• 🔄 Language-Agnostic Design: Unified graph schema across all supported languages

The system consists of two main components:

1. Multi-language Parser: Tree-sitter based parsing system that analyzes codebases and ingests data into Memgraph
2. RAG System (codebase_rag/): Interactive CLI for querying the stored knowledge graph

• 🌳 Tree-sitter Integration: Language-agnostic parsing using Tree-sitter grammars
• 📊 Graph Database: Memgraph for storing code structure as nodes and relationships
• 🤖 LLM Integration: Google Gemini for natural language processing
• 🔍 Code Analysis: Advanced AST traversal for extracting code elements across languages
• 🛠️ Query Tools: Specialized tools for graph querying and code retrieval
• ⚙️ Language Configuration: Configurable mappings for different programming languages

• Python 3.12+
• Docker & Docker Compose (for Memgraph)
• Google Gemini API key
• uv package manager

1. Clone the repository:

git clone https://github.com/vitali87/code-graph-rag.git
cd code-graph-rag

2. Install dependencies:

For basic Python support:

uv sync

For full multi-language support:

uv sync --extra treesitter-full

This installs Tree-sitter grammars for:

• Python (.py)
• JavaScript (.js, .jsx)
• TypeScript (.ts, .tsx)
• Rust (.rs)
• Go (.go)

3. Set up environment variables:

cp .env.example .env
# Edit .env with your configuration

Required environment variables:

GEMINI_API_KEY=your-api-key
GEMINI_MODEL_ID=gemini-2.5-pro
MODEL_CYPHER_ID=gemini-2.5-flash-lite-preview-06-17
MEMGRAPH_HOST=localhost
MEMGRAPH_PORT=7687

4. Start Memgraph database:

docker-compose up -d

Parse and ingest a multi-language repository into the knowledge graph:

For the first repository (clean start):

python -m codebase_rag.main --repo-path /path/to/repo1 --update-graph --clean

For additional repositories (preserve existing data):

python -m codebase_rag.main --repo-path /path/to/repo2 --update-graph
python -m codebase_rag.main --repo-path /path/to/repo3 --update-graph

Supported Languages: The system automatically detects and processes files based on extensions:

• Python: .py files
• JavaScript: .js, .jsx files
• TypeScript: .ts, .tsx files
• Rust: .rs files
• Go: .go files

Start the interactive RAG CLI:

python -m codebase_rag.main --repo-path /path/to/your/repo

Example queries (works across all supported languages):

• "Show me all classes that contain 'user' in their name"
• "Find functions related to database operations"
• "What methods does the User class have?"
• "Show me functions that handle authentication"
• "List all TypeScript components"
• "Find Rust structs and their methods"
• "Show me Go interfaces and implementations"

The knowledge graph uses the following node types and relationships:

• Project: Root node representing the entire repository
• Package: Language packages (Python: __init__.py, etc.)
• Module: Individual source code files (.py, .js, .ts, .rs, .go)
• Class: Class/Struct/Enum definitions across all languages
• Function: Module-level functions and standalone functions
• Method: Class methods and associated functions
• Folder: Regular directories
• File: All files (source code and others)
• ExternalPackage: External dependencies

• Python: function_definition, class_definition
• JavaScript/TypeScript: function_declaration, arrow_function, class_declaration
• Rust: function_item, struct_item, enum_item, impl_item
• Go: function_declaration, method_declaration, type_declaration

• CONTAINS_PACKAGE/MODULE/FILE/FOLDER: Hierarchical containment
• DEFINES: Module defines classes/functions
• DEFINES_METHOD: Class defines methods
• DEPENDS_ON_EXTERNAL: Project depends on external packages

Configuration is managed through environment variables and the config.py file:

MEMGRAPH_HOST = "localhost"
MEMGRAPH_PORT = 7687
GEMINI_MODEL_ID = "gemini-2.5-pro"  # Main RAG orchestrator model
MODEL_CYPHER_ID = "gemini-2.5-flash-lite-preview-06-17"  # Cypher generation model
TARGET_REPO_PATH = "."
GEMINI_API_KEY = "required"

code-graph-rag/
├── codebase_rag/              # RAG system package
│   ├── main.py                # CLI entry point
│   ├── config.py              # Configuration management
│   ├── graph_updater.py       # Tree-sitter based multi-language parser
│   ├── language_config.py     # Language-specific configurations
│   ├── prompts.py             # LLM prompts and schemas
│   ├── schemas.py             # Pydantic models
│   ├── services/              # Core services
│   │   └── llm.py             # Gemini LLM integration
│   └── tools/                 # RAG tools
│       ├── codebase_query.py  # Graph querying tool
│       ├── code_retrieval.py  # Code snippet retrieval
│       └── file_reader.py     # File content reading
├── docker-compose.yaml        # Memgraph setup
├── pyproject.toml            # Project dependencies & language extras
└── README.md                 # This file

• tree-sitter: Core Tree-sitter library for language-agnostic parsing
• tree-sitter-{language}: Language-specific grammars (Python, JS, TS, Rust, Go)
• pydantic-ai: AI agent framework for RAG orchestration
• pymgclient: Memgraph Python client for graph database operations
• loguru: Advanced logging with structured output
• python-dotenv: Environment variable management

• Python: Full support including nested functions, methods, classes, and package structure
• JavaScript/TypeScript: Functions, arrow functions, classes, and method definitions
• Rust: Functions, structs, enums, impl blocks, and associated functions
• Go: Functions, methods, type declarations, and struct definitions

# Basic Python-only support
uv sync

# Full multi-language support  
uv sync --extra treesitter-full

# Individual language support (if needed)
uv add tree-sitter-python tree-sitter-javascript tree-sitter-typescript tree-sitter-rust tree-sitter-go

The system uses a configuration-driven approach for language support. Each language is defined in codebase_rag/language_config.py with:

• File extensions: Which files to process
• AST node types: How to identify functions, classes, etc.
• Module structure: How modules/packages are organized
• Name extraction: How to extract names from AST nodes

Adding support for new languages requires only configuration changes, no code modifications.

1. Check Memgraph connection:

   • Ensure Docker containers are running: docker-compose ps
   • Verify Memgraph is accessible on port 7687

2. View database in Memgraph Lab:

   • Open http://localhost:3000
   • Connect to memgraph:7687

3. Enable debug logging:

   • The RAG orchestrator runs in debug mode by default
   • Check logs for detailed execution traces

1. Follow the established code structure
2. Keep files under 100 lines (as per user rules)
3. Use type annotations
4. Follow conventional commit messages
5. Use DRY principles

For issues or questions:

1. Check the logs for error details
2. Verify Memgraph connection
3. Ensure all environment variables are set
4. Review the graph schema matches your expectations