feat: harden agent verification and runtime contracts
Introduce two-gate checker settlement, explicit graph visibility and checkpoints, structured LLM reliability, centralized action specifications, and corresponding tests and design documentation.
This commit is contained in:
1 parent
ddfa420945
commit
cbe8f4f0c6
106 files changed
+11542
-4445
No files matched your search
@@ -0,0 +1,252 @@
|
||||
# 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
|
||||
import threading
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
|
||||
class FailureCategory(StrEnum):
|
||||
RATE_LIMIT = "rate_limit"
|
||||
PROVIDER_UNAVAILABLE = "provider_unavailable"
|
||||
TIMEOUT = "timeout"
|
||||
CONNECTION = "connection"
|
||||
AUTHENTICATION = "authentication"
|
||||
BAD_REQUEST = "bad_request"
|
||||
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", "quota", "resource_exhausted")
|
||||
_UNAVAILABLE_MARKERS = ("unavailable", "overloaded", "high demand", "service unavailable")
|
||||
_AUTH_MARKERS = ("unauthorized", "forbidden", "api key", "credential")
|
||||
|
||||
|
||||
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)
|
||||
|
||||
code = extract_status_code(error)
|
||||
message = str(error).lower()
|
||||
|
||||
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 "timed out" in message or "timeout" in message:
|
||||
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
|
||||
|
||||
def delay_for(self, attempt: int) -> float:
|
||||
"""Backoff delay before retry number ``attempt`` (1-based)."""
|
||||
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),
|
||||
}
|
||||
|
||||
|
||||
def retry_policy_for(category: FailureCategory) -> RetryPolicy:
|
||||
return _DEFAULT_RETRY_POLICIES.get(category, _DEFAULT_RETRY_POLICIES[FailureCategory.UNKNOWN])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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()
|
||||
}
|
||||
@@ -0,0 +1,194 @@
|
||||
# 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.
|
||||
|
||||
"""Unified structured-output parsing for LLM responses.
|
||||
|
||||
Single pipeline replacing the ad-hoc JSON extraction scattered across agents:
|
||||
strip thinking tags → extract the JSON payload (fenced or balanced-span) →
|
||||
parse → tolerant repair pass → optional schema validation. A failed parse
|
||||
returns a typed ``ParseFailure`` carrying the raw text and the reason — never
|
||||
a raw string masquerading as parsed data.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import json
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from artemis.llm.parser import extract_thinking_tags
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ParseFailure:
|
||||
"""A parse miss: keeps the raw text and the reason, explicitly typed."""
|
||||
|
||||
raw: str
|
||||
error: str
|
||||
|
||||
|
||||
class StructuredOutputError(ValueError):
|
||||
"""Raised when a required structured response could not be parsed."""
|
||||
|
||||
def __init__(self, failure: ParseFailure):
|
||||
super().__init__(failure.error)
|
||||
self.failure = failure
|
||||
|
||||
|
||||
def content_to_text(content: Any) -> str:
|
||||
"""Flatten a LangChain message content (str or block list) to plain text."""
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
parts: list[str] = []
|
||||
for item in content:
|
||||
if isinstance(item, str):
|
||||
parts.append(item)
|
||||
elif isinstance(item, dict) and item.get("type") != "thinking":
|
||||
text = item.get("text")
|
||||
if isinstance(text, str):
|
||||
parts.append(text)
|
||||
return "".join(parts)
|
||||
return "" if content is None else str(content)
|
||||
|
||||
|
||||
_FENCE_RE = re.compile(r"```(?:json)?\s*([\s\S]*?)```", re.IGNORECASE)
|
||||
|
||||
|
||||
def _balanced_span(text: str) -> str | None:
|
||||
"""Return the first balanced {...} or [...] span, string/escape aware."""
|
||||
start = None
|
||||
for i, ch in enumerate(text):
|
||||
if ch in "{[":
|
||||
start = i
|
||||
break
|
||||
if start is None:
|
||||
return None
|
||||
stack: list[str] = []
|
||||
in_string = False
|
||||
escaped = False
|
||||
for j in range(start, len(text)):
|
||||
ch = text[j]
|
||||
if in_string:
|
||||
if escaped:
|
||||
escaped = False
|
||||
elif ch == "\\":
|
||||
escaped = True
|
||||
elif ch == '"':
|
||||
in_string = False
|
||||
continue
|
||||
if ch == '"':
|
||||
in_string = True
|
||||
elif ch in "{[":
|
||||
stack.append(ch)
|
||||
elif ch in "}]":
|
||||
if not stack or (stack[-1], ch) not in {("{", "}"), ("[", "]")}:
|
||||
return None
|
||||
stack.pop()
|
||||
if not stack:
|
||||
return text[start : j + 1]
|
||||
return None
|
||||
|
||||
|
||||
def extract_json_candidate(text: str) -> str | None:
|
||||
"""Locate the JSON payload inside prose: fenced blocks first, then bare spans."""
|
||||
text = text.strip()
|
||||
for match in _FENCE_RE.finditer(text):
|
||||
span = _balanced_span(match.group(1).strip())
|
||||
if span:
|
||||
return span
|
||||
return _balanced_span(text)
|
||||
|
||||
|
||||
def repair_json(text: str) -> str:
|
||||
"""Remove comment and trailing-comma noise outside of strings."""
|
||||
out: list[str] = []
|
||||
i = 0
|
||||
n = len(text)
|
||||
in_string = False
|
||||
escaped = False
|
||||
while i < n:
|
||||
ch = text[i]
|
||||
if in_string:
|
||||
out.append(ch)
|
||||
if escaped:
|
||||
escaped = False
|
||||
elif ch == "\\":
|
||||
escaped = True
|
||||
elif ch == '"':
|
||||
in_string = False
|
||||
i += 1
|
||||
continue
|
||||
if ch == '"':
|
||||
in_string = True
|
||||
out.append(ch)
|
||||
i += 1
|
||||
continue
|
||||
if ch == "/" and i + 1 < n and text[i + 1] == "/":
|
||||
while i < n and text[i] != "\n":
|
||||
i += 1
|
||||
continue
|
||||
if ch == "/" and i + 1 < n and text[i + 1] == "*":
|
||||
i += 2
|
||||
while i + 1 < n and not (text[i] == "*" and text[i + 1] == "/"):
|
||||
i += 1
|
||||
i += 2
|
||||
continue
|
||||
if ch == ",":
|
||||
k = i + 1
|
||||
while k < n and text[k] in " \t\r\n":
|
||||
k += 1
|
||||
if k < n and text[k] in "]}":
|
||||
i += 1
|
||||
continue
|
||||
out.append(ch)
|
||||
i += 1
|
||||
return "".join(out)
|
||||
|
||||
|
||||
def parse_structured(text: Any, schema: type[BaseModel] | None = None) -> Any | ParseFailure:
|
||||
"""Parse an LLM response into JSON data (optionally schema-validated).
|
||||
|
||||
Returns the parsed object (or a validated ``schema`` instance), or a
|
||||
``ParseFailure`` describing why the text could not be parsed. Callers must
|
||||
check ``isinstance(result, ParseFailure)`` — the raw text is never
|
||||
silently passed through as a result.
|
||||
"""
|
||||
if not isinstance(text, str):
|
||||
text = content_to_text(text)
|
||||
if not text or not text.strip():
|
||||
return ParseFailure(raw=text or "", error="empty response")
|
||||
|
||||
cleaned, _ = extract_thinking_tags(text)
|
||||
candidate = extract_json_candidate(cleaned)
|
||||
if candidate is None:
|
||||
return ParseFailure(raw=text, error="no JSON object or array found in response")
|
||||
|
||||
last_error = ""
|
||||
for variant in (candidate, repair_json(candidate)):
|
||||
try:
|
||||
data = json.loads(variant)
|
||||
except json.JSONDecodeError as decode_error:
|
||||
last_error = str(decode_error)
|
||||
continue
|
||||
if schema is not None:
|
||||
try:
|
||||
return schema.model_validate(data)
|
||||
except ValidationError as validation_error:
|
||||
return ParseFailure(raw=text, error=f"schema validation failed: {validation_error}")
|
||||
return data
|
||||
return ParseFailure(raw=text, error=f"invalid JSON: {last_error}")
|
||||
Reference in new issue
Block a user