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

Add type-aware context expansion for search results

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-12-08 11:35:44 +02:00
parent dab04668bd
commit 9a5f8e4368
No known key found for this signature in database
20 changed files with 1259 additions and 413 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
CHANGELOG.md
View file

@ -37,9 +37,17 @@
- Inspector modal with keyboard navigation between pages
- CLI command: `haiku-rag visualize <chunk_id>`
- Requires `textual-image` dependency and terminal with image support
- **Type-Aware Context Expansion**: `expand_context()` now uses document structure for intelligent expansion
- Structural content (tables, code blocks, lists) expands to complete structures regardless of chunking
- Text content uses radius-based expansion via `text_context_radius` setting
- `max_context_items` and `max_context_chars` settings control expansion limits
- `SearchResult.format_for_agent()` method formats expanded results with metadata for LLM consumption
- **Inspector Context Modal**: Press `c` in the inspector to view expanded context for the selected chunk
### Changed
- **BREAKING: Config Renamed**: `context_chunk_radius` renamed to `text_context_radius`
- **BREAKING: `create_document()` API**: Removed `chunks` parameter
- `create_document()` now always processes content (converts, chunks, embeds)
- Use `import_document()` for pre-processed documents with custom chunks

4
docs/configuration/index.md
View file

@ -115,7 +115,9 @@ processing:
chunker: docling-local # docling-local or docling-serve
chunker_type: hybrid # hybrid or hierarchical
chunk_size: 256
context_chunk_radius: 0
text_context_radius: 0
max_context_items: 25
max_context_chars: 10000
chunking_tokenizer: "Qwen/Qwen3-Embedding-0.6B"
chunking_merge_peers: true
chunking_use_markdown_tables: false

26
docs/configuration/processing.md
View file

@ -10,7 +10,11 @@ Configure how documents are converted and chunked:
processing:
# Chunking configuration
chunk_size: 256 # Maximum tokens per chunk
context_chunk_radius: 0 # Context radius for chunk expansion
# Context expansion for search results
text_context_radius: 0 # Radius for text chunk expansion
max_context_items: 25 # Maximum items in expanded context
max_context_chars: 10000 # Maximum characters in expanded context
# Converter selection
converter: docling-local # docling-local or docling-serve
@ -133,20 +137,28 @@ processing:
- `false`: Tables as narrative text ("Value A, Column 2 = Value B")
- `true`: Tables as markdown (preserves table structure)
### Chunk Size and Context
### Chunk Size and Context Expansion
```yaml
processing:
# Chunk size for document processing
chunk_size: 256
# Number of adjacent chunks to include before/after retrieved chunks for context
# 0 = no expansion (default), 1 = include 1 chunk before and after, etc.
# When expanded chunks overlap or are adjacent, they are automatically merged
# into single chunks with continuous content to eliminate duplication
context_chunk_radius: 0
# Context expansion settings
# Controls how search results are expanded with surrounding content
text_context_radius: 0 # Chunks before/after to include for text content
max_context_items: 25 # Maximum doc items to include in expansion
max_context_chars: 10000 # Maximum characters in expanded content
```
Context expansion enriches search results with surrounding content from the source document:
- **text_context_radius**: For text content (paragraphs), includes N chunks 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.
- **max_context_chars**: Hard limit on total characters in expanded content.
Structural content (tables, code blocks, lists) uses type-aware expansion that automatically includes the complete structure regardless of how it was chunked. For example, if a table was split across multiple chunks, expansion retrieves the complete table.
## File Monitoring
Set directories to monitor for automatic indexing:

12
docs/inspector.md
View file

@ -36,6 +36,8 @@ Three panels display your data:
- `Tab` - Cycle between panels
- `↑` / `↓` - Navigate lists
- `/` - Open search modal
- `c` - Open context expansion modal (when viewing a chunk)
- `v` - Open visual grounding modal (when viewing a chunk)
- `q` - Quit
**Mouse:** Click to select, scroll to view content
@ -53,6 +55,16 @@ Press `/` to open the full-screen search modal:
Search uses hybrid (vector + full-text) search across all chunks. Content is rendered as markdown with syntax highlighting.
## Context Expansion
Press `c` while viewing a chunk to open the context expansion modal:
- Shows 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 `text_context_radius` setting
- Includes metadata like source document, content type, and relevance score
- Press `Esc` to close the modal
## Visual Grounding
Press `v` while viewing a chunk to open the visual grounding modal:

19
docs/python.md
View file

