tokenpak 1.0.0__py3-none-any.whl

tokenpak/__init__.py

@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: MIT
+"""TokenPak — Universal Content Compiler for LLMs.
+
+Public API surface for TokenPak v1.0.0.
+Formalizes importable classes for agent integrations, deployment, and testing.
+
+Quick start:
+    from tokenpak import TelemetryCollector, CacheManager, CompressionEngine, Budgeter
+
+Sub-package imports:
+    from tokenpak.telemetry import TelemetryCollector
+    from tokenpak.engines import CompactionEngine, HeuristicEngine
+    from tokenpak.registry import Block, BlockRegistry
+    from tokenpak.budgeter import Budgeter
+"""
+
+from __future__ import annotations
+
+__version__ = "1.0.0"
+__author__ = "Kevin Yang"
+__license__ = "MIT"
+__description__ = "Deterministic compression for multi-agent AI workflows"
+
+# ---------------------------------------------------------------------------
+# Telemetry
+# ---------------------------------------------------------------------------
+# ---------------------------------------------------------------------------
+# Sub-packages (for advanced use)
+# ---------------------------------------------------------------------------
+from tokenpak import agent, connectors, proxy
+
+# CompletionTracker: tracks per-completion cost, model, and latency
+from tokenpak.agent.telemetry.cost_tracker import CostTracker as CompletionTracker
+from tokenpak.budget import BudgetBlock
+
+# ---------------------------------------------------------------------------
+# Budgeting
+# ---------------------------------------------------------------------------
+from tokenpak.budgeter import Budgeter
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+from tokenpak.cli import main
+from tokenpak.engines import get_engine
+
+# ---------------------------------------------------------------------------
+# Compression / Compaction Engines
+# ---------------------------------------------------------------------------
+# CompressionEngine: abstract base for all compaction strategies
+from tokenpak.engines.base import CompactionEngine as CompressionEngine
+from tokenpak.engines.heuristic import HeuristicEngine
+from tokenpak.pack import CompiledResult, ContextPack, PackBlock, pack_prompt
+
+# ---------------------------------------------------------------------------
+# Content Blocks
+# ---------------------------------------------------------------------------
+from tokenpak.registry import Block, BlockRegistry
+
+# ---------------------------------------------------------------------------
+# Compile Reports
+# ---------------------------------------------------------------------------
+from tokenpak.report import Action, CompileReport, Decision
+
+# ---------------------------------------------------------------------------
+# Cache
+# ---------------------------------------------------------------------------
+# CacheManager: semantic cache store (get/set/hit-rate tracking)
+from tokenpak.telemetry.cache import CacheStore as CacheManager
+from tokenpak.telemetry.collector import TelemetryCollector
+
+# ---------------------------------------------------------------------------
+# Token Counting (Level 1 — single import, zero config)
+# ---------------------------------------------------------------------------
+from tokenpak.tokens import count_tokens
+from tokenpak.trace import (  # noqa: F401
+    TokenPakTrace,
+    TraceBuilder,
+    attach_trace_header,
+    attach_trace_envelope,
+    strip_trace,
+    strip_trace_header,
+    read_trace_header,
+    read_trace_envelope,
+    assert_no_leak,
+)
+
+
+# ---------------------------------------------------------------------------
+# Agent Handoff Protocol
+# ---------------------------------------------------------------------------
+from tokenpak.agent.agentic.handoff import (
+    HandoffBlock,
+    HandoffManager,
+    HandoffStatus,
+    HandoffWire as Handoff,
+    ContextRef,
+    TokenPak,
+)
+# ---------------------------------------------------------------------------
+# Agentic handoff protocol
+# ---------------------------------------------------------------------------
+from tokenpak.agent.agentic.handoff import (
+    ContextRef,
+    HandoffBlock,
+    HandoffManager,
+    HandoffStatus,
+    HandoffWire,
+    TokenPak,
+)
+# HandoffWire is the intended top-level "Handoff" API (pack-based wire format)
+# The internal Handoff dataclass (file-based) is available via
+# tokenpak.agent.agentic.handoff.Handoff
+Handoff = HandoffWire  # type: ignore
+
+# ---------------------------------------------------------------------------
+# Public API declaration
+# ---------------------------------------------------------------------------
+__all__ = [
+    # Metadata
+    "__version__",
+    "__author__",
+    "__license__",
+    "__description__",
+    # Telemetry
+    "TelemetryCollector",
+    "CompletionTracker",
+    # Cache
+    "CacheManager",
+    # Compression
+    "CompressionEngine",
+    "HeuristicEngine",
+    "get_engine",
+    # Content Blocks
+    "Block",
+    "BlockRegistry",
+    # Budgeting
+    "Budgeter",
+    "BudgetBlock",
+    # Compile Reports
+    "Action",
+    "CompileReport",
+    "Decision",
+    "ContextPack",
+    "PackBlock",
+    "CompiledResult",
+    # Incremental adoption helpers
+    "count_tokens",
+    "pack_prompt",
+    # Agent Handoff Protocol
+    "HandoffBlock",
+    "HandoffManager",
+    "HandoffStatus",
+    "Handoff",
+    "ContextRef",
+    "TokenPak",
+    # CLI
+    "main",
+    # Sub-packages
+    "connectors",
+    "agent",
+    "proxy",
+    # Agentic handoff protocol
+    "ContextRef",
+    "Handoff",
+    "HandoffBlock",
+    "HandoffManager",
+    "HandoffStatus",
+    "HandoffWire",
+    "TokenPak",
+]

