62da6086b8
Collapsing the caption text into `get_pictures_grouped` served the enrichment path, which uses it, but the multimodal reranker discards the second return value while still paying to read the column. That is the widest fan-out in the codebase, `limit * 10` candidates, and it previously projected self_ref and picture_data alone. `with_text` is opt-in and off by default, so the cheap projection is what a caller gets unless it asks for more. The reranker test asserts the projection as well as the query count, since a count alone would not notice the column coming back.
424 lines
16 KiB
Python
424 lines
16 KiB
Python
import base64
|
|
from collections.abc import Sequence
|
|
from typing import TYPE_CHECKING
|
|
|
|
from haiku.rag.store.models.chunk import Chunk, SearchResult, SearchType
|
|
from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX
|
|
|
|
if TYPE_CHECKING:
|
|
from PIL import Image as PILImage
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
|
|
async def search(
|
|
client: "HaikuRAG",
|
|
query: "str | bytes | PILImage.Image",
|
|
limit: int | None = None,
|
|
search_type: SearchType | None = None,
|
|
filter: str | None = None,
|
|
include_images: bool = True,
|
|
) -> list[SearchResult]:
|
|
"""Search for relevant chunks with optional reranking.
|
|
|
|
Args:
|
|
client: The HaikuRAG client (provides config + chunk repository).
|
|
query: Text (``str``) or image (``bytes`` / ``PIL.Image.Image``).
|
|
Image queries require a multimodal embedder and run vector-only.
|
|
limit: Maximum number of results to return. Defaults to config.search.limit.
|
|
search_type: "vector", "fts", or "hybrid".
|
|
Applicable only for text queries, where the default is "hybrid".
|
|
filter: Optional SQL WHERE clause to filter documents before searching chunks.
|
|
include_images: When True, populate ``SearchResult.image_data`` with
|
|
base64 picture bytes for picture-labeled chunks.
|
|
|
|
Returns:
|
|
List of SearchResult objects ordered by relevance.
|
|
"""
|
|
if limit is None:
|
|
limit = client._config.search.limit
|
|
|
|
if isinstance(query, str):
|
|
if search_type is None:
|
|
search_type = "hybrid"
|
|
|
|
reranker = client.reranker
|
|
|
|
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]
|
|
if client._config.reranking.multimodal:
|
|
await _attach_picture_data(client, chunks)
|
|
chunk_results = await reranker.rerank(query, chunks, top_n=limit)
|
|
else:
|
|
embedder = client.embedder
|
|
if not embedder.supports_images:
|
|
raise ValueError(
|
|
"Image queries require a multimodal embedder. Set "
|
|
"embeddings.model.multimodal: true on a vllm, voyageai, or cohere "
|
|
"model."
|
|
)
|
|
query_vector = await embedder.embed_image(query)
|
|
chunk_results = await client.chunk_repository.search(
|
|
query="",
|
|
limit=limit,
|
|
filter=filter,
|
|
query_vector=query_vector,
|
|
)
|
|
|
|
results = [SearchResult.from_chunk(chunk, score) for chunk, score in chunk_results]
|
|
results = _dedup_picture_chunks(results)
|
|
|
|
if include_images:
|
|
await _populate_image_data(client, results)
|
|
|
|
return results
|
|
|
|
|
|
async def _attach_picture_data(client: "HaikuRAG", chunks: list[Chunk]) -> None:
|
|
"""Attach picture bytes to synthetic picture chunks in-place, so a
|
|
multimodal reranker can score the pixels instead of just the chunk's
|
|
description text.
|
|
|
|
One query however many documents the candidates span, which matters here
|
|
more than anywhere: reranking fetches `limit * 10` candidates.
|
|
"""
|
|
by_doc: dict[str, list[tuple[Chunk, str]]] = {}
|
|
for chunk in chunks:
|
|
if chunk.document_id is None:
|
|
continue
|
|
refs = chunk.get_chunk_metadata().doc_item_refs
|
|
if len(refs) == 1 and refs[0].startswith(PICTURE_REF_PREFIX):
|
|
by_doc.setdefault(chunk.document_id, []).append((chunk, refs[0]))
|
|
|
|
bytes_by_document, _ = await client.document_item_repository.get_pictures_grouped(
|
|
{doc_id: [ref for _, ref in pairs] for doc_id, pairs in by_doc.items()}
|
|
)
|
|
for doc_id, doc_chunks in by_doc.items():
|
|
bytes_by_ref = bytes_by_document.get(doc_id, {})
|
|
for chunk, ref in doc_chunks:
|
|
data = bytes_by_ref.get(ref)
|
|
if data:
|
|
chunk._picture_data = data
|
|
|
|
|
|
def _dedup_picture_chunks(results: list[SearchResult]) -> list[SearchResult]:
|
|
"""Collapse duplicate picture-only chunks to one result per ``self_ref``.
|
|
|
|
A single picture can produce two chunks for the same self_ref: one whose
|
|
vector is the text embedding of the picture's description, and one whose
|
|
vector is the image embedding of the picture's bytes. Both can rank for
|
|
the same query. When two results share a single picture self_ref as
|
|
their only ref, keep the higher-scoring one. Wider chunks that span the
|
|
picture plus surrounding items pass through untouched.
|
|
"""
|
|
seen: dict[tuple[str | None, str], int] = {}
|
|
keep: list[bool] = [True] * len(results)
|
|
for i, r in enumerate(results):
|
|
if len(r.doc_item_refs) == 1 and r.doc_item_refs[0].startswith(
|
|
PICTURE_REF_PREFIX
|
|
):
|
|
key = (r.document_id, r.doc_item_refs[0])
|
|
prior = seen.get(key)
|
|
if prior is None:
|
|
seen[key] = i
|
|
elif r.score > results[prior].score:
|
|
keep[prior] = False
|
|
seen[key] = i
|
|
else:
|
|
keep[i] = False
|
|
return [r for r, k in zip(results, keep) if k]
|
|
|
|
|
|
async def _populate_image_data(client: "HaikuRAG", results: list[SearchResult]) -> None:
|
|
"""Attach base64 picture bytes to ``SearchResult.image_data`` in-place.
|
|
|
|
A result carries a picture when its refs include the picture directly, or
|
|
when they include the picture's caption — the common case where a prose
|
|
chunk carrying a figure's caption ranks while the picture is its own chunk.
|
|
Costs a fixed number of reads however many documents the result set spans.
|
|
"""
|
|
repo = client.document_item_repository
|
|
by_doc: dict[str, list[SearchResult]] = {}
|
|
for r in results:
|
|
if r.document_id and r.doc_item_refs:
|
|
by_doc.setdefault(r.document_id, []).append(r)
|
|
if not by_doc:
|
|
return
|
|
|
|
refs_by_document = {
|
|
doc_id: list({ref for r in doc_results for ref in r.doc_item_refs})
|
|
for doc_id, doc_results in by_doc.items()
|
|
}
|
|
captions_to_pictures = await repo.get_caption_picture_refs_grouped(refs_by_document)
|
|
|
|
# Which pictures each result wants, and which to fetch per document.
|
|
result_pictures: list[tuple[SearchResult, list[str]]] = []
|
|
wanted: dict[str, list[str]] = {}
|
|
for doc_id, doc_results in by_doc.items():
|
|
caption_to_picture = captions_to_pictures.get(doc_id, {})
|
|
seen: set[str] = set()
|
|
for r in doc_results:
|
|
pictures: list[str] = []
|
|
for ref in r.doc_item_refs:
|
|
picture = (
|
|
ref
|
|
if ref.startswith(PICTURE_REF_PREFIX)
|
|
else caption_to_picture.get(ref)
|
|
)
|
|
if picture and picture not in pictures:
|
|
pictures.append(picture)
|
|
if pictures:
|
|
result_pictures.append((r, pictures))
|
|
for picture in pictures:
|
|
if picture not in seen:
|
|
wanted.setdefault(doc_id, []).append(picture)
|
|
seen.add(picture)
|
|
if not wanted:
|
|
return
|
|
|
|
bytes_by_document, captions_by_document = await repo.get_pictures_grouped(
|
|
wanted, with_text=True
|
|
)
|
|
if not bytes_by_document:
|
|
return
|
|
|
|
for r, pictures in result_pictures:
|
|
bytes_by_ref = bytes_by_document.get(r.document_id or "", {})
|
|
captions_by_ref = captions_by_document.get(r.document_id or "", {})
|
|
attached: dict[str, str] = {}
|
|
captions: dict[str, str] = {}
|
|
for ref in pictures:
|
|
blob = bytes_by_ref.get(ref)
|
|
if blob:
|
|
attached[ref] = base64.b64encode(blob).decode("ascii")
|
|
caption = captions_by_ref.get(ref)
|
|
if caption:
|
|
captions[ref] = caption
|
|
if attached:
|
|
r.image_data = attached
|
|
if captions:
|
|
r.picture_captions = captions
|
|
|
|
|
|
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, window_for
|
|
|
|
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 = []
|
|
expandable = {
|
|
doc_id: doc_results
|
|
for doc_id, doc_results in document_groups.items()
|
|
if doc_id is not None and any(r.doc_item_refs for r in doc_results)
|
|
}
|
|
repo = client.document_item_repository
|
|
positions_by_document = await repo.resolve_refs_grouped(
|
|
{
|
|
doc_id: [ref for r in doc_results for ref in r.doc_item_refs]
|
|
for doc_id, doc_results in expandable.items()
|
|
}
|
|
)
|
|
windows = {
|
|
doc_id: window_for(positions)
|
|
for doc_id, positions in positions_by_document.items()
|
|
if positions
|
|
}
|
|
items_by_document = await repo.get_items_in_ranges(windows)
|
|
|
|
# In document_groups order: the score sort below is stable, so assembling
|
|
# expandable and passthrough documents in separate passes would reorder
|
|
# equal-scored results.
|
|
for doc_id, doc_results in document_groups.items():
|
|
if doc_id not in expandable:
|
|
expanded_results.extend(doc_results)
|
|
continue
|
|
expanded_results.extend(
|
|
expand_with_items(
|
|
doc_results,
|
|
max_chars,
|
|
positions_by_document.get(doc_id, {}),
|
|
items_by_document.get(doc_id, []),
|
|
)
|
|
)
|
|
|
|
expanded_results.sort(key=lambda r: r.score, reverse=True)
|
|
# image_data and picture_captions are preserved through expansion by
|
|
# expand_with_items — we deliberately do not re-attach bytes for refs
|
|
# introduced by section expansion, so the multimodal payload stays
|
|
# bounded by what was originally retrieved.
|
|
return expanded_results
|
|
|
|
|
|
async def visualize_chunk(
|
|
client: "HaikuRAG",
|
|
chunk: "Chunk | Sequence[Chunk]",
|
|
refs: list[str] | None = None,
|
|
expand: bool = True,
|
|
) -> list:
|
|
"""Render page images with bounding box highlights for one or more chunks.
|
|
|
|
When ``refs`` is given (the ``doc_item_refs`` of the citation, i.e. the
|
|
exact items the model saw), bounding boxes are resolved from them directly
|
|
so the visualization matches the cited context precisely. Otherwise, with
|
|
``expand=True`` (default) the chunks' context is re-expanded to recover the
|
|
surrounding section; with ``expand=False`` only the chunks' own items are
|
|
drawn, so the visualization shows just the retrieved chunk with no context.
|
|
|
|
The chunks' own items draw in a strong highlight; the remaining items draw
|
|
fainter, so the matched content stands out from its surrounding context.
|
|
Chunks from a different document than the first are ignored.
|
|
|
|
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
|
|
|
|
chunks = [chunk] if isinstance(chunk, Chunk) else list(chunk)
|
|
if not chunks:
|
|
return []
|
|
document_id = chunks[0].document_id
|
|
if not document_id:
|
|
return []
|
|
chunks = [c for c in chunks if c.document_id == document_id]
|
|
|
|
doc = await client.document_repository.get_docling_data(document_id)
|
|
if not doc:
|
|
return []
|
|
|
|
docling_doc = doc.get_docling_document()
|
|
if not docling_doc:
|
|
return []
|
|
|
|
matched_refs = {r for c in chunks for r in c.get_chunk_metadata().doc_item_refs}
|
|
|
|
if refs is not None:
|
|
all_refs = list(refs)
|
|
elif not expand:
|
|
# Chunk-only: draw just the retrieved chunks' own items, no context.
|
|
all_refs = list(matched_refs)
|
|
else:
|
|
# No stored context: re-expand the chunks to recover their section.
|
|
search_results = [
|
|
SearchResult(
|
|
content=c.content,
|
|
score=1.0,
|
|
chunk_id=c.id,
|
|
document_id=c.document_id,
|
|
doc_item_refs=meta.doc_item_refs,
|
|
page_numbers=meta.page_numbers,
|
|
)
|
|
for c in chunks
|
|
if (meta := c.get_chunk_metadata()).doc_item_refs
|
|
]
|
|
if search_results:
|
|
expanded = await expand_context(client, search_results)
|
|
all_refs = []
|
|
for result in expanded:
|
|
all_refs.extend(r for r in result.doc_item_refs if r not in all_refs)
|
|
if not all_refs:
|
|
all_refs = [r for sr in search_results for r in sr.doc_item_refs]
|
|
else:
|
|
all_refs = list(chunks[0].get_chunk_metadata().doc_item_refs)
|
|
|
|
matched_draw = [r for r in all_refs if r in matched_refs]
|
|
swept_refs = [r for r in all_refs if r not in matched_refs]
|
|
|
|
matched_boxes = ChunkMetadata(doc_item_refs=matched_draw).resolve_bounding_boxes(
|
|
docling_doc
|
|
)
|
|
swept_boxes = ChunkMetadata(doc_item_refs=swept_refs).resolve_bounding_boxes(
|
|
docling_doc
|
|
)
|
|
if not matched_boxes and not swept_boxes:
|
|
return []
|
|
|
|
# Group bounding boxes by page; swept boxes first so matched draw on top
|
|
boxes_by_page: dict[int, list] = {}
|
|
for bbox, is_matched in [(b, False) for b in swept_boxes] + [
|
|
(b, True) for b in matched_boxes
|
|
]:
|
|
if bbox.page_no not in boxes_by_page:
|
|
boxes_by_page[bbox.page_no] = []
|
|
boxes_by_page[bbox.page_no].append((bbox, is_matched))
|
|
|
|
# Load only the needed page images
|
|
pages_doc = await client.document_repository.get_pages_data(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, is_matched 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
|
|
|
|
if is_matched:
|
|
fill_color = (255, 150, 0, 55) # Orange, matched content
|
|
outline_color = (240, 130, 0, 150) # Orange outline
|
|
else:
|
|
fill_color = (255, 255, 0, 40) # Yellow, surrounding context
|
|
outline_color = (255, 165, 0, 100)
|
|
|
|
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
|