zettcode-engine 0.3.4__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- zettcode/_compat.py +35 -0
- zettcode/budget.py +24 -0
- zettcode/config.py +1088 -0
- zettcode/config_edit.py +457 -0
- zettcode/engine/__init__.py +39 -0
- zettcode/engine/agent/__init__.py +81 -0
- zettcode/engine/agent/base.py +60 -0
- zettcode/engine/agent/context.py +166 -0
- zettcode/engine/agent/exports.py +73 -0
- zettcode/engine/agent/interaction.py +54 -0
- zettcode/engine/agent/sessions.py +218 -0
- zettcode/engine/agent/settings.py +269 -0
- zettcode/engine/agent/turns.py +308 -0
- zettcode/engine/core.py +819 -0
- zettcode/engine/dialogue/__init__.py +6 -0
- zettcode/engine/dialogue/ask.py +113 -0
- zettcode/engine/dialogue/side.py +197 -0
- zettcode/engine/dialogue/title.py +80 -0
- zettcode/engine/export/__init__.py +421 -0
- zettcode/engine/export/trace/script.js +18 -0
- zettcode/engine/export/trace/style.css +136 -0
- zettcode/engine/export/trace/template.html +31 -0
- zettcode/engine/extensions/__init__.py +8 -0
- zettcode/engine/extensions/capabilities.py +68 -0
- zettcode/engine/extensions/compaction.py +259 -0
- zettcode/engine/extensions/context.py +521 -0
- zettcode/engine/extensions/files.py +101 -0
- zettcode/engine/extensions/guidance.py +17 -0
- zettcode/engine/extensions/harness.py +208 -0
- zettcode/engine/extensions/harness_switch.py +18 -0
- zettcode/engine/extensions/mcp_notice.py +49 -0
- zettcode/engine/extensions/relay.py +93 -0
- zettcode/engine/extensions/subagents.py +391 -0
- zettcode/engine/extensions/tool_output.py +153 -0
- zettcode/engine/extensions/truncation.py +156 -0
- zettcode/engine/extensions/usage.py +304 -0
- zettcode/engine/harness_paths.py +74 -0
- zettcode/engine/permission/__init__.py +197 -0
- zettcode/engine/permission/approval.py +66 -0
- zettcode/engine/permission/readonly.py +435 -0
- zettcode/engine/permission/sandbox.py +140 -0
- zettcode/engine/pipeline.py +191 -0
- zettcode/engine/plugins/__init__.py +72 -0
- zettcode/engine/plugins/commands.py +152 -0
- zettcode/engine/plugins/context.py +251 -0
- zettcode/engine/plugins/extension.py +183 -0
- zettcode/engine/plugins/groups.py +74 -0
- zettcode/engine/plugins/loader.py +367 -0
- zettcode/engine/plugins/mixins.py +156 -0
- zettcode/engine/plugins/plugin.py +190 -0
- zettcode/engine/plugins/services.py +68 -0
- zettcode/engine/plugins/settings.py +117 -0
- zettcode/engine/plugins/state.py +64 -0
- zettcode/engine/plugins/switches.py +553 -0
- zettcode/engine/py.typed +0 -0
- zettcode/engine/server_tools.py +179 -0
- zettcode/engine/storage/__init__.py +51 -0
- zettcode/engine/storage/layout.py +161 -0
- zettcode/engine/storage/ledger.py +659 -0
- zettcode/engine/storage/locks.py +182 -0
- zettcode/engine/storage/metadata.py +278 -0
- zettcode/engine/storage/records.py +373 -0
- zettcode/engine/storage/store.py +1056 -0
- zettcode/engine/tools/__init__.py +7 -0
- zettcode/engine/tools/boundary.py +56 -0
- zettcode/engine/tools/files.py +262 -0
- zettcode/engine/tools/images.py +32 -0
- zettcode/engine/tools/output.py +81 -0
- zettcode/engine/tools/shell.py +181 -0
- zettcode/engine/tools/workspace.py +271 -0
- zettcode/error_text.py +217 -0
- zettcode/filelock.py +125 -0
- zettcode/logging_setup.py +196 -0
- zettcode/mention.py +124 -0
- zettcode/modes.py +18 -0
- zettcode/paths.py +113 -0
- zettcode/plugin_tui.py +133 -0
- zettcode/processes.py +126 -0
- zettcode/prompt.py +94 -0
- zettcode/providers.py +215 -0
- zettcode/py.typed +0 -0
- zettcode/registry.py +63 -0
- zettcode/relay.py +267 -0
- zettcode/statusline_parts.py +84 -0
- zettcode/version.py +11 -0
- zettcode/wire.py +30 -0
- zettcode/wording.py +11 -0
- zettcode_engine-0.3.4.dist-info/METADATA +23 -0
- zettcode_engine-0.3.4.dist-info/RECORD +90 -0
- zettcode_engine-0.3.4.dist-info/WHEEL +4 -0
|
@@ -0,0 +1,659 @@
|
|
|
1
|
+
"""What a workspace's sessions have spent: tokens by model, by session, and by day.
|
|
2
|
+
|
|
3
|
+
Every model reply is stored with the usage its provider reported, so the totals ``/usage``
|
|
4
|
+
shows are read back from the conversation files rather than kept in a second place that could
|
|
5
|
+
disagree with them. A reader never builds a :class:`~zettcode.engine.storage.records.Session`
|
|
6
|
+
for this: it looks at one line at a time, parses only the lines that carry usage, and takes from
|
|
7
|
+
them the counters and the model's name.
|
|
8
|
+
|
|
9
|
+
**What is counted.** The sessions the reader works in. A subagent's session names an agent
|
|
10
|
+
(``agent_name`` in its header) and is left out: its tokens are real, but they are not part of
|
|
11
|
+
any conversation the reader had.
|
|
12
|
+
|
|
13
|
+
**A fork is not counted twice.** ``/fork`` copies the conversation into a new session together
|
|
14
|
+
with the usage of each reply, so adding every session up would charge the same answer to both.
|
|
15
|
+
A reply is identified by the request it answered, when it was written, the model that wrote it
|
|
16
|
+
and what it reported; the first session to hold it, in the order the sessions began, keeps it,
|
|
17
|
+
and a copy of it elsewhere is not counted. A fork therefore shows only what it spent after the
|
|
18
|
+
fork.
|
|
19
|
+
|
|
20
|
+
**Days** are the reader's own: a reply written at 23:30 UTC belongs to the day it was in the
|
|
21
|
+
reader's time zone, which is the day a calendar on their screen should put it on. The index
|
|
22
|
+
below keeps the moment, not the day, so changing time zone changes nothing it stored.
|
|
23
|
+
|
|
24
|
+
**The index.** Reading is a pass over every byte of every conversation, which is a fraction of a
|
|
25
|
+
second for tens of megabytes and many seconds for gigabytes, so what a pass finds is kept in
|
|
26
|
+
``usage.json`` beside the session folders. A conversation file is append-only, so the index
|
|
27
|
+
records, for each one, how far it was read and a fingerprint of the bytes at its start and just
|
|
28
|
+
before that point; the next read checks the fingerprints and reads only what came after. A file
|
|
29
|
+
that no longer matches (replaced, or edited in the part already read) is read again from the
|
|
30
|
+
beginning, and a damaged, foreign, or older index is ignored whole. The index is never the
|
|
31
|
+
truth, only a way to avoid re-reading it: deleting it costs one slow read, and it is written
|
|
32
|
+
best-effort, so a read-only store or two processes writing at once cost nothing but a re-read.
|
|
33
|
+
A rewrite that keeps both fingerprinted ends is not noticed; nothing in this program makes one.
|
|
34
|
+
|
|
35
|
+
**Entries.** The same pass counts a session's message records (the system prompt, each user message,
|
|
36
|
+
each reply and each tool result is one), which is what the sessions panel shows as its size. A line is
|
|
37
|
+
recognised by its ``"kind": "message"`` field without being parsed, so counting costs a substring search.
|
|
38
|
+
|
|
39
|
+
**The summary.** ``usage.json`` keeps every reply of every session, so it grows with the work. The sessions
|
|
40
|
+
panel needs a handful of totals for each of twenty sessions, so :meth:`LedgerReader.stats` keeps those in a
|
|
41
|
+
second, small file, ``sessions.json``, next to it: for each session the size and modification time of its
|
|
42
|
+
conversation file, and the totals found at that size. A session whose file still has that size and time is
|
|
43
|
+
answered from the summary alone, without opening ``usage.json`` or a conversation; any other is answered by the
|
|
44
|
+
ordinary incremental read, which then rewrites the summary. It is a cache like the index and is treated the
|
|
45
|
+
same way: damaged, foreign or older files are ignored whole, and writing it is best-effort.
|
|
46
|
+
"""
|
|
47
|
+
|
|
48
|
+
from __future__ import annotations
|
|
49
|
+
|
|
50
|
+
import json
|
|
51
|
+
import os
|
|
52
|
+
import zlib
|
|
53
|
+
from collections.abc import Iterable
|
|
54
|
+
from dataclasses import dataclass, field, replace
|
|
55
|
+
from datetime import date, datetime, timezone
|
|
56
|
+
from pathlib import Path
|
|
57
|
+
from typing import BinaryIO
|
|
58
|
+
|
|
59
|
+
from zett_agent.model import FIRST_CHUNK_LIMIT
|
|
60
|
+
|
|
61
|
+
from .layout import find_workspace
|
|
62
|
+
from .metadata import SESSION_METADATA_FILE
|
|
63
|
+
from .records import SAFE_SESSION_ID, STORE_FILE
|
|
64
|
+
|
|
65
|
+
#: The model a reply is filed under when its message does not say.
|
|
66
|
+
UNKNOWN_MODEL = "unknown"
|
|
67
|
+
|
|
68
|
+
#: The index, in the workspace's folder beside the session folders (which is where
|
|
69
|
+
#: ``metadata.jsonl`` already sits, so a session listing never mistakes it for a session).
|
|
70
|
+
INDEX_FILE = "usage.json"
|
|
71
|
+
|
|
72
|
+
#: Bumped when the index's shape changes; an index of another version is ignored. The fingerprint's width is part
|
|
73
|
+
#: of the shape: a stored tail that was taken over other bytes never matches, and re-reading the file is the whole cost.
|
|
74
|
+
INDEX_VERSION = 2
|
|
75
|
+
|
|
76
|
+
#: The per-session totals the sessions panel reads, beside the index, and the version of its shape.
|
|
77
|
+
SUMMARY_FILE = "sessions.json"
|
|
78
|
+
SUMMARY_VERSION = 1
|
|
79
|
+
|
|
80
|
+
#: What marks a stored line as a message record, with and without the space after the colon. Keys are written
|
|
81
|
+
#: in sorted order and a message's own text is escaped, so the pair never appears inside a message.
|
|
82
|
+
MESSAGE_MARKS = (b'"kind": "message"', b'"kind":"message"')
|
|
83
|
+
|
|
84
|
+
#: Bytes at the start of a file, and bytes before the read point, that fingerprint a file. The tail has to reach back
|
|
85
|
+
#: over one whole record, because an edit to a count inside the newest one has to be noticed: a store writes each line
|
|
86
|
+
#: with sorted keys, so the usage block is the last thing on it, and it carries two excerpts of the reply clipped at
|
|
87
|
+
#: ``FIRST_CHUNK_LIMIT`` characters each, a character taking up to four bytes. Raised from 64 when the runtime began
|
|
88
|
+
#: recording the first token's arrival and the block doubled; written as arithmetic, with room to spare, so that raising
|
|
89
|
+
#: the limit cannot silently outgrow it.
|
|
90
|
+
HEAD_BYTES = 256
|
|
91
|
+
TAIL_BYTES = 2 * FIRST_CHUNK_LIMIT * 4 + 1024
|
|
92
|
+
|
|
93
|
+
#: A quarter of an hour: every time zone's midnight falls on one, so a day is worked out once
|
|
94
|
+
#: for each and not once for each reply.
|
|
95
|
+
BUCKET_SECONDS = 900
|
|
96
|
+
|
|
97
|
+
#: One reply as the index keeps it: request id, stored timestamp, UTC seconds, model, then the
|
|
98
|
+
#: input, output, cache-read and reasoning counters.
|
|
99
|
+
Row = tuple[str, str, int, str, int, int, int, int]
|
|
100
|
+
|
|
101
|
+
|
|
102
|
+
@dataclass(slots=True)
|
|
103
|
+
class Tally:
|
|
104
|
+
"""Token counters added up over some replies.
|
|
105
|
+
|
|
106
|
+
``input_tokens`` already contains what came from the provider's cache and
|
|
107
|
+
``output_tokens`` already contains reasoning, as the provider reports them, so
|
|
108
|
+
:attr:`total` is their sum and nothing is counted twice.
|
|
109
|
+
"""
|
|
110
|
+
|
|
111
|
+
requests: int = 0
|
|
112
|
+
input_tokens: int = 0
|
|
113
|
+
output_tokens: int = 0
|
|
114
|
+
cache_read_tokens: int = 0
|
|
115
|
+
reasoning_tokens: int = 0
|
|
116
|
+
|
|
117
|
+
@property
|
|
118
|
+
def total(self) -> int:
|
|
119
|
+
"""Return everything sent plus everything generated."""
|
|
120
|
+
return self.input_tokens + self.output_tokens
|
|
121
|
+
|
|
122
|
+
@property
|
|
123
|
+
def cache_share(self) -> float | None:
|
|
124
|
+
"""Return the cached share of all input, or ``None`` before there is any."""
|
|
125
|
+
return self.cache_read_tokens / self.input_tokens if self.input_tokens else None
|
|
126
|
+
|
|
127
|
+
def add(self, other: Tally) -> None:
|
|
128
|
+
"""Add another tally's counters into this one."""
|
|
129
|
+
self.requests += other.requests
|
|
130
|
+
self.input_tokens += other.input_tokens
|
|
131
|
+
self.output_tokens += other.output_tokens
|
|
132
|
+
self.cache_read_tokens += other.cache_read_tokens
|
|
133
|
+
self.reasoning_tokens += other.reasoning_tokens
|
|
134
|
+
|
|
135
|
+
|
|
136
|
+
@dataclass(frozen=True, slots=True)
|
|
137
|
+
class SessionUsage:
|
|
138
|
+
"""What one session spent, how long it is, and when it began and last changed."""
|
|
139
|
+
|
|
140
|
+
session_id: str
|
|
141
|
+
title: str | None
|
|
142
|
+
started: datetime
|
|
143
|
+
updated: datetime
|
|
144
|
+
tally: Tally
|
|
145
|
+
models: tuple[str, ...]
|
|
146
|
+
entries: int = 0
|
|
147
|
+
|
|
148
|
+
|
|
149
|
+
@dataclass(frozen=True, slots=True)
|
|
150
|
+
class SessionStats:
|
|
151
|
+
"""The totals one row of the sessions panel shows.
|
|
152
|
+
|
|
153
|
+
Attributes:
|
|
154
|
+
session_id: The session.
|
|
155
|
+
entries: Message records its conversation file holds.
|
|
156
|
+
requests: Model replies it was charged for (a fork is not charged for the replies it copied).
|
|
157
|
+
input_tokens: Everything sent to the model, cache reads included.
|
|
158
|
+
output_tokens: Everything generated, reasoning included.
|
|
159
|
+
cache_read_tokens: The part of the input the provider served from its cache.
|
|
160
|
+
reasoning_tokens: The reasoning part of the output.
|
|
161
|
+
"""
|
|
162
|
+
|
|
163
|
+
session_id: str
|
|
164
|
+
entries: int = 0
|
|
165
|
+
requests: int = 0
|
|
166
|
+
input_tokens: int = 0
|
|
167
|
+
output_tokens: int = 0
|
|
168
|
+
cache_read_tokens: int = 0
|
|
169
|
+
reasoning_tokens: int = 0
|
|
170
|
+
|
|
171
|
+
@property
|
|
172
|
+
def total(self) -> int:
|
|
173
|
+
"""Return everything sent plus everything generated."""
|
|
174
|
+
return self.input_tokens + self.output_tokens
|
|
175
|
+
|
|
176
|
+
@property
|
|
177
|
+
def cache_share(self) -> float | None:
|
|
178
|
+
"""Return the cached share of all input, or ``None`` before there is any."""
|
|
179
|
+
return self.cache_read_tokens / self.input_tokens if self.input_tokens else None
|
|
180
|
+
|
|
181
|
+
|
|
182
|
+
@dataclass(slots=True)
|
|
183
|
+
class Ledger:
|
|
184
|
+
"""A workspace's usage three ways: by model, by session, and by day.
|
|
185
|
+
|
|
186
|
+
Attributes:
|
|
187
|
+
sessions: One row per counted session, last changed first. A session that has not
|
|
188
|
+
spent anything is a row of zeros, so the list is the list ``/resume`` shows.
|
|
189
|
+
models: Totals by model name.
|
|
190
|
+
days: Totals by calendar day, only for days something was spent.
|
|
191
|
+
total: Everything, which is also the sum of any one of the three.
|
|
192
|
+
"""
|
|
193
|
+
|
|
194
|
+
sessions: list[SessionUsage] = field(default_factory=list)
|
|
195
|
+
models: dict[str, Tally] = field(default_factory=dict)
|
|
196
|
+
days: dict[date, Tally] = field(default_factory=dict)
|
|
197
|
+
total: Tally = field(default_factory=Tally)
|
|
198
|
+
|
|
199
|
+
@property
|
|
200
|
+
def empty(self) -> bool:
|
|
201
|
+
"""Return whether no reply has been counted at all."""
|
|
202
|
+
return self.total.requests == 0
|
|
203
|
+
|
|
204
|
+
|
|
205
|
+
def parse_time(text: object) -> datetime | None:
|
|
206
|
+
"""Return a stored UTC timestamp as an aware datetime, or ``None`` when it is not one."""
|
|
207
|
+
if not isinstance(text, str):
|
|
208
|
+
return None
|
|
209
|
+
try:
|
|
210
|
+
# ``fromisoformat`` only learned the trailing ``Z`` in 3.11.
|
|
211
|
+
moment = datetime.fromisoformat(text.replace("Z", "+00:00"))
|
|
212
|
+
except ValueError:
|
|
213
|
+
return None
|
|
214
|
+
return moment if moment.tzinfo is not None else moment.replace(tzinfo=timezone.utc)
|
|
215
|
+
|
|
216
|
+
|
|
217
|
+
def _count(value: object) -> int:
|
|
218
|
+
"""Return a stored counter, reading anything that is not a whole number as zero."""
|
|
219
|
+
return value if isinstance(value, int) and not isinstance(value, bool) and value > 0 else 0
|
|
220
|
+
|
|
221
|
+
|
|
222
|
+
def _model_of(message: object) -> str:
|
|
223
|
+
"""Return the model a stored message names, which sits inside its first envelope."""
|
|
224
|
+
try:
|
|
225
|
+
name = message[0]["data"]["model"] # type: ignore[index]
|
|
226
|
+
except (IndexError, KeyError, TypeError):
|
|
227
|
+
return UNKNOWN_MODEL
|
|
228
|
+
return name if isinstance(name, str) and name else UNKNOWN_MODEL
|
|
229
|
+
|
|
230
|
+
|
|
231
|
+
def _row_of(line: str) -> Row | None:
|
|
232
|
+
"""Return the reply a stored line records, or ``None`` when it records none."""
|
|
233
|
+
# Most lines carry no usage; this keeps them from being parsed at all. A serializer that
|
|
234
|
+
# leaves out the space after the colon is read the same.
|
|
235
|
+
if '"usage":' not in line or '"usage": null' in line or '"usage":null' in line:
|
|
236
|
+
return None
|
|
237
|
+
try:
|
|
238
|
+
record = json.loads(line)
|
|
239
|
+
except ValueError:
|
|
240
|
+
return None
|
|
241
|
+
if not isinstance(record, dict) or record.get("kind") != "message":
|
|
242
|
+
return None
|
|
243
|
+
usage = record.get("usage")
|
|
244
|
+
moment = parse_time(record.get("created_at"))
|
|
245
|
+
if not isinstance(usage, dict) or moment is None:
|
|
246
|
+
return None
|
|
247
|
+
return (
|
|
248
|
+
str(record.get("request_id")),
|
|
249
|
+
str(record.get("created_at")),
|
|
250
|
+
int(moment.timestamp()),
|
|
251
|
+
_model_of(record.get("message")),
|
|
252
|
+
_count(usage.get("input_tokens")),
|
|
253
|
+
_count(usage.get("output_tokens")),
|
|
254
|
+
_count(usage.get("cache_read_tokens")),
|
|
255
|
+
_count(usage.get("reasoning_tokens")),
|
|
256
|
+
)
|
|
257
|
+
|
|
258
|
+
|
|
259
|
+
def _key_of(row: Row) -> tuple[str, str, str, int, int, int]:
|
|
260
|
+
"""Return what identifies a reply across sessions: a copy of it has the same key."""
|
|
261
|
+
return (row[0], row[1], row[3], row[4], row[5], row[6])
|
|
262
|
+
|
|
263
|
+
|
|
264
|
+
@dataclass(frozen=True, slots=True)
|
|
265
|
+
class _File:
|
|
266
|
+
"""What is known of one conversation file: its replies, and how far they were read.
|
|
267
|
+
|
|
268
|
+
Attributes:
|
|
269
|
+
size: Bytes the file had when it was looked at; with ``mtime_ns``, "nothing changed".
|
|
270
|
+
mtime_ns: Its modification time then.
|
|
271
|
+
offset: Bytes consumed, which stops before a final line that has no newline yet.
|
|
272
|
+
head: Fingerprint of the first bytes, so another file under the same name is noticed.
|
|
273
|
+
tail: Fingerprint of the bytes just before ``offset``, so an edit in what was read is.
|
|
274
|
+
started: The header's ``created_at``, as stored.
|
|
275
|
+
counted: False for a subagent's session, whose replies are never read.
|
|
276
|
+
rows: The replies found so far.
|
|
277
|
+
entries: The message records found so far.
|
|
278
|
+
"""
|
|
279
|
+
|
|
280
|
+
size: int
|
|
281
|
+
mtime_ns: int
|
|
282
|
+
offset: int
|
|
283
|
+
head: int
|
|
284
|
+
tail: int
|
|
285
|
+
started: str | None
|
|
286
|
+
counted: bool
|
|
287
|
+
rows: tuple[Row, ...]
|
|
288
|
+
entries: int = 0
|
|
289
|
+
|
|
290
|
+
|
|
291
|
+
def _fingerprint(handle: BinaryIO, start: int, end: int) -> int:
|
|
292
|
+
"""Return a checksum of the bytes from ``start`` to ``end``, clipped to the file's start."""
|
|
293
|
+
start = max(0, start)
|
|
294
|
+
handle.seek(start)
|
|
295
|
+
return zlib.crc32(handle.read(max(0, end - start)))
|
|
296
|
+
|
|
297
|
+
|
|
298
|
+
def _header_of(first: bytes) -> tuple[str | None, bool]:
|
|
299
|
+
"""Return when a session began, as stored, and whether it is one the reader had."""
|
|
300
|
+
try:
|
|
301
|
+
header = json.loads(first.decode("utf-8", errors="replace"))
|
|
302
|
+
except ValueError:
|
|
303
|
+
return None, True
|
|
304
|
+
if not isinstance(header, dict):
|
|
305
|
+
return None, True
|
|
306
|
+
created = header.get("created_at")
|
|
307
|
+
return (created if isinstance(created, str) else None), header.get("agent_name") is None
|
|
308
|
+
|
|
309
|
+
|
|
310
|
+
def _scan(path: Path, previous: _File | None) -> _File | None:
|
|
311
|
+
"""Read one conversation file, continuing from ``previous`` when the file only grew.
|
|
312
|
+
|
|
313
|
+
Returns ``None`` when the file cannot be read at all.
|
|
314
|
+
"""
|
|
315
|
+
try:
|
|
316
|
+
stat = path.stat()
|
|
317
|
+
with path.open("rb") as handle:
|
|
318
|
+
head = _fingerprint(handle, 0, HEAD_BYTES)
|
|
319
|
+
if (
|
|
320
|
+
previous is not None
|
|
321
|
+
and previous.head == head
|
|
322
|
+
and 0 < previous.offset <= stat.st_size
|
|
323
|
+
and _fingerprint(handle, previous.offset - TAIL_BYTES, previous.offset) == previous.tail
|
|
324
|
+
):
|
|
325
|
+
if not previous.counted:
|
|
326
|
+
return replace(previous, size=stat.st_size, mtime_ns=stat.st_mtime_ns)
|
|
327
|
+
offset, started, counted, rows = previous.offset, previous.started, True, list(previous.rows)
|
|
328
|
+
entries = previous.entries
|
|
329
|
+
else:
|
|
330
|
+
handle.seek(0)
|
|
331
|
+
first = handle.readline()
|
|
332
|
+
# A header with no newline yet is a session still being born: read nothing of it.
|
|
333
|
+
offset = len(first) if first.endswith(b"\n") else 0
|
|
334
|
+
started, counted = _header_of(first)
|
|
335
|
+
rows = []
|
|
336
|
+
entries = 0
|
|
337
|
+
if counted and offset:
|
|
338
|
+
handle.seek(offset)
|
|
339
|
+
for raw in handle:
|
|
340
|
+
# A line without its newline is still being written: leave it for next time.
|
|
341
|
+
if not raw.endswith(b"\n"):
|
|
342
|
+
break
|
|
343
|
+
offset += len(raw)
|
|
344
|
+
if any(mark in raw for mark in MESSAGE_MARKS):
|
|
345
|
+
entries += 1
|
|
346
|
+
row = _row_of(raw.decode("utf-8", errors="replace"))
|
|
347
|
+
if row is not None:
|
|
348
|
+
rows.append(row)
|
|
349
|
+
tail = _fingerprint(handle, offset - TAIL_BYTES, offset)
|
|
350
|
+
except OSError:
|
|
351
|
+
return None
|
|
352
|
+
return _File(stat.st_size, stat.st_mtime_ns, offset, head, tail, started, counted, tuple(rows), entries)
|
|
353
|
+
|
|
354
|
+
|
|
355
|
+
def _title_of(folder: Path) -> str | None:
|
|
356
|
+
"""Return the title a session's metadata file holds, if there is one."""
|
|
357
|
+
try:
|
|
358
|
+
data = json.loads((folder / SESSION_METADATA_FILE).read_text(encoding="utf-8"))
|
|
359
|
+
except (OSError, ValueError):
|
|
360
|
+
return None
|
|
361
|
+
title = data.get("title") if isinstance(data, dict) else None
|
|
362
|
+
return title if isinstance(title, str) and title.strip() else None
|
|
363
|
+
|
|
364
|
+
|
|
365
|
+
def _whole(value: object) -> bool:
|
|
366
|
+
"""Say whether a stored value is an integer (a boolean is not one here)."""
|
|
367
|
+
return isinstance(value, int) and not isinstance(value, bool)
|
|
368
|
+
|
|
369
|
+
|
|
370
|
+
def _row_from_index(row: object) -> Row:
|
|
371
|
+
"""Check one stored row's shape and return it as a tuple; raise ``ValueError`` if it is odd."""
|
|
372
|
+
if not isinstance(row, list) or len(row) != 8:
|
|
373
|
+
raise ValueError("a row has eight fields")
|
|
374
|
+
request, created, seconds, model, *counters = row
|
|
375
|
+
if not (isinstance(request, str) and isinstance(created, str) and isinstance(model, str)):
|
|
376
|
+
raise ValueError("a row's text fields are text")
|
|
377
|
+
if not all(_whole(number) for number in (seconds, *counters)):
|
|
378
|
+
raise ValueError("a row's counters are whole numbers")
|
|
379
|
+
return (request, created, seconds, model, *counters)
|
|
380
|
+
|
|
381
|
+
|
|
382
|
+
def _file_from_index(item: object) -> _File | None:
|
|
383
|
+
"""Rebuild one file's entry from the index, or ``None`` when it is not a well-formed one."""
|
|
384
|
+
if not isinstance(item, dict):
|
|
385
|
+
return None
|
|
386
|
+
try:
|
|
387
|
+
size, mtime_ns, offset, head, tail = (item[name] for name in ("size", "mtime_ns", "offset", "head", "tail"))
|
|
388
|
+
started, counted = item["started"], item["counted"]
|
|
389
|
+
entries = item["entries"]
|
|
390
|
+
if not all(_whole(number) for number in (size, mtime_ns, offset, head, tail, entries)) or entries < 0:
|
|
391
|
+
return None
|
|
392
|
+
if not isinstance(counted, bool) or not (started is None or isinstance(started, str)):
|
|
393
|
+
return None
|
|
394
|
+
rows = tuple(_row_from_index(row) for row in item["rows"])
|
|
395
|
+
except (KeyError, TypeError, ValueError):
|
|
396
|
+
return None
|
|
397
|
+
return _File(size, mtime_ns, offset, head, tail, started, counted, rows, entries)
|
|
398
|
+
|
|
399
|
+
|
|
400
|
+
class LedgerReader:
|
|
401
|
+
"""Reads a workspace's ledger, reading as little as it can.
|
|
402
|
+
|
|
403
|
+
Two things are remembered. In memory, each file's result is kept while its size and
|
|
404
|
+
modification time are as they were, so opening ``/usage`` again a minute later reads the
|
|
405
|
+
session being worked in and nothing else. On disk, the index (see the module's notes) lets a
|
|
406
|
+
program that has just started pick up where the last one stopped, and lets a file that has
|
|
407
|
+
grown be read from where it was last read instead of from its first byte.
|
|
408
|
+
"""
|
|
409
|
+
|
|
410
|
+
def __init__(self, root: str | Path, workspace: str | Path, *, index: bool = True) -> None:
|
|
411
|
+
"""Name the store and the workspace; nothing is read or created yet.
|
|
412
|
+
|
|
413
|
+
Args:
|
|
414
|
+
root: The directory holding one folder per workspace.
|
|
415
|
+
workspace: The project whose sessions are counted.
|
|
416
|
+
index: Keep the index on disk. ``False`` is the in-memory behaviour alone.
|
|
417
|
+
"""
|
|
418
|
+
self.directory = find_workspace(Path(root).expanduser(), workspace)
|
|
419
|
+
self.index_path = self.directory / INDEX_FILE
|
|
420
|
+
self.summary_path = self.directory / SUMMARY_FILE
|
|
421
|
+
self._use_index = index
|
|
422
|
+
self._files: dict[str, _File] = {}
|
|
423
|
+
self._loaded = False
|
|
424
|
+
|
|
425
|
+
def read(self) -> Ledger:
|
|
426
|
+
"""Return the ledger of every session the workspace holds right now."""
|
|
427
|
+
if not self._loaded:
|
|
428
|
+
self._loaded = True
|
|
429
|
+
if self._use_index:
|
|
430
|
+
self._files = self._load_index()
|
|
431
|
+
found: list[tuple[datetime, str, Path, _File]] = []
|
|
432
|
+
kept: dict[str, _File] = {}
|
|
433
|
+
changed = False
|
|
434
|
+
for path in sorted(self.directory.glob(f"*/{STORE_FILE}")):
|
|
435
|
+
session_id = path.parent.name
|
|
436
|
+
if not SAFE_SESSION_ID.match(session_id) or ".." in session_id:
|
|
437
|
+
continue
|
|
438
|
+
cached = self._files.get(session_id)
|
|
439
|
+
try:
|
|
440
|
+
stat = path.stat()
|
|
441
|
+
except OSError:
|
|
442
|
+
continue
|
|
443
|
+
if cached is not None and cached.size == stat.st_size and cached.mtime_ns == stat.st_mtime_ns:
|
|
444
|
+
read: _File | None = cached
|
|
445
|
+
else:
|
|
446
|
+
read = _scan(path, cached)
|
|
447
|
+
changed = True
|
|
448
|
+
if read is None:
|
|
449
|
+
continue
|
|
450
|
+
kept[session_id] = read
|
|
451
|
+
if read.counted:
|
|
452
|
+
began = parse_time(read.started) or datetime.fromtimestamp(0, tz=timezone.utc)
|
|
453
|
+
found.append((began, session_id, path, read))
|
|
454
|
+
# A session folder deleted since the last read is forgotten with it.
|
|
455
|
+
changed = changed or kept.keys() != self._files.keys()
|
|
456
|
+
self._files = kept
|
|
457
|
+
if changed and self._use_index and kept:
|
|
458
|
+
self._save_index(kept)
|
|
459
|
+
return _fold(found)
|
|
460
|
+
|
|
461
|
+
# -- the index ------------------------------------------------------------
|
|
462
|
+
def _load_index(self) -> dict[str, _File]:
|
|
463
|
+
"""Return the files the index on disk describes, or none when it cannot be trusted."""
|
|
464
|
+
try:
|
|
465
|
+
data = json.loads(self.index_path.read_text(encoding="utf-8"))
|
|
466
|
+
except (OSError, ValueError):
|
|
467
|
+
return {}
|
|
468
|
+
if not isinstance(data, dict) or data.get("version") != INDEX_VERSION:
|
|
469
|
+
return {}
|
|
470
|
+
sessions = data.get("sessions")
|
|
471
|
+
if not isinstance(sessions, dict):
|
|
472
|
+
return {}
|
|
473
|
+
files: dict[str, _File] = {}
|
|
474
|
+
for session_id, item in sessions.items():
|
|
475
|
+
entry = _file_from_index(item)
|
|
476
|
+
if entry is not None:
|
|
477
|
+
files[session_id] = entry
|
|
478
|
+
return files
|
|
479
|
+
|
|
480
|
+
def _save_index(self, files: dict[str, _File]) -> None:
|
|
481
|
+
"""Write the index whole and atomically; a store that cannot be written is left alone."""
|
|
482
|
+
payload = {
|
|
483
|
+
"version": INDEX_VERSION,
|
|
484
|
+
"sessions": {
|
|
485
|
+
session_id: {
|
|
486
|
+
"size": item.size,
|
|
487
|
+
"mtime_ns": item.mtime_ns,
|
|
488
|
+
"offset": item.offset,
|
|
489
|
+
"head": item.head,
|
|
490
|
+
"tail": item.tail,
|
|
491
|
+
"started": item.started,
|
|
492
|
+
"counted": item.counted,
|
|
493
|
+
"entries": item.entries,
|
|
494
|
+
"rows": [list(row) for row in item.rows],
|
|
495
|
+
}
|
|
496
|
+
for session_id, item in files.items()
|
|
497
|
+
},
|
|
498
|
+
}
|
|
499
|
+
self._write(self.index_path, payload)
|
|
500
|
+
|
|
501
|
+
def _write(self, path: Path, payload: dict[str, object]) -> None:
|
|
502
|
+
"""Write a cache file whole and atomically; a store that cannot be written is left alone."""
|
|
503
|
+
temporary = path.with_name(f"{path.name}.{os.getpid()}.tmp")
|
|
504
|
+
try:
|
|
505
|
+
temporary.write_text(json.dumps(payload, separators=(",", ":")), encoding="utf-8")
|
|
506
|
+
os.replace(temporary, path)
|
|
507
|
+
except OSError:
|
|
508
|
+
try:
|
|
509
|
+
temporary.unlink()
|
|
510
|
+
except OSError:
|
|
511
|
+
pass
|
|
512
|
+
|
|
513
|
+
# -- the summary ----------------------------------------------------------
|
|
514
|
+
def stats(self, session_ids: Iterable[str]) -> dict[str, SessionStats]:
|
|
515
|
+
"""Return the totals of the given sessions, reading as little as it can.
|
|
516
|
+
|
|
517
|
+
A session whose conversation file still has the size and modification time the summary recorded
|
|
518
|
+
is answered from the summary, so listing twenty sessions costs twenty ``stat`` calls however long
|
|
519
|
+
their conversations are. If any asked-for session has changed, or was never summarised, or a summarised
|
|
520
|
+
session has gone (which can move a reply from a fork back to the session that held it first), the
|
|
521
|
+
incremental read runs once and the summary is rewritten whole.
|
|
522
|
+
|
|
523
|
+
A session that is not counted (a subagent's) or cannot be read has no entry in the answer.
|
|
524
|
+
"""
|
|
525
|
+
wanted = [
|
|
526
|
+
session_id for session_id in session_ids if SAFE_SESSION_ID.match(session_id) and ".." not in session_id
|
|
527
|
+
]
|
|
528
|
+
summary = self._load_summary()
|
|
529
|
+
answer: dict[str, SessionStats] = {}
|
|
530
|
+
if summary and self._summary_is_current(summary):
|
|
531
|
+
for session_id in wanted:
|
|
532
|
+
entry = summary.get(session_id)
|
|
533
|
+
if entry is None:
|
|
534
|
+
break
|
|
535
|
+
if self._unchanged(session_id, entry):
|
|
536
|
+
answer[session_id] = entry[1]
|
|
537
|
+
continue
|
|
538
|
+
break
|
|
539
|
+
else:
|
|
540
|
+
return answer
|
|
541
|
+
ledger = self.read()
|
|
542
|
+
counted = {row.session_id: row for row in ledger.sessions}
|
|
543
|
+
fresh: dict[str, tuple[tuple[int, int], SessionStats]] = {}
|
|
544
|
+
for session_id, row in counted.items():
|
|
545
|
+
# The stamp is the one the read itself took, before it read: a file that grew since is then seen to
|
|
546
|
+
# differ, where a stamp taken now would pair the new size with the old totals.
|
|
547
|
+
scanned = self._files.get(session_id)
|
|
548
|
+
if scanned is not None:
|
|
549
|
+
fresh[session_id] = ((scanned.size, scanned.mtime_ns), _stats_of(row))
|
|
550
|
+
if self._use_index and fresh:
|
|
551
|
+
self._save_summary(fresh)
|
|
552
|
+
return {session_id: fresh[session_id][1] for session_id in wanted if session_id in fresh}
|
|
553
|
+
|
|
554
|
+
def _stamp(self, session_id: str) -> tuple[int, int] | None:
|
|
555
|
+
"""Return a conversation file's size and modification time, or ``None`` when it cannot be seen."""
|
|
556
|
+
try:
|
|
557
|
+
stat = (self.directory / session_id / STORE_FILE).stat()
|
|
558
|
+
except OSError:
|
|
559
|
+
return None
|
|
560
|
+
return stat.st_size, stat.st_mtime_ns
|
|
561
|
+
|
|
562
|
+
def _unchanged(self, session_id: str, entry: tuple[tuple[int, int], SessionStats]) -> bool:
|
|
563
|
+
return self._stamp(session_id) == entry[0]
|
|
564
|
+
|
|
565
|
+
def _summary_is_current(self, summary: dict[str, tuple[tuple[int, int], SessionStats]]) -> bool:
|
|
566
|
+
"""Say whether every session the summary names still has a conversation file."""
|
|
567
|
+
present = {path.parent.name for path in self.directory.glob(f"*/{STORE_FILE}")}
|
|
568
|
+
return summary.keys() <= present
|
|
569
|
+
|
|
570
|
+
def _load_summary(self) -> dict[str, tuple[tuple[int, int], SessionStats]]:
|
|
571
|
+
"""Return what the summary file holds, or nothing when it cannot be trusted."""
|
|
572
|
+
if not self._use_index:
|
|
573
|
+
return {}
|
|
574
|
+
try:
|
|
575
|
+
data = json.loads(self.summary_path.read_text(encoding="utf-8"))
|
|
576
|
+
except (OSError, ValueError):
|
|
577
|
+
return {}
|
|
578
|
+
if not isinstance(data, dict) or data.get("version") != SUMMARY_VERSION:
|
|
579
|
+
return {}
|
|
580
|
+
sessions = data.get("sessions")
|
|
581
|
+
if not isinstance(sessions, dict):
|
|
582
|
+
return {}
|
|
583
|
+
found: dict[str, tuple[tuple[int, int], SessionStats]] = {}
|
|
584
|
+
for session_id, item in sessions.items():
|
|
585
|
+
if not isinstance(session_id, str) or not isinstance(item, list) or len(item) != 8:
|
|
586
|
+
continue
|
|
587
|
+
if not all(_whole(number) and number >= 0 for number in item):
|
|
588
|
+
continue
|
|
589
|
+
size, mtime_ns, *counters = item
|
|
590
|
+
found[session_id] = ((size, mtime_ns), SessionStats(session_id, *counters))
|
|
591
|
+
return found
|
|
592
|
+
|
|
593
|
+
def _save_summary(self, fresh: dict[str, tuple[tuple[int, int], SessionStats]]) -> None:
|
|
594
|
+
"""Write the summary: for each session its stamp, then the seven numbers a row shows."""
|
|
595
|
+
payload = {
|
|
596
|
+
"version": SUMMARY_VERSION,
|
|
597
|
+
"sessions": {
|
|
598
|
+
session_id: [
|
|
599
|
+
stamp[0],
|
|
600
|
+
stamp[1],
|
|
601
|
+
stats.entries,
|
|
602
|
+
stats.requests,
|
|
603
|
+
stats.input_tokens,
|
|
604
|
+
stats.output_tokens,
|
|
605
|
+
stats.cache_read_tokens,
|
|
606
|
+
stats.reasoning_tokens,
|
|
607
|
+
]
|
|
608
|
+
for session_id, (stamp, stats) in fresh.items()
|
|
609
|
+
},
|
|
610
|
+
}
|
|
611
|
+
self._write(self.summary_path, payload)
|
|
612
|
+
|
|
613
|
+
|
|
614
|
+
def _fold(found: Iterable[tuple[datetime, str, Path, _File]]) -> Ledger:
|
|
615
|
+
"""Fold sessions, oldest first, into a ledger that charges each reply once."""
|
|
616
|
+
ledger = Ledger()
|
|
617
|
+
seen: set[tuple[str, str, str, int, int, int]] = set()
|
|
618
|
+
days: dict[int, date] = {}
|
|
619
|
+
for started, session_id, path, read in sorted(found, key=lambda item: (item[0], item[1])):
|
|
620
|
+
tally = Tally()
|
|
621
|
+
models: dict[str, None] = {}
|
|
622
|
+
for row in read.rows:
|
|
623
|
+
key = _key_of(row)
|
|
624
|
+
if key in seen:
|
|
625
|
+
continue
|
|
626
|
+
seen.add(key)
|
|
627
|
+
reply = Tally(1, row[4], row[5], row[6], row[7])
|
|
628
|
+
tally.add(reply)
|
|
629
|
+
ledger.models.setdefault(row[3], Tally()).add(reply)
|
|
630
|
+
bucket = row[2] // BUCKET_SECONDS
|
|
631
|
+
if bucket not in days:
|
|
632
|
+
moment = datetime.fromtimestamp(bucket * BUCKET_SECONDS, tz=timezone.utc)
|
|
633
|
+
days[bucket] = moment.astimezone().date()
|
|
634
|
+
ledger.days.setdefault(days[bucket], Tally()).add(reply)
|
|
635
|
+
models[row[3]] = None
|
|
636
|
+
ledger.total.add(tally)
|
|
637
|
+
try:
|
|
638
|
+
updated = datetime.fromtimestamp(path.stat().st_mtime, tz=timezone.utc)
|
|
639
|
+
except OSError:
|
|
640
|
+
updated = started
|
|
641
|
+
title = _title_of(path.parent)
|
|
642
|
+
ledger.sessions.append(SessionUsage(session_id, title, started, updated, tally, tuple(models), read.entries))
|
|
643
|
+
# The list answers "what was I working on", so it is ordered by the last change, not by the first word.
|
|
644
|
+
ledger.sessions.sort(key=lambda row: (row.updated, row.started, row.session_id), reverse=True)
|
|
645
|
+
return ledger
|
|
646
|
+
|
|
647
|
+
|
|
648
|
+
def _stats_of(row: SessionUsage) -> SessionStats:
|
|
649
|
+
"""Return the totals of one session's row of the ledger."""
|
|
650
|
+
tally = row.tally
|
|
651
|
+
return SessionStats(
|
|
652
|
+
row.session_id,
|
|
653
|
+
row.entries,
|
|
654
|
+
tally.requests,
|
|
655
|
+
tally.input_tokens,
|
|
656
|
+
tally.output_tokens,
|
|
657
|
+
tally.cache_read_tokens,
|
|
658
|
+
tally.reasoning_tokens,
|
|
659
|
+
)
|