64 lines
2.2 KiB
Python
64 lines
2.2 KiB
Python
from datetime import datetime, time, timedelta
|
|
from zoneinfo import ZoneInfo
|
|
|
|
from fastapi import HTTPException, status
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.core.config import settings
|
|
from app.models.daily_usage import DailyUsage
|
|
|
|
SHANGHAI_TZ = ZoneInfo("Asia/Shanghai")
|
|
|
|
|
|
class QuotaService:
|
|
def today(self):
|
|
return datetime.now(SHANGHAI_TZ).date()
|
|
|
|
def reset_at(self) -> datetime:
|
|
tomorrow = self.today() + timedelta(days=1)
|
|
return datetime.combine(tomorrow, time.min, tzinfo=SHANGHAI_TZ)
|
|
|
|
async def get_quota(self, db: AsyncSession, user_id: str) -> dict:
|
|
usage = await self._get_usage(db, user_id)
|
|
used = usage.count if usage else 0
|
|
return self._quota_payload(used)
|
|
|
|
async def consume(self, db: AsyncSession, user_id: str) -> dict:
|
|
usage = await self._get_usage(db, user_id)
|
|
if usage is None:
|
|
usage = DailyUsage(user_id=user_id, usage_date=self.today(), count=0)
|
|
db.add(usage)
|
|
await db.flush()
|
|
|
|
if usage.count >= settings.daily_reading_limit:
|
|
quota = self._quota_payload(usage.count)
|
|
raise HTTPException(
|
|
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
|
detail=f"今日 {settings.daily_reading_limit} 次解读机会已用完，请明天 0 点后再来。",
|
|
headers={
|
|
"X-RateLimit-Limit": str(quota["limit"]),
|
|
"X-RateLimit-Remaining": str(quota["remaining"]),
|
|
"X-RateLimit-Reset": quota["reset_at"].isoformat(),
|
|
},
|
|
)
|
|
|
|
usage.count += 1
|
|
await db.flush()
|
|
return self._quota_payload(usage.count)
|
|
|
|
async def _get_usage(self, db: AsyncSession, user_id: str) -> DailyUsage | None:
|
|
result = await db.execute(
|
|
select(DailyUsage).where(DailyUsage.user_id == user_id, DailyUsage.usage_date == self.today())
|
|
)
|
|
return result.scalar_one_or_none()
|
|
|
|
def _quota_payload(self, used: int) -> dict:
|
|
limit = settings.daily_reading_limit
|
|
return {
|
|
"limit": limit,
|
|
"used": used,
|
|
"remaining": max(0, limit - used),
|
|
"reset_at": self.reset_at(),
|
|
}
|