280 lines
10 KiB
Python
280 lines
10 KiB
Python
import os
|
||
import re
|
||
import time
|
||
import asyncio
|
||
import logging
|
||
from datetime import datetime, timezone
|
||
from typing import Dict, Any, List, Optional, Tuple
|
||
|
||
logger = logging.getLogger("AGYUsageMonitor")
|
||
|
||
def make_quota_bar(percent_remaining: float, length: int = 10) -> str:
|
||
"""
|
||
Generates an emoji progress bar for remaining quota:
|
||
>= 60%: 🟩 Green (Safe/High)
|
||
20% - 59%: 🟨 Yellow (Medium)
|
||
< 20%: 🟥 Red (Low)
|
||
"""
|
||
percent = max(0.0, min(100.0, float(percent_remaining)))
|
||
filled = int(round((percent / 100.0) * length))
|
||
filled = max(0, min(length, filled))
|
||
|
||
if percent >= 60:
|
||
fill_char = "🟩"
|
||
elif percent >= 20:
|
||
fill_char = "🟨"
|
||
else:
|
||
fill_char = "🟥"
|
||
|
||
empty_char = "⬜"
|
||
return f"{fill_char * filled}{empty_char * (length - filled)}"
|
||
|
||
def parse_iso_time_remaining(iso_str: str, is_fa: bool = True) -> str:
|
||
"""Calculates relative time remaining until reset timestamp."""
|
||
if not iso_str:
|
||
return "نامشخص" if is_fa else "Unknown"
|
||
try:
|
||
clean_iso = iso_str.strip().replace("Z", "+00:00")
|
||
target_dt = datetime.fromisoformat(clean_iso)
|
||
now_dt = datetime.now(timezone.utc)
|
||
diff_sec = int((target_dt - now_dt).total_seconds())
|
||
|
||
if diff_sec <= 0:
|
||
return "هماکنون در حال تمدید" if is_fa else "Resetting now"
|
||
|
||
days = diff_sec // 86400
|
||
hours = (diff_sec % 86400) // 3600
|
||
mins = (diff_sec % 3600) // 60
|
||
|
||
if is_fa:
|
||
parts = []
|
||
if days > 0:
|
||
parts.append(f"{days} روز")
|
||
if hours > 0:
|
||
parts.append(f"{hours} ساعت")
|
||
if mins > 0 or not parts:
|
||
parts.append(f"{mins} دقیقه")
|
||
return f"{' و '.join(parts)} دیگر"
|
||
else:
|
||
parts = []
|
||
if days > 0:
|
||
parts.append(f"{days}d")
|
||
if hours > 0:
|
||
parts.append(f"{hours}h")
|
||
parts.append(f"{mins}m")
|
||
return f"in {' '.join(parts)}"
|
||
except Exception as e:
|
||
return iso_str
|
||
|
||
def parse_agy_usage_output(usage_raw: str, credits_raw: str) -> Dict[str, Any]:
|
||
"""Parses raw stdout from 'agy --print /usage' and 'agy --print /credits'."""
|
||
data: Dict[str, Any] = {
|
||
"gemini_5h": None,
|
||
"gemini_5h_reset": "",
|
||
"gemini_weekly": None,
|
||
"gemini_weekly_reset": "",
|
||
"claude_5h": None,
|
||
"claude_5h_reset": "",
|
||
"claude_weekly": None,
|
||
"claude_weekly_reset": "",
|
||
"credits": "0",
|
||
"raw_text": usage_raw + "\n" + credits_raw,
|
||
}
|
||
|
||
lines = (usage_raw + "\n" + credits_raw).splitlines()
|
||
for line in lines:
|
||
line_clean = line.strip()
|
||
if not line_clean:
|
||
continue
|
||
|
||
# Check Gemini 5H
|
||
if re.search(r"Gemini.*Five Hour Limit Remaining", line_clean, re.I):
|
||
m = re.search(r"(\d+)%\s+(\S+)", line_clean)
|
||
if m:
|
||
data["gemini_5h"] = int(m.group(1))
|
||
data["gemini_5h_reset"] = m.group(2)
|
||
|
||
# Check Gemini Weekly
|
||
elif re.search(r"Gemini.*Weekly Limit Remaining", line_clean, re.I):
|
||
m = re.search(r"(\d+)%\s+(\S+)", line_clean)
|
||
if m:
|
||
data["gemini_weekly"] = int(m.group(1))
|
||
data["gemini_weekly_reset"] = m.group(2)
|
||
|
||
# Check Claude/GPT 5H
|
||
elif re.search(r"(Claude|GPT).*Five Hour Limit Remaining", line_clean, re.I):
|
||
m = re.search(r"(\d+)%\s+(\S+)", line_clean)
|
||
if m:
|
||
data["claude_5h"] = int(m.group(1))
|
||
data["claude_5h_reset"] = m.group(2)
|
||
|
||
# Check Claude/GPT Weekly
|
||
elif re.search(r"(Claude|GPT).*Weekly Limit Remaining", line_clean, re.I):
|
||
m = re.search(r"(\d+)%\s+(\S+)", line_clean)
|
||
if m:
|
||
data["claude_weekly"] = int(m.group(1))
|
||
data["claude_weekly_reset"] = m.group(2)
|
||
|
||
# Check credits
|
||
elif re.search(r"Remaining credits", line_clean, re.I):
|
||
m = re.search(r"Remaining credits\s+(\d+)", line_clean, re.I)
|
||
if m:
|
||
data["credits"] = m.group(1)
|
||
|
||
return data
|
||
|
||
async def fetch_and_render_usage_report(is_fa: bool = True, user_id: Optional[int] = None) -> str:
|
||
"""Executes agy /usage and /credits and formats with visual progress bars."""
|
||
from auth_manager import auth_manager
|
||
env = auth_manager.get_user_env(user_id) if user_id else os.environ.copy()
|
||
usage_raw = ""
|
||
credits_raw = ""
|
||
try:
|
||
proc1 = await asyncio.create_subprocess_exec(
|
||
"agy", "--print", "/usage",
|
||
stdout=asyncio.subprocess.PIPE,
|
||
stderr=asyncio.subprocess.PIPE,
|
||
env=env,
|
||
)
|
||
out1, _ = await proc1.communicate()
|
||
usage_raw = out1.decode("utf-8", errors="replace").strip()
|
||
|
||
proc2 = await asyncio.create_subprocess_exec(
|
||
"agy", "--print", "/credits",
|
||
stdout=asyncio.subprocess.PIPE,
|
||
stderr=asyncio.subprocess.PIPE,
|
||
env=env,
|
||
)
|
||
out2, _ = await proc2.communicate()
|
||
credits_raw = out2.decode("utf-8", errors="replace").strip()
|
||
except Exception as e:
|
||
logger.error(f"Failed to query AGY quota: {e}")
|
||
usage_raw = f"Error: {e}"
|
||
|
||
parsed = parse_agy_usage_output(usage_raw, credits_raw)
|
||
|
||
g_5h = parsed["gemini_5h"]
|
||
g_wk = parsed["gemini_weekly"]
|
||
c_5h = parsed["claude_5h"]
|
||
c_wk = parsed["claude_weekly"]
|
||
credits_val = parsed["credits"]
|
||
|
||
if g_5h is None and g_wk is None and c_5h is None:
|
||
# Fallback to pre-formatted raw output if parsing failed
|
||
if is_fa:
|
||
return (
|
||
"📊 <b>گزارش لحظهای سهمیه و وضعیت مصرف AGY</b>\n\n"
|
||
f"<pre>{usage_raw}\n\n{credits_raw}</pre>\n\n"
|
||
"💡 <i>سهمیههای ۵ ساعته و هفتگی به صورت خودکار تمدید میشوند.</i>"
|
||
)
|
||
else:
|
||
return (
|
||
"📊 <b>Real-time AGY Quota & Usage Report</b>\n\n"
|
||
f"<pre>{usage_raw}\n\n{credits_raw}</pre>\n\n"
|
||
"💡 <i>Quotas are reset on 5-hour and 7-day windows.</i>"
|
||
)
|
||
|
||
if is_fa:
|
||
lines = [
|
||
"📈 <b>گزارش وضعیت سهمیه و مصرف هوش مصنوعی (AGY)</b>",
|
||
"",
|
||
"✨ <b>۱. مدلهای Gemini (جمینای):</b>",
|
||
]
|
||
if g_5h is not None:
|
||
bar_5h = make_quota_bar(g_5h)
|
||
reset_5h = parse_iso_time_remaining(parsed["gemini_5h_reset"], is_fa=True)
|
||
lines.extend([
|
||
f" ⏳ <b>سهمیه ۵ ساعته:</b>",
|
||
f" {bar_5h} <b>{g_5h}% باقیمانده</b>",
|
||
f" • 🔄 <i>زمان تمدید:</i> <code>{reset_5h}</code>",
|
||
])
|
||
if g_wk is not None:
|
||
bar_wk = make_quota_bar(g_wk)
|
||
reset_wk = parse_iso_time_remaining(parsed["gemini_weekly_reset"], is_fa=True)
|
||
lines.extend([
|
||
f" 📅 <b>سهمیه هفتگی:</b>",
|
||
f" {bar_wk} <b>{g_wk}% باقیمانده</b>",
|
||
f" • 🔄 <i>زمان تمدید:</i> <code>{reset_wk}</code>",
|
||
])
|
||
|
||
lines.extend([
|
||
"",
|
||
"🌟 <b>۲. مدلهای Claude و GPT:</b>",
|
||
]
|
||
)
|
||
if c_5h is not None:
|
||
bar_c5 = make_quota_bar(c_5h)
|
||
reset_c5 = parse_iso_time_remaining(parsed["claude_5h_reset"], is_fa=True)
|
||
lines.extend([
|
||
f" ⏳ <b>سهمیه ۵ ساعته:</b>",
|
||
f" {bar_c5} <b>{c_5h}% باقیمانده</b>",
|
||
f" • 🔄 <i>زمان تمدید:</i> <code>{reset_c5}</code>",
|
||
])
|
||
if c_wk is not None:
|
||
bar_cwk = make_quota_bar(c_wk)
|
||
reset_cwk = parse_iso_time_remaining(parsed["claude_weekly_reset"], is_fa=True)
|
||
lines.extend([
|
||
f" 📅 <b>سهمیه هفتگی:</b>",
|
||
f" {bar_cwk} <b>{c_wk}% باقیمانده</b>",
|
||
f" • 🔄 <i>زمان تمدید:</i> <code>{reset_cwk}</code>",
|
||
])
|
||
|
||
lines.extend([
|
||
"",
|
||
f"🪙 <b>اعتبار مدلها (Model Credits):</b> <code>{credits_val}</code>",
|
||
"",
|
||
"💡 <i>سهمیههای مصرف به صورت خودکار در بازههای زمانی فوق بازنشانی میشوند.</i>"
|
||
])
|
||
else:
|
||
lines = [
|
||
"📈 <b>Real-Time AGY Quota & Usage Report</b>",
|
||
"",
|
||
"✨ <b>1. Gemini Models:</b>",
|
||
]
|
||
if g_5h is not None:
|
||
bar_5h = make_quota_bar(g_5h)
|
||
reset_5h = parse_iso_time_remaining(parsed["gemini_5h_reset"], is_fa=False)
|
||
lines.extend([
|
||
f" ⏳ <b>5-Hour Limit:</b>",
|
||
f" {bar_5h} <b>{g_5h}% Remaining</b>",
|
||
f" • 🔄 <i>Reset:</i> <code>{reset_5h}</code>",
|
||
])
|
||
if g_wk is not None:
|
||
bar_wk = make_quota_bar(g_wk)
|
||
reset_wk = parse_iso_time_remaining(parsed["gemini_weekly_reset"], is_fa=False)
|
||
lines.extend([
|
||
f" 📅 <b>Weekly Limit:</b>",
|
||
f" {bar_wk} <b>{g_wk}% Remaining</b>",
|
||
f" • 🔄 <i>Reset:</i> <code>{reset_wk}</code>",
|
||
])
|
||
|
||
lines.extend([
|
||
"",
|
||
"🌟 <b>2. Claude & GPT Models:</b>",
|
||
])
|
||
if c_5h is not None:
|
||
bar_c5 = make_quota_bar(c_5h)
|
||
reset_c5 = parse_iso_time_remaining(parsed["claude_5h_reset"], is_fa=False)
|
||
lines.extend([
|
||
f" ⏳ <b>5-Hour Limit:</b>",
|
||
f" {bar_c5} <b>{c_5h}% Remaining</b>",
|
||
f" • 🔄 <i>Reset:</i> <code>{reset_c5}</code>",
|
||
])
|
||
if c_wk is not None:
|
||
bar_cwk = make_quota_bar(c_wk)
|
||
reset_cwk = parse_iso_time_remaining(parsed["claude_weekly_reset"], is_fa=False)
|
||
lines.extend([
|
||
f" 📅 <b>Weekly Limit:</b>",
|
||
f" {bar_cwk} <b>{c_wk}% Remaining</b>",
|
||
f" • 🔄 <i>Reset:</i> <code>{reset_cwk}</code>",
|
||
])
|
||
|
||
lines.extend([
|
||
"",
|
||
f"🪙 <b>Model Credits:</b> <code>{credits_val}</code>",
|
||
"",
|
||
"💡 <i>Quotas automatically reset periodically in the windows shown above.</i>"
|
||
])
|
||
|
||
return "\n".join(lines)
|