tokenhelm 0.1.0rc1__py3-none-any.whl

tokenhelm/__init__.py

@@ -0,0 +1,82 @@
+"""TokenHelm — lightweight, framework-agnostic LLM token tracking and cost calculation.
+
+Public surface per ``specs/001-core-sdk/contracts/public-api.md``. Names exported here are the
+v0.x stable contract (Principle X). This slice (User Story 1) ships the OpenAI vertical;
+additional providers, loggers, storage, and streaming land in later phases.
+"""
+
+from __future__ import annotations
+
+# Errors
+# Extension-point interfaces (Principle VI)
+from .adapters.anthropic import AnthropicAdapter
+from .adapters.base import BaseAdapter, StreamAggregator
+from .adapters.gemini import GeminiAdapter
+from .adapters.ollama import OllamaAdapter
+from .adapters.openai import OpenAIAdapter
+from .core.calculator import CostCalculator
+from .core.errors import ProviderNotInstalledError, TokenHelmError
+
+# Core data model
+from .core.models import (
+    LLMCost,
+    LLMEvent,
+    LLMProvider,
+    LLMRequest,
+    LLMUsage,
+    RateEntry,
+)
+from .dispatch.base import EventDispatcher
+from .dispatch.default import DefaultEventDispatcher
+from .logging.base import Logger
+from .logging.console import ConsoleLogger
+from .logging.file import FileLogger
+from .logging.json import JSONLogger
+from .pricing.base import PricingProvider
+from .pricing.yaml_provider import YamlPricingProvider
+
+# Client + scope
+from .sdk.client import TokenHelm
+from .sdk.context import StreamSession, TraceScope
+from .storage.base import StorageBackend
+from .storage.memory import InMemoryStorageBackend
+
+__version__ = "0.1.0rc1"  # x-release-please-version
+
+__all__ = [
+    "__version__",
+    # client
+    "TokenHelm",
+    "TraceScope",
+    "StreamSession",
+    # event + enum
+    "LLMEvent",
+    "LLMProvider",
+    "LLMUsage",
+    "LLMCost",
+    "LLMRequest",
+    "RateEntry",
+    # interfaces
+    "BaseAdapter",
+    "StreamAggregator",
+    "PricingProvider",
+    "EventDispatcher",
+    "Logger",
+    "StorageBackend",
+    # built-in adapters
+    "OpenAIAdapter",
+    "AnthropicAdapter",
+    "GeminiAdapter",
+    "OllamaAdapter",
+    # default implementations
+    "YamlPricingProvider",
+    "DefaultEventDispatcher",
+    "ConsoleLogger",
+    "JSONLogger",
+    "FileLogger",
+    "InMemoryStorageBackend",
+    "CostCalculator",
+    # errors
+    "TokenHelmError",
+    "ProviderNotInstalledError",
+]

tokenhelm/adapters/__init__.py

@@ -0,0 +1,20 @@
+"""Provider adapters (extension point #1).
+
+``default_adapters()`` is the built-in adapter set the client registers by default. It lives
+here (not in ``core``) so the core stays decoupled from concrete adapter implementations —
+``core.extraction.UsageParser`` depends only on :class:`BaseAdapter`.
+"""
+
+from __future__ import annotations
+
+from .base import BaseAdapter
+
+
+def default_adapters() -> list[BaseAdapter]:
+    """Built-in adapters, in resolution order (US1 ships OpenAI; US2 adds the rest)."""
+    from .anthropic import AnthropicAdapter
+    from .gemini import GeminiAdapter
+    from .ollama import OllamaAdapter
+    from .openai import OpenAIAdapter
+
+    return [OpenAIAdapter(), AnthropicAdapter(), GeminiAdapter(), OllamaAdapter()]

tokenhelm/adapters/anthropic.py

