From 1b14a4643d48bb09ce3e9f71aafec8d8e6bd20b9 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 8 Dec 2025 12:20:55 +0200 Subject: [PATCH] Remove unused bounding box calculations --- haiku_rag_slim/haiku/rag/client.py | 50 ++----------------- .../haiku/rag/store/models/chunk.py | 5 +- tests/test_context_enhancement.py | 26 ---------- tests/test_search.py | 2 - 4 files changed, 6 insertions(+), 77 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 463e3c57..a890ca2d 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -881,7 +881,6 @@ class HaikuRAG: limit: int = 5, search_type: str = "hybrid", filter: str | None = None, - resolve_bounding_boxes: bool = False, ) -> list[SearchResult]: """Search for relevant chunks using the specified search method with optional reranking. @@ -890,7 +889,6 @@ 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 SearchResult objects ordered by relevance. @@ -909,34 +907,7 @@ class HaikuRAG: chunks = [chunk for chunk, _ in raw_results] chunk_results = await reranker.rerank(query, chunks, top_n=limit) - bounding_boxes_map: dict[str, list] | None = None - if resolve_bounding_boxes: - bounding_boxes_map = {} - doc_cache: dict[str, Document | None] = {} - - 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 - ) - - 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 + return [SearchResult.from_chunk(chunk, score) for chunk, score in chunk_results] async def expand_context( self, @@ -1165,8 +1136,6 @@ class HaikuRAG: Structural content (tables, code, lists) expands to complete structures. Text content uses radius-based expansion. """ - from haiku.rag.store.models.chunk import BoundingBox - all_items = list(docling_doc.iterate_items()) ref_to_index = { getattr(item, "self_ref", None): i @@ -1197,7 +1166,10 @@ class HaikuRAG: final_results: list[SearchResult] = [] for min_idx, max_idx, original_results in merged: - content_parts, refs, pages, labels, bboxes = [], [], set(), set(), [] + content_parts: list[str] = [] + refs: list[str] = [] + pages: set[int] = set() + labels: set[str] = set() for i in range(min_idx, max_idx + 1): item, _ = all_items[i] @@ -1215,16 +1187,6 @@ class HaikuRAG: for p in prov: if (page_no := getattr(p, "page_no", None)) is not None: pages.add(page_no) - if bbox := getattr(p, "bbox", None): - bboxes.append( - BoundingBox( - page_no=page_no or 0, - left=bbox.l, - top=bbox.t, - right=bbox.r, - bottom=bbox.b, - ) - ) # Merge headings preserving order all_headings: list[str] = [] @@ -1245,7 +1207,6 @@ class HaikuRAG: page_numbers=sorted(pages), headings=all_headings or None, labels=sorted(labels), - bounding_boxes=bboxes or None, ) ) @@ -1301,7 +1262,6 @@ class HaikuRAG: page_numbers=first.page_numbers, headings=first.headings, labels=first.labels, - bounding_boxes=first.bounding_boxes, ) ) diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index f3104811..881dc258 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/models/chunk.py @@ -121,16 +121,14 @@ class SearchResult(BaseModel): 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.""" + """Create from a Chunk.""" meta = chunk.get_chunk_metadata() return cls( content=chunk.content, @@ -143,7 +141,6 @@ class SearchResult(BaseModel): page_numbers=meta.page_numbers, headings=meta.headings, labels=meta.labels, - bounding_boxes=bounding_boxes, ) def format_for_agent(self) -> str: diff --git a/tests/test_context_enhancement.py b/tests/test_context_enhancement.py index c7b8d222..4b916f2c 100644 --- a/tests/test_context_enhancement.py +++ b/tests/test_context_enhancement.py @@ -688,29 +688,3 @@ First paragraph of results. assert len(r.content) > 0 # Score should be preserved (best score) assert r.score in [0.9, 0.8] - - -@pytest.mark.asyncio -async def test_expand_context_docling_preserves_bounding_boxes(temp_db_path): - """Test that expand_context preserves bounding boxes from DoclingDocument.""" - config = AppConfig() - config.processing.text_context_radius = 2 - - async with HaikuRAG(temp_db_path, config=config, create=True) as client: - doc = await client.create_document( - content="# Test\n\nSome content here.", - uri="test://bboxes", - ) - - assert doc.id is not None - chunks = await client.chunk_repository.get_by_document_id(doc.id) - - if chunks: - search_results = [SearchResult.from_chunk(chunks[0], 0.9)] - expanded = await client.expand_context(search_results) - - # Expanded results should exist - assert len(expanded) == 1 - # Bounding boxes may or may not be present depending on document - # but the field should be accessible - _ = expanded[0].bounding_boxes diff --git a/tests/test_search.py b/tests/test_search.py index bb023828..b8b8d012 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -237,8 +237,6 @@ async def test_search_graceful_degradation(temp_db_path): 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 == []