feat(ai): add guarded Claude fallback route
This commit is contained in:
@@ -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": {
|
||||
|
||||
Reference in New Issue
Block a user