@@ -0,0 +1,113 @@
+"""AnthropicAdapter — normalize Anthropic Messages responses (T034).
+
+Observe-don't-patch: reads attributes off the ``Message`` object the developer's own
+``anthropic`` client returns. Anthropic reports usage under ``usage.input_tokens`` /
+``output_tokens`` and does **not** supply a combined total (LLMUsage derives it). Prompt-cache
+counts (``cache_creation_input_tokens`` / ``cache_read_input_tokens``) are preserved in
+``LLMUsage.extra`` rather than folded into the core schema.
+"""
+
+from __future__ import annotations
+
+from ..core.models import LLMProvider, LLMUsage
+from .base import BaseAdapter, StreamAggregator
+
+_CACHE_FIELDS = ("cache_creation_input_tokens", "cache_read_input_tokens")
+
+_STREAM_EVENT_TYPES = frozenset(
+    {
+        "message_start",
+        "content_block_start",
+        "content_block_delta",
+        "content_block_stop",
+        "message_delta",
+        "message_stop",
+        "ping",
+    }
+)
+
+
+class AnthropicAdapter(BaseAdapter):
+    @property
+    def provider(self) -> LLMProvider:
+        return LLMProvider.ANTHROPIC
+
+    def identify(self, response: object) -> bool:
+        # Anthropic Message objects carry type == "message"; their usage uses input_tokens.
+        if getattr(response, "type", None) == "message":
+            return True
+        usage = getattr(response, "usage", None)
+        return (
+            usage is not None
+            and hasattr(usage, "input_tokens")
+            and hasattr(response, "stop_reason")
+        )
+
+    def extract_model(self, response: object) -> str:
+        return str(getattr(response, "model", "") or "")
+
+    def extract_usage(self, response: object) -> LLMUsage:
+        usage = getattr(response, "usage", None)
+        if usage is None:
+            return LLMUsage()
+        input_tokens = getattr(usage, "input_tokens", None)
+        output_tokens = getattr(usage, "output_tokens", None)
+        # Anthropic does not report a combined total; LLMUsage derives input + output.
+        return LLMUsage(
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            extra=_cache_extra(usage),
+        )
+
+    def identify_stream(self, chunk: object) -> bool:
+        return getattr(chunk, "type", None) in _STREAM_EVENT_TYPES
+
+    def new_stream_aggregator(self) -> StreamAggregator:
+        return _AnthropicStreamAggregator()
+
+
+def _cache_extra(usage: object) -> dict[str, int]:
+    extra: dict[str, int] = {}
+    for name in _CACHE_FIELDS:
+        value = getattr(usage, name, None)
+        if value is not None:
+            extra[name] = value
+    return extra
+
+
+class _AnthropicStreamAggregator(StreamAggregator):
+    """Anthropic streaming: input + model from ``message_start``; output from ``message_delta``.
+
+    ``message_start`` carries ``message.usage.input_tokens`` (and the model); subsequent
+    ``message_delta`` events carry the running ``usage.output_tokens``.
+    """
+
+    def __init__(self) -> None:
+        self._model = ""
+        self._input: int | None = None
+        self._output: int | None = None
+        self._extra: dict[str, int] = {}
+
+    def consume(self, chunk: object) -> None:
+        ctype = getattr(chunk, "type", None)
+        if ctype == "message_start":
+            message = getattr(chunk, "message", None)
+            if message is not None:
+                self._model = str(getattr(message, "model", "") or "")
+                usage = getattr(message, "usage", None)
+                if usage is not None:
+                    self._input = getattr(usage, "input_tokens", None)
+                    self._output = getattr(usage, "output_tokens", None)
+                    self._extra = _cache_extra(usage)
+        elif ctype == "message_delta":
+            usage = getattr(chunk, "usage", None)
+            if usage is not None:
+                output = getattr(usage, "output_tokens", None)
+                if output is not None:
+                    self._output = output
+
+    def model(self) -> str:
+        return self._model
+
+    def usage(self) -> LLMUsage:
+        return LLMUsage(input_tokens=self._input, output_tokens=self._output, extra=self._extra)

tokenhelm/adapters/base.py

@@ -0,0 +1,73 @@
+"""Extension point #1 — provider adapter contract (T006, T050).
+
+An adapter *observes* a response object the developer's own client returned; it never patches
+or wraps the provider SDK. Adapters identify a response by its shape (duck typing) so the core
+needs no provider SDK installed to normalize a response.
+
+Streaming (US4): each adapter owns its provider-specific chunk-aggregation logic behind a
+:class:`StreamAggregator`. The core tracker never understands streaming payloads — it only asks
+the aggregator for a final :class:`LLMUsage` and model. ``identify_stream`` recognizes a
+streaming *chunk* (whose shape may differ from a full response).
+"""
+
+from __future__ import annotations
+
+import abc
+
+from ..core.models import LLMProvider, LLMUsage
+
+
+class StreamAggregator(abc.ABC):
+    """Accumulates provider-specific stream chunks into a normalized usage/model.
+
+    One aggregator instance handles one stream. ``consume`` is called once per chunk (in
+    order); ``model`` and ``usage`` are read once at finalization. Implementations must
+    tolerate partial streams — returning whatever usage was seen so far, with missing counts
+    left as ``None`` (so the event flags ``usage_complete=False``).
+    """
+
+    @abc.abstractmethod
+    def consume(self, chunk: object) -> None:
+        """Accumulate one stream chunk. Must not raise on a chunk lacking usage."""
+
+    @abc.abstractmethod
+    def model(self) -> str:
+        """Best-known model id from the chunks seen so far."""
+
+    @abc.abstractmethod
+    def usage(self) -> LLMUsage:
+        """The aggregated usage so far (final once the stream is exhausted)."""
+
+
+class BaseAdapter(abc.ABC):
+    """Contract every provider adapter satisfies."""
+
+    @property
+    @abc.abstractmethod
+    def provider(self) -> LLMProvider:
+        """Which provider this adapter handles."""
+
+    @abc.abstractmethod
+    def identify(self, response: object) -> bool:
+        """True if this adapter can parse a completed (non-streaming) ``response``."""
+
+    @abc.abstractmethod
+    def extract_usage(self, response: object) -> LLMUsage:
+        """Normalize token usage into an :class:`LLMUsage`."""
+
+    @abc.abstractmethod
+    def extract_model(self, response: object) -> str:
+        """Read the model identifier reported by the response."""
+
+    # -- streaming (optional per adapter) --------------------------------------------
+
+    def identify_stream(self, chunk: object) -> bool:
+        """True if this adapter recognizes a streaming *chunk*. Defaults to ``identify``.
+
+        Override when a provider's streaming chunk shape differs from its full response.
+        """
+        return self.identify(chunk)
+
+    def new_stream_aggregator(self) -> StreamAggregator:
+        """Return a fresh :class:`StreamAggregator` for one stream. v0.1 default: unsupported."""
+        raise NotImplementedError(f"{type(self).__name__} does not support streaming aggregation.")

tokenhelm/adapters/gemini.py

