fix(providers): stop paying the prompt-cache write premium on one-shot fast-model calls (#11179)
This commit is contained in:
@@ -52,6 +52,7 @@ pub struct AnthropicFormatOptions {
|
||||
pub preserve_thinking_context: bool,
|
||||
pub thinking_disabled: bool,
|
||||
pub current_model: Option<String>,
|
||||
pub prompt_cache_disabled: bool,
|
||||
}
|
||||
|
||||
impl AnthropicFormatOptions {
|
||||
@@ -73,6 +74,7 @@ impl AnthropicFormatOptions {
|
||||
current_model: self
|
||||
.current_model
|
||||
.or_else(|| Some(model_config.model_name.clone())),
|
||||
prompt_cache_disabled: model_config.prompt_cache_disabled(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -427,6 +429,10 @@ fn format_messages_with_options(
|
||||
}));
|
||||
}
|
||||
|
||||
if options.prompt_cache_disabled {
|
||||
return anthropic_messages;
|
||||
}
|
||||
|
||||
// The last two user messages extend the cached prefix each turn.
|
||||
let mut user_count = 0;
|
||||
for message in anthropic_messages.iter_mut().rev() {
|
||||
@@ -463,7 +469,7 @@ fn anthropic_flavored_input_schema(input_schema: Arc<JsonObject>) -> Arc<JsonObj
|
||||
}
|
||||
|
||||
/// Convert internal Tool format to Anthropic's API tool specification
|
||||
pub fn format_tools(tools: &[Tool]) -> Vec<Value> {
|
||||
pub fn format_tools(tools: &[Tool], options: &AnthropicFormatOptions) -> Vec<Value> {
|
||||
let mut unique_tools = HashSet::new();
|
||||
let mut tool_specs = Vec::new();
|
||||
|
||||
@@ -477,6 +483,10 @@ pub fn format_tools(tools: &[Tool]) -> Vec<Value> {
|
||||
}
|
||||
}
|
||||
|
||||
if options.prompt_cache_disabled {
|
||||
return tool_specs;
|
||||
}
|
||||
|
||||
// Add "cache_control" to the last tool spec, if any. This means that all tool definitions,
|
||||
// will be cached as a single prefix.
|
||||
if let Some(last_tool) = tool_specs.last_mut() {
|
||||
@@ -490,7 +500,13 @@ pub fn format_tools(tools: &[Tool]) -> Vec<Value> {
|
||||
}
|
||||
|
||||
/// Convert system message to Anthropic's API system specification
|
||||
pub fn format_system(system: &str) -> Value {
|
||||
pub fn format_system(system: &str, options: &AnthropicFormatOptions) -> Value {
|
||||
if options.prompt_cache_disabled {
|
||||
return json!([{
|
||||
TYPE_FIELD: TEXT_TYPE,
|
||||
TEXT_TYPE: system
|
||||
}]);
|
||||
}
|
||||
json!([{
|
||||
TYPE_FIELD: TEXT_TYPE,
|
||||
TEXT_TYPE: system,
|
||||
@@ -765,8 +781,8 @@ pub fn create_request_for_model(
|
||||
) -> Result<Value> {
|
||||
let options = options.for_model(model_config);
|
||||
let anthropic_messages = format_messages_with_options(messages, &options);
|
||||
let tool_specs = format_tools(tools);
|
||||
let system_spec = format_system(system);
|
||||
let tool_specs = format_tools(tools, &options);
|
||||
let system_spec = format_system(system, &options);
|
||||
|
||||
if anthropic_messages.is_empty() {
|
||||
return Err(anyhow!("No valid messages to send to Anthropic API"));
|
||||
@@ -1341,9 +1357,7 @@ mod tests {
|
||||
&messages,
|
||||
&AnthropicFormatOptions {
|
||||
preserve_unsigned_thinking: true,
|
||||
preserve_thinking_context: false,
|
||||
thinking_disabled: false,
|
||||
current_model: None,
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
|
||||
@@ -1444,7 +1458,7 @@ mod tests {
|
||||
),
|
||||
];
|
||||
|
||||
let spec = format_tools(&tools);
|
||||
let spec = format_tools(&tools, &AnthropicFormatOptions::default());
|
||||
|
||||
assert_eq!(spec.len(), 2);
|
||||
assert_eq!(spec[0]["name"], "calculator");
|
||||
@@ -1459,7 +1473,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_system_to_anthropic_spec() {
|
||||
let system = "You are a helpful assistant.";
|
||||
let spec = format_system(system);
|
||||
let spec = format_system(system, &AnthropicFormatOptions::default());
|
||||
|
||||
assert!(spec.is_array());
|
||||
let spec_array = spec.as_array().unwrap();
|
||||
@@ -1635,8 +1649,7 @@ mod tests {
|
||||
AnthropicFormatOptions {
|
||||
preserve_unsigned_thinking: true,
|
||||
preserve_thinking_context: true,
|
||||
thinking_disabled: false,
|
||||
current_model: None,
|
||||
..Default::default()
|
||||
},
|
||||
)?;
|
||||
|
||||
@@ -2648,6 +2661,25 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disable_prompt_cache_removes_every_breakpoint() {
|
||||
let config = cfg("claude-sonnet-4-5").with_merged_request_params(
|
||||
std::collections::HashMap::from([(
|
||||
"disable_prompt_cache".to_string(),
|
||||
json!(true),
|
||||
)]),
|
||||
);
|
||||
let req = create_request_with_default_options(
|
||||
&config,
|
||||
"You are a summarizer.",
|
||||
&[Message::user().with_text("Summarize the conversation above.")],
|
||||
&sample_tools(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(!req.to_string().contains(CACHE_CONTROL_FIELD));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn breakpoints_land_on_the_last_content_block() {
|
||||
let messages = vec![Message::user()
|
||||
|
||||
@@ -4,7 +4,7 @@ use crate::formats::anthropic::{
|
||||
adaptive_output_effort, model_supports_temperature, thinking_block_is_stale,
|
||||
thinking_budget_tokens, thinking_type_for_provider, ThinkingType,
|
||||
};
|
||||
use crate::model::ModelConfig;
|
||||
use crate::model::{is_goose_internal_request_param, ModelConfig};
|
||||
|
||||
use crate::formats::openai::{
|
||||
extract_reasoning_effort, is_openai_responses_model, is_valid_function_name,
|
||||
@@ -584,6 +584,7 @@ pub fn create_request_for_provider(
|
||||
}
|
||||
|
||||
if CacheSemantics::for_model("databricks", &model_config.model_name).uses_explicit_breakpoints()
|
||||
&& !model_config.prompt_cache_disabled()
|
||||
{
|
||||
apply_chat_payload_breakpoints(&mut payload);
|
||||
}
|
||||
@@ -592,7 +593,7 @@ pub fn create_request_for_provider(
|
||||
if let Some(params) = &model_config.request_params {
|
||||
if let Some(obj) = payload.as_object_mut() {
|
||||
for (key, value) in params {
|
||||
if key == "thinking_effort" {
|
||||
if is_goose_internal_request_param(key) {
|
||||
continue;
|
||||
}
|
||||
obj.insert(key.clone(), value.clone());
|
||||
@@ -1255,6 +1256,30 @@ mod tests {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_create_request_one_shot_claude() -> anyhow::Result<()> {
|
||||
let model_config = ModelConfig::new("databricks-claude-sonnet-4-5")
|
||||
.with_merged_request_params(std::collections::HashMap::from([
|
||||
("anthropic_beta".to_string(), serde_json::json!(["ctx-1m"])),
|
||||
("disable_prompt_cache".to_string(), serde_json::json!(true)),
|
||||
]));
|
||||
let messages = vec![Message::user().with_text("Summarize the conversation above.")];
|
||||
|
||||
let request = create_request(
|
||||
&model_config,
|
||||
"system",
|
||||
&messages,
|
||||
&[],
|
||||
&ImageFormat::OpenAi,
|
||||
)?;
|
||||
|
||||
assert_eq!(request["anthropic_beta"], serde_json::json!(["ctx-1m"]));
|
||||
assert!(request.get("disable_prompt_cache").is_none());
|
||||
assert!(!request.to_string().contains("cache_control"));
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_create_request_adaptive_thinking_for_46_models() -> anyhow::Result<()> {
|
||||
let mut model_config = ModelConfig::new("databricks-claude-opus-4-6");
|
||||
|
||||
@@ -5,7 +5,7 @@ use crate::errors::ProviderError;
|
||||
use crate::images::{convert_image, detect_image_path, load_image_file, ImageFormat};
|
||||
use crate::json::{parse_tool_arguments, truncation_error_message};
|
||||
use crate::mcp_utils::extract_text_from_resource;
|
||||
use crate::model::ModelConfig;
|
||||
use crate::model::{is_goose_internal_request_param, ModelConfig};
|
||||
use crate::thinking::{
|
||||
split_think_blocks, ThinkFilter, ThinkingEffort, GEMINI_THOUGHT_SIGNATURE_KEY,
|
||||
};
|
||||
@@ -1727,7 +1727,7 @@ pub fn create_request_for_model_with_options(
|
||||
if let Some(params) = &model_config.request_params {
|
||||
if let Some(obj) = payload.as_object_mut() {
|
||||
for (key, value) in params {
|
||||
if key != "thinking_effort" && !is_reserved_request_param_key(key) {
|
||||
if !is_goose_internal_request_param(key) && !is_reserved_request_param_key(key) {
|
||||
obj.insert(key.clone(), value.clone());
|
||||
}
|
||||
}
|
||||
@@ -2839,6 +2839,9 @@ mod tests {
|
||||
("max_tokens".to_string(), json!(1)),
|
||||
("temperature".to_string(), json!(2.0)),
|
||||
("provider_custom".to_string(), json!("allowed")),
|
||||
("thinking_effort".to_string(), json!("high")),
|
||||
("disable_prompt_cache".to_string(), json!(true)),
|
||||
("preserve_thinking_context".to_string(), json!(true)),
|
||||
]);
|
||||
let model_config = test_model_config("glm-4.7")
|
||||
.with_max_tokens(Some(4096))
|
||||
@@ -2867,6 +2870,9 @@ mod tests {
|
||||
assert_eq!(request["max_tokens"], 1);
|
||||
assert_eq!(request["temperature"], 2.0);
|
||||
assert_eq!(request["provider_custom"], "allowed");
|
||||
assert!(request.get("thinking_effort").is_none());
|
||||
assert!(request.get("disable_prompt_cache").is_none());
|
||||
assert!(request.get("preserve_thinking_context").is_none());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -23,6 +23,18 @@ const INHERITED_SESSION_PARAM_KEYS: &[&str] = &[
|
||||
"preserve_unsigned_thinking",
|
||||
];
|
||||
|
||||
/// Request params goose consumes itself: formats that forward unknown params into
|
||||
/// the payload must skip these, or the provider gets an unrecognized wire parameter.
|
||||
pub fn is_goose_internal_request_param(key: &str) -> bool {
|
||||
matches!(
|
||||
key,
|
||||
"thinking_effort"
|
||||
| "disable_prompt_cache"
|
||||
| "preserve_thinking_context"
|
||||
| "preserve_unsigned_thinking"
|
||||
)
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct ModelConfig {
|
||||
pub model_name: String,
|
||||
@@ -300,6 +312,18 @@ impl ModelConfig {
|
||||
.and_then(|s| s.parse::<ThinkingEffort>().ok())
|
||||
}
|
||||
|
||||
pub fn with_prompt_cache_disabled(self) -> Self {
|
||||
self.with_merged_request_params(HashMap::from([(
|
||||
"disable_prompt_cache".to_string(),
|
||||
Value::Bool(true),
|
||||
)]))
|
||||
}
|
||||
|
||||
pub fn prompt_cache_disabled(&self) -> bool {
|
||||
self.request_param::<bool>("disable_prompt_cache")
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
pub fn request_param<T: for<'de> serde::Deserialize<'de>>(
|
||||
&self,
|
||||
request_key: &str,
|
||||
|
||||
@@ -375,8 +375,7 @@ fn format_options_for_provider(preserves_thinking: bool) -> AnthropicFormatOptio
|
||||
AnthropicFormatOptions {
|
||||
preserve_unsigned_thinking: preserves_thinking,
|
||||
preserve_thinking_context: preserves_thinking,
|
||||
thinking_disabled: false,
|
||||
current_model: None,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -112,6 +112,15 @@ pub async fn get_fast_model(
|
||||
}
|
||||
}
|
||||
|
||||
/// Fast tasks summarize a transcript or tool result that never recurs, so a prompt
|
||||
/// cache entry written for one can never be read back and only costs the
|
||||
/// cache-write premium.
|
||||
fn one_shot_model_config(model_config: ModelConfig) -> ModelConfig {
|
||||
model_config
|
||||
.with_thinking_effort(ThinkingEffort::Off)
|
||||
.with_prompt_cache_disabled()
|
||||
}
|
||||
|
||||
/// Run a completion for a lightweight "fast" task (session naming, compaction,
|
||||
/// summarization) using the provider's fast model, falling back to the supplied
|
||||
/// main `model_config` if the fast model errors.
|
||||
@@ -123,10 +132,11 @@ pub async fn complete_fast(
|
||||
messages: &[Message],
|
||||
tools: &[Tool],
|
||||
) -> Result<(Message, ProviderUsage), ProviderError> {
|
||||
let fast_model_config = get_fast_model(provider.get_name(), model_config)
|
||||
.await
|
||||
.map_err(|e| ProviderError::ExecutionError(e.to_string()))?
|
||||
.with_thinking_effort(ThinkingEffort::Off);
|
||||
let fast_model_config = one_shot_model_config(
|
||||
get_fast_model(provider.get_name(), model_config)
|
||||
.await
|
||||
.map_err(|e| ProviderError::ExecutionError(e.to_string()))?,
|
||||
);
|
||||
|
||||
match crate::session_context::with_session_id(
|
||||
Some(session_id.to_string()),
|
||||
@@ -142,9 +152,7 @@ pub async fn complete_fast(
|
||||
e,
|
||||
model_config.model_name
|
||||
);
|
||||
let fallback_config = model_config
|
||||
.clone()
|
||||
.with_thinking_effort(ThinkingEffort::Off);
|
||||
let fallback_config = one_shot_model_config(model_config.clone());
|
||||
crate::session_context::with_session_id(
|
||||
Some(session_id.to_string()),
|
||||
provider.complete(&fallback_config, system, messages, tools),
|
||||
@@ -261,6 +269,16 @@ fn parse_yaml_bool_config(key: &str, value: serde_yaml::Value) -> Result<bool> {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod one_shot_tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn prompt_cache_is_disabled() {
|
||||
assert!(one_shot_model_config(ModelConfig::new("claude-haiku-4-5")).prompt_cache_disabled());
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod azure_foundry_tests {
|
||||
use super::*;
|
||||
|
||||
@@ -237,7 +237,7 @@ impl BedrockProvider {
|
||||
let enabled = config
|
||||
.get_param::<bool>("BEDROCK_ENABLE_CACHING")
|
||||
.unwrap_or(false);
|
||||
enabled && model.model_name.contains("anthropic.claude")
|
||||
enabled && model.model_name.contains("anthropic.claude") && !model.prompt_cache_disabled()
|
||||
}
|
||||
|
||||
async fn post_mantle_streaming(
|
||||
@@ -1092,6 +1092,15 @@ mod tests {
|
||||
"Caching should be enabled for Claude models when BEDROCK_ENABLE_CACHING=true"
|
||||
);
|
||||
|
||||
let one_shot = model.with_merged_request_params(HashMap::from([(
|
||||
"disable_prompt_cache".to_string(),
|
||||
serde_json::json!(true),
|
||||
)]));
|
||||
assert!(
|
||||
!provider.should_enable_caching(&one_shot),
|
||||
"One-shot requests must not create cache points"
|
||||
);
|
||||
|
||||
std::env::remove_var("BEDROCK_ENABLE_CACHING");
|
||||
}
|
||||
|
||||
|
||||
@@ -249,7 +249,8 @@ impl Provider for LiteLLMProvider {
|
||||
false,
|
||||
)?;
|
||||
|
||||
if self.supports_cache_control(model_config).await {
|
||||
if !model_config.prompt_cache_disabled() && self.supports_cache_control(model_config).await
|
||||
{
|
||||
apply_chat_payload_breakpoints(&mut payload);
|
||||
}
|
||||
|
||||
|
||||
@@ -279,6 +279,7 @@ impl Provider for OpenRouterProvider {
|
||||
|
||||
if CacheSemantics::for_model(OPENROUTER_PROVIDER_NAME, &model_config.model_name)
|
||||
.uses_explicit_breakpoints()
|
||||
&& !model_config.prompt_cache_disabled()
|
||||
{
|
||||
apply_chat_payload_breakpoints(&mut payload);
|
||||
}
|
||||
|
||||
@@ -43,7 +43,8 @@ pub(crate) async fn generate_tool_title(
|
||||
let model_config = agent.model_config_for_session(session_id).await.ok()?;
|
||||
let fast_model_config = get_fast_model(provider.get_name(), &model_config)
|
||||
.await
|
||||
.ok()?;
|
||||
.ok()?
|
||||
.with_prompt_cache_disabled();
|
||||
let title = generate_tool_title_with_provider(
|
||||
provider.as_ref(),
|
||||
&fast_model_config,
|
||||
@@ -85,7 +86,8 @@ pub(crate) async fn generate_tool_chain_summary(
|
||||
let model_config = agent.model_config_for_session(session_id).await.ok()?;
|
||||
let fast_model_config = get_fast_model(provider.get_name(), &model_config)
|
||||
.await
|
||||
.ok()?;
|
||||
.ok()?
|
||||
.with_prompt_cache_disabled();
|
||||
let chain_summary = ToolChainSummary {
|
||||
summary: generate_tool_chain_summary_with_provider(
|
||||
provider.as_ref(),
|
||||
|
||||
Reference in New Issue
Block a user