From 9a5f8e43686d4547fba06d89a21c5959f8139b8b Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 8 Dec 2025 11:35:44 +0200 Subject: [PATCH] Add type-aware context expansion for search results --- CHANGELOG.md | 8 + docs/configuration/index.md | 4 +- docs/configuration/processing.md | 26 +- docs/inspector.md | 12 + docs/python.md | 19 +- evaluations/evaluations/benchmark.py | 4 +- haiku_rag_slim/haiku/rag/client.py | 165 +++- haiku_rag_slim/haiku/rag/config/models.py | 4 +- .../haiku/rag/graph/common/nodes.py | 5 +- .../haiku/rag/graph/common/prompts.py | 21 +- haiku_rag_slim/haiku/rag/inspector/app.py | 17 + .../rag/inspector/widgets/context_modal.py | 89 +++ .../rag/inspector/widgets/visual_modal.py | 1 + haiku_rag_slim/haiku/rag/qa/agent.py | 6 +- haiku_rag_slim/haiku/rag/qa/prompts.py | 21 +- .../haiku/rag/store/models/chunk.py | 63 ++ tests/test_chunk.py | 92 ++- tests/test_client.py | 371 +-------- tests/test_context_enhancement.py | 716 ++++++++++++++++++ tests/test_search.py | 28 + 20 files changed, 1259 insertions(+), 413 deletions(-) create mode 100644 haiku_rag_slim/haiku/rag/inspector/widgets/context_modal.py create mode 100644 tests/test_context_enhancement.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 283face8..d271b75c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -37,9 +37,17 @@ - Inspector modal with keyboard navigation between pages - CLI command: `haiku-rag visualize ` - 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 diff --git a/docs/configuration/index.md b/docs/configuration/index.md index 2d18adfd..01c6081e 100644 --- a/docs/configuration/index.md +++ b/docs/configuration/index.md @@ -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 diff --git a/docs/configuration/processing.md b/docs/configuration/processing.md index b3a8fbab..af298ed3 100644 --- a/docs/configuration/processing.md +++ b/docs/configuration/processing.md @@ -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: diff --git a/docs/inspector.md b/docs/inspector.md index f5df2637..406a36a2 100644 --- a/docs/inspector.md +++ b/docs/inspector.md @@ -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: diff --git a/docs/python.md b/docs/python.md index 51d99489..645813bc 100644 --- a/docs/python.md +++ b/docs/python.md @@ -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 diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index 754c4ef0..da7fbbc1 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -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, diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index ef28822b..463e3c57 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -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) diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 85417591..30c10943 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -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" diff --git a/haiku_rag_slim/haiku/rag/graph/common/nodes.py b/haiku_rag_slim/haiku/rag/graph/common/nodes.py index 5781b992..885036d0 100644 --- a/haiku_rag_slim/haiku/rag/graph/common/nodes.py +++ b/haiku_rag_slim/haiku/rag/graph/common/nodes.py @@ -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}" diff --git a/haiku_rag_slim/haiku/rag/graph/common/prompts.py b/haiku_rag_slim/haiku/rag/graph/common/prompts.py index 2a3a460c..2f17ad7f 100644 --- a/haiku_rag_slim/haiku/rag/graph/common/prompts.py +++ b/haiku_rag_slim/haiku/rag/graph/common/prompts.py @@ -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. diff --git a/haiku_rag_slim/haiku/rag/inspector/app.py b/haiku_rag_slim/haiku/rag/inspector/app.py index c6351d1d..dbdb62b2 100644 --- a/haiku_rag_slim/haiku/rag/inspector/app.py +++ b/haiku_rag_slim/haiku/rag/inspector/app.py @@ -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. diff --git a/haiku_rag_slim/haiku/rag/inspector/widgets/context_modal.py b/haiku_rag_slim/haiku/rag/inspector/widgets/context_modal.py new file mode 100644 index 00000000..b3340d82 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/inspector/widgets/context_modal.py @@ -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() diff --git a/haiku_rag_slim/haiku/rag/inspector/widgets/visual_modal.py b/haiku_rag_slim/haiku/rag/inspector/widgets/visual_modal.py index 50339eda..eb09ba76 100644 --- a/haiku_rag_slim/haiku/rag/inspector/widgets/visual_modal.py +++ b/haiku_rag_slim/haiku/rag/inspector/widgets/visual_modal.py @@ -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"), ] diff --git a/haiku_rag_slim/haiku/rag/qa/agent.py b/haiku_rag_slim/haiku/rag/qa/agent.py index 21bdb69e..c180e57f 100644 --- a/haiku_rag_slim/haiku/rag/qa/agent.py +++ b/haiku_rag_slim/haiku/rag/qa/agent.py @@ -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]]: diff --git a/haiku_rag_slim/haiku/rag/qa/prompts.py b/haiku_rag_slim/haiku/rag/qa/prompts.py index edc70cbc..e9aa5ac0 100644 --- a/haiku_rag_slim/haiku/rag/qa/prompts.py +++ b/haiku_rag_slim/haiku/rag/qa/prompts.py @@ -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 diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index 40c1c88d..f3104811 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/models/chunk.py @@ -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] diff --git a/tests/test_chunk.py b/tests/test_chunk.py index 3624d603..4c795440 100644 --- a/tests/test_chunk.py +++ b/tests/test_chunk.py @@ -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 diff --git a/tests/test_client.py b/tests/test_client.py index e7b9eadc..2a3ca4c3 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -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 diff --git a/tests/test_context_enhancement.py b/tests/test_context_enhancement.py new file mode 100644 index 00000000..c7b8d222 --- /dev/null +++ b/tests/test_context_enhancement.py @@ -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 diff --git a/tests/test_search.py b/tests/test_search.py index 841f77bb..bb023828 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -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()