Merge pull request #491 from ggozad/fix/image-embeddings-limit
Filter picture chunks at ingest and pool embedding/reranking HTTP clients
This commit is contained in:
commit
777ce193f8
14 changed files with 406 additions and 57 deletions
|
|
@ -1,6 +1,15 @@
|
||||||
# Changelog
|
# Changelog
|
||||||
## [Unreleased]
|
## [Unreleased]
|
||||||
|
|
||||||
|
### Changed
|
||||||
|
|
||||||
|
- Duplicate images within a document produce a single picture chunk.
|
||||||
|
- Pictures smaller than `processing.min_picture_size` pixels on the smaller side (default 64) no longer become picture chunks; `0` disables the filter.
|
||||||
|
|
||||||
|
### Fixed
|
||||||
|
|
||||||
|
- vLLM embedding and vLLM/Jina reranking reuse one HTTP client across requests instead of opening one per request.
|
||||||
|
|
||||||
## [0.64.0] - 2026-07-08
|
## [0.64.0] - 2026-07-08
|
||||||
|
|
||||||
### Added
|
### Added
|
||||||
|
|
|
||||||
|
|
@ -259,6 +259,8 @@ Per-image failures (404, timeout, oversized, unreadable) leave that picture as a
|
||||||
| `description` | on | yes | yes |
|
| `description` | on | yes | yes |
|
||||||
| `image` (default) | on | yes | no |
|
| `image` (default) | on | yes | no |
|
||||||
|
|
||||||
|
Not every picture becomes a picture chunk. Identical picture bytes within a document produce a single chunk, so a watermark or logo repeated on every page embeds once. Pictures smaller than `processing.min_picture_size` pixels on their smaller side (default 64, `0` disables) are skipped entirely. Filtered pictures keep their bytes in `document_items`, so context expansion and vision QA still see them.
|
||||||
|
|
||||||
Use `none` when you don't need picture content (e.g. very large reference manuals where RAM is tight). Use `description` to weave VLM-generated text into chunk content and keep bytes for later. Use `image` (default) to keep bytes without paying the VLM cost. The prompt is configurable under `prompts.picture_description`. See [Prompts](prompts.md).
|
Use `none` when you don't need picture content (e.g. very large reference manuals where RAM is tight). Use `description` to weave VLM-generated text into chunk content and keep bytes for later. Use `image` (default) to keep bytes without paying the VLM cost. The prompt is configurable under `prompts.picture_description`. See [Prompts](prompts.md).
|
||||||
|
|
||||||
```yaml
|
```yaml
|
||||||
|
|
@ -302,9 +304,9 @@ The Embedder column below is driven by `embeddings.model.multimodal`, not the pr
|
||||||
|---|---|---|---|
|
|---|---|---|---|
|
||||||
| `none` | any | text only (caption/surrounding) | none |
|
| `none` | any | text only (caption/surrounding) | none |
|
||||||
| `image` | text-only | text only (caption/surrounding) | none |
|
| `image` | text-only | text only (caption/surrounding) | none |
|
||||||
| `image` | multimodal | text only | one per picture, vector = image embedding |
|
| `image` | multimodal | text only | one per distinct picture, vector = image embedding |
|
||||||
| `description` | text-only | text + descriptions | none |
|
| `description` | text-only | text + descriptions | none |
|
||||||
| `description` | multimodal | text + descriptions | one per picture, vector = image embedding |
|
| `description` | multimodal | text + descriptions | one per distinct picture, vector = image embedding |
|
||||||
|
|
||||||
**What QA receives** at search time:
|
**What QA receives** at search time:
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -143,6 +143,16 @@ class HaikuRAG:
|
||||||
async def __aexit__(self, exc_type, exc_val, exc_tb): # noqa: ARG002
|
async def __aexit__(self, exc_type, exc_val, exc_tb): # noqa: ARG002
|
||||||
"""Async context manager exit."""
|
"""Async context manager exit."""
|
||||||
await self._await_vacuum_tasks()
|
await self._await_vacuum_tasks()
|
||||||
|
# Best-effort: __aexit__ may run during exception unwinding, and a
|
||||||
|
# raising close must not mask the original exception. The reranker is
|
||||||
|
# a cached_property — close it only if it was materialized.
|
||||||
|
try:
|
||||||
|
await self.embedder.aclose()
|
||||||
|
reranker = self.__dict__.get("reranker")
|
||||||
|
if reranker is not None:
|
||||||
|
await reranker.aclose()
|
||||||
|
except Exception:
|
||||||
|
logger.debug("Closing embedder/reranker failed on teardown", exc_info=True)
|
||||||
self.close()
|
self.close()
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,5 @@
|
||||||
import asyncio
|
import asyncio
|
||||||
|
import io
|
||||||
import logging
|
import logging
|
||||||
import tempfile
|
import tempfile
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
@ -14,7 +15,7 @@ from haiku.rag.store.models.chunk import Chunk
|
||||||
from haiku.rag.store.models.document_item import _picture_description_text
|
from haiku.rag.store.models.document_item import _picture_description_text
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from docling_core.types.doc.document import DoclingDocument
|
from docling_core.types.doc.document import DoclingDocument, PictureItem
|
||||||
|
|
||||||
from haiku.rag.embeddings import EmbedderWrapper
|
from haiku.rag.embeddings import EmbedderWrapper
|
||||||
|
|
||||||
|
|
@ -173,11 +174,13 @@ def _merge_picture_chunks(
|
||||||
text_chunks: list[Chunk],
|
text_chunks: list[Chunk],
|
||||||
document_id: str | None,
|
document_id: str | None,
|
||||||
existing_picture_data: dict[str, bytes] | None,
|
existing_picture_data: dict[str, bytes] | None,
|
||||||
|
min_picture_size: int,
|
||||||
) -> list[Chunk]:
|
) -> list[Chunk]:
|
||||||
picture_chunks = build_picture_chunks(
|
picture_chunks = build_picture_chunks(
|
||||||
docling_document,
|
docling_document,
|
||||||
document_id=document_id,
|
document_id=document_id,
|
||||||
existing_picture_data=existing_picture_data,
|
existing_picture_data=existing_picture_data,
|
||||||
|
min_picture_size=min_picture_size,
|
||||||
)
|
)
|
||||||
|
|
||||||
if not picture_chunks:
|
if not picture_chunks:
|
||||||
|
|
@ -235,22 +238,49 @@ async def chunk(
|
||||||
text_chunks,
|
text_chunks,
|
||||||
document_id,
|
document_id,
|
||||||
existing_picture_data,
|
existing_picture_data,
|
||||||
|
config.processing.min_picture_size,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _min_picture_side(picture: "PictureItem", data: bytes) -> float | None:
|
||||||
|
"""The picture's smaller pixel dimension, or None when it can't be
|
||||||
|
determined. Uses ``ImageRef.size`` when the live image is present; falls
|
||||||
|
back to a PIL header read of the bytes (rebuild path, where picture URIs
|
||||||
|
have been stripped)."""
|
||||||
|
if picture.image is not None:
|
||||||
|
return min(picture.image.size.width, picture.image.size.height)
|
||||||
|
|
||||||
|
from PIL import Image as PILImage
|
||||||
|
from PIL import UnidentifiedImageError
|
||||||
|
|
||||||
|
try:
|
||||||
|
with PILImage.open(io.BytesIO(data)) as img:
|
||||||
|
return min(img.size)
|
||||||
|
except UnidentifiedImageError:
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
def build_picture_chunks(
|
def build_picture_chunks(
|
||||||
docling_document: "DoclingDocument",
|
docling_document: "DoclingDocument",
|
||||||
*,
|
*,
|
||||||
document_id: str | None = None,
|
document_id: str | None = None,
|
||||||
existing_picture_data: dict[str, bytes] | None = None,
|
existing_picture_data: dict[str, bytes] | None = None,
|
||||||
|
min_picture_size: int = 0,
|
||||||
) -> list[Chunk]:
|
) -> list[Chunk]:
|
||||||
"""Emit one synthetic ``Chunk`` per ``PictureItem`` with available bytes.
|
"""Emit one synthetic ``Chunk`` per distinct ``PictureItem`` with available
|
||||||
|
bytes.
|
||||||
|
|
||||||
Bytes come from ``picture.image.uri`` (live data URI on a freshly-converted
|
Bytes come from ``picture.image.uri`` (live data URI on a freshly-converted
|
||||||
docling) or from ``existing_picture_data`` keyed by ``self_ref`` (snapshot
|
docling) or from ``existing_picture_data`` keyed by ``self_ref`` (snapshot
|
||||||
taken before a delete-and-re-extract cycle, when the live docling has had
|
taken before a delete-and-re-extract cycle, when the live docling has had
|
||||||
its picture URIs stripped). Pictures with no available bytes are skipped.
|
its picture URIs stripped). Pictures with no available bytes are skipped.
|
||||||
|
|
||||||
|
Pictures whose bytes were already seen in this document are skipped — the
|
||||||
|
first occurrence carries the chunk, so a watermark repeated on every page
|
||||||
|
embeds once. Pictures whose smaller side is under ``min_picture_size``
|
||||||
|
pixels are skipped entirely (``0`` disables the size filter; pictures
|
||||||
|
whose size can't be determined are kept).
|
||||||
|
|
||||||
The bytes ride on ``Chunk._picture_data`` (a PrivateAttr — not serialized)
|
The bytes ride on ``Chunk._picture_data`` (a PrivateAttr — not serialized)
|
||||||
so ``embed_chunks`` can route them through ``embed_image``. The
|
so ``embed_chunks`` can route them through ``embed_image``. The
|
||||||
``order`` field is left at its default (0); the caller (``chunk()``)
|
``order`` field is left at its default (0); the caller (``chunk()``)
|
||||||
|
|
@ -262,6 +292,7 @@ def build_picture_chunks(
|
||||||
)
|
)
|
||||||
|
|
||||||
existing = existing_picture_data or {}
|
existing = existing_picture_data or {}
|
||||||
|
seen: set[bytes] = set()
|
||||||
chunks: list[Chunk] = []
|
chunks: list[Chunk] = []
|
||||||
|
|
||||||
for picture in docling_document.pictures:
|
for picture in docling_document.pictures:
|
||||||
|
|
@ -271,6 +302,15 @@ def build_picture_chunks(
|
||||||
if picture_data is None:
|
if picture_data is None:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
|
if picture_data in seen:
|
||||||
|
continue
|
||||||
|
seen.add(picture_data)
|
||||||
|
|
||||||
|
if min_picture_size > 0:
|
||||||
|
side = _min_picture_side(picture, picture_data)
|
||||||
|
if side is not None and side < min_picture_size:
|
||||||
|
continue
|
||||||
|
|
||||||
text = extract_item_text(picture, docling_document) or ""
|
text = extract_item_text(picture, docling_document) or ""
|
||||||
|
|
||||||
page_numbers: list[int] = []
|
page_numbers: list[int] = []
|
||||||
|
|
|
||||||
|
|
@ -206,6 +206,11 @@ class ProcessingConfig(BaseModel):
|
||||||
- ``"image"``: docling generates picture images and stores them in
|
- ``"image"``: docling generates picture images and stores them in
|
||||||
``document_items.picture_data``; no VLM runs at ingest.
|
``document_items.picture_data``; no VLM runs at ingest.
|
||||||
"""
|
"""
|
||||||
|
min_picture_size: int = 64
|
||||||
|
"""Minimum pixel size (smaller side) for a picture to become a picture
|
||||||
|
chunk. Smaller pictures — icons, bullets, decorative graphics — are not
|
||||||
|
embedded or indexed; their bytes stay in ``document_items`` for context
|
||||||
|
expansion. ``0`` keeps all pictures."""
|
||||||
extract_pdf_attachments: bool = True
|
extract_pdf_attachments: bool = True
|
||||||
"""When a PDF carries `/EmbeddedFiles`, ingest each attachment as a separate
|
"""When a PDF carries `/EmbeddedFiles`, ingest each attachment as a separate
|
||||||
Document linked back to the wrapper via ``metadata.parent_uri``. Cap depth
|
Document linked back to the wrapper via ``metadata.parent_uri``. Cap depth
|
||||||
|
|
|
||||||
|
|
@ -73,6 +73,10 @@ class EmbedderWrapper:
|
||||||
"embeddings.model.multimodal: true on a vllm, voyageai, or cohere model."
|
"embeddings.model.multimodal: true on a vllm, voyageai, or cohere model."
|
||||||
)
|
)
|
||||||
|
|
||||||
|
async def aclose(self) -> None:
|
||||||
|
"""Release resources held by the embedder. No-op by default;
|
||||||
|
embedders that own an HTTP client override this."""
|
||||||
|
|
||||||
|
|
||||||
def _to_data_uri(image: "bytes | PILImage.Image") -> str:
|
def _to_data_uri(image: "bytes | PILImage.Image") -> str:
|
||||||
"""Render an image as a ``data:image/png;base64,...`` URI."""
|
"""Render an image as a ``data:image/png;base64,...`` URI."""
|
||||||
|
|
|
||||||
|
|
@ -38,6 +38,10 @@ class VLLMMultimodalEmbedder(EmbedderWrapper):
|
||||||
self._base_url = base_url.rstrip("/")
|
self._base_url = base_url.rstrip("/")
|
||||||
self._api_key = api_key
|
self._api_key = api_key
|
||||||
self._timeout = timeout
|
self._timeout = timeout
|
||||||
|
# One client reused across every request so the connection (and its
|
||||||
|
# name resolution) is established once and kept alive, rather than
|
||||||
|
# rebuilt per call.
|
||||||
|
self._client = httpx.AsyncClient(timeout=timeout)
|
||||||
|
|
||||||
def _headers(self) -> dict[str, str]:
|
def _headers(self) -> dict[str, str]:
|
||||||
headers = {"Content-Type": "application/json"}
|
headers = {"Content-Type": "application/json"}
|
||||||
|
|
@ -45,16 +49,18 @@ class VLLMMultimodalEmbedder(EmbedderWrapper):
|
||||||
headers["Authorization"] = f"Bearer {self._api_key}"
|
headers["Authorization"] = f"Bearer {self._api_key}"
|
||||||
return headers
|
return headers
|
||||||
|
|
||||||
|
async def aclose(self) -> None:
|
||||||
|
await self._client.aclose()
|
||||||
|
|
||||||
async def _post(self, body: dict[str, Any]) -> list[list[float]]:
|
async def _post(self, body: dict[str, Any]) -> list[list[float]]:
|
||||||
try:
|
try:
|
||||||
async with httpx.AsyncClient(timeout=self._timeout) as client:
|
response = await self._client.post(
|
||||||
response = await client.post(
|
f"{self._base_url}/embeddings",
|
||||||
f"{self._base_url}/embeddings",
|
json=body,
|
||||||
json=body,
|
headers=self._headers(),
|
||||||
headers=self._headers(),
|
)
|
||||||
)
|
response.raise_for_status()
|
||||||
response.raise_for_status()
|
payload = response.json()
|
||||||
payload = response.json()
|
|
||||||
except httpx.ConnectError as e:
|
except httpx.ConnectError as e:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"Could not connect to vLLM at {self._base_url}. "
|
f"Could not connect to vLLM at {self._base_url}. "
|
||||||
|
|
|
||||||
|
|
@ -18,3 +18,7 @@ class RerankerBase:
|
||||||
raise NotImplementedError(
|
raise NotImplementedError(
|
||||||
"Reranker is an abstract class. Please implement the _rerank method in a subclass."
|
"Reranker is an abstract class. Please implement the _rerank method in a subclass."
|
||||||
)
|
)
|
||||||
|
|
||||||
|
async def aclose(self) -> None:
|
||||||
|
"""Release resources held by the reranker. No-op by default;
|
||||||
|
rerankers that own an HTTP client override this."""
|
||||||
|
|
|
||||||
|
|
@ -14,34 +14,38 @@ class JinaReranker(RerankerBase):
|
||||||
self._api_key = os.environ.get("JINA_API_KEY")
|
self._api_key = os.environ.get("JINA_API_KEY")
|
||||||
if not self._api_key:
|
if not self._api_key:
|
||||||
raise ValueError("JINA_API_KEY environment variable required")
|
raise ValueError("JINA_API_KEY environment variable required")
|
||||||
|
# One client reused across rerank calls (connection kept alive).
|
||||||
|
self._client = httpx.AsyncClient()
|
||||||
|
|
||||||
|
async def aclose(self) -> None:
|
||||||
|
await self._client.aclose()
|
||||||
|
|
||||||
async def _rerank(
|
async def _rerank(
|
||||||
self, query: str, chunks: list[Chunk], top_n: int = 10
|
self, query: str, chunks: list[Chunk], top_n: int = 10
|
||||||
) -> list[tuple[Chunk, float]]:
|
) -> list[tuple[Chunk, float]]:
|
||||||
documents = [chunk.content for chunk in chunks]
|
documents = [chunk.content for chunk in chunks]
|
||||||
|
|
||||||
async with httpx.AsyncClient() as client:
|
response = await self._client.post(
|
||||||
response = await client.post(
|
"https://api.jina.ai/v1/rerank",
|
||||||
"https://api.jina.ai/v1/rerank",
|
json={
|
||||||
json={
|
"model": self._model,
|
||||||
"model": self._model,
|
"query": query,
|
||||||
"query": query,
|
"documents": documents,
|
||||||
"documents": documents,
|
"top_n": top_n,
|
||||||
"top_n": top_n,
|
},
|
||||||
},
|
headers={
|
||||||
headers={
|
"Authorization": f"Bearer {self._api_key}",
|
||||||
"Authorization": f"Bearer {self._api_key}",
|
"Content-Type": "application/json",
|
||||||
"Content-Type": "application/json",
|
},
|
||||||
},
|
)
|
||||||
)
|
response.raise_for_status()
|
||||||
response.raise_for_status()
|
|
||||||
|
|
||||||
result = response.json()
|
result = response.json()
|
||||||
|
|
||||||
scored_chunks = []
|
scored_chunks = []
|
||||||
for item in result.get("results", []):
|
for item in result.get("results", []):
|
||||||
index = item["index"]
|
index = item["index"]
|
||||||
score = item["relevance_score"]
|
score = item["relevance_score"]
|
||||||
scored_chunks.append((chunks[index], score))
|
scored_chunks.append((chunks[index], score))
|
||||||
|
|
||||||
return scored_chunks
|
return scored_chunks
|
||||||
|
|
|
||||||
|
|
@ -4,10 +4,15 @@ from haiku.rag.reranking.base import RerankerBase
|
||||||
from haiku.rag.store.models.chunk import Chunk
|
from haiku.rag.store.models.chunk import Chunk
|
||||||
|
|
||||||
|
|
||||||
class VLLMReranker(RerankerBase): # pragma: no cover
|
class VLLMReranker(RerankerBase):
|
||||||
def __init__(self, model: str, base_url: str):
|
def __init__(self, model: str, base_url: str):
|
||||||
self._model = model
|
self._model = model
|
||||||
self._base_url = base_url
|
self._base_url = base_url
|
||||||
|
# One client reused across rerank calls (connection kept alive).
|
||||||
|
self._client = httpx.AsyncClient()
|
||||||
|
|
||||||
|
async def aclose(self) -> None:
|
||||||
|
await self._client.aclose()
|
||||||
|
|
||||||
async def _rerank(
|
async def _rerank(
|
||||||
self, query: str, chunks: list[Chunk], top_n: int = 10
|
self, query: str, chunks: list[Chunk], top_n: int = 10
|
||||||
|
|
@ -15,26 +20,25 @@ class VLLMReranker(RerankerBase): # pragma: no cover
|
||||||
# Prepare documents for reranking
|
# Prepare documents for reranking
|
||||||
documents = [chunk.content for chunk in chunks]
|
documents = [chunk.content for chunk in chunks]
|
||||||
|
|
||||||
async with httpx.AsyncClient() as client:
|
response = await self._client.post(
|
||||||
response = await client.post(
|
f"{self._base_url}/v1/rerank",
|
||||||
f"{self._base_url}/v1/rerank",
|
json={"model": self._model, "query": query, "documents": documents},
|
||||||
json={"model": self._model, "query": query, "documents": documents},
|
headers={
|
||||||
headers={
|
"accept": "application/json",
|
||||||
"accept": "application/json",
|
"Content-Type": "application/json",
|
||||||
"Content-Type": "application/json",
|
},
|
||||||
},
|
)
|
||||||
)
|
response.raise_for_status()
|
||||||
response.raise_for_status()
|
|
||||||
|
|
||||||
result = response.json()
|
result = response.json()
|
||||||
|
|
||||||
# Extract scores and pair with chunks
|
# Extract scores and pair with chunks
|
||||||
scored_chunks = []
|
scored_chunks = []
|
||||||
for item in result.get("results", []):
|
for item in result.get("results", []):
|
||||||
index = item["index"]
|
index = item["index"]
|
||||||
score = item["relevance_score"]
|
score = item["relevance_score"]
|
||||||
scored_chunks.append((chunks[index], score))
|
scored_chunks.append((chunks[index], score))
|
||||||
|
|
||||||
# Sort by score (descending) and return top_n
|
# Sort by score (descending) and return top_n
|
||||||
scored_chunks.sort(key=lambda x: x[1], reverse=True)
|
scored_chunks.sort(key=lambda x: x[1], reverse=True)
|
||||||
return scored_chunks[:top_n]
|
return scored_chunks[:top_n]
|
||||||
|
|
|
||||||
|
|
@ -154,6 +154,29 @@ async def test_client_embedder_is_store_embedder(temp_db_path):
|
||||||
assert client.embedder is client.store.embedder
|
assert client.embedder is client.store.embedder
|
||||||
|
|
||||||
|
|
||||||
|
async def test_client_aexit_closes_embedder(temp_db_path, monkeypatch):
|
||||||
|
"""__aexit__ releases the embedder's HTTP resources. A never-accessed
|
||||||
|
reranker is not materialized just to be closed; an accessed-but-None
|
||||||
|
reranker (reranking disabled) is handled."""
|
||||||
|
from haiku.rag.client import HaikuRAG
|
||||||
|
|
||||||
|
closed = []
|
||||||
|
|
||||||
|
async def _record():
|
||||||
|
closed.append(True)
|
||||||
|
|
||||||
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||||
|
monkeypatch.setattr(client.embedder, "aclose", _record)
|
||||||
|
assert client.reranker is None # default config: reranking disabled
|
||||||
|
|
||||||
|
assert closed == [True]
|
||||||
|
|
||||||
|
async with HaikuRAG(temp_db_path) as client:
|
||||||
|
monkeypatch.setattr(client.embedder, "aclose", _record)
|
||||||
|
|
||||||
|
assert "reranker" not in client.__dict__
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.vcr()
|
@pytest.mark.vcr()
|
||||||
async def test_embed_chunks_basic(allow_model_requests):
|
async def test_embed_chunks_basic(allow_model_requests):
|
||||||
"""Test that embed_chunks generates embeddings for chunks."""
|
"""Test that embed_chunks generates embeddings for chunks."""
|
||||||
|
|
@ -408,6 +431,44 @@ async def test_vllm_embed_image_request_shape(monkeypatch):
|
||||||
assert url.startswith("data:image/png;base64,")
|
assert url.startswith("data:image/png;base64,")
|
||||||
|
|
||||||
|
|
||||||
|
async def test_vllm_reuses_pooled_client(monkeypatch):
|
||||||
|
"""The embedder builds one httpx client and reuses it across requests
|
||||||
|
instead of opening a fresh connection per call; aclose releases it."""
|
||||||
|
from haiku.rag.embeddings.vllm import VLLMMultimodalEmbedder
|
||||||
|
|
||||||
|
stats = {"constructed": 0, "closed": 0}
|
||||||
|
|
||||||
|
class FakeResponse:
|
||||||
|
def raise_for_status(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def json(self):
|
||||||
|
return {"data": [{"embedding": [0.1, 0.2]}]}
|
||||||
|
|
||||||
|
class FakeAsyncClient:
|
||||||
|
def __init__(self, *args, **kwargs):
|
||||||
|
stats["constructed"] += 1
|
||||||
|
|
||||||
|
async def post(self, url, json, headers):
|
||||||
|
return FakeResponse()
|
||||||
|
|
||||||
|
async def aclose(self):
|
||||||
|
stats["closed"] += 1
|
||||||
|
|
||||||
|
monkeypatch.setattr("httpx.AsyncClient", FakeAsyncClient)
|
||||||
|
|
||||||
|
embedder = VLLMMultimodalEmbedder(
|
||||||
|
model_name="x", vector_dim=2, base_url="http://localhost:8000/v1"
|
||||||
|
)
|
||||||
|
await embedder.embed_query("one")
|
||||||
|
await embedder.embed_query("two")
|
||||||
|
await embedder.embed_documents(["three", "four"])
|
||||||
|
|
||||||
|
assert stats["constructed"] == 1
|
||||||
|
await embedder.aclose()
|
||||||
|
assert stats["closed"] == 1
|
||||||
|
|
||||||
|
|
||||||
async def test_vllm_supports_images_flag():
|
async def test_vllm_supports_images_flag():
|
||||||
from haiku.rag.embeddings.vllm import VLLMMultimodalEmbedder
|
from haiku.rag.embeddings.vllm import VLLMMultimodalEmbedder
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -281,6 +281,8 @@ async def test_embed_only_preserves_picture_vectors(temp_db_path, monkeypatch):
|
||||||
model=EmbeddingModelConfig(provider="ollama", name="stub", vector_dim=4)
|
model=EmbeddingModelConfig(provider="ollama", name="stub", vector_dim=4)
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
# The fixture picture is 8x8; disable the size filter so it still chunks.
|
||||||
|
config.processing.min_picture_size = 0
|
||||||
|
|
||||||
docling_doc = _docling_doc_with_picture()
|
docling_doc = _docling_doc_with_picture()
|
||||||
|
|
||||||
|
|
@ -514,6 +516,124 @@ def test_build_picture_chunks_skips_pictures_without_bytes():
|
||||||
assert chunks == []
|
assert chunks == []
|
||||||
|
|
||||||
|
|
||||||
|
def _doc_with_picture_images(*images):
|
||||||
|
"""DoclingDocument with one paragraph and one PictureItem per PIL image."""
|
||||||
|
from docling_core.types.doc.document import DoclingDocument, ImageRef
|
||||||
|
from docling_core.types.doc.labels import DocItemLabel
|
||||||
|
|
||||||
|
doc = DoclingDocument(name="pics")
|
||||||
|
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Hello world")
|
||||||
|
for img in images:
|
||||||
|
doc.add_picture(image=ImageRef.from_pil(img, dpi=72))
|
||||||
|
return doc
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_picture_chunks_dedupes_identical_bytes():
|
||||||
|
"""Identical picture bytes within a document produce one chunk — the
|
||||||
|
first occurrence. A watermark repeated on every page embeds once."""
|
||||||
|
from PIL import Image as PILImageModule
|
||||||
|
|
||||||
|
from haiku.rag.client.processing import build_picture_chunks
|
||||||
|
|
||||||
|
red = PILImageModule.new("RGB", (100, 100), "red")
|
||||||
|
blue = PILImageModule.new("RGB", (100, 100), "blue")
|
||||||
|
doc = _doc_with_picture_images(red, red, blue, red)
|
||||||
|
|
||||||
|
chunks = build_picture_chunks(doc, document_id="doc-1")
|
||||||
|
|
||||||
|
refs = [c.metadata["doc_item_refs"][0] for c in chunks]
|
||||||
|
assert refs == ["#/pictures/0", "#/pictures/2"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_picture_chunks_skips_small_pictures():
|
||||||
|
"""Pictures whose smaller side is under min_picture_size are not chunked."""
|
||||||
|
from PIL import Image as PILImageModule
|
||||||
|
|
||||||
|
from haiku.rag.client.processing import build_picture_chunks
|
||||||
|
|
||||||
|
icon = PILImageModule.new("RGB", (16, 16), "red")
|
||||||
|
figure = PILImageModule.new("RGB", (100, 100), "blue")
|
||||||
|
banner = PILImageModule.new("RGB", (200, 16), "green")
|
||||||
|
doc = _doc_with_picture_images(icon, figure, banner)
|
||||||
|
|
||||||
|
chunks = build_picture_chunks(doc, document_id="doc-1", min_picture_size=64)
|
||||||
|
|
||||||
|
assert [c.metadata["doc_item_refs"][0] for c in chunks] == ["#/pictures/1"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_picture_chunks_measures_snapshot_bytes():
|
||||||
|
"""Rebuild path: picture.image is None, so size comes from a PIL header
|
||||||
|
read of the snapshot bytes — existing DBs shed small pictures on rebuild."""
|
||||||
|
import io
|
||||||
|
|
||||||
|
from PIL import Image as PILImageModule
|
||||||
|
|
||||||
|
from haiku.rag.client.processing import build_picture_chunks
|
||||||
|
|
||||||
|
icon_png = io.BytesIO()
|
||||||
|
PILImageModule.new("RGB", (16, 16), "red").save(icon_png, format="PNG")
|
||||||
|
doc = _doc_with_picture_images(PILImageModule.new("RGB", (16, 16), "red"))
|
||||||
|
for picture in doc.pictures:
|
||||||
|
picture.image = None
|
||||||
|
|
||||||
|
chunks = build_picture_chunks(
|
||||||
|
doc,
|
||||||
|
document_id="doc-1",
|
||||||
|
existing_picture_data={"#/pictures/0": icon_png.getvalue()},
|
||||||
|
min_picture_size=64,
|
||||||
|
)
|
||||||
|
assert chunks == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_picture_chunks_keeps_unmeasurable_bytes():
|
||||||
|
"""Bytes PIL can't parse are kept — the filter only drops what it can
|
||||||
|
measure."""
|
||||||
|
from haiku.rag.client.processing import build_picture_chunks
|
||||||
|
from tests.store.test_document_items import _docling_doc_with_picture
|
||||||
|
|
||||||
|
doc = _docling_doc_with_picture()
|
||||||
|
for picture in doc.pictures:
|
||||||
|
picture.image = None
|
||||||
|
|
||||||
|
chunks = build_picture_chunks(
|
||||||
|
doc,
|
||||||
|
document_id="doc-1",
|
||||||
|
existing_picture_data={"#/pictures/0": b"not-an-image"},
|
||||||
|
min_picture_size=64,
|
||||||
|
)
|
||||||
|
assert len(chunks) == 1
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_chunk_filters_small_pictures_by_config(monkeypatch):
|
||||||
|
"""``chunk()`` applies ``processing.min_picture_size`` — with the default
|
||||||
|
config, icon-sized pictures don't become picture chunks."""
|
||||||
|
from haiku.rag.client.processing import chunk
|
||||||
|
from haiku.rag.embeddings import EmbedderWrapper
|
||||||
|
from haiku.rag.store.models.chunk import Chunk
|
||||||
|
from tests.store.test_document_items import _docling_doc_with_picture
|
||||||
|
|
||||||
|
class StubMultimodalEmbedder(EmbedderWrapper):
|
||||||
|
supports_images = True
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__(embedder=None, vector_dim=4)
|
||||||
|
|
||||||
|
class StubChunker:
|
||||||
|
async def chunk(self, document):
|
||||||
|
return [Chunk(content="text", metadata={"doc_item_refs": ["#/texts/0"]})]
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"haiku.rag.chunkers.get_chunker", lambda *a, **kw: StubChunker()
|
||||||
|
)
|
||||||
|
|
||||||
|
doc = _docling_doc_with_picture() # 8x8 picture, below the 64px default
|
||||||
|
chunks = await chunk(AppConfig(), doc, embedder=StubMultimodalEmbedder())
|
||||||
|
|
||||||
|
assert [c.content for c in chunks] == ["text"]
|
||||||
|
assert not any("picture" in (c.metadata or {}).get("labels", []) for c in chunks)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_chunk_interleaves_picture_in_structural_order(monkeypatch):
|
async def test_chunk_interleaves_picture_in_structural_order(monkeypatch):
|
||||||
"""``chunk()`` merges text and picture chunks by their first
|
"""``chunk()`` merges text and picture chunks by their first
|
||||||
|
|
@ -554,7 +674,7 @@ async def test_chunk_interleaves_picture_in_structural_order(monkeypatch):
|
||||||
from docling_core.types.doc.labels import DocItemLabel
|
from docling_core.types.doc.labels import DocItemLabel
|
||||||
from PIL import Image as PILImageModule
|
from PIL import Image as PILImageModule
|
||||||
|
|
||||||
img = PILImageModule.new("RGB", (8, 8), "blue")
|
img = PILImageModule.new("RGB", (64, 64), "blue")
|
||||||
doc = DoclingDocument(name="ordered")
|
doc = DoclingDocument(name="ordered")
|
||||||
doc.add_text(label=DocItemLabel.PARAGRAPH, text="A")
|
doc.add_text(label=DocItemLabel.PARAGRAPH, text="A")
|
||||||
doc.add_text(label=DocItemLabel.PARAGRAPH, text="B")
|
doc.add_text(label=DocItemLabel.PARAGRAPH, text="B")
|
||||||
|
|
@ -676,6 +796,8 @@ async def test_ingest_emits_picture_chunks_with_multimodal_embedder(
|
||||||
model=EmbeddingModelConfig(provider="ollama", name="stub", vector_dim=4)
|
model=EmbeddingModelConfig(provider="ollama", name="stub", vector_dim=4)
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
# The fixture picture is 8x8; disable the size filter so it still chunks.
|
||||||
|
config.processing.min_picture_size = 0
|
||||||
|
|
||||||
async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
|
async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
|
||||||
chunks = await rag.chunk(docling_doc)
|
chunks = await rag.chunk(docling_doc)
|
||||||
|
|
|
||||||
|
|
@ -218,7 +218,7 @@ def test_merge_picture_chunks_no_pictures_returns_text_chunks():
|
||||||
doc = _doc_without_pictures()
|
doc = _doc_without_pictures()
|
||||||
text_chunks = [Chunk(content="a"), Chunk(content="b")]
|
text_chunks = [Chunk(content="a"), Chunk(content="b")]
|
||||||
|
|
||||||
result = _merge_picture_chunks(doc, text_chunks, None, None)
|
result = _merge_picture_chunks(doc, text_chunks, None, None, 0)
|
||||||
|
|
||||||
assert result is text_chunks
|
assert result is text_chunks
|
||||||
assert [c.order for c in result] == [0, 1]
|
assert [c.order for c in result] == [0, 1]
|
||||||
|
|
|
||||||
|
|
@ -274,6 +274,84 @@ class TestGetReranker:
|
||||||
assert getattr(result, attr) == value
|
assert getattr(result, attr) == value
|
||||||
|
|
||||||
|
|
||||||
|
class _PoolStats:
|
||||||
|
"""Fake httpx.AsyncClient factory counting constructions and closes."""
|
||||||
|
|
||||||
|
def __init__(self, response_json):
|
||||||
|
self.constructed = 0
|
||||||
|
self.closed = 0
|
||||||
|
stats = self
|
||||||
|
|
||||||
|
class FakeResponse:
|
||||||
|
def raise_for_status(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def json(self):
|
||||||
|
return response_json
|
||||||
|
|
||||||
|
class FakeAsyncClient:
|
||||||
|
def __init__(self, *args, **kwargs):
|
||||||
|
stats.constructed += 1
|
||||||
|
|
||||||
|
async def post(self, url, json, headers):
|
||||||
|
return FakeResponse()
|
||||||
|
|
||||||
|
async def aclose(self):
|
||||||
|
stats.closed += 1
|
||||||
|
|
||||||
|
self.client_class = FakeAsyncClient
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_vllm_reranker_reuses_pooled_client(monkeypatch):
|
||||||
|
"""One httpx client is built and reused across rerank calls; aclose
|
||||||
|
releases it."""
|
||||||
|
from haiku.rag.reranking.vllm import VLLMReranker
|
||||||
|
|
||||||
|
stats = _PoolStats({"results": [{"index": 0, "relevance_score": 0.9}]})
|
||||||
|
monkeypatch.setattr("httpx.AsyncClient", stats.client_class)
|
||||||
|
|
||||||
|
reranker = VLLMReranker(model="m", base_url="http://localhost:8000")
|
||||||
|
docs = [Chunk(content="a", order=0)]
|
||||||
|
await reranker.rerank("q", docs)
|
||||||
|
await reranker.rerank("q", docs)
|
||||||
|
|
||||||
|
assert stats.constructed == 1
|
||||||
|
await reranker.aclose()
|
||||||
|
assert stats.closed == 1
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_jina_reranker_reuses_pooled_client(monkeypatch):
|
||||||
|
"""One httpx client is built and reused across rerank calls; aclose
|
||||||
|
releases it."""
|
||||||
|
monkeypatch.setenv("JINA_API_KEY", "test-api-key")
|
||||||
|
from haiku.rag.reranking.jina import JinaReranker
|
||||||
|
|
||||||
|
stats = _PoolStats({"results": [{"index": 0, "relevance_score": 0.9}]})
|
||||||
|
monkeypatch.setattr("httpx.AsyncClient", stats.client_class)
|
||||||
|
|
||||||
|
reranker = JinaReranker("jina-reranker-v3")
|
||||||
|
docs = [Chunk(content="a", order=0)]
|
||||||
|
await reranker.rerank("q", docs)
|
||||||
|
await reranker.rerank("q", docs)
|
||||||
|
|
||||||
|
assert stats.constructed == 1
|
||||||
|
await reranker.aclose()
|
||||||
|
assert stats.closed == 1
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_reranker_base_aclose_is_noop():
|
||||||
|
"""Base aclose exists so client teardown can close any reranker."""
|
||||||
|
|
||||||
|
class Custom(RerankerBase):
|
||||||
|
async def _rerank(self, query, chunks, top_n=10):
|
||||||
|
return []
|
||||||
|
|
||||||
|
await Custom().aclose() # must not raise
|
||||||
|
|
||||||
|
|
||||||
def test_jina_reranker_missing_api_key(monkeypatch):
|
def test_jina_reranker_missing_api_key(monkeypatch):
|
||||||
monkeypatch.delenv("JINA_API_KEY", raising=False)
|
monkeypatch.delenv("JINA_API_KEY", raising=False)
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue