Incorporate additional docling meta in search(), expand_context() and relevant qa/research agents

.pre-commit-config.yaml

@@ -9,7 +9,7 @@ repos:
       - id: debug-statements
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.14.3
+    rev: v0.14.5
     hooks:
       # Run the linter.
       - id: ruff

haiku_rag_slim/haiku/rag/app.py

@@ -3,6 +3,7 @@ import json
 import logging
 from importlib.metadata import version as pkg_version
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 from rich.console import Console
 from rich.markdown import Markdown
@@ -24,8 +25,10 @@ from haiku.rag.graph.research.graph import build_research_graph
 from haiku.rag.graph.research.state import ResearchDeps, ResearchState
 from haiku.rag.mcp import create_mcp_server
 from haiku.rag.monitor import FileWatcher
-from haiku.rag.store.models.chunk import Chunk
 from haiku.rag.store.models.document import Document
+
+if TYPE_CHECKING:
+    from haiku.rag.store.models import SearchResult
 from haiku.rag.utils import format_bytes
 
 logger = logging.getLogger(__name__)
@@ -268,8 +271,8 @@ class HaikuRAGApp:
             if not results:
                 self.console.print("[yellow]No results found.[/yellow]")
                 return
-            for chunk, score in results:
-                self._rich_print_search_result(chunk, score)
+            for result in results:
+                self._rich_print_search_result(result)
 
     async def ask(
         self,
@@ -606,22 +609,25 @@ class HaikuRAGApp:
         self.console.print(content)
         self.console.rule()
 
-    def _rich_print_search_result(self, chunk: Chunk, score: float):
-        """Format a search result chunk for display."""
-        content = Markdown(chunk.content)
+    def _rich_print_search_result(self, result: "SearchResult"):
+        """Format a search result for display."""
+        content = Markdown(result.content)
         self.console.print(
-            f"[repr.attrib_name]document_id[/repr.attrib_name]: {chunk.document_id} "
-            f"[repr.attrib_name]score[/repr.attrib_name]: {score:.4f}"
+            f"[repr.attrib_name]document_id[/repr.attrib_name]: {result.document_id} "
+            f"[repr.attrib_name]score[/repr.attrib_name]: {result.score:.4f}"
         )
-        if chunk.document_uri:
+        if result.document_uri:
             self.console.print("[repr.attrib_name]document uri[/repr.attrib_name]:")
-            self.console.print(chunk.document_uri)
-        if chunk.document_title:
+            self.console.print(result.document_uri)
+        if result.document_title:
             self.console.print("[repr.attrib_name]document title[/repr.attrib_name]:")
-            self.console.print(chunk.document_title)
-        if chunk.document_meta:
-            self.console.print("[repr.attrib_name]document meta[/repr.attrib_name]:")
-            self.console.print(chunk.document_meta)
+            self.console.print(result.document_title)
+        if result.page_numbers:
+            self.console.print("[repr.attrib_name]pages[/repr.attrib_name]:")
+            self.console.print(", ".join(str(p) for p in result.page_numbers))
+        if result.headings:
+            self.console.print("[repr.attrib_name]headings[/repr.attrib_name]:")
+            self.console.print(" > ".join(result.headings))
         self.console.print("[repr.attrib_name]content[/repr.attrib_name]:")
         self.console.print(content)
         self.console.rule()

haiku_rag_slim/haiku/rag/client.py

@@ -16,7 +16,7 @@ from haiku.rag.config import AppConfig, Config
 from haiku.rag.converters import get_converter
 from haiku.rag.reranking import get_reranker
 from haiku.rag.store.engine import Store
-from haiku.rag.store.models.chunk import Chunk
+from haiku.rag.store.models.chunk import Chunk, SearchResult
 from haiku.rag.store.models.document import Document
 from haiku.rag.store.repositories.chunk import ChunkRepository
 from haiku.rag.store.repositories.document import DocumentRepository
@@ -132,26 +132,31 @@ class HaikuRAG:
             # Use converter to convert text
             converter = get_converter(self._config)
             docling_document = await converter.convert_text(content)
-            docling_json = docling_document.model_dump_json()
-            docling_version = docling_document.version
+
+            document = Document(
+                content=content,
+                uri=uri,
+                title=title,
+                metadata=metadata or {},
+                docling_document_json=docling_document.model_dump_json(),
+                docling_version=docling_document.version,
+            )
+
+            return await self.document_repository._create_and_chunk(
+                document, docling_document, chunks
+            )
         else:
             # Chunks already provided, no conversion needed
-            docling_document = None
-            docling_json = None
-            docling_version = None
+            document = Document(
+                content=content,
+                uri=uri,
+                title=title,
+                metadata=metadata or {},
+            )
 
-        document = Document(
-            content=content,
-            uri=uri,
-            title=title,
-            metadata=metadata or {},
-            docling_document_json=docling_json,
-            docling_version=docling_version,
-        )
-
-        return await self.document_repository._create_and_chunk(
-            document, docling_document, chunks
-        )
+            return await self.document_repository._create_and_chunk(
+                document, None, chunks
+            )
 
     async def create_document_from_source(
         self, source: str | Path, title: str | None = None, metadata: dict | None = None
@@ -553,7 +558,8 @@ class HaikuRAG:
         limit: int = 5,
         search_type: str = "hybrid",
         filter: str | None = None,
-    ) -> list[tuple[Chunk, float]]:
+        resolve_bounding_boxes: bool = False,
+    ) -> list[SearchResult]:
         """Search for relevant chunks using the specified search method with optional reranking.
 
         Args:
@@ -561,159 +567,377 @@ class HaikuRAG:
             limit: Maximum number of results to return.
             search_type: Type of search - "vector", "fts", or "hybrid" (default).
             filter: Optional SQL WHERE clause to filter documents before searching chunks.
+            resolve_bounding_boxes: Whether to resolve bounding boxes from DoclingDocument.
 
         Returns:
-            List of (chunk, score) tuples ordered by relevance.
+            List of SearchResult objects ordered by relevance.
         """
