Skip to content

Repository files navigation

Structure-Aware RAG on Databricks

Structure-aware RAG pipeline from governed documents to cited evidence

A portfolio-grade document intelligence pipeline that preserves the two things chunk-and-embed throws away: document structure and evidence provenance.

CI Python 3.10+ Tests Coverage License

Results · Deployment proof · Architecture · Quick start · Design decisions · Documentation

Portfolio snapshot

Capability Verifiable result
Public corpus Date-pinned 21 CFR Parts 11, 210, and 211 from eCFR
Databricks workflow Five serverless tasks completed successfully on Free Edition
Governed retrieval data 82 chunks; 60 relation rows (47 resolved, 13 out-of-corpus)
Relationship-aware evaluation 5/5 in-scope passes, 10/10 citation and evidence recall
Safety behavior 2/2 out-of-scope questions correctly abstained
Engineering quality 200 tests, 83% coverage, Ruff, mypy, package and CLI checks

Live evaluation

Live Databricks retrieval evaluation

The same reviewed set ran successfully on Databricks Free Edition with databricks-gte-large-en and a delta-sync Vector Search index. At the calibrated 0.66 score floor, the live baseline produced 2/5 question passes and 7/10 citation/evidence recall; persisted CFR edges improved it to 5/5 and 10/10, while both configurations abstained on 2/2 out-of-scope questions.

This panel is reproducibly generated from reviewed, authenticated Databricks CLI output by scripts/render_portfolio_evidence.py. It intentionally omits workspace URLs, user identities, run IDs, endpoint IDs, and storage paths.

Quick start

Python 3.10 or newer is required. Install the pinned development environment and run the complete local validation suite:

python -m pip install --upgrade pip==26.2
python -m pip install -r requirements-dev.txt
python scripts/validate.py

For a one-second, service-free walkthrough, run the synthetic pipeline. No Databricks workspace, model endpoint, or network access is required:

python pipelines/run_local_demo.py

For the real public-data demonstration, acquire the checksum-pinned eCFR subset and run the same pipeline locally:

python scripts/acquire_ecfr_demo.py --check
python scripts/acquire_ecfr_demo.py
python pipelines/run_public_demo.py

The public corpus contains 21 CFR Parts 11, 210, and 211 as of 2026-07-31. Complete raw XML stays under ignored data/raw/; the repository retains the official URLs, expected hashes, attribution, and a minimal structural fixture.

The committed public gold set currently contains five in-scope and two out-of-scope questions. On the deterministic local token-vector index:

Retrieval configuration Questions Citation recall Evidence recall Abstention
Vector + hierarchy baseline 2/5 6/10 6/10 2/2
Baseline + exact one-hop CFR edges 5/5 10/10 10/10 2/2

These are small public-demo results, not production accuracy claims. The relationship-aware path only follows explicit, exact, in-corpus CFR citations, adds referenced sections as separately citable evidence, and never traverses beyond one hop.

Live Databricks proof

Successful workflow

Open the complete five-task workflow capture

Successful five-task Databricks public-demo workflow

The portfolio-safe capture shows the complete serverless workflow succeeding: storage bootstrap, pinned-corpus verification, extraction and authoring, governed Delta ingestion plus index verification, and live retrieval evaluation. Workspace URLs, identities, resource IDs, and private data are excluded from the image.

Governed tables and index

Governed Databricks table and Vector Search state

How the exact one-hop path changes the evidence

One-hop CFR retrieval trace

For the reviewed Part 211 definitions question, the vector hit correctly finds the pointer in 21 CFR 211.3. The baseline stops there. The relationship-aware path resolves the observed § 210.3 citation against the pinned corpus and adds 21 CFR 210.3 as a separate, citable evidence chunk. Generated text is not used as evidence, and traversal stops after the single exact edge.

Those citations are also written as governed relation_edges rows containing source/target node IDs, the observed mention, corpus version, resolution status, and a deterministic content hash. The Databricks notebook writes both chunks and edges with scope-safe Delta snapshot merges: stale rows are deleted only inside the active project and collection, never by overwriting the shared table.


The problem

Here is a real paragraph from a service-desk procedure, and the chunk a fixed-window splitter produces from it:

Targets are suspended while a request is waiting on information from the requester. Time spent in the waiting state does not count toward the 8 business hour target.

Retrieved on its own, this chunk is unsearchable and unsafe.

Unsearchable, because it never says what it is about. It does not contain "response target", "SLA", or "escalation" — the words someone would actually type. Embed it and you get a vector for a sentence about suspension, floating free of the thing being suspended.

Unsafe, because the 8-hour target it qualifies is defined in the parent section. An agent who retrieves only this chunk knows there is an exception but not what it is an exception to.

Now the same chunk with its heading breadcrumb and its parent attached:

Response Targets > Exceptions to Response Targets
  Targets are suspended while a request is waiting on information from the
  requester. Time spent in the waiting state does not count toward the 8
  business hour target.

  surrounding context:
  [Response Targets]
  Standard requests receive a first response within 8 business hours. The clock
  starts when the request is acknowledged in the queue, not when it is submitted.

Same text. Now findable, and now correct. That difference is what this project is about, and the demo prints exactly this.

The second problem is quieter. Once a language model rewrites a passage to make it more retrievable, which text is the evidence? If the rewrite is what gets returned, the system is citing the model to itself. Getting that boundary wrong does not produce errors — it produces confident, fluent, unfalsifiable answers.


What it does

upstream source
    │
    ▼
governed landing zone ──────────── stable identity, auditable point-in-time state
    │
    ├─ 1  discover ─────────────── catalogue, classify, content-hash
    ├─ 2  plan ─────────────────── per-file capabilities; skips carry reasons
    ├─ 3  extract ──────────────── file-aware; one output contract
    ├─ 4  validate ─────────────── deterministic checks, before model spend
    ├─ 5  detect references ────── cross-document candidates
    ├─ 6  resolve references ───── corpus-wide; ambiguity stays unresolved
    ├─ 7  author ───────────────── provenance to front matter, prose to body
    ├─ 8  chunk ────────────────── heading-bounded, breadcrumb-carrying
    ├─ 9  enrich ───────────────── model normalisation behind a verified gate
    └─ 10 index ────────────────── merge to table, sync index from table
              │
              ▼
    retrieval ──────────────────── semantic search + hierarchical expansion
              │
              ▼
    context assembly ───────────── cited, provenance-labelled, budgeted

Two rules run through all of it

1. Generated text is a retrieval aid, never evidence.

Every chunk carries two text fields. retrieval_text is what gets embedded — it may be model-written and it carries the breadcrumb, because its only job is matching a query. evidence_text is the document's own words, and it is the only thing that may be quoted. When a model rewrites a passage, the rewrite becomes retrieval text and the original stays as evidence. The chunk gets easier to find without the rewrite ever becoming what the document said.

The gate is verified, not trusted. A model's self-reported groundedness is a claim; verify_grounding() checks the generated text against its source independently — every number, amount, and identifier must survive, and vocabulary must overlap — and the weaker of claimed and measured wins. A section that changes 500 EUR to 9999 EUR is marked unsupported and never reaches the index, no matter how confidently it was labelled.

2. Semantic search and hierarchical expansion are complementary.

Search finds candidates; it is indifferent to document structure. Expansion repairs context by walking the section tree rebuilt from chunk breadcrumbs, and attaching the parent and siblings within budget. Search alone returns fragments missing their conditions. Expansion alone has nowhere to start.


What the demo prints

See abridged local-demo output

Real output, abridged:

  discovered          7 file(s)
  planned             7 to extract, 0 skipped
  authored            7 knowledge-base document(s)
  raw chunks         28
  knowledge chunks   28 (25 indexable)
  references          3 resolved, 2 unresolved candidates
  quality         6 passed, 1 warning

Derivation:            Groundedness:
  observed      28       grounded          25
                         not_applicable     3

Excluded from the index: 3 chunk(s)
  [not_applicable] queue_naming_conventions.md > Scope
  [not_applicable] queue_naming_conventions.md > Source
  [not_applicable] queue_naming_conventions.md > Extraction Notes

procedures/escalation_handling.md
- escalation_handling.md
  - Response Targets
    - Priority Requests
    - Exceptions to Response Targets
    - Tier One
    - Tier Two

Then retrieval on the hard query, and the evaluation:

Query: 'does time waiting on the requester count toward the response target'
Status: found via vector, 1 chunk(s), ~152 tokens
Threshold in force: 0.15 (from the vector backend), best match 0.67

  hit rate         100.0%   right document retrieved
  MRR               1.000   rank of the right document
  evidence rate    100.0%   expected text present in evidence
  abstention rate  100.0%   correctly returned nothing
  quotable share   100.0%   chunks citable as source evidence

These synthetic results validate the local pipeline contracts; they are not an accuracy claim. The public eCFR evaluation above is the portfolio comparison, and docs/evaluation.md describes the denominators and failure modes each metric is designed to expose.

abstention rate is a first-class metric for the same reason. Optimising recall alone reliably produces a system that never says "I don't know", which scores well on every other metric and is the dangerous failure mode.


Design decisions worth defending

