Cache the system prompt and price cached tokens

hyperbolic/cli.py

@@ -122,7 +122,8 @@ def cmd_replay(args) -> int:
     rep = replay_run(args.run_dir, out_dir=args.out)
     print(f"replayed {args.run_dir}: {rep.steps} hours, {rep.calls} recorded calls, {rep.api_calls} API calls")
     print(f"  identical decisions and actions: {rep.identical_steps}/{rep.steps}")
-    print(f"  clamped hours {rep.clamped_steps}, fallback hours {rep.fallback_steps}, recorded cost ${rep.cost_usd:.4f}")
+    print(f"  clamped hours {rep.clamped_steps}, fallback hours {rep.fallback_steps}, recorded cost ${rep.cost_usd:.4f}, "
+          f"cache hit rate {_hit_rate(rep.cache_hit_rate)}")
     if rep.deterministic:
         print("  deterministic: yes")
         return 0
@@ -165,7 +166,8 @@ def cmd_compare(args) -> int:
         print(f"  {label}: {d}")
     r = s["agent_run"]
     print(f"  model calls {r['calls']}, clamped hours {r['clamped_hours']}, fallback hours {r['fallback_hours']}, "
-          f"cost ${r['cost_usd']:.4f} of ${r['budget_usd']:.2f}{' (budget stop hit)' if r['budget_stopped'] else ''}")
+          f"cost ${r['cost_usd']:.4f} of ${r['budget_usd']:.2f}{' (budget stop hit)' if r['budget_stopped'] else ''}"
+          + (f", cache hit rate {_hit_rate(r['cache_hit_rate'])}" if "cache_hit_rate" in r else ""))
     print(f"\nviewer data written to {args.out}")
     return 0
 
@@ -245,13 +247,21 @@ def _fmt(v):
     return "n/a" if v is None else f"{v:.4f}"
 
 
+def _hit_rate(v) -> str:
+    return "n/a" if v is None else f"{100 * v:.1f}%"
+
+
 def _print_result(res: dict) -> None:
     k = res["district_kpis"]
     parts = [f"{key}={k[key]:.4f}" for key in dc.HEADLINE_KPIS if k.get(key) is not None]
     print(f"  done: {res['steps']} hours{'' if res['complete'] else ' (partial)'}; " + ", ".join(parts))
     if res.get("calls"):
         print(f"  model calls {res['calls']}, clamped {res['clamped_steps']}, fallback {res['fallback_steps']}, "
               f"cost ${res['cost_usd']:.4f}{'; ' + res['budget_stop_reason'] if res['budget_stopped'] else ''}")
+        t = res.get("tokens")
+        if t:
+            print(f"  input tokens {t['input']} uncached, {t['cache_read']} cache reads, {t['cache_write']} cache "
+                  f"writes; output {t['output']}; cache hit rate {_hit_rate(res.get('cache_hit_rate'))}")
 
 
 def build_parser() -> argparse.ArgumentParser:

hyperbolic/live.py

@@ -408,6 +408,7 @@ def on_step(rec: dict) -> None:
                 "stop_reason": run.result.get("stop_reason"), "steps": run.result["steps"], "hours": run.hours,
                 "agent_run": str(run.agent_dir), "baseline_run": str(run.baseline_dir), "kpis": rows,
                 "cost_usd": run.result["cost_usd"], "calls": run.result["calls"],
+                "cache_hit_rate": run.result.get("cache_hit_rate"),
                 "clamped": run.result["clamped_steps"], "fallback": run.result["fallback_steps"],
                 "budget_stopped": run.result["budget_stopped"], "budget_stop_reason": run.result["budget_stop_reason"],
             })

hyperbolic/live_assets/live.js

