Merge branch 'simp/r2-state-b' into simp/r2-state
This commit is contained in:
Teknium
+234
-477
File diff suppressed because it is too large
Load Diff
+763
-1756
File diff suppressed because it is too large
Load Diff
+80
-169
@@ -13,66 +13,44 @@ from hermes_state_common import _COMPRESSION_CHILD_SQL, escape_like as _escape_l
|
||||
# caplog tests pin the "hermes_state" logger name.
|
||||
logger = logging.getLogger("hermes_state")
|
||||
|
||||
# ASCII controls (keeping \t \n \r for the whitespace collapse), then zero-width,
|
||||
# bidi override, object-replacement and interlinear-annotation code points.
|
||||
_TITLE_CONTROL_RE = re.compile(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]')
|
||||
_TITLE_INVISIBLE_RE = re.compile(r'[\u200b-\u200f\u2028-\u202e\u2060-\u2069\ufeff\ufffc\ufff9-\ufffb]')
|
||||
_NUMBERED_TITLE_RE = re.compile(r'^(.*?) #(\d+)$')
|
||||
|
||||
|
||||
class SessionTitlesMixin:
|
||||
"""Sanitizing, ranking auto/user titles, lineage-aware lookups."""
|
||||
|
||||
@classmethod
|
||||
def _title_rank(cls, source: Optional[str]) -> int:
|
||||
"""Rank a stored title_source.
|
||||
|
||||
NULL (pre-provenance rows) is indistinguishable from a manual ``/title``
|
||||
of that era, so it ranks as ``user``: auto-titling only ever fills
|
||||
genuinely empty legacy titles.
|
||||
"""
|
||||
"""Rank a stored title_source. NULL (pre-provenance rows) is indistinguishable
|
||||
from a manual ``/title`` of that era, so it ranks as ``user``."""
|
||||
if source is None:
|
||||
return cls._TITLE_SOURCE_RANK[cls.TITLE_SOURCE_USER]
|
||||
return cls._TITLE_SOURCE_RANK.get(str(source), 0)
|
||||
|
||||
@staticmethod
|
||||
def sanitize_title(title: Optional[str]) -> Optional[str]:
|
||||
"""Strip control/zero-width/bidi chars, collapse whitespace, normalize
|
||||
empty to None. Raises ValueError if longer than MAX_TITLE_LENGTH
|
||||
after cleaning."""
|
||||
"""Strip control/zero-width/bidi chars (and lone surrogates sqlite3 cannot
|
||||
bind), collapse whitespace, normalize empty to None. Raises ValueError if
|
||||
longer than MAX_TITLE_LENGTH after cleaning."""
|
||||
from hermes_state import SessionDB
|
||||
if not title:
|
||||
return None
|
||||
|
||||
# Lone surrogates cannot be bound by sqlite3 (UnicodeEncodeError).
|
||||
title = _sanitize_surrogates(title)
|
||||
|
||||
# ASCII controls, keeping \t \n \r so the whitespace collapse below
|
||||
# turns them into spaces.
|
||||
cleaned = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]', '', title)
|
||||
|
||||
# Zero-width, bidi override, object-replacement, interlinear annotation.
|
||||
cleaned = re.sub(
|
||||
r'[\u200b-\u200f\u2028-\u202e\u2060-\u2069\ufeff\ufffc\ufff9-\ufffb]',
|
||||
'', cleaned,
|
||||
)
|
||||
|
||||
cleaned = _TITLE_INVISIBLE_RE.sub('', _TITLE_CONTROL_RE.sub('', _sanitize_surrogates(title)))
|
||||
cleaned = re.sub(r'\s+', ' ', cleaned).strip()
|
||||
|
||||
if not cleaned:
|
||||
return None
|
||||
|
||||
if len(cleaned) > SessionDB.MAX_TITLE_LENGTH:
|
||||
raise ValueError(
|
||||
f"Title too long ({len(cleaned)} chars, max {SessionDB.MAX_TITLE_LENGTH})"
|
||||
)
|
||||
|
||||
raise ValueError(f"Title too long ({len(cleaned)} chars, max {SessionDB.MAX_TITLE_LENGTH})")
|
||||
return cleaned
|
||||
|
||||
def _is_compression_ancestor(
|
||||
self, conn, *, ancestor_id: str, descendant_id: str
|
||||
) -> bool:
|
||||
"""True if *ancestor_id* is a compression predecessor of *descendant_id*.
|
||||
|
||||
Uses the canonical continuation edge ``_COMPRESSION_CHILD_SQL`` (parent
|
||||
ended with ``end_reason = 'compression'`` and child started at/after its
|
||||
``ended_at``), which excludes delegate/branch children that also carry
|
||||
``parent_session_id``. One recursive CTE so the edge is defined once.
|
||||
"""
|
||||
def _is_compression_ancestor(self, conn, *, ancestor_id: str, descendant_id: str) -> bool:
|
||||
"""True if *ancestor_id* is a compression predecessor of *descendant_id*, via the
|
||||
canonical continuation edge ``_COMPRESSION_CHILD_SQL`` (excludes delegate/branch
|
||||
children that also carry ``parent_session_id``)."""
|
||||
if not ancestor_id or not descendant_id or ancestor_id == descendant_id:
|
||||
return False
|
||||
edge = _COMPRESSION_CHILD_SQL.format(a="child")
|
||||
@@ -93,23 +71,15 @@ class SessionTitlesMixin:
|
||||
).fetchone()
|
||||
return row is not None
|
||||
|
||||
def _set_session_title(
|
||||
self,
|
||||
session_id: str,
|
||||
title: str,
|
||||
*,
|
||||
source: str,
|
||||
) -> bool:
|
||||
def _set_session_title(self, session_id: str, title: str, *, source: str) -> bool:
|
||||
"""Write a title, enforcing provenance precedence.
|
||||
|
||||
A ``user`` write always lands. ``derived``/``llm`` land only when the
|
||||
row is untitled or holds strictly lower authority, so derived upgrades
|
||||
to llm exactly once, nothing overwrites a user name, and re-running the
|
||||
titler on an llm row is a no-op (stops sessions renaming themselves).
|
||||
No writer may move a hidden canonical Bot Chat off its title.
|
||||
|
||||
Read and write are one compare-and-swap in a single transaction, so a
|
||||
manual ``/title`` racing an in-flight generation is not clobbered.
|
||||
A ``user`` write always lands. ``derived``/``llm`` land only when the row is
|
||||
untitled or holds strictly lower authority (derived upgrades to llm exactly once,
|
||||
nothing overwrites a user name, re-running the titler on an llm row is a no-op).
|
||||
No writer may move a hidden canonical Bot Chat off its title. Read and write are
|
||||
one compare-and-swap transaction, so a manual ``/title`` racing an in-flight
|
||||
generation is not clobbered.
|
||||
"""
|
||||
title = self.sanitize_title(title)
|
||||
is_user = source == self.TITLE_SOURCE_USER
|
||||
@@ -117,18 +87,14 @@ class SessionTitlesMixin:
|
||||
|
||||
def _do(conn):
|
||||
current = conn.execute(
|
||||
"SELECT title, title_source, hidden FROM sessions WHERE id = ?",
|
||||
(session_id,),
|
||||
"SELECT title, title_source, hidden FROM sessions WHERE id = ?", (session_id,),
|
||||
).fetchone()
|
||||
if current is None:
|
||||
return 0
|
||||
# The canonical Bot Chat's NAME is its identity: Bot Mode resolves it
|
||||
# by exact-title lookup on every open, so a rename orphans the whole
|
||||
# conversation (next open mints an empty replacement and UNIQUE(title)
|
||||
# blocks renaming back). Refuse here, the single write path every
|
||||
# surface funnels through. Hidden is the discriminator: canonical
|
||||
# chats are born hidden; a visible session merely named "Bot Chat"
|
||||
# stays renameable. Provenance-blind so the auto-titler no-ops too.
|
||||
# The canonical Bot Chat's NAME is its identity (Bot Mode resolves it by
|
||||
# exact-title lookup on every open), so a rename orphans the conversation.
|
||||
# Hidden is the discriminator: canonical chats are born hidden; a visible
|
||||
# session merely named "Bot Chat" stays renameable. Provenance-blind.
|
||||
if (
|
||||
(current["title"] or "") == self.CANONICAL_BOT_CHAT_TITLE
|
||||
and bool(current["hidden"])
|
||||
@@ -141,100 +107,64 @@ class SessionTitlesMixin:
|
||||
"To start fresh, create a new bot instead."
|
||||
)
|
||||
return 0
|
||||
if not is_user and current["title"] is not None:
|
||||
if self._title_rank(current["title_source"]) >= new_rank:
|
||||
return 0
|
||||
|
||||
if not is_user and current["title"] is not None and self._title_rank(current["title_source"]) >= new_rank:
|
||||
return 0
|
||||
if title:
|
||||
cursor = conn.execute(
|
||||
"SELECT id FROM sessions WHERE title = ? AND id != ?",
|
||||
(title, session_id),
|
||||
)
|
||||
conflict = cursor.fetchone()
|
||||
conflict = conn.execute(
|
||||
"SELECT id FROM sessions WHERE title = ? AND id != ?", (title, session_id),
|
||||
).fetchone()
|
||||
if conflict:
|
||||
conflict_id = conflict["id"]
|
||||
# If the conflicting holder is a hidden compressed ancestor
|
||||
# of this continuation, the user cannot free the title, so
|
||||
# transfer it onto the tip. Uniqueness and lineage are kept.
|
||||
if self._is_compression_ancestor(
|
||||
conn, ancestor_id=conflict_id, descendant_id=session_id
|
||||
):
|
||||
conn.execute(
|
||||
"UPDATE sessions SET title = NULL WHERE id = ?",
|
||||
(conflict_id,),
|
||||
)
|
||||
# A hidden compressed ancestor holding the title cannot be freed by
|
||||
# the user, so transfer it onto the tip (uniqueness + lineage kept).
|
||||
if self._is_compression_ancestor(conn, ancestor_id=conflict_id, descendant_id=session_id):
|
||||
conn.execute("UPDATE sessions SET title = NULL WHERE id = ?", (conflict_id,))
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Title '{title}' is already in use by session {conflict_id}"
|
||||
)
|
||||
# CAS on the values just read (``IS`` is NULL-safe): a concurrent
|
||||
# write between the SELECT and here loses instead of being overwritten.
|
||||
raise ValueError(f"Title '{title}' is already in use by session {conflict_id}")
|
||||
# CAS on the values just read (``IS`` is NULL-safe): a concurrent write
|
||||
# between the SELECT and here loses instead of being overwritten.
|
||||
cursor = conn.execute(
|
||||
"UPDATE sessions SET title = ?, title_source = ? "
|
||||
"WHERE id = ? AND title IS ? AND title_source IS ?",
|
||||
(
|
||||
title,
|
||||
source if title else None,
|
||||
session_id,
|
||||
current["title"],
|
||||
current["title_source"],
|
||||
),
|
||||
(title, source if title else None, session_id, current["title"], current["title_source"]),
|
||||
)
|
||||
return cursor.rowcount
|
||||
|
||||
rowcount = self._execute_write(_do)
|
||||
return rowcount > 0
|
||||
return self._execute_write(_do) > 0
|
||||
|
||||
def set_session_title(self, session_id: str, title: str) -> bool:
|
||||
"""Set a title on the user's behalf (``user`` provenance; auto-titling
|
||||
never replaces it). Empty clears the title. Raises ValueError on a
|
||||
title conflict or validation failure. Automatic callers must use
|
||||
:meth:`set_auto_title`."""
|
||||
return self._set_session_title(
|
||||
session_id, title, source=self.TITLE_SOURCE_USER
|
||||
)
|
||||
"""Set a title on the user's behalf (``user`` provenance). Empty clears it.
|
||||
Raises ValueError on conflict or validation failure."""
|
||||
return self._set_session_title(session_id, title, source=self.TITLE_SOURCE_USER)
|
||||
|
||||
def set_auto_title(self, session_id: str, title: str, *, source: str) -> bool:
|
||||
"""Set an automatic title; False (untouched) when a higher-authority
|
||||
title already holds the row."""
|
||||
"""Set an automatic title; False (untouched) when a higher-authority title
|
||||
already holds the row."""
|
||||
if source not in (self.TITLE_SOURCE_DERIVED, self.TITLE_SOURCE_LLM):
|
||||
raise ValueError(f"invalid automatic title source: {source!r}")
|
||||
return self._set_session_title(session_id, title, source=source)
|
||||
|
||||
def set_auto_title_if_empty(self, session_id: str, title: str) -> bool:
|
||||
"""Back-compat shim (third-party plugins reference it by name); new
|
||||
code calls :meth:`set_auto_title` with an explicit source."""
|
||||
return self.set_auto_title(
|
||||
session_id, title, source=self.TITLE_SOURCE_LLM
|
||||
)
|
||||
"""Back-compat shim (third-party plugins reference it by name)."""
|
||||
return self.set_auto_title(session_id, title, source=self.TITLE_SOURCE_LLM)
|
||||
|
||||
def get_session_title(self, session_id: str) -> Optional[str]:
|
||||
"""Get the title for a session, or None."""
|
||||
with self._read_ctx() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT title FROM sessions WHERE id = ?", (session_id,)
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
row = self._read_one("SELECT title FROM sessions WHERE id = ?", (session_id,))
|
||||
return row["title"] if row else None
|
||||
|
||||
def get_session_title_source(self, session_id: str) -> Optional[str]:
|
||||
"""Get the provenance of a session's title, or None when untitled."""
|
||||
with self._read_ctx() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT title, title_source FROM sessions WHERE id = ?",
|
||||
(session_id,),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
row = self._read_one("SELECT title, title_source FROM sessions WHERE id = ?", (session_id,))
|
||||
if not row or row["title"] is None:
|
||||
return None
|
||||
return row["title_source"]
|
||||
|
||||
def set_session_title_source(self, session_id: str, source: str) -> bool:
|
||||
"""Overwrite a title's provenance without touching the text: a title
|
||||
copied across a compression rotation keeps the original's authority."""
|
||||
"""Overwrite a title's provenance without touching the text (a title copied
|
||||
across a compression rotation keeps the original's authority)."""
|
||||
if source not in self._TITLE_SOURCE_RANK:
|
||||
raise ValueError(f"invalid title source: {source!r}")
|
||||
|
||||
return self._write_rowcount(
|
||||
"UPDATE sessions SET title_source = ? "
|
||||
"WHERE id = ? AND title IS NOT NULL",
|
||||
@@ -243,63 +173,44 @@ class SessionTitlesMixin:
|
||||
|
||||
def get_session_by_title(self, title: str) -> Optional[Dict[str, Any]]:
|
||||
"""Look up a session by exact title. Returns session dict or None."""
|
||||
with self._read_ctx() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT s.*, "
|
||||
"COALESCE(sp.prompt, s.system_prompt) AS _system_prompt_resolved "
|
||||
"FROM sessions s "
|
||||
"LEFT JOIN system_prompts sp ON sp.hash = s.system_prompt_hash "
|
||||
"WHERE s.title = ?",
|
||||
(title,),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
row = self._read_one(
|
||||
"SELECT s.*, "
|
||||
"COALESCE(sp.prompt, s.system_prompt) AS _system_prompt_resolved "
|
||||
"FROM sessions s "
|
||||
"LEFT JOIN system_prompts sp ON sp.hash = s.system_prompt_hash "
|
||||
"WHERE s.title = ?",
|
||||
(title,),
|
||||
)
|
||||
return self._session_row_dict(row) if row else None
|
||||
|
||||
def resolve_session_by_title(self, title: str) -> Optional[str]:
|
||||
"""Resolve a title to a session ID, preferring the latest "title #N"
|
||||
continuation over the exact match."""
|
||||
exact = self.get_session_by_title(title)
|
||||
|
||||
# Escape LIKE wildcards so "%"/"_" in titles cannot false-match.
|
||||
escaped = _escape_like(title)
|
||||
with self._read_ctx() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT id, title, started_at FROM sessions "
|
||||
"WHERE title LIKE ? ESCAPE '\\' ORDER BY started_at DESC",
|
||||
(f"{escaped} #%",),
|
||||
)
|
||||
numbered = cursor.fetchall()
|
||||
|
||||
numbered = self._read_all(
|
||||
"SELECT id, title, started_at FROM sessions "
|
||||
"WHERE title LIKE ? ESCAPE '\\' ORDER BY started_at DESC",
|
||||
(f"{_escape_like(title)} #%",),
|
||||
)
|
||||
if numbered:
|
||||
return numbered[0]["id"]
|
||||
elif exact:
|
||||
return exact["id"]
|
||||
return None
|
||||
return exact["id"] if exact else None
|
||||
|
||||
def get_next_title_in_lineage(self, base_title: str) -> str:
|
||||
"""Next title in a lineage ("my session" → "my session #2"): strip any
|
||||
" #N" suffix, then increment the highest existing number."""
|
||||
match = re.match(r'^(.*?) #(\d+)$', base_title)
|
||||
if match:
|
||||
base = match.group(1)
|
||||
else:
|
||||
base = base_title
|
||||
|
||||
escaped = _escape_like(base)
|
||||
with self._read_ctx() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT title FROM sessions WHERE title = ? OR title LIKE ? ESCAPE '\\'",
|
||||
(base, f"{escaped} #%"),
|
||||
)
|
||||
existing = [row["title"] for row in cursor.fetchall()]
|
||||
|
||||
if not existing:
|
||||
"""Next title in a lineage ("my session" -> "my session #2"): strip any " #N"
|
||||
suffix, then increment the highest existing number."""
|
||||
match = _NUMBERED_TITLE_RE.match(base_title)
|
||||
base = match.group(1) if match else base_title
|
||||
rows = self._read_all(
|
||||
"SELECT title FROM sessions WHERE title = ? OR title LIKE ? ESCAPE '\\'",
|
||||
(base, f"{_escape_like(base)} #%"),
|
||||
)
|
||||
if not rows:
|
||||
return base
|
||||
|
||||
max_num = 1 # The unnumbered original counts as #1
|
||||
for t in existing:
|
||||
m = re.match(r'^.* #(\d+)$', t)
|
||||
max_num = 1 # the unnumbered original counts as #1
|
||||
for row in rows:
|
||||
m = re.match(r'^.* #(\d+)$', row["title"])
|
||||
if m:
|
||||
max_num = max(max_num, int(m.group(1)))
|
||||
|
||||
return f"{base} #{max_num + 1}"
|
||||
|
||||
+185
-348
@@ -14,24 +14,79 @@ from typing import Any, Dict, List, Optional, Tuple
|
||||
# caplog tests pin the "hermes_state" logger name.
|
||||
logger = logging.getLogger("hermes_state")
|
||||
|
||||
_TOKEN_UPDATE_ABSOLUTE_SQL = """UPDATE sessions SET
|
||||
input_tokens = ?,
|
||||
output_tokens = ?,
|
||||
cache_read_tokens = ?,
|
||||
cache_write_tokens = ?,
|
||||
reasoning_tokens = ?,
|
||||
estimated_cost_usd = COALESCE(?, 0),
|
||||
actual_cost_usd = CASE
|
||||
WHEN ? IS NULL THEN actual_cost_usd
|
||||
ELSE ?
|
||||
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 = ?
|
||||
WHERE id = ?"""
|
||||
|
||||
_TOKEN_UPDATE_DELTA_SQL = """UPDATE sessions SET
|
||||
input_tokens = input_tokens + ?,
|
||||
output_tokens = output_tokens + ?,
|
||||
cache_read_tokens = cache_read_tokens + ?,
|
||||
cache_write_tokens = cache_write_tokens + ?,
|
||||
reasoning_tokens = reasoning_tokens + ?,
|
||||
estimated_cost_usd = COALESCE(estimated_cost_usd, 0) + COALESCE(?, 0),
|
||||
actual_cost_usd = CASE
|
||||
WHEN ? IS NULL THEN actual_cost_usd
|
||||
ELSE COALESCE(actual_cost_usd, 0) + ?
|
||||
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 = COALESCE(api_call_count, 0) + ?
|
||||
WHERE id = ?"""
|
||||
|
||||
_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"""
|
||||
|
||||
|
||||
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,
|
||||
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 (stale ``Model:`` /
|
||||
``Provider:`` header) is rebuilt, like ``update_session_model``.
|
||||
"""
|
||||
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()
|
||||
|
||||
@@ -50,29 +105,20 @@ class SessionUsageMixin:
|
||||
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 and semantics as :meth:`update_token_counts`, applied
|
||||
asynchronously; cheap enough for the turn thread. After close() has
|
||||
stopped the writer, falls back to the synchronous path and may raise.
|
||||
"""
|
||||
"""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."""
|
||||
with self._token_queue_cond:
|
||||
thread = self._token_writer_thread
|
||||
writer_stopped = self._token_writer_stop and (
|
||||
thread is None or not thread.is_alive()
|
||||
)
|
||||
writer_stopped = self._token_writer_stop and (thread is None or not thread.is_alive())
|
||||
if not writer_stopped:
|
||||
self._token_queue.append((session_id, kwargs))
|
||||
if thread is None or not thread.is_alive():
|
||||
# Daemon so exit never hangs on accounting; the atexit hook
|
||||
# (registered once per instance) drains leftovers. Checking
|
||||
# ``not is_alive()`` rather than ``is None`` respawns a writer
|
||||
# that died from an unexpected escape, otherwise deltas
|
||||
# would pile up until a reader's flush drained them.
|
||||
# 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.
|
||||
thread = threading.Thread(
|
||||
target=self._token_writer_loop,
|
||||
name="session-db-token-writer",
|
||||
daemon=True,
|
||||
target=self._token_writer_loop, name="session-db-token-writer", daemon=True,
|
||||
)
|
||||
self._token_writer_thread = thread
|
||||
thread.start()
|
||||
@@ -88,18 +134,21 @@ class SessionUsageMixin:
|
||||
atexit.register(_drain_at_exit)
|
||||
self._token_queue_cond.notify_all()
|
||||
if writer_stopped:
|
||||
# close() ran (a stop-flagged but live writer still accepts; its
|
||||
# loop drains before exiting). Enqueueing now would drop the delta
|
||||
# silently — no writer, atexit hook gone — so apply inline and let a
|
||||
# closed-connection failure raise at the call site.
|
||||
# close() ran: enqueueing would drop the delta silently, so apply inline.
|
||||
self.update_token_counts(session_id, **kwargs)
|
||||
|
||||
def flush_token_counts(self, timeout: float = 5.0) -> bool:
|
||||
"""Block until every queued token delta has been applied.
|
||||
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()
|
||||
|
||||
False on timeout (callers then read totals stale by the queued deltas).
|
||||
Never raises: apply failures are logged by the writer.
|
||||
"""
|
||||
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
|
||||
@@ -107,20 +156,13 @@ 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 newer deltas committing
|
||||
# before older ones breaks last-non-None-wins / first-accounted-
|
||||
# route / COALESCE-backfill fields. Only a dead writer lets the
|
||||
# caller take leftovers; re-checked each wakeup because the writer
|
||||
# can exit mid-wait with deltas enqueued after its final check.
|
||||
# busy is claimed while draining so a concurrent flush cannot
|
||||
# report drained or pop a newer delta while this batch is
|
||||
# unapplied: 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", never "drain alongside".
|
||||
thread = self._token_writer_thread
|
||||
if (
|
||||
(thread is None or not thread.is_alive())
|
||||
and not self._token_writer_busy
|
||||
):
|
||||
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()
|
||||
@@ -130,12 +172,7 @@ class SessionUsageMixin:
|
||||
return False
|
||||
self._token_queue_cond.wait(remaining)
|
||||
if batch:
|
||||
try:
|
||||
self._apply_token_batch(batch)
|
||||
finally:
|
||||
with self._token_queue_cond:
|
||||
self._token_writer_busy = False
|
||||
self._token_queue_cond.notify_all()
|
||||
self._apply_claimed_batch(batch)
|
||||
return True
|
||||
|
||||
def _token_writer_loop(self) -> None:
|
||||
@@ -145,62 +182,44 @@ 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 same lock queue_token_counts() uses to decide
|
||||
# to spawn, 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
|
||||
# reads queue-then-busy and 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()
|
||||
try:
|
||||
self._apply_token_batch(batch)
|
||||
finally:
|
||||
with self._token_queue_cond:
|
||||
self._token_writer_busy = False
|
||||
self._token_queue_cond.notify_all()
|
||||
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 (callers cannot observe a
|
||||
# dead one); the merge is only an optimization.
|
||||
logger.warning(
|
||||
"async token accounting: coalesce failed, applying raw "
|
||||
"batch: %s", 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,
|
||||
)
|
||||
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 deltas never merge."""
|
||||
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,) + tuple(
|
||||
kwargs.get(f) for f in self._TOKEN_DELTA_ROUTE_FIELDS
|
||||
)
|
||||
key = (session_id,) + tuple(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:
|
||||
@@ -223,18 +242,16 @@ class SessionUsageMixin:
|
||||
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 and misorder/double-count.
|
||||
# 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.
|
||||
# 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:
|
||||
@@ -247,19 +264,13 @@ class SessionUsageMixin:
|
||||
)
|
||||
return
|
||||
self._token_queue_cond.wait(remaining)
|
||||
# busy BEFORE clearing the queue (same ordering as the writer loop),
|
||||
# or flush's lock-free fast path could see "empty and idle".
|
||||
# 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:
|
||||
try:
|
||||
self._apply_token_batch(batch)
|
||||
finally:
|
||||
with self._token_queue_cond:
|
||||
self._token_writer_busy = False
|
||||
self._token_queue_cond.notify_all()
|
||||
self._apply_claimed_batch(batch)
|
||||
|
||||
def _drain_token_queue_at_exit(self) -> None:
|
||||
try:
|
||||
@@ -268,126 +279,50 @@ class SessionUsageMixin:
|
||||
pass # never fatal at interpreter shutdown
|
||||
|
||||
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,
|
||||
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,
|
||||
) -> 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).
|
||||
"""
|
||||
# Ensure the row exists: under concurrent load the initial
|
||||
# create_session() may have failed on SQLite locking, and the UPDATE
|
||||
# would silently affect 0 rows.
|
||||
"""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)."""
|
||||
# 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)
|
||||
if absolute:
|
||||
sql = """UPDATE sessions SET
|
||||
input_tokens = ?,
|
||||
output_tokens = ?,
|
||||
cache_read_tokens = ?,
|
||||
cache_write_tokens = ?,
|
||||
reasoning_tokens = ?,
|
||||
estimated_cost_usd = COALESCE(?, 0),
|
||||
actual_cost_usd = CASE
|
||||
WHEN ? IS NULL THEN actual_cost_usd
|
||||
ELSE ?
|
||||
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 = ?
|
||||
WHERE id = ?"""
|
||||
else:
|
||||
sql = """UPDATE sessions SET
|
||||
input_tokens = input_tokens + ?,
|
||||
output_tokens = output_tokens + ?,
|
||||
cache_read_tokens = cache_read_tokens + ?,
|
||||
cache_write_tokens = cache_write_tokens + ?,
|
||||
reasoning_tokens = reasoning_tokens + ?,
|
||||
estimated_cost_usd = COALESCE(estimated_cost_usd, 0) + COALESCE(?, 0),
|
||||
actual_cost_usd = CASE
|
||||
WHEN ? IS NULL THEN actual_cost_usd
|
||||
ELSE COALESCE(actual_cost_usd, 0) + ?
|
||||
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 = COALESCE(api_call_count, 0) + ?
|
||||
WHERE id = ?"""
|
||||
has_accounted_usage = bool(
|
||||
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 or actual_cost_usd
|
||||
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,
|
||||
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,
|
||||
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. Each delta carries the route active at call time
|
||||
# and is recorded into session_model_usage keyed by it. Only the
|
||||
# incremental path records here: absolute cumulative updates cannot be
|
||||
# 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.
|
||||
record_model_usage = (not absolute) and (
|
||||
input_tokens or output_tokens or cache_read_tokens
|
||||
or cache_write_tokens or reasoning_tokens or api_call_count
|
||||
or estimated_cost_usd
|
||||
)
|
||||
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,),
|
||||
"SELECT model, billing_provider, api_call_count FROM sessions WHERE id = ?", (session_id,),
|
||||
).fetchone()
|
||||
existing_model = row["model"] if row is not None else None
|
||||
existing_provider = row["billing_provider"] if row is not None else None
|
||||
existing_api_calls = int((row["api_call_count"] if row is not None else 0) or 0)
|
||||
|
||||
# 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-provider 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 = (
|
||||
existing_api_calls == 0
|
||||
and has_accounted_usage
|
||||
@@ -406,146 +341,63 @@ class SessionUsageMixin:
|
||||
conn.execute(sql, params)
|
||||
if record_model_usage:
|
||||
self._record_model_usage(
|
||||
conn,
|
||||
session_id,
|
||||
model=model,
|
||||
billing_provider=billing_provider,
|
||||
billing_base_url=billing_base_url,
|
||||
billing_mode=billing_mode,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cache_read_tokens=cache_read_tokens,
|
||||
cache_write_tokens=cache_write_tokens,
|
||||
reasoning_tokens=reasoning_tokens,
|
||||
estimated_cost_usd=estimated_cost_usd,
|
||||
actual_cost_usd=actual_cost_usd,
|
||||
cost_status=cost_status,
|
||||
cost_source=cost_source,
|
||||
conn, session_id, model=model, billing_provider=billing_provider,
|
||||
billing_base_url=billing_base_url, billing_mode=billing_mode,
|
||||
input_tokens=input_tokens, output_tokens=output_tokens,
|
||||
cache_read_tokens=cache_read_tokens, cache_write_tokens=cache_write_tokens,
|
||||
reasoning_tokens=reasoning_tokens, estimated_cost_usd=estimated_cost_usd,
|
||||
actual_cost_usd=actual_cost_usd, cost_status=cost_status, cost_source=cost_source,
|
||||
api_call_count=api_call_count,
|
||||
)
|
||||
self._execute_write(_do)
|
||||
|
||||
def _record_model_usage(
|
||||
self,
|
||||
conn,
|
||||
session_id: str,
|
||||
*,
|
||||
model: Optional[str],
|
||||
billing_provider: Optional[str],
|
||||
billing_base_url: Optional[str],
|
||||
billing_mode: Optional[str],
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
cache_read_tokens: int,
|
||||
cache_write_tokens: int,
|
||||
reasoning_tokens: int,
|
||||
estimated_cost_usd: Optional[float],
|
||||
actual_cost_usd: Optional[float],
|
||||
cost_status: Optional[str],
|
||||
cost_source: Optional[str],
|
||||
api_call_count: int,
|
||||
task: str = "",
|
||||
self, conn, session_id: str, *, model: Optional[str], billing_provider: Optional[str],
|
||||
billing_base_url: Optional[str], billing_mode: Optional[str], input_tokens: int,
|
||||
output_tokens: int, cache_read_tokens: int, cache_write_tokens: int, reasoning_tokens: int,
|
||||
estimated_cost_usd: Optional[float], actual_cost_usd: Optional[float],
|
||||
cost_status: Optional[str], cost_source: Optional[str], api_call_count: int, task: str = "",
|
||||
) -> None:
|
||||
"""Accumulate a per-API-call usage delta into session_model_usage.
|
||||
|
||||
Runs inside the caller's write transaction, after the ``sessions``
|
||||
UPDATE, so per-model rows stay consistent with the summary row. A
|
||||
missing model/provider falls back to the session row (same COALESCE
|
||||
behaviour as the summary update). ``task`` is ``''`` for the main loop;
|
||||
auxiliary calls record their task name via :meth:`record_auxiliary_usage`.
|
||||
"""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.
|
||||
"""
|
||||
row = conn.execute(
|
||||
"SELECT model, billing_provider, billing_base_url, billing_mode "
|
||||
"FROM sessions WHERE id = ?",
|
||||
(session_id,),
|
||||
).fetchone()
|
||||
sess_model = row["model"] if row is not None else None
|
||||
sess_provider = row["billing_provider"] if row is not None else None
|
||||
sess_base_url = row["billing_base_url"] if row is not None else None
|
||||
sess_billing_mode = row["billing_mode"] if row is not None else None
|
||||
|
||||
# Aux rows must NOT inherit the main-loop route (vision on gemini while
|
||||
# the main loop runs anthropic); missing info stays 'unknown'/empty.
|
||||
if task:
|
||||
eff_model = model or "unknown"
|
||||
eff_provider = billing_provider or ""
|
||||
eff_base_url = billing_base_url or ""
|
||||
eff_billing_mode = billing_mode or ""
|
||||
else:
|
||||
eff_model = model or sess_model or "unknown"
|
||||
eff_provider = billing_provider or sess_provider or ""
|
||||
eff_base_url = billing_base_url or sess_base_url or ""
|
||||
eff_billing_mode = billing_mode or sess_billing_mode or ""
|
||||
sess = dict(row) if (row is not None and not task) else {}
|
||||
eff_model = model or sess.get("model") or "unknown"
|
||||
eff_provider = billing_provider or sess.get("billing_provider") or ""
|
||||
eff_base_url = billing_base_url or sess.get("billing_base_url") or ""
|
||||
eff_billing_mode = billing_mode or sess.get("billing_mode") or ""
|
||||
counts = [v or 0 for v in (input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, reasoning_tokens)]
|
||||
now = time.time()
|
||||
conn.execute(
|
||||
"""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""",
|
||||
_MODEL_USAGE_UPSERT_SQL,
|
||||
(
|
||||
session_id,
|
||||
eff_model,
|
||||
eff_provider,
|
||||
eff_base_url,
|
||||
eff_billing_mode,
|
||||
task or "",
|
||||
api_call_count or 0,
|
||||
input_tokens or 0,
|
||||
output_tokens or 0,
|
||||
cache_read_tokens or 0,
|
||||
cache_write_tokens or 0,
|
||||
reasoning_tokens or 0,
|
||||
float(estimated_cost_usd or 0.0),
|
||||
float(actual_cost_usd or 0.0),
|
||||
cost_status,
|
||||
cost_source,
|
||||
now,
|
||||
now,
|
||||
session_id, eff_model, eff_provider, eff_base_url, eff_billing_mode, 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,
|
||||
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, ...) against *session_id*.
|
||||
|
||||
Writes a per-(model, provider, task) delta into ``session_model_usage``
|
||||
WITHOUT touching the ``sessions`` summary row: the gateway overwrites
|
||||
session counters with absolute main-loop totals, so aux tokens there
|
||||
would be clobbered or double-counted. Insights read the union.
|
||||
``api_call_count`` may aggregate N calls (background-review forks).
|
||||
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:
|
||||
callers must never fail an aux call over accounting."""
|
||||
if not session_id or not task:
|
||||
return
|
||||
# FK to sessions.id: same INSERT OR IGNORE guard as update_token_counts.
|
||||
@@ -553,37 +405,23 @@ class SessionUsageMixin:
|
||||
|
||||
def _do(conn):
|
||||
self._record_model_usage(
|
||||
conn,
|
||||
session_id,
|
||||
model=model,
|
||||
billing_provider=billing_provider,
|
||||
billing_base_url=billing_base_url,
|
||||
billing_mode=None,
|
||||
input_tokens=input_tokens or 0,
|
||||
output_tokens=output_tokens or 0,
|
||||
cache_read_tokens=cache_read_tokens or 0,
|
||||
cache_write_tokens=cache_write_tokens or 0,
|
||||
reasoning_tokens=reasoning_tokens or 0,
|
||||
estimated_cost_usd=estimated_cost_usd,
|
||||
actual_cost_usd=None,
|
||||
cost_status=None,
|
||||
cost_source=None,
|
||||
api_call_count=(
|
||||
1 if api_call_count is None else int(api_call_count)
|
||||
),
|
||||
task=task,
|
||||
conn, session_id, model=model, billing_provider=billing_provider,
|
||||
billing_base_url=billing_base_url, billing_mode=None,
|
||||
input_tokens=input_tokens or 0, output_tokens=output_tokens or 0,
|
||||
cache_read_tokens=cache_read_tokens or 0, cache_write_tokens=cache_write_tokens or 0,
|
||||
reasoning_tokens=reasoning_tokens or 0, estimated_cost_usd=estimated_cost_usd,
|
||||
actual_cost_usd=None, cost_status=None, cost_source=None,
|
||||
api_call_count=1 if api_call_count is None else int(api_call_count), task=task,
|
||||
)
|
||||
self._execute_write(_do)
|
||||
|
||||
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, the same precedence a single row renders."""
|
||||
"""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),
|
||||
@@ -593,5 +431,4 @@ class SessionUsageMixin:
|
||||
""",
|
||||
params,
|
||||
)
|
||||
|
||||
return {"tokens": int(row[0] or 0), "cost_usd": float(row[1] or 0.0)}
|
||||
|
||||
Reference in New Issue
Block a user