diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index 4fad8361..7a074f5d 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -2,6 +2,7 @@ import asyncio import hashlib +import re from collections import OrderedDict from pathlib import Path @@ -14,6 +15,75 @@ from ..as_embedding import BaseAsEmbedding Miss = tuple[int, str, str] # (result_index, text, cache_key) _MAX_VECTOR_SPACE_ATTEMPTS = 3 +_PROVIDER_WIDE_STATUS_CODES = frozenset({401, 403, 404}) + +# Transport failures raised by the OpenAI SDK and httpx. They do not inherit from the built-in +# TimeoutError / ConnectionError / OSError, so the built-in check alone lets them slip into the +# generic handler and fan out into per-item requests. Matching on class names (across the MRO, +# so subclasses such as APITimeoutError are covered) keeps this module free of SDK imports. +_TRANSPORT_FAILURE_TYPES = frozenset( + { + "APIConnectionError", + "APITimeoutError", + "ConnectError", + "ConnectTimeout", + "PoolTimeout", + "ReadError", + "ReadTimeout", + "RemoteProtocolError", + "WriteError", + "WriteTimeout", + }, +) + +# The status a provider reports can live in more than one place. ``status_code`` is the +# OpenAI-compatible attribute, Google GenAI exposes an integer ``code``, httpx-style wrappers +# keep it on ``response``, and AgentScope's DashScope implementation raises a plain +# ``RuntimeError`` whose message embeds the structured response, so the status has to be read +# back out of the message. A string ``code`` names an error (``insufficient_quota``, +# ``rate_limit_exceeded``, ``InvalidApiKey``) and is never a status. +_STATUS_IN_MESSAGE = re.compile(r'"status_code"\s*:\s*(\d{3})') + +# The business error code travels the same way the status does: a string ``code`` on the +# OpenAI-compatible SDKs, ``error.code`` inside a structured body, and a serialized +# ``"code"`` in the message for AgentScope's DashScope ``RuntimeError``. Reading only the +# attribute loses the DashScope case and misclassifies a quota error as a rate limit. +_CODE_IN_MESSAGE = re.compile(r'"code"\s*:\s*"([^"]+)"') + +# A rejected batch can only be paid back one item at a time when the provider blames an input. +# A context-length rejection is the one status an over-long item explains by itself, so it is +# the only error allowed to fan out into per-item requests. +_INPUT_DEPENDENT_STATUS_CODES = frozenset({400}) + +# A 413 is ambiguous. Providers such as Ollama v0.40.3 normalize an embedding input-limit +# rejection to 413, but a bare 413 can also describe an unrelated request-body limit that no +# single input explains. It therefore only qualifies as input-dependent when the error itself +# names the input length, which is why it is checked through the marker list instead of being +# accepted outright. +_INPUT_DEPENDENT_AMBIGUOUS_STATUS_CODES = frozenset({413}) + +_INPUT_DEPENDENT_ERROR_TYPES = frozenset( + { + "ContextLengthError", + "ContextWindowExceededError", + "InputLengthError", + "TokenLimitError", + }, +) + +_INPUT_DEPENDENT_MARKERS = ( + "context length", + "context_length", + "context window", + "input length", + "input is too long", + "maximum context", + "reduce the length", + "token limit", + "too many tokens", + "too long", +) + @R.register("local") class LocalEmbeddingStore(BaseEmbeddingStore): @@ -162,7 +232,17 @@ class LocalEmbeddingStore(BaseEmbeddingStore): async def _compute_batch(self, batch: list[Miss], **kwargs) -> list[tuple[int, str, np.ndarray]]: texts = [text for _, text, _ in batch] - embeddings = await self._call_with_retry(texts, **kwargs) + # A provider-wide failure can end a request after some of its vectors were already + # recovered. Those vectors are still worth keeping, but they must not report the + # provider as healthy again: the unhealthy-to-healthy transition is what schedules + # missing-vector backfill, so a hidden outage would either trigger it while the + # provider is still unavailable or suppress a later genuine recovery. + terminal_failure: list[bool] = [] + embeddings = await self._call_with_retry( + texts, + terminal_failure=terminal_failure, + **kwargs, + ) if not embeddings or len(embeddings) != len(texts): return [] out: list[tuple[int, str, np.ndarray]] = [] @@ -178,22 +258,63 @@ class LocalEmbeddingStore(BaseEmbeddingStore): if bad_dims: details = ", ".join(f"{count} with dim {dim}" for dim, count in sorted(bad_dims.items())) self.logger.error(f"Embedding dimension mismatch in batch: expected {self.dimensions}; rejected {details}") - if out: + if out and not terminal_failure: self.is_healthy = True else: self.is_healthy = False return out - async def _call_with_retry(self, texts: list[str], **kwargs) -> list[list[float] | None] | None: + async def _call_with_retry( + self, + texts: list[str], + *, + isolated: bool = False, + terminal_failure: list[bool] | None = None, + **kwargs, + ) -> list[list[float] | None] | None: + """Answer a batch, retrying transient failures and isolating input ones. + + ``isolated`` marks a single-item attempt made while isolating a rejected batch. Such an + attempt keeps the store's transport backoff, because it is the last chance to keep that + one vector; every other request leaves transport retries to the embedding wrapper and the + provider SDK, which already own them. + + ``terminal_failure`` is an optional out-parameter. When the caller passes a list, + it is set to ``[True]`` if this call gave up on a failure that no remaining item + can explain (a provider-wide error, exhausted quota, or exhausted transient + retries). The per-item isolation loop uses it to stop instead of asking for every + remaining item and multiplying the rejected requests, and it forwards the same + flag to its caller so a request that recovered some vectors before the failure is + still not reported as healthy. + """ + + def stop() -> None: + """Record that this failure ends the work for the whole batch.""" + if terminal_failure is not None: + terminal_failure[:] = [True] + for attempt in range(self.max_retries): try: result = await self.as_embedding(texts, **kwargs) if result and len(result) == len(texts): return result - except (TimeoutError, ConnectionError, OSError): - if attempt < self.max_retries - 1: - await asyncio.sleep(2**attempt) except Exception as error: + if self._is_transport_failure(error): + # A connection failure or a timeout describes the endpoint, not any single + # item, so this request is never split per item. Which layer retries it is + # explicit: an isolated single-item attempt keeps the store's backoff, while + # any other request is left to the layers below, which already own that retry + # -- the OpenAI model retries a batch max_retries + 1 times and the SDK retries + # each of those attempts, so a batch they answer after 12 attempts used to + # reach the provider 36 times with this store retrying on top of them. + if isolated: + if attempt < self.max_retries - 1: + await asyncio.sleep(2**attempt) + continue + self.logger.exception("Embedding request failed: transport error, not splitting") + self.is_healthy = False + stop() + break if self._is_rate_limited(error): if attempt < self.max_retries - 1: delay = 2**attempt @@ -202,6 +323,7 @@ class LocalEmbeddingStore(BaseEmbeddingStore): continue self.logger.exception("Embedding request failed after exhausting rate-limit retries") self.is_healthy = False + stop() return None if ( self.quota_retry_delay is not None @@ -213,35 +335,227 @@ class LocalEmbeddingStore(BaseEmbeddingStore): ) await asyncio.sleep(self.quota_retry_delay) continue - self.logger.exception("Embedding request failed") - self.is_healthy = False - return None + if self._is_insufficient_quota(error): + # Quota is a service-level failure: no single item is to blame, so splitting the + # batch cannot isolate anything and would only multiply the rejected requests. + # This also covers the case where the quota retries above are exhausted. + self.logger.exception("Embedding request failed: quota exhausted") + self.is_healthy = False + stop() + return None + if self._is_provider_wide_failure(error): + # Credentials, permissions and server outages describe the endpoint, not any + # single item: isolation cannot fix them, so give up on the whole batch + # instead of multiplying the rejected requests. + # + # A 5xx is deliberately *not* retried here. The embedding wrapper and the + # provider SDK already retry transient server errors, and every extra layer + # that retries multiplies the attempts and the waiting time of the layers + # below it -- a batch the SDK answered after 12 attempts became 36 once this + # store retried on top of it. Retry ownership therefore stays with those + # layers: a caller that wants a longer budget configures it there instead of + # stacking another silent loop here. A transport failure is left to them for + # the same reason; it only skips the per-item split. + self.logger.exception("Embedding request failed: provider-wide error, not splitting") + stop() + break + if self._is_input_dependent_failure(error): + # A rejection a single input explains by itself is the one case the per-item + # split can resolve: isolating it keeps the remaining vectors usable. + self.logger.exception("Embedding request failed: input-dependent error") + self.is_healthy = False + if len(texts) > 1: + return await self._embed_individually( + texts, + terminal_failure=terminal_failure, + **kwargs, + ) + return None + # Anything else counts as provider-wide. Splitting cannot recover a vector unless + # the failing input is identified, so an unrecognized error is never answered + # with one request per item. + self.logger.exception("Embedding request failed: unrecognized error, not splitting") + stop() + break self.is_healthy = False + stop() return None + async def _embed_individually( + self, + texts: list[str], + *, + terminal_failure: list[bool] | None = None, + **kwargs, + ) -> list[list[float] | None]: + """Retry a rejected batch one item at a time so one bad item cannot drop the rest. + + Provider-side request errors, such as a per-item context-length limit, reject the + whole batch and would otherwise cost every vector in it. Isolating each item bounds + the damage to the item that actually failed and keeps the remaining vectors usable. + Each single-item request goes back through ``_call_with_retry`` with ``isolated=True`` so + a rate limit and a transient transport failure still get the store's backoff instead of + losing the vector; a provider-wide failure ends the work for the whole batch rather than + being asked for every remaining item, and a one-item request never splits again. + When the loop stops on a provider-wide failure, ``terminal_failure`` is set for the + caller so the vectors recovered before it are not read as a healthy provider. + """ + results: list[list[float] | None] = [] + for index, text in enumerate(texts): + terminal: list[bool] = [] + single = await self._call_with_retry( + [text], + isolated=True, + terminal_failure=terminal, + **kwargs, + ) + if single: + results.append(single[0]) + continue + self.logger.error(f"Embedding request failed for a single item (chars={len(text)})") + results.append(None) + if terminal and terminal[0]: + # No remaining item can explain a provider-wide failure, so asking for the + # rest would only multiply the rejected requests. Keep the vectors already + # recovered and leave the items that were never answered as None. The batch + # caller has to learn about it as well: a partial success must not report + # the provider as healthy again. + if terminal_failure is not None: + terminal_failure[:] = [True] + self.logger.error( + f"Stopping per-item isolation after {index + 1}/{len(texts)} items: " + "the failure is not specific to any input", + ) + results.extend([None] * (len(texts) - index - 1)) + break + return results + + @staticmethod + def _is_transport_failure(error: Exception) -> bool: + """Return whether an error describes a transport failure rather than a rejected input. + + Covers the built-in socket errors as well as the provider/httpx exceptions that describe + a connection failure or a timeout. Transport failures are provider-wide: no individual + input can be blamed, so the request is never split per item. Who retries them is explicit + in the transport branch of ``_call_with_retry``: an isolated single-item attempt keeps the + store's backoff because it is the last chance to keep that vector, while every other + request leaves the retry to the layers below, which already own it. + """ + if isinstance(error, (TimeoutError, ConnectionError, OSError)): + return True + return any(cls.__name__ in _TRANSPORT_FAILURE_TYPES for cls in type(error).__mro__) + + @staticmethod + def _http_status(error: Exception) -> int | None: + """Return the HTTP status an exception carries, across the provider SDK shapes. + + ``status_code`` covers the OpenAI-compatible SDKs, an integer ``code`` covers Google + GenAI, ``response.status_code`` covers httpx-style wrappers, and the structured + ``status_code`` embedded in the message covers AgentScope's DashScope ``RuntimeError``. + Only integers are accepted: a string ``code`` is a provider error name such as + ``insufficient_quota``, and reading it as a status would misclassify those errors. + """ + for attribute in ("status_code", "code"): + value = getattr(error, attribute, None) + if isinstance(value, int) and not isinstance(value, bool): + return value + value = getattr(getattr(error, "response", None), "status_code", None) + if isinstance(value, int) and not isinstance(value, bool): + return value + match = _STATUS_IN_MESSAGE.search(str(error)) + return int(match.group(1)) if match else None + + @staticmethod + def _is_provider_wide_failure(error: Exception) -> bool: + """Return whether an error describes the provider rather than a single input. + + Authentication, permission, missing model or route and server errors belong to the + endpoint, the credentials or the deployment: no individual input can be blamed, so + splitting the batch cannot isolate anything and would only multiply the rejected + requests. Rate limits and quota exhaustion have their own branches above and never + reach this classifier. A context-length ``400`` stays input-dependent on purpose, + because that is the one status an over-limit item can explain by itself. + """ + status = LocalEmbeddingStore._http_status(error) + if status is not None and (status in _PROVIDER_WIDE_STATUS_CODES or 500 <= status < 600): + return True + return type(error).__name__ in ( + "AuthenticationError", + "PermissionDeniedError", + "NotFoundError", + "InternalServerError", + ) + + @staticmethod + def _is_input_dependent_failure(error: Exception) -> bool: + """Return whether a single input can be blamed for the rejected request. + + Only a context-length rejection qualifies: it is the one error an over-long item + explains by itself, and therefore the only one a per-item retry can resolve. It can + arrive as a ``400`` or, from a provider that normalizes the rejection (Ollama reports + ``413``), as an ambiguous status that is accepted only when the error names the input + itself. Anything else stays on the provider-wide path, so a client never answers a + service failure with one request per item. + """ + status = LocalEmbeddingStore._http_status(error) + if status is not None and status not in _INPUT_DEPENDENT_STATUS_CODES: + if status not in _INPUT_DEPENDENT_AMBIGUOUS_STATUS_CODES: + return False + # The status alone does not name an input for an ambiguous code, so the error text + # has to carry the signal; otherwise this is a request-level rejection. + text = LocalEmbeddingStore._error_text(error) + return any(marker in text for marker in _INPUT_DEPENDENT_MARKERS) + if type(error).__name__ in _INPUT_DEPENDENT_ERROR_TYPES: + return True + text = LocalEmbeddingStore._error_text(error) + return any(marker in text for marker in _INPUT_DEPENDENT_MARKERS) + + @staticmethod + def _error_text(error: Exception) -> str: + """Flatten an exception and its OpenAI-compatible body into one lowercase string.""" + parts = [str(error)] + body = getattr(error, "body", None) + if isinstance(body, dict): + parts.append(str(body)) + return " ".join(parts).lower() + + @staticmethod + def _error_code(error: Exception) -> str | None: + """Return the provider's business error code, across the supported error shapes. + + ``code`` covers the OpenAI-compatible SDKs when the provider puts the error name + there, ``error.code`` inside ``body`` covers the structured OpenAI payload, and the + serialized message covers AgentScope's DashScope ``RuntimeError`` — the same shape + ``_http_status`` already reads the status back out of. Without the last case a + DashScope quota error is indistinguishable from an ordinary 429. + """ + code = getattr(error, "code", None) + if isinstance(code, str) and code: + return code + body = getattr(error, "body", None) + if isinstance(body, dict): + details = body.get("error", body) + if isinstance(details, dict): + inner = details.get("code") + if isinstance(inner, str) and inner: + return inner + match = _CODE_IN_MESSAGE.search(str(error)) + return match.group(1) if match else None + @staticmethod def _is_rate_limited(error: Exception) -> bool: """Recognize an OpenAI-compatible 429 response without importing a provider SDK.""" - if LocalEmbeddingStore._is_insufficient_quota(error): + code = LocalEmbeddingStore._error_code(error) + if code == "insufficient_quota": return False - if getattr(error, "status_code", None) == 429: + if code == "rate_limit_exceeded": return True - body = getattr(error, "body", None) - if not isinstance(body, dict): - return False - details = body.get("error", body) - return isinstance(details, dict) and details.get("code") == "rate_limit_exceeded" + return LocalEmbeddingStore._http_status(error) == 429 @staticmethod def _is_insufficient_quota(error: Exception) -> bool: """Recognize OpenAI-compatible quota errors without importing a provider SDK.""" - if getattr(error, "code", None) == "insufficient_quota": - return True - body = getattr(error, "body", None) - if not isinstance(body, dict): - return False - details = body.get("error", body) - return isinstance(details, dict) and details.get("code") == "insufficient_quota" + return LocalEmbeddingStore._error_code(error) == "insufficient_quota" def _validate_dim(self, emb: np.ndarray) -> bool: """Return whether an embedding exactly matches the configured dimension.""" diff --git a/reme/components/file_store/local_file_store.py b/reme/components/file_store/local_file_store.py index e5448f65..dcbe6177 100644 --- a/reme/components/file_store/local_file_store.py +++ b/reme/components/file_store/local_file_store.py @@ -43,6 +43,7 @@ _SEARCH_FILTER_RESERVED_KEYS = frozenset( "strict_date_filter", }, ) +_FAILED_CHUNK_PREVIEW = 5 @R.register("local") @@ -440,7 +441,12 @@ class LocalFileStore(BaseFileStore): if chunk.text and not self._embedding_dim_matches(chunk.embedding) ] if missing: - raise RuntimeError(f"embedding reindex incomplete: {len(missing)} chunks failed") + preview = ", ".join(missing[:_FAILED_CHUNK_PREVIEW]) + if len(missing) > _FAILED_CHUNK_PREVIEW: + preview += f" (+{len(missing) - _FAILED_CHUNK_PREVIEW} more)" + raise RuntimeError( + f"embedding reindex incomplete: {len(missing)} chunks failed: {preview}", + ) await self._finalize_embedding_reindex() await self._dump_owned_state() if self.embedding_store is not None: @@ -544,6 +550,7 @@ class LocalFileStore(BaseFileStore): return processed = 0 + filled = 0 batch_count = 0 embedding_started_at = time.monotonic() next_percent = _PROGRESS_LOG_PERCENT_STEP @@ -552,10 +559,11 @@ class LocalFileStore(BaseFileStore): await self.embedding_store.get_node_embeddings(batch) self._drop_stale_embeddings(batch, "backfill") processed += len(batch) + filled += sum(1 for chunk in batch if chunk.embedding is not None) batch_count += 1 next_percent = self._log_progress("embedding backfill", processed, total, next_percent) self.logger.info( - f"{self.name}: embedding batches complete: processed={processed}/{total}, " + f"{self.name}: embedding batches complete: processed={processed}/{total}, filled={filled}, " f"batches={batch_count}, elapsed={time.monotonic() - embedding_started_at:.3f}s", ) except asyncio.CancelledError: diff --git a/tests/unit/test_file_store_consistency.py b/tests/unit/test_file_store_consistency.py index b7750017..9e6af475 100644 --- a/tests/unit/test_file_store_consistency.py +++ b/tests/unit/test_file_store_consistency.py @@ -118,6 +118,16 @@ class RecoveringEmbeddingStore(CountingFakeEmbeddingStore): return await super().get_node_embeddings(nodes, **kwargs) +class PartialEmbeddingStore(FakeEmbeddingStore): + """Fake provider that leaves chunks whose text contains "beta" without a vector.""" + + async def get_node_embeddings(self, nodes: list[FileChunk], **_kwargs) -> list[FileChunk]: + for chunk_node in nodes: + if "beta" not in chunk_node.text: + chunk_node.embedding = self._embed(chunk_node.text) + return nodes + + class HealthCountingEmbeddingStore(FakeEmbeddingStore): """Fake provider that records eager health checks.""" @@ -838,6 +848,41 @@ def test_verified_resume_without_chunks_supersedes_inflight_health_check(): run(go()) +@pytest.mark.parametrize("store_factory", [_new_local_store, _new_faiss_store, _new_zvec_store]) +def test_embedding_reindex_error_names_the_chunks_without_vectors(store_factory): + """The final reindex error must name the chunks that still have no vector.""" + + async def go(): + with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp): + store = store_factory("t_reindex_failure_ids") + await store.start() + await set_chunks_with_graph( + store, + { + "good": chunk("good", "a.md", "alpha text"), + "bad": chunk("bad", "b.md", "beta text"), + }, + ) + store.embedding_store = PartialEmbeddingStore() + _ensure_zvec_collection(store) + if isinstance(store, FaissLocalFileStore): + store._rebuild_index() + elif isinstance(store, ZvecLocalFileStore): + store._rebuild_collection() + await store.require_embedding_rebuild() + + with pytest.raises(RuntimeError) as failure: + await store.reindex("embedding") + + message = str(failure.value) + assert "1 chunks failed" in message + assert "bad" in message + assert "good" not in message + await store.close() + + run(go()) + + @pytest.mark.parametrize("store_factory", [_new_local_store, _new_faiss_store, _new_zvec_store]) def test_embedding_reindex_is_explicit_and_keeps_vectors_disabled_until_success( store_factory, diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 51b7bf63..1d9ed6e7 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -3,9 +3,12 @@ # pylint: disable=protected-access import asyncio +import json from types import SimpleNamespace +import httpx import numpy as np +import openai from reme.components.as_embedding import DashScopeAsEmbedding, OllamaAsEmbedding, OpenAIAsEmbedding from reme.components.embedding_store.base_embedding_store import BaseEmbeddingStore @@ -126,6 +129,282 @@ class QuotaThenSuccessAsEmbedding: return [[1.0, 0.0] for _ in texts] +class ContextLengthError(Exception): + """OpenAI-compatible 400 raised when one item exceeds the provider's per-item window.""" + + status_code = 400 + body = {"error": {"message": "the input length exceeds the context length"}} + + +class PerItemLimitAsEmbedding: + """Reject a whole request when any single text is over the provider's per-item limit. + + Mirrors an OpenAI-compatible gateway: one over-limit item makes the provider answer + HTTP 400 for the entire batch, so a client that gives up on the batch loses every + vector in it. + """ + + dimensions = 2 + vector_space_id = "fakespace000" + max_chars = 32 + + def __init__(self): + self.batch_sizes: list[int] = [] + + def initialize_model(self): + """Mirror the real component's idempotent initialization hook.""" + + async def __call__(self, texts: list[str], **_kwargs): + self.batch_sizes.append(len(texts)) + if any(len(text) > self.max_chars for text in texts): + raise ContextLengthError("the input length exceeds the context length") + return [[1.0, 0.0] for _ in texts] + + +class AlwaysInsufficientQuotaAsEmbedding: + """Always fail with quota exhaustion, recording every requested batch size.""" + + dimensions = 2 + + def __init__(self): + self.batch_sizes: list[int] = [] + + async def __call__(self, texts: list[str], **_kwargs): + self.batch_sizes.append(len(texts)) + raise InsufficientQuotaError("quota exhausted") + + +class PerItemLimitThenTransientErrorAsEmbedding: + """Reject an over-limit batch, then fail the first single-item request once. + + Mirrors the review case: the batch is rejected because one item exceeds the per-item + window, and a sibling item then hits a transient provider failure while retried alone. + """ + + dimensions = 2 + vector_space_id = "fakespace000" + max_chars = 32 + + def __init__(self, transient: Exception): + self.transient = transient + self.batch_sizes: list[int] = [] + self.transient_failures = 0 + + def initialize_model(self): + """Mirror the real component's idempotent initialization hook.""" + + async def __call__(self, texts: list[str], **_kwargs): + self.batch_sizes.append(len(texts)) + if any(len(text) > self.max_chars for text in texts): + raise ContextLengthError("the input length exceeds the context length") + if len(texts) == 1 and self.transient_failures == 0: + self.transient_failures += 1 + raise self.transient + return [[1.0, 0.0] for _ in texts] + + +class PerItemLimitThenQuotaAsEmbedding: + """Reject an over-limit batch, recover the first isolated item, then exhaust the quota. + + Mirrors the review case: isolation starts, the first single-item request comes back, and + every later call then fails provider-wide. Asking for the remaining items would only + multiply the rejected requests, because no single item explains a quota failure. + """ + + dimensions = 2 + vector_space_id = "fakespace000" + max_chars = 32 + + def __init__(self): + self.batch_sizes: list[int] = [] + self.singles = 0 + + def initialize_model(self): + """Mirror the real component's idempotent initialization hook.""" + + async def __call__(self, texts: list[str], **_kwargs): + self.batch_sizes.append(len(texts)) + if any(len(text) > self.max_chars for text in texts): + raise ContextLengthError("the input length exceeds the context length") + self.singles += 1 + if self.singles > 1: + raise InsufficientQuotaError("quota exhausted") + return [[1.0, 0.0] for _ in texts] + + +class AuthenticationError(Exception): + """OpenAI-compatible 401 error used without importing the provider SDK.""" + + status_code = 401 + + +class PermissionDeniedError(Exception): + """OpenAI-compatible 403 error used without importing the provider SDK.""" + + status_code = 403 + + +class InternalServerError(Exception): + """OpenAI-compatible 500 error used without importing the provider SDK.""" + + status_code = 500 + + +class NotFoundError(Exception): + """OpenAI-compatible 404 error, e.g. a missing model or endpoint.""" + + status_code = 404 + + +class APIError(Exception): + """Base of the OpenAI SDK error hierarchy, reproduced without importing the SDK.""" + + +class APIConnectionError(APIError): + """OpenAI SDK connection failure: it does not inherit from the built-in ConnectionError.""" + + +class APITimeoutError(APIConnectionError): + """OpenAI SDK timeout, a subclass of APIConnectionError.""" + + +class AlwaysProviderWideFailureAsEmbedding: + """Always fail with the same provider-wide error, recording every requested batch size.""" + + dimensions = 2 + + def __init__(self, error: Exception): + self.error = error + self.batch_sizes: list[int] = [] + + async def __call__(self, texts: list[str], **_kwargs): + self.batch_sizes.append(len(texts)) + raise self.error + + +class TransientThenSuccessAsEmbedding: + """Fail the first `failures` calls with one error, then return valid embeddings.""" + + dimensions = 2 + vector_space_id = "fakespace000" + + def __init__(self, error: Exception, failures: int = 1): + self.error = error + self.failures = failures + self.batch_sizes: list[int] = [] + + def initialize_model(self): + """Mirror the real component's idempotent initialization hook.""" + + async def __call__(self, texts: list[str], **_kwargs): + self.batch_sizes.append(len(texts)) + if len(self.batch_sizes) <= self.failures: + raise self.error + return [[1.0, 0.0] for _ in texts] + + +_REAL_OPENAI_ASYNC_CLIENT = openai.AsyncClient +_OFFLINE_OPENAI_BASE_URL = "http://127.0.0.1:9/v1" + + +class CountingOpenAIWrapper(OpenAIAsEmbedding): + """Real OpenAI wrapper that records how often the store asked it for embeddings.""" + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + self.calls = 0 + + async def __call__(self, inputs, **kwargs): + self.calls += 1 + return await super().__call__(inputs, **kwargs) + + +class OfflineOpenAIEndpoint: + """Real openai.AsyncClient wired to an offline httpx transport that counts provider attempts. + + Gives the wrapper-and-SDK layers their real code path, so a regression test can assert how + many attempts the provider would see instead of trusting the store's own loop. + """ + + def __init__(self, handler, sdk_max_retries: int = 2) -> None: + self.requests = 0 + self.handler = handler + self.client = _REAL_OPENAI_ASYNC_CLIENT( + api_key="test-key", + base_url=_OFFLINE_OPENAI_BASE_URL, + http_client=httpx.AsyncClient(transport=httpx.MockTransport(self._handle)), + max_retries=sdk_max_retries, + ) + + def _handle(self, request): + self.requests += 1 + return self.handler(request) + + def install(self, monkeypatch) -> "OfflineOpenAIEndpoint": + """Route every openai.AsyncClient built during this test to the offline client.""" + monkeypatch.setattr(openai, "AsyncClient", lambda **kwargs: self.client) + return self + + +def offline_openai_endpoint(handler) -> OfflineOpenAIEndpoint: + """Build an offline endpoint whose transport follows ``handler``.""" + return OfflineOpenAIEndpoint(handler) + + +def openai_wrapper(name: str) -> CountingOpenAIWrapper: + """Build a real OpenAI embedding wrapper pointed at the offline endpoint.""" + return CountingOpenAIWrapper( + name=name, + model="text-embedding-v4", + backend="openai", + dimensions=2, + credential={"base_url": _OFFLINE_OPENAI_BASE_URL, "api_key": "test-key"}, + retry_delay=0, + pass_dimensions=False, + ) + + +def embedding_payload(request) -> httpx.Response: + """Answer an embeddings request the way the provider would.""" + body = json.loads(request.content.decode("utf-8")) + inputs = body["input"] + if isinstance(inputs, str): + inputs = [inputs] + return httpx.Response( + 200, + json={ + "object": "list", + "data": [ + {"object": "embedding", "index": position, "embedding": [0.0, 1.0]} for position in range(len(inputs)) + ], + "model": body.get("model", "text-embedding-v4"), + "usage": {"prompt_tokens": len(inputs), "total_tokens": len(inputs)}, + }, + ) + + +def context_length_rejection() -> httpx.Response: + """Answer a batch the way an OpenAI-compatible provider reports an over-long input.""" + return httpx.Response( + 400, + json={ + "error": { + "message": ( + "This model's maximum context length is 8191 tokens. " "Please reduce the length of the input." + ), + "type": "invalid_request_error", + "param": None, + "code": "context_length_exceeded", + }, + }, + ) + + +def read_timeout(request) -> httpx.Response: + """Fail one attempt the way a read timeout reaches the SDK.""" + raise httpx.ReadTimeout("simulated read timeout", request=request) + + class BadNodeEmbeddingStore(BaseEmbeddingStore): """Embedding store that returns wrong-dimensional vectors.""" @@ -730,3 +1009,821 @@ def test_start_ignores_cache_file_without_vector_space_tag(monkeypatch, tmp_path assert not store._cache run(go()) + + +def test_batch_failure_falls_back_to_per_item_embedding(): + """One over-limit item must cost one vector, not its whole batch.""" + + async def go(): + provider = PerItemLimitAsEmbedding() + store = LocalEmbeddingStore(name="t_local_embedding_poison_batch", enable_cache=False) + store.as_embedding = provider + + results = await store._call_with_retry(["ok", "x" * 64, "also ok"]) + + assert results == [[1.0, 0.0], None, [1.0, 0.0]] + assert provider.batch_sizes == [3, 1, 1, 1] + + run(go()) + + +def test_poison_chunk_keeps_its_batch_mates(): + """A rejected chunk must not clear the vectors of the chunks beside it.""" + + async def go(): + store = LocalEmbeddingStore(name="t_local_embedding_poison_nodes", enable_cache=False) + store.as_embedding = PerItemLimitAsEmbedding() + nodes = [EmbNode(text="first"), EmbNode(text="y" * 64), EmbNode(text="third")] + + await store.get_node_embeddings(nodes) + + assert nodes[0].embedding is not None + assert nodes[1].embedding is None + assert nodes[2].embedding is not None + + run(go()) + + +def test_single_item_failure_still_reports_a_failed_batch(): + """A one-item request has nothing to isolate, so the existing contract stands.""" + + async def go(): + store = LocalEmbeddingStore(name="t_local_embedding_single_poison", enable_cache=False) + store.as_embedding = PerItemLimitAsEmbedding() + + assert await store._call_with_retry(["x" * 64]) is None + + run(go()) + + +def test_rate_limited_batch_is_not_split_into_per_item_requests(monkeypatch): + """A 429 is a provider-level failure; splitting it would multiply the load.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + store = LocalEmbeddingStore(name="t_local_embedding_rate_limit_batch", max_retries=2) + embedding = AlwaysRateLimitedAsEmbedding() + store.as_embedding = embedding + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + assert await store._call_with_retry(["a", "b"]) is None + assert embedding.calls == 2 + assert sleeps == [1.0] + + run(go()) + + +def test_insufficient_quota_batch_is_not_split_into_per_item_requests(monkeypatch): + """Quota exhaustion is a service-level failure; splitting it would multiply the load.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + store = LocalEmbeddingStore(name="t_local_embedding_quota_batch", max_retries=2) + embedding = AlwaysInsufficientQuotaAsEmbedding() + store.as_embedding = embedding + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + assert await store._call_with_retry(["a", "b", "c"]) is None + assert embedding.batch_sizes == [3] + assert store.is_healthy is False + assert not sleeps + + run(go()) + + +def test_exhausted_quota_retries_still_do_not_split_the_batch(monkeypatch): + """Once the opt-in quota retries run out, the store gives up instead of splitting.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + store = LocalEmbeddingStore( + name="t_local_embedding_quota_batch_exhausted", + max_retries=3, + quota_retry_delay=5.0, + ) + embedding = AlwaysInsufficientQuotaAsEmbedding() + store.as_embedding = embedding + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + assert await store._call_with_retry(["a", "b"]) is None + assert embedding.batch_sizes == [2, 2, 2] + assert sleeps == [5.0, 5.0] + + run(go()) + + +def test_per_item_fallback_keeps_a_rate_limited_vector(monkeypatch): + """A 429 on one isolated item still gets the store's backoff and keeps its vector. + + Rate limiting is the failure class the store still owns: no lower layer reports it as a + provider-wide error worth its own budget, so the retry lives here. Connection errors and + timeouts are left to the wrapper and the SDK instead; see the transport tests above. + """ + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + provider = PerItemLimitThenTransientErrorAsEmbedding( + RateLimitError("Requests are too frequent"), + ) + store = LocalEmbeddingStore( + name="t_local_embedding_fallback_rate_limit", + max_retries=3, + enable_cache=False, + ) + store.as_embedding = provider + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + results = await store._call_with_retry(["ok", "x" * 64, "also ok"]) + + assert results == [[1.0, 0.0], None, [1.0, 0.0]] + assert provider.batch_sizes == [3, 1, 1, 1, 1] + assert provider.transient_failures == 1 + assert sleeps == [1.0] + + run(go()) + + +def test_provider_wide_credentials_errors_are_not_split_into_per_item_requests(): + """401/403 belong to the credentials: splitting ten items only multiplies the failures.""" + + async def go(): + for error in ( + AuthenticationError("invalid api key"), + PermissionDeniedError("no access to model"), + ): + provider = AlwaysProviderWideFailureAsEmbedding(error) + store = LocalEmbeddingStore(name="t_local_embedding_provider_wide", max_retries=3) + store.as_embedding = provider + + assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None + assert provider.batch_sizes == [10] + assert store.is_healthy is False + + run(go()) + + +def test_provider_wide_server_errors_are_not_retried_without_splitting(monkeypatch): + """A 500 is left to the layers that own retries, so this store asks exactly once.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + provider = AlwaysProviderWideFailureAsEmbedding(InternalServerError("upstream unavailable")) + store = LocalEmbeddingStore(name="t_local_embedding_provider_server", max_retries=3) + store.as_embedding = provider + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None + # One request only: the wrapper and the SDK already retry a 5xx, and retrying here + # as well multiplied their attempts (12 -> 36 in the reported reproduction). + assert provider.batch_sizes == [10] + assert sleeps == [] + assert store.is_healthy is False + + run(go()) + + +def test_sdk_transport_errors_are_provider_wide_and_not_retried_here(monkeypatch): + """SDK connection/timeout errors are provider-wide: never split, and not retried here. + + The wrapper and the SDK already retry a transport failure, so the store answers once and + leaves the retry budget to those layers instead of multiplying their attempts. + """ + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + for error in (APIConnectionError("connection reset"), APITimeoutError("timed out")): + sleeps.clear() + provider = AlwaysProviderWideFailureAsEmbedding(error) + store = LocalEmbeddingStore(name="t_local_embedding_sdk_transport", max_retries=3) + store.as_embedding = provider + + assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None + # One request only: the wrapper and the SDK already retry a transport failure + # (4 wrapper attempts, each retried by the SDK), so another loop here only + # multiplies them -- a persistent timeout used to cost 36 attempts instead of 12. + assert provider.batch_sizes == [10] + assert sleeps == [] + assert store.is_healthy is False + + run(go()) + + +def test_isolated_item_keeps_the_store_backoff_for_a_transport_error(monkeypatch): + """An isolated single-item attempt still gets the store's backoff and keeps its vector. + + This fake provider has no wrapper or SDK beneath it, so it pins the store's own budget: the + isolated attempt is the last chance to keep that one vector. A whole-batch transport failure + is left to the layers below instead; see the earlier transport tests. + """ + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + for transient in ( + APIConnectionError("connection reset"), + ConnectionError("connection reset"), + TimeoutError("request timed out"), + ): + sleeps.clear() + + provider = PerItemLimitThenTransientErrorAsEmbedding(transient) + store = LocalEmbeddingStore( + name="t_local_embedding_sdk_isolated", + max_retries=3, + enable_cache=False, + ) + store.as_embedding = provider + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + results = await store._call_with_retry(["ok", "x" * 64, "also ok"]) + + assert results == [[1.0, 0.0], None, [1.0, 0.0]] + assert provider.batch_sizes == [3, 1, 1, 1, 1] + assert sleeps == [1.0] + + run(go()) + + +def test_a_persistent_transport_error_is_not_multiplied_by_the_store(monkeypatch): + """A timeout is asked once: the wrapper and the SDK own the retry budget. + + The provider sees the 4 wrapper attempts x 3 SDK attempts those layers make, while the + store's max_retries=3 no longer multiplies them -- a persistent timeout used to cost 36 + provider attempts instead of 12. + """ + + async def go(): + endpoint = offline_openai_endpoint(read_timeout).install(monkeypatch) + provider = openai_wrapper("t_openai_transport_budget") + store = LocalEmbeddingStore( + name="t_local_embedding_sdk_transport_budget", + max_retries=3, + enable_cache=False, + ) + store.as_embedding = provider + + assert await store._call_with_retry([f"item{i}" for i in range(3)]) is None + assert provider.calls == 1 + assert endpoint.requests == 12 + assert store.is_healthy is False + + run(go()) + + +def test_a_transient_transport_error_still_recovers_through_the_sdk(monkeypatch): + """One timeout is answered by the SDK's own retry: the batch survives and never splits.""" + + async def go(): + state = {"timeouts": 1} + + def handler(request): + if state["timeouts"] > 0: + state["timeouts"] -= 1 + return read_timeout(request) + return embedding_payload(request) + + endpoint = offline_openai_endpoint(handler).install(monkeypatch) + provider = openai_wrapper("t_openai_transport_recover") + store = LocalEmbeddingStore( + name="t_local_embedding_sdk_transport_recover", + max_retries=3, + enable_cache=False, + ) + store.as_embedding = provider + + results = await store._call_with_retry([f"item{i}" for i in range(3)]) + + assert results == [[0.0, 1.0]] * 3 + assert provider.calls == 1 + assert endpoint.requests == 2 + + run(go()) + + +def test_an_isolated_item_still_recovers_through_the_sdk(monkeypatch): + """After a context-length rejection, a timeout on one remaining item is still recovered. + + The store asks the wrapper once per item; the retry that saves that item belongs to the SDK. + """ + + async def go(): + state = {"timeout_pending": True} + + def handler(request): + body = json.loads(request.content.decode("utf-8")) + inputs = body["input"] + if isinstance(inputs, str): + inputs = [inputs] + if any(len(item) > 32 for item in inputs): + return context_length_rejection() + if len(inputs) == 1 and state["timeout_pending"]: + state["timeout_pending"] = False + return read_timeout(request) + return embedding_payload(request) + + endpoint = offline_openai_endpoint(handler).install(monkeypatch) + provider = openai_wrapper("t_openai_transport_isolated") + store = LocalEmbeddingStore( + name="t_local_embedding_sdk_transport_isolated", + max_retries=3, + enable_cache=False, + ) + store.as_embedding = provider + + results = await store._call_with_retry(["ok", "x" * 64, "also ok"]) + + assert results == [[0.0, 1.0], None, [0.0, 1.0]] + assert provider.calls == 4 + assert endpoint.requests == 5 + + run(go()) + + +def test_provider_wide_statuses_exclude_input_dependent_bad_request(): + """404 is provider-wide, but an over-limit 400 must stay input-dependent.""" + + async def go(): + store = LocalEmbeddingStore(name="t_local_embedding_classifier") + for error in ( + AuthenticationError("bad key"), + PermissionDeniedError("no access"), + NotFoundError("model_not_found"), + InternalServerError("boom"), + ): + assert store._is_provider_wide_failure(error) is True + over_limit = ContextLengthError("the input length exceeds the context length") + assert store._is_provider_wide_failure(over_limit) is False + + run(go()) + + +def test_provider_wide_not_found_errors_are_not_split(monkeypatch): + """A missing model or endpoint is provider-wide: one request, then failure, never a split.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + provider = AlwaysProviderWideFailureAsEmbedding(NotFoundError("model_not_found")) + store = LocalEmbeddingStore(name="t_local_embedding_not_found", max_retries=3) + store.as_embedding = provider + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None + assert provider.batch_sizes == [10] + assert not sleeps + + run(go()) + + +class GenAiClientError(Exception): + """Google GenAI client error: the HTTP status is an integer ``code``, not ``status_code``.""" + + def __init__(self, code: int, message: str = "client error"): + super().__init__(message) + self.code = code + + +class GenAiServerError(Exception): + """Google GenAI server error: the HTTP status is an integer ``code``, not ``status_code``.""" + + def __init__(self, code: int, message: str = "server error"): + super().__init__(message) + self.code = code + + +class StringCodeError(Exception): + """OpenAI-compatible error whose ``code`` names the error instead of carrying a status.""" + + code = "invalid_api_key" + + +class BadRequestError(Exception): + """OpenAI-compatible 400 that names no input, such as a malformed payload.""" + + status_code = 400 + + +class ResponseLike: + """httpx-style wrapper that keeps the status on a ``response`` attribute.""" + + def __init__(self, status_code: int): + self.status_code = status_code + + +class WrappedTransportError(Exception): + """Error that only exposes its status through ``response.status_code``.""" + + def __init__(self, status_code: int): + super().__init__("wrapped") + self.response = ResponseLike(status_code) + + +class DashScopeRuntimeError(RuntimeError): + """AgentScope turns every non-200 DashScope response into a plain ``RuntimeError``. + + The structured response is embedded in the message instead of being kept as an attribute, + so a client that only reads exception attributes cannot classify it. + """ + + @staticmethod + def for_status(status_code: int, code: str = "InvalidApiKey") -> "DashScopeRuntimeError": + """Build the exception AgentScope raises for a non-200 DashScope response.""" + response = { + "status_code": status_code, + "request_id": "req-0", + "code": code, + "message": "request rejected", + "output": None, + "usage": None, + } + return DashScopeRuntimeError(f"Failed to get embedding from DashScope API: {json.dumps(response)}") + + +def test_http_status_reads_every_supported_exception_shape(): + """Read the status from status_code, an integer code, response.status_code or the message.""" + + store = LocalEmbeddingStore(name="t_local_embedding_status_shapes") + + assert store._http_status(AuthenticationError("bad key")) == 401 + assert store._http_status(NotFoundError("model_not_found")) == 404 + assert store._http_status(GenAiClientError(401)) == 401 + assert store._http_status(GenAiServerError(503)) == 503 + assert store._http_status(WrappedTransportError(404)) == 404 + assert store._http_status(DashScopeRuntimeError.for_status(401)) == 401 + assert store._http_status(ValueError("no status here")) is None + + +def test_http_status_ignores_a_string_code(): + """A string ``code`` names the error; reading it as a status would misclassify quota errors.""" + + store = LocalEmbeddingStore(name="t_local_embedding_status_string_code") + + assert store._http_status(StringCodeError("invalid key")) is None + assert store._is_insufficient_quota(StringCodeError("invalid key")) is False + assert store._http_status(InsufficientQuotaError("quota exhausted")) == 429 + assert store._is_insufficient_quota(InsufficientQuotaError("quota exhausted")) is True + + +def test_google_genai_service_failures_are_not_split_into_per_item_requests(monkeypatch): + """A Gemini 401 / 429 / 503 describes the service, so the batch is never split per item. + + A 503 is answered once: the wrapper and the SDK own the retry for a 5xx, and retrying + here as well multiplied their attempts. A 429 keeps its retry because the rate-limit + branch above this classifier is the one place that reacts to it. + """ + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + for error, expected_batches in ( + (GenAiClientError(401), [10]), + (GenAiServerError(503), [10]), + (GenAiClientError(429), [10, 10, 10]), + ): + sleeps.clear() + provider = AlwaysProviderWideFailureAsEmbedding(error) + store = LocalEmbeddingStore( + name="t_local_embedding_genai", + max_retries=3, + enable_cache=False, + ) + store.as_embedding = provider + + assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None + assert provider.batch_sizes == expected_batches + + run(go()) + + +def test_dashscope_runtime_errors_are_not_split_into_per_item_requests(monkeypatch): + """A DashScope failure arrives as a plain RuntimeError; the status lives in its message. + + As with the Gemini shapes, a 503 is answered once (the lower layers own that retry) and + a 429 keeps the retry from the rate-limit branch. + """ + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + for status, expected_batches in ((401, [10]), (503, [10]), (429, [10, 10, 10])): + sleeps.clear() + provider = AlwaysProviderWideFailureAsEmbedding(DashScopeRuntimeError.for_status(status)) + store = LocalEmbeddingStore( + name="t_local_embedding_dashscope", + max_retries=3, + enable_cache=False, + ) + store.as_embedding = provider + + assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None + assert provider.batch_sizes == expected_batches + + run(go()) + + +def test_unrecognized_errors_are_not_split_into_per_item_requests(): + """Only an input-dependent error may split; anything unrecognized stays a batch failure.""" + + async def go(): + provider = AlwaysProviderWideFailureAsEmbedding(RuntimeError("provider said no")) + store = LocalEmbeddingStore(name="t_local_embedding_unrecognized", max_retries=3) + store.as_embedding = provider + + assert await store._call_with_retry(["a", "b", "c"]) is None + assert provider.batch_sizes == [3] + + run(go()) + + +def test_input_dependent_classifier_pins_both_halves_of_the_boundary(): + """The isolation must keep working for the one error a single item can explain.""" + + store = LocalEmbeddingStore(name="t_local_embedding_input_dependent") + + assert ( + store._is_input_dependent_failure( + ContextLengthError("the input length exceeds the context length"), + ) + is True + ) + assert store._is_input_dependent_failure(ContextLengthError("too long")) is True + assert store._is_input_dependent_failure(AuthenticationError("bad key")) is False + assert store._is_input_dependent_failure(InternalServerError("boom")) is False + assert store._is_input_dependent_failure(GenAiClientError(400)) is False + assert store._is_input_dependent_failure(RuntimeError("provider said no")) is False + + +def test_bad_request_without_an_input_length_signal_is_not_split(): + """A bare 400 does not name an input, so the batch is not divided over it.""" + + async def go(): + provider = AlwaysProviderWideFailureAsEmbedding(BadRequestError("malformed payload")) + store = LocalEmbeddingStore(name="t_local_embedding_bare_400", max_retries=2) + store.as_embedding = provider + + assert await store._call_with_retry(["a", "b", "c"]) is None + assert provider.batch_sizes == [3] + + run(go()) + + +def test_isolation_stops_when_a_provider_wide_failure_hits(): + """Once isolation meets a provider-wide failure, the remaining items are not requested. + + The vectors already recovered stay in the result: only the items that were never answered + are dropped, and no further request is sent for them. + """ + + async def go(): + provider = PerItemLimitThenQuotaAsEmbedding() + store = LocalEmbeddingStore( + name="t_local_embedding_isolation_stops", + max_retries=2, + enable_cache=False, + ) + store.as_embedding = provider + + results = await store._call_with_retry(["ok", "x" * 64, "also ok", "still ok"]) + + assert results == [[1.0, 0.0], None, None, None] + # One rejected batch, the recovered first item, the over-limit item, then the quota + # failure. The fourth item is never requested once the quota failure is seen. + assert provider.batch_sizes == [4, 1, 1, 1] + + run(go()) + + +def test_partial_success_keeps_a_failed_provider_unhealthy(): + """Recovered vectors must not overwrite the health state of a provider-wide failure. + + Isolation can recover an early item and then meet a provider-wide failure such as + exhausted quota. ``_call_with_retry`` marks the store unhealthy and stops, but the + vectors recovered before it are still returned. Treating that partial success as a + healthy provider would hide the outage, and it would break the unhealthy-to-healthy + transition the file store uses to schedule missing-vector backfill. The existing + isolation test drives ``_call_with_retry`` directly, so it never sees the overwrite. + """ + + async def go(): + provider = PerItemLimitThenQuotaAsEmbedding() + store = LocalEmbeddingStore( + name="t_local_embedding_partial_health", + max_retries=2, + enable_cache=False, + ) + store.as_embedding = provider + + vectors = await store.get_embeddings(["ok", "x" * 64, "also ok", "still ok"]) + + # The vector recovered before the outage is kept ... + assert vectors[0] is not None + assert vectors[1] is None and vectors[2] is None and vectors[3] is None + # ... and the provider-wide failure still decides the health state. + assert store.is_healthy is False + # One rejected batch, the recovered first item, the over-limit item, then the quota + # failure: the fourth item is never requested. + assert provider.batch_sizes == [4, 1, 1, 1] + + run(go()) + + +def test_isolation_keeps_going_while_only_inputs_fail(): + """An input-dependent failure is local to its item, so isolation must continue. + + This pins the other half of the boundary: stopping on every isolated failure would lose + the vectors of the items that are perfectly embeddable on their own. + """ + + async def go(): + provider = PerItemLimitAsEmbedding() + store = LocalEmbeddingStore( + name="t_local_embedding_isolation_keeps_going", + max_retries=2, + enable_cache=False, + ) + store.as_embedding = provider + + results = await store._call_with_retry(["ok", "x" * 64, "also ok", "still ok"]) + + assert results == [[1.0, 0.0], None, [1.0, 0.0], [1.0, 0.0]] + assert provider.batch_sizes == [4, 1, 1, 1, 1] + + run(go()) + + +def test_dashscope_serialized_quota_errors_keep_their_quota_policy(monkeypatch): + """A serialized DashScope quota error must classify as quota, not as a rate limit. + + AgentScope embeds the structured response in the message, so the business error code has to + be read back out of it. Without that the opt-in quota retry policy is never applied and the + error is retried with the generic backoff instead. + """ + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + error = DashScopeRuntimeError.for_status(429, code="insufficient_quota") + assert LocalEmbeddingStore._is_insufficient_quota(error) is True + assert LocalEmbeddingStore._is_rate_limited(error) is False + + provider = AlwaysProviderWideFailureAsEmbedding(error) + store = LocalEmbeddingStore( + name="t_local_embedding_dashscope_quota", + max_retries=3, + quota_retry_delay=7.0, + enable_cache=False, + ) + store.as_embedding = provider + + assert await store._call_with_retry(["a", "b"]) is None + assert provider.batch_sizes == [2, 2, 2] + assert sleeps == [7.0, 7.0] + + run(go()) + + +class ContextLength413Error(Exception): + """An over-limit input reported as HTTP 413 instead of 400. + + Ollama v0.40.3 normalizes an embedding input-limit rejection to 413 and propagates that + status through the embedding endpoint, so the same over-long item that arrives as a 400 + from an OpenAI-compatible gateway is a 413 for that provider. + """ + + status_code = 413 + body = {"error": {"message": "the input length exceeds the context length"}} + + +class RequestBodyTooLargeError(Exception): + """A bare 413 that never names an input: the request body itself is too large.""" + + status_code = 413 + + +class PerItemLimitAs413Embedding: + """Reject a whole request with HTTP 413 when any single text is over the per-item limit.""" + + dimensions = 2 + vector_space_id = "fakespace000" + max_chars = 32 + + def __init__(self): + self.batch_sizes: list[int] = [] + + def initialize_model(self): + """Mirror the real component's idempotent initialization hook.""" + + async def __call__(self, texts: list[str], **_kwargs): + self.batch_sizes.append(len(texts)) + if any(len(text) > self.max_chars for text in texts): + raise ContextLength413Error("the input length exceeds the context length") + return [[1.0, 0.0] for _ in texts] + + +def test_context_length_rejection_reported_as_http_413_is_isolated(): + """An over-limit input carried by a 413 must still reach the per-item isolation path. + + Ollama normalizes an embedding input-limit rejection to 413, so a classifier that only + accepts 400 rejects the batch instead of isolating the item that caused it, and every + healthy vector in the request is lost. + """ + + async def go(): + provider = PerItemLimitAs413Embedding() + store = LocalEmbeddingStore( + name="t_local_embedding_413_isolated", + max_retries=2, + enable_cache=False, + ) + store.as_embedding = provider + + results = await store._call_with_retry(["ok", "x" * 64, "also ok"]) + + assert results == [[1.0, 0.0], None, [1.0, 0.0]] + assert provider.batch_sizes == [3, 1, 1, 1] + + run(go()) + + +def test_bare_http_413_without_an_input_signal_is_not_split(): + """A 413 that names no input stays provider-wide, so the batch is not divided over it. + + The status alone is ambiguous: it can describe an unrelated request-body limit, which no + single input explains, so only a 413 that blames the input itself may fan out per item. + """ + + async def go(): + provider = AlwaysProviderWideFailureAsEmbedding( + RequestBodyTooLargeError("request entity too large"), + ) + store = LocalEmbeddingStore(name="t_local_embedding_bare_413", max_retries=2) + store.as_embedding = provider + + assert await store._call_with_retry(["a", "b", "c"]) is None + assert provider.batch_sizes == [3] + + run(go()) + + +def test_input_dependent_classifier_accepts_a_413_only_with_an_input_signal(): + """Pin both halves of the 413 boundary, and that the 400 behaviour is unchanged.""" + + store = LocalEmbeddingStore(name="t_local_embedding_413_classifier") + + assert ( + store._is_input_dependent_failure( + ContextLength413Error("the input length exceeds the context length"), + ) + is True + ) + assert store._is_input_dependent_failure(ContextLength413Error("too long")) is True + assert store._is_input_dependent_failure(RequestBodyTooLargeError("request entity too large")) is False + assert store._is_input_dependent_failure(ContextLengthError("too long")) is True + assert store._is_input_dependent_failure(GenAiClientError(400)) is False