@@ -0,0 +1,92 @@
+"""GeminiAdapter — normalize Google Gemini responses (T033).
+
+Observe-don't-patch: reads attributes off the ``GenerateContentResponse`` object the
+developer's own ``google-genai`` client returns. Gemini reports usage under
+``usage_metadata`` (``prompt_token_count`` / ``candidates_token_count`` / ``total_token_count``)
+and the resolved model under ``model_version``.
+"""
+
+from __future__ import annotations
+
+from ..core.models import LLMProvider, LLMUsage
+from .base import BaseAdapter, StreamAggregator
+
+_GEMINI_EXTRA_FIELDS = (
+    "cached_content_token_count",
+    "thoughts_token_count",
+    "tool_use_prompt_token_count",
+)
+
+
+def _model_of(response: object) -> str:
+    return str(getattr(response, "model_version", None) or getattr(response, "model", "") or "")
+
+
+def _usage_from_metadata(meta: object) -> LLMUsage:
+    extra: dict[str, int] = {}
+    # Gemini may report cached / tool / thinking tokens; keep them out of the core schema.
+    for name in _GEMINI_EXTRA_FIELDS:
+        value = getattr(meta, name, None)
+        if value is not None:
+            extra[name] = value
+    return LLMUsage(
+        input_tokens=getattr(meta, "prompt_token_count", None),
+        output_tokens=getattr(meta, "candidates_token_count", None),
+        total_tokens=getattr(meta, "total_token_count", None),
+        extra=extra,
+    )
+
+
+class GeminiAdapter(BaseAdapter):
+    @property
+    def provider(self) -> LLMProvider:
+        return LLMProvider.GEMINI
+
+    def identify(self, response: object) -> bool:
+        # The ``usage_metadata`` container with Gemini's token-count fields is distinctive.
+        meta = getattr(response, "usage_metadata", None)
+        return meta is not None and (
+            hasattr(meta, "prompt_token_count") or hasattr(meta, "candidates_token_count")
+        )
+
+    def extract_model(self, response: object) -> str:
+        return _model_of(response)
+
+    def extract_usage(self, response: object) -> LLMUsage:
+        meta = getattr(response, "usage_metadata", None)
+        if meta is None:
+            return LLMUsage()
+        return _usage_from_metadata(meta)
+
+    def identify_stream(self, chunk: object) -> bool:
+        # Streaming chunks are GenerateContentResponse objects; earliest chunks may carry
+        # candidates before usage_metadata is populated.
+        return self.identify(chunk) or hasattr(chunk, "candidates")
+
+    def new_stream_aggregator(self) -> StreamAggregator:
+        return _GeminiStreamAggregator()
+
+
+class _GeminiStreamAggregator(StreamAggregator):
+    """Gemini streaming: keep the latest ``usage_metadata`` / ``model_version`` seen.
+
+    Gemini reports cumulative usage on each chunk, with the final chunk holding the totals.
+    """
+
+    def __init__(self) -> None:
+        self._model = ""
+        self._usage = LLMUsage()
+
+    def consume(self, chunk: object) -> None:
+        model = _model_of(chunk)
+        if model:
+            self._model = model
+        meta = getattr(chunk, "usage_metadata", None)
+        if meta is not None:
+            self._usage = _usage_from_metadata(meta)
+
+    def model(self) -> str:
+        return self._model
+
+    def usage(self) -> LLMUsage:
+        return self._usage

tokenhelm/adapters/ollama.py

@@ -0,0 +1,69 @@
+"""OllamaAdapter — normalize Ollama (local) responses (T035).
+
+Observe-don't-patch: reads attributes off the response object the developer's own ``ollama``
+client returns. Ollama reports usage as ``prompt_eval_count`` (input) and ``eval_count``
+(output); there is no combined total (LLMUsage derives it). Local inference is typically
+zero-cost — unlisted models simply resolve as "unpriced" (cost 0) without error.
+"""
+
+from __future__ import annotations
+
+from ..core.models import LLMProvider, LLMUsage
+from .base import BaseAdapter, StreamAggregator
+
+
+def _usage_of(response: object) -> LLMUsage:
+    # Ollama provides no combined total; LLMUsage derives input + output.
+    return LLMUsage(
+        input_tokens=getattr(response, "prompt_eval_count", None),
+        output_tokens=getattr(response, "eval_count", None),
+    )
+
+
+class OllamaAdapter(BaseAdapter):
+    @property
+    def provider(self) -> LLMProvider:
+        return LLMProvider.OLLAMA
+
+    def identify(self, response: object) -> bool:
+        # eval_count / prompt_eval_count are distinctive to Ollama responses.
+        return hasattr(response, "eval_count") or hasattr(response, "prompt_eval_count")
+
+    def extract_model(self, response: object) -> str:
+        return str(getattr(response, "model", "") or "")
+
+    def extract_usage(self, response: object) -> LLMUsage:
+        return _usage_of(response)
+
+    def identify_stream(self, chunk: object) -> bool:
+        # Intermediate stream chunks lack eval counts but carry model + done.
+        return self.identify(chunk) or (hasattr(chunk, "done") and hasattr(chunk, "model"))
+
+    def new_stream_aggregator(self) -> StreamAggregator:
+        return _OllamaStreamAggregator()
+
+
+class _OllamaStreamAggregator(StreamAggregator):
+    """Ollama streaming: the final ``done`` chunk carries prompt_eval_count / eval_count."""
+
+    def __init__(self) -> None:
+        self._model = ""
+        self._input: int | None = None
+        self._output: int | None = None
+
+    def consume(self, chunk: object) -> None:
+        model = getattr(chunk, "model", None)
+        if model:
+            self._model = str(model)
+        prompt_eval = getattr(chunk, "prompt_eval_count", None)
+        if prompt_eval is not None:
+            self._input = prompt_eval
+        eval_count = getattr(chunk, "eval_count", None)
+        if eval_count is not None:
+            self._output = eval_count
+
+    def model(self) -> str:
+        return self._model
+
+    def usage(self) -> LLMUsage:
+        return LLMUsage(input_tokens=self._input, output_tokens=self._output)

