robosage / monetization /token_manager.py
mgbam's picture
Update monetization/token_manager.py
4ae8332 verified
# token_manager.py - Manages monetization via Hugging Face tokens or Stripe
import os
import logging
logging.basicConfig(level=logging.INFO)
# Simple paywall using Hugging Face tokens or credits
class TokenManager:
def __init__(self, free_tier=True):
self.free_tier_enabled = free_tier
self.user_tokens = {} # user_id -> tokens
def register_user(self, user_id: str):
if user_id not in self.user_tokens:
self.user_tokens[user_id] = 5 if self.free_tier_enabled else 0
logging.info(f"[πŸ’³] Registered user '{user_id}' with {self.user_tokens[user_id]} tokens.")
def deduct_token(self, user_id: str) -> bool:
if self.user_tokens.get(user_id, 0) > 0:
self.user_tokens[user_id] -= 1
logging.info(f"[πŸ”»] Deducted token from '{user_id}'. Remaining: {self.user_tokens[user_id]}")
return True
else:
logging.warning(f"[🚫] No tokens left for '{user_id}'.")
return False
def add_tokens(self, user_id: str, amount: int):
self.user_tokens[user_id] = self.user_tokens.get(user_id, 0) + amount
logging.info(f"[βž•] Added {amount} tokens to '{user_id}'. Total: {self.user_tokens[user_id]}")
def check_balance(self, user_id: str) -> int:
return self.user_tokens.get(user_id, 0)
# Example
if __name__ == "__main__":
tm = TokenManager()
tm.register_user("user_001")
print("Balance:", tm.check_balance("user_001"))
tm.deduct_token("user_001")
print("Balance:", tm.check_balance("user_001"))