6cc1efd8f9
Voice replies had to come from the session's main model, so a heavy model made every spoken answer slow (#83894). A new auxiliary slot, voice_chat, picks the model for voice turns: - CLI/TUI voice mode, the Desktop voice conversation (prompt.submit voice_turn) and messaging voice notes mark the turn; agent/voice_turn_route.py binds the route after the prompt, row and preflight compaction were settled against the main model and undoes it before memory sync and background review, so the next typed turn and the review run on the main model. - A turn larger than the voice model's compression threshold stays on the main model with a one-time notice: switching to voice never compacts the conversation. - Usage lands in session_model_usage under task voice_chat; it never rewrites the session's recorded model/route, so resume keeps the model the user picked. - Slot appears in Desktop Settings -> Models, the dashboard, and `hermes model`; GPT-Live delegations keep using the session model.
436 lines
25 KiB
Python
436 lines
25 KiB
Python
"""Token/usage accounting mixin for SessionDB: the coalescing background token writer,
|
|
per-model usage rows, and billing-route columns. Writer thread state lives on the instance."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import atexit
|
|
import contextlib
|
|
import logging
|
|
import threading
|
|
import time
|
|
import weakref
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
# caplog tests pin the "hermes_state" logger name.
|
|
logger = logging.getLogger("hermes_state")
|
|
|
|
_TOKEN_COUNTERS = ("input_tokens", "output_tokens", "cache_read_tokens", "cache_write_tokens", "reasoning_tokens")
|
|
|
|
|
|
def _token_update_sql(delta: bool) -> str:
|
|
"""``UPDATE sessions`` for one usage report: *delta* adds to the stored counters (CLI
|
|
per-call path), otherwise sets them (gateway cumulative path). Cost/route columns
|
|
COALESCE-fill either way (statement text is pinned by the SQL trace harness)."""
|
|
def add(col: str) -> str: # "col + ?" / "COALESCE(col, 0) + ?" in delta mode, bare "?" otherwise
|
|
return f"{col} + ?" if delta else "?"
|
|
def add0(col: str) -> str:
|
|
return f"COALESCE({col}, 0) + ?" if delta else "?"
|
|
counters = "".join(f" {c} = {add(c)},\n" for c in _TOKEN_COUNTERS)
|
|
estimated = "COALESCE(estimated_cost_usd, 0) + COALESCE(?, 0)" if delta else "COALESCE(?, 0)"
|
|
return (
|
|
"UPDATE sessions SET\n" + counters
|
|
+ f""" estimated_cost_usd = {estimated},
|
|
actual_cost_usd = CASE
|
|
WHEN ? IS NULL THEN actual_cost_usd
|
|
ELSE {add0("actual_cost_usd")}
|
|
END,
|
|
cost_status = COALESCE(?, cost_status),
|
|
cost_source = COALESCE(?, cost_source),
|
|
pricing_version = COALESCE(?, pricing_version),
|
|
billing_provider = COALESCE(billing_provider, ?),
|
|
billing_base_url = COALESCE(billing_base_url, ?),
|
|
billing_mode = COALESCE(billing_mode, ?),
|
|
model = COALESCE(model, ?),
|
|
api_call_count = {add0("api_call_count")}
|
|
WHERE id = ?"""
|
|
)
|
|
|
|
|
|
_TOKEN_UPDATE_ABSOLUTE_SQL = _token_update_sql(delta=False)
|
|
_TOKEN_UPDATE_DELTA_SQL = _token_update_sql(delta=True)
|
|
|
|
_MODEL_USAGE_UPSERT_SQL = """INSERT INTO session_model_usage (
|
|
session_id, model, billing_provider, billing_base_url, billing_mode,
|
|
task, api_call_count, input_tokens, output_tokens,
|
|
cache_read_tokens, cache_write_tokens, reasoning_tokens,
|
|
estimated_cost_usd, actual_cost_usd, cost_status, cost_source,
|
|
first_seen, last_seen
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
ON CONFLICT(session_id, model, billing_provider, billing_base_url, billing_mode, task)
|
|
DO UPDATE SET
|
|
api_call_count = api_call_count + excluded.api_call_count,
|
|
input_tokens = input_tokens + excluded.input_tokens,
|
|
output_tokens = output_tokens + excluded.output_tokens,
|
|
cache_read_tokens = cache_read_tokens + excluded.cache_read_tokens,
|
|
cache_write_tokens = cache_write_tokens + excluded.cache_write_tokens,
|
|
reasoning_tokens = reasoning_tokens + excluded.reasoning_tokens,
|
|
estimated_cost_usd = estimated_cost_usd + excluded.estimated_cost_usd,
|
|
actual_cost_usd = actual_cost_usd + excluded.actual_cost_usd,
|
|
cost_status = COALESCE(excluded.cost_status, cost_status),
|
|
cost_source = COALESCE(excluded.cost_source, cost_source),
|
|
last_seen = excluded.last_seen"""
|
|
|
|
|
|
# Kwargs forwarded verbatim from update_token_counts / record_auxiliary_usage into
|
|
# _record_model_usage (the per-route attribution row).
|
|
_MODEL_USAGE_FIELDS = frozenset((
|
|
"model", "billing_provider", "billing_base_url", "billing_mode", "input_tokens", "output_tokens",
|
|
"cache_read_tokens", "cache_write_tokens", "reasoning_tokens", "estimated_cost_usd",
|
|
"actual_cost_usd", "cost_status", "cost_source", "api_call_count"))
|
|
|
|
|
|
class SessionUsageMixin:
|
|
"""Coalesced token writer, per-model usage rows, billing route."""
|
|
|
|
def update_session_billing_route(
|
|
self, session_id: str, *, provider: str, base_url: str, billing_mode: Optional[str] = None,
|
|
) -> None:
|
|
"""Unconditionally set the billing route (``update_token_counts`` only COALESCE-fills
|
|
NULLs) so the dashboard reflects the latest /model switch.
|
|
|
|
Route writers never touch the stored prompt; ``_stored_prompt_matches_runtime`` decides staleness.
|
|
|
|
See #48173, #48248.
|
|
"""
|
|
# Barrier against queued token deltas — see update_session_model.
|
|
self.flush_token_counts()
|
|
|
|
def _do(conn):
|
|
conn.execute("""UPDATE sessions SET
|
|
billing_provider = ?,
|
|
billing_base_url = ?,
|
|
billing_mode = COALESCE(?, billing_mode)
|
|
WHERE id = ?""", (provider, base_url, billing_mode, session_id))
|
|
self._execute_write(_do)
|
|
|
|
def queue_token_counts(self, session_id: str, **kwargs) -> None:
|
|
"""Enqueue a token/cost delta for the background writer (same kwargs as
|
|
:meth:`update_token_counts`). After close() stopped the writer, falls back to the
|
|
synchronous path and may raise."""
|
|
with self._token_queue_cond:
|
|
thread = self._token_writer_thread
|
|
writer_alive = thread is not None and thread.is_alive()
|
|
writer_stopped = self._token_writer_stop and not writer_alive
|
|
if not writer_stopped:
|
|
self._token_queue.append((session_id, kwargs))
|
|
if not writer_alive:
|
|
# Daemon so exit never hangs on accounting; the atexit hook drains
|
|
# leftovers. ``not is_alive()`` respawns a writer that died unexpectedly.
|
|
thread = threading.Thread(
|
|
target=self._token_writer_loop, name="session-db-token-writer", daemon=True)
|
|
self._token_writer_thread = thread
|
|
thread.start()
|
|
if self._token_atexit_hook is None:
|
|
self_ref = weakref.ref(self)
|
|
|
|
def _drain_at_exit() -> None:
|
|
db = self_ref()
|
|
if db is not None:
|
|
db._drain_token_queue_at_exit()
|
|
|
|
self._token_atexit_hook = _drain_at_exit
|
|
atexit.register(_drain_at_exit)
|
|
self._token_queue_cond.notify_all()
|
|
if writer_stopped:
|
|
# close() ran: enqueueing would drop the delta silently, so apply inline.
|
|
self.update_token_counts(session_id, **kwargs)
|
|
|
|
def _apply_claimed_batch(self, batch) -> None:
|
|
"""Apply a batch whose ``busy`` flag the caller already claimed, then release."""
|
|
try:
|
|
self._apply_token_batch(batch)
|
|
finally:
|
|
with self._token_queue_cond:
|
|
self._token_writer_busy = False
|
|
self._token_queue_cond.notify_all()
|
|
|
|
def flush_token_counts(self, timeout: float = 5.0) -> bool:
|
|
"""Block until every queued token delta has been applied. False on timeout (callers
|
|
then read totals stale by the queued deltas). Never raises."""
|
|
# Lock-free fast path: reads queue-then-busy (see ordering notes below).
|
|
if not self._token_queue and not self._token_writer_busy:
|
|
return True
|
|
batch = None
|
|
with self._token_queue_cond:
|
|
deadline = time.monotonic() + timeout
|
|
while self._token_queue or self._token_writer_busy:
|
|
# A live writer is authoritative even when stop-flagged: draining here would
|
|
# race its in-flight batch and reorder deltas (breaking last-non-None-wins /
|
|
# first-accounted-route / COALESCE-backfill fields). Only a dead writer lets
|
|
# the caller take leftovers; a claimed busy means "wait".
|
|
thread = self._token_writer_thread
|
|
if (thread is None or not thread.is_alive()) and not self._token_writer_busy:
|
|
self._token_writer_busy = True
|
|
batch = list(self._token_queue)
|
|
self._token_queue.clear()
|
|
break
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
return False
|
|
self._token_queue_cond.wait(remaining)
|
|
if batch:
|
|
self._apply_claimed_batch(batch)
|
|
return True
|
|
|
|
def _token_writer_loop(self) -> None:
|
|
while True:
|
|
with self._token_queue_cond:
|
|
idle_deadline = time.monotonic() + self._TOKEN_WRITER_IDLE_SECONDS
|
|
while not self._token_queue and not self._token_writer_stop:
|
|
remaining = idle_deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
# Retire under the lock queue_token_counts() spawns under, so no
|
|
# delta strands behind an exiting worker.
|
|
self._token_writer_thread = None
|
|
return
|
|
self._token_queue_cond.wait(remaining)
|
|
if not self._token_queue:
|
|
self._token_writer_thread = None
|
|
return # stop requested and fully drained
|
|
# busy BEFORE clearing the queue: flush's lock-free fast path must never see
|
|
# "empty and idle" while a popped batch is unapplied.
|
|
self._token_writer_busy = True
|
|
batch = list(self._token_queue)
|
|
self._token_queue.clear()
|
|
self._apply_claimed_batch(batch)
|
|
|
|
def _apply_token_batch(self, batch: List[Tuple[str, Dict[str, Any]]]) -> None:
|
|
"""Apply queued deltas in order, coalescing where safe. Never raises."""
|
|
try:
|
|
coalesced = self._coalesce_token_deltas(batch)
|
|
except Exception as exc:
|
|
# Coalescing must never kill the writer; the merge is only an optimization.
|
|
logger.warning("async token accounting: coalesce failed, applying raw batch: %s", exc)
|
|
coalesced = batch
|
|
for session_id, kwargs in coalesced:
|
|
try:
|
|
self.update_token_counts(session_id, **kwargs)
|
|
except Exception as exc:
|
|
# Accounting loss is logged, never raised into a turn.
|
|
logger.warning("async token accounting: apply failed (session=%s): %s", session_id, exc)
|
|
|
|
def _coalesce_token_deltas(self, batch: List[Tuple[str, Dict[str, Any]]]) -> List[Tuple[str, Dict[str, Any]]]:
|
|
"""Merge adjacent incremental deltas with an identical route, so ordering across
|
|
sessions and /model switches is preserved exactly. absolute=True never merges."""
|
|
groups: List[Tuple[Optional[tuple], str, Dict[str, Any]]] = []
|
|
for session_id, kwargs in batch:
|
|
key = None
|
|
if not kwargs.get("absolute"):
|
|
key = (session_id, *(kwargs.get(f) for f in self._TOKEN_DELTA_ROUTE_FIELDS))
|
|
if groups and key is not None and groups[-1][0] == key:
|
|
merged = groups[-1][2]
|
|
for f in self._TOKEN_DELTA_SUM_FIELDS:
|
|
merged[f] = merged.get(f, 0) + kwargs.get(f, 0)
|
|
for f in self._TOKEN_DELTA_COST_FIELDS:
|
|
value = kwargs.get(f)
|
|
if value is not None:
|
|
# All-None runs stay None so COALESCE keeps the stored value.
|
|
merged[f] = (merged.get(f) or 0.0) + value
|
|
else:
|
|
groups.append((key, session_id, dict(kwargs)))
|
|
return [(sid, kw) for _, sid, kw in groups]
|
|
|
|
def _stop_token_writer(self, join_timeout: float = 10.0) -> None:
|
|
"""Stop the writer thread and drain remaining deltas. Never raises."""
|
|
with self._token_queue_cond:
|
|
self._token_writer_stop = True
|
|
self._token_queue_cond.notify_all()
|
|
thread = self._token_writer_thread
|
|
if thread is not None and thread.is_alive():
|
|
thread.join(timeout=join_timeout)
|
|
if thread.is_alive():
|
|
# Writer stuck mid-apply: leave deltas unapplied rather than race it.
|
|
logger.warning(
|
|
"async token accounting: writer did not stop within %.0fs; "
|
|
"%d queued delta(s) not persisted", join_timeout, len(self._token_queue))
|
|
return
|
|
# Writer gone: apply leftovers synchronously under the same busy protocol. Wait out
|
|
# a flush caller-drain that already claimed busy — close() nulls the connection
|
|
# right after this returns and must not yank it mid-batch.
|
|
with self._token_queue_cond:
|
|
deadline = time.monotonic() + join_timeout
|
|
while self._token_writer_busy:
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
logger.warning(
|
|
"async token accounting: concurrent drain did not "
|
|
"finish within %.0fs; %d queued delta(s) not persisted",
|
|
join_timeout, len(self._token_queue))
|
|
return
|
|
self._token_queue_cond.wait(remaining)
|
|
# busy BEFORE clearing the queue (same ordering as the writer loop).
|
|
batch = list(self._token_queue)
|
|
if batch:
|
|
self._token_writer_busy = True
|
|
self._token_queue.clear()
|
|
if batch:
|
|
self._apply_claimed_batch(batch)
|
|
|
|
def _drain_token_queue_at_exit(self) -> None:
|
|
with contextlib.suppress(Exception): # never fatal at interpreter shutdown
|
|
self._stop_token_writer()
|
|
|
|
def update_token_counts(
|
|
self, session_id: str, input_tokens: int=0, output_tokens: int=0, model: str=None, cache_read_tokens: int=0,
|
|
cache_write_tokens: int=0, reasoning_tokens: int=0, estimated_cost_usd: Optional[float]=None,
|
|
actual_cost_usd: Optional[float]=None, cost_status: Optional[str]=None, cost_source: Optional[str]=None,
|
|
pricing_version: Optional[str]=None, billing_provider: Optional[str]=None, billing_base_url: Optional[str]=None,
|
|
billing_mode: Optional[str]=None, api_call_count: int=0, absolute: bool=False,
|
|
source: Optional[str]=None, task: str = "",
|
|
) -> None:
|
|
"""Update token counters and backfill model if unset. *absolute*=False increments
|
|
(per-API-call deltas, CLI path); *absolute*=True sets directly (gateway path,
|
|
where the cached agent holds cumulative totals). ``source`` is the session's real surface
|
|
for the row-existence guard; callers that don't know it leave the placeholder. ``task`` names a
|
|
per-turn route (``voice_chat``): its calls count in the session totals and land in
|
|
``session_model_usage`` under that task, but never become the session's recorded route."""
|
|
usage = {k: v for k, v in locals().items() if k in _MODEL_USAGE_FIELDS}
|
|
# Ensure the row exists: under concurrent load create_session() may have failed on
|
|
# locking, and the UPDATE would silently affect 0 rows. When this guard is the first
|
|
# writer it must carry the agent's real source: the turn lease treats an existing row as
|
|
# proof the create already happened, so the creator never returns to repair an anonymous
|
|
# ``unknown`` placeholder and the session stays a phantom for life (#111999).
|
|
self._insert_session_row(session_id, source or "unknown", model=model)
|
|
sql = _TOKEN_UPDATE_ABSOLUTE_SQL if absolute else _TOKEN_UPDATE_DELTA_SQL
|
|
has_usage = bool(input_tokens or output_tokens or cache_read_tokens or cache_write_tokens or reasoning_tokens
|
|
or api_call_count or estimated_cost_usd)
|
|
has_accounted_usage = bool(has_usage or actual_cost_usd)
|
|
params = (
|
|
input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, reasoning_tokens,
|
|
estimated_cost_usd, actual_cost_usd, actual_cost_usd, cost_status, cost_source, pricing_version,
|
|
billing_provider if has_accounted_usage else None,
|
|
billing_base_url if has_accounted_usage else None,
|
|
billing_mode if has_accounted_usage else None, model if has_accounted_usage else None,
|
|
api_call_count, session_id)
|
|
# Per-model attribution: the sessions row keeps one (model, provider) pair, so a
|
|
# mid-session /model switch would attribute every token to the initial model. Only
|
|
# the incremental path records here — absolute cumulative updates cannot be split
|
|
# back into routes; Insights reconciles the residual instead.
|
|
# ``update_token_counts`` is the single chokepoint every per-API-call delta flows through (CLI,
|
|
# gateway, cron, delegated runs — see conversation_loop / codex_runtime), and each call carries the
|
|
# model/provider *active at the time of that call*. Recording the per-call delta into
|
|
# session_model_usage keyed by the live model preserves an accurate per-model breakdown regardless
|
|
# of how many times the user switches. See #51607.
|
|
record_model_usage = (not absolute) and has_usage
|
|
|
|
def _do(conn):
|
|
row = conn.execute(
|
|
"SELECT model, billing_provider, api_call_count FROM sessions WHERE id = ?", (session_id,),
|
|
).fetchone()
|
|
existing = dict(row) if row is not None else {}
|
|
# create_session records the requested route before any API call. If that fails
|
|
# and fallback succeeds, the first accounted usage is the authoritative route;
|
|
# after that keep the row as is (one row cannot represent mixed usage).
|
|
first_accounted_route = (
|
|
not task and int(existing.get("api_call_count") or 0) == 0 and has_accounted_usage and bool(model)
|
|
and bool(billing_provider)
|
|
and (existing.get("model") != model or existing.get("billing_provider") != billing_provider)
|
|
)
|
|
if first_accounted_route:
|
|
conn.execute("""UPDATE sessions
|
|
SET model = ?, billing_provider = ?,
|
|
billing_base_url = ?, billing_mode = ?
|
|
WHERE id = ?""", (model, billing_provider, billing_base_url, billing_mode, session_id))
|
|
conn.execute(sql, params)
|
|
if record_model_usage:
|
|
self._record_model_usage(conn, session_id, task=task, **usage)
|
|
self._execute_write(_do)
|
|
|
|
def _record_model_usage(
|
|
self, conn, session_id: str, *, model: Optional[str]=None, billing_provider: Optional[str]=None,
|
|
billing_base_url: Optional[str]=None, billing_mode: Optional[str]=None, input_tokens: int=0,
|
|
output_tokens: int=0, cache_read_tokens: int=0, cache_write_tokens: int=0, reasoning_tokens: int=0,
|
|
estimated_cost_usd: Optional[float]=None, actual_cost_usd: Optional[float]=None,
|
|
cost_status: Optional[str]=None, cost_source: Optional[str]=None, api_call_count: int=0, task: str="",
|
|
) -> None:
|
|
"""Accumulate a per-API-call usage delta into session_model_usage, inside the caller's
|
|
write txn after the ``sessions`` UPDATE. A missing model/provider falls back to
|
|
the session row — except for aux rows (``task`` set), which must NOT inherit the
|
|
main-loop route (vision on gemini while the main loop runs anthropic): missing
|
|
info stays 'unknown'/empty.
|
|
|
|
``task`` distinguishes what kind of work consumed the tokens: ``''`` (empty) is the main agent loop;
|
|
auxiliary calls record their task name (``vision``, ``compression``, ``title_generation``, ...) via
|
|
:meth:`record_auxiliary_usage` (issue #23270).
|
|
"""
|
|
row = conn.execute(
|
|
"SELECT model, billing_provider, billing_base_url, billing_mode FROM sessions WHERE id = ?", (session_id,),
|
|
).fetchone()
|
|
sess = dict(row) if (row is not None and not task) else {}
|
|
counts = [v or 0 for v in (input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, reasoning_tokens)]
|
|
now = time.time()
|
|
conn.execute(_MODEL_USAGE_UPSERT_SQL, (
|
|
session_id, model or sess.get("model") or "unknown",
|
|
billing_provider or sess.get("billing_provider") or "",
|
|
billing_base_url or sess.get("billing_base_url") or "",
|
|
billing_mode or sess.get("billing_mode") or "", task or "", api_call_count or 0, *counts,
|
|
float(estimated_cost_usd or 0.0), float(actual_cost_usd or 0.0), cost_status, cost_source, now, now))
|
|
|
|
def record_auxiliary_usage(
|
|
self, session_id: str, task: str, *, model: Optional[str]=None, billing_provider: Optional[str]=None,
|
|
billing_base_url: Optional[str]=None, input_tokens: int=0, output_tokens: int=0, cache_read_tokens: int=0,
|
|
cache_write_tokens: int=0, reasoning_tokens: int=0, estimated_cost_usd: Optional[float]=None,
|
|
api_call_count: int=1,
|
|
) -> None:
|
|
"""Record an auxiliary LLM call's usage (vision, compression, title generation, ...)
|
|
as a per-(model, provider, task) delta in ``session_model_usage`` WITHOUT touching
|
|
the ``sessions`` summary row (the gateway overwrites those counters with absolute
|
|
main-loop totals). ``api_call_count`` may aggregate N calls. Best-effort.
|
|
|
|
See #23270.
|
|
Background-review forks record an aggregate of N fork API calls in one write with
|
|
``task='background_review'`` (issue #87250).
|
|
"""
|
|
usage = {k: v for k, v in locals().items() if k in _MODEL_USAGE_FIELDS}
|
|
if not session_id or not task:
|
|
return
|
|
usage["api_call_count"] = 1 if api_call_count is None else int(api_call_count)
|
|
# FK to sessions.id: same guard as update_token_counts; the aux path carries no surface, so
|
|
# the placeholder stays repairable by the creator's upsert (_insert_session_row).
|
|
self._insert_session_row(session_id, "unknown")
|
|
self._execute_write(lambda conn: self._record_model_usage(conn, session_id, task=task, **usage))
|
|
|
|
def auxiliary_usage_by_task(self, session_id: str) -> Dict[str, Dict[str, float]]:
|
|
"""Per-task auxiliary usage (``task != ''``: vision, compression, title_generation, ...) summed
|
|
over the session's compression lineage. Aux calls bill to the id the turn STARTED with while
|
|
compression mints child ids mid-turn, so a single-id read misses rows (#112848)."""
|
|
if not session_id:
|
|
return {}
|
|
chain = self._session_lineage_root_to_tip(session_id)
|
|
# An explicit ``/branch`` copy owns its spend: cut the walk at the nearest branch node so a
|
|
# resumed branch never absorbs aux rows billed to its source (hermes_state_messages does the same).
|
|
for i in range(len(chain) - 1, -1, -1):
|
|
if self._is_explicit_branch_session(chain[i]):
|
|
chain = chain[i:]
|
|
break
|
|
rows = self._read_all(
|
|
f"""SELECT task,
|
|
COALESCE(SUM(api_call_count), 0) AS api_calls,
|
|
COALESCE(SUM(input_tokens), 0) AS input_tokens,
|
|
COALESCE(SUM(output_tokens), 0) AS output_tokens,
|
|
COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens,
|
|
COALESCE(SUM(cache_write_tokens), 0) AS cache_write_tokens,
|
|
COALESCE(SUM(reasoning_tokens), 0) AS reasoning_tokens,
|
|
COALESCE(SUM(estimated_cost_usd), 0) AS estimated_cost_usd
|
|
FROM session_model_usage
|
|
WHERE session_id IN ({','.join('?' * len(chain))}) AND task != ''
|
|
GROUP BY task""",
|
|
chain,
|
|
)
|
|
return {row["task"]: {k: row[k] for k in row.keys() if k != "task"} for row in rows}
|
|
|
|
def usage_totals(self, *, min_message_count: int = 1, include_archived: bool = False) -> Dict[str, float]:
|
|
"""Tokens and spend across the whole store (one scan), so the sidebar total does not
|
|
shrink with paging. Spend prefers the billed figure over the estimate."""
|
|
where = ["parent_session_id IS NULL", "message_count >= ?"]
|
|
params: List[Any] = [min_message_count]
|
|
if not include_archived:
|
|
where.append("COALESCE(archived, 0) = 0")
|
|
row = self._read_one(f"""
|
|
SELECT COALESCE(SUM(COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)), 0),
|
|
COALESCE(SUM(COALESCE(actual_cost_usd, estimated_cost_usd, 0)), 0)
|
|
FROM sessions
|
|
WHERE {' AND '.join(where)}
|
|
""", params)
|
|
return {"tokens": int(row[0] or 0), "cost_usd": float(row[1] or 0.0)}
|