hermes-admin/hermes-agent
1
0
Fork 0
Code Issues Pull Requests Actions 28 Packages Projects Releases Wiki Activity

refactor(hermes_state_usage): compact docstrings/comments

Browse Source
This commit is contained in:
Teknium
2026-09-02 18:24:01 -07:00
parent c7a6c5ac39
commit ab5a64a7a4
1 changed files with 43 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

90
hermes_state_usage.py
View File

@@ -1,6 +1,5 @@
"""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 SessionDB instance."""
"""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
@@ -92,9 +91,9 @@ class SessionUsageMixin:
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. Also
nulls ``system_prompt`` so the cached snapshot header is rebuilt."""
"""Unconditionally set the billing route (``update_token_counts`` only COALESCE-fills
NULLs) so the dashboard reflects the latest /model switch; also nulls
``system_prompt`` so the cached snapshot header is rebuilt."""
# Barrier against queued token deltas — see update_session_model.
self.flush_token_counts()
@@ -114,8 +113,8 @@ class SessionUsageMixin:
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() has stopped the writer, falls back
to the synchronous path and may raise."""
: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()
@@ -124,8 +123,7 @@ class SessionUsageMixin:
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()`` (not ``is None``) respawns a writer
# that died from an unexpected escape.
# 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
@@ -155,8 +153,8 @@ class SessionUsageMixin:
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."""
"""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
@@ -164,11 +162,10 @@ class SessionUsageMixin:
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", never "drain alongside".
# 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
@@ -190,16 +187,16 @@ class SessionUsageMixin:
while not self._token_queue and not self._token_writer_stop:
remaining = idle_deadline - time.monotonic()
if remaining <= 0:
# Retire under the same lock queue_token_counts() uses to decide
# to spawn, so no delta strands behind an exiting worker.
# 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.
# 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()
@@ -255,9 +252,9 @@ class SessionUsageMixin:
"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.
# 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:
@@ -292,12 +289,12 @@ class SessionUsageMixin:
billing_base_url: Optional[str] = None, billing_mode: Optional[str] = None,
api_call_count: int = 0, absolute: bool = False,
) -> 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)."""
"""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)."""
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.
# Ensure the row exists: under concurrent load create_session() may have failed on
# locking, and the UPDATE would silently affect 0 rows.
self._insert_session_row(session_id, "unknown", model=model)
sql = _TOKEN_UPDATE_ABSOLUTE_SQL if absolute else _TOKEN_UPDATE_DELTA_SQL
has_usage = bool(
@@ -313,9 +310,9 @@ class SessionUsageMixin:
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.
# 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.
record_model_usage = (not absolute) and has_usage
def _do(conn):
@@ -323,9 +320,9 @@ class SessionUsageMixin:
"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).
# 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 = (
int(existing.get("api_call_count") or 0) == 0 and has_accounted_usage and bool(model)
and bool(billing_provider)
@@ -352,11 +349,11 @@ class SessionUsageMixin:
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 (same COALESCE behaviour as the summary update) — 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."""
"""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."""
row = conn.execute(
"SELECT model, billing_provider, billing_base_url, billing_mode "
"FROM sessions WHERE id = ?", (session_id,),
@@ -382,11 +379,10 @@ class SessionUsageMixin:
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:
callers must never fail an aux call over accounting."""
"""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."""
usage = {k: v for k, v in locals().items() if k in _MODEL_USAGE_FIELDS}
if not session_id or not task:
return
@@ -396,8 +392,8 @@ class SessionUsageMixin:
self._execute_write(lambda conn: self._record_model_usage(conn, session_id, task=task, **usage))
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."""
"""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: