Code-Graph-RAG: A Code Search Engine Powered by Knowledge Graphs 83

Code-Graph-RAG is a code search engine that integrates knowledge graphs with Retrieval-Augmented Generation (RAG) technology, aiming to improve the accuracy of code search.

Main Features:

• Supports multi-language codebases such as Python, JavaScript, TypeScript, Rust, and Go

• Cross-language Abstract Syntax Tree (AST) parsing based on Tree-sitter

• Uses Memgraph to store code structures as an interconnected graph

• Supports natural language queries against the codebase

• Converts natural language into Cypher query statements via Google Gemini

• Can directly retrieve the actual source code of functions/methods

• Capable of parsing pyproject.toml to analyze external dependencies

• Supports handling of nested functions and class hierarchies

• Adopts a unified graph schema design across all supported languages

Technical Principles:

• Knowledge Graph Construction: After analyzing the codebase, the large model extracts entities such as classes and functions along with their relationships (e.g., inheritance, dependencies) to build an interconnected graph. For example, when code documentation is input, the system splits the text into units, identifies code structures, and generates the corresponding graph.

• Graph-Enhanced Retrieval:

• Subgraph Retrieval: When a user asks a question, the system extracts key entities and expands associated paths in the graph (e.g., 3-hop neighborhoods) to generate a matching subgraph. For instance, when querying a function, it can link to related nodes such as classes and modules.

• Hybrid Retrieval Strategy: Combines vector retrieval (e.g., BERT embeddings) with graph embeddings (e.g., Node2Vec, TransE), and determines the final results through a weighted formula (e.g., Score=α⋅cos_sim(q,v)+β⋅graph_score(q,G)).

System Architecture:

• Multi-language Parser: A code parsing and data import system based on Tree-sitter

• RAG System: An interactive command-line query interface

Core Components Include:

• Tree-sitter Integration: Enables cross-language parsing

• Graph Database: Memgraph for storing nodes and relationships

• LLM Integration: Google Gemini handles natural language

• Code Analysis: Cross-language AST traversal to extract code elements

• Query tools: Dedicated graph query and code retrieval tools

• Language Configuration: Multi-language mapping parameters

Installation and Usage:

• Prerequisites: Python 3.12+, Docker and Docker Compose (Memgraph), Google Gemini API key, uv package manager

• Installation Steps: Clone the repository → Install dependencies → Configure environment variables → Start Memgraph

• Usage Workflow: Parse the codebase to generate a knowledge graph → Launch the RAG CLI for interactive queries

Application Scenarios:

• Code Analysis: Quickly understand code structure, such as querying classes with specific names or database-related functions

• Code Reuse: Obtain semantically matching code snippets through natural language descriptions of requirements, improving development efficiency

Project Address:

https://github.com/vitali87/code-graph-rag