fix(provider): improve embedding rate-limit resilience - #9430
Open
lxfight wants to merge 4 commits into
Open
Conversation
Contributor
There was a problem hiding this comment.
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>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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_minuteembedding_rate_limit_cooldownKeep 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 / 运行截图或测试结果
Result:
133 passed, 1 warning in 12.46sChecklist / 检查清单
😊 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.txtandpyproject.toml./ 我确保没有引入新依赖库,或者引入了新依赖库的同时将其添加到
requirements.txt和pyproject.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:
Bug Fixes:
Enhancements:
Documentation:
Tests: