diff --git a/api/adapters/embeddings/__init__.py b/api/adapters/embeddings/__init__.py index 64a48bb..daaa43a 100644 --- a/api/adapters/embeddings/__init__.py +++ b/api/adapters/embeddings/__init__.py @@ -1,6 +1,7 @@ from typing import TYPE_CHECKING from api.adapters.base import EmbeddingAdapter +from api.adapters.embeddings.errors import RateLimitError if TYPE_CHECKING: from api.models.config import EmbeddingConfig @@ -40,3 +41,52 @@ def get_embedding_adapter(config: "EmbeddingConfig") -> EmbeddingAdapter: ) raise ValueError(f"Unknown embedding provider: {provider!r}") + + +def same_embedding_shape(a: "EmbeddingConfig", b: "EmbeddingConfig") -> bool: + """Whether two embedding configs yield the same vector dimension, so a re-probe is needless. + + The dimension is fixed by provider + model + any output-dimensionality truncation, plus the + base URL (the same model *name* on a different server can be a different model). The rate limit, + API key, batch size, and concurrency don't change it — so an edit touching only those keeps the + dimension, and the vector store need not be re-sized from a fresh provider probe. + """ + return (a.provider, a.model, a.base_url, a.output_dimensionality) == ( + b.provider, + b.model, + b.base_url, + b.output_dimensionality, + ) + + +def resolve_dimensions( + embed: EmbeddingAdapter, + new_config: "EmbeddingConfig", + prior_config: "EmbeddingConfig | None", + prior_dimensions: int | None, +) -> int: + """The embedding vector size (to size the vector store), tolerant of a rate-limited provider. + + Probing the provider for the size (``embed.dimensions``) also confirms it is reachable and the + credentials work, and re-learns the size if it ever changed — so the probe runs first and a bad + key / unreachable host still fails fast. But a rate limit (HTTP 429) is transient, not an invalid + config: when it strikes and the embedding model shape is unchanged, the size cannot have changed, + so reuse the prior (still-live) store's known dimension instead of failing. That is what lets an + unrelated edit — notably raising ``max_rpm`` to relieve the 429 — go through while the provider + is throttled. Only a genuine model change while throttled truly can't be sized; that re-raises + for the caller to surface as a transient error (a 503 in the API apply, an error ack + rollback + in the worker reload). + + Shared by both config-apply call sites (:func:`api.services.config_service._build_adapters` and + :func:`worker.tasks.reload_adapters`) so the API and every worker size the store identically. + """ + try: + return embed.dimensions + except RateLimitError: + if ( + prior_dimensions is not None + and prior_config is not None + and same_embedding_shape(new_config, prior_config) + ): + return prior_dimensions + raise diff --git a/api/adapters/vector_store/pgvector.py b/api/adapters/vector_store/pgvector.py index 034aee8..43aba3c 100644 --- a/api/adapters/vector_store/pgvector.py +++ b/api/adapters/vector_store/pgvector.py @@ -181,6 +181,17 @@ def __init__(self, host: str, port: int, database: str, user: str, self._runner = _AsyncRunner() self._lock = asyncio.Lock() + @property + def dimensions(self) -> int: + """The fixed vector size this store was built for (its ``embedding vector(N)`` column). + + A construction-time constant — set once at boot / config-apply — so it is always available + without calling the embedding provider. That makes it the reliable source for sizing a + rebuilt store when the embedding model is unchanged (see + ``config_service._embedding_dimensions``), where probing the provider could hit a rate limit. + """ + return self._dimensions + # -- connection / schema bootstrap -------------------------------------- async def _bootstrap_schema(self) -> None: diff --git a/api/services/config_service.py b/api/services/config_service.py index 878618b..deb0a75 100644 --- a/api/services/config_service.py +++ b/api/services/config_service.py @@ -30,6 +30,8 @@ from fastapi import HTTPException from api.adapters.embeddings import get_embedding_adapter as resolve_embedding +from api.adapters.embeddings import resolve_dimensions +from api.adapters.embeddings.errors import RateLimitError from api.adapters.reranker import get_reranker as resolve_reranker from api.adapters.vector_store import get_vector_store as resolve_store from api.dependencies import ( @@ -165,7 +167,18 @@ def _build_adapters(config: AppConfig) -> tuple[Any, Any, Any]: unrelated ones (e.g. rotating the embedding API key). """ embed = resolve_embedding(config.embedding) - store = resolve_store(config.vector_store, embed.dimensions) + live_store = get_vector_store() + live_config = get_app_config() + # Size the store tolerantly: a rate-limited (429) dimension probe reuses the still-live store's + # known size when the embedding model is unchanged, so an unrelated edit (e.g. raising max_rpm to + # relieve the 429) isn't rejected. Only a model change that can't be probed re-raises -> 503 below. + dimensions = resolve_dimensions( + embed, + config.embedding, + live_config.embedding if live_config is not None else None, + live_store.dimensions if live_store is not None else None, + ) + store = resolve_store(config.vector_store, dimensions) reranker = _build_reranker_optional(config.reranker) return embed, store, reranker @@ -245,6 +258,13 @@ def _validate_and_build(merged: dict[str, Any]) -> tuple[Any, Any, Any, AppConfi embed, store, reranker = _build_adapters(new_config) except HTTPException: raise + except RateLimitError as exc: + # The provider throttled the dimension probe and the model changed, so the live size + # couldn't be reused. A transient limit, not an invalid config — tell the client to retry + # rather than rejecting a valid change (503, not 422). + raise HTTPException( + 503, f"Embedding provider is rate-limited; try again shortly: {exc}" + ) from exc except Exception as exc: # invalid config or unbuildable adapter raise HTTPException(422, f"Invalid configuration: {exc}") from exc return embed, store, reranker, new_config diff --git a/tests/unit/test_config_service.py b/tests/unit/test_config_service.py index 3b4273b..e4ef435 100644 --- a/tests/unit/test_config_service.py +++ b/tests/unit/test_config_service.py @@ -7,6 +7,7 @@ from fastapi import HTTPException from api import dependencies +from api.adapters.embeddings.errors import RateLimitError from api.models.config import AppConfig, EmbeddingConfig, RerankerConfig, VectorStoreConfig from api.services import config_reload as cr from api.services import config_service as cs @@ -17,8 +18,17 @@ class _FakeEmbed: dimensions = 384 +class _RateLimitedEmbed: + """Embedding adapter whose dimension probe hits the provider's rate limit (HTTP 429).""" + + @property + def dimensions(self) -> int: + raise RateLimitError("embedding provider HTTP 429: quota exceeded") + + class _FakeStore: - pass + def __init__(self, dimensions: int = 384) -> None: + self.dimensions = dimensions @pytest.fixture(autouse=True) @@ -281,6 +291,54 @@ def test_apply_config_preserves_masked_secret(tmp_path, monkeypatch): assert dependencies.get_app_config().embedding.api_key == "keep-me" +def test_apply_config_survives_rate_limited_probe_when_model_unchanged(tmp_path, monkeypatch): + # A throttled provider (HTTP 429) during the dimension probe must not fail a save that doesn't + # change the model: the dimension can't have changed, so the probe's 429 is caught and the live + # store's known size (768) is reused. This is the reported bug — editing max-requests/min was + # rejected because the probe hit Gemini's daily quota. + # + # The live embedding *adapter* is left unset (None) on purpose: it mirrors the real breakage, + # where a boot whose warm-up probe was itself rate-limited swallows the error and never registers + # an adapter. The reused size therefore comes from the live *store*, which always carries one. + live = AppConfig(embedding=EmbeddingConfig(provider="gemini", model="gemini-embedding-2", api_key="k")) + dependencies.set_app_config(live) + dependencies.set_vector_store(_FakeStore(dimensions=768)) # live store carries the known size + monkeypatch.setattr(cs, "resolve_embedding", lambda _c: _RateLimitedEmbed()) + monkeypatch.setattr(cs, "resolve_store", lambda _c, dims: _FakeStore(dims)) + monkeypatch.setattr(cs, "_config_path", lambda: tmp_path / "config.yaml") + + payload = AppConfig( + embedding=EmbeddingConfig( + provider="gemini", model="gemini-embedding-2", api_key=cs.SECRET_MASK, max_rpm=90 + ) + ) + result = cs.apply_config(payload) # must not raise on the 429 + + assert result["status"] == "applied" + assert dependencies.get_app_config().embedding.max_rpm == 90 # the throttle edit landed + assert dependencies.get_vector_store().dimensions == 768 # rebuilt store reused the known size + + +def test_apply_config_returns_503_when_model_change_probe_is_rate_limited(tmp_path, monkeypatch): + # Changing the model *while* the provider is throttled: the live store's size is for the OLD + # model, so the new model's dimension genuinely can't be reused — it must be probed, and that + # probe is rate-limited. Surface a transient 503 ("try again"), not a misleading 422. + live = AppConfig(embedding=EmbeddingConfig(provider="gemini", model="gemini-embedding-2", api_key="k")) + dependencies.set_app_config(live) + dependencies.set_vector_store(_FakeStore(dimensions=768)) + monkeypatch.setattr(cs, "resolve_embedding", lambda _c: _RateLimitedEmbed()) + monkeypatch.setattr(cs, "resolve_store", lambda _c, dims: _FakeStore(dims)) + monkeypatch.setattr(cs, "_config_path", lambda: tmp_path / "config.yaml") + + payload = AppConfig( + embedding=EmbeddingConfig(provider="gemini", model="gemini-embedding-99", api_key=cs.SECRET_MASK) + ) + with pytest.raises(HTTPException) as exc: + cs.apply_config(payload) + assert exc.value.status_code == 503 + assert "rate-limited" in exc.value.detail + + def test_list_ollama_models_uses_explicit_base_url(monkeypatch): monkeypatch.setattr( "api.adapters.llm_chat.list_ollama_models", lambda url: [f"model@{url}"] diff --git a/tests/unit/test_embedding_dimensions.py b/tests/unit/test_embedding_dimensions.py new file mode 100644 index 0000000..cb8dca0 --- /dev/null +++ b/tests/unit/test_embedding_dimensions.py @@ -0,0 +1,101 @@ +"""Unit tests for the shared, rate-limit-tolerant embedding-dimension helpers. + +``resolve_dimensions`` / ``same_embedding_shape`` back both config-apply call sites (the API's +``config_service._build_adapters`` and the worker's ``tasks.reload_adapters``), so the size a +throttled provider can't be probed for is reused from the live store rather than failing the save. +""" + +from __future__ import annotations + +import pytest + +from api.adapters.embeddings import resolve_dimensions, same_embedding_shape +from api.adapters.embeddings.errors import RateLimitError +from api.models.config import EmbeddingConfig + + +class _Embed: + """Embedding adapter whose dimension probe returns a fixed size.""" + + def __init__(self, dims: int) -> None: + self._dims = dims + + @property + def dimensions(self) -> int: + return self._dims + + +class _RateLimitedEmbed: + """Embedding adapter whose dimension probe hits the provider rate limit (HTTP 429).""" + + @property + def dimensions(self) -> int: + raise RateLimitError("embedding provider HTTP 429: quota exceeded") + + +def _cfg(**over: object) -> EmbeddingConfig: + base: dict = {"provider": "gemini", "model": "gemini-embedding-2", "api_key": "k"} + base.update(over) + return EmbeddingConfig(**base) + + +# ── same_embedding_shape ────────────────────────────────────────────────────── + + +def test_same_shape_ignores_non_dimension_fields(): + # Rate limit, API key, batch size, and concurrency don't affect the vector size. + a = _cfg(api_key="k1", max_rpm=50, batch_size=16, concurrency=4) + b = _cfg(api_key="k2", max_rpm=90, batch_size=32, concurrency=8) + assert same_embedding_shape(a, b) + + +@pytest.mark.parametrize( + "over", + [ + {"provider": "openai_compat"}, + {"model": "gemini-embedding-99"}, + {"base_url": "http://other:1234"}, + {"output_dimensionality": 1536}, + ], +) +def test_different_shape_when_a_dimension_field_changes(over): + assert not same_embedding_shape(_cfg(), _cfg(**over)) + + +# ── resolve_dimensions ──────────────────────────────────────────────────────── + + +def test_probes_when_not_rate_limited(): + # Happy path: the provider answers, so the freshly probed size wins and prior state is ignored + # (this is what self-heals a stale boot-fallback size once the provider is reachable again). + assert resolve_dimensions(_Embed(3072), _cfg(), prior_config=_cfg(), prior_dimensions=768) == 3072 + + +def test_reuses_prior_dimension_on_429_when_shape_unchanged(): + # 429 during the probe + unchanged model → reuse the prior store's known size instead of failing. + dims = resolve_dimensions( + _RateLimitedEmbed(), _cfg(max_rpm=90), prior_config=_cfg(max_rpm=50), prior_dimensions=768 + ) + assert dims == 768 + + +def test_reraises_on_429_when_model_changed(): + # 429 + a real model change → the new size genuinely can't be known; surface the limit. + with pytest.raises(RateLimitError): + resolve_dimensions( + _RateLimitedEmbed(), + _cfg(model="gemini-embedding-99"), + prior_config=_cfg(), + prior_dimensions=768, + ) + + +def test_reraises_on_429_when_no_prior_dimension(): + # 429 with nothing to reuse (no live store yet) → can't size it; surface the limit. + with pytest.raises(RateLimitError): + resolve_dimensions(_RateLimitedEmbed(), _cfg(), prior_config=_cfg(), prior_dimensions=None) + + +def test_reraises_on_429_when_no_prior_config(): + with pytest.raises(RateLimitError): + resolve_dimensions(_RateLimitedEmbed(), _cfg(), prior_config=None, prior_dimensions=768) diff --git a/tests/unit/test_worker_config_reload.py b/tests/unit/test_worker_config_reload.py index 4c8bc00..3cce8fe 100644 --- a/tests/unit/test_worker_config_reload.py +++ b/tests/unit/test_worker_config_reload.py @@ -3,9 +3,12 @@ from __future__ import annotations import json +from types import SimpleNamespace import pytest +from api.adapters.embeddings.errors import RateLimitError +from api.models.config import EmbeddingConfig from api.services import config_reload as cr from tests.unit.fake_redis import FakeRedis from worker import config_reload as wcr @@ -64,8 +67,17 @@ class _FakeEmbed: dimensions = 384 +class _RateLimitedEmbed: + """Embedding adapter whose dimension probe hits the provider rate limit (HTTP 429).""" + + @property + def dimensions(self) -> int: + raise RateLimitError("embedding provider HTTP 429: quota exceeded") + + class _FakeStore: - pass + def __init__(self, dimensions: int = 384) -> None: + self.dimensions = dimensions class _FakeConfig: @@ -81,7 +93,7 @@ def test_redis_client_builds_a_client(monkeypatch): def test_reload_adapters_clears_cache_and_rebuilds(monkeypatch): calls: list[str] = [] - monkeypatch.setattr(tasks, "reload_adapters", lambda: calls.append("reload")) + monkeypatch.setattr(tasks, "reload_adapters", lambda _prior: calls.append("reload")) wcr._reload_adapters() assert calls == ["reload"] @@ -113,6 +125,78 @@ def test_reload_adapters_rebuilds_singletons(monkeypatch): assert isinstance(tasks._vector_store_singleton, _FakeStore) +def test_reload_adapters_reuses_store_dim_when_rate_limited_and_shape_unchanged(monkeypatch): + # Mirrors the API apply: a 429 dimension probe during an unchanged-model reload (e.g. a raised + # max_rpm) must reuse the still-live store's size, not raise — a raise would error-ack and force + # the API to roll the whole config change back. + before = EmbeddingConfig(provider="gemini", model="gemini-embedding-2", api_key="k", max_rpm=50) + after = EmbeddingConfig(provider="gemini", model="gemini-embedding-2", api_key="k", max_rpm=90) + monkeypatch.setattr(tasks, "get_config", lambda: SimpleNamespace(embedding=after, vector_store=object())) + monkeypatch.setattr("api.adapters.embeddings.get_embedding_adapter", lambda _c: _RateLimitedEmbed()) + monkeypatch.setattr("api.adapters.vector_store.get_vector_store", lambda _c, d: _FakeStore(d)) + tasks._embedder_singleton = None + tasks._vector_store_singleton = _FakeStore(dimensions=768) # still-live store carries the size + + tasks.reload_adapters(before) # must not raise on the 429 + + assert tasks._vector_store_singleton.dimensions == 768 # rebuilt store reused the known size + + +def test_reload_adapters_defers_store_on_rate_limit_when_store_not_built(monkeypatch): + # The reported 409: a worker whose store singleton isn't built yet (continuously throttled since + # its last restart) gets an unchanged-model reload (a raised max_rpm). It has no prior dimension + # to reuse, but a 429 must NOT roll the save back — so it adopts the new embedder and defers the + # store to a lazy rebuild rather than raising (which would error-ack and force a rollback). + before = EmbeddingConfig(provider="gemini", model="gemini-embedding-2", api_key="k", max_rpm=50) + after = EmbeddingConfig(provider="gemini", model="gemini-embedding-2", api_key="k", max_rpm=90) + monkeypatch.setattr(tasks, "get_config", lambda: SimpleNamespace(embedding=after, vector_store=object())) + monkeypatch.setattr("api.adapters.embeddings.get_embedding_adapter", lambda _c: _RateLimitedEmbed()) + monkeypatch.setattr("api.adapters.vector_store.get_vector_store", lambda _c, d: _FakeStore(d)) + tasks._embedder_singleton = None + tasks._vector_store_singleton = None # never built (throttled since restart) + + tasks.reload_adapters(before) # must NOT raise on the 429 + + assert isinstance(tasks._embedder_singleton, _RateLimitedEmbed) # new config's embedder adopted + assert tasks._vector_store_singleton is None # store deferred for lazy rebuild + + +def test_reload_adapters_defers_store_on_rate_limit_when_model_changed(monkeypatch): + # A model change while throttled can't reuse the prior size (it was for the old model), but a 429 + # still must not roll the save back — defer the store to a lazy rebuild, which re-probes the new + # model's size once the provider is reachable. + before = EmbeddingConfig(provider="gemini", model="gemini-embedding-2", api_key="k") + after = EmbeddingConfig(provider="gemini", model="gemini-embedding-99", api_key="k") + monkeypatch.setattr(tasks, "get_config", lambda: SimpleNamespace(embedding=after, vector_store=object())) + monkeypatch.setattr("api.adapters.embeddings.get_embedding_adapter", lambda _c: _RateLimitedEmbed()) + monkeypatch.setattr("api.adapters.vector_store.get_vector_store", lambda _c, d: _FakeStore(d)) + tasks._embedder_singleton = None + tasks._vector_store_singleton = _FakeStore(dimensions=768) + + tasks.reload_adapters(before) # must NOT raise + + assert tasks._vector_store_singleton is None # deferred; the stale 768 is not kept for a new model + + +def test_reload_adapters_propagates_non_rate_limit_error(monkeypatch): + # Only a *rate limit* defers. A genuine build failure (a bad model / key — not a 429) must still + # propagate, so the listener error-acks and the API rolls back rather than silently accepting it. + class _BrokenEmbed: + @property + def dimensions(self) -> int: + raise ValueError("model 'ghost' not found") + + after = EmbeddingConfig(provider="gemini", model="ghost", api_key="k") + monkeypatch.setattr(tasks, "get_config", lambda: SimpleNamespace(embedding=after, vector_store=object())) + monkeypatch.setattr("api.adapters.embeddings.get_embedding_adapter", lambda _c: _BrokenEmbed()) + monkeypatch.setattr("api.adapters.vector_store.get_vector_store", lambda _c, d: _FakeStore(d)) + tasks._embedder_singleton = None + tasks._vector_store_singleton = _FakeStore(dimensions=768) + + with pytest.raises(ValueError): + tasks.reload_adapters(EmbeddingConfig(provider="gemini", model="gemini-embedding-2", api_key="k")) + + @pytest.fixture(autouse=True) def _reset_task_singletons(): yield diff --git a/worker/config_reload.py b/worker/config_reload.py index d82e557..68f09e7 100644 --- a/worker/config_reload.py +++ b/worker/config_reload.py @@ -47,7 +47,7 @@ def _reload_adapters() -> bool: before = get_config().embedding # snapshot before the cache is cleared get_config.cache_clear() - tasks.reload_adapters() + tasks.reload_adapters(before) # pass the prior embedding so a 429 can reuse the live dimension return get_config().embedding != before diff --git a/worker/tasks.py b/worker/tasks.py index 167d8d4..1025798 100644 --- a/worker/tasks.py +++ b/worker/tasks.py @@ -118,22 +118,45 @@ def _redis() -> Any: return _redis_singleton -def reload_adapters() -> None: +def reload_adapters(prior_embedding: Any = None) -> None: """Rebuild the embedder + vector-store singletons from the current config. - Called by the config hot-reload listener after ``get_config.cache_clear()`` so - a live config change takes effect without restarting the worker. Building here - (rather than nulling the singletons for lazy rebuild) surfaces a bad config - immediately, so the listener can ack an error and the API can roll back. + Called by the config hot-reload listener after ``get_config.cache_clear()`` so a live config + change takes effect without restarting the worker. Building here (rather than nulling the + singletons for lazy rebuild) surfaces a bad config immediately, so the listener can ack an error + and the API can roll back. + + A provider rate limit (HTTP 429) is the exception: it is transient, not a bad config, and must + never roll a config save back. ``prior_embedding`` (the embedding config from *before* this + reload, passed by the listener) lets the store keep its size when the model is unchanged — the + dimension can't have changed, so :func:`resolve_dimensions` reuses the live store's size. When + even that isn't possible (this worker hasn't built a store yet, or the model itself changed), the + store is *deferred*: the new embedder is still adopted (so a rotated key / raised ``max_rpm`` + takes effect) and the store is nulled for a lazy rebuild on next use, when the provider should be + reachable. Ingestion already pauses + retries on 429, so a deferred store is safe. Any *other* + build failure (a bad model / key — not a 429) still propagates, so the API rolls back. """ global _embedder_singleton, _vector_store_singleton - from api.adapters.embeddings import get_embedding_adapter + from api.adapters.embeddings import get_embedding_adapter, resolve_dimensions + from api.adapters.embeddings.errors import RateLimitError from api.adapters.vector_store import get_vector_store config = get_config() + prior_store = _vector_store_singleton # still the previous store; its size is the 429 fallback embedder = get_embedding_adapter(config.embedding) + try: + dimensions = resolve_dimensions( + embedder, + config.embedding, + prior_embedding, + prior_store.dimensions if prior_store is not None else None, + ) + except RateLimitError: + _embedder_singleton = embedder + _vector_store_singleton = None # defer: lazy rebuild re-probes when the provider is reachable + return _embedder_singleton = embedder - _vector_store_singleton = get_vector_store(config.vector_store, embedder.dimensions) + _vector_store_singleton = get_vector_store(config.vector_store, dimensions) # ---------------------------------------------------------------------------