What Is the Knowledge Graph Schema in Code-Graph-RAG?

The knowledge graph schema in Code-Graph-RAG is defined as a Protocol Buffers (protobuf) specification where the GraphCodeIndex message serves as the top-level container for strongly-typed Node entities and directed Relationship edges, creating a flattened, ID-based graph representation of codebases.

The vitali87/code-graph-rag project transforms raw source code into a queryable knowledge graph to support retrieval-augmented generation (RAG) systems. At its core, the system relies on a rigorous protobuf schema that standardizes how code entities—from high-level projects down to individual function calls—are encoded and interconnected. This schema ensures that language-specific abstract syntax trees (ASTs) can be unified into a common format for graph traversal and semantic search.

Core Schema Architecture

The knowledge graph schema centers on the GraphCodeIndex message defined in codec/schema.proto. This message acts as the root container for an entire codebase representation and contains two primary repeated fields:

• nodes – A list of generic Node objects that wrap strongly-typed entity messages.
• relationships – A list of Relationship objects that encode directed edges between nodes.

This design creates a flattened graph structure where nodes are referenced by unique identifiers rather than nested hierarchies, enabling efficient storage and traversal of large codebases.

Node Structure and Payloads

Each Node message utilizes a oneof payload field (lines 81–107 of codec/schema.proto) that can hold exactly one of several concrete entity types. This polymorphic approach allows the graph to store heterogeneous code elements—such as Project, Package, Folder, Module, Class, Function, and Method—within a uniform container.

The wrapper pattern decouples the graph topology from entity specifics, enabling the ingestion pipeline to handle new language constructs without altering the core graph structure.

Relationship Definitions

Relationships are defined by the Relationship message (lines 114–139 of codec/schema.proto), which establishes directed edges using source_id and target_id fields that reference node identifiers. Each relationship carries:

• A RelationshipType enum value specifying the semantic nature of the edge
• Optional properties stored as a protobuf Struct for metadata
• Textual labels (source_label, target_label) denoting the entity types involved

Entity Types and Primary Keys

The schema defines distinct protobuf messages for various code artifacts, each utilizing specific fields as primary identifiers:

Entity	Primary Key	Key Fields
Project	name	–
Package	qualified_name	name, path
Folder	path	name
File	path	name, extension
Module	qualified_name	name, path, decorators, flow_covered
ExternalModule	qualified_name	name, path
ModuleImplementation	qualified_name	name, path, implements_module (C++20)
ModuleInterface	qualified_name	name, path (C++20)
Function	qualified_name	name, docstring, line range, decorators, is_exported
Method	qualified_name	name, docstring, line range, decorators
Class	qualified_name	name, docstring, line range, decorators, is_exported
Interface	qualified_name	name, path, absolute_path
Enum	qualified_name	name, path, absolute_path
Type	qualified_name	name
Union	qualified_name	name

The consistent use of qualified_name for code elements and path for filesystem entities ensures reliable cross-referencing throughout the graph.

Relationship Types and Semantic Edges

The RelationshipType enum (lines 118–136 of codec/schema.proto) enumerates all supported edge semantics:

• Containment: CONTAINS_PACKAGE, CONTAINS_FOLDER, CONTAINS_FILE, CONTAINS_MODULE
• Definition: DEFINES (module to function/class/enum), DEFINES_METHOD (class to method)
• Invocation: CALLS (function/method to function/method)
• Inheritance: INHERITS, OVERRIDES
• Dependencies: IMPORTS, DEPENDS_ON_EXTERNAL

These relationship types enable precise graph traversals for use cases like "find all callers of this function" or "identify classes that inherit from a given base class."

Building and Serializing Graphs in Python

The generated Python bindings in codec/schema_pb2.py provide classes for constructing and manipulating graph instances programmatically.

Creating a Basic Graph Structure

from codec import schema_pb2 as sb

# Create nodes

project = sb.Node(project=sb.Project(name="my_project"))
folder = sb.Node(folder=sb.Folder(path="src/", name="src"))
file_node = sb.Node(
    file=sb.File(path="src/main.py", name="main.py", extension=".py")
)

