danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

replace fixed-radius expansion with section-bounded algorithm

Browse source 
Context expansion is now automatic and structure-aware. For structured
documents, expands within the section containing the match. For sections
that exceed the budget or are too small, expands item-by-item outward
skipping noise labels. Unstructured documents use budget-based outward
expansion. Results sorted by relevance score.
This commit is contained in:
Yiorgis Gozadinos 2026-04-15 12:39:41 +03:00
parent 364b1bc509
commit b6113bf8ab
No known key found for this signature in database
26 changed files with 575 additions and 968 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
CHANGELOG.md
View file

@ -4,12 +4,19 @@
### Added
- **Document items table**: Pre-extracted document items stored as individual rows with scalar indexes, enabling context expansion via indexed range queries (~2.5ms) instead of full DoclingDocument deserialization (~8.7s for large documents)
- **Section-bounded context expansion**: Expansion is now automatic and structure-aware — stays within section boundaries for structured documents, grows outward for unstructured ones. Noise labels (footnotes, page headers/footers) are filtered. Results without `doc_item_refs` pass through unexpanded.
### Changed
- **Database migration required**: Run `haiku-rag migrate` to populate `document_items` table for existing documents
- **Pin docling-core**: Upper bound added (`<2.72`) to prevent uncontrolled schema changes
### Removed
- **`context_radius` config**: Replaced by automatic section-bounded expansion. Context expansion no longer requires configuration.
- **DoclingDocument LRU cache**: No longer needed — the document_items table replaces in-memory caching for context expansion
- **`cachetools` dependency**: No longer used
## [0.39.0] - 2026-04-09
### Added

1
app/README.md
View file

@ -61,7 +61,6 @@ embeddings:
search:
limit: 10
context_radius: 1
```
See `haiku.rag.yaml.example` for all options.

1
app/haiku.rag.yaml.example
View file

@ -33,7 +33,6 @@ embeddings:
# Search settings
search:
limit: 5
context_radius: 0
# Provider settings
providers:

4
docs/apps.md
View file

@ -182,8 +182,8 @@ Search uses hybrid (vector + full-text) search across all chunks.
Press `c` while viewing a chunk to see the expanded context that would be provided to the QA agent:
- Type-aware expansion: tables, code blocks, and lists expand to their complete structures
- Text content expands based on `search.context_radius` setting
- Section-aware expansion: expands to fill the current document section
- Noise filtering: footnotes, page headers/footers excluded from structured documents
- Includes metadata like source document, content type, and relevance score
### Visual Grounding

3
docs/configuration/index.md
View file

@ -99,9 +99,8 @@ research:
search:
limit: 10 # Default number of results to return
context_radius: 0 # DocItems before/after to include for text content
max_context_items: 10 # Maximum items in expanded context
max_context_chars: 10000 # Maximum characters in expanded context
max_context_chars: 5000 # Maximum characters in expanded context
vector_index_metric: cosine # cosine, l2, or dot
vector_refine_factor: 30

8
docs/configuration/qa-research.md
View file

@ -7,17 +7,15 @@ Configure search behavior and context expansion:
```yaml
search:
limit: 10 # Default number of results to return
context_radius: 0 # DocItems before/after to include for text content
max_context_items: 10 # Maximum items in expanded context
max_context_chars: 10000 # Maximum characters in expanded context
max_context_chars: 5000 # Maximum characters in expanded context
```
- **limit**: Default number of search results to return when no limit is specified. Used by CLI, MCP server, QA, and research workflows. Default: 10
- **context_radius**: For text content (paragraphs), includes N DocItems before and after. Set to 0 to disable expansion (default).
- **max_context_items**: Limits how many document items (paragraphs, list items, etc.) can be included in expanded context. Default: 10.
- **max_context_chars**: Hard limit on total characters in expanded content. Default: 10000.
- **max_context_chars**: Hard limit on total characters in expanded content. Default: 5000.
Structural content (tables, code blocks, lists) uses type-aware expansion that automatically includes the complete structure regardless of how it was chunked.
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. 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`.

2
docs/configuration/storage.md
View file

@ -112,7 +112,7 @@ search:
vector_refine_factor: 30 # Re-ranking factor for accuracy
```
For search behavior settings (`limit`, `context_radius`, `max_context_items`, `max_context_chars`), see [QA and Research](qa-research.md#search-settings).
For search behavior settings (`limit`, `max_context_items`, `max_context_chars`), see [QA and Research](qa-research.md#search-settings).
- **vector_index_metric**: Distance metric for vector similarity:
- `cosine`: Cosine similarity (default, best for most embeddings)

16
docs/python.md
View file

@ -359,29 +359,27 @@ results = await client.search(
### Expanding Search Context
Expand search results with adjacent chunks for more complete context:
Expand search results with surrounding content from the document:
```python
# Get initial search results
search_results = await client.search("machine learning", limit=3)
# Expand search results with adjacent content from the source document
# Expand with section-bounded context
expanded_results = await client.expand_context(search_results)
# The expanded results contain chunks with combined content
for result in expanded_results:
print(f"Expanded content: {result.content}")
```
Context expansion uses your configuration settings:
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. 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.
- **search.context_radius**: For text content (paragraphs), includes N DocItems before and after
- **search.max_context_items**: Limits how many document items can be included
- **search.max_context_chars**: Hard limit on total characters
Configuration:
**Type-aware expansion**: Structural content (tables, code blocks, lists) automatically expands to include the complete structure, regardless of how it was split during chunking.
- **search.max_context_items**: Maximum items in expanded context. Default: 10.
- **search.max_context_chars**: Maximum characters in expanded context. Default: 5000.
**Smart Merging**: When expanded chunks overlap or are adjacent within the same document, they are automatically merged into single chunks with continuous content. This eliminates duplication and provides coherent text blocks. The merged chunk uses the highest relevance score from the original chunks.
**Smart Merging**: When expanded results overlap within the same document, they are automatically merged into a single result with continuous content and the highest relevance score.
## Question Answering

2
docs/tuning.md
View file

@ -26,7 +26,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-research.md#search-settings).
`context_radius` expands text chunks with neighboring document items. Structural content (tables, code blocks, lists) expands automatically to include the complete structure. This setting matters most with small `chunk_size` values, where individual chunks may lack sufficient context. `max_context_items` and `max_context_chars` cap 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_items` and `max_context_chars` cap expansion to prevent context bloat.
## Tuning Generation

1
evaluations/evaluations/benchmark.py
View file