tokenpak/__main__.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: MIT
+"""Allow `python -m tokenpak` to invoke the CLI."""
+
+from tokenpak.cli import main
+
+if __name__ == "__main__":
+    main()

tokenpak/adapters/__init__.py

@@ -0,0 +1,51 @@
+"""
+tokenpak.adapters — Unified SDK/framework adapter layer.
+
+All adapters share the ``TokenPakAdapter`` base contract:
+
+    prepare_request(request)  → normalised proxy dict
+    send(prepared_request)    → raw proxy response dict
+    parse_response(response)  → provider-native response dict
+    extract_tokens(response)  → token-usage summary dict
+
+Quick start
+-----------
+>>> from tokenpak.adapters import AnthropicAdapter, OpenAIAdapter
+>>> adapter = OpenAIAdapter(base_url="http://127.0.0.1:8767", api_key="sk-...")
+>>> response = adapter.call({"model": "gpt-4o", "messages": [...]})
+>>> usage = adapter.extract_tokens(response)
+
+See also
+--------
+- ``tokenpak.adapters.base``     — abstract base + exception hierarchy
+- ``tokenpak.adapters.anthropic`` — Anthropic Messages API
+- ``tokenpak.adapters.openai``    — OpenAI Chat Completions API
+- ``tokenpak.adapters.langchain`` — LangChain (ChatOpenAI / ChatAnthropic)
+- ``tokenpak.adapters.litellm``   — LiteLLM provider-agnostic routing
+"""
+
+from tokenpak.adapters.anthropic import AnthropicAdapter
+from tokenpak.adapters.base import (
+    TokenPakAdapter,
+    TokenPakAdapterError,
+    TokenPakAuthError,
+    TokenPakConfigError,
+    TokenPakTimeoutError,
+)
+from tokenpak.adapters.langchain import LangChainAdapter
+from tokenpak.adapters.litellm import LiteLLMAdapter
+from tokenpak.adapters.openai import OpenAIAdapter
+
+__all__ = [
+    # Base
+    "TokenPakAdapter",
+    "TokenPakAdapterError",
+    "TokenPakAuthError",
+    "TokenPakConfigError",
+    "TokenPakTimeoutError",
+    # Concrete adapters
+    "AnthropicAdapter",
+    "OpenAIAdapter",
+    "LangChainAdapter",
+    "LiteLLMAdapter",
+]

tokenpak/adapters/anthropic.py

