复用 ccswitch 统计 Codex token
This commit is contained in:
Codex
parent
7a48b50d5b
commit
2f5cfdb1c8
@ -1,11 +1,12 @@
|
||||
import argparse
|
||||
import asyncio
|
||||
import ctypes
|
||||
from datetime import date, timedelta
|
||||
from datetime import date, datetime, timedelta
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import queue
|
||||
import sqlite3
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
@ -24,8 +25,8 @@ DEFAULT_CHAR_UUID = "9f4af548-1a14-4fc1-8f5a-4b9f00000002"
|
||||
MAX_DISK_ROWS = 4
|
||||
MAX_DISK_TEXT = 124
|
||||
CODEX_DAILY_BAR_COUNT = 14
|
||||
CODEX_REFRESH_SECONDS = 60.0
|
||||
CODEX_USAGE_REFRESH_SECONDS = 300.0
|
||||
CODEX_REFRESH_SECONDS = 300.0
|
||||
CCSWITCH_CODEX_REFRESH_SECONDS = 60.0
|
||||
ZHIPU_REFRESH_SECONDS = 60.0
|
||||
ZHIPU_USAGE_REFRESH_SECONDS = 600.0
|
||||
ZHIPU_DAILY_BAR_COUNT = 14
|
||||
@ -45,6 +46,25 @@ def clamp_kbps(value: float) -> int:
|
||||
return max(0, min(999999, int(round(value))))
|
||||
|
||||
|
||||
def tokens_to_m10(tokens) -> int:
|
||||
if not isinstance(tokens, (int, float)) or tokens < 0:
|
||||
return -1
|
||||
return max(0, min(2147483647, int(round(tokens / 100000))))
|
||||
|
||||
|
||||
def daily_bars_from_token_map(tokens_by_day: dict[str, int], count: int) -> tuple[int, str]:
|
||||
today = date.today()
|
||||
wanted_days = [(today - timedelta(days=offset)).isoformat() for offset in range(count - 1, -1, -1)]
|
||||
values = [tokens_by_day.get(day, 0) for day in wanted_days]
|
||||
max_value = max(values) if values else 0
|
||||
if max_value > 0:
|
||||
percents = [max(0, min(100, int(round(value * 100 / max_value)))) for value in values]
|
||||
else:
|
||||
percents = [0 for _ in values]
|
||||
today_tokens_m10 = tokens_to_m10(tokens_by_day.get(today.isoformat(), 0))
|
||||
return today_tokens_m10, ":".join(str(value) for value in percents)
|
||||
|
||||
|
||||
class _PdhCounterValue(ctypes.Structure):
|
||||
_fields_ = [
|
||||
("status", wintypes.LONG),
|
||||
@ -421,9 +441,6 @@ class CodexRateLimitReader:
|
||||
self.refresh_seconds = refresh_seconds
|
||||
self.next_refresh = 0.0
|
||||
self.cached = (-1, -1, -1, -1, -1)
|
||||
self.usage_refresh_seconds = CODEX_USAGE_REFRESH_SECONDS
|
||||
self.next_usage_refresh = 0.0
|
||||
self.usage_cached = (-1, "")
|
||||
self.process: subprocess.Popen | None = None
|
||||
self.responses: queue.Queue[dict] = queue.Queue()
|
||||
self.next_id = 1
|
||||
@ -579,60 +596,144 @@ class CodexRateLimitReader:
|
||||
)
|
||||
return self.cached
|
||||
|
||||
@staticmethod
|
||||
def _tokens_to_m10(tokens) -> int:
|
||||
if not isinstance(tokens, (int, float)) or tokens < 0:
|
||||
return -1
|
||||
return max(0, min(2147483647, int(round(tokens / 100000))))
|
||||
class CcswitchCodexUsageReader:
|
||||
DEDUP_WINDOW_SECONDS = 10 * 60
|
||||
|
||||
def __init__(self, refresh_seconds: float = CCSWITCH_CODEX_REFRESH_SECONDS) -> None:
|
||||
self.refresh_seconds = refresh_seconds
|
||||
self.next_refresh = 0.0
|
||||
self.cached = (-1, "")
|
||||
|
||||
@staticmethod
|
||||
def _daily_bars_from_buckets(buckets: list[dict]) -> tuple[int, str]:
|
||||
today = date.today()
|
||||
wanted_days = [(today - timedelta(days=offset)).isoformat() for offset in range(CODEX_DAILY_BAR_COUNT - 1, -1, -1)]
|
||||
def _candidate_paths() -> list[Path]:
|
||||
paths: list[Path] = []
|
||||
override = os.environ.get("CCSWITCH_DB_PATH", "").strip()
|
||||
if override:
|
||||
paths.append(Path(override).expanduser())
|
||||
|
||||
home = Path.home()
|
||||
paths.append(home / ".cc-switch" / "cc-switch.db")
|
||||
|
||||
appdata = os.environ.get("APPDATA")
|
||||
if appdata:
|
||||
paths.append(Path(appdata) / "com.ccswitch.desktop" / "cc-switch.db")
|
||||
|
||||
local_appdata = os.environ.get("LOCALAPPDATA")
|
||||
if local_appdata:
|
||||
paths.append(Path(local_appdata) / "com.ccswitch.desktop" / "cc-switch.db")
|
||||
|
||||
unique: list[Path] = []
|
||||
seen: set[str] = set()
|
||||
for path in paths:
|
||||
key = str(path).lower()
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
unique.append(path)
|
||||
return unique
|
||||
|
||||
@classmethod
|
||||
def _find_db(cls) -> Path | None:
|
||||
for path in cls._candidate_paths():
|
||||
if path.is_file():
|
||||
return path
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _connect_readonly(path: Path) -> sqlite3.Connection:
|
||||
sqlite_path = urllib.parse.quote(str(path).replace("\\", "/"), safe="/:")
|
||||
con = sqlite3.connect(f"file:{sqlite_path}?mode=ro", uri=True, timeout=1.0)
|
||||
con.row_factory = sqlite3.Row
|
||||
return con
|
||||
|
||||
@classmethod
|
||||
def _effective_filter_sql(cls) -> str:
|
||||
return f"""
|
||||
NOT (
|
||||
COALESCE(l.data_source, 'proxy') IN ('session_log', 'codex_session', 'gemini_session', 'opencode_session')
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM proxy_request_logs proxy_dedup
|
||||
WHERE COALESCE(proxy_dedup.data_source, 'proxy') = 'proxy'
|
||||
AND proxy_dedup.app_type = l.app_type
|
||||
AND proxy_dedup.status_code >= 200
|
||||
AND proxy_dedup.status_code < 300
|
||||
AND proxy_dedup.input_tokens = l.input_tokens
|
||||
AND proxy_dedup.output_tokens = l.output_tokens
|
||||
AND proxy_dedup.cache_read_tokens = l.cache_read_tokens
|
||||
AND (
|
||||
proxy_dedup.cache_creation_tokens = l.cache_creation_tokens
|
||||
OR (
|
||||
l.cache_creation_tokens = 0
|
||||
AND COALESCE(l.data_source, 'proxy') IN ('codex_session', 'gemini_session', 'opencode_session')
|
||||
)
|
||||
)
|
||||
AND proxy_dedup.created_at BETWEEN
|
||||
l.created_at - {cls.DEDUP_WINDOW_SECONDS}
|
||||
AND l.created_at + {cls.DEDUP_WINDOW_SECONDS}
|
||||
AND (
|
||||
LOWER(proxy_dedup.model) = LOWER(l.model)
|
||||
OR LOWER(proxy_dedup.model) = 'unknown'
|
||||
OR LOWER(l.model) = 'unknown'
|
||||
)
|
||||
)
|
||||
)
|
||||
"""
|
||||
|
||||
def _query_daily_tokens(self, db_path: Path) -> dict[str, int]:
|
||||
start_day = date.today() - timedelta(days=CODEX_DAILY_BAR_COUNT - 1)
|
||||
start_ts = int(datetime.combine(start_day, datetime.min.time()).timestamp())
|
||||
real_total_sql = """
|
||||
(
|
||||
CASE
|
||||
WHEN COALESCE(l.input_tokens, 0) >= COALESCE(l.cache_read_tokens, 0)
|
||||
THEN COALESCE(l.input_tokens, 0) - COALESCE(l.cache_read_tokens, 0)
|
||||
ELSE COALESCE(l.input_tokens, 0)
|
||||
END
|
||||
+ COALESCE(l.output_tokens, 0)
|
||||
+ COALESCE(l.cache_creation_tokens, 0)
|
||||
+ COALESCE(l.cache_read_tokens, 0)
|
||||
)
|
||||
"""
|
||||
sql = f"""
|
||||
SELECT date(l.created_at, 'unixepoch', 'localtime') AS usage_day,
|
||||
SUM({real_total_sql}) AS real_total_tokens
|
||||
FROM proxy_request_logs l
|
||||
WHERE l.app_type = 'codex'
|
||||
AND l.status_code >= 200
|
||||
AND l.status_code < 300
|
||||
AND l.created_at >= ?
|
||||
AND {self._effective_filter_sql()}
|
||||
GROUP BY usage_day
|
||||
"""
|
||||
|
||||
tokens_by_day: dict[str, int] = {}
|
||||
|
||||
for bucket in buckets:
|
||||
if not isinstance(bucket, dict):
|
||||
continue
|
||||
start_date = bucket.get("startDate")
|
||||
tokens = bucket.get("tokens")
|
||||
if isinstance(start_date, str) and isinstance(tokens, (int, float)) and tokens >= 0:
|
||||
tokens_by_day[start_date] = int(tokens)
|
||||
|
||||
values = [tokens_by_day.get(day, 0) for day in wanted_days]
|
||||
max_value = max(values) if values else 0
|
||||
if max_value > 0:
|
||||
percents = [max(0, min(100, int(round(value * 100 / max_value)))) for value in values]
|
||||
else:
|
||||
percents = [0 for _ in values]
|
||||
|
||||
today_tokens_m10 = CodexRateLimitReader._tokens_to_m10(tokens_by_day.get(today.isoformat(), 0))
|
||||
return today_tokens_m10, ":".join(str(value) for value in percents)
|
||||
with self._connect_readonly(db_path) as con:
|
||||
for row in con.execute(sql, (start_ts,)):
|
||||
usage_day = row["usage_day"]
|
||||
tokens = row["real_total_tokens"]
|
||||
if isinstance(usage_day, str) and isinstance(tokens, (int, float)) and tokens >= 0:
|
||||
tokens_by_day[usage_day] = int(tokens)
|
||||
return tokens_by_day
|
||||
|
||||
def read_usage_snapshot(self) -> tuple[int, str]:
|
||||
now = time.monotonic()
|
||||
if now < self.next_usage_refresh:
|
||||
return self.usage_cached
|
||||
if now < self.next_refresh:
|
||||
return self.cached
|
||||
|
||||
if not self._ensure_started():
|
||||
self.next_usage_refresh = now + min(15.0, self.usage_refresh_seconds)
|
||||
return self.usage_cached
|
||||
db_path = self._find_db()
|
||||
if db_path is None:
|
||||
self.next_refresh = now + min(15.0, self.refresh_seconds)
|
||||
return self.cached
|
||||
|
||||
request_id = self._send("account/usage/read")
|
||||
if request_id is None:
|
||||
self.next_usage_refresh = now + min(15.0, self.usage_refresh_seconds)
|
||||
return self.usage_cached
|
||||
try:
|
||||
tokens_by_day = self._query_daily_tokens(db_path)
|
||||
except sqlite3.Error:
|
||||
self.next_refresh = now + min(15.0, self.refresh_seconds)
|
||||
return self.cached
|
||||
|
||||
message = self._wait_for(request_id, 8.0)
|
||||
result = message.get("result") if message else None
|
||||
buckets = result.get("dailyUsageBuckets") if isinstance(result, dict) else None
|
||||
if not isinstance(buckets, list) or not buckets:
|
||||
self.next_usage_refresh = now + min(15.0, self.usage_refresh_seconds)
|
||||
return self.usage_cached
|
||||
|
||||
self.next_usage_refresh = now + self.usage_refresh_seconds
|
||||
self.usage_cached = self._daily_bars_from_buckets(buckets)
|
||||
return self.usage_cached
|
||||
self.cached = daily_bars_from_token_map(tokens_by_day, CODEX_DAILY_BAR_COUNT)
|
||||
self.next_refresh = now + self.refresh_seconds
|
||||
return self.cached
|
||||
|
||||
|
||||
class ZhipuStatusReader:
|
||||
@ -824,6 +925,7 @@ class MetricsSampler:
|
||||
self.cpu_frequency = CpuFrequencyReader()
|
||||
self.disk_rates = LogicalDiskRateReader(self._disk_labels())
|
||||
self.codex = CodexRateLimitReader()
|
||||
self.codex_usage = CcswitchCodexUsageReader()
|
||||
self.zhipu = ZhipuStatusReader()
|
||||
psutil.cpu_percent(interval=None)
|
||||
|
||||
@ -927,7 +1029,7 @@ class MetricsSampler:
|
||||
|
||||
disk_count, disk_percent, disk_text = self._disk_summary(disk_rates)
|
||||
codex_5h, codex_week, codex_5h_reset, codex_week_reset, codex_credits = self.codex.read()
|
||||
codex_today_tokens_m10, codex_daily_bars = self.codex.read_usage_snapshot()
|
||||
codex_today_tokens_m10, codex_daily_bars = self.codex_usage.read_usage_snapshot()
|
||||
zhipu_status, zhipu_5h, zhipu_week, zhipu_5h_reset, zhipu_week_reset, zhipu_today_tokens_m10, zhipu_daily_bars = self.zhipu.read()
|
||||
|
||||
return (
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user