Spaces:
Sleeping
Sleeping
File size: 2,554 Bytes
1397957 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | from typing import Optional
from fastapi import HTTPException, Depends
from pydantic import BaseModel
from .auth import AuthUser, require_auth
from .supabase import get_client, is_enabled as supabase_enabled
from .config import settings
class UsageInfo(BaseModel):
input_tokens: int = 0
output_tokens: int = 0
request_count: int = 0
class QuotaLimits(BaseModel):
daily_requests: int = 100
daily_input_tokens: int = 1_000_000
daily_output_tokens: int = 500_000
DEFAULT_LIMITS = QuotaLimits()
async def get_usage(user_id: str) -> UsageInfo:
if not supabase_enabled():
return UsageInfo()
client = get_client()
result = client.rpc("get_opencode_usage", {"p_user_id": user_id}).execute()
if result.data and len(result.data) > 0:
row = result.data[0]
return UsageInfo(
input_tokens=row.get("input_tokens", 0),
output_tokens=row.get("output_tokens", 0),
request_count=row.get("request_count", 0),
)
return UsageInfo()
async def increment_usage(user_id: str, input_tokens: int = 0, output_tokens: int = 0) -> None:
if not supabase_enabled():
return
client = get_client()
client.rpc("increment_opencode_usage", {
"p_user_id": user_id,
"p_input_tokens": input_tokens,
"p_output_tokens": output_tokens,
}).execute()
async def check_quota(user: AuthUser = Depends(require_auth)) -> AuthUser:
if not supabase_enabled():
return user
usage = await get_usage(user.id)
limits = DEFAULT_LIMITS
if usage.request_count >= limits.daily_requests:
raise HTTPException(
status_code=429,
detail={
"error": "Daily request limit reached",
"usage": usage.model_dump(),
"limits": limits.model_dump(),
}
)
if usage.input_tokens >= limits.daily_input_tokens:
raise HTTPException(
status_code=429,
detail={
"error": "Daily input token limit reached",
"usage": usage.model_dump(),
"limits": limits.model_dump(),
}
)
if usage.output_tokens >= limits.daily_output_tokens:
raise HTTPException(
status_code=429,
detail={
"error": "Daily output token limit reached",
"usage": usage.model_dump(),
"limits": limits.model_dump(),
}
)
return user
|