tokenmeter-sdk 0.2.0__py3-none-any.whl

tokenmeter/__init__.py

@@ -0,0 +1,49 @@
+"""TokenMeter Python SDK.
+
+    import tokenmeter
+    tokenmeter.init(ingest_key="tl_live_...")
+    client = tokenmeter.wrap(OpenAI())
+
+    with tokenmeter.action("generate-report", customer_id="acme", feature="reports"):
+        client.chat.completions.create(...)
+"""
+from ._client import init
+from ._context import TraceContext, action, attach, current
+from ._diagnose import diagnose
+from ._preflight import GovernanceDenied, GovernanceUnavailable, PreflightDecision, governed, preflight
+from ._wrap import wrap, wrap_openai_compatible
+from ._read_api import ReadApiClient, ReadApiError, TokenMeterReadClient
+from ._topic import keyword_topic_classifier, latest_user_text
+
+__all__ = [
+    "init",
+    "wrap",
+    "wrap_openai_compatible",
+    "action",
+    "current",
+    "attach",
+    "TraceContext",
+    "diagnose",
+    "flush",
+    "preflight",
+    "governed",
+    "PreflightDecision",
+    "GovernanceDenied",
+    "GovernanceUnavailable",
+    "ReadApiClient",
+    "TokenMeterReadClient",
+    "ReadApiError",
+    "keyword_topic_classifier",
+    "latest_user_text",
+]
+
+__version__ = "0.2.0"
+
+
+def flush(deadline_s: float = 2.0) -> None:
+    """Best-effort synchronous flush. Never raises."""
+    from ._client import get_state
+
+    state = get_state()
+    if state is not None:
+        state.flusher.flush_best_effort(deadline_s=deadline_s)

tokenmeter/_buffer.py

@@ -0,0 +1,201 @@
+"""Ring buffer + background flush thread (PRD §7.1).
+
+Flush every 2s or 500 events, whichever first. Buffer full -> drop OLDEST,
+count the drop, ship the count in the next batch (never silent).
+"""
+from __future__ import annotations
+
+import atexit
+import collections
+import logging
+import threading
+import time
+from typing import Callable, Deque, List, Optional
+
+logger = logging.getLogger("tokenmeter")
+
+DEFAULT_MAXSIZE = 10_000
+DEFAULT_FLUSH_INTERVAL_S = 2.0
+DEFAULT_FLUSH_BATCH_SIZE = 500
+# DECISIONS.md §8: sustained rate limiting must be LOUD, not a silent spin.
+# 5 consecutive 429s at the default 2s flush interval is ~10s of sustained
+# rejection -- past one unlucky fixed-window boundary straddle, short of
+# letting the ~10k-event ring buffer fill silently.
+CONSECUTIVE_429_WARN_THRESHOLD = 5
+
+
+class RingBuffer:
+    """Thread-safe, fixed-capacity, drop-oldest-on-overflow queue.
+
+    `collections.deque(maxlen=N)` already evicts from the opposite end on
+    overflow (stdlib does it, rung 3) -- this class only adds the lock and the
+    dropped-event counter the eviction doesn't tell us about on its own.
+    """
+
+    def __init__(self, maxsize: int = DEFAULT_MAXSIZE) -> None:
+        self._dq: Deque[dict] = collections.deque(maxlen=maxsize)
+        self._lock = threading.Lock()
+        self._dropped = 0
+
+    def append(self, event: dict) -> None:
+        with self._lock:
+            if len(self._dq) >= (self._dq.maxlen or 0):
+                self._dropped += 1  # about to evict the oldest via maxlen
+            self._dq.append(event)
+
+    def __len__(self) -> int:
+        with self._lock:
+            return len(self._dq)
+
+    @property
+    def capacity(self) -> int:
+        return self._dq.maxlen or 0
+
+    def drain(self, max_items: Optional[int] = None) -> List[dict]:
+        """Pop up to max_items (or all) events. Returns [] if empty."""
+        with self._lock:
+            n = len(self._dq) if max_items is None else min(max_items, len(self._dq))
+            return [self._dq.popleft() for _ in range(n)]
+
+    def take_dropped(self) -> int:
+        """Reset and return the dropped-event count since the last call."""
+        with self._lock:
+            d, self._dropped = self._dropped, 0
+            return d
+
+    @property
+    def dropped(self) -> int:
+        """Peek the pending drop count without resetting it."""
+        with self._lock:
+            return self._dropped
+
+    def release_dropped(self, n: int) -> None:
+        """Subtract n from the pending drop count once it has actually been
+        ACKNOWLEDGED by the server, rather than zeroing it the moment it's
+        read (take_dropped()'s behavior, which is what Flusher used to call
+        unconditionally BEFORE knowing whether the send would even succeed --
+        a failed/retried send silently threw the count away, exactly wrong
+        during a sustained-429 run where every attempt discarded it again).
+        """
+        with self._lock:
+            self._dropped = max(0, self._dropped - n)
+
+    def requeue(self, events: List[dict]) -> None:
+        """Put a drained-but-unsent batch back, oldest-first, for a failed
+        send -- mirrors the TypeScript SDK's unshiftMany(). Without this,
+        drain() had already removed the events from the deque, and a raised
+        send() (network error, non-2xx) meant _flush_once()'s outer
+        try/except simply swallowed them: a real data-loss bug, not just a
+        counter bug. If the buffer filled back up while this batch was in
+        flight, the newest of what's being returned gives way (drop-oldest
+        semantics still hold; what's already in the buffer is newer)."""
+        with self._lock:
+            maxlen = self._dq.maxlen or 0
+            for ev in reversed(events):
+                if len(self._dq) >= maxlen:
+                    self._dropped += 1
+                    continue
+                self._dq.appendleft(ev)
+
+
+class Flusher:
+    """Background thread: flush when len(buffer) >= batch_size OR every interval_s."""
+
+    def __init__(
+        self,
+        buffer: RingBuffer,
+        send: Callable[[List[dict], int], None],
+        *,
+        interval_s: float = DEFAULT_FLUSH_INTERVAL_S,
+        batch_size: int = DEFAULT_FLUSH_BATCH_SIZE,
+    ) -> None:
+        self._buffer = buffer
+        self._send = send
+        self._interval_s = interval_s
+        self._batch_size = batch_size
+        self._wake = threading.Event()
+        self._stop = threading.Event()
+        self._thread = threading.Thread(target=self._run, name="tokenmeter-flusher", daemon=True)
+        self._atexit_registered = False
+        self._consecutive_429s = 0
+        self._warned_sustained_429 = False
+
+    def start(self) -> None:
+        self._thread.start()
+        if not self._atexit_registered:
+            atexit.register(self.flush_best_effort, deadline_s=2.0)
+            self._atexit_registered = True
+
+    def notify(self) -> None:
+        """Called on every append; wakes the flusher early once batch_size is hit."""
+        if len(self._buffer) >= self._batch_size:
+            self._wake.set()
+
+    def _run(self) -> None:
+        while not self._stop.is_set():
+            self._wake.wait(timeout=self._interval_s)
+            self._wake.clear()
+            self._flush_once()
+
+    def _flush_once(self) -> None:
+        events = self._buffer.drain()
+        if not events:
+            return
+        # Peek, don't consume: the count is only "communicated" once the
+        # server actually acknowledges it. take_dropped() used to be called
+        # unconditionally here, BEFORE knowing whether send() would even
+        # succeed -- a failed send meant the drop count it just zeroed was
+        # gone, silently under-reporting exactly the signal DECISIONS.md §8
+        # relies on to make loss visible (worst during a sustained-429 run,
+        # where every retry attempt reset it again).
+        dropped = self._buffer.dropped
+        try:
+            self._send(events, dropped)
+        except Exception as exc:  # SDK NEVER throws into host code -- swallow, log, move on.
+            # drain() already removed these from the buffer; without putting
+            # them back, a failed send silently discarded the whole batch --
+            # not just under-counted drops, an actual data-loss bug.
+            self._buffer.requeue(events)
+            self._note_failure(exc)
+            logger.debug("tokenmeter: flush failed", exc_info=True)
+            return
+        self._buffer.release_dropped(dropped)
+        self._consecutive_429s = 0
+        self._warned_sustained_429 = False
+
+    def _note_failure(self, exc: Exception) -> None:
+        """DECISIONS.md §8: sustained 429s must be loud, not a silent retry
+        spin. Mirrors diagnose()'s "print + log, never raise" approach --
+        one warning per incident (not one per attempt), reset once traffic
+        stops being rate-limited."""
+        code = getattr(exc, "code", None)
+        if code == 429:
+            self._consecutive_429s += 1
+        else:
+            self._consecutive_429s = 0
+            self._warned_sustained_429 = False
+            return
+        if self._consecutive_429s >= CONSECUTIVE_429_WARN_THRESHOLD and not self._warned_sustained_429:
+            self._warned_sustained_429 = True
+            msg = (
+                f"tokenmeter: {self._consecutive_429s} consecutive batches rate-limited (HTTP 429). "
+                "Events are being retried, not dropped by the server, but this SDK's local buffer "
+                f"(maxsize={self._buffer.capacity}) can still overflow under sustained pressure -- "
+                "watch the dropped-event count. If this persists, the ingest key's plan quota is "
+                "genuinely too low for this traffic. Run tokenmeter.diagnose() for more detail."
+            )
+            logger.warning(msg)
+            print("[tokenmeter] WARNING: " + msg)
+
+    def flush_best_effort(self, deadline_s: float = 2.0) -> None:
+        """Synchronous, best-effort, capped flush -- used by atexit and public flush()."""
+        deadline = time.monotonic() + deadline_s
+        try:
+            while len(self._buffer) > 0 and time.monotonic() < deadline:
+                self._flush_once()
+        except Exception:
+            logger.debug("tokenmeter: atexit flush failed", exc_info=True)
+
+    def stop(self) -> None:
+        self._stop.set()
+        self._wake.set()

tokenmeter/_client.py

@@ -0,0 +1,78 @@
+"""Global SDK state, set up once by tokenmeter.init()."""
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+from ._buffer import DEFAULT_FLUSH_BATCH_SIZE, DEFAULT_FLUSH_INTERVAL_S, DEFAULT_MAXSIZE, Flusher, RingBuffer
+from ._transport import Transport
+
+logger = logging.getLogger("tokenmeter")
+
+DEFAULT_ENDPOINT = "https://ingest.timewastevc.com"
+
+
+class _State:
+    def __init__(
+        self,
+        ingest_key: str,
+        endpoint: str = DEFAULT_ENDPOINT,
+        *,
+        buffer_size: int = DEFAULT_MAXSIZE,
+        flush_interval_s: float = DEFAULT_FLUSH_INTERVAL_S,
+        flush_batch_size: int = DEFAULT_FLUSH_BATCH_SIZE,
+        optimization: Optional[dict] = None,
+    ) -> None:
+        self.optimization = dict(optimization) if isinstance(optimization, dict) else {}
+        self.transport = Transport(ingest_key, endpoint)
+        self.buffer = RingBuffer(buffer_size)
+        self.flusher = Flusher(
+            self.buffer,
+            self.transport.send,
+            interval_s=flush_interval_s,
+            batch_size=flush_batch_size,
+        )
+        self.flusher.start()
+
+    def emit(self, event: dict) -> None:
+        # SDK NEVER throws into host code and NEVER adds latency to the caller --
+        # this is a memory append, not I/O, so it's already sub-microsecond; the
+        # try/except is only for the pathological case of an unpicklable/huge event.
+        try:
+            self.buffer.append(event)
+            self.flusher.notify()
+        except Exception:
+            logger.debug("tokenmeter: emit failed", exc_info=True)
+
+
+_state: Optional[_State] = None
+
+
+def init(
+    ingest_key: str,
+    endpoint: str = DEFAULT_ENDPOINT,
+    *,
+    router_url: Optional[str] = None,
+    governance_mode: str = "enforce",
+    governance_fail_mode: str = "closed",
+    optimization: Optional[dict] = None,
+    **kwargs,
+) -> None:
+    """router_url enables tokenmeter.preflight()/governed() (governance before direct provider calls)."""
+    global _state
+    from . import _preflight
+
+    _preflight.configure(ingest_key, router_url, mode=governance_mode, fail_mode=governance_fail_mode)
+    _state = _State(ingest_key, endpoint, optimization=optimization, **kwargs)
+
+
+def get_state() -> Optional[_State]:
+    return _state
+
+
+def _reset_for_tests() -> None:
+    """Test-only: drop global state so each test starts clean."""
+    global _state
+    if _state is not None:
+        _state.flusher.stop()
+    _state = None

