Compare commits

..

1 commit

Author SHA1 Message Date
Yiorgis Gozadinos
d2af1a2e62
Wait for drained jobs instead of sleeping in the breaker isolation test
Some checks failed
build-docs / build (push) Has been cancelled
Tests / lint (push) Has been cancelled
Tests / lint-frontend (push) Has been cancelled
Tests / URI paths (macos-latest, py3.13) (push) Has been cancelled
Tests / URI paths (macos-latest, py3.14) (push) Has been cancelled
Tests / URI paths (ubuntu-latest, py3.13) (push) Has been cancelled
Tests / URI paths (ubuntu-latest, py3.14) (push) Has been cancelled
Tests / URI paths (windows-latest, py3.13) (push) Has been cancelled
Tests / URI paths (windows-latest, py3.14) (push) Has been cancelled
build-docs / deploy (push) Has been cancelled
Tests / test (push) Has been cancelled
test_breaker_isolates_sources snapshotted job status after a fixed 0.2s
sleep and failed on a loaded runner while two good jobs were still in
flight. Poll for the drained jobs under a 5s deadline.

Assert the paused source's jobs are unattempted: the queued-uri assertion
alone passes with source isolation disabled, since a retried job returns
to QUEUED.
2026-09-07 14:42:03 +03:00
7 changed files with 67 additions and 202 deletions

View file

@ -51,8 +51,6 @@
- `haiku-rag mcp` covers the configured `lancedb.databases` set. `sources` on
`search_documents`, `search_documents_by_image` and `execute_code`; `source`
on `get_document`; an unknown name is a tool error. `DocumentInfo.source`.
- `footnote` items are no longer filtered from expanded context. The noise
labels are `page_header`, `page_footer` and `document_index`.
### Fixed
@ -61,8 +59,6 @@
- Past `analysis.code_timeout` a sandbox program starts no further host call.
Files served from memory and in-code `search()` / `list_documents()` were
not checked against the deadline.
- Context expansion keeps the item a result matched on when it carries a
noise label, and counts it toward the character budget.
### Removed

View file

@ -13,7 +13,7 @@ search:
- **limit**: Default number of search results to return when no limit is specified. Used by CLI, MCP server, and QA. Default: 5
- **max_context_chars**: Hard limit on total characters in expanded content. Default: 5000.
Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (page headers, page footers, table of contents). This naturally crosses into adjacent sections until the budget is filled. Picture and table matches are exempt: they return their enclosing section as-is and never cross section boundaries. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded.
Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (footnotes, page headers). This naturally crosses into adjacent sections until the budget is filled. Picture and table matches are exempt: they return their enclosing section as-is and never cross section boundaries. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded.
!!! note "Reranking behavior"
When a reranker is configured, search automatically retrieves 10x the requested limit, then reranks to return the final count. This improves result quality without requiring you to adjust `limit`.

View file

