prompt_cache_warmer/

lib.rs

//! # prompt-cache-warmer
//!
//! Pre-warm Anthropic prompt cache before user traffic.
//!
//! Anthropic charges 25% more on the first request that creates a cache
//! entry and 10% as much on subsequent reads. If user requests are slow or
//! expensive on the first hit of a new system prompt, you want that first
//! hit to be a cheap synthetic warmup, not a real user.
//!
//! This crate:
//!
//!   1. Takes your long system prompt (string or block list) and a model
//!      name.
//!   2. Inserts up to N `cache_control` breakpoints in the right places.
//!   3. Fires a tiny warmup call (`max_tokens = 8` by default).
//!   4. Optionally fires a second verification call and asserts
//!      `cache_read_input_tokens > 0`.
//!   5. Returns a [`WarmResult`] with timings, token counts, and estimated
//!      cost.
//!
//! ## Quick example
//!
//! ```
//! use prompt_cache_warmer::{Block, Usage, WarmCall, WarmRequest, WarmResponse, Warmer};
//!
//! // BYO transport: anything that implements `WarmCall`.
//! struct FakeClient;
//! impl WarmCall for FakeClient {
//!     type Error = std::convert::Infallible;
//!     fn call(&self, _req: &WarmRequest) -> Result<WarmResponse, Self::Error> {
//!         Ok(WarmResponse {
//!             usage: Usage {
//!                 input_tokens: 10,
//!                 output_tokens: 4,
//!                 cache_creation_input_tokens: 12_000,
//!                 cache_read_input_tokens: 0,
//!             },
//!         })
//!     }
//! }
//!
//! let warmer = Warmer::new(FakeClient);
//! let out = warmer
//!     .warm("claude-opus-4-7", "long system text")
//!     .unwrap();
//! assert_eq!(out.cache_creation_input_tokens, 12_000);
//! ```
//!
//! `Warmer` is generic over the transport, so you can plug in the real
//! Anthropic HTTP client, a Bedrock wrapper, or a fake for tests.

#![deny(missing_docs)]

use std::collections::HashMap;
use std::time::Instant;

// ---- public usage / response shapes ----------------------------------------

/// Token usage returned by the model on a single warm call.
///
/// This mirrors the four fields Anthropic returns on `message.usage`.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Usage {
    /// Regular (non-cache) input tokens.
    pub input_tokens: u64,
    /// Output tokens generated by the model.
    pub output_tokens: u64,
    /// Tokens billed at the cache-write multiplier on this call.
    pub cache_creation_input_tokens: u64,
    /// Tokens billed at the cache-read multiplier on this call.
    pub cache_read_input_tokens: u64,
}

/// One block of a system prompt.
///
/// Anthropic accepts a system prompt as either a string or a list of
/// `{type: "text", text, cache_control?}` blocks. We model the block form
/// directly; conversion from a plain string is handled by
/// [`to_system_blocks`] and the `From` impls.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Block {
    /// The block's text content.
    pub text: String,
    /// Optional `cache_control` marker. `Some(CacheControl::Ephemeral)`
    /// means the block ends a cacheable prefix.
    pub cache_control: Option<CacheControl>,
}

/// Cache control marker for a [`Block`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CacheControl {
    /// Anthropic's `{"type": "ephemeral"}` breakpoint.
    Ephemeral,
}

impl From<String> for Block {
    fn from(text: String) -> Self {
        Block {
            text,
            cache_control: None,
        }
    }
}

impl From<&str> for Block {
    fn from(text: &str) -> Self {
        Block {
            text: text.to_string(),
            cache_control: None,
        }
    }
}

