Build a Lightweight Graph-Based RAG System with LightRAG and Python

Table of Contents

The Real-World Problem: When Vector Search Misses the Big Picture

Standard vector search handles isolated lookups cleanly. Ask it for the definition of an error code, and it returns the right snippet in under 50 milliseconds. But it struggles when a question spans multiple systems and services.

Suppose you ask your internal documentation bot: “How does the Q3 auth update affect our Stripe webhook handlers?” A standard vector pipeline embeds your query and hunts for nearby text chunks. The problem? The answer does not live in one file. It spans PR discussions, an architecture RFC, and downstream worker handlers. Cosine similarity pulls back three disconnected paragraphs, misses the shared dependencies, and forces your LLM to guess the missing links.

Root Cause: Why Flat Embeddings Ignore System Relationships

Most RAG setups chop incoming documents into fixed 500-token chunks and embed each one independently. This creates two distinct blind spots:

Loss of entity topology: Chunking cuts off explicit links between components. A vector can tell you that “JWT” and “RS256” appear in the same paragraph. It cannot reliably trace how an invalid signature in AuthService bubbles up to BillingWorker three hops away.
Blindness to corpus-wide questions: Flat vectors excel at narrow lookups. They stumble on aggregate questions like “What are our three most common timeout causes across all 14 microservices?” No single chunk contains that overview.

Architectural Comparison: Vector RAG vs. GraphRAG vs. LightRAG

Knowledge graphs resolve these visibility gaps by mapping relationships directly into an entity graph. Here is how the three primary retrieval architectures differ in practice:

Architecture	Relationship Traversal	Indexing Cost (100k Tokens)	Incremental Updates
Standard Vector RAG	None (Isolated chunks)	~$0.002 (Embedding only)	Instant append
Microsoft GraphRAG	Hierarchical (Leiden communities)	~$8.00–$25.00+ (Extensive LLM passes)	Requires re-clustering graph hierarchy
LightRAG	Dual-level (Entity + Theme graphs)	~$0.15–$0.40 (Single-pass extraction)	Native document-level appends

Microsoft GraphRAG creates comprehensive, hierarchical community summaries across an entire dataset. That depth makes it great for macroeconomic analysis, but running recursive clustering over a 500-page engineering wiki can burn millions of tokens on a single build. More frustratingly, adding five new pull requests often means re-running that expensive clustering step from scratch.

LightRAG takes a leaner route. It extracts key entities and relationships in a single pass, then splits retrieval into two layers: low-level entity lookups and high-level thematic queries. You get relationship-aware answers without blowing through your API budget.

Implementation: Building a LightRAG Pipeline with Python

You can build a functional Graph RAG prototype with Python and an OpenAI-compatible endpoint. Follow these steps to configure storage, index relational data, and test dual-level queries.

Step 1: Install Dependencies

Create an isolated virtual environment and install the lightrag-hku package along with the official OpenAI SDK:

python3 -m venv venv
source venv/bin/activate
pip install lightrag-hku openai

Step 2: Initialize the Knowledge Graph Engine

Save the following script as graph_rag.py. We configure LightRAG with an asynchronous LLM wrapper for relationship extraction and a standard embedding function for node matching.

import os
import asyncio
from lightrag import LightRAG, QueryParam
from lightrag.llm import openai_complete_if_cache, openai_embedding
from lightrag.utils import EmbeddingFunc

# Set your API credentials
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "sk-your-key-here")

WORKING_DIR = "./lightrag_storage"
os.makedirs(WORKING_DIR, exist_ok=True)

async def llm_model_func(prompt, system_prompt=None, history_messages=[], **kwargs):
    return await openai_complete_if_cache(
        "gpt-4o-mini",
        prompt,
        system_prompt=system_prompt,
        history_messages=history_messages,
        **kwargs
    )

async def embedding_func(texts: list[str]):
    return await openai_embedding(
        texts,
        model="text-embedding-3-small"
    )

rag = LightRAG(
    working_dir=WORKING_DIR,
    llm_model_func=llm_model_func,
    embedding_func=EmbeddingFunc(
        embedding_dim=1536,
        max_token_size=8192,
        func=embedding_func
    )
)

Step 3: Index Interconnected Text

Next, feed interdependent architectural documentation into the engine. LightRAG extracts entities like AuthService and PaymentGateway, infers the relationships between them, and persists the graph locally.

async def index_docs():
    sample_doc = """
    AuthService handles OAuth2 tokens and session validation. In Q3, AuthService was updated
    to issue short-lived JWTs with asymmetric RS256 signing.
    
    PaymentGateway receives webhook events from Stripe. To verify event origins, PaymentGateway
    validates user identity via AuthService public keys before dispatching billing events.
    
    BillingService listens to PaymentGateway events. If AuthService rejects a session,
    BillingService flags the transaction for manual fraud review.
    """
    
    print("Indexing documentation and extracting graph nodes...")
    await rag.ainsert(sample_doc)
    print("Graph extraction complete.")

Step 4: Query Across Local and Global Contexts

LightRAG provides four distinct query modes tailored to different question types:

local: Traverses direct 1-hop connections around specific entities (e.g., “What signing algorithm does AuthService use?”).
global: Scans broader relationship clusters across the whole corpus (e.g., “Outline the fallback flow when authentication fails during payment processing.”).
hybrid: Merges entity graph traversal with standard vector retrieval for balanced coverage.
naive: Runs basic chunk-based vector search for direct A/B performance comparisons.

async def query_pipeline():
    question = "If AuthService rejects a session during a Stripe payment, what happens across all downstream services?"
    
    print("\n--- Querying with Hybrid Mode ---")
    response = await rag.aquery(
        question,
        param=QueryParam(mode="hybrid")
    )
    print(response)

async def main():
    await index_docs()
    await query_pipeline()

if __name__ == "__main__":
    asyncio.run(main())

Step 5: Run and Validate

Execute the script from your terminal:

python graph_rag.py

During the first run, LightRAG writes graph data to kv_store_*.json and graph_chunk_entity_relation.graphml inside your designated working directory.

Production Considerations

Moving from a local prototype to a production deployment requires addressing a few infrastructure trade-offs:

Graph storage limits: LightRAG defaults to in-memory NetworkX graphs saved as JSON. Once your dataset exceeds 50,000 entities, migrate to an external graph engine like Neo4j or Memgraph paired with a vector store such as Qdrant or Milvus.
Extraction cost control: Node extraction makes up over 85% of your ingestion bill. Running a lightweight model like gpt-4o-mini or a self-hosted Qwen 2.5 14B instance via Ollama keeps your per-document indexing costs under a fraction of a cent.
True incremental writes: Because LightRAG avoids hierarchical community re-indexing, calling rag.ainsert() on incoming updates merges new nodes and edges directly into the existing graph without locking your pipeline.