tokenmeter/_context.py

@@ -0,0 +1,131 @@
+"""Trace/action context propagation (PRD §11).
+
+contextvars.ContextVar survives asyncio task spawns (gather, create_task) and
+thread-pool submission natively as long as the submitted callable captures the
+context (contextvars.copy_context() / ThreadPoolExecutor + contextvars.copy_context
+via `contextvars.copy_context().run(...)`, or Python's own default propagation into
+`concurrent.futures.ThreadPoolExecutor.submit` -- which does NOT auto-propagate, so
+callers crossing a real OS thread boundary must use `attach(current())` explicitly,
+same escape hatch as a queue/serialize hop. This module provides that hatch.
+"""
+from __future__ import annotations
+
+import contextvars
+import dataclasses
+from typing import Optional
+
+from ._ids import uuid7_str
+
+
+@dataclasses.dataclass(frozen=True)
+class TraceContext:
+    trace_id: str
+    span_id: str
+    parent_span_id: Optional[str] = None
+    action: Optional[str] = None
+    customer_ref: Optional[str] = None
+    feature: Optional[str] = None
+    team: Optional[str] = None
+    environment: Optional[str] = None
+    project: Optional[str] = None
+    member_ref: Optional[str] = None
+    optimization: Optional[dict] = None
+
+    def to_dict(self) -> dict:
+        d = dataclasses.asdict(self)
+        # A function (e.g. optimization["topic_classifier"]) cannot cross a queue/serialize hop; the receiving side
+        # falls back to the classifier configured in its own tokenmeter.init().
+        if isinstance(d.get("optimization"), dict):
+            d["optimization"] = {k: v for k, v in d["optimization"].items() if not callable(v)}
+        return d
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "TraceContext":
+        return cls(**d)
+
+
+_current: "contextvars.ContextVar[Optional[TraceContext]]" = contextvars.ContextVar(
+    "tokenmeter_context", default=None
+)
+
+
+def current() -> Optional[TraceContext]:
+    """The active TraceContext, or None if no action() scope is open.
+
+    Serializable (dataclass of str fields) so it can cross a queue/serialize hop:
+    `payload = tokenmeter.current().to_dict()` ... `tokenmeter.attach(TraceContext.from_dict(payload))`.
+    """
+    return _current.get()
+
+
+def attach(ctx: Optional[TraceContext]) -> contextvars.Token:
+    """Escape hatch for thread pools and queue/serialize hops (PRD §11).
+
+    Returns the reset Token so the caller can restore the prior value if it owns
+    the context's lifetime (rarely needed -- a fresh thread/task has nothing to
+    restore to).
+    """
+    return _current.set(ctx)
+
+
+def _reset(token: contextvars.Token) -> None:
+    _current.reset(token)
+
+
+class action:
+    """Context manager. Nested calls form a span chain (parent_span_id chain)."""
+
+    def __init__(
+        self,
+        name: str,
+        *,
+        customer_id: Optional[str] = None,
+        customer_ref: Optional[str] = None,
+        feature: Optional[str] = None,
+        team: Optional[str] = None,
+        environment: Optional[str] = None,
+        project: Optional[str] = None,
+        member_ref: Optional[str] = None,
+        optimization: Optional[dict] = None,
+    ) -> None:
+        self.name = name
+        # accept either kwarg spelling used across the PRD text/examples
+        self.customer_ref = customer_ref or customer_id
+        self.feature = feature
+        self.team = team
+        self.environment = environment
+        self.project = project
+        self.member_ref = member_ref
+        self.optimization = dict(optimization) if isinstance(optimization, dict) else {}
+        self._token: Optional[contextvars.Token] = None
+
+    def __enter__(self) -> TraceContext:
+        parent = current()
+        ctx = TraceContext(
+            trace_id=parent.trace_id if parent else uuid7_str(),
+            span_id=uuid7_str(),
+            parent_span_id=parent.span_id if parent else None,
+            action=self.name,
+            customer_ref=self.customer_ref if self.customer_ref is not None else (parent.customer_ref if parent else None),
+            feature=self.feature if self.feature is not None else (parent.feature if parent else None),
+            team=self.team if self.team is not None else (parent.team if parent else None),
+            environment=self.environment if self.environment is not None else (parent.environment if parent else None),
+            project=self.project if self.project is not None else (parent.project if parent else None),
+            member_ref=self.member_ref if self.member_ref is not None else (parent.member_ref if parent else None),
+            optimization={**(parent.optimization or {} if parent else {}), **self.optimization},
+        )
+        self._token = attach(ctx)
+        return ctx
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        if self._token is not None:
+            _reset(self._token)
+        return None
+
+    # async support: `async with tokenmeter.action(...):` works identically since
+    # entering/exiting is synchronous context-var manipulation, not I/O.
+    async def __aenter__(self) -> TraceContext:
+        return self.__enter__()
+
+    async def __aexit__(self, exc_type, exc, tb) -> None:
+        self.__exit__(exc_type, exc, tb)