@ -367,7 +367,7 @@ for result in expanded_results:
print(f"Expanded content: {result.content}")
```
Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (page headers, page footers, table of contents). This naturally crosses into adjacent sections until the budget is filled. Picture and table matches are exempt: they return their enclosing section as-is and never cross section boundaries. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded.
Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (footnotes, page headers). This naturally crosses into adjacent sections until the budget is filled. Picture and table matches are exempt: they return their enclosing section as-is and never cross section boundaries. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded.
Configuration:

View file

@ -28,7 +28,7 @@ When configured, a cross-encoder reranker re-scores 10x the requested candidates
`limit` controls how many results reach the LLM. More candidates improve recall but increase token usage. See [Search Settings](configuration/qa.md#search-settings).
Context expansion is automatic and section-aware. Search results are expanded to include surrounding content from the same document section. For structured documents, expansion stays within section boundaries and filters noise (page headers, page footers and the table of contents). For unstructured documents, expansion grows outward until the character budget is filled. `max_context_chars` caps expansion to prevent context bloat.
Context expansion is automatic and section-aware. Search results are expanded to include surrounding content from the same document section. For structured documents, expansion stays within section boundaries and filters noise (footnotes, page headers). For unstructured documents, expansion grows outward until the character budget is filled. `max_context_chars` caps expansion to prevent context bloat.
## Tuning Generation
@ -89,7 +89,7 @@ Press `c` on a chunk to see the expanded context that would be fed to the RAG ca
- The expanded text. Section-aware expansion stays within section boundaries on structured documents and fills `max_context_chars` outward on unstructured ones.
- Source document, content type, and relevance score.
- Filtered noise. Page headers, page footers and the table of contents are excluded from structured documents.
- Filtered noise. Footnotes, page headers and footers are excluded from structured documents.
If `qa.model.vision = true` is set, the modal also renders the picture bytes attached to that chunk, so you see exactly what the vision model would receive.

View file

@ -27,19 +27,17 @@ For UNSTRUCTURED documents (no section headers):
In both cases:
- max_context_chars caps total characters per expanded result
- Noise labels (page_header, page_footer, document_index) are
excluded from content AND budget counting in structured documents,
except the items a result matched on
- Noise labels (footnote, page_header, page_footer, document_index) are
excluded from content AND budget counting in structured documents
- Results without doc_item_refs pass through unexpanded
"""
from collections.abc import Set as AbstractSet
from typing import Any
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.document_item import DocumentItem
_NOISE_LABELS = {"page_header", "page_footer", "document_index"}
_NOISE_LABELS = {"footnote", "page_header", "page_footer", "document_index"}
_SECTION_BOUNDARY_LABELS = {"section_header", "title"}
# Labels whose pertinent unit is the item plus its own section: expansion
@ -168,43 +166,37 @@ def _span_in_window(
return start < win_end and end > win_start
def _noise_positions(
items: list[DocumentItem], matched_positions: set[int]
) -> set[int]:
"""Positions skipped as noise; an item a result matched on never is."""
return {
item.position for item in items if item.label in _NOISE_LABELS
} - matched_positions
def _expand_outward(
items: list[DocumentItem],
center_idx: int,
max_chars: int,
noise: AbstractSet[int] = frozenset(),
skip_noise: bool = False,
lo_bound: int = 0,
hi_bound: int | None = None,
) -> tuple[int, int]:
"""Expand item-by-item outward from center until char budget is filled.
Items at ``noise`` positions are excluded from char counting.
When skip_noise is True, noise labels are excluded from char counting
(used in structured documents so footnotes don't consume budget).
lo_bound and hi_bound constrain expansion (e.g., to section edges).
"""
if hi_bound is None:
hi_bound = len(items) - 1
lo = hi = center_idx
char_count = len(items[center_idx].text)
center_is_noise = skip_noise and items[center_idx].label in _NOISE_LABELS
char_count = 0 if center_is_noise else len(items[center_idx].text)
while char_count < max_chars:
grew = False
if lo > lo_bound:
lo -= 1
if items[lo].position not in noise:
if not (skip_noise and items[lo].label in _NOISE_LABELS):
char_count += len(items[lo].text)
grew = True
if hi < hi_bound and char_count < max_chars:
hi += 1
if items[hi].position not in noise:
if not (skip_noise and items[hi].label in _NOISE_LABELS):
char_count += len(items[hi].text)
grew = True
if not grew:
@ -227,8 +219,6 @@ def _find_expansion_range(
if not has_sections:
return _expand_outward(items, center_idx, max_chars)
noise = _noise_positions(items, matched_positions)
# Build section spans: [(start_idx, end_idx), ...]
headers = [
i for i, item in enumerate(items) if item.label in _SECTION_BOUNDARY_LABELS
@ -251,7 +241,7 @@ def _find_expansion_range(
sec_chars = sum(
len(items[i].text)
for i in range(sec_start, sec_end + 1)
if items[i].position not in noise
if items[i].label not in _NOISE_LABELS
)
min_useful = int(max_chars * _MIN_SECTION_BUDGET_RATIO)
@ -263,7 +253,12 @@ def _find_expansion_range(
if sec_chars > max_chars:
# Section too large — expand outward bounded by section edges
return _expand_outward(
items, center_idx, max_chars, noise, lo_bound=sec_start, hi_bound=sec_end
items,
center_idx,
max_chars,
skip_noise=True,
lo_bound=sec_start,
hi_bound=sec_end,
)
# Picture/table hits stay section-bounded: their pertinent unit is the
@ -272,7 +267,7 @@ def _find_expansion_range(
return (items[sec_start].position, items[sec_end].position)
# Section too small (e.g., title+authors) — expand across boundaries
return _expand_outward(items, center_idx, max_chars, noise)
return _expand_outward(items, center_idx, max_chars, skip_noise=True)
def _group_lost_constituent(built: SearchResult, group: list[SearchResult]) -> bool:
@ -305,7 +300,7 @@ def _build_result(
range_end: int,
original_results: list[SearchResult],
pos_to_item: dict[int, DocumentItem],
noise: AbstractSet[int],
has_sections: bool,
max_chars: int,
) -> SearchResult:
"""Build one expanded result from the items in ``[range_start, range_end]``."""
@ -318,7 +313,9 @@ def _build_result(
for pos in range(range_start, range_end + 1):
item = pos_to_item.get(pos)
if item is None or pos in noise:
if item is None:
continue
if has_sections and item.label in _NOISE_LABELS:
continue
if item.text:
if content_parts:
@ -460,21 +457,8 @@ def expand_with_items(
ranges: list[tuple[int, int, SearchResult]] = []
passthrough: list[SearchResult] = []
def matched_positions(group: list[SearchResult]) -> set[int]:
return {
ref_positions[ref]
for result in group
for ref in result.doc_item_refs
if ref in ref_positions
}
def noise_for(group: list[SearchResult]) -> set[int]:
if not has_sections:
return set()
return _noise_positions(window_items, matched_positions(group))
for result in results:
matched = matched_positions([result])
matched = {ref_positions[r] for r in result.doc_item_refs if r in ref_positions}
if not matched:
passthrough.append(result)
continue
@ -489,7 +473,7 @@ def expand_with_items(
final_results: list[SearchResult] = []
for range_start, range_end, group in merged:
built = _build_result(
range_start, range_end, group, pos_to_item, noise_for(group), max_chars
range_start, range_end, group, pos_to_item, has_sections, max_chars
)
if len(group) > 1 and _group_lost_constituent(built, group):
# The merged window cannot afford every constituent's evidence:
@ -499,7 +483,7 @@ def expand_with_items(
lo, hi = constituent_range[id(result)]
final_results.append(
_build_result(
lo, hi, [result], pos_to_item, noise_for([result]), max_chars
lo, hi, [result], pos_to_item, has_sections, max_chars
)
)
continue

View file

@ -1054,16 +1054,23 @@ async def test_breaker_isolates_sources(client, jobs, sync):
for _ in range(10):
pool._breaker_for("bad").record_failure()
async def _good_jobs_drained():
while True:
done = await jobs.list_jobs(status=JobStatus.SUCCEEDED, limit=50)
if len(done) == 3:
return done
await asyncio.sleep(0.02)
await pool.start()
try:
await asyncio.sleep(0.2)
succeeded = await jobs.list_jobs(status=JobStatus.SUCCEEDED, limit=50)
succeeded = await asyncio.wait_for(_good_jobs_drained(), timeout=5.0)
queued = await jobs.list_jobs(status=JobStatus.QUEUED, limit=50)
finally:
await pool.stop()
assert {j.uri for j in succeeded} == {"g0", "g1", "g2"}
assert {j.uri for j in queued} == {"b0", "b1", "b2"}
assert [j.attempts for j in queued] == [0, 0, 0]
@pytest.mark.asyncio

View file

@ -103,20 +103,32 @@ class TestExpandOutward:
lo, hi = _expand_outward(items, 9, max_chars=999999)
assert hi == 9
def test_noise_positions_excluded_from_char_count(self):
def test_skip_noise_excludes_from_char_count(self):
items = [
_item(0, text="a" * 100),
_item(1, label="page_header", text="f" * 5000),
_item(1, label="footnote", text="f" * 5000),
_item(2, text="b" * 100),
_item(3, text="c" * 100),
_item(4, label="page_header", text="f" * 5000),
_item(4, label="footnote", text="f" * 5000),
_item(5, text="d" * 100),
]
lo, hi = _expand_outward(items, 2, max_chars=500, noise={1, 4})
# Noise (5000 chars each) does not count, so expansion passes it
lo, hi = _expand_outward(items, 2, max_chars=500, skip_noise=True)
# Footnotes (5000 chars each) should NOT count toward budget
# So we should expand past them
assert lo <= 0
assert hi >= 5
def test_noise_center_gets_zero_chars(self):
items = [
_item(0, text="a" * 200),
_item(1, label="document_index", text="x" * 10000),
_item(2, text="b" * 200),
]
lo, hi = _expand_outward(items, 1, max_chars=500, skip_noise=True)
# Center is noise, should start at 0 chars and expand outward
assert lo == 0
assert hi == 2
def test_respects_bounds(self):
items = [_item(i, text="x" * 100) for i in range(20)]
lo, hi = _expand_outward(items, 10, max_chars=999999, lo_bound=8, hi_bound=12)
@ -212,28 +224,16 @@ class TestFindExpansionRange:
items = [
_item(0, label="section_header", text="Section"),
_item(1, text="Real content." * 10),
_item(2, label="page_header", text="x" * 10000),
_item(2, label="footnote", text="x" * 10000),
_item(3, text="More content." * 10),
]
# Section non-noise chars: ~260 chars (items 0,1,3). Under 5000 budget.
# The header's 10000 chars do not count.
# The footnote's 10000 chars should NOT count.
lo, hi = _find_expansion_range(items, {1}, has_sections=True, max_chars=5000)
# Should return full section (it fits in budget excluding noise)
assert lo == 0
assert hi == 3
def test_matched_noise_item_counts_toward_section_chars(self):
items = [
_item(0, label="section_header", text="Contents"),
_item(1, text="Real content." * 10),
_item(2, label="document_index", text="x" * 10000),
_item(3, text="More content." * 10),
]
lo, hi = _find_expansion_range(items, {2}, has_sections=True, max_chars=5000)
# The matched index is 10000 chars: the section is over budget and
# expansion stays on the match.
assert (lo, hi) == (2, 2)
def test_items_before_first_header_form_section(self):
items = [
_item(0, text="Preamble text."),
@ -373,11 +373,12 @@ class TestExpandWithItems:
assert len(expanded) == 1
assert expanded[0].content == "original"
async def test_matched_noise_item_survives_expansion(self, temp_db_path):
"""A result that matched on a noise-labelled item keeps that item."""
async def test_noise_only_range_preserves_original(self, temp_db_path):
"""When noise filtering removes all content, original chunk is preserved."""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(temp_db_path, create=True) as rag:
# Structured document where the matched item's section has only noise
items = [
DocumentItem(
document_id="doc-1",
@ -420,134 +421,11 @@ class TestExpandWithItems:
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
assert "x" * 2000 in expanded[0].content
assert "#/texts/1" in expanded[0].doc_item_refs
async def test_footnotes_are_kept_in_expanded_content(self, temp_db_path):
from haiku.rag.client import HaikuRAG
async with HaikuRAG(temp_db_path, create=True) as rag:
items = [
DocumentItem(
document_id="doc-1",
position=0,
self_ref="#/texts/0",
label="section_header",
text="Chapter 1",
),
DocumentItem(
document_id="doc-1",
position=1,
self_ref="#/texts/1",
label="text",
text="Body paragraph. " * 80,
),
DocumentItem(
document_id="doc-1",
position=2,
self_ref="#/texts/2",
label="footnote",
text="1 See Smith v Jones, para 12.",
),
]
await rag.document_item_repository.create_items("doc-1", items)
result = SearchResult(
content="Body paragraph.",
score=0.9,
document_id="doc-1",
doc_item_refs=["#/texts/1"],
)
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
assert "See Smith v Jones" in expanded[0].content
assert "#/texts/2" in expanded[0].doc_item_refs
async def test_unmatched_noise_excluded_in_structured_document(self, temp_db_path):
from haiku.rag.client import HaikuRAG
async with HaikuRAG(temp_db_path, create=True) as rag:
items = [
DocumentItem(
document_id="doc-1",
position=0,
self_ref="#/texts/0",
label="section_header",
text="Chapter 1",
),
DocumentItem(
document_id="doc-1",
position=1,
self_ref="#/texts/1",
label="text",
text="Body paragraph. " * 80,
),
DocumentItem(
document_id="doc-1",
position=2,
self_ref="#/texts/2",
label="page_header",
text="RUNNING HEADER",
),
DocumentItem(
document_id="doc-1",
position=3,
self_ref="#/texts/3",
label="text",
text="Closing paragraph.",
),
]
await rag.document_item_repository.create_items("doc-1", items)
result = SearchResult(
content="Body paragraph.",
score=0.9,
document_id="doc-1",
doc_item_refs=["#/texts/1"],
)
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
assert "Closing paragraph." in expanded[0].content
assert "RUNNING HEADER" not in expanded[0].content
assert "#/texts/2" not in expanded[0].doc_item_refs
async def test_unstructured_document_keeps_noise_labels(self, temp_db_path):
from haiku.rag.client import HaikuRAG
async with HaikuRAG(temp_db_path, create=True) as rag:
items = [
DocumentItem(
document_id="doc-1",
position=0,
self_ref="#/texts/0",
label="page_header",
text="RUNNING HEADER",
),
DocumentItem(
document_id="doc-1",
position=1,
self_ref="#/texts/1",
label="text",
text="Body paragraph.",
),
]
await rag.document_item_repository.create_items("doc-1", items)
result = SearchResult(
content="Body paragraph.",
score=0.9,
document_id="doc-1",
doc_item_refs=["#/texts/1"],
)
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
assert "RUNNING HEADER" in expanded[0].content
# The TOC section's only non-header item is document_index (noise).
# The section_header "Table of Contents" has text but _expand_outward
# with skip_noise crosses into the Introduction section which has
# real content — so we get expanded content, not the fallback.
assert len(expanded[0].content) > 0
async def test_picture_expansion_stays_within_section_pages(self, temp_db_path):
from haiku.rag.client import HaikuRAG
@ -1681,6 +1559,6 @@ def test_build_result_skips_positions_with_no_item():
),
}
built = _build_result(0, 3, [original], pos_to_item, set(), 5000)
built = _build_result(0, 3, [original], pos_to_item, False, 5000)
assert built.content == "first\n\nlast"