复用 ccswitch 统计 Codex token
This commit is contained in:
parent
7a48b50d5b
commit
2f5cfdb1c8
@ -1,11 +1,12 @@
|
|||||||
import argparse
|
import argparse
|
||||||
import asyncio
|
import asyncio
|
||||||
import ctypes
|
import ctypes
|
||||||
from datetime import date, timedelta
|
from datetime import date, datetime, timedelta
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
import queue
|
import queue
|
||||||
|
import sqlite3
|
||||||
import subprocess
|
import subprocess
|
||||||
import sys
|
import sys
|
||||||
import threading
|
import threading
|
||||||
@ -24,8 +25,8 @@ DEFAULT_CHAR_UUID = "9f4af548-1a14-4fc1-8f5a-4b9f00000002"
|
|||||||
MAX_DISK_ROWS = 4
|
MAX_DISK_ROWS = 4
|
||||||
MAX_DISK_TEXT = 124
|
MAX_DISK_TEXT = 124
|
||||||
CODEX_DAILY_BAR_COUNT = 14
|
CODEX_DAILY_BAR_COUNT = 14
|
||||||
CODEX_REFRESH_SECONDS = 60.0
|
CODEX_REFRESH_SECONDS = 300.0
|
||||||
CODEX_USAGE_REFRESH_SECONDS = 300.0
|
CCSWITCH_CODEX_REFRESH_SECONDS = 60.0
|
||||||
ZHIPU_REFRESH_SECONDS = 60.0
|
ZHIPU_REFRESH_SECONDS = 60.0
|
||||||
ZHIPU_USAGE_REFRESH_SECONDS = 600.0
|
ZHIPU_USAGE_REFRESH_SECONDS = 600.0
|
||||||
ZHIPU_DAILY_BAR_COUNT = 14
|
ZHIPU_DAILY_BAR_COUNT = 14
|
||||||
@ -45,6 +46,25 @@ def clamp_kbps(value: float) -> int:
|
|||||||
return max(0, min(999999, int(round(value))))
|
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):
|
class _PdhCounterValue(ctypes.Structure):
|
||||||
_fields_ = [
|
_fields_ = [
|
||||||
("status", wintypes.LONG),
|
("status", wintypes.LONG),
|
||||||
@ -421,9 +441,6 @@ class CodexRateLimitReader:
|
|||||||
self.refresh_seconds = refresh_seconds
|
self.refresh_seconds = refresh_seconds
|
||||||
self.next_refresh = 0.0
|
self.next_refresh = 0.0
|
||||||
self.cached = (-1, -1, -1, -1, -1)
|
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.process: subprocess.Popen | None = None
|
||||||
self.responses: queue.Queue[dict] = queue.Queue()
|
self.responses: queue.Queue[dict] = queue.Queue()
|
||||||
self.next_id = 1
|
self.next_id = 1
|
||||||
@ -579,60 +596,144 @@ class CodexRateLimitReader:
|
|||||||
)
|
)
|
||||||
return self.cached
|
return self.cached
|
||||||
|
|
||||||
@staticmethod
|
class CcswitchCodexUsageReader:
|
||||||
def _tokens_to_m10(tokens) -> int:
|
DEDUP_WINDOW_SECONDS = 10 * 60
|
||||||
if not isinstance(tokens, (int, float)) or tokens < 0:
|
|
||||||
return -1
|
def __init__(self, refresh_seconds: float = CCSWITCH_CODEX_REFRESH_SECONDS) -> None:
|
||||||
return max(0, min(2147483647, int(round(tokens / 100000))))
|
self.refresh_seconds = refresh_seconds
|
||||||
|
self.next_refresh = 0.0
|
||||||
|
self.cached = (-1, "")
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _daily_bars_from_buckets(buckets: list[dict]) -> tuple[int, str]:
|
def _candidate_paths() -> list[Path]:
|
||||||
today = date.today()
|
paths: list[Path] = []
|
||||||
wanted_days = [(today - timedelta(days=offset)).isoformat() for offset in range(CODEX_DAILY_BAR_COUNT - 1, -1, -1)]
|
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] = {}
|
tokens_by_day: dict[str, int] = {}
|
||||||
|
with self._connect_readonly(db_path) as con:
|
||||||
for bucket in buckets:
|
for row in con.execute(sql, (start_ts,)):
|
||||||
if not isinstance(bucket, dict):
|
usage_day = row["usage_day"]
|
||||||
continue
|
tokens = row["real_total_tokens"]
|
||||||
start_date = bucket.get("startDate")
|
if isinstance(usage_day, str) and isinstance(tokens, (int, float)) and tokens >= 0:
|
||||||
tokens = bucket.get("tokens")
|
tokens_by_day[usage_day] = int(tokens)
|
||||||
if isinstance(start_date, str) and isinstance(tokens, (int, float)) and tokens >= 0:
|
return tokens_by_day
|
||||||
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)
|
|
||||||
|
|
||||||
def read_usage_snapshot(self) -> tuple[int, str]:
|
def read_usage_snapshot(self) -> tuple[int, str]:
|
||||||
now = time.monotonic()
|
now = time.monotonic()
|
||||||
if now < self.next_usage_refresh:
|
if now < self.next_refresh:
|
||||||
return self.usage_cached
|
return self.cached
|
||||||
|
|
||||||
if not self._ensure_started():
|
db_path = self._find_db()
|
||||||
self.next_usage_refresh = now + min(15.0, self.usage_refresh_seconds)
|
if db_path is None:
|
||||||
return self.usage_cached
|
self.next_refresh = now + min(15.0, self.refresh_seconds)
|
||||||
|
return self.cached
|
||||||
|
|
||||||
request_id = self._send("account/usage/read")
|
try:
|
||||||
if request_id is None:
|
tokens_by_day = self._query_daily_tokens(db_path)
|
||||||
self.next_usage_refresh = now + min(15.0, self.usage_refresh_seconds)
|
except sqlite3.Error:
|
||||||
return self.usage_cached
|
self.next_refresh = now + min(15.0, self.refresh_seconds)
|
||||||
|
return self.cached
|
||||||
|
|
||||||
message = self._wait_for(request_id, 8.0)
|
self.cached = daily_bars_from_token_map(tokens_by_day, CODEX_DAILY_BAR_COUNT)
|
||||||
result = message.get("result") if message else None
|
self.next_refresh = now + self.refresh_seconds
|
||||||
buckets = result.get("dailyUsageBuckets") if isinstance(result, dict) else None
|
return self.cached
|
||||||
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
|
|
||||||
|
|
||||||
|
|
||||||
class ZhipuStatusReader:
|
class ZhipuStatusReader:
|
||||||
@ -824,6 +925,7 @@ class MetricsSampler:
|
|||||||
self.cpu_frequency = CpuFrequencyReader()
|
self.cpu_frequency = CpuFrequencyReader()
|
||||||
self.disk_rates = LogicalDiskRateReader(self._disk_labels())
|
self.disk_rates = LogicalDiskRateReader(self._disk_labels())
|
||||||
self.codex = CodexRateLimitReader()
|
self.codex = CodexRateLimitReader()
|
||||||
|
self.codex_usage = CcswitchCodexUsageReader()
|
||||||
self.zhipu = ZhipuStatusReader()
|
self.zhipu = ZhipuStatusReader()
|
||||||
psutil.cpu_percent(interval=None)
|
psutil.cpu_percent(interval=None)
|
||||||
|
|
||||||
@ -927,7 +1029,7 @@ class MetricsSampler:
|
|||||||
|
|
||||||
disk_count, disk_percent, disk_text = self._disk_summary(disk_rates)
|
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_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()
|
zhipu_status, zhipu_5h, zhipu_week, zhipu_5h_reset, zhipu_week_reset, zhipu_today_tokens_m10, zhipu_daily_bars = self.zhipu.read()
|
||||||
|
|
||||||
return (
|
return (
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user