Graph-Code: A Graph-Based RAG System for Any Codebases

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 as well as editing capabilities.

combined.mp4

🛠️ Makefile Updates

Use the Makefile for:

• make install: Install project dependencies with full language support.
• make python: Install dependencies for Python only.
• make dev: Setup dev environment (install deps + pre-commit hooks).
• make test: Run all tests.
• make clean: Clean up build artifacts and cache.
• make help: Show available commands.

🚀 Features

• 🌍 Multi-Language Support:

  Language	Status	Extensions	Functions	Classes/Structs	Modules	Package Detection	Additional Features
  ✅ Python	Fully Supported	.py	✅	✅	✅	__init__.py	Type inference, decorators, nested functions
  ✅ JavaScript	Fully Supported	.js, .jsx	✅	✅	✅	-	ES6 modules, CommonJS, prototype methods, object methods, arrow functions
  ✅ TypeScript	Fully Supported	.ts, .tsx	✅	✅	✅	-	Interfaces, type aliases, enums, namespaces, ES6/CommonJS modules
  ✅ C++	Fully Supported	.cpp, .h, .hpp, .cc, .cxx, .hxx, .hh, .ixx, .cppm, .ccm	✅	✅ (classes/structs/unions/enums)	✅	CMakeLists.txt, Makefile	Constructors, destructors, operator overloading, templates, lambdas, C++20 modules, namespaces
  ✅ Lua	Fully Supported	.lua	✅	✅ (tables/modules)	✅	-	Local/global functions, metatables, closures, coroutines
  ✅ Rust	Fully Supported	.rs	✅	✅ (structs/enums)	✅	-	impl blocks, associated functions
  ✅ Java	Fully Supported	.java	✅	✅ (classes/interfaces/enums)	✅	package declarations	Generics, annotations, modern features (records/sealed classes), concurrency, reflection
  🚧 Go	In Development	.go	✅	✅ (structs)	✅	-	Methods, type declarations
  🚧 Scala	In Development	.scala, .sc	✅	✅ (classes/objects/traits)	✅	package declarations	Case classes, objects
  🚧 C#	In Development	.cs	-	-	-	-	Classes, interfaces, generics (planned)

• 🌳 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: Supports both cloud models (Google Gemini), local models (Ollama), and OpenAI models for natural language to Cypher translation

• 🤖 OpenAI Integration: Leverage OpenAI models to enhance AI functionalities.

• 📝 Code Snippet Retrieval: Retrieves actual source code snippets for found functions/methods

• ✍️ Advanced File Editing: Surgical code replacement with AST-based function targeting, visual diff previews, and exact code block modifications

• ⚡️ Shell Command Execution: Can execute terminal commands for tasks like running tests or using CLI tools.

• 🚀 Interactive Code Optimization: AI-powered codebase optimization with language-specific best practices and interactive approval workflow

• 📚 Reference-Guided Optimization: Use your own coding standards and architectural documents to guide optimization suggestions

• 🔗 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

🏗️ Architecture

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

📋 Prerequisites

• Python 3.12+
• Docker & Docker Compose (for Memgraph)
• cmake (required for building pymgclient dependency)
• For cloud models: Google Gemini API key
• For local models: Ollama installed and running
• uv package manager

Installing cmake

On macOS:

brew install cmake

On Linux (Ubuntu/Debian):

sudo apt-get update
sudo apt-get install cmake

On Linux (CentOS/RHEL):

sudo yum install cmake
# or on newer versions:
sudo dnf install cmake

🛠️ Installation

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

For development (including tests and pre-commit hooks):

make dev

This installs all dependencies and sets up pre-commit hooks automatically.

This installs Tree-sitter grammars for all supported languages (see Multi-Language Support section).

3. Set up environment variables:

cp .env.example .env
# Edit .env with your configuration (see options below)

Configuration Options

Option 1: Cloud Models (Gemini)

# .env file
GEMINI_API_KEY=your_gemini_api_key_here

Get your free API key from Google AI Studio.

Option 2: OpenAI Models

# .env file
OPENAI_API_KEY=your_openai_api_key_here

Option 3: Local Models (Ollama)

# .env file
LOCAL_MODEL_ENDPOINT=http://localhost:11434/v1
LOCAL_ORCHESTRATOR_MODEL_ID=llama3
LOCAL_CYPHER_MODEL_ID=llama3
LOCAL_MODEL_API_KEY=ollama

Install and run Ollama:

# Install Ollama (macOS/Linux)
curl -fsSL https://ollama.ai/install.sh | sh

# Pull required models
ollama pull llama3
# Or try other models like:
# ollama pull llama3.1
# ollama pull mistral
# ollama pull codellama

# Ollama will automatically start serving on localhost:11434

Note: Local models provide privacy and no API costs, but may have lower accuracy compared to cloud models like Gemini.

