Adding Graph Retrieval to a RAG Pipeline With Python

The majority of retrieval-augmented generation (RAG) examples start as follows: document segmentation; create embedding representations for segmented document sections ("chunks"); store those representations in a vector database and use them to "retrieve" the "best matching" document section(s) for any user-inputted search query.

That is a good starting point. However, there is one significant limitation to the process described above: not all questions relate solely to similarity. Some questions relate to relationships. For example:

Which team owns the API that Checkout Service uses?

While a vector search system could potentially identify documents related to "Checkout Service," "Payment API," and "Platform Team," the critical piece of information is the relationship between these concepts:

Checkout Service --DEPENDS ON--> Payment API
Platform Team    --OWNS-->       Payment API

This is when graph retrieval becomes a value-added element.

Unlike retrieving only documents based on similarity to a query, a graph retrieval layer can retrieve connected Facts from a graph, providing relationship-aware context to a large language model (LLM) while allowing the LLM to generate the final answer. In other words, instead of simply providing additional context to an LLM in terms of a specific set of retrieved documents, a graph retrieval layer would be able to provide structured context.

In this tutorial, we will develop a basic graph retrieval layer to enhance a RAG pipeline utilizing Python and NetworkX. We will load relationship triples from a JSON File, build a directed graph, retrieve the local neighborhood of any entity matched by the retrieval layer, and convert that graph-based context into a format suitable for an LLM Prompt.

Why Graph Retrieval Helps

Most traditional RAG methods are effective for questions that have answers composed within a couple of data "chunks."

Examples include:

• What database does Payment API utilize? As an example, if you have documentation that states, "The Payment API utilizes PostgreSQL," then a vector search should normally identify it.

However, as soon as you are looking for information about related facts (i.e., multiple connected pieces of fact), plain similarity search may not provide sufficient results:

• Who owns the API that the Checkout Service relies upon? In order to determine who owns the API that the Checkout Service relies upon, you need to perform at least two link traversals:
  • Checkout Service   ->     Payment API
  • Platform Team       ->      Payment API

Since this creates a relationship map (a graph), systems, teams, APIs, and databases all represent nodes in the graph, while their relationships all represent the edges in the graph. As long as this structure is in place, retrieval can follow edges (relationships) rather than simply rely on semantic similarities.

This does not mean that graph retrieval replaces vector search completely. A hybrid approach has been found to be more successful in most production environments; specifically, the vector search identifies relevant text, while the graph retrieval identifies relevant relationships among those texts, and ultimately the final answer will be determined using both.

Therefore, our main emphasis in this article will be solely on the graph retrieval layer.

What We Are Building

We are creating a short Python application:

• Loading data (relationship triples) from file "triples.json".
• Building a directional graph using NetworkX.
• Identifying all entities mentioned within the user's input query.
• Retrieving related graph information.
• Converting that related information to form context for the LLM prompt.
• Writing the processed contextual information to "graph_context.json."

Output will appear as follows:

{
"question": "which team is responsible for owning the api Checkout Service depends on?",
"context": [
	"Checkout Service depends_on payment api.",
	"Platform Team OWNS Payment API."
	]
}

Then you could pass that created contextual information into an LLM prompt portion of a RAG prompt.

Project Setup

Create a new project folder:

mkdir graph-retrieval-rag

cd graph-retrieval-rag

Create a requirements.txt file:

networkx==3.3

Install the dependency:

pip install -r requirements.txt

Create a file named 'triples.json':

[
  {
    "subject": "Checkout Service",
    "relation": "DEPENDS_ON",
    "object": "Payment API",
    "source_text": "Checkout Service depends on Payment API during payment authorization."
  },
  {
    "subject": "Payment API",
    "relation": "USES",
    "object": "PostgreSQL",
    "source_text": "Payment API uses PostgreSQL for transaction metadata."
  },
  {
    "subject": "Platform Team",
    "relation": "OWNS",
    "object": "Payment API",
    "source_text": "Platform Team owns Payment API."
  },
  {
    "subject": "Recommendation Service",
    "relation": "CALLS",
    "object": "Catalog API",
    "source_text": "Recommendation Service calls Catalog API to retrieve product details."
  },
  {
    "subject": "Catalog API",
    "relation": "USES",
    "object": "Elasticsearch",
    "source_text": "Catalog API uses Elasticsearch for product search."
  },
  {
    "subject": "Search Team",
    "relation": "OWNS",
    "object": "Catalog API",
    "source_text": "Search Team owns Catalog API."
  }
]

