Skip to content

fix(provider): improve embedding rate-limit resilience - #9430

Open
lxfight wants to merge 4 commits into
AstrBotDevs:masterfrom
lxfight:fix/embedding-rate-limit-resilience
Open

fix(provider): improve embedding rate-limit resilience#9430
lxfight wants to merge 4 commits into
AstrBotDevs:masterfrom
lxfight:fix/embedding-rate-limit-resilience

Conversation

@lxfight

@lxfight lxfight commented Jul 28, 2026

Copy link
Copy Markdown
Member

Improve Embedding Provider reliability when processing knowledge-base uploads and queries under API rate limits or transient failures.

The previous batch implementation could return embeddings out of input order and retried requests without coordinating concurrent callers. Rate-limited providers could therefore continue receiving requests while cooling down, causing repeated failures.

Modifications / 改动点

  • Preserve input order when concurrent embedding batches finish out of order.

  • Add provider-scoped request pacing with configurable requests-per-minute limits.

  • Handle Retry-After, transient HTTP statuses, network failures, exponential backoff, and retry jitter.

  • Preserve structured HTTP status metadata for Gemini, NVIDIA, and Ollama Embedding providers.

  • Avoid retrying permanent client errors such as HTTP 401.

  • Route single embedding operations through the shared retry and pacing path.

  • Classify failures as embedding errors only when embedding generation fails, preserving existing storage error behavior.

  • Add localized configuration metadata for:

    • embedding_max_requests_per_minute
    • embedding_rate_limit_cooldown
  • Keep existing and third-party Embedding Provider configurations unlimited by default unless a limit is explicitly configured.

  • Add focused regression coverage for pacing, cooldowns, retries, provider adapters, ordering, and knowledge-base error boundaries.

  • Introduce no new dependencies.

  • This is NOT a breaking change. / 这不是一个破坏性变更。

Screenshots or Test Results / 运行截图或测试结果

uv run pytest \
  tests/unit/test_embedding_provider_batch.py \
  tests/unit/test_faiss_vec_db.py \
  tests/unit/test_kb_upload_atomicity.py \
  tests/unit/test_knowledge_base_service_contract.py \
  tests/test_openai_embedding_source.py \
  tests/test_fastapi_v1_dashboard.py -q

Result:

133 passed, 1 warning in 12.46s

Checklist / 检查清单

  • 😊 If there are new features added in the PR, I have discussed it with the authors through issues/emails, etc.
    / 如果 PR 中有新加入的功能,已经通过 Issue / 邮件等方式和作者讨论过。

  • 👀 My changes have been well-tested, and "Verification Steps" and "Screenshots" have been provided above.
    / 我的更改经过了良好的测试,并已在上方提供了“验证步骤”和“运行截图”

  • 🤓 I have ensured that no new dependencies are introduced, OR if new dependencies are introduced, they have been added to the appropriate locations in requirements.txt and pyproject.toml.
    / 我确保没有引入新依赖库,或者引入了新依赖库的同时将其添加到 requirements.txtpyproject.toml 文件相应位置。

  • 😮 My changes do not introduce malicious code.
    / 我的更改没有引入恶意代码。

Summary by Sourcery

Improve embedding provider resilience and observability under API rate limits by adding shared throttled retry logic, configurable pacing, and clearer error handling for knowledge-base operations.

New Features:

  • Add provider-scoped throttling for embedding requests with configurable per-minute limits and cooldown windows.

Bug Fixes:

  • Ensure embedding batches preserve input order and classify embedding failures separately from storage errors.

Enhancements:

  • Route single embedding calls through the shared batched retry and pacing path to improve reliability under rate limits and transient errors.
  • Preserve HTTP status and Retry-After metadata from Gemini, NVIDIA, and Ollama embedding providers for smarter retry behavior.
  • Surface user-friendly knowledge-base upload errors when embedding generation fails.

Documentation:

  • Document new embedding rate-limit configuration options in localized dashboard metadata.

Tests:

  • Add regression tests covering embedding batch retries, provider-level pacing, cooldown handling, status preservation, and knowledge-base error boundaries.

@dosubot dosubot Bot added size:L This PR changes 100-499 lines, ignoring generated files. area:provider The bug / feature is about AI Provider, Models, LLM Agent, LLM Agent Runner. feature:knowledge-base The bug / feature is about knowledge base labels Jul 28, 2026

@sourcery-ai sourcery-ai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey - I've found 1 issue

Prompt for AI Agents
Please address the comments from this code review:

## Individual Comments

### Comment 1
<location path="astrbot/core/provider/provider.py" line_range="327" />
<code_context>
             pass


