use serde_json::{json, Value};
use super::types::OptimizerConfig;
pub fn optimize(body: &mut Value, config: &OptimizerConfig) {
if !config.thinking_optimizer {
return;
}
let model = match body.get("model").and_then(|m| m.as_str()) {
Some(model) => model.to_lowercase(),
None => return,
};
if model.contains("haiku") {
return;
}
if model.contains("opus-4-6") || model.contains("sonnet-4-6") {
body["thinking"] = json!({"type": "adaptive"});
body["output_config"] = json!({"effort": "max"});
append_beta(body, "context-1m-2025-08-07");
return;
}
let max_tokens = body
.get("max_tokens")
.and_then(|value| value.as_u64())
.unwrap_or(16_384);
let budget_target = max_tokens.saturating_sub(1);
let thinking_type = body
.get("thinking")
.and_then(|thinking| thinking.get("type"))
.and_then(|value| value.as_str());
match thinking_type {
None | Some("disabled") => {
body["thinking"] = json!({
"type": "enabled",
"budget_tokens": budget_target,
});
}
Some("enabled") => {
let current_budget = body
.get("thinking")
.and_then(|thinking| thinking.get("budget_tokens"))
.and_then(|value| value.as_u64())
.unwrap_or(0);
if current_budget < budget_target {
body["thinking"]["budget_tokens"] = json!(budget_target);
}
}
Some(_) => {}
}
append_beta(body, "interleaved-thinking-2025-05-14");
}
fn append_beta(body: &mut Value, beta: &str) {
match body.get("anthropic_beta") {
Some(Value::Array(existing))
if existing.iter().any(|value| value.as_str() == Some(beta)) => {}
Some(Value::Array(_)) => body["anthropic_beta"]
.as_array_mut()
.unwrap()
.push(json!(beta)),
_ => body["anthropic_beta"] = json!([beta]),
}
}
#[cfg(test)]
mod tests {
use super::*;
fn enabled_config() -> OptimizerConfig {
OptimizerConfig {
enabled: true,
thinking_optimizer: true,
cache_injection: true,
cache_ttl: "1h".to_string(),
}
}
#[test]
fn adaptive_models_use_adaptive_thinking() {
let mut body = json!({
"model": "anthropic.claude-opus-4-6-20250514-v1:0",
"max_tokens": 128,
});
optimize(&mut body, &enabled_config());
assert_eq!(body["thinking"]["type"], "adaptive");
assert_eq!(body["output_config"]["effort"], "max");
assert!(body["anthropic_beta"]
.as_array()
.unwrap()
.iter()
.any(|value| value == "context-1m-2025-08-07"));
}
#[test]
fn legacy_models_enable_thinking_budget() {
let mut body = json!({
"model": "anthropic.claude-sonnet-4-5-20250514-v1:0",
"max_tokens": 32,
});
optimize(&mut body, &enabled_config());
assert_eq!(body["thinking"]["type"], "enabled");
assert_eq!(body["thinking"]["budget_tokens"], 31);
assert!(body["anthropic_beta"]
.as_array()
.unwrap()
.iter()
.any(|value| value == "interleaved-thinking-2025-05-14"));
}
}