@@ -315,7 +315,8 @@
     table.appendChild(tb);
     box.appendChild(table);
     var facts = el("p", "live-note", "API spend " + money(ev.cost_usd) + ", " + ev.calls + " calls, " + ev.clamped + " clamped hours, " +
-      ev.fallback + " fallback hours" + (ev.budget_stopped ? ", budget stop hit" : "") + ". Recorded in " + ev.agent_run + "; replay it with `hyperbolic replay`.");
+      ev.fallback + " fallback hours" + (ev.cache_hit_rate != null ? ", cache hit rate " + fmt(100 * ev.cache_hit_rate, 1) + "%" : "") +
+      (ev.budget_stopped ? ", budget stop hit" : "") + ". Recorded in " + ev.agent_run + "; replay it with `hyperbolic replay`.");
     box.appendChild(facts);
     var pub = button("live-publish", "Publish to docs/data", "primary");
     var note = el("p", "live-note"); note.id = "live-published";

hyperbolic/model.py

@@ -25,10 +25,13 @@
 
 KEY_ENV_VARS = ("HYPERBOLIC_ANTHROPIC_KEY", "TRIAGE_ANTHROPIC_KEY")
 
-# USD per million tokens (input, output).
+# USD per million tokens. input is uncached input; a cache write (5 minute TTL) bills at 1.25x
+# the input price and a cache read at 0.1x.
 MODELS = {
-    "haiku": {"id": "claude-haiku-5-5", "input_per_mtok": 0.10, "output_per_mtok": 0.50},
-    "sonnet": {"id": "claude-sonnet-5-5", "input_per_mtok": 2.00, "output_per_mtok": 10.00},
+    "haiku": {"id": "claude-haiku-5-5", "input_per_mtok": 0.10, "cache_write_per_mtok": 0.125,
+              "cache_read_per_mtok": 0.01, "output_per_mtok": 0.50},
+    "sonnet": {"id": "claude-sonnet-5-5", "input_per_mtok": 2.00, "cache_write_per_mtok": 2.50,
+               "cache_read_per_mtok": 0.20, "output_per_mtok": 10.00},
 }
 DEFAULT_MODEL = "haiku"
 DEFAULT_BUDGET_USD = 2.00
@@ -55,15 +58,35 @@ def request_sha256(request: dict) -> str:
     return hashlib.sha256(blob.encode()).hexdigest()
 
 
-def call_cost_usd(model_key: str, input_tokens: int, output_tokens: int) -> float:
+def call_cost_usd(model_key: str, input_tokens: int, output_tokens: int,
+                  cache_read_tokens: int = 0, cache_write_tokens: int = 0) -> float:
+    """input_tokens is uncached input only, as the API reports it beside the two cache counts."""
     p = MODELS[model_key]
-    return (input_tokens * p["input_per_mtok"] + output_tokens * p["output_per_mtok"]) / 1_000_000
+    return (input_tokens * p["input_per_mtok"] + cache_write_tokens * p["cache_write_per_mtok"]
+            + cache_read_tokens * p["cache_read_per_mtok"] + output_tokens * p["output_per_mtok"]) / 1_000_000
+
+
+def cache_hit_rate(input_tokens: int, cache_read_tokens: int, cache_write_tokens: int) -> float | None:
+    """Share of all input tokens served from the cache; None when there was no input."""
+    total = input_tokens + cache_read_tokens + cache_write_tokens
+    return None if total == 0 else cache_read_tokens / total
+
+
+def _cached_prefix_chars(request: dict) -> int:
+    """Characters up to and including the last block marked for caching (tools render before system)."""
+    system = request.get("system")
+    blocks = [*request.get("tools", []), *(system if isinstance(system, list) else [])]
+    marked = [i for i, b in enumerate(blocks) if b.get("cache_control")]
+    return len(json.dumps(blocks[: marked[-1] + 1], ensure_ascii=False)) if marked else 0
 
 
 def worst_case_cost_usd(model_key: str, request: dict) -> float:
-    """Upper bound for one call: about 3 characters per input token, full max_tokens output."""
+    """Upper bound for one call: about 3 characters per input token, full max_tokens output, and a
+    cached prefix priced as a cache write, the dearest way it can bill."""
     chars = len(json.dumps(request, ensure_ascii=False))
-    return call_cost_usd(model_key, math.ceil(chars / 3), int(request.get("max_tokens", DEFAULT_MAX_TOKENS)))
+    prefix = _cached_prefix_chars(request)
+    return call_cost_usd(model_key, math.ceil((chars - prefix) / 3), int(request.get("max_tokens", DEFAULT_MAX_TOKENS)),
+                         cache_write_tokens=math.ceil(prefix / 3))
 
 
 @dataclass
