# Beyond Flat RAG: Structure-Aware Graph Expansion for Multi-Hop Reasoning with GraphSAGE > Source: > Published: 2026-09-11 21:06:01+00:00 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 && 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.