-        # Get reranker if available
         reranker = get_reranker(config=self._config)
 
         if reranker is None:
-            # No reranking - return direct search results
-            return await self.chunk_repository.search(query, limit, search_type, filter)
+            chunk_results = await self.chunk_repository.search(
+                query, limit, search_type, filter
+            )
+        else:
+            search_limit = limit * 10
+            raw_results = await self.chunk_repository.search(
+                query, search_limit, search_type, filter
+            )
+            chunks = [chunk for chunk, _ in raw_results]
+            chunk_results = await reranker.rerank(query, chunks, top_n=limit)
 
-        # Get more initial results (10X) for reranking
-        search_limit = limit * 10
-        search_results = await self.chunk_repository.search(
-            query, search_limit, search_type, filter
-        )
+        bounding_boxes_map: dict[str, list] | None = None
+        if resolve_bounding_boxes:
+            bounding_boxes_map = {}
+            doc_cache: dict[str, Document | None] = {}
 
-        # Apply reranking
-        chunks = [chunk for chunk, _ in search_results]
-        reranked_results = await reranker.rerank(query, chunks, top_n=limit)
+            for chunk, _ in chunk_results:
+                if chunk.document_id and chunk.id:
+                    if chunk.document_id not in doc_cache:
+                        doc_cache[chunk.document_id] = await self.get_document_by_id(
+                            chunk.document_id
+                        )
 
-        # Return reranked results with scores from reranker
-        return reranked_results
+                    doc = doc_cache[chunk.document_id]
+                    if doc:
+                        docling_doc = doc.get_docling_document()
+                        if docling_doc:
+                            meta = chunk.get_chunk_metadata()
+                            bounding_boxes_map[chunk.id] = meta.resolve_bounding_boxes(
+                                docling_doc
+                            )
+
+        results = []
+        for chunk, score in chunk_results:
+            bboxes = None
+            if bounding_boxes_map and chunk.id:
+                bboxes = bounding_boxes_map.get(chunk.id)
+            results.append(SearchResult.from_chunk(chunk, score, bboxes))
+        return results
 
     async def expand_context(
         self,
-        search_results: list[tuple[Chunk, float]],
+        search_results: list[SearchResult],
         radius: int | None = None,
-    ) -> list[tuple[Chunk, float]]:
-        """Expand search results with adjacent chunks, merging overlapping chunks.
+    ) -> list[SearchResult]:
+        """Expand search results with adjacent content from the source document.
+
+        When DoclingDocument is available and results have doc_item_refs, expands
+        by finding adjacent DocItems with accurate bounding boxes and metadata.
+        Otherwise, falls back to chunk-based expansion using adjacent chunks.
 
         Args:
-            search_results: List of (chunk, score) tuples from search.
-            radius: Number of adjacent chunks to include before/after each chunk.
+            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 (chunk, score) tuples with expanded and merged context chunks.
+            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
 
-        # Group chunks by document_id to handle merging within documents
-        document_groups = {}
-        for chunk, score in search_results:
-            doc_id = chunk.document_id
+        # Group by document_id for efficient processing
+        document_groups: dict[str | None, list[SearchResult]] = {}
+        for result in search_results:
+            doc_id = result.document_id
             if doc_id not in document_groups:
                 document_groups[doc_id] = []
-            document_groups[doc_id].append((chunk, score))
+            document_groups[doc_id].append(result)
 
-        results = []
+        expanded_results = []
 
-        for doc_id, doc_chunks in document_groups.items():
-            # Get all expanded ranges for this document
-            expanded_ranges = []
-            for chunk, score in doc_chunks:
-                adjacent_chunks = await self.chunk_repository.get_adjacent_chunks(
-                    chunk, radius
+        for doc_id, doc_results in document_groups.items():
+            if doc_id is None:
+                expanded_results.extend(doc_results)
+                continue
+
+            # Fetch the document to get DoclingDocument
+            doc = await self.get_document_by_id(doc_id)
+            if doc is None:
+                expanded_results.extend(doc_results)
+                continue
+
+            docling_doc = doc.get_docling_document()
+
+            # Check if we can use DoclingDocument-based expansion
+            has_docling = docling_doc is not None
+            has_refs = any(r.doc_item_refs for r in doc_results)
+
+            if has_docling and has_refs:
+                # Use DoclingDocument-based expansion
+                expanded = await self._expand_with_docling(
+                    doc_results, docling_doc, radius
                 )
+                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)
 
-                all_chunks = adjacent_chunks + [chunk]
+        return expanded_results
 
-                # Get the range of orders for this expanded chunk
-                orders = [c.order for c in all_chunks]
-                min_order = min(orders)
-                max_order = max(orders)
+    async def _expand_with_docling(
+        self,
+        results: list[SearchResult],
+        docling_doc,
+        radius: int,
+    ) -> list[SearchResult]:
+        """Expand results using DoclingDocument structure."""
+        from haiku.rag.store.models.chunk import BoundingBox
 