/// The request payload handed to a [`WarmCall`] transport.
///
/// This is intentionally minimal; we leave model-specific extras (top-p,
/// metadata, etc.) to the caller's transport wrapper.
#[derive(Debug, Clone)]
pub struct WarmRequest {
    /// Anthropic model id (e.g. `"claude-opus-4-7"`).
    pub model: String,
    /// System prompt, as one or more cacheable blocks.
    pub system_blocks: Vec<Block>,
    /// User/assistant messages. Defaults to a single `"ok"` ping if the
    /// caller doesn't supply any.
    pub messages: Vec<Message>,
    /// Optional tool definitions (`name`, `input_schema`, ...). Opaque to
    /// this crate; pass them straight through.
    pub tools: Vec<Tool>,
    /// `max_tokens` for the warm call. Defaults to 8.
    pub max_tokens: u32,
}

/// A single chat message (role + content text).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Message {
    /// `"user"` or `"assistant"`.
    pub role: String,
    /// Free-form text content.
    pub content: String,
}

/// An opaque tool definition. The transport owns serialization.
///
/// We only carry the tool name here so [`Warmer::warm`] callers can
/// inspect what got sent; richer tool shapes live in the transport.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Tool {
    /// Tool name (matches Anthropic's `tools[].name`).
    pub name: String,
}

/// Response shape returned by a [`WarmCall`] transport.
#[derive(Debug, Clone)]
pub struct WarmResponse {
    /// Token usage for this call.
    pub usage: Usage,
}

/// A transport that can execute a [`WarmRequest`] and return a
/// [`WarmResponse`].
///
/// Implement this for the Anthropic SDK, a Bedrock client, or a fake.
pub trait WarmCall {
    /// Transport-specific error type.
    type Error: std::error::Error;
    /// Execute the warm request.
    fn call(&self, req: &WarmRequest) -> Result<WarmResponse, Self::Error>;
}

// ---- block helpers ---------------------------------------------------------

/// Coerce a system arg into the Anthropic block list shape.
///
/// ```
/// use prompt_cache_warmer::{to_system_blocks, Block};
/// let blocks = to_system_blocks("hi");
/// assert_eq!(blocks, vec![Block::from("hi")]);
/// ```
pub fn to_system_blocks(system: &str) -> Vec<Block> {
    vec![Block::from(system)]
}

/// Add up to `n` ephemeral `cache_control` markers, evenly spaced and
/// ending with the last block.
///
/// Anthropic caps the number of breakpoints at 4 per request; we cap at 4
/// here too. Blocks that already carry a `cache_control` are preserved.
///
/// ```
/// use prompt_cache_warmer::{add_cache_breakpoints, Block, CacheControl};
/// let blocks = vec![Block::from("a"), Block::from("b")];
/// let out = add_cache_breakpoints(&blocks, 1);
/// assert_eq!(out[0].cache_control, None);
/// assert_eq!(out[1].cache_control, Some(CacheControl::Ephemeral));
/// ```
pub fn add_cache_breakpoints(blocks: &[Block], n: usize) -> Vec<Block> {
    if n == 0 || blocks.is_empty() {
        return blocks.to_vec();
    }

    let capped = n.min(4);
    let len = blocks.len();
    let positions: std::collections::HashSet<usize> = if capped >= len {
        (0..len).collect()
    } else {
        let step = len as f64 / capped as f64;
        (0..capped)
            .map(|i| (((i + 1) as f64 * step) as usize).saturating_sub(1))
            .collect()
    };

    blocks
        .iter()
        .enumerate()
        .map(|(i, b)| {
            let mut nb = b.clone();
            if positions.contains(&i) && nb.cache_control.is_none() {
                nb.cache_control = Some(CacheControl::Ephemeral);
            }
            nb
        })
        .collect()
}

// ---- pricing ---------------------------------------------------------------

/// Per-million-token list price for a single model.
///
/// Multipliers ([`CACHE_WRITE_MULTIPLIER`], [`CACHE_READ_MULTIPLIER`]) are
/// applied at cost-estimate time, not stored here.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ModelPrice {
    /// Per-1M-token regular input price (USD).
    pub input: f64,
    /// Per-1M-token output price (USD).
    pub output: f64,
}

