Cache Breakpoint Planner — Python source

"""Cache Breakpoint Planner — find the blocks a set of prompts share and
place cache breakpoints where they pay.

Language: Python (3.10+, standard library only)
Port of src/lib/cacheBreakpointPlanner.ts (the canonical TypeScript
implementation). javascript.js in this set carries the same port; this file
mirrors it for Python.
Tool page: https://dev.cosmolabs.org/tools/cache-breakpoint-planner
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import List, Sequence

#: Cached reads bill at ~0.1x — the saving on the cached share is ~90%.
CACHE_READ_DISCOUNT: float = 0.1


@dataclass(frozen=True)
class PromptSession:
    """One prompt session: an id plus its ordered blocks."""

    id: str
    blocks: List[str] = field(default_factory=list)


@dataclass(frozen=True)
class Breakpoint:
    """Place the cache breakpoint AFTER this block index (0-based)."""

    after_block: int
    label: str
    reason: str
    cached_tokens: int


@dataclass(frozen=True)
class PerSessionRow:
    id: str
    total_tokens: int
    unique_tokens: int
    cached_ratio: float


@dataclass(frozen=True)
class BreakpointPlan:
    prefix_blocks: List[str] = field(default_factory=list)
    prefix_tokens: int = 0
    suffix_blocks: List[str] = field(default_factory=list)
    suffix_tokens: int = 0
    breakpoints: List[Breakpoint] = field(default_factory=list)
    per_session: List[PerSessionRow] = field(default_factory=list)
    #: Estimated cost saving across the sessions vs no caching (0-1).
    estimated_savings: float = 0.0
    warnings: List[str] = field(default_factory=list)


def _tok(text: str) -> int:
    """The `type: 'prose'` path of the tokenEstimator, inlined: every
    non-empty line costs max(1, round(length / 4)) tokens; empty text is 0."""
    if not text:
        return 0
    tokens = 0
    for line in text.splitlines():
        if line:
            tokens += max(1, round(len(line) / 4))
    return tokens


def plan_breakpoints(sessions: Sequence[PromptSession]) -> BreakpointPlan:
    """Plan cache breakpoints for a set of prompt sessions.

    Finds the common leading/trailing blocks across every session and places
    breakpoints where the cache pays: everything stable before the
    breakpoint, everything per-request after it.
    """
    warnings: List[str] = []
    valid = [s for s in sessions if isinstance(s.blocks, list)]

    if not valid:
        return BreakpointPlan(
            warnings=["No sessions given — paste at least two prompts to compare."]
        )
    if len(valid) == 1:
        warnings.append("Only one session — a prefix needs at least two prompts to detect.")

    # Common leading blocks by position.
    shortest = min(len(s.blocks) for s in valid)
    prefix_end = 0
    while prefix_end < shortest and all(
        s.blocks[prefix_end] == valid[0].blocks[prefix_end] for s in valid
    ):
        prefix_end += 1

    # Common trailing blocks, matched from each session's own tail, never
    # overlapping the prefix.
    suffix_len = 0
    while suffix_len < shortest - prefix_end and all(
        s.blocks[len(s.blocks) - 1 - suffix_len]
        == valid[0].blocks[len(valid[0].blocks) - 1 - suffix_len]
        for s in valid
    ):
        suffix_len += 1

    prefix_blocks = valid[0].blocks[:prefix_end]
    suffix_blocks = valid[0].blocks[len(valid[0].blocks) - suffix_len:] if suffix_len > 0 else []
    prefix_tokens = _tok("\n".join(prefix_blocks))
    suffix_tokens = _tok("\n".join(suffix_blocks))

    breakpoints: List[Breakpoint] = []
    if prefix_blocks:
        breakpoints.append(
            Breakpoint(
                after_block=prefix_end - 1,
                label="after the shared prefix",
                reason=(
                    f"{len(prefix_blocks)} block(s) identical across every session — "
                    "cache once, hit on every request."
                ),
                cached_tokens=prefix_tokens,
            )
        )
    if suffix_len > 0:
        breakpoints.append(
            Breakpoint(
                after_block=-1,  # terminal: the shared tail sits at the end
                label="shared tail",
                reason=(
                    f"{suffix_len} trailing block(s) also identical — extend the cache "
                    "segment or accept the re-read."
                ),
                cached_tokens=suffix_tokens,
            )
        )
    if not breakpoints:
        warnings.append(
            "No shared leading or trailing blocks — nothing to cache across these sessions."
        )

    per_session: List[PerSessionRow] = []
    for s in valid:
        total_tokens = _tok("\n".join(s.blocks))
        unique_tokens = max(total_tokens - prefix_tokens - suffix_tokens, 0)
        cached_ratio = (
            min((prefix_tokens + suffix_tokens) / total_tokens, 1.0)
            if total_tokens > 0
            else 0.0
        )
        per_session.append(
            PerSessionRow(s.id, total_tokens, unique_tokens, cached_ratio)
        )

    avg_total = sum(p.total_tokens for p in per_session) / len(per_session)
    cached_share = (
        min((prefix_tokens + suffix_tokens) / avg_total, 1.0) if avg_total > 0 else 0.0
    )
    estimated_savings = cached_share * (1 - CACHE_READ_DISCOUNT)

    return BreakpointPlan(
        prefix_blocks=prefix_blocks,
        prefix_tokens=prefix_tokens,
        suffix_blocks=suffix_blocks,
        suffix_tokens=suffix_tokens,
        breakpoints=breakpoints,
        per_session=per_session,
        estimated_savings=estimated_savings,
        warnings=warnings,
    )