diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index abcd1b09..7537fe4b 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -862,89 +862,17 @@ class HaikuRAG: search_type: str = "hybrid", filter: str | None = None, ) -> list[SearchResult]: - """Search for relevant chunks using the specified search method with optional reranking. + from haiku.rag.client.search import search - Args: - query: The search query string. - limit: Maximum number of results to return. Defaults to config.search.default_limit. - search_type: Type of search - "vector", "fts", or "hybrid" (default). - filter: Optional SQL WHERE clause to filter documents before searching chunks. - - Returns: - List of SearchResult objects ordered by relevance. - """ - if limit is None: - limit = self._config.search.limit - - reranker = get_reranker(config=self._config) - - if reranker is None: - 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) - - return [SearchResult.from_chunk(chunk, score) for chunk, score in chunk_results] + return await search(self, query, limit, search_type, filter) async def expand_context( self, search_results: list[SearchResult], ) -> list[SearchResult]: - """Expand search results with surrounding content from the document. + from haiku.rag.client.search import expand_context - Uses the document_items table for section-bounded expansion. - See haiku.rag.context for the algorithm description. - - Results without doc_item_refs pass through unexpanded. This happens - when chunks were created without docling metadata (e.g., custom chunks - passed to import_document). - - Args: - search_results: List of SearchResult objects from search. - - Returns: - List of SearchResult objects with expanded content. - """ - from haiku.rag.context import expand_with_items - - max_chars = self._config.search.max_context_chars - - # 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(result) - - expanded_results = [] - - for doc_id, doc_results in document_groups.items(): - if doc_id is None: - expanded_results.extend(doc_results) - continue - - has_refs = any(r.doc_item_refs for r in doc_results) - if not has_refs: - expanded_results.extend(doc_results) - continue - - expanded = await expand_with_items( - self.document_item_repository, - doc_id, - doc_results, - max_chars, - ) - expanded_results.extend(expanded) - - expanded_results.sort(key=lambda r: r.score, reverse=True) - return expanded_results + return await expand_context(self, search_results) async def ask( self, @@ -1076,113 +1004,9 @@ class HaikuRAG: ) async def visualize_chunk(self, chunk: Chunk) -> list: - """Render page images with bounding box highlights for a chunk. + from haiku.rag.client.search import visualize_chunk - Expands the chunk's context to find the full section, then resolves - bounding boxes from all items in the expanded range. This ensures - visualization covers all pages the expanded content spans. - - Args: - chunk: The chunk to visualize. - - Returns: - List of PIL Image objects, one per page with bounding boxes. - Empty list if no bounding boxes or page images available. - """ - from copy import deepcopy - - from PIL import ImageDraw - - from haiku.rag.store.models.chunk import ChunkMetadata - - # Get the document structure (from cache if available) - if not chunk.document_id: - return [] - - doc = await self.document_repository.get_docling_data(chunk.document_id) - if not doc: - return [] - - docling_doc = doc.get_docling_document() - if not docling_doc: - return [] - - # Expand context to get all doc_item_refs in the section - chunk_meta = chunk.get_chunk_metadata() - if chunk_meta.doc_item_refs: - search_result = SearchResult( - content=chunk.content, - score=1.0, - chunk_id=chunk.id, - document_id=chunk.document_id, - doc_item_refs=chunk_meta.doc_item_refs, - page_numbers=chunk_meta.page_numbers, - ) - expanded = await self.expand_context([search_result]) - refs = expanded[0].doc_item_refs if expanded else chunk_meta.doc_item_refs - meta = ChunkMetadata(doc_item_refs=refs) - else: - meta = chunk_meta - bounding_boxes = meta.resolve_bounding_boxes(docling_doc) - if not bounding_boxes: - return [] - - # Group bounding boxes by page - boxes_by_page: dict[int, list] = {} - for bbox in bounding_boxes: - if bbox.page_no not in boxes_by_page: - boxes_by_page[bbox.page_no] = [] - boxes_by_page[bbox.page_no].append(bbox) - - # Load only the needed page images - pages_doc = await self.document_repository.get_pages_data(chunk.document_id) - if not pages_doc: - return [] - page_images = pages_doc.get_page_images(list(boxes_by_page.keys())) - - # Render each page with its bounding boxes - images = [] - for page_no in sorted(boxes_by_page.keys()): - if page_no not in page_images: - continue - - page = page_images[page_no] - if page.image is None or page.image.pil_image is None: - continue - - pil_image = page.image.pil_image - page_height = page.size.height - - # Calculate scale factor (image pixels vs document coordinates) - scale_x = pil_image.width / page.size.width - scale_y = pil_image.height / page.size.height - - # Draw bounding boxes - image = deepcopy(pil_image) - draw = ImageDraw.Draw(image, "RGBA") - - for bbox in boxes_by_page[page_no]: - # Convert from document coordinates to image coordinates - # Document coords are bottom-left origin, PIL uses top-left - x0 = bbox.left * scale_x - y0 = (page_height - bbox.top) * scale_y - x1 = bbox.right * scale_x - y1 = (page_height - bbox.bottom) * scale_y - - # Ensure proper ordering (y0 should be less than y1 for PIL) - if y0 > y1: - y0, y1 = y1, y0 - - # Draw filled rectangle with transparency - fill_color = (255, 255, 0, 40) # Yellow with transparency - outline_color = (255, 165, 0, 100) # Orange outline - - draw.rectangle([(x0, y0), (x1, y1)], fill=fill_color, outline=None) - draw.rectangle([(x0, y0), (x1, y1)], outline=outline_color, width=1) - - images.append(image) - - return images + return await visualize_chunk(self, chunk) async def rebuild_database( self, mode: RebuildMode = RebuildMode.FULL diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py new file mode 100644 index 00000000..10af0621 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -0,0 +1,196 @@ +from typing import TYPE_CHECKING + +from haiku.rag.reranking import get_reranker +from haiku.rag.store.models.chunk import Chunk, SearchResult + +if TYPE_CHECKING: + from haiku.rag.client import HaikuRAG + + +async def search( + client: "HaikuRAG", + query: str, + limit: int | None = None, + search_type: str = "hybrid", + filter: str | None = None, +) -> list[SearchResult]: + """Search for relevant chunks with optional reranking. + + Args: + client: The HaikuRAG client (provides config + chunk repository). + query: The search query string. + limit: Maximum number of results to return. Defaults to config.search.limit. + search_type: Type of search - "vector", "fts", or "hybrid" (default). + filter: Optional SQL WHERE clause to filter documents before searching chunks. + + Returns: + List of SearchResult objects ordered by relevance. + """ + if limit is None: + limit = client._config.search.limit + + reranker = get_reranker(config=client._config) + + if reranker is None: + chunk_results = await client.chunk_repository.search( + query, limit, search_type, filter + ) + else: + search_limit = limit * 10 + raw_results = await client.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) + + return [SearchResult.from_chunk(chunk, score) for chunk, score in chunk_results] + + +async def expand_context( + client: "HaikuRAG", + search_results: list[SearchResult], +) -> list[SearchResult]: + """Expand search results with surrounding content from the document. + + Uses the document_items table for section-bounded expansion. + See haiku.rag.context for the algorithm description. + + Results without doc_item_refs pass through unexpanded. This happens when + chunks were created without docling metadata (e.g., custom chunks passed + to import_document). + """ + from haiku.rag.context import expand_with_items + + max_chars = client._config.search.max_context_chars + + # 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(result) + + expanded_results = [] + + for doc_id, doc_results in document_groups.items(): + if doc_id is None: + expanded_results.extend(doc_results) + continue + + has_refs = any(r.doc_item_refs for r in doc_results) + if not has_refs: + expanded_results.extend(doc_results) + continue + + expanded = await expand_with_items( + client.document_item_repository, + doc_id, + doc_results, + max_chars, + ) + expanded_results.extend(expanded) + + expanded_results.sort(key=lambda r: r.score, reverse=True) + return expanded_results + + +async def visualize_chunk(client: "HaikuRAG", chunk: Chunk) -> list: + """Render page images with bounding box highlights for a chunk. + + Expands the chunk's context to find the full section, then resolves + bounding boxes from all items in the expanded range. This ensures + visualization covers all pages the expanded content spans. + + Returns a list of PIL Image objects, one per page with bounding boxes. + Empty list if no bounding boxes or page images available. + """ + from copy import deepcopy + + from PIL import ImageDraw + + from haiku.rag.store.models.chunk import ChunkMetadata + + if not chunk.document_id: + return [] + + doc = await client.document_repository.get_docling_data(chunk.document_id) + if not doc: + return [] + + docling_doc = doc.get_docling_document() + if not docling_doc: + return [] + + # Expand context to get all doc_item_refs in the section + chunk_meta = chunk.get_chunk_metadata() + if chunk_meta.doc_item_refs: + search_result = SearchResult( + content=chunk.content, + score=1.0, + chunk_id=chunk.id, + document_id=chunk.document_id, + doc_item_refs=chunk_meta.doc_item_refs, + page_numbers=chunk_meta.page_numbers, + ) + expanded = await expand_context(client, [search_result]) + refs = expanded[0].doc_item_refs if expanded else chunk_meta.doc_item_refs + meta = ChunkMetadata(doc_item_refs=refs) + else: + meta = chunk_meta + bounding_boxes = meta.resolve_bounding_boxes(docling_doc) + if not bounding_boxes: + return [] + + # Group bounding boxes by page + boxes_by_page: dict[int, list] = {} + for bbox in bounding_boxes: + if bbox.page_no not in boxes_by_page: + boxes_by_page[bbox.page_no] = [] + boxes_by_page[bbox.page_no].append(bbox) + + # Load only the needed page images + pages_doc = await client.document_repository.get_pages_data(chunk.document_id) + if not pages_doc: + return [] + page_images = pages_doc.get_page_images(list(boxes_by_page.keys())) + + # Render each page with its bounding boxes + images = [] + for page_no in sorted(boxes_by_page.keys()): + if page_no not in page_images: + continue + + page = page_images[page_no] + if page.image is None or page.image.pil_image is None: + continue + + pil_image = page.image.pil_image + page_height = page.size.height + + # Scale factor: image pixels vs document coordinates + scale_x = pil_image.width / page.size.width + scale_y = pil_image.height / page.size.height + + image = deepcopy(pil_image) + draw = ImageDraw.Draw(image, "RGBA") + + for bbox in boxes_by_page[page_no]: + # Document coords are bottom-left origin; PIL uses top-left + x0 = bbox.left * scale_x + y0 = (page_height - bbox.top) * scale_y + x1 = bbox.right * scale_x + y1 = (page_height - bbox.bottom) * scale_y + + if y0 > y1: + y0, y1 = y1, y0 + + fill_color = (255, 255, 0, 40) # Yellow with transparency + outline_color = (255, 165, 0, 100) # Orange outline + + draw.rectangle([(x0, y0), (x1, y1)], fill=fill_color, outline=None) + draw.rectangle([(x0, y0), (x1, y1)], outline=outline_color, width=1) + + images.append(image) + + return images