+        # Build index of all DocItems for expansion
+        all_items = list(docling_doc.iterate_items())
+        ref_to_index: dict[str, int] = {}
+
+        # Map refs to indices
+        for i, (_, item) in enumerate(all_items):
+            self_ref = getattr(item, "self_ref", None)
+            if self_ref:
+                ref_to_index[self_ref] = i
+
+        expanded_results = []
+
+        for result in results:
+            if not result.doc_item_refs:
+                expanded_results.append(result)
+                continue
+
+            # Find indices of all refs in this result
+            indices = []
+            for ref in result.doc_item_refs:
+                if ref in ref_to_index:
+                    indices.append(ref_to_index[ref])
+
+            if not indices:
+                expanded_results.append(result)
+                continue
+
+            # Expand range
+            min_idx = max(0, min(indices) - radius)
+            max_idx = min(len(all_items) - 1, max(indices) + radius)
+
+            # Collect expanded DocItems
+            expanded_content_parts = []
+            expanded_refs = []
+            expanded_page_numbers: set[int] = set()
+            expanded_labels: set[str] = set()
+            expanded_bboxes: list[BoundingBox] = []
+
+            for i in range(min_idx, max_idx + 1):
+                _, item = all_items[i]
+
+                # Get content
+                text = getattr(item, "text", None)
+                if text:
+                    expanded_content_parts.append(text)
+
+                # Get self_ref
+                self_ref = getattr(item, "self_ref", None)
+                if self_ref:
+                    expanded_refs.append(self_ref)
+
+                # Get label
+                label = getattr(item, "label", None)
+                if label:
+                    expanded_labels.add(
+                        str(label.value) if hasattr(label, "value") else str(label)
+                    )
+
+                # Get provenance (page numbers and bounding boxes)
+                prov = getattr(item, "prov", None)
+                if prov:
+                    for prov_item in prov:
+                        page_no = getattr(prov_item, "page_no", None)
+                        if page_no is not None:
+                            expanded_page_numbers.add(page_no)
+
+                        bbox = getattr(prov_item, "bbox", None)
+                        if bbox is not None:
+                            expanded_bboxes.append(
+                                BoundingBox(
+                                    page_no=page_no or 0,
+                                    left=bbox.l,
+                                    top=bbox.t,
+                                    right=bbox.r,
+                                    bottom=bbox.b,
+                                )
+                            )
+
+            expanded_results.append(
+                SearchResult(
+                    content="\n\n".join(expanded_content_parts),
+                    score=result.score,
+                    chunk_id=result.chunk_id,
+                    document_id=result.document_id,
+                    document_uri=result.document_uri,
+                    document_title=result.document_title,
+                    doc_item_refs=expanded_refs,
+                    page_numbers=sorted(expanded_page_numbers),
+                    headings=result.headings,
+                    labels=sorted(expanded_labels),
+                    bounding_boxes=expanded_bboxes if expanded_bboxes else None,
+                )
+            )
+
+        return expanded_results
+
+    async def _expand_with_chunks(
+        self,
+        doc_id: str,
+        results: list[SearchResult],
+        radius: int,
+    ) -> list[SearchResult]:
+        """Expand results using chunk-based adjacency."""
+        # Fetch all chunks for this document
+        all_chunks = await self.chunk_repository.get_by_document_id(doc_id)
+        if not all_chunks:
+            return results
+
+        # Build content -> chunk mapping and order -> chunk mapping
+        content_to_chunk = {c.content: c for c in all_chunks}
+        chunk_by_order = {c.order: c for c in all_chunks}
+        max_order = max(chunk_by_order.keys())
+        min_order = min(chunk_by_order.keys())
+
+        # Build expanded ranges for merging
+        expanded_ranges = []
+
+        for result in results:
+            # Find matching chunk by content
+            matching_chunk = content_to_chunk.get(result.content)
+            if matching_chunk is None:
                 expanded_ranges.append(
                     {
-                        "original_chunk": chunk,
-                        "score": score,
-                        "min_order": min_order,
-                        "max_order": max_order,
-                        "all_chunks": sorted(all_chunks, key=lambda c: c.order),
+                        "original_result": result,
+                        "score": result.score,
+                        "min_order": -1,
+                        "max_order": -1,
+                        "chunks": [],
                     }
                 )
+                continue
 
-            # Merge overlapping/adjacent ranges
-            merged_ranges = self._merge_overlapping_ranges(expanded_ranges)
+            # Calculate range
+            start_order = max(min_order, matching_chunk.order - radius)
+            end_order = min(max_order, matching_chunk.order + radius)
 
-            # Create merged chunks
-            for merged_range in merged_ranges:
-                combined_content_parts = [c.content for c in merged_range["all_chunks"]]
+            range_chunks = [
+                chunk_by_order[o]
+                for o in range(start_order, end_order + 1)
+                if o in chunk_by_order
+            ]
 
-                # Use the first original chunk for metadata
-                original_chunk = merged_range["original_chunks"][0]
+            expanded_ranges.append(
+                {
+                    "original_result": result,
+                    "score": result.score,
+                    "min_order": start_order,
+                    "max_order": end_order,
+                    "chunks": range_chunks,
+                }
+            )
 