@ -49,7 +49,6 @@ def build_experiment_metadata(
"embedder_dim": config.embeddings.model.vector_dim,
"chunk_size": config.processing.chunk_size,
"search_limit": config.search.limit,
"context_radius": config.search.context_radius,
"max_context_items": config.search.max_context_items,
"max_context_chars": config.search.max_context_chars,
"rerank_provider": config.reranking.model.provider

1
evaluations/tests/test_benchmark.py
View file

@ -31,7 +31,6 @@ class TestBuildExperimentMetadata:
assert result["embedder_dim"] == config.embeddings.model.vector_dim
assert result["chunk_size"] == config.processing.chunk_size
assert result["search_limit"] == config.search.limit
assert result["context_radius"] == config.search.context_radius
assert result["qa_provider"] == config.qa.model.provider
assert result["qa_model"] == config.qa.model.name
assert "judge_provider" not in result

374
haiku_rag_slim/haiku/rag/client.py
View file

@ -1090,24 +1090,23 @@ class HaikuRAG:
self,
search_results: list[SearchResult],
) -> list[SearchResult]:
"""Expand search results with adjacent content from the source document.
"""Expand search results with surrounding content from the document.
When DoclingDocument is available and results have doc_item_refs, expands
by finding adjacent DocItems with accurate bounding boxes and metadata.
Otherwise, falls back to chunk-based expansion using adjacent chunks.
Uses the document_items table for section-bounded expansion.
See haiku.rag.context for the algorithm description.
Expansion is type-aware based on content:
- Tables, code blocks, and lists expand to include complete structures
- Text content uses the configured radius (search.context_radius)
- Expansion is limited by search.max_context_items and search.max_context_chars
Results without doc_item_refs pass through unexpanded. This happens
when chunks were created without docling metadata (e.g., custom chunks
passed to import_document).
Args:
search_results: List of SearchResult objects from search.
Returns:
List of SearchResult objects with expanded content and resolved provenance.
List of SearchResult objects with expanded content.
"""
radius = self._config.search.context_radius
from haiku.rag.context import expand_with_items
max_items = self._config.search.max_context_items
max_chars = self._config.search.max_context_chars
@ -1127,347 +1126,22 @@ class HaikuRAG:
continue
has_refs = any(r.doc_item_refs for r in doc_results)
docling_doc = None
if not has_refs:
expanded_results.extend(doc_results)
continue
if has_refs:
# Only load docling data when refs exist (skips content blob)
doc = await self.document_repository.get_docling_data(doc_id)
if doc is not None:
docling_doc = doc.get_docling_document()
if docling_doc is not None and has_refs:
# Use DoclingDocument-based expansion
expanded = await self._expand_with_docling(
doc_results,
docling_doc,
radius,
max_items,
max_chars,
)
expanded_results.extend(expanded)
else:
# Fall back to chunk-based expansion (always uses fixed radius)
if radius > 0:
expanded = await self._expand_with_chunks(
doc_id, doc_results, radius
)
expanded_results.extend(expanded)
else:
expanded_results.extend(doc_results)
expanded = await expand_with_items(
self.document_item_repository,
doc_id,
doc_results,
max_items,
max_chars,
)
expanded_results.extend(expanded)
expanded_results.sort(key=lambda r: r.score, reverse=True)
return expanded_results
def _merge_ranges(
self, ranges: list[tuple[int, int, SearchResult]]
) -> list[tuple[int, int, list[SearchResult]]]:
"""Merge overlapping or adjacent ranges."""
if not ranges:
return []
sorted_ranges = sorted(ranges, key=lambda x: x[0])
merged: list[tuple[int, int, list[SearchResult]]] = []
cur_min, cur_max, cur_results = (
sorted_ranges[0][0],
sorted_ranges[0][1],
[sorted_ranges[0][2]],
)
for min_idx, max_idx, result in sorted_ranges[1:]:
if cur_max >= min_idx - 1: # Overlapping or adjacent
cur_max = max(cur_max, max_idx)
cur_results.append(result)
else:
merged.append((cur_min, cur_max, cur_results))
cur_min, cur_max, cur_results = min_idx, max_idx, [result]
merged.append((cur_min, cur_max, cur_results))
return merged
# Label groups for type-aware expansion
_STRUCTURAL_LABELS = {
"table",
"code",
"list_item",
"form",
"key_value_region",
"field_region",
}
def _extract_item_text(self, item, docling_doc) -> str | None:
"""Extract text content from a DocItem.
Handles different item types:
- TextItem, SectionHeaderItem, etc.: Use .text attribute
- TableItem: Use export_to_markdown() for table content
- PictureItem: Use export_to_markdown() with PLACEHOLDER mode to avoid base64
"""
from docling_core.types.doc.base import ImageRefMode
from docling_core.types.doc.document import PictureItem
# Try simple text attribute first (works for most items)
if text := getattr(item, "text", None):
return text
# For pictures: use PLACEHOLDER mode to avoid base64 images in content.
# This still includes VLM descriptions (annotations) and captions.
if isinstance(item, PictureItem):
return item.export_to_markdown(
docling_doc,
image_mode=ImageRefMode.PLACEHOLDER,
image_placeholder="",
)
# For tables and other items with export_to_markdown
if hasattr(item, "export_to_markdown"):
try:
return item.export_to_markdown(docling_doc)
except Exception:
pass
# Fallback for items with captions
if caption := getattr(item, "caption", None):
if hasattr(caption, "text"):
return caption.text
return None
def _get_item_label(self, item) -> str | None:
"""Extract label string from a DocItem."""
label = getattr(item, "label", None)
if label is None:
return None
return str(label.value) if hasattr(label, "value") else str(label)
def _compute_type_aware_range(
self,
all_items: list,
indices: list[int],
radius: int,
max_items: int,
max_chars: int,
) -> tuple[int, int]:
"""Compute expansion range based on content type with limits.
For structural content (tables, code, lists), expands to include complete
structures. For text, uses the configured radius. Applies hybrid limits.
"""
if not indices:
return (0, 0)
min_idx = min(indices)
max_idx = max(indices)
# Determine the primary label type from matched items
labels_in_chunk = set()
for idx in indices:
item, _ = all_items[idx]
if label := self._get_item_label(item):
labels_in_chunk.add(label)
# Check if we have structural content
is_structural = bool(labels_in_chunk & self._STRUCTURAL_LABELS)
if is_structural:
# Expand to complete structure boundaries
# Expand backwards to find structure start
while min_idx > 0:
prev_item, _ = all_items[min_idx - 1]
prev_label = self._get_item_label(prev_item)
if prev_label in labels_in_chunk & self._STRUCTURAL_LABELS:
min_idx -= 1
else:
break
# Expand forwards to find structure end
while max_idx < len(all_items) - 1:
next_item, _ = all_items[max_idx + 1]
next_label = self._get_item_label(next_item)
if next_label in labels_in_chunk & self._STRUCTURAL_LABELS:
max_idx += 1
else:
break
else:
# Text content: use radius-based expansion
min_idx = max(0, min_idx - radius)
max_idx = min(len(all_items) - 1, max_idx + radius)
# Apply hybrid limits
# First check item count hard limit
if max_idx - min_idx + 1 > max_items:
# Center the window around original indices
original_center = (min(indices) + max(indices)) // 2
half_items = max_items // 2
min_idx = max(0, original_center - half_items)
max_idx = min(len(all_items) - 1, min_idx + max_items - 1)
# Then check character soft limit (but keep at least original items)
char_count = 0
effective_max = min_idx
for i in range(min_idx, max_idx + 1):
item, _ = all_items[i]
text = getattr(item, "text", "") or ""
char_count += len(text)
effective_max = i
# Once we've included original items, check char limit
if i >= max(indices) and char_count > max_chars:
break
max_idx = effective_max
return (min_idx, max_idx)
async def _expand_with_docling(
self,
results: list[SearchResult],
docling_doc,
radius: int,
max_items: int,
max_chars: int,
) -> list[SearchResult]:
"""Expand results using DoclingDocument structure.
Structural content (tables, code, lists) expands to complete structures.
Text content uses radius-based expansion.
"""
all_items = list(docling_doc.iterate_items())
ref_to_index = {
getattr(item, "self_ref", None): i
for i, (item, _) in enumerate(all_items)
if getattr(item, "self_ref", None)
}
# Compute expanded ranges
ranges: list[tuple[int, int, SearchResult]] = []
passthrough: list[SearchResult] = []
for result in results:
indices = [
ref_to_index[r] for r in result.doc_item_refs if r in ref_to_index
]
if not indices:
passthrough.append(result)
continue
min_idx, max_idx = self._compute_type_aware_range(
all_items, indices, radius, max_items, max_chars
)
ranges.append((min_idx, max_idx, result))
# Merge overlapping ranges
merged = self._merge_ranges(ranges)
final_results: list[SearchResult] = []
for min_idx, max_idx, original_results in merged:
content_parts: list[str] = []
refs: list[str] = []
pages: set[int] = set()
labels: set[str] = set()
for i in range(min_idx, max_idx + 1):
item, _ = all_items[i]
# Extract text content - handle different item types
text = self._extract_item_text(item, docling_doc)
if text:
content_parts.append(text)
if self_ref := getattr(item, "self_ref", None):
refs.append(self_ref)
if label := getattr(item, "label", None):
labels.add(
str(label.value) if hasattr(label, "value") else str(label)
)
if prov := getattr(item, "prov", None):
for p in prov:
if (page_no := getattr(p, "page_no", None)) is not None:
pages.add(page_no)
# Merge headings preserving order
all_headings: list[str] = []
for r in original_results:
if r.headings:
all_headings.extend(h for h in r.headings if h not in all_headings)
first = original_results[0]
final_results.append(
SearchResult(
content="\n\n".join(content_parts),
score=max(r.score for r in original_results),
chunk_id=first.chunk_id,
document_id=first.document_id,
document_uri=first.document_uri,
document_title=first.document_title,
doc_item_refs=refs,
page_numbers=sorted(pages),
headings=all_headings or None,
labels=sorted(labels),
)
)
return final_results + passthrough
async def _expand_with_chunks(
self,
doc_id: str,
results: list[SearchResult],
radius: int,
) -> list[SearchResult]:
"""Expand results using chunk-based adjacency."""
# Build ranges from result orders
ranges: list[tuple[int, int, SearchResult]] = []
passthrough: list[SearchResult] = []
for result in results:
if result.chunk_id is None:
passthrough.append(result)
continue
start = result.order - radius
end = result.order + radius
ranges.append((start, end, result))
if not ranges:
return results
# Compute the full order range needed and fetch only those chunks
all_starts = [s for s, _, _ in ranges]
all_ends = [e for _, e, _ in ranges]
range_min = min(all_starts)
range_max = max(all_ends)
chunks_in_range = await self.chunk_repository.get_chunks_in_range(
doc_id, range_min, range_max
)
if not chunks_in_range:
return results
chunk_by_order = {c.order: c for c in chunks_in_range}
# Merge and build results
final_results: list[SearchResult] = []
for min_idx, max_idx, original_results in self._merge_ranges(ranges):
# Collect chunks in order
merged_chunks = [
chunk_by_order[o]
for o in range(min_idx, max_idx + 1)
if o in chunk_by_order
]
first = original_results[0]
final_results.append(
SearchResult(
content="".join(c.content for c in merged_chunks),
score=max(r.score for r in original_results),
chunk_id=first.chunk_id,
document_id=first.document_id,
document_uri=first.document_uri,
document_title=first.document_title,
doc_item_refs=first.doc_item_refs,
page_numbers=first.page_numbers,
headings=first.headings,
labels=first.labels,
)
)
return final_results + passthrough
async def ask(
self,
question: str,
@ -1794,18 +1468,12 @@ class HaikuRAG:
Used by RECHUNK and FULL modes after the chunks table has been cleared.
"""
from haiku.rag.store.engine import DocumentRecord
from haiku.rag.store.models.document import invalidate_docling_document_cache
if not documents:
return
now = datetime.now().isoformat()
# Invalidate cache for all documents being updated
for doc in documents:
if doc.id:
invalidate_docling_document_cache(doc.id)
# Batch update documents using merge_insert (single LanceDB version)
doc_records = []
for doc in documents:

