tokenprof 0.1.0__py3-none-any.whl

tokenprof/__init__.py

@@ -0,0 +1,16 @@
+"""tokenprof: a profiler for the context window."""
+
+from tokenprof.attribute import profile_request
+from tokenprof.diff import diff_turns
+from tokenprof.types import Category, Profile, Segment, Turn
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Category",
+    "Profile",
+    "Segment",
+    "Turn",
+    "profile_request",
+    "diff_turns",
+]

tokenprof/adapters/__init__.py

@@ -0,0 +1,45 @@
+"""Adapter registry.
+
+Adding a framework means adding a module here and one line to ADAPTERS.
+That is the intended contribution surface.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from tokenprof.adapters.anthropic_messages import AnthropicMessagesAdapter
+from tokenprof.adapters.base import Adapter
+from tokenprof.adapters.openai_chat import OpenAIChatAdapter
+
+ADAPTERS: list[Adapter] = [
+    AnthropicMessagesAdapter(),
+    OpenAIChatAdapter(),
+]
+
+
+def get_adapter(name: str) -> Adapter:
+    for a in ADAPTERS:
+        if a.name == name:
+            return a
+    known = ", ".join(a.name for a in ADAPTERS)
+    raise KeyError(f"unknown adapter {name!r}. known adapters: {known}")
+
+
+def detect(request: dict[str, Any]) -> Adapter:
+    """Pick an adapter by payload shape.
+
+    Order matters: Anthropic is checked first because its signals are
+    positive (top-level system, input_schema) while the OpenAI adapter
+    accepts the general case.
+    """
+    for a in ADAPTERS:
+        if a.matches(request):
+            return a
+    raise ValueError(
+        "could not detect a provider from this payload. "
+        "Pass --provider explicitly, or open an issue with a redacted sample."
+    )
+
+
+__all__ = ["ADAPTERS", "Adapter", "detect", "get_adapter"]

tokenprof/adapters/anthropic_messages.py

@@ -0,0 +1,136 @@
+"""Anthropic Messages API request payloads."""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Any
+
+from tokenprof.adapters.base import measure, stringify
+from tokenprof.tokenizer import Tokenizer
+from tokenprof.types import Category, Segment
+
+
+class AnthropicMessagesAdapter:
+    name = "anthropic_messages"
+
+    def matches(self, request: dict[str, Any]) -> bool:
+        if not isinstance(request.get("messages"), list):
+            return False
+        if isinstance(request.get("system"), (str, list)):
+            return True
+        for spec in request.get("tools") or []:
+            if isinstance(spec, dict) and "input_schema" in spec:
+                return True
+        return False
+
+    def model_of(self, request: dict[str, Any]) -> str:
+        return str(request.get("model", ""))
+
+    def segments(self, request: dict[str, Any], tok: Tokenizer) -> list[Segment]:
+        # Emitted in the order the provider actually assembles the prompt:
+        # tools, then system, then messages. Cache prefixes are computed over
+        # that order, so getting it wrong would put the breakpoint in the
+        # wrong place.
+        out: list[Segment] = []
+        marked: list[int] = []
+
+        for spec in request.get("tools") or []:
+            name = spec.get("name", "?") if isinstance(spec, dict) else "?"
+            t, c, dg = measure(stringify(spec), tok)
+            if isinstance(spec, dict) and spec.get("cache_control"):
+                marked.append(len(out))
+            out.append(Segment(Category.TOOL_SCHEMA, t, c, name=name, digest=dg))
+
+        system = request.get("system")
+        if system is not None:
+            if isinstance(system, list):
+                for i, block in enumerate(system):
+                    text = block.get("text", "") if isinstance(block, dict) else str(block)
+                    t, c, dg = measure(text, tok)
+                    if not (t or c):
+                        continue
+                    if isinstance(block, dict) and block.get("cache_control"):
+                        marked.append(len(out))
+                    out.append(
+                        Segment(Category.SYSTEM, t, c, name=f"block[{i}]", index=i, digest=dg)
+                    )
+            else:
+                t, c, dg = measure(str(system), tok)
+                if t or c:
+                    out.append(Segment(Category.SYSTEM, t, c, name="system", digest=dg))
+
+        messages = request.get("messages") or []
+        last_user = _last_user_index(messages)
+
+        for i, msg in enumerate(messages):
+            if not isinstance(msg, dict):
+                continue
+            role = msg.get("role", "")
+            content = msg.get("content")
+
+            # A single user message can carry several tool_result blocks. They
+            # are attributed individually, because "which tool is flooding the
+            # context" is the question worth answering.
+            for cat, name, text, is_marked in _split_content(content, role):
+                t, c, dg = measure(text, tok)
+                if t == 0 and c == 0:
+                    continue
+                if cat is None:
+                    cat = _role_category(role, i == last_user)
+                if is_marked:
+                    marked.append(len(out))
+                out.append(Segment(cat, t, c, name=name, index=i, digest=dg))
+
+        # A cache_control marker caches the whole prefix up to and including
+        # that block, so only the last marker matters for how much is cached.
+        if marked:
+            cutoff = max(marked)
+            out = [
+                dataclasses.replace(s, cached=True) if i <= cutoff else s for i, s in enumerate(out)
+            ]
+        return out
+
+
+def _role_category(role: str, is_last_user: bool) -> Category:
+    if role == "user":
+        return Category.CURRENT_USER if is_last_user else Category.HISTORY_USER
+    if role == "assistant":
+        return Category.HISTORY_ASSISTANT
+    return Category.OTHER
+
+
+def _last_user_index(messages: list[Any]) -> int:
+    for i in range(len(messages) - 1, -1, -1):
+        m = messages[i]
+        if isinstance(m, dict) and m.get("role") == "user":
+            return i
+    return -1
+
+
+def _split_content(content: Any, role: str) -> list[tuple[Category | None, str, str, bool]]:
+    """Break message content into (category_override, name, text, cache_marked)."""
+    if content is None:
+        return []
+    if isinstance(content, str):
+        return [(None, "", content, False)]
+    if not isinstance(content, list):
+        return [(None, "", stringify(content), False)]
+
+    out: list[tuple[Category | None, str, str, bool]] = []
+    for block in content:
+        if not isinstance(block, dict):
+            out.append((None, "", str(block), False))
+            continue
+        marked = bool(block.get("cache_control"))
+        btype = block.get("type")
+        if btype == "tool_result":
+            name = str(block.get("tool_use_id") or "tool_result")
+            out.append((Category.TOOL_RESULT, name, stringify(block.get("content")), marked))
+        elif btype == "tool_use":
+            name = str(block.get("name") or "tool_use")
+            out.append((Category.HISTORY_ASSISTANT, name, stringify(block), marked))
+        elif btype == "text":
+            out.append((None, "", block.get("text", ""), marked))
+        else:
+            out.append((None, str(btype or ""), stringify(block), marked))
+    return out

tokenprof/adapters/base.py

@@ -0,0 +1,59 @@
+"""Adapter contract.
+
+An adapter turns one provider's request payload into a list of Segments.
+That is the whole interface. Keeping it this narrow is what makes a new
+adapter an afternoon of work instead of a refactor, and it is why the core
+has no dependency on any framework.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from typing import Any, Protocol
+
+from tokenprof.tokenizer import Tokenizer
+from tokenprof.types import Segment
+
+
+class Adapter(Protocol):
+    #: Short identifier used by --provider and in reports.
+    name: str
+
+    def matches(self, request: dict[str, Any]) -> bool:
+        """True if this adapter understands the payload shape."""
+        ...
+
+    def model_of(self, request: dict[str, Any]) -> str: ...
+
+    def segments(self, request: dict[str, Any], tok: Tokenizer) -> list[Segment]: ...
+
+
+def measure(text: str, tok: Tokenizer) -> tuple[int, int, str]:
+    """Return (tokens, chars, digest) for a string.
+
+    The digest is what makes cache analysis possible: two turns can be walked
+    segment by segment to find exactly where the prompt prefix stops being
+    byte-identical, which is the difference between a cache hit and paying
+    full price.
+    """
+    return tok.count(text), len(text), digest_of(text)
+
+
+def digest_of(text: str) -> str:
+    return hashlib.blake2b(text.encode("utf-8", "replace"), digest_size=8).hexdigest()
+
+
+def stringify(value: Any) -> str:
+    """Flatten a value the way a provider serializes it into the prompt.
+
+    Tool schemas and structured content reach the model as JSON, so counting
+    the JSON is closer to the truth than counting a Python repr. Not exact,
+    since providers wrap this in their own formatting, and the wrapper is
+    small relative to the payload.
+    """
+    if value is None:
+        return ""
+    if isinstance(value, str):
+        return value
+    return json.dumps(value, ensure_ascii=False, separators=(",", ":"))

