Initial release: aircontext

scripts/core/token_estimator.py

@@ -0,0 +1,53 @@
+"""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)