Augmenting the RAG System on AstroLlama

Enhancing the ingest software in AstroLlama to load Astronomy books

As we saw last time, RAG means Retrieval Augmented Generation which allows you to extend your model to include additional information, essentially giving it a library of text to search to provide through a database of text “chunks” Simply, if a keyword is detected in a chunk it is provided to the LLM in it’s context window which allows it to provide that info to the user. Along with Model Context Protocol (MCP) this is how LLMs have their knowledge base improved without additional training.

In my case I have a very large library of scanned Astronomy Books (I went totally manic on scanning books at one time, to get rid of a large paper library) I note since I am using my personal library for copyright reasons I cannot share my RAG library, just the tools to build your own.

Below I asked Claude to summarize a rather lengthy development session implementing extensive RAG for AstroLlama.

The Basic Problem

A RAG system needs to break documents into chunks, embed them into a vector database, and at query time pull the most semantically relevant chunks into the LLM’s context window. The naive approach — split every 500 characters, embed, done — works for clean text but falls apart quickly with astronomy PDFs. Books are a mix of typeset text, embedded equations, tables, star charts, and telescope images. A fixed character split will happily cleave a sentence in half, separate a figure caption from its image, or merge two unrelated entries. You also lose all the images entirely.

The new ingestion subsystem in AstroLlama tries to do better on all three counts: smarter text extraction, semantic chunking, and image extraction with preserved links back to the source material.

Two Processing Paths

The first thing the ingester does for any PDF is ask: is this a scanned document, or does it have real embedded text?

It does this by sampling the first few pages with pdfplumber and counting the average characters extracted. Below 50 characters per page, it assumes the PDF is an image scan (common with older astronomy texts digitized from print). Above that, it treats it as a text-based document.

Path 1 — Text PDFs → Marker

For text-based PDFs, the pipeline runs Marker, a document conversion model that uses a combination of layout detection (Surya), OCR, and a fine-tuned model (Texify) for mathematical notation. The output is clean Markdown that preserves headings, lists, tables, and equations far better than raw pdfplumber extraction. Marker also extracts embedded figures as separate image files. The resulting Markdown goes straight into the chunker. Images are saved to data/images/<pdf_stem>/ and their relative paths are stored in chunk metadata.

Path 2 — Scanned PDFs → unstructured hi_res

For scanned pages, Marker can’t help much — there’s no embedded text to work with. Instead the pipeline uses unstructured with the hi_res strategy, which renders each page at 200 DPI and runs a YOLOX layout detection model to classify regions: title, narrative text, figure, table, and so on.

Text regions get their text extracted and assembled. Figure regions get their bounding boxes handed to PyMuPDF (fitz), which renders and crops those exact rectangles at 2× resolution and saves them as PNG files. This means even scanned-document figures end up in data/images/ with predictable names and can be linked back to the chunks from the same page.

Semantic Chunking with Chonkie

Whether the text came from Marker or unstructured, it goes through Chonkie’s SemanticChunker. Rather than splitting on character count, it uses sentence embeddings to find natural semantic boundaries — grouping sentences that are topically cohesive and splitting where the subject changes. For an encyclopedia this works particularly well because entries tend to be self-contained topics; the chunker keeps them together rather than splitting mid-entry.

Chonkie is declared mandatory — the ingester exits immediately with a clear error if it’s not installed. The silent-fallback-to-character-splitting approach was tried first and produced noticeably worse retrieval.

Default settings: 512 tokens per chunk, similarity threshold 0.5. These can be tuned per-document with --chunk-size and adjusting the threshold in code.

ChromaDB Metadata

Each chunk stored in ChromaDB carries more than just the text. The metadata schema:

source        — absolute path to the original PDF
source_title  — stem of the filename (used for display)
chunk         — index within the document
page_num      — page number if known (0 = whole-document, e.g. from Marker)
has_images    — boolean flag
images        — JSON-encoded list of relative paths under data/images/

The images field being a JSON string is a ChromaDB constraint — it only accepts scalar metadata values. At retrieval time the application decodes it back to a list.

Surface Images in the Chat UI

The FastAPI app mounts data/images/ as a static directory at /images, so any saved figure is immediately web-accessible at http://localhost:8080/images/<stem>/p3_fig0.png.

The retriever has a query_with_metadata() method (alongside the original query() for backwards compatibility) that returns full metadata with each chunk. The tool orchestrator calls this instead of the bare query, injects figure references like [Figure: /images/...] into the LLM’s system context so the model knows images exist and can reference them, and emits tool_image SSE events for each associated figure so the frontend renders them inline alongside the response text.

Dedicated Ingest Container

Marker’s models (Surya + Texify) weigh around 1.5 GB, and unstructured-inference’s YOLOX layout model adds another ~400 MB. Since we don’t need any of that in the app or MCP containers that serve live requests, I created a dedicated ingest Docker service separate from the main application. In case you’re not familiar with Docker it allows you to create virtual machines in “containers” that can be configured using simple files and rebuilt as needed. AstroLlama now runs completely on containers to simplify running it. On Windows just install Docker Desktop and run the ./start.ps1 script.

The service uses a Docker Compose profile so it never starts with a normal docker compose up:

# First-time build
docker compose --profile ingest build ingest

# Test extraction on the first 20 pages, produce an HTML report
docker compose --profile ingest run --rm ingest \
  python scripts/test_pdf_ingest.py \
  --pdf “/books/Encyclopedia of Astronomy and Astrophysics.pdf” \
  --pages 1-20

# Full ingest into ChromaDB
docker compose --profile ingest run --rm ingest \
  python scripts/ingest.py \
  --source “/books/Encyclopedia of Astronomy and Astrophysics.pdf” \
  --extract-images

Model weights are cached in a named Docker volume (marker_cache) so they survive container removal and don’t re-download on every ingest run. The book library is mounted read-only via a BOOKS_PATH env variable in .env.

The HTML Test Report

Before committing a whole book to ChromaDB it’s useful to validate extraction quality. scripts/test_pdf_ingest.py processes a page range and writes a self-contained HTML report showing, for each page, the detected layout elements, extracted images as inline thumbnails, and the Chonkie chunks with their token counts. This lets you catch problems — a scanned page that needed the layout path, OCR noise in figure captions, chunk boundaries that split a table — before ingesting the full document.

What’s Still Rough

The main limitation right now is that Marker produces one big Markdown string for the whole document rather than per-page output, so image-to-chunk association after Marker is coarse (all images get attributed to the whole document rather than the specific page they came from). Unstructured gives better per-page association but is slower and less clean on text quality. A future improvement would be to use Marker’s internal page-level data structures to do finer-grained association. The other known issue is that for very large books (1000+ pages), unstructured’s hi_res path is slow even on GPU — plan for 2–5 seconds per page. Marker is faster but still not instant. Batch ingestion of a whole library is an overnight job, not interactive.