tokenprof/adapters/openai_chat.py

@@ -0,0 +1,107 @@
+"""OpenAI Chat Completions request payloads."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from tokenprof.adapters.base import measure, stringify
+from tokenprof.tokenizer import Tokenizer
+from tokenprof.types import Category, Segment
+
+
+class OpenAIChatAdapter:
+    name = "openai_chat"
+
+    def matches(self, request: dict[str, Any]) -> bool:
+        """Discriminate against Anthropic, which also uses a "messages" list.
+
+        Two signals only one of them has: Anthropic puts the system prompt at
+        the top level, and its tool specs carry "input_schema". OpenAI wraps
+        tool specs in "function" and keeps the system prompt inside messages.
+        """
+        if not isinstance(request.get("messages"), list):
+            return False
+        if isinstance(request.get("system"), (str, list)):
+            return False
+        for spec in request.get("tools") or []:
+            if isinstance(spec, dict):
+                if "input_schema" in spec:
+                    return False
+                if "function" in spec:
+                    return True
+        return True
+
+    def model_of(self, request: dict[str, Any]) -> str:
+        return str(request.get("model", ""))
+
+    def segments(self, request: dict[str, Any], tok: Tokenizer) -> list[Segment]:
+        out: list[Segment] = []
+
+        for spec in request.get("tools") or []:
+            fn = spec.get("function", spec) if isinstance(spec, dict) else {}
+            name = fn.get("name", "?") if isinstance(fn, dict) else "?"
+            t, c, dg = measure(stringify(spec), tok)
+            out.append(Segment(Category.TOOL_SCHEMA, t, c, name=name, digest=dg))
+
+        messages = request.get("messages") or []
+        last_user = _last_user_index(messages)
+
+        for i, msg in enumerate(messages):
+            if not isinstance(msg, dict):
+                continue
+            role = msg.get("role", "")
+            text = _content_text(msg.get("content"))
+
+            if role == "assistant" and msg.get("tool_calls"):
+                text += stringify(msg["tool_calls"])
+
+            t, c, dg = measure(text, tok)
+            if t == 0 and c == 0:
+                continue
+
+            if role == "system" or role == "developer":
+                out.append(Segment(Category.SYSTEM, t, c, name=role, index=i, digest=dg))
+            elif role == "tool":
+                out.append(
+                    Segment(
+                        Category.TOOL_RESULT,
+                        t,
+                        c,
+                        name=str(msg.get("name") or msg.get("tool_call_id") or "tool"),
+                        index=i,
+                        digest=dg,
+                    )
+                )
+            elif role == "assistant":
+                out.append(Segment(Category.HISTORY_ASSISTANT, t, c, index=i, digest=dg))
+            elif role == "user":
+                cat = Category.CURRENT_USER if i == last_user else Category.HISTORY_USER
+                out.append(Segment(cat, t, c, index=i, digest=dg))
+            else:
+                out.append(Segment(Category.OTHER, t, c, name=role, index=i, digest=dg))
+
+        return out
+
+
+def _last_user_index(messages: list[Any]) -> int:
+    for i in range(len(messages) - 1, -1, -1):
+        m = messages[i]
+        if isinstance(m, dict) and m.get("role") == "user":
+            return i
+    return -1
+
+
+def _content_text(content: Any) -> str:
+    if content is None:
+        return ""
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        parts = []
+        for block in content:
+            if isinstance(block, dict):
+                parts.append(block.get("text") or stringify(block))
+            else:
+                parts.append(str(block))
+        return "".join(parts)
+    return stringify(content)

tokenprof/attribute.py

@@ -0,0 +1,93 @@
+"""Turn a raw request payload into a Turn."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterable, Iterator
+from typing import Any
+
+from tokenprof.adapters import Adapter, detect, get_adapter
+from tokenprof.tokenizer import Tokenizer, get_tokenizer
+from tokenprof.types import Profile, Turn
+
+
+def profile_request(
+    request: dict[str, Any],
+    *,
+    provider: str | None = None,
+    tokenizer: Tokenizer | None = None,
+    index: int = 0,
+) -> Turn:
+    """Profile a single request payload."""
+    adapter: Adapter = get_adapter(provider) if provider else detect(request)
+    model = adapter.model_of(request)
+    tok = tokenizer or get_tokenizer(model)
+    return Turn(
+        segments=adapter.segments(request, tok),
+        model=model,
+        provider=adapter.name,
+        index=index,
+        tokenizer=tok.name,
+    )
+
+
+def profile_stream(
+    records: Iterable[dict[str, Any]],
+    *,
+    provider: str | None = None,
+    tokenizer: Tokenizer | None = None,
+) -> Profile:
+    """Profile a sequence of records into one Profile.
+
+    The tokenizer is resolved once from the first record and reused, because
+    building a tiktoken encoding per turn dominates runtime on long sessions.
+    """
+    profile = Profile()
+    tok = tokenizer
+    for i, rec in enumerate(records):
+        request = unwrap(rec)
+        if tok is None:
+            adapter = get_adapter(provider) if provider else detect(request)
+            tok = get_tokenizer(adapter.model_of(request))
+        profile.turns.append(profile_request(request, provider=provider, tokenizer=tok, index=i))
+    return profile
+
+
+def unwrap(record: dict[str, Any]) -> dict[str, Any]:
+    """Accept either a bare request payload or a wrapper around one.
+
+    Recorders tend to store {"request": {...}, "ts": ...}. Both shapes are
+    accepted so nobody has to reshape their logs before getting an answer.
+    """
+    for key in ("request", "body", "payload", "kwargs"):
+        inner = record.get(key)
+        if isinstance(inner, dict) and "messages" in inner:
+            return inner
+    return record
+
+
+def read_jsonl(path: str) -> Iterator[dict[str, Any]]:
+    """Read a .jsonl file, skipping blank lines.
+
+    A malformed line raises with its line number rather than being dropped
+    silently, because a profile that quietly ignored half your session is
+    worse than one that refuses to run.
+    """
+    import sys
+
+    handle = sys.stdin if path == "-" else open(path, encoding="utf-8")
+    try:
+        for lineno, line in enumerate(handle, 1):
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                obj = json.loads(line)
+            except json.JSONDecodeError as exc:
+                raise ValueError(f"{path}:{lineno}: invalid JSON: {exc}") from exc
+            if not isinstance(obj, dict):
+                raise ValueError(f"{path}:{lineno}: expected a JSON object")
+            yield obj
+    finally:
+        if handle is not sys.stdin:
+            handle.close()

