GraphRAG on Lakebase

Knowledge-graph-augmented RAG served entirely from Lakebase (managed Postgres). Classic RAG retrieves text chunks by vector similarity and misses the relationships between them. GraphRAG adds a knowledge graph so retrieval can traverse relationships, not just match text — then answers from the expanded context.

Architecture

Question
   │
   ▼
pgvector HNSW (cosine)  ──►  seed nodes (semantic entry points)
   │
   ▼
recursive CTE over edges ──►  k-hop expansion (bounded by max_hops)
   │
   ▼
blended rank (seed_similarity × 0.5^hop) ──►  context ──►  model-serving answer

Everything lives in Lakebase: nodes + edges relational tables (typed, JSONB props) as the graph store, a recursive CTE (WITH RECURSIVE) for traversal, and pgvector for the semantic seed. A generic Databricks model-serving endpoint supplies embeddings and answer synthesis — swap in whichever provider your workspace has registered.

The graph schema

Three tables under a graph schema hold everything (sql/schema.sql):

Table	Key columns	Notes
`graph.nodes`	`node_id TEXT PK`, `node_type`, `name`, `props JSONB`, `path LTREE`	Business-key ids (`product:122`, `supplier:S2`); `props` holds type-specific attributes; optional `ltree` path for hierarchy rollups (`us.northeast.boston`). Trigram + GiST indexes for fuzzy name and path lookups.
`graph.edges`	PK `(src_id, dst_id, rel)`, `props JSONB`	The composite key dedupes relationships. Relations: `SUPPLIED_BY`, `LOCATED_IN`, `BELONGS_TO`, `SUBSTITUTE_FOR`, `SURGES_IN`. Indexed both directions (`edges_src_rel_idx`, `edges_dst_rel_idx`).
`graph.node_embeddings`	`node_id TEXT PK`, `embedding VECTOR(1024)`	1024-dim to match `databricks-gte-large-en`. HNSW index: `vector_cosine_ops` with `m = 16, ef_construction = 64`.

Cross-stack interop (`gold_triplets`)

A knowledge graph is portable between engines when they agree on one row shape. sql/gold_triplets_mapping.sql projects graph.nodes / graph.edges into a portable 8-column triplet contract — subject_id, subject_type, predicate, object_id, object_type, confidence, source_method, source_agent — so the same graph can be emitted to or consumed from another graph stack without migrating tables. It is a column contract, not a dependency: the underlying tables are unchanged and nothing new is imported. Provenance in an edge’s props wins where present; otherwise source_method is derived from the relation and confidence defaults to 1.0, with the numeric cast guarded so a malformed props value becomes NULL rather than 1.0, so a consumer filtering on high confidence cannot silently ingest malformed rows as gold. Undirected relations (SUBSTITUTE_FOR) are stored once with sorted endpoints and made bidirectional only at query time, so the view emits both directions for them — a directed consumer would otherwise miss the reverse.

How retrieval works

The retrieval SQL is three stages in one statement:

Semantic seed — an HNSW ANN scan finds the entry nodes closest to the question embedding, ordered by embedding <=> :query_embedding (pgvector’s <=> is cosine distance, so similarity = 1 - distance). Cosine is used because embeddings are direction-, not magnitude-, meaningful.
Graph expansion — a WITH RECURSIVE walk follows edges out from each seed up to :max_hops (typically 2 — “two degrees of separation”). Edges are first materialized in both directions (UNION ALL of forward and reverse) so an undirected relation like SUBSTITUTE_FOR traverses either way. A per-path visited array guards against cycles.
Blended ranking — each reachable node is scored
```
graph_score = MAX(seed_similarity × 0.5 ^ hop)
```
The 0.5 ^ hop decay halves a node’s contribution per hop, so a direct neighbor of a strong seed outranks a distant one; MAX means a node reachable from several seeds keeps its strongest path. The top 25 by graph_score (with the relationship types traversed) become the LLM context.

This is the GraphRAG win: a supplier or substitute that shares no keywords with the question still surfaces because it’s one hop from a semantically-matched node — something flat vector RAG never retrieves.

Building the graph from documents

The example above starts from structured dims. When the source is documents, graph_upstream.py covers the indexing phase that comes first: parse with ai_parse_document, chunk into passages with stable ids, extract typed triples with ai_query under a constrained schema, resolve entities, and emit the gold_triplets contract. The two steps needing a workspace live in sql/upstream_ai_functions.sql.

