Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lxfight wants to merge 4 commits into
base: master
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
4ba7b1e
fix(provider): add embedding retry and rate-limit handling
lxfight Jul 28, 2026
5fbebdf
fix(kb): preserve embedding and storage error boundaries
lxfight Jul 28, 2026
ee0037b
test(provider): cover embedding retry and pacing
lxfight Jul 28, 2026
5b2c374
fix(provider): honor cooldown without rpm limit
lxfight Jul 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions astrbot/core/config/default.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1841,6 +1841,8 @@
"embedding_model": "",
"embedding_dimensions": 1024,
"embedding_dimensions_mode": "auto",
"embedding_max_requests_per_minute": 120,
"embedding_rate_limit_cooldown": 120,
"timeout": 20,
"proxy": "",
},
Expand All @@ -1855,6 +1857,8 @@
"embedding_api_base": "",
"embedding_model": "gemini-embedding-exp-03-07",
"embedding_dimensions": 768,
"embedding_max_requests_per_minute": 120,
"embedding_rate_limit_cooldown": 120,
"timeout": 20,
"proxy": "",
},
Expand All @@ -1870,6 +1874,8 @@
"embedding_model": "nvidia/llama-nemotron-embed-1b-v2",
"input_type": "passage",
"embedding_dimensions": 1024,
"embedding_max_requests_per_minute": 120,
"embedding_rate_limit_cooldown": 120,
"timeout": 20,
"proxy": "",
},
Expand All @@ -1883,6 +1889,8 @@
"embedding_api_base": "http://localhost:11434",
"embedding_model": "nomic-embed-text",
"embedding_dimensions": 768,
"embedding_max_requests_per_minute": 0,
"embedding_rate_limit_cooldown": 120,
"timeout": 60,
"proxy": "",
},
Expand Down Expand Up @@ -2267,6 +2275,16 @@
"description": "API Base URL",
"type": "string",
},
"embedding_max_requests_per_minute": {
"description": "Embedding 每分钟最大请求数",
"type": "int",
"hint": "限制同一 Embedding Provider 的批量向量化请求速率。设置为 0 表示不限制；远程 API 建议按服务商限额填写。",
},
"embedding_rate_limit_cooldown": {
"description": "Embedding 限速冷却上限（秒）",
"type": "int",
"hint": "当 Embedding API 返回 Retry-After 或限速错误时，同一 Provider 后续请求的最大等待秒数。",
},
"volcengine_cluster": {
"type": "string",
"description": "火山引擎集群",
Expand Down
31 changes: 22 additions & 9 deletions astrbot/core/db/vec_db/faiss_impl/vec_db.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ async def insert(
metadata = metadata or {}
str_id = id or str(uuid.uuid4()) # 使用 UUID 作为原始 ID

vector = await self.embedding_provider.get_embedding(content)
vector = await self.embedding_provider.get_embedding_with_retry(content)
vector = np.array(vector, dtype=np.float32)

# 使用 DocumentStorage 的方法插入文档
Expand Down Expand Up @@ -109,13 +109,26 @@ async def insert_batch(

start = time.time()
logger.debug(f"Generating embeddings for {len(contents)} contents...")
vectors = await self.embedding_provider.get_embeddings_batch(
contents,
batch_size=batch_size,
tasks_limit=tasks_limit,
max_retries=max_retries,
progress_callback=progress_callback,
)
try:
vectors = await self.embedding_provider.get_embeddings_batch(
contents,
batch_size=batch_size,
tasks_limit=tasks_limit,
max_retries=max_retries,
progress_callback=progress_callback,
)
except KnowledgeBaseUploadError:
raise
except Exception as exc:
raise KnowledgeBaseUploadError(
stage="embedding",
user_message=(
"向量化失败：调用 Embedding API 生成向量时出错。"
"请检查 Embedding 服务是否可用、是否触发限流，"
"并尝试降低批量大小或并发数后重试。"
),
details={"cause": str(exc)},
) from exc
end = time.time()
logger.debug(
f"Generated embeddings for {len(contents)} contents in {end - start:.2f} seconds.",
Expand Down Expand Up @@ -217,7 +230,7 @@ async def retrieve(
List[Result]: 查询结果

"""
embedding = await self.embedding_provider.get_embedding(query)
embedding = await self.embedding_provider.get_embedding_with_retry(query)
scores, indices = await self.embedding_storage.search(
vector=np.array(embedding).astype("float32"),
k=fetch_k if metadata_filters else k,
Expand Down
Loading
Loading