Spaces:
Running
Running
File size: 2,660 Bytes
afd56bc | 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 93 94 95 96 97 98 99 100 | from datetime import datetime, timezone
from sqlalchemy.orm import Session
from core.subscription.db import SessionLocal
from core.subscription.models import UserUsage, User, UsageLog
def get_or_create_usage(db: Session, user_id: str) -> UserUsage:
user = db.query(User).filter(User.clerk_id == user_id).first()
if not user:
user = User(clerk_id=user_id, tier="free")
db.add(user)
db.commit()
usage = db.query(UserUsage).filter(UserUsage.user_id == user_id).first()
if not usage:
usage = UserUsage(user_id=user_id)
db.add(usage)
db.commit()
db.refresh(usage)
return usage
def check_and_reset_limits(usage: UserUsage, db: Session):
now = datetime.now(timezone.utc)
changed = False
# Reset dzienny wizarda
if usage.last_reset_daily.date() < now.date():
usage.wizard_iterations_today = 0
usage.last_reset_daily = now
changed = True
# Reset miesięczny tokenów
if (
usage.last_reset_monthly.month != now.month
or usage.last_reset_monthly.year != now.year
):
usage.tokens_used_month = 0
usage.last_reset_monthly = now
changed = True
if changed:
db.commit()
def increment_wizard_iteration(user_id: str) -> int:
db = SessionLocal()
try:
usage = get_or_create_usage(db, user_id)
check_and_reset_limits(usage, db)
usage.wizard_iterations_today += 1
# Log audytowy
log = UsageLog(user_id=user_id, action_type="wizard_iteration", tokens_cost=0)
db.add(log)
db.commit()
return usage.wizard_iterations_today
finally:
db.close()
def get_wizard_iterations(user_id: str) -> int:
db = SessionLocal()
try:
usage = get_or_create_usage(db, user_id)
check_and_reset_limits(usage, db)
return usage.wizard_iterations_today
finally:
db.close()
def increment_tokens(user_id: str, tokens: int, action_type: str = "llm_call") -> int:
db = SessionLocal()
try:
usage = get_or_create_usage(db, user_id)
check_and_reset_limits(usage, db)
usage.tokens_used_month += tokens
# Log audytowy
log = UsageLog(user_id=user_id, action_type=action_type, tokens_cost=tokens)
db.add(log)
db.commit()
return usage.tokens_used_month
finally:
db.close()
def get_used_tokens(user_id: str) -> int:
db = SessionLocal()
try:
usage = get_or_create_usage(db, user_id)
check_and_reset_limits(usage, db)
return usage.tokens_used_month
finally:
db.close()
|