7be71ebbac
Retrieval scores are each database's own rank arithmetic; the databases in a selection share an embedder, so similarity in that one space is the signal comparable across databases by construction. Measured product to product against score ordering: +8.3 to +16.6pp recall@5 across five cells on two corpora, flat in collection count and corpus shape where score ordering dips with both, closing roughly 60% of the gap to a reranker; order-sensitivity residual 0.00pp in every cell. Exact ties collapse from 51-81% of candidates to under 1%. Full-text-only searches keep retrieval-score order, having no query vector. The vector column already travels with every search result, so the similarity costs no additional transfer; per-chunk embeddings are materialized only for the federated path that reads them.
718 lines
27 KiB
Python
718 lines
27 KiB
Python
import base64
|
|
from collections.abc import Sequence
|
|
from typing import TYPE_CHECKING
|
|
|
|
from haiku.rag.store.models.chunk import (
|
|
Chunk,
|
|
SearchResult,
|
|
SearchType,
|
|
qualified_id,
|
|
)
|
|
from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX
|
|
from haiku.rag.utils import gather_all
|
|
|
|
if TYPE_CHECKING:
|
|
from PIL import Image as PILImage
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.client.session import FederatedSession, SingleDatabaseSession
|
|
|
|
|
|
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
|
|
|
|
resolved = _resolved_search_type(query, search_type)
|
|
# The repository embeds late, so a filter matching nothing never embeds.
|
|
query_vector = (
|
|
None if isinstance(query, str) else await _embed_query(client, query, resolved)
|
|
)
|
|
candidates = await client.chunk_repository.search(
|
|
query=query if isinstance(query, str) else "",
|
|
limit=_fetch_limit(client, query, limit),
|
|
search_type=resolved,
|
|
filter=filter,
|
|
query_vector=query_vector,
|
|
)
|
|
chunk_results = await _rank(client, query, candidates, limit)
|
|
|
|
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 search_sources(
|
|
client: "HaikuRAG",
|
|
query: "str | bytes | PILImage.Image",
|
|
limit: int | None = None,
|
|
search_type: SearchType | None = None,
|
|
filter: str | None = None,
|
|
include_images: bool = True,
|
|
sources: list[str] | None = None,
|
|
) -> list[SearchResult]:
|
|
"""Search several databases and fuse their results into one ranked list.
|
|
|
|
Fetch, fuse, truncate, then enrich: enrichment runs on the survivors through
|
|
the database each came from, so its cost is that of a single-database
|
|
search.
|
|
"""
|
|
if limit is None:
|
|
limit = client._config.search.limit
|
|
|
|
names = list(client.source_names) if sources is None else list(sources)
|
|
if not names:
|
|
return []
|
|
selected = await client.clients_for(names)
|
|
if len(selected) == 1:
|
|
# One database is an ordinary search: it keeps the database's own
|
|
# hybrid scores and embeds late.
|
|
return await selected[0].search(
|
|
query, limit, search_type, filter, include_images
|
|
)
|
|
resolved = _resolved_search_type(query, search_type)
|
|
if resolved != "fts":
|
|
client._require_one_embedder(selected)
|
|
|
|
# One over-fetch decision, one query vector and one reranker for the whole
|
|
# set: the databases share an embedder.
|
|
fetch_limit = _fetch_limit(client, query, limit)
|
|
query_vector = await _embed_query(selected[0], query, resolved)
|
|
text = query if isinstance(query, str) else ""
|
|
# Embeddings are read only by cosine fusion: a reranker scores the union
|
|
# itself, and its 10x over-fetch would materialize them for nothing.
|
|
uses_cosine = query_vector is not None and (
|
|
not isinstance(query, str) or client.reranker is None
|
|
)
|
|
per_source = await gather_all(
|
|
*(
|
|
c.chunk_repository.search(
|
|
query=text,
|
|
limit=fetch_limit,
|
|
search_type=resolved,
|
|
filter=filter,
|
|
query_vector=query_vector,
|
|
with_vectors=uses_cosine,
|
|
)
|
|
for c in selected
|
|
)
|
|
)
|
|
|
|
ranked = await _fuse(
|
|
client, selected, query, per_source, limit, query_vector=query_vector
|
|
)
|
|
|
|
results: list[SearchResult] = []
|
|
for owner, chunk, score in ranked:
|
|
result = SearchResult.from_chunk(chunk, score)
|
|
result.source = owner.source
|
|
results.append(result)
|
|
results = _dedup_picture_chunks(results)
|
|
|
|
if include_images:
|
|
by_owner: dict[str, list[SearchResult]] = {}
|
|
for result in results:
|
|
if result.source:
|
|
by_owner.setdefault(result.source, []).append(result)
|
|
owners = await client.clients_for(list(by_owner))
|
|
await gather_all(
|
|
*(
|
|
_populate_image_data(owner, by_owner[name])
|
|
for name, owner in zip(by_owner, owners, strict=True)
|
|
)
|
|
)
|
|
|
|
return results
|
|
|
|
|
|
async def _fuse(
|
|
federator: "HaikuRAG",
|
|
clients: list["HaikuRAG"],
|
|
query: "str | bytes | PILImage.Image",
|
|
per_source: list[list[tuple[Chunk, float]]],
|
|
limit: int,
|
|
query_vector: list[float] | None = None,
|
|
) -> list[tuple["HaikuRAG", Chunk, float]]:
|
|
"""One ranked list from several, keeping each candidate's owner.
|
|
|
|
A configured reranker scores the union directly, which is what makes ranking
|
|
across databases tractable: it compares query against document and does not
|
|
care where a candidate came from. Without one, the union is ordered by
|
|
cosine similarity to the query vector: the databases in a selection share an
|
|
embedder, so similarity in that one space is the signal that is comparable
|
|
across databases by construction, where retrieval scores are each database's
|
|
own rank arithmetic. A search with no query vector (full-text) orders by the
|
|
retrieval score instead. In both, ties resolve by within-database rank — the
|
|
candidate nothing in its own database beat wins — and only a tie on both
|
|
falls to configured order. The returned score is the one the union was
|
|
ordered by, so downstream re-sorts (context expansion) preserve this order.
|
|
"""
|
|
owned = [
|
|
(client, chunk, score)
|
|
for client, candidates in zip(clients, per_source, strict=True)
|
|
for chunk, score in candidates
|
|
]
|
|
if not owned:
|
|
return []
|
|
|
|
# The reranker interface takes a text query, so an image query skips it,
|
|
# and the check precedes `reranker`, which builds the reranker on first
|
|
# access.
|
|
if isinstance(query, str):
|
|
reranker = federator.reranker
|
|
if reranker is not None:
|
|
chunks = [chunk for _, chunk, _ in owned]
|
|
if federator._config.reranking.multimodal:
|
|
await gather_all(
|
|
*(
|
|
_attach_picture_data(
|
|
c, [chunk for owner, chunk, _ in owned if owner is c]
|
|
)
|
|
for c in clients
|
|
)
|
|
)
|
|
reranked = await reranker.rerank(query, chunks, top_n=limit)
|
|
# Identity, since chunk ids repeat between copies of a database.
|
|
owner_of = {id(chunk): client for client, chunk, _ in owned}
|
|
if any(id(chunk) not in owner_of for chunk, _ in reranked):
|
|
raise ValueError(
|
|
f"{type(reranker).__name__} returned chunks that are not the "
|
|
"ones it was given, so the database each came from is lost; "
|
|
"a reranker must return objects from the list passed to it"
|
|
)
|
|
return [(owner_of[id(chunk)], chunk, score) for chunk, score in reranked]
|
|
|
|
scored: list[tuple[float, float, HaikuRAG, Chunk]] = []
|
|
for client, candidates in zip(clients, per_source, strict=True):
|
|
for rank, (chunk, score) in enumerate(candidates):
|
|
scored.append((1.0 / (_RRF_K + rank + 1), score, client, chunk))
|
|
|
|
embeddings = [chunk.embedding for _, _, _, chunk in scored]
|
|
if query_vector is not None and all(e is not None for e in embeddings):
|
|
similarities = _cosine_to(query_vector, embeddings) # ty: ignore[invalid-argument-type]
|
|
scored = [
|
|
(rank_score, similarity, client, chunk)
|
|
for (rank_score, _, client, chunk), similarity in zip(
|
|
scored, similarities, strict=True
|
|
)
|
|
]
|
|
scored.sort(key=lambda item: (item[1], item[0]), reverse=True)
|
|
return [(client, chunk, score) for _, score, client, chunk in scored[:limit]]
|
|
|
|
|
|
def _cosine_to(query_vector: list[float], embeddings: list[list[float]]) -> list[float]:
|
|
"""Cosine similarity of each embedding to the query vector.
|
|
|
|
A zero-norm vector has no direction, so its similarity is 0 rather than a
|
|
division error.
|
|
"""
|
|
import numpy as np
|
|
|
|
query = np.asarray(query_vector, dtype=np.float32)
|
|
matrix = np.asarray(embeddings, dtype=np.float32)
|
|
norms = np.linalg.norm(matrix, axis=1) * np.linalg.norm(query)
|
|
with np.errstate(divide="ignore", invalid="ignore"):
|
|
similarities = np.where(norms > 0, matrix @ query / norms, 0.0)
|
|
return [float(s) for s in similarities]
|
|
|
|
|
|
# Reciprocal rank fusion's smoothing constant, the value the literature uses.
|
|
_RRF_K = 60
|
|
|
|
|
|
# Candidates per requested result when a reranker will re-order them.
|
|
_RERANK_OVERFETCH = 10
|
|
|
|
|
|
def _fetch_limit(
|
|
client: "HaikuRAG",
|
|
query: "str | bytes | PILImage.Image",
|
|
limit: int,
|
|
) -> int:
|
|
"""How many candidates to fetch per database.
|
|
|
|
Only a text query with a reranker over-fetches: an image query keeps its
|
|
vector ranking, and the type is checked before `reranker`, which loads model
|
|
weights for a local one on first access.
|
|
"""
|
|
if not isinstance(query, str):
|
|
return limit
|
|
return limit * _RERANK_OVERFETCH if client.reranker else limit
|
|
|
|
|
|
def _resolved_search_type(
|
|
query: "str | bytes | PILImage.Image", search_type: SearchType | None
|
|
) -> SearchType:
|
|
"""The search actually run for this query.
|
|
|
|
An image query has no text to match against, so it is vector-only whatever
|
|
the caller asked for; a text query defaults to hybrid.
|
|
"""
|
|
if not isinstance(query, str):
|
|
return "vector"
|
|
return search_type or "hybrid"
|
|
|
|
|
|
async def _embed_query(
|
|
client: "HaikuRAG", query: "str | bytes | PILImage.Image", search_type: SearchType
|
|
) -> list[float] | None:
|
|
"""The query as a vector, or None when the search needs no vector.
|
|
|
|
The caller computes it once for however many databases the search covers:
|
|
the databases in a selection share an embedder. `search_type` is the
|
|
resolved one, so only a text query ever reaches this as full-text.
|
|
"""
|
|
if search_type == "fts":
|
|
return None
|
|
if isinstance(query, str):
|
|
return await client.embedder.embed_query(query)
|
|
|
|
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."
|
|
)
|
|
return await embedder.embed_image(query)
|
|
|
|
|
|
async def _rank(
|
|
client: "HaikuRAG",
|
|
query: "str | bytes | PILImage.Image",
|
|
candidates: list[tuple[Chunk, float]],
|
|
limit: int,
|
|
) -> list[tuple[Chunk, float]]:
|
|
"""Order candidates and cut them to `limit`.
|
|
|
|
The reranker interface takes a text query, so an image query keeps the
|
|
vector ranking. Its type is checked before `client.reranker`, which builds
|
|
the reranker on first access and loads model weights for a local one.
|
|
"""
|
|
if not isinstance(query, str):
|
|
return candidates[:limit]
|
|
|
|
reranker = client.reranker
|
|
if reranker is None:
|
|
return candidates[:limit]
|
|
|
|
chunks = [chunk for chunk, _ in candidates]
|
|
if client._config.reranking.multimodal:
|
|
await _attach_picture_data(client, chunks)
|
|
return await reranker.rerank(query, chunks, top_n=limit)
|
|
|
|
|
|
async def _attach_picture_data(client: "HaikuRAG", chunks: list[Chunk]) -> None:
|
|
"""Attach picture bytes to synthetic picture chunks in-place; a multimodal
|
|
reranker scores the pixels.
|
|
|
|
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``.
|
|
|
|
Keyed by database as well: copies of a database hold the same document id,
|
|
and each copy's picture is its own result.
|
|
|
|
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 | 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.source, 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_sources(
|
|
federated: "FederatedSession",
|
|
search_results: list[SearchResult],
|
|
) -> list[SearchResult]:
|
|
"""Expand results drawn from several databases, each through its own.
|
|
|
|
A result naming no database passes through unexpanded: it cannot be placed,
|
|
which is the case for results a caller built by hand.
|
|
"""
|
|
by_source: dict[str, list[SearchResult]] = {}
|
|
unsourced: list[SearchResult] = []
|
|
for result in search_results:
|
|
if result.source:
|
|
by_source.setdefault(result.source, []).append(result)
|
|
else:
|
|
unsourced.append(result)
|
|
names = list(by_source)
|
|
sessions = await federated.sessions_for(names)
|
|
expanded_groups = await gather_all(
|
|
*(
|
|
expand_context(session, by_source[name])
|
|
for name, session in zip(names, sessions, strict=True)
|
|
)
|
|
)
|
|
merged = unsourced + [r for group in expanded_groups for r in group]
|
|
# Grouping by database must not become the tiebreak: fused scores tie often,
|
|
# so equal scores keep the order they were fused in.
|
|
arrival = {
|
|
qualified_id(result.source, result.chunk_id): rank
|
|
for rank, result in enumerate(search_results)
|
|
if result.chunk_id
|
|
}
|
|
|
|
def fused_rank(result: SearchResult) -> int:
|
|
return min(
|
|
(
|
|
arrival[key]
|
|
for cid in (result.chunk_id, *result.chunk_ids)
|
|
if (key := qualified_id(result.source, cid)) in arrival
|
|
),
|
|
default=len(arrival),
|
|
)
|
|
|
|
merged.sort(key=lambda r: (-r.score, fused_rank(r)))
|
|
return merged
|
|
|
|
|
|
async def expand_context(
|
|
session: "SingleDatabaseSession",
|
|
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 = session.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 = session.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(
|
|
session: "SingleDatabaseSession",
|
|
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 session.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(session, 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 session.document_repository.get_pages_data(document_id)
|
|
if not pages_doc:
|
|
return []
|
|
page_images = pages_doc.get_page_images(list(boxes_by_page.keys()))
|
|
|
|
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
|