diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c6d5446a..030f8438 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index d40622a2..eed988e8 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/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() diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 523515b2..739e6e62 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/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( diff --git a/haiku_rag_slim/haiku/rag/graph/common/nodes.py b/haiku_rag_slim/haiku/rag/graph/common/nodes.py index 7aee94a7..62720abd 100644 --- a/haiku_rag_slim/haiku/rag/graph/common/nodes.py +++ b/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}" diff --git a/haiku_rag_slim/haiku/rag/graph/common/prompts.py b/haiku_rag_slim/haiku/rag/graph/common/prompts.py index ed10ec98..3b497511 100644 --- a/haiku_rag_slim/haiku/rag/graph/common/prompts.py +++ b/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. diff --git a/haiku_rag_slim/haiku/rag/graph/research/prompts.py b/haiku_rag_slim/haiku/rag/graph/research/prompts.py index 5b90696e..540019ee 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/prompts.py +++ b/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. diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index c6c3f97e..90dd30ba 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/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 [] diff --git a/haiku_rag_slim/haiku/rag/qa/agent.py b/haiku_rag_slim/haiku/rag/qa/agent.py index 2f308008..854c531b 100644 --- a/haiku_rag_slim/haiku/rag/qa/agent.py +++ b/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: diff --git a/haiku_rag_slim/haiku/rag/qa/prompts.py b/haiku_rag_slim/haiku/rag/qa/prompts.py index b25b36af..80b1cb01 100644 --- a/haiku_rag_slim/haiku/rag/qa/prompts.py +++ b/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 """ diff --git a/haiku_rag_slim/haiku/rag/store/models/__init__.py b/haiku_rag_slim/haiku/rag/store/models/__init__.py index f01bd304..d1cc810b 100644 --- a/haiku_rag_slim/haiku/rag/store/models/__init__.py +++ b/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", +] diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index 576ffd80..40c1c88d 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/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, + ) diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py b/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py index 0ece23d1..ef8661fb 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py +++ b/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) diff --git a/tests/test_client.py b/tests/test_client.py index 1d258aec..90342ee8 100644 --- a/tests/test_client.py +++ b/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 diff --git a/tests/test_filter.py b/tests/test_filter.py index cc9c4918..936253a0 100644 --- a/tests/test_filter.py +++ b/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 diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 389f11bc..8c60b558 100644 --- a/tests/test_mcp.py +++ b/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() diff --git a/tests/test_search.py b/tests/test_search.py index 6436d88f..2144fdae 100644 --- a/tests/test_search.py +++ b/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()