tokenprof/cache.py

@@ -0,0 +1,108 @@
+"""Prompt cache analysis.
+
+Providers cache a *prefix*. If the first N tokens of a request are
+byte-identical to the previous one, you pay a fraction for them. One volatile
+value near the front (a timestamp, a session id, a re-ordered tool list) moves
+the break point to zero and you quietly pay full price on every turn while
+your dashboard still reports that caching is enabled.
+
+This module finds that break point by walking two turns segment by segment
+and comparing content digests.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from tokenprof.types import Profile, Segment, Turn
+
+
+@dataclass(frozen=True)
+class PrefixBreak:
+    """Where one turn stopped matching the previous one."""
+
+    before_index: int
+    after_index: int
+    #: Segments that matched, counted from the start of the prompt.
+    stable_segments: int
+    stable_tokens: int
+    total_tokens: int
+    #: The first segment that differed, from the later turn. None when the
+    #: whole earlier prompt is still a prefix of the later one, which is the
+    #: healthy case for an append-only conversation.
+    breaking: Segment | None = None
+
+    @property
+    def stable_share(self) -> float:
+        return self.stable_tokens / self.total_tokens if self.total_tokens else 0.0
+
+    @property
+    def healthy(self) -> bool:
+        """True when nothing in the shared region changed."""
+        return self.breaking is None
+
+
+def prefix_break(before: Turn, after: Turn) -> PrefixBreak:
+    """Compare two turns and find where the cacheable prefix ends."""
+    n = 0
+    breaking: Segment | None = None
+    for a, b in zip(before.segments, after.segments):
+        if a.digest and a.digest == b.digest:
+            n += 1
+            continue
+        breaking = b
+        break
+
+    stable = sum(s.tokens for s in after.segments[:n])
+    return PrefixBreak(
+        before_index=before.index,
+        after_index=after.index,
+        stable_segments=n,
+        stable_tokens=stable,
+        total_tokens=after.total_tokens,
+        breaking=breaking,
+    )
+
+
+@dataclass
+class CacheReport:
+    breaks: list[PrefixBreak]
+    #: Tokens the provider was explicitly told to cache, summed over turns.
+    marked_tokens: int = 0
+    #: Tokens that were re-sent unchanged and could have been cached.
+    reusable_tokens: int = 0
+    #: Tokens re-sent after the prefix broke, so full price every turn.
+    rebuilt_tokens: int = 0
+
+    @property
+    def worst(self) -> PrefixBreak | None:
+        unhealthy = [b for b in self.breaks if not b.healthy]
+        if not unhealthy:
+            return None
+        return min(unhealthy, key=lambda b: b.stable_share)
+
+    @property
+    def thrashing(self) -> bool:
+        """True when the prefix breaks early enough to lose most of the win."""
+        return any(not b.healthy and b.stable_share < 0.5 for b in self.breaks)
+
+    def offenders(self) -> dict[str, int]:
+        """How often each segment is the one that broke the prefix."""
+        out: dict[str, int] = {}
+        for b in self.breaks:
+            if b.breaking is None:
+                continue
+            key = f"{b.breaking.category.value}:{b.breaking.name or '(unnamed)'}"
+            out[key] = out.get(key, 0) + 1
+        return dict(sorted(out.items(), key=lambda kv: kv[1], reverse=True))
+
+
+def analyze(profile: Profile) -> CacheReport:
+    report = CacheReport(breaks=[])
+    for before, after in zip(profile.turns, profile.turns[1:]):
+        b = prefix_break(before, after)
+        report.breaks.append(b)
+        report.reusable_tokens += b.stable_tokens
+        report.rebuilt_tokens += b.total_tokens - b.stable_tokens
+    report.marked_tokens = sum(t.cached_tokens() for t in profile.turns)
+    return report