+class EmbeddingProviderError(Exception):
+    """Represent an Embedding API failure with HTTP response metadata.
+
</code_context>
<issue_to_address>
**issue (complexity):** Consider refactoring the new embedding retry/backoff logic by either using or removing EmbeddingProviderError, extracting nested helpers into private methods, and centralizing retry policy into a dedicated helper to keep get_embeddings_batch simpler.

You can keep all the new functionality but reduce complexity and coupling by:

1. **Either wire `EmbeddingProviderError` into the retry path or remove it for now**

Right now it adds surface area but is unused. If you want to keep it, make `process_batch` wrap raw exceptions in `EmbeddingProviderError` so callers can access status/response:

```python
class EmbeddingProviderError(Exception):
    ...

# Inside get_embeddings_batch.process_batch
except Exception as e:
    last_error = e
    if (
        attempt >= max_retries - 1
        or not _is_retryable_embedding_error(e)
    ):
        # Preserve the original exception as cause
        raise EmbeddingProviderError(
            f"批次 {batch_idx} 处理失败,共尝试 {attempt + 1} 次: {e!s}",
            status_code=_get_status_code(e),
            response=getattr(e, "response", None),
        ) from e
```

If you don’t plan to use it in this PR (e.g., no one inspects `status_code` / `response`), removing the class now and reintroducing it when needed will simplify the mental model.

2. **Move the inner helper functions to private instance methods**

The nested helpers in `get_embeddings_batch` make the method a “god function.” Pulling them out as private methods keeps the main flow readable without changing behavior:

```python
class EmbeddingProvider(AbstractProvider):
    ...

    def _get_status_code_from_error(self, error: BaseException) -> int | None:
        for attr in ("status_code", "status", "code"):
            value = getattr(error, attr, None)
            if isinstance(value, int):
                return value
        response = getattr(error, "response", None)
        if response is not None:
            status_code = getattr(response, "status_code", None)
            if isinstance(status_code, int):
                return status_code
        return None

    def _get_retry_after_seconds_from_error(self, error: BaseException) -> float | None:
        # move body of _get_retry_after_seconds here unchanged
        ...

    def _is_retryable_embedding_error(self, error: BaseException) -> bool:
        status_code = self._get_status_code_from_error(error)
        ...
```

Then `get_embeddings_batch` becomes simpler:

```python
status_code = self._get_status_code_from_error(e)
retry_after = self._get_retry_after_seconds_from_error(e)
if not self._is_retryable_embedding_error(e) or attempt >= max_retries - 1:
    break
```

This also makes it easy to unit test these pieces independently.

3. **Extract the retry/backoff decision into a dedicated helper**

Right now `_retry_delay_seconds`, `_delay_embedding_requests`, and the status-code checks are intertwined inside `get_embeddings_batch`. You can consolidate the retry/backoff policy into a single method, so `process_batch` doesn’t need to know the details:

```python
class EmbeddingProvider(AbstractProvider):
    ...

    def _compute_embedding_retry_delay(
        self,
        attempt: int,
        error: BaseException,
        retry_status_codes: set[int],
        retry_backoff_max_s: float,
    ) -> float | None:
        """Return delay in seconds, or None if we should not retry."""
        if not self._is_retryable_embedding_error(error):
            return None

        retry_after = self._get_retry_after_seconds_from_error(error)
        if retry_after is not None:
            delay = min(retry_after, self._embedding_rate_limit_cooldown_max_s)
        else:
            base = min(float(2**attempt), retry_backoff_max_s)
            jitter = random.uniform(0.0, min(1.0, base))
            delay = min(base + jitter, retry_backoff_max_s)

        status_code = self._get_status_code_from_error(error)
        if (
            retry_after is not None
            or status_code in retry_status_codes
            or (status_code is not None and 500 <= status_code <= 599)
        ):
            self._delay_embedding_requests(delay)

        return delay
```

Then `process_batch` reads much clearer:

```python
for attempt in range(max_retries):
    try:
        await self._wait_for_embedding_request_slot()
        batch_embeddings = await self.get_embeddings(batch_texts)
        ...
        return
    except Exception as e:
        last_error = e
        delay = self._compute_embedding_retry_delay(
            attempt=attempt,
            error=e,
            retry_status_codes=retry_status_codes,
            retry_backoff_max_s=retry_backoff_max_s,
        )
        if delay is None or attempt >= max_retries - 1:
            break
        logger.warning(
            "Embedding batch %s failed (attempt %s/%s): %s; retrying in %.2fs",
            batch_idx,
            attempt + 1,
            max_retries,
            e,
            delay,
        )
        await asyncio.sleep(delay)
```

Now `get_embeddings_batch` focuses on batching & concurrency, while HTTP parsing and backoff policy are encapsulated in small, testable methods.
</issue_to_address>

Sourcery is free for open source - if you like our reviews please consider sharing them ✨
Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.

Comment thread astrbot/core/provider/provider.py
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:provider The bug / feature is about AI Provider, Models, LLM Agent, LLM Agent Runner. feature:knowledge-base The bug / feature is about knowledge base size:L This PR changes 100-499 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant