Visual-citation RAG for M&A / investment data rooms. Ask a question, get a grounded answer with the exact region highlighted on the original scanned page — retrieved over page images, not lossy OCR text, verified before it's shown, and costed in dollars on every response.
Ask a question → the relevant region resolves on the page → a grounded, cited answer with its cost. · ▶ full-quality MP4
Why retrieve over page images instead of OCR text? In real filings the answer lives in the layout — a number in a table cell, a signature block, a stamp, a figure in an infographic. Text-only RAG flattens the page and throws that structure away. Atlas keeps the page intact and lets the model look at it.
- What it does
- What makes it different
- System architecture
- The query pipeline (LangGraph)
- Visual citations — the headline feature
- Hybrid retrieval
- Ingestion
- Observability & cost
- Engineering decisions & trade-offs
- Measured on the demo corpus
- Tech stack
- Quick start
- Project structure
- Testing
- Roadmap
- License
Atlas ingests a heterogeneous data room — scanned PDFs, contracts, financial tables, charts — and answers analyst questions the way a diligence associate would: by pointing at the page.
- Ingest — render every page to an image, embed it with ColPali (late-interaction visual retrieval) on HF ZeroGPU, embed the extracted text with bge-m3, and pull named entities with GLiNER into a Neo4j knowledge graph.
- Retrieve — fuse three channels (visual + dense + graph) with Reciprocal Rank Fusion, then re-rank with a cross-encoder.
- Answer — Gemini 3.1 Flash Lite reads the actual page images and returns an answer with a
per-claim
{page, quote}. - Verify — a second pass re-checks every claim against its cited page and refuses, re-retrieves, or escalates to a human rather than hallucinate.
- Cite — the ColPali patch-similarity map is the highlight: the answer ships with a cropped page region + bounding box + heatmap.
- Account — every LLM call and every pipeline node is traced and costed; each response carries a
meta.costbreakdown in USD.
The demo corpus is a real public SEC S-1 filing (Uber Technologies, 407 pages) fetched from EDGAR.
| Atlas | Typical "chat with PDF" | |
|---|---|---|
| Retrieval unit | The rendered page image (ColPali late interaction) | OCR'd text chunks |
| Layout awareness | Preserved — tables, figures, stamps, signatures | Flattened / lost |
| Citations | Pixel-level: crop + bbox + heatmap on the scan | Page number, at best |
| Groundedness | Verifier gate → refuse / re-retrieve / escalate | Answer as-is |
| Retrieval channels | Visual + dense + entity-graph, RRF-fused + reranked | Single dense channel |
| Ops | Per-call $ + latency, budget guard, Langfuse/Grafana/ledger | None |
This isn't a wrapper around an LLM — it's the retrieval, verification, and observability machinery a real diligence tool would need, built from scratch behind clean interfaces.
flowchart LR
subgraph Ingest["Ingestion — idempotent, sha256-deduped"]
PDF["S-1 / data-room PDF"] --> R["PyMuPDF render<br/>page to PNG"]
R --> CP["ColPali embed<br/>(HF ZeroGPU)"]
R --> TX["Text + word-bbox<br/>extract and chunk"]
TX --> D["bge-m3 dense"]
TX --> NER["GLiNER entities<br/>(windowed)"]
end
subgraph Stores["Stores"]
PG[("Postgres + pgvector<br/>multivectors · pooled ANN<br/>dense · cost ledger")]
NEO[("Neo4j<br/>entity graph")]
RD[("Redis<br/>dedup · cache · HITL")]
end
CP --> PG
D --> PG
NER --> NEO
subgraph Query["Query — LangGraph"]
Q["Analyst question"] --> API["FastAPI"]
API --> GR["Planner → Retrieve → Answer → Verify"]
end
PG --> GR
NEO --> GR
RD --> GR
GR --> ANS["Answer + visual citation<br/>+ meta.cost"]
GR -. traces .-> OBS["OTel → Langfuse<br/>Prometheus / Grafana"]
The control plane stays laptop-light. ColPali (a 3B PaliGemma-based model) runs remotely on ZeroGPU behind a small Gradio Space; MaxSim scoring and citation heatmaps are pure numpy, so torch is not in the core install. Gemini is a hosted API. Only Postgres and Redis need to run locally (Neo4j, Prometheus and Grafana are opt-in compose profiles).
The pipeline is a StateGraph. The interesting part is the gate after the verifier — it balances
answer recall against API spend instead of blindly retrying:
flowchart TD
P["planner"] --> R["retriever<br/>visual + dense + graph → RRF → rerank"]
R --> A["answerer<br/>Gemini reads page images →<br/>answer + per-claim citations"]
A --> V{"verifier<br/>re-checks each claim vs cited page image"}
V -->|"grounded"| F["finalize<br/>answer + visual citation + meta.cost"]
V -->|"unverified — retries remain"| RP["replan · widen retrieval"]
V -->|"unverified — budget spent"| H["hitl · escalate to human review"]
V -->|"no evidence found"| F
RP --> R
H --> F
- Grounded → return the answer with its citation and cost.
- Drafted but unverified, retries remaining → re-plan (widen retrieval via
extra_k) and try again — up tomax_retries. - Still unverified after the retry budget → escalate to a Redis-backed human review queue (
/review/queue,/review/resolve); the reviewer's decision is logged as eval data. - Model found nothing → refuse immediately — no wasted retry loops, no wasted tokens.
ColPali is a late-interaction model: each page is stored as a multivector — one embedding per image patch — and each query as one embedding per token. The relevance score is MaxSim:
score(query, page) = Σ_i max_j ( q_i · p_j ) # sum over query tokens of best-matching patch
The elegant part: the per-patch term of that same dot product is the citation heatmap — no extra model, no separate "highlighter":
relevance(patch_j) = max_i ( q_i · p_j ) → reshape to (ny, nx) grid → upsample → threshold → bbox
So the highlight drawn on the scan is computed locally in numpy from vectors already in the database. Click an answer in the Streamlit demo and the exact region lights up on the original page.
Three independent channels are fused, then reranked with a cross-encoder:
flowchart TD
Q["Query"] --> VE["ColPali query embed<br/>(ZeroGPU)"]
Q --> DE["bge-m3 dense embed"]
Q --> GE["GLiNER query entities"]
VE --> VIS["Visual channel<br/>pooled ANN → exact MaxSim"]
DE --> DEN["Dense channel<br/>pgvector cosine"]
GE --> GRAPH["Graph channel<br/>Neo4j entity + 1-hop co-occurrence"]
VIS --> RRF["Reciprocal Rank Fusion"]
DEN --> RRF
GRAPH --> RRF
RRF --> RR["bge-reranker-v2-m3<br/>cross-encoder"]
RR --> OUT["top-k candidates → answerer"]
Scaling the visual channel. Exact MaxSim over every page's full multivector is accurate but grows
linearly with the corpus. Atlas stores a mean-pooled, L2-normalised single vector per page in a
pgvector HNSW index, uses it for a coarse ANN shortlist (top-visual_coarse_n), then runs exact
MaxSim only on the shortlist — turning a linear scan into a sub-linear one, with a full-scan
fallback if the pooled table isn't backfilled yet. On the 407-page corpus this took the visual scan
from 734 ms → 181 ms (ADR 0001).
The graph channel extracts entities from the query with GLiNER, matches them against the Neo4j graph, and expands one co-occurrence hop — so "who advised on the offering?" can reach the underwriting section through the entities it mentions, even when the page shares few keywords with the question. If Neo4j is unavailable the channel fast-fails and degrades to visual + dense.
Ingestion is idempotent (sha256 + Redis SETNX) and atomic (a partially-ingested document is
deleted on failure, and the dedup lock is always released in a finally):
sequenceDiagram
participant U as atlas ingest
participant P as Pipeline
participant R as Redis
participant Z as ColPali Space · ZeroGPU
participant DB as Postgres+pgvector
participant N as Neo4j
U->>P: ingest_pdf(path)
P->>R: SETNX sha256 lock (dedup)
P->>P: PyMuPDF render pages to PNG
P->>P: extract text + word bboxes, chunk
P->>Z: embed_pages(images)
Z-->>P: multivectors + patch grids
P->>DB: store multivectors + pooled ANN vectors
P->>P: bge-m3 dense embeds
P->>DB: store dense vectors + chunks
P->>P: GLiNER windowed NER
P->>N: upsert entity graph (best-effort)
P->>R: release lock (finally)
Note over P,DB: atomic — partial doc deleted on failure
Entity extraction runs GLiNER over overlapping ~1,200-char windows and dedupes, instead of truncating the page — a long prospectus page yields dozens of entities the naïve single-pass approach would silently drop.
Cost and latency are first-class, not an afterthought — this is a system designed to run on ~$3 of Gemini credit without surprises.
- Per LLM call: read
usage_metadata(input / output / thinking tokens) → price it frompricing.yaml→ write a row to the Postgres cost ledger (llm_calls,request_costs). - Per node: each pipeline node opens an OpenTelemetry span (GenAI semantic conventions) with latency; ColPali calls also record ZeroGPU seconds.
- Three sinks: Langfuse (trace-level $/latency + dashboards), Prometheus + Grafana
(provisioned dashboard,
observabilityprofile), and the SQL ledger for rollups and cost-per-answer in the eval report. - Guardrails: per-request USD cap, a hard daily budget (
max_usd_per_day), and a Redis answer cache keyed on(question, index_version)to avoid paying twice for the same question.
Every /query response includes:
| Decision | Why | Trade-off / notes |
|---|---|---|
| Retrieve over page images (ColPali) | Preserves layout; enables pixel citations | Heavier per-page store than text (ADR 0002) |
| ColPali on HF ZeroGPU, not local | 16 GB Mac can't host a 3B VLM; keeps torch out of core | Network hop + cold starts; downscale images before upload |
| MaxSim + heatmaps in numpy | No torch in the control plane; citations are "free" | Re-implements late interaction by hand (unit-tested) |
| Gemini 3.1 Flash Lite | Cheap enough to run a verifier pass on every answer | function_calling structured output (json-schema returned empty on this model) |
| Pooled-ANN two-stage visual search | Linear MaxSim doesn't scale | Approximate shortlist; exact rescoring recovers precision |
| Verifier re-retrieve loop, cost-aware | Recall without runaway spend | Adds latency on hard questions; capped by max_retries |
| Postgres + pgvector for everything | One store for dense, multivector, pooled ANN, and the ledger | VectorChord / Qdrant is the scale path (documented) |
Full write-ups live in docs/architecture.md and docs/adr/.
Observed during development on the Uber Technologies S-1/A (407 pages) — illustrative, not a published benchmark. The eval harness (
make eval) reports recall@k, nDCG@10, citation precision, groundedness, refusal calibration, Cohen's κ vs. human labels, and cost-per-answer.
| Metric | Value |
|---|---|
| Corpus | Uber Technologies S-1/A — 407 pages |
| Visual scan latency | 734 ms → 181 ms (pooled ANN vs. full MaxSim) |
| Entity graph | ~2,040 entities · 395 pages · ~25k co-occurrence edges |
| Cost per answered query | ~$0.002 (answerer + verifier, Gemini 3.1 Flash Lite) |
| Automated tests | 40 passing |
Orchestration LangGraph · LangChain · VLM Gemini 3.1 Flash Lite (Google AI Studio) ·
Visual retrieval ColPali vidore/colpali-v1.2-hf on HF ZeroGPU · Dense bge-m3 · Rerank
bge-reranker-v2-m3 · NER GLiNER · Stores PostgreSQL + pgvector, Redis, Neo4j · API/UI
FastAPI, Streamlit · Docs PyMuPDF, Playwright · Observability OpenTelemetry, Langfuse,
Prometheus + Grafana · Tooling Python 3.11, uv, ruff, pytest.
Prerequisites: Python ≥ 3.11, uv, Docker, a Google AI Studio API key, and a Hugging Face token with ZeroGPU access (HF Pro).
cp .env.example .env # fill GOOGLE_API_KEY, HF_TOKEN, COLPALI_SPACE_ID (LANGFUSE_* optional)
make install # uv sync (core + local models + ingest + graph) + playwright chromium
make up # Postgres+pgvector + Redis (make up-obs adds Prometheus+Grafana)
# Deploy the ColPali ZeroGPU Space once — see services/colpali_space/README.md — then:
make download CIK=0001543151 # fetch + render Uber's S-1 to data/raw/
make ingest PDF=data/raw/<file>.pdf
make serve # FastAPI on :8000
make ui # Streamlit demo on :8501Ask a question:
curl -s localhost:8000/query -d '{"question":"Who are the underwriters?"}' \
| jq '{answer, citations, cost: .meta.cost}'Optional V2 (GraphRAG + ANN):
make neo4j # start Neo4j (compose --profile v2)
uv run atlas backfill-pooled # pooled ANN vectors for existing pages
uv run atlas backfill-graph # load pages/entities into the Neo4j graphEvaluate & test:
make eval # recall@k, nDCG, citation precision, groundedness, κ, cost/answer → report
make test # 40 tests (Gemini + ColPali Space mocked — no keys needed)All secrets live in
.env, which is git-ignored. Only the blank.env.exampleis committed.
atlas/
├── src/atlas/
│ ├── models/ # colpali (zerogpu+local), maxsim (numpy), gemini, dense, reranker, ner
│ ├── ingest/ # download (EDGAR→PDF), render, extract, pipeline, kafka_consumer (V2 stub)
│ ├── index/ # pgvector_store, graph_store (Neo4j), schema
│ ├── retrieve/ # visual, dense, graph, fusion (RRF), rerank, hybrid
│ ├── pipeline/ # LangGraph: build.py + nodes/{planner,retriever,answerer,verifier,finalize,...}
│ ├── citations/ # heatmap.py (patch relevance → bbox), render_citation.py
│ ├── observability/ # tracing (OTel+Langfuse), cost, metrics (Prometheus), ledger, instrument
│ ├── hitl/ # Redis human-review queue (V2)
│ ├── eval/ # dataset, metrics, judge (Gemini + Cohen's κ), run_eval, report
│ ├── api/ # FastAPI app + routes (query, ingest, review, metrics)
│ └── cli.py # `atlas` command (download-s1, ingest, serve, backfill-*, seed-eval, eval)
├── services/colpali_space/ # Gradio ZeroGPU Space: embed_pages / embed_query
├── infra/ # postgres/init.sql, prometheus, grafana dashboards, k8s (stub)
├── docs/ # architecture.md + adr/
├── ui/streamlit_app.py # demo UI with the cost/latency strip
└── tests/ # unit + integration (external seams mocked)
make test # uv run pytest — 40 tests
make lint # ruff check + format --checkTests mock the external seams (ColPali Space, Gemini, Postgres, Redis) so the real orchestration is exercised — answer-vs-refuse-vs-escalate routing, RRF fusion, MaxSim math, heatmap→bbox, cost accounting — without needing a deployed Space or an API key. The integration suite drives the full LangGraph, including the verifier loop actually enqueuing to the HITL queue.
- V1 — shipped & running end-to-end. ZeroGPU ColPali · dense + RRF + rerank · Gemini answerer · verifier/refusal · pixel-level visual citations · SEC S-1 ingest · full cost/latency observability · Streamlit demo · eval harness with Cohen's κ.
- V2 — implemented. Windowed NER (entity-loss fix) · pooled-ANN two-stage visual search · Neo4j GraphRAG as a third retrieval channel · verifier re-retrieve loop + HITL escalation queue.
- V3 — scaffolded behind interfaces. Multi-tenant RLS + access-scoped citations · MNPI / conflict governance · Kafka ingest consumer · self-hosted vLLM provider · A/B retrieval routing · K8s manifests.
MIT © Ayush Gupta