admin_adjust_user_quota.py

#!/usr/bin/env python3
"""
管理员工具：调整用户的每日调用次数配额

用法:
    python admin_adjust_user_quota.py <username> <action>

示例:
    # 重置用户的今日调用次数为0（相当于给他4次新配额）
    python admin_adjust_user_quota.py learnmlf reset
    
    # 设置用户今日已用次数为1（剩余3次）
    python admin_adjust_user_quota.py learnmlf set 1
    
    # 查看用户当前状态
    python admin_adjust_user_quota.py learnmlf show
    
    # 给用户额外增加2次配额（即今日可用6次）
    python admin_adjust_user_quota.py learnmlf add 2
"""

import os
import sys
from datetime import datetime
from app import DatasetManager, MAX_DAILY_CALLS

# 配置
HF_TOKEN = os.environ.get("HF_TOKEN", "")
DATASET_REPO_ID = "learnmlf/video-evaluations"

def show_usage():
    """显示使用说明"""
    print(__doc__)
    sys.exit(1)

def show_user_status(dm: DatasetManager, username: str):
    """显示用户当前状态"""
    calls_today = dm.get_user_calls_today(username)
    remaining = max(0, MAX_DAILY_CALLS - calls_today)
    today = datetime.now().strftime("%Y-%m-%d")
    
    print("\n" + "="*60)
    print(f"📊 用户状态: {username}")
    print("="*60)
    print(f"日期: {today}")
    print(f"已用次数: {calls_today}/{MAX_DAILY_CALLS}")
    print(f"剩余次数: {remaining}")
    print(f"状态: {'✅ 可用' if remaining > 0 else '❌ 已用完'}")
    print("="*60)

def reset_user_quota(dm: DatasetManager, username: str):
    """重置用户今日配额（设为0）"""
    today = datetime.now().strftime("%Y-%m-%d")
    count_file = f"call_counts/{today}/{username}.json"
    
    import json
    import tempfile
    
    count_data = {
        'user': username,
        'date': today,
        'count': 0,
        'last_updated': datetime.now().isoformat(),
        'note': 'Reset by admin'
    }
    
    with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False, encoding='utf-8') as f:
        json.dump(count_data, f, ensure_ascii=False, indent=2)
        temp_path = f.name
    
    try:
        dm.api.upload_file(
            path_or_fileobj=temp_path,
            path_in_repo=count_file,
            repo_id=DATASET_REPO_ID,
            repo_type="dataset",
            token=HF_TOKEN
        )
        print(f"\n✅ 用户 {username} 的今日调用次数已重置为 0")
        print(f"   → 剩余次数: {MAX_DAILY_CALLS}")
    finally:
        os.unlink(temp_path)

def set_user_quota(dm: DatasetManager, username: str, count: int):
    """设置用户今日已用次数"""
    if count < 0:
        print(f"❌ 错误: 次数不能为负数")
        sys.exit(1)
    
    today = datetime.now().strftime("%Y-%m-%d")
    count_file = f"call_counts/{today}/{username}.json"
    
    import json
    import tempfile
    
    count_data = {
        'user': username,
        'date': today,
        'count': count,
        'last_updated': datetime.now().isoformat(),
        'note': f'Set to {count} by admin'
    }
    
    with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False, encoding='utf-8') as f:
        json.dump(count_data, f, ensure_ascii=False, indent=2)
        temp_path = f.name
    
    try:
        dm.api.upload_file(
            path_or_fileobj=temp_path,
            path_in_repo=count_file,
            repo_id=DATASET_REPO_ID,
            repo_type="dataset",
            token=HF_TOKEN
        )
        remaining = max(0, MAX_DAILY_CALLS - count)
        print(f"\n✅ 用户 {username} 的今日调用次数已设置为 {count}")
        print(f"   → 剩余次数: {remaining}")
    finally:
        os.unlink(temp_path)

