216 lines
8.9 KiB
Python
216 lines
8.9 KiB
Python
"""Stream-json reader and accounting for one `claude -p` call: pure functions, no IO."""
|
|
import datetime as _dt
|
|
import json
|
|
import re
|
|
import time
|
|
|
|
TTL_1H, TTL_5M = 3600, 300
|
|
CACHE_HIT_RATIO = 0.9
|
|
MU_FIELDS = ("inputTokens", "outputTokens", "cacheReadInputTokens", "cacheCreationInputTokens")
|
|
TOKEN_KEYS = ("input", "output", "cache_creation", "cache_read", "cache_creation_1h", "cache_creation_5m")
|
|
|
|
|
|
def iso(t: float | None = None) -> str:
|
|
"""ISO 8601 with the local offset, from epoch seconds (now if not given)."""
|
|
if t is None:
|
|
t = time.time()
|
|
return _dt.datetime.fromtimestamp(t).astimezone().isoformat(timespec="seconds")
|
|
|
|
|
|
def parse_iso(s) -> float | None:
|
|
"""Epoch seconds of an ISO string written by `iso`, or None."""
|
|
if not s or not isinstance(s, str):
|
|
return None
|
|
try:
|
|
return _dt.datetime.fromisoformat(s).timestamp()
|
|
except ValueError:
|
|
return None
|
|
|
|
|
|
class StreamState:
|
|
"""What one `claude -p --output-format stream-json --verbose` stream told us so far."""
|
|
|
|
def __init__(self):
|
|
self.init = None # the system/init event
|
|
self.session_id = None
|
|
self.msg_ids: list = []
|
|
self._ids: set = set()
|
|
self.first_usage = None # usage of the first assistant event (cache verdict)
|
|
self.first_at = None # its arrival time (the base prefix read)
|
|
self.last_usage = None # usage of the last assistant event (base context)
|
|
self.last_assistant_at = None
|
|
self.result = None # the result event
|
|
self.refused = None # text of a model-restricted warning
|
|
self.events = 0
|
|
|
|
def feed(self, ev: dict, t: float | None) -> None:
|
|
if not isinstance(ev, dict):
|
|
return
|
|
self.events += 1
|
|
ty = ev.get("type")
|
|
if ty == "system":
|
|
st = ev.get("subtype")
|
|
if st == "init" and self.init is None:
|
|
self.init = ev
|
|
self.session_id = ev.get("session_id") or self.session_id
|
|
elif (st == "informational" and ev.get("level") == "warning"
|
|
and "is restricted" in str(ev.get("content", ""))):
|
|
self.refused = str(ev.get("content"))
|
|
elif ty == "assistant":
|
|
m = ev.get("message") or {}
|
|
mid = m.get("id")
|
|
if mid not in self._ids:
|
|
self._ids.add(mid)
|
|
self.msg_ids.append(mid)
|
|
u = m.get("usage") or {}
|
|
if self.first_usage is None:
|
|
self.first_usage, self.first_at = u, t
|
|
self.last_usage, self.last_assistant_at = u, t
|
|
elif ty == "result":
|
|
self.result = ev
|
|
self.session_id = ev.get("session_id") or self.session_id
|
|
|
|
@property
|
|
def turns(self) -> int:
|
|
return len(self.msg_ids)
|
|
|
|
|
|
def context_tokens(usage: dict | None) -> int:
|
|
"""Prompt-side tokens of one request: input + cache creation + cache read."""
|
|
if not usage:
|
|
return 0
|
|
return (int(usage.get("input_tokens") or 0) + int(usage.get("cache_creation_input_tokens") or 0)
|
|
+ int(usage.get("cache_read_input_tokens") or 0))
|
|
|
|
|
|
_API_ERROR_TEXT = ( # checked in this order, after the status code
|
|
("prompt_too_long", re.compile(r"prompt is too long|prompt too long|prompt_too_long", re.I)),
|
|
("usage_limit", re.compile(r"usage limit|rate[ _-]?limit|hit your limit|api error:?\s*429", re.I)),
|
|
("overloaded", re.compile(r"overloaded|api error:?\s*529", re.I)),
|
|
("api_error", re.compile(r"api error", re.I)),
|
|
)
|
|
|
|
|
|
def api_error_reason(res: dict | None) -> tuple[str, str] | None:
|
|
"""(reason, detail) for a result that carries an API error, else None. `reason` is
|
|
`usage_limit` (429), `overloaded` (529), `prompt_too_long` or `api_error`."""
|
|
if not res:
|
|
return None
|
|
status = res.get("api_error_status")
|
|
try:
|
|
status = int(status) if status not in (None, "") else None
|
|
except (TypeError, ValueError):
|
|
status = None
|
|
if not res.get("is_error") and status is None:
|
|
return None
|
|
parts = [res.get("result"), res.get("error"), res.get("errors")]
|
|
text = " | ".join(p if isinstance(p, str) else json.dumps(p) for p in parts if p)
|
|
detail = "api_error_status %s: %s" % (status, text[:200] if text else "(no text)")
|
|
if status == 429:
|
|
return "usage_limit", detail
|
|
if status == 529:
|
|
return "overloaded", detail
|
|
for reason, rx in _API_ERROR_TEXT:
|
|
if rx.search(text):
|
|
return reason, detail
|
|
if status is not None:
|
|
return "api_error", detail
|
|
return None
|
|
|
|
|
|
def classify(ss: StreamState, exit_code, stopped: bool) -> tuple[str, str | None]:
|
|
"""(status, reason) of a finished call; the first matching case wins."""
|
|
res = ss.result
|
|
match (bool(stopped), ss.refused, res, api_error_reason(res), exit_code):
|
|
case (True, _, _, _, _):
|
|
return "stopped", "stopped"
|
|
case (_, str(), _, _, _):
|
|
return "error", "model_refused"
|
|
case (_, _, None, _, code):
|
|
return "error", "exit_code" if code else "no_result"
|
|
case (_, _, _, (reason, _), _):
|
|
return "error", reason
|
|
case (_, _, dict() as r, _, _) if r.get("is_error"):
|
|
return "error", "is_error"
|
|
case (_, _, _, _, code) if code:
|
|
return "error", "exit_code"
|
|
case _:
|
|
return "ok", None
|
|
|
|
|
|
def usage_from_result(result: dict, before_mu: dict, cost_before) -> dict:
|
|
"""Tokens (with the 1h/5m split) from `result.usage`; `models_used` = modelUsage entries that
|
|
grew against the session's totals before the call, largest first; cost delta against the
|
|
session's cost before the call. Costs are rounded here, once, at 6 decimals."""
|
|
result = result or {}
|
|
usage = result.get("usage") or {}
|
|
cc = usage.get("cache_creation") or {}
|
|
tokens = {
|
|
"input": int(usage.get("input_tokens") or 0),
|
|
"output": int(usage.get("output_tokens") or 0),
|
|
"cache_creation": int(usage.get("cache_creation_input_tokens") or 0),
|
|
"cache_read": int(usage.get("cache_read_input_tokens") or 0),
|
|
"cache_creation_1h": int(cc.get("ephemeral_1h_input_tokens") or 0),
|
|
"cache_creation_5m": int(cc.get("ephemeral_5m_input_tokens") or 0),
|
|
}
|
|
after = result.get("modelUsage") or {}
|
|
before_mu = before_mu or {}
|
|
delta = {}
|
|
for m, v in after.items():
|
|
b = before_mu.get(m) or {}
|
|
delta[m] = {f: int((v or {}).get(f) or 0) - int(b.get(f) or 0) for f in MU_FIELDS}
|
|
grew = [m for m, d in delta.items() if any(x > 0 for x in d.values())]
|
|
models_used = sorted(grew, key=lambda m: (-sum(delta[m].values()), m))
|
|
sums = {f: sum(d[f] for d in delta.values()) for f in MU_FIELDS}
|
|
mismatch = (sums["inputTokens"], sums["outputTokens"], sums["cacheCreationInputTokens"],
|
|
sums["cacheReadInputTokens"]) != (tokens["input"], tokens["output"],
|
|
tokens["cache_creation"], tokens["cache_read"])
|
|
total = result.get("total_cost_usd")
|
|
if total is None:
|
|
cost_delta = None
|
|
else:
|
|
cost_delta = round(float(total) - float(cost_before or 0.0), 6)
|
|
total = round(float(total), 6)
|
|
return {"tokens": tokens, "models_used": models_used, "mismatch": mismatch,
|
|
"cost_delta": cost_delta, "total": total, "model_usage": after}
|
|
|
|
|
|
def cache_verdict(first_usage: dict | None, base_ctx) -> str | None:
|
|
"""The first request's cache read against the base's context tokens."""
|
|
if first_usage is None:
|
|
return None
|
|
cr = int(first_usage.get("cache_read_input_tokens") or 0)
|
|
if base_ctx and cr >= CACHE_HIT_RATIO * float(base_ctx):
|
|
return "hit"
|
|
return "partial" if cr > 0 else "miss"
|
|
|
|
|
|
def infer_ttl(tokens: dict | None, prev):
|
|
"""The TTL of a base's cache entries from one call's tokens: 1h written -> 3600, only 5m
|
|
written -> 300, otherwise (read only, or nothing) `prev`."""
|
|
if not tokens:
|
|
return prev
|
|
if int(tokens.get("cache_creation_1h") or 0) > 0:
|
|
return TTL_1H
|
|
if int(tokens.get("cache_creation_5m") or 0) > 0:
|
|
return TTL_5M
|
|
return prev
|
|
|
|
|
|
def wrote_only_5m(tokens: dict | None) -> bool:
|
|
"""The call wrote 5-minute cache entries and no 1-hour ones."""
|
|
return (bool(tokens) and int(tokens.get("cache_creation_5m") or 0) > 0
|
|
and int(tokens.get("cache_creation_1h") or 0) == 0)
|
|
|
|
|
|
def warmth(last_hit_ts: float | None, ttl_s, now: float) -> dict:
|
|
"""warm until `margin` before last_hit + ttl, expiring until last_hit + ttl, then cold;
|
|
margin = max(120, ttl/10). No hit at all is cold."""
|
|
if last_hit_ts is None:
|
|
return {"base_state": "cold", "warm_until": None, "ttl_left_s": None}
|
|
ttl = float(ttl_s or TTL_1H)
|
|
margin = max(120.0, ttl / 10.0)
|
|
wu = float(last_hit_ts) + ttl
|
|
state = "warm" if now < wu - margin else ("expiring" if now < wu else "cold")
|
|
return {"base_state": state, "warm_until": iso(wu), "ttl_left_s": int(max(0.0, wu - now))}
|