4. Start Memgraph database:

docker-compose up -d

🎯 Usage

The Graph-Code system offers four main modes of operation:

1. Parse & Ingest: Build knowledge graph from your codebase
2. Interactive Query: Ask questions about your code in natural language
3. Export & Analyze: Export graph data for programmatic analysis
4. AI Optimization: Get AI-powered optimization suggestions for your code.
5. Editing: Perform surgical code replacements and modifications with precise targeting.

Step 1: Parse a Repository

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

For the first repository (clean start):

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

For additional repositories (preserve existing data):

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

The system automatically detects and processes files for all supported languages (see Multi-Language Support section).

Step 2: Query the Codebase

Start the interactive RAG CLI:

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

Specify Custom Models:

# Use specific local models
python -m codebase_rag.main start --repo-path /path/to/your/repo \
  --orchestrator-model llama3.1 \
  --cypher-model codellama

# Use specific Gemini models
python -m codebase_rag.main start --repo-path /path/to/your/repo \
  --orchestrator-model gemini-2.0-flash-thinking-exp-01-21 \
  --cypher-model gemini-2.5-flash-lite-preview-06-17

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"
• "Find all C++ operator overloads in the Matrix class"
• "Show me C++ template functions with their specializations"
• "List all C++ namespaces and their contained classes"
• "Find C++ lambda expressions used in algorithms"
• "Add logging to all database connection functions"
• "Refactor the User class to use dependency injection"
• "Convert these Python functions to async/await pattern"
• "Add error handling to authentication methods"
• "Optimize this function for better performance"

Step 3: Export Graph Data (New!)

For programmatic access and integration with other tools, you can export the entire knowledge graph to JSON:

Export during graph update:

python -m codebase_rag.main start --repo-path /path/to/repo --update-graph --clean -o my_graph.json

Export existing graph without updating:

python -m codebase_rag.main export -o my_graph.json

Working with exported data:

from codebase_rag.graph_loader import load_graph

# Load the exported graph
graph = load_graph("my_graph.json")

# Get summary statistics
summary = graph.summary()
print(f"Total nodes: {summary['total_nodes']}")
print(f"Total relationships: {summary['total_relationships']}")

# Find specific node types
functions = graph.find_nodes_by_label("Function")
classes = graph.find_nodes_by_label("Class")

# Analyze relationships
for func in functions[:5]:
    relationships = graph.get_relationships_for_node(func.node_id)
    print(f"Function {func.properties['name']} has {len(relationships)} relationships")

Example analysis script:

python examples/graph_export_example.py my_graph.json

This provides a reliable, programmatic way to access your codebase structure without LLM restrictions, perfect for:

• Integration with other tools
• Custom analysis scripts
• Building documentation generators
• Creating code metrics dashboards

Step 4: Code Optimization (New!)

For AI-powered codebase optimization with best practices guidance:

Basic optimization for a specific language:

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

Optimization with reference documentation:

python -m codebase_rag.main optimize python \
  --repo-path /path/to/your/repo \
  --reference-document /path/to/best_practices.md

Using specific models for optimization:

python -m codebase_rag.main optimize javascript \
  --repo-path /path/to/frontend \
  --orchestrator-model gemini-2.0-flash-thinking-exp-01-21

Supported Languages for Optimization: All supported languages: python, javascript, typescript, rust, go, java, scala, cpp

How It Works:

1. Analysis Phase: The agent analyzes your codebase structure using the knowledge graph
2. Pattern Recognition: Identifies common anti-patterns, performance issues, and improvement opportunities
3. Best Practices Application: Applies language-specific best practices and patterns
4. Interactive Approval: Presents each optimization suggestion for your approval before implementation
5. Guided Implementation: Implements approved changes with detailed explanations

Example Optimization Session:

Starting python optimization session...
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ The agent will analyze your python codebase and propose specific          ┃
┃ optimizations. You'll be asked to approve each suggestion before          ┃
┃ implementation. Type 'exit' or 'quit' to end the session.                 ┃
┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛

🔍 Analyzing codebase structure...
📊 Found 23 Python modules with potential optimizations

💡 Optimization Suggestion #1:
   File: src/data_processor.py
   Issue: Using list comprehension in a loop can be optimized
   Suggestion: Replace with generator expression for memory efficiency

   [y/n] Do you approve this optimization?

Reference Document Support: You can provide reference documentation (like coding standards, architectural guidelines, or best practices documents) to guide the optimization process:

# Use company coding standards
python -m codebase_rag.main optimize python \
  --reference-document ./docs/coding_standards.md

# Use architectural guidelines
python -m codebase_rag.main optimize java \
  --reference-document ./ARCHITECTURE.md

# Use performance best practices
python -m codebase_rag.main optimize rust \
  --reference-document ./docs/performance_guide.md