/// Lookup table mapping model id -> [`ModelPrice`].
pub type PriceTable = HashMap<&'static str, ModelPrice>;

/// Cache-write multiplier applied to `cache_creation_input_tokens`.
pub const CACHE_WRITE_MULTIPLIER: f64 = 1.25;
/// Cache-read multiplier applied to `cache_read_input_tokens`.
pub const CACHE_READ_MULTIPLIER: f64 = 0.10;

/// Built-in best-effort pricing as of 2026-Q2. Override with
/// [`Warmer::with_prices`] for unsupported models.
pub fn default_prices() -> PriceTable {
    let mut t: PriceTable = HashMap::new();
    t.insert(
        "claude-opus-4-7",
        ModelPrice {
            input: 15.0,
            output: 75.0,
        },
    );
    t.insert(
        "claude-opus-4-6",
        ModelPrice {
            input: 15.0,
            output: 75.0,
        },
    );
    t.insert(
        "claude-sonnet-4-6",
        ModelPrice {
            input: 3.0,
            output: 15.0,
        },
    );
    t.insert(
        "claude-haiku-4-5",
        ModelPrice {
            input: 0.80,
            output: 4.0,
        },
    );
    t
}

// ---- result type -----------------------------------------------------------

/// Result of a single [`Warmer::warm`] (or [`Warmer::warm_verified`]) call.
#[derive(Debug, Clone, PartialEq)]
pub struct WarmResult {
    /// The model id passed in.
    pub model: String,
    /// `cache_creation_input_tokens` from the first call.
    pub cache_creation_input_tokens: u64,
    /// `cache_read_input_tokens` from the first call.
    pub cache_read_input_tokens: u64,
    /// `input_tokens` from the first call.
    pub input_tokens: u64,
    /// `output_tokens` from the first call.
    pub output_tokens: u64,
    /// Wall-clock latency of the first (warm) call, in milliseconds.
    pub latency_ms_warm: u64,
    /// `cache_read_input_tokens` from the optional verification call.
    pub verified_hit_tokens: Option<u64>,
    /// Wall-clock latency of the optional verification call, in
    /// milliseconds.
    pub latency_ms_verify: Option<u64>,
    /// Estimated USD cost of the warm call, or `None` if the model isn't
    /// in the price table.
    pub cost_usd: Option<f64>,
}

// ---- main type -------------------------------------------------------------

/// Caller-facing entry point. Generic over a [`WarmCall`] transport.
pub struct Warmer<C: WarmCall> {
    client: C,
    prices: PriceTable,
}

/// High-level input for [`Warmer::warm`]/[`Warmer::warm_verified`].
///
/// Use [`WarmInput::new`] for the common case (string system prompt + the
/// default "ok" ping). For richer cases, build the struct directly.
#[derive(Debug, Clone)]
pub struct WarmInput {
    /// Anthropic model id.
    pub model: String,
    /// System prompt blocks.
    pub system_blocks: Vec<Block>,
    /// Optional user/assistant messages. If empty, a single user `"ok"`
    /// message is sent.
    pub messages: Vec<Message>,
    /// Optional tool list.
    pub tools: Vec<Tool>,
    /// `max_tokens` for the warm call. Defaults to 8.
    pub max_tokens: u32,
    /// Number of `cache_control` breakpoints to inject (capped at 4).
    pub breakpoints: usize,
    /// Text used for the default user ping when `messages` is empty.
    pub ping_text: String,
}

impl WarmInput {
    /// Shortcut: a string system prompt with default options.
    pub fn new(model: impl Into<String>, system: impl AsRef<str>) -> Self {
        WarmInput {
            model: model.into(),
            system_blocks: to_system_blocks(system.as_ref()),
            messages: Vec::new(),
            tools: Vec::new(),
            max_tokens: 8,
            breakpoints: 1,
            ping_text: "ok".to_string(),
        }
    }
}