This file is intentionally small. In a real system, these triples could come from entity extraction, service catalogs, architecture documents, source code metadata, or cloud inventory.

Step 1: Load Relationship Triples

Create a file named graph_retriever.py.

import json 


def load_triples(path): 
  with open(path, "r", encoding="utf-8") as file: 
    return json.load(file)

Each triple has four fields:

  subject
  relation
  object
  source_text

The source_text field matters because graph context should remain traceable. When the model uses a retrieved relationship, the system should know where that relationship came from.

Step 2: Build a Directed Graph

Next, we will create a graph from these triples by importing the NetworkX library for Python, which can create different graph types (for example, directed or undirected).

import networkx as nx


def build_graph(triples):
    graph = nx.DiGraph()

    for triple in triples:
        subject = triple["subject"]
        object_ = triple["object"]
        relation = triple["relation"]

        graph.add_node(subject)
        graph.add_node(object_)

        graph.add_edge(
            subject,
            object_,
            relation=relation,
            source_text=triple["source_text"],
        )

    return graph

A directed graph is important because relationship direction changes meaning.

This:

Checkout Service --DEPENDS_ON--> Payment API

does not mean the same thing as:

Payment API --DEPENDS_ON--> Checkout Service

Direction matters for dependency analysis, ownership lookup, impact analysis, and many GraphRAG workflows.

Step 3: Find Query Entities

We're going to find entities mentioned in a question using simple string searching as the means of finding them.

def find_query_entities(graph, question):
    question_lower = question.lower()
    matches = []

    for node in graph.nodes:
        if node.lower() in question_lower:
            matches.append(node)

    return matches

The code here does not represent a full NLP solution. This is simply a straightforward baseline.

In production, you would probably want to do things like aliasing, entity linking, fuzzy matching, etc., or train your own model. In order to normalize the way that a question is phrased versus how something is named in the graph. For instance, "checkout" might be used to ask about the "Checkout Service.

Step 4: Retrieve a Graph Neighborhood

After identifying the graph element that will answer a question, identify all other elements connected to it.

One method to do this would be to gather together both outgoing and incoming edges at each identified node.

This could be done with the following code snippet:

def retrieve_neighborhood(graph, seed_entities, max_hops=2):
    visited_nodes = set(seed_entities)
    frontier = set(seed_entities)
    retrieved_edges = []

    for _ in range(max_hops):
        next_frontier = set()

        for node in frontier:
            for neighbor in graph.successors(node):
                edge_data = graph.get_edge_data(node, neighbor)
                retrieved_edges.append((node, edge_data["relation"], neighbor, edge_data))
                next_frontier.add(neighbor)

            for neighbor in graph.predecessors(node):
                edge_data = graph.get_edge_data(neighbor, node)
                retrieved_edges.append((neighbor, edge_data["relation"], node, edge_data))
                next_frontier.add(neighbor)

        frontier = next_frontier - visited_nodes
        visited_nodes.update(next_frontier)

    return deduplicate_edges(retrieved_edges)

  
# Adding a helper function to remove duplicate edges:

def deduplicate_edges(edges):
    seen = set()
    unique_edges = []

    for source, relation, target, edge_data in edges:
        key = (source, relation, target)

        if key not in seen:
            seen.add(key)
            unique_edges.append((source, relation, target, edge_data))

    return unique_edges

The max_hops parameter controls how far retrieval travels from the original query entity.

For example:

max_hops=1 retrieves immediate neighbors.
max_hops=2 retrieves neighbors of neighbors.

While higher values give you a larger amount of additional knowledge about your subject matter, this may also introduce a lot of "noise" into your results.

Step 5: Convert Graph Results Into LLM Context

The retrieved graph edges are structured data. Before sending them to an LLM, convert them into short, readable facts.

def format_edges_as_context(edges):
    context = []

    for source, relation, target, edge_data in edges:
        context.append(
            {
                "fact": f"{source} {relation} {target}.",
                "source_text": edge_data["source_text"],
            }
        )

    return context

This gives the prompt two useful forms of evidence:

