Featured
Building production-ready GraphRAG systems with Kedro
A demo project to manage the graph, vector index, and relational store
Standard Retrieval-Augmented Generation (RAG) works well for factual lookup. For example, if you use semantic similarity to retrieve a document, you’ll usually find the correct passage. However, it struggles when the lookup requires reasoning across relationships, say if the answer depends on facts that are distributed across records that share little lexical or semantic similarity with each other.
As an illustration for healthcare, a search might want to discover the risk of a particular medication to patients with a certain condition. A typical query might be: “For a chronic kidney disease patient, which medication may not be safe and interact with their poor kidney function?”
A dense retriever might return the passages whose embeddings sit the closest to the query vector, like a chunk about chronic kidney disease and another about drugs with contraindications for renal patients. But to respond to this question as accurately as possible, we’d have to gather information from separate records, like a diagnosis note, a lab report, or a drug prescription. None of these directly references the others, and there is no structure that encodes the intersection of all of these pieces of data.
GraphRAG addresses this by pairing semantic search with a knowledge graph that represents the structure explicitly: entities such as conditions, medications, and test results become nodes, and the relationships between them become edges. Before passing context to the language model, the system traverses the relationships connected to the retrieved entities so the model can reason across linked nodes rather than treating each document in isolation.
In this article, we describe how we used Kedro to build a demo of a simple reference implementation.
The basics
In our example, the GraphRAG system maintains three representations of the same underlying data, which need to stay consistent with each other as source data changes. The representations are:
- a knowledge graph for relationship traversal
- a vector index for semantic retrieval
- a relational store for aggregate statistics.
Data engineering matters as much as retrieval because of the need for consistency. Inconsistent data arises, for example, when data teams use a collection of scripts, run in undocumented order, with no guarantee that the graph, the vector index, and the analytics store are produced from the same source run. When the knowledge graph becomes stale relative to the embeddings, the retrieval and graph traversal return contradictory context. As a result, the model resolves the inconsistency silently, and may resolve it incorrectly.
Incremental updates compound the challenge. Rebuilding the full graph each time new records arrive is impractical at production scale. A production system needs targeted, merge-capable updates that leave unchanged nodes and edges untouched while incorporating new information: a classic data engineering pipeline problem.
An example pipeline
Using Kedro, we built a five-stage pipeline demo over a synthetic healthcare dataset of 55,500 patient records covering conditions, medications, insurers, admission types, diagnostic results, and demographics. The full source is on the Kedro Academy GitHub repository. While this example is based on healthcare data, none of this structure is specific to healthcare, and the example is deliberately simple to keep it self-contained and easy to run. The entity extraction, graph building patterns, and incremental merge logic apply to any domain with typed entities and relationships.
The pipeline stages are as follows:
- Clean and summarize: The first stage normalizes raw records, removes duplicates, and computes summary statistics per entity. We used patient counts per condition, average billing amounts per admission type, and result distributions across demographic groups. This stage stores them in SQLite and reuses them as retrieval context in later stages.
- Build the knowledge graph: The second stage extracts entities and typed relationships. Nodes represent conditions, medications, insurers, admission types, test results, and blood types. Edges carry a relationship label (
TREATED_WITH,COVERED_BY,ADMITTED_AS,SHOWS_RESULT,ASSOCIATED_WITH) and a weight recording how frequently each relationship appears in the data. The stage produces a NetworkX JSON graph with 30 nodes and 120 weighted edges, and an interactive D3.js visualization of the entity network. - Incrementally update the graph: Rather than rebuilding from scratch on each data refresh, the next stage merges new records into the existing graph using create-or-update semantics. It adds new entities, increments edge weights for existing relationships, and leaves unchanged portions untouched. In this demo, each run of the stage merges a fixed batch of 5,000 records rather than the full corpus, and it updates only the knowledge graph. We designed it to model a real production constraint: while rebuilding a relatively small graph like this is cheap and resolves in a few seconds, a production-sized knowledge graph based on a larger, more complex set of data would take much longer.
- Generate vector embeddings: For each entity, the next stage generates a natural-language summary of its characteristics and connections, then embeds those summaries using OpenAI’s
text-embedding-3-smallmodel. This stage stores the embeddings in ChromaDB for semantic similarity search at query time. - Answer questions: The final stage runs an AI agent backed by GPT-4o. When a query arrives, the agent embeds it and queries ChromaDB for the most similar entity summaries. Starting from those results, it traverses one hop across the knowledge graph to collect neighboring entities and the relationships between them. That combined context (semantic matches plus explicit graph structure) goes to the model. Standard RAG asks, “What is most similar to this query?” This stage adds, “What else is connected to those results?” This is a hybrid GraphRAG pattern: vector retrieval followed by graph traversal.
All five stages write to three different backends, and Kedro’s Data Catalog manages loading and saving for all of them from a single configuration file. Every dataset (the NetworkX graph, the ChromaDB collection, the SQLite statistics table, the LangChain prompt template) is declared in catalog.yml so there is one place to see where each input and output lives and there are no separate scripts to reconcile manually. A full default run builds all three stores from the same source data. The incremental update refreshes only the knowledge graph; keeping the vector index and statistics in step with it would require extending their pipelines in the same way.
Get QuantumBlack, AI by McKinsey’s stories in your inbox
Join Medium for free to get updates from this writer.
Each stage is an independent Kedro pipeline with explicit inputs and outputs. The graph construction logic, the embedding model, and the retrieval strategy sit in separate modules. Swapping the embedding model requires changing one node and its catalog entry, while the graph and statistics stages are unaffected. Kedro-Viz renders the full data lineage across all five stages, and a published Kedro-Viz build can be viewed without any installation on the viewer’s side.
This approach also means it is possible to treat generative AI components as first-class Data Catalog entries, rather than scripts or loose objects spread around the project. It makes them easier to manage and keep track of, and they can be used as input or output for nodes like any other entry. Through the use of Kedro, generative AI datasets receive easy credentials management and follow the same load/save contract as other Kedro datasets.
The graph_update pipeline is excluded from the default run and must be invoked explicitly. It writes to the same catalog entry as graph_construction, so running both in the same default pipeline would create a DAG conflict. Full rebuilds and incremental updates are mutually exclusive operations, so making that explicit in the pipeline registry prevents silent errors.
Choosing between standard RAG and GraphRAG
GraphRAG is not a replacement for standard RAG. For factual retrieval, a vector index is simpler to build, cheaper to operate, and easier to debug.
The case for a knowledge graph becomes clear when questions require multi-hop reasoning: finding conditions that share treatments, identifying which medications cluster around a particular insurer, or following relationship chains across several entity types. Healthcare illustrates the pattern, but the same logic applies wherever the connections between entities carry as much meaning as the entities themselves, for example, in areas as diverse as finance, legal, or supply chain management.
A reasonable approach is to begin with vector search and add a knowledge graph once relationship-based reasoning becomes a core requirement rather than an occasional edge case.
Resources
If you want to go deeper on GraphRAG itself: the concept was popularized by Microsoft Research in From Local to Global: A Graph RAG Approach to Query-Focused Summarization (Edge et al., 2024), with an accompanying open-source implementation at microsoft/graphrag. For a broader view of the field, Graph Retrieval-Augmented Generation: A Survey (Peng et al., 2024) covers the range of approaches beyond the pattern used here. Neo4j also has a practical, framework-agnostic tutorial on building a RAG system on a knowledge graph.
Kedro is an open-source project supported by a full-time development team. It is part of the LF AI & Data Foundation, with a neutral, vendor-independent governance model that supports long-term sustainability and community-driven development. Beyond the core framework, Kedro is supported by an active open-source community. Contributors and users share patterns, extensions, and real-world experience. A growing list of blog posts, videos, and projects that use Kedro can be found in the awesome-kedro GitHub repository.
Kedro is also backed by extensive documentation. This includes tutorials, example projects, coffee chats, and code snippets ranging from introductory features to more advanced use cases. Teams can get started quickly with a ready-to-use project template, which includes sample datasets and enables users to train their first model with just a few commands from the terminal. If you check it out, drop us a comment below to let us know how it goes!
This article was authored by Laura Couto and Alice Cima with additional insights from Bruce Philp and Joseph Perkins. The editor was Jo Stichbury.