tokenmeter/_diagnose.py

@@ -0,0 +1,57 @@
+"""tokenmeter.diagnose() -- finding M4: a typo'd ingest key must never be silent.
+
+Everywhere else the SDK swallows errors (that's the whole point of "never throws
+into host code"). diagnose() is the one deliberate exception: it is called on
+purpose, synchronously, specifically to make failures LOUD -- it prints a clear
+verdict and returns a structured result, and it is the documented answer to
+"why do I see zero data in my dashboard."
+"""
+from __future__ import annotations
+
+import urllib.error
+
+from ._client import DEFAULT_ENDPOINT, get_state
+from ._transport import Transport
+
+_STATUS_HINTS = {
+    401: "invalid_key: the ingest key is wrong, revoked, or malformed. Check tokenmeter.init(ingest_key=...).",
+    403: "key_revoked: this key (or its org) was revoked/purged. Issue a new key from the dashboard.",
+    413: "batch_too_large: unrelated to diagnose() itself -- your real traffic is sending oversized batches.",
+    429: "rate_limited: traffic looks abuse-shaped (well over the plan quota), not ordinary overage -- ordinary overage is accepted (DECISIONS.md §8). A single 429 is retried and not data loss; SUSTAINED 429s can be, since this SDK's local buffer is bounded -- watch for the sustained-429 warning this SDK logs.",
+    503: "unavailable: TokenMeter ingest is temporarily down. Events buffer locally and retry.",
+}
+
+
+def diagnose(ingest_key: str = None, endpoint: str = None) -> dict:
+    """Synchronously pings the Ingest API and returns
+    {"ok": bool, "status": int|None, "message": str}. Never raises. Prints the
+    verdict too, because the whole point is to be found by a human staring at a
+    silent dashboard, not by someone who already knew to check a return value.
+    """
+    state = get_state()
+    key = ingest_key or (state.transport.ingest_key if state else None)
+    url = endpoint or (state.transport.endpoint.rsplit("/v1/", 1)[0] if state else DEFAULT_ENDPOINT)
+
+    if not key:
+        result = {"ok": False, "status": None, "message": "no ingest key: call tokenmeter.init(ingest_key=...) first, or pass one to diagnose()."}
+        _report(result)
+        return result
+
+    try:
+        pong = Transport(key, url).ping()
+        result = {"ok": True, "status": pong["status"], "message": f"tokenmeter: ingest key valid, endpoint reachable ({url})."}
+    except urllib.error.HTTPError as e:
+        hint = _STATUS_HINTS.get(e.code, f"unexpected HTTP {e.code}")
+        result = {"ok": False, "status": e.code, "message": f"tokenmeter: {hint}"}
+    except urllib.error.URLError as e:
+        result = {"ok": False, "status": None, "message": f"tokenmeter: could not reach {url}: {e.reason}"}
+    except Exception as e:  # still must never raise into host code
+        result = {"ok": False, "status": None, "message": f"tokenmeter: diagnose() failed unexpectedly: {e!r}"}
+
+    _report(result)
+    return result
+
+
+def _report(result: dict) -> None:
+    prefix = "[tokenmeter diagnose] OK: " if result["ok"] else "[tokenmeter diagnose] FAILED: "
+    print(prefix + result["message"])

