68 lines
2 KiB
Python
68 lines
2 KiB
Python
"""Sum the per-block token statistics from an opencode session trace.
|
|
|
|
Each conversation block in a trace ends with a line like::
|
|
|
|
│ tokens 18 in 134 out 63628 cached 0 written $0.0000 tool-calls
|
|
|
|
This script adds up the *in*, *out* and *cached* columns across all blocks
|
|
and prints the three totals; the written-token count, the price and the
|
|
tool-call indicator are ignored.
|
|
|
|
Usage:
|
|
|
|
python3 scripts/count_tokens.py opencode/000_opencode_session_....txt
|
|
"""
|
|
|
|
import argparse
|
|
import re
|
|
from collections.abc import Iterable
|
|
from pathlib import Path
|
|
|
|
# The stat line is preceded by a box-drawing character and ends either with
|
|
# "tool-calls" or "stop"; only the numeric columns matter.
|
|
STAT_RE = re.compile(r"tokens (\d+) in (\d+) out (\d+) cached (\d+) written")
|
|
|
|
|
|
def sum_tokens(lines: Iterable[str]) -> tuple[int, int, int]:
|
|
"""Add up the in/out/cached token counts of all stat lines.
|
|
|
|
Args:
|
|
lines: Iterable of trace lines.
|
|
|
|
Returns:
|
|
A ``(tokens_in, tokens_out, cached)`` triple of sums.
|
|
"""
|
|
tokens_in = 0
|
|
tokens_out = 0
|
|
cached = 0
|
|
for line in lines:
|
|
match = STAT_RE.search(line)
|
|
if not match:
|
|
continue
|
|
tokens_in += int(match.group(1))
|
|
tokens_out += int(match.group(2))
|
|
cached += int(match.group(3))
|
|
return tokens_in, tokens_out, cached
|
|
|
|
|
|
def main() -> None:
|
|
"""Parse arguments, read the trace file and print the totals."""
|
|
parser = argparse.ArgumentParser(
|
|
description="Sum the token statistics from an opencode session trace."
|
|
)
|
|
parser.add_argument("trace_file", type=Path, help="path to the trace file")
|
|
args = parser.parse_args()
|
|
|
|
if not args.trace_file.is_file():
|
|
parser.error(f"no such file: {args.trace_file}")
|
|
|
|
with args.trace_file.open(encoding="utf-8", errors="replace") as trace:
|
|
tokens_in, tokens_out, cached = sum_tokens(trace)
|
|
|
|
print(f"tokens in: {tokens_in}")
|
|
print(f"tokens out: {tokens_out}")
|
|
print(f"cached: {cached}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|