3
haiku_rag_slim/haiku/rag/config/models.py
View file

@ -174,9 +174,8 @@ class ProcessingConfig(BaseModel):
class SearchConfig(BaseModel):
limit: int = 10
context_radius: int = 0
max_context_items: int = 10
max_context_chars: int = 10000
max_context_chars: int = 5000
vector_index_metric: Literal["cosine", "l2", "dot"] = "cosine"
vector_refine_factor: int = 30

249
haiku_rag_slim/haiku/rag/context.py Normal file
View file

@ -0,0 +1,249 @@
"""Section-bounded context expansion for search results.
Expands search results with surrounding content from the document using
the document_items table. The algorithm adapts to document structure:
For STRUCTURED documents (containing section_header or title labels):
1. Resolve matched doc_item_refs to positions in the items table
2. Find section boundaries around each match (section_header/title labels)
3. If the section fits within the budget, include it entirely
4. If the section exceeds the budget, OR the section is too small (under
20% of max_context_chars), expand item-by-item from the match center
outward, skipping noise labels. This lets small sections (e.g., a
title+authors area) grow into the next section's content.
5. Merge overlapping ranges from multiple results in the same document
For UNSTRUCTURED documents (no section headers):
Expand outward item-by-item from the match center until the character
budget is filled. No noise filtering (unstructured docs typically only
have text items).
In both cases:
- max_context_chars caps total characters per expanded result
- max_context_items caps total items per expanded result
- 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 haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.document_item import DocumentItem
from haiku.rag.store.repositories.document_item import DocumentItemRepository
_NOISE_LABELS = {"footnote", "page_header", "page_footer", "document_index"}
_SECTION_BOUNDARY_LABELS = {"section_header", "title"}
# Sections with fewer chars than this fraction of max_context_chars are
# considered too small — expansion falls through to item-by-item outward
# growth, which naturally crosses into adjacent sections.
_MIN_SECTION_BUDGET_RATIO = 0.2
def _merge_ranges(
ranges: list[tuple[int, int, SearchResult]],
) -> list[tuple[int, int, list[SearchResult]]]:
"""Merge overlapping or adjacent ranges."""
if not ranges:
return []
sorted_ranges = sorted(ranges, key=lambda x: x[0])
merged: list[tuple[int, int, list[SearchResult]]] = []
cur_min, cur_max, cur_results = (
sorted_ranges[0][0],
sorted_ranges[0][1],
[sorted_ranges[0][2]],
)
for min_idx, max_idx, result in sorted_ranges[1:]:
if cur_max >= min_idx - 1: # Overlapping or adjacent
cur_max = max(cur_max, max_idx)
cur_results.append(result)
else:
merged.append((cur_min, cur_max, cur_results))
cur_min, cur_max, cur_results = min_idx, max_idx, [result]
merged.append((cur_min, cur_max, cur_results))
return merged
def _expand_outward(
items: list[DocumentItem],
center_idx: int,
max_items: int,
max_chars: int,
skip_noise: bool = False,
) -> tuple[int, int]:
"""Expand item-by-item outward from center until budget is filled.
When skip_noise is True, noise labels are excluded from char counting
(used in structured documents so footnotes don't consume budget).
"""
lo = hi = center_idx
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 and hi - lo + 1 < max_items:
grew = False
if lo > 0:
lo -= 1
if not (skip_noise and items[lo].label in _NOISE_LABELS):
char_count += len(items[lo].text)
grew = True
if hi < len(items) - 1 and char_count < max_chars:
hi += 1
if not (skip_noise and items[hi].label in _NOISE_LABELS):
char_count += len(items[hi].text)
grew = True
if not grew:
break
return (items[lo].position, items[hi].position)
def _find_expansion_range(
items: list[DocumentItem],
matched_positions: set[int],
has_sections: bool,
max_items: int,
max_chars: int,
) -> tuple[int, int]:
"""Find the expansion range for matched positions within a window of items."""
pos_to_idx = {item.position: i for i, item in enumerate(items)}
matched_indices = sorted(pos_to_idx[p] for p in matched_positions)
center_idx = matched_indices[len(matched_indices) // 2]
if not has_sections:
return _expand_outward(items, center_idx, max_items, max_chars)
# Build section spans: [(start_idx, end_idx), ...]
headers = [
i for i, item in enumerate(items) if item.label in _SECTION_BOUNDARY_LABELS
]
sections: list[tuple[int, int]] = []
if headers[0] > 0:
sections.append((0, headers[0] - 1))
for j, h in enumerate(headers):
end = headers[j + 1] - 1 if j + 1 < len(headers) else len(items) - 1
sections.append((h, end))
# Find which section contains the center match
current = 0
for j, (start, end) in enumerate(sections):
if start <= center_idx <= end:
current = j
break
sec_start, sec_end = sections[current]
sec_chars = sum(
len(items[i].text)
for i in range(sec_start, sec_end + 1)
if items[i].label not in _NOISE_LABELS
)
# Section fits nicely in the budget — return it as-is
min_useful = int(max_chars * _MIN_SECTION_BUDGET_RATIO)
if min_useful <= sec_chars <= max_chars and sec_end - sec_start + 1 <= max_items:
return (items[sec_start].position, items[sec_end].position)
# Section is too large or too small — expand item-by-item from center.
# For too-large sections this stays within budget.
# For too-small sections (e.g., title+authors) this naturally grows
# into adjacent sections until the budget is filled.
return _expand_outward(items, center_idx, max_items, max_chars, skip_noise=True)
async def expand_with_items(
document_item_repository: DocumentItemRepository,
document_id: str,
results: list[SearchResult],
max_items: int,
max_chars: int,
) -> list[SearchResult]:
"""Expand results using the document_items table."""
all_refs = []
for result in results:
all_refs.extend(result.doc_item_refs)
ref_positions = await document_item_repository.resolve_refs(document_id, all_refs)
if not ref_positions:
return results
# Fetch all items for the document to ensure we always detect section
# structure correctly, regardless of where the match falls.
item_count = await document_item_repository.get_item_count(document_id)
window_items = await document_item_repository.get_items_in_range(
document_id, 0, item_count
)
if not window_items:
return results
has_sections = any(item.label in _SECTION_BOUNDARY_LABELS for item in window_items)
# Compute expansion ranges per result
ranges: list[tuple[int, int, SearchResult]] = []
passthrough: list[SearchResult] = []
for result in results:
matched = {ref_positions[r] for r in result.doc_item_refs if r in ref_positions}
if not matched:
passthrough.append(result)
continue
lo, hi = _find_expansion_range(
window_items, matched, has_sections, max_items, max_chars
)
ranges.append((lo, hi, result))
merged = _merge_ranges(ranges)
# Build results from the window items
pos_to_item = {item.position: item for item in window_items}
final_results: list[SearchResult] = []
for range_start, range_end, original_results in merged:
content_parts: list[str] = []
refs: list[str] = []
pages: set[int] = set()
labels: set[str] = set()
for pos in range(range_start, range_end + 1):
item = pos_to_item.get(pos)
if item is None:
continue
if has_sections and item.label in _NOISE_LABELS:
continue
if item.text:
content_parts.append(item.text)
refs.append(item.self_ref)
if item.label:
labels.add(item.label)
pages.update(item.page_numbers)
all_headings: list[str] = []
for r in original_results:
if r.headings:
all_headings.extend(h for h in r.headings if h not in all_headings)
first = original_results[0]
# If noise filtering removed all content, preserve the original
expanded_content = "\n\n".join(content_parts)
if not expanded_content:
expanded_content = first.content
final_results.append(
SearchResult(
content=expanded_content,
score=max(r.score for r in original_results),
chunk_id=first.chunk_id,
document_id=first.document_id,
document_uri=first.document_uri,
document_title=first.document_title,
doc_item_refs=refs or first.doc_item_refs,
page_numbers=sorted(pages) or first.page_numbers,
headings=all_headings or None,
labels=sorted(labels) or first.labels,
)
)
return final_results + passthrough

42
haiku_rag_slim/haiku/rag/store/models/document.py
View file

@ -2,7 +2,6 @@ import json
from datetime import datetime
from typing import TYPE_CHECKING
from cachetools import LRUCache
from pydantic import BaseModel, Field
from haiku.rag.store.compression import compress_docling_split, decompress_json
@ -11,38 +10,6 @@ if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument, PageItem
_docling_document_cache: LRUCache[str, "DoclingDocument"] = LRUCache(maxsize=100)
def _validate_without_pages(compressed_data: bytes) -> "DoclingDocument":
"""Decompress and validate DoclingDocument."""
from docling_core.types.doc.document import DoclingDocument
json_str = decompress_json(compressed_data)
return DoclingDocument.model_validate_json(json_str)
def _get_cached_docling_document(
document_id: str, compressed_data: bytes
) -> "DoclingDocument":
"""Get or parse DoclingDocument with LRU caching by document ID.
Strips page images before validation for performance — cached documents
do not contain page data.
"""
if document_id in _docling_document_cache:
return _docling_document_cache[document_id]
doc = _validate_without_pages(compressed_data)
_docling_document_cache[document_id] = doc
return doc
def invalidate_docling_document_cache(document_id: str) -> None:
"""Remove a document from the DoclingDocument cache."""
_docling_document_cache.pop(document_id, None)
class Document(BaseModel):
"""
Represents a document with an ID, content, and metadata.
@ -73,19 +40,16 @@ class Document(BaseModel):
def get_docling_document(self) -> "DoclingDocument | None":
"""Parse and return the stored DoclingDocument (without page images).
Uses LRU cache (keyed by document ID) to avoid repeated parsing.
Returns:
The parsed DoclingDocument, or None if not stored.
"""
if self.docling_document is None:
return None
# No caching for documents without ID
if self.id is None:
return _validate_without_pages(self.docling_document)
from docling_core.types.doc.document import DoclingDocument
return _get_cached_docling_document(self.id, self.docling_document)
json_str = decompress_json(self.docling_document)
return DoclingDocument.model_validate_json(json_str)
def get_page_images(self, page_numbers: list[int]) -> "dict[int, PageItem]":
"""Decompress and return page images for the requested page numbers.

7
haiku_rag_slim/haiku/rag/store/repositories/document.py
View file

@ -152,13 +152,9 @@ class DocumentRepository:
async def update(self, entity: Document) -> Document:
"""Update an existing document."""
self.store._assert_writable()
from haiku.rag.store.models.document import invalidate_docling_document_cache
assert entity.id, "Document ID is required for update"
# Invalidate cache before update
invalidate_docling_document_cache(entity.id)
# Update timestamp
now = datetime.now().isoformat()
entity.updated_at = datetime.fromisoformat(now)
@ -184,15 +180,12 @@ class DocumentRepository:
async def delete(self, entity_id: str) -> bool:
"""Delete a document by its ID."""
self.store._assert_writable()
from haiku.rag.store.models.document import invalidate_docling_document_cache
# Check if document exists
doc = await self.get_by_id(entity_id)
if doc is None:
return False
invalidate_docling_document_cache(entity_id)
# Delete associated chunks and items first
await self.chunk_repository.delete_by_document_id(entity_id)
await self.document_item_repository.delete_by_document_id(entity_id)

1
haiku_rag_slim/pyproject.toml
View file

@ -22,7 +22,6 @@ classifiers = [
]
dependencies = [
"cachetools>=7.0.5",
"docling-core>=2.71.0,<2.72",
"haiku.skills>=0.14.0",
"httpx>=0.28.1",

62
tests/cassettes/test_context_enhancement/test_expand_context_keeps_separate_non_overlapping.yaml
View file

File diff suppressed because one or more lines are too long

58
tests/cassettes/test_context_enhancement/test_expand_context_merges_overlapping_chunks.yaml
View file

File diff suppressed because one or more lines are too long

94
tests/cassettes/test_context_enhancement/test_expand_context_multiple_documents.yaml
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_context_enhancement/test_expand_context_radius_zero.yaml
View file

File diff suppressed because one or more lines are too long

0
tests/cassettes/test_context_enhancement/test_text_expansion_uses_radius.yaml → tests/cassettes/test_context_enhancement/test_text_expansion_includes_surrounding.yaml
View file

238
tests/test_context.py Normal file
View file

@ -0,0 +1,238 @@
from haiku.rag.context import (
_expand_outward,
_find_expansion_range,
_merge_ranges,
)
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.document_item import DocumentItem
def _item(
position: int, label: str = "text", text: str = "", pages: list[int] | None = None
) -> DocumentItem:
return DocumentItem(
document_id="doc-1",
position=position,
self_ref=f"#/texts/{position}",
label=label,
text=text or f"Text for item {position}.",
page_numbers=pages or [1],
)
def _result(score: float = 0.5, refs: list[str] | None = None) -> SearchResult:
return SearchResult(
content="original",
score=score,
document_id="doc-1",
doc_item_refs=refs or [],
)
class TestMergeRanges:
def test_empty(self):
assert _merge_ranges([]) == []
def test_no_overlap(self):
r1, r2 = _result(), _result()
merged = _merge_ranges([(0, 5, r1), (10, 15, r2)])
assert len(merged) == 2
assert merged[0] == (0, 5, [r1])
assert merged[1] == (10, 15, [r2])
def test_overlapping(self):
r1, r2 = _result(0.9), _result(0.8)
merged = _merge_ranges([(0, 10, r1), (5, 15, r2)])
assert len(merged) == 1
assert merged[0] == (0, 15, [r1, r2])
def test_adjacent(self):
r1, r2 = _result(), _result()
merged = _merge_ranges([(0, 5, r1), (6, 10, r2)])
assert len(merged) == 1
assert merged[0] == (0, 10, [r1, r2])
def test_sorts_by_position(self):
r1, r2 = _result(), _result()
merged = _merge_ranges([(10, 15, r1), (0, 5, r2)])
assert merged[0][0] == 0
assert merged[1][0] == 10
class TestExpandOutward:
def test_basic_expansion(self):
items = [_item(i, text=f"{'x' * 100}") for i in range(10)]
lo, hi = _expand_outward(items, 5, max_items=10, max_chars=500)
assert lo <= 5
assert hi >= 5
total = sum(
len(items[i].text)
for i in range(lo, hi + 1)
if items[i].position >= lo and items[i].position <= hi
)
# Should be around 500 chars (may overshoot by one item)
assert total >= 400
def test_respects_max_items(self):
items = [_item(i, text="x") for i in range(100)]
lo, hi = _expand_outward(items, 50, max_items=5, max_chars=999999)
count = hi - lo + 1
# May overshoot by 1-2 items due to alternating expansion
assert count <= 7
def test_respects_max_chars(self):
items = [_item(i, text=f"{'x' * 200}") for i in range(20)]
lo, hi = _expand_outward(items, 10, max_items=999, max_chars=500)
total = sum(
len(items[i].text)
for i in range(lo, hi + 1)
if items[i].position >= lo and items[i].position <= hi
)
# Should be near 500, may overshoot by one item (~200 chars)
assert total <= 900
def test_center_at_start(self):
items = [_item(i) for i in range(10)]
lo, hi = _expand_outward(items, 0, max_items=5, max_chars=999999)
assert lo == 0
def test_center_at_end(self):
items = [_item(i) for i in range(10)]
lo, hi = _expand_outward(items, 9, max_items=5, max_chars=999999)
assert hi == 9
def test_skip_noise_excludes_from_char_count(self):
items = [
_item(0, text="a" * 100),
_item(1, label="footnote", text="f" * 5000),
_item(2, text="b" * 100),
_item(3, text="c" * 100),
_item(4, label="footnote", text="f" * 5000),
_item(5, text="d" * 100),
]
lo, hi = _expand_outward(items, 2, max_items=10, 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_items=10, max_chars=500, skip_noise=True)
# Center is noise, should start at 0 chars and expand outward
assert lo == 0
assert hi == 2
class TestFindExpansionRange:
def _structured_items(self):
"""Document with two sections, each over min_useful (1000 chars)."""
return [
_item(0, label="section_header", text="Introduction"),
_item(1, text="First paragraph. " * 40), # ~680 chars
_item(2, text="Second paragraph. " * 40), # ~720 chars
_item(3, label="footnote", text="Some footnote."),
_item(4, label="section_header", text="Methods"),
_item(5, text="Methods paragraph one. " * 40), # ~920 chars
_item(6, text="Methods paragraph two. " * 40), # ~920 chars
]
def test_structured_returns_section(self):
items = self._structured_items()
lo, hi = _find_expansion_range(
items, {1}, has_sections=True, max_items=20, max_chars=5000
)
# Should return the Introduction section (items 0-3)
assert lo == 0
assert hi == 3
def test_structured_different_section(self):
items = self._structured_items()
lo, hi = _find_expansion_range(
items, {5}, has_sections=True, max_items=20, max_chars=5000
)
# Should return the Methods section (items 4-6)
assert lo == 4
assert hi == 6
def test_structured_large_section_falls_back_to_outward(self):
items = [
_item(0, label="section_header", text="Big Section"),
] + [_item(i, text="x" * 1000) for i in range(1, 20)]
# Section has 19 * 1000 = 19000 chars, way over 5000 budget
lo, hi = _find_expansion_range(
items, {10}, has_sections=True, max_items=50, max_chars=5000
)
# Should NOT return the full section
total = sum(
len(items[i].text) for i in range(lo, hi + 1) if items[i].position >= lo
)
assert total < 10000
def test_structured_small_section_expands_outward(self):
items = [
_item(0, label="title", text="Paper Title"),
_item(1, text="Author names"),
_item(2, label="section_header", text="Abstract"),
_item(3, text="Abstract content. " * 50),
_item(4, label="section_header", text="Introduction"),
_item(5, text="Intro content. " * 50),
]
# Title section (items 0-1) is tiny (~25 chars) < 20% of 5000
lo, hi = _find_expansion_range(
items, {0}, has_sections=True, max_items=20, max_chars=5000
)
# Should expand past the title section into the abstract
assert hi >= 3
def test_unstructured_expands_outward(self):
items = [_item(i, text=f"Paragraph {i}. " * 10) for i in range(10)]
lo, hi = _find_expansion_range(
items, {5}, has_sections=False, max_items=20, max_chars=5000
)
assert lo < 5
assert hi > 5
def test_multiple_matched_positions_uses_center(self):
items = [_item(i, text="x" * 100) for i in range(20)]
# Match at positions 3 and 7, center should be index for position 5 (median)
lo, hi = _find_expansion_range(
items, {3, 7}, has_sections=False, max_items=5, max_chars=999999
)
center = (lo + hi) // 2
# Center should be around position 5
assert 3 <= center <= 7
def test_noise_excluded_from_section_char_count(self):
items = [
_item(0, label="section_header", text="Section"),
_item(1, text="Real content." * 10),
_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 footnote's 10000 chars should NOT count.
lo, hi = _find_expansion_range(
items, {1}, has_sections=True, max_items=20, max_chars=5000
)
# Should return full section (it fits in budget excluding noise)
assert lo == 0
assert hi == 3
def test_items_before_first_header_form_section(self):
items = [
_item(0, text="Preamble text."),
_item(1, text="More preamble."),
_item(2, label="section_header", text="First Section"),
_item(3, text="Section content."),
]
lo, hi = _find_expansion_range(
items, {0}, has_sections=True, max_items=20, max_chars=5000
)
# Match is in preamble section (items 0-1), which is small
# Should expand outward into the first section
assert hi >= 2

238
tests/test_context_enhancement.py
View file

@ -5,7 +5,6 @@ from docling_core.types.doc.labels import DocItemLabel
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models import SearchResult
from haiku.rag.store.models.chunk import Chunk
async def create_document_with_docling(
@ -224,11 +223,10 @@ async def test_code_expansion_includes_adjacent_blocks(
@pytest.mark.vcr()
async def test_text_expansion_uses_radius(temp_db_path):
"""Text content expansion should use radius, not structural boundaries."""
async def test_text_expansion_includes_surrounding(temp_db_path):
"""Text content expansion should include surrounding paragraphs."""
config = AppConfig()
config.processing.chunk_size = 32
config.search.context_radius = 1 # Small radius
# Create a document with longer paragraphs that will split
doc = DoclingDocument(name="text_test")
@ -260,8 +258,7 @@ async def test_text_expansion_uses_radius(temp_db_path):
original = results[0]
expanded = await client.expand_context(results)
# With radius=1, expansion should include adjacent paragraphs
# Content length should be >= original (may add adjacent content)
# Expansion should include adjacent paragraphs
assert len(expanded[0].content) >= len(original.content)
@ -335,210 +332,61 @@ async def test_max_items_limit_caps_expansion(temp_db_path):
assert item_count <= 2, f"Expected at most 2 items, got {item_count}"
@pytest.mark.vcr()
async def test_expand_context_radius_zero(temp_db_path):
"""Test expand_context with radius 0 returns original results."""
# Default config has context_radius=0
async def test_expand_context_single_item_document(temp_db_path):
"""Test expand_context with a single-item document."""
from haiku.rag.store.models.document import Document
docling_doc = DoclingDocument(name="simple")
docling_doc.add_text(label=DocItemLabel.PARAGRAPH, text="Simple test content")
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content="Simple test content")
document = Document(content="Simple test content")
document.set_docling(docling_doc)
doc = await client._store_document_with_chunks(document, [], docling_doc)
assert doc.id is not None
chunks = await client.chunk_repository.get_by_document_id(doc.id)
search_results = [SearchResult.from_chunk(chunks[0], 0.9)]
# Create a search result with a doc_item_ref pointing to the item
items = await client.document_item_repository.get_items_in_range(doc.id, 0, 10)
assert len(items) > 0
search_results = [
SearchResult(
content="Simple test content",
score=0.9,
document_id=doc.id,
doc_item_refs=[items[0].self_ref],
)
]
expanded_results = await client.expand_context(search_results)
# Should return exactly the same results
assert len(expanded_results) == 1
assert expanded_results[0].content == search_results[0].content
assert expanded_results[0].score == search_results[0].score
assert expanded_results[0].score == 0.9
assert "Simple test content" in expanded_results[0].content
@pytest.mark.vcr()
async def test_expand_context_multiple_documents(temp_db_path):
"""Test expand_context with results from multiple documents."""
config = AppConfig()
config.search.context_radius = 1
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
# Create first document with manual chunks
docling_doc1 = DoclingDocument(name="doc1")
docling_doc1.add_text(label=DocItemLabel.TEXT, text="Doc1 content")
doc1_chunks = [
Chunk(content="Doc1 Part A", order=0),
Chunk(content="Doc1 Part B", order=1),
Chunk(content="Doc1 Part C", order=2),
]
doc1 = await client.import_document(
docling_document=docling_doc1, chunks=doc1_chunks, uri="doc1.txt"
)
# Create second document with manual chunks
docling_doc2 = DoclingDocument(name="doc2")
docling_doc2.add_text(label=DocItemLabel.TEXT, text="Doc2 content")
doc2_chunks = [
Chunk(content="Doc2 Section X", order=0),
Chunk(content="Doc2 Section Y", order=1),
]
doc2 = await client.import_document(
docling_document=docling_doc2, chunks=doc2_chunks, uri="doc2.txt"
)
assert doc1.id is not None
assert doc2.id is not None
chunks1 = await client.chunk_repository.get_by_document_id(doc1.id)
chunks2 = await client.chunk_repository.get_by_document_id(doc2.id)
# Get middle chunk from doc1 (order=1) and first chunk from doc2 (order=0)
chunk1 = next(c for c in chunks1 if c.order == 1)
chunk2 = next(c for c in chunks2 if c.order == 0)
async def test_expand_context_no_refs_passes_through(temp_db_path):
"""Results without doc_item_refs pass through unexpanded."""
async with HaikuRAG(temp_db_path, create=True) as client:
# A search result with no doc_item_refs should pass through as-is
search_results = [
SearchResult.from_chunk(chunk1, 0.8),
SearchResult.from_chunk(chunk2, 0.7),
SearchResult(
content="Some chunk content",
score=0.8,
document_id="some-doc",
doc_item_refs=[],
)
]
expanded_results = await client.expand_context(search_results)
expanded = await client.expand_context(search_results)
assert len(expanded_results) == 2
# Check first expanded result (should include chunks 0,1,2 from doc1)
expanded1 = expanded_results[0]
assert expanded1.score == 0.8
assert "Doc1 Part A" in expanded1.content
assert "Doc1 Part B" in expanded1.content
assert "Doc1 Part C" in expanded1.content
# Check second expanded result (should include chunks 0,1 from doc2)
expanded2 = expanded_results[1]
assert expanded2.score == 0.7
assert "Doc2 Section X" in expanded2.content
assert "Doc2 Section Y" in expanded2.content
@pytest.mark.vcr()
async def test_expand_context_merges_overlapping_chunks(temp_db_path):
"""Test that overlapping expanded chunks are merged into one."""
config = AppConfig()
config.search.context_radius = 1
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
# Create document with 5 chunks
docling_doc = DoclingDocument(name="test")
docling_doc.add_text(label=DocItemLabel.TEXT, text="Full document content")
manual_chunks = [
Chunk(content="Chunk 0", order=0),
Chunk(content="Chunk 1", order=1),
Chunk(content="Chunk 2", order=2),
Chunk(content="Chunk 3", order=3),
Chunk(content="Chunk 4", order=4),
]
doc = await client.import_document(
docling_document=docling_doc, chunks=manual_chunks
)
assert doc.id is not None
chunks = await client.chunk_repository.get_by_document_id(doc.id)
# Get adjacent chunks (orders 1 and 2) - these will overlap when expanded
chunk1 = next(c for c in chunks if c.order == 1)
chunk2 = next(c for c in chunks if c.order == 2)
# With radius=1:
# chunk1 expanded would be [0,1,2]
# chunk2 expanded would be [1,2,3]
# These should merge into one chunk containing [0,1,2,3]
search_results = [
SearchResult.from_chunk(chunk1, 0.8),
SearchResult.from_chunk(chunk2, 0.7),
]
expanded_results = await client.expand_context(search_results)
# Should have only 1 merged result instead of 2 overlapping ones
assert len(expanded_results) == 1
merged = expanded_results[0]
# Should contain all chunks from 0 to 3
assert "Chunk 0" in merged.content
assert "Chunk 1" in merged.content
assert "Chunk 2" in merged.content
assert "Chunk 3" in merged.content
assert "Chunk 4" not in merged.content # Should not include chunk 4
# Should use the higher score (0.8)
assert merged.score == 0.8
@pytest.mark.vcr()
async def test_expand_context_keeps_separate_non_overlapping(temp_db_path):
"""Test that non-overlapping expanded chunks remain separate."""
config = AppConfig()
config.search.context_radius = 1
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
# Create document with chunks far apart
docling_doc = DoclingDocument(name="test")
docling_doc.add_text(label=DocItemLabel.TEXT, text="Full document content")
manual_chunks = [
Chunk(content="Chunk 0", order=0),
Chunk(content="Chunk 1", order=1),
Chunk(content="Chunk 2", order=2),
Chunk(content="Chunk 5", order=5), # Gap here
Chunk(content="Chunk 6", order=6),
Chunk(content="Chunk 7", order=7),
]
doc = await client.import_document(
docling_document=docling_doc, chunks=manual_chunks
)
assert doc.id is not None
chunks = await client.chunk_repository.get_by_document_id(doc.id)
# Get chunks by index - they will have sequential orders 0,1,2,3,4,5
# So get chunk with order=0 and chunk with order=5 (far enough apart)
chunk0 = next(c for c in chunks if c.order == 0) # Content: "Chunk 0"
chunk5 = next(
c for c in chunks if c.order == 5
) # Content: "Chunk 7" but now at order 5
# chunk0 expanded: [0,1] with radius=1 (orders 0,1)
# chunk5 expanded: [4,5] with radius=1 (orders 4,5)
search_results = [
SearchResult.from_chunk(chunk0, 0.8),
SearchResult.from_chunk(chunk5, 0.7),
]
expanded_results = await client.expand_context(search_results)
# Should have 2 separate results
assert len(expanded_results) == 2
# Sort by score to ensure predictable order
expanded_results.sort(key=lambda x: x.score, reverse=True)
chunk0_expanded = expanded_results[0]
chunk5_expanded = expanded_results[1]
# First chunk (order=0) expanded should contain orders [0,1]
# Content should be "Chunk 0" + "Chunk 1"
assert "Chunk 0" in chunk0_expanded.content
assert "Chunk 1" in chunk0_expanded.content
assert "Chunk 5" not in chunk0_expanded.content
assert chunk0_expanded.score == 0.8
# Second chunk (order=5) expanded should contain orders [4,5]
# Content should be "Chunk 6" (order 4) + "Chunk 7" (order 5)
assert "Chunk 6" in chunk5_expanded.content
assert "Chunk 7" in chunk5_expanded.content
assert "Chunk 0" not in chunk5_expanded.content
assert chunk5_expanded.score == 0.7
assert len(expanded) == 1
assert expanded[0].content == "Some chunk content"
assert expanded[0].score == 0.8
@pytest.mark.vcr()
async def test_expand_context_with_docling_merges_overlapping(temp_db_path):
"""Test that expand_context with DoclingDocument merges overlapping results."""
config = AppConfig()
config.search.context_radius = 3
markdown_content = """# Chapter 1
@ -593,7 +441,6 @@ This is paragraph four about topic C.
async def test_expand_context_docling_merges_metadata(temp_db_path):
"""Test that expand_context properly merges metadata from multiple results."""
config = AppConfig()
config.search.context_radius = 10
markdown_content = """# Introduction
@ -676,7 +523,6 @@ async def test_expand_context_no_base64_images(temp_db_path):
base64 image data from leaking into the expanded content.
"""
config = AppConfig()
config.search.context_radius = 5 # Expand enough to include pictures
docling_doc = create_picture_document()
@ -713,7 +559,6 @@ async def test_expand_context_no_base64_images_docling_local(temp_db_path):
config.processing.converter = "docling-local"
config.processing.chunker = "docling-local"
config.processing.conversion_options.do_ocr = False
config.search.context_radius = 5
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
pdf_path = Path(__file__).parent / "data" / "doclaynet.pdf"
@ -747,7 +592,6 @@ async def test_expand_context_no_base64_images_docling_serve(temp_db_path):
config = AppConfig()
config.processing.converter = "docling-serve"
config.processing.chunker = "docling-serve"
config.search.context_radius = 5
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
pdf_path = Path(__file__).parent / "data" / "doclaynet.pdf"

87
tests/test_document.py
View file

@ -142,93 +142,6 @@ def test_document_get_docling_document_none():
assert document.get_docling_document() is None
def test_document_get_docling_document_caching():
"""Test that get_docling_document uses LRU cache keyed by document ID."""
from haiku.rag.store.models.document import (
_docling_document_cache,
invalidate_docling_document_cache,
)
doc_json = {
"name": "test_doc",
"texts": [
{
"self_ref": "#/texts/0",
"text": "Test text",
"orig": "Test text",
"label": "paragraph",
},
],
"tables": [],
"pictures": [],
"groups": [],
"body": {"self_ref": "#/body", "children": []},
"furniture": {"self_ref": "#/furniture", "children": []},
}
import json
from haiku.rag.store.compression import compress_json
compressed = compress_json(json.dumps(doc_json))
# Clear cache to get clean state
_docling_document_cache.clear()
document = Document(
id="test-doc-id", content="Test content", docling_document=compressed
)
# First call - not in cache
assert "test-doc-id" not in _docling_document_cache
doc1 = document.get_docling_document()
assert "test-doc-id" in _docling_document_cache
# Second call - cache hit, same object
doc2 = document.get_docling_document()
assert doc1 is doc2
# Invalidation removes from cache
invalidate_docling_document_cache("test-doc-id")
assert "test-doc-id" not in _docling_document_cache
def test_document_get_docling_document_no_id_no_cache():
"""Test that documents without ID don't use cache."""
from haiku.rag.store.models.document import _docling_document_cache
doc_json = {
"name": "test_doc",
"texts": [],
"tables": [],
"pictures": [],
"groups": [],
"body": {"self_ref": "#/body", "children": []},
"furniture": {"self_ref": "#/furniture", "children": []},
}
import json
from haiku.rag.store.compression import compress_json
compressed = compress_json(json.dumps(doc_json))
# Clear cache
_docling_document_cache.clear()
# Document without ID
document = Document(content="Test content", docling_document=compressed)
doc1 = document.get_docling_document()
doc2 = document.get_docling_document()
# Cache should remain empty (no ID to cache by)
assert len(_docling_document_cache) == 0
# Each call parses fresh (different objects)
assert doc1 is not doc2
def test_set_docling_splits_structure_and_pages():
"""set_docling stores structure and pages separately."""
import json

2
uv.lock
View file

@ -1503,7 +1503,6 @@ name = "haiku-rag-slim"
version = "0.40.0"
source = { editable = "haiku_rag_slim" }
dependencies = [
{ name = "cachetools" },
{ name = "docling-core" },
{ name = "haiku-skills" },
{ name = "httpx" },
@ -1568,7 +1567,6 @@ zeroentropy = [
[package.metadata]
requires-dist = [
{ name = "cachetools", specifier = ">=7.0.5" },
{ name = "cohere", marker = "extra == 'cohere'", specifier = ">=5.21.1" },
{ name = "docling", marker = "extra == 'docling'", specifier = ">=2.84.0" },
{ name = "docling-core", specifier = ">=2.71.0,<2.72" },