diff --git a/src/providers.rs b/src/providers.rs index 16c778f..e8e3dec 100644 --- a/src/providers.rs +++ b/src/providers.rs @@ -1028,23 +1028,65 @@ pub fn anthropic_messages(req: &ChatRequest) -> Vec { messages } +/// Prompt caching: a breakpoint on the last block of the last message, so the whole history +/// prefix is cached too — history is append-only, so each step of a turn reads what the last +/// step wrote. Anthropic accepts at most four breakpoints; we use three (system, tools, tail). +fn cache_tail(mut messages: Vec) -> Vec { + if let Some(last) = messages.last_mut() { + match last.get_mut("content") { + Some(Value::Array(blocks)) => { + if let Some(b) = blocks.last_mut() { + if b.is_object() { + b["cache_control"] = json!({"type": "ephemeral"}); + } + } + } + Some(Value::String(text)) => { + let t = std::mem::take(text); + last["content"] = json!([{"type": "text", "text": t, "cache_control": {"type": "ephemeral"}}]); + } + _ => {} + } + } + messages +} + +/// (input_tokens, cache_read_input_tokens, cache_creation_input_tokens) summed over the +/// session's Claude answers: the cache hit rate is the number that tells whether the +/// breakpoints work (see `anybot` diagnostics / make doctor). +static CACHE_STATS: std::sync::Mutex<(u64, u64, u64)> = std::sync::Mutex::new((0, 0, 0)); + +pub fn cache_stats() -> (u64, u64, u64) { + *CACHE_STATS.lock().unwrap_or_else(|e| e.into_inner()) +} + fn stream_anthropic(req: &ChatRequest, cancel: &AtomicBool, send: &dyn Fn(StreamEvent)) -> Result { let mut body = json!({ "model": req.model, "max_tokens": MAX_OUTPUT_TOKENS_ANTHROPIC, - "messages": anthropic_messages(req), + "messages": [], "stream": true, }); let system = system_with_summaries(req); if !system.trim().is_empty() { - body["system"] = json!(system); + // Prompt caching: system + tool schemas are most of every request and stay identical + // between the steps of a turn, so both get an ephemeral breakpoint (Anthropic caches + // the prefix up to each marked block; the cache is written on the first step of a turn + // and read on the rest, which is where a computer-use turn's tokens live). + body["system"] = json!([{"type": "text", "text": system, "cache_control": {"type": "ephemeral"}}]); } if claude_adaptive(&req.model) { body["thinking"] = json!({"type": "adaptive", "display": "summarized"}); } if !req.tools.is_empty() { - body["tools"] = json!(req.tools.iter().map(|t| json!({"name": t["name"], "description": t["description"], "input_schema": t["parameters"]})).collect::>()); + let mut tools: Vec = req.tools.iter().map(|t| json!({"name": t["name"], "description": t["description"], "input_schema": t["parameters"]})).collect(); + // One breakpoint on the last tool marks the whole (stable, order-fixed) tools block. + if let Some(last) = tools.last_mut() { + last["cache_control"] = json!({"type": "ephemeral"}); + } + body["tools"] = json!(tools); } + body["messages"] = json!(cache_tail(anthropic_messages(req))); let mut r = agent() .post(&format!("{}/messages", req.base_url)) .header("Content-Type", "application/json") @@ -1132,6 +1174,15 @@ fn stream_anthropic(req: &ChatRequest, cancel: &AtomicBool, send: &dyn Fn(Stream } } } + "message_start" => { + if let Some(u) = v["message"]["usage"].as_object() { + let g = |k: &str| u.get(k).and_then(|x| x.as_u64()).unwrap_or(0); + let mut s = CACHE_STATS.lock().unwrap_or_else(|e| e.into_inner()); + s.0 += g("input_tokens"); + s.1 += g("cache_read_input_tokens"); + s.2 += g("cache_creation_input_tokens"); + } + } "message_delta" => { if v["delta"]["stop_reason"] == "refusal" { out.text.push_str("\n\n(模型拒絕回答這個要求。)"); @@ -1654,4 +1705,18 @@ mod tests { assert_eq!(msgs[i + 1]["role"], "tool"); assert_eq!(msgs[i + 2]["role"], "tool"); } + + #[test] + fn cache_breakpoints_mark_the_stable_prefix() { + // cache_tail marks the last block of the last message, in both content shapes. + let m = cache_tail(vec![json!({"role": "user", "content": "hi"})]); + assert_eq!(m[0]["role"], "user"); + assert_eq!(m[0]["content"][0]["cache_control"]["type"], "ephemeral"); + let m = cache_tail(vec![json!({"role": "user", "content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}]})]); + assert_eq!(m[0]["content"][1]["cache_control"]["type"], "ephemeral"); + assert!(m[0]["content"][0]["cache_control"].is_null()); + assert_eq!(cache_tail(vec![]).len(), 0); + let m = cache_tail(vec![json!({"role": "user", "content": []})]); + assert!(m[0]["content"].as_array().unwrap().is_empty()); + } }