tokenmeter/_ids.py

@@ -0,0 +1,45 @@
+"""UUIDv7 (RFC 9562) event ids: time-ordered, client-generated, dedup key upstream."""
+from __future__ import annotations
+
+import os
+import time
+import uuid
+
+
+def uuid7() -> uuid.UUID:
+    """48-bit unix ms timestamp + version/variant bits + 74 random bits."""
+    unix_ts_ms = int(time.time() * 1000) & 0xFFFFFFFFFFFF  # 48 bits
+    rand = os.urandom(10)  # 12 bits rand_a + 62 bits rand_b = 74 bits, rounded up to 10 bytes
+
+    rand_a = int.from_bytes(rand[0:2], "big") & 0x0FFF  # 12 bits
+    rand_b = int.from_bytes(rand[2:10], "big") & 0x3FFFFFFFFFFFFFFF  # 62 bits
+
+    b = bytearray(16)
+    b[0:6] = unix_ts_ms.to_bytes(6, "big")
+    b[6] = 0x70 | (rand_a >> 8)  # version 7 in top nibble
+    b[7] = rand_a & 0xFF
+    b[8] = 0x80 | (rand_b >> 56) & 0x3F  # variant bits 10
+    b[9:16] = ((rand_b >> 0) & 0xFFFFFFFFFFFFFF).to_bytes(7, "big")
+    return uuid.UUID(bytes=bytes(b))
+
+
+def uuid7_str() -> str:
+    return str(uuid7())
+
+
+def demo() -> None:
+    ids = [uuid7() for _ in range(1000)]
+    strs = [str(i) for i in ids]
+    assert len(set(strs)) == 1000, "collision in 1000 draws"
+    for i in ids:
+        assert i.version == 7
+    # Timestamp prefix (first 12 hex chars = 48-bit ms) must be non-decreasing across
+    # calls in issuance order -- the property that makes these dedup/sort-friendly.
+    # (Random tail bits are NOT required to sort within the same millisecond.)
+    ts_prefixes = [s.replace("-", "")[:12] for s in strs]
+    assert ts_prefixes == sorted(ts_prefixes), "uuid7 timestamp prefix must be time-ordered"
+    print("uuid7 demo: ok")
+
+
+if __name__ == "__main__":
+    demo()