Uh oh!
There was an error while loading. Please reload this page.
Add Anthropic prompt cache controls#5
Uh oh!
There was an error while loading. Please reload this page.
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,370 @@ | ||
| //! Anthropic Messages API provider with explicit prompt-cache breakpoints. | ||
| //! | ||
| //! Unlike OpenAI-compatible APIs, Anthropic enables prompt caching by attaching | ||
| //! `{"type":"ephemeral"}` as `cache_control` to a system/content block. This | ||
| //! adapter turns TinyInference's cacheable prompt segments into that wire shape | ||
| //! and maps the provider's cache usage counters back into [`Usage`]. | ||
| use async_trait::async_trait; | ||
| use serde_json::{Value, json}; | ||
| use std::time::Duration; | ||
| use crate::message::{AssistantMessage, ContentBlock, Message}; | ||
| use crate::model::{ChatModel, ModelProfile, ModelRequest, ModelResponse}; | ||
| use crate::usage::Usage; | ||
| use crate::{Error, Result}; | ||
| const DEFAULT_BASE_URL: &str = "https://api.anthropic.com/v1"; | ||
| const DEFAULT_MODEL: &str = "claude-sonnet-4-6"; | ||
| const ANTHROPIC_VERSION: &str = "2023-06-01"; | ||
| /// A chat model backed by Anthropic's native Messages API. | ||
| pub struct AnthropicModel { | ||
| client: reqwest::Client, | ||
| api_key: String, | ||
| base_url: String, | ||
| model: String, | ||
| profile: ModelProfile, | ||
| } | ||
| impl std::fmt::Debug for AnthropicModel { | ||
| fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { | ||
| formatter | ||
| .debug_struct("AnthropicModel") | ||
| .field("client", &self.client) | ||
| .field("api_key", &"[redacted]") | ||
| .field("base_url", &self.base_url) | ||
| .field("model", &self.model) | ||
| .field("profile", &self.profile) | ||
| .finish() | ||
| } | ||
| } | ||
| impl AnthropicModel { | ||
| /// Creates an Anthropic model using the default Messages API endpoint. | ||
| pub fn new(api_key: impl Into<String>) -> Self { | ||
| Self::with_base_url(api_key, DEFAULT_BASE_URL) | ||
| } | ||
| /// Creates an Anthropic model targeting a Messages-API-compatible endpoint. | ||
| pub fn with_base_url(api_key: impl Into<String>, base_url: impl Into<String>) -> Self { | ||
| let model = DEFAULT_MODEL.to_string(); | ||
| Self { | ||
| client: reqwest::Client::new(), | ||
| api_key: api_key.into(), | ||
| base_url: base_url.into().trim_end_matches('/').to_string(), | ||
| profile: ModelProfile { | ||
| provider: Some("anthropic".to_string()), | ||
| model: Some(model.clone()), | ||
| ..ModelProfile::default() | ||
| }, | ||
| model, | ||
| } | ||
| } | ||
| /// Overrides the default model id used when a request does not specify one. | ||
| pub fn with_model(mut self, model: impl Into<String>) -> Self { | ||
| self.model = model.into(); | ||
| self.profile.model = Some(self.model.clone()); | ||
| self | ||
| } | ||
| /// Reads `ANTHROPIC_API_KEY`, plus optional `ANTHROPIC_BASE_URL` and | ||
| /// `ANTHROPIC_MODEL`, from the environment. | ||
| /// | ||
| /// # Errors | ||
| /// | ||
| /// Returns [`Error::Model`] when `ANTHROPIC_API_KEY` is not set. | ||
| pub fn from_env() -> Result<Self> { | ||
senamakel marked this conversation as resolved.
Uh oh!There was an error while loading. Please reload this page. | ||
| let key = std::env::var("ANTHROPIC_API_KEY") | ||
| .map_err(|_| Error::Model("ANTHROPIC_API_KEY is not set".to_string()))?; | ||
| let base_url = | ||
| std::env::var("ANTHROPIC_BASE_URL").unwrap_or_else(|_| DEFAULT_BASE_URL.into()); | ||
| let model = std::env::var("ANTHROPIC_MODEL").unwrap_or_else(|_| DEFAULT_MODEL.into()); | ||
| Ok(Self::with_base_url(key, base_url).with_model(model)) | ||
| } | ||
| fn endpoint(&self) -> String { | ||
| if self.base_url.ends_with("/messages") { | ||
| self.base_url.clone() | ||
| } else { | ||
| format!("{}/messages", self.base_url) | ||
| } | ||
| } | ||
| } | ||
| /// Builds the native Messages API body. Kept pure so cache-control placement is | ||
| /// exhaustively testable without a network server. | ||
| pub(crate) fn request_body(request: &ModelRequest, default_model: &str) -> Value { | ||
| let cache_enabled = request | ||
| .cache_policy | ||
| .as_ref() | ||
| .is_some_and(|policy| policy.protect_prompt_prefix) | ||
| && request | ||
| .cache_segments | ||
| .iter() | ||
| .any(|segment| segment.cacheable); | ||
| let mut system = Vec::new(); | ||
| let mut messages = Vec::new(); | ||
| for message in &request.messages { | ||
| match message { | ||
| Message::System(system_message) => { | ||
| system.extend(content_blocks(&system_message.content)) | ||
| } | ||
| Message::User(user_message) => messages.push(json!({ | ||
| "role": "user", | ||
| "content": content_blocks(&user_message.content), | ||
| })), | ||
| Message::Assistant(assistant_message) => { | ||
| let mut content = content_blocks(&assistant_message.content); | ||
| content.extend(assistant_message.tool_calls.iter().map(|call| { | ||
| json!({ | ||
| "type": "tool_use", | ||
| "id": call.id, | ||
| "name": call.name, | ||
| "input": call.arguments, | ||
There was a problem hiding this comment. Serialize tool-call arguments as a JSON object for Anthropic Anthropic's Messages API expects [RULE] wrong-argument-type-for-wire · | ||
| }) | ||
| })); | ||
| messages.push(json!({ "role": "assistant", "content": content })); | ||
| } | ||
| Message::Tool(tool_message) => messages.push(json!({ | ||
| "role": "user", | ||
| "content": [{ | ||
| "type": "tool_result", | ||
| "tool_use_id": tool_message.tool_call_id, | ||
| "content": content_blocks(&tool_message.content), | ||
| }], | ||
| })), | ||
| } | ||
| } | ||
| if cache_enabled { | ||
| let block = system.last_mut().or_else(|| { | ||
| messages | ||
| .first_mut() | ||
| .and_then(|message| message["content"].as_array_mut()) | ||
| .and_then(|blocks| blocks.first_mut()) | ||
| }); | ||
| if let Some(block) = block { | ||
| block["cache_control"] = json!({ "type": "ephemeral" }); | ||
| } | ||
| } | ||
| let mut body = json!({ | ||
| "model": request.model.as_deref().unwrap_or(default_model), | ||
| "max_tokens": request.max_tokens.unwrap_or(1024), | ||
| "messages": messages, | ||
| }); | ||
senamakel marked this conversation as resolved.
Uh oh!There was an error while loading. Please reload this page. | ||
| if !system.is_empty() { | ||
| body["system"] = Value::Array(system); | ||
| } | ||
| if let Some(temperature) = request.temperature { | ||
| body["temperature"] = json!(temperature); | ||
| } | ||
| if let Some(top_p) = request.top_p { | ||
| body["top_p"] = json!(top_p); | ||
| } | ||
| if !request.stop_sequences.is_empty() { | ||
| body["stop_sequences"] = json!(request.stop_sequences); | ||
| } | ||
| body | ||
| } | ||
| fn content_blocks(content: &[ContentBlock]) -> Vec<Value> { | ||
| content | ||
| .iter() | ||
| .filter_map(|block| match block { | ||
| ContentBlock::Text(text) => Some(json!({ "type": "text", "text": text })), | ||
| ContentBlock::Json(value) => Some(json!({ "type": "text", "text": value.to_string() })), | ||
| ContentBlock::Thinking { text, .. } => Some(json!({ "type": "text", "text": text })), | ||
| ContentBlock::RedactedThinking { data } => { | ||
| Some(json!({ "type": "text", "text": data })) | ||
| } | ||
| ContentBlock::Image(_) | ContentBlock::ProviderExtension(_) => None, | ||
| }) | ||
| .collect() | ||
| } | ||
| fn parse_response(body: Value) -> Result<ModelResponse> { | ||
| let text = body["content"] | ||
| .as_array() | ||
| .into_iter() | ||
| .flatten() | ||
| .filter_map(|block| { | ||
| (block["type"].as_str() == Some("text")) | ||
| .then(|| block["text"].as_str()) | ||
| .flatten() | ||
| }) | ||
| .collect::<String>(); | ||
| let usage = body.get("usage").map(|usage| { | ||
| let uncached_input_tokens = usage["input_tokens"].as_u64().unwrap_or(0); | ||
| let cache_read_tokens = usage["cache_read_input_tokens"].as_u64().unwrap_or(0); | ||
| let cache_creation_tokens = usage["cache_creation_input_tokens"].as_u64().unwrap_or(0); | ||
| let input_tokens = uncached_input_tokens + cache_read_tokens + cache_creation_tokens; | ||
| let output_tokens = usage["output_tokens"].as_u64().unwrap_or(0); | ||
| Usage { | ||
| input_tokens, | ||
| output_tokens, | ||
| total_tokens: input_tokens + output_tokens, | ||
| cache_read_tokens, | ||
| cache_creation_tokens, | ||
| ..Usage::default() | ||
| } | ||
| }); | ||
| Ok(ModelResponse { | ||
There was a problem hiding this comment. Parse tool_use blocks from Anthropic responses
[RULE] dropped-tool-calls · | ||
| message: AssistantMessage { | ||
| id: body["id"].as_str().map(str::to_string), | ||
| content: vec![ContentBlock::Text(text)], | ||
| tool_calls: Vec::new(), | ||
| usage, | ||
| }, | ||
| usage, | ||
| finish_reason: body["stop_reason"].as_str().map(str::to_string), | ||
| raw: Some(body), | ||
| resolved_model: None, | ||
| continue_turn: None, | ||
| served_from_cache: false, | ||
| }) | ||
| } | ||
| #[async_trait] | ||
| impl<State: Send + Sync> ChatModel<State> for AnthropicModel { | ||
| fn profile(&self) -> Option<&ModelProfile> { | ||
| Some(&self.profile) | ||
| } | ||
| fn cache_identity(&self) -> Option<String> { | ||
| Some(format!("anthropic:{}:{}", self.base_url, self.model)) | ||
| } | ||
| async fn invoke(&self, _state: &State, request: ModelRequest) -> Result<ModelResponse> { | ||
| let request_builder = self | ||
| .client | ||
| .post(self.endpoint()) | ||
| .header("x-api-key", &self.api_key) | ||
| .header("anthropic-version", ANTHROPIC_VERSION) | ||
| .json(&request_body(&request, &self.model)); | ||
| let request_builder = match request.timeout_ms { | ||
| Some(timeout_ms) => request_builder.timeout(Duration::from_millis(timeout_ms)), | ||
| None => request_builder, | ||
| }; | ||
| let response = request_builder | ||
| .send() | ||
senamakel marked this conversation as resolved.
Uh oh!There was an error while loading. Please reload this page. | ||
| .await | ||
| .map_err(|error| Error::Model(format!("anthropic request failed: {error}")))?; | ||
| let status = response.status(); | ||
| let body: Value = response | ||
| .json() | ||
| .await | ||
| .map_err(|error| Error::Model(format!("anthropic response was not JSON: {error}")))?; | ||
| if !status.is_success() { | ||
| return Err(Error::Model(format!( | ||
| "anthropic returned HTTP {status}: {}", | ||
| body["error"]["message"].as_str().unwrap_or("unknown error") | ||
| ))); | ||
Comment on lines
+259
to
+262
There was a problem hiding this comment.
When Anthropic returns a routine 429 or transient 5xx response, this collapses the failure into AGENTS.md reference: AGENTS.md:L38-L40 Useful? React with 👍 / 👎. | ||
| } | ||
| parse_response(body) | ||
| } | ||
| } | ||
| #[cfg(test)] | ||
| mod test { | ||
| use super::*; | ||
| use crate::cache::CachePolicy; | ||
| use crate::model::{PromptSegment, SegmentRole}; | ||
| #[test] | ||
| fn cacheable_system_prefix_becomes_an_anthropic_cache_breakpoint() { | ||
| let request = ModelRequest::new(vec![ | ||
| Message::system("stable instructions"), | ||
| Message::user("hello"), | ||
| ]) | ||
| .with_cache_segments(vec![PromptSegment { | ||
| id: "system".into(), | ||
| role: SegmentRole::System, | ||
| cacheable: true, | ||
| }]) | ||
| .with_cache_policy(CachePolicy { | ||
| protect_prompt_prefix: true, | ||
| ..CachePolicy::default() | ||
| }); | ||
| let body = request_body(&request, "test-model"); | ||
| assert_eq!( | ||
| body["system"][0]["cache_control"], | ||
| json!({ "type": "ephemeral" }) | ||
| ); | ||
| assert_eq!(body["messages"][0]["role"], "user"); | ||
| } | ||
| #[test] | ||
| fn cacheable_user_prefix_becomes_a_content_block_breakpoint() { | ||
| let request = ModelRequest::new(vec![Message::user("stable context")]) | ||
| .with_cache_segments(vec![PromptSegment { | ||
| id: "history".into(), | ||
| role: SegmentRole::History, | ||
| cacheable: true, | ||
| }]) | ||
| .with_cache_policy(CachePolicy { | ||
| protect_prompt_prefix: true, | ||
| ..CachePolicy::default() | ||
| }); | ||
| let body = request_body(&request, "test-model"); | ||
| assert_eq!(body["messages"][0]["content"][0]["type"], "text"); | ||
| assert_eq!(body["messages"][0]["content"][0]["text"], "stable context"); | ||
| assert_eq!( | ||
| body["messages"][0]["content"][0]["cache_control"], | ||
| json!({ "type": "ephemeral" }) | ||
| ); | ||
| } | ||
| #[test] | ||
| fn tool_results_use_anthropic_tool_result_blocks() { | ||
| let request = ModelRequest::new(vec![Message::Tool(crate::message::ToolMessage { | ||
| tool_call_id: "tool_1".into(), | ||
| content: vec![ContentBlock::Text("42".into())], | ||
| trusted_verbatim: false, | ||
| artifact: None, | ||
| })]); | ||
| let body = request_body(&request, "test-model"); | ||
| assert_eq!(body["messages"][0]["role"], "user"); | ||
| assert_eq!(body["messages"][0]["content"][0]["type"], "tool_result"); | ||
| assert_eq!(body["messages"][0]["content"][0]["tool_use_id"], "tool_1"); | ||
| } | ||
| #[test] | ||
| fn default_model_is_current() { | ||
| assert_eq!(AnthropicModel::new("key").model, "claude-sonnet-4-6"); | ||
| } | ||
| #[test] | ||
| fn cache_usage_is_mapped_from_anthropic_response() { | ||
| let response = parse_response(json!({ | ||
| "id": "msg_1", "content": [{ "type": "text", "text": "hello" }], "stop_reason": "end_turn", | ||
| "usage": { "input_tokens": 100, "output_tokens": 5, "cache_read_input_tokens": 90, "cache_creation_input_tokens": 10 } | ||
| })).unwrap(); | ||
| assert_eq!(response.text(), "hello"); | ||
| let usage = response.usage.unwrap(); | ||
| assert_eq!(usage.input_tokens, 200); | ||
| assert_eq!(usage.total_tokens, 205); | ||
| assert_eq!(usage.cache_read_tokens, 90); | ||
| assert_eq!(usage.cache_creation_tokens, 10); | ||
| } | ||
| #[test] | ||
| fn request_body_forwards_generation_controls() { | ||
| let request = ModelRequest::new(vec![Message::user("hello")]) | ||
| .with_temperature(0.2) | ||
| .with_top_p(0.8) | ||
| .with_stop_sequences(["END"]); | ||
| let body = request_body(&request, "test-model"); | ||
| assert_eq!(body["temperature"], 0.2); | ||
| assert_eq!(body["top_p"], 0.8); | ||
| assert_eq!(body["stop_sequences"], json!(["END"])); | ||
| } | ||
| #[test] | ||
| fn debug_redacts_the_api_key() { | ||
| let model = AnthropicModel::new("secret-api-key"); | ||
| let debug = format!("{model:?}"); | ||
| assert!(debug.contains("[redacted]")); | ||
| assert!(!debug.contains("secret-api-key")); | ||
| } | ||
| } | ||
Oops, something went wrong.
Uh oh!
There was an error while loading. Please reload this page.
Oops, something went wrong.
Uh oh!
There was an error while loading. Please reload this page.
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.