diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 38529774..82c0910e 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -143,6 +143,13 @@ class HaikuRAG: async def __aexit__(self, exc_type, exc_val, exc_tb): # noqa: ARG002 """Async context manager exit.""" await self._await_vacuum_tasks() + # Release the embedder's pooled HTTP client (if any). Best-effort like + # the vacuum drain: __aexit__ runs during exception unwinding, so a + # raising close here must not mask the original exception. + try: + await self.embedder.aclose() + except Exception: + logger.debug("Embedder aclose failed on teardown", exc_info=True) self.close() return False diff --git a/haiku_rag_slim/haiku/rag/config/__init__.py b/haiku_rag_slim/haiku/rag/config/__init__.py index c2b79f0d..a992d3c8 100644 --- a/haiku_rag_slim/haiku/rag/config/__init__.py +++ b/haiku_rag_slim/haiku/rag/config/__init__.py @@ -12,6 +12,7 @@ from haiku.rag.config.models import ( CircuitBreakerConfig, ConversionOptions, DoclingServeConfig, + EmbeddingHTTPConfig, EmbeddingModelConfig, EmbeddingsConfig, FSSourceConfig, @@ -42,6 +43,7 @@ __all__ = [ "CircuitBreakerConfig", "ConversionOptions", "DoclingServeConfig", + "EmbeddingHTTPConfig", "EmbeddingModelConfig", "EmbeddingsConfig", "FSSourceConfig", diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index b9f2734c..008efd99 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -68,9 +68,29 @@ class LanceDBConfig(BaseModel): storage_options: dict[str, str] = Field(default_factory=dict) +class EmbeddingHTTPConfig(BaseModel): + """HTTP transport tuning for embedding requests to OpenAI-compatible + servers (vllm, openai, ollama). + + Applied to the pooled ``httpx.AsyncClient`` the embedder reuses across + every request, so a connection — and its name resolution — is established + once and kept warm. Providers whose SDKs manage their own transport + (voyageai, cohere, sentence-transformers) ignore these settings. + + ``max_connections`` bounds concurrent in-flight requests; keep it >= any + future embedding concurrency so the pool is never the limiter. + """ + + timeout_s: float = 60.0 + max_connections: int = 16 + max_keepalive_connections: int = 16 + keepalive_expiry_s: float = 300.0 + + class EmbeddingsConfig(BaseModel): model: EmbeddingModelConfig = Field(default_factory=EmbeddingModelConfig) batch_size: int = 512 + http: EmbeddingHTTPConfig = Field(default_factory=EmbeddingHTTPConfig) class RerankingConfig(BaseModel): diff --git a/haiku_rag_slim/haiku/rag/embeddings/__init__.py b/haiku_rag_slim/haiku/rag/embeddings/__init__.py index f4ea1af7..b61bb39d 100644 --- a/haiku_rag_slim/haiku/rag/embeddings/__init__.py +++ b/haiku_rag_slim/haiku/rag/embeddings/__init__.py @@ -1,7 +1,9 @@ import base64 +import hashlib import io from typing import TYPE_CHECKING, Any +import httpx from pydantic_ai.embeddings import Embedder from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel from pydantic_ai.providers.ollama import OllamaProvider @@ -12,10 +14,26 @@ from haiku.rag.config import AppConfig, Config if TYPE_CHECKING: from PIL import Image as PILImage - from haiku.rag.config.models import EmbeddingModelConfig + from haiku.rag.config.models import EmbeddingHTTPConfig, EmbeddingModelConfig from haiku.rag.store.models.chunk import Chunk +def build_http_client(http: "EmbeddingHTTPConfig") -> httpx.AsyncClient: + """A pooled ``httpx.AsyncClient`` configured from ``EmbeddingHTTPConfig``. + + Shared across an embedder's requests so a connection (and its name + resolution) is established once and kept warm rather than rebuilt per call. + """ + return httpx.AsyncClient( + timeout=httpx.Timeout(http.timeout_s), + limits=httpx.Limits( + max_connections=http.max_connections, + max_keepalive_connections=http.max_keepalive_connections, + keepalive_expiry=http.keepalive_expiry_s, + ), + ) + + ImageInput = "bytes | PILImage.Image" @@ -34,11 +52,17 @@ class EmbedderWrapper: embedder: Embedder | None, vector_dim: int, supports_images: bool | None = None, + *, + owned_http_client: "Any | None" = None, ): self._embedder = embedder self._vector_dim = vector_dim if supports_images is not None: self.supports_images = supports_images + # An httpx.AsyncClient this wrapper built and must close on teardown — + # e.g. the pooled client passed to an openai/ollama provider. None when + # the underlying SDK owns its own transport. + self._owned_http_client = owned_http_client @property def vector_dim(self) -> int: @@ -73,6 +97,15 @@ class EmbedderWrapper: "embeddings.model.multimodal: true on a vllm, voyageai, or cohere model." ) + async def aclose(self) -> None: + """Release any resources held by the embedder. Closes a pooled HTTP + client this wrapper owns (openai/ollama); a no-op otherwise. Lets + callers tear down uniformly regardless of embedder type. Subclasses + that own their own client (e.g. vLLM) override this.""" + if self._owned_http_client is not None: + await self._owned_http_client.aclose() + self._owned_http_client = None + def _to_data_uri(image: "bytes | PILImage.Image") -> str: """Render an image as a ``data:image/png;base64,...`` URI.""" @@ -152,8 +185,27 @@ async def embed_chunks( "embeddings.model.multimodal: true on a vllm, voyageai, or cohere " "model, or omit picture chunks." ) + # Identical image bytes embed to identical vectors, so embed each + # distinct image once and reuse the result for every chunk that shares + # it. A document that repeats one figure across many pages (header, + # watermark, logo) collapses from one request per occurrence to one + # per unique image. Keyed by a FIPS-safe content hash; order is + # preserved because we append one vector per chunk in chunk order. + embedding_cache: dict[bytes, list[float]] = {} for chunk in picture_chunks: - picture_embeddings.append(await embedder.embed_image(chunk._picture_data)) + data = chunk._picture_data + key = ( + hashlib.sha256(data, usedforsecurity=False).digest() + if isinstance(data, bytes | bytearray) + else None + ) + if key is not None and (cached := embedding_cache.get(key)) is not None: + picture_embeddings.append(cached) + continue + embedding = await embedder.embed_image(data) + if key is not None: + embedding_cache[key] = embedding + picture_embeddings.append(embedding) text_iter = iter(text_embeddings) picture_iter = iter(picture_embeddings) @@ -187,32 +239,36 @@ def get_embedder(config: AppConfig = Config) -> EmbedderWrapper: An embedder instance configured according to the config. """ embedding_model = config.embeddings.model + http = config.embeddings.http provider = embedding_model.provider model_name = embedding_model.name vector_dim = embedding_model.vector_dim if embedding_model.multimodal: - return _get_multimodal_embedder(embedding_model) + return _get_multimodal_embedder(embedding_model, http) if provider == "ollama": # Use model-level base_url if set, otherwise fall back to providers config base_url = embedding_model.base_url or config.providers.ollama.base_url if not base_url.rstrip("/").endswith("/v1"): base_url = base_url.rstrip("/") + "/v1" + client = build_http_client(http) model = OpenAIEmbeddingModel( model_name, - provider=OllamaProvider(base_url=base_url), + provider=OllamaProvider(base_url=base_url, http_client=client), ) - return EmbedderWrapper(Embedder(model), vector_dim) + return EmbedderWrapper(Embedder(model), vector_dim, owned_http_client=client) if provider == "openai": + client = build_http_client(http) + provider_kwargs: dict[str, Any] = {"http_client": client} if embedding_model.base_url: - model = OpenAIEmbeddingModel( - model_name, - provider=OpenAIProvider(base_url=embedding_model.base_url), - ) - return EmbedderWrapper(Embedder(model), vector_dim) - return EmbedderWrapper(Embedder(f"openai:{model_name}"), vector_dim) + provider_kwargs["base_url"] = embedding_model.base_url + model = OpenAIEmbeddingModel( + model_name, + provider=OpenAIProvider(**provider_kwargs), + ) + return EmbedderWrapper(Embedder(model), vector_dim, owned_http_client=client) if provider == "voyageai": return EmbedderWrapper(Embedder(f"voyageai:{model_name}"), vector_dim) @@ -230,7 +286,7 @@ def get_embedder(config: AppConfig = Config) -> EmbedderWrapper: base_url = _vllm_base_url(embedding_model.base_url) return VLLMMultimodalEmbedder( - model_name, vector_dim, base_url=base_url, supports_images=False + model_name, vector_dim, base_url=base_url, http=http, supports_images=False ) raise ValueError(f"Unsupported embedding provider: {provider}") @@ -245,6 +301,7 @@ def _vllm_base_url(base_url: str | None) -> str: def _get_multimodal_embedder( embedding_model: "EmbeddingModelConfig", + http: "EmbeddingHTTPConfig", ) -> EmbedderWrapper: """Build an image-capable embedder for providers that support multimodal. @@ -260,7 +317,7 @@ def _get_multimodal_embedder( base_url = _vllm_base_url(embedding_model.base_url) return VLLMMultimodalEmbedder( - model_name, vector_dim, base_url=base_url, supports_images=True + model_name, vector_dim, base_url=base_url, http=http, supports_images=True ) if provider == "voyageai": diff --git a/haiku_rag_slim/haiku/rag/embeddings/vllm.py b/haiku_rag_slim/haiku/rag/embeddings/vllm.py index fd49eab8..aed54761 100644 --- a/haiku_rag_slim/haiku/rag/embeddings/vllm.py +++ b/haiku_rag_slim/haiku/rag/embeddings/vllm.py @@ -11,11 +11,13 @@ Models like ``Qwen/Qwen3-VL-Embedding-8B`` and ``jinaai/jina-embeddings-v4`` ship with chat templates that map both shapes into a shared vector space. """ +import asyncio from typing import TYPE_CHECKING, Any import httpx -from haiku.rag.embeddings import EmbedderWrapper, _to_data_uri +from haiku.rag.config import EmbeddingHTTPConfig +from haiku.rag.embeddings import EmbedderWrapper, _to_data_uri, build_http_client if TYPE_CHECKING: from PIL import Image as PILImage @@ -28,7 +30,7 @@ class VLLMMultimodalEmbedder(EmbedderWrapper): vector_dim: int, base_url: str, api_key: str | None = None, - timeout: float = 60.0, + http: EmbeddingHTTPConfig | None = None, supports_images: bool = True, ): super().__init__( @@ -37,7 +39,16 @@ class VLLMMultimodalEmbedder(EmbedderWrapper): self._model_name = model_name self._base_url = base_url.rstrip("/") self._api_key = api_key - self._timeout = timeout + # Connection-pool + timeout tuning from config (defaults preserve the + # historical 60s / 16-connection behavior when constructed directly). + self._http = http or EmbeddingHTTPConfig() + # One pooled client reused across every request (text, query, image) so + # a connection — and its name resolution — is established once and kept + # warm, rather than a fresh connect per call. Built lazily inside a + # running loop; the lock ensures concurrent first-callers create only + # one. + self._client: httpx.AsyncClient | None = None + self._client_lock = asyncio.Lock() def _headers(self) -> dict[str, str]: headers = {"Content-Type": "application/json"} @@ -45,16 +56,30 @@ class VLLMMultimodalEmbedder(EmbedderWrapper): headers["Authorization"] = f"Bearer {self._api_key}" return headers + async def _get_client(self) -> httpx.AsyncClient: + if self._client is None: + async with self._client_lock: + if self._client is None: + self._client = build_http_client(self._http) + return self._client + + async def aclose(self) -> None: + """Close the pooled HTTP client. Idempotent and safe to call on + teardown even when no request was ever made (the client is lazy).""" + if self._client is not None: + await self._client.aclose() + self._client = None + async def _post(self, body: dict[str, Any]) -> list[list[float]]: + client = await self._get_client() 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 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}. " @@ -62,7 +87,8 @@ class VLLMMultimodalEmbedder(EmbedderWrapper): ) from e except httpx.TimeoutException as e: raise ValueError( - f"Request to vLLM timed out after {self._timeout}s. Error: {e}" + f"Request to vLLM timed out after {self._http.timeout_s}s. " + f"Error: {e}" ) from e except httpx.HTTPStatusError as e: if e.response.status_code == 401: diff --git a/tests/test_embedder.py b/tests/test_embedder.py index db8787ee..742b71be 100644 --- a/tests/test_embedder.py +++ b/tests/test_embedder.py @@ -228,6 +228,77 @@ async def test_embed_chunks_picture_with_text_only_embedder_raises(): await embed_chunks([chunk], get_embedder(config), config) +class _ImageStubEmbedder(EmbedderWrapper): + """Multimodal stub that records each image it is asked to embed and + returns a distinct vector per unique payload.""" + + def __init__(self): + super().__init__(embedder=None, vector_dim=4, supports_images=True) + self.embedded: list[bytes] = [] + self._vectors: dict[bytes, list[float]] = {} + + async def embed_documents(self, texts): + return [[0.1] * 4 for _ in texts] + + async def embed_image(self, image): + self.embedded.append(image) + # Deterministic, payload-specific vector so callers can prove that a + # reused (deduped) vector really came from the matching image. + vec = self._vectors.setdefault(image, [float(len(self._vectors))] * 4) + return list(vec) + + +async def test_embed_chunks_dedupes_identical_pictures(): + """Identical image bytes are embedded once and the vector is reused for + every chunk that shares them, preserving order.""" + img_a = b"\x89PNG\r\n\x1a\nAAAA" + img_b = b"\x89PNG\r\n\x1a\nBBBB" + # Order: a, b, a, a, b — 5 chunks, 2 unique images. + payloads = [img_a, img_b, img_a, img_a, img_b] + chunks = [] + for i, data in enumerate(payloads): + c = Chunk(id=f"pic{i}", content="x", order=i) + c._picture_data = data + chunks.append(c) + + embedder = _ImageStubEmbedder() + embedded = await embed_chunks(chunks, embedder, AppConfig()) + + # Only the two unique images hit embed_image, in first-seen order. + assert embedder.embedded == [img_a, img_b] + # Every chunk gets a vector, and duplicates share the right one. + vecs = [c.embedding for c in embedded] + assert vecs[0] == vecs[2] == vecs[3] # all img_a + assert vecs[1] == vecs[4] # all img_b + assert vecs[0] != vecs[1] + + +async def test_embed_chunks_non_bytes_picture_not_deduped(): + """Picture data that isn't bytes bypasses the content-hash cache and is + embedded per-occurrence (defensive path; build_picture_chunks yields bytes + in practice).""" + payload = "not-bytes-sentinel" + chunks = [] + for i in range(2): + c = Chunk(id=f"pic{i}", content="x", order=i) + c._picture_data = payload + chunks.append(c) + + embedder = _ImageStubEmbedder() + embedded = await embed_chunks(chunks, embedder, AppConfig()) + + # No dedup for unhashable/non-bytes payloads: both occurrences embed. + assert embedder.embedded == [payload, payload] + assert embedded[0].embedding is not None + assert embedded[1].embedding is not None + + +async def test_embedder_aclose_without_owned_client_is_noop(): + """Base aclose is a no-op when the wrapper owns no HTTP client.""" + embedder = EmbedderWrapper(embedder=None, vector_dim=4) + await embedder.aclose() # must not raise + + async def test_embed_chunks_respects_configured_batch_size(monkeypatch): """`embeddings.batch_size` controls how `embed_chunks` slices its input. @@ -408,6 +479,185 @@ 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.""" + from haiku.rag.embeddings.vllm import VLLMMultimodalEmbedder + + constructed: list[object] = [] + + 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): + constructed.append(self) + self.closed = False + + async def post(self, url, json, headers): + return FakeResponse() + + async def aclose(self): + self.closed = True + + 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 len(constructed) == 1 # one pooled client, reused + await embedder.aclose() + assert constructed[0].closed is True + + +async def test_vllm_aclose_without_request_is_noop(monkeypatch): + """aclose is safe when no request was ever made (client is lazy).""" + from haiku.rag.embeddings.vllm import VLLMMultimodalEmbedder + + def fail(*args, **kwargs): + raise AssertionError("no client should be built before a request") + + monkeypatch.setattr("httpx.AsyncClient", fail) + + embedder = VLLMMultimodalEmbedder( + model_name="x", vector_dim=2, base_url="http://localhost:8000/v1" + ) + await embedder.aclose() # must not raise or construct a client + + +def test_embeddings_http_config_defaults(): + """Defaults preserve the historical transport behavior.""" + http = EmbeddingsConfig().http + assert http.timeout_s == 60.0 + assert http.max_connections == 16 + assert http.max_keepalive_connections == 16 + assert http.keepalive_expiry_s == 300.0 + + +async def test_vllm_client_built_from_http_config(monkeypatch): + """The pooled client's timeout and pool limits come from + EmbeddingHTTPConfig, not hardcoded constants.""" + from haiku.rag.config import EmbeddingHTTPConfig + from haiku.rag.embeddings.vllm import VLLMMultimodalEmbedder + + captured: dict = {} + + class FakeResponse: + def raise_for_status(self): + pass + + def json(self): + return {"data": [{"embedding": [0.1]}]} + + class FakeAsyncClient: + def __init__(self, *args, **kwargs): + captured["timeout"] = kwargs.get("timeout") + captured["limits"] = kwargs.get("limits") + + async def post(self, url, json, headers): + return FakeResponse() + + monkeypatch.setattr("httpx.AsyncClient", FakeAsyncClient) + + http = EmbeddingHTTPConfig( + timeout_s=12.0, + max_connections=3, + max_keepalive_connections=2, + keepalive_expiry_s=45.0, + ) + embedder = VLLMMultimodalEmbedder( + model_name="x", vector_dim=2, base_url="http://localhost:8000/v1", http=http + ) + await embedder.embed_query("hi") + + assert captured["timeout"].read == 12.0 + assert captured["limits"].max_connections == 3 + assert captured["limits"].max_keepalive_connections == 2 + assert captured["limits"].keepalive_expiry == 45.0 + + +async def test_openai_embedder_owns_and_closes_pooled_client(monkeypatch): + """The openai flow builds one pooled client, hands it to the provider, and + the wrapper closes it on aclose().""" + import haiku.rag.embeddings as emb + + closed = {"value": False} + + class FakeClient: + async def aclose(self): + closed["value"] = True + + fake_client = FakeClient() + captured: dict = {} + + class FakeProvider: + def __init__(self, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr(emb, "build_http_client", lambda http: fake_client) + monkeypatch.setattr(emb, "OpenAIProvider", FakeProvider) + monkeypatch.setattr( + emb, "OpenAIEmbeddingModel", lambda name, provider: ("model", name, provider) + ) + monkeypatch.setattr(emb, "Embedder", lambda model: ("embedder", model)) + + config = AppConfig( + embeddings=EmbeddingsConfig( + model=EmbeddingModelConfig( + provider="openai", name="text-embedding-3-small", vector_dim=1536 + ) + ) + ) + embedder = emb.get_embedder(config) + + # The same pooled client is passed to the provider and owned by the wrapper. + assert captured["http_client"] is fake_client + await embedder.aclose() + assert closed["value"] is True + + +async def test_openai_embedder_forwards_base_url(monkeypatch): + """A configured base_url is forwarded to the OpenAI provider (OpenAI- + compatible servers like vLLM/LM Studio).""" + import haiku.rag.embeddings as emb + + class FakeClient: + async def aclose(self): + pass + + captured: dict = {} + + class FakeProvider: + def __init__(self, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr(emb, "build_http_client", lambda http: FakeClient()) + monkeypatch.setattr(emb, "OpenAIProvider", FakeProvider) + monkeypatch.setattr(emb, "OpenAIEmbeddingModel", lambda name, provider: object()) + monkeypatch.setattr(emb, "Embedder", lambda model: object()) + + config = AppConfig( + embeddings=EmbeddingsConfig( + model=EmbeddingModelConfig( + provider="openai", + name="text-embedding-3-small", + vector_dim=1536, + base_url="http://vllm:8000/v1", + ) + ) + ) + emb.get_embedder(config) + assert captured["base_url"] == "http://vllm:8000/v1" + + async def test_vllm_supports_images_flag(): from haiku.rag.embeddings.vllm import VLLMMultimodalEmbedder