mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
feat: enable Anthropic prompt caching via automatic cache_control injection (#660)
* feat(llm): add Anthropic prompt caching and cache token tracking - Inject cache_control via additional_params for Claude models in rig_adapter - Add cache_read_input_tokens and cache_creation_input_tokens to CompletionResponse and ToolCompletionResponse - Extract cached_input_tokens from rig-core unified Usage - Add is_anthropic_model() detection helper with provider prefix support - Log prompt cache hits at debug level (consistent with response_cache) - Add 7 unit tests for cache injection and model detection - Update all mock providers and test fixtures with new fields * feat(cost): apply 90% cache discount to prompt-cached tokens in CostGuard - Add cache_read_input_tokens to TokenUsage so cache counts flow from CompletionResponse through the reasoning layer to the dispatcher - Update CostGuard::record_llm_call() to accept cache_read_input_tokens: cached tokens are billed at 10% of the normal input rate - Thread cache_read_input_tokens from dispatcher into CostGuard - Add test_cache_discount_reduces_cost verifying exact savings match 90% of input cost for fully-cached requests - Update all existing test callers with zero-cache parameter * refactor(cache): scope cache_control to Anthropic backend and validate model support - Replace model-name-based is_anthropic_model() with explicit enable_prompt_cache flag on RigAdapter, set only for the direct Anthropic backend via with_prompt_cache(true) - Add supports_prompt_cache() to validate model names per Anthropic docs: only Claude 3+ models support caching; claude-2 and claude-instant are excluded to prevent 400 errors - Warn when caching is enabled but model does not support it - Replace is_anthropic_model tests with flag-based and model validation tests * fix(cache): validate model at construction and propagate cache metrics through proxy - Move supports_prompt_cache() check into with_prompt_cache() so unsupported models are detected once at construction, not per request - Add cache_read_input_tokens and cache_creation_input_tokens to ProxyCompletionResponse and ProxyToolCompletionResponse with serde(default) for backward compatibility - Pass cache metrics through orchestrator proxy instead of zeroing - Use claude-opus-4-6 in cache discount test to match Anthropic semantics * feat(llm): add configurable cache retention with write surcharge - Add CacheRetention enum (none/short/long) to AnthropicDirectConfig - Parse ANTHROPIC_CACHE_RETENTION env var (default: short) - Inject TTL-aware cache_control (short=5m ephemeral, long=1h) - Extract cache_creation_input_tokens from raw Anthropic response - Add cache_write_multiplier() to LlmProvider trait (1.25x short, 2.0x long) - Pipe dynamic write multiplier through dispatcher to CostGuard - Add TokenUsage.cache_creation_input_tokens field - Add tests for Long TTL injection, 5m and 1h write surcharges - Document ANTHROPIC_CACHE_RETENTION in .env.example * docs: fix stale cache_retention field comment * fix: resolve CI failures after upstream merge - Add missing cost_per_token arg to cache test callsites - Apply cargo fmt to long lines in tests and tracing macros * fix: address Copilot review feedback - Use saturating_add for cache token sum to prevent u32 overflow - Tighten supports_prompt_cache to explicitly match claude-3+/claude-4+ and named families (claude-sonnet/claude-opus/claude-haiku) * fix: adapt prompt caching to registry architecture and add missing cache fields - Resolve merge conflicts: adapt CacheRetention and cache injection to the declarative provider registry (RegistryProviderConfig replaces AnthropicDirectConfig) - Parse ANTHROPIC_CACHE_RETENTION env var in create_anthropic_from_registry() - Use Anthropic automatic caching via top-level cache_control in additional_params (rig-core #[serde(flatten)] places it at request root) - Add cache_read/creation_input_tokens fields to all mock LlmProviders added on main after PR #291 branched (response_cache, dispatcher, provider_chaos, trace_llm) - Suppress clippy::too_many_arguments on record_llm_call and build_rig_request - Add regression tests for cache injection (short/long/none) and cache_write_multiplier values Co-Authored-By: Canvinus <[email protected]> * fix: delegate cache_write_multiplier through provider wrappers and make cache_read_discount configurable The 6 decorator providers (Retry, CircuitBreaker, Failover, SmartRouting, CachedProvider, RecordingLlm) did not delegate cache_write_multiplier() to their inner provider, causing it to always return 1.0 instead of the actual 1.25x/2.0x from RigAdapter. This fix adds delegation for both cache_write_multiplier() and the new cache_read_discount() method. Also makes the cache read discount per-provider instead of hardcoding Anthropic's 90% discount (÷10). OpenAI uses 50% (÷2), so the discount is now returned by each provider via the LlmProvider trait. Addresses review feedback on PR #660. Co-Authored-By: Claude Opus 4.6 <[email protected]> * style: cargo fmt Co-Authored-By: Claude Opus 4.6 <[email protected]> * test: add CacheRetention FromStr/Display unit tests Tests cover primary values, aliases (off/disabled/5m/ephemeral/1h), case-insensitivity, invalid input error, and Display round-trip. Addresses Copilot review feedback on PR #660. Co-Authored-By: Claude Opus 4.6 <[email protected]> --------- Co-authored-by: Andrey <[email protected]> Co-authored-by: Andrey Gruzdev <[email protected]> Co-authored-by: Claude Opus 4.6 <[email protected]>
This commit is contained in:
co-authored by
Andrey
Andrey Gruzdev
Claude Opus 4.6
parent
633b234e44
commit
424a0366a9
@@ -57,6 +57,17 @@ NEARAI_AUTH_URL=https://private.near.ai
|
|||||||
# LLM_BASE_URL=https://api.fireworks.ai/inference/v1
|
# LLM_BASE_URL=https://api.fireworks.ai/inference/v1
|
||||||
# LLM_API_KEY=fw_...
|
# LLM_API_KEY=fw_...
|
||||||
|
|
||||||
|
# === Anthropic Direct ===
|
||||||
|
# LLM_BACKEND=anthropic
|
||||||
|
# ANTHROPIC_MODEL=claude-sonnet-4-6
|
||||||
|
# ANTHROPIC_API_KEY=sk-ant-...
|
||||||
|
# ANTHROPIC_BASE_URL=https://api.anthropic.com # default
|
||||||
|
# Prompt cache retention — controls Anthropic server-side prompt caching:
|
||||||
|
# none = disabled (no cache_control injected)
|
||||||
|
# short = 5-minute TTL, 1.25× (125%) write surcharge (default)
|
||||||
|
# long = 1-hour TTL, 2.0× (200%) write surcharge
|
||||||
|
# ANTHROPIC_CACHE_RETENTION=short
|
||||||
|
|
||||||
# For full provider setup guide see docs/LLM_PROVIDERS.md
|
# For full provider setup guide see docs/LLM_PROVIDERS.md
|
||||||
|
|
||||||
# Channel Configuration
|
# Channel Configuration
|
||||||
|
|||||||
+230
-11
@@ -151,21 +151,46 @@ impl CostGuard {
|
|||||||
/// Record a completed LLM action: its token costs and the action timestamp.
|
/// Record a completed LLM action: its token costs and the action timestamp.
|
||||||
///
|
///
|
||||||
/// Call this AFTER an LLM call completes so that costs are tracked.
|
/// Call this AFTER an LLM call completes so that costs are tracked.
|
||||||
|
/// - `cache_read_input_tokens`: tokens served from cache.
|
||||||
|
/// - `cache_creation_input_tokens`: tokens written to cache.
|
||||||
|
/// - `cache_read_discount`: divisor for cache-read cost (e.g. 10 for Anthropic 90% off, 2 for OpenAI 50% off).
|
||||||
|
/// - `cache_write_multiplier`: cost multiplier for cache writes (1.25 for 5m, 2.0 for 1h).
|
||||||
///
|
///
|
||||||
/// When `cost_per_token` is `Some`, those rates are used directly (provider-
|
/// When `cost_per_token` is `Some`, those rates are used directly (provider-
|
||||||
/// sourced pricing). When `None`, falls back to the static `costs::model_cost`
|
/// sourced pricing). When `None`, falls back to the static `costs::model_cost`
|
||||||
/// lookup table, then `costs::default_cost`.
|
/// lookup table, then `costs::default_cost`.
|
||||||
|
#[allow(clippy::too_many_arguments)]
|
||||||
pub async fn record_llm_call(
|
pub async fn record_llm_call(
|
||||||
&self,
|
&self,
|
||||||
model: &str,
|
model: &str,
|
||||||
input_tokens: u32,
|
input_tokens: u32,
|
||||||
output_tokens: u32,
|
output_tokens: u32,
|
||||||
|
cache_read_input_tokens: u32,
|
||||||
|
cache_creation_input_tokens: u32,
|
||||||
|
cache_read_discount: Decimal,
|
||||||
|
cache_write_multiplier: Decimal,
|
||||||
cost_per_token: Option<(Decimal, Decimal)>,
|
cost_per_token: Option<(Decimal, Decimal)>,
|
||||||
) -> Decimal {
|
) -> Decimal {
|
||||||
let (input_rate, output_rate) = cost_per_token
|
let (input_rate, output_rate) = cost_per_token
|
||||||
.unwrap_or_else(|| costs::model_cost(model).unwrap_or_else(costs::default_cost));
|
.unwrap_or_else(|| costs::model_cost(model).unwrap_or_else(costs::default_cost));
|
||||||
let cost =
|
// Cached read tokens cost input_rate / cache_read_discount (provider-specific).
|
||||||
input_rate * Decimal::from(input_tokens) + output_rate * Decimal::from(output_tokens);
|
// Cached write tokens cost write_multiplier × input_rate (e.g. 1.25× for 5m, 2× for 1h).
|
||||||
|
// Uncached tokens = total input - cache reads - cache writes.
|
||||||
|
let cached_total = cache_read_input_tokens.saturating_add(cache_creation_input_tokens);
|
||||||
|
let uncached_input = input_tokens.saturating_sub(cached_total);
|
||||||
|
let effective_discount = if cache_read_discount.is_zero() {
|
||||||
|
Decimal::ONE
|
||||||
|
} else {
|
||||||
|
cache_read_discount
|
||||||
|
};
|
||||||
|
let cache_read_cost =
|
||||||
|
input_rate * Decimal::from(cache_read_input_tokens) / effective_discount;
|
||||||
|
let cache_write_cost =
|
||||||
|
input_rate * Decimal::from(cache_creation_input_tokens) * cache_write_multiplier;
|
||||||
|
let cost = input_rate * Decimal::from(uncached_input)
|
||||||
|
+ cache_read_cost
|
||||||
|
+ cache_write_cost
|
||||||
|
+ output_rate * Decimal::from(output_tokens);
|
||||||
|
|
||||||
// Update daily cost (reset if new day)
|
// Update daily cost (reset if new day)
|
||||||
{
|
{
|
||||||
@@ -267,7 +292,16 @@ mod tests {
|
|||||||
|
|
||||||
// Record a big call, still allowed
|
// Record a big call, still allowed
|
||||||
guard
|
guard
|
||||||
.record_llm_call("gpt-4o", 100_000, 100_000, None)
|
.record_llm_call(
|
||||||
|
"gpt-4o",
|
||||||
|
100_000,
|
||||||
|
100_000,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
Decimal::ONE,
|
||||||
|
Decimal::ONE,
|
||||||
|
None,
|
||||||
|
)
|
||||||
.await;
|
.await;
|
||||||
assert!(guard.check_allowed().await.is_ok());
|
assert!(guard.check_allowed().await.is_ok());
|
||||||
}
|
}
|
||||||
@@ -285,7 +319,18 @@ mod tests {
|
|||||||
// Record a call that costs more than $0.01
|
// Record a call that costs more than $0.01
|
||||||
// gpt-4o: input=$0.0000025/tok, output=$0.00001/tok
|
// gpt-4o: input=$0.0000025/tok, output=$0.00001/tok
|
||||||
// 10000 input + 10000 output = $0.025 + $0.10 = $0.125
|
// 10000 input + 10000 output = $0.025 + $0.10 = $0.125
|
||||||
guard.record_llm_call("gpt-4o", 10_000, 10_000, None).await;
|
guard
|
||||||
|
.record_llm_call(
|
||||||
|
"gpt-4o",
|
||||||
|
10_000,
|
||||||
|
10_000,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
Decimal::ONE,
|
||||||
|
Decimal::ONE,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
// Now should be blocked
|
// Now should be blocked
|
||||||
let result = guard.check_allowed().await;
|
let result = guard.check_allowed().await;
|
||||||
@@ -308,7 +353,9 @@ mod tests {
|
|||||||
// First 3 actions allowed
|
// First 3 actions allowed
|
||||||
for _ in 0..3 {
|
for _ in 0..3 {
|
||||||
assert!(guard.check_allowed().await.is_ok());
|
assert!(guard.check_allowed().await.is_ok());
|
||||||
guard.record_llm_call("gpt-4o", 10, 10, None).await;
|
guard
|
||||||
|
.record_llm_call("gpt-4o", 10, 10, 0, 0, Decimal::ONE, Decimal::ONE, None)
|
||||||
|
.await;
|
||||||
}
|
}
|
||||||
|
|
||||||
// 4th should be blocked
|
// 4th should be blocked
|
||||||
@@ -329,7 +376,9 @@ mod tests {
|
|||||||
|
|
||||||
assert_eq!(guard.daily_spend().await, Decimal::ZERO);
|
assert_eq!(guard.daily_spend().await, Decimal::ZERO);
|
||||||
|
|
||||||
let cost = guard.record_llm_call("gpt-4o", 1000, 500, None).await;
|
let cost = guard
|
||||||
|
.record_llm_call("gpt-4o", 1000, 500, 0, 0, Decimal::ONE, Decimal::ONE, None)
|
||||||
|
.await;
|
||||||
assert!(cost > Decimal::ZERO);
|
assert!(cost > Decimal::ZERO);
|
||||||
assert_eq!(guard.daily_spend().await, cost);
|
assert_eq!(guard.daily_spend().await, cost);
|
||||||
}
|
}
|
||||||
@@ -340,8 +389,12 @@ mod tests {
|
|||||||
|
|
||||||
assert_eq!(guard.actions_this_hour().await, 0);
|
assert_eq!(guard.actions_this_hour().await, 0);
|
||||||
|
|
||||||
guard.record_llm_call("gpt-4o", 10, 10, None).await;
|
guard
|
||||||
guard.record_llm_call("gpt-4o", 10, 10, None).await;
|
.record_llm_call("gpt-4o", 10, 10, 0, 0, Decimal::ONE, Decimal::ONE, None)
|
||||||
|
.await;
|
||||||
|
guard
|
||||||
|
.record_llm_call("gpt-4o", 10, 10, 0, 0, Decimal::ONE, Decimal::ONE, None)
|
||||||
|
.await;
|
||||||
|
|
||||||
assert_eq!(guard.actions_this_hour().await, 2);
|
assert_eq!(guard.actions_this_hour().await, 2);
|
||||||
}
|
}
|
||||||
@@ -378,10 +431,23 @@ mod tests {
|
|||||||
assert!(guard.model_usage().await.is_empty());
|
assert!(guard.model_usage().await.is_empty());
|
||||||
|
|
||||||
// Record calls for two different models
|
// Record calls for two different models
|
||||||
guard.record_llm_call("gpt-4o", 1000, 500, None).await;
|
|
||||||
guard.record_llm_call("gpt-4o", 2000, 1000, None).await;
|
|
||||||
guard
|
guard
|
||||||
.record_llm_call("claude-3-5-sonnet-20241022", 500, 200, None)
|
.record_llm_call("gpt-4o", 1000, 500, 0, 0, Decimal::ONE, Decimal::ONE, None)
|
||||||
|
.await;
|
||||||
|
guard
|
||||||
|
.record_llm_call("gpt-4o", 2000, 1000, 0, 0, Decimal::ONE, Decimal::ONE, None)
|
||||||
|
.await;
|
||||||
|
guard
|
||||||
|
.record_llm_call(
|
||||||
|
"claude-3-5-sonnet-20241022",
|
||||||
|
500,
|
||||||
|
200,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
Decimal::ONE,
|
||||||
|
Decimal::ONE,
|
||||||
|
None,
|
||||||
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
let usage = guard.model_usage().await;
|
let usage = guard.model_usage().await;
|
||||||
@@ -402,4 +468,157 @@ mod tests {
|
|||||||
// Costs should differ since models have different pricing
|
// Costs should differ since models have different pricing
|
||||||
assert_ne!(gpt.cost, claude.cost);
|
assert_ne!(gpt.cost, claude.cost);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_cache_discount_reduces_cost() {
|
||||||
|
let guard = CostGuard::new(CostGuardConfig::default());
|
||||||
|
|
||||||
|
// Full price: 1000 input + 500 output, no cache
|
||||||
|
let full_cost = guard
|
||||||
|
.record_llm_call(
|
||||||
|
"claude-opus-4-6",
|
||||||
|
1000,
|
||||||
|
500,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
Decimal::ONE,
|
||||||
|
Decimal::ONE,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
let guard2 = CostGuard::new(CostGuardConfig::default());
|
||||||
|
|
||||||
|
// Same tokens but all input cached (90% discount on input)
|
||||||
|
let cached_cost = guard2
|
||||||
|
.record_llm_call(
|
||||||
|
"claude-opus-4-6",
|
||||||
|
1000,
|
||||||
|
500,
|
||||||
|
1000,
|
||||||
|
0,
|
||||||
|
dec!(10),
|
||||||
|
Decimal::ONE,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
// Cached cost must be strictly less than full cost
|
||||||
|
assert!(
|
||||||
|
cached_cost < full_cost,
|
||||||
|
"cached_cost ({}) should be less than full_cost ({})",
|
||||||
|
cached_cost,
|
||||||
|
full_cost
|
||||||
|
);
|
||||||
|
|
||||||
|
// The difference should be exactly 90% of the input cost
|
||||||
|
let (input_rate, _) = costs::model_cost("claude-opus-4-6").unwrap();
|
||||||
|
let expected_savings = input_rate * Decimal::from(1000u32) * dec!(9) / dec!(10);
|
||||||
|
let actual_savings = full_cost - cached_cost;
|
||||||
|
assert_eq!(
|
||||||
|
actual_savings, expected_savings,
|
||||||
|
"savings should be 90% of input cost for fully-cached request"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_cache_write_surcharge_increases_cost() {
|
||||||
|
let guard = CostGuard::new(CostGuardConfig::default());
|
||||||
|
|
||||||
|
// Full price: 1000 input + 500 output, no cache activity
|
||||||
|
let full_cost = guard
|
||||||
|
.record_llm_call(
|
||||||
|
"claude-opus-4-6",
|
||||||
|
1000,
|
||||||
|
500,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
Decimal::ONE,
|
||||||
|
Decimal::ONE,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
let guard2 = CostGuard::new(CostGuardConfig::default());
|
||||||
|
|
||||||
|
// Same tokens, but all input tokens are cache writes (1.25x surcharge for 5m TTL)
|
||||||
|
let short_multiplier = Decimal::new(125, 2); // 1.25
|
||||||
|
let write_cost = guard2
|
||||||
|
.record_llm_call(
|
||||||
|
"claude-opus-4-6",
|
||||||
|
1000,
|
||||||
|
500,
|
||||||
|
0,
|
||||||
|
1000,
|
||||||
|
Decimal::ONE,
|
||||||
|
short_multiplier,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
// Write cost must be strictly greater than full cost
|
||||||
|
assert!(
|
||||||
|
write_cost > full_cost,
|
||||||
|
"write_cost ({}) should be greater than full_cost ({})",
|
||||||
|
write_cost,
|
||||||
|
full_cost
|
||||||
|
);
|
||||||
|
|
||||||
|
// The difference should be exactly 25% of the input cost
|
||||||
|
let (input_rate, _) = costs::model_cost("claude-opus-4-6").unwrap();
|
||||||
|
let expected_surcharge = input_rate * Decimal::from(1000u32) * dec!(0.25);
|
||||||
|
let actual_surcharge = write_cost - full_cost;
|
||||||
|
assert_eq!(
|
||||||
|
actual_surcharge, expected_surcharge,
|
||||||
|
"surcharge should be 25% of input cost for 5m cache writes"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_cache_write_surcharge_long_ttl() {
|
||||||
|
let guard = CostGuard::new(CostGuardConfig::default());
|
||||||
|
|
||||||
|
// Full price: 1000 input + 500 output
|
||||||
|
let full_cost = guard
|
||||||
|
.record_llm_call(
|
||||||
|
"claude-opus-4-6",
|
||||||
|
1000,
|
||||||
|
500,
|
||||||
|
0,
|
||||||
|
0,
|
||||||
|
Decimal::ONE,
|
||||||
|
Decimal::ONE,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
let guard2 = CostGuard::new(CostGuardConfig::default());
|
||||||
|
|
||||||
|
// All input tokens are cache writes with 2.0x multiplier (1h TTL)
|
||||||
|
let long_multiplier = Decimal::TWO;
|
||||||
|
let write_cost = guard2
|
||||||
|
.record_llm_call(
|
||||||
|
"claude-opus-4-6",
|
||||||
|
1000,
|
||||||
|
500,
|
||||||
|
0,
|
||||||
|
1000,
|
||||||
|
Decimal::ONE,
|
||||||
|
long_multiplier,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
// Write cost > full cost
|
||||||
|
assert!(write_cost > full_cost);
|
||||||
|
|
||||||
|
// Surcharge should be 100% of input cost (2.0x - 1.0x = 1.0x)
|
||||||
|
let (input_rate, _) = costs::model_cost("claude-opus-4-6").unwrap();
|
||||||
|
let expected_surcharge = input_rate * Decimal::from(1000u32);
|
||||||
|
let actual_surcharge = write_cost - full_cost;
|
||||||
|
assert_eq!(
|
||||||
|
actual_surcharge, expected_surcharge,
|
||||||
|
"surcharge should be 100% of input cost for 1h cache writes"
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -278,12 +278,18 @@ impl Agent {
|
|||||||
|
|
||||||
// Record cost and track token usage
|
// Record cost and track token usage
|
||||||
let model_name = self.llm().active_model_name();
|
let model_name = self.llm().active_model_name();
|
||||||
|
let read_discount = self.llm().cache_read_discount();
|
||||||
|
let write_multiplier = self.llm().cache_write_multiplier();
|
||||||
let call_cost = self
|
let call_cost = self
|
||||||
.cost_guard()
|
.cost_guard()
|
||||||
.record_llm_call(
|
.record_llm_call(
|
||||||
&model_name,
|
&model_name,
|
||||||
output.usage.input_tokens,
|
output.usage.input_tokens,
|
||||||
output.usage.output_tokens,
|
output.usage.output_tokens,
|
||||||
|
output.usage.cache_read_input_tokens,
|
||||||
|
output.usage.cache_creation_input_tokens,
|
||||||
|
read_discount,
|
||||||
|
write_multiplier,
|
||||||
Some(self.llm().cost_per_token()),
|
Some(self.llm().cost_per_token()),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
@@ -1062,6 +1068,8 @@ mod tests {
|
|||||||
input_tokens: 0,
|
input_tokens: 0,
|
||||||
output_tokens: 0,
|
output_tokens: 0,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1075,6 +1083,8 @@ mod tests {
|
|||||||
input_tokens: 0,
|
input_tokens: 0,
|
||||||
output_tokens: 0,
|
output_tokens: 0,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1635,6 +1645,8 @@ mod tests {
|
|||||||
input_tokens: 0,
|
input_tokens: 0,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1650,6 +1662,8 @@ mod tests {
|
|||||||
input_tokens: 0,
|
input_tokens: 0,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
// Tools available: always call one.
|
// Tools available: always call one.
|
||||||
@@ -1663,6 +1677,8 @@ mod tests {
|
|||||||
input_tokens: 0,
|
input_tokens: 0,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
finish_reason: FinishReason::ToolUse,
|
finish_reason: FinishReason::ToolUse,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -1787,6 +1803,8 @@ mod tests {
|
|||||||
input_tokens: 0,
|
input_tokens: 0,
|
||||||
output_tokens: 2,
|
output_tokens: 2,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1801,6 +1819,8 @@ mod tests {
|
|||||||
input_tokens: 0,
|
input_tokens: 0,
|
||||||
output_tokens: 2,
|
output_tokens: 2,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
// Always call a tool that does not exist in the registry.
|
// Always call a tool that does not exist in the registry.
|
||||||
@@ -1814,6 +1834,8 @@ mod tests {
|
|||||||
input_tokens: 0,
|
input_tokens: 0,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
finish_reason: FinishReason::ToolUse,
|
finish_reason: FinishReason::ToolUse,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -9,6 +9,50 @@ use crate::llm::registry::{ProviderProtocol, ProviderRegistry};
|
|||||||
use crate::llm::session::SessionConfig;
|
use crate::llm::session::SessionConfig;
|
||||||
use crate::settings::Settings;
|
use crate::settings::Settings;
|
||||||
|
|
||||||
|
/// Prompt cache retention policy for Anthropic.
|
||||||
|
///
|
||||||
|
/// Controls Anthropic's automatic prompt caching via a top-level
|
||||||
|
/// `cache_control` field injected through rig-core's `additional_params`.
|
||||||
|
/// - `None` — caching disabled, no `cache_control` injected.
|
||||||
|
/// - `Short` — 5-minute TTL (default), `{"type": "ephemeral"}`, 1.25× write surcharge.
|
||||||
|
/// - `Long` — 1-hour TTL, `{"type": "ephemeral", "ttl": "1h"}`, 2× write surcharge.
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||||
|
pub enum CacheRetention {
|
||||||
|
/// No prompt caching.
|
||||||
|
None,
|
||||||
|
/// 5-minute TTL (default). Write cost: 1.25× base input.
|
||||||
|
#[default]
|
||||||
|
Short,
|
||||||
|
/// 1-hour TTL. Write cost: 2× base input.
|
||||||
|
Long,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::str::FromStr for CacheRetention {
|
||||||
|
type Err = String;
|
||||||
|
|
||||||
|
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||||
|
match s.to_lowercase().as_str() {
|
||||||
|
"none" | "off" | "disabled" => Ok(Self::None),
|
||||||
|
"short" | "5m" | "ephemeral" => Ok(Self::Short),
|
||||||
|
"long" | "1h" => Ok(Self::Long),
|
||||||
|
_ => Err(format!(
|
||||||
|
"invalid cache retention '{}', expected one of: none, short, long",
|
||||||
|
s
|
||||||
|
)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::fmt::Display for CacheRetention {
|
||||||
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
|
match self {
|
||||||
|
Self::None => write!(f, "none"),
|
||||||
|
Self::Short => write!(f, "short"),
|
||||||
|
Self::Long => write!(f, "long"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Resolved configuration for a registry-based provider.
|
/// Resolved configuration for a registry-based provider.
|
||||||
///
|
///
|
||||||
/// This single struct replaces what used to be five separate config types
|
/// This single struct replaces what used to be five separate config types
|
||||||
@@ -755,4 +799,86 @@ mod tests {
|
|||||||
std::env::remove_var("LLM_BACKEND");
|
std::env::remove_var("LLM_BACKEND");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn cache_retention_from_str_primary_values() {
|
||||||
|
assert_eq!(
|
||||||
|
"none".parse::<CacheRetention>().unwrap(),
|
||||||
|
CacheRetention::None
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
"short".parse::<CacheRetention>().unwrap(),
|
||||||
|
CacheRetention::Short
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
"long".parse::<CacheRetention>().unwrap(),
|
||||||
|
CacheRetention::Long
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn cache_retention_from_str_aliases() {
|
||||||
|
assert_eq!(
|
||||||
|
"off".parse::<CacheRetention>().unwrap(),
|
||||||
|
CacheRetention::None
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
"disabled".parse::<CacheRetention>().unwrap(),
|
||||||
|
CacheRetention::None
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
"5m".parse::<CacheRetention>().unwrap(),
|
||||||
|
CacheRetention::Short
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
"ephemeral".parse::<CacheRetention>().unwrap(),
|
||||||
|
CacheRetention::Short
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
"1h".parse::<CacheRetention>().unwrap(),
|
||||||
|
CacheRetention::Long
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn cache_retention_from_str_case_insensitive() {
|
||||||
|
assert_eq!(
|
||||||
|
"NONE".parse::<CacheRetention>().unwrap(),
|
||||||
|
CacheRetention::None
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
"Short".parse::<CacheRetention>().unwrap(),
|
||||||
|
CacheRetention::Short
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
"LONG".parse::<CacheRetention>().unwrap(),
|
||||||
|
CacheRetention::Long
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
"Ephemeral".parse::<CacheRetention>().unwrap(),
|
||||||
|
CacheRetention::Short
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn cache_retention_from_str_invalid() {
|
||||||
|
let err = "bogus".parse::<CacheRetention>().unwrap_err();
|
||||||
|
assert!(
|
||||||
|
err.contains("bogus"),
|
||||||
|
"error should mention the invalid value"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn cache_retention_display_round_trip() {
|
||||||
|
for variant in [
|
||||||
|
CacheRetention::None,
|
||||||
|
CacheRetention::Short,
|
||||||
|
CacheRetention::Long,
|
||||||
|
] {
|
||||||
|
let s = variant.to_string();
|
||||||
|
let parsed: CacheRetention = s.parse().unwrap();
|
||||||
|
assert_eq!(parsed, variant, "round-trip failed for {s}");
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+1
-1
@@ -36,7 +36,7 @@ pub use self::database::{DatabaseBackend, DatabaseConfig, SslMode, default_libsq
|
|||||||
pub use self::embeddings::EmbeddingsConfig;
|
pub use self::embeddings::EmbeddingsConfig;
|
||||||
pub use self::heartbeat::HeartbeatConfig;
|
pub use self::heartbeat::HeartbeatConfig;
|
||||||
pub use self::hygiene::HygieneConfig;
|
pub use self::hygiene::HygieneConfig;
|
||||||
pub use self::llm::{LlmConfig, NearAiConfig, RegistryProviderConfig};
|
pub use self::llm::{CacheRetention, LlmConfig, NearAiConfig, RegistryProviderConfig};
|
||||||
pub use self::routines::RoutineConfig;
|
pub use self::routines::RoutineConfig;
|
||||||
pub use self::safety::SafetyConfig;
|
pub use self::safety::SafetyConfig;
|
||||||
pub use self::sandbox::{ClaudeCodeConfig, SandboxModeConfig};
|
pub use self::sandbox::{ClaudeCodeConfig, SandboxModeConfig};
|
||||||
|
|||||||
@@ -245,6 +245,14 @@ impl LlmProvider for CircuitBreakerProvider {
|
|||||||
self.inner.cost_per_token()
|
self.inner.cost_per_token()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn cache_write_multiplier(&self) -> Decimal {
|
||||||
|
self.inner.cache_write_multiplier()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cache_read_discount(&self) -> Decimal {
|
||||||
|
self.inner.cache_read_discount()
|
||||||
|
}
|
||||||
|
|
||||||
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||||
self.check_allowed().await?;
|
self.check_allowed().await?;
|
||||||
match self.inner.complete(request).await {
|
match self.inner.complete(request).await {
|
||||||
|
|||||||
@@ -296,6 +296,14 @@ impl LlmProvider for FailoverProvider {
|
|||||||
self.providers[self.last_used.load(Ordering::Relaxed)].cost_per_token()
|
self.providers[self.last_used.load(Ordering::Relaxed)].cost_per_token()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn cache_write_multiplier(&self) -> Decimal {
|
||||||
|
self.providers[self.last_used.load(Ordering::Relaxed)].cache_write_multiplier()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cache_read_discount(&self) -> Decimal {
|
||||||
|
self.providers[self.last_used.load(Ordering::Relaxed)].cache_read_discount()
|
||||||
|
}
|
||||||
|
|
||||||
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||||
let (provider_idx, response) = self
|
let (provider_idx, response) = self
|
||||||
.try_providers(|provider| {
|
.try_providers(|provider| {
|
||||||
@@ -404,6 +412,8 @@ mod tests {
|
|||||||
input_tokens: 10,
|
input_tokens: 10,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
}))),
|
}))),
|
||||||
tool_complete_result: Mutex::new(Some(Ok(ToolCompletionResponse {
|
tool_complete_result: Mutex::new(Some(Ok(ToolCompletionResponse {
|
||||||
content: Some(content.to_string()),
|
content: Some(content.to_string()),
|
||||||
@@ -411,6 +421,8 @@ mod tests {
|
|||||||
input_tokens: 10,
|
input_tokens: 10,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
}))),
|
}))),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -792,6 +804,8 @@ mod tests {
|
|||||||
input_tokens: 10,
|
input_tokens: 10,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -817,6 +831,8 @@ mod tests {
|
|||||||
input_tokens: 10,
|
input_tokens: 10,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
+29
-1
@@ -177,6 +177,8 @@ fn create_openai_compat_from_registry(
|
|||||||
fn create_anthropic_from_registry(
|
fn create_anthropic_from_registry(
|
||||||
config: &RegistryProviderConfig,
|
config: &RegistryProviderConfig,
|
||||||
) -> Result<Arc<dyn LlmProvider>, LlmError> {
|
) -> Result<Arc<dyn LlmProvider>, LlmError> {
|
||||||
|
use crate::config::CacheRetention;
|
||||||
|
use crate::config::helpers::optional_env;
|
||||||
use rig::providers::anthropic;
|
use rig::providers::anthropic;
|
||||||
|
|
||||||
let api_key = config
|
let api_key = config
|
||||||
@@ -200,8 +202,32 @@ fn create_anthropic_from_registry(
|
|||||||
reason: format!("Failed to create Anthropic client: {e}"),
|
reason: format!("Failed to create Anthropic client: {e}"),
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
|
// Resolve prompt cache retention from env (default: Short).
|
||||||
|
// Injects top-level cache_control via additional_params for Anthropic
|
||||||
|
// automatic caching (the API auto-places the breakpoint at the last
|
||||||
|
// cacheable block).
|
||||||
|
let cache_retention: CacheRetention = optional_env("ANTHROPIC_CACHE_RETENTION")
|
||||||
|
.ok()
|
||||||
|
.flatten()
|
||||||
|
.and_then(|val| match val.parse::<CacheRetention>() {
|
||||||
|
Ok(r) => Some(r),
|
||||||
|
Err(e) => {
|
||||||
|
tracing::warn!("Invalid ANTHROPIC_CACHE_RETENTION: {e}; defaulting to short");
|
||||||
|
None
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
let model = client.completion_model(&config.model);
|
let model = client.completion_model(&config.model);
|
||||||
|
|
||||||
|
if cache_retention != CacheRetention::None {
|
||||||
|
tracing::info!(
|
||||||
|
model = %config.model,
|
||||||
|
retention = %cache_retention,
|
||||||
|
"Anthropic automatic prompt caching enabled"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
tracing::info!(
|
tracing::info!(
|
||||||
provider = %config.provider_id,
|
provider = %config.provider_id,
|
||||||
model = %config.model,
|
model = %config.model,
|
||||||
@@ -209,7 +235,9 @@ fn create_anthropic_from_registry(
|
|||||||
"Using Anthropic provider"
|
"Using Anthropic provider"
|
||||||
);
|
);
|
||||||
|
|
||||||
Ok(Arc::new(RigAdapter::new(model, &config.model)))
|
Ok(Arc::new(
|
||||||
|
RigAdapter::new(model, &config.model).with_cache_retention(cache_retention),
|
||||||
|
))
|
||||||
}
|
}
|
||||||
|
|
||||||
fn create_ollama_from_registry(
|
fn create_ollama_from_registry(
|
||||||
|
|||||||
@@ -499,6 +499,8 @@ impl LlmProvider for NearAiChatProvider {
|
|||||||
finish_reason,
|
finish_reason,
|
||||||
input_tokens,
|
input_tokens,
|
||||||
output_tokens,
|
output_tokens,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -604,6 +606,8 @@ impl LlmProvider for NearAiChatProvider {
|
|||||||
finish_reason,
|
finish_reason,
|
||||||
input_tokens,
|
input_tokens,
|
||||||
output_tokens,
|
output_tokens,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -153,6 +153,12 @@ pub struct CompletionResponse {
|
|||||||
pub input_tokens: u32,
|
pub input_tokens: u32,
|
||||||
pub output_tokens: u32,
|
pub output_tokens: u32,
|
||||||
pub finish_reason: FinishReason,
|
pub finish_reason: FinishReason,
|
||||||
|
/// Tokens read from the provider's server-side prompt cache (Anthropic).
|
||||||
|
/// Zero when caching is not supported or on a cache miss.
|
||||||
|
pub cache_read_input_tokens: u32,
|
||||||
|
/// Tokens written to the provider's server-side prompt cache (Anthropic).
|
||||||
|
/// Zero when caching is not supported or no new prefix was cached.
|
||||||
|
pub cache_creation_input_tokens: u32,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Why the completion finished.
|
/// Why the completion finished.
|
||||||
@@ -254,6 +260,10 @@ pub struct ToolCompletionResponse {
|
|||||||
pub input_tokens: u32,
|
pub input_tokens: u32,
|
||||||
pub output_tokens: u32,
|
pub output_tokens: u32,
|
||||||
pub finish_reason: FinishReason,
|
pub finish_reason: FinishReason,
|
||||||
|
/// Tokens read from the provider's server-side prompt cache (Anthropic).
|
||||||
|
pub cache_read_input_tokens: u32,
|
||||||
|
/// Tokens written to the provider's server-side prompt cache (Anthropic).
|
||||||
|
pub cache_creation_input_tokens: u32,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Metadata about a model returned by the provider's API.
|
/// Metadata about a model returned by the provider's API.
|
||||||
@@ -328,6 +338,23 @@ pub trait LlmProvider: Send + Sync {
|
|||||||
let (input_cost, output_cost) = self.cost_per_token();
|
let (input_cost, output_cost) = self.cost_per_token();
|
||||||
input_cost * Decimal::from(input_tokens) + output_cost * Decimal::from(output_tokens)
|
input_cost * Decimal::from(input_tokens) + output_cost * Decimal::from(output_tokens)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Cost multiplier for cache-creation tokens (Anthropic prompt caching).
|
||||||
|
///
|
||||||
|
/// Returns `1.0` by default (no surcharge). Anthropic providers return
|
||||||
|
/// `1.25` for 5-minute TTL or `2.0` for 1-hour TTL.
|
||||||
|
fn cache_write_multiplier(&self) -> Decimal {
|
||||||
|
Decimal::ONE
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Discount divisor for cache-read tokens.
|
||||||
|
///
|
||||||
|
/// Cached-read cost = `input_rate / cache_read_discount()`.
|
||||||
|
/// Returns `1` by default (no discount). Anthropic returns `10` (90% off),
|
||||||
|
/// OpenAI would return `2` (50% off).
|
||||||
|
fn cache_read_discount(&self) -> Decimal {
|
||||||
|
Decimal::ONE
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Sanitize a message list to ensure tool_use / tool_result integrity.
|
/// Sanitize a message list to ensure tool_use / tool_result integrity.
|
||||||
|
|||||||
@@ -292,6 +292,10 @@ pub struct ToolSelection {
|
|||||||
pub struct TokenUsage {
|
pub struct TokenUsage {
|
||||||
pub input_tokens: u32,
|
pub input_tokens: u32,
|
||||||
pub output_tokens: u32,
|
pub output_tokens: u32,
|
||||||
|
/// Tokens served from the provider's server-side prompt cache (Anthropic).
|
||||||
|
pub cache_read_input_tokens: u32,
|
||||||
|
/// Tokens written to the provider's prompt cache (Anthropic).
|
||||||
|
pub cache_creation_input_tokens: u32,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl TokenUsage {
|
impl TokenUsage {
|
||||||
@@ -434,6 +438,8 @@ impl Reasoning {
|
|||||||
let usage = TokenUsage {
|
let usage = TokenUsage {
|
||||||
input_tokens: response.input_tokens,
|
input_tokens: response.input_tokens,
|
||||||
output_tokens: response.output_tokens,
|
output_tokens: response.output_tokens,
|
||||||
|
cache_read_input_tokens: response.cache_read_input_tokens,
|
||||||
|
cache_creation_input_tokens: response.cache_creation_input_tokens,
|
||||||
};
|
};
|
||||||
Ok((clean_response(&response.content), usage))
|
Ok((clean_response(&response.content), usage))
|
||||||
}
|
}
|
||||||
@@ -612,6 +618,8 @@ Respond in JSON format:
|
|||||||
let usage = TokenUsage {
|
let usage = TokenUsage {
|
||||||
input_tokens: response.input_tokens,
|
input_tokens: response.input_tokens,
|
||||||
output_tokens: response.output_tokens,
|
output_tokens: response.output_tokens,
|
||||||
|
cache_read_input_tokens: response.cache_read_input_tokens,
|
||||||
|
cache_creation_input_tokens: response.cache_creation_input_tokens,
|
||||||
};
|
};
|
||||||
|
|
||||||
// If there were tool calls, return them for execution
|
// If there were tool calls, return them for execution
|
||||||
@@ -690,6 +698,8 @@ Respond in JSON format:
|
|||||||
usage: TokenUsage {
|
usage: TokenUsage {
|
||||||
input_tokens: response.input_tokens,
|
input_tokens: response.input_tokens,
|
||||||
output_tokens: response.output_tokens,
|
output_tokens: response.output_tokens,
|
||||||
|
cache_read_input_tokens: response.cache_read_input_tokens,
|
||||||
|
cache_creation_input_tokens: response.cache_creation_input_tokens,
|
||||||
},
|
},
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -461,6 +461,14 @@ impl LlmProvider for RecordingLlm {
|
|||||||
self.inner.cost_per_token()
|
self.inner.cost_per_token()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn cache_write_multiplier(&self) -> Decimal {
|
||||||
|
self.inner.cache_write_multiplier()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cache_read_discount(&self) -> Decimal {
|
||||||
|
self.inner.cache_read_discount()
|
||||||
|
}
|
||||||
|
|
||||||
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||||
let (hint, tool_results) = self.capture_new_messages(&request.messages).await;
|
let (hint, tool_results) = self.capture_new_messages(&request.messages).await;
|
||||||
let response = self.inner.complete(request).await?;
|
let response = self.inner.complete(request).await?;
|
||||||
|
|||||||
@@ -181,6 +181,14 @@ impl LlmProvider for CachedProvider {
|
|||||||
self.inner.cost_per_token()
|
self.inner.cost_per_token()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn cache_write_multiplier(&self) -> Decimal {
|
||||||
|
self.inner.cache_write_multiplier()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cache_read_discount(&self) -> Decimal {
|
||||||
|
self.inner.cache_read_discount()
|
||||||
|
}
|
||||||
|
|
||||||
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||||
let effective_model = self.inner.effective_model_name(request.model.as_deref());
|
let effective_model = self.inner.effective_model_name(request.model.as_deref());
|
||||||
let key = cache_key(&effective_model, &request);
|
let key = cache_key(&effective_model, &request);
|
||||||
@@ -352,6 +360,8 @@ mod tests {
|
|||||||
input_tokens: 1,
|
input_tokens: 1,
|
||||||
output_tokens: 1,
|
output_tokens: 1,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -365,6 +375,8 @@ mod tests {
|
|||||||
input_tokens: 1,
|
input_tokens: 1,
|
||||||
output_tokens: 1,
|
output_tokens: 1,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -109,6 +109,14 @@ impl LlmProvider for RetryProvider {
|
|||||||
self.inner.cost_per_token()
|
self.inner.cost_per_token()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn cache_write_multiplier(&self) -> Decimal {
|
||||||
|
self.inner.cache_write_multiplier()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cache_read_discount(&self) -> Decimal {
|
||||||
|
self.inner.cache_read_discount()
|
||||||
|
}
|
||||||
|
|
||||||
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||||
let mut last_error: Option<LlmError> = None;
|
let mut last_error: Option<LlmError> = None;
|
||||||
|
|
||||||
|
|||||||
+255
-5
@@ -3,6 +3,7 @@
|
|||||||
//! This lets us use any rig-core provider (OpenAI, Anthropic, Ollama, etc.) as an
|
//! This lets us use any rig-core provider (OpenAI, Anthropic, Ollama, etc.) as an
|
||||||
//! `Arc<dyn LlmProvider>` without changing any of the agent, reasoning, or tool code.
|
//! `Arc<dyn LlmProvider>` without changing any of the agent, reasoning, or tool code.
|
||||||
|
|
||||||
|
use crate::config::CacheRetention;
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use rig::OneOrMany;
|
use rig::OneOrMany;
|
||||||
use rig::completion::{
|
use rig::completion::{
|
||||||
@@ -14,6 +15,7 @@ use rig::message::{
|
|||||||
ToolResultContent, UserContent,
|
ToolResultContent, UserContent,
|
||||||
};
|
};
|
||||||
use rust_decimal::Decimal;
|
use rust_decimal::Decimal;
|
||||||
|
use rust_decimal_macros::dec;
|
||||||
use serde::Serialize;
|
use serde::Serialize;
|
||||||
use serde::de::DeserializeOwned;
|
use serde::de::DeserializeOwned;
|
||||||
use serde_json::Value as JsonValue;
|
use serde_json::Value as JsonValue;
|
||||||
@@ -34,6 +36,11 @@ pub struct RigAdapter<M: CompletionModel> {
|
|||||||
model_name: String,
|
model_name: String,
|
||||||
input_cost: Decimal,
|
input_cost: Decimal,
|
||||||
output_cost: Decimal,
|
output_cost: Decimal,
|
||||||
|
/// Prompt cache retention policy (Anthropic only).
|
||||||
|
/// When not `CacheRetention::None`, injects top-level `cache_control`
|
||||||
|
/// via `additional_params` for Anthropic automatic caching. Also controls
|
||||||
|
/// the cost multiplier for cache-creation tokens.
|
||||||
|
cache_retention: CacheRetention,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<M: CompletionModel> RigAdapter<M> {
|
impl<M: CompletionModel> RigAdapter<M> {
|
||||||
@@ -47,8 +54,35 @@ impl<M: CompletionModel> RigAdapter<M> {
|
|||||||
model_name: name,
|
model_name: name,
|
||||||
input_cost,
|
input_cost,
|
||||||
output_cost,
|
output_cost,
|
||||||
|
cache_retention: CacheRetention::None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Set Anthropic prompt cache retention policy.
|
||||||
|
///
|
||||||
|
/// Controls both cache injection and cost tracking:
|
||||||
|
/// - `None` — no caching, no surcharge (1.0×).
|
||||||
|
/// - `Short` — 5-minute TTL via `{"type": "ephemeral"}`, 1.25× write surcharge.
|
||||||
|
/// - `Long` — 1-hour TTL via `{"type": "ephemeral", "ttl": "1h"}`, 2.0× write surcharge.
|
||||||
|
///
|
||||||
|
/// Cache injection uses Anthropic's **automatic caching** — a top-level
|
||||||
|
/// `cache_control` field in `additional_params` that gets `#[serde(flatten)]`'d
|
||||||
|
/// into the request body by rig-core.
|
||||||
|
///
|
||||||
|
/// If the configured model does not support caching (e.g. claude-2),
|
||||||
|
/// a warning is logged once at construction and caching is disabled.
|
||||||
|
pub fn with_cache_retention(mut self, retention: CacheRetention) -> Self {
|
||||||
|
if retention != CacheRetention::None && !supports_prompt_cache(&self.model_name) {
|
||||||
|
tracing::warn!(
|
||||||
|
model = %self.model_name,
|
||||||
|
"Prompt caching requested but model does not support it; disabling"
|
||||||
|
);
|
||||||
|
self.cache_retention = CacheRetention::None;
|
||||||
|
} else {
|
||||||
|
self.cache_retention = retention;
|
||||||
|
}
|
||||||
|
self
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// -- Type conversion helpers --
|
// -- Type conversion helpers --
|
||||||
@@ -360,7 +394,44 @@ fn saturate_u32(val: u64) -> u32 {
|
|||||||
val.min(u32::MAX as u64) as u32
|
val.min(u32::MAX as u64) as u32
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Returns `true` if the model supports Anthropic prompt caching.
|
||||||
|
///
|
||||||
|
/// Per Anthropic docs, only Claude 3+ models support prompt caching.
|
||||||
|
/// Unsupported: claude-2, claude-2.1, claude-instant-*.
|
||||||
|
fn supports_prompt_cache(name: &str) -> bool {
|
||||||
|
let lower = name.to_lowercase();
|
||||||
|
// Strip optional provider prefix (e.g. "anthropic/claude-...")
|
||||||
|
let model = lower.strip_prefix("anthropic/").unwrap_or(&lower);
|
||||||
|
// Only Claude 3+ families support prompt caching
|
||||||
|
model.starts_with("claude-3")
|
||||||
|
|| model.starts_with("claude-4")
|
||||||
|
|| model.starts_with("claude-sonnet")
|
||||||
|
|| model.starts_with("claude-opus")
|
||||||
|
|| model.starts_with("claude-haiku")
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Extract `cache_creation_input_tokens` from the raw provider response.
|
||||||
|
///
|
||||||
|
/// Rig-core's unified `Usage` does not surface this field, but Anthropic's raw
|
||||||
|
/// response includes it at `usage.cache_creation_input_tokens`. We serialize the
|
||||||
|
/// raw response to JSON and attempt to read the value.
|
||||||
|
fn extract_cache_creation<T: Serialize>(raw: &T) -> u32 {
|
||||||
|
serde_json::to_value(raw)
|
||||||
|
.ok()
|
||||||
|
.and_then(|v| v.get("usage")?.get("cache_creation_input_tokens")?.as_u64())
|
||||||
|
.map(|n| n.min(u32::MAX as u64) as u32)
|
||||||
|
.unwrap_or(0)
|
||||||
|
}
|
||||||
|
|
||||||
/// Build a rig-core CompletionRequest from our internal types.
|
/// Build a rig-core CompletionRequest from our internal types.
|
||||||
|
///
|
||||||
|
/// When `cache_retention` is not `None`, injects a top-level `cache_control`
|
||||||
|
/// field via `additional_params`. Rig-core's `AnthropicCompletionRequest`
|
||||||
|
/// uses `#[serde(flatten)]` on `additional_params`, so the field lands at
|
||||||
|
/// the request root — which is exactly what Anthropic's **automatic caching**
|
||||||
|
/// expects. The API auto-places the cache breakpoint at the last cacheable
|
||||||
|
/// block and moves it forward as conversations grow.
|
||||||
|
#[allow(clippy::too_many_arguments)]
|
||||||
fn build_rig_request(
|
fn build_rig_request(
|
||||||
preamble: Option<String>,
|
preamble: Option<String>,
|
||||||
mut history: Vec<RigMessage>,
|
mut history: Vec<RigMessage>,
|
||||||
@@ -368,6 +439,7 @@ fn build_rig_request(
|
|||||||
tool_choice: Option<RigToolChoice>,
|
tool_choice: Option<RigToolChoice>,
|
||||||
temperature: Option<f32>,
|
temperature: Option<f32>,
|
||||||
max_tokens: Option<u32>,
|
max_tokens: Option<u32>,
|
||||||
|
cache_retention: CacheRetention,
|
||||||
) -> Result<RigRequest, LlmError> {
|
) -> Result<RigRequest, LlmError> {
|
||||||
// rig-core requires at least one message in chat_history
|
// rig-core requires at least one message in chat_history
|
||||||
if history.is_empty() {
|
if history.is_empty() {
|
||||||
@@ -379,6 +451,17 @@ fn build_rig_request(
|
|||||||
reason: format!("Failed to build chat history: {}", e),
|
reason: format!("Failed to build chat history: {}", e),
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
|
// Inject top-level cache_control for Anthropic automatic prompt caching.
|
||||||
|
let additional_params = match cache_retention {
|
||||||
|
CacheRetention::None => None,
|
||||||
|
CacheRetention::Short => Some(serde_json::json!({
|
||||||
|
"cache_control": {"type": "ephemeral"}
|
||||||
|
})),
|
||||||
|
CacheRetention::Long => Some(serde_json::json!({
|
||||||
|
"cache_control": {"type": "ephemeral", "ttl": "1h"}
|
||||||
|
})),
|
||||||
|
};
|
||||||
|
|
||||||
Ok(RigRequest {
|
Ok(RigRequest {
|
||||||
preamble,
|
preamble,
|
||||||
chat_history,
|
chat_history,
|
||||||
@@ -387,7 +470,7 @@ fn build_rig_request(
|
|||||||
temperature: temperature.map(|t| t as f64),
|
temperature: temperature.map(|t| t as f64),
|
||||||
max_tokens: max_tokens.map(|t| t as u64),
|
max_tokens: max_tokens.map(|t| t as u64),
|
||||||
tool_choice,
|
tool_choice,
|
||||||
additional_params: None,
|
additional_params,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -405,6 +488,22 @@ where
|
|||||||
(self.input_cost, self.output_cost)
|
(self.input_cost, self.output_cost)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn cache_write_multiplier(&self) -> Decimal {
|
||||||
|
match self.cache_retention {
|
||||||
|
CacheRetention::None => Decimal::ONE,
|
||||||
|
CacheRetention::Short => Decimal::new(125, 2), // 1.25× (125% of input rate)
|
||||||
|
CacheRetention::Long => Decimal::TWO, // 2.0× (200% of input rate)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cache_read_discount(&self) -> Decimal {
|
||||||
|
if self.cache_retention != CacheRetention::None {
|
||||||
|
dec!(10) // Anthropic: 90% discount (cost = input_rate / 10)
|
||||||
|
} else {
|
||||||
|
Decimal::ONE
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||||
if let Some(requested_model) = request.model.as_deref()
|
if let Some(requested_model) = request.model.as_deref()
|
||||||
&& requested_model != self.model_name.as_str()
|
&& requested_model != self.model_name.as_str()
|
||||||
@@ -427,6 +526,7 @@ where
|
|||||||
None,
|
None,
|
||||||
request.temperature,
|
request.temperature,
|
||||||
request.max_tokens,
|
request.max_tokens,
|
||||||
|
self.cache_retention,
|
||||||
)?;
|
)?;
|
||||||
|
|
||||||
let response =
|
let response =
|
||||||
@@ -440,12 +540,26 @@ where
|
|||||||
|
|
||||||
let (text, _tool_calls, finish) = extract_response(&response.choice, &response.usage);
|
let (text, _tool_calls, finish) = extract_response(&response.choice, &response.usage);
|
||||||
|
|
||||||
Ok(CompletionResponse {
|
let resp = CompletionResponse {
|
||||||
content: text.unwrap_or_default(),
|
content: text.unwrap_or_default(),
|
||||||
input_tokens: saturate_u32(response.usage.input_tokens),
|
input_tokens: saturate_u32(response.usage.input_tokens),
|
||||||
output_tokens: saturate_u32(response.usage.output_tokens),
|
output_tokens: saturate_u32(response.usage.output_tokens),
|
||||||
finish_reason: finish,
|
finish_reason: finish,
|
||||||
})
|
cache_read_input_tokens: saturate_u32(response.usage.cached_input_tokens),
|
||||||
|
cache_creation_input_tokens: extract_cache_creation(&response.raw_response),
|
||||||
|
};
|
||||||
|
|
||||||
|
if resp.cache_read_input_tokens > 0 {
|
||||||
|
tracing::debug!(
|
||||||
|
model = %self.model_name,
|
||||||
|
input = resp.input_tokens,
|
||||||
|
output = resp.output_tokens,
|
||||||
|
cache_read = resp.cache_read_input_tokens,
|
||||||
|
"prompt cache hit",
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(resp)
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn complete_with_tools(
|
async fn complete_with_tools(
|
||||||
@@ -478,6 +592,7 @@ where
|
|||||||
tool_choice,
|
tool_choice,
|
||||||
request.temperature,
|
request.temperature,
|
||||||
request.max_tokens,
|
request.max_tokens,
|
||||||
|
self.cache_retention,
|
||||||
)?;
|
)?;
|
||||||
|
|
||||||
let response =
|
let response =
|
||||||
@@ -504,13 +619,27 @@ where
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(ToolCompletionResponse {
|
let resp = ToolCompletionResponse {
|
||||||
content: text,
|
content: text,
|
||||||
tool_calls,
|
tool_calls,
|
||||||
input_tokens: saturate_u32(response.usage.input_tokens),
|
input_tokens: saturate_u32(response.usage.input_tokens),
|
||||||
output_tokens: saturate_u32(response.usage.output_tokens),
|
output_tokens: saturate_u32(response.usage.output_tokens),
|
||||||
finish_reason: finish,
|
finish_reason: finish,
|
||||||
})
|
cache_read_input_tokens: saturate_u32(response.usage.cached_input_tokens),
|
||||||
|
cache_creation_input_tokens: extract_cache_creation(&response.raw_response),
|
||||||
|
};
|
||||||
|
|
||||||
|
if resp.cache_read_input_tokens > 0 {
|
||||||
|
tracing::debug!(
|
||||||
|
model = %self.model_name,
|
||||||
|
input = resp.input_tokens,
|
||||||
|
output = resp.output_tokens,
|
||||||
|
cache_read = resp.cache_read_input_tokens,
|
||||||
|
"prompt cache hit",
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(resp)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn active_model_name(&self) -> String {
|
fn active_model_name(&self) -> String {
|
||||||
@@ -869,4 +998,125 @@ mod tests {
|
|||||||
let known = HashSet::from(["echo".to_string()]);
|
let known = HashSet::from(["echo".to_string()]);
|
||||||
assert_eq!(normalize_tool_name("other_tool", &known), "other_tool");
|
assert_eq!(normalize_tool_name("other_tool", &known), "other_tool");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_build_rig_request_injects_cache_control_short() {
|
||||||
|
let req = build_rig_request(
|
||||||
|
Some("You are helpful.".to_string()),
|
||||||
|
vec![RigMessage::user("Hello")],
|
||||||
|
Vec::new(),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
CacheRetention::Short,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let params = req
|
||||||
|
.additional_params
|
||||||
|
.expect("should have additional_params for Short retention");
|
||||||
|
assert_eq!(params["cache_control"]["type"], "ephemeral");
|
||||||
|
assert!(
|
||||||
|
params["cache_control"].get("ttl").is_none(),
|
||||||
|
"Short retention should not include ttl"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_build_rig_request_injects_cache_control_long() {
|
||||||
|
let req = build_rig_request(
|
||||||
|
Some("You are helpful.".to_string()),
|
||||||
|
vec![RigMessage::user("Hello")],
|
||||||
|
Vec::new(),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
CacheRetention::Long,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let params = req
|
||||||
|
.additional_params
|
||||||
|
.expect("should have additional_params for Long retention");
|
||||||
|
assert_eq!(params["cache_control"]["type"], "ephemeral");
|
||||||
|
assert_eq!(params["cache_control"]["ttl"], "1h");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_build_rig_request_no_cache_control_when_none() {
|
||||||
|
let req = build_rig_request(
|
||||||
|
Some("You are helpful.".to_string()),
|
||||||
|
vec![RigMessage::user("Hello")],
|
||||||
|
Vec::new(),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
CacheRetention::None,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
req.additional_params.is_none(),
|
||||||
|
"additional_params should be None when cache is disabled"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Verify that the multiplier match arms in `RigAdapter::cache_write_multiplier`
|
||||||
|
/// produce the expected values. We use a standalone helper because constructing
|
||||||
|
/// a real `RigAdapter` requires a rig `Model` (which needs network/provider setup).
|
||||||
|
/// The helper mirrors the same match expression — if the impl drifts, the
|
||||||
|
/// `test_build_rig_request_*` tests will still catch regressions end-to-end.
|
||||||
|
#[test]
|
||||||
|
fn test_cache_write_multiplier_values() {
|
||||||
|
use rust_decimal::Decimal;
|
||||||
|
// None → 1.0× (no surcharge)
|
||||||
|
assert_eq!(
|
||||||
|
cache_write_multiplier_for(CacheRetention::None),
|
||||||
|
Decimal::ONE
|
||||||
|
);
|
||||||
|
// Short → 1.25× (25% surcharge)
|
||||||
|
assert_eq!(
|
||||||
|
cache_write_multiplier_for(CacheRetention::Short),
|
||||||
|
Decimal::new(125, 2)
|
||||||
|
);
|
||||||
|
// Long → 2.0× (100% surcharge)
|
||||||
|
assert_eq!(
|
||||||
|
cache_write_multiplier_for(CacheRetention::Long),
|
||||||
|
Decimal::TWO
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cache_write_multiplier_for(retention: CacheRetention) -> rust_decimal::Decimal {
|
||||||
|
match retention {
|
||||||
|
CacheRetention::None => rust_decimal::Decimal::ONE,
|
||||||
|
CacheRetention::Short => rust_decimal::Decimal::new(125, 2),
|
||||||
|
CacheRetention::Long => rust_decimal::Decimal::TWO,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// -- supports_prompt_cache tests --
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_supports_prompt_cache_supported_models() {
|
||||||
|
// All Claude 3+ models per Anthropic docs
|
||||||
|
assert!(supports_prompt_cache("claude-opus-4-6"));
|
||||||
|
assert!(supports_prompt_cache("claude-sonnet-4-6"));
|
||||||
|
assert!(supports_prompt_cache("claude-sonnet-4"));
|
||||||
|
assert!(supports_prompt_cache("claude-haiku-4-5"));
|
||||||
|
assert!(supports_prompt_cache("claude-3-5-sonnet-20241022"));
|
||||||
|
assert!(supports_prompt_cache("claude-haiku-3"));
|
||||||
|
assert!(supports_prompt_cache("Claude-Opus-4-5")); // case-insensitive
|
||||||
|
assert!(supports_prompt_cache("anthropic/claude-sonnet-4-6")); // provider prefix
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_supports_prompt_cache_unsupported_models() {
|
||||||
|
// Legacy Claude models that predate caching
|
||||||
|
assert!(!supports_prompt_cache("claude-2"));
|
||||||
|
assert!(!supports_prompt_cache("claude-2.1"));
|
||||||
|
assert!(!supports_prompt_cache("claude-instant-1.2"));
|
||||||
|
// Non-Claude models
|
||||||
|
assert!(!supports_prompt_cache("gpt-4o"));
|
||||||
|
assert!(!supports_prompt_cache("llama3"));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -857,6 +857,14 @@ impl LlmProvider for SmartRoutingProvider {
|
|||||||
self.primary.cost_per_token()
|
self.primary.cost_per_token()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn cache_write_multiplier(&self) -> Decimal {
|
||||||
|
self.primary.cache_write_multiplier()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cache_read_discount(&self) -> Decimal {
|
||||||
|
self.primary.cache_read_discount()
|
||||||
|
}
|
||||||
|
|
||||||
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||||
self.stats.total_requests.fetch_add(1, Ordering::Relaxed);
|
self.stats.total_requests.fetch_add(1, Ordering::Relaxed);
|
||||||
|
|
||||||
@@ -1471,6 +1479,8 @@ mod tests {
|
|||||||
input_tokens: 10,
|
input_tokens: 10,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
finish_reason: crate::llm::FinishReason::Stop,
|
finish_reason: crate::llm::FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
};
|
};
|
||||||
assert!(SmartRoutingProvider::response_is_uncertain(&response));
|
assert!(SmartRoutingProvider::response_is_uncertain(&response));
|
||||||
}
|
}
|
||||||
@@ -1482,6 +1492,8 @@ mod tests {
|
|||||||
input_tokens: 10,
|
input_tokens: 10,
|
||||||
output_tokens: 0,
|
output_tokens: 0,
|
||||||
finish_reason: crate::llm::FinishReason::Stop,
|
finish_reason: crate::llm::FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
};
|
};
|
||||||
assert!(SmartRoutingProvider::response_is_uncertain(&response));
|
assert!(SmartRoutingProvider::response_is_uncertain(&response));
|
||||||
}
|
}
|
||||||
@@ -1493,6 +1505,8 @@ mod tests {
|
|||||||
input_tokens: 10,
|
input_tokens: 10,
|
||||||
output_tokens: 1,
|
output_tokens: 1,
|
||||||
finish_reason: crate::llm::FinishReason::Stop,
|
finish_reason: crate::llm::FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
};
|
};
|
||||||
assert!(!SmartRoutingProvider::response_is_uncertain(&response));
|
assert!(!SmartRoutingProvider::response_is_uncertain(&response));
|
||||||
}
|
}
|
||||||
@@ -1505,6 +1519,8 @@ mod tests {
|
|||||||
input_tokens: 10,
|
input_tokens: 10,
|
||||||
output_tokens: 20,
|
output_tokens: 20,
|
||||||
finish_reason: crate::llm::FinishReason::Stop,
|
finish_reason: crate::llm::FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
};
|
};
|
||||||
assert!(!SmartRoutingProvider::response_is_uncertain(&response));
|
assert!(!SmartRoutingProvider::response_is_uncertain(&response));
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -160,6 +160,8 @@ async fn llm_complete(
|
|||||||
input_tokens: resp.input_tokens,
|
input_tokens: resp.input_tokens,
|
||||||
output_tokens: resp.output_tokens,
|
output_tokens: resp.output_tokens,
|
||||||
finish_reason: format_finish_reason(resp.finish_reason),
|
finish_reason: format_finish_reason(resp.finish_reason),
|
||||||
|
cache_read_input_tokens: resp.cache_read_input_tokens,
|
||||||
|
cache_creation_input_tokens: resp.cache_creation_input_tokens,
|
||||||
}))
|
}))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -189,6 +191,8 @@ async fn llm_complete_with_tools(
|
|||||||
input_tokens: resp.input_tokens,
|
input_tokens: resp.input_tokens,
|
||||||
output_tokens: resp.output_tokens,
|
output_tokens: resp.output_tokens,
|
||||||
finish_reason: format_finish_reason(resp.finish_reason),
|
finish_reason: format_finish_reason(resp.finish_reason),
|
||||||
|
cache_read_input_tokens: resp.cache_read_input_tokens,
|
||||||
|
cache_creation_input_tokens: resp.cache_creation_input_tokens,
|
||||||
}))
|
}))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -174,6 +174,8 @@ impl LlmProvider for StubLlm {
|
|||||||
input_tokens: 10,
|
input_tokens: 10,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -191,6 +193,8 @@ impl LlmProvider for StubLlm {
|
|||||||
input_tokens: 10,
|
input_tokens: 10,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -52,6 +52,10 @@ pub struct ProxyCompletionResponse {
|
|||||||
pub input_tokens: u32,
|
pub input_tokens: u32,
|
||||||
pub output_tokens: u32,
|
pub output_tokens: u32,
|
||||||
pub finish_reason: String,
|
pub finish_reason: String,
|
||||||
|
#[serde(default)]
|
||||||
|
pub cache_read_input_tokens: u32,
|
||||||
|
#[serde(default)]
|
||||||
|
pub cache_creation_input_tokens: u32,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Serialize, Deserialize)]
|
#[derive(Debug, Serialize, Deserialize)]
|
||||||
@@ -71,6 +75,10 @@ pub struct ProxyToolCompletionResponse {
|
|||||||
pub input_tokens: u32,
|
pub input_tokens: u32,
|
||||||
pub output_tokens: u32,
|
pub output_tokens: u32,
|
||||||
pub finish_reason: String,
|
pub finish_reason: String,
|
||||||
|
#[serde(default)]
|
||||||
|
pub cache_read_input_tokens: u32,
|
||||||
|
#[serde(default)]
|
||||||
|
pub cache_creation_input_tokens: u32,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Completion result for the worker to report when done.
|
/// Completion result for the worker to report when done.
|
||||||
@@ -227,6 +235,8 @@ impl WorkerHttpClient {
|
|||||||
input_tokens: proxy_resp.input_tokens,
|
input_tokens: proxy_resp.input_tokens,
|
||||||
output_tokens: proxy_resp.output_tokens,
|
output_tokens: proxy_resp.output_tokens,
|
||||||
finish_reason: parse_finish_reason(&proxy_resp.finish_reason),
|
finish_reason: parse_finish_reason(&proxy_resp.finish_reason),
|
||||||
|
cache_read_input_tokens: proxy_resp.cache_read_input_tokens,
|
||||||
|
cache_creation_input_tokens: proxy_resp.cache_creation_input_tokens,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -254,6 +264,8 @@ impl WorkerHttpClient {
|
|||||||
input_tokens: proxy_resp.input_tokens,
|
input_tokens: proxy_resp.input_tokens,
|
||||||
output_tokens: proxy_resp.output_tokens,
|
output_tokens: proxy_resp.output_tokens,
|
||||||
finish_reason: parse_finish_reason(&proxy_resp.finish_reason),
|
finish_reason: parse_finish_reason(&proxy_resp.finish_reason),
|
||||||
|
cache_read_input_tokens: proxy_resp.cache_read_input_tokens,
|
||||||
|
cache_creation_input_tokens: proxy_resp.cache_creation_input_tokens,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -70,6 +70,8 @@ impl LlmProvider for MockLlmProvider {
|
|||||||
input_tokens: 10,
|
input_tokens: 10,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -95,6 +97,8 @@ impl LlmProvider for MockLlmProvider {
|
|||||||
input_tokens: 15,
|
input_tokens: 15,
|
||||||
output_tokens: 8,
|
output_tokens: 8,
|
||||||
finish_reason: FinishReason::ToolUse,
|
finish_reason: FinishReason::ToolUse,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
} else {
|
} else {
|
||||||
Ok(ToolCompletionResponse {
|
Ok(ToolCompletionResponse {
|
||||||
@@ -103,6 +107,8 @@ impl LlmProvider for MockLlmProvider {
|
|||||||
input_tokens: 10,
|
input_tokens: 10,
|
||||||
output_tokens: 4,
|
output_tokens: 4,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -141,6 +147,8 @@ impl LlmProvider for FixedModelProvider {
|
|||||||
input_tokens: 10,
|
input_tokens: 10,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -154,6 +162,8 @@ impl LlmProvider for FixedModelProvider {
|
|||||||
input_tokens: 10,
|
input_tokens: 10,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -88,6 +88,8 @@ impl LlmProvider for FlakeyProvider {
|
|||||||
input_tokens: 10,
|
input_tokens: 10,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -115,6 +117,8 @@ impl LlmProvider for FlakeyProvider {
|
|||||||
input_tokens: 10,
|
input_tokens: 10,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -192,6 +196,8 @@ impl LlmProvider for GarbageProvider {
|
|||||||
input_tokens: 0,
|
input_tokens: 0,
|
||||||
output_tokens: 0,
|
output_tokens: 0,
|
||||||
finish_reason: FinishReason::Unknown,
|
finish_reason: FinishReason::Unknown,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -206,6 +212,8 @@ impl LlmProvider for GarbageProvider {
|
|||||||
input_tokens: 0,
|
input_tokens: 0,
|
||||||
output_tokens: 0,
|
output_tokens: 0,
|
||||||
finish_reason: FinishReason::Unknown,
|
finish_reason: FinishReason::Unknown,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -248,6 +256,8 @@ impl LlmProvider for ReliableProvider {
|
|||||||
input_tokens: 10,
|
input_tokens: 10,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -262,6 +272,8 @@ impl LlmProvider for ReliableProvider {
|
|||||||
input_tokens: 10,
|
input_tokens: 10,
|
||||||
output_tokens: 5,
|
output_tokens: 5,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -578,6 +578,8 @@ impl LlmProvider for TraceLlm {
|
|||||||
input_tokens,
|
input_tokens,
|
||||||
output_tokens,
|
output_tokens,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
}),
|
}),
|
||||||
TraceResponse::ToolCalls { .. } => Err(LlmError::RequestFailed {
|
TraceResponse::ToolCalls { .. } => Err(LlmError::RequestFailed {
|
||||||
provider: self.model_name.clone(),
|
provider: self.model_name.clone(),
|
||||||
@@ -610,6 +612,8 @@ impl LlmProvider for TraceLlm {
|
|||||||
input_tokens,
|
input_tokens,
|
||||||
output_tokens,
|
output_tokens,
|
||||||
finish_reason: FinishReason::Stop,
|
finish_reason: FinishReason::Stop,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
}),
|
}),
|
||||||
TraceResponse::ToolCalls {
|
TraceResponse::ToolCalls {
|
||||||
tool_calls,
|
tool_calls,
|
||||||
@@ -630,6 +634,8 @@ impl LlmProvider for TraceLlm {
|
|||||||
input_tokens,
|
input_tokens,
|
||||||
output_tokens,
|
output_tokens,
|
||||||
finish_reason: FinishReason::ToolUse,
|
finish_reason: FinishReason::ToolUse,
|
||||||
|
cache_read_input_tokens: 0,
|
||||||
|
cache_creation_input_tokens: 0,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
TraceResponse::UserInput { .. } => Err(LlmError::RequestFailed {
|
TraceResponse::UserInput { .. } => Err(LlmError::RequestFailed {
|
||||||
|
|||||||
Reference in New Issue
Block a user