From 2163ab1aa3cd74c74abe7cf7f832a29b581bee4a Mon Sep 17 00:00:00 2001 From: Joe-Kneeland <104656007+Joe-Kneeland@users.noreply.github.com> Date: Tue, 18 Aug 2026 08:51:53 -0600 Subject: [PATCH] fix(llm): forward provider_params.timeout to the OpenAI-compatible embedding client (#1024) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(llm): forward provider_params.timeout to the OpenAI-compatible embedding client #832 and #903 added a configurable request timeout for the LLM registry and the Gemini embedding client respectively, but the OpenAI-compatible embedding client (src/embedding_client.py) was never wired up. It constructed AsyncOpenAI with no timeout at all, so a stalled socket against a slow or contended OpenAI-compatible backend (e.g. a self-hosted embedding model under load) wedges the deriver worker's event loop indefinitely — the exact failure #785/#903 describe, just via a code path #903 didn't cover. EmbeddingModelConfig now carries provider_params through from resolve_embedding_model_config, mirroring how resolve_model_config already does it for ModelConfig, and the OpenAI branch of _EmbeddingClient.__init__ extracts `timeout` via the existing request_timeout_from_extra_params helper. Unset stays unset — no existing behavior changes. Reproduced and verified against a real self-hosted deployment (local Ollama backend under load): before this fix, a single stuck embedding call blocked all deriver queue processing for 20+ minutes with no error logged, twice in one session. * fix(embedding): use first-class timeout on embedding model config provider_params is the LLM per-request escape hatch; embedding timeouts are client-construction knobs and belong next to max_batch_size. Wire the field for OpenAI and Gemini, omit the OpenAI kwarg when unset so the SDK default stays, and keep Gemini's 10-minute floor when unset. * test(embedding): live coverage for first-class embedding timeout Exercise EmbeddingModelConfig.timeout on one representative OpenAI and Gemini model: configured timeout lands on the SDK client, and a near-zero timeout aborts before the provider answers. --------- Co-authored-by: Aakash Kattelu --- config.toml.example | 2 + docs/v3/contributing/configuration.mdx | 8 ++ src/config.py | 19 +++ src/embedding_client.py | 22 +-- tests/live_llm/README.md | 2 +- tests/live_llm/conftest.py | 12 +- tests/live_llm/test_live_embeddings.py | 84 +++++++++++ tests/llm/test_embedding_client.py | 185 ++++++++++++++++++++++++- 8 files changed, 317 insertions(+), 17 deletions(-) diff --git a/config.toml.example b/config.toml.example index c5517999..60d9092a 100644 --- a/config.toml.example +++ b/config.toml.example @@ -77,6 +77,8 @@ model = "text-embedding-3-small" # Optional provider request input cap. Useful for OpenAI-compatible embedding # APIs with smaller limits, such as DashScope text-embedding-v4. # max_batch_size = 10 +# Optional client HTTP timeout in seconds (OpenAI + Gemini). +# timeout = 90.0 # Optional module-level endpoint overrides # [embedding.model_config.overrides] diff --git a/docs/v3/contributing/configuration.mdx b/docs/v3/contributing/configuration.mdx index c02a4ebb..3c210e71 100644 --- a/docs/v3/contributing/configuration.mdx +++ b/docs/v3/contributing/configuration.mdx @@ -267,6 +267,7 @@ EMBEDDING_MAX_TOKENS_PER_REQUEST=300000 EMBEDDING_MODEL_CONFIG__TRANSPORT=openai # openai, gemini EMBEDDING_MODEL_CONFIG__MODEL=text-embedding-3-small EMBEDDING_MODEL_CONFIG__MAX_BATCH_SIZE=10 # optional per-request input cap +EMBEDDING_MODEL_CONFIG__TIMEOUT=90.0 # optional client HTTP timeout (seconds) # Optional endpoint overrides EMBEDDING_MODEL_CONFIG__OVERRIDES__BASE_URL=http://localhost:8000/v1 @@ -279,6 +280,13 @@ document a per-request limit. Set it when an OpenAI-compatible embedding provider accepts fewer inputs per request, such as DashScope `text-embedding-v4` with a limit of 10. +`EMBEDDING_MODEL_CONFIG__TIMEOUT` is an optional client HTTP timeout in +seconds. OpenAI-compatible transports receive it as the SDK `timeout` kwarg +(omitted when unset, so the SDK default applies). Gemini converts it to +milliseconds on `http_options.timeout`, and keeps its existing 10-minute +default when unset. The value is validated at config load the same way as +LLM `provider_params.timeout` (positive, finite number of seconds). + Forwarding `dimensions=` to OpenAI-compatible providers is controlled by `EMBEDDING_MODEL_CONFIG__DIMENSIONS_MODE`: - `auto` (default): forwards `dimensions=` when **the operator has explicitly set `EMBEDDING_VECTOR_DIMENSIONS`** — provenance, not value — and the configured model is not on the known-rejecting list (currently `text-embedding-ada-002`). Explicit `EMBEDDING_VECTOR_DIMENSIONS=1536` *does* trigger the forward; this is how `text-embedding-3-large` truncation to 1536 is expressed. Deployments that leave the setting unset get their existing behavior (`dimensions=` is not forwarded). diff --git a/src/config.py b/src/config.py index d8100694..092cb6ee 100644 --- a/src/config.py +++ b/src/config.py @@ -394,6 +394,15 @@ class ConfiguredEmbeddingModelSettings(BaseModel): dimensions_mode: EmbeddingDimensionsMode = "auto" encoding_format_mode: EmbeddingEncodingFormatMode = "auto" max_batch_size: Annotated[int, Field(gt=0)] | None = None + # Client HTTP timeout in seconds. OpenAI receives seconds; Gemini converts to ms. + timeout: float | None = None + + @field_validator("timeout", mode="before") + @classmethod + def _validate_timeout(cls, v: Any) -> float | None: + if v is None: + return None + return coerce_provider_timeout(v) @model_validator(mode="before") @classmethod @@ -431,6 +440,15 @@ class EmbeddingModelConfig(BaseModel): api_key: str | None = None base_url: str | None = None max_batch_size: Annotated[int, Field(gt=0)] | None = None + # Client HTTP timeout in seconds. OpenAI receives seconds; Gemini converts to ms. + timeout: float | None = None + + @field_validator("timeout", mode="before") + @classmethod + def _validate_timeout(cls, v: Any) -> float | None: + if v is None: + return None + return coerce_provider_timeout(v) @model_validator(mode="before") @classmethod @@ -556,6 +574,7 @@ def resolve_embedding_model_config( api_key=api_key, base_url=configured.overrides.base_url, max_batch_size=configured.max_batch_size, + timeout=configured.timeout, ) diff --git a/src/embedding_client.py b/src/embedding_client.py index 08a4ff04..3cee4484 100644 --- a/src/embedding_client.py +++ b/src/embedding_client.py @@ -195,13 +195,13 @@ class _EmbeddingClient: from google import genai from google.genai import types as genai_types - # 10-minute HTTP timeout, in lockstep with the LLM registry's Gemini - # client (`src/llm/registry.py:_build_gemini_http_options`). Without - # this, a stalled Gemini embedding socket wedges the deriver worker - # exactly the way #785 describes for the LLM client. + # Default 10-minute HTTP timeout matches the LLM registry Gemini client. + timeout_ms = ( + int(config.timeout * 1000) if config.timeout is not None else 600_000 + ) http_options = genai_types.HttpOptions( base_url=config.base_url, - timeout=600_000, + timeout=timeout_ms, ) self.client: genai.Client | AsyncOpenAI = genai.Client( api_key=config.api_key, @@ -216,10 +216,14 @@ class _EmbeddingClient: raise ValueError("OpenAI API key is required") from openai import AsyncOpenAI - self.client = AsyncOpenAI( - api_key=config.api_key, - base_url=config.base_url, - ) + # Omit timeout when unset so the OpenAI SDK keeps its own default. + client_kwargs: dict[str, Any] = { + "api_key": config.api_key, + "base_url": config.base_url, + } + if config.timeout is not None: + client_kwargs["timeout"] = config.timeout + self.client = AsyncOpenAI(**client_kwargs) self.max_embedding_tokens = max_input_tokens self.max_batch_size = config.max_batch_size or 2048 diff --git a/tests/live_llm/README.md b/tests/live_llm/README.md index 6da7dfb4..442049bc 100644 --- a/tests/live_llm/README.md +++ b/tests/live_llm/README.md @@ -68,5 +68,5 @@ Coverage by provider: - OpenAI transport → OpenRouter non-reasoning models (e.g. `inception/mercury-2`): non-chat / diffusion architectures must stay on `max_tokens`, no `reasoning_effort`, tool-calling parameter-schema compatibility is the canary for exotic OR-served providers - Gemini 2.5/3.0 classes: structured outputs, cached-content reuse, thought signatures, multi-turn tool replay - Gemini 3.1 class: thinking and tool replay coverage by default; structured-output/caching coverage should only be added once Google documents support for that path -- Embeddings (`test_live_embeddings.py`): single embed, batched embed, batch-vs-single alignment, and chunk-to-id mapping for every configured embedding model. `gemini-embedding-2*` is the reason this exists — those models collapse a list of bare strings into one document (#745), and only a live call catches it +- Embeddings (`test_live_embeddings.py`): single embed, batched embed, batch-vs-single alignment, and chunk-to-id mapping for every configured embedding model. `gemini-embedding-2*` is the reason this exists — those models collapse a list of bare strings into one document (#745), and only a live call catches it. Also covers first-class `EmbeddingModelConfig.timeout` plumbing (one representative model per transport): configured timeout lands on the SDK client, and a near-zero timeout aborts before the provider answers - OpenAI-compatible embedding providers (e.g. OpenRouter's `google/gemini-embedding-001`): the #932 surface. Those providers reject a base64 embedding request outright (HTTP 400) or answer HTTP 200 with empty data, so the whole matrix fails without `encoding_format="float"`. Real OpenAI accepts base64 happily, so only a third-party provider catches it. Note that OpenRouter load-balances across upstreams, so the base64 failure is per-attempt rather than guaranteed: a retry can land on an endpoint that accepts it. `test_live_openai_float_encoding_matches_base64` covers the other side, that the float switch must not move vectors on real OpenAI diff --git a/tests/live_llm/conftest.py b/tests/live_llm/conftest.py index eb7b8dbb..9c58bb87 100644 --- a/tests/live_llm/conftest.py +++ b/tests/live_llm/conftest.py @@ -66,14 +66,23 @@ def require_embedding_key(spec: LiveEmbeddingSpec) -> str: return key +_EMBEDDING_CONFIG_OVERRIDE_KEYS = frozenset({"timeout", "max_batch_size"}) + + def make_embedding_client( spec: LiveEmbeddingSpec, **overrides: Any ) -> _EmbeddingClient: """Build a live embedding client for one matrix entry. Bypasses the `EmbeddingClient` singleton so each spec gets its own client - without mutating global settings. + without mutating global settings. `timeout` and `max_batch_size` land on + `EmbeddingModelConfig`; remaining kwargs go to `_EmbeddingClient`. """ + config_overrides = { + key: overrides.pop(key) + for key in _EMBEDDING_CONFIG_OVERRIDE_KEYS + if key in overrides + } kwargs: dict[str, Any] = { "vector_dimensions": spec.dimensions, "max_input_tokens": 2048, @@ -90,6 +99,7 @@ def make_embedding_client( model=spec.model, api_key=require_embedding_key(spec), base_url=spec.base_url, + **config_overrides, ), **kwargs, ) diff --git a/tests/live_llm/test_live_embeddings.py b/tests/live_llm/test_live_embeddings.py index c8af34b8..6ef2b048 100644 --- a/tests/live_llm/test_live_embeddings.py +++ b/tests/live_llm/test_live_embeddings.py @@ -1,10 +1,15 @@ from __future__ import annotations +import time from typing import Any, cast +import httpx +import openai import pytest from openai import AsyncOpenAI +from src.config import EmbeddingTransport + from .conftest import cosine_similarity, make_embedding_client from .embedding_matrix import LiveEmbeddingSpec, get_live_embedding_specs @@ -25,6 +30,52 @@ OPENAI_NATIVE_SPECS = tuple( spec for spec in ALL_SPECS if spec.family == "openai_embedding" ) +GENEROUS_TIMEOUT_SECONDS = 120.0 +TIGHT_TIMEOUT_SECONDS = 0.01 +# Well under the client defaults; generous enough to absorb SDK retries. +TIGHT_TIMEOUT_WALL_CLOCK_LIMIT_SECONDS = 30 + +EMBEDDING_TIMEOUT_EXCEPTIONS: dict[ + EmbeddingTransport, tuple[type[BaseException], ...] +] = { + "openai": (openai.APITimeoutError,), + # google-genai raises httpx or aiohttp timeouts depending on its transport; + # aiohttp surfaces as asyncio.TimeoutError (== builtins.TimeoutError). + "gemini": (httpx.TimeoutException, TimeoutError), +} + +TRANSPORT_MARKS = { + "openai": pytest.mark.requires_openai, + "gemini": pytest.mark.requires_gemini, +} + + +def representative_embedding_specs() -> list[Any]: + """One spec per transport — timeout plumbing is client-level, not model-level.""" + params: list[Any] = [] + for transport in ("openai", "gemini"): + specs = get_live_embedding_specs(transport=transport) + if not specs: + continue + # Prefer the native family over openai-compatible proxies. + family = f"{transport}_embedding" + native = next((s for s in specs if s.family == family), specs[0]) + params.append( + pytest.param(native, marks=TRANSPORT_MARKS[transport], id=native.id) + ) + return params + + +def assert_embedding_timeout_on_client( + client: Any, transport: EmbeddingTransport, timeout_seconds: float +) -> None: + if transport == "gemini": + http_options = client.client._api_client._http_options + assert http_options.timeout == int(timeout_seconds * 1000) + return + openai_client = cast(AsyncOpenAI, client.client) + assert openai_client.timeout == timeout_seconds + @pytest.mark.asyncio @pytest.mark.parametrize("spec", ALL_SPECS, ids=lambda spec: spec.id) @@ -160,3 +211,36 @@ async def test_live_gemini_batch_embed_survives_batch_split( assert len(embeddings) == len(BATCH_TEXTS) assert len({tuple(embedding) for embedding in embeddings}) == len(BATCH_TEXTS) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("spec", representative_embedding_specs()) +async def test_live_embedding_timeout_reaches_the_client( + spec: LiveEmbeddingSpec, +) -> None: + """Configured embedding timeout lands on the provider SDK client.""" + client = make_embedding_client(spec, timeout=GENEROUS_TIMEOUT_SECONDS) + + embedding = await client.embed(BATCH_TEXTS[0]) + + assert len(embedding) == spec.dimensions + assert_embedding_timeout_on_client(client, spec.transport, GENEROUS_TIMEOUT_SECONDS) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("spec", representative_embedding_specs()) +async def test_live_tight_embedding_timeout_aborts_request( + spec: LiveEmbeddingSpec, +) -> None: + """A near-zero embedding timeout aborts before the provider can answer.""" + client = make_embedding_client(spec, timeout=TIGHT_TIMEOUT_SECONDS) + + started = time.monotonic() + with pytest.raises(EMBEDDING_TIMEOUT_EXCEPTIONS[spec.transport]): + await client.embed(BATCH_TEXTS[0]) + elapsed = time.monotonic() - started + + assert elapsed < TIGHT_TIMEOUT_WALL_CLOCK_LIMIT_SECONDS, ( + f"tight embedding timeout took {elapsed:.1f}s — client timeout " + f"likely not applied" + ) diff --git a/tests/llm/test_embedding_client.py b/tests/llm/test_embedding_client.py index b50a66f0..9df80be1 100644 --- a/tests/llm/test_embedding_client.py +++ b/tests/llm/test_embedding_client.py @@ -65,7 +65,13 @@ async def test_openai_embedding_client_uses_configured_model_and_dimensions( fake_embeddings = FakeOpenAIEmbeddingsAPI([0.1] * 8) class FakeOpenAIClient: - def __init__(self, *, api_key: str | None, base_url: str | None) -> None: + def __init__( + self, + *, + api_key: str | None, + base_url: str | None, + timeout: float | None = None, + ) -> None: self.api_key: str | None = api_key self.base_url: str | None = base_url self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings @@ -104,7 +110,13 @@ async def test_openai_embedding_client_rejects_dimension_mismatch( fake_embeddings = FakeOpenAIEmbeddingsAPI([0.1] * 7) class FakeOpenAIClient: - def __init__(self, *, api_key: str | None, base_url: str | None) -> None: + def __init__( + self, + *, + api_key: str | None, + base_url: str | None, + timeout: float | None = None, + ) -> None: self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings monkeypatch.setattr("openai.AsyncOpenAI", FakeOpenAIClient) @@ -220,6 +232,118 @@ async def test_gemini_embedding_client_keeps_timeout_without_base_url( assert gemini_client.http_options.timeout == 600_000 +@pytest.mark.asyncio +async def test_openai_embedding_client_forwards_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Configured embedding timeout reaches the OpenAI-compatible client.""" + + class FakeOpenAIClient: + def __init__( + self, + *, + api_key: str | None, + base_url: str | None, + timeout: float | None = None, + ) -> None: + self.api_key: str | None = api_key + self.base_url: str | None = base_url + self.timeout: float | None = timeout + self.embeddings: FakeOpenAIEmbeddingsAPI = FakeOpenAIEmbeddingsAPI( + [0.1] * 8 + ) + + monkeypatch.setattr("openai.AsyncOpenAI", FakeOpenAIClient) + + client = _EmbeddingClient( + EmbeddingModelConfig( + transport="openai", + model="text-embedding-3-small", + api_key="test-key", + timeout=45, + ), + vector_dimensions=8, + max_input_tokens=8192, + max_tokens_per_request=300_000, + send_dimensions=False, + ) + + openai_client = cast(Any, client.client) + assert openai_client.timeout == 45.0 + + +@pytest.mark.asyncio +async def test_openai_embedding_client_omits_timeout_when_unset( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Unset timeout omits the kwarg so the OpenAI SDK keeps its default.""" + + missing = object() + + class FakeOpenAIClient: + def __init__( + self, + *, + api_key: str | None, + base_url: str | None, + timeout: object = missing, + ) -> None: + self.api_key: str | None = api_key + self.base_url: str | None = base_url + self.timeout: object = timeout + self.embeddings: FakeOpenAIEmbeddingsAPI = FakeOpenAIEmbeddingsAPI( + [0.1] * 8 + ) + + monkeypatch.setattr("openai.AsyncOpenAI", FakeOpenAIClient) + + client = _EmbeddingClient( + EmbeddingModelConfig( + transport="openai", + model="text-embedding-3-small", + api_key="test-key", + ), + vector_dimensions=8, + max_input_tokens=8192, + max_tokens_per_request=300_000, + send_dimensions=False, + ) + + openai_client = cast(Any, client.client) + assert openai_client.timeout is missing + + +@pytest.mark.asyncio +async def test_gemini_embedding_client_forwards_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Configured embedding timeout reaches Gemini as milliseconds.""" + + class FakeGeminiClient: + def __init__(self, *, api_key: str | None, http_options: Any) -> None: + self.api_key: str | None = api_key + self.http_options: Any = http_options + self.aio: Any = SimpleNamespace(models=SimpleNamespace()) + + monkeypatch.setattr("google.genai.Client", FakeGeminiClient) + + client = _EmbeddingClient( + EmbeddingModelConfig( + transport="gemini", + model="gemini-embedding-001", + api_key="gemini-key", + timeout=45, + ), + vector_dimensions=8, + max_input_tokens=4096, + max_tokens_per_request=300_000, + send_dimensions=False, + ) + + gemini_client = cast(Any, client.client) + assert gemini_client.http_options.timeout == 45_000 + + def _build_openai_client( monkeypatch: pytest.MonkeyPatch, *, @@ -233,7 +357,13 @@ def _build_openai_client( fake_embeddings = FakeOpenAIEmbeddingsAPI(embedding) class FakeOpenAIClient: - def __init__(self, *, api_key: str | None, base_url: str | None) -> None: + def __init__( + self, + *, + api_key: str | None, + base_url: str | None, + timeout: float | None = None, + ) -> None: self.api_key: str | None = api_key self.base_url: str | None = base_url self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings @@ -707,7 +837,13 @@ async def test_simple_batch_embed_respects_token_budget_per_request( fake_embeddings = FakeOpenAIEmbeddingsAPI([0.5] * 4) class FakeOpenAIClient: - def __init__(self, *, api_key: str | None, base_url: str | None) -> None: + def __init__( + self, + *, + api_key: str | None, + base_url: str | None, + timeout: float | None = None, + ) -> None: self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings monkeypatch.setattr("openai.AsyncOpenAI", FakeOpenAIClient) @@ -745,7 +881,13 @@ async def test_simple_batch_embed_rejects_oversized_input( fake_embeddings = FakeOpenAIEmbeddingsAPI([0.1] * 4) class FakeOpenAIClient: - def __init__(self, *, api_key: str | None, base_url: str | None) -> None: + def __init__( + self, + *, + api_key: str | None, + base_url: str | None, + timeout: float | None = None, + ) -> None: self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings monkeypatch.setattr("openai.AsyncOpenAI", FakeOpenAIClient) @@ -775,7 +917,13 @@ def test_prepare_chunks_returns_ordered_chunks( fake_embeddings = FakeOpenAIEmbeddingsAPI([0.1] * 4) class FakeOpenAIClient: - def __init__(self, *, api_key: str | None, base_url: str | None) -> None: + def __init__( + self, + *, + api_key: str | None, + base_url: str | None, + timeout: float | None = None, + ) -> None: self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings monkeypatch.setattr("openai.AsyncOpenAI", FakeOpenAIClient) @@ -818,6 +966,31 @@ def test_embedding_model_config_parses_max_batch_size_from_env( assert resolved.max_batch_size == 10 +def test_embedding_model_config_parses_timeout_from_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + s = _build_embedding_settings( + {"EMBEDDING_MODEL_CONFIG__TIMEOUT": "90.0"}, + monkeypatch, + ) + + assert s.MODEL_CONFIG.timeout == 90.0 + + resolved = resolve_embedding_model_config(s.MODEL_CONFIG) + assert resolved.timeout == 90.0 + + +def test_embedding_model_config_rejects_invalid_timeout() -> None: + with pytest.raises( + ValueError, match=r"provider_params\.timeout must be a positive number" + ): + EmbeddingModelConfig( + transport="openai", + model="text-embedding-3-small", + timeout=-1, + ) + + @pytest.mark.asyncio async def test_gemini_process_batch_wraps_contents_as_content_part( monkeypatch: pytest.MonkeyPatch,