Remove unused bounding box calculations

This commit is contained in:
Yiorgis Gozadinos 2025-12-08 12:20:55 +02:00
parent 9a5f8e4368
commit 1b14a4643d
No known key found for this signature in database
4 changed files with 6 additions and 77 deletions

View file

@ -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,
)
)

View file

@ -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:

View file

@ -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

View file

@ -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 == []