1. A normalized graph fact.
2. The original source sentence.

That makes the final answer easier to debug and easier to cite.

Step 6: Build a RAG Prompt

Now create a simple prompt builder.

def build_prompt(question, graph_context):
    context_lines = []

    for item in graph_context:
        context_lines.append(f"- {item['fact']} Source: {item['source_text']}")

    context_text = "\n".join(context_lines)

    return f"""
You are answering a question using graph-retrieved context.

Use only the facts provided in the context. If the context does not contain enough information, say so.

Question:
{question}

Context:
{context_text}

Answer:
""".strip()

This prompt is intentionally restrictive. The model should use the source instead of guessing.

A good production application would also include a format for citations, rules for how to write responses, and fallback behavior.

Complete Example

Here is the full code:

import json
import networkx as nx


def load_triples(path):
    with open(path, "r", encoding="utf-8") as file:
        return json.load(file)


def build_graph(triples):
    graph = nx.DiGraph()

    for triple in triples:
        subject = triple["subject"]
        object_ = triple["object"]
        relation = triple["relation"]

        graph.add_node(subject)
        graph.add_node(object_)

        graph.add_edge(
            subject,
            object_,
            relation=relation,
            source_text=triple["source_text"],
        )

    return graph


def find_query_entities(graph, question):
    question_lower = question.lower()
    matches = []

    for node in graph.nodes:
        if node.lower() in question_lower:
            matches.append(node)

    return matches


def retrieve_neighborhood(graph, seed_entities, max_hops=2):
    visited_nodes = set(seed_entities)
    frontier = set(seed_entities)
    retrieved_edges = []

    for _ in range(max_hops):
        next_frontier = set()

        for node in frontier:
            for neighbor in graph.successors(node):
                edge_data = graph.get_edge_data(node, neighbor)
                retrieved_edges.append((node, edge_data["relation"], neighbor, edge_data))
                next_frontier.add(neighbor)

            for neighbor in graph.predecessors(node):
                edge_data = graph.get_edge_data(neighbor, node)
                retrieved_edges.append((neighbor, edge_data["relation"], node, edge_data))
                next_frontier.add(neighbor)

        frontier = next_frontier - visited_nodes
        visited_nodes.update(next_frontier)

    return deduplicate_edges(retrieved_edges)


def deduplicate_edges(edges):
    seen = set()
    unique_edges = []

    for source, relation, target, edge_data in edges:
        key = (source, relation, target)

        if key not in seen:
            seen.add(key)
            unique_edges.append((source, relation, target, edge_data))

    return unique_edges


def format_edges_as_context(edges):
    context = []

    for source, relation, target, edge_data in edges:
        context.append(
            {
                "fact": f"{source} {relation} {target}.",
                "source_text": edge_data["source_text"],
            }
        )

    return context


def build_prompt(question, graph_context):
    context_lines = []

    for item in graph_context:
        context_lines.append(f"- {item['fact']} Source: {item['source_text']}")

    context_text = "\n".join(context_lines)

    return f"""
You are answering a question using graph-retrieved context.

Use only the facts provided in the context. If the context does not contain enough information, say so.

Question:
{question}

Context:
{context_text}

Answer:
""".strip()


def write_json(path, data):
    with open(path, "w", encoding="utf-8") as file:
        json.dump(data, file, indent=2)


if __name__ == "__main__":
    question = "Which team owns the API that Checkout Service depends on?"

    triples = load_triples("triples.json")
    graph = build_graph(triples)

    seed_entities = find_query_entities(graph, question)

    if not seed_entities:
        print("No matching graph entities found in the question.")
        raise SystemExit(0)

    retrieved_edges = retrieve_neighborhood(graph, seed_entities, max_hops=2)
    graph_context = format_edges_as_context(retrieved_edges)
    prompt = build_prompt(question, graph_context)

    result = {
        "question": question,
        "seed_entities": seed_entities,
        "graph_context": graph_context,
        "prompt": prompt,
    }

    write_json("graph_context.json", result)

    print("Seed entities:")
    print(seed_entities)

    print("\nGraph context:")
    for item in graph_context:
        print("-", item["fact"])

    print("\nPrompt:")
    print(prompt)

Run it:

python graph_retriever.py

Expected output:

Seed entities:
['Checkout Service']

