54 lines
1.8 KiB
Python
54 lines
1.8 KiB
Python
"""Token-usage estimation for the active chain of a JSONL transcript.
|
|
|
|
v0.1 ships only `char_div_3_5` — count characters of the active chain's content
|
|
fields and divide by 3.5. This is a rough but cheap heuristic; tiktoken-based
|
|
estimation can be added later behind the same interface.
|
|
"""
|
|
from __future__ import annotations
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from .jsonl_ops import find_active_chain, load_messages
|
|
|
|
|
|
def _char_count(content: Any) -> int:
|
|
if content is None:
|
|
return 0
|
|
if isinstance(content, str):
|
|
return len(content)
|
|
if isinstance(content, list):
|
|
total = 0
|
|
for item in content:
|
|
if isinstance(item, dict):
|
|
if "text" in item and isinstance(item["text"], str):
|
|
total += len(item["text"])
|
|
elif "input" in item:
|
|
total += len(json.dumps(item.get("input", {}), ensure_ascii=False))
|
|
elif "content" in item:
|
|
total += _char_count(item["content"])
|
|
else:
|
|
total += len(str(item))
|
|
return total
|
|
if isinstance(content, dict):
|
|
return _char_count(content.get("content"))
|
|
return len(str(content))
|
|
|
|
|
|
def _message_chars(msg: dict[str, Any]) -> int:
|
|
inner = msg.get("message")
|
|
if isinstance(inner, dict):
|
|
return _char_count(inner.get("content"))
|
|
if "tool_use_id" in msg and "content" in msg:
|
|
return _char_count(msg.get("content"))
|
|
return 0
|
|
|
|
|
|
def estimate_active_chain_tokens(transcript: Path, *, method: str = "char_div_3_5") -> int:
|
|
if method != "char_div_3_5":
|
|
raise ValueError(f"Unsupported estimate_method: {method}")
|
|
msgs = load_messages(transcript)
|
|
chain = find_active_chain(msgs)
|
|
chars = sum(_message_chars(m) for m in chain)
|
|
return int(chars / 3.5)
|