A governed landing zone, not reads against the upstream system. Stable identity (a document's id survives re-uploads and edits), auditable point-in-time state (the manifest records the content hash actually processed), and one access boundary instead of credentials for every source. Transfer is a connector concern — see sources.py.

Provenance in front matter, never in the body. Body prose becomes chunks, so an authoring format that writes metadata as ## Source sections puts a near-identical block into every document in the corpus. Those blocks then compete with real content for every query: short, topically generic, and as numerous as the corpus. The chunker still recognises and excludes them, because migrated corpora have them — that is what the three not_applicable chunks above are.

The index syncs from the table, not from the pipeline. The table is authoritative and the index is a disposable projection. Rebuilding an index never re-runs a model call, and "why is this in the index?" is a table query.

The similarity threshold comes from the backend. Embedding cosine puts a good match near 0.8; TF-IDF cosine on a short query rarely passes 0.4. A threshold hard-coded in the retrieval layer is wrong for every backend but the one it was tuned against — silently returning noise on one and nothing at all on the other. Each backend declares its own and retrieval asks.

Skips are decisions, not omissions. Every skipped file carries a reason. 900 of 1,000 files skipped because content is unchanged is a healthy incremental run; 900 skipped because an extractor is missing is an incident. In a total count those are identical.

Retrieval logs record candidates and selections. A log holding only what was returned cannot distinguish a retrieval miss (nothing relevant found) from a selection problem (the right chunk ranked fourth, budget cut it at three). Those have opposite fixes.

Everything service-dependent sits behind a protocol. LLMClient, VectorIndex, TableStore, SourceConnector each have a working local implementation and a Databricks adapter. That is why the pipeline runs offline and why 200 tests need no workspace — not a testing compromise, a structural property.


Layout

Browse the package and repository map
src/structured_kb/
  contracts.py       stage-to-stage dataclasses; Derivation and Groundedness
  config.py          env-driven settings; placeholder defaults only
  sources.py         landing-zone connectors
  discovery.py       walk, classify, hash, manifest
  planning.py        per-file capabilities, incremental skip, run state
  extraction/        markdown, delimited, json + registry
  quality.py         deterministic checks; issues vs warnings
  relations.py       reference detection and conservative resolution
  authoring.py       front-matter provenance contract
  chunking.py        heading-bounded chunker + SectionTree
  enrichment/        prompts, LLM clients, the verified groundedness gate
  indexing.py        TableStore + VectorIndex (local and Databricks)
  retrieval.py       search, expansion, budget, context assembly
  evaluation.py      hit rate, MRR, evidence rate, abstention rate
  observability.py   processing and retrieval telemetry
  demo.py            synthetic corpus and golden queries
  pipeline.py        orchestration
  cli.py             demo / discover / config / query

notebooks/           Databricks source-format notebooks (no output cells)
pipelines/           run_local_demo.py
tests/               200 tests, no service dependencies
docs/                architecture, data contracts, retrieval, evaluation, roadmap

Usage

Python 3.10 or newer is required. Python 3.10 is the supported floor because it matches both the CI matrix and Databricks Runtime 13.3 LTS; CI also exercises Python 3.12.

python -m pip install --upgrade pip==26.2
python -m pip install -r requirements-dev.txt

# Exact local equivalent of CI: lint, types, tests with an 82% coverage floor,
# package build, dependency check, compilation, CLI help, and local demo.
python scripts/validate.py

structured-kb discover ./data/landing         # read-only: what would a run do?
structured-kb config                          # resolve and validate settings
structured-kb query ./data/landing "your question"

Against a real workspace: copy .env.example to .env, fill it in, then run notebooks/01_extract_and_author.py and notebooks/02_ingest_and_index.py.

Databricks public-demo deployment

databricks.yml packages the project as a demo-only Asset Bundle. Its job creates a schema and Volume inside an existing demo catalog, verifies three staged eCFR files against the checksum-pinned manifest, authors the corpus, merges governed chunk and relation rows, syncs Vector Search, and runs the reviewed evaluation. Endpoint names are required deployment inputs; no credential or workspace identifier is committed.

See docs/databricks_demo_runbook.md for validate, plan, deploy, run, evidence, and teardown commands. Run databricks bundle plan before the first deployment; local CI remains service-free even though the public demo has also completed in an isolated live workspace. Start with dry_run=true — it reports every action and writes nothing.

The core package has no dependencies. requests and databricks-vectorsearch are needed only for the remote adapters.


What this is not

  • Not a production retrieval engine. LocalVectorIndex is TF-IDF: no embeddings, no ANN, no semantic generalisation. It exists so the layers above it — filtering, thresholds, expansion, budgeting, provenance — are fully exercised without a service.
  • Not a complete extractor set. Markdown, delimited, and JSON are implemented. PDF, DOCX, XLSX, and PPTX are declared in extraction.NOT_IMPLEMENTED so the planner reports the gap instead of crashing on it.
  • Not redistributing raw public source files. The repository includes a checksum-pinned acquisition manifest, official eCFR URLs, attribution, and a minimal structural fixture. Full raw XML is downloaded into ignored data/raw/ so the demo stays reproducible without committing a data snapshot.
  • No OCR, no reranking, no incremental reference resolution. Documented in docs/roadmap.md rather than implied.

Documentation

architecture.md Stage-by-stage flow and why the boundaries fall where they do
data-contracts.md Every dataclass and table column, with provenance semantics
retrieval.md Search, expansion, thresholds, budgets, context assembly
evaluation.md The four metrics, current numbers, and the known miss
roadmap.md What is deliberately absent, and what closing each gap needs

Licensed MIT.

About

Structure-aware RAG on Databricks with governed Delta data, exact CFR relationship expansion, and reproducible public-data evaluation.

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages