Beyond Flat RAG: Structure-Aware Graph Expansion for Multi-Hop Reasoning with GraphSAGE

A working example project called GraphSAGE-GraphRAG demonstrates that structure-aware retrieval using a Neo4j knowledge graph and an offline-trained GraphSAGE GNN can solve multi-hop reasoning that flat vector RAG fails on. In a 16-document SkyLink Airlines flight-operations corpus, flat RAG retrieved four lexically similar chunks (inc-2209, inc-2209-rca, inc-1187, team-groundops) and incorrectly named Team Ground Operations as both the cause of incident INC-2209 and the team to page, while GraphRAG surfaced seed entities N882AA, INC-2209, and Flight 202 by graph structure and reached the capa-turnaround runbook that names the correct team. The project cites the SAGE (Structure Aware Graph Expansion) framework and an arXiv paper at arxiv.org/abs/2602.16964, with code hosted at github.com/h-swathi-shenoy/graphsage-graphrag.

When one chunk isn’t enough GraphSAGE walks the chain that flat RAG can’t see GraphSAGE-GraphRAG is a working example of Structure-Aware Retrieval solving the multi-hop retrieval problem that flat vector RAG cannot. While traditional RAG treats document chunks as isolated units ranked purely by semantic similarity, this project builds a knowledge graph in Neo4j and trains a GraphSAGE GNN on it offline, so every node’s embedding absorbs its neighborhood before a single question is ever asked. This approach echoes the ideas behind the excellent SAGE Structure Aware Graph Expansion framework and the research paper it’s built on: https://www.arxiv.org/abs/2602.16964 https://www.arxiv.org/abs/2602.16964 . https://github.com/h-swathi-shenoy/graphsage-graphrag/tree/main https://github.com/h-swathi-shenoy/graphsage-graphrag/tree/main Flat RAG works by embedding every chunk of text independently and retrieving whichever chunks are semantically closest to the question. That works fine when the answer lives inside one document. It breaks down the moment the answer is a chain of facts spread across multiple documents , none of which reference each other in a way that text embeddings can pick up. The Dataset This project ships with a small, self-contained example knowledge base 16 short documents describing the flight operations of a fictional airline, SkyLink Airlines. It’s deliberately small so the multi-hop chain is easy to follow by hand, but it’s structured the way a real operations knowledge base is: a mix of entities flights, aircraft, systems, teams, policies and incidents that reference them indirectly. What’s in the corpus: • Flights & aircraft: Flight 202 ORD → LAX , flown by Aircraft N882AA , which is shared across two other routes on rotation • Systems: ramp-scheduler gate/ground scheduling, 8-slot-per-hour cap , crew-scheduler pilot/crew assignment, FAA duty-time rules • Teams: Team Flight Operations , Team Ground Operations , Team Airport Operations , Team Crew Scheduling each owns a different piece of the puzzle • Policies & changes: SP-118 a staffing-policy change , PL-FAA-7 a safety policy it violated • Incidents: INC-2209 the flight delay and its root-cause analysis INC-2209-rca , plus an unrelated earlier incident INC-1187 for contrast • Runbooks: capa-turnaround , capa-crewhold the “what to do” documents that operators actually consult Each document is written the way real operational documentation is written: focused on its own system or event, with almost no cross-referencing. The ramp-scheduler doc never mentions Flight 202. The incident report never explains why the scheduler ran out of capacity. That’s the whole point the corpus is designed so that no single chunk answers a root-cause question, forcing retrieval to either get lucky on lexical overlap flat RAG or actually traverse relationships GraphRAG . Lets look at the following concrete example: Flight-delay incident INC-2209 When asked “Which team should be paged for INC-2209, and which team actually caused it?” and F lat vector RAG retrieves the four chunks that lexically resemble the question inc-2209 , inc-2209-rca , inc-1187 , team-groundops because they all contain words like “team,” “incident,” and “caused.” It never retrieves capa-turnaround , the runbook that actually names the correct team to page, because that document shares almost no vocabulary with the question. The result isn’t just incomplete it’s factually wrong: the LLM confidently names Team Ground Operations for both halves of the question, when in reality that team caused the incident but isn’t who gets paged. GraphRAG starts from the same question but asks a different kind of question of the data: not “what sounds similar?” but “what is this incident connected to?” By attaching the question as a virtual node and running it through GraphSAGE, it surfaces seed entities based on graph structure N882AA , INC-2209 , Flight 202 , then walks two hops in Neo4j to reach ramp-scheduler → OWNED BY → Team Airport Operations and, separately, SP-118 → TRIGGERED BY → Team Ground Operations . Neither of those two facts lives in the same document, and neither is reachable by similarity search alone — only by following the relationship edges between them. That’s the core failure mode this diagram illustrates: flat RAG is a similarity search, not a reasoning process. It fails silently with no low-confidence signal whenever the correct answer depends on a connection between two facts that don’t happen to use similar words. Getting Started The full source is on GitHub, and you can reproduce every screenshot in this article on your own machine in about two minutes of actual runtime plus setup : 1. git clone <repo-url && cd graphsage-graphrag2. docker compose up -d Neo4j 5.26 on 7474 browser / 7687 bolt . Starts Neo4j locally3. python3 -m venv .venv .4. /.venv/bin/pip install -r requirements.txt5. ./.venv/bin/python 01 ingest.py build the graph ~1 min, Bedrock calls 6. ./.venv/bin/python 02 train graphsage.py train the GNN ~20 s, CPU 7. ./.venv/bin/python 03 ask.py answer the demo questions There are three stages, and each one is a standalone script you can run and inspect independently: ingest the documents into a graph, train GraphSAGE on that graph, then retrieve using both the graph structure and the trained embeddings. The rest of this section walks through each stage in the order it runs, using the actual numbers and output from this project’s own pipeline — not simplified pseudocode. The first step turns the 16 flight-ops documents present in corpus.py into a graph instead of a flat list of chunks. For each document: 1. The raw text is stored and embedded as a Chunk this is what flat vector RAG would use on its own . 2. An LLM Nova Lite reads the same text and extracts entities Flight 202 , ramp-scheduler , Team Ground Operations and relationships between them DEPENDS ON , OWNED BY , TRIGGERED BY , … . 3. Those entities and relationships are merged into the graph, so the same entity mentioned in two different documents becomes one connected node — not two disconnected mentions. Why this matters: This is the step that creates the connections flat RAG can never see. ramp-scheduler appears in one document and Team Airport Operations in another — a chunk-similarity search treats them as unrelated text. But here, once extracted into the graph, they’re one hop apart. One deliberate simplification worth knowing about: entities are matched by lowercasing their name Flight 202 → flight 202 . That’s enough for a 16-document demo, but it’s the first thing that breaks at real-world scale — two spellings of the same entity become two disconnected nodes. This is a known limitation, not an oversight. The result, on this corpus: 16 chunks → 33 entities, 54 relationships, 70 mentions. Small enough to inspect by hand in the Neo4j browser, large enough that no single chunk contains the full INC-2209 chain. Each entity also gets its own embedding built from its name, type, description, and the text of chunks that mention it. That embedding is the input feature GraphSAGE trains on next it’s separate from the chunk embeddings used for plain vector search Under the hood: the Cypher cypherCREATE CONSTRAINT chunk id IF NOT EXISTS FOR c:Chunk REQUIRE c.id IS UNIQUECREATE CONSTRAINT entity key IF NOT EXISTS FOR e:Entity REQUIRE e.key IS UNIQUECREATE VECTOR INDEX chunk vec IF NOT EXISTSFOR c:Chunk ON c.embedding OPTIONS {indexConfig: { 'vector.dimensions': 1024, 'vector.similarity function': 'cosine'}}That vector index is what powers the flat RAG baseline in this project it’s the same Neo4j instance doing both jobs, which is what makes the — compare flag a fair, apples-to-apples test. Every extracted entity is merged, not created this is what turns “Team Ground Operations” mentioned in three different documents into one node instead of three duplicates: cypherMERGE n:Entity {key: $key} ON CREATE SET n.name = $name, n.type = $type, n.description = $descON MATCH SET n.description = CASEWHEN size coalesce n.description,'' < size $desc THEN $desc ELSE n.description ENDWITH n MATCH c:Chunk {id: $doc} MERGE c - :MENTIONS - n key is the entity's lowercased name this single line is the entity-resolution strategy for the whole project, and the reason it's flagged as the first thing to improve at scale. Run the following query, right after ingestion, to see the exact chain the article is built around: cypherMATCH p = :Entity {key:'flight 202'} - :RELATES TO 1..3 - :Entity {key:'ramp-scheduler'} RETURN p LIMIT 5That's the query that proves the connection exists in the graph before GraphSAGE ever runs GraphSAGE's job next section is to make that same connection discoverable without writing this exact Cypher by hand, i.e., from a natural-language question instead of a known entity pair. At this point the graph exists in Neo4j, but it’s just structure no model has learned anything from it yet. This step trains GraphSAGE so that every node’s embedding absorbs information from its neighbors, not just its own text. Why not just use the entity embeddings from Stage 1 directly? Because those embeddings only capture what a document says about an entity they know nothing about what that entity is How it works? Each node repeatedly updates its embedding by combining its own vector with the average of its neighbors’ vectors, through a small learned transformation stack that a couple of times, and a node’s final embedding reflects everything within 2 hops of it. Training is unsupervised there are no labels here, just the graph’s own structure. The signal comes from random walks: nodes that co-occur on a short walk are treated as should be similar, and the model is pushed to agree. Why GraphSAGE specifically? GraphSAGE learns aggregator functions . Given any nodes neighbors, even a node that didn’t exist at training time, it can produce a sensible embedding on the spot. That property is the entire reason the retrieval step next section works at all — a user’s question is exactly the kind of “new node” GraphSAGE was built to handle. Once, the training is completed, the trained embeddings are written back as a ‘.sage’ property on each Chunk/Entity node in Neo4j itself, not a separate file, though graph cache.pkl / sage model.pt are also saved locally as a cache/checkpoint for retraining without re-hitting Neo4j. By this point there are two separate things sitting in Neo4j: the original chunks with their text embeddings Stage 1 , and a graph where every entity’s embedding now reflects its neighbors Stage 2 . Retrieval is where both get used together . Think of it as three questions asked in sequence: 1. “Which text sounds like this question?” : The ordinary vector search. Titan embeds the question, Neo4j’s vector index returns the top few chunks whose wording is closest. This is flat RAG, and it’s still run here not thrown away, just no longer the only path. 2. “Which entities is this question structurally close to?” : This is the part that makes GraphRAG different, and the trickiest idea in the whole project. The question doesn’t exist in the graph. So it’s temporarily inserted as a new node, connected to the handful of existing nodes it resembles most in raw text, and then run through the already-trained GraphSAGE model the same one from Stage 2, no retraining involved. Why this works at all: GraphSAGE never memorized “the question” it learned a general rule for how to combine a node with its neighbors . Apply that same rule to a brand-new node, and it produces a sensible, structure-aware embedding on the spot aggegrating feature of GraphSAGE . That’s the entire point of using GraphSAGE which only knows the exact nodes it was trained on. The output of this step is a short list of seed entities not because they sound like the question, but because they sit in the right neighborhood of the graph. In practice, these often look nothing like a keyword match: on the flight-delay example, the top seed for “why was Flight 202 delayed” is the aircraft, N882AA not because the question mentions the tail number, but because that’s structurally where the answer lives. 3. “What is each seed entity connected to?” starting from those seed entities, the graph is walked outward one or two relationship hops, collecting every entity and relationship along the way. This is the step that reaches ramp-scheduler , Staffing Policy SP-118 , and Team Ground Operations none of which were retrieved by step 1, because none of them share vocabulary with the original question. Then everything is handed to the LLM together: The relationships found in step 3 as a simple list “ ramp-scheduler owned by Team Airport Operations ” , plus the actual document text from steps 1 and 3, with an explicit instruction to use the relationships to connect facts across documents. The relationships are what let the model reason across chunks instead of just summarizing whichever one ranked highest. Now I believe you have realized that flat RAG answers based on what the question sounds like” . But , this retriever answers based on “ what the question is connected to” and those turn out to be very different sets of documents whenever the real answer requires more than one hop. TL; DR The whole system is three scripts, run in order, each one handing something concrete to the next: Stage 1 : Ingest 01 ingest.py turns 16 plain-text documents into a graph. Nova Lite reads each document and extracts entities and relationships; those get merged into Neo4j. This runs once . Output: a graph of 33 entities, 54 relationships, 70 chunk-to-entity mentions sitting permanently in Neo4j. Stage 2: Train 02 train graphsage.py reads that same graph back out, flattens it into 49 nodes and 113 edges, and trains GraphSAGE on it unsupervised. This also runs once or periodically, if new documents are added later . The important thing it produces isn’t a file it’s two things: a trained model sage model.pt , the reusable aggregator functions and a 128-dimensional .sage embedding written onto every node in Neo4j itself , so the graph now carries both its structure and its trained representation together. Stage 3: Retrieve 03 ask.py / retriever.py is the only stage that runs per question. It doesn’t touch the documents or retrain anything it reuses the trained model from Stage 2 and the graph from Stage 1 to a embed the question, b run it through GraphSAGE as a temporary “virtual node” to get structure-aware seed entities, and walk 2 hops from those seeds in Neo4j to pull in the connected facts. demo compare trace — View Image https://jumpshare.com/s/3H9Ii7iJ1zK3DviyKF1f This project is a small, concrete step toward retrieval systems that are both smarter and more trustworthy. By moving structural reasoning offline extracting an entity graph and training GraphSAGE on it before any question is ever asked the expensive work of discovering relationships is separated from the latency-sensitive work of answering a query. What’s left at inference time is a single forward pass and a Cypher walk: fast enough to run per-question, deterministic enough to explain afterward. Unlike a pure vector-similarity retriever, where all you can observe is what got retrieved, this pipeline is traceable end-to-end: which node the question attached to, which seed entities GraphSAGE ranked and why, which relationship edges were traversed, and which chunks backed the final answer. Every hop in the chain is inspectable including the one hop flat RAG silently drops. That difference showed up concretely, not theoretically: asked the same question, flat RAG named the same team for two different roles paged and at-fault because both answers were anchored to the same top-ranked chunk. GraphRAG separated them, because it was reasoning over structure, not just wording. As retrieval systems move into domains where a wrong answer isn’t just annoying but costly incident response, compliance, operational root-causing it retrieved something relevant stops being good enough. The question becomes why it retrieved that, and whether the chain of reasoning behind the answer can be checked. This project doesn’t promise a bigger model or a cleverer prompt. It promises a graph, an inductive embedding, and a traversal you can point to. And in multi-hop retrieval, that’s the difference between a system that sounds right and one you can actually verify Thank you for reading 😊 Github Link: https://github.com/h-swathi-shenoy/graphsage-graphrag/tree/main https://github.com/h-swathi-shenoy/graphsage-graphrag/tree/main Beyond Flat RAG: Structure-Aware Graph Expansion for Multi-Hop Reasoning with GraphSAGE https://blog.devgenius.io/beyond-flat-rag-structure-aware-graph-expansion-for-multi-hop-reasoning-with-graphsage-4f780c1277d8 was originally published in Dev Genius https://blog.devgenius.io on Medium, where people are continuing the conversation by highlighting and responding to this story.