RAG Chunk Selector

Select retrieval-augmented generation (RAG) chunks while avoiding duplicate text and overloading one document.

Requirements

• Define select_chunks(chunks, k, max_per_doc=2).
• Each chunk is a dict with id, doc_id, text, and score.
• Return selected chunk IDs.
• Higher score wins.
• Ties sort by id alphabetically.
• Skip chunks whose normalized text was already selected.
• Normalized text is lowercase with whitespace collapsed.
• Select at most max_per_doc chunks per document.
• Return an empty list when k <= 0 or max_per_doc <= 0.

Example

chunks = [
    {"id": "a", "doc_id": "d1", "text": "Hello world", "score": 0.9},
    {"id": "b", "doc_id": "d2", "text": "hello   world", "score": 0.8},
]
assert select_chunks(chunks, 2) == ["a"]

Constraints

• Don't mutate input chunks.
• Use deterministic ordering.
• Empty or whitespace-only text can still be selected once.
• All fields are present. Chunk IDs are unique, case-sensitive ASCII strings; document IDs and text are strings. Scores are finite numbers, and k and max_per_doc are integers.
• Return IDs in rank order. Scan that order greedily until k chunks are selected or candidates run out. A skipped chunk consumes neither a document slot nor a normalized-text slot.

For example, with a one-chunk document limit, select a top-ranked chunk from d1, skip a second d1 chunk with text "shared fact", then allow a lower-ranked d2 chunk with the same text. The skipped text was never selected. This policy enforces diversity constraints; it does not maximize the sum of scores over all possible selections.