chenmingxuan
e1ea099592
Install and update artemis-use with task diagnostics and plugin guidance. Preserve cross-device file receipts, review canonical reports, guard stage transitions and reduce small history summaries. Improve exploration lifecycle, request recovery and browser live connections.
309 lines
12 KiB
Python
309 lines
12 KiB
Python
# Copyright 2026 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Failure classification, retry policies, and circuit breaking for LLM calls.
|
|
|
|
This is the single source of truth for how ARTEMIS reacts to a failed model
|
|
call. Every provider/transport exception is converted into a ``Failure``
|
|
carrying explicit recovery decisions, and every retry loop in the codebase is
|
|
expected to consult ``retry_policy_for`` instead of inventing its own rules.
|
|
|
|
The taxonomy was promoted from ``artemis.agents.video_analyzer.reliability``
|
|
(which now extends it with media-specific categories) so that the same
|
|
classification and circuit-breaker quality applies to every agent call, not
|
|
just the video pipeline.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from enum import StrEnum
|
|
import random
|
|
from functools import lru_cache
|
|
import threading
|
|
import time
|
|
from typing import Any
|
|
|
|
from artemis.llm.stream_guard import ToolArgumentOverflow
|
|
|
|
|
|
class FailureCategory(StrEnum):
|
|
RATE_LIMIT = "rate_limit"
|
|
PROVIDER_UNAVAILABLE = "provider_unavailable"
|
|
TIMEOUT = "timeout"
|
|
CONNECTION = "connection"
|
|
AUTHENTICATION = "authentication"
|
|
BAD_REQUEST = "bad_request"
|
|
INVALID_OUTPUT = "invalid_output"
|
|
CANCELLED = "cancelled"
|
|
UNKNOWN = "unknown"
|
|
|
|
|
|
#: Categories that count toward circuit-breaker state: transient provider-side
|
|
#: trouble where hammering the endpoint makes things worse.
|
|
TRANSIENT_CATEGORIES: frozenset[str] = frozenset(
|
|
{
|
|
FailureCategory.RATE_LIMIT.value,
|
|
FailureCategory.PROVIDER_UNAVAILABLE.value,
|
|
FailureCategory.TIMEOUT.value,
|
|
FailureCategory.CONNECTION.value,
|
|
}
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Failure:
|
|
category: FailureCategory
|
|
retryable: bool
|
|
should_fallback: bool
|
|
|
|
|
|
def extract_status_code(error: BaseException) -> int | None:
|
|
"""Best-effort extraction of an HTTP-ish status code from an exception."""
|
|
code = getattr(error, "code", None) or getattr(error, "status_code", None)
|
|
try:
|
|
return int(code) if code is not None else None
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
|
|
_RATE_LIMIT_MARKERS = ("rate limit", "too many requests", "quota", "resource_exhausted")
|
|
_UNAVAILABLE_MARKERS = ("unavailable", "overloaded", "high demand", "service unavailable")
|
|
_AUTH_MARKERS = ("unauthorized", "forbidden", "api key", "credential")
|
|
# "deadline exceeded"/"deadline_exceeded" cover google-genai SDK timeouts that
|
|
# surface without an HTTP status code attached.
|
|
_TIMEOUT_MARKERS = ("timed out", "timeout", "deadline exceeded", "deadline_exceeded")
|
|
_TOOL_HISTORY_ERRORS = (
|
|
"no tool call found for function call output",
|
|
"no tool output found for function call",
|
|
"tool_call_id not found",
|
|
)
|
|
|
|
|
|
def classify_failure(error: BaseException) -> Failure:
|
|
"""Convert a provider/transport exception into explicit recovery decisions."""
|
|
if isinstance(error, KeyboardInterrupt) or type(error).__name__ == "CancelledError":
|
|
return Failure(FailureCategory.CANCELLED, False, False)
|
|
if isinstance(error, ToolArgumentOverflow):
|
|
return Failure(FailureCategory.INVALID_OUTPUT, True, True)
|
|
|
|
code = extract_status_code(error)
|
|
message = str(error).lower()
|
|
|
|
# Compatible Responses endpoints can report input validation inside an
|
|
# HTTP 200 stream as APIError without a status code. Retrying the same
|
|
# broken call/output history cannot repair it, even on a fallback model.
|
|
if any(marker in message for marker in _TOOL_HISTORY_ERRORS):
|
|
return Failure(FailureCategory.BAD_REQUEST, False, False)
|
|
|
|
if code == 429 or any(marker in message for marker in _RATE_LIMIT_MARKERS):
|
|
return Failure(FailureCategory.RATE_LIMIT, True, True)
|
|
if code in {500, 502, 503, 504} or any(marker in message for marker in _UNAVAILABLE_MARKERS):
|
|
return Failure(FailureCategory.PROVIDER_UNAVAILABLE, True, True)
|
|
if isinstance(error, TimeoutError) or any(marker in message for marker in _TIMEOUT_MARKERS):
|
|
return Failure(FailureCategory.TIMEOUT, True, True)
|
|
if isinstance(error, (ConnectionError, OSError)):
|
|
return Failure(FailureCategory.CONNECTION, True, True)
|
|
if code in {401, 403} or any(marker in message for marker in _AUTH_MARKERS):
|
|
# Not retryable against the same endpoint, but a fallback endpoint may
|
|
# use a different (working) credential.
|
|
return Failure(FailureCategory.AUTHENTICATION, False, True)
|
|
if code in {400, 404, 413, 415, 422}:
|
|
# Programmer/request errors: retrying is useless and falling back only
|
|
# hides the bug inside a different model's output.
|
|
return Failure(FailureCategory.BAD_REQUEST, False, False)
|
|
return Failure(FailureCategory.UNKNOWN, True, True)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Typed terminal errors
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class LLMCallError(Exception):
|
|
"""Base class for terminal LLM call failures raised by the gateway."""
|
|
|
|
def __init__(self, message: str, *, failure: Failure, cause: BaseException | None = None):
|
|
super().__init__(message)
|
|
self.failure = failure
|
|
if cause is not None:
|
|
self.__cause__ = cause
|
|
|
|
|
|
class LLMPermanentError(LLMCallError):
|
|
"""A non-retryable failure (auth, bad request). Retrying cannot help."""
|
|
|
|
|
|
class LLMExhaustedError(LLMCallError):
|
|
"""Retryable failures exhausted every attempt and the pause deadline."""
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Retry policies
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RetryPolicy:
|
|
"""How many attempts a failure category deserves and how to space them."""
|
|
|
|
max_attempts: int
|
|
base_delay: float
|
|
max_delay: float
|
|
fixed: bool = False
|
|
|
|
def delay_for(self, attempt: int) -> float:
|
|
"""Backoff delay before retry number ``attempt`` (1-based)."""
|
|
if self.fixed:
|
|
return self.base_delay
|
|
delay = min(self.max_delay, self.base_delay * (2 ** max(0, attempt - 1)))
|
|
return delay + random.uniform(0.1, 0.5)
|
|
|
|
|
|
_DEFAULT_RETRY_POLICIES: dict[FailureCategory, RetryPolicy] = {
|
|
FailureCategory.RATE_LIMIT: RetryPolicy(max_attempts=4, base_delay=10.0, max_delay=60.0),
|
|
FailureCategory.PROVIDER_UNAVAILABLE: RetryPolicy(
|
|
max_attempts=4, base_delay=5.0, max_delay=30.0
|
|
),
|
|
FailureCategory.TIMEOUT: RetryPolicy(max_attempts=3, base_delay=2.0, max_delay=15.0),
|
|
FailureCategory.CONNECTION: RetryPolicy(max_attempts=3, base_delay=2.0, max_delay=15.0),
|
|
FailureCategory.UNKNOWN: RetryPolicy(max_attempts=2, base_delay=1.0, max_delay=5.0),
|
|
# Non-retryable categories get a single attempt by definition.
|
|
FailureCategory.AUTHENTICATION: RetryPolicy(max_attempts=1, base_delay=0.0, max_delay=0.0),
|
|
FailureCategory.BAD_REQUEST: RetryPolicy(max_attempts=1, base_delay=0.0, max_delay=0.0),
|
|
FailureCategory.CANCELLED: RetryPolicy(max_attempts=1, base_delay=0.0, max_delay=0.0),
|
|
}
|
|
|
|
|
|
@lru_cache(maxsize=32)
|
|
def _configured_policy(retries: int, interval: float) -> RetryPolicy:
|
|
return RetryPolicy(retries + 1, interval, interval, fixed=True)
|
|
|
|
|
|
def retry_policy_for(category: FailureCategory | str) -> RetryPolicy:
|
|
"""Look up the retry policy for a failure category.
|
|
|
|
Accepts the generic ``FailureCategory`` or any string-valued category from
|
|
a domain extension (e.g. the video pipeline's ``VideoFailureCategory``).
|
|
Extension-only categories fall back to the UNKNOWN policy, so callers must
|
|
gate on ``failure.retryable`` before consulting attempt counts.
|
|
"""
|
|
if not isinstance(category, FailureCategory):
|
|
try:
|
|
category = FailureCategory(str(category))
|
|
except ValueError:
|
|
category = FailureCategory.UNKNOWN
|
|
if category in {
|
|
FailureCategory.AUTHENTICATION,
|
|
FailureCategory.BAD_REQUEST,
|
|
FailureCategory.CANCELLED,
|
|
}:
|
|
return _DEFAULT_RETRY_POLICIES[category]
|
|
from artemis.config.settings import settings
|
|
|
|
if category is FailureCategory.INVALID_OUTPUT:
|
|
return _configured_policy(
|
|
min(1, settings.OPENAI_MAX_RETRIES), settings.OPENAI_RETRY_INTERVAL_SECONDS
|
|
)
|
|
if category is FailureCategory.RATE_LIMIT:
|
|
return _configured_policy(
|
|
settings.OPENAI_MAX_RETRIES, settings.OPENAI_RATE_LIMIT_RETRY_SECONDS
|
|
)
|
|
return _configured_policy(settings.OPENAI_MAX_RETRIES, settings.OPENAI_RETRY_INTERVAL_SECONDS)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Circuit breaker
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class _CircuitState:
|
|
consecutive_failures: int = 0
|
|
open_until: float = 0.0
|
|
half_open_in_flight: bool = False
|
|
|
|
|
|
class CircuitBreaker:
|
|
"""Small circuit breaker keyed by provider/model identity.
|
|
|
|
Only transient categories (see ``TRANSIENT_CATEGORIES``) trip the breaker;
|
|
programmer errors and auth failures never open it. ``failure`` may be any
|
|
object with a ``category`` attribute whose value matches this module's
|
|
category strings, so domain extensions (e.g. the video pipeline's failure
|
|
type) plug in without conversion.
|
|
"""
|
|
|
|
def __init__(self, threshold: int = 3, cooldown_seconds: float = 60.0) -> None:
|
|
self.threshold = max(1, int(threshold))
|
|
self.cooldown_seconds = max(0.1, float(cooldown_seconds))
|
|
self._states: dict[str, _CircuitState] = {}
|
|
self._lock = threading.RLock()
|
|
|
|
def allow(self, key: str, *, now: float | None = None) -> bool:
|
|
current = time.time() if now is None else now
|
|
with self._lock:
|
|
state = self._states.get(key)
|
|
if state is None:
|
|
return True
|
|
if state.open_until:
|
|
if current < state.open_until or state.half_open_in_flight:
|
|
return False
|
|
# Half-open: permit exactly one trial. Other concurrent calls
|
|
# continue to bypass the primary until it succeeds or fails.
|
|
state.half_open_in_flight = True
|
|
state.consecutive_failures = max(0, self.threshold - 1)
|
|
return True
|
|
return not state.half_open_in_flight
|
|
|
|
def open_remaining(self, key: str, *, now: float | None = None) -> float:
|
|
"""Seconds until the circuit for ``key`` half-opens (0 if not open)."""
|
|
current = time.time() if now is None else now
|
|
with self._lock:
|
|
state = self._states.get(key)
|
|
if state is None or not state.open_until:
|
|
return 0.0
|
|
return max(0.0, state.open_until - current)
|
|
|
|
def record_success(self, key: str) -> None:
|
|
with self._lock:
|
|
self._states.pop(key, None)
|
|
|
|
def record_failure(self, key: str, failure: Any, *, now: float | None = None) -> None:
|
|
category = str(getattr(failure, "category", ""))
|
|
if category not in TRANSIENT_CATEGORIES:
|
|
with self._lock:
|
|
state = self._states.get(key)
|
|
if state is not None and state.half_open_in_flight:
|
|
self._states.pop(key, None)
|
|
return
|
|
current = time.time() if now is None else now
|
|
with self._lock:
|
|
state = self._states.setdefault(key, _CircuitState())
|
|
state.half_open_in_flight = False
|
|
state.consecutive_failures += 1
|
|
if state.consecutive_failures >= self.threshold:
|
|
state.open_until = current + self.cooldown_seconds
|
|
|
|
def snapshot(self) -> dict[str, dict[str, Any]]:
|
|
with self._lock:
|
|
return {
|
|
key: {
|
|
"consecutive_failures": state.consecutive_failures,
|
|
"open_until": state.open_until,
|
|
"half_open_in_flight": state.half_open_in_flight,
|
|
}
|
|
for key, state in self._states.items()
|
|
}
|