@@ -99,6 +122,9 @@ class CallResult:
     model: str | None = None
     latency_s: float = 0.0
     extra: dict[str, Any] = field(default_factory=dict)
+    # input_tokens counts uncached input only; these are the cached prefix read or written.
+    cache_read_input_tokens: int = 0
+    cache_creation_input_tokens: int = 0
 
 
 class Supervisor(Protocol):
@@ -135,12 +161,13 @@ def call(self, step: int, request: dict) -> CallResult:
         usage = resp.usage
         result = CallResult(
             text=None,
-            input_tokens=int(usage.input_tokens or 0) + int(getattr(usage, "cache_read_input_tokens", 0) or 0)
-            + int(getattr(usage, "cache_creation_input_tokens", 0) or 0),
+            input_tokens=int(usage.input_tokens or 0),
             output_tokens=int(usage.output_tokens or 0),
             stop_reason=resp.stop_reason,
             model=resp.model,
             latency_s=latency,
+            cache_read_input_tokens=int(getattr(usage, "cache_read_input_tokens", 0) or 0),
+            cache_creation_input_tokens=int(getattr(usage, "cache_creation_input_tokens", 0) or 0),
         )
         if resp.stop_reason == "refusal":
             result.error = "model declined the request"
@@ -156,6 +183,20 @@ class ReplayMismatch(RuntimeError):
     pass
 
 
+def _recorded_result(reply: dict) -> CallResult:
+    # Runs recorded before prompt caching carry no cache counts; their input_tokens is all input.
+    return CallResult(
+        text=reply.get("text"),
+        input_tokens=reply.get("input_tokens", 0),
+        output_tokens=reply.get("output_tokens", 0),
+        stop_reason=reply.get("stop_reason"),
+        error=reply.get("error"),
+        model=reply.get("model"),
+        cache_read_input_tokens=reply.get("cache_read_input_tokens", 0),
+        cache_creation_input_tokens=reply.get("cache_creation_input_tokens", 0),
+    )
+
+
 class ReplaySupervisor:
     """Serves recorded replies. Fails loudly if the request differs from the recording."""
 
@@ -177,15 +218,7 @@ def call(self, step: int, request: dict) -> CallResult:
         sha = request_sha256(request)
         if sha != rec["request_sha256"]:
             raise ReplayMismatch(f"step {step}: request differs from recording")
-        reply = rec["reply"]
-        return CallResult(
-            text=reply.get("text"),
-            input_tokens=reply.get("input_tokens", 0),
-            output_tokens=reply.get("output_tokens", 0),
-            stop_reason=reply.get("stop_reason"),
-            error=reply.get("error"),
-            model=reply.get("model"),
-        )
+        return _recorded_result(rec["reply"])
 
 
 class RecordedReplySupervisor:
@@ -203,15 +236,7 @@ def call(self, step: int, request: dict) -> CallResult:
         rec = self._inner._by_step.get(step)
         if rec is None:
             return CallResult(None, error="no recorded reply for this hour")
-        reply = rec["reply"]
-        return CallResult(
-            text=reply.get("text"),
-            input_tokens=reply.get("input_tokens", 0),
-            output_tokens=reply.get("output_tokens", 0),
-            stop_reason=reply.get("stop_reason"),
-            error=reply.get("error"),
-            model=reply.get("model"),
-        )
+        return _recorded_result(rec["reply"])
 
 
 @dataclass(frozen=True)

hyperbolic/prompt.py

@@ -1105,12 +1105,25 @@ def user_message(obs: dict) -> str:
     )
 
 