The agent will incorporate the guidance from your reference documents when suggesting optimizations, ensuring they align with your project's standards and architectural decisions.

Common CLI Arguments:

• --orchestrator-model: Specify model for main operations
• --cypher-model: Specify model for graph queries
• --repo-path: Path to repository (defaults to current directory)
• --reference-document: Path to reference documentation (optimization only)

📊 Graph Schema

The knowledge graph uses the following node types and relationships:

Node Types

• Project: Root node representing the entire repository
• Package: Language packages (Python: __init__.py, etc.)
• Module: Individual source code files (.py, .js, .jsx, .ts, .tsx, .rs, .go, .scala, .sc, .java)
• 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

Language-Specific Mappings

• Python: function_definition, class_definition
• JavaScript/TypeScript: function_declaration, arrow_function, class_declaration
• C++: function_definition, template_declaration, lambda_expression, class_specifier, struct_specifier, union_specifier, enum_specifier
• Rust: function_item, struct_item, enum_item, impl_item
• Go: function_declaration, method_declaration, type_declaration
• Scala: function_definition, class_definition, object_definition, trait_definition
• Java: method_declaration, class_declaration, interface_declaration, enum_declaration

Relationships

• CONTAINS_PACKAGE: Project or Package contains Package nodes
• CONTAINS_FOLDER: Project, Package, or Folder contains Folder nodes
• CONTAINS_FILE: Project, Package, or Folder contains File nodes
• CONTAINS_MODULE: Project, Package, or Folder contains Module nodes
• DEFINES: Module defines classes/functions
• DEFINES_METHOD: Class defines methods
• DEPENDS_ON_EXTERNAL: Project depends on external packages
• CALLS: Function or Method calls other functions/methods

🔧 Configuration

Configuration is managed through environment variables in .env file:

Gemini (Cloud) Configuration

• GEMINI_API_KEY: Required when using Google models.
• GEMINI_MODEL_ID: Main model for orchestration (default: gemini-2.5-pro)
• MODEL_CYPHER_ID: Model for Cypher generation (default: gemini-2.5-flash-lite-preview-06-17)

Local Models Configuration