-                merged_chunk = Chunk(
-                    id=original_chunk.id,
-                    document_id=original_chunk.document_id,
-                    content="".join(combined_content_parts),
-                    metadata=original_chunk.metadata,
-                    document_uri=original_chunk.document_uri,
-                    document_title=original_chunk.document_title,
-                    document_meta=original_chunk.document_meta,
+        # Merge overlapping ranges
+        merged_ranges = self._merge_chunk_ranges(expanded_ranges)
+
+        # Convert to SearchResults
+        expanded_results = []
+        for merged in merged_ranges:
+            if not merged["chunks"]:
+                # No chunks found, return original
+                expanded_results.append(merged["original_results"][0])
+                continue
+
+            combined_content = "".join(c.content for c in merged["chunks"])
+            original = merged["original_results"][0]
+            best_score = max(merged["scores"])
+
+            expanded_results.append(
+                SearchResult(
+                    content=combined_content,
+                    score=best_score,
+                    chunk_id=original.chunk_id,
+                    document_id=original.document_id,
+                    document_uri=original.document_uri,
+                    document_title=original.document_title,
+                    doc_item_refs=original.doc_item_refs,
+                    page_numbers=original.page_numbers,
+                    headings=original.headings,
+                    labels=original.labels,
+                    bounding_boxes=original.bounding_boxes,
                 )
+            )
 
-                # Use the highest score from merged chunks
-                best_score = max(merged_range["scores"])
-                results.append((merged_chunk, best_score))
+        return expanded_results
 
-        return results
+    def _merge_chunk_ranges(self, expanded_ranges: list[dict]) -> list[dict]:
+        """Merge overlapping or adjacent chunk ranges."""
+        # Filter out ranges without chunks
+        valid_ranges = [r for r in expanded_ranges if r["chunks"]]
+        invalid_ranges = [r for r in expanded_ranges if not r["chunks"]]
 
-    def _merge_overlapping_ranges(self, expanded_ranges):
-        """Merge overlapping or adjacent expanded ranges."""
-        if not expanded_ranges:
-            return []
+        if not valid_ranges:
+            return [
+                {
+                    "original_results": [r["original_result"]],
+                    "scores": [r["score"]],
+                    "chunks": [],
+                }
+                for r in invalid_ranges
+            ]
 
         # Sort by min_order
-        sorted_ranges = sorted(expanded_ranges, key=lambda x: x["min_order"])
+        sorted_ranges = sorted(valid_ranges, key=lambda x: x["min_order"])
         merged = []
 
         current = {
             "min_order": sorted_ranges[0]["min_order"],
             "max_order": sorted_ranges[0]["max_order"],
-            "original_chunks": [sorted_ranges[0]["original_chunk"]],
+            "original_results": [sorted_ranges[0]["original_result"]],
             "scores": [sorted_ranges[0]["score"]],
-            "all_chunks": sorted_ranges[0]["all_chunks"],
+            "chunks": sorted_ranges[0]["chunks"],
         }
 
         for range_info in sorted_ranges[1:]:
-            # Check if ranges overlap or are adjacent (max_order + 1 >= min_order)
+            # Check if ranges overlap or are adjacent
             if current["max_order"] >= range_info["min_order"] - 1:
                 # Merge ranges
                 current["max_order"] = max(
                     current["max_order"], range_info["max_order"]
                 )
-                current["original_chunks"].append(range_info["original_chunk"])
+                current["original_results"].append(range_info["original_result"])
                 current["scores"].append(range_info["score"])
 
-                # Merge all_chunks and deduplicate by order
-                all_chunks_dict = {}
-                for chunk in current["all_chunks"] + range_info["all_chunks"]:
-                    order = chunk.order
-                    all_chunks_dict[order] = chunk
-                current["all_chunks"] = [
-                    all_chunks_dict[order] for order in sorted(all_chunks_dict.keys())
-                ]
+                # Merge chunks and deduplicate by order
+                chunks_dict = {c.order: c for c in current["chunks"]}
+                for chunk in range_info["chunks"]:
+                    chunks_dict[chunk.order] = chunk
+                current["chunks"] = [chunks_dict[o] for o in sorted(chunks_dict.keys())]
             else:
-                # No overlap, add current to merged and start new
                 merged.append(current)
                 current = {
                     "min_order": range_info["min_order"],
                     "max_order": range_info["max_order"],
-                    "original_chunks": [range_info["original_chunk"]],
+                    "original_results": [range_info["original_result"]],
                     "scores": [range_info["score"]],
-                    "all_chunks": range_info["all_chunks"],
+                    "chunks": range_info["chunks"],
                 }
 
-        # Add the last range
         merged.append(current)
+
+        # Add back invalid ranges
+        for r in invalid_ranges:
+            merged.append(
+                {
+                    "original_results": [r["original_result"]],
+                    "scores": [r["score"]],
+                    "chunks": [],
+                }
+            )
+
         return merged
 
     async def ask(

haiku_rag_slim/haiku/rag/graph/common/nodes.py

@@ -103,8 +103,8 @@ def create_plan_node[AgentDepsT: GraphAgentDeps](
                 ctx2: RunContext[AgentDepsT], query: str, limit: int = 6
             ) -> str:
                 results = await ctx2.deps.client.search(query, limit=limit)
-                expanded = await ctx2.deps.client.expand_context(results)
-                return "\n\n".join(chunk.content for chunk, _ in expanded)
+                results = await ctx2.deps.client.expand_context(results)
+                return "\n\n".join(r.content for r in results)
 
             # Tool is registered via decorator above
             _ = gather_context
@@ -228,17 +228,24 @@ async def _do_search[AgentDepsT: GraphAgentDeps](
     async def search_and_answer(
         ctx2: RunContext[AgentDepsT], query: str, limit: int = 5
     ) -> str:
-        search_results = await ctx2.deps.client.search(query, limit=limit)
-        expanded = await ctx2.deps.client.expand_context(search_results)
+        results = await ctx2.deps.client.search(query, limit=limit)
+        results = await ctx2.deps.client.expand_context(results)
 
-        entries: list[dict[str, Any]] = [
-            {
-                "text": chunk.content,
-                "score": score,
-                "document_uri": (chunk.document_title or chunk.document_uri or ""),
+        entries: list[dict[str, Any]] = []
+        for r in results:
+            entry: dict[str, Any] = {
+                "text": r.content,
+                "score": r.score,
+                "document_uri": (r.document_uri or ""),
             }
-            for chunk, score in expanded
-        ]
+            if r.document_title:
+                entry["document_title"] = r.document_title
+            if r.page_numbers:
+                entry["page_numbers"] = r.page_numbers
+            if r.headings:
+                entry["headings"] = r.headings
+            entries.append(entry)
+
         if not entries:
             return f"No relevant information found in the knowledge base for: {query}"
 

haiku_rag_slim/haiku/rag/graph/common/prompts.py

@@ -27,15 +27,20 @@ Tasks:
 
 Tool usage:
 - Always call search_and_answer before drafting any answer.
-- The tool returns snippets with verbatim `text`, a relevance `score`, and the
-  originating document identifier (document title if available, otherwise URI).
+- The tool returns snippets with:
+  - `text`: verbatim content
+  - `score`: relevance score
+  - `document_uri`: full path to the source document
+  - `document_title`: title if available
+  - `page_numbers`: list of page numbers where content appears (if available)
+  - `headings`: section heading hierarchy (if available)
 - You may call the tool multiple times to refine or broaden context, but do not
   exceed 3 total calls. Favor precision over volume.
 - Use scores to prioritize evidence, but include only the minimal subset of
   snippet texts (verbatim) in SearchAnswer.context (typically 1-4).
-- Set SearchAnswer.sources to the corresponding document identifiers for the
-  snippets you used (title if available, otherwise URI; one per snippet; same
-  order as context). Context must be text-only.
+- Set SearchAnswer.sources to include document_uri, page numbers, and headings
+  for each snippet used. Format: "document_uri (p. X, Section: Y)" or just
+  "document_uri" if no page/heading info. One source per context snippet.
 - If no relevant information is found, clearly say so and return an empty
   context list and sources list.
 

haiku_rag_slim/haiku/rag/graph/research/prompts.py

@@ -87,7 +87,11 @@ Report guidelines (map to output fields):
 - conclusions: 2–4 bullets that follow logically from findings.
 - recommendations: 2–5 actionable bullets tied to findings.
 - limitations: 1–3 bullets describing key constraints or uncertainties.
-- sources_summary: 2–4 sentences summarizing sources used and their reliability.
+- sources_summary: List specific sources used with document paths, page numbers,
+  and section headings where available. Format each as:
+  "- /path/to/document.pdf (p. 5, Section: Introduction)" or
+  "- /path/to/file.md (Section: Getting Started)"
+  Include one bullet per distinct source document.
 
 Style:
 - Base all content solely on the collected evidence.

haiku_rag_slim/haiku/rag/mcp.py

@@ -7,12 +7,7 @@ from pydantic import BaseModel
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import AppConfig, Config
 from haiku.rag.graph.research.models import ResearchReport
-
-
-class SearchResult(BaseModel):
-    document_id: str
-    content: str
-    score: float
+from haiku.rag.store.models import SearchResult
 
 
 class DocumentResult(BaseModel):
@@ -87,22 +82,7 @@ def create_mcp_server(db_path: Path, config: AppConfig = Config) -> FastMCP:
         """Search the RAG system for documents using hybrid search (vector similarity + full-text search)."""
         try:
             async with HaikuRAG(db_path, config=config) as rag:
-                results = await rag.search(query, limit)
-
-                search_results = []
-                for chunk, score in results:
-                    assert chunk.document_id is not None, (
-                        "Chunk document_id should not be None in search results"
-                    )
-                    search_results.append(
-                        SearchResult(
-                            document_id=chunk.document_id,
-                            content=chunk.content,
-                            score=score,
-                        )
-                    )
-
-                return search_results
+                return await rag.search(query, limit)
         except Exception:
             return []
 

haiku_rag_slim/haiku/rag/qa/agent.py

@@ -8,11 +8,20 @@ from haiku.rag.graph.common import get_model
 from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT, QA_SYSTEM_PROMPT_WITH_CITATIONS
 
 
-class SearchResult(BaseModel):
+class ToolSearchResult(BaseModel):
+    """Search result model exposed to the LLM tool."""
+
     content: str = Field(description="The document text content")
     score: float = Field(description="Relevance score (higher is more relevant)")
-    document_uri: str = Field(
-        description="Source title (if available) or URI/path of the document"
+    document_uri: str = Field(description="The URI/path of the source document")
+    document_title: str | None = Field(
+        default=None, description="The title of the document (if available)"
+    )
+    page_numbers: list[int] = Field(
+        default=[], description="Page numbers where this content appears"
+    )
+    headings: list[str] | None = Field(
+        default=None, description="Section heading hierarchy for this content"
     )
 
 
@@ -50,18 +59,21 @@ class QuestionAnswerAgent:
             ctx: RunContext[Dependencies],
             query: str,
             limit: int = 5,
-        ) -> list[SearchResult]:
+        ) -> list[ToolSearchResult]:
             """Search the knowledge base for relevant documents."""
-            search_results = await ctx.deps.client.search(query, limit=limit)
-            expanded_results = await ctx.deps.client.expand_context(search_results)
+            results = await ctx.deps.client.search(query, limit=limit)
+            results = await ctx.deps.client.expand_context(results)
 
             return [
-                SearchResult(
-                    content=chunk.content,
-                    score=score,
-                    document_uri=(chunk.document_title or chunk.document_uri or ""),
+                ToolSearchResult(
+                    content=r.content,
+                    score=r.score,
+                    document_uri=(r.document_uri or ""),
+                    document_title=r.document_title,
+                    page_numbers=r.page_numbers,
+                    headings=r.headings,
                 )
-                for chunk, score in expanded_results
+                for r in results
             ]
 
     async def answer(self, question: str) -> str:

haiku_rag_slim/haiku/rag/qa/prompts.py

@@ -44,17 +44,22 @@ Guidelines:
 
 Citation Format:
 After your answer, include a "Citations:" section that lists:
-- The document title (if available) or URI from each search result used
-- A brief excerpt (first 50-100 characters) of the content that supported your answer
-- Format: "Citations:\n- [document title or URI]: [content_excerpt]..."
+- The document URI (from the document_uri field) - always include the full path
+- The document title if available (from the document_title field)
+- Page number(s) if available (from the page_numbers field)
+- Section heading if available (from the headings field)
+- A VERBATIM excerpt (copy-paste exact text, first 50-100 characters) from the content field - do NOT summarize or paraphrase
 
 Example response format:
 [Your answer here]
 
 Citations:
-- /path/to/document1.pdf: "This document explains that AFMAN stands for Air Force Manual..."
-- /path/to/document2.pdf: "The manual provides guidance on military procedures and..."
+
+- **/docs/user-guide.pdf** - "User Manual" (p. 5, Section: Introduction)
+  The system requires Python 3.10 or higher to run properly...
+
+- **/reports/quarterly-analysis.md** (pp. 12-13, Section: Results)
+  Revenue increased by 15% compared to the previous quarter...
 
 Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents.
-/no_think
 """

haiku_rag_slim/haiku/rag/store/models/__init__.py

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

haiku_rag_slim/haiku/rag/store/models/chunk.py

@@ -6,6 +6,16 @@ if TYPE_CHECKING:
     from docling_core.types.doc.document import DocItem, DoclingDocument
 
 
+class BoundingBox(BaseModel):
+    """Bounding box coordinates for visual grounding."""
+
+    page_no: int
+    left: float
+    top: float
+    right: float
+    bottom: float
+
+
 class ChunkMetadata(BaseModel):
     """
     Structured metadata for a chunk, including DoclingDocument references.
@@ -46,6 +56,37 @@ class ChunkMetadata(BaseModel):
                 continue
         return doc_items
 
+    def resolve_bounding_boxes(
+        self, docling_document: "DoclingDocument"
+    ) -> list[BoundingBox]:
+        """Resolve doc_item_refs to bounding boxes for visual grounding.
+
+        Args:
+            docling_document: The parent DoclingDocument containing the items.
+
+        Returns:
+            List of BoundingBox objects from resolved DocItems' provenance.
+        """
+        bounding_boxes = []
+        for doc_item in self.resolve_doc_items(docling_document):
+            prov = getattr(doc_item, "prov", None)
+            if not prov:
+                continue
+            for prov_item in prov:
+                bbox = getattr(prov_item, "bbox", None)
+                if bbox is None:
+                    continue
+                bounding_boxes.append(
+                    BoundingBox(
+                        page_no=prov_item.page_no,
+                        left=bbox.l,
+                        top=bbox.t,
+                        right=bbox.r,
+                        bottom=bbox.b,
+                    )
+                )
+        return bounding_boxes
+
 
 class Chunk(BaseModel):
     """
@@ -65,3 +106,42 @@ class Chunk(BaseModel):
     def get_chunk_metadata(self) -> ChunkMetadata:
         """Parse metadata dict into structured ChunkMetadata."""
         return ChunkMetadata.model_validate(self.metadata)
+
+
+class SearchResult(BaseModel):
+    """Search result with optional provenance information for citations."""
+
+    content: str
+    score: float
+    chunk_id: str | None = None
+    document_id: str | None = None
+    document_uri: str | None = None
+    document_title: str | None = None
+    doc_item_refs: list[str] = []
+    page_numbers: list[int] = []
+    headings: list[str] | None = None
+    labels: list[str] = []
+    bounding_boxes: list[BoundingBox] | None = None
+
+    @classmethod
+    def from_chunk(
+        cls,
+        chunk: "Chunk",
+        score: float,
+        bounding_boxes: list[BoundingBox] | None = None,
+    ) -> "SearchResult":
+        """Create from a Chunk with optional bounding boxes."""
+        meta = chunk.get_chunk_metadata()
+        return cls(
+            content=chunk.content,
+            score=score,
+            chunk_id=chunk.id,
+            document_id=chunk.document_id,
+            document_uri=chunk.document_uri,
+            document_title=chunk.document_title,
+            doc_item_refs=meta.doc_item_refs,
+            page_numbers=meta.page_numbers,
+            headings=meta.headings,
+            labels=meta.labels,
+            bounding_boxes=bounding_boxes,
+        )

haiku_rag_slim/haiku/rag/store/upgrades/__init__.py

@@ -6,16 +6,6 @@ from packaging.version import Version, parse
 
 from haiku.rag.store.engine import Store
 
-from .v0_9_3 import upgrade_fts_phrase as upgrade_0_9_3_fts  # noqa: E402
-from .v0_9_3 import upgrade_order as upgrade_0_9_3_order  # noqa: E402
-from .v0_10_1 import upgrade_add_title as upgrade_0_10_1_add_title  # noqa: E402
-from .v0_19_6 import (  # noqa: E402
-    upgrade_embeddings_model_config as upgrade_0_19_6_embeddings,
-)
-from .v0_20_0 import (
-    upgrade_add_docling_document as upgrade_0_20_0_docling,  # noqa: E402
-)
-
 logger = logging.getLogger(__name__)
 
 
@@ -63,6 +53,20 @@ def run_pending_upgrades(store: Store, from_version: str, to_version: str) -> No
         logger.info("Completed upgrade %s", step.version)
 
 
+# Import upgrade modules AFTER Upgrade class is defined to avoid circular imports
+# ruff: noqa: E402, I001
+from haiku.rag.store.upgrades.v0_9_3 import upgrade_fts_phrase as upgrade_0_9_3_fts
+from haiku.rag.store.upgrades.v0_9_3 import upgrade_order as upgrade_0_9_3_order
+from haiku.rag.store.upgrades.v0_10_1 import (
+    upgrade_add_title as upgrade_0_10_1_add_title,
+)
+from haiku.rag.store.upgrades.v0_19_6 import (
+    upgrade_embeddings_model_config as upgrade_0_19_6_embeddings,
+)
+from haiku.rag.store.upgrades.v0_20_0 import (
+    upgrade_add_docling_document as upgrade_0_20_0_docling,
+)
+
 upgrades.append(upgrade_0_9_3_order)
 upgrades.append(upgrade_0_9_3_fts)
 upgrades.append(upgrade_0_10_1_add_title)

tests/test_client.py

@@ -641,19 +641,19 @@ async def test_client_search(temp_db_path):
         results = await client.search("Python programming", limit=3)
 
         assert len(results) > 0
-        assert all(len(result) == 2 for result in results)
-
-        # Verify first result is from the Python document (doc1)
-        first_chunk, _ = results[0]
-        assert first_chunk.document_id == doc1.id
+        # Verify results are SearchResult objects with expected fields
+        first_result = results[0]
+        assert first_result.content
+        assert first_result.score >= 0
+        assert first_result.document_id == doc1.id
 
         # Test search with different query
         ml_results = await client.search("machine learning data", limit=2)
         assert len(ml_results) > 0
 
         # Verify first result is from the machine learning document (doc2)
-        first_ml_chunk, _ = ml_results[0]
-        assert first_ml_chunk.document_id == doc2.id
+        first_ml_result = ml_results[0]
+        assert first_ml_result.document_id == doc2.id
 
         # Test search with limit parameter
         limited_results = await client.search("programming", limit=1)
@@ -782,6 +782,8 @@ async def test_client_ask_with_cite(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:
@@ -808,50 +810,55 @@ async def test_client_expand_context(temp_db_path):
         chunks = await client.chunk_repository.get_by_document_id(doc.id)
         assert len(chunks) == 5
 
-        # Find the middle chunk (order=2)
+        # Find the middle chunk (order=2) and convert to SearchResult
         middle_chunk = next(c for c in chunks if c.order == 2)
-        search_results = [(middle_chunk, 0.8)]
+        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_chunk, score = expanded_results[0]
+        expanded = expanded_results[0]
 
-        # Check that the expanded chunk has combined content and preserves title/uri
-        assert expanded_chunk.id == middle_chunk.id
-        assert score == 0.8
-        assert "Chunk 2 content" in expanded_chunk.content
-        assert expanded_chunk.document_title == "test_doc_title"
-        assert expanded_chunk.document_uri == "test_doc.txt"
+        # 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_chunk.content
-        assert "Chunk 1 content" in expanded_chunk.content
-        assert "Chunk 2 content" in expanded_chunk.content
-        assert "Chunk 3 content" in expanded_chunk.content
-        assert "Chunk 4 content" in expanded_chunk.content
+        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."""
+    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 = [(chunks[0], 0.9)]
+        search_results = [SearchResult.from_chunk(chunks[0], 0.9)]
         expanded_results = await client.expand_context(search_results, radius=0)
 
         # Should return exactly the same results
-        assert expanded_results == search_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
@@ -882,23 +889,24 @@ async def test_client_expand_context_multiple_chunks(temp_db_path):
         chunk1 = next(c for c in chunks1 if c.order == 1)
         chunk2 = next(c for c in chunks2 if c.order == 0)
 
-        search_results = [(chunk1, 0.8), (chunk2, 0.7)]
+        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, score1 = expanded_results[0]
-        assert expanded1.id == chunk1.id
-        assert score1 == 0.8
+        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, score2 = expanded_results[1]
-        assert expanded2.id == chunk2.id
-        assert score2 == 0.7
+        expanded2 = expanded_results[1]
+        assert expanded2.score == 0.7
         assert "Doc2 Section X" in expanded2.content
         assert "Doc2 Section Y" in expanded2.content
 
@@ -906,6 +914,8 @@ async def test_client_expand_context_multiple_chunks(temp_db_path):
 @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 = [
@@ -931,28 +941,33 @@ async def test_client_expand_context_merges_overlapping_chunks(temp_db_path):
         # 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 = [(chunk1, 0.8), (chunk2, 0.7)]
+        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_chunk, score = expanded_results[0]
+        merged = expanded_results[0]
 
         # Should contain all chunks from 0 to 3
-        assert "Chunk 0" in merged_chunk.content
-        assert "Chunk 1" in merged_chunk.content
-        assert "Chunk 2" in merged_chunk.content
-        assert "Chunk 3" in merged_chunk.content
-        assert "Chunk 4" not in merged_chunk.content  # Should not include chunk 4
+        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 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 = [
@@ -980,17 +995,20 @@ async def test_client_expand_context_keeps_separate_non_overlapping(temp_db_path
 
         # chunk0 expanded: [0,1] with radius=1 (orders 0,1)
         # chunk5 expanded: [4,5] with radius=1 (orders 4,5)
-        search_results = [(chunk0, 0.8), (chunk5, 0.7)]
+        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[1], reverse=True)
+        expanded_results.sort(key=lambda x: x.score, reverse=True)
 
-        chunk0_expanded, score1 = expanded_results[0]
-        chunk5_expanded, score2 = expanded_results[1]
+        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"
@@ -999,14 +1017,14 @@ async def test_client_expand_context_keeps_separate_non_overlapping(temp_db_path
         assert (
             "Chunk 5" not in chunk0_expanded.content
         )  # Should not have chunk 7 content
-        assert score1 == 0.8
+        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 score2 == 0.7
+        assert chunk5_expanded.score == 0.7
 
 
 @pytest.mark.asyncio

tests/test_filter.py

@@ -24,17 +24,17 @@ async def test_search_with_uri_filter(temp_db_path):
             "tutorial", limit=5, filter="uri LIKE '%example.com%'"
         )
         assert len(results) > 0
-        for chunk, _ in results:
-            assert chunk.document_uri is not None
-            assert "example.com" in chunk.document_uri
+        for result in results:
+            assert result.document_uri is not None
+            assert "example.com" in result.document_uri
 
         # Filter by exact URI
         results = await client.search(
             "tutorial", limit=5, filter="uri = 'https://other.com/java.html'"
         )
         assert len(results) > 0
-        for chunk, _ in results:
-            assert chunk.document_uri == "https://other.com/java.html"
+        for result in results:
+            assert result.document_uri == "https://other.com/java.html"
 
 
 @pytest.mark.asyncio
@@ -58,9 +58,9 @@ async def test_search_with_title_filter(temp_db_path):
             "programming", limit=5, filter="title LIKE '%Python%'"
         )
         assert len(results) > 0
-        for chunk, _ in results:
-            assert chunk.document_title is not None
-            assert "Python" in chunk.document_title
+        for result in results:
+            assert result.document_title is not None
+            assert "Python" in result.document_title
 
 
 @pytest.mark.asyncio
@@ -89,11 +89,11 @@ async def test_search_with_combined_filters(temp_db_path):
             "AI", limit=5, filter="uri LIKE '%arxiv%' AND title LIKE '%Machine%'"
         )
         assert len(results) > 0
-        for chunk, _ in results:
-            assert chunk.document_uri is not None
-            assert chunk.document_title is not None
-            assert "arxiv" in chunk.document_uri
-            assert "Machine" in chunk.document_title
+        for result in results:
+            assert result.document_uri is not None
+            assert result.document_title is not None
+            assert "arxiv" in result.document_uri
+            assert "Machine" in result.document_title
 
         # Filter with OR condition (if supported)
         results = await client.search(
@@ -159,15 +159,15 @@ async def test_search_filter_with_all_search_types(temp_db_path):
             filter="uri LIKE '%ai.example%'",
         )
         assert len(results) > 0
-        for chunk, _ in results:
-            assert chunk.document_uri is not None
-            assert "ai.example" in chunk.document_uri
+        for result in results:
+            assert result.document_uri is not None
+            assert "ai.example" in result.document_uri
 
         # Test FTS search with filter
         results = await client.search(
             "learning", limit=5, search_type="fts", filter="title = 'ML Guide'"
         )
-        assert all(chunk.document_title == "ML Guide" for chunk, _ in results)
+        assert all(r.document_title == "ML Guide" for r in results)
 
         # Test hybrid search with filter (default)
         results = await client.search(
@@ -176,6 +176,6 @@ async def test_search_filter_with_all_search_types(temp_db_path):
             search_type="hybrid",
             filter="uri LIKE '%other.com%'",
         )
-        for chunk, _ in results:
-            assert chunk.document_uri is not None
-            assert "other.com" in chunk.document_uri
+        for result in results:
+            assert result.document_uri is not None
+            assert "other.com" in result.document_uri

tests/test_mcp.py

@@ -123,11 +123,12 @@ async def test_mcp_search_documents():
         db_path = Path(temp_dir) / "test.lancedb"
         mcp = create_mcp_server(db_path)
 
-        from haiku.rag.store.models.chunk import Chunk
+        from haiku.rag.store.models import SearchResult
 
-        mock_chunk1 = Chunk(content="Result 1", document_id="doc1")
-        mock_chunk2 = Chunk(content="Result 2", document_id="doc2")
-        mock_results = [(mock_chunk1, 0.9), (mock_chunk2, 0.8)]
+        mock_results = [
+            SearchResult(content="Result 1", score=0.9, document_id="doc1"),
+            SearchResult(content="Result 2", score=0.8, document_id="doc2"),
+        ]
 
         with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class:
             mock_rag = AsyncMock()

tests/test_search.py

@@ -3,6 +3,7 @@ from datasets import Dataset
 
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import Config
+from haiku.rag.store.models import SearchResult
 
 
 @pytest.mark.asyncio
@@ -184,3 +185,62 @@ async def test_search_score_types(temp_db_path):
             )
 
     client.close()
+
+
+@pytest.mark.asyncio
+async def test_search_returns_search_result(temp_db_path):
+    """Test that client.search() returns SearchResult with provenance info."""
+    client = HaikuRAG(db_path=temp_db_path, config=Config)
+
+    await client.create_document(
+        content="Machine learning models can classify images with high accuracy.",
+        uri="https://example.com/ml.html",
+        title="ML Guide",
+    )
+
+    results = await client.search("machine learning", limit=3)
+
+    assert len(results) > 0
+    result = results[0]
+    assert isinstance(result, SearchResult)
+    assert result.content
+    assert result.score > 0
+    assert result.document_uri == "https://example.com/ml.html"
+    assert result.document_title == "ML Guide"
+    # page_numbers and headings come from chunk metadata
+    assert isinstance(result.page_numbers, list)
+    assert isinstance(result.labels, list)
+
+    client.close()
+
+
+@pytest.mark.asyncio
+async def test_search_graceful_degradation(temp_db_path):
+    """Test search works when docling data is unavailable."""
+    from haiku.rag.store.models import Chunk
+
+    client = HaikuRAG(db_path=temp_db_path, config=Config)
+
+    # Create document with custom chunks (no docling document)
+    custom_chunks = [
+        Chunk(content="Custom chunk without docling metadata", metadata={}),
+    ]
+    await client.create_document(
+        content="Document with custom chunks",
+        uri="https://example.com/custom.html",
+        chunks=custom_chunks,
+    )
+
+    results = await client.search("custom chunk", limit=3)
+
+    assert len(results) > 0
+    result = results[0]
+    assert isinstance(result, SearchResult)
+    assert result.content
+    # Bounding boxes should be None when docling is unavailable
+    assert result.bounding_boxes is None
+    # Metadata defaults should still work
+    assert result.page_numbers == []
+    assert result.labels == []
+
+    client.close()