feat(ai): add guarded Claude fallback route

This commit is contained in:
ogt
2026-07-15 15:57:24 +08:00
parent f57845ba5e
commit 5c5eb48091
56 changed files with 4942 additions and 760 deletions

View File

@@ -1,5 +1,5 @@
"""
AI Rate Limiter - Gemini API 用量閥值控制
AI Rate Limiter - paid AI provider usage and spend guard
=========================================
防止最終付費備援用量暴衝;超過任一硬閘即 no-write fail closed。
@@ -79,8 +79,8 @@ COST_LIMITS = {
}
@dataclass(frozen=True)
class GeminiPricingPolicy:
"""Verified Standard text pricing used by the paid-fallback guard."""
class PaidProviderPricingPolicy:
"""Reviewed exact-model pricing used by the paid-fallback guard."""
model: str
input_usd_per_million: float
@@ -98,8 +98,8 @@ class GeminiPricingPolicy:
# Unknown models are intentionally absent and therefore fail closed. Updating a
# model requires an explicit, reviewed pricing policy update in the same change.
GEMINI_PRICING_POLICIES: dict[str, GeminiPricingPolicy] = {
"gemini-2.5-flash-lite": GeminiPricingPolicy(
GEMINI_PRICING_POLICIES: dict[str, PaidProviderPricingPolicy] = {
"gemini-2.5-flash-lite": PaidProviderPricingPolicy(
model="gemini-2.5-flash-lite",
input_usd_per_million=0.10,
output_usd_per_million=0.40,
@@ -112,11 +112,45 @@ GEMINI_PRICING_POLICIES: dict[str, GeminiPricingPolicy] = {
),
}
# Sonnet 5 had an introductory $2/$10 price through 2026-08-31. The guard
# deliberately reserves the higher published standard $3/$15 rate so a
# promotion expiry cannot under-reserve spend without a source change.
CLAUDE_PRICING_POLICIES: dict[str, PaidProviderPricingPolicy] = {
"claude-sonnet-5": PaidProviderPricingPolicy(
model="claude-sonnet-5",
input_usd_per_million=3.0,
output_usd_per_million=15.0,
source="https://platform.claude.com/docs/en/about-claude/models/overview",
version="anthropic-standard-ceiling-2026-07-15",
checked_at="2026-07-15",
),
}
def get_gemini_pricing_policy(model: str) -> GeminiPricingPolicy | None:
# Backward-compatible public type name used by existing Gemini tests/imports.
GeminiPricingPolicy = PaidProviderPricingPolicy
def get_gemini_pricing_policy(model: str) -> PaidProviderPricingPolicy | None:
"""Return a verified exact-model policy; no prefix or family matching."""
return GEMINI_PRICING_POLICIES.get(str(model).strip())
def get_claude_pricing_policy(model: str) -> PaidProviderPricingPolicy | None:
"""Return the reviewed exact-model Anthropic policy; aliases fail closed."""
return CLAUDE_PRICING_POLICIES.get(str(model).strip())
def get_paid_provider_pricing_policy(
provider: str,
model: str,
) -> PaidProviderPricingPolicy | None:
"""Resolve exact provider/model pricing without family or prefix matching."""
if provider == "gemini":
return get_gemini_pricing_policy(model)
if provider == "claude":
return get_claude_pricing_policy(model)
return None
# Redis Keys
REDIS_KEY_PREFIX = "ai_rate:"
RPM_KEY = f"{REDIS_KEY_PREFIX}rpm:{{provider}}"
@@ -156,7 +190,7 @@ class AIGenerationReservation:
estimated_completion_tokens: int
estimated_cost_usd: float
date: str
model: str = "gemini-2.5-flash-lite"
model: str = ""
pricing_source: str = ""
pricing_version: str = ""
pricing_checked_at: str = ""
@@ -201,7 +235,7 @@ local function receipt(status, reason)
terminal_state = 'pending_finalize'
end
redis.call('HSET', KEYS[9],
'schema_version', 'gemini_generation_receipt_v1',
'schema_version', 'ai_paid_generation_receipt_v1',
'provider', ARGV[9],
'receipt_id', ARGV[10],
'status', status,
@@ -235,7 +269,7 @@ end
-- One paid provider call is permitted for one canonical run_id; trace_id and
-- work_item_id remain mandatory receipt correlation. The claim and the
-- receipt/budget transition live in the same Lua execution so concurrent
-- workers cannot double-reserve or double-call Gemini.
-- workers cannot double-reserve or double-call a paid provider.
local existing_receipt_id = redis.call('GET', KEYS[10])
if existing_receipt_id then
return {0, 'idempotent_replay', existing_receipt_id}
@@ -400,7 +434,7 @@ elseif ARGV[6] == 'authentication_rejected' then
end
redis.call('HSET', KEYS[8],
'schema_version', 'ai_provider_authentication_status_v1',
'provider', 'gemini',
'provider', ARGV[10],
'authenticated', authentication_state,
'authentication_verified', success and 'true' or 'false',
'verification_status', verification_status,
@@ -497,12 +531,12 @@ class AIRateLimiter:
return hashlib.sha256(payload).hexdigest()
@staticmethod
def _estimate_gemini_usage(
def _estimate_paid_usage(
prompt: str,
max_output_tokens: int,
pricing: GeminiPricingPolicy,
pricing: PaidProviderPricingPolicy,
) -> tuple[int, int, float]:
"""Conservatively reserve Gemini tokens and spend before generation."""
"""Conservatively reserve paid-provider tokens and spend before generation."""
# One token per UTF-8 byte deliberately over-reserves mixed zh-TW/English
# prompts. Completion reserves the configured hard output ceiling.
prompt_tokens = max(1, len(prompt.encode("utf-8")))
@@ -512,23 +546,30 @@ class AIRateLimiter:
@staticmethod
def _provider_limits(provider: str) -> tuple[dict[str, int], dict[str, float]]:
"""Use one settings-backed Gemini quota source across router and health."""
"""Use one settings-backed paid-provider quota source across runtime/readback."""
limits = dict(RATE_LIMITS[provider])
costs = dict(COST_LIMITS.get(provider, {}))
if provider != "gemini":
if provider not in {"claude", "gemini"}:
return limits, costs
from src.core.config import settings
prefix = provider.upper()
limits.update(
rpm=int(settings.GEMINI_RPM_LIMIT),
daily_requests=int(settings.GEMINI_DAILY_QUOTA),
daily_tokens=int(settings.GEMINI_DAILY_TOKEN_LIMIT),
rpm=int(getattr(settings, f"{prefix}_RPM_LIMIT")),
daily_requests=int(getattr(settings, f"{prefix}_DAILY_QUOTA")),
daily_tokens=int(getattr(settings, f"{prefix}_DAILY_TOKEN_LIMIT")),
)
costs.update(
daily_cost_usd=float(settings.GEMINI_DAILY_COST_LIMIT_USD),
total_cost_usd=float(settings.GEMINI_TOTAL_COST_LIMIT_USD),
alert_threshold_usd=float(settings.GEMINI_COST_ALERT_THRESHOLD_USD),
daily_cost_usd=float(
getattr(settings, f"{prefix}_DAILY_COST_LIMIT_USD")
),
total_cost_usd=float(
getattr(settings, f"{prefix}_TOTAL_COST_LIMIT_USD")
),
alert_threshold_usd=float(
getattr(settings, f"{prefix}_COST_ALERT_THRESHOLD_USD")
),
)
return limits, costs
@@ -556,7 +597,7 @@ class AIRateLimiter:
receipt_id=receipt_id,
)
mapping = {
"schema_version": "gemini_generation_receipt_v1",
"schema_version": "ai_paid_generation_receipt_v1",
"provider": provider,
"receipt_id": receipt_id,
"status": "blocked",
@@ -625,14 +666,16 @@ class AIRateLimiter:
Redis or configuration failure denies the call. A blocked attempt still
gets a non-secret receipt when Redis is available.
"""
if provider != "gemini":
raise ValueError("reserve_generation currently supports gemini only")
if provider not in {"claude", "gemini"}:
raise ValueError("reserve_generation requires a supported paid provider")
from src.core.config import settings
from src.utils.timezone import now_taipei_iso
today = self._get_today()
model_name = str(model if model is not None else settings.GEMINI_MODEL).strip()
provider_prefix = provider.upper()
configured_model = getattr(settings, f"{provider_prefix}_MODEL")
model_name = str(model if model is not None else configured_model).strip()
trace_id, run_id, work_item_id = self._safe_receipt_context(context)
if not (trace_id and run_id and work_item_id):
receipt_id = str(uuid4())
@@ -671,7 +714,7 @@ class AIRateLimiter:
run_id,
)
receipt_id = f"gen-{identity_hash[:32]}"
pricing = get_gemini_pricing_policy(model_name)
pricing = get_paid_provider_pricing_policy(provider, model_name)
if pricing is None:
receipt_written = await self._write_blocked_generation_receipt(
provider=provider,
@@ -708,9 +751,9 @@ class AIRateLimiter:
output_limit = int(
max_output_tokens
if max_output_tokens is not None
else settings.GEMINI_MAX_OUTPUT_TOKENS
else getattr(settings, f"{provider_prefix}_MAX_OUTPUT_TOKENS")
)
prompt_tokens, completion_tokens, estimated_cost = self._estimate_gemini_usage(
prompt_tokens, completion_tokens, estimated_cost = self._estimate_paid_usage(
prompt,
output_limit,
pricing,
@@ -747,7 +790,7 @@ class AIRateLimiter:
or costs["alert_threshold_usd"] <= 0
or costs["alert_threshold_usd"] > costs["total_cost_usd"]
):
raise ValueError("invalid_gemini_cost_guard_limits")
raise ValueError(f"invalid_{provider}_cost_guard_limits")
redis = await self._get_redis()
if redis is None:
@@ -897,7 +940,10 @@ class AIRateLimiter:
from src.utils.timezone import now_taipei_iso
try:
pricing = get_gemini_pricing_policy(reservation.model)
pricing = get_paid_provider_pricing_policy(
reservation.provider,
reservation.model,
)
if (
pricing is None
or reservation.pricing_source != pricing.source
@@ -906,7 +952,7 @@ class AIRateLimiter:
or reservation.input_usd_per_million != pricing.input_usd_per_million
or reservation.output_usd_per_million != pricing.output_usd_per_million
):
raise ValueError("gemini_pricing_reservation_mismatch")
raise ValueError("paid_provider_pricing_reservation_mismatch")
redis = await self._get_redis()
if redis is None:
@@ -925,7 +971,10 @@ class AIRateLimiter:
else 0
)
daily_ttl = self._seconds_until_tomorrow()
auth_status_ttl = int(settings.GEMINI_USAGE_RECEIPT_TTL_SECONDS)
provider_prefix = reservation.provider.upper()
auth_status_ttl = int(
getattr(settings, f"{provider_prefix}_USAGE_RECEIPT_TTL_SECONDS")
)
receipt_key = GENERATION_RECEIPT_KEY.format(
provider=reservation.provider,
receipt_id=reservation.receipt_id,
@@ -962,7 +1011,13 @@ class AIRateLimiter:
str(error_code)[:120],
auth_status_ttl,
daily_ttl,
int(settings.GEMINI_FAILURE_COOLDOWN_SECONDS),
int(
getattr(
settings,
f"{provider_prefix}_FAILURE_COOLDOWN_SECONDS",
)
),
reservation.provider,
)
result_code = int(result[0])
raw_status = result[1]
@@ -1010,9 +1065,10 @@ class AIRateLimiter:
Returns:
tuple[bool, str | None]: (是否允許, 拒絕原因)
"""
if provider == "gemini":
if provider in {"claude", "gemini"}:
logger.error(
"gemini_legacy_rate_gate_rejected",
"paid_provider_legacy_rate_gate_rejected",
provider=provider,
reason="reserve_generation_required",
)
return False, "reserve_generation_required"
@@ -1318,7 +1374,9 @@ class AIRateLimiter:
"""Read public-safe authentication evidence; configuration is separate."""
from src.core.config import settings
configured = bool(settings.GEMINI_API_KEY) if provider == "gemini" else False
configured = bool(
getattr(settings, f"{provider.upper()}_API_KEY", "")
) if provider in {"claude", "gemini"} else False
base = {
"provider": provider,
"configured": configured,
@@ -1462,11 +1520,13 @@ class AIRateLimiter:
pricing_info: dict[str, Any] = {}
pricing_missing = False
if provider == "gemini":
if provider in {"claude", "gemini"}:
from src.core.config import settings
configured_model = str(settings.GEMINI_MODEL).strip()
pricing = get_gemini_pricing_policy(configured_model)
configured_model = str(
getattr(settings, f"{provider.upper()}_MODEL")
).strip()
pricing = get_paid_provider_pricing_policy(provider, configured_model)
pricing_missing = pricing is None
pricing_info = {
"pricing_policy": {