Spaces:
Running
Running
File size: 1,583 Bytes
a31f556 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 | import sqlite3
import uuid
import logging
from datetime import datetime, timezone
from backend.services.usb_monitor import get_db_path
from backend.services.usb_vault import KeyDomain
def track_token_usage(domain: KeyDomain, tokens: int):
"""Logs token usage for a specific KeyDomain."""
try:
db_path = get_db_path()
with sqlite3.connect(db_path) as conn:
conn.execute('PRAGMA journal_mode=WAL')
conn.execute(
"INSERT INTO key_usage (id, domain, tokens_used, timestamp) VALUES (?, ?, ?, ?)",
(str(uuid.uuid4()), domain.value, tokens, datetime.now(timezone.utc).isoformat())
)
conn.commit()
except Exception as e:
logging.error(f"Failed to track token usage for {domain.value}: {e}")
def get_today_usage_by_domain():
"""Returns aggregated token usage per domain for the current day."""
try:
db_path = get_db_path()
today = datetime.now(timezone.utc).strftime('%Y-%m-%d')
with sqlite3.connect(db_path) as conn:
conn.execute('PRAGMA journal_mode=WAL')
cursor = conn.cursor()
# SQLite DATE() extracts YYYY-MM-DD from ISO format string
cursor.execute(
"SELECT domain, SUM(tokens_used) FROM key_usage WHERE DATE(timestamp) = ? GROUP BY domain",
(today,)
)
rows = cursor.fetchall()
return {row[0]: row[1] for row in rows}
except Exception as e:
logging.error(f"Failed to fetch key usage: {e}")
return {}
|