281 lines
10 KiB
Python
281 lines
10 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""
|
||
|
|
token_estimator.py — deterministic, local token estimation for doc-hygiene.
|
||
|
|
|
||
|
|
Produces the `raw_tokens` count that feeds each report entry's
|
||
|
|
`token_estimate` field (see the frozen report schema:
|
||
|
|
`openspec/specs/report-schema/spec.md`). This is the deterministic seam
|
||
|
|
required by invariant #6 — NO model, NO network/API call, ever.
|
||
|
|
|
||
|
|
Accuracy goal
|
||
|
|
-------------
|
||
|
|
This is BALLPARK relative ranking of documentation bloat, not billing-grade
|
||
|
|
tokenization. A swappable backend lets a heavier (more accurate) tokenizer be
|
||
|
|
used when available, while a zero-dependency heuristic guarantees the seam
|
||
|
|
always works.
|
||
|
|
|
||
|
|
Backends (pluggable via the `TokenEstimator` ABC)
|
||
|
|
-------------------------------------------------
|
||
|
|
- `HeuristicEstimator` — ALWAYS available. Pure-Python ~4-chars/token rule
|
||
|
|
(`ceil(len(text) / 4)`). Zero third-party dependencies.
|
||
|
|
- `TiktokenEstimator` — OPTIONAL accuracy upgrade. Uses tiktoken's
|
||
|
|
`o200k_base` encoding. Only selected when tiktoken is importable AND its
|
||
|
|
vocab is vendored in the local cache (no cold-cache network download).
|
||
|
|
|
||
|
|
`default_estimator()` performs the selection and NEVER raises — it always
|
||
|
|
returns a working estimator, falling back to the heuristic when tiktoken or
|
||
|
|
its vendored vocab is unavailable. The active backend is introspectable via
|
||
|
|
`.name` so the `check` skill can record which tokenizer produced the counts.
|
||
|
|
|
||
|
|
Offline / determinism guarantee (invariant #6)
|
||
|
|
----------------------------------------------
|
||
|
|
tiktoken has no "offline only" flag: calling `get_encoding("o200k_base")`
|
||
|
|
with a cold cache silently fetches the vocab over the network. We therefore
|
||
|
|
DETECT BEFORE FETCH — we set `TIKTOKEN_CACHE_DIR` to a path inside the plugin
|
||
|
|
and check that the expected cache file already exists *before* importing the
|
||
|
|
encoding. If it is absent we fall back to the heuristic rather than triggering
|
||
|
|
a download.
|
||
|
|
|
||
|
|
CACHE-FILENAME GOTCHA: tiktoken names its cache file by
|
||
|
|
`sha1(blobpath_url).hexdigest()` — NOT by the sha256 content hash and NOT by a
|
||
|
|
human-readable name. For `o200k_base` the URL is
|
||
|
|
`https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken`,
|
||
|
|
which hashes to the filename:
|
||
|
|
|
||
|
|
fb374d419588a4632f3f557e76b4b70aebbca790
|
||
|
|
|
||
|
|
A vendored vocab MUST live at `<cache_dir>/fb374d419588a4632f3f557e76b4b70aebbca790`
|
||
|
|
or the cache never hits and tiktoken falls back to a network fetch.
|
||
|
|
|
||
|
|
Pre-warming the vendored vocab (one-time, online, optional)
|
||
|
|
-----------------------------------------------------------
|
||
|
|
TIKTOKEN_CACHE_DIR=scripts/tiktoken_cache \
|
||
|
|
python -c "import tiktoken; tiktoken.get_encoding('o200k_base')"
|
||
|
|
|
||
|
|
The ~2 MB vocab blob is intentionally NOT committed to git. When it is absent
|
||
|
|
the estimator runs on the heuristic backend; this keeps the repo lean and the
|
||
|
|
runtime offline-safe.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import hashlib
|
||
|
|
import json
|
||
|
|
import math
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
from abc import ABC, abstractmethod
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Optional
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# Constants — tiktoken vendored-cache resolution
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
|
||
|
|
# The encoding we standardise on (current OpenAI tokenizer; close enough for a
|
||
|
|
# relative ballpark of Claude-context bloat — accuracy goal is ranking, not
|
||
|
|
# billing).
|
||
|
|
_O200K_BASE_NAME = "o200k_base"
|
||
|
|
_O200K_BASE_BLOBPATH = (
|
||
|
|
"https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken"
|
||
|
|
)
|
||
|
|
|
||
|
|
# tiktoken's on-disk cache file is named sha1(blobpath_url).hexdigest().
|
||
|
|
# Pre-computed here so the offline pre-check needs no tiktoken import.
|
||
|
|
_O200K_BASE_CACHE_FILENAME = hashlib.sha1(
|
||
|
|
_O200K_BASE_BLOBPATH.encode()
|
||
|
|
).hexdigest() # == "fb374d419588a4632f3f557e76b4b70aebbca790"
|
||
|
|
|
||
|
|
# Default vendored-cache directory inside the plugin (sibling of this file).
|
||
|
|
_DEFAULT_CACHE_DIR = Path(__file__).resolve().parent / "tiktoken_cache"
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# Estimator abstraction
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
|
||
|
|
class TokenEstimator(ABC):
|
||
|
|
"""Pluggable, deterministic token-count backend.
|
||
|
|
|
||
|
|
Subclasses implement `estimate(text)`; the file convenience and report
|
||
|
|
wrapper are shared. Every backend exposes `.name` so callers can record
|
||
|
|
which tokenizer produced a count.
|
||
|
|
"""
|
||
|
|
|
||
|
|
#: Short, stable backend identifier (e.g. "heuristic", "tiktoken:o200k_base").
|
||
|
|
name: str = "abstract"
|
||
|
|
|
||
|
|
@abstractmethod
|
||
|
|
def estimate(self, text: str) -> int:
|
||
|
|
"""Return the token count for *text* (>= 0). Deterministic, no I/O."""
|
||
|
|
raise NotImplementedError
|
||
|
|
|
||
|
|
def estimate_file(self, path: Path) -> int:
|
||
|
|
"""Return the token count of *path*'s UTF-8 contents.
|
||
|
|
|
||
|
|
Decode errors are tolerated (``errors="replace"``) to match the
|
||
|
|
scanner's read convention, so arbitrary/binary-ish docs never raise.
|
||
|
|
A missing/unreadable file yields 0.
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
text = Path(path).read_text(encoding="utf-8", errors="replace")
|
||
|
|
except OSError:
|
||
|
|
return 0
|
||
|
|
return self.estimate(text)
|
||
|
|
|
||
|
|
def estimate_for_report(self, text: str) -> dict:
|
||
|
|
"""Return a schema-shaped ``token_estimate`` object for *text*.
|
||
|
|
|
||
|
|
v1 populates only `raw_tokens` (the required field). The
|
||
|
|
injection-frequency weighting fields are the v2 bonus and are emitted
|
||
|
|
as `null` here so the shape is explicit for the `check` skill.
|
||
|
|
See `openspec/specs/report-schema/spec.md` (Per-Entry Token Estimate).
|
||
|
|
"""
|
||
|
|
return {
|
||
|
|
"raw_tokens": self.estimate(text),
|
||
|
|
"injection_frequency": None,
|
||
|
|
"weighted_tokens": None,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
class HeuristicEstimator(TokenEstimator):
|
||
|
|
"""Zero-dependency ~4-chars/token estimator: ``ceil(len(text) / 4)``.
|
||
|
|
|
||
|
|
Boundaries: ``""`` -> 0, 4 chars -> 1, 5 chars -> 2. Always available;
|
||
|
|
this is the guaranteed fallback that keeps the deterministic seam working
|
||
|
|
with no third-party packages installed.
|
||
|
|
"""
|
||
|
|
|
||
|
|
name = "heuristic"
|
||
|
|
|
||
|
|
#: Average characters per token for the heuristic. ~4 is the long-standing
|
||
|
|
#: rule-of-thumb for English prose / markdown.
|
||
|
|
CHARS_PER_TOKEN = 4
|
||
|
|
|
||
|
|
def estimate(self, text: str) -> int:
|
||
|
|
if not text:
|
||
|
|
return 0
|
||
|
|
return math.ceil(len(text) / self.CHARS_PER_TOKEN)
|
||
|
|
|
||
|
|
|
||
|
|
class TiktokenEstimator(TokenEstimator):
|
||
|
|
"""Accuracy-upgrade backend wrapping a tiktoken encoding.
|
||
|
|
|
||
|
|
Constructed with an already-loaded encoding object so this class performs
|
||
|
|
no import or cache resolution itself (that lives in `default_estimator`,
|
||
|
|
which guarantees the encoding came from the vendored cache, not a network
|
||
|
|
fetch). `disallowed_special=()` ensures arbitrary markdown containing
|
||
|
|
special-token-looking substrings (e.g. ``<|endoftext|>``) never raises.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, encoding, encoding_name: str = _O200K_BASE_NAME) -> None:
|
||
|
|
self._encoding = encoding
|
||
|
|
self.name = f"tiktoken:{encoding_name}"
|
||
|
|
|
||
|
|
def estimate(self, text: str) -> int:
|
||
|
|
if not text:
|
||
|
|
return 0
|
||
|
|
return len(self._encoding.encode(text, disallowed_special=()))
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# Backend selection (never raises)
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
|
||
|
|
def _vendored_vocab_present(cache_dir: Path) -> bool:
|
||
|
|
"""Return True iff the o200k_base vocab is already in *cache_dir*.
|
||
|
|
|
||
|
|
Pure filesystem check by the sha1-of-blobpath filename — see the
|
||
|
|
CACHE-FILENAME GOTCHA in the module docstring. This is the detect-before-
|
||
|
|
fetch gate that keeps the seam offline (invariant #6).
|
||
|
|
"""
|
||
|
|
return (cache_dir / _O200K_BASE_CACHE_FILENAME).is_file()
|
||
|
|
|
||
|
|
|
||
|
|
def default_estimator(cache_dir: Optional[Path] = None) -> TokenEstimator:
|
||
|
|
"""Return the best available working estimator. NEVER raises.
|
||
|
|
|
||
|
|
Selection order:
|
||
|
|
1. `TiktokenEstimator` — only if tiktoken is importable AND the
|
||
|
|
o200k_base vocab is vendored in *cache_dir* (so loading it triggers
|
||
|
|
no network fetch).
|
||
|
|
2. `HeuristicEstimator` — the always-available fallback otherwise.
|
||
|
|
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
cache_dir:
|
||
|
|
Directory to resolve the vendored tiktoken vocab from. Defaults to
|
||
|
|
the plugin's `scripts/tiktoken_cache/`. Injectable so tests can point
|
||
|
|
at an empty dir to force (and assert) the heuristic fallback without
|
||
|
|
any network access.
|
||
|
|
"""
|
||
|
|
resolved_cache = Path(cache_dir) if cache_dir is not None else _DEFAULT_CACHE_DIR
|
||
|
|
|
||
|
|
# Detect-before-fetch: bail to heuristic unless the vocab is already local.
|
||
|
|
if not _vendored_vocab_present(resolved_cache):
|
||
|
|
return HeuristicEstimator()
|
||
|
|
|
||
|
|
try:
|
||
|
|
# Point tiktoken at the vendored cache BEFORE importing/using it.
|
||
|
|
os.environ["TIKTOKEN_CACHE_DIR"] = str(resolved_cache)
|
||
|
|
import tiktoken # noqa: WPS433 (optional dependency, imported lazily)
|
||
|
|
|
||
|
|
encoding = tiktoken.get_encoding(_O200K_BASE_NAME)
|
||
|
|
return TiktokenEstimator(encoding, _O200K_BASE_NAME)
|
||
|
|
except Exception:
|
||
|
|
# Import error, corrupt cache, or any other failure — fall back.
|
||
|
|
# The seam must always yield a working estimator (invariant #6).
|
||
|
|
return HeuristicEstimator()
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
# CLI entry point
|
||
|
|
# ---------------------------------------------------------------------------
|
||
|
|
|
||
|
|
def main(argv: Optional[list] = None) -> int:
|
||
|
|
import argparse
|
||
|
|
|
||
|
|
parser = argparse.ArgumentParser(
|
||
|
|
description=(
|
||
|
|
"doc-hygiene token estimator — emits a JSON token estimate for a "
|
||
|
|
"file (deterministic, no model, no network)."
|
||
|
|
)
|
||
|
|
)
|
||
|
|
parser.add_argument("path", help="Path to the file to estimate.")
|
||
|
|
parser.add_argument(
|
||
|
|
"--cache-dir",
|
||
|
|
default=None,
|
||
|
|
help="Vendored tiktoken cache dir (default: scripts/tiktoken_cache/).",
|
||
|
|
)
|
||
|
|
args = parser.parse_args(argv)
|
||
|
|
|
||
|
|
target = Path(args.path)
|
||
|
|
if not target.is_file():
|
||
|
|
print(
|
||
|
|
json.dumps({"error": f"not a file: {args.path}"}, indent=2),
|
||
|
|
file=sys.stderr,
|
||
|
|
)
|
||
|
|
return 1
|
||
|
|
|
||
|
|
estimator = default_estimator(
|
||
|
|
cache_dir=Path(args.cache_dir) if args.cache_dir else None
|
||
|
|
)
|
||
|
|
raw_tokens = estimator.estimate_file(target)
|
||
|
|
|
||
|
|
output = {
|
||
|
|
"path": str(target),
|
||
|
|
"backend": estimator.name,
|
||
|
|
"token_estimate": {
|
||
|
|
"raw_tokens": raw_tokens,
|
||
|
|
"injection_frequency": None,
|
||
|
|
"weighted_tokens": None,
|
||
|
|
},
|
||
|
|
}
|
||
|
|
print(json.dumps(output, indent=2))
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
sys.exit(main())
|