@@ -0,0 +1,258 @@
+"""
+TokenPak Anthropic SDK Adapter
+
+Routes Anthropic ``messages.create`` requests through the TokenPak proxy.
+Preserves the full Anthropic Messages API shape on both request and response
+so callers require zero code changes beyond swapping in this adapter.
+
+Request format handled:
+  {
+    "model": "claude-3-5-sonnet-20241022",
+    "max_tokens": 4096,
+    "system": "...",
+    "messages": [{"role": "user", "content": "..."}],
+    ...
+  }
+
+Token extraction
+----------------
+Anthropic reports exact usage in every non-streaming response:
+  ``usage.input_tokens``, ``usage.output_tokens``,
+  ``usage.cache_read_input_tokens``, ``usage.cache_creation_input_tokens``
+
+Error handling
+--------------
+- 401/403 → TokenPakAuthError
+- timeout → TokenPakTimeoutError
+- provider error block in response → TokenPakAdapterError
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+from typing import Any
+
+try:
+    import requests as _requests
+    _REQUESTS_AVAILABLE = True
+except ImportError:
+    _REQUESTS_AVAILABLE = False
+
+from tokenpak.adapters.base import (
+    TokenPakAdapter,
+    TokenPakAdapterError,
+    TokenPakAuthError,
+    TokenPakConfigError,
+    TokenPakTimeoutError,
+)
+
+_log = logging.getLogger("tokenpak.adapters.anthropic")
+
+# Required fields in every request
+_REQUIRED_FIELDS = frozenset({"model", "messages", "max_tokens"})
+
+# Anthropic stop_reason → canonical finish_reason
+_STOP_REASON_MAP: dict[str, str] = {
+    "end_turn": "stop",
+    "max_tokens": "max_tokens",
+    "stop_sequence": "stop_sequence",
+    "tool_use": "tool_use",
+}
+
+
+class AnthropicAdapter(TokenPakAdapter):
+    """TokenPak adapter for the Anthropic Messages API.
+
+    Usage
+    -----
+    >>> adapter = AnthropicAdapter(
+    ...     base_url="http://127.0.0.1:8767",
+    ...     api_key="sk-ant-...",
+    ... )
+    >>> response = adapter.call({
+    ...     "model": "claude-3-5-sonnet-20241022",
+    ...     "max_tokens": 1024,
+    ...     "messages": [{"role": "user", "content": "Hello"}],
+    ... })
+    >>> tokens = adapter.extract_tokens(response)
+    """
+
+    provider_name: str = "anthropic"
+
+    PROXY_PATH: str = "/v1/messages"
+    ANTHROPIC_VERSION: str = "2023-06-01"
+
+    # ── prepare_request ───────────────────────────────────────────────────
+
+    def prepare_request(self, request: dict[str, Any]) -> dict[str, Any]:
+        """Validate and normalise an Anthropic request.
+
+        Validation:
+        - ``model``, ``messages``, ``max_tokens`` must be present
+        - ``messages`` must be a non-empty list
+        - Each message must have ``role`` (str) and ``content`` (str | list)
+
+        Normalisation:
+        - Strips unknown top-level keys that would be rejected by the proxy
+          (conservative: passes through everything; proxy decides)
+        - Adds ``stream: false`` default if not specified
+        """
+        missing = _REQUIRED_FIELDS - request.keys()
+        if missing:
+            raise TokenPakConfigError(
+                f"AnthropicAdapter.prepare_request: missing required fields: {sorted(missing)}"
+            )
+
+        messages = request.get("messages")
+        if not isinstance(messages, list) or not messages:
+            raise TokenPakConfigError(
+                "AnthropicAdapter.prepare_request: 'messages' must be a non-empty list."
+            )
+
+        for i, msg in enumerate(messages):
+            if not isinstance(msg, dict):
+                raise TokenPakConfigError(
+                    f"AnthropicAdapter.prepare_request: messages[{i}] must be a dict."
+                )
+            if "role" not in msg or "content" not in msg:
+                raise TokenPakConfigError(
+                    f"AnthropicAdapter.prepare_request: messages[{i}] must have 'role' and 'content'."
+                )
+
+        prepared = dict(request)
+        prepared.setdefault("stream", False)
+
+        self.logger.debug(
+            "prepare_request model=%s messages=%d",
+            prepared.get("model"),
+            len(messages),
+        )
+        return prepared
+
+    # ── send ──────────────────────────────────────────────────────────────
+
+    def send(self, prepared_request: dict[str, Any]) -> dict[str, Any]:
+        """POST to ``{base_url}/v1/messages`` through the proxy."""
+        if not _REQUESTS_AVAILABLE:
+            raise TokenPakAdapterError(
+                "AnthropicAdapter.send: 'requests' package is required. "
+                "Install with: pip install requests"
+            )
+
+        url = f"{self.base_url}{self.PROXY_PATH}"
+        headers = {
+            "Content-Type": "application/json",
+            "x-api-key": self.api_key,
+            "anthropic-version": self.ANTHROPIC_VERSION,
+        }
+
+        self.logger.debug("send POST %s model=%s", url, prepared_request.get("model"))
+        t0 = time.monotonic()
+
+        try:
+            resp = _requests.post(
+                url,
+                headers=headers,
+                data=json.dumps(prepared_request, ensure_ascii=False).encode("utf-8"),
+                timeout=self.timeout_s,
+            )
+        except _requests.exceptions.Timeout as exc:
+            raise TokenPakTimeoutError(
+                f"AnthropicAdapter.send: request timed out after {self.timeout_s}s."
+            ) from exc
+        except _requests.exceptions.RequestException as exc:
+            raise TokenPakAdapterError(
+                f"AnthropicAdapter.send: HTTP transport error: {exc}"
+            ) from exc
+
+        elapsed_ms = (time.monotonic() - t0) * 1000
+        self.logger.info("send complete status=%d elapsed_ms=%.1f", resp.status_code, elapsed_ms)
+
+        if resp.status_code in (401, 403):
+            raise TokenPakAuthError(
+                f"AnthropicAdapter.send: authentication failed (HTTP {resp.status_code}).",
+                status_code=resp.status_code,
+            )
+
+        if not resp.ok:
+            try:
+                err_body = resp.json()
+            except Exception:
+                err_body = resp.text
+            raise TokenPakAdapterError(
+                f"AnthropicAdapter.send: proxy returned HTTP {resp.status_code}.",
+                status_code=resp.status_code,
+                raw=err_body,
+            )
+
+        try:
+            return resp.json()
+        except Exception as exc:
+            raise TokenPakAdapterError(
+                "AnthropicAdapter.send: response body is not valid JSON."
+            ) from exc
+
+    # ── parse_response ────────────────────────────────────────────────────
+
+    def parse_response(self, response: dict[str, Any]) -> dict[str, Any]:
+        """Validate proxy response and surface provider errors.
+
+        Returns the response unchanged if valid; raises ``TokenPakAdapterError``
+        when the response contains an Anthropic error block.
+        """
+        error = response.get("error")
+        if error:
+            if isinstance(error, dict):
+                msg = error.get("message", str(error))
+                err_type = error.get("type", "unknown")
+            else:
+                msg, err_type = str(error), "unknown"
+            raise TokenPakAdapterError(
+                f"AnthropicAdapter.parse_response: provider error [{err_type}]: {msg}",
+                raw=response,
+            )
+
+        if "content" not in response and "type" not in response:
+            self.logger.warning(
+                "parse_response: response missing 'content' and 'type' — "
+                "may be malformed: %s",
+                list(response.keys()),
+            )
+
+        return response
+
+    # ── extract_tokens ────────────────────────────────────────────────────
+
+    def extract_tokens(self, response: dict[str, Any]) -> dict[str, Any]:
+        """Extract Anthropic usage block.
+
+        Returns zeros with a warning if ``usage`` is absent (e.g. streaming
+        partial chunks where billing data hasn't arrived yet).
+        """
+        usage = response.get("usage", {})
+        if not usage:
+            self.logger.warning(
+                "extract_tokens: no 'usage' block in response — returning zeros."
+            )
+            return {
+                "input_tokens": 0,
+                "output_tokens": 0,
+                "cache_read": 0,
+                "cache_write": 0,
+                "total": 0,
+            }
+
+        input_tokens = int(usage.get("input_tokens", 0))
+        output_tokens = int(usage.get("output_tokens", 0))
+        cache_read = int(usage.get("cache_read_input_tokens", 0))
+        cache_write = int(usage.get("cache_creation_input_tokens", 0))
+
+        return {
+            "input_tokens": input_tokens,
+            "output_tokens": output_tokens,
+            "cache_read": cache_read,
+            "cache_write": cache_write,
+            "total": input_tokens + output_tokens,
+        }

tokenpak/adapters/base.py

@@ -0,0 +1,212 @@
+"""
+TokenPak Unified Adapter Base — TokenPakAdapter
+
+All SDK/framework adapters (Anthropic, OpenAI, LangChain, LiteLLM, etc.)
+inherit from this class and implement the four lifecycle hooks.
+
+Lifecycle
+---------
+1. ``prepare_request(request)``  — validate & normalise SDK request → proxy format
+2. ``send(prepared_request)``    — POST to TokenPak proxy, handle errors/timeouts
+3. ``parse_response(response)``  — proxy response → SDK-native format
+4. ``extract_tokens(response)``  — pull token-usage counts for budgeting
+
+Error Handling Contract
+-----------------------
+All concrete adapters MUST:
+- Wrap HTTP errors in ``TokenPakAdapterError``
+- Wrap timeout errors in ``TokenPakTimeoutError``
+- Wrap auth/config errors in ``TokenPakConfigError``
+- Never raise bare ``requests.exceptions.*`` or provider SDK exceptions
+  directly — always translate to the canonical hierarchy.
+
+Logging
+-------
+Use ``self.logger`` (standard ``logging.Logger``) for all log output.
+Log levels:
+- DEBUG   : request/response body summaries (no credentials)
+- INFO    : round-trip timings, cache-hit notices
+- WARNING : retried requests, degraded fallbacks
+- ERROR   : terminal failures before raising exceptions
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from abc import ABC, abstractmethod
+from typing import Any
+
+# ── Canonical exception hierarchy ─────────────────────────────────────────
+
+class TokenPakAdapterError(Exception):
+    """Base exception for all TokenPak adapter errors."""
+
+    def __init__(self, message: str, status_code: int | None = None, raw: Any = None) -> None:
+        super().__init__(message)
+        self.status_code = status_code
+        self.raw = raw
+
+
+class TokenPakTimeoutError(TokenPakAdapterError):
+    """Raised when the proxy does not respond within ``timeout_s`` seconds."""
+
+
+class TokenPakConfigError(TokenPakAdapterError):
+    """Raised when required configuration (base_url, api_key) is missing/invalid."""
+
+
+class TokenPakAuthError(TokenPakAdapterError):
+    """Raised on 401/403 responses from the proxy."""
+
+
+# ── Base adapter ───────────────────────────────────────────────────────────
+
+class TokenPakAdapter(ABC):
+    """Abstract base class for all TokenPak SDK/framework adapters.
+
+    Parameters
+    ----------
+    base_url:
+        TokenPak proxy endpoint, e.g. ``"http://127.0.0.1:8767"``.
+        Must not have a trailing slash.
+    api_key:
+        Provider API key forwarded transparently through the proxy.
+    timeout_s:
+        Request timeout in seconds.  Defaults to 120.
+    """
+
+    #: Subclasses set this to a stable, lowercase identifier (e.g. "anthropic").
+    provider_name: str = "unknown"
+
+    DEFAULT_TIMEOUT_S: float = 120.0
+
+    def __init__(self, base_url: str, api_key: str, timeout_s: float | None = None) -> None:
+        if not base_url:
+            raise TokenPakConfigError("base_url must not be empty.")
+        if not api_key:
+            raise TokenPakConfigError("api_key must not be empty.")
+
+        self.base_url = base_url.rstrip("/")
+        self.api_key = api_key
+        self.timeout_s = timeout_s if timeout_s is not None else self.DEFAULT_TIMEOUT_S
+        self.logger = logging.getLogger(f"tokenpak.adapters.{self.provider_name}")
+
+    # ── Lifecycle hooks ────────────────────────────────────────────────────
+
+    @abstractmethod
+    def prepare_request(self, request: dict[str, Any]) -> dict[str, Any]:
+        """Validate and normalise an SDK request dict into proxy format.
+
+        Parameters
+        ----------
+        request:
+            Raw dict as the SDK caller would build it (model, messages, etc.).
+
+        Returns
+        -------
+        dict
+            Proxy-ready request dict.  Must contain at least ``"model"``
+            and ``"messages"`` (or the provider's equivalent).
+
+        Raises
+        ------
+        TokenPakConfigError
+            If required fields are missing or have invalid types.
+        """
+
+    @abstractmethod
+    def send(self, prepared_request: dict[str, Any]) -> dict[str, Any]:
+        """POST *prepared_request* to the TokenPak proxy and return the response.
+
+        Parameters
+        ----------
+        prepared_request:
+            The dict returned by ``prepare_request``.
+
+        Returns
+        -------
+        dict
+            Raw response dict from the proxy.
+
+        Raises
+        ------
+        TokenPakTimeoutError
+            If the request exceeds ``self.timeout_s``.
+        TokenPakAuthError
+            On 401/403 responses.
+        TokenPakAdapterError
+            For all other HTTP/transport errors.
+        """
+
+    @abstractmethod
+    def parse_response(self, response: dict[str, Any]) -> dict[str, Any]:
+        """Convert a raw proxy response into the provider's native SDK format.
+
+        Parameters
+        ----------
+        response:
+            Raw dict returned by ``send``.
+
+        Returns
+        -------
+        dict
+            Response shaped exactly as the provider's SDK would return it,
+            so callers can switch to TokenPak without code changes.
+
+        Raises
+        ------
+        TokenPakAdapterError
+            If the response is malformed or contains a provider error.
+        """
+
+    @abstractmethod
+    def extract_tokens(self, response: dict[str, Any]) -> dict[str, Any]:
+        """Extract token usage counts from a response.
+
+        Parameters
+        ----------
+        response:
+            Either the raw proxy dict or the parsed SDK-format dict.
+
+        Returns
+        -------
+        dict with keys:
+            - ``input_tokens``  (int) — billed input tokens
+            - ``output_tokens`` (int) — billed output tokens
+            - ``cache_read``    (int) — tokens served from cache (0 if N/A)
+            - ``cache_write``   (int) — tokens written to cache (0 if N/A)
+            - ``total``         (int) — ``input_tokens + output_tokens``
+        """
+
+    # ── Convenience helpers ────────────────────────────────────────────────
+
+    def call(self, request: dict[str, Any]) -> dict[str, Any]:
+        """Full pipeline: prepare → send → parse_response.
+
+        Returns the provider-native response dict.  Logs round-trip time
+        at INFO level.
+
+        Parameters
+        ----------
+        request:
+            Raw SDK-style request dict.
+        """
+        t0 = time.monotonic()
+        prepared = self.prepare_request(request)
+        raw_response = self.send(prepared)
+        parsed = self.parse_response(raw_response)
+        elapsed_ms = (time.monotonic() - t0) * 1000
+        self.logger.info(
+            "call complete provider=%s model=%s elapsed_ms=%.1f",
+            self.provider_name,
+            request.get("model", "?"),
+            elapsed_ms,
+        )
+        return parsed
+
+    def __repr__(self) -> str:
+        return (
+            f"<{type(self).__name__} provider={self.provider_name!r}"
+            f" base_url={self.base_url!r}>"
+        )