#!/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 `/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())