diff --git a/astrbot/core/config/default.py b/astrbot/core/config/default.py index f8c46c26e4..c915d9fb10 100644 --- a/astrbot/core/config/default.py +++ b/astrbot/core/config/default.py @@ -2439,6 +2439,11 @@ def get_local_permission_defaults(system: str | None = None) -> dict: "options": ["auto", "always", "never"], "hint": "控制是否在 OpenAI 兼容 Embedding 请求中发送 dimensions 参数。auto 会仅对官方 OpenAI embedding-3 模型自动发送；第三方兼容 API 如需该参数可改为 always，报错时改为 never。", }, + "embedding_max_batch_items": { + "description": "单次请求最大文本数", + "type": "int", + "hint": "留空则自动判断。当嵌入服务限制单次请求可携带的文本数量时（如百炼 text-embedding-v3/v4 最多 10 条），填写该上限可避免请求被拒绝；已知上限的服务无需填写。", + }, "embedding_model": { "description": "嵌入模型", "type": "string", diff --git a/astrbot/core/db/vec_db/faiss_impl/vec_db.py b/astrbot/core/db/vec_db/faiss_impl/vec_db.py index 13fde6c0ef..594d4b2669 100644 --- a/astrbot/core/db/vec_db/faiss_impl/vec_db.py +++ b/astrbot/core/db/vec_db/faiss_impl/vec_db.py @@ -6,6 +6,7 @@ from astrbot import logger from astrbot.core.exceptions import KnowledgeBaseUploadError from astrbot.core.provider.provider import EmbeddingProvider, RerankProvider +from astrbot.core.utils.error_redaction import redact_sensitive_text from ..base import BaseVecDB, Result from .document_storage import DocumentStorage @@ -125,13 +126,38 @@ async def insert_batch( start = time.time() logger.debug(f"Generating embeddings for {len(contents)} contents...") - vectors = await self.embedding_provider.get_embeddings_batch( - embedding_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( + embedding_contents, + batch_size=batch_size, + tasks_limit=tasks_limit, + max_retries=max_retries, + progress_callback=progress_callback, + ) + except KnowledgeBaseUploadError: + raise + except Exception as exc: + # Attach the embedding stage here: otherwise the raw provider error + # escapes unlabelled and gets reported downstream as a storage + # failure, hiding the real cause (see kb_helper's stage mapping). + # Provider errors can quote the request they failed on, API key + # included ("Incorrect API key provided: sk-..."), and this text + # reaches the upload log and the dashboard's failure list, so redact + # it first. The summary shown to users is bounded; the log keeps the + # redacted error in full. + cause = redact_sensitive_text(str(exc)).strip() or type(exc).__name__ + summary = cause if len(cause) <= 300 else cause[:300] + "…" + raise KnowledgeBaseUploadError( + stage="embedding", + user_message=f"向量化失败：调用嵌入模型时出错。原因：{summary}", + details={ + "cause": cause, + "error_type": type(exc).__name__, + "provider": type(self.embedding_provider).__name__, + "content_count": content_count, + "batch_size": batch_size, + }, + ) from exc end = time.time() logger.debug( f"Generated embeddings for {len(contents)} contents in {end - start:.2f} seconds.", diff --git a/astrbot/core/knowledge_base/kb_helper.py b/astrbot/core/knowledge_base/kb_helper.py index 8a2fb32825..b2e76a994b 100644 --- a/astrbot/core/knowledge_base/kb_helper.py +++ b/astrbot/core/knowledge_base/kb_helper.py @@ -19,6 +19,7 @@ from astrbot.core.provider.provider import ( Provider as LLMProvider, ) +from astrbot.core.utils.error_redaction import redact_sensitive_text from .chunking.base import BaseChunker from .chunking.markdown import MarkdownChunker @@ -432,7 +433,9 @@ async def embedding_progress_callback(current, total) -> None: details={ "file_name": file_name, "doc_id": doc_id, - "cause": str(exc), + # Redacted: the cause is logged, and a provider error can + # quote the request it failed on, API key included. + "cause": redact_sensitive_text(str(exc)), }, ) from exc diff --git a/astrbot/core/provider/embedding_batch_limits.py b/astrbot/core/provider/embedding_batch_limits.py new file mode 100644 index 0000000000..276b9d0d14 --- /dev/null +++ b/astrbot/core/provider/embedding_batch_limits.py @@ -0,0 +1,85 @@ +"""Known per-request input limits of embedding services. + +Every embedding request in AstrBot funnels through +``EmbeddingProvider.get_embeddings_batch``. Some services reject a request that +carries more than N inputs with a deterministic client error, which retrying can +never overcome; the limit therefore has to be known *before* the request is +built. This module holds the small amount of metadata AstrBot actually knows +about those limits, so adapters can declare it cheaply instead of each carrying +its own copy. + +Kept dependency-free on purpose: ``provider/sources/dashscope_embedding_source`` +imports the dashscope SDK and registers an adapter at import time, which makes it +a poor place to share a plain lookup table from. +""" + +# Aliyun Model Studio (DashScope) caps the number of inputs accepted by one +# native embedding request. Exceeding it returns HTTP 400 +# "batch size is invalid, it should not be larger than 10." +_DASHSCOPE_MAX_ITEMS: dict[str, int] = { + "text-embedding-v1": 25, + "text-embedding-v2": 25, + "text-embedding-v3": 10, + "text-embedding-v4": 10, +} +# Conservative fallback for DashScope embedding models whose limit is not +# documented per generation (e.g. the multimodal ones). Under-estimating only +# costs extra requests; over-estimating reproduces the failure. +DASHSCOPE_DEFAULT_MAX_ITEMS = 10 + +# Model families served by DashScope that are embeddings but carry no per-model +# entry above. Matched on substrings because these names carry version suffixes. +_DASHSCOPE_EMBEDDING_NAME_HINTS = ( + "multimodal-embedding", + "vl-embedding", + "tongyi-embedding-vision", +) + +# Hosts serving the DashScope API, including the OpenAI-compatible mode which +# generic OpenAI endpoints are commonly pointed at. +_DASHSCOPE_HOSTS = frozenset( + { + "dashscope.aliyuncs.com", + "dashscope-intl.aliyuncs.com", + } +) + + +def dashscope_max_batch_items(model: str | None) -> int | None: + """Return the DashScope input limit for ``model``, or None if not applicable. + + None means "this is not a DashScope embedding model" -- deliberately not + "an embedding model with an unknown limit", so that a model name belonging + to another vendor is never clamped by this table. + """ + if not model: + return None + # Deployment prefixes are common ("models/text-embedding-v4"). + name = str(model).strip().lower().rsplit("/", 1)[-1] + if name in _DASHSCOPE_MAX_ITEMS: + return _DASHSCOPE_MAX_ITEMS[name] + if name.startswith("text-embedding-v"): + # A text-embedding generation newer than the table: assume the current + # (stricter) generation's limit rather than the legacy 25. + return DASHSCOPE_DEFAULT_MAX_ITEMS + if any(hint in name for hint in _DASHSCOPE_EMBEDDING_NAME_HINTS): + return DASHSCOPE_DEFAULT_MAX_ITEMS + return None + + +def is_dashscope_host(hostname: str | None) -> bool: + """Whether ``hostname`` points at the DashScope API.""" + if not hostname: + return False + return hostname.strip().lower() in _DASHSCOPE_HOSTS + + +def combine_caps(*caps: int | None) -> int | None: + """Combine declared limits into a single effective limit. + + Only known limits participate, and the result is their minimum: a limit + declared by the adapter can never be raised by a looser config value or by + a second detection path. + """ + known = [cap for cap in caps if isinstance(cap, int) and cap > 0] + return min(known) if known else None diff --git a/astrbot/core/provider/provider.py b/astrbot/core/provider/provider.py index 5e1274c24c..cb64849815 100644 --- a/astrbot/core/provider/provider.py +++ b/astrbot/core/provider/provider.py @@ -1,9 +1,11 @@ import abc import asyncio import os +import re from collections.abc import AsyncGenerator from typing import Literal, TypeAlias, Union +from astrbot import logger from astrbot.core.agent.message import ContentPart, Message, is_checkpoint_message from astrbot.core.agent.tool import ToolSet from astrbot.core.provider.entities import ( @@ -25,6 +27,129 @@ ] +# --------------------------------------------------------------------------- +# Embedding batch-size rejections +# +# A request refused because it carried too many inputs fails deterministically: +# the same request will be refused again however often it is retried. Such a +# failure has to be recognised so the caller shrinks the request (which is what +# adaptive splitting does) instead of spending the retry budget on it. All of +# the matching below is deliberately narrow -- a misread makes AstrBot split a +# batch that was never too large, so unrelated failures (token limits, rate +# limits, bad credentials) must never qualify. +# --------------------------------------------------------------------------- + +# Client error codes a provider may use to report "too many inputs in one request". +_BATCH_SIZE_STATUS_CODES = frozenset({400, 413, 414, 422}) + +# 4xx codes that are still worth retrying (timeouts / conflicts / throttling). +_RETRYABLE_CLIENT_STATUS_CODES = frozenset({408, 409, 425, 429}) + +_BATCH_SIZE_PATTERNS = tuple( + re.compile(pattern, re.IGNORECASE) + for pattern in ( + r"batch[\s_-]*(?:size|length|count|limit)", + r"too many (?:inputs?|texts?|items?|sentences?|contents?|documents?|chunks?|rows?)", + r"input (?:array|list|batch)[^.]{0,32}too (?:long|large|many)", + r"(?:maximum|max)\s+(?:number of\s+)?(?:inputs?|texts?|items?|rows?)", + r"range of inputs?(?: length)?[^.]{0,24}\[\s*\d+\s*,\s*\d+\s*\]", + # zh + r"批量[\s\S]{0,8}(?:超|不能|不得|最多|限制|过大)", + r"(?:单次|每次|一次)[\s\S]{0,8}(?:最多|不超过|不能超过|不得超过)\s*\d+", + ) +) + +# Per-item content limits: splitting the request would not help and would mask +# the real cause, so these disqualify a message even if it also mentions a size. +_BATCH_SIZE_EXCLUSIONS = tuple( + re.compile(pattern, re.IGNORECASE) + for pattern in ( + r"context[\s_-]*(?:length|window)", + r"max_tokens|maximum (?:output )?tokens|token limit|too many tokens", + r"rate limit|too many requests", + r"文本(?:过长|长度)|输入(?:过长|长度超)|内容(?:过长)", + ) +) + +_STATUS_PATTERN = re.compile( + r"(?:http|status(?:\s*code)?)\s*[:=]?\s*(\d{3})\b", + re.IGNORECASE, +) + +# "should not be larger than 10", "maximum of 10", "[1, 10]", "最多 10 条" +_LIMIT_HINT_PATTERNS = tuple( + re.compile(pattern, re.IGNORECASE) + for pattern in ( + r"(?:larger|greater|more|bigger)\s+than\s*[:：]?\s*(\d+)", + r"(?:at most|no more than|maximum of|max)\s*[:：]?\s*(\d+)", + r"\[\s*\d+\s*,\s*(\d+)\s*\]", + r"(?:不超过|最多|不能超过|不得超过)\s*(\d+)", + ) +) + + +def _exception_status_code(exc: BaseException) -> int | None: + """Best-effort HTTP status code for a provider failure. + + Covers both shapes used in this repo: SDK errors carrying ``status_code`` + (``openai.BadRequestError``) and plain ``Exception``s whose text embeds + ``(HTTP 400)`` (the DashScope adapter). + """ + for attr in ("status_code", "http_status", "status"): + value = getattr(exc, attr, None) + if isinstance(value, int) and 100 <= value < 600: + return value + response = getattr(exc, "response", None) + value = getattr(response, "status_code", None) + if isinstance(value, int) and 100 <= value < 600: + return value + match = _STATUS_PATTERN.search(str(exc)) + if match: + code = int(match.group(1)) + if 100 <= code < 600: + return code + return None + + +def _looks_like_batch_size_error(exc: BaseException) -> bool: + """Whether ``exc`` says the request carried too many inputs.""" + message = str(exc) + if not message: + return False + if any(pattern.search(message) for pattern in _BATCH_SIZE_EXCLUSIONS): + return False + status = _exception_status_code(exc) + if status is not None and status not in _BATCH_SIZE_STATUS_CODES: + return False + return any(pattern.search(message) for pattern in _BATCH_SIZE_PATTERNS) + + +def _extract_batch_limit(exc: BaseException, attempted: int) -> int | None: + """Read the provider's stated input limit out of a rejection message. + + Only a value that is a plausible limit for the request just refused + (``1 <= limit < attempted``) is returned, so a number quoted for any other + reason can never become the new batch size. + """ + message = str(exc) + for pattern in _LIMIT_HINT_PATTERNS: + for match in pattern.finditer(message): + limit = int(match.group(1)) + if 1 <= limit < attempted: + return limit + return None + + +def _is_retryable_error(exc: BaseException) -> bool: + """Whether re-issuing the identical request could plausibly succeed.""" + status = _exception_status_code(exc) + if status is None: + return True + if status >= 500 or status in _RETRYABLE_CLIENT_STATUS_CODES: + return True + return not 400 <= status < 500 + + class AbstractProvider(abc.ABC): """Provider Abstract Class""" @@ -326,6 +451,9 @@ def __init__(self, provider_config: dict, provider_settings: dict) -> None: super().__init__(provider_config) self.provider_config = provider_config self.provider_settings = provider_settings + # Upper bound on the inputs per request, once a provider has rejected a + # batch and stated it. None until then (see get_embeddings_batch). + self._learned_batch_size: int | None = None @abc.abstractmethod async def get_embedding(self, text: str) -> list[float]: @@ -345,6 +473,38 @@ def get_dim(self) -> int: async def test(self) -> None: await self.get_embedding("astrbot") + def get_max_batch_size(self) -> int | None: + """单次嵌入请求允许携带的最大文本数量。 + + Returning ``None`` (the default) means the limit is unknown, which keeps + the caller-supplied ``batch_size`` untouched. Adapters whose service + documents a fixed limit should override this; a generic OpenAI-compatible + endpoint whose limit cannot be inferred can declare it through the + ``embedding_max_batch_items`` provider config key instead. + + This is only a *hint* used to avoid a doomed request up front -- + ``get_embeddings_batch`` additionally recovers from a rejection at + runtime, so a wrong or missing value cannot break batching. + """ + raw = (getattr(self, "provider_config", None) or {}).get( + "embedding_max_batch_items" + ) + if raw is None or raw == "": + return None + try: + value = int(raw) + except (TypeError, ValueError): + logger.warning( + f"embedding_max_batch_items is not a valid integer: '{raw}', ignored." + ) + return None + if value <= 0: + logger.warning( + f"embedding_max_batch_items must be positive, got {value}, ignored." + ) + return None + return value + async def get_embeddings_batch( self, texts: list[str], @@ -357,7 +517,7 @@ async def get_embeddings_batch( Args: texts: 文本列表 - batch_size: 每批处理的文本数量 + batch_size: 每批处理的文本数量；超过提供商单次请求上限时会被自动压低 tasks_limit: 并发任务数量限制 max_retries: 失败时的最大重试次数 progress_callback: 进度回调函数，接收参数 (current, total) @@ -366,32 +526,92 @@ async def get_embeddings_batch( 向量列表 """ + if not texts: + return [] + + batch_size = max(1, int(batch_size)) + # A request that carries too many inputs is refused deterministically, so + # clamp to whatever limit the provider declares before building any work. + declared_cap = self.get_max_batch_size() + if declared_cap is not None and batch_size > declared_cap: + logger.debug( + f"[{type(self).__name__}] batch_size {batch_size} exceeds the " + f"provider limit {declared_cap}, using {declared_cap}." + ) + batch_size = declared_cap + + # A single rejection reveals the real limit, so record it and reuse it + # for the remaining batches of this call (and later calls on this + # instance). Learned values always come from a message that *stated* a + # limit; a bare rejection only tightens the size for the current call, + # so an unrelated error misread as a size rejection cannot permanently + # shrink this provider's batches. + learned_cap = getattr(self, "_learned_batch_size", None) + cap_state = {"size": batch_size} + if isinstance(learned_cap, int) and 0 < learned_cap < cap_state["size"]: + cap_state["size"] = learned_cap + semaphore = asyncio.Semaphore(tasks_limit) batch_results: dict[int, list[list[float]]] = {} failed_batches: list[tuple[int, list[str]]] = [] completed_count = 0 total_count = len(texts) + def note_limit(size: int) -> None: + if size < cap_state["size"]: + cap_state["size"] = size + + async def embed_slice(slice_texts: list[str]) -> list[list[float]]: + """Embed one slice, splitting it when the provider refuses its size.""" + if len(slice_texts) > cap_state["size"]: + # A limit learned earlier in this call already rules this out. + return await split_and_embed(slice_texts) + for attempt in range(max_retries): + try: + return await self.get_embeddings(slice_texts) + except Exception as e: + if _looks_like_batch_size_error(e) and len(slice_texts) > 1: + limit = _extract_batch_limit(e, len(slice_texts)) + if limit is not None: + self._learned_batch_size = limit + note_limit(limit) + else: + note_limit(len(slice_texts) // 2) + logger.debug( + f"[{type(self).__name__}] provider rejected {len(slice_texts)} " + f"inputs, retrying in smaller batches: {e!s}" + ) + return await split_and_embed(slice_texts) + if attempt == max_retries - 1 or not _is_retryable_error(e): + raise + # 等待一段时间后重试，使用指数退避 + await asyncio.sleep(2**attempt) + + async def split_and_embed(slice_texts: list[str]) -> list[list[float]]: + # Only ever called with len(slice_texts) > cap_state["size"], so a + # known limit is honoured exactly (32 inputs against a limit of 10 + # become 10/10/10/2, not four rounds of halving) and an unknown one + # halves. Every piece is strictly smaller than the slice, so the + # recursion terminates, and a single rejected text re-raises the + # provider's own error instead of splitting forever. + step = max(1, min(cap_state["size"], len(slice_texts) - 1)) + embeddings: list[list[float]] = [] + for start in range(0, len(slice_texts), step): + embeddings.extend(await embed_slice(slice_texts[start : start + step])) + return embeddings + async def process_batch(batch_idx: int, batch_texts: list[str]) -> None: nonlocal completed_count async with semaphore: - for attempt in range(max_retries): - try: - batch_embeddings = await self.get_embeddings(batch_texts) - batch_results[batch_idx] = batch_embeddings - completed_count += len(batch_texts) - if progress_callback: - await progress_callback(completed_count, total_count) - return - except Exception as e: - if attempt == max_retries - 1: - # 最后一次重试失败，记录失败的批次 - failed_batches.append((batch_idx, batch_texts)) - raise Exception( - f"批次 {batch_idx} 处理失败，已重试 {max_retries} 次: {e!s}", - ) - # 等待一段时间后重试，使用指数退避 - await asyncio.sleep(2**attempt) + try: + batch_embeddings = await embed_slice(batch_texts) + except Exception as e: + failed_batches.append((batch_idx, batch_texts)) + raise Exception(f"批次 {batch_idx} 处理失败: {e!s}") from e + batch_results[batch_idx] = batch_embeddings + completed_count += len(batch_texts) + if progress_callback: + await progress_callback(completed_count, total_count) tasks = [] for i in range(0, len(texts), batch_size): diff --git a/astrbot/core/provider/sources/dashscope_embedding_source.py b/astrbot/core/provider/sources/dashscope_embedding_source.py index 56fd14fb38..965c76ecbb 100644 --- a/astrbot/core/provider/sources/dashscope_embedding_source.py +++ b/astrbot/core/provider/sources/dashscope_embedding_source.py @@ -6,6 +6,11 @@ from astrbot import logger +from ..embedding_batch_limits import ( + DASHSCOPE_DEFAULT_MAX_ITEMS, + combine_caps, + dashscope_max_batch_items, +) from ..entities import ProviderType from ..provider import EmbeddingProvider from ..register import register_provider_adapter @@ -140,3 +145,16 @@ def get_dim(self) -> int: f"'{self.provider_config['embedding_dimensions']}', ignored." ) return 0 + + def get_max_batch_size(self) -> int | None: + """DashScope rejects a request carrying more inputs than the model allows. + + The limit is per model generation (see ``embedding_batch_limits``); + everything this adapter serves is a DashScope embedding model, so an + unrecognised name falls back to the conservative default rather than to + "unknown". A user-declared limit can only lower the result. + """ + return combine_caps( + dashscope_max_batch_items(self.model) or DASHSCOPE_DEFAULT_MAX_ITEMS, + super().get_max_batch_size(), + ) diff --git a/astrbot/core/provider/sources/openai_embedding_source.py b/astrbot/core/provider/sources/openai_embedding_source.py index 8f29017f4e..964353fc84 100644 --- a/astrbot/core/provider/sources/openai_embedding_source.py +++ b/astrbot/core/provider/sources/openai_embedding_source.py @@ -6,6 +6,11 @@ from astrbot import logger +from ..embedding_batch_limits import ( + combine_caps, + dashscope_max_batch_items, + is_dashscope_host, +) from ..entities import ProviderType from ..provider import EmbeddingProvider from ..register import register_provider_adapter @@ -122,6 +127,34 @@ def get_dim(self) -> int: ) return 0 + def get_max_batch_size(self) -> int | None: + """Declared per-request input limit, when the endpoint reveals one. + + Some services reached through this generic adapter cap the number of + inputs per request (DashScope returns HTTP 400 for more than 10 inputs + on text-embedding-v3/v4). The adapter cannot know that from its own + identity, so, like ``embedding_dimensions``, the limit is inferred from + the configured base URL host; anything else falls back to the + ``embedding_max_batch_items`` config key. + """ + detected = None + try: + api_base = _normalize_api_base( + self.provider_config.get( + "embedding_api_base", "https://api.openai.com/v1" + ) + or "https://api.openai.com/v1" + ) + hostname = urlparse(api_base).hostname + except (ValueError, AttributeError): + hostname = None + if is_dashscope_host(hostname): + detected = dashscope_max_batch_items( + getattr(self, "model", None) + or self.provider_config.get("embedding_model") + ) + return combine_caps(detected, super().get_max_batch_size()) + async def terminate(self): if self.client: await self.client.close() diff --git a/astrbot/core/utils/error_redaction.py b/astrbot/core/utils/error_redaction.py index dcab07ac58..559cddafb1 100644 --- a/astrbot/core/utils/error_redaction.py +++ b/astrbot/core/utils/error_redaction.py @@ -11,7 +11,8 @@ r"(?i)(?P(?P['\"])authorization(?P=kq)\s*:\s*)(?P['\"])bearer\s+[^'\"]+(?P=vq)" ) _QUERY_FIELD_PATTERN = re.compile( - rf"(?i)(?P{_SECRET_KEYS}\s*=\s*)(?P[^&'\" ]+)" + rf"(?i)(?P{_SECRET_KEYS}\s*=\s*)" + r"(?P(?P['\"])[^'\"]+(?P=vq)|[^&'\" ]+)" ) _QUERY_PARAM_PATTERN = re.compile( r"(?i)(?P[?&](?:api_?key|key|access_?token|auth_?token)=)(?P[^&'\" ]+)" @@ -20,7 +21,10 @@ r"(?i)(?P\bauthorization\s*:\s*bearer\s+)(?P[A-Za-z0-9._\-]+)" ) _BEARER_PATTERN = re.compile(r"(?i)(?P\bbearer\s+)(?P[A-Za-z0-9._\-]+)") -_SK_PATTERN = re.compile(r"\bsk-[A-Za-z0-9]{16,}\b") +# Vendor keys are "sk-" plus a long body that may itself be dash-segmented, +# e.g. "sk-proj-..." (OpenAI) or "sk-ant-api03-..." (Anthropic). Requiring a +# minimum length keeps short placeholders like "sk-test" untouched. +_SK_PATTERN = re.compile(r"\bsk-[A-Za-z0-9_-]{16,}\b") def _redact_json_field(match: re.Match[str]) -> str: diff --git a/tests/unit/test_embedding_batch_split.py b/tests/unit/test_embedding_batch_split.py new file mode 100644 index 0000000000..ea39eeb6b9 --- /dev/null +++ b/tests/unit/test_embedding_batch_split.py @@ -0,0 +1,516 @@ +"""Embedding requests must survive a provider's per-request input limit. + +Regression coverage for the knowledge-base upload failure where a provider +rejects a batch that carries too many inputs (DashScope text-embedding-v3/v4 +return HTTP 400 "batch size is invalid, it should not be larger than 10.") and +the retry loop can never recover, aborting the whole upload. +""" + +from pathlib import Path + +import pytest + +from astrbot.core.db.vec_db.faiss_impl.vec_db import FaissVecDB +from astrbot.core.exceptions import KnowledgeBaseUploadError +from astrbot.core.provider.embedding_batch_limits import ( + combine_caps, + dashscope_max_batch_items, + is_dashscope_host, +) +from astrbot.core.provider.provider import ( + EmbeddingProvider, + _extract_batch_limit, + _is_retryable_error, + _looks_like_batch_size_error, +) + + +class LimitedEmbeddingProvider(EmbeddingProvider): + """Provider that refuses any request carrying more than ``limit`` inputs.""" + + def __init__(self, limit: int, *, state_the_limit: bool = True) -> None: + super().__init__({}, {}) + self.limit = limit + self.state_the_limit = state_the_limit + self.calls: list[list[str]] = [] + + async def get_embedding(self, text: str) -> list[float]: + return (await self.get_embeddings([text]))[0] + + async def get_embeddings(self, text: list[str]) -> list[list[float]]: + self.calls.append(list(text)) + if len(text) > self.limit: + hint = ( + f"batch size is invalid, it should not be larger than {self.limit}." + if self.state_the_limit + else "batch size is invalid." + ) + raise Exception( + f"DashScope Embedding API request failed (HTTP 400): " + f"InvalidParameter - {hint}" + ) + return [[float(item)] for item in text] + + def get_dim(self) -> int: + return 1 + + +def _texts(count: int) -> list[str]: + return [str(i) for i in range(count)] + + +def _embedding_lengths(provider: LimitedEmbeddingProvider) -> list[int]: + return [len(call) for call in provider.calls] + + +# --------------------------------------------------------------------------- +# Declared limits +# --------------------------------------------------------------------------- + + +class DeclaredCapProvider(LimitedEmbeddingProvider): + def get_max_batch_size(self) -> int | None: + return self.limit + + +@pytest.mark.asyncio +async def test_declared_limit_prevents_the_oversized_request() -> None: + provider = DeclaredCapProvider(10) + + embeddings = await provider.get_embeddings_batch( + _texts(25), batch_size=32, tasks_limit=2 + ) + + assert embeddings == [[float(i)] for i in range(25)] + # Clamped up front, so no request is ever refused. + assert max(_embedding_lengths(provider)) == 10 + + +@pytest.mark.asyncio +async def test_provider_config_can_declare_the_limit() -> None: + provider = LimitedEmbeddingProvider(10) + provider.provider_config = {"embedding_max_batch_items": 4} + + assert provider.get_max_batch_size() == 4 + + embeddings = await provider.get_embeddings_batch(_texts(12), batch_size=32) + assert len(embeddings) == 12 + assert max(_embedding_lengths(provider)) == 4 + + +@pytest.mark.asyncio +async def test_invalid_declared_limit_falls_back_to_unknown() -> None: + provider = LimitedEmbeddingProvider(10) + provider.provider_config = {"embedding_max_batch_items": "many"} + + assert provider.get_max_batch_size() is None + + +# --------------------------------------------------------------------------- +# Adaptive splitting +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_rejected_batch_is_split_and_results_stay_ordered() -> None: + provider = LimitedEmbeddingProvider(10) + + embeddings = await provider.get_embeddings_batch( + _texts(32), batch_size=32, tasks_limit=1 + ) + + assert embeddings == [[float(i)] for i in range(32)] + assert _embedding_lengths(provider)[0] == 32 # the doomed request + assert max(_embedding_lengths(provider)) == 32 + + +@pytest.mark.asyncio +async def test_stated_limit_is_honoured_exactly() -> None: + provider = LimitedEmbeddingProvider(10) + + await provider.get_embeddings_batch(_texts(32), batch_size=32, tasks_limit=1) + + # The provider stated "not be larger than 10", so the retry uses 10/10/10/2 + # rather than halving down to 16/8. + assert provider._learned_batch_size == 10 + assert _embedding_lengths(provider) == [32, 10, 10, 10, 2] + + +@pytest.mark.asyncio +async def test_unstated_limit_halves_until_it_fits() -> None: + provider = LimitedEmbeddingProvider(3, state_the_limit=False) + + embeddings = await provider.get_embeddings_batch( + _texts(8), batch_size=8, tasks_limit=1 + ) + + assert embeddings == [[float(i)] for i in range(8)] + # Nothing was stated, so the reduced size must not outlive this call. + assert provider._learned_batch_size is None + assert _embedding_lengths(provider)[0] == 8 + assert max(_embedding_lengths(provider)) == 8 + + +@pytest.mark.asyncio +async def test_learned_limit_is_reused_by_later_calls() -> None: + provider = LimitedEmbeddingProvider(10) + + await provider.get_embeddings_batch(_texts(32), batch_size=32, tasks_limit=1) + provider.calls.clear() + + embeddings = await provider.get_embeddings_batch( + _texts(20), batch_size=32, tasks_limit=1 + ) + + assert len(embeddings) == 20 + assert _embedding_lengths(provider) == [10, 10] # no second wasted attempt at 32 + + +@pytest.mark.asyncio +async def test_progress_callback_reports_the_full_batch() -> None: + provider = LimitedEmbeddingProvider(10) + seen: list[tuple[int, int]] = [] + + async def callback(current: int, total: int) -> None: + seen.append((current, total)) + + await provider.get_embeddings_batch( + _texts(32), batch_size=32, tasks_limit=1, progress_callback=callback + ) + + assert seen == [(32, 32)] + + +@pytest.mark.asyncio +async def test_retryable_failures_are_still_retried() -> None: + class FlakyProvider(LimitedEmbeddingProvider): + def __init__(self) -> None: + super().__init__(100) + self.failures = 2 + + async def get_embeddings(self, text: list[str]) -> list[list[float]]: + self.calls.append(list(text)) + if self.failures: + self.failures -= 1 + raise Exception("HTTP 429: rate limit exceeded, please retry") + return [[float(item)] for item in text] + + provider = FlakyProvider() + + embeddings = await provider.get_embeddings_batch(_texts(2), batch_size=2) + + assert embeddings == [[0.0], [1.0]] + assert len(provider.calls) == 1 + 2 # the two throttled attempts plus the success + + +@pytest.mark.asyncio +async def test_a_single_rejected_text_surfaces_the_provider_error() -> None: + provider = LimitedEmbeddingProvider(0) + + with pytest.raises(Exception, match="有 1 个批次处理失败.*batch size is invalid"): + await provider.get_embeddings_batch(["only"], batch_size=1, tasks_limit=1) + + +# --------------------------------------------------------------------------- +# Failure classification +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "message", + [ + "DashScope Embedding API request failed (HTTP 400): InvalidParameter - " + "batch size is invalid, it should not be larger than 10.", + "Error code: 400 - {'error': {'message': 'Too many inputs provided'}}", + "HTTP 422: batch length exceeds the maximum of 25", + "请求失败：单次最多 10 条", + ], +) +def test_batch_size_rejections_are_recognised(message: str) -> None: + assert _looks_like_batch_size_error(Exception(message)) + + +@pytest.mark.parametrize( + "message", + [ + # Per-item content limits: splitting the request cannot help. + "HTTP 400: context length must not be larger than 8192", + "HTTP 400: max_tokens is invalid, batch size is fine otherwise", + "HTTP 400: 文本过长，单次批量过大", + # Not size related at all. + "HTTP 401: Incorrect API key provided", + "HTTP 404: model not found", + "HTTP 429: too many requests, batch size limit reached", + "HTTP 500: internal server error", + ], +) +def test_unrelated_failures_are_not_mistaken_for_size_rejections(message: str) -> None: + assert not _looks_like_batch_size_error(Exception(message)) + + +def test_sdk_errors_carrying_a_status_code_are_recognised() -> None: + """The OpenAI-compatible path raises ``openai.BadRequestError``.""" + import httpx + from openai import BadRequestError + + def bad_request(message: str) -> BadRequestError: + return BadRequestError( + message, + response=httpx.Response( + 400, request=httpx.Request("POST", "https://example.com/v1/embeddings") + ), + body=None, + ) + + error = bad_request( + "Error code: 400 - {'error': {'message': 'batch size is invalid, " + "it should not be larger than 10.'}}" + ) + + assert _looks_like_batch_size_error(error) + assert _extract_batch_limit(error, attempted=32) == 10 + assert not _is_retryable_error(error) + assert not _looks_like_batch_size_error(bad_request("Error code: 400 - bad key")) + + +def test_stated_limit_is_extracted_only_when_plausible() -> None: + exc = Exception("HTTP 400: it should not be larger than 10.") + + assert _extract_batch_limit(exc, attempted=32) == 10 + # A limit that is not smaller than the request just refused cannot explain it. + assert _extract_batch_limit(exc, attempted=10) is None + assert ( + _extract_batch_limit(Exception("HTTP 400: bad request"), attempted=32) is None + ) + + +@pytest.mark.parametrize( + ("message", "retryable"), + [ + ("HTTP 429: rate limit", True), + ("HTTP 503: service unavailable", True), + ("connection reset by peer", True), + ("HTTP 400: invalid request", False), + ("HTTP 401: invalid api key", False), + ], +) +def test_retryability(message: str, retryable: bool) -> None: + assert _is_retryable_error(Exception(message)) is retryable + + +# --------------------------------------------------------------------------- +# Adapter-declared limits +# --------------------------------------------------------------------------- + + +def test_dashscope_limits_are_per_model() -> None: + assert dashscope_max_batch_items("text-embedding-v4") == 10 + assert dashscope_max_batch_items("text-embedding-v3") == 10 + assert dashscope_max_batch_items("text-embedding-v2") == 25 + assert dashscope_max_batch_items("qwen3-vl-embedding") == 10 + assert dashscope_max_batch_items("text-embedding-3-small") is None + assert dashscope_max_batch_items(None) is None + + +def test_dashscope_host_detection() -> None: + assert is_dashscope_host("dashscope.aliyuncs.com") + assert is_dashscope_host("dashscope-intl.aliyuncs.com") + assert not is_dashscope_host("api.openai.com") + assert not is_dashscope_host(None) + + +def test_combined_caps_only_ever_lower() -> None: + assert combine_caps(None, 10) == 10 + assert combine_caps(25, 10) == 10 + assert combine_caps(None, None) is None + + +def test_dashscope_adapter_declares_the_native_limit() -> None: + from astrbot.core.provider.sources.dashscope_embedding_source import ( + DashScopeEmbeddingProvider, + ) + + def make(model: str, **extra): + provider = DashScopeEmbeddingProvider.__new__(DashScopeEmbeddingProvider) + provider.model = model + provider.provider_config = {"embedding_model": model, **extra} + return provider + + assert make("text-embedding-v4").get_max_batch_size() == 10 + assert make("text-embedding-v2").get_max_batch_size() == 25 + # The multimodal models have no documented per-model limit: stay conservative. + assert make("qwen3-vl-embedding").get_max_batch_size() == 10 + # A user-declared limit can only lower the result. + assert ( + make("text-embedding-v2", embedding_max_batch_items=4).get_max_batch_size() == 4 + ) + + +def test_openai_compatible_adapter_sniffs_dashscope_hosts() -> None: + from astrbot.core.provider.sources.openai_embedding_source import ( + OpenAIEmbeddingProvider, + ) + + def make(api_base: str, model: str, **extra): + provider = OpenAIEmbeddingProvider.__new__(OpenAIEmbeddingProvider) + provider.model = model + provider.provider_config = { + "embedding_api_base": api_base, + "embedding_model": model, + **extra, + } + return provider + + dashscope_openai_mode = "https://dashscope.aliyuncs.com/compatible-mode/v1" + assert make(dashscope_openai_mode, "text-embedding-v4").get_max_batch_size() == 10 + # Any other OpenAI-compatible gateway is unknown unless configured. + assert ( + make("https://api.openai.com/v1", "text-embedding-3-small").get_max_batch_size() + is None + ) + assert ( + make( + "https://api.openai.com/v1", + "text-embedding-3-small", + embedding_max_batch_items=8, + ).get_max_batch_size() + == 8 + ) + + +@pytest.mark.asyncio +async def test_openai_compatible_dashscope_endpoint_is_clamped_end_to_end() -> None: + """The second reported setup: DashScope through the OpenAI-compatible mode.""" + from unittest.mock import AsyncMock, MagicMock + + from astrbot.core.provider.sources.openai_embedding_source import ( + OpenAIEmbeddingProvider, + ) + + provider = OpenAIEmbeddingProvider.__new__(OpenAIEmbeddingProvider) + provider.model = "text-embedding-v4" + provider.provider_config = { + "embedding_api_base": "https://dashscope.aliyuncs.com/compatible-mode/v1", + "embedding_model": "text-embedding-v4", + } + received: list[int] = [] + + async def create(*, input, model): # noqa: A002 - mirrors the SDK signature + received.append(len(input)) + response = MagicMock() + response.data = [MagicMock(embedding=[0.0]) for _ in input] + return response + + provider.client = MagicMock() + provider.client.embeddings.create = AsyncMock(side_effect=create) + + embeddings = await provider.get_embeddings_batch( + _texts(25), batch_size=32, tasks_limit=3 + ) + + assert len(embeddings) == 25 + assert received == [10, 10, 5] # never 32, and no rejected request + + +# --------------------------------------------------------------------------- +# Persistence into a real vector store +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_insert_batch_succeeds_against_a_limited_provider(tmp_path: Path) -> None: + """End-to-end: a knowledge-base insert must not fail because of the limit.""" + dim = 4 + + class CappedProvider(LimitedEmbeddingProvider): + async def get_embeddings(self, text: list[str]) -> list[list[float]]: + if len(text) > self.limit: + return await LimitedEmbeddingProvider.get_embeddings(self, text) + return [[0.1, 0.2, 0.3, 0.4] for _ in text] + + def get_dim(self) -> int: + return dim + + provider = CappedProvider(10) + vec_db = FaissVecDB( + doc_store_path=str(tmp_path / "doc.db"), + index_store_path=str(tmp_path / "index.faiss"), + embedding_provider=provider, + ) + await vec_db.initialize() + try: + contents = _texts(40) + await vec_db.insert_batch( + contents=contents, + metadatas=[{"kb_id": "kb-1"} for _ in contents], + ids=[f"id-{i}" for i in range(len(contents))], + batch_size=32, # the dashboard default that triggered the bug + ) + + assert await vec_db.count_documents() == 40 + assert max(_embedding_lengths(provider)) == 32 # only the refused attempt + finally: + await vec_db.close() + + +@pytest.mark.asyncio +async def test_embedding_failure_is_reported_as_an_embedding_error( + tmp_path: Path, +) -> None: + """A provider failure must not be relabelled as a storage failure.""" + vec_db = _vec_db_whose_provider_raises("HTTP 401: invalid api key") + + with pytest.raises(KnowledgeBaseUploadError) as exc_info: + await FaissVecDB.insert_batch( + vec_db, + contents=["chunk-1"], + metadatas=[{}], + ids=["doc-1"], + ) + + error = exc_info.value + assert error.stage == "embedding" + assert "向量化失败" in error.user_message + assert "invalid api key" in error.user_message + assert error.details["cause"] == "HTTP 401: invalid api key" + + +@pytest.mark.asyncio +async def test_provider_error_text_is_redacted() -> None: + """Provider errors quote the request, so the API key must not be carried on. + + The message reaches the upload log and the dashboard's failure list. + """ + key = "sk-proj-abcdefghijklmnopqrstuvwxyz012345" + vec_db = _vec_db_whose_provider_raises( + f"HTTP 401: Incorrect API key provided: {key}. You can find your API key " + f"at https://platform.openai.com/account/api-keys." + ) + + with pytest.raises(KnowledgeBaseUploadError) as exc_info: + await FaissVecDB.insert_batch( + vec_db, + contents=["chunk-1"], + metadatas=[{}], + ids=["doc-1"], + ) + + error = exc_info.value + assert key not in error.user_message + assert key not in error.details["cause"] + assert "[REDACTED]" in error.user_message + # Everything that is not a secret is still reported. + assert "Incorrect API key provided" in error.user_message + + +def _vec_db_whose_provider_raises(message: str) -> FaissVecDB: + vec_db = FaissVecDB.__new__(FaissVecDB) + vec_db.embedding_provider = LimitedEmbeddingProvider(10) + + async def _call(*args, **kwargs): + raise Exception(message) + + vec_db.embedding_provider.get_embeddings_batch = _call + vec_db.document_storage = None + vec_db.embedding_storage = None + return vec_db diff --git a/tests/unit/test_error_redaction.py b/tests/unit/test_error_redaction.py new file mode 100644 index 0000000000..e3b0c8d233 --- /dev/null +++ b/tests/unit/test_error_redaction.py @@ -0,0 +1,68 @@ +"""Error text reaching logs and the dashboard must not carry credentials.""" + +import pytest + +from astrbot.core.utils.error_redaction import redact_sensitive_text, safe_error + + +@pytest.mark.parametrize( + "key", + [ + "sk-abcdefghijklmnopqrstuvwxyz012345", # legacy OpenAI + "sk-proj-abcdefghijklmnopqrstuvwxyz012345", # project keys + "sk-ant-api03-abcdefghijklmnopqrstuvwxyz", # Anthropic + "sk-1234567890abcdef_-1234567890", # dashsk-_ bodies + ], +) +def test_vendor_keys_are_redacted(key: str) -> None: + redacted = redact_sensitive_text(f"401 Incorrect API key provided: {key}") + + assert key not in redacted + assert "[REDACTED]" in redacted + + +@pytest.mark.parametrize( + "message", + [ + 'api_key="supersecretvalue"', + "api_key=supersecretvalue&model=x", + "https://example.com/v1?key=supersecretvalue", + "Authorization: Bearer abcdefghijklmnop", + "bearer abcdefghijklmnop", + 'authorization: "Bearer abcdefghijklmnop"', + ], +) +def test_named_credentials_are_redacted(message: str) -> None: + redacted = redact_sensitive_text(message) + + assert "supersecretvalue" not in redacted + assert "abcdefghijklmnop" not in redacted + assert "[REDACTED]" in redacted + + +def test_short_placeholders_are_left_alone() -> None: + """A short stand-in is not a credential and should stay readable in logs.""" + message = "using access_token=sk-test for the sandbox" + + assert ( + redact_sensitive_text(message) + == "using access_token=[REDACTED] for the sandbox" + ) + + +def test_safe_error_keeps_the_prefix_and_redacts_the_body() -> None: + assert ( + safe_error( + "request failed: ", + Exception( + "401 Incorrect API key provided: sk-proj-abcdefghijklmnopqrstuvwxyz012345" + ), + ) + == "request failed: 401 Incorrect API key provided: [REDACTED]" + ) + + +def test_safe_error_can_be_asked_not_to_redact() -> None: + assert safe_error("", "api_key=supersecretvalue", redact=False) == ( + "api_key=supersecretvalue" + ) diff --git a/tests/unit/test_kb_upload_atomicity.py b/tests/unit/test_kb_upload_atomicity.py index c7ad2d750b..539ce6ae23 100644 --- a/tests/unit/test_kb_upload_atomicity.py +++ b/tests/unit/test_kb_upload_atomicity.py @@ -349,6 +349,59 @@ async def fake_save_media(**kwargs): assert not media_file.exists() +@pytest.mark.asyncio +async def test_storage_failure_redacts_secrets_from_details( + tmp_path: Path, + stub_provider_manager_module, +) -> None: + """The cause is logged, so it must not carry the provider's API key.""" + KBHelper = _import_kb_helper() + + helper = KBHelper.__new__(KBHelper) + helper.kb = KnowledgeBase( + kb_name="Test KB", + description="", + embedding_provider_id="emb", + ) + helper.kb_db = MagicMock() + helper.vec_db = AsyncMock() + helper.kb_medias_dir = tmp_path / "medias" + helper.kb_medias_dir.mkdir() + helper.chunker = AsyncMock() + helper.chunker.chunk = AsyncMock(return_value=["hello world"]) + helper._save_media = AsyncMock(return_value=None) + helper.vec_db.insert_batch.side_effect = RuntimeError( + "401 Incorrect API key provided: sk-proj-abcdefghijklmnopqrstuvwxyz012345" + ) + helper.vec_db.delete_documents = AsyncMock() + helper.kb_db.get_db = _successful_get_db(_session_with_begin()) + + parse_result = MagicMock() + parse_result.text = "hello world" + parse_result.media = [] + + with ( + patch( + "astrbot.core.knowledge_base.kb_helper.select_parser", + new=AsyncMock( + return_value=MagicMock(parse=AsyncMock(return_value=parse_result)), + ), + ), + patch.object(helper, "_ensure_vec_db", new=AsyncMock()), + pytest.raises(KnowledgeBaseUploadError) as exc_info, + ): + await helper.upload_document( + file_name="demo.txt", + file_content=b"hello world", + file_type="txt", + ) + + assert exc_info.value.stage == "storage" + cause = exc_info.value.details["cause"] + assert "sk-proj-abcdefghijklmnopqrstuvwxyz012345" not in cause + assert "[REDACTED]" in cause + + @pytest.mark.asyncio @pytest.mark.parametrize( ("file_name", "file_type"),