diff --git a/CHANGELOG.md b/CHANGELOG.md index a121d02e..7e8786ad 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,15 @@ # Changelog ## [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 ### Added diff --git a/docs/configuration/processing.md b/docs/configuration/processing.md index 8e50e174..d59ce8aa 100644 --- a/docs/configuration/processing.md +++ b/docs/configuration/processing.md @@ -259,6 +259,8 @@ Per-image failures (404, timeout, oversized, unreadable) leave that picture as a | `description` | on | yes | yes | | `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). ```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 | | `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` | 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: diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 38529774..910ef6a0 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -143,6 +143,16 @@ class HaikuRAG: async def __aexit__(self, exc_type, exc_val, exc_tb): # noqa: ARG002 """Async context manager exit.""" 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() return False diff --git a/haiku_rag_slim/haiku/rag/client/processing.py b/haiku_rag_slim/haiku/rag/client/processing.py index 198f7736..87f59ff3 100644 --- a/haiku_rag_slim/haiku/rag/client/processing.py +++ b/haiku_rag_slim/haiku/rag/client/processing.py @@ -1,4 +1,5 @@ import asyncio +import io import logging import tempfile 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 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 @@ -173,11 +174,13 @@ def _merge_picture_chunks( text_chunks: list[Chunk], document_id: str | None, existing_picture_data: dict[str, bytes] | None, + min_picture_size: int, ) -> list[Chunk]: picture_chunks = build_picture_chunks( docling_document, document_id=document_id, existing_picture_data=existing_picture_data, + min_picture_size=min_picture_size, ) if not picture_chunks: @@ -235,22 +238,49 @@ async def chunk( text_chunks, document_id, 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( docling_document: "DoclingDocument", *, document_id: str | None = None, existing_picture_data: dict[str, bytes] | None = None, + min_picture_size: int = 0, ) -> 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 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 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) so ``embed_chunks`` can route them through ``embed_image``. The ``order`` field is left at its default (0); the caller (``chunk()``) @@ -262,6 +292,7 @@ def build_picture_chunks( ) existing = existing_picture_data or {} + seen: set[bytes] = set() chunks: list[Chunk] = [] for picture in docling_document.pictures: @@ -271,6 +302,15 @@ def build_picture_chunks( if picture_data is None: 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 "" page_numbers: list[int] = [] diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index b9f2734c..c2e79804 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -206,6 +206,11 @@ class ProcessingConfig(BaseModel): - ``"image"``: docling generates picture images and stores them in ``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 """When a PDF carries `/EmbeddedFiles`, ingest each attachment as a separate Document linked back to the wrapper via ``metadata.parent_uri``. Cap depth diff --git a/haiku_rag_slim/haiku/rag/embeddings/__init__.py b/haiku_rag_slim/haiku/rag/embeddings/__init__.py index f4ea1af7..2b548745 100644 --- a/haiku_rag_slim/haiku/rag/embeddings/__init__.py +++ b/haiku_rag_slim/haiku/rag/embeddings/__init__.py @@ -73,6 +73,10 @@ class EmbedderWrapper: "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: """Render an image as a ``data:image/png;base64,...`` URI.""" diff --git a/haiku_rag_slim/haiku/rag/embeddings/vllm.py b/haiku_rag_slim/haiku/rag/embeddings/vllm.py index fd49eab8..1da1555e 100644 --- a/haiku_rag_slim/haiku/rag/embeddings/vllm.py +++ b/haiku_rag_slim/haiku/rag/embeddings/vllm.py @@ -38,6 +38,10 @@ class VLLMMultimodalEmbedder(EmbedderWrapper): self._base_url = base_url.rstrip("/") self._api_key = api_key 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]: headers = {"Content-Type": "application/json"} @@ -45,16 +49,18 @@ class VLLMMultimodalEmbedder(EmbedderWrapper): headers["Authorization"] = f"Bearer {self._api_key}" return headers + async def aclose(self) -> None: + await self._client.aclose() + async def _post(self, body: dict[str, Any]) -> list[list[float]]: try: - async with httpx.AsyncClient(timeout=self._timeout) as client: - response = await client.post( - f"{self._base_url}/embeddings", - json=body, - headers=self._headers(), - ) - response.raise_for_status() - payload = response.json() + response = await self._client.post( + f"{self._base_url}/embeddings", + json=body, + headers=self._headers(), + ) + response.raise_for_status() + payload = response.json() except httpx.ConnectError as e: raise ValueError( f"Could not connect to vLLM at {self._base_url}. " diff --git a/haiku_rag_slim/haiku/rag/reranking/base.py b/haiku_rag_slim/haiku/rag/reranking/base.py index 6c31d986..703ed912 100644 --- a/haiku_rag_slim/haiku/rag/reranking/base.py +++ b/haiku_rag_slim/haiku/rag/reranking/base.py @@ -18,3 +18,7 @@ class RerankerBase: raise NotImplementedError( "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.""" diff --git a/haiku_rag_slim/haiku/rag/reranking/jina.py b/haiku_rag_slim/haiku/rag/reranking/jina.py index 0144edb3..ea5bff2d 100644 --- a/haiku_rag_slim/haiku/rag/reranking/jina.py +++ b/haiku_rag_slim/haiku/rag/reranking/jina.py @@ -14,34 +14,38 @@ class JinaReranker(RerankerBase): self._api_key = os.environ.get("JINA_API_KEY") if not self._api_key: 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( self, query: str, chunks: list[Chunk], top_n: int = 10 ) -> list[tuple[Chunk, float]]: documents = [chunk.content for chunk in chunks] - async with httpx.AsyncClient() as client: - response = await client.post( - "https://api.jina.ai/v1/rerank", - json={ - "model": self._model, - "query": query, - "documents": documents, - "top_n": top_n, - }, - headers={ - "Authorization": f"Bearer {self._api_key}", - "Content-Type": "application/json", - }, - ) - response.raise_for_status() + response = await self._client.post( + "https://api.jina.ai/v1/rerank", + json={ + "model": self._model, + "query": query, + "documents": documents, + "top_n": top_n, + }, + headers={ + "Authorization": f"Bearer {self._api_key}", + "Content-Type": "application/json", + }, + ) + response.raise_for_status() - result = response.json() + result = response.json() - scored_chunks = [] - for item in result.get("results", []): - index = item["index"] - score = item["relevance_score"] - scored_chunks.append((chunks[index], score)) + scored_chunks = [] + for item in result.get("results", []): + index = item["index"] + score = item["relevance_score"] + scored_chunks.append((chunks[index], score)) - return scored_chunks + return scored_chunks diff --git a/haiku_rag_slim/haiku/rag/reranking/vllm.py b/haiku_rag_slim/haiku/rag/reranking/vllm.py index 96271219..67909c8d 100644 --- a/haiku_rag_slim/haiku/rag/reranking/vllm.py +++ b/haiku_rag_slim/haiku/rag/reranking/vllm.py @@ -4,10 +4,15 @@ from haiku.rag.reranking.base import RerankerBase 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): self._model = model 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( self, query: str, chunks: list[Chunk], top_n: int = 10 @@ -15,26 +20,25 @@ class VLLMReranker(RerankerBase): # pragma: no cover # Prepare documents for reranking documents = [chunk.content for chunk in chunks] - async with httpx.AsyncClient() as client: - response = await client.post( - f"{self._base_url}/v1/rerank", - json={"model": self._model, "query": query, "documents": documents}, - headers={ - "accept": "application/json", - "Content-Type": "application/json", - }, - ) - response.raise_for_status() + response = await self._client.post( + f"{self._base_url}/v1/rerank", + json={"model": self._model, "query": query, "documents": documents}, + headers={ + "accept": "application/json", + "Content-Type": "application/json", + }, + ) + response.raise_for_status() - result = response.json() + result = response.json() - # Extract scores and pair with chunks - scored_chunks = [] - for item in result.get("results", []): - index = item["index"] - score = item["relevance_score"] - scored_chunks.append((chunks[index], score)) + # Extract scores and pair with chunks + scored_chunks = [] + for item in result.get("results", []): + index = item["index"] + score = item["relevance_score"] + scored_chunks.append((chunks[index], score)) - # Sort by score (descending) and return top_n - scored_chunks.sort(key=lambda x: x[1], reverse=True) - return scored_chunks[:top_n] + # Sort by score (descending) and return top_n + scored_chunks.sort(key=lambda x: x[1], reverse=True) + return scored_chunks[:top_n] diff --git a/tests/test_embedder.py b/tests/test_embedder.py index db8787ee..bcb82da1 100644 --- a/tests/test_embedder.py +++ b/tests/test_embedder.py @@ -154,6 +154,29 @@ async def test_client_embedder_is_store_embedder(temp_db_path): 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() async def test_embed_chunks_basic(allow_model_requests): """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,") +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(): from haiku.rag.embeddings.vllm import VLLMMultimodalEmbedder diff --git a/tests/test_picture_in_context.py b/tests/test_picture_in_context.py index 6c7ab2ea..7e2bd500 100644 --- a/tests/test_picture_in_context.py +++ b/tests/test_picture_in_context.py @@ -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) ) ) + # 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() @@ -514,6 +516,124 @@ def test_build_picture_chunks_skips_pictures_without_bytes(): 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 async def test_chunk_interleaves_picture_in_structural_order(monkeypatch): """``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 PIL import Image as PILImageModule - img = PILImageModule.new("RGB", (8, 8), "blue") + img = PILImageModule.new("RGB", (64, 64), "blue") doc = DoclingDocument(name="ordered") doc.add_text(label=DocItemLabel.PARAGRAPH, text="A") 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) ) ) + # 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: chunks = await rag.chunk(docling_doc) diff --git a/tests/test_processing.py b/tests/test_processing.py index 571f9581..8e7133d2 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -218,7 +218,7 @@ def test_merge_picture_chunks_no_pictures_returns_text_chunks(): doc = _doc_without_pictures() 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 [c.order for c in result] == [0, 1] diff --git a/tests/test_reranker.py b/tests/test_reranker.py index e67d3f3c..30dc6171 100644 --- a/tests/test_reranker.py +++ b/tests/test_reranker.py @@ -274,6 +274,84 @@ class TestGetReranker: 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): monkeypatch.delenv("JINA_API_KEY", raising=False)