186 lines
10 KiB
Plaintext
186 lines
10 KiB
Plaintext
@staticmethod
|
|
def _usage_record_values(record: dict[str, Any]) -> tuple:
|
|
"""Validate a metadata-only event; never retain response or prompt bodies."""
|
|
if not isinstance(record, dict):
|
|
raise ValueError("model usage event must be a mapping")
|
|
|
|
def text(name: str, maximum: int = 256) -> str:
|
|
value = str(record.get(name) or "").strip()
|
|
if len(value) > maximum:
|
|
raise ValueError(f"model usage {name} is too long")
|
|
return value
|
|
|
|
def integer(name: str, *, nullable: bool = False, default: int = 0) -> int | None:
|
|
value = record.get(name)
|
|
if value is None:
|
|
return None if nullable else default
|
|
if isinstance(value, bool) or not isinstance(value, int):
|
|
raise ValueError(f"model usage {name} must be an integer or null")
|
|
if value < 0 or value > 9_223_372_036_854_775_807:
|
|
raise ValueError(f"model usage {name} is out of range")
|
|
return value
|
|
|
|
event_id = text("event_id")
|
|
if not event_id:
|
|
raise ValueError("model usage event_id is required")
|
|
purpose, role, state = text("purpose"), text("role"), text("status")
|
|
if purpose not in {"chat", "guard", "knowledge"}:
|
|
raise ValueError("model usage purpose is invalid")
|
|
if role not in {"answer", "judge"}:
|
|
raise ValueError("model usage role is invalid")
|
|
if state not in {"success", "error"}:
|
|
raise ValueError("model usage status is invalid")
|
|
attempt = integer("attempt", default=1)
|
|
if not attempt:
|
|
raise ValueError("model usage attempt must be positive")
|
|
created = record.get("created_at")
|
|
if created:
|
|
try:
|
|
stamp = datetime.fromisoformat(str(created).replace("Z", "+00:00"))
|
|
except (TypeError, ValueError) as exc:
|
|
raise ValueError("model usage created_at must be an ISO timestamp") from exc
|
|
if stamp.tzinfo is None:
|
|
raise ValueError("model usage created_at requires a timezone")
|
|
stamp = stamp.astimezone(timezone.utc)
|
|
else:
|
|
stamp = datetime.now(timezone.utc)
|
|
return (
|
|
event_id, text("request_id"), text("task_id"), text("tenant_id"),
|
|
integer("desktop_account_id", nullable=True), text("provider_id"),
|
|
text("provider_name", 512), text("model", 512), text("kind", 80),
|
|
purpose, role, attempt, state,
|
|
*(integer(name, nullable=True) for name in (
|
|
"input_tokens", "output_tokens", "total_tokens",
|
|
"cached_input_tokens", "reasoning_tokens")),
|
|
integer("latency_ms"), stamp.isoformat(timespec="milliseconds"),
|
|
)
|
|
|
|
def record_model_usage(self, record: dict[str, Any]) -> bool:
|
|
"""Append one actual upstream attempt, idempotently by event_id.
|
|
|
|
Raise on invalid data/storage failures so the caller can retry or warn;
|
|
False means this event already exists. No token values are estimated.
|
|
"""
|
|
return bool(self.record_model_usage_batch([record]))
|
|
|
|
def record_model_usage_batch(self, records: list[dict[str, Any]]) -> int:
|
|
"""Atomically persist a request's attempts; return newly inserted rows."""
|
|
values = [self._usage_record_values(record) for record in records]
|
|
if not values:
|
|
return 0
|
|
with self.connect() as db:
|
|
db.execute("BEGIN")
|
|
before = db.total_changes
|
|
db.executemany(
|
|
"""INSERT INTO model_usage_events (
|
|
event_id,request_id,task_id,tenant_id,desktop_account_id,
|
|
provider_id,provider_name,model,kind,purpose,role,attempt,status,
|
|
input_tokens,output_tokens,total_tokens,cached_input_tokens,
|
|
reasoning_tokens,latency_ms,created_at)
|
|
VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
|
|
ON CONFLICT(event_id) DO NOTHING""",
|
|
values,
|
|
)
|
|
inserted = db.total_changes - before
|
|
return inserted
|
|
|
|
@staticmethod
|
|
def _model_usage_window(days: int) -> tuple[int, str, str]:
|
|
days = max(1, min(int(days or 7), 90))
|
|
today = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0)
|
|
since = today - timedelta(days=days - 1)
|
|
until = today + timedelta(days=1)
|
|
return days, since.isoformat(timespec="milliseconds"), until.isoformat(timespec="milliseconds")
|
|
|
|
@staticmethod
|
|
def _model_usage_where(since: str, until: str, tenant_id: Any,
|
|
purpose: str, provider_id: str, model: str) -> tuple[str, list[Any]]:
|
|
tenant_clause, tenant_params = tenant_filter(tenant_id)
|
|
clauses = ["created_at >= ?", "created_at < ?", tenant_clause]
|
|
params: list[Any] = [since, until, *tenant_params]
|
|
for column, value in (("purpose", purpose), ("provider_id", provider_id), ("model", model)):
|
|
if value:
|
|
clauses.append(f"{column} = ?")
|
|
params.append(str(value))
|
|
return " AND ".join(clauses), params
|
|
|
|
@staticmethod
|
|
def _model_usage_aggregation_sql() -> str:
|
|
fields = ["COUNT(*) AS request_count", "COUNT(total_tokens) AS known_count",
|
|
"COUNT(*) - COUNT(total_tokens) AS missing_count"]
|
|
for name in ("input", "output", "total", "cached_input", "reasoning"):
|
|
column = name + "_tokens"
|
|
# With calls but no reported values, NULL is not the same as zero.
|
|
fields.append(f"CASE WHEN COUNT(*) = 0 THEN 0 ELSE SUM({column}) END AS {column}")
|
|
if name != "total":
|
|
fields.append(f"COUNT({column}) AS {name}_known_count")
|
|
return ", ".join(fields)
|
|
|
|
def model_usage_stats(self, days: int = 7, tenant_id: Any = "default",
|
|
purpose: str = "", provider_id: str = "", model: str = "") -> dict[str, Any]:
|
|
"""Actual upstream attempts, including retries/judges/errors; UTC days."""
|
|
days, since, until = self._model_usage_window(days)
|
|
where, params = self._model_usage_where(since, until, tenant_id, purpose, provider_id, model)
|
|
aggregates = self._model_usage_aggregation_sql()
|
|
with self.connect() as db:
|
|
db.execute("BEGIN")
|
|
summary = dict(db.execute(
|
|
f"SELECT {aggregates} FROM model_usage_events WHERE {where}", params).fetchone())
|
|
by_model = [dict(row) for row in db.execute(
|
|
f"""SELECT provider_id, MAX(provider_name) AS provider_name, model, {aggregates}
|
|
FROM model_usage_events WHERE {where}
|
|
GROUP BY provider_id, model ORDER BY request_count DESC, provider_id, model""", params)]
|
|
by_purpose = [dict(row) for row in db.execute(
|
|
f"""SELECT purpose, {aggregates} FROM model_usage_events WHERE {where}
|
|
GROUP BY purpose ORDER BY purpose""", params)]
|
|
daily_rows = {row["date"]: dict(row) for row in db.execute(
|
|
f"""SELECT substr(created_at, 1, 10) AS date, {aggregates}
|
|
FROM model_usage_events WHERE {where} GROUP BY date ORDER BY date""", params)}
|
|
# Old orchestration logs are a separate coverage indication, never
|
|
# backfilled into token events. datetime() normalises historical
|
|
# timestamps that may have been written with a local UTC offset.
|
|
tenant_clause, tenant_params = tenant_filter(tenant_id, "calls.tenant_id")
|
|
legacy_where = ["datetime(calls.created_at) >= datetime(?)",
|
|
"datetime(calls.created_at) < datetime(?)", tenant_clause]
|
|
legacy_params: list[Any] = [since, until, *tenant_params]
|
|
if purpose:
|
|
legacy_where.append("calls.purpose = ?")
|
|
legacy_params.append(purpose)
|
|
legacy_where.append("""NOT EXISTS (
|
|
SELECT 1 FROM model_usage_events AS usage
|
|
WHERE calls.task_id != '' AND usage.task_id != '' AND usage.task_id = calls.task_id
|
|
AND (usage.tenant_id = calls.tenant_id OR
|
|
(usage.tenant_id IN ('', 'default') AND calls.tenant_id IN ('', 'default')))
|
|
)""")
|
|
legacy_unmetered = db.execute(
|
|
f"SELECT COUNT(*) FROM model_calls AS calls WHERE {' AND '.join(legacy_where)}",
|
|
legacy_params).fetchone()[0]
|
|
zero = {key: 0 for key in summary}
|
|
first = datetime.fromisoformat(since)
|
|
daily = []
|
|
for index in range(days):
|
|
date = (first + timedelta(days=index)).date().isoformat()
|
|
daily.append(daily_rows.get(date, {"date": date, **zero}))
|
|
return {
|
|
"days": days, "since": since, "timezone": "UTC", "summary": summary,
|
|
"by_model": by_model, "by_purpose": by_purpose, "daily": daily,
|
|
"legacy_unmetered_calls": int(legacy_unmetered),
|
|
"legacy_unmetered_scope": "account_time_purpose",
|
|
"legacy_unmetered_note": "旧调用按账号、UTC日期和用途统计；没有task_id的旧记录无法关联，亦无法按模型或服务商准确筛选。",
|
|
}
|
|
|
|
def list_model_usage(self, days: int = 7, limit: int = 50, offset: int = 0,
|
|
tenant_id: Any = "default", purpose: str = "",
|
|
provider_id: str = "", model: str = "") -> dict[str, Any]:
|
|
"""Paginated metadata only, retaining NULL for unavailable usage."""
|
|
_, since, until = self._model_usage_window(days)
|
|
limit, offset = max(1, min(int(limit or 50), 200)), max(0, int(offset or 0))
|
|
where, params = self._model_usage_where(since, until, tenant_id, purpose, provider_id, model)
|
|
with self.connect() as db:
|
|
db.execute("BEGIN")
|
|
total = db.execute(f"SELECT COUNT(*) FROM model_usage_events WHERE {where}", params).fetchone()[0]
|
|
rows = db.execute(
|
|
f"""SELECT * FROM model_usage_events WHERE {where}
|
|
ORDER BY created_at DESC, event_id DESC LIMIT ? OFFSET ?""", [*params, limit, offset]).fetchall()
|
|
return {"total": int(total), "items": [dict(row) for row in rows]}
|