+def system_text(request: dict) -> str:
+    """The system prompt's text, whether it was sent as a plain string or as cached blocks."""
+    system = request["system"]
+    return system if isinstance(system, str) else "".join(b["text"] for b in system)
+
+
 def build_request(obs: dict, cfg: EnvelopeConfig, buildings: tuple[str, ...], model_key: str,
                   max_tokens: int = DEFAULT_MAX_TOKENS, effort: str = DEFAULT_EFFORT,
-                  prompt_version: int = 1) -> dict:
+                  prompt_version: int = 1, prompt_cache: bool = False) -> dict:
+    """prompt_cache marks the system prompt as the cached prefix. It depends only on the envelope
+    config and the building names, so it is the same bytes every hour of a run; everything that
+    changes by the hour is in the user message after it. Off by default so runs recorded before
+    prompt caching rebuild the request they sent."""
     system = {1: system_prompt, 2: system_prompt_v2, 3: system_prompt_v3,
               4: system_prompt_v4, 5: system_prompt_v5, 6: system_prompt_v6,
               7: system_prompt_v7, 8: system_prompt_v8, 9: system_prompt_v9}[prompt_version](cfg, buildings)
+    if prompt_cache:
+        # Default 5 minute TTL: hourly calls run back to back, seconds apart.
+        system = [{"type": "text", "text": system, "cache_control": {"type": "ephemeral"}}]
     return {
         "model": MODELS[model_key]["id"],
         "max_tokens": max_tokens,

hyperbolic/runner.py

@@ -3,8 +3,10 @@
 A run directory holds everything needed to replay it with zero model or simulator calls:
   manifest.json  settings, dataset, envelope config
   steps.jsonl    per hour: district state, actions applied, outcome, decision
-  calls.jsonl    per model call: full request, reply, tokens, cost
-  result.json    CityLearn KPIs and counts
+  calls.jsonl    per model call: full request, reply, tokens, cost. The reply's input_tokens is
+                 uncached input, beside cache_read_input_tokens and cache_creation_input_tokens
+                 (runs from before prompt caching record no cache counts and no cache reads).
+  result.json    CityLearn KPIs and counts, token totals and the cache hit rate
 """
 
 from __future__ import annotations
@@ -25,6 +27,7 @@
     ReplayMismatch,
     ReplaySupervisor,
     Supervisor,
+    cache_hit_rate,
     call_cost_usd,
     request_sha256,
     worst_case_cost_usd,
@@ -170,11 +173,14 @@ class AgentController:
     max_tokens: int = 2000
     effort: str = "low"
     prompt_version: int = PROMPT_VERSION
+    # New runs cache the system prompt; replay passes the manifest's value (absent before caching).
+    prompt_cache: bool = True
     last: Decision | None = None
     lookahead: Lookahead = field(default_factory=Lookahead)
     peaks: PeakTracker | None = None
     mode: str = "agent"
     counts: dict = field(default_factory=lambda: {"calls": 0, "clamped": 0, "fallback": 0, "budget_skipped": 0})
+    tokens: dict = field(default_factory=lambda: {"input": 0, "cache_read": 0, "cache_write": 0, "output": 0})
 
     def describe(self) -> dict:
         return {
@@ -185,6 +191,7 @@ def describe(self) -> dict:
             "max_tokens": self.max_tokens,
             "effort": self.effort,
             "prompt_version": self.prompt_version,
+            "prompt_cache": self.prompt_cache,
         }
 
     def decide(self, step: int, state: dict, writer: RunWriter) -> tuple[Decision, dict | None]:
@@ -198,7 +205,7 @@ def decide(self, step: int, state: dict, writer: RunWriter) -> tuple[Decision, d
         buildings = tuple(state["buildings"])
         obs = observe(self.prompt_version, state, ctx, self.last, self.lookahead, self.envelope)
         request = build_request(obs, self.envelope.config, buildings, self.supervisor.model_key,
-                                self.max_tokens, self.effort, self.prompt_version)
+                                self.max_tokens, self.effort, self.prompt_version, self.prompt_cache)
         worst = worst_case_cost_usd(self.supervisor.model_key, request)
         call_info = None
         if not self.budget.allows(worst):
@@ -211,9 +218,13 @@ def decide(self, step: int, state: dict, writer: RunWriter) -> tuple[Decision, d
                 raise
             except Exception as exc:  # any client failure is a missing reply, never a crash
                 res = CallResult(None, error=f"client error: {type(exc).__name__}")
-            cost = call_cost_usd(self.supervisor.model_key, res.input_tokens, res.output_tokens)
+            cost = call_cost_usd(self.supervisor.model_key, res.input_tokens, res.output_tokens,
+                                 res.cache_read_input_tokens, res.cache_creation_input_tokens)
             self.budget.charge(cost)
             self.counts["calls"] += 1
+            for k, v in (("input", res.input_tokens), ("cache_read", res.cache_read_input_tokens),
+                         ("cache_write", res.cache_creation_input_tokens), ("output", res.output_tokens)):
+                self.tokens[k] += v
             sha = request_sha256(request)
             writer.call({
                 "step": step,
@@ -222,6 +233,8 @@ def decide(self, step: int, state: dict, writer: RunWriter) -> tuple[Decision, d
                 "reply": {
                     "text": res.text,
                     "input_tokens": res.input_tokens,
+                    "cache_read_input_tokens": res.cache_read_input_tokens,
+                    "cache_creation_input_tokens": res.cache_creation_input_tokens,
                     "output_tokens": res.output_tokens,
                     "stop_reason": res.stop_reason,
                     "error": res.error,
@@ -351,6 +364,11 @@ def run_session(
         "budget_stop_reason": budget.stop_reason if budget else None,
         "finished_at": dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds"),
     }
+    tokens = getattr(controller, "tokens", None)
+    if tokens is not None:
+        result["tokens"] = dict(tokens)
+        rate = cache_hit_rate(tokens["input"], tokens["cache_read"], tokens["cache_write"])
+        result["cache_hit_rate"] = None if rate is None else round(rate, 4)
     if stopped:
         result["stopped"] = True
         result["stop_reason"] = f"stopped by the user after {step} of {episode_hours} hours"
@@ -410,6 +428,7 @@ class ReplayReport:
     clamped_steps: int
     fallback_steps: int
     error: str | None = None
+    cache_hit_rate: float | None = None
 
     @property
     def deterministic(self) -> bool:
@@ -432,6 +451,8 @@ def replay_run(run_dir: Path, out_dir: Path | None = None, log=lambda *_: None)
             effort=sup["effort"],
             # Runs recorded before prompt versions existed used version 1.
             prompt_version=int(sup.get("prompt_version", 1)),
+            # Runs recorded before prompt caching sent the system prompt as a plain string.
+            prompt_cache=bool(sup.get("prompt_cache", False)),
         )
     writer: RunWriter = RunWriter(out_dir) if out_dir is not None else NullWriter()
     error = None
@@ -454,6 +475,9 @@ def replay_run(run_dir: Path, out_dir: Path | None = None, log=lambda *_: None)
         mismatched.append(min(len(replayed), len(sim.steps)))
     calls = load_jsonl(run_dir / "calls.jsonl")
     counts = getattr(controller, "counts", {})
+    replies = [c.get("reply") or {} for c in calls]
+    hit_rate = cache_hit_rate(*(sum(r.get(k, 0) or 0 for r in replies) for k in
+                                ("input_tokens", "cache_read_input_tokens", "cache_creation_input_tokens")))
     return ReplayReport(
         steps=len(sim.steps),
         identical_steps=identical,
@@ -464,4 +488,5 @@ def replay_run(run_dir: Path, out_dir: Path | None = None, log=lambda *_: None)
         clamped_steps=counts.get("clamped", 0),
         fallback_steps=counts.get("fallback", 0),
         error=error,
+        cache_hit_rate=hit_rate,
     )

hyperbolic/viewer.py

@@ -197,6 +197,8 @@ def build_scenario(baseline_dir: Path, agent_dir: Path) -> dict:
             "clamped_hours": sum(1 for d in decisions if d["status"] == "clamped"),
             "fallback_hours": sum(1 for d in decisions if d["status"] == "fallback"),
             "budget_stopped": bool(ar.get("budget_stopped")),
+            # Only runs recorded with prompt caching carry it, so older viewer data is unchanged.
+            **({"cache_hit_rate": ar["cache_hit_rate"]} if "cache_hit_rate" in ar else {}),
             "envelope": am.get("envelope"),
             **_trim_end_count(am, decisions),
         },