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

Merge pull request #344 from ggozad/feat/docling-document-load-performace

Browse source 
Fast context expansion.
This commit is contained in:
Yiorgis Gozadinos 2026-04-16 12:33:52 +03:00 committed by GitHub
commit 969c45fe5c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
40 changed files with 1523 additions and 1018 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
CHANGELOG.md
View file

@ -1,6 +1,25 @@
# Changelog
## [Unreleased]
### 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
- **`max_searches` default**: Raised from 3 to 5 — faster expansion makes additional searches inexpensive
- **Improved QA prompt**: Stronger instruction to refuse answering from tangentially related content
- **Improved judge prompt**: Asymmetric evaluation — generated answers that are more comprehensive than expected are not penalized
### 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

1
docs/configuration/index.md
View file

@ -99,7 +99,6 @@ 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
vector_index_metric: cosine # cosine, l2, or dot

4
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
```
- **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.
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: 10000.
**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

25
evaluations/evaluations/evaluators/judge.py
View file

@ -4,27 +4,28 @@ from pydantic_ai import Agent
from haiku.rag.config.models import AppConfig, ModelConfig
from haiku.rag.utils import get_model
ANSWER_EQUIVALENCE_RUBRIC = """You are evaluating whether two answers to the same question are semantically equivalent.
ANSWER_EQUIVALENCE_RUBRIC = """You are evaluating whether a generated answer is equivalent to an expected answer for a given question.
EVALUATION CRITERIA:
Rate as EQUIVALENT if:
✓ Both answers contain the same core factual information
✓ Both directly address the question asked
✓ The generated answer contains the core factual information from the expected answer
✓ The generated answer directly addresses the question asked
✓ The key claims and conclusions are consistent
✓ Any additional detail in one answer doesn't contradict the other
✓ The generated answer may include additional correct details not in the expected answer — this is fine
Rate as NOT EQUIVALENT if:
✗ Factual contradictions exist between the answers
✗ One answer fails to address the core question
✗ Key information is missing that changes the meaning
✗ The answers lead to different conclusions or implications
✗ The generated answer contradicts facts in the expected answer
✗ The generated answer fails to address the core question
✗ Key information from the expected answer is missing in a way that changes the meaning
✗ The answers lead to different conclusions or actions
GUIDELINES:
- Ignore minor differences in phrasing, style, or formatting
- Focus on semantic meaning rather than exact wording
- Consider both answers correct if they convey the same essential information
- The evaluation is asymmetric: judge the generated answer against the expected answer, not the other way around
- A generated answer that is MORE detailed or comprehensive than the expected answer is EQUIVALENT, as long as it doesn't contradict it
- If the expected answer is incomplete or narrow, do not penalize the generated answer for being broader
- Ignore differences in phrasing, style, or formatting
- Focus on whether a user would get the correct guidance from the generated answer
- Be tolerant of different levels of detail if the core answer is preserved
- Evaluate based on what a person asking this question would need to know
"""

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

2
haiku_rag_slim/haiku/rag/agents/qa/prompts.py
View file

@ -32,8 +32,8 @@ Guidelines:
- Base answers strictly on retrieved content - do not use external knowledge
- Use the Source and Type metadata to understand context
- If multiple results are relevant, synthesize them coherently
- If information is insufficient, say: "I cannot find enough information in the knowledge base to answer this question."
- Be concise and direct - avoid elaboration unless asked
- Results are ordered by relevance, with rank 1 being most relevant
- If the search tool tells you the search limit is reached, stop searching immediately and answer with what you have
- If the retrieved documents do not directly address the question, say: "I cannot find enough information in the knowledge base to answer this question." Do not guess or infer an answer from tangentially related content.
"""

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

@ -20,12 +20,12 @@ from haiku.rag.reranking import get_reranker
from haiku.rag.store.engine import Store
from haiku.rag.store.models.chunk import Chunk, SearchResult
from haiku.rag.store.models.document import Document
from haiku.rag.store.models.document_item import extract_items
from haiku.rag.store.repositories.chunk import ChunkRepository
from haiku.rag.store.repositories.document import (
DocumentRepository,
_escape_sql_string,
)
from haiku.rag.store.repositories.document import DocumentRepository
from haiku.rag.store.repositories.document_item import DocumentItemRepository
from haiku.rag.store.repositories.settings import SettingsRepository
from haiku.rag.utils import escape_sql_string
if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
@ -96,6 +96,7 @@ class HaikuRAG:
)
self.document_repository = DocumentRepository(self.store)
self.chunk_repository = ChunkRepository(self.store)
self.document_item_repository = DocumentItemRepository(self.store)
@property
def is_read_only(self) -> bool:
@ -354,6 +355,7 @@ class HaikuRAG:
self,
document: Document,
chunks: list[Chunk],
docling_document: "DoclingDocument",
) -> Document:
"""Store a document with chunks, embedding any that lack embeddings.
@ -362,6 +364,7 @@ class HaikuRAG:
Args:
document: The document to store (will be created).
chunks: Chunks to store (will be embedded if lacking embeddings).
docling_document: The DoclingDocument to extract items from.
Returns:
The created Document instance with ID set.
@ -389,6 +392,10 @@ class HaikuRAG:
# Batch create all chunks in a single operation
await self.chunk_repository.create(chunks)
# Extract and store document items for context expansion
items = extract_items(created_doc.id, docling_document)
await self.document_item_repository.create_items(created_doc.id, items)
# Vacuum old versions in background (non-blocking) if auto_vacuum enabled
if self._config.storage.auto_vacuum:
asyncio.create_task(self.store.vacuum())
@ -403,6 +410,7 @@ class HaikuRAG:
self,
document: Document,
chunks: list[Chunk],
docling_document: "DoclingDocument | None" = None,
) -> Document:
"""Update a document and replace its chunks, embedding any that lack embeddings.
@ -411,6 +419,8 @@ class HaikuRAG:
Args:
document: The document to update (must have ID set).
chunks: Chunks to replace existing (will be embedded if lacking embeddings).
docling_document: The DoclingDocument to extract items from.
When None, existing items are preserved.
Returns:
The updated Document instance.
@ -441,6 +451,14 @@ class HaikuRAG:
# Batch create all chunks in a single operation
await self.chunk_repository.create(chunks)
# Replace document items when a new DoclingDocument is provided
if docling_document is not None:
await self.document_item_repository.delete_by_document_id(
updated_doc.id
)
items = extract_items(updated_doc.id, docling_document)
await self.document_item_repository.create_items(updated_doc.id, items)
# Vacuum old versions in background (non-blocking) if auto_vacuum enabled
if self._config.storage.auto_vacuum:
asyncio.create_task(self.store.vacuum())
@ -498,7 +516,9 @@ class HaikuRAG:
document.set_docling(docling_document)
# Store document and chunks
return await self._store_document_with_chunks(document, embedded_chunks)
return await self._store_document_with_chunks(
document, embedded_chunks, docling_document
)
async def import_document(
self,
@ -536,7 +556,9 @@ class HaikuRAG:
)
document.set_docling(docling_document)
return await self._store_document_with_chunks(document, chunks)
return await self._store_document_with_chunks(
document, chunks, docling_document
)
async def create_document_from_source(
self, source: str | Path, title: str | None = None, metadata: dict | None = None
@ -677,7 +699,7 @@ class HaikuRAG:
docling_document, stored_content
)
return await self._update_document_with_chunks(
existing_doc, embedded_chunks
existing_doc, embedded_chunks, docling_document
)
else:
# Create new document
@ -690,7 +712,9 @@ class HaikuRAG:
metadata=metadata,
)
document.set_docling(docling_document)
return await self._store_document_with_chunks(document, embedded_chunks)
return await self._store_document_with_chunks(
document, embedded_chunks, docling_document
)
async def _create_or_update_document_from_url(
self, url: str, title: str | None = None, metadata: dict | None = None
@ -790,7 +814,7 @@ class HaikuRAG:
docling_document, stored_content
)
return await self._update_document_with_chunks(
existing_doc, embedded_chunks
existing_doc, embedded_chunks, docling_document
)
else:
# Create new document
@ -803,7 +827,9 @@ class HaikuRAG:
metadata=metadata,
)
document.set_docling(docling_document)
return await self._store_document_with_chunks(document, embedded_chunks)
return await self._store_document_with_chunks(
document, embedded_chunks, docling_document
)
def _get_extension_from_content_type_or_url(
self, url: str, content_type: str
@ -882,7 +908,7 @@ class HaikuRAG:
if doc:
return doc
safe_input = _escape_sql_string(id_or_title)
safe_input = escape_sql_string(id_or_title)
docs = await self.list_documents(filter=f"title = '{safe_input}'")
if docs and docs[0].id:
return await self.get_document_by_id(docs[0].id)
@ -956,7 +982,9 @@ class HaikuRAG:
elif content is not None:
existing_doc.content = content
return await self._update_document_with_chunks(existing_doc, chunks)
return await self._update_document_with_chunks(
existing_doc, chunks, docling_document
)
# DoclingDocument provided without chunks - chunk and embed using primitives
if docling_document is not None:
@ -966,7 +994,7 @@ class HaikuRAG:
new_chunks = await self.chunk(docling_document)
embedded_chunks = await embed_chunks(new_chunks, self._config)
return await self._update_document_with_chunks(
existing_doc, embedded_chunks
existing_doc, embedded_chunks, docling_document
)
# Content provided without chunks - convert, chunk, and embed using primitives
@ -977,7 +1005,9 @@ class HaikuRAG:
new_chunks = await self.chunk(converted_docling)
embedded_chunks = await embed_chunks(new_chunks, self._config)
return await self._update_document_with_chunks(existing_doc, embedded_chunks)
return await self._update_document_with_chunks(
existing_doc, embedded_chunks, converted_docling
)
async def delete_document(self, document_id: str) -> bool:
"""Delete a document by its ID."""
@ -1058,24 +1088,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
@ -1095,347 +1124,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,
@ -1758,21 +1462,16 @@ class HaikuRAG:
"""Batch write documents and chunks during rebuild.
This performs two writes: one for all document updates, one for all chunks.
Also repopulates document items from the stored docling document.
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:
@ -1800,6 +1499,15 @@ class HaikuRAG:
if chunks:
await self.chunk_repository.create(chunks)
# Repopulate document items from stored docling data
for doc in documents:
assert doc.id is not None
docling_doc = doc.get_docling_document()
if docling_doc is not None:
await self.document_item_repository.delete_by_document_id(doc.id)
items = extract_items(doc.id, docling_doc)
await self.document_item_repository.create_items(doc.id, items)
async def _rebuild_rechunk(
self, documents: list[Document]
) -> AsyncGenerator[str, None]:

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

@ -80,7 +80,7 @@ class QAConfig(BaseModel):
temperature=0.3,
)
)
max_searches: int = 3
max_searches: int = 5
class ResearchConfig(BaseModel):
@ -174,7 +174,6 @@ class ProcessingConfig(BaseModel):
class SearchConfig(BaseModel):
limit: int = 10
context_radius: int = 0
max_context_items: int = 10
max_context_chars: int = 10000
vector_index_metric: Literal["cosine", "l2", "dot"] = "cosine"

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

@ -0,0 +1,255 @@
"""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 a window of items around matched positions. The margin must be
# wide enough to find section boundaries (the nearest section_header/title
# above and below the match).
all_positions = sorted(ref_positions.values())
window_margin = max_items * 10
window_start = max(0, min(all_positions) - window_margin)
window_end = max(all_positions) + window_margin
window_items = await document_item_repository.get_items_in_range(
document_id, window_start, window_end
)
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]
# Expansion should never return less content than the original chunk.
# This can happen when item texts are fragmented (e.g., docling splits
# formatted HTML list items into many small text nodes).
expanded_content = "\n\n".join(content_parts)
if len(expanded_content) < len(first.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

33
haiku_rag_slim/haiku/rag/store/engine.py
View file

@ -107,6 +107,15 @@ def create_chunk_model(vector_dim: int):
return ChunkRecord
class DocumentItemRecord(LanceModel):
document_id: str
position: int
self_ref: str
label: str = Field(default="")
text: str = Field(default="")
page_numbers: str = Field(default="[]")
class SettingsRecord(LanceModel):
id: str = Field(default="settings")
settings: str = Field(default="{}")
@ -256,6 +265,7 @@ class Store:
for table in [
self.documents_table,
self.chunks_table,
self.document_items_table,
self.settings_table,
]:
table.optimize(cleanup_older_than=retention)
@ -358,7 +368,7 @@ class Store:
def _init_tables(self):
"""Initialize database tables (create if they don't exist)."""
existing_tables = self.db.list_tables().tables
required_tables = {"documents", "chunks", "settings"}
required_tables = {"documents", "chunks", "document_items", "settings"}
missing_tables = required_tables - set(existing_tables)
if missing_tables and self._read_only:
@ -385,6 +395,23 @@ class Store:
"content_fts", replace=True, with_position=True, remove_stop_words=False
)
# Create or open document_items table
if "document_items" in existing_tables:
self.document_items_table = self.db.open_table("document_items")
else:
self.document_items_table = self.db.create_table(
"document_items", schema=DocumentItemRecord
)
self.document_items_table.create_scalar_index(
"document_id", index_type="BTREE", replace=True
)
self.document_items_table.create_scalar_index(
"position", index_type="BTREE", replace=True
)
self.document_items_table.create_scalar_index(
"self_ref", index_type="BTREE", replace=True
)
# Create or open settings table
if "settings" in existing_tables:
self.settings_table = self.db.open_table("settings")
@ -528,6 +555,7 @@ class Store:
return {
"documents": int(self.documents_table.version),
"chunks": int(self.chunks_table.version),
"document_items": int(self.document_items_table.version),
"settings": int(self.settings_table.version),
}
@ -540,6 +568,7 @@ class Store:
self._assert_writable()
self.documents_table.restore(int(versions["documents"]))
self.chunks_table.restore(int(versions["chunks"]))
self.document_items_table.restore(int(versions["document_items"]))
self.settings_table.restore(int(versions["settings"]))
return True
@ -569,6 +598,7 @@ class Store:
tables = [
("documents", self.documents_table),
("chunks", self.chunks_table),
("document_items", self.document_items_table),
("settings", self.settings_table),
]
@ -620,6 +650,7 @@ class Store:
table_map = {
"documents": self.documents_table,
"chunks": self.chunks_table,
"document_items": self.document_items_table,
"settings": self.settings_table,
}
table = table_map.get(table_name)

2
haiku_rag_slim/haiku/rag/store/models/__init__.py
View file

@ -1,10 +1,12 @@
from .chunk import BoundingBox, Chunk, ChunkMetadata, SearchResult
from .document import Document
from .document_item import DocumentItem
__all__ = [
"BoundingBox",
"Chunk",
"ChunkMetadata",
"Document",
"DocumentItem",
"SearchResult",
]

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.

88
haiku_rag_slim/haiku/rag/store/models/document_item.py Normal file
View file

@ -0,0 +1,88 @@
from typing import TYPE_CHECKING
from pydantic import BaseModel
if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument, NodeItem
class DocumentItem(BaseModel):
document_id: str
position: int
self_ref: str
label: str = ""
text: str = ""
page_numbers: list[int] = []
def extract_item_text(item: "NodeItem", docling_doc: "DoclingDocument") -> 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, TableItem
if text := getattr(item, "text", None):
return text
if isinstance(item, PictureItem):
return item.export_to_markdown(
docling_doc,
image_mode=ImageRefMode.PLACEHOLDER,
image_placeholder="",
)
if isinstance(item, TableItem):
try:
return item.export_to_markdown(docling_doc)
except Exception:
pass
if caption := getattr(item, "caption", None):
if hasattr(caption, "text"):
return caption.text
return None
def extract_items(
document_id: str, docling_doc: "DoclingDocument"
) -> list[DocumentItem]:
"""Extract document items from a DoclingDocument for the items table.
Runs iterate_items() and extracts the fields needed for context expansion:
self_ref, label, pre-rendered text, and page numbers from provenance.
Items are stored as docling produces them — container items (e.g., list_item)
may have empty text with content in their children.
"""
items: list[DocumentItem] = []
for position, (item, _level) in enumerate(docling_doc.iterate_items()):
label = getattr(item, "label", None)
label_str = str(label.value) if hasattr(label, "value") else str(label or "")
text = extract_item_text(item, docling_doc) or ""
page_numbers: list[int] = []
if prov := getattr(item, "prov", None):
for p in prov:
page_no = getattr(p, "page_no", None)
if page_no is not None and page_no not in page_numbers:
page_numbers.append(page_no)
items.append(
DocumentItem(
document_id=document_id,
position=position,
self_ref=item.self_ref,
label=label_str,
text=text,
page_numbers=sorted(page_numbers),
)
)
return items

2
haiku_rag_slim/haiku/rag/store/repositories/__init__.py
View file

@ -1,9 +1,11 @@
from haiku.rag.store.repositories.chunk import ChunkRepository
from haiku.rag.store.repositories.document import DocumentRepository
from haiku.rag.store.repositories.document_item import DocumentItemRepository
from haiku.rag.store.repositories.settings import SettingsRepository
__all__ = [
"ChunkRepository",
"DocumentItemRepository",
"DocumentRepository",
"SettingsRepository",
]

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

@ -4,11 +4,7 @@ from uuid import uuid4
from haiku.rag.store.engine import DocumentRecord, Store, get_documents_arrow_schema
from haiku.rag.store.models.document import Document
def _escape_sql_string(value: str) -> str:
"""Escape single quotes in SQL string literals."""
return value.replace("'", "''")
from haiku.rag.utils import escape_sql_string
class DocumentRepository:
@ -17,6 +13,7 @@ class DocumentRepository:
def __init__(self, store: Store) -> None:
self.store = store
self._chunk_repository = None
self._document_item_repository = None
@property
def chunk_repository(self):
@ -27,6 +24,17 @@ class DocumentRepository:
self._chunk_repository = ChunkRepository(self.store)
return self._chunk_repository
@property
def document_item_repository(self):
"""Lazy-load DocumentItemRepository when needed."""
if self._document_item_repository is None:
from haiku.rag.store.repositories.document_item import (
DocumentItemRepository,
)
self._document_item_repository = DocumentItemRepository(self.store)
return self._document_item_repository
def _record_to_document(self, record: DocumentRecord) -> Document:
"""Convert a DocumentRecord to a Document model."""
return Document(
@ -79,7 +87,7 @@ class DocumentRepository:
async def get_by_id(self, entity_id: str) -> Document | None:
"""Get a document by its ID."""
safe_id = _escape_sql_string(entity_id)
safe_id = escape_sql_string(entity_id)
results = list(
self.store.documents_table.search()
.where(f"id = '{safe_id}'")
@ -96,7 +104,7 @@ class DocumentRepository:
async def get_docling_data(self, entity_id: str) -> Document | None:
"""Get a document with only docling data loaded (skips content blob)."""
safe_id = _escape_sql_string(entity_id)
safe_id = escape_sql_string(entity_id)
results = list(
self.store.documents_table.search()
.select(self._DOCLING_COLUMNS)
@ -118,7 +126,7 @@ class DocumentRepository:
async def get_pages_data(self, entity_id: str) -> Document | None:
"""Get a document with only page image data loaded."""
safe_id = _escape_sql_string(entity_id)
safe_id = escape_sql_string(entity_id)
results = list(
self.store.documents_table.search()
.select(["id", "docling_pages"])
@ -140,19 +148,15 @@ 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)
# Update the record
safe_id = _escape_sql_string(entity.id)
safe_id = escape_sql_string(entity.id)
self.store.documents_table.update(
where=f"id = '{safe_id}'",
values={
@ -172,21 +176,18 @@ 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 cache before delete
invalidate_docling_document_cache(entity_id)
# Delete associated chunks first
# 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)
# Delete the document
safe_id = _escape_sql_string(entity_id)
safe_id = escape_sql_string(entity_id)
self.store.documents_table.delete(f"id = '{safe_id}'")
return True
@ -256,7 +257,7 @@ class DocumentRepository:
async def get_by_uri(self, uri: str) -> Document | None:
"""Get a document by its URI."""
escaped_uri = _escape_sql_string(uri)
escaped_uri = escape_sql_string(uri)
results = list(
self.store.documents_table.search()
.where(f"uri = '{escaped_uri}'")
@ -272,8 +273,23 @@ class DocumentRepository:
async def delete_all(self) -> None:
"""Delete all documents from the database."""
self.store._assert_writable()
# Delete all chunks first
from haiku.rag.store.engine import DocumentItemRecord
# Delete all chunks and items first
await self.chunk_repository.delete_all()
self.store.db.drop_table("document_items")
self.store.document_items_table = self.store.db.create_table(
"document_items", schema=DocumentItemRecord
)
self.store.document_items_table.create_scalar_index(
"document_id", index_type="BTREE", replace=True
)
self.store.document_items_table.create_scalar_index(
"position", index_type="BTREE", replace=True
)
self.store.document_items_table.create_scalar_index(
"self_ref", index_type="BTREE", replace=True
)
# Get count before deletion
count = len(

86
haiku_rag_slim/haiku/rag/store/repositories/document_item.py Normal file
View file

@ -0,0 +1,86 @@
import json
from haiku.rag.store.engine import DocumentItemRecord, Store
from haiku.rag.store.models.document_item import DocumentItem
from haiku.rag.utils import escape_sql_string
class DocumentItemRepository:
"""Repository for DocumentItem operations."""
def __init__(self, store: Store) -> None:
self.store = store
def _record_to_item(self, row: dict) -> DocumentItem:
return DocumentItem(
document_id=row["document_id"],
position=row["position"],
self_ref=row["self_ref"],
label=row.get("label", ""),
text=row.get("text", ""),
page_numbers=json.loads(row.get("page_numbers", "[]")),
)
async def create_items(self, document_id: str, items: list[DocumentItem]) -> None:
"""Bulk insert items for a document."""
if not items:
return
self.store._assert_writable()
records = [
DocumentItemRecord(
document_id=document_id,
position=item.position,
self_ref=item.self_ref,
label=item.label,
text=item.text,
page_numbers=json.dumps(item.page_numbers),
)
for item in items
]
self.store.document_items_table.add(records)
async def get_items_in_range(
self, document_id: str, start: int, end: int
) -> list[DocumentItem]:
"""Get items for a document within a position range (inclusive)."""
safe_id = escape_sql_string(document_id)
rows = (
self.store.document_items_table.search()
.where(
f"document_id = '{safe_id}' "
f"AND position >= {start} AND position <= {end}"
)
.to_list()
)
items = [self._record_to_item(row) for row in rows]
items.sort(key=lambda x: x.position)
return items
async def resolve_refs(self, document_id: str, refs: list[str]) -> dict[str, int]:
"""Resolve self_refs to positions. Returns {self_ref: position}."""
if not refs:
return {}
safe_id = escape_sql_string(document_id)
refs_sql = ", ".join(f"'{escape_sql_string(r)}'" for r in refs)
rows = (
self.store.document_items_table.search()
.select(["self_ref", "position"])
.where(f"document_id = '{safe_id}' AND self_ref IN ({refs_sql})")
.to_list()
)
return {row["self_ref"]: row["position"] for row in rows}
async def get_item_count(self, document_id: str) -> int:
"""Count items for a document."""
safe_id = escape_sql_string(document_id)
return self.store.document_items_table.count_rows(
filter=f"document_id = '{safe_id}'"
)
async def delete_by_document_id(self, document_id: str) -> None:
"""Delete all items for a document."""
self.store._assert_writable()
safe_id = escape_sql_string(document_id)
self.store.document_items_table.delete(f"document_id = '{safe_id}'")

4
haiku_rag_slim/haiku/rag/store/upgrades/__init__.py
View file

@ -81,8 +81,12 @@ from haiku.rag.store.upgrades.v0_25_0 import (
from haiku.rag.store.upgrades.v0_38_0 import (
upgrade_split_pages_zstd as upgrade_0_38_0_split_pages,
)
from haiku.rag.store.upgrades.v0_40_0 import (
upgrade_populate_document_items as upgrade_0_40_0_document_items,
)
upgrades.append(upgrade_0_20_0_docling)
upgrades.append(upgrade_0_23_1_contextualize)
upgrades.append(upgrade_0_25_0_compress)
upgrades.append(upgrade_0_38_0_split_pages)
upgrades.append(upgrade_0_40_0_document_items)

98
haiku_rag_slim/haiku/rag/store/upgrades/v0_40_0.py Normal file
View file

@ -0,0 +1,98 @@
import json
import logging
from haiku.rag.store.engine import DocumentItemRecord, Store
from haiku.rag.store.upgrades import Upgrade
from haiku.rag.utils import escape_sql_string
logger = logging.getLogger(__name__)
def _apply_populate_document_items(store: Store) -> None: # pragma: no cover
"""Populate document_items table from existing docling documents."""
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.store.compression import decompress_json
from haiku.rag.store.models.document_item import extract_items
# Get all document IDs that have docling data
ids = [
row["id"]
for row in store.documents_table.search().select(["id"]).to_arrow().to_pylist()
]
if not ids:
logger.info("No documents to migrate")
return
total = len(ids)
logger.info("Populating document_items for %d documents", total)
migrated = 0
skipped = 0
for idx, doc_id in enumerate(ids, 1):
# Load only docling data
safe_id = escape_sql_string(doc_id)
rows = (
store.documents_table.search()
.select(["id", "docling_document"])
.where(f"id = '{safe_id}'")
.limit(1)
.to_list()
)
if not rows:
skipped += 1
continue
row = rows[0]
docling_blob = row.get("docling_document")
if not docling_blob or not isinstance(docling_blob, bytes):
skipped += 1
continue
try:
json_str = decompress_json(docling_blob)
docling_doc = DoclingDocument.model_validate_json(json_str)
items = extract_items(doc_id, docling_doc)
if items:
records = [
DocumentItemRecord(
document_id=item.document_id,
position=item.position,
self_ref=item.self_ref,
label=item.label,
text=item.text,
page_numbers=json.dumps(item.page_numbers),
)
for item in items
]
store.document_items_table.add(records)
migrated += 1
if idx % 10 == 0 or idx == total:
logger.info(
"Progress: %d/%d documents (%d migrated, %d skipped)",
idx,
total,
migrated,
skipped,
)
except Exception:
logger.warning("Failed to extract items for document %s", doc_id)
skipped += 1
logger.info(
"Migration complete: %d migrated, %d skipped out of %d",
migrated,
skipped,
total,
)
upgrade_populate_document_items = Upgrade(
version="0.40.0",
apply=_apply_populate_document_items,
description="Populate document_items table for context expansion",
)

5
haiku_rag_slim/haiku/rag/utils.py
View file

@ -452,6 +452,11 @@ def build_prompt(base_prompt: str, config: "AppConfig") -> str:
return base_prompt
def escape_sql_string(value: str) -> str:
"""Escape single quotes in SQL string literals."""
return value.replace("'", "''")
def get_package_versions() -> dict[str, str]:
"""Get versions of haiku.rag and its dependencies.

5
haiku_rag_slim/pyproject.toml
View file

@ -2,7 +2,7 @@
name = "haiku.rag-slim"
description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling - Minimal dependencies"
version = "0.39.0"
version = "0.40.0"
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
license = { text = "MIT" }
readme = { file = "README.md", content-type = "text/markdown" }
@ -22,8 +22,7 @@ classifiers = [
]
dependencies = [
"cachetools>=7.0.5",
"docling-core>=2.71.0",
"docling-core>=2.71.0,<2.72",
"haiku.skills>=0.14.0",
"httpx>=0.28.1",
"jinja2>=3.1.0",

4
pyproject.toml
View file

@ -2,7 +2,7 @@
name = "haiku.rag"
description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling"
version = "0.39.0"
version = "0.40.0"
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
license = { text = "MIT" }
readme = { file = "README.md", content-type = "text/markdown" }
@ -30,7 +30,7 @@ classifiers = [
]
dependencies = [
"haiku.rag-slim[docling,voyageai,mxbai,cohere,zeroentropy,tui]==0.39.0",
"haiku.rag-slim[docling,voyageai,mxbai,cohere,zeroentropy,tui]==0.40.0",
]
[project.scripts]

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

351
tests/store/test_document_items.py Normal file
View file

@ -0,0 +1,351 @@
import pytest
from haiku.rag.client import HaikuRAG
from haiku.rag.store.engine import Store
from haiku.rag.store.models.document_item import (
DocumentItem,
extract_item_text,
extract_items,
)
from haiku.rag.store.repositories.document_item import DocumentItemRepository
def _make_docling_doc():
"""Create a DoclingDocument with mixed item types for testing."""
from docling_core.types.doc.document import DoclingDocument, TableData
from docling_core.types.doc.labels import DocItemLabel
doc = DoclingDocument(name="test")
doc.add_text(label=DocItemLabel.SECTION_HEADER, text="Introduction")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="This is the first paragraph.")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="This is the second paragraph.")
doc.add_table(data=TableData(num_rows=2, num_cols=2, table_cells=[]))
doc.add_text(label=DocItemLabel.SECTION_HEADER, text="Conclusion")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Final thoughts here.")
return doc
class TestExtractItems:
def test_extracts_all_items(self):
doc = _make_docling_doc()
items = extract_items("doc-1", doc)
assert len(items) == 6
assert all(item.document_id == "doc-1" for item in items)
assert [item.position for item in items] == [0, 1, 2, 3, 4, 5]
def test_extracts_labels(self):
doc = _make_docling_doc()
items = extract_items("doc-1", doc)
assert items[0].label == "section_header"
assert items[1].label == "paragraph"
assert items[3].label == "table"
assert items[4].label == "section_header"
def test_extracts_text(self):
doc = _make_docling_doc()
items = extract_items("doc-1", doc)
assert items[0].text == "Introduction"
assert items[1].text == "This is the first paragraph."
assert items[5].text == "Final thoughts here."
def test_extracts_self_refs(self):
doc = _make_docling_doc()
items = extract_items("doc-1", doc)
assert all(item.self_ref.startswith("#/") for item in items)
def test_table_gets_markdown_text(self):
doc = _make_docling_doc()
items = extract_items("doc-1", doc)
table_item = items[3]
assert table_item.label == "table"
# Table should have some text from export_to_markdown
assert isinstance(table_item.text, str)
class TestExtractItemText:
def test_text_item(self):
from docling_core.types.doc.document import DoclingDocument
from docling_core.types.doc.labels import DocItemLabel
doc = DoclingDocument(name="test")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Hello world")
item, _ = next(iter(doc.iterate_items()))
assert extract_item_text(item, doc) == "Hello world"
def test_returns_none_for_empty_item(self):
from docling_core.types.doc.document import DoclingDocument
doc = DoclingDocument(name="test")
# An empty doc has no items to extract text from
items = extract_items("doc-1", doc)
assert items == []
@pytest.mark.asyncio
class TestDocumentItemRepository:
async def test_create_and_get_range(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag:
repo = DocumentItemRepository(rag.store)
items = [
DocumentItem(
document_id="doc-1",
position=i,
self_ref=f"#/texts/{i}",
label="paragraph",
text=f"Item {i}",
page_numbers=[1],
)
for i in range(10)
]
await repo.create_items("doc-1", items)
result = await repo.get_items_in_range("doc-1", 3, 7)
assert len(result) == 5
assert result[0].position == 3
assert result[-1].position == 7
assert result[0].text == "Item 3"
async def test_resolve_refs(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag:
repo = DocumentItemRepository(rag.store)
items = [
DocumentItem(
document_id="doc-1",
position=i,
self_ref=f"#/texts/{i}",
label="paragraph",
text=f"Item {i}",
)
for i in range(10)
]
await repo.create_items("doc-1", items)
refs = await repo.resolve_refs(
"doc-1", ["#/texts/2", "#/texts/7", "#/texts/999"]
)
assert refs == {"#/texts/2": 2, "#/texts/7": 7}
async def test_get_item_count(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag:
repo = DocumentItemRepository(rag.store)
items = [
DocumentItem(
document_id="doc-1",
position=i,
self_ref=f"#/texts/{i}",
label="paragraph",
text=f"Item {i}",
)
for i in range(15)
]
await repo.create_items("doc-1", items)
assert await repo.get_item_count("doc-1") == 15
assert await repo.get_item_count("nonexistent") == 0
async def test_delete_by_document_id(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag:
repo = DocumentItemRepository(rag.store)
for doc_id in ["doc-1", "doc-2"]:
items = [
DocumentItem(
document_id=doc_id,
position=i,
self_ref=f"#/texts/{i}",
label="paragraph",
text=f"Item {i}",
)
for i in range(5)
]
await repo.create_items(doc_id, items)
await repo.delete_by_document_id("doc-1")
assert await repo.get_item_count("doc-1") == 0
assert await repo.get_item_count("doc-2") == 5
async def test_empty_refs_returns_empty(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag:
repo = DocumentItemRepository(rag.store)
assert await repo.resolve_refs("doc-1", []) == {}
async def test_items_sorted_by_position(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag:
repo = DocumentItemRepository(rag.store)
# Insert in reverse order
items = [
DocumentItem(
document_id="doc-1",
position=9 - i,
self_ref=f"#/texts/{9 - i}",
label="paragraph",
text=f"Item {9 - i}",
)
for i in range(10)
]
await repo.create_items("doc-1", items)
result = await repo.get_items_in_range("doc-1", 0, 9)
positions = [item.position for item in result]
assert positions == sorted(positions)
@pytest.mark.asyncio
class TestDocumentItemPopulation:
async def test_store_document_populates_items(self, temp_db_path):
"""Test that _store_document_with_chunks populates items when given a docling_document."""
from haiku.rag.store.models.document import Document
docling_doc = _make_docling_doc()
async with HaikuRAG(temp_db_path, create=True) as rag:
document = Document(
content="test content",
uri="test://doc",
)
document.set_docling(docling_doc)
# Use _store_document_with_chunks directly with empty chunks
# to avoid needing embeddings
created = await rag._store_document_with_chunks(document, [], docling_doc)
assert created.id is not None
count = await rag.document_item_repository.get_item_count(created.id)
assert count == 6
items = await rag.document_item_repository.get_items_in_range(
created.id, 0, count
)
assert items[0].label == "section_header"
assert items[0].text == "Introduction"
assert items[1].label == "paragraph"
async def test_update_document_replaces_items(self, temp_db_path):
"""Test that _update_document_with_chunks replaces items."""
from docling_core.types.doc.document import DoclingDocument
from docling_core.types.doc.labels import DocItemLabel
from haiku.rag.store.models.document import Document
docling_doc = _make_docling_doc()
async with HaikuRAG(temp_db_path, create=True) as rag:
document = Document(
content="test content",
uri="test://doc",
)
document.set_docling(docling_doc)
created = await rag._store_document_with_chunks(document, [], docling_doc)
assert created.id is not None
assert await rag.document_item_repository.get_item_count(created.id) == 6
# Update with a simpler document
new_doc = DoclingDocument(name="updated")
new_doc.add_text(label=DocItemLabel.PARAGRAPH, text="Only one item now.")
created.set_docling(new_doc)
await rag._update_document_with_chunks(created, [], new_doc)
assert await rag.document_item_repository.get_item_count(created.id) == 1
async def test_delete_document_cascades_items(self, temp_db_path):
"""Test that deleting a document also deletes its items."""
from haiku.rag.store.models.document import Document
docling_doc = _make_docling_doc()
async with HaikuRAG(temp_db_path, create=True) as rag:
document = Document(
content="test content",
uri="test://doc",
)
document.set_docling(docling_doc)
created = await rag._store_document_with_chunks(document, [], docling_doc)
assert created.id is not None
assert await rag.document_item_repository.get_item_count(created.id) == 6
await rag.delete_document(created.id)
assert await rag.document_item_repository.get_item_count(created.id) == 0
class TestDocumentItemMigration:
def test_migration_populates_items_for_existing_documents(self, temp_db_path):
"""Test that the v0.40.0 migration populates items for pre-existing documents."""
from haiku.rag.store.compression import compress_docling_split
from haiku.rag.store.engine import DocumentRecord
docling_doc = _make_docling_doc()
json_str = docling_doc.model_dump_json()
structure, pages = compress_docling_split(json_str)
# Create a database at a pre-migration version with a document
store = Store(temp_db_path, create=True, skip_migration_check=True)
store.set_haiku_version("0.39.0")
doc_record = DocumentRecord(
id="test-doc-1",
content="test content",
uri="test://doc",
docling_document=structure,
docling_pages=pages,
docling_version=docling_doc.version,
)
store.documents_table.add([doc_record])
# Verify no items exist yet
assert store.document_items_table.count_rows() == 0
store.close()
# Re-open with skip_migration_check and run migration
store = Store(temp_db_path, skip_migration_check=True)
applied = store.migrate()
# Should have applied the v0.40.0 migration
assert any("document_items" in desc for desc in applied)
# Items should now exist
item_count = store.document_items_table.count_rows(
filter="document_id = 'test-doc-1'"
)
assert item_count == 6
# Verify item content
items = (
store.document_items_table.search()
.where("document_id = 'test-doc-1'")
.to_list()
)
labels = {row["label"] for row in items}
assert "section_header" in labels
assert "paragraph" in labels
assert "table" in labels
store.close()
def test_migration_skips_documents_without_docling(self, temp_db_path):
"""Test that migration handles documents without docling data."""
from haiku.rag.store.engine import DocumentRecord
store = Store(temp_db_path, create=True, skip_migration_check=True)
store.set_haiku_version("0.39.0")
doc_record = DocumentRecord(
id="no-docling",
content="plain text document",
)
store.documents_table.add([doc_record])
store.close()
store = Store(temp_db_path, skip_migration_check=True)
store.migrate()
# No items should have been created
assert store.document_items_table.count_rows() == 0
store.close()

8
tests/test_client.py
View file

@ -1502,12 +1502,12 @@ async def test_client_convert_with_html_format(temp_db_path):
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_sql_injection_is_blocked_with_escaping(temp_db_path):
"""SQL injection is blocked when using _escape_sql_string.
"""SQL injection is blocked when using escape_sql_string.
This test verifies that _escape_sql_string properly prevents SQL injection
This test verifies that escape_sql_string properly prevents SQL injection
by escaping single quotes in user input.
"""
from haiku.rag.store.repositories.document import _escape_sql_string
from haiku.rag.utils import escape_sql_string
async with HaikuRAG(temp_db_path, create=True) as client:
# Create documents
@ -1529,7 +1529,7 @@ async def test_sql_injection_is_blocked_with_escaping(temp_db_path):
# With proper escaping, single quotes become double quotes
# so the filter becomes: title = 'x'' OR title LIKE ''%'
# which searches for a literal title containing the injection string
safe_payload = _escape_sql_string(injection_payload)
safe_payload = escape_sql_string(injection_payload)
docs = await client.list_documents(filter=f"title = '{safe_payload}'")
# Should find 0 documents (injection is escaped, searching for literal string)

384
tests/test_context.py Normal file
View file

@ -0,0 +1,384 @@
import pytest
from haiku.rag.context import (
_expand_outward,
_find_expansion_range,
_merge_ranges,
expand_with_items,
)
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
@pytest.mark.asyncio
class TestExpandWithItems:
async def test_unresolvable_refs_returns_original(self, temp_db_path):
from haiku.rag.client import HaikuRAG
from haiku.rag.store.models.document import Document
async with HaikuRAG(temp_db_path, create=True) as rag:
doc = await rag._store_document_with_chunks(
Document(content="test"),
[],
__import__(
"docling_core.types.doc.document", fromlist=["DoclingDocument"]
).DoclingDocument(name="t"),
)
result = SearchResult(
content="original",
score=0.9,
document_id=doc.id,
doc_item_refs=["#/texts/999999"],
)
expanded = await expand_with_items(
rag.document_item_repository, doc.id, [result], 10, 5000
)
assert len(expanded) == 1
assert expanded[0].content == "original"
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",
position=0,
self_ref="#/texts/0",
label="section_header",
text="Table of Contents",
),
DocumentItem(
document_id="doc-1",
position=1,
self_ref="#/texts/1",
label="document_index",
text="x" * 2000,
),
DocumentItem(
document_id="doc-1",
position=2,
self_ref="#/texts/2",
label="section_header",
text="Introduction",
),
DocumentItem(
document_id="doc-1",
position=3,
self_ref="#/texts/3",
label="text",
text="Intro content. " * 100,
),
]
await rag.document_item_repository.create_items("doc-1", items)
result = SearchResult(
content="original chunk content",
score=0.9,
document_id="doc-1",
doc_item_refs=["#/texts/1"],
)
expanded = await expand_with_items(
rag.document_item_repository, "doc-1", [result], 10, 5000
)
assert len(expanded) == 1
# 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_fragmented_items_preserve_chunk(self, temp_db_path):
"""When items are fragmented (e.g., list_item children), the original
chunk content is preserved if expansion produces less text."""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(temp_db_path, create=True) as rag:
# Simulate docling's list_item structure: container with empty text,
# children with tiny fragments
items = [
DocumentItem(
document_id="doc-1",
position=0,
self_ref="#/texts/0",
label="section_header",
text="Steps",
),
DocumentItem(
document_id="doc-1",
position=1,
self_ref="#/texts/1",
label="list_item",
text="",
),
DocumentItem(
document_id="doc-1",
position=2,
self_ref="#/texts/2",
label="text",
text="Click",
),
DocumentItem(
document_id="doc-1",
position=3,
self_ref="#/texts/3",
label="text",
text="+",
),
DocumentItem(
document_id="doc-1",
position=4,
self_ref="#/texts/4",
label="text",
text="Add a New Service",
),
]
await rag.document_item_repository.create_items("doc-1", items)
# The chunk had properly assembled content from the chunker
result = SearchResult(
content="1. Click + Add a New Service in the dashboard.",
score=0.9,
document_id="doc-1",
doc_item_refs=["#/texts/1", "#/texts/2", "#/texts/3", "#/texts/4"],
)
expanded = await expand_with_items(
rag.document_item_repository, "doc-1", [result], 10, 5000
)
assert len(expanded) == 1
# Expansion produces "Steps\n\nClick\n\n+\n\nAdd a New Service" = 38 chars
# which is less than the chunk's 46 chars — fallback preserves the chunk
assert expanded[0].content == result.content

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

3
tests/test_info.py
View file

@ -5,6 +5,7 @@ import pytest
from haiku.rag.app import HaikuRAGApp
from haiku.rag.config.models import AppConfig, LanceDBConfig
from haiku.rag.store.engine import DocumentItemRecord
@pytest.mark.asyncio
@ -33,6 +34,7 @@ async def test_app_info_outputs(temp_db_path, capsys):
settings_tbl = db.create_table("settings", schema=SettingsRecord)
docs_tbl = db.create_table("documents", schema=DocumentRecord)
chunks_tbl = db.create_table("chunks", schema=ChunkRecord)
db.create_table("document_items", schema=DocumentItemRecord)
# Insert one of each - using the new config format
settings_tbl.add(
@ -113,6 +115,7 @@ async def test_app_info_with_vector_index(temp_db_path, capsys):
settings_tbl = db.create_table("settings", schema=SettingsRecord)
docs_tbl = db.create_table("documents", schema=DocumentRecord)
chunks_tbl = db.create_table("chunks", schema=ChunkRecord)
db.create_table("document_items", schema=DocumentItemRecord)
# Insert settings
settings_tbl.add(

8
uv.lock
View file

@ -1418,7 +1418,7 @@ wheels = [
[[package]]
name = "haiku-rag"
version = "0.39.0"
version = "0.40.0"
source = { editable = "." }
dependencies = [
{ name = "haiku-rag-slim", extra = ["cohere", "docling", "mxbai", "tui", "voyageai", "zeroentropy"] },
@ -1500,10 +1500,9 @@ requires-dist = [
[[package]]
name = "haiku-rag-slim"
version = "0.39.0"
version = "0.40.0"
source = { editable = "haiku_rag_slim" }
dependencies = [
{ name = "cachetools" },
{ name = "docling-core" },
{ name = "haiku-skills" },
{ name = "httpx" },
@ -1568,10 +1567,9 @@ 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" },
{ name = "docling-core", specifier = ">=2.71.0,<2.72" },
{ name = "haiku-skills", specifier = ">=0.14.0" },
{ name = "httpx", specifier = ">=0.28.1" },
{ name = "jinja2", specifier = ">=3.1.0" },