Two lessons from running it against a live workspace are worth carrying into your own build:

Entity resolution is the load-bearing step, and Jaccard alone is not enough. The same organization appears as “Acme”, “Acme Foods” and “ACME Foods Inc.”; each spelling would otherwise become its own node and fragment the graph precisely where multi-hop retrieval needs it joined. Jaccard trigram similarity scores “Acme Foods” against “Acme Foods Incorporated” at just 0.42 — below any sane threshold — because it penalizes length difference, and legal suffixes are the most common way one entity is written two ways. resolve_entities() therefore also applies a containment measure, the shape of pg_trgm’s word_similarity(), which scores that pair 0.92. Every merge is written back as a SAME_AS edge, so the resolution is auditable in the graph.

Pin the relation vocabulary, not just the entity types. Left free, the model invents a predicate per sentence: a live run over two short documents produced SUPPLIES_TO_RETAILERS_ACROSS, OPERATES_DISTRIBUTION_CENTER_IN, IS_LOCATED_IN and BELONGS_TO_CATEGORY — the last two near-misses for this schema’s LOCATED_IN and BELONGS_TO. Each spelling becomes a distinct edge label, and typed-path traversal degrades toward an untyped walk. Constrain the predicate list in the extraction prompt and pass allowed_predicates=DEFAULT_PREDICATES, on_unknown="drop" as a backstop.

Building the graph safely

assemble_graph() in graph_build.py is a pure function that turns rows + LLM enrichment into (nodes, edges). Its add_edge() only adds an edge if both endpoints exist — so hallucinated substitute ids or orphaned supply rows are dropped and logged rather than creating dangling references. Undirected SUBSTITUTE_FOR edges are stored once (endpoints sorted) to avoid duplicates.

The retrieval and build logic is validated entirely offline by smoketest/graphrag_logic_smoketest.py — 82 assertions (on DuckDB, no Lakebase or model endpoint needed) covering the semantic seed, graph expansion surfacing context flat RAG misses, the dangling-edge guard, 0.5^hop score decay, the max_hops depth bound, and the seed_floor distractor guard.

Deploy with Asset Bundles

Prerequisites: a Databricks workspace with Lakebase (Autoscaling) enabled, a model-serving embeddings endpoint and a chat endpoint registered, plus the databricks CLI and uv.

cd agents/graphrag
databricks bundle deploy -t dev \
  --var lakebase_database="projects/<project>/branches/<branch>/databases/<id>"
databricks bundle run graphrag_build -t dev

The bundle deploys notebooks/graphrag_build_and_query.py as a job that assembles a small example supply-chain graph, embeds its nodes, writes to Lakebase, and queries it. The two Lakebase I/O cells are scaffolding you complete (the Postgres connection is workspace-specific); the retrieval logic itself is fully validated offline by the smoke test:

cd agents/graphrag
uv run --python 3.11 --with duckdb --with numpy smoketest/graphrag_logic_smoketest.py

Configuration and tuning

Variable / setting	Purpose
`lakebase_database`	Full Lakebase database resource path (required).
`max_hops`	Traversal depth from the seeds. `2` is the sweet spot; higher pulls in more distant (and lower-scored) context at the cost of a wider recursive walk.
`seed_floor`	Minimum cosine similarity a semantic seed must clear before it enters the graph walk. Cosine ranges `[-1, 1]`, so the default `-1.0` keeps every seed (identical to no floor). Raise it on distractor-heavy corpora, where weak seeds bridge into unrelated subgraphs and dilute precision. Calibrate to your embedding model’s observed range rather than a fixed constant: measured live with `databricks-gte-large-en`, this graph’s similarities span `0.38`-`0.71`, so `0.3` is a no-op and `0.555` is where the distractors drop. Required bind — Postgres has no server-side default for a named parameter, so SQL callers must pass it (`-1.0` for the old behavior); the Python twin defaults it. A floor above every similarity empties the seed CTE and returns zero rows, so the caller should retry at `-1.0` or decline to answer.
`VECTOR(1024)`	Embedding dimension — must match your embeddings endpoint (1024 for `databricks-gte-large-en`).
HNSW `m` / `ef_construction`	Index build quality vs. speed (`16` / `64` here). Raise for higher recall on larger graphs.
Embeddings endpoint	Model-serving endpoint used to embed nodes and questions.
Chat endpoint	Model-serving endpoint used to synthesize the final answer.

Edit this page →