tokenhelm/adapters/openai.py

@@ -0,0 +1,81 @@
+"""OpenAIAdapter — normalize OpenAI responses (T024).
+
+Observe-don't-patch: this reads attributes off whatever response object the developer's own
+OpenAI client returned. It identifies responses by shape (duck typing), so the core does not
+need the ``openai`` package installed to normalize a response. Supports both the Chat
+Completions shape (``usage.prompt_tokens`` / ``completion_tokens``) and the Responses API
+shape (``usage.input_tokens`` / ``output_tokens``).
+"""
+
+from __future__ import annotations
+
+from ..core.models import LLMProvider, LLMUsage
+from .base import BaseAdapter, StreamAggregator
+
+# OpenAI response objects carry an ``object`` discriminator, e.g. "chat.completion",
+# "response", or "text_completion". We use it (when present) to disambiguate from other
+# providers that also expose input/output token counts (notably Anthropic).
+_OPENAI_OBJECT_PREFIXES = ("chat.completion", "response", "text_completion")
+
+
+class OpenAIAdapter(BaseAdapter):
+    @property
+    def provider(self) -> LLMProvider:
+        return LLMProvider.OPENAI
+
+    def identify(self, response: object) -> bool:
+        obj = getattr(response, "object", None)
+        if isinstance(obj, str) and obj.startswith(_OPENAI_OBJECT_PREFIXES):
+            return True
+        usage = getattr(response, "usage", None)
+        # Chat Completions usage has the OpenAI-specific ``prompt_tokens`` field.
+        return usage is not None and hasattr(usage, "prompt_tokens")
+
+    def extract_model(self, response: object) -> str:
+        return str(getattr(response, "model", "") or "")
+
+    def extract_usage(self, response: object) -> LLMUsage:
+        usage = getattr(response, "usage", None)
+        if usage is None:
+            return LLMUsage()
+        return _usage_from_openai(usage)
+
+    def new_stream_aggregator(self) -> StreamAggregator:
+        return _OpenAIStreamAggregator()
+
+
+def _usage_from_openai(usage: object) -> LLMUsage:
+    """Normalize an OpenAI ``usage`` object (Chat Completions or Responses API naming)."""
+    input_tokens = getattr(usage, "prompt_tokens", None)
+    output_tokens = getattr(usage, "completion_tokens", None)
+    if input_tokens is None and output_tokens is None:
+        input_tokens = getattr(usage, "input_tokens", None)
+        output_tokens = getattr(usage, "output_tokens", None)
+    total_tokens = getattr(usage, "total_tokens", None)
+    return LLMUsage(
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        total_tokens=total_tokens,
+    )
+
+
+class _OpenAIStreamAggregator(StreamAggregator):
+    """OpenAI streaming: final usage arrives in a late chunk's ``usage`` (include_usage)."""
+
+    def __init__(self) -> None:
+        self._model = ""
+        self._usage = LLMUsage()
+
+    def consume(self, chunk: object) -> None:
+        model = getattr(chunk, "model", None)
+        if model:
+            self._model = str(model)
+        usage = getattr(chunk, "usage", None)
+        if usage is not None:
+            self._usage = _usage_from_openai(usage)
+
+    def model(self) -> str:
+        return self._model
+
+    def usage(self) -> LLMUsage:
+        return self._usage