• LOCAL_MODEL_ENDPOINT: Ollama endpoint (default: http://localhost:11434/v1)
• LOCAL_ORCHESTRATOR_MODEL_ID: Model for main RAG orchestration (default: llama3)
• LOCAL_CYPHER_MODEL_ID: Model for Cypher query generation (default: llama3)
• LOCAL_MODEL_API_KEY: API key for local models (default: ollama)

Other Settings

• MEMGRAPH_HOST: Memgraph hostname (default: localhost)
• MEMGRAPH_PORT: Memgraph port (default: 7687)
• TARGET_REPO_PATH: Default repository path (default: .)

Key Dependencies

• tree-sitter: Core Tree-sitter library for language-agnostic parsing
• tree-sitter-{language}: Language-specific grammars (Python, JS, TS, Rust, Go, Scala, Java)
• 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

🤖 Agentic Workflow & Tools

The agent is designed with a deliberate workflow to ensure it acts with context and precision, especially when modifying the file system.

Core Tools

The agent has access to a suite of tools to understand and interact with the codebase:

• query_codebase_knowledge_graph: The primary tool for understanding the repository. It queries the graph database to find files, functions, classes, and their relationships based on natural language.
• get_code_snippet: Retrieves the exact source code for a specific function or class.
• read_file_content: Reads the entire content of a specified file.
• create_new_file: Creates a new file with specified content.
• replace_code_surgically: Surgically replaces specific code blocks in files. Requires exact target code and replacement. Only modifies the specified block, leaving rest of file unchanged. True surgical patching.
• execute_shell_command: Executes a shell command in the project's environment.

Intelligent and Safe File Editing

The agent uses AST-based function targeting with Tree-sitter for precise code modifications. Features include:

• Visual diff preview before changes
• Surgical patching that only modifies target code blocks
• Multi-language support across all supported languages
• Security sandbox preventing edits outside project directory
• Smart function matching with qualified names and line numbers

🌍 Multi-Language Support

Language-Specific Features

• Python: Full support including nested functions, methods, classes, decorators, type hints, and package structure
• JavaScript: ES6 modules, CommonJS modules, prototype-based methods, object methods, arrow functions, classes, and JSX support
• TypeScript: All JavaScript features plus interfaces, type aliases, enums, namespaces, generics, and advanced type inference
• C++: Comprehensive support including functions, classes, structs, unions, enums, constructors, destructors, operator overloading, templates, lambdas, namespaces, C++20 modules, inheritance, method calls, and modern C++ features
• Lua: Functions, local/global variables, tables, metatables, closures, coroutines, and object-oriented patterns
• Rust: Functions, structs, enums, impl blocks, traits, and associated functions
• Go: Functions, methods, type declarations, interfaces, and struct definitions
• Scala: Functions, methods, classes, objects, traits, case classes, implicits, and Scala 3 syntax
• Java: Methods, constructors, classes, interfaces, enums, annotations, generics, modern features (records, sealed classes, switch expressions), concurrency patterns, reflection, and enterprise frameworks

Adding New Languages

Graph-Code makes it easy to add support for any language that has a Tree-sitter grammar. The system automatically handles grammar compilation and integration.

⚠️ Recommendation: While you can add languages yourself, we recommend waiting for official full support to ensure optimal parsing quality, comprehensive feature coverage, and robust integration. The languages marked as "In Development" above will receive dedicated optimization and testing.

💡 Request Support: If you want a specific language to be officially supported, please submit an issue with your language request.

Quick Start: Add a Language

Use the built-in language management tool to add any Tree-sitter supported language:

# Add a language using the standard tree-sitter repository
python -m codebase_rag.tools.language add-grammar <language-name>

# Examples:
python -m codebase_rag.tools.language add-grammar c-sharp
python -m codebase_rag.tools.language add-grammar php
python -m codebase_rag.tools.language add-grammar ruby
python -m codebase_rag.tools.language add-grammar kotlin

Custom Grammar Repositories

For languages hosted outside the standard tree-sitter organization:

# Add a language from a custom repository
python -m codebase_rag.tools.language add-grammar --grammar-url https://github.com/custom/tree-sitter-mylang

What Happens Automatically

When you add a language, the tool automatically:

1. Downloads the Grammar: Clones the tree-sitter grammar repository as a git submodule
2. Detects Configuration: Auto-extracts language metadata from tree-sitter.json
3. Analyzes Node Types: Automatically identifies AST node types for:
   • Functions/methods (method_declaration, function_definition, etc.)
   • Classes/structs (class_declaration, struct_declaration, etc.)
   • Modules/files (compilation_unit, source_file, etc.)
   • Function calls (call_expression, method_invocation, etc.)
4. Compiles Bindings: Builds Python bindings from the grammar source
5. Updates Configuration: Adds the language to codebase_rag/language_config.py
6. Enables Parsing: Makes the language immediately available for codebase analysis

Example: Adding C# Support

$ python -m codebase_rag.tools.language add-grammar c-sharp
🔍 Using default tree-sitter URL: https://github.com/tree-sitter/tree-sitter-c-sharp
🔄 Adding submodule from https://github.com/tree-sitter/tree-sitter-c-sharp...
✅ Successfully added submodule at grammars/tree-sitter-c-sharp
Auto-detected language: c-sharp
Auto-detected file extensions: ['cs']
Auto-detected node types:
Functions: ['destructor_declaration', 'method_declaration', 'constructor_declaration']
Classes: ['struct_declaration', 'enum_declaration', 'interface_declaration', 'class_declaration']
Modules: ['compilation_unit', 'file_scoped_namespace_declaration', 'namespace_declaration']
Calls: ['invocation_expression']

✅ Language 'c-sharp' has been added to the configuration!
📝 Updated codebase_rag/language_config.py

Managing Languages

# List all configured languages
python -m codebase_rag.tools.language list-languages

# Remove a language (this also removes the git submodule unless --keep-submodule is specified)
python -m codebase_rag.tools.language remove-language <language-name>

Language Configuration

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

"language-name": LanguageConfig(
    name="language-name",
    file_extensions=[".ext1", ".ext2"],
    function_node_types=["function_declaration", "method_declaration"],
    class_node_types=["class_declaration", "struct_declaration"],
    module_node_types=["compilation_unit", "source_file"],
    call_node_types=["call_expression", "method_invocation"],
),

Troubleshooting

Grammar not found: If the automatic URL doesn't work, use a custom URL:

python -m codebase_rag.tools.language add-grammar --grammar-url https://github.com/custom/tree-sitter-mylang

Version incompatibility: If you get "Incompatible Language version" errors, update your tree-sitter package:

uv add tree-sitter@latest

Missing node types: The tool automatically detects common node patterns, but you can manually adjust the configuration in language_config.py if needed.

📦 Building a binary

You can build a binary of the application using the build_binary.py script. This script uses PyInstaller to package the application and its dependencies into a single executable.

python build_binary.py

The resulting binary will be located in the dist directory.

🐛 Debugging

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. For local models:

   • Verify Ollama is running: ollama list
   • Check if models are downloaded: ollama pull llama3
   • Test Ollama API: curl http://localhost:11434/v1/models
   • Check Ollama logs: ollama logs

🤝 Contributing

Please see CONTRIBUTING.md for detailed contribution guidelines.

Good first PRs are from TODO issues.

🙋‍♂️ Support

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

Star History