MM-7: PDF chunk-level retrieval

[Atharva-Kanherkar]

Problem

Storing a PDF produces one embedding for the entire document. A single 768-dim vector captures the overall gist but cannot represent individual words, sentences, or pages. This means keyword-style queries like "medicine" or "prescription" score no better than noise (~0.5 similarity) even when the PDF literally contains those words.

This was observed directly:

• Stored a prescription PDF with content label "my prescription"
• Queried "prescription" → ranked 9th (sim=0.563), behind unrelated WAV files and smoke tests
• Queried "medicine" → ranked 16th (sim=0.496), below the noise floor

The root cause is architectural: one vector per document cannot support content-level retrieval.

What every RAG system does about this

Split the document into chunks, embed each chunk separately, retrieve at chunk granularity, deduplicate by parent document before returning results.

Proposed approach

1. Text extraction + chunking

For each PDF:

1. Extract text per page (using pymupdf or pdfplumber)
2. If a page is short enough for Gemini's token limit (~8000 tokens), embed the whole page as one chunk
3. If a page is too long, split at paragraph boundaries into sub-chunks
4. Each chunk becomes its own record

2. Storage model

Use Option A: Chunks are records. Each chunk is a SemanticMemory with a parent_id metadata field linking back to the original document record.

• Fits the existing store contract — no changes to BaseStore
• The parent document record can still hold the whole-document multimodal embedding (for visual/layout similarity)
• Chunks hold text-extracted content (for keyword/content retrieval)
• Both are searchable through the same retriever

3. Retrieval changes

Add a deduplication step to the retriever between ranking and returning:

• Group results by parent_id (if present)
• Keep the highest-scoring chunk as the representative for each parent
• Return deduplicated results

This means a query for "medicine" would match the specific page/paragraph that mentions medicine, and the result would surface the parent document with that chunk's similarity score.

4. Optional preprocessing layer (not in scope for this issue, but worth designing for)

Text extraction handles the keyword retrieval gap for PDFs. But for images and other non-text media, there's no text to extract. A separate preprocessing layer could use a vision/multimodal LLM (e.g., Gemini Flash) to generate text descriptions of media before embedding:

image → Gemini Flash "describe this image" → "a prescription label showing amoxicillin 500mg" → embed text

This is not part of the chunking issue. It's a separate concern:

• Chunking = split existing text into smaller pieces for better retrieval granularity. Cheap, local, no LLM.
• Preprocessing = generate new text from non-text media. Requires an LLM call per item, costs money at scale.

The architecture should keep these as separate, optional layers. Chunking should work without preprocessing, and preprocessing should work without chunking. If both are present, a preprocessed image description becomes just another chunk that the retriever handles normally.

The key design constraint: the LLM should never be in the retrieval path (per-query). Preprocessing is one-time at storage time. Retrieval stays embedding-only.

5. Scope decisions to lock

• Text-only chunks vs multimodal page chunks: Text extraction is simpler and sufficient for keyword retrieval. Multimodal page embeddings (render page to image + text jointly) preserve visual layout but are more expensive. Start with text-only?
• Chunk granularity: Page-level is the natural split for PDFs. Paragraph-level gives better precision but more vectors. Start with page-level?
• Parent record: Should the original whole-document embedding still be stored alongside chunks? It's useful for "find documents like this one" queries but adds storage.
• Overlap: Should adjacent chunks share some overlapping text (e.g., last 2 sentences of previous chunk prepended)? Improves retrieval for concepts that span page boundaries.

6. Repo impact

• New: text extraction utility (likely utils/pdf.py or utils/chunking.py)
• New: dependency on a PDF text extraction library
• Modified: stores/semantic_store.py — chunk-aware storage path for PDFs
• Modified: retrieval/retriever.py — deduplication step
• Modified: models/base.py or record metadata — parent_id field
• New: tests for chunk extraction, chunk storage, deduplicated retrieval

7. What this does NOT cover

• Audio/video chunking for content retrieval (existing time-based chunking handles embedding limits, not content retrieval)
• Replacing the multimodal embedding path — chunks complement it, not replace it
• LLM-based preprocessing of images/media (described above as a future optional layer)
• Non-PDF document types (DOCX, HTML, etc.)
• LLM-based reranking — the current weighted ranker (relevance + recency + importance) is free and stays in the retrieval path

Context

This emerged from testing the Step 5 multimodal pipeline. The multimodal embedding model works correctly for semantic/visual similarity. But for content-level retrieval ("find the document that mentions X"), single-vector-per-document is insufficient regardless of the embedding model. Chunking is the standard solution.

References

• Observation from live testing: prescription PDF ranked 9th/16th for direct keyword queries
• Gemini embedding docs confirm 8192 token limit per text input
• Current chunking code: utils/embeddings.py lines 379-413 (audio/video time-based chunking only)