impl<C: WarmCall> Warmer<C> {
    /// Build a [`Warmer`] using the default price table.
    pub fn new(client: C) -> Self {
        Warmer {
            client,
            prices: default_prices(),
        }
    }

    /// Build a [`Warmer`] with a caller-supplied price table.
    pub fn with_prices(client: C, prices: PriceTable) -> Self {
        Warmer { client, prices }
    }

    /// Borrow the underlying transport. Useful when the transport carries
    /// its own state (call log, metrics, etc.) you want to inspect.
    pub fn client(&self) -> &C {
        &self.client
    }

    /// Convenience shortcut around [`Warmer::warm_with`] using the default
    /// [`WarmInput::new`] options.
    pub fn warm(
        &self,
        model: impl Into<String>,
        system: impl AsRef<str>,
    ) -> Result<WarmResult, C::Error> {
        self.warm_with(WarmInput::new(model, system))
    }

    /// Convenience shortcut that also runs a verification call.
    pub fn warm_verified(
        &self,
        model: impl Into<String>,
        system: impl AsRef<str>,
    ) -> Result<WarmResult, C::Error> {
        self.warm_with_verified(WarmInput::new(model, system))
    }

    /// Fire a single warm call.
    pub fn warm_with(&self, input: WarmInput) -> Result<WarmResult, C::Error> {
        self.warm_inner(input, false)
    }

    /// Fire a warm call, then a second call, and record the second call's
    /// `cache_read_input_tokens` on the result.
    pub fn warm_with_verified(&self, input: WarmInput) -> Result<WarmResult, C::Error> {
        self.warm_inner(input, true)
    }

    fn warm_inner(&self, input: WarmInput, verify: bool) -> Result<WarmResult, C::Error> {
        let WarmInput {
            model,
            system_blocks,
            messages,
            tools,
            max_tokens,
            breakpoints,
            ping_text,
        } = input;

        let system_blocks = add_cache_breakpoints(&system_blocks, breakpoints);
        let messages = if messages.is_empty() {
            vec![Message {
                role: "user".to_string(),
                content: ping_text,
            }]
        } else {
            messages
        };

        let request = WarmRequest {
            model: model.clone(),
            system_blocks,
            messages,
            tools,
            max_tokens,
        };

        let t0 = Instant::now();
        let resp1 = self.client.call(&request)?;
        let warm_ms = t0.elapsed().as_millis() as u64;
        let usage1 = resp1.usage;

        let (verified, verify_ms) = if verify {
            let t1 = Instant::now();
            let resp2 = self.client.call(&request)?;
            let verify_ms = t1.elapsed().as_millis() as u64;
            (Some(resp2.usage.cache_read_input_tokens), Some(verify_ms))
        } else {
            (None, None)
        };

        Ok(WarmResult {
            cost_usd: self.estimate_cost(&model, &usage1),
            model,
            cache_creation_input_tokens: usage1.cache_creation_input_tokens,
            cache_read_input_tokens: usage1.cache_read_input_tokens,
            input_tokens: usage1.input_tokens,
            output_tokens: usage1.output_tokens,
            latency_ms_warm: warm_ms,
            verified_hit_tokens: verified,
            latency_ms_verify: verify_ms,
        })
    }

    fn estimate_cost(&self, model: &str, usage: &Usage) -> Option<f64> {
        let price = self.prices.get(model)?;
        let per_in = price.input / 1_000_000.0;
        let per_out = price.output / 1_000_000.0;
        let regular_in = usage.input_tokens as f64;
        let write_in = usage.cache_creation_input_tokens as f64;
        let read_in = usage.cache_read_input_tokens as f64;
        let out = usage.output_tokens as f64;
        Some(
            regular_in * per_in
                + write_in * per_in * CACHE_WRITE_MULTIPLIER
                + read_in * per_in * CACHE_READ_MULTIPLIER
                + out * per_out,
        )
    }
}