From 5a99877fddfb033d811e5ac4f6398a71fc80e0ee Mon Sep 17 00:00:00 2001 From: "zhanghe.1216" <@bytedance.com> Date: Mon, 12 Oct 2026 20:07:45 +0800 Subject: [PATCH 01/16] refactor(vlm): centralize retry handling --- openviking/models/vlm/__init__.py | 2 + openviking/models/vlm/backends/litellm_vlm.py | 30 +-- openviking/models/vlm/backends/openai_vlm.py | 32 +-- .../models/vlm/backends/volcengine_vlm.py | 32 ++- openviking/models/vlm/base.py | 11 +- openviking/models/vlm/retry_vlm.py | 186 ++++++++++++++++++ openviking/utils/model_retry.py | 64 +++++- openviking_cli/utils/config/vlm_config.py | 4 + tests/unit/test_extra_headers_vlm.py | 2 +- tests/unit/test_kimi_glm_vlm.py | 4 +- tests/unit/test_model_retry.py | 33 ++++ tests/unit/test_retry_vlm.py | 167 ++++++++++++++++ 12 files changed, 490 insertions(+), 77 deletions(-) create mode 100644 openviking/models/vlm/retry_vlm.py create mode 100644 tests/unit/test_retry_vlm.py diff --git a/openviking/models/vlm/__init__.py b/openviking/models/vlm/__init__.py index 8d761ec5ea..b4f3f265a2 100644 --- a/openviking/models/vlm/__init__.py +++ b/openviking/models/vlm/__init__.py @@ -16,10 +16,12 @@ VLMFactory, ) from .registry import get_all_provider_names, is_valid_provider +from .retry_vlm import RetryVLM __all__ = [ "VLMBase", "VLMFactory", + "RetryVLM", "FailoverVLM", "MultiCredentialVLM", "AllCredentialsFailedError", diff --git a/openviking/models/vlm/backends/litellm_vlm.py b/openviking/models/vlm/backends/litellm_vlm.py index ea8ea7dbc7..699ad343ff 100644 --- a/openviking/models/vlm/backends/litellm_vlm.py +++ b/openviking/models/vlm/backends/litellm_vlm.py @@ -18,7 +18,6 @@ from openviking.telemetry import tracer from openviking.utils.message_format import format_messages, sanitize_openai_messages -from openviking.utils.model_retry import retry_async, retry_sync from openviking.utils.multimodal import redact_image_data_urls from openviking_cli.utils import get_logger @@ -293,6 +292,7 @@ def _build_kwargs( "messages": sanitize_openai_messages(messages), "temperature": self.temperature, "timeout": self.timeout, + "num_retries": 0, } effective_max_tokens = max_tokens if max_tokens is not None else self.max_tokens if effective_max_tokens is not None: @@ -475,12 +475,7 @@ def _call() -> Union[str, VLMResponse]: return self._build_vlm_response(response, has_tools=True) return self._clean_response(self._extract_content_from_response(response)) - return retry_sync( - _call, - max_retries=self.max_retries, - logger=logger, - operation_name="LiteLLM VLM completion", - ) + return _call() async def get_completion_async( self, @@ -511,12 +506,7 @@ async def _call() -> Union[str, VLMResponse]: return self._build_vlm_response(response, has_tools=True) return self._clean_response(self._extract_content_from_response(response)) - return await retry_async( - _call, - max_retries=self.max_retries, - logger=logger, - operation_name="LiteLLM VLM async completion", - ) + return await _call() def get_vision_completion( self, @@ -539,12 +529,7 @@ def _call() -> Union[str, VLMResponse]: return self._build_vlm_response(response, has_tools=True) return self._clean_response(self._extract_content_from_response(response)) - return retry_sync( - _call, - max_retries=self.max_retries, - logger=logger, - operation_name="LiteLLM VLM vision completion", - ) + return _call() async def get_vision_completion_async( self, @@ -567,12 +552,7 @@ async def _call() -> Union[str, VLMResponse]: return self._build_vlm_response(response, has_tools=True) return self._clean_response(self._extract_content_from_response(response)) - return await retry_async( - _call, - max_retries=self.max_retries, - logger=logger, - operation_name="LiteLLM VLM async vision completion", - ) + return await _call() def _update_token_usage_from_response( self, diff --git a/openviking/models/vlm/backends/openai_vlm.py b/openviking/models/vlm/backends/openai_vlm.py index 6d3f86b851..7f6023eadf 100644 --- a/openviking/models/vlm/backends/openai_vlm.py +++ b/openviking/models/vlm/backends/openai_vlm.py @@ -27,8 +27,6 @@ except ImportError: openai = None -from openviking.utils.model_retry import retry_async, retry_sync - from ..base import ToolCall, VLMBase, VLMResponse from ..registry import DEFAULT_AZURE_API_VERSION @@ -67,7 +65,7 @@ def _build_openai_client_kwargs( else: kwargs = {"api_key": api_key, "base_url": api_base} kwargs["timeout"] = timeout - # OpenViking owns provider retry/backoff via retry_sync/retry_async. + # RetryVLM owns retry/backoff; the SDK performs exactly one request. kwargs["max_retries"] = 0 if extra_headers: kwargs["default_headers"] = extra_headers @@ -358,12 +356,7 @@ def _call() -> Union[str, VLMResponse]: return self._build_vlm_response(response, has_tools=True) return self._extract_completion_content(response, elapsed) - return retry_sync( - _call, - max_retries=self.max_retries, - logger=logger, - operation_name="OpenAI VLM completion", - ) + return _call() async def get_completion_async( self, @@ -396,12 +389,7 @@ async def _call() -> Union[str, VLMResponse]: + format_messages(redact_image_data_urls(kwargs.get("messages", []))) ) - return await retry_async( - _call, - max_retries=self.max_retries, - logger=logger, - operation_name="OpenAI VLM async completion", - ) + return await _call() def _detect_image_format(self, data: bytes) -> str: """Detect image format from magic bytes. @@ -479,12 +467,7 @@ def _call() -> Union[str, VLMResponse]: return self._build_vlm_response(response, has_tools=True) return self._extract_completion_content(response, elapsed) - return retry_sync( - _call, - max_retries=self.max_retries, - logger=logger, - operation_name="OpenAI VLM vision completion", - ) + return _call() async def get_vision_completion_async( self, @@ -511,9 +494,4 @@ async def _call() -> Union[str, VLMResponse]: return self._build_vlm_response(response, has_tools=True) return await self._extract_completion_content_async(response, elapsed) - return await retry_async( - _call, - max_retries=self.max_retries, - logger=logger, - operation_name="OpenAI VLM async vision completion", - ) + return await _call() diff --git a/openviking/models/vlm/backends/volcengine_vlm.py b/openviking/models/vlm/backends/volcengine_vlm.py index 099cff78e1..977b0b225a 100644 --- a/openviking/models/vlm/backends/volcengine_vlm.py +++ b/openviking/models/vlm/backends/volcengine_vlm.py @@ -15,7 +15,6 @@ ) from openviking.telemetry import tracer from openviking.utils.message_format import format_messages, sanitize_openai_messages -from openviking.utils.model_retry import retry_async from openviking.utils.multimodal import redact_image_data_urls from openviking_cli.utils import get_logger @@ -120,12 +119,12 @@ def get_client(self): raise ImportError( "Please install volcenginesdkarkruntime: pip install volcenginesdkarkruntime" ) - kwargs = dict( - api_key=self.api_key, - base_url=self.api_base, - timeout=self.timeout, - max_retries=0, - ) + kwargs = { + "api_key": self.api_key, + "base_url": self.api_base, + "timeout": self.timeout, + "max_retries": 0, + } http_client = create_optional_sync_httpx_client( self.api_base, timeout=self.timeout, @@ -143,12 +142,12 @@ def _build_async_client(self): raise ImportError( "Please install volcenginesdkarkruntime: pip install volcenginesdkarkruntime" ) - kwargs = dict( - api_key=self.api_key, - base_url=self.api_base, - timeout=self.timeout, - max_retries=0, - ) + kwargs = { + "api_key": self.api_key, + "base_url": self.api_base, + "timeout": self.timeout, + "max_retries": 0, + } http_client = create_optional_async_httpx_client( self.api_base, timeout=self.timeout, @@ -293,12 +292,7 @@ async def _call() -> Union[str, VLMResponse]: tracer.info(f"message.content={content}") return content - return await retry_async( - _call, - max_retries=self.max_retries, - logger=logger, - operation_name="VolcEngine VLM async completion", - ) + return await _call() def _detect_image_format(self, data: bytes) -> str: """Detect image format from magic bytes. diff --git a/openviking/models/vlm/base.py b/openviking/models/vlm/base.py index ec52d012d4..70743d11cc 100644 --- a/openviking/models/vlm/base.py +++ b/openviking/models/vlm/base.py @@ -346,8 +346,8 @@ class VLMFactory: """VLM factory class, creates corresponding VLM instance based on config""" @staticmethod - def create(config: Dict[str, Any]) -> VLMBase: - """Create VLM instance + def create_provider(config: Dict[str, Any]) -> VLMBase: + """Create one concrete VLM provider without an outer retry decorator. Args: config: VLM config, must contain 'provider' field @@ -391,6 +391,13 @@ def create(config: Dict[str, Any]) -> VLMBase: return LiteLLMVLMProvider(config) + @staticmethod + def create(config: Dict[str, Any]) -> VLMBase: + """Create one VLM target with the standard retry owner.""" + from .retry_vlm import RetryVLM + + return RetryVLM(VLMFactory.create_provider(config)) + def _annotate_vlm_error(exc: Exception, vlm_instance: "VLMBase") -> None: """Attach model and api_base info to an exception for better error diagnostics. diff --git a/openviking/models/vlm/retry_vlm.py b/openviking/models/vlm/retry_vlm.py new file mode 100644 index 0000000000..7828ad428f --- /dev/null +++ b/openviking/models/vlm/retry_vlm.py @@ -0,0 +1,186 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Retry decorator for ordinary VLM requests.""" + +from pathlib import Path +from typing import Any, Dict, List, Optional, Union + +from openviking.utils.model_retry import retry_async, retry_sync +from openviking_cli.utils import get_logger + +from .base import VLMBase, VLMResponse + +logger = get_logger(__name__) + + +class RetryVLM(VLMBase): + """Own ordinary request retries for one concrete VLM target.""" + + def __init__(self, inner: VLMBase): + if isinstance(inner, RetryVLM): + raise ValueError("VLM is already wrapped with RetryVLM") + super().__init__(inner.config) + self.inner = inner + + def __getattr__(self, name: str) -> Any: + return getattr(self.inner, name) + + def __setattr__(self, name: str, value: Any) -> None: + inner = self.__dict__.get("inner") + if ( + inner is not None + and name not in self.__dict__ + and not hasattr(type(self), name) + and hasattr(inner, name) + ): + setattr(inner, name, value) + return + object.__setattr__(self, name, value) + + @property + def token_tracker(self): + return self.inner.token_tracker + + def get_token_usage(self) -> Dict[str, Any]: + return self.inner.get_token_usage() + + def reset_token_usage(self) -> None: + self.inner.reset_token_usage() + + def update_token_usage(self, *args, **kwargs) -> None: + self.inner.update_token_usage(*args, **kwargs) + + def record_failed_call(self, *, duration_seconds: float, error: Exception) -> None: + self.inner.record_failed_call(duration_seconds=duration_seconds, error=error) + + def get_completion( + self, + prompt: str = "", + thinking: Optional[bool] = None, + tools: Optional[List[Dict[str, Any]]] = None, + tool_choice: Optional[str] = None, + messages: Optional[List[Dict[str, Any]]] = None, + **kwargs: Any, + ) -> Union[str, VLMResponse]: + return retry_sync( + lambda: self.inner.get_completion( + prompt=prompt, + thinking=thinking, + tools=tools, + tool_choice=tool_choice, + messages=messages, + **kwargs, + ), + max_retries=self.max_retries, + logger=logger, + operation_name=f"{self.provider} VLM completion", + ) + + async def get_completion_async( + self, + prompt: str = "", + thinking: Optional[bool] = None, + tools: Optional[List[Dict[str, Any]]] = None, + tool_choice: Optional[str] = None, + messages: Optional[List[Dict[str, Any]]] = None, + max_tokens: Optional[int] = None, + **kwargs: Any, + ) -> Union[str, VLMResponse]: + return await retry_async( + lambda: self.inner.get_completion_async( + prompt=prompt, + thinking=thinking, + tools=tools, + tool_choice=tool_choice, + messages=messages, + max_tokens=max_tokens, + **kwargs, + ), + max_retries=self.max_retries, + logger=logger, + operation_name=f"{self.provider} VLM async completion", + ) + + def get_vision_completion( + self, + prompt: str = "", + images: Optional[List[Union[str, Path, bytes]]] = None, + thinking: Optional[bool] = None, + tools: Optional[List[Dict[str, Any]]] = None, + tool_choice: Optional[str] = None, + messages: Optional[List[Dict[str, Any]]] = None, + **kwargs: Any, + ) -> Union[str, VLMResponse]: + return retry_sync( + lambda: self.inner.get_vision_completion( + prompt=prompt, + images=images, + thinking=thinking, + tools=tools, + tool_choice=tool_choice, + messages=messages, + **kwargs, + ), + max_retries=self.max_retries, + logger=logger, + operation_name=f"{self.provider} VLM vision completion", + ) + + async def get_vision_completion_async( + self, + prompt: str = "", + images: Optional[List[Union[str, Path, bytes]]] = None, + thinking: Optional[bool] = None, + tools: Optional[List[Dict[str, Any]]] = None, + tool_choice: Optional[str] = None, + messages: Optional[List[Dict[str, Any]]] = None, + **kwargs: Any, + ) -> Union[str, VLMResponse]: + return await retry_async( + lambda: self.inner.get_vision_completion_async( + prompt=prompt, + images=images, + thinking=thinking, + tools=tools, + tool_choice=tool_choice, + messages=messages, + **kwargs, + ), + max_retries=self.max_retries, + logger=logger, + operation_name=f"{self.provider} VLM async vision completion", + ) + + def supports_media( + self, + *, + media_type: str, + filename: str, + size_bytes: int, + ) -> bool: + return self.inner.supports_media( + media_type=media_type, + filename=filename, + size_bytes=size_bytes, + ) + + async def get_media_completion_async( + self, + *, + prompt: str, + media_path: Path, + filename: str, + media_type: str, + ) -> str: + return await self.inner.get_media_completion_async( + prompt=prompt, + media_path=media_path, + filename=filename, + media_type=media_type, + ) + + def is_available(self) -> bool: + return self.inner.is_available() + + def close(self) -> None: + self.inner.close() diff --git a/openviking/utils/model_retry.py b/openviking/utils/model_retry.py index 2caf88c82c..0c55ef2116 100644 --- a/openviking/utils/model_retry.py +++ b/openviking/utils/model_retry.py @@ -204,6 +204,53 @@ def _pattern_matches(text_lower: str, text_compact: str, pattern: str) -> bool: return pattern in text_lower or pattern in text_compact +def _extract_structured_error_details(error: BaseException) -> tuple[list[int], list[str]]: + """Extract trusted status and provider code fields from an exception chain.""" + statuses: list[int] = [] + codes: list[str] = [] + seen: set[int] = set() + current: BaseException | None = error + + while current is not None and id(current) not in seen: + seen.add(id(current)) + candidates = [current, getattr(current, "response", None)] + for candidate in candidates: + if candidate is None: + continue + status = getattr(candidate, "status_code", None) + try: + status_int = int(status) + except (TypeError, ValueError): + status_int = None + if status_int is not None and 100 <= status_int <= 599: + statuses.append(status_int) + + for attr in ("error_code", "code", "type"): + value = getattr(current, attr, None) + if isinstance(value, str) and value: + codes.append(value) + elif attr == "code" and isinstance(value, int) and 100 <= value <= 599: + statuses.append(value) + + body = getattr(current, "body", None) + if isinstance(body, dict): + body_candidates = [body] + nested = body.get("error") + if isinstance(nested, dict): + body_candidates.append(nested) + for item in body_candidates: + for key in ("error_code", "code", "type"): + value = item.get(key) + if isinstance(value, str) and value: + codes.append(value) + elif key == "code" and isinstance(value, int) and 100 <= value <= 599: + statuses.append(value) + + current = current.__cause__ or current.__context__ + + return statuses, codes + + def classify_api_error(error: Exception) -> str: """Classify an API error into one of the ERROR_CLASS_* categories. @@ -234,7 +281,8 @@ def classify_api_error(error: Exception) -> str: if exc is not None and isinstance(exc, _PERMANENT_IO_ERRORS): return ERROR_CLASS_PERMANENT - texts = [str(error)] + structured_statuses, structured_codes = _extract_structured_error_details(error) + texts = [str(error), *structured_codes] if error.__cause__ is not None: texts.append(str(error.__cause__)) @@ -254,6 +302,15 @@ def classify_api_error(error: Exception) -> str: if _pattern_matches(text_lower, text_compact, pattern): return ERROR_CLASS_CONTENT_SAFETY + if 413 in structured_statuses: + return ERROR_CLASS_INPUT_TOO_LARGE + + if 400 in structured_statuses: + return ERROR_CLASS_PERMANENT + + if any(status in (401, 403) for status in structured_statuses): + return ERROR_CLASS_AUTH + for text in texts: text_lower = text.lower() text_compact = text_lower.replace(" ", "") @@ -276,6 +333,11 @@ def classify_api_error(error: Exception) -> str: if pattern in text_lower: return ERROR_CLASS_QUOTA_EXCEEDED + if 429 in structured_statuses or any( + status in (500, 502, 503, 504) for status in structured_statuses + ): + return ERROR_CLASS_TRANSIENT + for text in texts: text_lower = text.lower() text_compact = text_lower.replace(" ", "") diff --git a/openviking_cli/utils/config/vlm_config.py b/openviking_cli/utils/config/vlm_config.py index 832d9a53e0..7bb88532b4 100644 --- a/openviking_cli/utils/config/vlm_config.py +++ b/openviking_cli/utils/config/vlm_config.py @@ -45,6 +45,10 @@ def _bind_token_usage_tracker(instance: Any, tracker: Any) -> None: return if hasattr(instance, "_token_tracker"): instance._token_tracker = tracker + inner = getattr(instance, "inner", None) + if inner is not None: + _bind_token_usage_tracker(inner, tracker) + return if hasattr(instance, "_vlm_instances"): for child in instance._vlm_instances: _bind_token_usage_tracker(child, tracker) diff --git a/tests/unit/test_extra_headers_vlm.py b/tests/unit/test_extra_headers_vlm.py index 1eaebcb53f..3d313f7e83 100644 --- a/tests/unit/test_extra_headers_vlm.py +++ b/tests/unit/test_extra_headers_vlm.py @@ -354,7 +354,7 @@ def test_volcengine_sync_client_applies_timeout_and_disables_sdk_retries( (400, "Total tokens of multi-modal content and text exceed max message tokens.", 1), (400, "The parameter model is invalid.", 1), (401, "Unauthorized", 1), - (503, "Service Unavailable", 6), + (503, "Service Unavailable", 1), ], ) @patch("volcenginesdkarkruntime.AsyncArk") diff --git a/tests/unit/test_kimi_glm_vlm.py b/tests/unit/test_kimi_glm_vlm.py index a33d11f192..a4e63e4217 100644 --- a/tests/unit/test_kimi_glm_vlm.py +++ b/tests/unit/test_kimi_glm_vlm.py @@ -71,8 +71,8 @@ def test_vlm_factory_routes_first_class_kimi_and_glm_providers(): kimi_vlm = VLMFactory.create({"provider": "kimi", "api_key": "kimi-key", "model": "kimi-code"}) glm_vlm = VLMFactory.create({"provider": "glm", "api_key": "glm-key", "model": "glm-4.6v"}) - assert kimi_vlm.__class__.__name__ == "KimiVLM" - assert glm_vlm.__class__.__name__ == "GLMVLM" + assert kimi_vlm.inner.__class__.__name__ == "KimiVLM" + assert glm_vlm.inner.__class__.__name__ == "GLMVLM" def test_vlm_config_uses_canonical_provider_names(): diff --git a/tests/unit/test_model_retry.py b/tests/unit/test_model_retry.py index cedf5a2d87..e35725f650 100644 --- a/tests/unit/test_model_retry.py +++ b/tests/unit/test_model_retry.py @@ -77,6 +77,11 @@ def __init__(self, *, status_code=None, error_code=None, code=None, body=None): self.body = body +class _Response: + def __init__(self, status_code): + self.status_code = status_code + + def test_extract_metric_error_code_prefers_structured_provider_code(): assert ( extract_metric_error_code(_ProviderError(status_code=429, code="RateLimitExceeded")) @@ -96,6 +101,34 @@ def test_extract_metric_error_code_uses_safe_fallbacks_only(): assert extract_metric_error_code(RuntimeError("request_id=not-a-metric-label")) == "unknown" +@pytest.mark.parametrize( + ("status_code", "expected"), + [ + (400, ERROR_CLASS_PERMANENT), + (401, ERROR_CLASS_AUTH), + (403, ERROR_CLASS_AUTH), + (413, ERROR_CLASS_INPUT_TOO_LARGE), + (429, ERROR_CLASS_TRANSIENT), + (503, ERROR_CLASS_TRANSIENT), + ], +) +def test_classify_api_error_prefers_structured_response_status(status_code, expected): + error = RuntimeError("provider request failed without status text") + error.response = _Response(status_code) + + assert classify_api_error(error) == expected + + +def test_structured_quota_code_takes_precedence_over_429_status(): + error = _ProviderError(status_code=429, code="AccountQuotaExceeded") + + assert classify_api_error(error) == ERROR_CLASS_QUOTA_EXCEEDED + + +def test_unstructured_number_is_not_treated_as_http_status(): + assert classify_api_error(RuntimeError("request id abc503xyz")) == ERROR_CLASS_UNKNOWN + + def test_classify_all_credentials_failed_prefers_transient_over_auth(): err = AllCredentialsFailedError( [ diff --git a/tests/unit/test_retry_vlm.py b/tests/unit/test_retry_vlm.py new file mode 100644 index 0000000000..734752b5e0 --- /dev/null +++ b/tests/unit/test_retry_vlm.py @@ -0,0 +1,167 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +import pytest + +from openviking.models.vlm.base import VLMBase, VLMFactory +from openviking.models.vlm.retry_vlm import RetryVLM + + +class ProviderError(RuntimeError): + def __init__(self, status_code: int): + super().__init__("provider request failed") + self.status_code = status_code + + +class FakeVLM(VLMBase): + def __init__(self, outcomes, *, media_outcomes=None): + super().__init__({"provider": "fake", "model": "fake-model", "max_retries": 2}) + self.outcomes = list(outcomes) + self.media_outcomes = list(media_outcomes or ["media"]) + self.calls = [] + self.closed = False + + def _next(self, method, kwargs): + self.calls.append((method, kwargs)) + outcome = self.outcomes.pop(0) + if isinstance(outcome, Exception): + raise outcome + return outcome + + def get_completion(self, **kwargs): + return self._next("text", kwargs) + + async def get_completion_async(self, **kwargs): + return self._next("async_text", kwargs) + + def get_vision_completion(self, **kwargs): + return self._next("vision", kwargs) + + async def get_vision_completion_async(self, **kwargs): + return self._next("async_vision", kwargs) + + def supports_media(self, **kwargs): + return True + + async def get_media_completion_async(self, **kwargs): + self.calls.append(("media", kwargs)) + outcome = self.media_outcomes.pop(0) + if isinstance(outcome, Exception): + raise outcome + return outcome + + def close(self): + self.closed = True + + +@pytest.fixture(autouse=True) +def no_retry_sleep(monkeypatch): + async def _no_sleep(_delay): + return None + + monkeypatch.setattr("openviking.utils.model_retry.time.sleep", lambda _delay: None) + monkeypatch.setattr("openviking.utils.model_retry.asyncio.sleep", _no_sleep) + + +@pytest.mark.parametrize( + ("method", "expected_method"), + [ + ("get_completion", "text"), + ("get_vision_completion", "vision"), + ], +) +def test_retry_vlm_retries_sync_transient_failures(method, expected_method): + inner = FakeVLM([ProviderError(503), "ok"]) + vlm = RetryVLM(inner) + + result = getattr(vlm, method)(prompt="hello") + + assert result == "ok" + assert [name for name, _kwargs in inner.calls] == [expected_method, expected_method] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("method", "expected_method"), + [ + ("get_completion_async", "async_text"), + ("get_vision_completion_async", "async_vision"), + ], +) +async def test_retry_vlm_retries_async_transient_failures(method, expected_method): + inner = FakeVLM([ProviderError(503), "ok"]) + vlm = RetryVLM(inner) + + result = await getattr(vlm, method)(prompt="hello") + + assert result == "ok" + assert [name for name, _kwargs in inner.calls] == [expected_method, expected_method] + + +@pytest.mark.parametrize("status_code", [400, 401]) +def test_retry_vlm_does_not_retry_permanent_or_auth_errors(status_code): + error = ProviderError(status_code) + inner = FakeVLM([error, "unexpected"]) + + with pytest.raises(ProviderError) as raised: + RetryVLM(inner).get_completion(prompt="hello") + + assert raised.value is error + assert len(inner.calls) == 1 + + +def test_retry_vlm_honors_zero_retries(): + inner = FakeVLM([ProviderError(503), "unexpected"]) + inner.max_retries = 0 + vlm = RetryVLM(inner) + vlm.max_retries = 0 + + with pytest.raises(ProviderError): + vlm.get_completion(prompt="hello") + + assert len(inner.calls) == 1 + + +@pytest.mark.asyncio +async def test_retry_vlm_forwards_extension_kwargs_and_does_not_retry_media(tmp_path): + media_error = ProviderError(503) + inner = FakeVLM(["ok"], media_outcomes=[media_error, "unexpected"]) + vlm = RetryVLM(inner) + + assert await vlm.get_completion_async(prompt="hello", reasoning_effort="high") == "ok" + assert inner.calls[0][1]["reasoning_effort"] == "high" + + media_path = tmp_path / "clip.mp4" + media_path.write_bytes(b"media") + with pytest.raises(ProviderError) as raised: + await vlm.get_media_completion_async( + prompt="describe", + media_path=Path(media_path), + filename="clip.mp4", + media_type="video/mp4", + ) + + assert raised.value is media_error + assert [name for name, _kwargs in inner.calls].count("media") == 1 + + +def test_retry_vlm_delegates_token_tracker_and_close(): + inner = FakeVLM(["ok"]) + vlm = RetryVLM(inner) + + assert vlm.token_tracker is inner.token_tracker + vlm.close() + assert inner.closed is True + + +def test_vlm_factory_exposes_raw_and_standard_creation(): + config = {"provider": "kimi", "api_key": "key", "model": "kimi-code"} + + raw = VLMFactory.create_provider(config) + standard = VLMFactory.create(config) + + assert raw.__class__.__name__ == "KimiVLM" + assert isinstance(standard, RetryVLM) + assert standard.inner.__class__.__name__ == "KimiVLM" From 390f1821b28f0dbe2d4d4c78eb5e30a4c018c051 Mon Sep 17 00:00:00 2001 From: "zhanghe.1216" <@bytedance.com> Date: Mon, 12 Oct 2026 20:33:28 +0800 Subject: [PATCH 02/16] =?UTF-8?q?test(vlm):=20=E8=A1=A5=E5=85=85=E7=BB=9F?= =?UTF-8?q?=E4=B8=80=E9=87=8D=E8=AF=95=E8=BE=B9=E7=95=8C=E5=9B=9E=E5=BD=92?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/unit/test_codex_vlm.py | 6 +++- tests/unit/test_retry_vlm.py | 64 +++++++++++++++++++++++++++++++++++- 2 files changed, 68 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_codex_vlm.py b/tests/unit/test_codex_vlm.py index 0808af5657..2b1c108a24 100644 --- a/tests/unit/test_codex_vlm.py +++ b/tests/unit/test_codex_vlm.py @@ -219,10 +219,14 @@ async def test_codex_async_stream_uses_text_deltas_when_completed_output_is_none @patch("openviking.models.vlm.backends.codex_auth.has_codex_auth_available", return_value=True) def test_vlm_config_accepts_codex_without_api_key(_mock_auth_available): + from openviking.models.vlm.retry_vlm import RetryVLM + config = VLMConfig(provider="openai-codex", model="gpt-5.3-codex") assert config.is_available() is True - assert config.get_vlm_instance().__class__.__name__ == "CodexVLM" + instance = config.get_vlm_instance() + assert isinstance(instance, RetryVLM) + assert instance.inner.__class__.__name__ == "CodexVLM" @patch("openviking.models.vlm.backends.codex_auth.has_codex_auth_available", return_value=True) diff --git a/tests/unit/test_retry_vlm.py b/tests/unit/test_retry_vlm.py index 734752b5e0..c3312d884c 100644 --- a/tests/unit/test_retry_vlm.py +++ b/tests/unit/test_retry_vlm.py @@ -1,12 +1,14 @@ # Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. # SPDX-License-Identifier: Apache-2.0 +import asyncio from pathlib import Path import pytest -from openviking.models.vlm.base import VLMBase, VLMFactory +from openviking.models.vlm.base import MultiCredentialVLM, VLMBase, VLMFactory from openviking.models.vlm.retry_vlm import RetryVLM +from openviking_cli.utils.config.vlm_config import VLMConfig class ProviderError(RuntimeError): @@ -124,6 +126,38 @@ def test_retry_vlm_honors_zero_retries(): assert len(inner.calls) == 1 +def test_retry_vlm_uses_max_retries_as_additional_attempts(): + errors = [ProviderError(503) for _ in range(3)] + inner = FakeVLM(errors) + + with pytest.raises(ProviderError) as raised: + RetryVLM(inner).get_completion(prompt="hello") + + assert raised.value is errors[-1] + assert len(inner.calls) == 3 + + +@pytest.mark.asyncio +async def test_retry_vlm_does_not_retry_cancelled_request(): + started = asyncio.Event() + + class BlockingVLM(FakeVLM): + async def get_completion_async(self, **kwargs): + self.calls.append(("async_text", kwargs)) + started.set() + await asyncio.Event().wait() + + inner = BlockingVLM([]) + task = asyncio.create_task(RetryVLM(inner).get_completion_async(prompt="hello")) + await started.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert len(inner.calls) == 1 + + @pytest.mark.asyncio async def test_retry_vlm_forwards_extension_kwargs_and_does_not_retry_media(tmp_path): media_error = ProviderError(503) @@ -165,3 +199,31 @@ def test_vlm_factory_exposes_raw_and_standard_creation(): assert raw.__class__.__name__ == "KimiVLM" assert isinstance(standard, RetryVLM) assert standard.inner.__class__.__name__ == "KimiVLM" + + +def test_vlm_config_wraps_each_credential_inside_failover(): + config = VLMConfig( + model="kimi-code", + credentials=[ + { + "id": "primary", + "provider": "kimi", + "api_key": "primary-key", + }, + { + "id": "backup", + "provider": "glm", + "api_key": "backup-key", + "model": "glm-4.6v", + }, + ], + ) + + instance = config.get_vlm_instance() + + assert isinstance(instance, MultiCredentialVLM) + assert all(isinstance(child, RetryVLM) for child in instance._vlm_instances) + assert [child.inner.__class__.__name__ for child in instance._vlm_instances] == [ + "KimiVLM", + "GLMVLM", + ] From 929a437d28f04a48ecc4382ef6d3f45add1fa52e Mon Sep 17 00:00:00 2001 From: "zhanghe.1216" <@bytedance.com> Date: Mon, 12 Oct 2026 21:05:03 +0800 Subject: [PATCH 03/16] =?UTF-8?q?refactor(vlm):=20=E6=94=B6=E7=B4=A7?= =?UTF-8?q?=E7=BB=9F=E4=B8=80=E9=87=8D=E8=AF=95=E5=B1=82=E6=94=B9=E5=8A=A8?= =?UTF-8?q?=E8=8C=83=E5=9B=B4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 撤回全局错误分类和无关 API 扩张，使本次改造仅聚焦 VLM 重试职责收拢。 --- openviking/models/vlm/__init__.py | 2 - .../models/vlm/backends/volcengine_vlm.py | 24 +++---- openviking/models/vlm/base.py | 4 +- openviking/models/vlm/retry_vlm.py | 12 ---- openviking/utils/model_retry.py | 64 +------------------ tests/unit/test_extra_headers_vlm.py | 25 ++++---- tests/unit/test_model_retry.py | 33 ---------- tests/unit/test_retry_vlm.py | 11 ++-- 8 files changed, 34 insertions(+), 141 deletions(-) diff --git a/openviking/models/vlm/__init__.py b/openviking/models/vlm/__init__.py index b4f3f265a2..8d761ec5ea 100644 --- a/openviking/models/vlm/__init__.py +++ b/openviking/models/vlm/__init__.py @@ -16,12 +16,10 @@ VLMFactory, ) from .registry import get_all_provider_names, is_valid_provider -from .retry_vlm import RetryVLM __all__ = [ "VLMBase", "VLMFactory", - "RetryVLM", "FailoverVLM", "MultiCredentialVLM", "AllCredentialsFailedError", diff --git a/openviking/models/vlm/backends/volcengine_vlm.py b/openviking/models/vlm/backends/volcengine_vlm.py index 977b0b225a..e51bbc054f 100644 --- a/openviking/models/vlm/backends/volcengine_vlm.py +++ b/openviking/models/vlm/backends/volcengine_vlm.py @@ -119,12 +119,12 @@ def get_client(self): raise ImportError( "Please install volcenginesdkarkruntime: pip install volcenginesdkarkruntime" ) - kwargs = { - "api_key": self.api_key, - "base_url": self.api_base, - "timeout": self.timeout, - "max_retries": 0, - } + kwargs = dict( + api_key=self.api_key, + base_url=self.api_base, + timeout=self.timeout, + max_retries=0, + ) http_client = create_optional_sync_httpx_client( self.api_base, timeout=self.timeout, @@ -142,12 +142,12 @@ def _build_async_client(self): raise ImportError( "Please install volcenginesdkarkruntime: pip install volcenginesdkarkruntime" ) - kwargs = { - "api_key": self.api_key, - "base_url": self.api_base, - "timeout": self.timeout, - "max_retries": 0, - } + kwargs = dict( + api_key=self.api_key, + base_url=self.api_base, + timeout=self.timeout, + max_retries=0, + ) http_client = create_optional_async_httpx_client( self.api_base, timeout=self.timeout, diff --git a/openviking/models/vlm/base.py b/openviking/models/vlm/base.py index 70743d11cc..930b41ceea 100644 --- a/openviking/models/vlm/base.py +++ b/openviking/models/vlm/base.py @@ -346,7 +346,7 @@ class VLMFactory: """VLM factory class, creates corresponding VLM instance based on config""" @staticmethod - def create_provider(config: Dict[str, Any]) -> VLMBase: + def _create_provider(config: Dict[str, Any]) -> VLMBase: """Create one concrete VLM provider without an outer retry decorator. Args: @@ -396,7 +396,7 @@ def create(config: Dict[str, Any]) -> VLMBase: """Create one VLM target with the standard retry owner.""" from .retry_vlm import RetryVLM - return RetryVLM(VLMFactory.create_provider(config)) + return RetryVLM(VLMFactory._create_provider(config)) def _annotate_vlm_error(exc: Exception, vlm_instance: "VLMBase") -> None: diff --git a/openviking/models/vlm/retry_vlm.py b/openviking/models/vlm/retry_vlm.py index 7828ad428f..d8b01d7086 100644 --- a/openviking/models/vlm/retry_vlm.py +++ b/openviking/models/vlm/retry_vlm.py @@ -25,18 +25,6 @@ def __init__(self, inner: VLMBase): def __getattr__(self, name: str) -> Any: return getattr(self.inner, name) - def __setattr__(self, name: str, value: Any) -> None: - inner = self.__dict__.get("inner") - if ( - inner is not None - and name not in self.__dict__ - and not hasattr(type(self), name) - and hasattr(inner, name) - ): - setattr(inner, name, value) - return - object.__setattr__(self, name, value) - @property def token_tracker(self): return self.inner.token_tracker diff --git a/openviking/utils/model_retry.py b/openviking/utils/model_retry.py index 0c55ef2116..2caf88c82c 100644 --- a/openviking/utils/model_retry.py +++ b/openviking/utils/model_retry.py @@ -204,53 +204,6 @@ def _pattern_matches(text_lower: str, text_compact: str, pattern: str) -> bool: return pattern in text_lower or pattern in text_compact -def _extract_structured_error_details(error: BaseException) -> tuple[list[int], list[str]]: - """Extract trusted status and provider code fields from an exception chain.""" - statuses: list[int] = [] - codes: list[str] = [] - seen: set[int] = set() - current: BaseException | None = error - - while current is not None and id(current) not in seen: - seen.add(id(current)) - candidates = [current, getattr(current, "response", None)] - for candidate in candidates: - if candidate is None: - continue - status = getattr(candidate, "status_code", None) - try: - status_int = int(status) - except (TypeError, ValueError): - status_int = None - if status_int is not None and 100 <= status_int <= 599: - statuses.append(status_int) - - for attr in ("error_code", "code", "type"): - value = getattr(current, attr, None) - if isinstance(value, str) and value: - codes.append(value) - elif attr == "code" and isinstance(value, int) and 100 <= value <= 599: - statuses.append(value) - - body = getattr(current, "body", None) - if isinstance(body, dict): - body_candidates = [body] - nested = body.get("error") - if isinstance(nested, dict): - body_candidates.append(nested) - for item in body_candidates: - for key in ("error_code", "code", "type"): - value = item.get(key) - if isinstance(value, str) and value: - codes.append(value) - elif key == "code" and isinstance(value, int) and 100 <= value <= 599: - statuses.append(value) - - current = current.__cause__ or current.__context__ - - return statuses, codes - - def classify_api_error(error: Exception) -> str: """Classify an API error into one of the ERROR_CLASS_* categories. @@ -281,8 +234,7 @@ def classify_api_error(error: Exception) -> str: if exc is not None and isinstance(exc, _PERMANENT_IO_ERRORS): return ERROR_CLASS_PERMANENT - structured_statuses, structured_codes = _extract_structured_error_details(error) - texts = [str(error), *structured_codes] + texts = [str(error)] if error.__cause__ is not None: texts.append(str(error.__cause__)) @@ -302,15 +254,6 @@ def classify_api_error(error: Exception) -> str: if _pattern_matches(text_lower, text_compact, pattern): return ERROR_CLASS_CONTENT_SAFETY - if 413 in structured_statuses: - return ERROR_CLASS_INPUT_TOO_LARGE - - if 400 in structured_statuses: - return ERROR_CLASS_PERMANENT - - if any(status in (401, 403) for status in structured_statuses): - return ERROR_CLASS_AUTH - for text in texts: text_lower = text.lower() text_compact = text_lower.replace(" ", "") @@ -333,11 +276,6 @@ def classify_api_error(error: Exception) -> str: if pattern in text_lower: return ERROR_CLASS_QUOTA_EXCEEDED - if 429 in structured_statuses or any( - status in (500, 502, 503, 504) for status in structured_statuses - ): - return ERROR_CLASS_TRANSIENT - for text in texts: text_lower = text.lower() text_compact = text_lower.replace(" ", "") diff --git a/tests/unit/test_extra_headers_vlm.py b/tests/unit/test_extra_headers_vlm.py index 3d313f7e83..d8c86108e6 100644 --- a/tests/unit/test_extra_headers_vlm.py +++ b/tests/unit/test_extra_headers_vlm.py @@ -349,17 +349,17 @@ def test_volcengine_sync_client_applies_timeout_and_disables_sdk_retries( assert call_kwargs["max_retries"] == 0 @pytest.mark.parametrize( - ("status_code", "message", "expected_attempts"), + ("status_code", "message"), [ - (400, "Total tokens of multi-modal content and text exceed max message tokens.", 1), - (400, "The parameter model is invalid.", 1), - (401, "Unauthorized", 1), - (503, "Service Unavailable", 1), + (400, "Total tokens of multi-modal content and text exceed max message tokens."), + (400, "The parameter model is invalid."), + (401, "Unauthorized"), + (503, "Service Unavailable"), ], ) @patch("volcenginesdkarkruntime.AsyncArk") - def test_volcengine_async_client_uses_shared_retry_policy( - self, mock_async_ark_class, status_code, message, expected_attempts + def test_volcengine_async_client_performs_single_request( + self, mock_async_ark_class, status_code, message ): error = ArkAPIStatusError( f"Error code: {status_code} - {message}", @@ -382,12 +382,11 @@ def test_volcengine_async_client_uses_shared_retry_policy( } ) - with patch("openviking.utils.model_retry.asyncio.sleep", new=AsyncMock()): - with pytest.raises(ArkAPIStatusError) as raised: - asyncio.run(vlm.get_completion_async("hello")) + with pytest.raises(ArkAPIStatusError) as raised: + asyncio.run(vlm.get_completion_async("hello")) assert raised.value is error - assert mock_client.chat.completions.create.await_count == expected_attempts + assert mock_client.chat.completions.create.await_count == 1 call_kwargs = mock_async_ark_class.call_args.kwargs assert call_kwargs["timeout"] == 12.0 assert call_kwargs["max_retries"] == 0 @@ -531,7 +530,7 @@ async def test_completion_body_overrides_configured_defaults( api_key="sk-test", http_client=httpx.Client(transport=httpx.MockTransport(send)), ) as client: - monkeypatch.setattr(vlm, "get_client", lambda: client) + monkeypatch.setattr(vlm.inner, "get_client", lambda: client) assert vlm.get_completion("hello") == "ok" assert vlm.get_vision_completion("describe", images=[b"\x89PNG\r\n\x1a\n"]) == "ok" @@ -539,7 +538,7 @@ async def test_completion_body_overrides_configured_defaults( api_key="sk-test", http_client=httpx.AsyncClient(transport=httpx.MockTransport(send)), ) as async_client: - monkeypatch.setattr(vlm, "get_async_client", lambda: async_client) + monkeypatch.setattr(vlm.inner, "get_async_client", lambda: async_client) assert await vlm.get_completion_async("hello") == "ok" assert ( await vlm.get_vision_completion_async("describe", images=[b"\x89PNG\r\n\x1a\n"]) diff --git a/tests/unit/test_model_retry.py b/tests/unit/test_model_retry.py index e35725f650..cedf5a2d87 100644 --- a/tests/unit/test_model_retry.py +++ b/tests/unit/test_model_retry.py @@ -77,11 +77,6 @@ def __init__(self, *, status_code=None, error_code=None, code=None, body=None): self.body = body -class _Response: - def __init__(self, status_code): - self.status_code = status_code - - def test_extract_metric_error_code_prefers_structured_provider_code(): assert ( extract_metric_error_code(_ProviderError(status_code=429, code="RateLimitExceeded")) @@ -101,34 +96,6 @@ def test_extract_metric_error_code_uses_safe_fallbacks_only(): assert extract_metric_error_code(RuntimeError("request_id=not-a-metric-label")) == "unknown" -@pytest.mark.parametrize( - ("status_code", "expected"), - [ - (400, ERROR_CLASS_PERMANENT), - (401, ERROR_CLASS_AUTH), - (403, ERROR_CLASS_AUTH), - (413, ERROR_CLASS_INPUT_TOO_LARGE), - (429, ERROR_CLASS_TRANSIENT), - (503, ERROR_CLASS_TRANSIENT), - ], -) -def test_classify_api_error_prefers_structured_response_status(status_code, expected): - error = RuntimeError("provider request failed without status text") - error.response = _Response(status_code) - - assert classify_api_error(error) == expected - - -def test_structured_quota_code_takes_precedence_over_429_status(): - error = _ProviderError(status_code=429, code="AccountQuotaExceeded") - - assert classify_api_error(error) == ERROR_CLASS_QUOTA_EXCEEDED - - -def test_unstructured_number_is_not_treated_as_http_status(): - assert classify_api_error(RuntimeError("request id abc503xyz")) == ERROR_CLASS_UNKNOWN - - def test_classify_all_credentials_failed_prefers_transient_over_auth(): err = AllCredentialsFailedError( [ diff --git a/tests/unit/test_retry_vlm.py b/tests/unit/test_retry_vlm.py index c3312d884c..c14a21e025 100644 --- a/tests/unit/test_retry_vlm.py +++ b/tests/unit/test_retry_vlm.py @@ -13,7 +13,12 @@ class ProviderError(RuntimeError): def __init__(self, status_code: int): - super().__init__("provider request failed") + messages = { + 400: "400 invalid parameter", + 401: "401 unauthorized", + 503: "503 service unavailable", + } + super().__init__(messages[status_code]) self.status_code = status_code @@ -190,13 +195,11 @@ def test_retry_vlm_delegates_token_tracker_and_close(): assert inner.closed is True -def test_vlm_factory_exposes_raw_and_standard_creation(): +def test_vlm_factory_wraps_provider_with_standard_retry(): config = {"provider": "kimi", "api_key": "key", "model": "kimi-code"} - raw = VLMFactory.create_provider(config) standard = VLMFactory.create(config) - assert raw.__class__.__name__ == "KimiVLM" assert isinstance(standard, RetryVLM) assert standard.inner.__class__.__name__ == "KimiVLM" From 8c9c8ef32f6053a4e2d0316c1b236b1de5671cf5 Mon Sep 17 00:00:00 2001 From: "zhanghe.1216" <@bytedance.com> Date: Tue, 13 Oct 2026 15:22:10 +0800 Subject: [PATCH 04/16] =?UTF-8?q?refactor(embedding):=20=E7=BB=9F=E4=B8=80?= =?UTF-8?q?=E6=A8=A1=E5=9E=8B=E9=87=8D=E8=AF=95=E8=81=8C=E8=B4=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 将重试决策集中到 RetryEmbedder，避免 Provider 与 SDK 重试叠乘，同时保持组合模型和账号资源语义。 --- openviking/config/embedding.py | 3 + openviking/models/embedder/__init__.py | 2 + openviking/models/embedder/base.py | 100 +++--- .../models/embedder/cohere_embedders.py | 8 +- .../models/embedder/dashscope_embedders.py | 10 +- .../models/embedder/gemini_embedders.py | 21 +- openviking/models/embedder/jina_embedders.py | 6 +- .../models/embedder/litellm_embedders.py | 6 +- openviking/models/embedder/local_embedders.py | 2 +- .../models/embedder/minimax_embedders.py | 27 +- .../models/embedder/openai_embedders.py | 4 +- openviking/models/embedder/retry_embedder.py | 93 ++++++ .../models/embedder/vikingdb_embedders.py | 12 +- .../models/embedder/volcengine_embedders.py | 14 +- .../models/embedder/voyage_embedders.py | 6 +- .../utils/config/embedding_config.py | 8 +- tests/unit/test_embedding_config_voyage.py | 5 +- tests/unit/test_encoding_format_embedding.py | 17 +- tests/unit/test_litellm_embedder.py | 8 +- tests/unit/test_retry_embedder.py | 297 ++++++++++++++++++ 20 files changed, 518 insertions(+), 131 deletions(-) create mode 100644 openviking/models/embedder/retry_embedder.py create mode 100644 tests/unit/test_retry_embedder.py diff --git a/openviking/config/embedding.py b/openviking/config/embedding.py index d683970976..656bfaa735 100644 --- a/openviking/config/embedding.py +++ b/openviking/config/embedding.py @@ -13,6 +13,7 @@ from openviking.config.scope import ScopeKind from openviking.config.vector import AccountVectorConfigResolver, VectorRuntimeSettings from openviking.models.embedder.base import EmbedderBase +from openviking.models.embedder.retry_embedder import RetryEmbedder from openviking.models.vlm.token_usage import TokenUsageTracker from openviking.service.task_tracker_concurrency import KeyedAsyncLockPool from openviking.utils.circuit_breaker import CircuitBreaker, CircuitBreakerOpen @@ -114,6 +115,8 @@ def _bind_account_state( child = getattr(embedder, name, None) if child is not None: children.append(child) + if isinstance(embedder, RetryEmbedder): + children.append(embedder.inner) for child in children: _bind_account_state(child, semaphore, tracker, account_id) diff --git a/openviking/models/embedder/__init__.py b/openviking/models/embedder/__init__.py index 229c35b945..8b8670ee0f 100644 --- a/openviking/models/embedder/__init__.py +++ b/openviking/models/embedder/__init__.py @@ -44,6 +44,7 @@ LiteLLMDenseEmbedder = None # litellm not installed from openviking.models.embedder.minimax_embedders import MinimaxDenseEmbedder from openviking.models.embedder.openai_embedders import OpenAIDenseEmbedder +from openviking.models.embedder.retry_embedder import RetryEmbedder from openviking.models.embedder.vikingdb_embedders import ( VikingDBDenseEmbedder, VikingDBHybridEmbedder, @@ -69,6 +70,7 @@ "HybridEmbedderBase", "CompositeHybridEmbedder", "FailoverEmbedder", + "RetryEmbedder", "AllCredentialsFailedError", # Google Gemini implementations "GeminiDenseEmbedder", diff --git a/openviking/models/embedder/base.py b/openviking/models/embedder/base.py index efb9ec1369..329ede0b84 100644 --- a/openviking/models/embedder/base.py +++ b/openviking/models/embedder/base.py @@ -21,8 +21,6 @@ from openviking.utils.model_retry import ( OrderedCredentialSwitcher, classify_api_error, - retry_async, - retry_sync, ) from openviking_cli.utils import get_logger @@ -377,66 +375,58 @@ def close(self): """Release resources, subclasses can override as needed""" pass - def _run_with_retry(self, func: Callable[[], T], *, logger=None, operation_name: str) -> T: - def _wrapped() -> T: - previous_started_at = self._active_call_started_at - self._active_call_started_at = time.monotonic() - try: - return func() - finally: - self._active_call_started_at = previous_started_at - - return retry_sync( - _wrapped, - max_retries=self.max_retries, - logger=logger, - operation_name=operation_name, - ) + def _run_once( + self, + func: Callable[[], T], + *, + logger=None, + operation_name: str = "embedding", + ) -> T: + """Run one provider request with per-attempt timing state.""" + previous_started_at = self._active_call_started_at + self._active_call_started_at = time.monotonic() + try: + return func() + finally: + self._active_call_started_at = previous_started_at - async def _run_with_async_retry( + async def _run_once_async( self, func: Callable[[], Awaitable[T]], *, logger=None, operation_name: str, ) -> T: - async def _wrapped() -> T: - semaphore = getattr(self, "_account_semaphore", None) - if semaphore is None: - semaphore = _get_async_embed_semaphore(self.max_concurrent) - wait_started = time.monotonic() - await semaphore.acquire() - wait_elapsed = time.monotonic() - wait_started - telemetry = get_current_telemetry() - telemetry.set("embedding.async.max_concurrent", self.max_concurrent) - telemetry.set("embedding.async.wait_ms", round(wait_elapsed * 1000, 3)) - - started = time.monotonic() - previous_started_at = self._active_call_started_at - self._active_call_started_at = started - try: - return await func() - finally: - elapsed = time.monotonic() - started - telemetry.set("embedding.async.duration_ms", round(elapsed * 1000, 3)) - if logger and elapsed >= 3.0: - logger.warning( - "%s slow call provider=%s model=%s wait_ms=%.2f duration_ms=%.2f", - operation_name, - self.provider, - self.model_name, - wait_elapsed * 1000, - elapsed * 1000, - ) - self._active_call_started_at = previous_started_at - semaphore.release() - - return await retry_async( - _wrapped, - max_retries=self.max_retries, - logger=logger, - operation_name=operation_name, - ) + """Run one async provider request with concurrency and telemetry guards.""" + semaphore = getattr(self, "_account_semaphore", None) + if semaphore is None: + semaphore = _get_async_embed_semaphore(self.max_concurrent) + wait_started = time.monotonic() + await semaphore.acquire() + wait_elapsed = time.monotonic() - wait_started + telemetry = get_current_telemetry() + telemetry.set("embedding.async.max_concurrent", self.max_concurrent) + telemetry.set("embedding.async.wait_ms", round(wait_elapsed * 1000, 3)) + + started = time.monotonic() + previous_started_at = self._active_call_started_at + self._active_call_started_at = started + try: + return await func() + finally: + elapsed = time.monotonic() - started + telemetry.set("embedding.async.duration_ms", round(elapsed * 1000, 3)) + if logger and elapsed >= 3.0: + logger.warning( + "%s slow call provider=%s model=%s wait_ms=%.2f duration_ms=%.2f", + operation_name, + self.provider, + self.model_name, + wait_elapsed * 1000, + elapsed * 1000, + ) + self._active_call_started_at = previous_started_at + semaphore.release() @property def is_dense(self) -> bool: diff --git a/openviking/models/embedder/cohere_embedders.py b/openviking/models/embedder/cohere_embedders.py index f31e676c63..43bdabe86c 100644 --- a/openviking/models/embedder/cohere_embedders.py +++ b/openviking/models/embedder/cohere_embedders.py @@ -135,7 +135,11 @@ def _normalize_vector(self, vector: List[float]) -> List[float]: def embed(self, text: str, is_query: bool = False) -> EmbedResult: input_type = "search_query" if is_query else "search_document" try: - vectors = self._call_api([text], input_type) + vectors = self._run_once( + lambda: self._call_api([text], input_type), + logger=logger, + operation_name="Cohere embedding", + ) result = EmbedResult(dense_vector=self._normalize_vector(vectors[0])) # Estimate token usage estimated_tokens = self._estimate_tokens(text) @@ -161,7 +165,7 @@ async def _call() -> EmbedResult: return EmbedResult(dense_vector=self._normalize_vector(vectors[0])) try: - result = await self._run_with_async_retry( + result = await self._run_once_async( _call, logger=logger, operation_name="Cohere async embedding", diff --git a/openviking/models/embedder/dashscope_embedders.py b/openviking/models/embedder/dashscope_embedders.py index 0ad338b191..44e3d3c959 100644 --- a/openviking/models/embedder/dashscope_embedders.py +++ b/openviking/models/embedder/dashscope_embedders.py @@ -89,6 +89,7 @@ def __init__( self._openai_client = openai.OpenAI( api_key=self.api_key, base_url=f"{self.api_base}/compatible-mode/v1", + max_retries=0, ) # Multimodal mode: httpx self._httpx_client = httpx.Client( @@ -208,6 +209,7 @@ def _get_async_openai_client(self) -> openai.AsyncOpenAI: lambda: openai.AsyncOpenAI( api_key=self.api_key, base_url=f"{self.api_base}/compatible-mode/v1", + max_retries=0, ) ) @@ -246,7 +248,7 @@ def _call() -> EmbedResult: return EmbedResult(dense_vector=truncate_and_normalize(vector, self._dimension)) try: - return self._run_with_retry(_call, logger=logger, operation_name="DashScope embedding") + return self._run_once(_call, logger=logger, operation_name="DashScope embedding") except Exception as e: raise RuntimeError(f"DashScope embedding failed: {e}") from e @@ -279,7 +281,7 @@ async def _call() -> EmbedResult: return EmbedResult(dense_vector=truncate_and_normalize(vector, self._dimension)) try: - return await self._run_with_async_retry( + return await self._run_once_async( _call, logger=logger, operation_name="DashScope async embedding" ) except Exception as e: @@ -318,7 +320,7 @@ def _call() -> EmbedResult: return EmbedResult(dense_vector=truncate_and_normalize(vector, self._dimension)) try: - return self._run_with_retry( + return self._run_once( _call, logger=logger, operation_name="DashScope content embedding" ) except Exception as e: @@ -348,7 +350,7 @@ async def _call() -> EmbedResult: return EmbedResult(dense_vector=truncate_and_normalize(vector, self._dimension)) try: - return await self._run_with_async_retry( + return await self._run_once_async( _call, logger=logger, operation_name="DashScope async content embedding" ) except Exception as e: diff --git a/openviking/models/embedder/gemini_embedders.py b/openviking/models/embedder/gemini_embedders.py index 3ab1cd9377..501b74f661 100644 --- a/openviking/models/embedder/gemini_embedders.py +++ b/openviking/models/embedder/gemini_embedders.py @@ -139,12 +139,7 @@ def __init__( self._client_kwargs: Dict[str, Any] = {"api_key": api_key} if _HTTP_RETRY_AVAILABLE: self._client_kwargs["http_options"] = HttpOptions( - retry_options=HttpRetryOptions( - attempts=max(self.max_retries + 1, 1), - initial_delay=0.5, - max_delay=8.0, - exp_base=2.0, - ) + retry_options=HttpRetryOptions(attempts=1) ) self.client = genai.Client(**self._client_kwargs) self.task_type = task_type @@ -213,14 +208,10 @@ def _call() -> EmbedResult: return EmbedResult(dense_vector=vector) try: - result = ( - _call() - if _HTTP_RETRY_AVAILABLE - else self._run_with_retry( - _call, - logger=logger, - operation_name="Gemini embedding", - ) + result = self._run_once( + _call, + logger=logger, + operation_name="Gemini embedding", ) # Estimate token usage estimated_tokens = self._estimate_tokens(text) @@ -262,7 +253,7 @@ async def _call() -> EmbedResult: return EmbedResult(dense_vector=vector) try: - result = await self._run_with_async_retry( + result = await self._run_once_async( _call, logger=logger, operation_name="Gemini async embedding", diff --git a/openviking/models/embedder/jina_embedders.py b/openviking/models/embedder/jina_embedders.py index 91bbf69114..74629dd81c 100644 --- a/openviking/models/embedder/jina_embedders.py +++ b/openviking/models/embedder/jina_embedders.py @@ -126,6 +126,7 @@ def __init__( self.client = openai.OpenAI( api_key=self.api_key, base_url=self.api_base, + max_retries=0, ) self._async_client_cache = LoopScopedAsyncClientCache() @@ -168,6 +169,7 @@ def _get_async_client(self): lambda: openai.AsyncOpenAI( api_key=self.api_key, base_url=self.api_base, + max_retries=0, ) ) @@ -202,7 +204,7 @@ def _call() -> EmbedResult: return EmbedResult(dense_vector=vector) try: - result = self._run_with_retry( + result = self._run_once( _call, logger=logger, operation_name="Jina embedding", @@ -230,7 +232,7 @@ async def _call() -> EmbedResult: return EmbedResult(dense_vector=response.data[0].embedding) try: - result = await self._run_with_async_retry( + result = await self._run_once_async( _call, logger=logger, operation_name="Jina async embedding", diff --git a/openviking/models/embedder/litellm_embedders.py b/openviking/models/embedder/litellm_embedders.py index 2630ee0b89..2c0365ed88 100644 --- a/openviking/models/embedder/litellm_embedders.py +++ b/openviking/models/embedder/litellm_embedders.py @@ -101,7 +101,7 @@ def _truncate_vector(self, vector: List[float]) -> List[float]: def _build_kwargs(self, is_query: bool = False) -> Dict[str, Any]: """Build kwargs dict for litellm.embedding() call.""" - kwargs: Dict[str, Any] = {"model": self.model_name} + kwargs: Dict[str, Any] = {"model": self.model_name, "num_retries": 0} if self.api_key: kwargs["api_key"] = self.api_key @@ -191,7 +191,7 @@ def _call() -> EmbedResult: return EmbedResult(dense_vector=vector) try: - return self._run_with_retry( + return self._run_once( _call, logger=logger, operation_name="LiteLLM embedding", @@ -210,7 +210,7 @@ async def _call() -> EmbedResult: return EmbedResult(dense_vector=vector) try: - return await self._run_with_async_retry( + return await self._run_once_async( _call, logger=logger, operation_name="LiteLLM async embedding", diff --git a/openviking/models/embedder/local_embedders.py b/openviking/models/embedder/local_embedders.py index 350f5b8e0d..d4cb56eab1 100644 --- a/openviking/models/embedder/local_embedders.py +++ b/openviking/models/embedder/local_embedders.py @@ -202,7 +202,7 @@ def embed(self, text: str, is_query: bool = False) -> EmbedResult: formatted = self._format_text(text, is_query=is_query) try: - result = self._run_with_retry( + result = self._run_once( lambda: self._embed_formatted_text(formatted), logger=logger, operation_name="local embedding", diff --git a/openviking/models/embedder/minimax_embedders.py b/openviking/models/embedder/minimax_embedders.py index f9ce0b6245..9c46f38ea2 100644 --- a/openviking/models/embedder/minimax_embedders.py +++ b/openviking/models/embedder/minimax_embedders.py @@ -6,8 +6,6 @@ import httpx import requests -from requests.adapters import HTTPAdapter -from urllib3.util.retry import Retry from openviking.models.embedder.base import DenseEmbedderBase, EmbedResult from openviking.utils.async_client_cache import LoopScopedAsyncClientCache @@ -78,8 +76,7 @@ def __init__( if not self.api_key: raise ValueError("api_key is required for MiniMax embedder") - # Initialize session with retry logic - self.session = self._create_session() + self.session = requests.Session() self._async_client_cache = LoopScopedAsyncClientCache() # Auto-detect dimension if not provided @@ -90,20 +87,6 @@ def __init__( logger.warning(f"Failed to detect MiniMax dimension: {e}. Defaulting to 1536.") self._dimension = 1536 - def _create_session(self) -> requests.Session: - """Create a requests session with retry logic""" - session = requests.Session() - retry_strategy = Retry( - total=self.max_retries, - backoff_factor=0.5, - status_forcelist=[429, 500, 502, 503, 504], - allowed_methods=["POST"], - ) - adapter = HTTPAdapter(max_retries=retry_strategy) - session.mount("https://", adapter) - session.mount("http://", adapter) - return session - def _detect_dimension(self) -> int: """Detect dimension by making an actual API call""" result = self.embed("test") @@ -200,7 +183,11 @@ async def _call_api_async(self, texts: List[str], is_query: bool = False) -> Lis def embed(self, text: str, is_query: bool = False) -> EmbedResult: """Perform dense embedding on text""" - vectors = self._call_api([text], is_query=is_query) + vectors = self._run_once( + lambda: self._call_api([text], is_query=is_query), + logger=logger, + operation_name="MiniMax embedding", + ) result = EmbedResult(dense_vector=vectors[0]) # Estimate token usage estimated_tokens = self._estimate_tokens(text) @@ -217,7 +204,7 @@ async def _call() -> EmbedResult: vectors = await self._call_api_async([text], is_query=is_query) return EmbedResult(dense_vector=vectors[0]) - result = await self._run_with_async_retry( + result = await self._run_once_async( _call, logger=logger, operation_name="MiniMax async embedding", diff --git a/openviking/models/embedder/openai_embedders.py b/openviking/models/embedder/openai_embedders.py index c4c122d201..eec981220d 100644 --- a/openviking/models/embedder/openai_embedders.py +++ b/openviking/models/embedder/openai_embedders.py @@ -390,7 +390,7 @@ def _call() -> EmbedResult: return EmbedResult(dense_vector=vector) try: - return self._run_with_retry( + return self._run_once( _call, logger=logger, operation_name="OpenAI embedding", @@ -412,7 +412,7 @@ async def _call() -> EmbedResult: return EmbedResult(dense_vector=self._truncate_vector(response.data[0].embedding)) try: - return await self._run_with_async_retry( + return await self._run_once_async( _call, logger=logger, operation_name="OpenAI async embedding", diff --git a/openviking/models/embedder/retry_embedder.py b/openviking/models/embedder/retry_embedder.py new file mode 100644 index 0000000000..b7b04aa435 --- /dev/null +++ b/openviking/models/embedder/retry_embedder.py @@ -0,0 +1,93 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Retry decorator for one concrete embedding provider.""" + +from typing import Any, Dict + +from openviking.models.embedder.base import EmbedderBase, EmbeddingInput, EmbedResult +from openviking.utils.model_retry import retry_async, retry_sync +from openviking_cli.utils import get_logger + +logger = get_logger(__name__) + + +class RetryEmbedder(EmbedderBase): + """Own retries for one concrete embedding target.""" + + def __init__(self, inner: EmbedderBase): + if isinstance(inner, RetryEmbedder): + raise ValueError("Embedder is already wrapped with RetryEmbedder") + super().__init__(inner.model_name, inner.config) + self.inner = inner + self.provider = inner.provider + + @property + def query_param(self) -> Any: + return getattr(self.inner, "query_param", None) + + @property + def document_param(self) -> Any: + return getattr(self.inner, "document_param", None) + + def embed( + self, + content: EmbeddingInput, + is_query: bool = False, + **kwargs: Any, + ) -> EmbedResult: + return retry_sync( + lambda: self.inner.embed(content, is_query=is_query, **kwargs), + max_retries=self.max_retries, + logger=logger, + operation_name=f"{self.provider} embedding", + ) + + async def embed_async( + self, + content: EmbeddingInput, + is_query: bool = False, + **kwargs: Any, + ) -> EmbedResult: + return await retry_async( + lambda: self.inner.embed_async(content, is_query=is_query, **kwargs), + max_retries=self.max_retries, + logger=logger, + operation_name=f"{self.provider} async embedding", + ) + + def prepare_embedding_input(self, content: EmbeddingInput) -> EmbeddingInput: + return self.inner.prepare_embedding_input(content) + + @property + def supports_multimodal(self) -> bool: + return self.inner.supports_multimodal + + @property + def is_dense(self) -> bool: + return self.inner.is_dense + + @property + def is_sparse(self) -> bool: + return self.inner.is_sparse + + @property + def is_hybrid(self) -> bool: + return self.inner.is_hybrid + + def get_dimension(self) -> int: + dimension_getter = getattr(self.inner, "get_dimension", None) + if dimension_getter is None: + raise AttributeError( + f"{type(self.inner).__name__} has no get_dimension " + "(sparse embedders have no fixed dimension)" + ) + return dimension_getter() + + def close(self) -> None: + self.inner.close() + + def get_token_usage(self) -> Dict[str, Any]: + return self.inner.get_token_usage() + + def reset_token_usage(self) -> None: + self.inner.reset_token_usage() diff --git a/openviking/models/embedder/vikingdb_embedders.py b/openviking/models/embedder/vikingdb_embedders.py index 5feff642b0..c98b41efa9 100644 --- a/openviking/models/embedder/vikingdb_embedders.py +++ b/openviking/models/embedder/vikingdb_embedders.py @@ -197,7 +197,7 @@ def _call() -> EmbedResult: return EmbedResult(dense_vector=dense_vector) - result = self._run_with_retry( + result = self._run_once( _call, logger=logger, operation_name="VikingDB embedding", @@ -224,7 +224,7 @@ async def _call() -> EmbedResult: dense_vector = self._truncate_and_normalize(item["dense_embedding"], self.dimension) return EmbedResult(dense_vector=dense_vector) - result = await self._run_with_async_retry( + result = await self._run_once_async( _call, logger=logger, operation_name="VikingDB async embedding", @@ -277,7 +277,7 @@ def _call() -> EmbedResult: return EmbedResult(sparse_vector=sparse_vector) - result = self._run_with_retry( + result = self._run_once( _call, logger=logger, operation_name="VikingDB sparse embedding", @@ -306,7 +306,7 @@ async def _call() -> EmbedResult: sparse_vector = self._process_sparse_embedding(item["sparse_embedding"]) return EmbedResult(sparse_vector=sparse_vector) - result = await self._run_with_async_retry( + result = await self._run_once_async( _call, logger=logger, operation_name="VikingDB async sparse embedding", @@ -383,7 +383,7 @@ def _call() -> EmbedResult: return EmbedResult(dense_vector=dense_vector, sparse_vector=sparse_vector) - result = self._run_with_retry( + result = self._run_once( _call, logger=logger, operation_name="VikingDB hybrid embedding", @@ -419,7 +419,7 @@ async def _call() -> EmbedResult: sparse_vector = self._process_sparse_embedding(item["sparse_embedding"]) return EmbedResult(dense_vector=dense_vector, sparse_vector=sparse_vector) - result = await self._run_with_async_retry( + result = await self._run_once_async( _call, logger=logger, operation_name="VikingDB async hybrid embedding", diff --git a/openviking/models/embedder/volcengine_embedders.py b/openviking/models/embedder/volcengine_embedders.py index 6149caa779..86b9c0fc09 100644 --- a/openviking/models/embedder/volcengine_embedders.py +++ b/openviking/models/embedder/volcengine_embedders.py @@ -81,6 +81,7 @@ async def _measure_embedding_call_async( def _create_sync_ark(kwargs: Dict[str, Any]): client_kwargs = dict(kwargs) + client_kwargs["max_retries"] = 0 http_client = create_optional_sync_httpx_client( client_kwargs.get("base_url"), timeout=60.0, @@ -92,6 +93,7 @@ def _create_sync_ark(kwargs: Dict[str, Any]): def _create_async_ark(kwargs: Dict[str, Any]): client_kwargs = dict(kwargs) + client_kwargs["max_retries"] = 0 http_client = create_optional_async_httpx_client( client_kwargs.get("base_url"), timeout=60.0, @@ -302,7 +304,7 @@ def _embed_call(): return EmbedResult(dense_vector=vector) try: - return self._run_with_retry( + return self._run_once( _embed_call, logger=logger, operation_name="Volcengine embedding", @@ -350,7 +352,7 @@ async def _embed_call() -> EmbedResult: return EmbedResult(dense_vector=truncate_and_normalize(vector, self.dimension)) try: - return await self._run_with_async_retry( + return await self._run_once_async( _embed_call, logger=logger, operation_name="Volcengine async embedding", @@ -471,7 +473,7 @@ def _embed_call(): return EmbedResult(sparse_vector=process_sparse_embedding(sparse_vector)) try: - return self._run_with_retry( + return self._run_once( _embed_call, logger=logger, operation_name="Volcengine sparse embedding", @@ -505,7 +507,7 @@ async def _embed_call() -> EmbedResult: return EmbedResult(sparse_vector=process_sparse_embedding(sparse_vector)) try: - return await self._run_with_async_retry( + return await self._run_once_async( _embed_call, logger=logger, operation_name="Volcengine async sparse embedding", @@ -656,7 +658,7 @@ def _embed_call(): ) try: - return self._run_with_retry( + return self._run_once( _embed_call, logger=logger, operation_name="Volcengine hybrid embedding", @@ -694,7 +696,7 @@ async def _embed_call() -> EmbedResult: ) try: - return await self._run_with_async_retry( + return await self._run_once_async( _embed_call, logger=logger, operation_name="Volcengine async hybrid embedding", diff --git a/openviking/models/embedder/voyage_embedders.py b/openviking/models/embedder/voyage_embedders.py index e47310b0ed..67708efd3f 100644 --- a/openviking/models/embedder/voyage_embedders.py +++ b/openviking/models/embedder/voyage_embedders.py @@ -82,6 +82,7 @@ def __init__( self.client = openai.OpenAI( api_key=self.api_key, base_url=self.api_base, + max_retries=0, ) self._async_client_cache = LoopScopedAsyncClientCache() @@ -98,6 +99,7 @@ def _get_async_client(self): lambda: openai.AsyncOpenAI( api_key=self.api_key, base_url=self.api_base, + max_retries=0, ) ) @@ -110,7 +112,7 @@ def _call() -> EmbedResult: return EmbedResult(dense_vector=vector) try: - result = self._run_with_retry( + result = self._run_once( _call, logger=logger, operation_name="Voyage embedding", @@ -137,7 +139,7 @@ async def _call() -> EmbedResult: return EmbedResult(dense_vector=response.data[0].embedding) try: - result = await self._run_with_async_retry( + result = await self._run_once_async( _call, logger=logger, operation_name="Voyage async embedding", diff --git a/openviking_cli/utils/config/embedding_config.py b/openviking_cli/utils/config/embedding_config.py index 723d306398..5f773c86d5 100644 --- a/openviking_cli/utils/config/embedding_config.py +++ b/openviking_cli/utils/config/embedding_config.py @@ -744,6 +744,7 @@ def _create_embedder( LocalDenseEmbedder, MinimaxDenseEmbedder, OpenAIDenseEmbedder, + RetryEmbedder, VikingDBDenseEmbedder, VikingDBHybridEmbedder, VikingDBSparseEmbedder, @@ -1021,7 +1022,12 @@ def _create_embedder( "google-genai is not installed. Install it with: pip install openviking[gemini]" ) params = param_builder(config) - return embedder_class(**params) + embedder = embedder_class(**params) + if provider == "local": + return embedder + if isinstance(embedder, RetryEmbedder): + return embedder + return RetryEmbedder(embedder) def get_embedder(self): """Get embedder instance based on configuration. diff --git a/tests/unit/test_embedding_config_voyage.py b/tests/unit/test_embedding_config_voyage.py index eef0f0f274..701cc3633f 100644 --- a/tests/unit/test_embedding_config_voyage.py +++ b/tests/unit/test_embedding_config_voyage.py @@ -4,7 +4,7 @@ import pytest -from openviking.models.embedder import VoyageDenseEmbedder +from openviking.models.embedder import RetryEmbedder, VoyageDenseEmbedder from openviking_cli.utils.config.embedding_config import EmbeddingConfig, EmbeddingModelConfig @@ -48,5 +48,6 @@ def test_voyage_get_embedder_returns_voyage_dense_embedder(): ) embedder = config.get_embedder() - assert isinstance(embedder, VoyageDenseEmbedder) + assert isinstance(embedder, RetryEmbedder) + assert isinstance(embedder.inner, VoyageDenseEmbedder) assert embedder.get_dimension() == 1024 diff --git a/tests/unit/test_encoding_format_embedding.py b/tests/unit/test_encoding_format_embedding.py index 278a3b7df7..520b5d2f44 100644 --- a/tests/unit/test_encoding_format_embedding.py +++ b/tests/unit/test_encoding_format_embedding.py @@ -15,7 +15,7 @@ import pytest from pydantic import ValidationError -from openviking.models.embedder import OpenAIDenseEmbedder +from openviking.models.embedder import OpenAIDenseEmbedder, RetryEmbedder from openviking_cli.utils.config.embedding_config import EmbeddingConfig, EmbeddingModelConfig @@ -109,8 +109,9 @@ def test_openai_factory_forwards(self, mock_openai_class): encoding_format="float", ) embedder = EmbeddingConfig(dense=cfg)._create_embedder("openai", "dense", cfg) - assert isinstance(embedder, OpenAIDenseEmbedder) - assert embedder.encoding_format == "float" + assert isinstance(embedder, RetryEmbedder) + assert isinstance(embedder.inner, OpenAIDenseEmbedder) + assert embedder.inner.encoding_format == "float" @patch("openai.AzureOpenAI") def test_azure_factory_forwards(self, mock_azure_class): @@ -125,8 +126,9 @@ def test_azure_factory_forwards(self, mock_azure_class): encoding_format="float", ) embedder = EmbeddingConfig(dense=cfg)._create_embedder("azure", "dense", cfg) - assert isinstance(embedder, OpenAIDenseEmbedder) - assert embedder.encoding_format == "float" + assert isinstance(embedder, RetryEmbedder) + assert isinstance(embedder.inner, OpenAIDenseEmbedder) + assert embedder.inner.encoding_format == "float" @patch("openai.OpenAI") def test_factory_omits_when_unset(self, mock_openai_class): @@ -138,5 +140,6 @@ def test_factory_omits_when_unset(self, mock_openai_class): api_key="sk-test", ) embedder = EmbeddingConfig(dense=cfg)._create_embedder("openai", "dense", cfg) - assert isinstance(embedder, OpenAIDenseEmbedder) - assert embedder.encoding_format is None + assert isinstance(embedder, RetryEmbedder) + assert isinstance(embedder.inner, OpenAIDenseEmbedder) + assert embedder.inner.encoding_format is None diff --git a/tests/unit/test_litellm_embedder.py b/tests/unit/test_litellm_embedder.py index aded09e9cc..cc341897d1 100644 --- a/tests/unit/test_litellm_embedder.py +++ b/tests/unit/test_litellm_embedder.py @@ -206,6 +206,7 @@ def test_factory_creates_litellm_embedder(self, mock_litellm): """EmbeddingConfig factory should create LiteLLMDenseEmbedder for provider='litellm'.""" mock_litellm.embedding.return_value = _mock_litellm_response() + from openviking.models.embedder import RetryEmbedder from openviking.models.embedder.litellm_embedders import LiteLLMDenseEmbedder cfg = EmbeddingModelConfig( @@ -217,10 +218,11 @@ def test_factory_creates_litellm_embedder(self, mock_litellm): ) embedder = EmbeddingConfig(dense=cfg)._create_embedder("litellm", "dense", cfg) - assert isinstance(embedder, LiteLLMDenseEmbedder) + assert isinstance(embedder, RetryEmbedder) + assert isinstance(embedder.inner, LiteLLMDenseEmbedder) assert embedder.model_name == "openai/text-embedding-3-small" - assert embedder.api_key == "test-key" - assert embedder.api_base == "https://openrouter.ai/api/v1" + assert embedder.inner.api_key == "test-key" + assert embedder.inner.api_base == "https://openrouter.ai/api/v1" @patch("openviking.models.embedder.litellm_embedders.litellm") def test_factory_forwards_query_document_params(self, mock_litellm): diff --git a/tests/unit/test_retry_embedder.py b/tests/unit/test_retry_embedder.py new file mode 100644 index 0000000000..6ce38860bd --- /dev/null +++ b/tests/unit/test_retry_embedder.py @@ -0,0 +1,297 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio + +import pytest + +from openviking.concurrency import AsyncSemaphore +from openviking.config.embedding import _bind_account_state +from openviking.models.embedder.base import ( + CompositeHybridEmbedder, + DenseEmbedderBase, + EmbedResult, + FailoverEmbedder, + SparseEmbedderBase, +) +from openviking.models.embedder.retry_embedder import RetryEmbedder +from openviking.models.vlm.token_usage import TokenUsageTracker +from openviking_cli.utils.config.embedding_config import ( + EmbeddingConfig, + EmbeddingCredential, + EmbeddingModelConfig, +) + + +class ProviderError(RuntimeError): + def __init__(self, status_code: int): + super().__init__(f"HTTP {status_code}") + self.status_code = status_code + + +class FakeDenseEmbedder(DenseEmbedderBase): + def __init__(self, outcomes, *, max_retries: int = 2): + super().__init__( + "fake-dense", + {"provider": "fake", "max_retries": max_retries, "max_concurrent": 1}, + ) + self.provider = "fake" + self.outcomes = list(outcomes) + self.calls = 0 + self.closed = False + + def _next(self): + self.calls += 1 + outcome = self.outcomes.pop(0) + if isinstance(outcome, BaseException): + raise outcome + return EmbedResult(dense_vector=outcome) + + def embed(self, content, is_query=False, **kwargs): + return self._run_once(self._next) + + async def embed_async(self, content, is_query=False, **kwargs): + return await self._run_once_async( + self._next_async, + operation_name="fake embedding", + ) + + async def _next_async(self): + return self._next() + + def get_dimension(self): + return 2 + + @property + def supports_multimodal(self): + return True + + def close(self): + self.closed = True + + +class FakeSparseEmbedder(SparseEmbedderBase): + def __init__(self, outcomes, *, max_retries: int = 2): + super().__init__("fake-sparse", {"provider": "fake", "max_retries": max_retries}) + self.provider = "fake" + self.outcomes = list(outcomes) + self.calls = 0 + + def _next(self): + self.calls += 1 + outcome = self.outcomes.pop(0) + if isinstance(outcome, BaseException): + raise outcome + return EmbedResult(sparse_vector=outcome) + + def embed(self, content, is_query=False, **kwargs): + return self._next() + + async def embed_async(self, content, is_query=False, **kwargs): + return self._next() + + +@pytest.fixture(autouse=True) +def no_retry_sleep(monkeypatch): + monkeypatch.setattr("openviking.utils.model_retry.time.sleep", lambda _delay: None) + + +def test_retry_embedder_retries_transient_and_preserves_capabilities(): + inner = FakeDenseEmbedder([ProviderError(503), [0.1, 0.2]]) + embedder = RetryEmbedder(inner) + + result = embedder.embed("query", is_query=True) + + assert result.dense_vector == [0.1, 0.2] + assert inner.calls == 2 + assert embedder.provider == "fake" + assert embedder.get_dimension() == 2 + assert embedder.is_dense is True + assert embedder.is_sparse is False + assert embedder.is_hybrid is False + assert embedder.supports_multimodal is True + + +@pytest.mark.parametrize("status_code", [400, 401]) +def test_retry_embedder_does_not_retry_permanent_or_auth_errors(status_code): + error = ProviderError(status_code) + inner = FakeDenseEmbedder([error, [0.1, 0.2]]) + + with pytest.raises(ProviderError) as raised: + RetryEmbedder(inner).embed("query") + + assert raised.value is error + assert inner.calls == 1 + + +def test_retry_embedder_uses_max_retries_as_additional_attempts(): + errors = [ProviderError(503) for _ in range(3)] + inner = FakeDenseEmbedder(errors, max_retries=2) + + with pytest.raises(ProviderError) as raised: + RetryEmbedder(inner).embed("query") + + assert raised.value is errors[-1] + assert inner.calls == 3 + + +def test_retry_embedder_sparse_dimension_remains_unsupported(): + embedder = RetryEmbedder(FakeSparseEmbedder([{"term": 1.0}])) + + assert embedder.is_sparse is True + with pytest.raises(AttributeError, match="no get_dimension"): + embedder.get_dimension() + + +@pytest.mark.asyncio +async def test_retry_releases_semaphore_before_backoff(monkeypatch): + inner = FakeDenseEmbedder([ProviderError(503), [0.1, 0.2]], max_retries=1) + semaphore = AsyncSemaphore(1) + inner._account_semaphore = semaphore + observed_values = [] + + async def observe_backoff(_delay): + observed_values.append(semaphore._value) + + monkeypatch.setattr("openviking.utils.model_retry.asyncio.sleep", observe_backoff) + + result = await RetryEmbedder(inner).embed_async("query") + + assert result.dense_vector == [0.1, 0.2] + assert inner.calls == 2 + assert observed_values == [1] + assert semaphore._value == 1 + + +@pytest.mark.asyncio +async def test_cancelled_attempt_releases_semaphore(): + started = asyncio.Event() + + class BlockingDenseEmbedder(FakeDenseEmbedder): + async def _next_async(self): + self.calls += 1 + started.set() + await asyncio.Event().wait() + + inner = BlockingDenseEmbedder([], max_retries=2) + semaphore = AsyncSemaphore(1) + inner._account_semaphore = semaphore + task = asyncio.create_task(RetryEmbedder(inner).embed_async("query")) + await started.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert inner.calls == 1 + assert semaphore._value == 1 + + +@pytest.mark.asyncio +async def test_composite_retries_only_failed_branch(monkeypatch): + async def no_sleep(_delay): + return None + + monkeypatch.setattr("openviking.utils.model_retry.asyncio.sleep", no_sleep) + dense = FakeDenseEmbedder([[0.1, 0.2]]) + sparse = FakeSparseEmbedder([ProviderError(503), {"term": 1.0}], max_retries=1) + composite = CompositeHybridEmbedder(RetryEmbedder(dense), RetryEmbedder(sparse)) + + result = await composite.embed_async("query") + + assert result.dense_vector == [0.1, 0.2] + assert result.sparse_vector == {"term": 1.0} + assert dense.calls == 1 + assert sparse.calls == 2 + + +def test_failover_runs_after_current_credential_retry_exhaustion(): + primary = FakeDenseEmbedder([ProviderError(503), ProviderError(503)], max_retries=1) + backup = FakeDenseEmbedder([[0.3, 0.4]], max_retries=1) + failover = FailoverEmbedder( + [RetryEmbedder(primary), RetryEmbedder(backup)], + ["primary", "backup"], + ) + + result = failover.embed("query") + + assert result.dense_vector == [0.3, 0.4] + assert primary.calls == 2 + assert backup.calls == 1 + + +def test_account_state_reaches_retry_inner_and_close_is_delegated_once(): + inner = FakeDenseEmbedder([[0.1, 0.2]]) + embedder = RetryEmbedder(inner) + semaphore = AsyncSemaphore(2) + tracker = TokenUsageTracker() + + _bind_account_state(embedder, semaphore, tracker, "account-a") + + assert embedder._account_semaphore is semaphore + assert inner._account_semaphore is semaphore + assert inner._token_tracker is tracker + assert inner._account_id == "account-a" + embedder.close() + assert inner.closed is True + + +def test_factory_wraps_single_remote_provider(monkeypatch): + monkeypatch.setattr("openai.OpenAI", lambda **_kwargs: object()) + config = EmbeddingConfig( + dense=EmbeddingModelConfig( + provider="openai", + model="text-embedding-3-small", + api_key="key", + dimension=2, + ) + ) + + embedder = config.get_embedder() + + assert isinstance(embedder, RetryEmbedder) + assert embedder.inner.__class__.__name__ == "OpenAIDenseEmbedder" + + +def test_factory_places_retry_inside_failover(monkeypatch): + monkeypatch.setattr("openai.OpenAI", lambda **_kwargs: object()) + config = EmbeddingConfig( + dense=EmbeddingModelConfig( + provider="openai", + model="text-embedding-3-small", + dimension=2, + credentials=[ + EmbeddingCredential(id="primary", api_key="key-a"), + EmbeddingCredential(id="backup", api_key="key-b"), + ], + ) + ) + + embedder = config.get_embedder() + + assert isinstance(embedder, FailoverEmbedder) + assert all(isinstance(child, RetryEmbedder) for child in embedder._embedders) + + +def test_factory_places_retry_inside_composite(monkeypatch): + monkeypatch.setattr("openai.OpenAI", lambda **_kwargs: object()) + monkeypatch.setattr("volcenginesdkarkruntime.Ark", lambda **_kwargs: object()) + config = EmbeddingConfig( + dense=EmbeddingModelConfig( + provider="openai", + model="text-embedding-3-small", + api_key="openai-key", + dimension=2, + ), + sparse=EmbeddingModelConfig( + provider="volcengine", + model="doubao-embedding", + api_key="ark-key", + ), + ) + + embedder = config.get_embedder() + + assert isinstance(embedder, CompositeHybridEmbedder) + assert isinstance(embedder.dense_embedder, RetryEmbedder) + assert isinstance(embedder.sparse_embedder, RetryEmbedder) From ba4701e8994d91c8d22c139f3f39b20d10576816 Mon Sep 17 00:00:00 2001 From: "zhanghe.1216" <@bytedance.com> Date: Tue, 13 Oct 2026 20:14:59 +0800 Subject: [PATCH 05/16] =?UTF-8?q?fix(embedding):=20=E6=81=A2=E5=A4=8D=20Ge?= =?UTF-8?q?mini=20429=20=E9=87=8D=E8=AF=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../models/embedder/gemini_embedders.py | 2 +- tests/unit/test_gemini_embedder.py | 33 ++++++++++++++++++- 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/openviking/models/embedder/gemini_embedders.py b/openviking/models/embedder/gemini_embedders.py index 501b74f661..a81dbf6653 100644 --- a/openviking/models/embedder/gemini_embedders.py +++ b/openviking/models/embedder/gemini_embedders.py @@ -52,7 +52,7 @@ 401: "Invalid API key. Verify your GOOGLE_API_KEY or api_key in config.", 403: "Permission denied. API key may lack access to this model.", 404: "Model not found: '{model}'. Check spelling (e.g. 'gemini-embedding-2-preview').", - 429: "Quota exceeded. Wait and retry, or increase your Google API quota.", + 429: "Resource exhausted. Wait and retry.", 500: "Gemini service error (Google-side). Retry after a delay.", 503: "Gemini service unavailable. Retry after a delay.", } diff --git a/tests/unit/test_gemini_embedder.py b/tests/unit/test_gemini_embedder.py index d0ed804b46..ff33eefb9d 100644 --- a/tests/unit/test_gemini_embedder.py +++ b/tests/unit/test_gemini_embedder.py @@ -143,6 +143,37 @@ def test_embed_raises_runtime_error_on_api_error(self, mock_client_class): with pytest.raises(RuntimeError, match="Gemini embedding failed"): embedder.embed("hello") + @patch("openviking.models.embedder.gemini_embedders.genai.Client") + def test_retry_embedder_retries_transient_429(self, mock_client_class, monkeypatch): + from google.genai.errors import APIError + + from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder + from openviking.models.embedder.retry_embedder import RetryEmbedder + + monkeypatch.setattr("openviking.utils.model_retry.time.sleep", lambda _delay: None) + mock_client = mock_client_class.return_value + mock_response = MagicMock() + mock_response.status_code = 429 + mock_client.models.embed_content.side_effect = [ + APIError( + 429, + {"error": {"message": "Rate limit exceeded. Please retry."}}, + response=mock_response, + ), + _make_mock_result([[0.1, 0.2, 0.3]]), + ] + inner = GeminiDenseEmbedder( + "gemini-embedding-2-preview", + api_key="key", + dimension=3, + config={"max_retries": 1}, + ) + + result = RetryEmbedder(inner).embed("hello") + + assert result.dense_vector == [0.1, 0.2, 0.3] + assert mock_client.models.embed_content.call_count == 2 + @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_embed_empty_string_returns_zero_vector(self, mock_client_class): from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder @@ -198,7 +229,7 @@ class TestGeminiErrorMessages: (401, "Invalid API key"), (403, "Permission denied"), (404, "Model not found"), - (429, "Quota exceeded"), + (429, "Resource exhausted"), (500, "service error"), ], ) From c86fabd25d5d54a0d86bec55bc7f2cf4415bf4af Mon Sep 17 00:00:00 2001 From: "zhanghe.1216" <@bytedance.com> Date: Wed, 14 Oct 2026 15:07:11 +0800 Subject: [PATCH 06/16] =?UTF-8?q?refactor(models):=20=E7=BB=9F=E4=B8=80?= =?UTF-8?q?=E6=A8=A1=E5=9E=8B=E9=94=99=E8=AF=AF=E5=88=86=E7=B1=BB=E4=B8=8E?= =?UTF-8?q?=E9=87=8D=E8=AF=95=E8=BE=B9=E7=95=8C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../models/embedder/cohere_embedders.py | 3 + .../models/embedder/dashscope_embedders.py | 5 + .../models/embedder/gemini_embedders.py | 3 + openviking/models/embedder/jina_embedders.py | 3 + .../models/embedder/litellm_embedders.py | 3 + .../models/embedder/minimax_embedders.py | 21 +- .../models/embedder/openai_embedders.py | 5 +- openviking/models/embedder/retry_embedder.py | 2 + .../models/embedder/vikingdb_embedders.py | 22 +- .../models/embedder/volcengine_embedders.py | 16 +- .../models/embedder/voyage_embedders.py | 3 + openviking/models/model_error.py | 370 ++++++++++++++++++ openviking/models/vlm/backends/litellm_vlm.py | 5 + openviking/models/vlm/backends/openai_vlm.py | 5 + .../models/vlm/backends/volcengine_vlm.py | 29 +- openviking/models/vlm/retry_vlm.py | 4 + openviking/utils/model_retry.py | 62 ++- tests/unit/test_cohere_embedder.py | 3 +- tests/unit/test_extra_headers_vlm.py | 22 +- tests/unit/test_gemini_embedder.py | 10 +- tests/unit/test_jina_embedder.py | 9 +- tests/unit/test_model_error.py | 294 ++++++++++++++ tests/unit/test_model_retry.py | 15 +- tests/unit/test_retry_embedder.py | 11 +- tests/unit/test_retry_vlm.py | 11 +- tests/unit/test_voyage_embedder.py | 3 +- 26 files changed, 882 insertions(+), 57 deletions(-) create mode 100644 openviking/models/model_error.py create mode 100644 tests/unit/test_model_error.py diff --git a/openviking/models/embedder/cohere_embedders.py b/openviking/models/embedder/cohere_embedders.py index 43bdabe86c..b3f502faef 100644 --- a/openviking/models/embedder/cohere_embedders.py +++ b/openviking/models/embedder/cohere_embedders.py @@ -12,6 +12,7 @@ import httpx from openviking.models.embedder.base import DenseEmbedderBase, EmbedResult, truncate_and_normalize +from openviking.models.model_error import model_error_boundary from openviking.utils.async_client_cache import LoopScopedAsyncClientCache from openviking_cli.utils import get_logger @@ -132,6 +133,7 @@ def _normalize_vector(self, vector: List[float]) -> List[float]: return truncate_and_normalize(vector, self._dimension) return vector + @model_error_boundary def embed(self, text: str, is_query: bool = False) -> EmbedResult: input_type = "search_query" if is_query else "search_document" try: @@ -157,6 +159,7 @@ def embed(self, text: str, is_query: bool = False) -> EmbedResult: except Exception as e: raise RuntimeError(f"Cohere embedding failed: {e}") from e + @model_error_boundary async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: input_type = "search_query" if is_query else "search_document" diff --git a/openviking/models/embedder/dashscope_embedders.py b/openviking/models/embedder/dashscope_embedders.py index 44e3d3c959..b29b940805 100644 --- a/openviking/models/embedder/dashscope_embedders.py +++ b/openviking/models/embedder/dashscope_embedders.py @@ -17,6 +17,7 @@ EmbedResult, truncate_and_normalize, ) +from openviking.models.model_error import model_error_boundary from openviking.telemetry import get_current_telemetry from openviking.utils.async_client_cache import LoopScopedAsyncClientCache from openviking_cli.utils.logger import default_logger as logger @@ -225,6 +226,7 @@ def _get_async_httpx_client(self) -> httpx.AsyncClient: # embed # ------------------------------------------------------------------ + @model_error_boundary def embed(self, text: EmbeddingInput, is_query: bool = False) -> EmbedResult: if isinstance(text, list): contents = self._to_dashscope_contents(text) @@ -256,6 +258,7 @@ def _call() -> EmbedResult: # embed_async # ------------------------------------------------------------------ + @model_error_boundary async def embed_async(self, text: EmbeddingInput, is_query: bool = False) -> EmbedResult: if isinstance(text, list): contents = self._to_dashscope_contents(text) @@ -291,6 +294,7 @@ async def _call() -> EmbedResult: # embed_content — multimodal content (text + image URLs) # ------------------------------------------------------------------ + @model_error_boundary def embed_content( self, contents: List[Dict[str, str]], *, enable_fusion: Optional[bool] = None ) -> EmbedResult: @@ -326,6 +330,7 @@ def _call() -> EmbedResult: except Exception as e: raise RuntimeError(f"DashScope content embedding failed: {e}") from e + @model_error_boundary async def embed_content_async( self, contents: List[Dict[str, str]], *, enable_fusion: Optional[bool] = None ) -> EmbedResult: diff --git a/openviking/models/embedder/gemini_embedders.py b/openviking/models/embedder/gemini_embedders.py index a81dbf6653..1963edb586 100644 --- a/openviking/models/embedder/gemini_embedders.py +++ b/openviking/models/embedder/gemini_embedders.py @@ -20,6 +20,7 @@ EmbedResult, truncate_and_normalize, ) +from openviking.models.model_error import model_error_boundary from openviking_cli.utils import get_logger logger = get_logger(__name__) @@ -184,6 +185,7 @@ def __repr__(self) -> str: f"task_type={self.task_type!r})" ) + @model_error_boundary def embed( self, text: str, @@ -225,6 +227,7 @@ def _call() -> EmbedResult: except (APIError, ClientError) as e: _raise_api_error(e, self.model_name) + @model_error_boundary async def embed_async( self, text: str, diff --git a/openviking/models/embedder/jina_embedders.py b/openviking/models/embedder/jina_embedders.py index 74629dd81c..2a00ec491c 100644 --- a/openviking/models/embedder/jina_embedders.py +++ b/openviking/models/embedder/jina_embedders.py @@ -10,6 +10,7 @@ DenseEmbedderBase, EmbedResult, ) +from openviking.models.model_error import model_error_boundary from openviking.utils.async_client_cache import LoopScopedAsyncClientCache from openviking_cli.utils import get_logger @@ -183,6 +184,7 @@ def _raise_task_error(self, error: openai.APIError) -> None: f"to a valid task type for this model. API details: {error.message}" ) from error + @model_error_boundary def embed(self, text: str, is_query: bool = False) -> EmbedResult: """Perform dense embedding on text @@ -224,6 +226,7 @@ def _call() -> EmbedResult: except Exception as e: raise RuntimeError(f"Embedding failed: {str(e)}") from e + @model_error_boundary async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: client = self._get_async_client() diff --git a/openviking/models/embedder/litellm_embedders.py b/openviking/models/embedder/litellm_embedders.py index 2c0365ed88..88c039c872 100644 --- a/openviking/models/embedder/litellm_embedders.py +++ b/openviking/models/embedder/litellm_embedders.py @@ -12,6 +12,7 @@ import litellm from openviking.models.embedder.base import DenseEmbedderBase, EmbedResult +from openviking.models.model_error import model_error_boundary from openviking.telemetry import get_current_telemetry from openviking_cli.utils import get_logger @@ -166,6 +167,7 @@ def _usage_value(key: str, default: int = 0) -> int: completion_tokens=output_tokens, ) + @model_error_boundary def embed(self, text: str, is_query: bool = False) -> EmbedResult: """Perform dense embedding on text via litellm. @@ -199,6 +201,7 @@ def _call() -> EmbedResult: except Exception as e: raise RuntimeError(f"LiteLLM embedding failed: {e}") from e + @model_error_boundary async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: async def _call() -> EmbedResult: kwargs = self._build_kwargs(is_query=is_query) diff --git a/openviking/models/embedder/minimax_embedders.py b/openviking/models/embedder/minimax_embedders.py index 9c46f38ea2..c397124d19 100644 --- a/openviking/models/embedder/minimax_embedders.py +++ b/openviking/models/embedder/minimax_embedders.py @@ -8,6 +8,7 @@ import requests from openviking.models.embedder.base import DenseEmbedderBase, EmbedResult +from openviking.models.model_error import ModelError, ModelErrorKind, model_error_boundary from openviking.utils.async_client_cache import LoopScopedAsyncClientCache from openviking_cli.utils.logger import default_logger as logger @@ -112,7 +113,12 @@ def _call_api(self, texts: List[str], is_query: bool = False) -> List[List[float # Check for business error code base_resp = data.get("base_resp", {}) if base_resp.get("status_code") != 0: - raise RuntimeError(f"MiniMax API error: {base_resp.get('status_msg')}") + raise ModelError( + ModelErrorKind.UNKNOWN, + f"MiniMax API error: {base_resp.get('status_msg')}", + provider="minimax", + code=str(base_resp.get("status_code")), + ) vectors = data.get("vectors", []) if not vectors: @@ -120,6 +126,8 @@ def _call_api(self, texts: List[str], is_query: bool = False) -> List[List[float return vectors + except ModelError: + raise except requests.exceptions.RequestException as e: raise RuntimeError(f"MiniMax network error: {str(e)}") from e except Exception as e: @@ -169,18 +177,26 @@ async def _call_api_async(self, texts: List[str], is_query: bool = False) -> Lis base_resp = data.get("base_resp", {}) if base_resp.get("status_code") != 0: - raise RuntimeError(f"MiniMax API error: {base_resp.get('status_msg')}") + raise ModelError( + ModelErrorKind.UNKNOWN, + f"MiniMax API error: {base_resp.get('status_msg')}", + provider="minimax", + code=str(base_resp.get("status_code")), + ) vectors = data.get("vectors", []) if not vectors: raise RuntimeError("MiniMax API returned empty vectors") return vectors + except ModelError: + raise except httpx.HTTPError as e: raise RuntimeError(f"MiniMax network error: {str(e)}") from e except Exception as e: raise RuntimeError(f"MiniMax embedding failed: {str(e)}") from e + @model_error_boundary def embed(self, text: str, is_query: bool = False) -> EmbedResult: """Perform dense embedding on text""" vectors = self._run_once( @@ -199,6 +215,7 @@ def embed(self, text: str, is_query: bool = False) -> EmbedResult: ) return result + @model_error_boundary async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: async def _call() -> EmbedResult: vectors = await self._call_api_async([text], is_query=is_query) diff --git a/openviking/models/embedder/openai_embedders.py b/openviking/models/embedder/openai_embedders.py index eec981220d..a0a76e9e8f 100644 --- a/openviking/models/embedder/openai_embedders.py +++ b/openviking/models/embedder/openai_embedders.py @@ -6,7 +6,6 @@ import httpx import openai -from openviking_cli.utils import get_logger from openviking.models.embedder.base import ( DenseEmbedderBase, @@ -15,6 +14,7 @@ HybridEmbedderBase, SparseEmbedderBase, ) +from openviking.models.model_error import model_error_boundary from openviking.models.network import ( create_optional_async_httpx_client, create_optional_sync_httpx_client, @@ -22,6 +22,7 @@ from openviking.models.vlm.registry import DEFAULT_AZURE_API_VERSION from openviking.telemetry import get_current_telemetry from openviking.utils.async_client_cache import LoopScopedAsyncClientCache +from openviking_cli.utils import get_logger logger = get_logger(__name__) @@ -362,6 +363,7 @@ def _build_async_client(): return self._async_client_cache.get(_build_async_client) + @model_error_boundary def embed(self, text: EmbeddingInput | List[str], is_query: bool = False) -> EmbedResult: """Perform dense embedding on text, text batches, or multimodal content @@ -400,6 +402,7 @@ def _call() -> EmbedResult: except Exception as e: raise RuntimeError(f"Embedding failed: {str(e)}") from e + @model_error_boundary async def embed_async(self, text: EmbeddingInput | List[str], is_query: bool = False) -> EmbedResult: client = self._get_async_client() embedding_input = self._prepare_embedding_input(text) diff --git a/openviking/models/embedder/retry_embedder.py b/openviking/models/embedder/retry_embedder.py index b7b04aa435..d93be2bd57 100644 --- a/openviking/models/embedder/retry_embedder.py +++ b/openviking/models/embedder/retry_embedder.py @@ -40,6 +40,7 @@ def embed( max_retries=self.max_retries, logger=logger, operation_name=f"{self.provider} embedding", + model_errors_only=True, ) async def embed_async( @@ -53,6 +54,7 @@ async def embed_async( max_retries=self.max_retries, logger=logger, operation_name=f"{self.provider} async embedding", + model_errors_only=True, ) def prepare_embedding_input(self, content: EmbeddingInput) -> EmbeddingInput: diff --git a/openviking/models/embedder/vikingdb_embedders.py b/openviking/models/embedder/vikingdb_embedders.py index c98b41efa9..631c728e2b 100644 --- a/openviking/models/embedder/vikingdb_embedders.py +++ b/openviking/models/embedder/vikingdb_embedders.py @@ -12,6 +12,7 @@ HybridEmbedderBase, SparseEmbedderBase, ) +from openviking.models.model_error import model_error_boundary, model_error_from_status from openviking.storage.vectordb.collection.volcengine_clients import ( DEFAULT_TIMEOUT, ClientForDataApi, @@ -65,10 +66,12 @@ def _call_api( try: response = self.client.do_req("POST", path, req_body=req_body) if response.status_code != 200: - logger.warning( - f"VikingDB API returned bad code: {response.status_code}, message: {response.text}" + raise model_error_from_status( + response.status_code, + f"VikingDB API returned bad code: {response.status_code}, " + f"message: {response.text}", + provider="vikingdb", ) - return [] result = response.json() return result.get("result", {}).get("data", []) @@ -102,12 +105,11 @@ async def _call_api_async( content=req.body, ) if response.status_code != 200: - logger.warning( - "VikingDB API returned bad code: %s, message: %s", + raise model_error_from_status( response.status_code, - response.text, + f"VikingDB API returned bad code: {response.status_code}, message: {response.text}", + provider="vikingdb", ) - return [] result = response.json() return result.get("result", {}).get("data", []) @@ -182,6 +184,7 @@ def _resolve_input_type(self, is_query: bool) -> Optional[str]: return self.document_param return None + @model_error_boundary def embed(self, text: str, is_query: bool = False) -> EmbedResult: input_type = self._resolve_input_type(is_query) @@ -212,6 +215,7 @@ def _call() -> EmbedResult: ) return result + @model_error_boundary async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: async def _call() -> EmbedResult: results = await self._call_api_async([text], dense_model=self.dense_model) @@ -264,6 +268,7 @@ def __init__( "version": model_version, } + @model_error_boundary def embed(self, text: str, is_query: bool = False) -> EmbedResult: def _call() -> EmbedResult: results = self._call_api([text], sparse_model=self.sparse_model) @@ -292,6 +297,7 @@ def _call() -> EmbedResult: ) return result + @model_error_boundary async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: async def _call() -> EmbedResult: results = await self._call_api_async([text], sparse_model=self.sparse_model) @@ -360,6 +366,7 @@ def _resolve_input_type(self, is_query: bool) -> Optional[str]: return self.document_param return None + @model_error_boundary def embed(self, text: str, is_query: bool = False) -> EmbedResult: input_type = self._resolve_input_type(is_query) @@ -398,6 +405,7 @@ def _call() -> EmbedResult: ) return result + @model_error_boundary async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: async def _call() -> EmbedResult: results = await self._call_api_async( diff --git a/openviking/models/embedder/volcengine_embedders.py b/openviking/models/embedder/volcengine_embedders.py index 86b9c0fc09..b13b360729 100644 --- a/openviking/models/embedder/volcengine_embedders.py +++ b/openviking/models/embedder/volcengine_embedders.py @@ -3,12 +3,11 @@ """Volcengine Embedder Implementation""" import time - from typing import Any, Awaitable, Callable, Dict, List, Optional, TypeVar import volcenginesdkarkruntime -from openviking_cli.utils.logger import default_logger as logger +from openviking.metrics.datasources import EmbeddingEventDataSource from openviking.models.embedder.base import ( DenseEmbedderBase, EmbeddingInput, @@ -18,15 +17,16 @@ extract_text_from_content, truncate_and_normalize, ) +from openviking.models.model_error import model_error_boundary from openviking.models.network import ( create_optional_async_httpx_client, create_optional_sync_httpx_client, ) -from openviking.telemetry import get_current_telemetry -from openviking.metrics.datasources import EmbeddingEventDataSource from openviking.observability.context import get_root_observability_context -from openviking.utils.model_retry import extract_metric_error_code +from openviking.telemetry import get_current_telemetry from openviking.utils.async_client_cache import LoopScopedAsyncClientCache +from openviking.utils.model_retry import extract_metric_error_code +from openviking_cli.utils.logger import default_logger as logger VOLCENGINE_CLIENT_REQUEST_ID_HEADER = "X-Client-Request-Id" VOLCENGINE_CLIENT_REQUEST_ID = "ToB-direct,OpenViking_Service,openviking-service_cn-beijing" @@ -254,6 +254,7 @@ def _usage_value(key: str, default: int = 0) -> int: duration_seconds=duration_seconds, ) + @model_error_boundary def embed(self, content: "EmbeddingInput", is_query: bool = False) -> EmbedResult: """Perform dense embedding on text or multimodal content @@ -317,6 +318,7 @@ def _get_async_client(self): lambda: _create_async_ark(self._ark_kwargs) ) + @model_error_boundary async def embed_async(self, content: "EmbeddingInput", is_query: bool = False) -> EmbedResult: client = self._get_async_client() @@ -439,6 +441,7 @@ def _usage_value(key: str, default: int = 0) -> int: duration_seconds=duration_seconds, ) + @model_error_boundary def embed(self, content: "EmbeddingInput", is_query: bool = False) -> EmbedResult: """Perform sparse embedding on text or multimodal content @@ -486,6 +489,7 @@ def _get_async_client(self): lambda: _create_async_ark(self._ark_kwargs) ) + @model_error_boundary async def embed_async(self, content: "EmbeddingInput", is_query: bool = False) -> EmbedResult: client = self._get_async_client() @@ -619,6 +623,7 @@ def supports_multimodal(self) -> bool: """ return False + @model_error_boundary def embed(self, content: "EmbeddingInput", is_query: bool = False) -> EmbedResult: """Perform hybrid embedding on text or multimodal content @@ -671,6 +676,7 @@ def _get_async_client(self): lambda: _create_async_ark(self._ark_kwargs) ) + @model_error_boundary async def embed_async(self, content: "EmbeddingInput", is_query: bool = False) -> EmbedResult: client = self._get_async_client() diff --git a/openviking/models/embedder/voyage_embedders.py b/openviking/models/embedder/voyage_embedders.py index 67708efd3f..2c856a2759 100644 --- a/openviking/models/embedder/voyage_embedders.py +++ b/openviking/models/embedder/voyage_embedders.py @@ -7,6 +7,7 @@ import openai from openviking.models.embedder.base import DenseEmbedderBase, EmbedResult +from openviking.models.model_error import model_error_boundary from openviking.utils.async_client_cache import LoopScopedAsyncClientCache from openviking_cli.utils import get_logger @@ -103,6 +104,7 @@ def _get_async_client(self): ) ) + @model_error_boundary def embed(self, text: str, is_query: bool = False) -> EmbedResult: """Perform dense embedding on text.""" @@ -131,6 +133,7 @@ def _call() -> EmbedResult: except Exception as e: raise RuntimeError(f"Embedding failed: {str(e)}") from e + @model_error_boundary async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: client = self._get_async_client() diff --git a/openviking/models/model_error.py b/openviking/models/model_error.py new file mode 100644 index 0000000000..2128ce212d --- /dev/null +++ b/openviking/models/model_error.py @@ -0,0 +1,370 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Provider-neutral model API errors.""" + +from __future__ import annotations + +import functools +import inspect +from enum import Enum +from typing import Any, Callable, TypeVar, cast + + +class ModelErrorKind(str, Enum): + TRANSIENT_SERVICE_ERROR = "transient_service_error" + TRANSIENT_NETWORK_ERROR = "transient_network_error" + RATE_LIMIT_ERROR = "rate_limit_error" + AUTHENTICATION_ERROR = "authentication_error" + PERMISSION_DENIED = "permission_denied" + QUOTA_EXHAUSTED = "quota_exhausted" + MODEL_NOT_FOUND = "model_not_found" + UNSUPPORTED = "unsupported" + INVALID_REQUEST = "invalid_request" + INPUT_TOO_LARGE = "input_too_large" + CONTENT_SAFETY = "content_safety" + UNKNOWN = "unknown" + + +class ModelError(Exception): + """A structured model API failure reported by a provider boundary.""" + + def __init__( + self, + kind: ModelErrorKind, + message: str, + *, + provider: str | None = None, + status_code: int | None = None, + code: str | None = None, + ) -> None: + super().__init__(message) + self.kind = kind + self.provider = provider + self.status_code = status_code + self.code = code + + +_CODE_KINDS = { + # Authentication and authorization. + "APIKEYINVALID": ModelErrorKind.AUTHENTICATION_ERROR, + "INVALIDAPIKEY": ModelErrorKind.AUTHENTICATION_ERROR, + "UNAUTHENTICATED": ModelErrorKind.AUTHENTICATION_ERROR, + "AUTHENTICATIONERROR": ModelErrorKind.AUTHENTICATION_ERROR, + "PERMISSIONDENIED": ModelErrorKind.PERMISSION_DENIED, + "ACCESSDENIED": ModelErrorKind.PERMISSION_DENIED, + # Quota and billing exhaustion. + "INSUFFICIENTQUOTA": ModelErrorKind.QUOTA_EXHAUSTED, + "ACCOUNTQUOTAEXCEEDED": ModelErrorKind.QUOTA_EXHAUSTED, + "MODELQUOTAEXCEEDED": ModelErrorKind.QUOTA_EXHAUSTED, + "QUOTAEXCEEDED": ModelErrorKind.QUOTA_EXHAUSTED, + "BILLINGHARDLIMITREACHED": ModelErrorKind.QUOTA_EXHAUSTED, + "ACCOUNTOVERDUE": ModelErrorKind.QUOTA_EXHAUSTED, + # Request rate limits. + "RATELIMITEXCEEDED": ModelErrorKind.RATE_LIMIT_ERROR, + "REQUESTLIMITEXCEEDED": ModelErrorKind.RATE_LIMIT_ERROR, + "REQUESTBURSTTOOFAST": ModelErrorKind.RATE_LIMIT_ERROR, + "TOOMANYREQUESTS": ModelErrorKind.RATE_LIMIT_ERROR, + "MODELACCOUNTRPMRATELIMITEXCEEDED": ModelErrorKind.RATE_LIMIT_ERROR, + "MODELACCOUNTTPMRATELIMITEXCEEDED": ModelErrorKind.RATE_LIMIT_ERROR, + "MODELACCOUNTTOTALTOKENRATELIMITEXCEEDED": ModelErrorKind.RATE_LIMIT_ERROR, + # Request shape and model capabilities. + "CONTEXTLENGTHEXCEEDED": ModelErrorKind.INPUT_TOO_LARGE, + "INPUTTEXTEXCEEDSLIMIT": ModelErrorKind.INPUT_TOO_LARGE, + "CONTENTTOOLARGE": ModelErrorKind.INPUT_TOO_LARGE, + "CONTENTFILTER": ModelErrorKind.CONTENT_SAFETY, + "CONTENTPOLICYVIOLATION": ModelErrorKind.CONTENT_SAFETY, + "SAFETY": ModelErrorKind.CONTENT_SAFETY, + "MODELNOTFOUND": ModelErrorKind.MODEL_NOT_FOUND, + "INVALIDMODEL": ModelErrorKind.MODEL_NOT_FOUND, + "INVALIDENDPOINTORMODELNOTFOUND": ModelErrorKind.MODEL_NOT_FOUND, + "UNSUPPORTEDPARAMETER": ModelErrorKind.UNSUPPORTED, + "UNSUPPORTEDVALUE": ModelErrorKind.UNSUPPORTED, +} + +_TRANSIENT_SERVICE_STATUS = {409, 500, 502, 503} +_TRANSIENT_NETWORK_STATUS = {408, 504} + + +def _exception_chain(error: BaseException) -> list[BaseException]: + chain: list[BaseException] = [] + seen: set[int] = set() + current: BaseException | None = error + while current is not None and id(current) not in seen: + chain.append(current) + seen.add(id(current)) + current = current.__cause__ or current.__context__ + return chain + + +def _normalized_code(value: object) -> str | None: + if not isinstance(value, str) or not value: + return None + return "".join(character for character in value.upper() if character.isalnum()) + + +def _body(error: BaseException) -> dict[str, Any]: + value = getattr(error, "body", None) + if not isinstance(value, dict): + response = getattr(error, "response", None) + try: + value = response.json() if response is not None else None + except Exception: + value = None + if isinstance(value, dict): + nested = value.get("error") + if isinstance(nested, dict): + return {**value, **nested} + return value + return {} + + +def _error_code(error: BaseException) -> str | None: + body = _body(error) + for value in ( + getattr(error, "error_code", None), + getattr(error, "code", None), + body.get("error_code"), + body.get("code"), + body.get("type"), + ): + if isinstance(value, str) and value: + return value + return None + + +def _status_code(error: BaseException) -> int | None: + for value in ( + getattr(error, "status_code", None), + getattr(getattr(error, "response", None), "status_code", None), + getattr(error, "code", None), + ): + if isinstance(value, int) and 100 <= value <= 599: + return value + return None + + +def _structured_reasons(value: object) -> set[str]: + reasons: set[str] = set() + if isinstance(value, dict): + for key, item in value.items(): + if key in {"reason", "status"} and isinstance(item, str): + normalized = _normalized_code(item) + if normalized: + reasons.add(normalized) + reasons.update(_structured_reasons(item)) + elif isinstance(value, (list, tuple)): + for item in value: + reasons.update(_structured_reasons(item)) + elif hasattr(value, "model_dump"): + reasons.update(_structured_reasons(value.model_dump())) + return reasons + + +def _kind_from_code(error: BaseException) -> ModelErrorKind | None: + normalized = _normalized_code(_error_code(error)) + if normalized in _CODE_KINDS: + return _CODE_KINDS[normalized] + + details = getattr(error, "details", None) + reasons = _structured_reasons(details) + if reasons & {"QUOTAEXCEEDED", "BILLINGDISABLED"}: + return ModelErrorKind.QUOTA_EXHAUSTED + if reasons & {"RATELIMITEXCEEDED", "SERVINGCAPACITYEXCEEDED"}: + return ModelErrorKind.RATE_LIMIT_ERROR + if reasons & {"APIKEYINVALID", "UNAUTHENTICATED"}: + return ModelErrorKind.AUTHENTICATION_ERROR + return None + + +def _kind_from_status(status_code: int | None) -> ModelErrorKind | None: + if status_code in _TRANSIENT_SERVICE_STATUS: + return ModelErrorKind.TRANSIENT_SERVICE_ERROR + if status_code in _TRANSIENT_NETWORK_STATUS: + return ModelErrorKind.TRANSIENT_NETWORK_ERROR + if status_code == 401: + return ModelErrorKind.AUTHENTICATION_ERROR + if status_code == 403: + return ModelErrorKind.PERMISSION_DENIED + if status_code == 405 or status_code == 501: + return ModelErrorKind.UNSUPPORTED + if status_code == 413: + return ModelErrorKind.INPUT_TOO_LARGE + if status_code == 429: + return ModelErrorKind.RATE_LIMIT_ERROR + if status_code in {400, 422}: + return ModelErrorKind.INVALID_REQUEST + if status_code is not None and status_code >= 500: + return ModelErrorKind.TRANSIENT_SERVICE_ERROR + return None + + +def _sdk_kind(error: BaseException) -> tuple[ModelErrorKind, str] | None: + module = type(error).__module__ + name = type(error).__name__ + mro_names = {base.__name__ for base in type(error).__mro__} + + if isinstance(error, TimeoutError): + return ModelErrorKind.TRANSIENT_NETWORK_ERROR, "network" + if isinstance(error, ConnectionError): + return ModelErrorKind.TRANSIENT_NETWORK_ERROR, "network" + + if module.startswith("openai"): + provider = "openai" + class_kinds = { + "APITimeoutError": ModelErrorKind.TRANSIENT_NETWORK_ERROR, + "APIConnectionError": ModelErrorKind.TRANSIENT_NETWORK_ERROR, + "AuthenticationError": ModelErrorKind.AUTHENTICATION_ERROR, + "PermissionDeniedError": ModelErrorKind.PERMISSION_DENIED, + "NotFoundError": ModelErrorKind.MODEL_NOT_FOUND, + "InternalServerError": ModelErrorKind.TRANSIENT_SERVICE_ERROR, + } + kind = ( + class_kinds.get(name) + or _kind_from_code(error) + or _kind_from_status(_status_code(error)) + ) + return kind or ModelErrorKind.UNKNOWN, provider + + if module.startswith("volcenginesdkarkruntime"): + provider = "volcengine" + class_kinds = { + "ArkAPITimeoutError": ModelErrorKind.TRANSIENT_NETWORK_ERROR, + "ArkAPIConnectionError": ModelErrorKind.TRANSIENT_NETWORK_ERROR, + "ArkAuthenticationError": ModelErrorKind.AUTHENTICATION_ERROR, + "ArkPermissionDeniedError": ModelErrorKind.PERMISSION_DENIED, + "ArkNotFoundError": ModelErrorKind.MODEL_NOT_FOUND, + "ArkInternalServerError": ModelErrorKind.TRANSIENT_SERVICE_ERROR, + "ArkContentFilterFinishReasonError": ModelErrorKind.CONTENT_SAFETY, + } + kind = ( + class_kinds.get(name) + or _kind_from_code(error) + or _kind_from_status(_status_code(error)) + ) + return kind or ModelErrorKind.UNKNOWN, provider + + if module.startswith("litellm"): + provider = "litellm" + class_kinds = { + "Timeout": ModelErrorKind.TRANSIENT_NETWORK_ERROR, + "APIConnectionError": ModelErrorKind.TRANSIENT_NETWORK_ERROR, + "RateLimitError": ModelErrorKind.RATE_LIMIT_ERROR, + "AuthenticationError": ModelErrorKind.AUTHENTICATION_ERROR, + "PermissionDeniedError": ModelErrorKind.PERMISSION_DENIED, + "NotFoundError": ModelErrorKind.MODEL_NOT_FOUND, + "BadRequestError": ModelErrorKind.INVALID_REQUEST, + "ContentPolicyViolationError": ModelErrorKind.CONTENT_SAFETY, + "ContextWindowExceededError": ModelErrorKind.INPUT_TOO_LARGE, + "ServiceUnavailableError": ModelErrorKind.TRANSIENT_SERVICE_ERROR, + "InternalServerError": ModelErrorKind.TRANSIENT_SERVICE_ERROR, + } + return _kind_from_code(error) or class_kinds.get(name, ModelErrorKind.UNKNOWN), provider + + if module.startswith("google.genai"): + kind = _kind_from_code(error) + status_code = _status_code(error) + # RESOURCE_EXHAUSTED is ambiguous unless Google supplies a reason. + if status_code == 429 and kind is None: + kind = ModelErrorKind.UNKNOWN + return kind or _kind_from_status(status_code) or ModelErrorKind.UNKNOWN, "gemini" + + if module.startswith("httpx"): + if name == "HTTPStatusError": + kind = _kind_from_code(error) or _kind_from_status(_status_code(error)) + elif "TimeoutException" in mro_names: + kind = ModelErrorKind.TRANSIENT_NETWORK_ERROR + elif mro_names & {"NetworkError", "TransportError", "ProtocolError", "ProxyError"}: + kind = ModelErrorKind.TRANSIENT_NETWORK_ERROR + else: + kind = None + return kind or ModelErrorKind.UNKNOWN, "http" + + if module.startswith("requests"): + if name == "HTTPError": + kind = _kind_from_code(error) or _kind_from_status(_status_code(error)) + elif mro_names & {"Timeout", "ConnectTimeout", "ReadTimeout", "ConnectionError"}: + kind = ModelErrorKind.TRANSIENT_NETWORK_ERROR + else: + kind = None + return kind or ModelErrorKind.UNKNOWN, "http" + + return None + + +def model_error_from_status( + status_code: int, + message: str, + *, + provider: str, + code: str | None = None, +) -> ModelError: + """Build a model error from a structured HTTP or provider status.""" + normalized = _normalized_code(code) + kind = _CODE_KINDS.get(normalized) if normalized else None + return ModelError( + kind or _kind_from_status(status_code) or ModelErrorKind.UNKNOWN, + message, + provider=provider, + status_code=status_code, + code=code, + ) + + +def model_error_from_exception(error: BaseException) -> ModelError | None: + """Convert a recognized SDK/API exception without inspecting free-form text.""" + for item in _exception_chain(error): + if isinstance(item, ModelError): + return item + classified = _sdk_kind(item) + if classified is None: + continue + kind, provider = classified + return ModelError( + kind, + str(error), + provider=provider, + status_code=_status_code(item), + code=_error_code(item), + ) + return None + + +F = TypeVar("F", bound=Callable[..., Any]) + + +def model_error_boundary(func: F) -> F: + """Translate recognized provider failures at a public provider boundary.""" + + if inspect.iscoroutinefunction(func): + + @functools.wraps(func) + async def _async_wrapper(*args: Any, **kwargs: Any) -> Any: + try: + return await func(*args, **kwargs) + except Exception as error: + converted = model_error_from_exception(error) + if converted is None or converted is error: + raise + if args and isinstance(getattr(args[0], "provider", None), str): + converted.provider = args[0].provider + if any(item is converted for item in _exception_chain(error)): + raise converted + raise converted from error + + return cast(F, _async_wrapper) + + @functools.wraps(func) + def _sync_wrapper(*args: Any, **kwargs: Any) -> Any: + try: + return func(*args, **kwargs) + except Exception as error: + converted = model_error_from_exception(error) + if converted is None or converted is error: + raise + if args and isinstance(getattr(args[0], "provider", None), str): + converted.provider = args[0].provider + if any(item is converted for item in _exception_chain(error)): + raise converted + raise converted from error + + return cast(F, _sync_wrapper) diff --git a/openviking/models/vlm/backends/litellm_vlm.py b/openviking/models/vlm/backends/litellm_vlm.py index 699ad343ff..917046e6e6 100644 --- a/openviking/models/vlm/backends/litellm_vlm.py +++ b/openviking/models/vlm/backends/litellm_vlm.py @@ -16,6 +16,7 @@ from litellm import acompletion, completion from litellm.types.utils import all_litellm_params +from openviking.models.model_error import model_error_boundary from openviking.telemetry import tracer from openviking.utils.message_format import format_messages, sanitize_openai_messages from openviking.utils.multimodal import redact_image_data_urls @@ -454,6 +455,7 @@ def _build_vision_kwargs( kwargs_messages = [{"role": "user", "content": content}] return self._build_kwargs(model, kwargs_messages, tools, tool_choice, thinking=thinking) + @model_error_boundary def get_completion( self, prompt: str = "", @@ -477,6 +479,7 @@ def _call() -> Union[str, VLMResponse]: return _call() + @model_error_boundary async def get_completion_async( self, prompt: str = "", @@ -508,6 +511,7 @@ async def _call() -> Union[str, VLMResponse]: return await _call() + @model_error_boundary def get_vision_completion( self, prompt: str = "", @@ -531,6 +535,7 @@ def _call() -> Union[str, VLMResponse]: return _call() + @model_error_boundary async def get_vision_completion_async( self, prompt: str = "", diff --git a/openviking/models/vlm/backends/openai_vlm.py b/openviking/models/vlm/backends/openai_vlm.py index 7f6023eadf..fec0d336f4 100644 --- a/openviking/models/vlm/backends/openai_vlm.py +++ b/openviking/models/vlm/backends/openai_vlm.py @@ -12,6 +12,7 @@ import httpx +from openviking.models.model_error import model_error_boundary from openviking.models.network import ( create_optional_async_httpx_client, create_optional_sync_httpx_client, @@ -334,6 +335,7 @@ async def _extract_completion_content_async(self, response, elapsed: float) -> s content = self._extract_content_from_response(response) return self._clean_response(content) + @model_error_boundary def get_completion( self, prompt: str = "", @@ -358,6 +360,7 @@ def _call() -> Union[str, VLMResponse]: return _call() + @model_error_boundary async def get_completion_async( self, prompt: str = "", @@ -442,6 +445,7 @@ def _prepare_image(self, image: Union[str, Path, bytes]) -> Dict[str, Any]: } return {"type": "image_url", "image_url": {"url": image}} + @model_error_boundary def get_vision_completion( self, prompt: str = "", @@ -469,6 +473,7 @@ def _call() -> Union[str, VLMResponse]: return _call() + @model_error_boundary async def get_vision_completion_async( self, prompt: str = "", diff --git a/openviking/models/vlm/backends/volcengine_vlm.py b/openviking/models/vlm/backends/volcengine_vlm.py index e51bbc054f..4f6c5f3be9 100644 --- a/openviking/models/vlm/backends/volcengine_vlm.py +++ b/openviking/models/vlm/backends/volcengine_vlm.py @@ -9,6 +9,7 @@ from pathlib import Path from typing import Any, Dict, List, Optional, Union +from openviking.models.model_error import model_error_boundary from openviking.models.network import ( create_optional_async_httpx_client, create_optional_sync_httpx_client, @@ -119,12 +120,12 @@ def get_client(self): raise ImportError( "Please install volcenginesdkarkruntime: pip install volcenginesdkarkruntime" ) - kwargs = dict( - api_key=self.api_key, - base_url=self.api_base, - timeout=self.timeout, - max_retries=0, - ) + kwargs = { + "api_key": self.api_key, + "base_url": self.api_base, + "timeout": self.timeout, + "max_retries": 0, + } http_client = create_optional_sync_httpx_client( self.api_base, timeout=self.timeout, @@ -142,12 +143,12 @@ def _build_async_client(self): raise ImportError( "Please install volcenginesdkarkruntime: pip install volcenginesdkarkruntime" ) - kwargs = dict( - api_key=self.api_key, - base_url=self.api_base, - timeout=self.timeout, - max_retries=0, - ) + kwargs = { + "api_key": self.api_key, + "base_url": self.api_base, + "timeout": self.timeout, + "max_retries": 0, + } http_client = create_optional_async_httpx_client( self.api_base, timeout=self.timeout, @@ -189,6 +190,7 @@ async def get_media_completion_async( media_type=media_type, ) + @model_error_boundary def get_completion( self, prompt: str = "", @@ -231,6 +233,7 @@ def get_completion( return result return self._clean_response(str(result)) + @model_error_boundary async def get_completion_async( self, prompt: str = "", @@ -408,6 +411,7 @@ def _prepare_image(self, image: Union[str, Path, bytes]) -> Dict[str, Any]: else: return {"type": "image_url", "image_url": {"url": image}} + @model_error_boundary def get_vision_completion( self, prompt: str = "", @@ -458,6 +462,7 @@ def get_vision_completion( return result return self._clean_response(str(result)) + @model_error_boundary async def get_vision_completion_async( self, prompt: str = "", diff --git a/openviking/models/vlm/retry_vlm.py b/openviking/models/vlm/retry_vlm.py index d8b01d7086..3c68fd8102 100644 --- a/openviking/models/vlm/retry_vlm.py +++ b/openviking/models/vlm/retry_vlm.py @@ -62,6 +62,7 @@ def get_completion( max_retries=self.max_retries, logger=logger, operation_name=f"{self.provider} VLM completion", + model_errors_only=True, ) async def get_completion_async( @@ -87,6 +88,7 @@ async def get_completion_async( max_retries=self.max_retries, logger=logger, operation_name=f"{self.provider} VLM async completion", + model_errors_only=True, ) def get_vision_completion( @@ -112,6 +114,7 @@ def get_vision_completion( max_retries=self.max_retries, logger=logger, operation_name=f"{self.provider} VLM vision completion", + model_errors_only=True, ) async def get_vision_completion_async( @@ -137,6 +140,7 @@ async def get_vision_completion_async( max_retries=self.max_retries, logger=logger, operation_name=f"{self.provider} VLM async vision completion", + model_errors_only=True, ) def supports_media( diff --git a/openviking/utils/model_retry.py b/openviking/utils/model_retry.py index 2caf88c82c..cb34396c81 100644 --- a/openviking/utils/model_retry.py +++ b/openviking/utils/model_retry.py @@ -8,6 +8,7 @@ import time from typing import Awaitable, Callable, TypeVar +from openviking.models.model_error import ModelErrorKind, model_error_from_exception from openviking.pyagfs.exceptions import AGFSNotADirectoryError from openviking.utils.exceptions import AllCredentialsFailedError @@ -230,6 +231,23 @@ def classify_api_error(error: Exception) -> str: return ERROR_CLASS_AUTH return ERROR_CLASS_UNKNOWN + model_error = model_error_from_exception(error) + if model_error is not None: + return { + ModelErrorKind.TRANSIENT_SERVICE_ERROR: ERROR_CLASS_TRANSIENT, + ModelErrorKind.TRANSIENT_NETWORK_ERROR: ERROR_CLASS_TRANSIENT, + ModelErrorKind.RATE_LIMIT_ERROR: ERROR_CLASS_TRANSIENT, + ModelErrorKind.AUTHENTICATION_ERROR: ERROR_CLASS_AUTH, + ModelErrorKind.PERMISSION_DENIED: ERROR_CLASS_AUTH, + ModelErrorKind.QUOTA_EXHAUSTED: ERROR_CLASS_QUOTA_EXCEEDED, + ModelErrorKind.MODEL_NOT_FOUND: ERROR_CLASS_PERMANENT, + ModelErrorKind.UNSUPPORTED: ERROR_CLASS_PERMANENT, + ModelErrorKind.INVALID_REQUEST: ERROR_CLASS_PERMANENT, + ModelErrorKind.INPUT_TOO_LARGE: ERROR_CLASS_INPUT_TOO_LARGE, + ModelErrorKind.CONTENT_SAFETY: ERROR_CLASS_CONTENT_SAFETY, + ModelErrorKind.UNKNOWN: ERROR_CLASS_UNKNOWN, + }[model_error.kind] + for exc in (error, getattr(error, "__cause__", None)): if exc is not None and isinstance(exc, _PERMANENT_IO_ERRORS): return ERROR_CLASS_PERMANENT @@ -393,6 +411,28 @@ def _compute_delay( return delay +def _retry_limit( + error: Exception, + *, + max_retries: int, + is_retryable: Callable[[Exception], bool], + model_errors_only: bool, +) -> int: + model_error = model_error_from_exception(error) + if model_error is not None: + if model_error.kind in { + ModelErrorKind.TRANSIENT_SERVICE_ERROR, + ModelErrorKind.TRANSIENT_NETWORK_ERROR, + ModelErrorKind.RATE_LIMIT_ERROR, + ModelErrorKind.UNKNOWN, + }: + return max_retries + return 0 + if model_errors_only: + return 0 + return max_retries if is_retryable(error) else 0 + + def retry_sync( func: Callable[[], T], *, @@ -403,6 +443,7 @@ def retry_sync( is_retryable: Callable[[Exception], bool] = is_retryable_api_error, logger=None, operation_name: str = "operation", + model_errors_only: bool = False, ) -> T: """Retry a sync function on known transient errors.""" attempt = 0 @@ -411,7 +452,13 @@ def retry_sync( try: return func() except Exception as e: - if max_retries <= 0 or attempt >= max_retries or not is_retryable(e): + retry_limit = _retry_limit( + e, + max_retries=max_retries, + is_retryable=is_retryable, + model_errors_only=model_errors_only, + ) + if attempt >= retry_limit: raise delay = _compute_delay( @@ -425,7 +472,7 @@ def retry_sync( "%s failed with retryable error (retry %d/%d): %s; retrying in %.2fs", operation_name, attempt + 1, - max_retries, + retry_limit, e, delay, ) @@ -443,6 +490,7 @@ async def retry_async( is_retryable: Callable[[Exception], bool] = is_retryable_api_error, logger=None, operation_name: str = "operation", + model_errors_only: bool = False, ) -> T: """Retry an async function on known transient errors.""" attempt = 0 @@ -451,7 +499,13 @@ async def retry_async( try: return await func() except Exception as e: - if max_retries <= 0 or attempt >= max_retries or not is_retryable(e): + retry_limit = _retry_limit( + e, + max_retries=max_retries, + is_retryable=is_retryable, + model_errors_only=model_errors_only, + ) + if attempt >= retry_limit: raise delay = _compute_delay( @@ -465,7 +519,7 @@ async def retry_async( "%s failed with retryable error (retry %d/%d): %s; retrying in %.2fs", operation_name, attempt + 1, - max_retries, + retry_limit, e, delay, ) diff --git a/tests/unit/test_cohere_embedder.py b/tests/unit/test_cohere_embedder.py index 27aee11fc0..82dbc2c090 100644 --- a/tests/unit/test_cohere_embedder.py +++ b/tests/unit/test_cohere_embedder.py @@ -12,6 +12,7 @@ COHERE_MODEL_DIMENSIONS, get_cohere_model_default_dimension, ) +from openviking.models.model_error import ModelError class TestCohereDenseEmbedder: @@ -167,7 +168,7 @@ def test_embed_api_error(self, mock_client_class): api_key="bad-key", ) - with pytest.raises(RuntimeError, match="Cohere API error"): + with pytest.raises(ModelError, match="Cohere API error"): embedder.embed("Hello world") @patch("openviking.models.embedder.cohere_embedders.httpx.Client") diff --git a/tests/unit/test_extra_headers_vlm.py b/tests/unit/test_extra_headers_vlm.py index d8c86108e6..62b2ede2f5 100644 --- a/tests/unit/test_extra_headers_vlm.py +++ b/tests/unit/test_extra_headers_vlm.py @@ -14,6 +14,7 @@ from volcenginesdkarkruntime import Ark, AsyncArk from volcenginesdkarkruntime._exceptions import ArkAPIStatusError +from openviking.models.model_error import ModelError, ModelErrorKind from openviking.models.vlm.backends.litellm_vlm import ( LiteLLMVLMProvider, detect_provider_by_model, @@ -349,17 +350,21 @@ def test_volcengine_sync_client_applies_timeout_and_disables_sdk_retries( assert call_kwargs["max_retries"] == 0 @pytest.mark.parametrize( - ("status_code", "message"), + ("status_code", "message", "expected_kind"), [ - (400, "Total tokens of multi-modal content and text exceed max message tokens."), - (400, "The parameter model is invalid."), - (401, "Unauthorized"), - (503, "Service Unavailable"), + ( + 400, + "Total tokens of multi-modal content and text exceed max message tokens.", + ModelErrorKind.INVALID_REQUEST, + ), + (400, "The parameter model is invalid.", ModelErrorKind.INVALID_REQUEST), + (401, "Unauthorized", ModelErrorKind.AUTHENTICATION_ERROR), + (503, "Service Unavailable", ModelErrorKind.TRANSIENT_SERVICE_ERROR), ], ) @patch("volcenginesdkarkruntime.AsyncArk") def test_volcengine_async_client_performs_single_request( - self, mock_async_ark_class, status_code, message + self, mock_async_ark_class, status_code, message, expected_kind ): error = ArkAPIStatusError( f"Error code: {status_code} - {message}", @@ -382,10 +387,11 @@ def test_volcengine_async_client_performs_single_request( } ) - with pytest.raises(ArkAPIStatusError) as raised: + with pytest.raises(ModelError) as raised: asyncio.run(vlm.get_completion_async("hello")) - assert raised.value is error + assert raised.value.kind == expected_kind + assert raised.value.__cause__ is error assert mock_client.chat.completions.create.await_count == 1 call_kwargs = mock_async_ark_class.call_args.kwargs assert call_kwargs["timeout"] == 12.0 diff --git a/tests/unit/test_gemini_embedder.py b/tests/unit/test_gemini_embedder.py index ff33eefb9d..c4b74f4749 100644 --- a/tests/unit/test_gemini_embedder.py +++ b/tests/unit/test_gemini_embedder.py @@ -9,6 +9,8 @@ import pytest +from openviking.models.model_error import ModelError + def _make_mock_embedding(values): emb = MagicMock() @@ -130,7 +132,7 @@ def test_embed_passes_task_type_in_config(self, mock_client_class): assert kwargs["config"].task_type == "RETRIEVAL_QUERY" @patch("openviking.models.embedder.gemini_embedders.genai.Client") - def test_embed_raises_runtime_error_on_api_error(self, mock_client_class): + def test_embed_raises_model_error_on_api_error(self, mock_client_class): from google.genai.errors import APIError from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder @@ -140,7 +142,7 @@ def test_embed_raises_runtime_error_on_api_error(self, mock_client_class): mock_response.status_code = 401 mock_client.models.embed_content.side_effect = APIError(401, {}, response=mock_response) embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="key") - with pytest.raises(RuntimeError, match="Gemini embedding failed"): + with pytest.raises(ModelError, match="Gemini embedding failed"): embedder.embed("hello") @patch("openviking.models.embedder.gemini_embedders.genai.Client") @@ -244,7 +246,7 @@ def test_error_code_hint(self, mock_client, code, match): mock_resp.status_code = code mock.models.embed_content.side_effect = APIError(code, {}, response=mock_resp) embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="k") - with pytest.raises(RuntimeError, match=match): + with pytest.raises(ModelError, match=match): embedder.embed("hello") @patch("openviking.models.embedder.gemini_embedders.genai.Client") @@ -258,7 +260,7 @@ def test_error_message_includes_http_code(self, mock_client): mock_resp.status_code = 404 mock.models.embed_content.side_effect = APIError(404, {}, response=mock_resp) embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="k") - with pytest.raises(RuntimeError, match="HTTP 404"): + with pytest.raises(ModelError, match="HTTP 404"): embedder.embed("hello") diff --git a/tests/unit/test_jina_embedder.py b/tests/unit/test_jina_embedder.py index 0f2266482d..7cc36689f3 100644 --- a/tests/unit/test_jina_embedder.py +++ b/tests/unit/test_jina_embedder.py @@ -10,6 +10,7 @@ from openviking.models.embedder.jina_embedders import ( JINA_MODEL_DIMENSIONS, ) +from openviking.models.model_error import ModelError from openviking_cli.utils.config.embedding_config import EmbeddingModelConfig @@ -240,7 +241,7 @@ def test_embed_api_error(self, mock_openai_class): api_key="test-api-key", ) - with pytest.raises(RuntimeError, match="Jina API error"): + with pytest.raises(ModelError, match="Jina API error"): embedder.embed("Hello world") def test_build_extra_body_none(self): @@ -290,7 +291,7 @@ def test_dimension_validation_within_range(self): @patch("openviking.models.embedder.jina_embedders.openai.OpenAI") def test_422_task_error_actionable_message(self, mock_openai_class): - """422 error mentioning 'task' should produce actionable RuntimeError.""" + """422 error mentioning 'task' should preserve the actionable provider message.""" import openai mock_client = MagicMock() @@ -308,7 +309,7 @@ def test_422_task_error_actionable_message(self, mock_openai_class): api_key="test-key", ) - with pytest.raises(RuntimeError, match="query_param.*document_param"): + with pytest.raises(ModelError, match="query_param.*document_param"): embedder.embed("hello") @patch("openviking.models.embedder.jina_embedders.openai.OpenAI") @@ -331,7 +332,7 @@ def test_non_422_error_passthrough(self, mock_openai_class): api_key="test-key", ) - with pytest.raises(RuntimeError, match="Jina API error"): + with pytest.raises(ModelError, match="Jina API error"): embedder.embed("hello") def test_code_model_dimensions(self): diff --git a/tests/unit/test_model_error.py b/tests/unit/test_model_error.py new file mode 100644 index 0000000000..ac5461ae5c --- /dev/null +++ b/tests/unit/test_model_error.py @@ -0,0 +1,294 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio + +import httpx +import openai +import pytest +from google.genai.errors import ClientError +from litellm import exceptions as litellm_exceptions +from volcenginesdkarkruntime import _exceptions as ark_exceptions + +from openviking.models.model_error import ( + ModelError, + ModelErrorKind, + model_error_boundary, + model_error_from_exception, + model_error_from_status, +) +from openviking.utils.model_retry import retry_async, retry_sync + + +def _response(status_code: int) -> httpx.Response: + return httpx.Response(status_code, request=httpx.Request("POST", "https://api.example.test")) + + +def _http_status_error(status_code: int, body: dict) -> httpx.HTTPStatusError: + response = httpx.Response( + status_code, + request=httpx.Request("POST", "https://api.example.test"), + json=body, + ) + return httpx.HTTPStatusError( + "provider request failed", + request=response.request, + response=response, + ) + + +@pytest.mark.parametrize( + ("error", "expected"), + [ + ( + openai.RateLimitError( + "quota", + response=_response(429), + body={"error": {"code": "insufficient_quota"}}, + ), + ModelErrorKind.QUOTA_EXHAUSTED, + ), + ( + openai.RateLimitError( + "rate limited", + response=_response(429), + body={"error": {"code": "rate_limit_exceeded"}}, + ), + ModelErrorKind.RATE_LIMIT_ERROR, + ), + ( + openai.APIConnectionError(request=httpx.Request("POST", "https://api.openai.com")), + ModelErrorKind.TRANSIENT_NETWORK_ERROR, + ), + ( + ark_exceptions.ArkRateLimitError( + "quota", + response=_response(429), + body={"error": {"code": "AccountQuotaExceeded"}}, + request_id="request-id", + ), + ModelErrorKind.QUOTA_EXHAUSTED, + ), + ( + ark_exceptions.ArkContentFilterFinishReasonError(), + ModelErrorKind.CONTENT_SAFETY, + ), + ( + litellm_exceptions.ContextWindowExceededError( + "too large", + model="model", + llm_provider="openai", + ), + ModelErrorKind.INPUT_TOO_LARGE, + ), + ( + litellm_exceptions.ServiceUnavailableError( + "unavailable", + model="model", + llm_provider="openai", + ), + ModelErrorKind.TRANSIENT_SERVICE_ERROR, + ), + ( + httpx.ReadTimeout("timed out"), + ModelErrorKind.TRANSIENT_NETWORK_ERROR, + ), + ( + _http_status_error(429, {"code": "AccountQuotaExceeded"}), + ModelErrorKind.QUOTA_EXHAUSTED, + ), + ( + openai.NotFoundError( + "model not found", + response=_response(404), + body={"error": {"code": "model_not_found"}}, + ), + ModelErrorKind.MODEL_NOT_FOUND, + ), + ( + _http_status_error(404, {}), + ModelErrorKind.UNKNOWN, + ), + ], +) +def test_sdk_errors_map_to_model_error_kind(error, expected): + converted = model_error_from_exception(error) + + assert converted is not None + assert converted.kind == expected + + +def test_plain_http_404_requires_structured_model_code(): + converted = model_error_from_status( + 404, + "not found", + provider="http-provider", + ) + + assert converted.kind == ModelErrorKind.UNKNOWN + + +@pytest.mark.parametrize( + ("reason", "expected"), + [ + ("QUOTA_EXCEEDED", ModelErrorKind.QUOTA_EXHAUSTED), + ("RATE_LIMIT_EXCEEDED", ModelErrorKind.RATE_LIMIT_ERROR), + (None, ModelErrorKind.UNKNOWN), + ], +) +def test_gemini_resource_exhausted_requires_structured_reason(reason, expected): + details = [] if reason is None else [{"reason": reason}] + error = ClientError( + 429, + { + "error": { + "code": 429, + "message": "resource exhausted", + "status": "RESOURCE_EXHAUSTED", + "details": details, + } + }, + ) + + converted = model_error_from_exception(error) + + assert converted is not None + assert converted.kind == expected + + +def test_provider_boundary_preserves_original_sdk_error_as_cause(): + sdk_error = openai.BadRequestError( + "bad request", + response=_response(400), + body={"error": {"code": "invalid_request_error"}}, + ) + + @model_error_boundary + def _call(): + raise sdk_error + + with pytest.raises(ModelError) as raised: + _call() + + assert raised.value.kind == ModelErrorKind.INVALID_REQUEST + assert raised.value.__cause__ is sdk_error + + +def test_provider_boundary_reports_actual_provider_name(): + class _Provider: + provider = "jina" + + @model_error_boundary + def call(self): + raise openai.APIConnectionError( + request=httpx.Request("POST", "https://api.jina.ai") + ) + + with pytest.raises(ModelError) as raised: + _Provider().call() + + assert raised.value.provider == "jina" + + +def test_provider_boundary_does_not_convert_local_program_error(): + error = ValueError("HTTP 503 is data, not an API failure") + + @model_error_boundary + def _call(): + raise error + + with pytest.raises(ValueError) as raised: + _call() + + assert raised.value is error + + +def test_unknown_model_error_uses_standard_retry_budget(): + attempts = 0 + + def _call(): + nonlocal attempts + attempts += 1 + raise ModelError(ModelErrorKind.UNKNOWN, "unclassified API error") + + with pytest.raises(ModelError): + retry_sync( + _call, + max_retries=5, + base_delay=0, + jitter=False, + model_errors_only=True, + ) + + assert attempts == 6 + + +@pytest.mark.parametrize( + "kind", + [ + ModelErrorKind.AUTHENTICATION_ERROR, + ModelErrorKind.PERMISSION_DENIED, + ModelErrorKind.QUOTA_EXHAUSTED, + ModelErrorKind.MODEL_NOT_FOUND, + ModelErrorKind.UNSUPPORTED, + ModelErrorKind.INVALID_REQUEST, + ModelErrorKind.INPUT_TOO_LARGE, + ModelErrorKind.CONTENT_SAFETY, + ], +) +def test_strict_retry_propagates_unhandled_model_error_unchanged(kind): + attempts = 0 + error = ModelError(kind, "handled by another layer") + + def _call(): + nonlocal attempts + attempts += 1 + raise error + + with pytest.raises(ModelError) as raised: + retry_sync( + _call, + max_retries=5, + base_delay=0, + jitter=False, + model_errors_only=True, + ) + + assert raised.value is error + assert attempts == 1 + + +def test_strict_retry_does_not_parse_local_error_message(): + attempts = 0 + + def _call(): + nonlocal attempts + attempts += 1 + raise RuntimeError("HTTP 503") + + with pytest.raises(RuntimeError): + retry_sync( + _call, + max_retries=5, + base_delay=0, + jitter=False, + model_errors_only=True, + ) + + assert attempts == 1 + + +@pytest.mark.asyncio +async def test_provider_boundary_and_retry_preserve_cancellation(): + attempts = 0 + + @model_error_boundary + async def _call(): + nonlocal attempts + attempts += 1 + raise asyncio.CancelledError + + with pytest.raises(asyncio.CancelledError): + await retry_async(_call, max_retries=5, model_errors_only=True) + + assert attempts == 1 diff --git a/tests/unit/test_model_retry.py b/tests/unit/test_model_retry.py index cedf5a2d87..1717c723cf 100644 --- a/tests/unit/test_model_retry.py +++ b/tests/unit/test_model_retry.py @@ -4,6 +4,7 @@ import pytest +from openviking.models.model_error import ModelError, ModelErrorKind from openviking.pyagfs.exceptions import ( AGFSClientError, AGFSConnectionError, @@ -99,8 +100,18 @@ def test_extract_metric_error_code_uses_safe_fallbacks_only(): def test_classify_all_credentials_failed_prefers_transient_over_auth(): err = AllCredentialsFailedError( [ - ("cred-a", ERROR_CLASS_AUTH, RuntimeError("401 Unauthorized"), 0), - ("cred-b", ERROR_CLASS_TRANSIENT, RuntimeError("500 server error"), 1), + ( + "cred-a", + ERROR_CLASS_AUTH, + ModelError(ModelErrorKind.AUTHENTICATION_ERROR, "unauthorized"), + 0, + ), + ( + "cred-b", + ERROR_CLASS_TRANSIENT, + ModelError(ModelErrorKind.TRANSIENT_SERVICE_ERROR, "server error"), + 1, + ), ] ) assert classify_api_error(err) == ERROR_CLASS_TRANSIENT diff --git a/tests/unit/test_retry_embedder.py b/tests/unit/test_retry_embedder.py index 6ce38860bd..29a5009d57 100644 --- a/tests/unit/test_retry_embedder.py +++ b/tests/unit/test_retry_embedder.py @@ -15,6 +15,7 @@ SparseEmbedderBase, ) from openviking.models.embedder.retry_embedder import RetryEmbedder +from openviking.models.model_error import ModelError, ModelErrorKind from openviking.models.vlm.token_usage import TokenUsageTracker from openviking_cli.utils.config.embedding_config import ( EmbeddingConfig, @@ -23,10 +24,14 @@ ) -class ProviderError(RuntimeError): +class ProviderError(ModelError): def __init__(self, status_code: int): - super().__init__(f"HTTP {status_code}") - self.status_code = status_code + kinds = { + 400: ModelErrorKind.INVALID_REQUEST, + 401: ModelErrorKind.AUTHENTICATION_ERROR, + 503: ModelErrorKind.TRANSIENT_SERVICE_ERROR, + } + super().__init__(kinds[status_code], f"HTTP {status_code}", status_code=status_code) class FakeDenseEmbedder(DenseEmbedderBase): diff --git a/tests/unit/test_retry_vlm.py b/tests/unit/test_retry_vlm.py index c14a21e025..28265ff28e 100644 --- a/tests/unit/test_retry_vlm.py +++ b/tests/unit/test_retry_vlm.py @@ -6,20 +6,25 @@ import pytest +from openviking.models.model_error import ModelError, ModelErrorKind from openviking.models.vlm.base import MultiCredentialVLM, VLMBase, VLMFactory from openviking.models.vlm.retry_vlm import RetryVLM from openviking_cli.utils.config.vlm_config import VLMConfig -class ProviderError(RuntimeError): +class ProviderError(ModelError): def __init__(self, status_code: int): messages = { 400: "400 invalid parameter", 401: "401 unauthorized", 503: "503 service unavailable", } - super().__init__(messages[status_code]) - self.status_code = status_code + kinds = { + 400: ModelErrorKind.INVALID_REQUEST, + 401: ModelErrorKind.AUTHENTICATION_ERROR, + 503: ModelErrorKind.TRANSIENT_SERVICE_ERROR, + } + super().__init__(kinds[status_code], messages[status_code], status_code=status_code) class FakeVLM(VLMBase): diff --git a/tests/unit/test_voyage_embedder.py b/tests/unit/test_voyage_embedder.py index fffcbad894..72b68586da 100644 --- a/tests/unit/test_voyage_embedder.py +++ b/tests/unit/test_voyage_embedder.py @@ -8,6 +8,7 @@ from openviking.models.embedder import VoyageDenseEmbedder from openviking.models.embedder.voyage_embedders import VOYAGE_MODEL_DIMENSIONS +from openviking.models.model_error import ModelError class TestVoyageDenseEmbedder: @@ -114,5 +115,5 @@ def test_embed_api_error(self, mock_openai_class): api_key="voyage-key", ) - with pytest.raises(RuntimeError, match="Voyage API error"): + with pytest.raises(ModelError, match="Voyage API error"): embedder.embed("Hello world") From 95ca04eb7b53b38bbe57c076a936a2442e64db84 Mon Sep 17 00:00:00 2001 From: "zhanghe.1216" <@bytedance.com> Date: Wed, 14 Oct 2026 16:13:58 +0800 Subject: [PATCH 07/16] =?UTF-8?q?perf(retrieval):=20=E5=9C=A8=E7=BA=BF?= =?UTF-8?q?=E6=A3=80=E7=B4=A2=E7=A6=81=E7=94=A8=E6=A8=A1=E5=9E=8B=E9=87=8D?= =?UTF-8?q?=E8=AF=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- openviking/models/embedder/retry_embedder.py | 6 ++-- openviking/models/vlm/retry_vlm.py | 10 +++--- openviking/service/search_service.py | 3 ++ openviking/utils/model_retry.py | 28 +++++++++++++++- tests/service/test_search_service_image.py | 20 ++++++++++++ tests/unit/test_model_retry.py | 34 ++++++++++++++++++++ tests/unit/test_retry_embedder.py | 14 ++++++++ tests/unit/test_retry_vlm.py | 14 ++++++++ 8 files changed, 120 insertions(+), 9 deletions(-) diff --git a/openviking/models/embedder/retry_embedder.py b/openviking/models/embedder/retry_embedder.py index d93be2bd57..b176c52145 100644 --- a/openviking/models/embedder/retry_embedder.py +++ b/openviking/models/embedder/retry_embedder.py @@ -5,7 +5,7 @@ from typing import Any, Dict from openviking.models.embedder.base import EmbedderBase, EmbeddingInput, EmbedResult -from openviking.utils.model_retry import retry_async, retry_sync +from openviking.utils.model_retry import effective_max_retries, retry_async, retry_sync from openviking_cli.utils import get_logger logger = get_logger(__name__) @@ -37,7 +37,7 @@ def embed( ) -> EmbedResult: return retry_sync( lambda: self.inner.embed(content, is_query=is_query, **kwargs), - max_retries=self.max_retries, + max_retries=effective_max_retries(self.max_retries), logger=logger, operation_name=f"{self.provider} embedding", model_errors_only=True, @@ -51,7 +51,7 @@ async def embed_async( ) -> EmbedResult: return await retry_async( lambda: self.inner.embed_async(content, is_query=is_query, **kwargs), - max_retries=self.max_retries, + max_retries=effective_max_retries(self.max_retries), logger=logger, operation_name=f"{self.provider} async embedding", model_errors_only=True, diff --git a/openviking/models/vlm/retry_vlm.py b/openviking/models/vlm/retry_vlm.py index 3c68fd8102..89a0598eac 100644 --- a/openviking/models/vlm/retry_vlm.py +++ b/openviking/models/vlm/retry_vlm.py @@ -5,7 +5,7 @@ from pathlib import Path from typing import Any, Dict, List, Optional, Union -from openviking.utils.model_retry import retry_async, retry_sync +from openviking.utils.model_retry import effective_max_retries, retry_async, retry_sync from openviking_cli.utils import get_logger from .base import VLMBase, VLMResponse @@ -59,7 +59,7 @@ def get_completion( messages=messages, **kwargs, ), - max_retries=self.max_retries, + max_retries=effective_max_retries(self.max_retries), logger=logger, operation_name=f"{self.provider} VLM completion", model_errors_only=True, @@ -85,7 +85,7 @@ async def get_completion_async( max_tokens=max_tokens, **kwargs, ), - max_retries=self.max_retries, + max_retries=effective_max_retries(self.max_retries), logger=logger, operation_name=f"{self.provider} VLM async completion", model_errors_only=True, @@ -111,7 +111,7 @@ def get_vision_completion( messages=messages, **kwargs, ), - max_retries=self.max_retries, + max_retries=effective_max_retries(self.max_retries), logger=logger, operation_name=f"{self.provider} VLM vision completion", model_errors_only=True, @@ -137,7 +137,7 @@ async def get_vision_completion_async( messages=messages, **kwargs, ), - max_retries=self.max_retries, + max_retries=effective_max_retries(self.max_retries), logger=logger, operation_name=f"{self.provider} VLM async vision completion", model_errors_only=True, diff --git a/openviking/service/search_service.py b/openviking/service/search_service.py index 840d72bc62..59e7e3d50e 100644 --- a/openviking/service/search_service.py +++ b/openviking/service/search_service.py @@ -17,6 +17,7 @@ is_http_url, is_viking_uri, ) +from openviking.utils.model_retry import model_retry_scope from openviking_cli.exceptions import InvalidArgumentError, NotInitializedError from openviking_cli.retrieve import ContextType from openviking_cli.utils import get_logger @@ -109,6 +110,7 @@ async def _resolve_image_url( "image_url must be a data:image base64 URI, http(s) URL, or viking:// URI." ) + @model_retry_scope(0) async def search( self, query: str, @@ -167,6 +169,7 @@ async def search( ) return result + @model_retry_scope(0) async def find( self, query: str, diff --git a/openviking/utils/model_retry.py b/openviking/utils/model_retry.py index cb34396c81..7dab79fbe6 100644 --- a/openviking/utils/model_retry.py +++ b/openviking/utils/model_retry.py @@ -6,7 +6,9 @@ import re import threading import time -from typing import Awaitable, Callable, TypeVar +from contextlib import asynccontextmanager +from contextvars import ContextVar +from typing import AsyncIterator, Awaitable, Callable, TypeVar from openviking.models.model_error import ModelErrorKind, model_error_from_exception from openviking.pyagfs.exceptions import AGFSNotADirectoryError @@ -16,6 +18,30 @@ T = TypeVar("T") +_max_retries_override: ContextVar[int | None] = ContextVar( + "model_max_retries_override", + default=None, +) + + +@asynccontextmanager +async def model_retry_scope(max_retries: int) -> AsyncIterator[None]: + """Temporarily override model retries for the current async context.""" + if max_retries < 0: + raise ValueError("max_retries must be non-negative") + token = _max_retries_override.set(max_retries) + try: + yield + finally: + _max_retries_override.reset(token) + + +def effective_max_retries(configured_max_retries: int) -> int: + """Return the request override when set, otherwise the configured value.""" + override = _max_retries_override.get() + return configured_max_retries if override is None else override + + # Error classification categories returned by classify_api_error() ERROR_CLASS_PERMANENT = "permanent" # request-level 4xx (e.g. 400 invalid parameter) ERROR_CLASS_AUTH = "auth" # credential-level 401/403 (key invalid / no permission / overdue) diff --git a/tests/service/test_search_service_image.py b/tests/service/test_search_service_image.py index 50bead176d..1f877b29d0 100644 --- a/tests/service/test_search_service_image.py +++ b/tests/service/test_search_service_image.py @@ -2,12 +2,14 @@ from openviking.server.identity import RequestContext, Role from openviking.service.search_service import SearchService +from openviking.utils.model_retry import effective_max_retries from openviking_cli.session.user_id import UserIdentifier class FakeVikingFS: def __init__(self): self.image_url = None + self.retry_values = [] async def read_file_bytes(self, uri, ctx=None): assert uri == "viking://resources/cat.png" @@ -15,6 +17,11 @@ async def read_file_bytes(self, uri, ctx=None): async def find(self, **kwargs): self.image_url = kwargs["image_url"] + self.retry_values.append(effective_max_retries(3)) + return {"ok": True} + + async def search(self, **kwargs): + self.retry_values.append(effective_max_retries(3)) return {"ok": True} @@ -32,3 +39,16 @@ async def test_search_service_resolves_viking_image_url_to_data_uri(): assert result == {"ok": True} assert fs.image_url.startswith("data:image/png;base64,") + + +@pytest.mark.asyncio +async def test_search_and_find_disable_model_retries_for_the_call(): + fs = FakeVikingFS() + service = SearchService(fs) + ctx = RequestContext(user=UserIdentifier("acc", "user"), role=Role.USER) + + await service.find(query="find", ctx=ctx) + await service.search(query="search", ctx=ctx) + + assert fs.retry_values == [0, 0] + assert effective_max_retries(3) == 3 diff --git a/tests/unit/test_model_retry.py b/tests/unit/test_model_retry.py index 1717c723cf..dd2160bc4f 100644 --- a/tests/unit/test_model_retry.py +++ b/tests/unit/test_model_retry.py @@ -2,6 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 """Tests for shared model retry helpers.""" +import asyncio + import pytest from openviking.models.model_error import ModelError, ModelErrorKind @@ -21,7 +23,9 @@ ERROR_CLASS_TRANSIENT, ERROR_CLASS_UNKNOWN, classify_api_error, + effective_max_retries, extract_metric_error_code, + model_retry_scope, retry_async, retry_sync, ) @@ -69,6 +73,36 @@ def test_classify_api_error_recognizes_request_burst_too_fast(): assert classify_api_error(RuntimeError("RequestBurstTooFast")) == ERROR_CLASS_TRANSIENT +@pytest.mark.asyncio +async def test_model_retry_scope_overrides_and_restores_configured_value(): + assert effective_max_retries(3) == 3 + + async with model_retry_scope(0): + assert effective_max_retries(3) == 0 + + assert effective_max_retries(3) == 3 + + +@pytest.mark.asyncio +async def test_model_retry_scope_is_isolated_between_tasks(): + scoped_started = asyncio.Event() + default_checked = asyncio.Event() + + async def scoped_call(): + async with model_retry_scope(0): + scoped_started.set() + await default_checked.wait() + return effective_max_retries(3) + + async def default_call(): + await scoped_started.wait() + retries = effective_max_retries(3) + default_checked.set() + return retries + + assert await asyncio.gather(scoped_call(), default_call()) == [0, 3] + + class _ProviderError(RuntimeError): def __init__(self, *, status_code=None, error_code=None, code=None, body=None): super().__init__("provider request failed") diff --git a/tests/unit/test_retry_embedder.py b/tests/unit/test_retry_embedder.py index 29a5009d57..b4c3c263b1 100644 --- a/tests/unit/test_retry_embedder.py +++ b/tests/unit/test_retry_embedder.py @@ -17,6 +17,7 @@ from openviking.models.embedder.retry_embedder import RetryEmbedder from openviking.models.model_error import ModelError, ModelErrorKind from openviking.models.vlm.token_usage import TokenUsageTracker +from openviking.utils.model_retry import model_retry_scope from openviking_cli.utils.config.embedding_config import ( EmbeddingConfig, EmbeddingCredential, @@ -168,6 +169,19 @@ async def observe_backoff(_delay): assert semaphore._value == 1 +@pytest.mark.asyncio +async def test_retry_embedder_honors_call_scope_without_mutating_instance(): + inner = FakeDenseEmbedder([ProviderError(503), [0.1, 0.2]], max_retries=2) + embedder = RetryEmbedder(inner) + + async with model_retry_scope(0): + with pytest.raises(ProviderError): + await embedder.embed_async("query") + + assert inner.calls == 1 + assert embedder.max_retries == 2 + + @pytest.mark.asyncio async def test_cancelled_attempt_releases_semaphore(): started = asyncio.Event() diff --git a/tests/unit/test_retry_vlm.py b/tests/unit/test_retry_vlm.py index 28265ff28e..e4b3eabe2d 100644 --- a/tests/unit/test_retry_vlm.py +++ b/tests/unit/test_retry_vlm.py @@ -9,6 +9,7 @@ from openviking.models.model_error import ModelError, ModelErrorKind from openviking.models.vlm.base import MultiCredentialVLM, VLMBase, VLMFactory from openviking.models.vlm.retry_vlm import RetryVLM +from openviking.utils.model_retry import model_retry_scope from openviking_cli.utils.config.vlm_config import VLMConfig @@ -147,6 +148,19 @@ def test_retry_vlm_uses_max_retries_as_additional_attempts(): assert len(inner.calls) == 3 +@pytest.mark.asyncio +async def test_retry_vlm_honors_call_scope_without_mutating_instance(): + inner = FakeVLM([ProviderError(503), "unexpected"]) + vlm = RetryVLM(inner) + + async with model_retry_scope(0): + with pytest.raises(ProviderError): + await vlm.get_completion_async(prompt="hello") + + assert len(inner.calls) == 1 + assert vlm.max_retries == 2 + + @pytest.mark.asyncio async def test_retry_vlm_does_not_retry_cancelled_request(): started = asyncio.Event() From dedfc3a098cb4f6bcf0ccb01db6a6c14b372890e Mon Sep 17 00:00:00 2001 From: "zhanghe.1216" <@bytedance.com> Date: Wed, 14 Oct 2026 16:33:55 +0800 Subject: [PATCH 08/16] =?UTF-8?q?test(api):=20=E4=BF=AE=E6=AD=A3=E8=B5=84?= =?UTF-8?q?=E6=BA=90=E4=B8=8E=E5=A4=8D=E5=88=B6=E7=94=A8=E4=BE=8B=E7=A8=B3?= =?UTF-8?q?=E5=AE=9A=E6=80=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/api_test/filesystem/test_fs_cp.py | 4 ---- .../resources_retrieval/test_build_error_handling.py | 3 ++- .../test_build_error_handling_slow.py | 3 ++- 3 files changed, 4 insertions(+), 6 deletions(-) diff --git a/tests/api_test/filesystem/test_fs_cp.py b/tests/api_test/filesystem/test_fs_cp.py index b99ec2a177..552f58f0da 100644 --- a/tests/api_test/filesystem/test_fs_cp.py +++ b/tests/api_test/filesystem/test_fs_cp.py @@ -153,10 +153,6 @@ def test_cp_rejects_file_directory_type_conflicts(self, api_client, source_is_di child_content = f"directory child remains {suffix}" try: assert api_client.fs_mkdir(directory_uri).status_code == 200 - # mkdir queues a parent refresh; finish it before preparing files so - # this test exercises type validation, not a transient refresh lock. - settled = api_client.system_wait(timeout=30) - assert settled.status_code == 200, settled.text for uri, content in ((file_uri, file_content), (child_uri, child_content)): write = api_client.fs_write(uri, content, mode="create", wait=True) assert write.status_code == 200, write.text diff --git a/tests/api_test/scenarios/resources_retrieval/test_build_error_handling.py b/tests/api_test/scenarios/resources_retrieval/test_build_error_handling.py index d33157ecab..6d2fc85d5e 100644 --- a/tests/api_test/scenarios/resources_retrieval/test_build_error_handling.py +++ b/tests/api_test/scenarios/resources_retrieval/test_build_error_handling.py @@ -141,10 +141,11 @@ def test_error_corrupted_zip(self, api_client): try: response = api_client.add_resource(path=zip_path, wait=True) - assert response.status_code == 500, f"损坏ZIP应返回 500, 实际: {response.status_code}" + assert response.status_code == 400, f"损坏ZIP应返回 400, 实际: {response.status_code}" data = response.json() assert data.get("status") == "error", f"损坏ZIP应返回 error, 实际: {data.get('status')}" + assert data.get("error", {}).get("code") == "INVALID_ARGUMENT" print("✓ TC-E16 损坏的ZIP文件处理通过") finally: diff --git a/tests/api_test/scenarios/resources_retrieval_slow/test_build_error_handling_slow.py b/tests/api_test/scenarios/resources_retrieval_slow/test_build_error_handling_slow.py index 8836fa5f13..1200437f04 100644 --- a/tests/api_test/scenarios/resources_retrieval_slow/test_build_error_handling_slow.py +++ b/tests/api_test/scenarios/resources_retrieval_slow/test_build_error_handling_slow.py @@ -379,10 +379,11 @@ def test_error_corrupted_zip(self, api_client): try: response = api_client.add_resource(path=zip_path, wait=True) - assert response.status_code == 500 + assert response.status_code == 400 data = response.json() if data.get("status") == "error": + assert data.get("error", {}).get("code") == "INVALID_ARGUMENT" error_msg = _extract_error_message(data).lower() assert ( "zip" in error_msg From 3bb74f94486fe95d3d5a910b921571ccdf1fff7f Mon Sep 17 00:00:00 2001 From: "zhanghe.1216" <@bytedance.com> Date: Wed, 14 Oct 2026 16:49:37 +0800 Subject: [PATCH 09/16] =?UTF-8?q?fix(queue):=20=E9=99=90=E5=88=B6=E6=A8=A1?= =?UTF-8?q?=E5=9E=8B=E9=94=99=E8=AF=AF=E9=87=8D=E6=96=B0=E5=85=A5=E9=98=9F?= =?UTF-8?q?=E8=8C=83=E5=9B=B4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- openviking/storage/collection_schemas.py | 16 ++++++- openviking/storage/queuefs/retry_policy.py | 16 +++++++ .../storage/queuefs/semantic_processor.py | 21 ++++++++-- tests/storage/test_collection_schemas.py | 42 ++++++++++++++++++- tests/storage/test_retry_policy.py | 32 ++++++++++++++ ...antic_processor_permanent_storage_error.py | 27 +++++++++++- 6 files changed, 145 insertions(+), 9 deletions(-) create mode 100644 tests/storage/test_retry_policy.py diff --git a/openviking/storage/collection_schemas.py b/openviking/storage/collection_schemas.py index a71a37eed2..a5d2f010a0 100644 --- a/openviking/storage/collection_schemas.py +++ b/openviking/storage/collection_schemas.py @@ -35,7 +35,10 @@ ) from openviking.storage.queuefs.named_queue import DequeueHandlerBase from openviking.storage.queuefs.process_result import ProcessResult -from openviking.storage.queuefs.retry_policy import MAX_FAILURE_REQUEUE_COUNT +from openviking.storage.queuefs.retry_policy import ( + MAX_FAILURE_REQUEUE_COUNT, + is_queue_retryable, +) from openviking.storage.vector_ids import vector_record_id from openviking.storage.viking_vector_index_backend import ( VIKINGDB_CONTENT_MAX_SIZE, @@ -821,7 +824,16 @@ async def on_dequeue(self, data: Optional[Dict[str, Any]]) -> ProcessResult: request_failed_message = error_msg return ProcessResult.failed(error_msg) - # Transient or unknown — re-enqueue for retry + if not is_queue_retryable(embed_err): + execute_status = "error" + self._log_embedding_error(logging.ERROR, error_msg, embedding_msg) + self._merge_request_stats( + embedding_msg.telemetry_id, + error_count=1, + ) + request_failed_message = error_msg + return ProcessResult.failed(error_msg) + self._log_embedding_error(logging.WARNING, error_msg, embedding_msg) execute_status = "requeued" if self._vikingdb.has_queue_manager: diff --git a/openviking/storage/queuefs/retry_policy.py b/openviking/storage/queuefs/retry_policy.py index 9214830913..bf07586772 100644 --- a/openviking/storage/queuefs/retry_policy.py +++ b/openviking/storage/queuefs/retry_policy.py @@ -2,4 +2,20 @@ # SPDX-License-Identifier: AGPL-3.0 """Shared limits for durable queue retry policies.""" +from openviking.models.model_error import ModelErrorKind, model_error_from_exception + MAX_FAILURE_REQUEUE_COUNT = 2 + +_QUEUE_RETRYABLE_MODEL_ERROR_KINDS = frozenset( + { + ModelErrorKind.TRANSIENT_SERVICE_ERROR, + ModelErrorKind.TRANSIENT_NETWORK_ERROR, + ModelErrorKind.RATE_LIMIT_ERROR, + } +) + + +def is_queue_retryable(error: BaseException) -> bool: + """Return whether a normalized model error may be retried by a durable queue.""" + model_error = model_error_from_exception(error) + return model_error is not None and model_error.kind in _QUEUE_RETRYABLE_MODEL_ERROR_KINDS diff --git a/openviking/storage/queuefs/semantic_processor.py b/openviking/storage/queuefs/semantic_processor.py index 58930f4b62..8106fac473 100644 --- a/openviking/storage/queuefs/semantic_processor.py +++ b/openviking/storage/queuefs/semantic_processor.py @@ -51,7 +51,10 @@ from openviking.storage.index_action import FieldPatch from openviking.storage.queuefs.named_queue import DequeueHandlerBase from openviking.storage.queuefs.process_result import ProcessResult -from openviking.storage.queuefs.retry_policy import MAX_FAILURE_REQUEUE_COUNT +from openviking.storage.queuefs.retry_policy import ( + MAX_FAILURE_REQUEUE_COUNT, + is_queue_retryable, +) from openviking.storage.queuefs.semantic_executor import SemanticTreeExecutor, SemanticTreeStats from openviking.storage.queuefs.semantic_lock import SemanticLockScope from openviking.storage.queuefs.semantic_msg import SemanticMsg, build_semantic_coalesce_key @@ -834,8 +837,7 @@ async def on_dequeue( if msg is not None: await self._cleanup_local_artifact(msg) return ProcessResult.failed(str(e)) - else: - # Transient or unknown — re-enqueue for retry + elif is_queue_retryable(e): execute_status = "requeued" logger.warning( f"Transient API error processing semantic message, re-enqueueing: {e}", @@ -853,6 +855,19 @@ async def on_dequeue( execute_status = "error" return result return ProcessResult.failed(str(e)) + else: + execute_status = "error" + logger.error( + f"Non-retryable error processing semantic message, dropping: {e}", + exc_info=True, + ) + if msg is not None: + self._merge_request_stats(msg.telemetry_id, error_count=1) + get_request_wait_tracker().mark_semantic_failed( + msg.telemetry_id, msg.id, str(e) + ) + await self._cleanup_local_artifact(msg) + return ProcessResult.failed(str(e)) finally: if msg is not None and execute_started_at is not None: tracker = get_request_wait_tracker() diff --git a/tests/storage/test_collection_schemas.py b/tests/storage/test_collection_schemas.py index 630fe96dc2..3c66e5b2c4 100644 --- a/tests/storage/test_collection_schemas.py +++ b/tests/storage/test_collection_schemas.py @@ -13,6 +13,7 @@ import requests from openviking.models.embedder.base import DenseEmbedderBase, EmbedResult +from openviking.models.model_error import ModelError, ModelErrorKind from openviking.server.identity import RequestContext, Role, UserIdentifier from openviking.service.resource_service import ResourceService from openviking.storage.collection_schemas import ( @@ -883,7 +884,10 @@ async def enqueue_embedding_msg(self, msg): class _TransientErrorEmbedder(_DummyEmbedder): async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: - raise RuntimeError("500 Internal Server Error") + raise ModelError( + ModelErrorKind.TRANSIENT_SERVICE_ERROR, + "500 Internal Server Error", + ) vikingdb = _QueueingVikingDB() monkeypatch.setattr( @@ -929,7 +933,10 @@ async def enqueue_embedding_msg(self, msg): class _TransientErrorEmbedder(_DummyEmbedder): async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: self.calls += 1 - raise RuntimeError("500 Internal Server Error") + raise ModelError( + ModelErrorKind.TRANSIENT_SERVICE_ERROR, + "500 Internal Server Error", + ) embedder = _TransientErrorEmbedder() vikingdb = _QueueingVikingDB() @@ -965,6 +972,37 @@ async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: assert "failure requeue limit reached" in results[-1].error +@pytest.mark.asyncio +async def test_embedding_unknown_model_error_fails_without_reenqueue(monkeypatch): + class _QueueingVikingDB: + is_closing = False + has_queue_manager = True + + def __init__(self): + self.enqueued = [] + + async def enqueue_embedding_msg(self, msg): + self.enqueued.append(msg) + + class _UnknownErrorEmbedder(_DummyEmbedder): + async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: + raise ModelError(ModelErrorKind.UNKNOWN, "unclassified provider error") + + vikingdb = _QueueingVikingDB() + monkeypatch.setattr( + "openviking_cli.utils.config.get_openviking_config", + lambda: _DummyConfig(_UnknownErrorEmbedder()), + ) + handler = TextEmbeddingHandler(vikingdb) + + result = await handler.on_dequeue(_build_queue_payload()) + + assert result.outcome is ProcessOutcome.FAILED + assert result.error is not None + assert "unclassified provider error" in result.error + assert vikingdb.enqueued == [] + + @pytest.mark.asyncio async def test_embedding_auth_error_fails_terminally_without_reenqueue(monkeypatch): """A credential (401/403) failure must fail terminally, not re-enqueue: an diff --git a/tests/storage/test_retry_policy.py b/tests/storage/test_retry_policy.py new file mode 100644 index 0000000000..5ae3a7cee4 --- /dev/null +++ b/tests/storage/test_retry_policy.py @@ -0,0 +1,32 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 + +import pytest + +from openviking.models.model_error import ModelError, ModelErrorKind +from openviking.storage.queuefs.retry_policy import is_queue_retryable + + +@pytest.mark.parametrize( + ("kind", "expected"), + [ + (ModelErrorKind.TRANSIENT_SERVICE_ERROR, True), + (ModelErrorKind.TRANSIENT_NETWORK_ERROR, True), + (ModelErrorKind.RATE_LIMIT_ERROR, True), + (ModelErrorKind.AUTHENTICATION_ERROR, False), + (ModelErrorKind.PERMISSION_DENIED, False), + (ModelErrorKind.QUOTA_EXHAUSTED, False), + (ModelErrorKind.MODEL_NOT_FOUND, False), + (ModelErrorKind.UNSUPPORTED, False), + (ModelErrorKind.INVALID_REQUEST, False), + (ModelErrorKind.INPUT_TOO_LARGE, False), + (ModelErrorKind.CONTENT_SAFETY, False), + (ModelErrorKind.UNKNOWN, False), + ], +) +def test_queue_retryable_model_error_kinds(kind, expected): + assert is_queue_retryable(ModelError(kind, "model request failed")) is expected + + +def test_unrecognized_error_is_not_queue_retryable(): + assert is_queue_retryable(RuntimeError("500 Internal Server Error")) is False diff --git a/tests/storage/test_semantic_processor_permanent_storage_error.py b/tests/storage/test_semantic_processor_permanent_storage_error.py index 999f0e12e4..9af2053c99 100644 --- a/tests/storage/test_semantic_processor_permanent_storage_error.py +++ b/tests/storage/test_semantic_processor_permanent_storage_error.py @@ -10,6 +10,7 @@ import pytest +from openviking.models.model_error import ModelError, ModelErrorKind from openviking.pyagfs.exceptions import AGFSNotADirectoryError from openviking.service.task_queue_middleware import TaskWorkQueueMiddleware from openviking.service.task_work_index import TaskWorkIndex @@ -219,7 +220,13 @@ async def test_malformed_resources_do_not_pause_following_valid_work(semantic_en @pytest.mark.asyncio @pytest.mark.parametrize( ("error", "trips_breaker"), - [(TimeoutError("API timeout"), True), (LockAcquisitionError("tree is locked"), False)], + [ + ( + ModelError(ModelErrorKind.TRANSIENT_NETWORK_ERROR, "API timeout"), + True, + ), + (LockAcquisitionError("tree is locked"), False), + ], ) async def test_retryable_errors_keep_existing_behavior( semantic_env, monkeypatch, error, trips_breaker @@ -254,7 +261,7 @@ async def test_retryable_errors_keep_existing_behavior( async def test_semantic_failure_requeue_limit_is_terminal(semantic_env): env = semantic_env msg = env.message(failure_requeue_count=MAX_FAILURE_REQUEUE_COUNT) - error = TimeoutError("API timeout") + error = ModelError(ModelErrorKind.TRANSIENT_SERVICE_ERROR, "API timeout") env.dag.run.side_effect = error result = await asyncio.wait_for(env.processor.on_dequeue(msg.to_dict()), timeout=1) @@ -272,6 +279,22 @@ async def test_semantic_failure_requeue_limit_is_terminal(semantic_env): } +@pytest.mark.asyncio +async def test_semantic_unknown_model_error_fails_without_requeue(semantic_env): + env = semantic_env + msg = env.message() + error = ModelError(ModelErrorKind.UNKNOWN, "unclassified provider error") + env.dag.run.side_effect = error + + result = await asyncio.wait_for(env.processor.on_dequeue(msg.to_dict()), timeout=1) + + assert result.outcome is ProcessOutcome.FAILED + assert result.error == str(error) + env.retry_queue.enqueue.assert_not_awaited() + await _assert_request_failed(env, msg, error) + env.processor._account_breaker(msg.account_id).check() + + @pytest.mark.asyncio async def test_semantic_failure_requeue_stops_after_three_attempts(semantic_env): env = semantic_env From 3614a8d095224a3874963b7f601cad7b10859323 Mon Sep 17 00:00:00 2001 From: "zhanghe.1216" <@bytedance.com> Date: Wed, 14 Oct 2026 17:36:12 +0800 Subject: [PATCH 10/16] =?UTF-8?q?fix(queue):=20=E5=88=A0=E9=99=A4=E6=A8=A1?= =?UTF-8?q?=E5=9E=8B=E9=94=99=E8=AF=AF=E9=87=8D=E6=96=B0=E5=85=A5=E9=98=9F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- openviking/storage/collection_schemas.py | 53 +------- openviking/storage/queuefs/retry_policy.py | 21 --- .../storage/queuefs/semantic_processor.py | 38 +----- tests/storage/test_collection_schemas.py | 123 ++---------------- tests/storage/test_memory_semantic_stall.py | 26 ++-- tests/storage/test_retry_policy.py | 32 ----- ...antic_processor_permanent_storage_error.py | 90 +++---------- 7 files changed, 37 insertions(+), 346 deletions(-) delete mode 100644 openviking/storage/queuefs/retry_policy.py delete mode 100644 tests/storage/test_retry_policy.py diff --git a/openviking/storage/collection_schemas.py b/openviking/storage/collection_schemas.py index a5d2f010a0..25baed1552 100644 --- a/openviking/storage/collection_schemas.py +++ b/openviking/storage/collection_schemas.py @@ -35,10 +35,6 @@ ) from openviking.storage.queuefs.named_queue import DequeueHandlerBase from openviking.storage.queuefs.process_result import ProcessResult -from openviking.storage.queuefs.retry_policy import ( - MAX_FAILURE_REQUEUE_COUNT, - is_queue_retryable, -) from openviking.storage.vector_ids import vector_record_id from openviking.storage.viking_vector_index_backend import ( VIKINGDB_CONTENT_MAX_SIZE, @@ -824,56 +820,9 @@ async def on_dequeue(self, data: Optional[Dict[str, Any]]) -> ProcessResult: request_failed_message = error_msg return ProcessResult.failed(error_msg) - if not is_queue_retryable(embed_err): - execute_status = "error" - self._log_embedding_error(logging.ERROR, error_msg, embedding_msg) - self._merge_request_stats( - embedding_msg.telemetry_id, - error_count=1, - ) - request_failed_message = error_msg - return ProcessResult.failed(error_msg) - - self._log_embedding_error(logging.WARNING, error_msg, embedding_msg) - execute_status = "requeued" - if self._vikingdb.has_queue_manager: - if embedding_msg.failure_requeue_count >= MAX_FAILURE_REQUEUE_COUNT: - terminal_error = ( - f"{error_msg} (failure requeue limit reached: " - f"{MAX_FAILURE_REQUEUE_COUNT})" - ) - self._log_embedding_error( - logging.ERROR, terminal_error, embedding_msg - ) - self._merge_request_stats(embedding_msg.telemetry_id, error_count=1) - execute_status = "error" - request_failed_message = terminal_error - return ProcessResult.failed(terminal_error) - embedding_msg.failure_requeue_count += 1 - try: - await self._reenqueue_embedding_msg(embedding_msg) - self._merge_request_stats( - embedding_msg.telemetry_id, - requeue_count=1, - ) - get_request_wait_tracker().record_embedding_requeue( - embedding_msg.telemetry_id - ) - logger.info( - "Re-enqueued embedding message after transient error: %s", - self._embedding_delivery_log_context(embedding_msg), - ) - return ProcessResult.requeued() - except Exception as requeue_err: - logger.error( - self._embedding_error_msg( - embedding_msg, - f"Failed to re-enqueue message: {requeue_err}", - ) - ) - self._merge_request_stats(embedding_msg.telemetry_id, error_count=1) execute_status = "error" + self._log_embedding_error(logging.ERROR, error_msg, embedding_msg) request_failed_message = error_msg return ProcessResult.failed(error_msg) diff --git a/openviking/storage/queuefs/retry_policy.py b/openviking/storage/queuefs/retry_policy.py deleted file mode 100644 index bf07586772..0000000000 --- a/openviking/storage/queuefs/retry_policy.py +++ /dev/null @@ -1,21 +0,0 @@ -# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. -# SPDX-License-Identifier: AGPL-3.0 -"""Shared limits for durable queue retry policies.""" - -from openviking.models.model_error import ModelErrorKind, model_error_from_exception - -MAX_FAILURE_REQUEUE_COUNT = 2 - -_QUEUE_RETRYABLE_MODEL_ERROR_KINDS = frozenset( - { - ModelErrorKind.TRANSIENT_SERVICE_ERROR, - ModelErrorKind.TRANSIENT_NETWORK_ERROR, - ModelErrorKind.RATE_LIMIT_ERROR, - } -) - - -def is_queue_retryable(error: BaseException) -> bool: - """Return whether a normalized model error may be retried by a durable queue.""" - model_error = model_error_from_exception(error) - return model_error is not None and model_error.kind in _QUEUE_RETRYABLE_MODEL_ERROR_KINDS diff --git a/openviking/storage/queuefs/semantic_processor.py b/openviking/storage/queuefs/semantic_processor.py index 8106fac473..8352f40bf1 100644 --- a/openviking/storage/queuefs/semantic_processor.py +++ b/openviking/storage/queuefs/semantic_processor.py @@ -51,10 +51,6 @@ from openviking.storage.index_action import FieldPatch from openviking.storage.queuefs.named_queue import DequeueHandlerBase from openviking.storage.queuefs.process_result import ProcessResult -from openviking.storage.queuefs.retry_policy import ( - MAX_FAILURE_REQUEUE_COUNT, - is_queue_retryable, -) from openviking.storage.queuefs.semantic_executor import SemanticTreeExecutor, SemanticTreeStats from openviking.storage.queuefs.semantic_lock import SemanticLockScope from openviking.storage.queuefs.semantic_msg import SemanticMsg, build_semantic_coalesce_key @@ -330,20 +326,7 @@ async def _requeue_semantic_msg_after_error( error: Exception, *, work: SemanticMessageWork, - consume_failure_budget: bool = True, ) -> ProcessResult: - if consume_failure_budget and msg.failure_requeue_count >= MAX_FAILURE_REQUEUE_COUNT: - terminal_error = f"{error} (failure requeue limit reached: {MAX_FAILURE_REQUEUE_COUNT})" - logger.error("Semantic message reached failure requeue limit: %s", msg.uri) - self._merge_request_stats(msg.telemetry_id, error_count=1) - get_request_wait_tracker().mark_semantic_failed( - msg.telemetry_id, msg.id, terminal_error - ) - await self._cleanup_local_artifact(msg) - return ProcessResult.failed(terminal_error) - - if consume_failure_budget: - msg.failure_requeue_count += 1 try: await work.reenqueue_after_error() self._merge_request_stats(msg.telemetry_id, requeue_count=1) @@ -796,7 +779,6 @@ async def on_dequeue( msg, e, work=work, - consume_failure_budget=False, ) if result.error is not None: execute_status = "error" @@ -837,28 +819,10 @@ async def on_dequeue( if msg is not None: await self._cleanup_local_artifact(msg) return ProcessResult.failed(str(e)) - elif is_queue_retryable(e): - execute_status = "requeued" - logger.warning( - f"Transient API error processing semantic message, re-enqueueing: {e}", - exc_info=True, - ) - if breaker is not None: - breaker.record_failure(e) - if msg is not None and work is not None: - result = await self._requeue_semantic_msg_after_error( - msg, - e, - work=work, - ) - if result.error is not None: - execute_status = "error" - return result - return ProcessResult.failed(str(e)) else: execute_status = "error" logger.error( - f"Non-retryable error processing semantic message, dropping: {e}", + f"Error processing semantic message, not re-enqueueing: {e}", exc_info=True, ) if msg is not None: diff --git a/tests/storage/test_collection_schemas.py b/tests/storage/test_collection_schemas.py index 3c66e5b2c4..9e11623dec 100644 --- a/tests/storage/test_collection_schemas.py +++ b/tests/storage/test_collection_schemas.py @@ -33,7 +33,6 @@ from openviking.storage.index_action import IndexAction from openviking.storage.queuefs.embedding_msg import EmbeddingMsg from openviking.storage.queuefs.process_result import ProcessOutcome -from openviking.storage.queuefs.retry_policy import MAX_FAILURE_REQUEUE_COUNT from openviking.storage.vector_ids import vector_record_id from openviking.storage.vectordb import engine as vectordb_engine from openviking.storage.vectordb.collection.collection import Collection @@ -863,117 +862,15 @@ async def enqueue_embedding_msg(self, msg): @pytest.mark.asyncio @pytest.mark.parametrize( - ("failure_requeue_count", "expected_outcome"), + "kind", [ - (0, ProcessOutcome.REQUEUED), - (MAX_FAILURE_REQUEUE_COUNT, ProcessOutcome.FAILED), + ModelErrorKind.TRANSIENT_SERVICE_ERROR, + ModelErrorKind.TRANSIENT_NETWORK_ERROR, + ModelErrorKind.RATE_LIMIT_ERROR, + ModelErrorKind.UNKNOWN, ], ) -async def test_embedding_failure_requeue_limit( - monkeypatch, failure_requeue_count, expected_outcome -): - class _QueueingVikingDB: - is_closing = False - has_queue_manager = True - - def __init__(self): - self.enqueued = [] - - async def enqueue_embedding_msg(self, msg): - self.enqueued.append(msg) - - class _TransientErrorEmbedder(_DummyEmbedder): - async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: - raise ModelError( - ModelErrorKind.TRANSIENT_SERVICE_ERROR, - "500 Internal Server Error", - ) - - vikingdb = _QueueingVikingDB() - monkeypatch.setattr( - "openviking_cli.utils.config.get_openviking_config", - lambda: _DummyConfig(_TransientErrorEmbedder()), - ) - handler = TextEmbeddingHandler(vikingdb) - msg = EmbeddingMsg( - message="hello", - context_data={ - "id": "id-1", - "uri": "viking://resources/sample", - "account_id": "default", - "abstract": "sample", - }, - failure_requeue_count=failure_requeue_count, - ) - - result = await handler.on_dequeue({"data": msg.to_json()}) - - assert result.outcome is expected_outcome - if expected_outcome is ProcessOutcome.REQUEUED: - assert len(vikingdb.enqueued) == 1 - assert vikingdb.enqueued[0].failure_requeue_count == 1 - else: - assert vikingdb.enqueued == [] - assert result.error is not None - assert "failure requeue limit reached" in result.error - - -@pytest.mark.asyncio -async def test_embedding_failure_requeue_stops_after_three_attempts(monkeypatch): - class _QueueingVikingDB: - is_closing = False - has_queue_manager = True - - def __init__(self): - self.enqueued = [] - - async def enqueue_embedding_msg(self, msg): - self.enqueued.append(msg) - - class _TransientErrorEmbedder(_DummyEmbedder): - async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: - self.calls += 1 - raise ModelError( - ModelErrorKind.TRANSIENT_SERVICE_ERROR, - "500 Internal Server Error", - ) - - embedder = _TransientErrorEmbedder() - vikingdb = _QueueingVikingDB() - monkeypatch.setattr( - "openviking_cli.utils.config.get_openviking_config", - lambda: _DummyConfig(embedder), - ) - handler = TextEmbeddingHandler(vikingdb) - msg = EmbeddingMsg( - message="hello", - context_data={ - "id": "id-1", - "uri": "viking://resources/sample", - "account_id": "default", - "abstract": "sample", - }, - ) - - results = [] - for _ in range(3): - results.append(await handler.on_dequeue({"data": msg.to_json()})) - if results[-1].outcome is ProcessOutcome.REQUEUED: - msg = vikingdb.enqueued[-1] - - assert [result.outcome for result in results] == [ - ProcessOutcome.REQUEUED, - ProcessOutcome.REQUEUED, - ProcessOutcome.FAILED, - ] - assert [msg.failure_requeue_count for msg in vikingdb.enqueued] == [1, 2] - assert embedder.calls == 3 - assert results[-1].error is not None - assert "failure requeue limit reached" in results[-1].error - - -@pytest.mark.asyncio -async def test_embedding_unknown_model_error_fails_without_reenqueue(monkeypatch): +async def test_embedding_model_error_fails_without_reenqueue(monkeypatch, kind): class _QueueingVikingDB: is_closing = False has_queue_manager = True @@ -984,14 +881,14 @@ def __init__(self): async def enqueue_embedding_msg(self, msg): self.enqueued.append(msg) - class _UnknownErrorEmbedder(_DummyEmbedder): + class _ModelErrorEmbedder(_DummyEmbedder): async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: - raise ModelError(ModelErrorKind.UNKNOWN, "unclassified provider error") + raise ModelError(kind, "model request failed") vikingdb = _QueueingVikingDB() monkeypatch.setattr( "openviking_cli.utils.config.get_openviking_config", - lambda: _DummyConfig(_UnknownErrorEmbedder()), + lambda: _DummyConfig(_ModelErrorEmbedder()), ) handler = TextEmbeddingHandler(vikingdb) @@ -999,7 +896,7 @@ async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: assert result.outcome is ProcessOutcome.FAILED assert result.error is not None - assert "unclassified provider error" in result.error + assert "model request failed" in result.error assert vikingdb.enqueued == [] diff --git a/tests/storage/test_memory_semantic_stall.py b/tests/storage/test_memory_semantic_stall.py index 99979cf562..1e682013ba 100644 --- a/tests/storage/test_memory_semantic_stall.py +++ b/tests/storage/test_memory_semantic_stall.py @@ -54,9 +54,6 @@ def _build_data(msg: SemanticMsg) -> dict: @pytest.mark.parametrize("context_type", ["resource", "memory", "skill"]) async def test_retry_cancellation_keeps_skill_wait_isolated(monkeypatch, context_type): processor = _processor() - processor._circuit_breaker = SimpleNamespace( - check=MagicMock(side_effect=CircuitBreakerOpen), retry_after=0 - ) entered, release = asyncio.Event(), asyncio.Event() written = [] @@ -71,6 +68,9 @@ async def enqueue(msg): lambda: SimpleNamespace(SEMANTIC="Semantic", get_queue=lambda _: queue), ) msg = _make_msg(context_type=context_type, telemetry_id=str(uuid4())) + processor._circuit_breakers[msg.account_id] = SimpleNamespace( + check=MagicMock(side_effect=CircuitBreakerOpen), retry_after=0 + ) worker = asyncio.create_task(processor.on_dequeue(msg.to_dict())) try: await asyncio.wait_for(entered.wait(), 1) @@ -167,14 +167,8 @@ async def test_memory_ls_error_returns_failed(): @pytest.mark.asyncio -async def test_memory_ls_transient_error_requeues(): - """Transient errors during ls() re-enqueue the msg and increment requeue count. - - A 500-class error wrapped by the processor's `raise RuntimeError(...) from e` - is classified as `transient`. The outer on_dequeue() path must call - _reenqueue_semantic_msg(), bump requeue_count, and return REQUEUED without - a terminal error. - """ +async def test_memory_ls_transient_error_returns_failed(): + """Queue processing does not retry an error after the owning layer exhausts it.""" processor = _processor() fake_fs = MagicMock() @@ -184,8 +178,6 @@ async def test_memory_ls_transient_error_requeues(): msg = _make_msg(telemetry_id="tel-1") data = _build_data(msg) - reenqueue_mock = AsyncMock() - with ( patch( "openviking.storage.queuefs.semantic_processor.get_viking_fs", @@ -195,14 +187,14 @@ async def test_memory_ls_transient_error_requeues(): "openviking.storage.queuefs.semantic_work.resolve_telemetry", return_value=None, ), - patch.object(processor, "_reenqueue_semantic_msg", new=reenqueue_mock), ): result = await processor.on_dequeue(data) - assert result.outcome is ProcessOutcome.REQUEUED + assert result.outcome is ProcessOutcome.FAILED assert result.value is None - assert result.error is None - reenqueue_mock.assert_awaited_once() + assert result.error == ( + f"Failed to list memory directory {msg.uri}: 500 Internal Server Error" + ) @pytest.mark.asyncio diff --git a/tests/storage/test_retry_policy.py b/tests/storage/test_retry_policy.py deleted file mode 100644 index 5ae3a7cee4..0000000000 --- a/tests/storage/test_retry_policy.py +++ /dev/null @@ -1,32 +0,0 @@ -# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. -# SPDX-License-Identifier: AGPL-3.0 - -import pytest - -from openviking.models.model_error import ModelError, ModelErrorKind -from openviking.storage.queuefs.retry_policy import is_queue_retryable - - -@pytest.mark.parametrize( - ("kind", "expected"), - [ - (ModelErrorKind.TRANSIENT_SERVICE_ERROR, True), - (ModelErrorKind.TRANSIENT_NETWORK_ERROR, True), - (ModelErrorKind.RATE_LIMIT_ERROR, True), - (ModelErrorKind.AUTHENTICATION_ERROR, False), - (ModelErrorKind.PERMISSION_DENIED, False), - (ModelErrorKind.QUOTA_EXHAUSTED, False), - (ModelErrorKind.MODEL_NOT_FOUND, False), - (ModelErrorKind.UNSUPPORTED, False), - (ModelErrorKind.INVALID_REQUEST, False), - (ModelErrorKind.INPUT_TOO_LARGE, False), - (ModelErrorKind.CONTENT_SAFETY, False), - (ModelErrorKind.UNKNOWN, False), - ], -) -def test_queue_retryable_model_error_kinds(kind, expected): - assert is_queue_retryable(ModelError(kind, "model request failed")) is expected - - -def test_unrecognized_error_is_not_queue_retryable(): - assert is_queue_retryable(RuntimeError("500 Internal Server Error")) is False diff --git a/tests/storage/test_semantic_processor_permanent_storage_error.py b/tests/storage/test_semantic_processor_permanent_storage_error.py index 9af2053c99..5cdbafeac9 100644 --- a/tests/storage/test_semantic_processor_permanent_storage_error.py +++ b/tests/storage/test_semantic_processor_permanent_storage_error.py @@ -17,7 +17,6 @@ from openviking.storage.errors import LockAcquisitionError from openviking.storage.queuefs.named_queue import NamedQueue from openviking.storage.queuefs.process_result import ProcessOutcome -from openviking.storage.queuefs.retry_policy import MAX_FAILURE_REQUEUE_COUNT from openviking.storage.queuefs.semantic_executor import SemanticTreeStats from openviking.storage.queuefs.semantic_msg import SemanticMsg from openviking.storage.queuefs.semantic_processor import SemanticProcessor @@ -218,26 +217,14 @@ async def test_malformed_resources_do_not_pause_following_valid_work(semantic_en @pytest.mark.asyncio -@pytest.mark.parametrize( - ("error", "trips_breaker"), - [ - ( - ModelError(ModelErrorKind.TRANSIENT_NETWORK_ERROR, "API timeout"), - True, - ), - (LockAcquisitionError("tree is locked"), False), - ], -) -async def test_retryable_errors_keep_existing_behavior( - semantic_env, monkeypatch, error, trips_breaker -): +async def test_lock_error_requeues_without_consuming_failure_budget(semantic_env, monkeypatch): env = semantic_env sleep = AsyncMock() monkeypatch.setattr("openviking.storage.queuefs.semantic_processor.asyncio.sleep", sleep) msg = env.message() breaker = CircuitBreaker(failure_threshold=1) env.processor._circuit_breakers[msg.account_id] = breaker - env.dag.run.side_effect = error + env.dag.run.side_effect = LockAcquisitionError("tree is locked") result = await asyncio.wait_for(env.processor.on_dequeue(msg.to_dict()), timeout=1) @@ -245,45 +232,27 @@ async def test_retryable_errors_keep_existing_behavior( env.retry_queue.enqueue.assert_awaited_once() requeued = env.retry_queue.enqueue.call_args.args[0] assert requeued.id == msg.id - assert requeued.failure_requeue_count == (0 if isinstance(error, LockAcquisitionError) else 1) + assert requeued.failure_requeue_count == 0 assert not env.tracker.is_complete(msg.telemetry_id) assert env.tracker.build_queue_status(msg.telemetry_id)["Semantic"]["requeue_count"] == 1 - if trips_breaker: - with pytest.raises(CircuitBreakerOpen): - breaker.check() - sleep.assert_awaited_once() - else: - breaker.check() - sleep.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_semantic_failure_requeue_limit_is_terminal(semantic_env): - env = semantic_env - msg = env.message(failure_requeue_count=MAX_FAILURE_REQUEUE_COUNT) - error = ModelError(ModelErrorKind.TRANSIENT_SERVICE_ERROR, "API timeout") - env.dag.run.side_effect = error - - result = await asyncio.wait_for(env.processor.on_dequeue(msg.to_dict()), timeout=1) - - terminal_error = f"{error} (failure requeue limit reached: {MAX_FAILURE_REQUEUE_COUNT})" - assert result.outcome is ProcessOutcome.FAILED - assert result.error == terminal_error - env.retry_queue.enqueue.assert_not_awaited() - await asyncio.wait_for(env.tracker.wait_for_request(msg.telemetry_id), timeout=1) - assert env.tracker.build_queue_status(msg.telemetry_id)["Semantic"] == { - "processed": 0, - "requeue_count": 0, - "error_count": 1, - "errors": [{"message": terminal_error}], - } + breaker.check() + sleep.assert_not_awaited() @pytest.mark.asyncio -async def test_semantic_unknown_model_error_fails_without_requeue(semantic_env): +@pytest.mark.parametrize( + "kind", + [ + ModelErrorKind.TRANSIENT_SERVICE_ERROR, + ModelErrorKind.TRANSIENT_NETWORK_ERROR, + ModelErrorKind.RATE_LIMIT_ERROR, + ModelErrorKind.UNKNOWN, + ], +) +async def test_semantic_model_error_fails_without_requeue(semantic_env, kind): env = semantic_env msg = env.message() - error = ModelError(ModelErrorKind.UNKNOWN, "unclassified provider error") + error = ModelError(kind, "model request failed") env.dag.run.side_effect = error result = await asyncio.wait_for(env.processor.on_dequeue(msg.to_dict()), timeout=1) @@ -295,33 +264,6 @@ async def test_semantic_unknown_model_error_fails_without_requeue(semantic_env): env.processor._account_breaker(msg.account_id).check() -@pytest.mark.asyncio -async def test_semantic_failure_requeue_stops_after_three_attempts(semantic_env): - env = semantic_env - msg = env.message() - work = SimpleNamespace(reenqueue_after_error=AsyncMock()) - error = TimeoutError("API timeout") - - results = [] - for _ in range(3): - msg = SemanticMsg.from_dict(msg.to_dict()) - results.append( - await env.processor._requeue_semantic_msg_after_error( - msg, - error, - work=work, - ) - ) - - assert [result.outcome for result in results] == [ - ProcessOutcome.REQUEUED, - ProcessOutcome.REQUEUED, - ProcessOutcome.FAILED, - ] - assert msg.failure_requeue_count == MAX_FAILURE_REQUEUE_COUNT - assert work.reenqueue_after_error.await_count == MAX_FAILURE_REQUEUE_COUNT - - @pytest.mark.asyncio async def test_permanent_api_error_still_opens_breaker(semantic_env): env = semantic_env From bc3e5510a57e3c8c582af2dc34e9a3c2310a4f1f Mon Sep 17 00:00:00 2001 From: qin-ctx Date: Wed, 14 Oct 2026 19:11:39 +0800 Subject: [PATCH 11/16] =?UTF-8?q?test(models):=20=E5=B0=86=E9=94=99?= =?UTF-8?q?=E8=AF=AF=E4=B8=8E=E9=87=8D=E8=AF=95=E8=A6=86=E7=9B=96=E6=94=B6?= =?UTF-8?q?=E6=95=9B=E5=88=B0=E7=8E=B0=E6=9C=89=E5=A5=91=E7=BA=A6=E6=B5=8B?= =?UTF-8?q?=E8=AF=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/service/test_search_service_image.py | 29 +- tests/storage/test_collection_schemas.py | 13 +- ...antic_processor_permanent_storage_error.py | 13 +- tests/unit/test_gemini_embedder.py | 30 +- tests/unit/test_jina_embedder.py | 34 +- tests/unit/test_model_error.py | 294 ------------------ tests/unit/test_model_retry.py | 66 ++-- tests/unit/test_retry_embedder.py | 37 +-- tests/unit/test_retry_vlm.py | 26 +- 9 files changed, 91 insertions(+), 451 deletions(-) delete mode 100644 tests/unit/test_model_error.py diff --git a/tests/service/test_search_service_image.py b/tests/service/test_search_service_image.py index 1f877b29d0..be3ff50c72 100644 --- a/tests/service/test_search_service_image.py +++ b/tests/service/test_search_service_image.py @@ -1,3 +1,5 @@ +import asyncio + import pytest from openviking.server.identity import RequestContext, Role @@ -9,7 +11,7 @@ class FakeVikingFS: def __init__(self): self.image_url = None - self.retry_values = [] + self.retry_values = {} async def read_file_bytes(self, uri, ctx=None): assert uri == "viking://resources/cat.png" @@ -17,38 +19,31 @@ async def read_file_bytes(self, uri, ctx=None): async def find(self, **kwargs): self.image_url = kwargs["image_url"] - self.retry_values.append(effective_max_retries(3)) + await asyncio.sleep(0) # Keep the search scope active during the concurrent call. + self.retry_values["find"] = effective_max_retries(3) return {"ok": True} async def search(self, **kwargs): - self.retry_values.append(effective_max_retries(3)) + self.retry_values["search"] = effective_max_retries(3) return {"ok": True} @pytest.mark.asyncio -async def test_search_service_resolves_viking_image_url_to_data_uri(): +async def test_image_search_resolves_uri_and_isolates_retry_scope(): fs = FakeVikingFS() service = SearchService(fs) ctx = RequestContext(user=UserIdentifier("acc", "user"), role=Role.USER) - result = await service.find( - query="", - image_url="viking://resources/cat.png", - ctx=ctx, + result, _ = await asyncio.gather( + service.find(query="", image_url="viking://resources/cat.png", ctx=ctx), + fs.search(query="concurrent background work"), ) assert result == {"ok": True} assert fs.image_url.startswith("data:image/png;base64,") + assert fs.retry_values == {"find": 0, "search": 3} - -@pytest.mark.asyncio -async def test_search_and_find_disable_model_retries_for_the_call(): - fs = FakeVikingFS() - service = SearchService(fs) - ctx = RequestContext(user=UserIdentifier("acc", "user"), role=Role.USER) - - await service.find(query="find", ctx=ctx) await service.search(query="search", ctx=ctx) - assert fs.retry_values == [0, 0] + assert fs.retry_values == {"find": 0, "search": 0} assert effective_max_retries(3) == 3 diff --git a/tests/storage/test_collection_schemas.py b/tests/storage/test_collection_schemas.py index 9e11623dec..7fe27c7491 100644 --- a/tests/storage/test_collection_schemas.py +++ b/tests/storage/test_collection_schemas.py @@ -861,16 +861,7 @@ async def enqueue_embedding_msg(self, msg): @pytest.mark.asyncio -@pytest.mark.parametrize( - "kind", - [ - ModelErrorKind.TRANSIENT_SERVICE_ERROR, - ModelErrorKind.TRANSIENT_NETWORK_ERROR, - ModelErrorKind.RATE_LIMIT_ERROR, - ModelErrorKind.UNKNOWN, - ], -) -async def test_embedding_model_error_fails_without_reenqueue(monkeypatch, kind): +async def test_embedding_model_error_fails_without_reenqueue(monkeypatch): class _QueueingVikingDB: is_closing = False has_queue_manager = True @@ -883,7 +874,7 @@ async def enqueue_embedding_msg(self, msg): class _ModelErrorEmbedder(_DummyEmbedder): async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: - raise ModelError(kind, "model request failed") + raise ModelError(ModelErrorKind.TRANSIENT_SERVICE_ERROR, "model request failed") vikingdb = _QueueingVikingDB() monkeypatch.setattr( diff --git a/tests/storage/test_semantic_processor_permanent_storage_error.py b/tests/storage/test_semantic_processor_permanent_storage_error.py index 5cdbafeac9..9582e724b2 100644 --- a/tests/storage/test_semantic_processor_permanent_storage_error.py +++ b/tests/storage/test_semantic_processor_permanent_storage_error.py @@ -240,19 +240,10 @@ async def test_lock_error_requeues_without_consuming_failure_budget(semantic_env @pytest.mark.asyncio -@pytest.mark.parametrize( - "kind", - [ - ModelErrorKind.TRANSIENT_SERVICE_ERROR, - ModelErrorKind.TRANSIENT_NETWORK_ERROR, - ModelErrorKind.RATE_LIMIT_ERROR, - ModelErrorKind.UNKNOWN, - ], -) -async def test_semantic_model_error_fails_without_requeue(semantic_env, kind): +async def test_semantic_model_error_fails_without_requeue(semantic_env): env = semantic_env msg = env.message() - error = ModelError(kind, "model request failed") + error = ModelError(ModelErrorKind.TRANSIENT_SERVICE_ERROR, "model request failed") env.dag.run.side_effect = error result = await asyncio.wait_for(env.processor.on_dequeue(msg.to_dict()), timeout=1) diff --git a/tests/unit/test_gemini_embedder.py b/tests/unit/test_gemini_embedder.py index c4b74f4749..e5a5154fe4 100644 --- a/tests/unit/test_gemini_embedder.py +++ b/tests/unit/test_gemini_embedder.py @@ -9,7 +9,7 @@ import pytest -from openviking.models.model_error import ModelError +from openviking.models.model_error import ModelError, ModelErrorKind def _make_mock_embedding(values): @@ -226,29 +226,33 @@ def test_dimension_too_high_raises(self, mock_client): class TestGeminiErrorMessages: @pytest.mark.parametrize( - "code,match", + "reason,expected_kind", [ - (401, "Invalid API key"), - (403, "Permission denied"), - (404, "Model not found"), - (429, "Resource exhausted"), - (500, "service error"), + ("QUOTA_EXCEEDED", ModelErrorKind.QUOTA_EXHAUSTED), + ("RATE_LIMIT_EXCEEDED", ModelErrorKind.RATE_LIMIT_ERROR), + (None, ModelErrorKind.UNKNOWN), ], ) @patch("openviking.models.embedder.gemini_embedders.genai.Client") - def test_error_code_hint(self, mock_client, code, match): - from google.genai.errors import APIError + def test_resource_exhausted_uses_structured_reason(self, mock_client, reason, expected_kind): + from google.genai.errors import ClientError from openviking.models.embedder.gemini_embedders import GeminiDenseEmbedder mock = mock_client.return_value - mock_resp = MagicMock() - mock_resp.status_code = code - mock.models.embed_content.side_effect = APIError(code, {}, response=mock_resp) + error = ClientError( + 429, + {"error": {"details": [] if reason is None else [{"reason": reason}]}}, + ) + mock.models.embed_content.side_effect = error embedder = GeminiDenseEmbedder("gemini-embedding-2-preview", api_key="k") - with pytest.raises(ModelError, match=match): + with pytest.raises(ModelError, match="Resource exhausted") as raised: embedder.embed("hello") + assert raised.value.kind is expected_kind + assert raised.value.status_code == 429 + assert raised.value.__cause__.__cause__ is error + @patch("openviking.models.embedder.gemini_embedders.genai.Client") def test_error_message_includes_http_code(self, mock_client): from google.genai.errors import APIError diff --git a/tests/unit/test_jina_embedder.py b/tests/unit/test_jina_embedder.py index 7cc36689f3..c7ebe64d75 100644 --- a/tests/unit/test_jina_embedder.py +++ b/tests/unit/test_jina_embedder.py @@ -10,7 +10,7 @@ from openviking.models.embedder.jina_embedders import ( JINA_MODEL_DIMENSIONS, ) -from openviking.models.model_error import ModelError +from openviking.models.model_error import ModelError, ModelErrorKind from openviking_cli.utils.config.embedding_config import EmbeddingModelConfig @@ -230,20 +230,25 @@ def test_embed_api_error(self, mock_openai_class): mock_client = MagicMock() mock_openai_class.return_value = mock_client - mock_client.embeddings.create.side_effect = openai.APIError( + error = openai.APIError( message="Test API error", request=MagicMock(), body=None, ) + mock_client.embeddings.create.side_effect = error embedder = JinaDenseEmbedder( model_name="jina-embeddings-v5-text-small", api_key="test-api-key", ) - with pytest.raises(ModelError, match="Jina API error"): + with pytest.raises(ModelError, match="Jina API error") as raised: embedder.embed("Hello world") + assert raised.value.kind is ModelErrorKind.UNKNOWN + assert raised.value.provider == "jina" + assert raised.value.__cause__.__cause__ is error + def test_build_extra_body_none(self): """Test _build_extra_body returns None when no params set""" embedder = JinaDenseEmbedder( @@ -312,29 +317,6 @@ def test_422_task_error_actionable_message(self, mock_openai_class): with pytest.raises(ModelError, match="query_param.*document_param"): embedder.embed("hello") - @patch("openviking.models.embedder.jina_embedders.openai.OpenAI") - def test_non_422_error_passthrough(self, mock_openai_class): - """Non-422 API errors should use the generic 'Jina API error' message.""" - import openai - - mock_client = MagicMock() - mock_openai_class.return_value = mock_client - - error = openai.APIError( - message="Internal server error", - request=MagicMock(), - body=None, - ) - mock_client.embeddings.create.side_effect = error - - embedder = JinaDenseEmbedder( - model_name="jina-embeddings-v5-text-small", - api_key="test-key", - ) - - with pytest.raises(ModelError, match="Jina API error"): - embedder.embed("hello") - def test_code_model_dimensions(self): """Code models should have correct default dimensions.""" embedder_1_5b = JinaDenseEmbedder( diff --git a/tests/unit/test_model_error.py b/tests/unit/test_model_error.py deleted file mode 100644 index ac5461ae5c..0000000000 --- a/tests/unit/test_model_error.py +++ /dev/null @@ -1,294 +0,0 @@ -# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. -# SPDX-License-Identifier: Apache-2.0 - -import asyncio - -import httpx -import openai -import pytest -from google.genai.errors import ClientError -from litellm import exceptions as litellm_exceptions -from volcenginesdkarkruntime import _exceptions as ark_exceptions - -from openviking.models.model_error import ( - ModelError, - ModelErrorKind, - model_error_boundary, - model_error_from_exception, - model_error_from_status, -) -from openviking.utils.model_retry import retry_async, retry_sync - - -def _response(status_code: int) -> httpx.Response: - return httpx.Response(status_code, request=httpx.Request("POST", "https://api.example.test")) - - -def _http_status_error(status_code: int, body: dict) -> httpx.HTTPStatusError: - response = httpx.Response( - status_code, - request=httpx.Request("POST", "https://api.example.test"), - json=body, - ) - return httpx.HTTPStatusError( - "provider request failed", - request=response.request, - response=response, - ) - - -@pytest.mark.parametrize( - ("error", "expected"), - [ - ( - openai.RateLimitError( - "quota", - response=_response(429), - body={"error": {"code": "insufficient_quota"}}, - ), - ModelErrorKind.QUOTA_EXHAUSTED, - ), - ( - openai.RateLimitError( - "rate limited", - response=_response(429), - body={"error": {"code": "rate_limit_exceeded"}}, - ), - ModelErrorKind.RATE_LIMIT_ERROR, - ), - ( - openai.APIConnectionError(request=httpx.Request("POST", "https://api.openai.com")), - ModelErrorKind.TRANSIENT_NETWORK_ERROR, - ), - ( - ark_exceptions.ArkRateLimitError( - "quota", - response=_response(429), - body={"error": {"code": "AccountQuotaExceeded"}}, - request_id="request-id", - ), - ModelErrorKind.QUOTA_EXHAUSTED, - ), - ( - ark_exceptions.ArkContentFilterFinishReasonError(), - ModelErrorKind.CONTENT_SAFETY, - ), - ( - litellm_exceptions.ContextWindowExceededError( - "too large", - model="model", - llm_provider="openai", - ), - ModelErrorKind.INPUT_TOO_LARGE, - ), - ( - litellm_exceptions.ServiceUnavailableError( - "unavailable", - model="model", - llm_provider="openai", - ), - ModelErrorKind.TRANSIENT_SERVICE_ERROR, - ), - ( - httpx.ReadTimeout("timed out"), - ModelErrorKind.TRANSIENT_NETWORK_ERROR, - ), - ( - _http_status_error(429, {"code": "AccountQuotaExceeded"}), - ModelErrorKind.QUOTA_EXHAUSTED, - ), - ( - openai.NotFoundError( - "model not found", - response=_response(404), - body={"error": {"code": "model_not_found"}}, - ), - ModelErrorKind.MODEL_NOT_FOUND, - ), - ( - _http_status_error(404, {}), - ModelErrorKind.UNKNOWN, - ), - ], -) -def test_sdk_errors_map_to_model_error_kind(error, expected): - converted = model_error_from_exception(error) - - assert converted is not None - assert converted.kind == expected - - -def test_plain_http_404_requires_structured_model_code(): - converted = model_error_from_status( - 404, - "not found", - provider="http-provider", - ) - - assert converted.kind == ModelErrorKind.UNKNOWN - - -@pytest.mark.parametrize( - ("reason", "expected"), - [ - ("QUOTA_EXCEEDED", ModelErrorKind.QUOTA_EXHAUSTED), - ("RATE_LIMIT_EXCEEDED", ModelErrorKind.RATE_LIMIT_ERROR), - (None, ModelErrorKind.UNKNOWN), - ], -) -def test_gemini_resource_exhausted_requires_structured_reason(reason, expected): - details = [] if reason is None else [{"reason": reason}] - error = ClientError( - 429, - { - "error": { - "code": 429, - "message": "resource exhausted", - "status": "RESOURCE_EXHAUSTED", - "details": details, - } - }, - ) - - converted = model_error_from_exception(error) - - assert converted is not None - assert converted.kind == expected - - -def test_provider_boundary_preserves_original_sdk_error_as_cause(): - sdk_error = openai.BadRequestError( - "bad request", - response=_response(400), - body={"error": {"code": "invalid_request_error"}}, - ) - - @model_error_boundary - def _call(): - raise sdk_error - - with pytest.raises(ModelError) as raised: - _call() - - assert raised.value.kind == ModelErrorKind.INVALID_REQUEST - assert raised.value.__cause__ is sdk_error - - -def test_provider_boundary_reports_actual_provider_name(): - class _Provider: - provider = "jina" - - @model_error_boundary - def call(self): - raise openai.APIConnectionError( - request=httpx.Request("POST", "https://api.jina.ai") - ) - - with pytest.raises(ModelError) as raised: - _Provider().call() - - assert raised.value.provider == "jina" - - -def test_provider_boundary_does_not_convert_local_program_error(): - error = ValueError("HTTP 503 is data, not an API failure") - - @model_error_boundary - def _call(): - raise error - - with pytest.raises(ValueError) as raised: - _call() - - assert raised.value is error - - -def test_unknown_model_error_uses_standard_retry_budget(): - attempts = 0 - - def _call(): - nonlocal attempts - attempts += 1 - raise ModelError(ModelErrorKind.UNKNOWN, "unclassified API error") - - with pytest.raises(ModelError): - retry_sync( - _call, - max_retries=5, - base_delay=0, - jitter=False, - model_errors_only=True, - ) - - assert attempts == 6 - - -@pytest.mark.parametrize( - "kind", - [ - ModelErrorKind.AUTHENTICATION_ERROR, - ModelErrorKind.PERMISSION_DENIED, - ModelErrorKind.QUOTA_EXHAUSTED, - ModelErrorKind.MODEL_NOT_FOUND, - ModelErrorKind.UNSUPPORTED, - ModelErrorKind.INVALID_REQUEST, - ModelErrorKind.INPUT_TOO_LARGE, - ModelErrorKind.CONTENT_SAFETY, - ], -) -def test_strict_retry_propagates_unhandled_model_error_unchanged(kind): - attempts = 0 - error = ModelError(kind, "handled by another layer") - - def _call(): - nonlocal attempts - attempts += 1 - raise error - - with pytest.raises(ModelError) as raised: - retry_sync( - _call, - max_retries=5, - base_delay=0, - jitter=False, - model_errors_only=True, - ) - - assert raised.value is error - assert attempts == 1 - - -def test_strict_retry_does_not_parse_local_error_message(): - attempts = 0 - - def _call(): - nonlocal attempts - attempts += 1 - raise RuntimeError("HTTP 503") - - with pytest.raises(RuntimeError): - retry_sync( - _call, - max_retries=5, - base_delay=0, - jitter=False, - model_errors_only=True, - ) - - assert attempts == 1 - - -@pytest.mark.asyncio -async def test_provider_boundary_and_retry_preserve_cancellation(): - attempts = 0 - - @model_error_boundary - async def _call(): - nonlocal attempts - attempts += 1 - raise asyncio.CancelledError - - with pytest.raises(asyncio.CancelledError): - await retry_async(_call, max_retries=5, model_errors_only=True) - - assert attempts == 1 diff --git a/tests/unit/test_model_retry.py b/tests/unit/test_model_retry.py index dd2160bc4f..596b989ad7 100644 --- a/tests/unit/test_model_retry.py +++ b/tests/unit/test_model_retry.py @@ -2,11 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 """Tests for shared model retry helpers.""" -import asyncio - +import httpx +import openai import pytest -from openviking.models.model_error import ModelError, ModelErrorKind +from openviking.models.model_error import ModelError, ModelErrorKind, model_error_boundary from openviking.pyagfs.exceptions import ( AGFSClientError, AGFSConnectionError, @@ -23,9 +23,7 @@ ERROR_CLASS_TRANSIENT, ERROR_CLASS_UNKNOWN, classify_api_error, - effective_max_retries, extract_metric_error_code, - model_retry_scope, retry_async, retry_sync, ) @@ -73,36 +71,6 @@ def test_classify_api_error_recognizes_request_burst_too_fast(): assert classify_api_error(RuntimeError("RequestBurstTooFast")) == ERROR_CLASS_TRANSIENT -@pytest.mark.asyncio -async def test_model_retry_scope_overrides_and_restores_configured_value(): - assert effective_max_retries(3) == 3 - - async with model_retry_scope(0): - assert effective_max_retries(3) == 0 - - assert effective_max_retries(3) == 3 - - -@pytest.mark.asyncio -async def test_model_retry_scope_is_isolated_between_tasks(): - scoped_started = asyncio.Event() - default_checked = asyncio.Event() - - async def scoped_call(): - async with model_retry_scope(0): - scoped_started.set() - await default_checked.wait() - return effective_max_retries(3) - - async def default_call(): - await scoped_started.wait() - retries = effective_max_retries(3) - default_checked.set() - return retries - - assert await asyncio.gather(scoped_call(), default_call()) == [0, 3] - - class _ProviderError(RuntimeError): def __init__(self, *, status_code=None, error_code=None, code=None, body=None): super().__init__("provider request failed") @@ -165,16 +133,19 @@ def _call(): @pytest.mark.asyncio -async def test_retry_async_does_not_retry_unknown_error(): +async def test_retry_async_preserves_local_error_without_parsing_message(): attempts = {"count": 0} + error = RuntimeError("HTTP 503 is input data, not an API failure") + @model_error_boundary async def _call(): attempts["count"] += 1 - raise RuntimeError("some unexpected validation failure") + raise error - with pytest.raises(RuntimeError): - await retry_async(_call, max_retries=3) + with pytest.raises(RuntimeError) as raised: + await retry_async(_call, max_retries=3, model_errors_only=True) + assert raised.value is error assert attempts["count"] == 1 @@ -268,16 +239,25 @@ def test_content_safety_takes_precedence_over_400(): def test_retry_sync_does_not_retry_quota_exceeded(): - """Quota-exceeded errors should NOT be retried.""" + """Structured quota exhaustion takes precedence over the retryable HTTP 429.""" attempts = {"count": 0} + sdk_error = openai.RateLimitError( + "rate limited", + response=httpx.Response(429, request=httpx.Request("POST", "https://example.test")), + body={"error": {"code": "insufficient_quota"}}, + ) + @model_error_boundary def _call(): attempts["count"] += 1 - raise RuntimeError("AccountQuotaExceeded") + raise sdk_error - with pytest.raises(RuntimeError, match="AccountQuotaExceeded"): - retry_sync(_call, max_retries=5) + with pytest.raises(ModelError) as raised: + retry_sync(_call, max_retries=5, model_errors_only=True) + assert raised.value.kind is ModelErrorKind.QUOTA_EXHAUSTED + assert raised.value.__cause__ is sdk_error + assert classify_api_error(raised.value) == ERROR_CLASS_QUOTA_EXCEEDED assert attempts["count"] == 1 diff --git a/tests/unit/test_retry_embedder.py b/tests/unit/test_retry_embedder.py index b4c3c263b1..3d4d755dfd 100644 --- a/tests/unit/test_retry_embedder.py +++ b/tests/unit/test_retry_embedder.py @@ -130,11 +130,11 @@ def test_retry_embedder_does_not_retry_permanent_or_auth_errors(status_code): assert inner.calls == 1 -def test_retry_embedder_uses_max_retries_as_additional_attempts(): - errors = [ProviderError(503) for _ in range(3)] +def test_retry_embedder_unknown_api_error_exhausts_configured_budget(): + errors = [ModelError(ModelErrorKind.UNKNOWN, "unclassified API error") for _ in range(3)] inner = FakeDenseEmbedder(errors, max_retries=2) - with pytest.raises(ProviderError) as raised: + with pytest.raises(ModelError) as raised: RetryEmbedder(inner).embed("query") assert raised.value is errors[-1] @@ -150,8 +150,10 @@ def test_retry_embedder_sparse_dimension_remains_unsupported(): @pytest.mark.asyncio -async def test_retry_releases_semaphore_before_backoff(monkeypatch): - inner = FakeDenseEmbedder([ProviderError(503), [0.1, 0.2]], max_retries=1) +async def test_scoped_failure_restores_retries_and_releases_semaphore(monkeypatch): + error = ProviderError(503) + inner = FakeDenseEmbedder([error, ProviderError(503), [0.1, 0.2]], max_retries=1) + embedder = RetryEmbedder(inner) semaphore = AsyncSemaphore(1) inner._account_semaphore = semaphore observed_values = [] @@ -161,27 +163,22 @@ async def observe_backoff(_delay): monkeypatch.setattr("openviking.utils.model_retry.asyncio.sleep", observe_backoff) - result = await RetryEmbedder(inner).embed_async("query") + async with model_retry_scope(0): + with pytest.raises(ProviderError) as raised: + await embedder.embed_async("query") + assert raised.value is error + assert inner.calls == 1 + assert observed_values == [] + assert semaphore._value == 1 + + result = await embedder.embed_async("query") assert result.dense_vector == [0.1, 0.2] - assert inner.calls == 2 + assert inner.calls == 3 assert observed_values == [1] assert semaphore._value == 1 -@pytest.mark.asyncio -async def test_retry_embedder_honors_call_scope_without_mutating_instance(): - inner = FakeDenseEmbedder([ProviderError(503), [0.1, 0.2]], max_retries=2) - embedder = RetryEmbedder(inner) - - async with model_retry_scope(0): - with pytest.raises(ProviderError): - await embedder.embed_async("query") - - assert inner.calls == 1 - assert embedder.max_retries == 2 - - @pytest.mark.asyncio async def test_cancelled_attempt_releases_semaphore(): started = asyncio.Event() diff --git a/tests/unit/test_retry_vlm.py b/tests/unit/test_retry_vlm.py index e4b3eabe2d..eb42ba917c 100644 --- a/tests/unit/test_retry_vlm.py +++ b/tests/unit/test_retry_vlm.py @@ -103,14 +103,21 @@ def test_retry_vlm_retries_sync_transient_failures(method, expected_method): ("get_vision_completion_async", "async_vision"), ], ) -async def test_retry_vlm_retries_async_transient_failures(method, expected_method): - inner = FakeVLM([ProviderError(503), "ok"]) +async def test_retry_vlm_scoped_failure_restores_async_retries(method, expected_method): + error = ProviderError(503) + inner = FakeVLM([error, ProviderError(503), "ok"]) vlm = RetryVLM(inner) + async with model_retry_scope(0): + with pytest.raises(ProviderError) as raised: + await getattr(vlm, method)(prompt="hello") + assert raised.value is error + assert len(inner.calls) == 1 + result = await getattr(vlm, method)(prompt="hello") assert result == "ok" - assert [name for name, _kwargs in inner.calls] == [expected_method, expected_method] + assert [name for name, _kwargs in inner.calls] == [expected_method] * 3 @pytest.mark.parametrize("status_code", [400, 401]) @@ -148,19 +155,6 @@ def test_retry_vlm_uses_max_retries_as_additional_attempts(): assert len(inner.calls) == 3 -@pytest.mark.asyncio -async def test_retry_vlm_honors_call_scope_without_mutating_instance(): - inner = FakeVLM([ProviderError(503), "unexpected"]) - vlm = RetryVLM(inner) - - async with model_retry_scope(0): - with pytest.raises(ProviderError): - await vlm.get_completion_async(prompt="hello") - - assert len(inner.calls) == 1 - assert vlm.max_retries == 2 - - @pytest.mark.asyncio async def test_retry_vlm_does_not_retry_cancelled_request(): started = asyncio.Event() From 4da8fb78f6c77c773618651a4c61e3cfe4607b27 Mon Sep 17 00:00:00 2001 From: "zhanghe.1216" <@bytedance.com> Date: Wed, 14 Oct 2026 20:25:58 +0800 Subject: [PATCH 12/16] =?UTF-8?q?fix(models):=20=E7=BB=9F=E4=B8=80?= =?UTF-8?q?=E9=94=99=E8=AF=AF=E8=AF=AD=E4=B9=89=E5=B9=B6=E4=BF=AE=E5=A4=8D?= =?UTF-8?q?=E6=A8=A1=E5=9E=8B=E6=95=85=E9=9A=9C=E8=BD=AC=E7=A7=BB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 直接基于 ModelErrorKind 驱动重试、凭证切换与熔断策略，确保模型不存在时切换备用配置，并优先解析 Provider 结构化业务错误。 --- bot/vikingbot/compile/ops/common.py | 4 +- bot/vikingbot/compile/ops/reduce.py | 4 +- openviking/config/embedding.py | 17 +-- openviking/models/embedder/base.py | 26 ++-- openviking/models/model_error.py | 18 ++- openviking/models/vlm/base.py | 41 +++--- openviking/storage/collection_schemas.py | 29 ++--- .../storage/queuefs/semantic_processor.py | 38 ++++-- openviking/utils/circuit_breaker.py | 24 +++- openviking/utils/exceptions.py | 39 +++++- openviking/utils/model_retry.py | 114 ++++++++--------- tests/storage/test_collection_schemas.py | 18 +-- ...antic_processor_permanent_storage_error.py | 9 +- ...est_failover_embedder_permanent_advance.py | 29 +++-- tests/unit/test_model_retry.py | 9 +- .../unit/test_ordered_credential_switcher.py | 32 +++-- tests/unit/test_retry_embedder.py | 17 +++ tests/unit/test_retry_vlm.py | 119 ++++++++++++++++++ 18 files changed, 388 insertions(+), 199 deletions(-) diff --git a/bot/vikingbot/compile/ops/common.py b/bot/vikingbot/compile/ops/common.py index 360ec2363b..6c49864df1 100644 --- a/bot/vikingbot/compile/ops/common.py +++ b/bot/vikingbot/compile/ops/common.py @@ -11,7 +11,7 @@ from typing import TYPE_CHECKING from openviking.core.namespace import relative_uri_path -from openviking.utils.model_retry import ERROR_CLASS_INPUT_TOO_LARGE, classify_api_error +from openviking.models.model_error import ModelErrorKind, model_error_kind from vikingbot.compile import file_ops from vikingbot.compile.hashing import content_hash, digest from vikingbot.compile.pipeline_io import MAX_JOB_ATTEMPTS @@ -347,7 +347,7 @@ async def job(runtime: Pipeline, name, inputs, work): except OSError as exc: if ( attempt == MAX_JOB_ATTEMPTS - 1 - or classify_api_error(exc) == ERROR_CLASS_INPUT_TOO_LARGE + or model_error_kind(exc) is ModelErrorKind.INPUT_TOO_LARGE ): raise runtime.metrics["job_retries"] += 1 diff --git a/bot/vikingbot/compile/ops/reduce.py b/bot/vikingbot/compile/ops/reduce.py index 1c4b548321..eb7af8d550 100644 --- a/bot/vikingbot/compile/ops/reduce.py +++ b/bot/vikingbot/compile/ops/reduce.py @@ -11,7 +11,7 @@ from typing import TYPE_CHECKING from openviking.core.namespace import relative_uri_path -from openviking.utils.model_retry import ERROR_CLASS_INPUT_TOO_LARGE, classify_api_error +from openviking.models.model_error import ModelErrorKind, model_error_kind from vikingbot.compile import file_ops from vikingbot.compile.hashing import content_hash, digest from vikingbot.compile.ops import common @@ -317,7 +317,7 @@ def validate_files(response): break except ModelCallError as exc: if ( - classify_api_error(exc) != ERROR_CLASS_INPUT_TOO_LARGE + model_error_kind(exc) is not ModelErrorKind.INPUT_TOO_LARGE or depth == 3 or stage == "map" ): diff --git a/openviking/config/embedding.py b/openviking/config/embedding.py index 656bfaa735..29ca34af28 100644 --- a/openviking/config/embedding.py +++ b/openviking/config/embedding.py @@ -17,12 +17,6 @@ from openviking.models.vlm.token_usage import TokenUsageTracker from openviking.service.task_tracker_concurrency import KeyedAsyncLockPool from openviking.utils.circuit_breaker import CircuitBreaker, CircuitBreakerOpen -from openviking.utils.model_retry import ( - ERROR_CLASS_AUTH, - ERROR_CLASS_INPUT_TOO_LARGE, - ERROR_CLASS_PERMANENT, - classify_api_error, -) from openviking_cli.utils.logger import get_logger logger = get_logger(__name__) @@ -260,16 +254,7 @@ async def _execute(self, resource, content, is_query): f"got {len(result.dense_vector)}" ) except Exception as exc: - # The breaker is shared by every call on this account, so only - # provider-health failures may trip it. Request-level rejections - # (400 / input too large) are specific to one input, and auth - # errors are surfaced per call instead. - if classify_api_error(exc) not in { - ERROR_CLASS_AUTH, - ERROR_CLASS_INPUT_TOO_LARGE, - ERROR_CLASS_PERMANENT, - }: - resource.breaker.record_failure(exc) + resource.breaker.record_failure(exc) raise resource.breaker.record_success() return result diff --git a/openviking/models/embedder/base.py b/openviking/models/embedder/base.py index 329ede0b84..a8e21c2350 100644 --- a/openviking/models/embedder/base.py +++ b/openviking/models/embedder/base.py @@ -12,16 +12,14 @@ from threading import Lock from typing import Any, AsyncIterator, Awaitable, Callable, Dict, List, Optional, TypeVar, Union +from openviking.models.model_error import model_error_kind from openviking.telemetry import get_current_telemetry from openviking.utils.embedding_input import ( resolve_embedding_max_input_tokens, truncate_embedding_input, ) from openviking.utils.exceptions import AllCredentialsFailedError -from openviking.utils.model_retry import ( - OrderedCredentialSwitcher, - classify_api_error, -) +from openviking.utils.model_retry import OrderedCredentialSwitcher from openviking_cli.utils import get_logger T = TypeVar("T") @@ -856,16 +854,18 @@ def _embed_with_failover(self, method_name: str, *args, **kwargs) -> EmbedResult self._switcher.commit_success(idx) return result except Exception as exc: - error_class = classify_api_error(exc) - aggregated_errors.append((credential_id, error_class, exc, idx)) + kind = model_error_kind(exc) - if self._switcher.is_fail_fast(error_class): + if self._switcher.is_fail_fast(kind): # Request-level failure: re-raise the original exception; # trying other credentials is useless. raise + aggregated_errors.append((credential_id, kind, exc, idx)) logger.warning( - f"Credential {credential_id} failed with {error_class}, trying next credential" + "Credential %s failed with %s, trying next credential", + credential_id, + kind.value, ) raise AllCredentialsFailedError(aggregated_errors) @@ -901,16 +901,18 @@ async def _embed_with_failover_async(self, method_name: str, *args, **kwargs) -> self._switcher.commit_success(idx) return result except Exception as exc: - error_class = classify_api_error(exc) - aggregated_errors.append((credential_id, error_class, exc, idx)) + kind = model_error_kind(exc) - if self._switcher.is_fail_fast(error_class): + if self._switcher.is_fail_fast(kind): # Request-level failure: re-raise the original exception; # trying other credentials is useless. raise + aggregated_errors.append((credential_id, kind, exc, idx)) logger.warning( - f"Credential {credential_id} failed with {error_class}, trying next credential" + "Credential %s failed with %s, trying next credential", + credential_id, + kind.value, ) raise AllCredentialsFailedError(aggregated_errors) diff --git a/openviking/models/model_error.py b/openviking/models/model_error.py index 2128ce212d..de2b3f5c9e 100644 --- a/openviking/models/model_error.py +++ b/openviking/models/model_error.py @@ -219,8 +219,8 @@ def _sdk_kind(error: BaseException) -> tuple[ModelErrorKind, str] | None: "InternalServerError": ModelErrorKind.TRANSIENT_SERVICE_ERROR, } kind = ( - class_kinds.get(name) - or _kind_from_code(error) + _kind_from_code(error) + or class_kinds.get(name) or _kind_from_status(_status_code(error)) ) return kind or ModelErrorKind.UNKNOWN, provider @@ -237,8 +237,8 @@ def _sdk_kind(error: BaseException) -> tuple[ModelErrorKind, str] | None: "ArkContentFilterFinishReasonError": ModelErrorKind.CONTENT_SAFETY, } kind = ( - class_kinds.get(name) - or _kind_from_code(error) + _kind_from_code(error) + or class_kinds.get(name) or _kind_from_status(_status_code(error)) ) return kind or ModelErrorKind.UNKNOWN, provider @@ -329,6 +329,16 @@ def model_error_from_exception(error: BaseException) -> ModelError | None: return None +def model_error_kind(error: BaseException) -> ModelErrorKind | None: + """Return the structured model kind carried by an exception chain.""" + for item in _exception_chain(error): + kind = getattr(item, "kind", None) + if isinstance(kind, ModelErrorKind): + return kind + converted = model_error_from_exception(error) + return converted.kind if converted is not None else None + + F = TypeVar("F", bound=Callable[..., Any]) diff --git a/openviking/models/vlm/base.py b/openviking/models/vlm/base.py index 930b41ceea..deb07cb633 100644 --- a/openviking/models/vlm/base.py +++ b/openviking/models/vlm/base.py @@ -9,11 +9,11 @@ from pathlib import Path from typing import Any, AsyncIterator, Callable, Dict, List, Optional, Union +from openviking.models.model_error import model_error_kind from openviking.utils.exceptions import AllCredentialsFailedError from openviking.utils.model_retry import ( OrderedCredentialSwitcher, PrimaryBackupSwitcher, - classify_api_error, extract_metric_error_code, ) from openviking_cli.utils import get_logger @@ -880,17 +880,19 @@ def _get_completion_with_failover(self, method_name: str, *args, **kwargs): return result except Exception as exc: _annotate_vlm_error(exc, vlm_instance) - error_class = classify_api_error(exc) - aggregated_errors.append((credential_id, error_class, exc, idx)) + kind = model_error_kind(exc) - if self._switcher.is_fail_fast(error_class): + if self._switcher.is_fail_fast(kind): # Request-level failure (400 / input too large / content # safety): re-raise the original exception so callers can # react to its type; trying other credentials is useless. raise + aggregated_errors.append((credential_id, kind, exc, idx)) self._logger.warning( - f"Credential {credential_id} failed with {error_class}, trying next credential" + "Credential %s failed with %s, trying next credential", + credential_id, + kind.value, ) raise AllCredentialsFailedError(aggregated_errors) @@ -927,16 +929,18 @@ async def _get_completion_with_failover_async(self, method_name: str, *args, **k return result except Exception as exc: _annotate_vlm_error(exc, vlm_instance) - error_class = classify_api_error(exc) - aggregated_errors.append((credential_id, error_class, exc, idx)) + kind = model_error_kind(exc) - if self._switcher.is_fail_fast(error_class): + if self._switcher.is_fail_fast(kind): # Request-level failure: re-raise the original exception; # trying other credentials is useless. raise + aggregated_errors.append((credential_id, kind, exc, idx)) self._logger.warning( - f"Credential {credential_id} failed with {error_class}, trying next credential" + "Credential %s failed with %s, trying next credential", + credential_id, + kind.value, ) raise AllCredentialsFailedError(aggregated_errors) @@ -972,13 +976,14 @@ async def stream_with_failover( if emitted: raise - error_class = classify_api_error(exc) - aggregated_errors.append((credential_id, error_class, exc, idx)) - if self._switcher.is_fail_fast(error_class): + kind = model_error_kind(exc) + if self._switcher.is_fail_fast(kind): raise - + aggregated_errors.append((credential_id, kind, exc, idx)) self._logger.warning( - f"Credential {credential_id} failed with {error_class}, trying next credential" + "Credential %s failed with %s, trying next credential", + credential_id, + kind.value, ) raise AllCredentialsFailedError(aggregated_errors) @@ -1117,16 +1122,16 @@ async def get_media_completion_async( return result except Exception as error: _annotate_vlm_error(error, vlm) - error_class = classify_api_error(error) - errors.append((credential_id, error_class, error, idx)) + kind = model_error_kind(error) if idx == start: active_credential_failed = True - if self._switcher.is_fail_fast(error_class): + if self._switcher.is_fail_fast(kind): raise + errors.append((credential_id, kind, error, idx)) self._logger.warning( "Credential %s failed media understanding with %s, trying next credential", credential_id, - error_class, + kind.value, ) if attempted: diff --git a/openviking/storage/collection_schemas.py b/openviking/storage/collection_schemas.py index 25baed1552..73b9b2d789 100644 --- a/openviking/storage/collection_schemas.py +++ b/openviking/storage/collection_schemas.py @@ -19,6 +19,7 @@ from openviking.core.context import ContextType, ResourceContentType from openviking.models.embedder.base import embed_compat +from openviking.models.model_error import ModelErrorKind, model_error_kind from openviking.server.identity import RequestContext, Role from openviking.service.task_tracker_concurrency import run_to_completion from openviking.storage.acl import ACL_GRANT_FIELDS, ACL_MODE_FIELD, AclMode @@ -43,16 +44,8 @@ ) from openviking.telemetry import bind_telemetry, resolve_telemetry from openviking.telemetry.request_wait_tracker import get_request_wait_tracker -from openviking.utils.circuit_breaker import ( - CircuitBreakerOpen, - classify_api_error, -) +from openviking.utils.circuit_breaker import CircuitBreakerOpen from openviking.utils.log_correlation import log_correlation -from openviking.utils.model_retry import ( - ERROR_CLASS_AUTH, - ERROR_CLASS_INPUT_TOO_LARGE, - ERROR_CLASS_PERMANENT, -) from openviking.utils.tags import merge_search_tags from openviking.utils.time_utils import get_current_timestamp from openviking_cli.session.user_id import UserIdentifier @@ -782,32 +775,40 @@ async def on_dequeue(self, data: Optional[Dict[str, Any]]) -> ProcessResult: embedding_msg, f"Failed to generate embedding: {embed_err}", ) - error_class = classify_api_error(embed_err) + kind = model_error_kind(embed_err) try: from openviking.metrics.datasources import EmbeddingEventDataSource EmbeddingEventDataSource.record_error( - error_code=str(error_class or "unknown"), + error_code=kind.value if kind is not None else "unknown", account_id=embedding_msg.context_data.get("account_id"), ) except Exception: pass - if error_class == ERROR_CLASS_INPUT_TOO_LARGE: + if kind is ModelErrorKind.INPUT_TOO_LARGE: execute_status = "error" self._log_embedding_error(logging.ERROR, error_msg, embedding_msg) self._merge_request_stats(embedding_msg.telemetry_id, error_count=1) request_failed_message = error_msg return ProcessResult.failed(error_msg) - if error_class == ERROR_CLASS_PERMANENT: + if kind in { + ModelErrorKind.MODEL_NOT_FOUND, + ModelErrorKind.UNSUPPORTED, + ModelErrorKind.INVALID_REQUEST, + ModelErrorKind.CONTENT_SAFETY, + }: execute_status = "error" self._log_embedding_error(logging.CRITICAL, error_msg, embedding_msg) self._merge_request_stats(embedding_msg.telemetry_id, error_count=1) request_failed_message = error_msg return ProcessResult.failed(error_msg) - if error_class == ERROR_CLASS_AUTH: + if kind in { + ModelErrorKind.AUTHENTICATION_ERROR, + ModelErrorKind.PERMISSION_DENIED, + }: execute_status = "error" # Bad/expired credential: retrying cannot succeed. Fail # terminally instead of re-enqueueing, which would cycle diff --git a/openviking/storage/queuefs/semantic_processor.py b/openviking/storage/queuefs/semantic_processor.py index 8352f40bf1..4163c32164 100644 --- a/openviking/storage/queuefs/semantic_processor.py +++ b/openviking/storage/queuefs/semantic_processor.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, List, Optional, Set, Tuple from openviking.core.namespace import classify_uri +from openviking.models.model_error import ModelErrorKind, model_error_kind from openviking.observability.context import ( bind_root_observability_context, reset_root_observability_context, @@ -64,13 +65,8 @@ ) from openviking.telemetry.request_wait_tracker import get_request_wait_tracker from openviking.telemetry.span_models import create_root_span_attributes -from openviking.utils.circuit_breaker import ( - CircuitBreaker, - CircuitBreakerOpen, - classify_api_error, -) +from openviking.utils.circuit_breaker import CircuitBreaker, CircuitBreakerOpen from openviking.utils.ingest_options import IngestOptions -from openviking.utils.model_retry import ERROR_CLASS_INPUT_TOO_LARGE, ERROR_CLASS_PERMANENT from openviking_cli.session.user_id import UserIdentifier from openviking_cli.utils import VikingURI from openviking_cli.utils.config import get_openviking_config @@ -81,6 +77,18 @@ logger = get_logger(__name__) +_PERMANENT_STORAGE_ERRORS = ( + FileNotFoundError, + PermissionError, + IsADirectoryError, + NotADirectoryError, + AGFSNotADirectoryError, +) + + +def _is_permanent_storage_error(error: BaseException) -> bool: + return any(isinstance(item, _PERMANENT_STORAGE_ERRORS) for item in (error, error.__cause__)) + class RequestQueueStats: processed: int = 0 @@ -785,8 +793,8 @@ async def on_dequeue( return result return ProcessResult.failed(str(e)) - error_class = classify_api_error(e) - if error_class == ERROR_CLASS_INPUT_TOO_LARGE: + kind = model_error_kind(e) + if kind is ModelErrorKind.INPUT_TOO_LARGE: execute_status = "error" logger.error( f"Input too large processing semantic message, dropping: {e}", @@ -800,16 +808,18 @@ async def on_dequeue( if msg is not None: await self._cleanup_local_artifact(msg) return ProcessResult.failed(str(e)) - elif error_class == ERROR_CLASS_PERMANENT: + elif _is_permanent_storage_error(e) or kind in { + ModelErrorKind.MODEL_NOT_FOUND, + ModelErrorKind.UNSUPPORTED, + ModelErrorKind.INVALID_REQUEST, + ModelErrorKind.CONTENT_SAFETY, + }: execute_status = "error" logger.critical( f"Permanent error processing semantic message, dropping: {e}", exc_info=True, ) - # A malformed filesystem target does not indicate an API outage. - if breaker is not None and not any( - isinstance(exc, AGFSNotADirectoryError) for exc in (e, e.__cause__) - ): + if breaker is not None and kind is not None: breaker.record_failure(e) if msg is not None: self._merge_request_stats(msg.telemetry_id, error_count=1) @@ -821,6 +831,8 @@ async def on_dequeue( return ProcessResult.failed(str(e)) else: execute_status = "error" + if breaker is not None and kind is not None: + breaker.record_failure(e) logger.error( f"Error processing semantic message, not re-enqueueing: {e}", exc_info=True, diff --git a/openviking/utils/circuit_breaker.py b/openviking/utils/circuit_breaker.py index 9af434a075..8305aa22ac 100644 --- a/openviking/utils/circuit_breaker.py +++ b/openviking/utils/circuit_breaker.py @@ -7,6 +7,7 @@ import threading import time +from openviking.models.model_error import ModelErrorKind, model_error_kind from openviking.utils.model_retry import ( ERROR_CLASS_AUTH, ERROR_CLASS_INPUT_TOO_LARGE, @@ -95,9 +96,19 @@ def record_success(self) -> None: def record_failure(self, error: Exception) -> None: """Record a failed API call. May trip the breaker.""" - error_class = classify_api_error(error) - if error_class == ERROR_CLASS_INPUT_TOO_LARGE: - logger.info(f"Circuit breaker ignoring row-specific input error: {error}") + kind = model_error_kind(error) + legacy_class = classify_api_error(error) if kind is None else None + if kind in { + ModelErrorKind.AUTHENTICATION_ERROR, + ModelErrorKind.PERMISSION_DENIED, + ModelErrorKind.MODEL_NOT_FOUND, + ModelErrorKind.UNSUPPORTED, + ModelErrorKind.INVALID_REQUEST, + ModelErrorKind.INPUT_TOO_LARGE, + ModelErrorKind.CONTENT_SAFETY, + } or legacy_class == ERROR_CLASS_INPUT_TOO_LARGE: + label = kind.value if kind is not None else legacy_class + logger.info("Circuit breaker ignoring %s error: %s", label, error) return with self._lock: @@ -115,14 +126,15 @@ def record_failure(self, error: Exception) -> None: ) return - if error_class in ( + if kind is ModelErrorKind.QUOTA_EXHAUSTED or legacy_class in { ERROR_CLASS_PERMANENT, ERROR_CLASS_AUTH, ERROR_CLASS_QUOTA_EXCEEDED, - ): + }: self._state = _STATE_OPEN self._current_reset_timeout = self._base_reset_timeout - logger.info(f"Circuit breaker tripped immediately on {error_class} error: {error}") + label = kind.value if kind is not None else legacy_class + logger.info("Circuit breaker tripped immediately on %s: %s", label, error) return if self._failure_count >= self._failure_threshold: diff --git a/openviking/utils/exceptions.py b/openviking/utils/exceptions.py index d3b8444b24..0bf1119761 100644 --- a/openviking/utils/exceptions.py +++ b/openviking/utils/exceptions.py @@ -2,6 +2,8 @@ # SPDX-License-Identifier: AGPL-3.0 """Common exception helpers.""" +from openviking.models.model_error import ModelErrorKind + HTTP_STATUS_TO_ERROR_CODE = { 400: "INVALID_ARGUMENT", 401: "UNAUTHENTICATED", @@ -34,15 +36,42 @@ def error_code_from_http_status(status: int | None, default: str = "INVALID_ARGU class AllCredentialsFailedError(Exception): """Raised when all credentials in the chain have failed.""" - def __init__(self, errors: list[tuple[str, str, Exception, int]]): + def __init__( + self, + errors: list[tuple[str, ModelErrorKind | str, Exception, int]], + ): """Initialize the error with a list of credential failures. Args: - errors: List of tuples containing (credential_id, error_class, exception, attempts) + errors: List of tuples containing + (credential_id, model_error_kind, exception, credential_index) """ - self.errors = errors + self.errors = [ + ( + credential_id, + kind if isinstance(kind, ModelErrorKind) else ModelErrorKind.UNKNOWN, + error, + credential_index, + ) + for credential_id, kind, error, credential_index in errors + ] + kinds = [kind for _, kind, _, _ in self.errors] + self.kind = self._aggregate_kind(kinds) message = "All credentials failed:\n" + "\n".join( - f" - {cred_id}: {error_class} - {exc}" - for cred_id, error_class, exc, attempts in errors + f" - {cred_id}: {kind.value} - {exc}" for cred_id, kind, exc, _ in self.errors ) super().__init__(message) + + @staticmethod + def _aggregate_kind(kinds: list[ModelErrorKind]) -> ModelErrorKind: + for kind in ( + ModelErrorKind.TRANSIENT_SERVICE_ERROR, + ModelErrorKind.TRANSIENT_NETWORK_ERROR, + ModelErrorKind.RATE_LIMIT_ERROR, + ModelErrorKind.QUOTA_EXHAUSTED, + ): + if kind in kinds: + return kind + if kinds and all(kind is kinds[0] for kind in kinds): + return kinds[0] + return ModelErrorKind.UNKNOWN diff --git a/openviking/utils/model_retry.py b/openviking/utils/model_retry.py index 7dab79fbe6..5fe9a9a18b 100644 --- a/openviking/utils/model_retry.py +++ b/openviking/utils/model_retry.py @@ -10,9 +10,12 @@ from contextvars import ContextVar from typing import AsyncIterator, Awaitable, Callable, TypeVar -from openviking.models.model_error import ModelErrorKind, model_error_from_exception +from openviking.models.model_error import ( + ModelErrorKind, + model_error_from_exception, + model_error_kind, +) from openviking.pyagfs.exceptions import AGFSNotADirectoryError -from openviking.utils.exceptions import AllCredentialsFailedError logger = logging.getLogger(__name__) @@ -244,36 +247,9 @@ def classify_api_error(error: Exception) -> str: - ``quota_exceeded`` is checked before ``transient`` because quota errors typically include "429" / "TooManyRequests" which would otherwise match the transient category. - - an aggregated ``AllCredentialsFailedError`` is classified from its - per-credential classes, not its concatenated message. + Structured model errors are intentionally handled by their callers using + ``ModelErrorKind`` and are not collapsed into this legacy string taxonomy. """ - if isinstance(error, AllCredentialsFailedError): - classes = [ec for (_cid, ec, _exc, _idx) in error.errors if ec] - if ERROR_CLASS_TRANSIENT in classes: - return ERROR_CLASS_TRANSIENT - if ERROR_CLASS_QUOTA_EXCEEDED in classes: - return ERROR_CLASS_QUOTA_EXCEEDED - if classes and all(ec == ERROR_CLASS_AUTH for ec in classes): - return ERROR_CLASS_AUTH - return ERROR_CLASS_UNKNOWN - - model_error = model_error_from_exception(error) - if model_error is not None: - return { - ModelErrorKind.TRANSIENT_SERVICE_ERROR: ERROR_CLASS_TRANSIENT, - ModelErrorKind.TRANSIENT_NETWORK_ERROR: ERROR_CLASS_TRANSIENT, - ModelErrorKind.RATE_LIMIT_ERROR: ERROR_CLASS_TRANSIENT, - ModelErrorKind.AUTHENTICATION_ERROR: ERROR_CLASS_AUTH, - ModelErrorKind.PERMISSION_DENIED: ERROR_CLASS_AUTH, - ModelErrorKind.QUOTA_EXHAUSTED: ERROR_CLASS_QUOTA_EXCEEDED, - ModelErrorKind.MODEL_NOT_FOUND: ERROR_CLASS_PERMANENT, - ModelErrorKind.UNSUPPORTED: ERROR_CLASS_PERMANENT, - ModelErrorKind.INVALID_REQUEST: ERROR_CLASS_PERMANENT, - ModelErrorKind.INPUT_TOO_LARGE: ERROR_CLASS_INPUT_TOO_LARGE, - ModelErrorKind.CONTENT_SAFETY: ERROR_CLASS_CONTENT_SAFETY, - ModelErrorKind.UNKNOWN: ERROR_CLASS_UNKNOWN, - }[model_error.kind] - for exc in (error, getattr(error, "__cause__", None)): if exc is not None and isinstance(exc, _PERMANENT_IO_ERRORS): return ERROR_CLASS_PERMANENT @@ -332,6 +308,14 @@ def classify_api_error(error: Exception) -> str: def is_retryable_api_error(error: Exception) -> bool: """Return True if the error should be retried.""" + kind = model_error_kind(error) + if kind is not None: + return kind in { + ModelErrorKind.TRANSIENT_SERVICE_ERROR, + ModelErrorKind.TRANSIENT_NETWORK_ERROR, + ModelErrorKind.RATE_LIMIT_ERROR, + ModelErrorKind.UNKNOWN, + } return classify_api_error(error) == ERROR_CLASS_TRANSIENT @@ -556,8 +540,8 @@ async def retry_async( class PrimaryBackupSwitcher: """Thread-safe primary/backup switcher with automatic failback logic. - When an error of type ERROR_CLASS_PERMANENT or ERROR_CLASS_QUOTA_EXCEEDED occurs, - switches to backup immediately. Then, after either: + When a credential- or provider-specific model error occurs, switches to + backup immediately. Then, after either: - 10 minutes have passed, OR - 200 requests have been made to backup it will attempt to failback to primary. If failback fails, it switches back @@ -615,18 +599,29 @@ def record_primary_success(self) -> None: def record_primary_failure(self, error: Exception) -> bool: """Record a primary failure. Returns True if should switch to backup. - Switches to backup immediately for ERROR_CLASS_PERMANENT, - ERROR_CLASS_AUTH or ERROR_CLASS_QUOTA_EXCEEDED. + Request-specific errors and unknown failures do not switch models. """ - error_class = classify_api_error(error) - if error_class in ( - ERROR_CLASS_PERMANENT, - ERROR_CLASS_AUTH, - ERROR_CLASS_QUOTA_EXCEEDED, - ): + kind = model_error_kind(error) + should_switch = kind in { + ModelErrorKind.TRANSIENT_SERVICE_ERROR, + ModelErrorKind.TRANSIENT_NETWORK_ERROR, + ModelErrorKind.RATE_LIMIT_ERROR, + ModelErrorKind.AUTHENTICATION_ERROR, + ModelErrorKind.PERMISSION_DENIED, + ModelErrorKind.QUOTA_EXHAUSTED, + ModelErrorKind.MODEL_NOT_FOUND, + } + if kind is None: + should_switch = classify_api_error(error) in { + ERROR_CLASS_PERMANENT, + ERROR_CLASS_AUTH, + ERROR_CLASS_QUOTA_EXCEEDED, + } + if should_switch: with self._lock: if not self._using_backup: - logger.warning(f"Primary failed with {error_class}, switching to backup") + label = kind.value if kind is not None else classify_api_error(error) + logger.warning("Primary failed with %s, switching to backup", label) self._using_backup = True # Always reset timer and counter when we fail (whether initial fail or failback fail) self._switch_to_backup_time = time.monotonic() @@ -650,8 +645,7 @@ def is_using_backup(self) -> bool: class OrderedCredentialSwitcher: """Thread-safe ordered N-credential switcher with hierarchical failback. - Supports ordered failover across multiple credentials. When a credential fails - with quota_exceeded or permanent error, it advances to the next credential. + Supports ordered failover across multiple credentials. After failback thresholds are met, it attempts to move back to a higher-priority credential (one step at a time, not all the way back to index 0). @@ -672,17 +666,8 @@ def __init__( failback_request_count: Number of requests after which to attempt failback Note: - Failure handling is driven by the error class (see - ``classify_api_error``): - - - request-level errors (``permanent`` 400 / ``input_too_large`` / - ``content_safety``) fail fast: the same request fails on every - credential of the same model, so switching is useless. - - credential-level ``auth`` errors (401/403) advance to the next - credential in multi-credential mode; the last (or single) - credential fails fast. - - ``quota_exceeded`` (and ``transient`` once its retries are - exhausted) and ``unknown`` advance to the next credential. + Request-specific and unknown model errors fail fast. Credential, + model-routing, quota, and exhausted transient failures advance. """ if n < 1: raise ValueError("Number of credentials must be >= 1") @@ -746,18 +731,19 @@ def on_success(self, idx: int) -> None: self._active_request_count += 1 @staticmethod - def is_fail_fast(error_class: str) -> bool: + def is_fail_fast(kind: ModelErrorKind | None) -> bool: """Whether an error is request-level and must not try other credentials. - Request-level errors (400 parameter error, input too large, content - safety) fail on every credential of the same model, so the caller should - re-raise immediately instead of cycling through credentials. + Unknown failures remain within the retry layer and never expand into a + credential or model failover. """ - return error_class in ( - ERROR_CLASS_PERMANENT, - ERROR_CLASS_INPUT_TOO_LARGE, - ERROR_CLASS_CONTENT_SAFETY, - ) + return kind is None or kind in { + ModelErrorKind.UNSUPPORTED, + ModelErrorKind.INVALID_REQUEST, + ModelErrorKind.INPUT_TOO_LARGE, + ModelErrorKind.CONTENT_SAFETY, + ModelErrorKind.UNKNOWN, + } def commit_success(self, idx: int) -> None: """Record that credential ``idx`` successfully served a request. diff --git a/tests/storage/test_collection_schemas.py b/tests/storage/test_collection_schemas.py index 7fe27c7491..c0f149c014 100644 --- a/tests/storage/test_collection_schemas.py +++ b/tests/storage/test_collection_schemas.py @@ -910,7 +910,11 @@ async def enqueue_embedding_msg(self, msg): class _AuthErrorEmbedder(_DummyEmbedder): async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: - raise RuntimeError("Error code: 401 - {'code': 'AuthenticationError'} Unauthorized") + raise ModelError( + ModelErrorKind.AUTHENTICATION_ERROR, + "Error code: 401 - {'code': 'AuthenticationError'} Unauthorized", + status_code=401, + ) vikingdb = _QueueingVikingDB() monkeypatch.setattr( @@ -942,9 +946,11 @@ class _RejectFirstEmbedder(_DummyEmbedder): async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: if self.calls == 0: self.calls += 1 - raise RuntimeError( + raise ModelError( + ModelErrorKind.INVALID_REQUEST, "Error code: 400 - {'error': {'code': 'InvalidParameter', " - "'message': 'Invalid base64 image url'}}" + "'message': 'Invalid base64 image url'}}", + status_code=400, ) return self.embed(text, is_query=is_query) @@ -956,7 +962,7 @@ async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: handler = TextEmbeddingHandler(SimpleNamespace(is_closing=False, has_queue_manager=True)) provider = handler._embedding_provider - with pytest.raises(RuntimeError, match="400"): + with pytest.raises(ModelError, match="400"): await provider.embed("default", "bad input", is_query=is_query) (await handler.breaker()).check() # would raise CircuitBreakerOpen if tripped @@ -1131,10 +1137,6 @@ async def embed_async(self, text: str, is_query: bool = False) -> EmbedResult: "openviking.metrics.datasources.EmbeddingEventDataSource.record_error", staticmethod(lambda **kwargs: captured.update(kwargs)), ) - monkeypatch.setattr( - "openviking.storage.collection_schemas.classify_api_error", - lambda _err: "unknown", - ) handler = TextEmbeddingHandler(_DummyVikingDB()) await handler.on_dequeue(_build_queue_payload_for_account("acct-embed-error")) diff --git a/tests/storage/test_semantic_processor_permanent_storage_error.py b/tests/storage/test_semantic_processor_permanent_storage_error.py index 9582e724b2..af31296adc 100644 --- a/tests/storage/test_semantic_processor_permanent_storage_error.py +++ b/tests/storage/test_semantic_processor_permanent_storage_error.py @@ -256,10 +256,10 @@ async def test_semantic_model_error_fails_without_requeue(semantic_env): @pytest.mark.asyncio -async def test_permanent_api_error_still_opens_breaker(semantic_env): +async def test_invalid_request_does_not_open_breaker(semantic_env): env = semantic_env msg = env.message() - error = RuntimeError("400 invalid model parameter") + error = ModelError(ModelErrorKind.INVALID_REQUEST, "400 invalid model parameter") env.dag.run.side_effect = error result = await asyncio.wait_for(env.processor.on_dequeue(msg.to_dict()), timeout=1) @@ -268,8 +268,9 @@ async def test_permanent_api_error_still_opens_breaker(semantic_env): assert result.error == str(error) env.retry_queue.enqueue.assert_not_awaited() await _assert_request_failed(env, msg, error) - with pytest.raises(CircuitBreakerOpen): - env.processor._account_breaker(msg.account_id).check() + breaker = env.processor._account_breaker(msg.account_id) + breaker.check() + assert breaker._failure_count == 0 @pytest.mark.asyncio diff --git a/tests/unit/test_failover_embedder_permanent_advance.py b/tests/unit/test_failover_embedder_permanent_advance.py index c423d1c20a..664862ad2d 100644 --- a/tests/unit/test_failover_embedder_permanent_advance.py +++ b/tests/unit/test_failover_embedder_permanent_advance.py @@ -17,6 +17,7 @@ import pytest from openviking.models.embedder.base import EmbedderBase, EmbedResult, FailoverEmbedder +from openviking.models.model_error import ModelError, ModelErrorKind from openviking.utils.exceptions import AllCredentialsFailedError @@ -47,16 +48,22 @@ async def embed_async(self, content, is_query: bool = False) -> EmbedResult: def _make_400_error() -> Exception: - """Construct a request-level 400 parameter error (PERMANENT, fail-fast).""" - return RuntimeError( + """Construct a request-level invalid request error.""" + return ModelError( + ModelErrorKind.INVALID_REQUEST, "Error code: 400 - {'error': {'message': " - "'The parameter `model` specified in the request are not valid'}}" + "'The parameter `model` specified in the request are not valid'}}", + status_code=400, ) def _make_401_error() -> Exception: - """Construct a credential-level 401 auth error (AUTH, advances in multi-credential).""" - return RuntimeError("Error code: 401 - {'error': {'message': 'Incorrect API key provided'}}") + """Construct a credential-level auth error.""" + return ModelError( + ModelErrorKind.AUTHENTICATION_ERROR, + "Error code: 401 - {'error': {'message': 'Incorrect API key provided'}}", + status_code=401, + ) def test_auth_error_on_primary_advances_to_backup(): @@ -106,7 +113,7 @@ def test_permanent_400_fails_fast_without_trying_backup(): credential_ids=["primary", "backup"], ) - with pytest.raises(RuntimeError, match="400"): + with pytest.raises(ModelError, match="400"): fe.embed("hello") assert primary.calls == 1 @@ -202,12 +209,18 @@ def test_fast_failover_commits_new_active_index(): def test_content_safety_does_not_cycle_the_ring(): """Request-level content-safety errors fail fast and must not try others.""" - cs = _StubEmbedder("cs", error=RuntimeError("content_filter: sensitive content detected")) + cs = _StubEmbedder( + "cs", + error=ModelError( + ModelErrorKind.CONTENT_SAFETY, + "content_filter: sensitive content detected", + ), + ) backup = _StubEmbedder("backup") fe = FailoverEmbedder(embedders=[cs, backup], credential_ids=["cs", "backup"]) - with pytest.raises(RuntimeError, match="content_filter"): + with pytest.raises(ModelError, match="content_filter"): fe.embed("hello") assert cs.calls == 1 diff --git a/tests/unit/test_model_retry.py b/tests/unit/test_model_retry.py index 596b989ad7..ff4ab99765 100644 --- a/tests/unit/test_model_retry.py +++ b/tests/unit/test_model_retry.py @@ -99,24 +99,24 @@ def test_extract_metric_error_code_uses_safe_fallbacks_only(): assert extract_metric_error_code(RuntimeError("request_id=not-a-metric-label")) == "unknown" -def test_classify_all_credentials_failed_prefers_transient_over_auth(): +def test_all_credentials_failed_prefers_transient_over_auth(): err = AllCredentialsFailedError( [ ( "cred-a", - ERROR_CLASS_AUTH, + ModelErrorKind.AUTHENTICATION_ERROR, ModelError(ModelErrorKind.AUTHENTICATION_ERROR, "unauthorized"), 0, ), ( "cred-b", - ERROR_CLASS_TRANSIENT, + ModelErrorKind.TRANSIENT_SERVICE_ERROR, ModelError(ModelErrorKind.TRANSIENT_SERVICE_ERROR, "server error"), 1, ), ] ) - assert classify_api_error(err) == ERROR_CLASS_TRANSIENT + assert err.kind is ModelErrorKind.TRANSIENT_SERVICE_ERROR def test_retry_sync_retries_transient_error_until_success(): @@ -257,7 +257,6 @@ def _call(): assert raised.value.kind is ModelErrorKind.QUOTA_EXHAUSTED assert raised.value.__cause__ is sdk_error - assert classify_api_error(raised.value) == ERROR_CLASS_QUOTA_EXCEEDED assert attempts["count"] == 1 diff --git a/tests/unit/test_ordered_credential_switcher.py b/tests/unit/test_ordered_credential_switcher.py index 0e05dae945..05c55618d4 100644 --- a/tests/unit/test_ordered_credential_switcher.py +++ b/tests/unit/test_ordered_credential_switcher.py @@ -5,16 +5,8 @@ import pytest -from openviking.utils.model_retry import ( - ERROR_CLASS_AUTH, - ERROR_CLASS_CONTENT_SAFETY, - ERROR_CLASS_INPUT_TOO_LARGE, - ERROR_CLASS_PERMANENT, - ERROR_CLASS_QUOTA_EXCEEDED, - ERROR_CLASS_TRANSIENT, - ERROR_CLASS_UNKNOWN, - OrderedCredentialSwitcher, -) +from openviking.models.model_error import ModelErrorKind +from openviking.utils.model_retry import OrderedCredentialSwitcher class TestOrderedCredentialSwitcher: @@ -41,14 +33,18 @@ def test_invalid_initialization(self): OrderedCredentialSwitcher(n=-1) def test_is_fail_fast_classification(self): - """Request-level errors are fail-fast; credential/quota/transient are not.""" - assert OrderedCredentialSwitcher.is_fail_fast(ERROR_CLASS_PERMANENT) is True - assert OrderedCredentialSwitcher.is_fail_fast(ERROR_CLASS_INPUT_TOO_LARGE) is True - assert OrderedCredentialSwitcher.is_fail_fast(ERROR_CLASS_CONTENT_SAFETY) is True - assert OrderedCredentialSwitcher.is_fail_fast(ERROR_CLASS_AUTH) is False - assert OrderedCredentialSwitcher.is_fail_fast(ERROR_CLASS_QUOTA_EXCEEDED) is False - assert OrderedCredentialSwitcher.is_fail_fast(ERROR_CLASS_TRANSIENT) is False - assert OrderedCredentialSwitcher.is_fail_fast(ERROR_CLASS_UNKNOWN) is False + """Request/unknown errors fail fast; model-routing failures may advance.""" + assert OrderedCredentialSwitcher.is_fail_fast(ModelErrorKind.INVALID_REQUEST) is True + assert OrderedCredentialSwitcher.is_fail_fast(ModelErrorKind.INPUT_TOO_LARGE) is True + assert OrderedCredentialSwitcher.is_fail_fast(ModelErrorKind.CONTENT_SAFETY) is True + assert OrderedCredentialSwitcher.is_fail_fast(ModelErrorKind.UNKNOWN) is True + assert OrderedCredentialSwitcher.is_fail_fast(None) is True + assert OrderedCredentialSwitcher.is_fail_fast(ModelErrorKind.AUTHENTICATION_ERROR) is False + assert OrderedCredentialSwitcher.is_fail_fast(ModelErrorKind.QUOTA_EXHAUSTED) is False + assert ( + OrderedCredentialSwitcher.is_fail_fast(ModelErrorKind.TRANSIENT_SERVICE_ERROR) is False + ) + assert OrderedCredentialSwitcher.is_fail_fast(ModelErrorKind.MODEL_NOT_FOUND) is False def test_commit_success_different_index_fast_failover(self): """commit_success on a different index commits it as the new active one.""" diff --git a/tests/unit/test_retry_embedder.py b/tests/unit/test_retry_embedder.py index 3d4d755dfd..18b5231c63 100644 --- a/tests/unit/test_retry_embedder.py +++ b/tests/unit/test_retry_embedder.py @@ -236,6 +236,23 @@ def test_failover_runs_after_current_credential_retry_exhaustion(): assert backup.calls == 1 +def test_unknown_error_retries_without_expanding_to_credential_failover(): + errors = [ModelError(ModelErrorKind.UNKNOWN, "unclassified API error") for _ in range(3)] + primary = FakeDenseEmbedder(errors, max_retries=2) + backup = FakeDenseEmbedder([[0.3, 0.4]], max_retries=2) + failover = FailoverEmbedder( + [RetryEmbedder(primary), RetryEmbedder(backup)], + ["primary", "backup"], + ) + + with pytest.raises(ModelError) as raised: + failover.embed("query") + + assert raised.value is errors[-1] + assert primary.calls == 3 + assert backup.calls == 0 + + def test_account_state_reaches_retry_inner_and_close_is_delegated_once(): inner = FakeDenseEmbedder([[0.1, 0.2]]) embedder = RetryEmbedder(inner) diff --git a/tests/unit/test_retry_vlm.py b/tests/unit/test_retry_vlm.py index eb42ba917c..9e55db8de5 100644 --- a/tests/unit/test_retry_vlm.py +++ b/tests/unit/test_retry_vlm.py @@ -4,11 +4,14 @@ import asyncio from pathlib import Path +import httpx +import openai import pytest from openviking.models.model_error import ModelError, ModelErrorKind from openviking.models.vlm.base import MultiCredentialVLM, VLMBase, VLMFactory from openviking.models.vlm.retry_vlm import RetryVLM +from openviking.utils.circuit_breaker import CircuitBreaker from openviking.utils.model_retry import model_retry_scope from openviking_cli.utils.config.vlm_config import VLMConfig @@ -243,3 +246,119 @@ def test_vlm_config_wraps_each_credential_inside_failover(): "KimiVLM", "GLMVLM", ] + + +def _openai_client(handler, *, base_url="https://example.test/v1"): + return openai.OpenAI( + api_key="test-key", + base_url=base_url, + max_retries=0, + http_client=httpx.Client(transport=httpx.MockTransport(handler)), + ) + + +def test_vlm_config_fails_over_when_primary_model_is_not_found(): + calls = {"primary": 0, "backup": 0} + + def primary_handler(request): + calls["primary"] += 1 + assert request.url.host == "primary.example.test" + return httpx.Response( + 404, + json={ + "error": { + "message": "The model does not exist", + "type": "invalid_request_error", + "code": "model_not_found", + } + }, + ) + + def backup_handler(request): + calls["backup"] += 1 + assert request.url.host == "backup.example.test" + return httpx.Response( + 200, + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1, + "model": "backup-model", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "backup response"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + + config = VLMConfig( + model="primary-model", + max_retries=2, + credentials=[ + { + "id": "primary", + "provider": "openai", + "api_key": "primary-key", + "api_base": "https://primary.example.test/v1", + "model": "missing-model", + }, + { + "id": "backup", + "provider": "openai", + "api_key": "backup-key", + "api_base": "https://backup.example.test/v1", + "model": "backup-model", + }, + ], + ) + instance = config.get_vlm_instance() + primary = instance._vlm_instances[0].inner + backup = instance._vlm_instances[1].inner + primary._sync_client = _openai_client(primary_handler, base_url=primary.api_base) + backup._sync_client = _openai_client(backup_handler, base_url=backup.api_base) + + assert instance.get_completion(prompt="hello") == "backup response" + assert calls == {"primary": 1, "backup": 1} + assert instance.active_credential_id == "backup" + + +def test_openai_500_context_length_code_is_terminal_and_does_not_trip_breaker(): + calls = 0 + + def handler(_request): + nonlocal calls + calls += 1 + return httpx.Response( + 500, + json={ + "error": { + "message": "Maximum context length exceeded", + "type": "server_error", + "code": "context_length_exceeded", + } + }, + ) + + config = VLMConfig( + model="test-model", + provider="openai", + api_key="test-key", + max_retries=3, + ) + instance = config.get_vlm_instance() + instance.inner._sync_client = _openai_client(handler) + + with pytest.raises(ModelError) as raised: + instance.get_completion(prompt="too long") + + assert raised.value.kind is ModelErrorKind.INPUT_TOO_LARGE + assert calls == 1 + + breaker = CircuitBreaker(failure_threshold=1) + breaker.record_failure(raised.value) + breaker.check() + assert breaker._failure_count == 0 From 0fde1b511e5f28b8293bb0c853de78ae59bd695d Mon Sep 17 00:00:00 2001 From: "zhanghe.1216" <@bytedance.com> Date: Wed, 14 Oct 2026 22:03:27 +0800 Subject: [PATCH 13/16] =?UTF-8?q?fix(cli):=20=E5=AF=B9=E9=BD=90=E7=AD=89?= =?UTF-8?q?=E5=BE=85=E5=91=BD=E4=BB=A4=20timeout=20=E7=A7=BB=E9=99=A4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/cli/conftest.py | 6 ------ tests/cli/test_cli_filesystem.py | 5 ++++- tests/cli/test_cli_system.py | 13 ------------- 3 files changed, 4 insertions(+), 20 deletions(-) diff --git a/tests/cli/conftest.py b/tests/cli/conftest.py index 35ff6a0b78..cd2f69fad8 100644 --- a/tests/cli/conftest.py +++ b/tests/cli/conftest.py @@ -298,9 +298,7 @@ def _write_cli_config(): CLI_CONFIG_PATH = _write_cli_config() -ADD_RESOURCE_WAIT_TIMEOUT = "300" ADD_RESOURCE_COMMAND_TIMEOUT = 360 -WRITE_WAIT_TIMEOUT = "300" WRITE_COMMAND_TIMEOUT = 360 @@ -493,8 +491,6 @@ def ov_add_resource(path, to_uri, *extra_args, attempts=15, interval=10): to_uri, *extra_args, "--wait", - "--timeout", - ADD_RESOURCE_WAIT_TIMEOUT, "-o", "json", ] @@ -572,8 +568,6 @@ def ov_write(uri, content, *extra_args): content, *extra_args, "--wait", - "--timeout", - WRITE_WAIT_TIMEOUT, "-o", "json", ], diff --git a/tests/cli/test_cli_filesystem.py b/tests/cli/test_cli_filesystem.py index 1a9cee45b5..9649ba9c43 100644 --- a/tests/cli/test_cli_filesystem.py +++ b/tests/cli/test_cli_filesystem.py @@ -127,7 +127,10 @@ def test_mv_directory(self, test_dir_uri): f.write("move me") temp_path = f.name try: - ov_add_resource(temp_path, src_uri) + add_result = ov_add_resource(temp_path, src_uri) + assert add_result["exit_code"] == 0, ( + f"add-resource for mv failed: {add_result['stderr'][:300]}" + ) finally: os.unlink(temp_path) r = ov_mv(src_uri, dst_uri) diff --git a/tests/cli/test_cli_system.py b/tests/cli/test_cli_system.py index 46d6bc43f9..029716cc0e 100644 --- a/tests/cli/test_cli_system.py +++ b/tests/cli/test_cli_system.py @@ -31,16 +31,3 @@ def test_status(self): assert data is not None, f"status should return JSON, got stdout: {r['stdout'][:200]}" assert data.get("ok") is True, f"Expected ok=true, got {data.get('ok')}" assert "result" in data, "'result' field should exist in status response" - - -class TestSystemWait: - @pytest.mark.skip( - reason="CLI --timeout passed as query param but server only reads from request body, " - "causing infinite wait when queue is not empty. " - "Re-enable after CLI/Server timeout parameter mismatch is fixed." - ) - def test_wait_with_timeout(self): - r = ov(["wait", "--timeout", "30", "-o", "json"], timeout=120) - assert r["exit_code"] == 0, ( - f"ov wait should exit 0, got {r['exit_code']}: {r['stderr'][:300]}" - ) From 0ff397db56d75a23d6d992cdf3b9e0672253d77f Mon Sep 17 00:00:00 2001 From: "zhanghe.1216" <@bytedance.com> Date: Thu, 15 Oct 2026 11:56:03 +0800 Subject: [PATCH 14/16] =?UTF-8?q?fix(vikingbot):=20=E4=BF=9D=E7=95=99?= =?UTF-8?q?=E6=A8=A1=E5=9E=8B=E8=BE=93=E5=85=A5=E8=BF=87=E9=95=BF=E9=94=99?= =?UTF-8?q?=E8=AF=AF=E8=AF=AD=E4=B9=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- bot/tests/test_compile_model_errors.py | 102 +++++++++++++++++++++++++ bot/vikingbot/compile/pipeline_io.py | 7 ++ bot/vikingbot/providers/base.py | 3 + bot/vikingbot/providers/vlm_adapter.py | 3 + 4 files changed, 115 insertions(+) create mode 100644 bot/tests/test_compile_model_errors.py diff --git a/bot/tests/test_compile_model_errors.py b/bot/tests/test_compile_model_errors.py new file mode 100644 index 0000000000..89d42050c8 --- /dev/null +++ b/bot/tests/test_compile_model_errors.py @@ -0,0 +1,102 @@ +from collections import Counter +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from vikingbot.compile.ops import reduce as reduce_op +from vikingbot.compile.pipeline_io import JsonModel, ModelCallError +from vikingbot.compile.plan import Node, Transform +from vikingbot.compile.results import FileDraft, FileResponse, Group, Record +from vikingbot.providers.vlm_adapter import VLMProviderAdapter + +from openviking.models.model_error import ModelError, ModelErrorKind, model_error_kind + + +@pytest.mark.asyncio +async def test_compile_preserves_vlm_input_too_large_error_kind(): + error = ModelError( + ModelErrorKind.INPUT_TOO_LARGE, + "Error code: 400 - context_length_exceeded", + provider="openai", + status_code=400, + code="context_length_exceeded", + ) + vlm = SimpleNamespace( + provider="openai", + thinking=None, + get_completion_async=AsyncMock(side_effect=error), + ) + provider = VLMProviderAdapter( + vlm, + default_model="test-model", + langfuse_client=SimpleNamespace(enabled=False, _client=None), + ) + files = SimpleNamespace(put=AsyncMock()) + model = JsonModel( + provider, + "test-model", + 0.0, + files, + SimpleNamespace(merge_input_chars=16_000, pipeline_call_seconds=5), + {}, + Counter(), + ) + + with pytest.raises(ModelCallError) as raised: + await model.call("reduce", [], [], max_tokens=128) + + assert model_error_kind(raised.value) is ModelErrorKind.INPUT_TOO_LARGE + assert vlm.get_completion_async.await_count == 1 + + +@pytest.mark.asyncio +async def test_reduce_compresses_after_preserved_input_too_large_error(monkeypatch): + original = Record("original", "payload/original", ["e1"], "original", {}, []) + compressed = Record("compressed", "payload/compressed", ["e1"], "compressed", {}, []) + overflow = ModelCallError("reduce: context_length_exceeded") + overflow.__cause__ = ModelError( + ModelErrorKind.INPUT_TOO_LARGE, + "Error code: 400 - context_length_exceeded", + ) + response = FileResponse( + files=[FileDraft(path="result.md", content="summary", inputs=["compressed"])] + ) + ask = AsyncMock(side_effect=[overflow, response]) + runtime = SimpleNamespace( + contract=SimpleNamespace( + reduce=Transform( + instructions="Summarize the inputs.", + output="files", + execution="direct", + ), + distinguish={}, + ), + target="viking://resources/output", + evidence={"e1": {"uri": "viking://resources/source.md"}}, + catalog={}, + records={"original": original}, + system="system", + files=SimpleNamespace(), + model=SimpleNamespace(ask=ask), + ) + payload = AsyncMock(side_effect=lambda _runtime, record: {"id": record.record_id}) + pack = AsyncMock(return_value=[[original]]) + combine = AsyncMock(return_value=[compressed]) + save_files = AsyncMock(return_value=["viking://resources/output/result.md"]) + monkeypatch.setattr(reduce_op.common, "payload", payload) + monkeypatch.setattr(reduce_op.common, "pack", pack) + monkeypatch.setattr(reduce_op.common, "transform", combine) + monkeypatch.setattr(reduce_op.file_ops, "save_files", save_files) + + result = await reduce_op.reduce_group( + runtime, + Node("summary", "reduce", "groups", task="reduce"), + Group("group-1", [original]), + ) + + assert result == ["viking://resources/output/result.md"] + assert ask.await_count == 2 + assert ask.await_args_list[0].args[2]["inputs"] == [{"id": "original"}] + assert ask.await_args_list[1].args[2]["inputs"] == [{"id": "compressed"}] + pack.assert_awaited_once() + combine.assert_awaited_once() diff --git a/bot/vikingbot/compile/pipeline_io.py b/bot/vikingbot/compile/pipeline_io.py index 3aef9ea87e..c51d185b6a 100644 --- a/bot/vikingbot/compile/pipeline_io.py +++ b/bot/vikingbot/compile/pipeline_io.py @@ -16,6 +16,7 @@ from loguru import logger from pydantic import BaseModel, ValidationError +from openviking.models.model_error import ModelError from openviking.utils.token_estimation import estimate_text_tokens from openviking_cli.exceptions import NotFoundError from vikingbot.compile.hashing import digest @@ -318,6 +319,12 @@ async def call(self, stage, messages, tools, *, max_tokens): self.usage[name] = self.usage.get(name, 0) + value self.metrics[f"{stage}_{name}"] += value if response.finish_reason == "error": + error_kind = getattr(response, "error_kind", None) + if error_kind is not None: + raise ModelError( + error_kind, + response.content or "Model transport failure; see provider diagnostics", + ) raise OSError( response.content or "Model transport failure; see provider diagnostics" ) diff --git a/bot/vikingbot/providers/base.py b/bot/vikingbot/providers/base.py index 159b1b1c90..17365eb949 100644 --- a/bot/vikingbot/providers/base.py +++ b/bot/vikingbot/providers/base.py @@ -7,6 +7,8 @@ import json_repair +from openviking.models.model_error import ModelErrorKind + @dataclass class ToolCallRequest: @@ -27,6 +29,7 @@ class LLMResponse: finish_reason: str = "stop" usage: dict[str, int] = field(default_factory=dict) reasoning_content: str | None = None # Kimi, DeepSeek-R1 etc. + error_kind: ModelErrorKind | None = None @property def has_tool_calls(self) -> bool: diff --git a/bot/vikingbot/providers/vlm_adapter.py b/bot/vikingbot/providers/vlm_adapter.py index bd5d0fe49a..f47575a32c 100644 --- a/bot/vikingbot/providers/vlm_adapter.py +++ b/bot/vikingbot/providers/vlm_adapter.py @@ -15,6 +15,7 @@ from loguru import logger +from openviking.models.model_error import model_error_kind from openviking.models.vlm.backends.volcengine_vlm import build_volcengine_request_headers from openviking.utils.model_retry import is_retryable_rate_limit_error, rate_limit_retry_delay from openviking.utils.multimodal import redact_image_data_urls @@ -336,6 +337,7 @@ async def chat( return LLMResponse( content=f"Error calling LLM in VLM Adapter: {str(e)}", finish_reason="error", + error_kind=model_error_kind(e), ) async def chat_stream( @@ -389,6 +391,7 @@ async def chat_stream( response=LLMResponse( content=f"Error calling LLM in VLM Adapter stream: {str(exc)}", finish_reason="error", + error_kind=model_error_kind(exc), ), ) From 67a0a6015fc66c9c2d3752c5ed5d6bb1552e63ab Mon Sep 17 00:00:00 2001 From: qin-ctx Date: Thu, 15 Oct 2026 14:23:16 +0800 Subject: [PATCH 15/16] =?UTF-8?q?fix(models):=20=E4=BF=9D=E7=95=99?= =?UTF-8?q?=E8=BE=93=E5=85=A5=E8=B6=85=E8=BF=87=E6=89=B9=E6=AC=A1=E5=AE=B9?= =?UTF-8?q?=E9=87=8F=E7=9A=84=20500=20=E9=94=99=E8=AF=AF=E5=88=86=E7=B1=BB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- openviking/models/model_error.py | 20 +++++++---- tests/unit/test_retry_vlm.py | 60 +++++++++++++++++++++++--------- 2 files changed, 58 insertions(+), 22 deletions(-) diff --git a/openviking/models/model_error.py b/openviking/models/model_error.py index de2b3f5c9e..dccbe36476 100644 --- a/openviking/models/model_error.py +++ b/openviking/models/model_error.py @@ -6,6 +6,7 @@ import functools import inspect +import re from enum import Enum from typing import Any, Callable, TypeVar, cast @@ -218,11 +219,18 @@ def _sdk_kind(error: BaseException) -> tuple[ModelErrorKind, str] | None: "NotFoundError": ModelErrorKind.MODEL_NOT_FOUND, "InternalServerError": ModelErrorKind.TRANSIENT_SERVICE_ERROR, } - kind = ( - _kind_from_code(error) - or class_kinds.get(name) - or _kind_from_status(_status_code(error)) - ) + kind = _kind_from_code(error) + if kind is None and _status_code(error) == 500: + # llama.cpp reports inputs above its physical batch capacity as + # HTTP 500 without a business code. Retrying the same input cannot help. + message = _body(error).get("message") + if isinstance(message, str) and re.fullmatch( + r"input \(\d+ tokens\) is too large to process\. " + r"increase the physical batch size \(current batch size: \d+\)", + message, + ): + kind = ModelErrorKind.INPUT_TOO_LARGE + kind = kind or class_kinds.get(name) or _kind_from_status(_status_code(error)) return kind or ModelErrorKind.UNKNOWN, provider if module.startswith("volcenginesdkarkruntime"): @@ -311,7 +319,7 @@ def model_error_from_status( def model_error_from_exception(error: BaseException) -> ModelError | None: - """Convert a recognized SDK/API exception without inspecting free-form text.""" + """Convert a recognized SDK/API exception at the provider boundary.""" for item in _exception_chain(error): if isinstance(item, ModelError): return item diff --git a/tests/unit/test_retry_vlm.py b/tests/unit/test_retry_vlm.py index 9e55db8de5..5f18ab0f26 100644 --- a/tests/unit/test_retry_vlm.py +++ b/tests/unit/test_retry_vlm.py @@ -11,7 +11,7 @@ from openviking.models.model_error import ModelError, ModelErrorKind from openviking.models.vlm.base import MultiCredentialVLM, VLMBase, VLMFactory from openviking.models.vlm.retry_vlm import RetryVLM -from openviking.utils.circuit_breaker import CircuitBreaker +from openviking.utils.circuit_breaker import CircuitBreaker, CircuitBreakerOpen from openviking.utils.model_retry import model_retry_scope from openviking_cli.utils.config.vlm_config import VLMConfig @@ -326,22 +326,46 @@ def backup_handler(request): assert instance.active_credential_id == "backup" -def test_openai_500_context_length_code_is_terminal_and_does_not_trip_breaker(): +@pytest.mark.parametrize( + ("error_body", "expected_kind", "expected_attempts"), + [ + ( + { + "message": "Maximum context length exceeded", + "type": "server_error", + "code": "context_length_exceeded", + }, + ModelErrorKind.INPUT_TOO_LARGE, + 1, + ), + ( + { + "message": ( + "input (8525 tokens) is too large to process. increase the physical batch size " + "(current batch size: 2048)" + ), + "type": "server_error", + "code": 500, + }, + ModelErrorKind.INPUT_TOO_LARGE, + 1, + ), + ( + {"message": "Internal server error", "type": "server_error", "code": 500}, + ModelErrorKind.TRANSIENT_SERVICE_ERROR, + 4, + ), + ], +) +def test_openai_500_retry_and_breaker_follow_error_kind( + error_body, expected_kind, expected_attempts +): calls = 0 def handler(_request): nonlocal calls calls += 1 - return httpx.Response( - 500, - json={ - "error": { - "message": "Maximum context length exceeded", - "type": "server_error", - "code": "context_length_exceeded", - } - }, - ) + return httpx.Response(500, json={"error": error_body}) config = VLMConfig( model="test-model", @@ -355,10 +379,14 @@ def handler(_request): with pytest.raises(ModelError) as raised: instance.get_completion(prompt="too long") - assert raised.value.kind is ModelErrorKind.INPUT_TOO_LARGE - assert calls == 1 + assert raised.value.kind is expected_kind + assert calls == expected_attempts breaker = CircuitBreaker(failure_threshold=1) breaker.record_failure(raised.value) - breaker.check() - assert breaker._failure_count == 0 + if expected_kind is ModelErrorKind.INPUT_TOO_LARGE: + breaker.check() + assert breaker._failure_count == 0 + else: + with pytest.raises(CircuitBreakerOpen): + breaker.check() From 3c14fa923b369b666f81847c0bc473a2eea95bac Mon Sep 17 00:00:00 2001 From: qin-ctx Date: Thu, 15 Oct 2026 14:42:44 +0800 Subject: [PATCH 16/16] =?UTF-8?q?test(models):=20=E7=B2=BE=E7=AE=80?= =?UTF-8?q?=E9=87=8D=E5=A4=8D=E6=B5=8B=E8=AF=95=E5=B9=B6=E4=BF=9D=E7=95=99?= =?UTF-8?q?=E5=85=B3=E9=94=AE=E5=9B=9E=E5=BD=92=E8=A6=86=E7=9B=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- bot/tests/test_compile_model_errors.py | 102 ----------- tests/unit/test_model_retry.py | 51 ++++-- tests/unit/test_retry_embedder.py | 30 +-- tests/unit/test_retry_vlm.py | 173 +++--------------- .../unit/test_vikingbot_vlm_adapter_retry.py | 23 ++- 5 files changed, 88 insertions(+), 291 deletions(-) delete mode 100644 bot/tests/test_compile_model_errors.py diff --git a/bot/tests/test_compile_model_errors.py b/bot/tests/test_compile_model_errors.py deleted file mode 100644 index 89d42050c8..0000000000 --- a/bot/tests/test_compile_model_errors.py +++ /dev/null @@ -1,102 +0,0 @@ -from collections import Counter -from types import SimpleNamespace -from unittest.mock import AsyncMock - -import pytest -from vikingbot.compile.ops import reduce as reduce_op -from vikingbot.compile.pipeline_io import JsonModel, ModelCallError -from vikingbot.compile.plan import Node, Transform -from vikingbot.compile.results import FileDraft, FileResponse, Group, Record -from vikingbot.providers.vlm_adapter import VLMProviderAdapter - -from openviking.models.model_error import ModelError, ModelErrorKind, model_error_kind - - -@pytest.mark.asyncio -async def test_compile_preserves_vlm_input_too_large_error_kind(): - error = ModelError( - ModelErrorKind.INPUT_TOO_LARGE, - "Error code: 400 - context_length_exceeded", - provider="openai", - status_code=400, - code="context_length_exceeded", - ) - vlm = SimpleNamespace( - provider="openai", - thinking=None, - get_completion_async=AsyncMock(side_effect=error), - ) - provider = VLMProviderAdapter( - vlm, - default_model="test-model", - langfuse_client=SimpleNamespace(enabled=False, _client=None), - ) - files = SimpleNamespace(put=AsyncMock()) - model = JsonModel( - provider, - "test-model", - 0.0, - files, - SimpleNamespace(merge_input_chars=16_000, pipeline_call_seconds=5), - {}, - Counter(), - ) - - with pytest.raises(ModelCallError) as raised: - await model.call("reduce", [], [], max_tokens=128) - - assert model_error_kind(raised.value) is ModelErrorKind.INPUT_TOO_LARGE - assert vlm.get_completion_async.await_count == 1 - - -@pytest.mark.asyncio -async def test_reduce_compresses_after_preserved_input_too_large_error(monkeypatch): - original = Record("original", "payload/original", ["e1"], "original", {}, []) - compressed = Record("compressed", "payload/compressed", ["e1"], "compressed", {}, []) - overflow = ModelCallError("reduce: context_length_exceeded") - overflow.__cause__ = ModelError( - ModelErrorKind.INPUT_TOO_LARGE, - "Error code: 400 - context_length_exceeded", - ) - response = FileResponse( - files=[FileDraft(path="result.md", content="summary", inputs=["compressed"])] - ) - ask = AsyncMock(side_effect=[overflow, response]) - runtime = SimpleNamespace( - contract=SimpleNamespace( - reduce=Transform( - instructions="Summarize the inputs.", - output="files", - execution="direct", - ), - distinguish={}, - ), - target="viking://resources/output", - evidence={"e1": {"uri": "viking://resources/source.md"}}, - catalog={}, - records={"original": original}, - system="system", - files=SimpleNamespace(), - model=SimpleNamespace(ask=ask), - ) - payload = AsyncMock(side_effect=lambda _runtime, record: {"id": record.record_id}) - pack = AsyncMock(return_value=[[original]]) - combine = AsyncMock(return_value=[compressed]) - save_files = AsyncMock(return_value=["viking://resources/output/result.md"]) - monkeypatch.setattr(reduce_op.common, "payload", payload) - monkeypatch.setattr(reduce_op.common, "pack", pack) - monkeypatch.setattr(reduce_op.common, "transform", combine) - monkeypatch.setattr(reduce_op.file_ops, "save_files", save_files) - - result = await reduce_op.reduce_group( - runtime, - Node("summary", "reduce", "groups", task="reduce"), - Group("group-1", [original]), - ) - - assert result == ["viking://resources/output/result.md"] - assert ask.await_count == 2 - assert ask.await_args_list[0].args[2]["inputs"] == [{"id": "original"}] - assert ask.await_args_list[1].args[2]["inputs"] == [{"id": "compressed"}] - pack.assert_awaited_once() - combine.assert_awaited_once() diff --git a/tests/unit/test_model_retry.py b/tests/unit/test_model_retry.py index ff4ab99765..35945c47ce 100644 --- a/tests/unit/test_model_retry.py +++ b/tests/unit/test_model_retry.py @@ -119,16 +119,23 @@ def test_all_credentials_failed_prefers_transient_over_auth(): assert err.kind is ModelErrorKind.TRANSIENT_SERVICE_ERROR -def test_retry_sync_retries_transient_error_until_success(): +def test_retry_sync_retries_transient_error_until_success(monkeypatch): attempts = {"count": 0} + monkeypatch.setattr("openviking.utils.model_retry.time.sleep", lambda _delay: None) + sdk_error = openai.InternalServerError( + "Internal server error", + response=httpx.Response(500, request=httpx.Request("POST", "https://example.test")), + body={"error": {"code": 500, "type": "server_error", "message": "Internal server error"}}, + ) + @model_error_boundary def _call(): attempts["count"] += 1 if attempts["count"] < 3: - raise RuntimeError("429 TooManyRequests") + raise sdk_error return "ok" - assert retry_sync(_call, max_retries=3) == "ok" + assert retry_sync(_call, max_retries=3, model_errors_only=True) == "ok" assert attempts["count"] == 3 @@ -329,30 +336,44 @@ def test_20015_inside_request_id_is_not_input_too_large(): @pytest.mark.parametrize( - "message", + "body", [ - "expected maxLength: 50000, actual: 75000", - "Error code: 400 - InvalidParameter: Total tokens of multi-modal content " - "and text exceed max message tokens.", + {"code": "context_length_exceeded", "message": "Maximum context length exceeded"}, + { + "code": 500, + "type": "server_error", + "message": ( + "input (8525 tokens) is too large to process. increase the physical batch size " + "(current batch size: 2048)" + ), + }, ], ) -def test_input_too_large_does_not_retry_or_trip_breaker(message): +def test_input_too_large_does_not_retry_or_trip_breaker(body, monkeypatch): from openviking.utils.circuit_breaker import CircuitBreaker breaker = CircuitBreaker() attempts = {"count": 0} + monkeypatch.setattr("openviking.utils.model_retry.time.sleep", lambda _delay: None) + sdk_error = openai.InternalServerError( + "Input rejected", + response=httpx.Response(500, request=httpx.Request("POST", "https://example.test")), + body={"error": body}, + ) + @model_error_boundary def _call(): attempts["count"] += 1 - error = RuntimeError(message) - breaker.record_failure(error) - raise error + raise sdk_error - with pytest.raises(RuntimeError) as raised: - retry_sync(_call, max_retries=5) + # Repeated bad inputs must not block a subsequent valid request. + for _ in range(5): + with pytest.raises(ModelError) as raised: + retry_sync(_call, max_retries=3, model_errors_only=True) + assert raised.value.kind is ModelErrorKind.INPUT_TOO_LARGE + breaker.record_failure(raised.value) - assert str(raised.value) == message - assert attempts["count"] == 1 + assert attempts["count"] == 5 breaker.check() diff --git a/tests/unit/test_retry_embedder.py b/tests/unit/test_retry_embedder.py index 18b5231c63..9851e7138c 100644 --- a/tests/unit/test_retry_embedder.py +++ b/tests/unit/test_retry_embedder.py @@ -130,15 +130,20 @@ def test_retry_embedder_does_not_retry_permanent_or_auth_errors(status_code): assert inner.calls == 1 -def test_retry_embedder_unknown_api_error_exhausts_configured_budget(): +def test_unknown_api_error_exhausts_retry_budget_without_credential_failover(): errors = [ModelError(ModelErrorKind.UNKNOWN, "unclassified API error") for _ in range(3)] - inner = FakeDenseEmbedder(errors, max_retries=2) + primary = FakeDenseEmbedder(errors, max_retries=2) + backup = FakeDenseEmbedder([[0.3, 0.4]], max_retries=2) + failover = FailoverEmbedder( + [RetryEmbedder(primary), RetryEmbedder(backup)], ["primary", "backup"] + ) with pytest.raises(ModelError) as raised: - RetryEmbedder(inner).embed("query") + failover.embed("query") assert raised.value is errors[-1] - assert inner.calls == 3 + assert primary.calls == 3 + assert backup.calls == 0 def test_retry_embedder_sparse_dimension_remains_unsupported(): @@ -236,23 +241,6 @@ def test_failover_runs_after_current_credential_retry_exhaustion(): assert backup.calls == 1 -def test_unknown_error_retries_without_expanding_to_credential_failover(): - errors = [ModelError(ModelErrorKind.UNKNOWN, "unclassified API error") for _ in range(3)] - primary = FakeDenseEmbedder(errors, max_retries=2) - backup = FakeDenseEmbedder([[0.3, 0.4]], max_retries=2) - failover = FailoverEmbedder( - [RetryEmbedder(primary), RetryEmbedder(backup)], - ["primary", "backup"], - ) - - with pytest.raises(ModelError) as raised: - failover.embed("query") - - assert raised.value is errors[-1] - assert primary.calls == 3 - assert backup.calls == 0 - - def test_account_state_reaches_retry_inner_and_close_is_delegated_once(): inner = FakeDenseEmbedder([[0.1, 0.2]]) embedder = RetryEmbedder(inner) diff --git a/tests/unit/test_retry_vlm.py b/tests/unit/test_retry_vlm.py index 5f18ab0f26..5b43d81499 100644 --- a/tests/unit/test_retry_vlm.py +++ b/tests/unit/test_retry_vlm.py @@ -9,9 +9,8 @@ import pytest from openviking.models.model_error import ModelError, ModelErrorKind -from openviking.models.vlm.base import MultiCredentialVLM, VLMBase, VLMFactory +from openviking.models.vlm.base import VLMBase from openviking.models.vlm.retry_vlm import RetryVLM -from openviking.utils.circuit_breaker import CircuitBreaker, CircuitBreakerOpen from openviking.utils.model_retry import model_retry_scope from openviking_cli.utils.config.vlm_config import VLMConfig @@ -211,79 +210,20 @@ def test_retry_vlm_delegates_token_tracker_and_close(): assert inner.closed is True -def test_vlm_factory_wraps_provider_with_standard_retry(): - config = {"provider": "kimi", "api_key": "key", "model": "kimi-code"} +def test_vlm_config_switches_model_after_primary_not_found(): + calls = [] - standard = VLMFactory.create(config) - - assert isinstance(standard, RetryVLM) - assert standard.inner.__class__.__name__ == "KimiVLM" - - -def test_vlm_config_wraps_each_credential_inside_failover(): - config = VLMConfig( - model="kimi-code", - credentials=[ - { - "id": "primary", - "provider": "kimi", - "api_key": "primary-key", - }, - { - "id": "backup", - "provider": "glm", - "api_key": "backup-key", - "model": "glm-4.6v", - }, - ], - ) - - instance = config.get_vlm_instance() - - assert isinstance(instance, MultiCredentialVLM) - assert all(isinstance(child, RetryVLM) for child in instance._vlm_instances) - assert [child.inner.__class__.__name__ for child in instance._vlm_instances] == [ - "KimiVLM", - "GLMVLM", - ] - - -def _openai_client(handler, *, base_url="https://example.test/v1"): - return openai.OpenAI( - api_key="test-key", - base_url=base_url, - max_retries=0, - http_client=httpx.Client(transport=httpx.MockTransport(handler)), - ) - - -def test_vlm_config_fails_over_when_primary_model_is_not_found(): - calls = {"primary": 0, "backup": 0} - - def primary_handler(request): - calls["primary"] += 1 - assert request.url.host == "primary.example.test" - return httpx.Response( - 404, - json={ - "error": { - "message": "The model does not exist", - "type": "invalid_request_error", - "code": "model_not_found", - } - }, - ) - - def backup_handler(request): - calls["backup"] += 1 + def handler(request): + calls.append(request.url.host) + if request.url.host == "primary.example.test": + return httpx.Response( + 404, + json={"error": {"code": "model_not_found", "message": "Model does not exist"}}, + ) assert request.url.host == "backup.example.test" return httpx.Response( 200, json={ - "id": "chatcmpl-test", - "object": "chat.completion", - "created": 1, - "model": "backup-model", "choices": [ { "index": 0, @@ -291,24 +231,21 @@ def backup_handler(request): "finish_reason": "stop", } ], - "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, }, ) config = VLMConfig( - model="primary-model", + model="missing-model", + provider="openai", max_retries=2, credentials=[ { "id": "primary", - "provider": "openai", "api_key": "primary-key", "api_base": "https://primary.example.test/v1", - "model": "missing-model", }, { "id": "backup", - "provider": "openai", "api_key": "backup-key", "api_base": "https://backup.example.test/v1", "model": "backup-model", @@ -316,77 +253,17 @@ def backup_handler(request): ], ) instance = config.get_vlm_instance() - primary = instance._vlm_instances[0].inner - backup = instance._vlm_instances[1].inner - primary._sync_client = _openai_client(primary_handler, base_url=primary.api_base) - backup._sync_client = _openai_client(backup_handler, base_url=backup.api_base) - - assert instance.get_completion(prompt="hello") == "backup response" - assert calls == {"primary": 1, "backup": 1} - assert instance.active_credential_id == "backup" - - -@pytest.mark.parametrize( - ("error_body", "expected_kind", "expected_attempts"), - [ - ( - { - "message": "Maximum context length exceeded", - "type": "server_error", - "code": "context_length_exceeded", - }, - ModelErrorKind.INPUT_TOO_LARGE, - 1, - ), - ( - { - "message": ( - "input (8525 tokens) is too large to process. increase the physical batch size " - "(current batch size: 2048)" - ), - "type": "server_error", - "code": 500, - }, - ModelErrorKind.INPUT_TOO_LARGE, - 1, - ), - ( - {"message": "Internal server error", "type": "server_error", "code": 500}, - ModelErrorKind.TRANSIENT_SERVICE_ERROR, - 4, - ), - ], -) -def test_openai_500_retry_and_breaker_follow_error_kind( - error_body, expected_kind, expected_attempts -): - calls = 0 - - def handler(_request): - nonlocal calls - calls += 1 - return httpx.Response(500, json={"error": error_body}) + for credential in instance._vlm_instances: + credential.inner._sync_client = openai.OpenAI( + api_key="test-key", + base_url=credential.inner.api_base, + max_retries=0, + http_client=httpx.Client(transport=httpx.MockTransport(handler)), + ) - config = VLMConfig( - model="test-model", - provider="openai", - api_key="test-key", - max_retries=3, - ) - instance = config.get_vlm_instance() - instance.inner._sync_client = _openai_client(handler) - - with pytest.raises(ModelError) as raised: - instance.get_completion(prompt="too long") - - assert raised.value.kind is expected_kind - assert calls == expected_attempts - - breaker = CircuitBreaker(failure_threshold=1) - breaker.record_failure(raised.value) - if expected_kind is ModelErrorKind.INPUT_TOO_LARGE: - breaker.check() - assert breaker._failure_count == 0 - else: - with pytest.raises(CircuitBreakerOpen): - breaker.check() + try: + assert instance.get_completion(prompt="hello") == "backup response" + assert calls == ["primary.example.test", "backup.example.test"] + assert instance.active_credential_id == "backup" + finally: + instance.close() diff --git a/tests/unit/test_vikingbot_vlm_adapter_retry.py b/tests/unit/test_vikingbot_vlm_adapter_retry.py index 4bf3c9790b..b05a399329 100644 --- a/tests/unit/test_vikingbot_vlm_adapter_retry.py +++ b/tests/unit/test_vikingbot_vlm_adapter_retry.py @@ -1,8 +1,11 @@ +from collections import Counter from types import SimpleNamespace +from unittest.mock import AsyncMock import httpx import pytest import vikingbot.providers.vlm_adapter as vlm_adapter +from vikingbot.compile.pipeline_io import JsonModel, ModelCallError from vikingbot.providers.vlm_adapter import VLMProviderAdapter from volcenginesdkarkruntime._exceptions import ( ArkAPITimeoutError, @@ -10,6 +13,7 @@ ArkRateLimitError, ) +from openviking.models.model_error import ModelError, ModelErrorKind, model_error_kind from openviking.models.vlm.backends.openai_vlm import OpenAIVLM from openviking.models.vlm.backends.volcengine_vlm import ( VOLCENGINE_CLIENT_REQUEST_ID, @@ -121,19 +125,28 @@ async def _sleep(delay: float): @pytest.mark.asyncio -async def test_chat_does_not_retry_errors_without_rate_limit_markers(monkeypatch): +async def test_chat_preserves_input_error_for_compile_without_retry(monkeypatch): async def _sleep(_delay: float): raise AssertionError("non-retryable errors must not sleep/retry") monkeypatch.setattr(vlm_adapter.asyncio, "sleep", _sleep) - fake_vlm = _FakeVLM([RuntimeError("AuthenticationError Unauthorized")]) + fake_vlm = _FakeVLM([ModelError(ModelErrorKind.INPUT_TOO_LARGE, "Input rejected")]) adapter = VLMProviderAdapter(fake_vlm, "test-model", langfuse_client=_DisabledLangfuse()) + model = JsonModel( + adapter, + "test-model", + 0.0, + SimpleNamespace(put=AsyncMock()), + SimpleNamespace(merge_input_chars=16_000, pipeline_call_seconds=5), + {}, + Counter(), + ) - response = await adapter.chat(messages=[{"role": "user", "content": "hello"}]) + with pytest.raises(ModelCallError) as raised: + await model.call("reduce", [{"role": "user", "content": "hello"}], [], max_tokens=128) - assert response.finish_reason == "error" - assert "AuthenticationError" in response.content + assert model_error_kind(raised.value) is ModelErrorKind.INPUT_TOO_LARGE assert fake_vlm.calls == 1