tokenhelm/core/__init__.py

@@ -0,0 +1 @@
+"""Core: data model, tracker, extraction, calculator, config, errors."""

tokenhelm/core/calculator.py

@@ -0,0 +1,61 @@
+"""Cost calculation and currency formatting (T021).
+
+``CostCalculator`` depends ONLY on the :class:`PricingProvider` interface (Decision 3). A
+``None`` rate lookup yields an unpriced ``LLMCost`` (``priced=False``, zero cost) — never an
+error (FR-013).
+"""
+
+from __future__ import annotations
+
+from decimal import Decimal
+from typing import TYPE_CHECKING
+
+from .config import DEFAULT_CURRENCY
+from .models import LLMCost, LLMProvider, LLMUsage
+
+if TYPE_CHECKING:  # pragma: no cover - typing only
+    from ..pricing.base import PricingProvider
+
+_PER_MILLION = Decimal(1_000_000)
+
+
+class CostCalculator:
+    """Compute :class:`LLMCost` from usage + a pricing source."""
+
+    def __init__(self, pricing: PricingProvider, currency: str = DEFAULT_CURRENCY) -> None:
+        self._pricing = pricing
+        self.currency = currency
+
+    def compute(self, provider: LLMProvider, model: str, usage: LLMUsage) -> LLMCost:
+        rate = self._pricing.get_rates(provider, model)
+        if rate is None:
+            return LLMCost(currency=self.currency, priced=False)
+
+        input_tokens = usage.input_tokens or 0
+        output_tokens = usage.output_tokens or 0
+        input_cost = (Decimal(input_tokens) / _PER_MILLION) * rate.input_rate
+        output_cost = (Decimal(output_tokens) / _PER_MILLION) * rate.output_rate
+        return LLMCost(
+            input_cost=input_cost,
+            output_cost=output_cost,
+            total_cost=input_cost + output_cost,
+            currency=self.currency,
+            priced=True,
+        )
+
+
+class CurrencyFormatter:
+    """Human-readable money formatting for loggers/CLI output."""
+
+    _SYMBOLS = {"USD": "$", "EUR": "€", "GBP": "£", "JPY": "¥"}
+
+    def __init__(self, currency: str = DEFAULT_CURRENCY, *, places: int = 6) -> None:
+        self.currency = currency
+        self.places = places
+
+    def format(self, amount: Decimal) -> str:
+        symbol = self._SYMBOLS.get(self.currency, "")
+        value = f"{amount:.{self.places}f}"
+        if symbol:
+            return f"{symbol}{value}"
+        return f"{value} {self.currency}"

tokenhelm/core/config.py

@@ -0,0 +1,14 @@
+"""Configuration constants and helpers (T012).
+
+Kept intentionally small for v0.1 — the client holds the live configuration; this module
+provides shared defaults so other modules don't hardcode them.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+DEFAULT_CURRENCY = "USD"
+
+#: Bundled pricing file shipped with the package (``tokenhelm/data/pricing.yaml``).
+DEFAULT_PRICING_PATH = Path(__file__).resolve().parent.parent / "data" / "pricing.yaml"

tokenhelm/core/errors.py

@@ -0,0 +1,24 @@
+"""Error types for TokenHelm (T013).
+
+Tracking never raises on *missing data* — missing usage or missing pricing are represented
+as flags on the event (FR-013). These exceptions cover structural problems only.
+"""
+
+from __future__ import annotations
+
+
+class TokenHelmError(Exception):
+    """Base error, e.g. no adapter recognized a response object."""
+
+
+class ProviderNotInstalledError(TokenHelmError):
+    """A matching provider adapter needs an optional extra that isn't installed."""
+
+    def __init__(self, provider: str, extra: str | None = None) -> None:
+        extra = extra or provider
+        super().__init__(
+            f"The '{provider}' provider requires the optional extra. "
+            f'Install it with: pip install "tokenhelm[{extra}]"'
+        )
+        self.provider = provider
+        self.extra = extra