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()