From c4cfd885935443b7961a25f9ce13782eeb4f9797 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Mon, 31 Aug 2026 19:57:31 +0300 Subject: [PATCH 1/7] feat: add Anthropic prompt cache controls Co-authored-by: Medulla --- .../tinyinference/src/providers/anthropic.rs | 242 ++++++++++++++++++ crates/tinyinference/src/providers/mod.rs | 7 +- 2 files changed, 246 insertions(+), 3 deletions(-) create mode 100644 crates/tinyinference/src/providers/anthropic.rs diff --git a/crates/tinyinference/src/providers/anthropic.rs b/crates/tinyinference/src/providers/anthropic.rs new file mode 100644 index 0000000..6b6ee87 --- /dev/null +++ b/crates/tinyinference/src/providers/anthropic.rs @@ -0,0 +1,242 @@ +//! 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 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-20250514"; +const ANTHROPIC_VERSION: &str = "2023-06-01"; + +/// A chat model backed by Anthropic's native Messages API. +#[derive(Debug)] +pub struct AnthropicModel { + client: reqwest::Client, + api_key: String, + base_url: String, + model: String, + profile: ModelProfile, +} + +impl AnthropicModel { + /// Creates an Anthropic model using the default Messages API endpoint. + pub fn new(api_key: impl Into) -> 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, base_url: impl Into) -> 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()), + tool_calling: true, + streaming: true, + ..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) -> 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. + pub fn from_env() -> Result { + 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 { + let role = match message { + Message::System(system_message) => { + system.push(text_block( + system_message + .content + .iter() + .filter_map(ContentBlock::as_text) + .collect(), + )); + continue; + } + Message::User(_) | Message::Tool(_) => "user", + Message::Assistant(_) => "assistant", + }; + messages.push(json!({ "role": role, "content": message.text() })); + } + if cache_enabled { + if let Some(block) = system.last_mut() { + block["cache_control"] = json!({ "type": "ephemeral" }); + } else if let Some(message) = messages.first_mut() { + message["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, + }); + if !system.is_empty() { + body["system"] = Value::Array(system); + } + body +} + +fn text_block(text: String) -> Value { + json!({ "type": "text", "text": text }) +} + +fn parse_response(body: Value) -> Result { + 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::(); + let usage = body.get("usage").map(|usage| Usage { + input_tokens: usage["input_tokens"].as_u64().unwrap_or(0), + output_tokens: usage["output_tokens"].as_u64().unwrap_or(0), + cache_read_tokens: usage["cache_read_input_tokens"].as_u64().unwrap_or(0), + cache_creation_tokens: usage["cache_creation_input_tokens"].as_u64().unwrap_or(0), + ..Usage::default() + }); + Ok(ModelResponse { + 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 ChatModel for AnthropicModel { + fn profile(&self) -> Option<&ModelProfile> { + Some(&self.profile) + } + + fn cache_identity(&self) -> Option { + Some(format!("anthropic:{}:{}", self.base_url, self.model)) + } + + async fn invoke(&self, _state: &State, request: ModelRequest) -> Result { + let response = self + .client + .post(self.endpoint()) + .header("x-api-key", &self.api_key) + .header("anthropic-version", ANTHROPIC_VERSION) + .json(&request_body(&request, &self.model)) + .send() + .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") + ))); + } + 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 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.cache_read_tokens, 90); + assert_eq!(usage.cache_creation_tokens, 10); + } +} diff --git a/crates/tinyinference/src/providers/mod.rs b/crates/tinyinference/src/providers/mod.rs index 5e37487..a0e8f48 100644 --- a/crates/tinyinference/src/providers/mod.rs +++ b/crates/tinyinference/src/providers/mod.rs @@ -14,12 +14,13 @@ //! |---|---| //! | [`MockModel`] | Implemented — deterministic, no network | //! | [`openai`] (and OpenAI-compatible endpoints) | Implemented | +//! | [`anthropic`] (Messages API, including prompt caching) | Implemented | //! //! [`MockModel`] is always compiled and needs no network, keeping the default //! build offline and deterministic. The [`openai`] module is always compiled //! too (it pulls no extra dependencies) and additionally serves every //! OpenAI-compatible endpoint (Ollama, DeepSeek, Groq, xAI, OpenRouter, -//! Together, Mistral, and Anthropic's OpenAI-compat endpoint) through the same +//! Together, and Mistral) through the same //! Chat Completions wire format. The default build stays offline anyway: the //! adapter only touches the network when invoked, and the live tests //! early-return without `OPENAI_API_KEY`. @@ -29,7 +30,7 @@ //! //! ```text //! pub mod openai; // always compiled -//! // #[cfg(feature = "anthropic")] pub mod anthropic; +//! pub mod anthropic; // always compiled //! // #[cfg(feature = "ollama")] pub mod ollama; //! ``` @@ -39,8 +40,8 @@ mod types; // The OpenAI Chat Completions adapter is always compiled; it also serves every // OpenAI-compatible endpoint. Providers with a different wire protocol would be // added behind their own Cargo feature. +pub mod anthropic; pub mod openai; -// #[cfg(feature = "anthropic")] pub mod anthropic; // #[cfg(feature = "ollama")] pub mod ollama; pub use types::*; From 767aea422f8094147375e05eac4b6bf8101ae088 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Mon, 31 Aug 2026 19:59:07 +0300 Subject: [PATCH 2/7] fix: advertise supported Anthropic capabilities Co-authored-by: Medulla --- crates/tinyinference/src/providers/anthropic.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/crates/tinyinference/src/providers/anthropic.rs b/crates/tinyinference/src/providers/anthropic.rs index 6b6ee87..34a683b 100644 --- a/crates/tinyinference/src/providers/anthropic.rs +++ b/crates/tinyinference/src/providers/anthropic.rs @@ -43,8 +43,6 @@ impl AnthropicModel { profile: ModelProfile { provider: Some("anthropic".to_string()), model: Some(model.clone()), - tool_calling: true, - streaming: true, ..ModelProfile::default() }, model, From 68042a0e9e3c4d8cb558372e3e6c4879160ef5fd Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Mon, 31 Aug 2026 20:41:21 +0300 Subject: [PATCH 3/7] fix: redact Anthropic credentials in debug output Co-authored-by: Medulla --- .../tinyinference/src/providers/anthropic.rs | 22 ++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/crates/tinyinference/src/providers/anthropic.rs b/crates/tinyinference/src/providers/anthropic.rs index 34a683b..2b68400 100644 --- a/crates/tinyinference/src/providers/anthropic.rs +++ b/crates/tinyinference/src/providers/anthropic.rs @@ -18,7 +18,6 @@ const DEFAULT_MODEL: &str = "claude-sonnet-4-20250514"; const ANTHROPIC_VERSION: &str = "2023-06-01"; /// A chat model backed by Anthropic's native Messages API. -#[derive(Debug)] pub struct AnthropicModel { client: reqwest::Client, api_key: String, @@ -27,6 +26,19 @@ pub struct AnthropicModel { 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) -> Self { @@ -237,4 +249,12 @@ mod test { assert_eq!(usage.cache_read_tokens, 90); assert_eq!(usage.cache_creation_tokens, 10); } + + #[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")); + } } From 1010801209c5ab92f6724744e95e0759466ca843 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Mon, 31 Aug 2026 20:43:05 +0300 Subject: [PATCH 4/7] fix: correct Anthropic cache breakpoint fallback Co-authored-by: Medulla --- .../tinyinference/src/providers/anthropic.rs | 33 +++++++++++++++++-- 1 file changed, 31 insertions(+), 2 deletions(-) diff --git a/crates/tinyinference/src/providers/anthropic.rs b/crates/tinyinference/src/providers/anthropic.rs index 2b68400..f6abbd5 100644 --- a/crates/tinyinference/src/providers/anthropic.rs +++ b/crates/tinyinference/src/providers/anthropic.rs @@ -14,7 +14,7 @@ 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-20250514"; +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. @@ -122,7 +122,11 @@ pub(crate) fn request_body(request: &ModelRequest, default_model: &str) -> Value if let Some(block) = system.last_mut() { block["cache_control"] = json!({ "type": "ephemeral" }); } else if let Some(message) = messages.first_mut() { - message["cache_control"] = json!({ "type": "ephemeral" }); + message["content"] = json!([{ + "type": "text", + "text": message["content"].as_str().unwrap_or_default(), + "cache_control": { "type": "ephemeral" }, + }]); } } let mut body = json!({ @@ -238,6 +242,31 @@ mod test { 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]["cache_control"], + json!({ "type": "ephemeral" }) + ); + } + + #[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!({ From 83b70f8232b11f6111ace0517b73eff01df7eb0d Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Mon, 31 Aug 2026 20:46:37 +0300 Subject: [PATCH 5/7] fix: preserve Anthropic message content blocks Co-authored-by: Medulla --- .../tinyinference/src/providers/anthropic.rs | 90 ++++++++++++++----- 1 file changed, 67 insertions(+), 23 deletions(-) diff --git a/crates/tinyinference/src/providers/anthropic.rs b/crates/tinyinference/src/providers/anthropic.rs index f6abbd5..2abb43f 100644 --- a/crates/tinyinference/src/providers/anthropic.rs +++ b/crates/tinyinference/src/providers/anthropic.rs @@ -102,31 +102,49 @@ pub(crate) fn request_body(request: &ModelRequest, default_model: &str) -> Value let mut system = Vec::new(); let mut messages = Vec::new(); for message in &request.messages { - let role = match message { + match message { Message::System(system_message) => { - system.push(text_block( - system_message - .content - .iter() - .filter_map(ContentBlock::as_text) - .collect(), - )); - continue; + system.push(Value::Array(content_blocks(&system_message.content))) } - Message::User(_) | Message::Tool(_) => "user", - Message::Assistant(_) => "assistant", - }; - messages.push(json!({ "role": role, "content": message.text() })); + 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, + }) + })); + 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 { - if let Some(block) = system.last_mut() { + let block = system + .last_mut() + .and_then(Value::as_array_mut) + .and_then(|blocks| blocks.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" }); - } else if let Some(message) = messages.first_mut() { - message["content"] = json!([{ - "type": "text", - "text": message["content"].as_str().unwrap_or_default(), - "cache_control": { "type": "ephemeral" }, - }]); } } let mut body = json!({ @@ -140,8 +158,19 @@ pub(crate) fn request_body(request: &ModelRequest, default_model: &str) -> Value body } -fn text_block(text: String) -> Value { - json!({ "type": "text", "text": text }) +fn content_blocks(content: &[ContentBlock]) -> Vec { + 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 { @@ -236,7 +265,7 @@ mod test { }); let body = request_body(&request, "test-model"); assert_eq!( - body["system"][0]["cache_control"], + body["system"][0][0]["cache_control"], json!({ "type": "ephemeral" }) ); assert_eq!(body["messages"][0]["role"], "user"); @@ -256,12 +285,27 @@ mod test { }); 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"); From ab2144f54e981ff3c22cc9eaa96a901471d04f38 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Mon, 31 Aug 2026 20:48:12 +0300 Subject: [PATCH 6/7] fix: preserve Anthropic system cache blocks Co-authored-by: Medulla --- .../tinyinference/src/providers/anthropic.rs | 20 ++++++++----------- 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/crates/tinyinference/src/providers/anthropic.rs b/crates/tinyinference/src/providers/anthropic.rs index 2abb43f..ec10c98 100644 --- a/crates/tinyinference/src/providers/anthropic.rs +++ b/crates/tinyinference/src/providers/anthropic.rs @@ -104,7 +104,7 @@ pub(crate) fn request_body(request: &ModelRequest, default_model: &str) -> Value for message in &request.messages { match message { Message::System(system_message) => { - system.push(Value::Array(content_blocks(&system_message.content))) + system.extend(content_blocks(&system_message.content)) } Message::User(user_message) => messages.push(json!({ "role": "user", @@ -133,16 +133,12 @@ pub(crate) fn request_body(request: &ModelRequest, default_model: &str) -> Value } } if cache_enabled { - let block = system - .last_mut() - .and_then(Value::as_array_mut) - .and_then(|blocks| blocks.last_mut()) - .or_else(|| { - messages - .first_mut() - .and_then(|message| message["content"].as_array_mut()) - .and_then(|blocks| blocks.first_mut()) - }); + 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" }); } @@ -265,7 +261,7 @@ mod test { }); let body = request_body(&request, "test-model"); assert_eq!( - body["system"][0][0]["cache_control"], + body["system"][0]["cache_control"], json!({ "type": "ephemeral" }) ); assert_eq!(body["messages"][0]["role"], "user"); From d2e377ae19a6d27785927f4564526653ae8a8906 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Mon, 31 Aug 2026 20:50:22 +0300 Subject: [PATCH 7/7] fix: honor Anthropic request and usage contracts Co-authored-by: Medulla --- .../tinyinference/src/providers/anthropic.rs | 57 ++++++++++++++++--- 1 file changed, 49 insertions(+), 8 deletions(-) diff --git a/crates/tinyinference/src/providers/anthropic.rs b/crates/tinyinference/src/providers/anthropic.rs index ec10c98..092ec47 100644 --- a/crates/tinyinference/src/providers/anthropic.rs +++ b/crates/tinyinference/src/providers/anthropic.rs @@ -7,6 +7,7 @@ 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}; @@ -70,6 +71,10 @@ impl AnthropicModel { /// 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 { let key = std::env::var("ANTHROPIC_API_KEY") .map_err(|_| Error::Model("ANTHROPIC_API_KEY is not set".to_string()))?; @@ -151,6 +156,15 @@ pub(crate) fn request_body(request: &ModelRequest, default_model: &str) -> Value 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 } @@ -180,12 +194,20 @@ fn parse_response(body: Value) -> Result { .flatten() }) .collect::(); - let usage = body.get("usage").map(|usage| Usage { - input_tokens: usage["input_tokens"].as_u64().unwrap_or(0), - output_tokens: usage["output_tokens"].as_u64().unwrap_or(0), - cache_read_tokens: usage["cache_read_input_tokens"].as_u64().unwrap_or(0), - cache_creation_tokens: usage["cache_creation_input_tokens"].as_u64().unwrap_or(0), - ..Usage::default() + 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 { message: AssistantMessage { @@ -214,12 +236,17 @@ impl ChatModel for AnthropicModel { } async fn invoke(&self, _state: &State, request: ModelRequest) -> Result { - let response = self + 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)) + .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() .await .map_err(|error| Error::Model(format!("anthropic request failed: {error}")))?; @@ -315,10 +342,24 @@ mod test { })).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");