Antigravity Python API を本番で落とさない — リトライ・タイムアウト・サーキットブレーカー実装ガイド

# timeouts.py — モデル種別に応じたタイムアウト戦略
from dataclasses import dataclass
import httpx
 
@dataclass(frozen=True)
class ModelTimeouts:
    connect: float
    read: float
    write: float
    pool: float
 
    def to_httpx(self) -> httpx.Timeout:
        return httpx.Timeout(
            connect=self.connect,
            read=self.read,
            write=self.write,
            pool=self.pool,
        )
 
# レイテンシ特性が違うので、同じ client を使い回さないのが重要
TIMEOUTS = {
    # Flash 系は軽量で基本的に速い。read が 30 秒超えたら異常
    "gemini-3-flash": ModelTimeouts(connect=5, read=30, write=10, pool=2),
    # Pro 系は長考することがあるので read を緩める
    "gemini-3-pro": ModelTimeouts(connect=5, read=120, write=10, pool=2),
    # Thinking モードはさらに余裕を持たせる
    "gemini-3-pro-thinking": ModelTimeouts(connect=5, read=240, write=10, pool=2),
}
 
def client_for(model: str) -> httpx.AsyncClient:
    t = TIMEOUTS.get(model)
    if t is None:
        raise ValueError(f"Unknown model: {model}")
    return httpx.AsyncClient(
        timeout=t.to_httpx(),
        limits=httpx.Limits(max_connections=100, max_keepalive_connections=20),
        # HTTP/2 を使うと多重化が効いてレイテンシが下がる
        http2=True,
    )

# handler.py — リクエスト全体のデッドラインを設ける
async def handle_user_request(user_input: str, model: str) -> str:
    async with asyncio.timeout(45):  # ユーザーを 45 秒以上待たせない
        async with client_for(model) as client:
            request = _build_request(user_input, model)
            response = await call_with_retry(
                client, request, max_attempts=3, total_budget_sec=30
            )
            return response.json()["candidates"][0]["content"]["parts"][0]["text"]

# circuit.py — シンプルな非同期サーキットブレーカー
from __future__ import annotations
import asyncio
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Awaitable, Callable, TypeVar
 
from ai_errors import RetryableAIError
 
T = TypeVar("T")
 
class State(Enum):
    CLOSED = "closed"
    OPEN = "open"
    HALF_OPEN = "half_open"
 
@dataclass
class CircuitBreaker:
    failure_threshold: int = 5          # 連続失敗この数で OPEN
    recovery_timeout_sec: float = 30.0  # OPEN 後にこの秒数経過で HALF_OPEN
    state: State = State.CLOSED
    _failure_count: int = 0
    _opened_at: float | None = None
    _lock: asyncio.Lock = field(default_factory=asyncio.Lock)
 
    async def call(self, fn: Callable[[], Awaitable[T]]) -> T:
        async with self._lock:
            if self.state is State.OPEN:
                # 回復時間が経過したか確認
                assert self._opened_at is not None
                if time.monotonic() - self._opened_at >= self.recovery_timeout_sec:
                    self.state = State.HALF_OPEN
                else:
                    raise RetryableAIError("Circuit is OPEN — fast-failing")
 
        try:
            result = await fn()
        except Exception:
            async with self._lock:
                self._failure_count += 1
                if (
                    self.state is State.HALF_OPEN
                    or self._failure_count >= self.failure_threshold
                ):
                    self.state = State.OPEN
                    self._opened_at = time.monotonic()
            raise
        else:
            async with self._lock:
                self._failure_count = 0
                self.state = State.CLOSED
                self._opened_at = None
            return result

# 実際の使用例
breaker = CircuitBreaker(failure_threshold=5, recovery_timeout_sec=30.0)
 
async def protected_call() -> str:
    async def _do() -> str:
        async with client_for("gemini-3-pro") as client:
            request = _build_request(user_input, "gemini-3-pro")
            response = await call_with_retry(client, request)
            return response.json()["candidates"][0]["content"]["parts"][0]["text"]
 
    return await breaker.call(_do)

# streaming.py — 部分応答チェックポイント付きのストリーミング
from dataclasses import dataclass, field
from typing import AsyncIterator
from google import genai
from google.genai import types
 
@dataclass
class StreamCheckpoint:
    produced: str = ""
    finished: bool = False
    error: Exception | None = field(default=None, repr=False)
 
async def stream_with_checkpoint(
    client: genai.Client, model: str, prompt: str, checkpoint: StreamCheckpoint
) -> AsyncIterator[str]:
    """ストリームを返しつつ、生成済みテキストを checkpoint に蓄積する。"""
    try:
        async for chunk in await client.aio.models.generate_content_stream(
            model=model, contents=prompt,
        ):
            text = chunk.text or ""
            checkpoint.produced += text
            yield text
        checkpoint.finished = True
    except Exception as exc:
        checkpoint.error = exc
        raise
 
