Prompt caching — 10× cheaper repeated calls

from pathlib import Path
import sys

ROOT = (
    Path.cwd().resolve().parent
    if Path.cwd().name == "notebooks"
    else Path.cwd().resolve()
)
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

import os
from shipit_agent.llms import AnthropicChatLLM
from shipit_agent.models import Message

llm = AnthropicChatLLM(
    model="claude-opus-4-1",
    api_key="offline-demo-key",   # not used unless we actually call .complete()
    prompt_caching=True,          # default on for Anthropic-family adapters
)
llm.prompt_caching

import json
from shipit_agent import FunctionTool


def get_weather(city: str) -> str:
    """Return the weather for a city."""
    return "sunny"


def search(query: str) -> str:
    """Search the web."""
    return "results"


tools = [
    FunctionTool.from_callable(get_weather).schema(),
    FunctionTool.from_callable(search).schema(),
]

req = llm._build_request_kwargs(
    messages=[Message(role="user", content="What's the weather in Paris?")],
    tools=tools,
    system_prompt="You are a helpful weather assistant.",
)

print("system block carries cache_control:")
print(json.dumps(req["system"], indent=2))
print("\nlast tool carries cache_control:", "cache_control" in req["tools"][-1])
print("first tool does NOT:", "cache_control" not in req["tools"][0])

llm_off = AnthropicChatLLM(model="claude-opus-4-1", api_key="x", prompt_caching=False)
req_off = llm_off._build_request_kwargs(
    messages=[Message(role="user", content="hi")],
    tools=tools,
    system_prompt="You are a helpful weather assistant.",
)
print("system is a plain string:", isinstance(req_off["system"], str))
print("no tool marked:", all("cache_control" not in t for t in req_off["tools"]))

from shipit_agent.llms import LiteLLMChatLLM
from shipit_agent.llms.litellm_adapter import _apply_prompt_caching

payload_messages = [
    {"role": "system", "content": "You are a helpful weather assistant."},
    {"role": "user", "content": "What's the weather in Paris?"},
]
marked_messages, marked_tools = _apply_prompt_caching(payload_messages, tools)

print("system promoted to cached block form:")
print(json.dumps(marked_messages[0], indent=2))
print("\nlast tool marked:", "cache_control" in marked_tools[-1])

lite = LiteLLMChatLLM(model="anthropic/claude-opus-4-1", prompt_caching=True)
print("LiteLLM prompt_caching:", lite.prompt_caching)

from shipit_agent import CostTracker

# Two calls: call #1 writes the prefix into the cache; call #2 reads it back cheap.
first_call_usage = {
    "input_tokens": 1200,
    "output_tokens": 80,
    "cache_creation_input_tokens": 1100,   # stable prefix written to cache
    "cache_read_input_tokens": 0,
}
second_call_usage = {
    "input_tokens": 120,                   # only the new user turn is full price
    "output_tokens": 80,
    "cache_creation_input_tokens": 0,
    "cache_read_input_tokens": 1100,       # prefix served from cache (~10% cost)
}

tracker = CostTracker()
print("CostTracker reads these keys via as_hooks():")
print("  cache_read_input_tokens  ->", second_call_usage["cache_read_input_tokens"])
print("  cache_creation_input_tokens ->", first_call_usage["cache_creation_input_tokens"])
print("\nWiring: Agent(llm=llm, hooks=tracker.as_hooks()) — cache reads are billed cheaper.")

if os.getenv("ANTHROPIC_API_KEY"):
    from shipit_agent import Agent

    live = AnthropicChatLLM(model="claude-opus-4-1", prompt_caching=True)
    agent = Agent(llm=live, tools=[FunctionTool.from_callable(get_weather)])
    r1 = agent.run("Weather in Paris?")   # cache creation
    r2 = agent.run("Weather in Lyon?")    # cache read on the shared prefix
    print("run 1 usage:", r1.usage)
    print("run 2 usage:", r2.usage)
else:
    print("Set ANTHROPIC_API_KEY to run the live cached call. Skipping (offline).")