@ -309,20 +309,23 @@ Expand search results with adjacent chunks for more complete context:
# Get initial search results
search_results = await client.search("machine learning", limit=3)
# Expand with adjacent chunks using config setting
# Expand search results with adjacent content from the source document
expanded_results = await client.expand_context(search_results)
# Or specify a custom radius
expanded_results = await client.expand_context(search_results, radius=2)
# The expanded results contain chunks with combined content from adjacent chunks
# The expanded results contain chunks with combined content
for result in expanded_results:
print(f"Expanded content: {result.content}") # Now includes before/after chunks
print(f"Expanded content: {result.content}")
```
**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.
Context expansion uses your configuration settings:
This is automatically used by the QA system when `processing.context_chunk_radius > 0` (configured in `haiku.rag.yaml`) to provide better answers with more complete context.
- **text_context_radius**: For text content (paragraphs), includes N chunks before and after
- **max_context_items**: Limits how many document items can be included
- **max_context_chars**: Hard limit on total characters
**Type-aware expansion**: Structural content (tables, code blocks, lists) automatically expands to include the complete structure, regardless of how it was split during chunking.
**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.
## Question Answering

4
evaluations/evaluations/benchmark.py
View file

@ -45,7 +45,9 @@ def build_experiment_metadata(
"embedder_model": config.embeddings.model.name,
"embedder_dim": config.embeddings.model.vector_dim,
"chunk_size": config.processing.chunk_size,
"context_chunk_radius": config.processing.context_chunk_radius,
"text_context_radius": config.processing.text_context_radius,
"max_context_items": config.processing.max_context_items,
"max_context_chars": config.processing.max_context_chars,
"rerank_provider": config.reranking.model.provider
if config.reranking.model
else None,

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

@ -941,7 +941,6 @@ class HaikuRAG:
async def expand_context(
self,
search_results: list[SearchResult],
radius: int | None = None,
) -> list[SearchResult]:
"""Expand search results with adjacent content from the source document.
@ -949,18 +948,20 @@ class HaikuRAG:
by finding adjacent DocItems with accurate bounding boxes and metadata.
Otherwise, falls back to chunk-based expansion using adjacent chunks.
Expansion is type-aware based on content:
- Tables, code blocks, and lists expand to include complete structures
- Text content uses the configured radius (text_context_radius)
- Expansion is limited by max_context_items and max_context_chars
Args:
search_results: List of SearchResult objects from search.
radius: Number of adjacent items to include before/after.
If None, uses config.processing.context_chunk_radius.
Returns:
List of SearchResult objects with expanded content and resolved provenance.
"""
if radius is None:
radius = self._config.processing.context_chunk_radius
if radius == 0:
return search_results
radius = self._config.processing.text_context_radius
max_items = self._config.processing.max_context_items
max_chars = self._config.processing.max_context_chars
# Group by document_id for efficient processing
document_groups: dict[str | None, list[SearchResult]] = {}
@ -992,13 +993,22 @@ class HaikuRAG:
if has_docling and has_refs:
# Use DoclingDocument-based expansion
expanded = await self._expand_with_docling(
doc_results, docling_doc, radius
doc_results,
docling_doc,
radius,
max_items,
max_chars,
)
expanded_results.extend(expanded)
else:
# Fall back to chunk-based expansion
expanded = await self._expand_with_chunks(doc_id, doc_results, radius)
expanded_results.extend(expanded)
# 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)
return expanded_results
@ -1028,13 +1038,133 @@ class HaikuRAG:
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"}
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 caption if available
"""
# Try simple text attribute first (works for most items)
if text := getattr(item, "text", None):
return text
# For tables, export as markdown
if hasattr(item, "export_to_markdown"):
try:
return item.export_to_markdown(docling_doc)
except Exception:
pass
# For pictures/charts, try to get caption
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."""
"""Expand results using DoclingDocument structure.
Structural content (tables, code, lists) expands to complete structures.
Text content uses radius-based expansion.
"""
from haiku.rag.store.models.chunk import BoundingBox
all_items = list(docling_doc.iterate_items())
@ -1055,8 +1185,11 @@ class HaikuRAG:
if not indices:
passthrough.append(result)
continue
min_idx = max(0, min(indices) - radius)
max_idx = min(len(all_items) - 1, max(indices) + radius)
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
@ -1068,7 +1201,9 @@ class HaikuRAG:
for i in range(min_idx, max_idx + 1):
item, _ = all_items[i]
if text := getattr(item, "text", None):
# 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)

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

@ -110,7 +110,9 @@ class ConversionOptions(BaseModel):
class ProcessingConfig(BaseModel):
chunk_size: int = 256
context_chunk_radius: int = 0
text_context_radius: int = 0
max_context_items: int = 25
max_context_chars: int = 10000
converter: str = "docling-local"
chunker: str = "docling-local"
chunker_type: str = "hybrid"

5
haiku_rag_slim/haiku/rag/graph/common/nodes.py
View file

@ -239,9 +239,8 @@ async def _do_search[AgentDepsT: GraphAgentDeps](
# Store results for citation resolution
ctx2.deps.search_results = results
parts = []
for r in results:
parts.append(f"[{r.chunk_id}] (score: {r.score:.2f}) {r.content}")
# Format with metadata for agent context
parts = [r.format_for_agent() for r in results]
if not parts:
return f"No relevant information found in the knowledge base for: {query}"

21
haiku_rag_slim/haiku/rag/graph/common/prompts.py
View file

@ -27,14 +27,31 @@ Process:
4. Provide a concise answer based strictly on the retrieved content.
The search tool returns results like:
[9bde5847-44c9-400a-8997-0e6b65babf92] (score: 0.85) Content text here...
[d5a63c82-cb40-439f-9b2e-de7d177829b7] (score: 0.72) More content...
[9bde5847-44c9-400a-8997-0e6b65babf92] (score: 0.85)
Source: "Document Title" > Section > Subsection
Type: paragraph
Content:
The actual text content here...
[d5a63c82-cb40-439f-9b2e-de7d177829b7] (score: 0.72)
Source: "Another Document"
Type: table
Content:
| Column 1 | Column 2 |
...
Each result includes:
- chunk_id in brackets and relevance score
- Source: document title and section hierarchy (when available)
- Type: content type like paragraph, table, code, list_item (when available)
- Content: the actual text
IMPORTANT: In cited_chunks, use the EXACT, COMPLETE chunk ID (the full UUID).
Do NOT truncate or shorten chunk IDs.
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 so clearly.
- Be concise and direct; avoid meta commentary about the process.

17
haiku_rag_slim/haiku/rag/inspector/app.py
View file

@ -69,6 +69,7 @@ class InspectorApp(App): # type: ignore[misc] # pragma: no cover
Binding("q", "quit", "Quit", show=True),
Binding("/", "search", "Search", show=True),
Binding("v", "show_visual", "Visual", show=True),
Binding("c", "show_context", "Context", show=True),
]
def __init__(self, db_path: Path):
@ -192,6 +193,22 @@ class InspectorApp(App): # type: ignore[misc] # pragma: no cover
await self.push_screen(VisualGroundingModal(chunk=chunk, client=self.client))
async def action_show_context(self) -> None:
"""Show how the currently selected chunk would be formatted for agents."""
if not self.client:
return
chunk_list = self.query_one(ChunkList)
idx = chunk_list.list_view.index
if idx is None or idx >= len(chunk_list.chunks):
return
chunk = chunk_list.chunks[idx]
from haiku.rag.inspector.widgets.context_modal import ContextModal
await self.push_screen(ContextModal(chunk=chunk, client=self.client))
def run_inspector(db_path: Path | None = None) -> None: # pragma: no cover
"""Run the inspector TUI.

89
haiku_rag_slim/haiku/rag/inspector/widgets/context_modal.py Normal file
View file

@ -0,0 +1,89 @@
from typing import TYPE_CHECKING
from textual.app import ComposeResult
from textual.binding import Binding
from textual.containers import VerticalScroll
from textual.screen import Screen
from textual.widgets import Markdown, Static
from haiku.rag.store.models import SearchResult
if TYPE_CHECKING:
from haiku.rag.client import HaikuRAG
from haiku.rag.store.models import Chunk
class ContextModal(Screen): # pragma: no cover
"""Modal screen for displaying how a chunk appears to agents."""
BINDINGS = [
Binding("escape", "dismiss", "Close", show=True),
Binding("c", "dismiss", "Close", show=True),
]
CSS = """
ContextModal {
background: $surface;
layout: vertical;
}
#context-header {
dock: top;
height: auto;
padding: 1;
}
#context-content {
height: 1fr;
width: 100%;
padding: 1;
}
#context-content Markdown {
width: 100%;
}
"""
def __init__(self, chunk: "Chunk", client: "HaikuRAG"):
super().__init__()
self.chunk = chunk
self.client = client
self._content_widget = Markdown("Loading...")
def compose(self) -> ComposeResult:
yield Static("[bold]Agent Context Format[/bold]", id="context-header")
with VerticalScroll(id="context-content"):
yield self._content_widget
async def on_mount(self) -> None:
"""Load and display the expanded context."""
# Create a SearchResult from the chunk
chunk_meta = self.chunk.get_chunk_metadata()
search_result = SearchResult(
content=self.chunk.content,
score=0.0,
chunk_id=self.chunk.id,
document_id=self.chunk.document_id,
document_uri=self.chunk.document_uri,
document_title=self.chunk.document_title,
doc_item_refs=chunk_meta.doc_item_refs,
page_numbers=chunk_meta.page_numbers,
headings=chunk_meta.headings,
labels=chunk_meta.labels,
)
# Expand context using the client (this is what agents actually receive)
expanded_results = await self.client.expand_context([search_result])
expanded = expanded_results[0] if expanded_results else search_result
formatted = expanded.format_for_agent()
content = (
"*This is how the chunk appears to agents after context expansion:*\n\n---\n\n"
f"{formatted}"
)
await self._content_widget.update(content)
async def action_dismiss(self, result=None) -> None:
self.app.pop_screen()

1
haiku_rag_slim/haiku/rag/inspector/widgets/visual_modal.py
View file

@ -20,6 +20,7 @@ class VisualGroundingModal(Screen): # pragma: no cover
BINDINGS = [
Binding("escape", "dismiss", "Close", show=True),
Binding("v", "dismiss", "Close", show=True),
Binding("left", "prev_page", "Previous Page"),
Binding("right", "next_page", "Next Page"),
]

6
haiku_rag_slim/haiku/rag/qa/agent.py
View file

@ -50,10 +50,8 @@ class QuestionAnswerAgent:
results = await ctx.deps.client.expand_context(results)
# Store results for citation resolution
ctx.deps.search_results = results
# Format with chunk IDs
parts = []
for r in results:
parts.append(f"[{r.chunk_id}] (score: {r.score:.2f}) {r.content}")
# Format with metadata for agent context
parts = [r.format_for_agent() for r in results]
return "\n\n".join(parts) if parts else "No results found."
async def answer(self, question: str) -> tuple[str, list[Citation]]:

21
haiku_rag_slim/haiku/rag/qa/prompts.py
View file

@ -7,13 +7,30 @@ Process:
4. Provide a concise answer based strictly on the retrieved content
The search tool returns results like:
[chunk_abc123] (score: 0.85) Content text here...
[chunk_def456] (score: 0.72) More content...
[chunk_abc123] (score: 0.85)
Source: "Document Title" > Section > Subsection
Type: paragraph
Content:
The actual text content here...
[chunk_def456] (score: 0.72)
Source: "Another Document"
Type: table
Content:
| Column 1 | Column 2 |
...
Each result includes:
- chunk_id in brackets and relevance score
- Source: document title and section hierarchy (when available)
- Type: content type like paragraph, table, code, list_item (when available)
- Content: the actual text
In your response, include the chunk IDs you used in cited_chunks.
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

63
haiku_rag_slim/haiku/rag/store/models/chunk.py
View file

@ -145,3 +145,66 @@ class SearchResult(BaseModel):
labels=meta.labels,
bounding_boxes=bounding_boxes,
)
def format_for_agent(self) -> str:
"""Format this search result for inclusion in agent context.
Produces a structured format with metadata that helps LLMs understand
the source and nature of the content.
"""
parts = [f"[{self.chunk_id}] (score: {self.score:.2f})"]
# Document source info
source_parts = []
if self.document_title:
source_parts.append(f'"{self.document_title}"')
if self.headings:
source_parts.append(" > ".join(self.headings))
if source_parts:
parts.append(f"Source: {' > '.join(source_parts)}")
# Content type (use primary label if available)
if self.labels:
primary_label = self._get_primary_label()
if primary_label:
parts.append(f"Type: {primary_label}")
# The actual content
parts.append(f"Content:\n{self.content}")
return "\n".join(parts)
def _get_primary_label(self) -> str | None:
"""Get the most significant label for display.
Prioritizes structural labels over text labels.
"""
if not self.labels:
return None
# Priority order: structural > contextual > text
priority = {
"table": 1,
"code": 2,
"form": 3,
"key_value_region": 4,
"list_item": 5,
"formula": 6,
"chart": 7,
"picture": 8,
"caption": 9,
"footnote": 10,
"section_header": 11,
"title": 12,
}
# Find highest priority label
best_label = None
best_priority = float("inf")
for label in self.labels:
if label in priority and priority[label] < best_priority:
best_label = label
best_priority = priority[label]
# Return best structural/special label, or first label if all are text
return best_label if best_label else self.labels[0]

92
tests/test_chunk.py
View file

@ -3,7 +3,7 @@ from datasets import Dataset
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.store.models.chunk import Chunk, ChunkMetadata
from haiku.rag.store.models.chunk import Chunk, ChunkMetadata, SearchResult
@pytest.mark.asyncio
@ -203,3 +203,93 @@ def test_chunk_metadata_resolve_empty_refs():
doc_items = chunk_meta.resolve_doc_items(docling_doc)
assert doc_items == []
def test_search_result_format_for_agent_full():
"""Test format_for_agent with all metadata present."""
result = SearchResult(
content="This is the chunk content about elections.",
score=0.85,
chunk_id="chunk-123",
document_id="doc-456",
document_uri="file:///docs/report.pdf",
document_title="Annual Report 2024",
headings=["Chapter 1", "Section 1.1", "Elections"],
labels=["paragraph", "table"],
page_numbers=[1, 2],
)
formatted = result.format_for_agent()
assert "[chunk-123]" in formatted
assert "(score: 0.85)" in formatted
assert (
'Source: "Annual Report 2024" > Chapter 1 > Section 1.1 > Elections'
in formatted
)
assert "Type: table" in formatted # table has higher priority than paragraph
assert "Content:\nThis is the chunk content about elections." in formatted
def test_search_result_format_for_agent_minimal():
"""Test format_for_agent with minimal metadata."""
result = SearchResult(
content="Some content here.",
score=0.72,
chunk_id="chunk-abc",
)
formatted = result.format_for_agent()
assert "[chunk-abc]" in formatted
assert "(score: 0.72)" in formatted
assert "Source:" not in formatted # No title or headings
assert "Type:" not in formatted # No labels
assert "Content:\nSome content here." in formatted
def test_search_result_format_for_agent_title_only():
"""Test format_for_agent with only document title."""
result = SearchResult(
content="Content text.",
score=0.60,
chunk_id="chunk-xyz",
document_title="My Document",
)
formatted = result.format_for_agent()
assert 'Source: "My Document"' in formatted
def test_search_result_format_for_agent_headings_only():
"""Test format_for_agent with only headings (no title)."""
result = SearchResult(
content="Content text.",
score=0.60,
chunk_id="chunk-xyz",
headings=["Introduction", "Background"],
)
formatted = result.format_for_agent()
assert "Source: Introduction > Background" in formatted
def test_search_result_get_primary_label():
"""Test _get_primary_label prioritization."""
# Table takes priority over text labels
result = SearchResult(content="x", score=0.5, labels=["paragraph", "table", "text"])
assert result._get_primary_label() == "table"
# Code takes priority over list_item
result = SearchResult(content="x", score=0.5, labels=["list_item", "code"])
assert result._get_primary_label() == "code"
# Text labels fall through to first
result = SearchResult(content="x", score=0.5, labels=["paragraph", "text"])
assert result._get_primary_label() == "paragraph"
# Empty labels
result = SearchResult(content="x", score=0.5, labels=[])
assert result._get_primary_label() is None

371
tests/test_client.py
View file

@ -758,384 +758,19 @@ async def test_client_ask(monkeypatch, temp_db_path):
@pytest.mark.asyncio
async def test_client_expand_context(temp_db_path):
"""Test expanding search results with adjacent chunks."""
from haiku.rag.store.models import SearchResult
# Mock Config to have CONTEXT_CHUNK_RADIUS = 2
with patch("haiku.rag.client.Config.processing.context_chunk_radius", 2):
async with HaikuRAG(temp_db_path, create=True) as client:
# Create chunks manually with precomputed embeddings to avoid network
dim = client.chunk_repository.embedder._vector_dim
z = [0.0] * dim
manual_chunks = [
Chunk(content="Chunk 0 content", order=0, embedding=z),
Chunk(content="Chunk 1 content", order=1, embedding=z),
Chunk(content="Chunk 2 content", order=2, embedding=z),
Chunk(content="Chunk 3 content", order=3, embedding=z),
Chunk(content="Chunk 4 content", order=4, embedding=z),
]
doc = await client.import_document(
content="Full document content",
chunks=manual_chunks,
uri="test_doc.txt",
title="test_doc_title",
)
# Get all chunks for the document
assert doc.id is not None
chunks = await client.chunk_repository.get_by_document_id(doc.id)
assert len(chunks) == 5
# Find the middle chunk (order=2) and convert to SearchResult
middle_chunk = next(c for c in chunks if c.order == 2)
search_results = [SearchResult.from_chunk(middle_chunk, 0.8)]
# Test expand_context with radius=2 and document title preserved
expanded_results = await client.expand_context(search_results, radius=2)
assert len(expanded_results) == 1
expanded = expanded_results[0]
# Check that the expanded result has combined content and preserves title/uri
assert expanded.score == 0.8
assert "Chunk 2 content" in expanded.content
assert expanded.document_title == "test_doc_title"
assert expanded.document_uri == "test_doc.txt"
# Should include all chunks (radius=2 from chunk 2 = chunks 0,1,2,3,4)
assert "Chunk 0 content" in expanded.content
assert "Chunk 1 content" in expanded.content
assert "Chunk 2 content" in expanded.content
assert "Chunk 3 content" in expanded.content
assert "Chunk 4 content" in expanded.content
@pytest.mark.asyncio
async def test_client_expand_context_radius_zero(temp_db_path):
"""Test expand_context with radius 0 returns original results."""
"""Test that expand_context method exists and works with basic input."""
from haiku.rag.store.models import SearchResult
async with HaikuRAG(temp_db_path, create=True) as client:
# Create a simple document
doc = await client.create_document(content="Simple test content")
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)]
expanded_results = await client.expand_context(search_results, radius=0)
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
@pytest.mark.asyncio
async def test_client_expand_context_multiple_chunks(temp_db_path):
"""Test expand_context with multiple search results."""
from haiku.rag.store.models import SearchResult
with patch("haiku.rag.client.Config.processing.context_chunk_radius", 1):
async with HaikuRAG(temp_db_path, create=True) as client:
# Create first document with manual chunks
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(
content="Doc1 content", chunks=doc1_chunks, uri="doc1.txt"
)
# Create second document with manual chunks
doc2_chunks = [
Chunk(content="Doc2 Section X", order=0),
Chunk(content="Doc2 Section Y", order=1),
]
doc2 = await client.import_document(
content="Doc2 content", 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)
search_results = [
SearchResult.from_chunk(chunk1, 0.8),
SearchResult.from_chunk(chunk2, 0.7),
]
expanded_results = await client.expand_context(search_results, radius=1)
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.asyncio
async def test_client_expand_context_merges_overlapping_chunks(temp_db_path):
"""Test that overlapping expanded chunks are merged into one."""
from haiku.rag.store.models import SearchResult
async with HaikuRAG(temp_db_path, create=True) as client:
# Create document with 5 chunks
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(
content="Full document content", 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, radius=1)
# 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.asyncio
async def test_client_expand_context_keeps_separate_non_overlapping(temp_db_path):
"""Test that non-overlapping expanded chunks remain separate."""
from haiku.rag.store.models import SearchResult
async with HaikuRAG(temp_db_path, create=True) as client:
# Create document with chunks far apart
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(
content="Full document content", 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, radius=1)
# 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
) # Should not have chunk 7 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
@pytest.mark.asyncio
async def test_client_expand_context_with_docling_merges_overlapping(temp_db_path):
"""Test that expand_context with DoclingDocument merges overlapping results."""
from haiku.rag.store.models import SearchResult
# Create a document with structured content that will have doc_item_refs
markdown_content = """# Chapter 1
This is paragraph one about topic A.
This is paragraph two about topic A continued.
This is paragraph three about topic B.
# Chapter 2
This is paragraph four about topic C.
"""
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content=markdown_content,
uri="test://structured",
)
assert doc.id is not None
assert doc.docling_document_json is not None
# Get chunks which should have doc_item_refs
chunks = await client.chunk_repository.get_by_document_id(doc.id)
assert len(chunks) >= 1
# Find chunks that have doc_item_refs (from docling chunking)
chunks_with_refs = [c for c in chunks if c.get_chunk_metadata().doc_item_refs]
if len(chunks_with_refs) >= 2:
# Create search results from adjacent chunks
search_results = [
SearchResult.from_chunk(chunks_with_refs[0], 0.9),
SearchResult.from_chunk(chunks_with_refs[1], 0.8),
]
# Expand with radius that should cause overlap
expanded = await client.expand_context(search_results, radius=3)
# If chunks were adjacent, they should be merged
# The expanded results should have merged metadata
assert len(expanded) >= 1
# Check that expanded result has page_numbers populated
for r in expanded:
# Should have doc_item_refs from expansion
assert r.doc_item_refs is not None
@pytest.mark.asyncio
async def test_client_expand_context_docling_merges_metadata(temp_db_path):
"""Test that expand_context properly merges metadata from multiple results."""
from haiku.rag.store.models import SearchResult
markdown_content = """# Introduction
First paragraph of introduction.
Second paragraph of introduction.
# Methods
First paragraph of methods section.
Second paragraph of methods section.
# Results
First paragraph of results.
"""
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content=markdown_content,
uri="test://metadata-merge",
)
assert doc.id is not None
chunks = await client.chunk_repository.get_by_document_id(doc.id)
chunks_with_refs = [c for c in chunks if c.get_chunk_metadata().doc_item_refs]
if len(chunks_with_refs) >= 2:
# Get chunks with different headings if possible
chunk1 = chunks_with_refs[0]
chunk2 = chunks_with_refs[-1] # Last chunk likely has different heading
search_results = [
SearchResult.from_chunk(chunk1, 0.9),
SearchResult.from_chunk(chunk2, 0.8),
]
# Expand with large radius to potentially merge
expanded = await client.expand_context(search_results, radius=10)
# Check that results have proper structure
for r in expanded:
# Content should be non-empty
assert len(r.content) > 0
# Score should be preserved (best score)
assert r.score in [0.9, 0.8]
@pytest.mark.asyncio
async def test_client_expand_context_docling_preserves_bounding_boxes(temp_db_path):
"""Test that expand_context preserves bounding boxes from DoclingDocument."""
from haiku.rag.store.models import SearchResult
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="# Test\n\nSome content here.",
uri="test://bboxes",
)
assert doc.id is not None
chunks = await client.chunk_repository.get_by_document_id(doc.id)
if chunks:
search_results = [SearchResult.from_chunk(chunks[0], 0.9)]
expanded = await client.expand_context(search_results, radius=2)
# Expanded results should exist
assert len(expanded) == 1
# Bounding boxes may or may not be present depending on document
# but the field should be accessible
_ = expanded[0].bounding_boxes
assert expanded_results[0].score == 0.9
@pytest.mark.asyncio

716
tests/test_context_enhancement.py Normal file
View file

@ -0,0 +1,716 @@
import pytest
from docling_core.types.doc.document import DoclingDocument, TableData
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(
client: HaikuRAG, docling_doc: DoclingDocument, title: str
):
"""Helper to create a document from a DoclingDocument using import_document."""
chunks = await client.chunk(docling_doc)
embedded_chunks = await client._ensure_chunks_embedded(chunks)
return await client.import_document(
chunks=embedded_chunks,
title=title,
docling_document_json=docling_doc.model_dump_json(),
docling_version=docling_doc.version,
)
def create_table_document() -> DoclingDocument:
"""Create a document with a table that will be split across chunks."""
doc = DoclingDocument(name="table_test")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Introduction paragraph.")
doc.add_heading(text="Employee Data", level=1)
# Create a table with enough content to span multiple chunks
table_data = TableData(num_cols=3, num_rows=0)
table_data.add_row(["Name", "Age", "City"])
table_data.add_row(["Alice Smith", "30", "New York"])
table_data.add_row(["Bob Johnson", "25", "Los Angeles"])
table_data.add_row(["Charlie Brown", "35", "Chicago"])
table_data.add_row(["Diana Ross", "28", "Miami"])
doc.add_table(data=table_data)
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Conclusion paragraph.")
return doc
def create_list_document() -> DoclingDocument:
"""Create a document with list items that will be split across chunks."""
doc = DoclingDocument(name="list_test")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Shopping list for the week:")
doc.add_list_item(
text="Fresh organic apples from the farmers market", enumerated=False
)
doc.add_list_item(text="Ripe yellow bananas for smoothies", enumerated=False)
doc.add_list_item(text="Valencia oranges for fresh juice", enumerated=False)
doc.add_list_item(text="Seedless red grapes as healthy snack", enumerated=False)
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Remember to bring reusable bags.")
return doc
def create_code_document() -> DoclingDocument:
"""Create a document with adjacent code blocks that will be split."""
doc = DoclingDocument(name="code_test")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Here are several code snippets:")
# Multiple adjacent code blocks - type-aware expansion should group them
doc.add_text(label=DocItemLabel.CODE, text="# Part 1: Setup\nimport os\nimport sys")
doc.add_text(
label=DocItemLabel.CODE, text='# Part 2: Config\nCONFIG = {"debug": True}'
)
doc.add_text(
label=DocItemLabel.CODE, text="# Part 3: Main\ndef main():\n print(CONFIG)"
)
doc.add_text(label=DocItemLabel.PARAGRAPH, text="End of code examples.")
return doc
@pytest.fixture
def small_chunk_config() -> AppConfig:
"""Config with small chunk size to force splitting."""
config = AppConfig()
config.processing.chunk_size = 32
config.processing.max_context_items = 25
config.processing.max_context_chars = 10000
return config
@pytest.mark.asyncio
async def test_table_expansion_includes_split_rows(temp_db_path, small_chunk_config):
"""Verify that table expansion retrieves rows that were split into different chunks."""
docling_doc = create_table_document()
async with HaikuRAG(temp_db_path, config=small_chunk_config, create=True) as client:
doc = await create_document_with_docling(client, docling_doc, "Table Test")
assert doc.id is not None
# Search for table content
results = await client.search("Alice Smith New York employee", limit=5)
table_results = [r for r in results if "table" in r.labels]
assert len(table_results) > 0, (
f"No table results. Labels: {[r.labels for r in results]}"
)
original = table_results[0]
# Verify the original chunk does NOT contain all table data
# (proving we need expansion)
original_has_all = all(
name in original.content for name in ["Alice", "Bob", "Charlie", "Diana"]
)
# Expand context
expanded = await client.expand_context(table_results[:1])
assert len(expanded) == 1
expanded_content = expanded[0].content
# After expansion, we should have the complete table
assert "Alice" in expanded_content
assert "Bob" in expanded_content
assert "Charlie" in expanded_content
assert "Diana" in expanded_content
# Verify expansion actually added content (unless chunk already had everything)
if not original_has_all:
assert len(expanded_content) > len(original.content), (
"Expansion should have added content"
)
@pytest.mark.asyncio
async def test_list_expansion_includes_split_items(temp_db_path, small_chunk_config):
"""Verify that list expansion retrieves items that were split into different chunks."""
docling_doc = create_list_document()
async with HaikuRAG(temp_db_path, config=small_chunk_config, create=True) as client:
doc = await create_document_with_docling(client, docling_doc, "List Test")
assert doc.id is not None
# Search for a list item
results = await client.search("grapes healthy snack", limit=5)
list_results = [r for r in results if "list_item" in r.labels]
assert len(list_results) > 0, (
f"No list results. Labels: {[r.labels for r in results]}"
)
original = list_results[0]
# Check what the original chunk contains
original_items = sum(
1
for item in ["apples", "bananas", "oranges", "grapes"]
if item in original.content.lower()
)
# Expand context
expanded = await client.expand_context(list_results[:1])
assert len(expanded) == 1
expanded_content = expanded[0].content.lower()
# Count items after expansion
expanded_items = sum(
1
for item in ["apples", "bananas", "oranges", "grapes"]
if item in expanded_content
)
# Expansion should include at least as many items (more if split)
assert expanded_items >= original_items
# If original didn't have all items, expansion should have added some
if original_items < 4:
assert expanded_items > original_items, (
f"Expansion should have added items. Original: {original_items}, Expanded: {expanded_items}"
)
@pytest.mark.asyncio
async def test_code_expansion_includes_adjacent_blocks(
temp_db_path, small_chunk_config
):
"""Verify that code expansion retrieves adjacent code blocks split across chunks."""
docling_doc = create_code_document()
async with HaikuRAG(temp_db_path, config=small_chunk_config, create=True) as client:
doc = await create_document_with_docling(client, docling_doc, "Code Test")
assert doc.id is not None
# Search for middle code block (Part 2)
results = await client.search("CONFIG debug True", limit=5)
code_results = [r for r in results if "code" in r.labels]
assert len(code_results) > 0, (
f"No code results. Labels: {[r.labels for r in results]}"
)
original = code_results[0]
# Check what parts the original chunk has
original_parts = sum(
1 for part in ["Part 1", "Part 2", "Part 3"] if part in original.content
)
# Expand context
expanded = await client.expand_context(code_results[:1])
assert len(expanded) == 1
expanded_content = expanded[0].content
# Count parts after expansion
expanded_parts = sum(
1 for part in ["Part 1", "Part 2", "Part 3"] if part in expanded_content
)
# Expansion should include at least as many parts
assert expanded_parts >= original_parts
# If original didn't have all parts, expansion should have added some
if original_parts < 3:
assert expanded_parts > original_parts, (
f"Expansion should have added code blocks. Original: {original_parts}, Expanded: {expanded_parts}"
)
@pytest.mark.asyncio
async def test_text_expansion_uses_radius(temp_db_path):
"""Text content expansion should use radius, not structural boundaries."""
config = AppConfig()
config.processing.chunk_size = 32
config.processing.text_context_radius = 1 # Small radius
# Create a document with longer paragraphs that will split
doc = DoclingDocument(name="text_test")
doc.add_text(
label=DocItemLabel.PARAGRAPH,
text="First paragraph with enough content to be its own chunk in the document.",
)
doc.add_text(
label=DocItemLabel.PARAGRAPH,
text="Second paragraph contains different information about software testing.",
)
doc.add_text(
label=DocItemLabel.PARAGRAPH,
text="Third paragraph discusses various topics and provides more details.",
)
doc.add_text(
label=DocItemLabel.PARAGRAPH,
text="Fourth paragraph concludes the document with final thoughts.",
)
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
document = await create_document_with_docling(client, doc, "Text Only")
assert document.id is not None
# Search for second paragraph
results = await client.search("software testing", limit=1)
assert len(results) > 0
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)
assert len(expanded[0].content) >= len(original.content)
@pytest.mark.asyncio
async def test_expansion_preserves_metadata(temp_db_path, small_chunk_config):
"""Expansion should preserve document metadata."""
docling_doc = create_table_document()
async with HaikuRAG(temp_db_path, config=small_chunk_config, create=True) as client:
doc = await create_document_with_docling(client, docling_doc, "Metadata Test")
assert doc.id is not None
results = await client.search("Introduction paragraph", limit=1)
assert len(results) > 0
expanded = await client.expand_context(results)
assert expanded[0].document_title == "Metadata Test"
assert expanded[0].chunk_id == results[0].chunk_id
assert expanded[0].document_id == results[0].document_id
@pytest.mark.asyncio
async def test_format_for_agent_output(temp_db_path, small_chunk_config):
"""format_for_agent should include source, type, and content sections."""
docling_doc = create_table_document()
async with HaikuRAG(temp_db_path, config=small_chunk_config, create=True) as client:
doc = await create_document_with_docling(client, docling_doc, "Format Test")
assert doc.id is not None
results = await client.search("Alice Smith employee data", limit=5)
table_results = [r for r in results if "table" in r.labels]
assert len(table_results) > 0
expanded = await client.expand_context(table_results[:1])
formatted = expanded[0].format_for_agent()
# Check format structure
assert "score:" in formatted
assert 'Source: "Format Test"' in formatted
assert "Type: table" in formatted
assert "Content:" in formatted
@pytest.mark.asyncio
async def test_max_items_limit_caps_expansion(temp_db_path):
"""Expansion should respect max_context_items limit."""
config = AppConfig()
config.processing.chunk_size = 32
config.processing.max_context_items = 2 # Very restrictive
docling_doc = create_list_document()
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
doc = await create_document_with_docling(client, docling_doc, "Limit Test")
assert doc.id is not None
results = await client.search("grapes", limit=1)
assert len(results) > 0
expanded = await client.expand_context(results)
# With max_items=2, expansion should be limited
content = expanded[0].content.lower()
item_count = sum(
1 for item in ["apples", "bananas", "oranges", "grapes"] if item in content
)
# Should have at most 2 items (the limit)
assert item_count <= 2, f"Expected at most 2 items, got {item_count}"
@pytest.mark.asyncio
async def test_search_result_get_primary_label():
"""Test _get_primary_label prioritizes structural labels correctly."""
# Table should be prioritized
result = SearchResult(
content="test",
score=0.5,
chunk_id="c1",
document_id="d1",
labels=["paragraph", "table", "text"],
)
assert result._get_primary_label() == "table"
# Code should be prioritized over paragraph
result = SearchResult(
content="test",
score=0.5,
chunk_id="c2",
document_id="d2",
labels=["paragraph", "code"],
)
assert result._get_primary_label() == "code"
# list_item should be prioritized
result = SearchResult(
content="test",
score=0.5,
chunk_id="c3",
document_id="d3",
labels=["text", "list_item"],
)
assert result._get_primary_label() == "list_item"
# Returns first label when no priority match
result = SearchResult(
content="test",
score=0.5,
chunk_id="c4",
document_id="d4",
labels=["paragraph", "text"],
)
assert result._get_primary_label() == "paragraph"
# Returns None for empty labels
result = SearchResult(
content="test",
score=0.5,
chunk_id="c5",
document_id="d5",
labels=[],
)
assert result._get_primary_label() is None
@pytest.mark.asyncio
async def test_expand_context_radius_zero(temp_db_path):
"""Test expand_context with radius 0 returns original results."""
# Default config has text_context_radius=0
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content="Simple test content")
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)]
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
@pytest.mark.asyncio
async def test_expand_context_multiple_documents(temp_db_path):
"""Test expand_context with results from multiple documents."""
config = AppConfig()
config.processing.text_context_radius = 1
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
# Create first document with manual chunks
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(
content="Doc1 content", chunks=doc1_chunks, uri="doc1.txt"
)
# Create second document with manual chunks
doc2_chunks = [
Chunk(content="Doc2 Section X", order=0),
Chunk(content="Doc2 Section Y", order=1),
]
doc2 = await client.import_document(
content="Doc2 content", 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)
search_results = [
SearchResult.from_chunk(chunk1, 0.8),
SearchResult.from_chunk(chunk2, 0.7),
]
expanded_results = 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.asyncio
async def test_expand_context_merges_overlapping_chunks(temp_db_path):
"""Test that overlapping expanded chunks are merged into one."""
config = AppConfig()
config.processing.text_context_radius = 1
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
# Create document with 5 chunks
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(
content="Full document content", 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.asyncio
async def test_expand_context_keeps_separate_non_overlapping(temp_db_path):
"""Test that non-overlapping expanded chunks remain separate."""
config = AppConfig()
config.processing.text_context_radius = 1
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
# Create document with chunks far apart
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(
content="Full document content", 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
@pytest.mark.asyncio
async def test_expand_context_with_docling_merges_overlapping(temp_db_path):
"""Test that expand_context with DoclingDocument merges overlapping results."""
config = AppConfig()
config.processing.text_context_radius = 3
markdown_content = """# Chapter 1
This is paragraph one about topic A.
This is paragraph two about topic A continued.
This is paragraph three about topic B.
# Chapter 2
This is paragraph four about topic C.
"""
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
doc = await client.create_document(
content=markdown_content,
uri="test://structured",
)
assert doc.id is not None
assert doc.docling_document_json is not None
# Get chunks which should have doc_item_refs
chunks = await client.chunk_repository.get_by_document_id(doc.id)
assert len(chunks) >= 1
# Find chunks that have doc_item_refs (from docling chunking)
chunks_with_refs = [c for c in chunks if c.get_chunk_metadata().doc_item_refs]
if len(chunks_with_refs) >= 2:
# Create search results from adjacent chunks
search_results = [
SearchResult.from_chunk(chunks_with_refs[0], 0.9),
SearchResult.from_chunk(chunks_with_refs[1], 0.8),
]
# Expand with configured radius that should cause overlap
expanded = await client.expand_context(search_results)
# If chunks were adjacent, they should be merged
# The expanded results should have merged metadata
assert len(expanded) >= 1
# Check that expanded result has page_numbers populated
for r in expanded:
# Should have doc_item_refs from expansion
assert r.doc_item_refs is not None
@pytest.mark.asyncio
async def test_expand_context_docling_merges_metadata(temp_db_path):
"""Test that expand_context properly merges metadata from multiple results."""
config = AppConfig()
config.processing.text_context_radius = 10
markdown_content = """# Introduction
First paragraph of introduction.
Second paragraph of introduction.
# Methods
First paragraph of methods section.
Second paragraph of methods section.
# Results
First paragraph of results.
"""
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
doc = await client.create_document(
content=markdown_content,
uri="test://metadata-merge",
)
assert doc.id is not None
chunks = await client.chunk_repository.get_by_document_id(doc.id)
chunks_with_refs = [c for c in chunks if c.get_chunk_metadata().doc_item_refs]
if len(chunks_with_refs) >= 2:
# Get chunks with different headings if possible
chunk1 = chunks_with_refs[0]
chunk2 = chunks_with_refs[-1] # Last chunk likely has different heading
search_results = [
SearchResult.from_chunk(chunk1, 0.9),
SearchResult.from_chunk(chunk2, 0.8),
]
# Expand with large radius to potentially merge
expanded = await client.expand_context(search_results)
# Check that results have proper structure
for r in expanded:
# Content should be non-empty
assert len(r.content) > 0
# Score should be preserved (best score)
assert r.score in [0.9, 0.8]
@pytest.mark.asyncio
async def test_expand_context_docling_preserves_bounding_boxes(temp_db_path):
"""Test that expand_context preserves bounding boxes from DoclingDocument."""
config = AppConfig()
config.processing.text_context_radius = 2
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
doc = await client.create_document(
content="# Test\n\nSome content here.",
uri="test://bboxes",
)
assert doc.id is not None
chunks = await client.chunk_repository.get_by_document_id(doc.id)
if chunks:
search_results = [SearchResult.from_chunk(chunks[0], 0.9)]
expanded = await client.expand_context(search_results)
# Expanded results should exist
assert len(expanded) == 1
# Bounding boxes may or may not be present depending on document
# but the field should be accessible
_ = expanded[0].bounding_boxes

28
tests/test_search.py
View file

@ -244,3 +244,31 @@ async def test_search_graceful_degradation(temp_db_path):
assert result.labels == []
client.close()
@pytest.mark.asyncio
async def test_search_result_format_includes_metadata(temp_db_path):
"""Test that formatted search results include document metadata."""
async with HaikuRAG(temp_db_path, create=True) as client:
await client.create_document(
content="Important information about machine learning algorithms.",
title="ML Guide",
uri="https://example.com/ml-guide",
)
results = await client.search("machine learning", limit=1)
assert len(results) > 0
formatted = results[0].format_for_agent()
# Should include chunk ID and score
assert "[" in formatted and "]" in formatted
assert "score:" in formatted
# Should include document title in Source
assert "ML Guide" in formatted
assert "Source:" in formatted
# Should include content
assert "Content:" in formatted
assert "machine learning" in formatted.lower()