def add_user_quota(dm: DatasetManager, username: str, extra: int):
    """给用户增加额外配额（减少已用次数）"""
    if extra <= 0:
        print(f"❌ 错误: 额外次数必须大于0")
        sys.exit(1)
    
    current = dm.get_user_calls_today(username)
    new_count = current - extra  # 减少已用次数 = 增加剩余次数（允许负数）
    
    today = datetime.now().strftime("%Y-%m-%d")
    count_file = f"call_counts/{today}/{username}.json"
    
    import json
    import tempfile
    
    count_data = {
        'user': username,
        'date': today,
        'count': new_count,
        'last_updated': datetime.now().isoformat(),
        'note': f'Added {extra} extra calls by admin (was {current})'
    }
    
    with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False, encoding='utf-8') as f:
        json.dump(count_data, f, ensure_ascii=False, indent=2)
        temp_path = f.name
    
    try:
        dm.api.upload_file(
            path_or_fileobj=temp_path,
            path_in_repo=count_file,
            repo_id=DATASET_REPO_ID,
            repo_type="dataset",
            token=HF_TOKEN
        )
        actual_extra = current - new_count
        new_remaining = MAX_DAILY_CALLS - new_count
        print(f"\n✅ 已给用户 {username} 增加 {actual_extra} 次额外配额")
        print(f"   之前: {current}/{MAX_DAILY_CALLS} (剩余 {MAX_DAILY_CALLS - current})")
        print(f"   现在: {new_count}/{MAX_DAILY_CALLS} (剩余 {new_remaining})")
        print(f"   → 总配额: {new_remaining} 次")
    finally:
        os.unlink(temp_path)

def main():
    if not HF_TOKEN:
        print("❌ 错误: 请设置 HF_TOKEN 环境变量")
        print("   export HF_TOKEN='your_token_here'")
        sys.exit(1)
    
    if len(sys.argv) < 3:
        show_usage()
    
    username = sys.argv[1]
    action = sys.argv[2].lower()
    
    # 初始化 DatasetManager
    dm = DatasetManager(DATASET_REPO_ID, HF_TOKEN)
    
    print(f"\n🔧 管理员工具 - 用户配额管理")
    print(f"Dataset: {DATASET_REPO_ID}")
    print(f"日期: {datetime.now().strftime('%Y-%m-%d')}")
    print(f"标准配额: {MAX_DAILY_CALLS} 次/天")
    
    if action == "show":
        # 查看状态
        show_user_status(dm, username)
    
    elif action == "reset":
        # 重置为0
        print(f"\n正在重置用户 {username} 的配额...")
        show_user_status(dm, username)
        
        confirm = input(f"\n确认重置为 0 吗？ (yes/no): ")
        if confirm.lower() == 'yes':
            reset_user_quota(dm, username)
            show_user_status(dm, username)
        else:
            print("❌ 操作已取消")
    
    elif action == "set":
        # 设置为指定值
        if len(sys.argv) < 4:
            print("❌ 错误: 缺少次数参数")
            print("   用法: python admin_adjust_user_quota.py <username> set <count>")
            sys.exit(1)
        
        try:
            count = int(sys.argv[3])
        except ValueError:
            print(f"❌ 错误: 无效的数字 '{sys.argv[3]}'")
            sys.exit(1)
        
        print(f"\n正在设置用户 {username} 的已用次数为 {count}...")
        show_user_status(dm, username)
        
        confirm = input(f"\n确认设置为 {count} 吗？ (yes/no): ")
        if confirm.lower() == 'yes':
            set_user_quota(dm, username, count)
            show_user_status(dm, username)
        else:
            print("❌ 操作已取消")
    
    elif action == "add":
        # 增加额外配额
        if len(sys.argv) < 4:
            print("❌ 错误: 缺少额外次数参数")
            print("   用法: python admin_adjust_user_quota.py <username> add <extra_count>")
            sys.exit(1)
        
        try:
            extra = int(sys.argv[3])
        except ValueError:
            print(f"❌ 错误: 无效的数字 '{sys.argv[3]}'")
            sys.exit(1)
        
        print(f"\n正在给用户 {username} 增加 {extra} 次额外配额...")
        show_user_status(dm, username)
        
        confirm = input(f"\n确认增加 {extra} 次配额吗？ (yes/no): ")
        if confirm.lower() == 'yes':
            add_user_quota(dm, username, extra)
            show_user_status(dm, username)
        else:
            print("❌ 操作已取消")
    
    else:
        print(f"❌ 错误: 未知操作 '{action}'")
        print("   支持的操作: show, reset, set, add")
        show_usage()

if __name__ == "__main__":
    main()