Graph context:
- Checkout Service DEPENDS_ON Payment API.
- Payment API USES PostgreSQL.
- Platform Team OWNS Payment API.

Prompt:
You are answering a question using graph-retrieved context.

Use only the facts provided in the context. If the context does not contain enough information, say so.

Question:
Which team owns the API that Checkout Service depends on?

Context:
- Checkout Service DEPENDS_ON Payment API. Source: Checkout Service depends on Payment API during payment authorization.
- Payment API USES PostgreSQL. Source: Payment API uses PostgreSQL for transaction metadata.
- Platform Team OWNS Payment API. Source: Platform Team owns Payment API.

Answer:

The answer is now straightforward:

Platform Team owns Payment API, which Checkout Service depends on.

The important part is not that an LLM can phrase the answer. The important part is that retrieval found the connected facts needed to support it.

What This Adds to a RAG Pipeline

A standard RAG pipeline often looks like this:

Documents
   |
Chunking
   |
Embeddings
   |
Vector Search
   |
  LLM

A graph-enhanced pipeline adds another retrieval path:

Question
   |
Query entity detection
   |
Graph neighborhood retrieval
   |
Graph context formatting
   |
  LLM

In a hybrid system, both paths can work together:

Question
   |
   |---- Vector search -> relevant text chunks
   |
   |---- Graph search  -> connected facts
   |
Combined context
   |
  LLM

This is useful because vector search and graph retrieval solve different problems.

Vector search is good at finding semantically similar text. Graph retrieval is good at following explicit relationships.

Production Considerations

The graph retrieval layer used in this example is very rudimentary. A production-quality implementation will have many additional "guardrails."

Normalized entities are necessary for performing entity matching. Exact string matches will miss most common aliases (i.e., checkout, checkout-service, etc.); however, production-quality implementations will use canonical IDs/aliases/entity linking.

Graph retrieval has limits. While there is no limit to how large a graph can grow in terms of returned context, returning too much context can be detrimental. To limit the amount of context returned, implement boundaries, including but not limited to hop count, edge type, node type, confidence value, and source data freshness.

All relationships in the graph are not equal in reliability. For example, a dependency established by a service catalog entry may be a more reliable relationship than a dependency inferred from an older document. Therefore, store your confidence level, timestamp, and other source metadata values along with every edge.

Not all relationships are equally reliable. Contextual information, in addition to just what was retrieved from the graph, should also contain evidence that supports or refutes the answer. As a result of providing evidence with contextual information, when someone reviews your answer, they will be able to easily verify the answer, and if something does go wrong during development, it will make it easier to debug.

Hybrid approaches are best suited for using graph retrieval. While graph retrieval can help the model understand why certain pieces of information relate to each other, document chunks can provide much more detail about those related pieces. When retrieving both types of information, the model will likely receive more complete contextual information.

Evaluation is mandatory. You will need to create realistic test questions and determine if the sub-graph retrieved from the database includes enough factual information for answering those test questions. Never assume that graph retrieval will improve performance unless you measure it.

Key Takeaways

Graph retrieval adds relational contextual knowledge to a RAG framework.

A simple directed graph can answer many complex questions involving multi-step logic that would be difficult to accomplish with plain similarity-based search.

One key design decision in developing an effective graph-retrieval system is what determines the success of the retrieval layer: does it identify the correct entities and their associated relationships?

Graph retrieval works well when the question involves dependency, ownership, impact assessment, lineage, and/or multi-step reasoning.

In addition to graph retrieval, build a production-quality retrieval system that includes source tracking, entity normalization, confidence scoring, and evaluation.

Try It Yourself

Add these triples to triples.json

{
  "subject": "Billing Worker",
  "relation": "PUBLISHES_TO",
  "object": "Kafka",
  "source_text": "Billing Worker publishes payment events to Kafka."
},
{
  "subject": "Data Platform Team",
  "relation": "OWNS",
  "object": "Kafka",
  "source_text": "Data Platform Team owns Kafka."
}

Then change the question to:

Which team owns the system that Billing Worker publishes events to?

Your goal is to retrieve:

Billing Worker PUBLISHES_TO Kafka.
Data Platform Team OWNS Kafka.

Once that works, try lowering max_hops from 2 to 1 and compare the retrieved context. This is a simple way to see why hop count matters in graph retrieval.