# Create relationships

rel_proj_folder = sb.Relationship(
    type=sb.Relationship.CONTAINS_FOLDER,
    source_id=project.project.name,
    target_id=folder.folder.path,
    source_label="Project",
    target_label="Folder",
)
rel_folder_file = sb.Relationship(
    type=sb.Relationship.CONTAINS_FILE,
    source_id=folder.folder.path,
    target_id=file_node.file.path,
    source_label="Folder",
    target_label="File",
)

# Assemble the graph

graph = sb.GraphCodeIndex(
    nodes=[project, folder, file_node],
    relationships=[rel_proj_folder, rel_folder_file],
)

# Serialize to protobuf binary

with open("graph_index.pb", "wb") as f:
    f.write(graph.SerializeToString())

Adding Functions and Call Relationships

# Create a function node

func = sb.Node(
    function=sb.Function(
        qualified_name="src.main.main",
        name="main",
        docstring="Entry point",
        start_line=1,
        end_line=10,
        decorators=[],
        is_exported=True,
    )
)

# Create a CALLS relationship

call_rel = sb.Relationship(
    type=sb.Relationship.CALLS,
    source_id="src.main.main",      # caller

    target_id="src.utils.helper",   # callee

    source_label="Function",
    target_label="Function",
)

# Append to existing graph

graph.nodes.append(func)
graph.relationships.append(call_rel)

Schema Implementation and Consumption

Several modules interact with the knowledge graph schema during the ingestion and query pipeline:

• codec/schema.proto – The canonical protobuf definition specifying all messages, enums, and fields.
• codec/schema_pb2.py – Auto-generated Python classes providing serialization and deserialization capabilities.
• codebase_rag/graph_loader.py – Handles loading serialized GraphCodeIndex protobuf files from disk into memory.
• codebase_rag/graph_updater.py – Implements logic for mutating the graph, including node addition and relationship updates during incremental indexing.
• codebase_rag/schema_builder.py – Contains utility functions that translate language-specific AST nodes into the standardized protobuf entity messages.

Summary

• The knowledge graph schema in Code-Graph-RAG is implemented in codec/schema.proto using Protocol Buffers.
• The GraphCodeIndex message serves as the root container for lists of Node and Relationship objects.
• Nodes use a oneof payload pattern to wrap specific entity types (Project, Function, Class, etc.), each identified by qualified_name or path.
• Relationships are directed edges with typed enums (CONTAINS_*, DEFINES, CALLS, INHERITS) connecting nodes by ID.
• The generated schema_pb2.py bindings enable Python code to construct, serialize, and deserialize graph indices for storage and retrieval.

Frequently Asked Questions

What file format defines the knowledge graph schema in Code-Graph-RAG?

The schema is defined in codec/schema.proto as a Protocol Buffers (protobuf) specification. This file declares all entity messages, the GraphCodeIndex container, and the RelationshipType enum. The Python implementation uses auto-generated bindings from this file located at codec/schema_pb2.py.

How does the schema handle different programming language constructs uniformly?

The schema employs a polymorphic node wrapper where the generic Node message contains a oneof field that can hold any specific entity type (Function, Class, Module, etc.). This allows the graph to store heterogeneous code elements from multiple languages within a single, uniform structure while preserving type safety through the protobuf definition.

What fields are used to identify nodes uniquely in the graph?

Nodes are identified by primary key fields specific to their entity type: qualified_name for code symbols (functions, classes, modules) and path for filesystem elements (files, folders). These identifiers populate the source_id and target_id fields in Relationship messages to create edges without data duplication.

Which modules are responsible for building and loading the knowledge graph?

The graph construction and persistence pipeline relies on codebase_rag/schema_builder.py for AST-to-protobuf translation, codebase_rag/graph_updater.py for incremental graph modifications, and codebase_rag/graph_loader.py for deserializing stored GraphCodeIndex protobuf files into queryable in-memory structures.

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