async def resume_if_interrupted(
    client: genai.Client, model: str, original_prompt: str
) -> str:
    cp = StreamCheckpoint()
    try:
        async for _ in stream_with_checkpoint(client, model, original_prompt, cp):
            pass
        return cp.produced
    except Exception:
        # 中断した場合、続きだけ生成させる
        if not cp.produced:
            raise  # 1 トークンも出ていなかったので普通にリトライ
        resume_prompt = (
            f"{original_prompt}\n\n"
            f"[以下は先ほど途中まで生成した内容です。この続きから、"
            f"重複なく自然に続けて出力してください。]\n\n{cp.produced}"
        )
        cp2 = StreamCheckpoint()
        async for _ in stream_with_checkpoint(client, model, resume_prompt, cp2):
            pass
        return cp.produced + cp2.produced

# budget.py — RPM と JPY の両方で絞るリミッタ
import asyncio
import time
from dataclasses import dataclass
 
@dataclass
class BudgetBucket:
    """RPM と日次コスト上限を同時に管理するバケット。"""
    rpm_limit: int
    daily_yen_limit: float
    _tokens: float = 0.0
    _spent_today_yen: float = 0.0
    _last_refill: float = 0.0
    _lock: asyncio.Lock = None
 
    def __post_init__(self):
        self._tokens = float(self.rpm_limit)
        self._last_refill = time.monotonic()
        self._lock = asyncio.Lock()
 
    async def acquire(self, estimated_yen: float) -> None:
        """1 リクエスト分の枠と予算を確保する。待機が必要なら待つ。"""
        async with self._lock:
            # RPM の補充(1 分で rpm_limit トークンが入る)
            now = time.monotonic()
            elapsed = now - self._last_refill
            self._tokens = min(
                self.rpm_limit,
                self._tokens + elapsed * (self.rpm_limit / 60.0),
            )
            self._last_refill = now
 
            # コスト超過チェックは待っても解決しないので即 raise
            if self._spent_today_yen + estimated_yen > self.daily_yen_limit:
                raise RuntimeError(
                    f"Daily budget exceeded: "
                    f"{self._spent_today_yen:.1f} + {estimated_yen:.1f} "
                    f"> {self.daily_yen_limit:.1f}"
                )
 
            # RPM は待てば回復する
            while self._tokens < 1.0:
                sleep_for = (1.0 - self._tokens) * 60.0 / self.rpm_limit
                await asyncio.sleep(sleep_for)
                now = time.monotonic()
                elapsed = now - self._last_refill
                self._tokens = min(
                    self.rpm_limit,
                    self._tokens + elapsed * (self.rpm_limit / 60.0),
                )
                self._last_refill = now
 
            self._tokens -= 1.0
            self._spent_today_yen += estimated_yen

# idempotency.py — 重複実行を防ぐ最小セット
import hashlib
from typing import Awaitable, Callable, TypeVar
 
T = TypeVar("T")
 
class IdempotencyStore:
    """実運用では Redis や DB を使うが、概念だけを示す簡易版。"""
    def __init__(self):
        self._cache: dict[str, object] = {}
 
    async def run_once(self, key: str, fn: Callable[[], Awaitable[T]]) -> T:
        if key in self._cache:
            return self._cache[key]  # type: ignore[return-value]
        result = await fn()
        self._cache[key] = result
        return result
 
def request_idempotency_key(user_id: str, payload: str) -> str:
    return hashlib.sha256(f"{user_id}:{payload}".encode()).hexdigest()

# test_resilience.py — 本番相当の失敗注入テスト
import pytest
import respx
import httpx
from tenacity import RetryError
 
from resilient_client import call_with_retry
 
@pytest.mark.asyncio
@respx.mock
async def test_retries_on_503_then_succeeds():
    route = respx.post("https://example.test/gen").mock(
        side_effect=[
            httpx.Response(503, text="Service Unavailable"),
            httpx.Response(503, text="Service Unavailable"),
            httpx.Response(200, json={"ok": True}),
        ]
    )
    async with httpx.AsyncClient() as client:
        req = client.build_request("POST", "https://example.test/gen")
        res = await call_with_retry(client, req, max_attempts=5)
    assert res.status_code == 200
    assert route.call_count == 3
 
@pytest.mark.asyncio
@respx.mock
async def test_no_retry_on_400():
    respx.post("https://example.test/gen").mock(
        return_value=httpx.Response(400, text="Bad Request")
    )
    async with httpx.AsyncClient() as client:
        req = client.build_request("POST", "https://example.test/gen")
        with pytest.raises(Exception):
            await call_with_retry(client, req, max_attempts=5)