gitops/optimclaw
1
0
Fork 0
mirror of https://github.com/outbackdingo/optimclaw.git synced 2026-08-25 14:53:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: enable Anthropic prompt caching via automatic cache_control injection (#660)

Browse Source 
* 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:
Illia Polosukhin
2026-03-07 09:10:05 +00:00
committed by GitHub
co-authored by Andrey Andrey Gruzdev Claude Opus 4.6
parent 633b234e44
commit 424a0366a9
22 changed files with 831 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.env.example
+11
View File
@@ -57,6 +57,17 @@ NEARAI_AUTH_URL=https://private.near.ai
# LLM_BASE_URL=https://api.fireworks.ai/inference/v1
# 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
# Channel Configuration
src/agent/cost_guard.rs
+230 -11
View File
@@ -151,21 +151,46 @@ impl CostGuard {
/// Record a completed LLM action: its token costs and the action timestamp.
///
/// 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-
/// sourced pricing). When `None`, falls back to the static `costs::model_cost`
/// lookup table, then `costs::default_cost`.
#[allow(clippy::too_many_arguments)]
pub async fn record_llm_call(
&self,
model: &str,
input_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)>,
) -> Decimal {
let (input_rate, output_rate) = cost_per_token
.unwrap_or_else(|| costs::model_cost(model).unwrap_or_else(costs::default_cost));
let cost =
input_rate * Decimal::from(input_tokens) + output_rate * Decimal::from(output_tokens);
// Cached read tokens cost input_rate / cache_read_discount (provider-specific).
// 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)
{
@@ -267,7 +292,16 @@ mod tests {
// Record a big call, still allowed
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;
assert!(guard.check_allowed().await.is_ok());
}
@@ -285,7 +319,18 @@ mod tests {
// Record a call that costs more than $0.01
// gpt-4o: input=$0.0000025/tok, output=$0.00001/tok
// 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
let result = guard.check_allowed().await;
@@ -308,7 +353,9 @@ mod tests {
// First 3 actions allowed
for _ in 0..3 {
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
@@ -329,7 +376,9 @@ mod tests {
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_eq!(guard.daily_spend().await, cost);
}
@@ -340,8 +389,12 @@ mod tests {
assert_eq!(guard.actions_this_hour().await, 0);
guard.record_llm_call("gpt-4o", 10, 10, None).await;
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;
guard
.record_llm_call("gpt-4o", 10, 10, 0, 0, Decimal::ONE, Decimal::ONE, None)
.await;
assert_eq!(guard.actions_this_hour().await, 2);
}
@@ -378,10 +431,23 @@ mod tests {
assert!(guard.model_usage().await.is_empty());
// 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
.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;
let usage = guard.model_usage().await;
@@ -402,4 +468,157 @@ mod tests {
// Costs should differ since models have different pricing
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"
);
}
}
src/agent/dispatcher.rs
+22
View File
@@ -278,12 +278,18 @@ impl Agent {
// Record cost and track token usage
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
.cost_guard()
.record_llm_call(
&model_name,
output.usage.input_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()),
)
.await;
@@ -1062,6 +1068,8 @@ mod tests {
input_tokens: 0,
output_tokens: 0,
finish_reason: FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
}
@@ -1075,6 +1083,8 @@ mod tests {
input_tokens: 0,
output_tokens: 0,
finish_reason: FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
}
}
@@ -1635,6 +1645,8 @@ mod tests {
input_tokens: 0,
output_tokens: 5,
finish_reason: FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
}
@@ -1650,6 +1662,8 @@ mod tests {
input_tokens: 0,
output_tokens: 5,
finish_reason: FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
});
}
// Tools available: always call one.
@@ -1663,6 +1677,8 @@ mod tests {
input_tokens: 0,
output_tokens: 5,
finish_reason: FinishReason::ToolUse,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
}
}
@@ -1787,6 +1803,8 @@ mod tests {
input_tokens: 0,
output_tokens: 2,
finish_reason: FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
}
@@ -1801,6 +1819,8 @@ mod tests {
input_tokens: 0,
output_tokens: 2,
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.
@@ -1814,6 +1834,8 @@ mod tests {
input_tokens: 0,
output_tokens: 5,
finish_reason: FinishReason::ToolUse,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
}
}
src/config/llm.rs
+126
View File
@@ -9,6 +9,50 @@ use crate::llm::registry::{ProviderProtocol, ProviderRegistry};
use crate::llm::session::SessionConfig;
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.
///
/// This single struct replaces what used to be five separate config types
@@ -755,4 +799,86 @@ mod tests {
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}");
}
}
}
src/config/mod.rs
+1 -1
View File
@@ -36,7 +36,7 @@ pub use self::database::{DatabaseBackend, DatabaseConfig, SslMode, default_libsq
pub use self::embeddings::EmbeddingsConfig;
pub use self::heartbeat::HeartbeatConfig;
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::safety::SafetyConfig;
pub use self::sandbox::{ClaudeCodeConfig, SandboxModeConfig};
src/llm/circuit_breaker.rs
+8
View File
@@ -245,6 +245,14 @@ impl LlmProvider for CircuitBreakerProvider {
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> {
self.check_allowed().await?;
match self.inner.complete(request).await {
src/llm/failover.rs
+16
View File
@@ -296,6 +296,14 @@ impl LlmProvider for FailoverProvider {
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> {
let (provider_idx, response) = self
.try_providers(|provider| {
@@ -404,6 +412,8 @@ mod tests {
input_tokens: 10,
output_tokens: 5,
finish_reason: FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
}))),
tool_complete_result: Mutex::new(Some(Ok(ToolCompletionResponse {
content: Some(content.to_string()),
@@ -411,6 +421,8 @@ mod tests {
input_tokens: 10,
output_tokens: 5,
finish_reason: FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
}))),
}
}
@@ -792,6 +804,8 @@ mod tests {
input_tokens: 10,
output_tokens: 5,
finish_reason: FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
}
@@ -817,6 +831,8 @@ mod tests {
input_tokens: 10,
output_tokens: 5,
finish_reason: FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
}
src/llm/mod.rs
+29 -1
View File
@@ -177,6 +177,8 @@ fn create_openai_compat_from_registry(
fn create_anthropic_from_registry(
config: &RegistryProviderConfig,
) -> Result<Arc<dyn LlmProvider>, LlmError> {
use crate::config::CacheRetention;
use crate::config::helpers::optional_env;
use rig::providers::anthropic;
let api_key = config
@@ -200,8 +202,32 @@ fn create_anthropic_from_registry(
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);
if cache_retention != CacheRetention::None {
tracing::info!(
model = %config.model,
retention = %cache_retention,
"Anthropic automatic prompt caching enabled"
);
}
tracing::info!(
provider = %config.provider_id,
model = %config.model,
@@ -209,7 +235,9 @@ fn create_anthropic_from_registry(
"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(
src/llm/nearai_chat.rs
+4
View File
@@ -499,6 +499,8 @@ impl LlmProvider for NearAiChatProvider {
finish_reason,
input_tokens,
output_tokens,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
}
@@ -604,6 +606,8 @@ impl LlmProvider for NearAiChatProvider {
finish_reason,
input_tokens,
output_tokens,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
}
src/llm/provider.rs
+27
View File
@@ -153,6 +153,12 @@ pub struct CompletionResponse {
pub input_tokens: u32,
pub output_tokens: u32,
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.
@@ -254,6 +260,10 @@ pub struct ToolCompletionResponse {
pub input_tokens: u32,
pub output_tokens: u32,
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.
@@ -328,6 +338,23 @@ pub trait LlmProvider: Send + Sync {
let (input_cost, output_cost) = self.cost_per_token();
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.
src/llm/reasoning.rs
+10
View File
@@ -292,6 +292,10 @@ pub struct ToolSelection {
pub struct TokenUsage {
pub input_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 {
@@ -434,6 +438,8 @@ impl Reasoning {
let usage = TokenUsage {
input_tokens: response.input_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))
}
@@ -612,6 +618,8 @@ Respond in JSON format:
let usage = TokenUsage {
input_tokens: response.input_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
@@ -690,6 +698,8 @@ Respond in JSON format:
usage: TokenUsage {
input_tokens: response.input_tokens,
output_tokens: response.output_tokens,
cache_read_input_tokens: response.cache_read_input_tokens,
cache_creation_input_tokens: response.cache_creation_input_tokens,
},
})
}
src/llm/recording.rs
+8
View File
@@ -461,6 +461,14 @@ impl LlmProvider for RecordingLlm {
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> {
let (hint, tool_results) = self.capture_new_messages(&request.messages).await;
let response = self.inner.complete(request).await?;
src/llm/response_cache.rs
+12
View File
@@ -181,6 +181,14 @@ impl LlmProvider for CachedProvider {
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> {
let effective_model = self.inner.effective_model_name(request.model.as_deref());
let key = cache_key(&effective_model, &request);
@@ -352,6 +360,8 @@ mod tests {
input_tokens: 1,
output_tokens: 1,
finish_reason: FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
}
@@ -365,6 +375,8 @@ mod tests {
input_tokens: 1,
output_tokens: 1,
finish_reason: FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
}
}
src/llm/retry.rs
+8
View File
@@ -109,6 +109,14 @@ impl LlmProvider for RetryProvider {
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> {
let mut last_error: Option<LlmError> = None;
src/llm/rig_adapter.rs
+255 -5
View File
@@ -3,6 +3,7 @@
//! 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.
use crate::config::CacheRetention;
use async_trait::async_trait;
use rig::OneOrMany;
use rig::completion::{
@@ -14,6 +15,7 @@ use rig::message::{
ToolResultContent, UserContent,
};
use rust_decimal::Decimal;
use rust_decimal_macros::dec;
use serde::Serialize;
use serde::de::DeserializeOwned;
use serde_json::Value as JsonValue;
@@ -34,6 +36,11 @@ pub struct RigAdapter<M: CompletionModel> {
model_name: String,
input_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> {
@@ -47,8 +54,35 @@ impl<M: CompletionModel> RigAdapter<M> {
model_name: name,
input_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 --
@@ -360,7 +394,44 @@ fn saturate_u32(val: u64) -> 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.
///
/// 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(
preamble: Option<String>,
mut history: Vec<RigMessage>,
@@ -368,6 +439,7 @@ fn build_rig_request(
tool_choice: Option<RigToolChoice>,
temperature: Option<f32>,
max_tokens: Option<u32>,
cache_retention: CacheRetention,
) -> Result<RigRequest, LlmError> {
// rig-core requires at least one message in chat_history
if history.is_empty() {
@@ -379,6 +451,17 @@ fn build_rig_request(
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 {
preamble,
chat_history,
@@ -387,7 +470,7 @@ fn build_rig_request(
temperature: temperature.map(|t| t as f64),
max_tokens: max_tokens.map(|t| t as u64),
tool_choice,
additional_params: None,
additional_params,
})
}
@@ -405,6 +488,22 @@ where
(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> {
if let Some(requested_model) = request.model.as_deref()
&& requested_model != self.model_name.as_str()
@@ -427,6 +526,7 @@ where
None,
request.temperature,
request.max_tokens,
self.cache_retention,
)?;
let response =
@@ -440,12 +540,26 @@ where
let (text, _tool_calls, finish) = extract_response(&response.choice, &response.usage);
Ok(CompletionResponse {
let resp = CompletionResponse {
content: text.unwrap_or_default(),
input_tokens: saturate_u32(response.usage.input_tokens),
output_tokens: saturate_u32(response.usage.output_tokens),
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(
@@ -478,6 +592,7 @@ where
tool_choice,
request.temperature,
request.max_tokens,
self.cache_retention,
)?;
let response =
@@ -504,13 +619,27 @@ where
}
}
Ok(ToolCompletionResponse {
let resp = ToolCompletionResponse {
content: text,
tool_calls,
input_tokens: saturate_u32(response.usage.input_tokens),
output_tokens: saturate_u32(response.usage.output_tokens),
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 {
@@ -869,4 +998,125 @@ mod tests {
let known = HashSet::from(["echo".to_string()]);
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"));
}
}
src/llm/smart_routing.rs
+16
View File
@@ -857,6 +857,14 @@ impl LlmProvider for SmartRoutingProvider {
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> {
self.stats.total_requests.fetch_add(1, Ordering::Relaxed);
@@ -1471,6 +1479,8 @@ mod tests {
input_tokens: 10,
output_tokens: 5,
finish_reason: crate::llm::FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
};
assert!(SmartRoutingProvider::response_is_uncertain(&response));
}
@@ -1482,6 +1492,8 @@ mod tests {
input_tokens: 10,
output_tokens: 0,
finish_reason: crate::llm::FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
};
assert!(SmartRoutingProvider::response_is_uncertain(&response));
}
@@ -1493,6 +1505,8 @@ mod tests {
input_tokens: 10,
output_tokens: 1,
finish_reason: crate::llm::FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
};
assert!(!SmartRoutingProvider::response_is_uncertain(&response));
}
@@ -1505,6 +1519,8 @@ mod tests {
input_tokens: 10,
output_tokens: 20,
finish_reason: crate::llm::FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
};
assert!(!SmartRoutingProvider::response_is_uncertain(&response));
}
src/orchestrator/api.rs
+4
View File
@@ -160,6 +160,8 @@ async fn llm_complete(
input_tokens: resp.input_tokens,
output_tokens: resp.output_tokens,
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,
output_tokens: resp.output_tokens,
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,
}))
}
src/testing.rs
+4
View File
@@ -174,6 +174,8 @@ impl LlmProvider for StubLlm {
input_tokens: 10,
output_tokens: 5,
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,
output_tokens: 5,
finish_reason: FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
}
}
src/worker/api.rs
+12
View File
@@ -52,6 +52,10 @@ pub struct ProxyCompletionResponse {
pub input_tokens: u32,
pub output_tokens: u32,
pub finish_reason: String,
#[serde(default)]
pub cache_read_input_tokens: u32,
#[serde(default)]
pub cache_creation_input_tokens: u32,
}
#[derive(Debug, Serialize, Deserialize)]
@@ -71,6 +75,10 @@ pub struct ProxyToolCompletionResponse {
pub input_tokens: u32,
pub output_tokens: u32,
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.
@@ -227,6 +235,8 @@ impl WorkerHttpClient {
input_tokens: proxy_resp.input_tokens,
output_tokens: proxy_resp.output_tokens,
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,
output_tokens: proxy_resp.output_tokens,
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,
})
}
tests/openai_compat_integration.rs
+10
View File
@@ -70,6 +70,8 @@ impl LlmProvider for MockLlmProvider {
input_tokens: 10,
output_tokens: 5,
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,
output_tokens: 8,
finish_reason: FinishReason::ToolUse,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
} else {
Ok(ToolCompletionResponse {
@@ -103,6 +107,8 @@ impl LlmProvider for MockLlmProvider {
input_tokens: 10,
output_tokens: 4,
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,
output_tokens: 5,
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,
output_tokens: 5,
finish_reason: FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
}
tests/provider_chaos.rs
+12
View File
@@ -88,6 +88,8 @@ impl LlmProvider for FlakeyProvider {
input_tokens: 10,
output_tokens: 5,
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,
output_tokens: 5,
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,
output_tokens: 0,
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,
output_tokens: 0,
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,
output_tokens: 5,
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,
output_tokens: 5,
finish_reason: FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
}
}
tests/support/trace_llm.rs
+6
View File
@@ -578,6 +578,8 @@ impl LlmProvider for TraceLlm {
input_tokens,
output_tokens,
finish_reason: FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
}),
TraceResponse::ToolCalls { .. } => Err(LlmError::RequestFailed {
provider: self.model_name.clone(),
@@ -610,6 +612,8 @@ impl LlmProvider for TraceLlm {
input_tokens,
output_tokens,
finish_reason: FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
}),
TraceResponse::ToolCalls {
tool_calls,
@@ -630,6 +634,8 @@ impl LlmProvider for TraceLlm {
input_tokens,
output_tokens,
finish_reason: FinishReason::ToolUse,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
}
TraceResponse::UserInput { .. } => Err(LlmError::RequestFailed {