mirror of
https://github.com/katanemo/plano.git
synced 2026-07-26 17:01:04 +02:00
feat(routing): automatic prompt caching + a per-session routing budget (#982)
* feat(routing): cache-aware routing and zero-config prompt caching
Reconcile prompt caching with intelligent routing: zero-config implicit
session affinity, cache_control preservation/injection across transforms,
cache-adjusted routing economics, and a response-driven cache-hit feedback loop.
* refactor(caching): model-aware cache markers, drop cache-aware routing
* feat(routing): session stickiness cache-regret cost gate
* refactor(routing): expose per-model input and cached rates
* fix(docker): patch libssh2 CVEs by upgrading after package install
* feat: record counterfactual route on vetoed session switches
Add opt-in session_stickiness.record_counterfactual that emits the
plano.switch.counterfactual_route OTEL attribute when the cache-regret
gate retains the previous model, capturing the route it would have taken
had the switch been allowed. Telemetry only; the candidate is never
dispatched. Includes schema, demo config, and a getting-started guide.
* refactor: group prompt-caching and session-stickiness logic into cohesive modules
Extract the two concerns interleaved in the LLM handler into dedicated
sibling modules for readability: prompt_caching.rs (cache-marker injection)
and session_stickiness.rs (session key/prefix-hash resolution, pin lookup,
cache-regret gate, and pin planning). handlers/llm/mod.rs is now a thin
orchestrator that calls them at clear phase boundaries. Behavior-preserving.
* refactor(routing): unify session routing into a cumulative switch budget
Replace observed_cache_hit + per-request threshold with time/provider-TTL
warmth and a per-session switch budget in a single session_router::route()
used by both the full-proxy and /routing decision paths.
* refactor(routing): move switch budget to routing.routing_budget and rename telemetry
* refactor(routing): drop negative-cost budget refund from routing_budget
* refactor(routing): express routing_budget as a percentage overhead cap
Replace the absolute seed_usd pool with max_overhead_pct (whole-number
percent) measured against the session's running never-switch baseline:
allow a paid switch only while cumulative switch spend stays within
max_overhead_pct% of what staying on the anchor would have cost. Track
baseline_usd and switch_spend_usd on the binding (monotonic, no refunds).
* refactor(observability): align routing-budget telemetry with the overhead cap
Rename the switch-decision reasons (within_budget/over_budget -> within_cap/
over_cap) and the session span attributes to match the percentage overhead-cap
model: budget_remaining_in_usd -> overhead_pct + switch_spend_in_usd +
baseline_in_usd, and switch threshold_in_usd -> overhead_ceiling_in_usd.
* feat(observability): emit per-request and cumulative session cost
Price each turn from the catalog rates and surface it as llm.usage.{input,
output,total}_cost_usd on the (llm) span, then accumulate into a conversation-
level session_cost_usd on the binding and emit plano.session.total_cost_in_usd
on the routing span. Cache-creation tokens are priced at the plain input rate,
and the OpenAI vs Anthropic prompt-token convention (cached folded in vs
reported separately) is captured at parse time so the cost math is correct for
both. Rates are resolved request-side since the response path is synchronous.
* refactor(routing): price the never-switch baseline against the session default model
Distinguish the session's default_model (the model it started on, i.e. what it
would have cost by never switching) from anchor_model (the model that handled
the latest request, which the session is warm on). The never-switch baseline now
grows at the default model's cached rate rather than the drifting anchor's, so
the overhead-cap denominator stays true to "% above never-switching"; switch cost
is still measured against the current anchor. Also rename estimate_context_tokens
-> actual_context_tokens (and est_context_tokens -> context_tokens) since the
router prefers the provider's real prompt-token count when the session is warm.
* feat(routing): track bounded route history and price returns to still-warm models
Record a bounded (LRU, capped) per-model visit history on the session binding
and use it to sharpen the switch-cost estimate: a return to a model still within
its cache window re-reads only the tokens appended since its last visit (at its
cached rate) instead of the whole context at the uncached rate, so an A->B->A
switch is no longer over-charged as a full re-ingest. History is carried through
the response-side refresh (refining the anchor's entry from real usage) and
cleared on prefix drift. Adds the plano.switch.candidate_warm_tokens span
attribute.
* refactor(routing): tidy test names and rename pin events to binding events
Shorten run-on test names and replace stale "pin" vocabulary in the session
binding metric (brightstaff_session_pin_events_total -> _binding_events_total),
dropping lifecycle event labels that are no longer emitted.
* feat(routing): price switch cost against uncached rates when caching is off
The overhead-cap gate previously always priced the anchor at its cached rate,
assuming a warm provider cache. When prompt caching is disabled there is no
warm cache to lose, so both the switch cost and the never-switch baseline are
now priced at the plain uncached input rate:
switch_cost = context_tokens x (candidate_uncached - anchor_uncached) / 1M.
* fix(brightstaff): bound pricing-feed fetch with timeouts and warn when session state runs without a tenant header
Co-authored-by: Cursor <cursoragent@cursor.com>
* test(routing): validate the session overhead cap end-to-end across a multi-turn session
Co-authored-by: Cursor <cursoragent@cursor.com>
* docs(demos): fix stale pinning semantics and add routing_budget demo script
Co-authored-by: Cursor <cursoragent@cursor.com>
---------
Co-authored-by: Adil Hafeez <adil.hafeez@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Musa
GitHub
parent
80bb044857
commit
844f08bda7
38 changed files with 4639 additions and 316 deletions
181
crates/brightstaff/src/affinity.rs
Normal file
181
crates/brightstaff/src/affinity.rs
Normal file
|
|
@ -0,0 +1,181 @@
|
|||
//! Implicit session-affinity key derivation.
|
||||
//!
|
||||
//! When a request carries no explicit `X-Model-Affinity` header, Plano derives a
|
||||
//! stable session key from the parts of the prompt that repeat verbatim at the head
|
||||
//! of every turn — the same bytes the provider's prompt cache is keyed on:
|
||||
//!
|
||||
//! ```text
|
||||
//! session_key = hash(system + tools + first_user_message)
|
||||
//! prefix_hash = hash(system + tools)
|
||||
//! ```
|
||||
//!
|
||||
//! The session key is constant for the life of a conversation (history grows at the
|
||||
//! tail, not the head), so turns 2+ reuse the same pin without any client changes.
|
||||
//! The prefix hash covers only the fully-stable segment and is stored with the pin
|
||||
//! for drift detection: if it changes, the provider cache is already lost and
|
||||
//! re-routing fresh is safe.
|
||||
//!
|
||||
//! Only salted hashes are ever stored — never prompt content.
|
||||
|
||||
use hermesllm::apis::openai::{Message, Role};
|
||||
|
||||
/// Salt folded into every hash so stored keys can't be trivially correlated with
|
||||
/// prompt content across systems. Deterministic across processes/replicas so a
|
||||
/// shared Redis session cache keys consistently.
|
||||
const HASH_SALT: &str = "plano-affinity-v1";
|
||||
|
||||
/// Prefix distinguishing derived keys from client-supplied `X-Model-Affinity` ids.
|
||||
const IMPLICIT_KEY_PREFIX: &str = "implicit:";
|
||||
|
||||
/// Derived affinity identifiers for one request.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ImplicitAffinity {
|
||||
/// Session-cache key: `implicit:{hex}` over system + tools + first user message.
|
||||
pub session_key: String,
|
||||
/// Hash of the stable prefix only (system + tools), for drift detection.
|
||||
pub prefix_hash: u64,
|
||||
}
|
||||
|
||||
/// FNV-1a 64-bit — stable across processes and Rust versions (unlike `DefaultHasher`),
|
||||
/// dependency-free, and plenty for cache keying (collisions merely over-pin, which is
|
||||
/// cache-friendly).
|
||||
fn fnv1a64(chunks: &[&str]) -> u64 {
|
||||
const OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;
|
||||
const PRIME: u64 = 0x0000_0100_0000_01b3;
|
||||
let mut hash = OFFSET_BASIS;
|
||||
let mut feed = |bytes: &[u8]| {
|
||||
for &b in bytes {
|
||||
hash ^= b as u64;
|
||||
hash = hash.wrapping_mul(PRIME);
|
||||
}
|
||||
// Field separator so ("ab","c") and ("a","bc") hash differently.
|
||||
hash ^= 0x1f;
|
||||
hash = hash.wrapping_mul(PRIME);
|
||||
};
|
||||
feed(HASH_SALT.as_bytes());
|
||||
for chunk in chunks {
|
||||
feed(chunk.as_bytes());
|
||||
}
|
||||
hash
|
||||
}
|
||||
|
||||
/// Derive the implicit affinity key from parsed request messages and tool names.
|
||||
///
|
||||
/// Returns `None` when there is no user message to anchor on (nothing distinguishes
|
||||
/// the conversation, so pinning would be meaningless).
|
||||
pub fn derive_implicit_affinity(
|
||||
messages: &[Message],
|
||||
tool_names: Option<&[String]>,
|
||||
tenant_id: Option<&str>,
|
||||
) -> Option<ImplicitAffinity> {
|
||||
let system_text: String = messages
|
||||
.iter()
|
||||
.filter(|m| matches!(m.role, Role::System | Role::Developer))
|
||||
.filter_map(|m| m.content.as_ref().map(|c| c.to_string()))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
let first_user = messages
|
||||
.iter()
|
||||
.find(|m| matches!(m.role, Role::User))
|
||||
.and_then(|m| m.content.as_ref().map(|c| c.to_string()))?;
|
||||
|
||||
let tools_text = tool_names.map(|names| names.join(",")).unwrap_or_default();
|
||||
let tenant = tenant_id.unwrap_or_default();
|
||||
|
||||
let prefix_hash = fnv1a64(&[tenant, &system_text, &tools_text]);
|
||||
let session_hash = fnv1a64(&[tenant, &system_text, &tools_text, &first_user]);
|
||||
|
||||
Some(ImplicitAffinity {
|
||||
session_key: format!("{IMPLICIT_KEY_PREFIX}{session_hash:016x}"),
|
||||
prefix_hash,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use hermesllm::apis::openai::MessageContent;
|
||||
|
||||
fn msg(role: Role, text: &str) -> Message {
|
||||
Message {
|
||||
role,
|
||||
content: Some(MessageContent::Text(text.to_string())),
|
||||
name: None,
|
||||
tool_call_id: None,
|
||||
tool_calls: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn key_is_stable_as_history_grows() {
|
||||
let turn1 = vec![
|
||||
msg(Role::System, "You are a coding agent."),
|
||||
msg(Role::User, "Fix the bug in main.rs"),
|
||||
];
|
||||
let turn5 = vec![
|
||||
msg(Role::System, "You are a coding agent."),
|
||||
msg(Role::User, "Fix the bug in main.rs"),
|
||||
msg(Role::Assistant, "Done. Anything else?"),
|
||||
msg(Role::User, "Now add tests"),
|
||||
msg(Role::Assistant, "Added."),
|
||||
];
|
||||
|
||||
let a1 = derive_implicit_affinity(&turn1, None, None).unwrap();
|
||||
let a5 = derive_implicit_affinity(&turn5, None, None).unwrap();
|
||||
assert_eq!(a1.session_key, a5.session_key);
|
||||
assert_eq!(a1.prefix_hash, a5.prefix_hash);
|
||||
assert!(a1.session_key.starts_with("implicit:"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn different_first_user_message_yields_different_session() {
|
||||
let base = msg(Role::System, "You are a coding agent.");
|
||||
let a = derive_implicit_affinity(
|
||||
&[base.clone(), msg(Role::User, "conversation A")],
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
let b = derive_implicit_affinity(&[base, msg(Role::User, "conversation B")], None, None)
|
||||
.unwrap();
|
||||
// Same stable prefix, different conversations.
|
||||
assert_eq!(a.prefix_hash, b.prefix_hash);
|
||||
assert_ne!(a.session_key, b.session_key);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn changed_system_prompt_changes_prefix_hash() {
|
||||
let a = derive_implicit_affinity(
|
||||
&[msg(Role::System, "v1 prompt"), msg(Role::User, "hi")],
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
let b = derive_implicit_affinity(
|
||||
&[msg(Role::System, "v2 prompt"), msg(Role::User, "hi")],
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.unwrap();
|
||||
assert_ne!(a.prefix_hash, b.prefix_hash);
|
||||
assert_ne!(a.session_key, b.session_key);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tools_and_tenant_are_part_of_the_key() {
|
||||
let messages = [msg(Role::System, "s"), msg(Role::User, "u")];
|
||||
let plain = derive_implicit_affinity(&messages, None, None).unwrap();
|
||||
let with_tools =
|
||||
derive_implicit_affinity(&messages, Some(&["get_weather".to_string()]), None).unwrap();
|
||||
let with_tenant = derive_implicit_affinity(&messages, None, Some("acme")).unwrap();
|
||||
assert_ne!(plain.session_key, with_tools.session_key);
|
||||
assert_ne!(plain.session_key, with_tenant.session_key);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_user_message_yields_none() {
|
||||
assert!(derive_implicit_affinity(&[msg(Role::System, "s")], None, None).is_none());
|
||||
assert!(derive_implicit_affinity(&[], None, None).is_none());
|
||||
}
|
||||
}
|
||||
|
|
@ -1,7 +1,10 @@
|
|||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use common::configuration::{Agent, FilterPipeline, Listener, ModelAlias, SpanAttributes};
|
||||
use common::configuration::{
|
||||
Agent, EffectivePromptCaching, EffectiveRoutingBudget, FilterPipeline, Listener, ModelAlias,
|
||||
SpanAttributes,
|
||||
};
|
||||
use common::llm_providers::LlmProviders;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
|
|
@ -31,4 +34,10 @@ pub struct AppState {
|
|||
/// When false, agentic signal analysis is skipped on LLM responses to save CPU.
|
||||
/// Controlled by `overrides.disable_signals` in plano config.
|
||||
pub signals_enabled: bool,
|
||||
/// Instance-wide automatic prompt-caching settings, resolved once from the
|
||||
/// top-level `prompt_caching` config. Disabled by default (opt-in).
|
||||
pub prompt_caching: EffectivePromptCaching,
|
||||
/// Per-session model-switch cost gate, resolved from `routing.routing_budget`.
|
||||
/// Independent of prompt caching; `None` when not configured (off by default).
|
||||
pub routing_budget: Option<EffectiveRoutingBudget>,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1239,6 +1239,7 @@ impl ArchFunctionHandler {
|
|||
total_tokens: 0,
|
||||
prompt_tokens_details: None,
|
||||
completion_tokens_details: None,
|
||||
..Default::default()
|
||||
},
|
||||
system_fingerprint: None,
|
||||
service_tier: None,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,9 @@
|
|||
use bytes::Bytes;
|
||||
use common::configuration::{FilterPipeline, ModelAlias};
|
||||
use common::consts::{ARCH_IS_STREAMING_HEADER, ARCH_PROVIDER_HINT_HEADER, MODEL_AFFINITY_HEADER};
|
||||
use common::consts::{
|
||||
ARCH_IS_STREAMING_HEADER, ARCH_PROVIDER_HINT_HEADER, MODEL_AFFINITY_HEADER, PLANO_CACHE_HEADER,
|
||||
PLANO_PREFIX_HASH_HEADER,
|
||||
};
|
||||
use common::llm_providers::LlmProviders;
|
||||
use hermesllm::apis::openai::Message;
|
||||
use hermesllm::apis::openai_responses::InputParam;
|
||||
|
|
@ -19,6 +22,8 @@ use tokio::sync::RwLock;
|
|||
use tracing::{debug, info, info_span, warn, Instrument};
|
||||
|
||||
pub(crate) mod model_selection;
|
||||
pub(crate) mod prompt_caching;
|
||||
pub(crate) mod session_router;
|
||||
|
||||
use crate::app_state::AppState;
|
||||
use crate::handlers::agents::pipeline::PipelineProcessor;
|
||||
|
|
@ -31,7 +36,7 @@ use crate::state::{
|
|||
};
|
||||
use crate::streaming::{
|
||||
create_streaming_response, create_streaming_response_with_output_filter, truncate_message,
|
||||
LlmMetricsCtx, ObservableStreamProcessor, StreamProcessor,
|
||||
LlmMetricsCtx, ObservableStreamProcessor, SessionUpdateCtx, StreamProcessor,
|
||||
};
|
||||
use crate::tracing::{
|
||||
collect_custom_trace_attributes, llm as tracing_llm, operation_component,
|
||||
|
|
@ -112,8 +117,7 @@ async fn llm_chat_inner(
|
|||
}
|
||||
}
|
||||
|
||||
// Session pinning: extract session ID and check cache before routing
|
||||
let session_id: Option<String> = request_headers
|
||||
let explicit_session_id: Option<String> = request_headers
|
||||
.get(MODEL_AFFINITY_HEADER)
|
||||
.and_then(|h| h.to_str().ok())
|
||||
.map(|s| s.to_string());
|
||||
|
|
@ -123,40 +127,17 @@ async fn llm_chat_inner(
|
|||
.and_then(|hdr| request_headers.get(hdr))
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(|s| s.to_string());
|
||||
let cached_route = if let Some(ref sid) = session_id {
|
||||
state
|
||||
.orchestrator_service
|
||||
.get_cached_route(sid, tenant_id.as_deref())
|
||||
.await
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let (pinned_model, pinned_route_name): (Option<String>, Option<String>) = match cached_route {
|
||||
Some(c) => (Some(c.model_name), c.route_name),
|
||||
None => (None, None),
|
||||
};
|
||||
|
||||
// Record session id on the LLM span for the observability console.
|
||||
if let Some(ref sid) = session_id {
|
||||
get_active_span(|span| {
|
||||
span.set_attribute(opentelemetry::KeyValue::new(
|
||||
tracing_plano::SESSION_ID,
|
||||
sid.clone(),
|
||||
));
|
||||
});
|
||||
}
|
||||
if let Some(ref route_name) = pinned_route_name {
|
||||
get_active_span(|span| {
|
||||
span.set_attribute(opentelemetry::KeyValue::new(
|
||||
tracing_plano::ROUTE_NAME,
|
||||
route_name.clone(),
|
||||
));
|
||||
});
|
||||
}
|
||||
// `X-Plano-Cache: off` disables implicit pinning + marker injection for one call.
|
||||
let cache_off_for_request = request_headers
|
||||
.get(PLANO_CACHE_HEADER)
|
||||
.and_then(|h| h.to_str().ok())
|
||||
.is_some_and(|v| v.eq_ignore_ascii_case("off"));
|
||||
|
||||
let full_qualified_llm_provider_url = format!("{}{}", state.llm_provider_url, request_path);
|
||||
|
||||
// --- Phase 1: Parse and validate the incoming request ---
|
||||
// (Parsing happens before session-key derivation: the implicit affinity key is
|
||||
// computed from the request body's stable prompt prefix.)
|
||||
let parsed = match parse_and_validate_request(
|
||||
request,
|
||||
&request_path,
|
||||
|
|
@ -187,6 +168,14 @@ async fn llm_chat_inner(
|
|||
provider_id,
|
||||
} = parsed;
|
||||
|
||||
// Prompt caching is configured once for the whole instance and never influences
|
||||
// which model routing selects — it only keeps a conversation on the same
|
||||
// model/provider so the upstream prompt cache stays warm across turns. Note:
|
||||
// session affinity and pin lookup are derived later (Phase 2a), *after* input
|
||||
// filters and Responses-API state merge, so the prefix hash reflects the exact
|
||||
// stable prefix sent upstream rather than the pre-filter request body.
|
||||
let prompt_caching = state.prompt_caching;
|
||||
|
||||
// Record LLM-specific span attributes
|
||||
let span = tracing::Span::current();
|
||||
if let Some(temp) = temperature {
|
||||
|
|
@ -283,6 +272,18 @@ async fn llm_chat_inner(
|
|||
*bad_request.status_mut() = StatusCode::BAD_REQUEST;
|
||||
return Ok(bad_request);
|
||||
}
|
||||
|
||||
// Auto-inject prompt-cache markers (grouped in `prompt_caching`). No-op unless
|
||||
// caching is enabled for this request and the model needs explicit markers.
|
||||
prompt_caching::inject_cache_markers(
|
||||
&mut client_request,
|
||||
provider_id,
|
||||
&model_name_only,
|
||||
&upstream_api,
|
||||
&prompt_caching,
|
||||
cache_off_for_request,
|
||||
&alias_resolved_model,
|
||||
);
|
||||
}
|
||||
|
||||
// --- Phase 2: Resolve conversation state (v1/responses API) ---
|
||||
|
|
@ -301,6 +302,40 @@ async fn llm_chat_inner(
|
|||
Err(response) => return Ok(response),
|
||||
};
|
||||
|
||||
// --- Phase 2a: Session affinity (see `session_router`) ---
|
||||
// Derived here, after input filters and Responses-API state merge, so the
|
||||
// implicit session key and prefix hash are computed over the exact stable
|
||||
// prefix (system + tools + first user message) that is actually sent upstream —
|
||||
// the same bytes the provider's prompt cache is keyed on. The prefix hash is
|
||||
// derived even when caching is off, so the `x-plano-prefix-hash` RING_HASH
|
||||
// replica-stickiness header still works.
|
||||
let routing_budget = state.routing_budget;
|
||||
// Derive the implicit session key when either prompt-caching affinity or the
|
||||
// routing budget is active — the budget needs a session anchor even with caching off.
|
||||
let implicit_affinity_enabled = prompt_caching.session_affinity || routing_budget.is_some();
|
||||
let request_messages = client_request.get_messages();
|
||||
let session_router::SessionResolution {
|
||||
request_prefix_hash,
|
||||
session_id,
|
||||
} = session_router::resolve_session(
|
||||
explicit_session_id,
|
||||
&request_messages,
|
||||
tool_names.as_deref(),
|
||||
tenant_id.as_deref(),
|
||||
implicit_affinity_enabled,
|
||||
cache_off_for_request,
|
||||
);
|
||||
|
||||
// Record session id on the LLM span for the observability console.
|
||||
if let Some(ref sid) = session_id {
|
||||
get_active_span(|span| {
|
||||
span.set_attribute(opentelemetry::KeyValue::new(
|
||||
tracing_plano::SESSION_ID,
|
||||
sid.clone(),
|
||||
));
|
||||
});
|
||||
}
|
||||
|
||||
// Serialize request for upstream BEFORE router consumes it
|
||||
let client_request_bytes_for_upstream: Bytes =
|
||||
match ProviderRequestType::to_bytes(&client_request) {
|
||||
|
|
@ -313,76 +348,86 @@ async fn llm_chat_inner(
|
|||
}
|
||||
};
|
||||
|
||||
// --- Phase 3: Route the request (or use pinned model from session cache) ---
|
||||
let resolved_model = if let Some(cached_model) = pinned_model {
|
||||
info!(
|
||||
session_id = %session_id.as_deref().unwrap_or(""),
|
||||
model = %cached_model,
|
||||
"using pinned routing decision from cache"
|
||||
);
|
||||
cached_model
|
||||
// Context size for the switch-cost estimate, computed before the router consumes
|
||||
// the request. Only needed when stickiness could act on a session.
|
||||
let context_tokens: u64 = if session_id.is_some() && routing_budget.is_some() {
|
||||
session_router::actual_context_tokens(&request_messages, &model_name_only)
|
||||
} else {
|
||||
let routing_span = info_span!(
|
||||
"routing",
|
||||
component = "routing",
|
||||
http.method = "POST",
|
||||
http.target = %request_path,
|
||||
model.requested = %model_from_request,
|
||||
model.alias_resolved = %alias_resolved_model,
|
||||
route.selected_model = tracing::field::Empty,
|
||||
routing.determination_ms = tracing::field::Empty,
|
||||
);
|
||||
let routing_result = match async {
|
||||
set_service_name(operation_component::ROUTING);
|
||||
router_chat_get_upstream_model(
|
||||
Arc::clone(&state.orchestrator_service),
|
||||
client_request,
|
||||
&request_path,
|
||||
&request_id,
|
||||
inline_routing_preferences,
|
||||
)
|
||||
.await
|
||||
}
|
||||
.instrument(routing_span)
|
||||
.await
|
||||
{
|
||||
Ok(result) => result,
|
||||
Err(err) => {
|
||||
let mut internal_error = Response::new(full(err.message));
|
||||
*internal_error.status_mut() = err.status_code;
|
||||
return Ok(internal_error);
|
||||
}
|
||||
};
|
||||
|
||||
let (router_selected_model, route_name) =
|
||||
(routing_result.model_name, routing_result.route_name);
|
||||
let model = if router_selected_model != "none" {
|
||||
router_selected_model
|
||||
} else {
|
||||
alias_resolved_model.clone()
|
||||
};
|
||||
|
||||
// Record route name on the LLM span (only when the orchestrator produced one).
|
||||
if let Some(ref rn) = route_name {
|
||||
if !rn.is_empty() && rn != "none" {
|
||||
get_active_span(|span| {
|
||||
span.set_attribute(opentelemetry::KeyValue::new(
|
||||
tracing_plano::ROUTE_NAME,
|
||||
rn.clone(),
|
||||
));
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ref sid) = session_id {
|
||||
state
|
||||
.orchestrator_service
|
||||
.cache_route(sid.clone(), tenant_id.as_deref(), model.clone(), route_name)
|
||||
.await;
|
||||
}
|
||||
|
||||
model
|
||||
0
|
||||
};
|
||||
|
||||
// --- Phase 3: Route the request (quality), then apply the session-cache decision ---
|
||||
// Routing stays cache-blind: the quality router always picks a candidate. The
|
||||
// session router then honors it or sticks to the warm anchor (see `session_router`).
|
||||
let routing_span = info_span!(
|
||||
"routing",
|
||||
component = "routing",
|
||||
http.method = "POST",
|
||||
http.target = %request_path,
|
||||
model.requested = %model_from_request,
|
||||
model.alias_resolved = %alias_resolved_model,
|
||||
route.selected_model = tracing::field::Empty,
|
||||
routing.determination_ms = tracing::field::Empty,
|
||||
);
|
||||
let routing_result = match async {
|
||||
set_service_name(operation_component::ROUTING);
|
||||
router_chat_get_upstream_model(
|
||||
Arc::clone(&state.orchestrator_service),
|
||||
client_request,
|
||||
&request_path,
|
||||
&request_id,
|
||||
inline_routing_preferences,
|
||||
)
|
||||
.await
|
||||
}
|
||||
.instrument(routing_span)
|
||||
.await
|
||||
{
|
||||
Ok(result) => result,
|
||||
Err(err) => {
|
||||
let mut internal_error = Response::new(full(err.message));
|
||||
*internal_error.status_mut() = err.status_code;
|
||||
return Ok(internal_error);
|
||||
}
|
||||
};
|
||||
|
||||
let candidate_model = if routing_result.model_name != "none" {
|
||||
routing_result.model_name
|
||||
} else {
|
||||
alias_resolved_model.clone()
|
||||
};
|
||||
let candidate_route = routing_result.route_name;
|
||||
|
||||
let decision = session_router::route(
|
||||
&state.orchestrator_service,
|
||||
routing_budget.as_ref(),
|
||||
session_router::RouteFacts {
|
||||
session_id: session_id.as_deref(),
|
||||
tenant_id: tenant_id.as_deref(),
|
||||
prefix_hash: request_prefix_hash,
|
||||
context_tokens,
|
||||
candidate_model: &candidate_model,
|
||||
candidate_route: candidate_route.as_deref(),
|
||||
caching_enabled: prompt_caching.enabled,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let resolved_model = decision.model;
|
||||
let resolved_route_name = decision.route_name;
|
||||
|
||||
// Record route name on the LLM span (only when a real route was produced).
|
||||
if let Some(ref rn) = resolved_route_name {
|
||||
if !rn.is_empty() && rn != "none" {
|
||||
get_active_span(|span| {
|
||||
span.set_attribute(opentelemetry::KeyValue::new(
|
||||
tracing_plano::ROUTE_NAME,
|
||||
rn.clone(),
|
||||
));
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
tracing::Span::current().record(tracing_llm::MODEL_NAME, resolved_model.as_str());
|
||||
|
||||
// Record the provider (derived from the `provider/model` prefix) so
|
||||
|
|
@ -398,6 +443,57 @@ async fn llm_chat_inner(
|
|||
});
|
||||
}
|
||||
|
||||
// Resolve the dispatched model's catalog rates now (this side is async; the
|
||||
// response path that prices the turn is synchronous). `None` when no cost feed is
|
||||
// configured → no per-request/session cost is computed.
|
||||
let cost_rates = if session_id.is_some() {
|
||||
state
|
||||
.orchestrator_service
|
||||
.model_rates(&resolved_model)
|
||||
.await
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let cache_read_discount = state
|
||||
.routing_budget
|
||||
.as_ref()
|
||||
.map(|b| b.cache_read_discount)
|
||||
.unwrap_or(common::configuration::DEFAULT_CACHE_READ_DISCOUNT);
|
||||
|
||||
// Response-side refresh: update `last_used` + the context-size estimate from the
|
||||
// real response. The routing decision itself was already persisted by `route()`.
|
||||
let session_update_ctx: Option<SessionUpdateCtx> =
|
||||
session_id.as_ref().map(|sid| SessionUpdateCtx {
|
||||
orchestrator: Arc::clone(&state.orchestrator_service),
|
||||
session_id: sid.clone(),
|
||||
tenant_id: tenant_id.clone(),
|
||||
anchor_model: resolved_model.clone(),
|
||||
default_model: decision.default_model.clone(),
|
||||
route_name: resolved_route_name.clone(),
|
||||
prefix_hash: request_prefix_hash,
|
||||
baseline_usd: decision.baseline_usd,
|
||||
switch_spend_usd: decision.switch_spend_usd,
|
||||
switches: decision.switches,
|
||||
history: decision.history.clone(),
|
||||
session_cost_usd: decision.session_cost_usd,
|
||||
cost_rates,
|
||||
cache_read_discount,
|
||||
context_tokens: decision.cached_tokens,
|
||||
gc_ttl: decision.gc_ttl,
|
||||
});
|
||||
|
||||
// Forward the prefix hash so self-hosted multi-replica backends can do
|
||||
// KV-aware replica stickiness (consistent-hash on this header at the
|
||||
// cluster layer).
|
||||
if let Some(hash) = request_prefix_hash {
|
||||
if let Ok(val) = header::HeaderValue::from_str(&format!("{hash:016x}")) {
|
||||
request_headers.insert(
|
||||
header::HeaderName::from_static(PLANO_PREFIX_HASH_HEADER),
|
||||
val,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// --- Phase 4: Forward to upstream and stream back ---
|
||||
send_upstream(
|
||||
&state.http_client,
|
||||
|
|
@ -415,6 +511,7 @@ async fn llm_chat_inner(
|
|||
state.state_storage.clone(),
|
||||
request_id,
|
||||
&state.filter_pipeline,
|
||||
session_update_ctx,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
|
@ -692,6 +789,7 @@ async fn send_upstream(
|
|||
state_storage: Option<Arc<dyn StateStorage>>,
|
||||
request_id: String,
|
||||
filter_pipeline: &Arc<FilterPipeline>,
|
||||
session_update_ctx: Option<SessionUpdateCtx>,
|
||||
) -> Result<Response<BoxBody<Bytes, hyper::Error>>, hyper::Error> {
|
||||
let span_name = if model_from_request == resolved_model {
|
||||
format!("POST {} {}", request_path, resolved_model)
|
||||
|
|
@ -807,7 +905,7 @@ async fn send_upstream(
|
|||
let byte_stream = llm_response.bytes_stream();
|
||||
|
||||
// Create base processor for metrics and tracing
|
||||
let base_processor = ObservableStreamProcessor::new(
|
||||
let mut base_processor = ObservableStreamProcessor::new(
|
||||
operation_component::LLM,
|
||||
span_name,
|
||||
request_start_time,
|
||||
|
|
@ -818,6 +916,13 @@ async fn send_upstream(
|
|||
model: metric_model.clone(),
|
||||
upstream_status: upstream_status.as_u16(),
|
||||
});
|
||||
// Only refresh the session binding for successful responses — errors carry no
|
||||
// usage block and shouldn't reset the warmth clock.
|
||||
if upstream_status.is_success() {
|
||||
if let Some(update_ctx) = session_update_ctx {
|
||||
base_processor = base_processor.with_session_update(update_ctx);
|
||||
}
|
||||
}
|
||||
|
||||
let output_filter_request_headers = if filter_pipeline.has_output_filters() {
|
||||
Some(request_headers.clone())
|
||||
|
|
|
|||
70
crates/brightstaff/src/handlers/llm/prompt_caching.rs
Normal file
70
crates/brightstaff/src/handlers/llm/prompt_caching.rs
Normal file
|
|
@ -0,0 +1,70 @@
|
|||
//! Prompt-caching request handling for the LLM path.
|
||||
//!
|
||||
//! This module owns the "make the upstream provider's prompt cache work" concern:
|
||||
//! resolving the correct cache-marking strategy for the `(gateway × model family ×
|
||||
//! upstream API)` combination and injecting the markers into the outbound request.
|
||||
//! It never influences routing — see [`super::session_router`] for the session-cache
|
||||
//! lookup and the routing-budget switch-cost concern.
|
||||
|
||||
use common::configuration::EffectivePromptCaching;
|
||||
use hermesllm::clients::SupportedUpstreamAPIs;
|
||||
use hermesllm::{cache_marker_strategy, CacheMarkerStrategy, ProviderId, ProviderRequestType};
|
||||
use tracing::debug;
|
||||
|
||||
/// Auto-inject prompt-cache markers into `client_request`.
|
||||
///
|
||||
/// The strategy is resolved from `(gateway × model family × upstream API)` so that
|
||||
/// Anthropic-family models cache whether they arrive over the native Messages API or
|
||||
/// an OpenAI-compatible gateway (DigitalOcean, OpenRouter), while OpenAI-family models
|
||||
/// (which cache automatically) and unimplemented backends are left untouched.
|
||||
///
|
||||
/// A no-op when caching is disabled, `inject_cache_control` is off, the request opted
|
||||
/// out (`X-Plano-Cache: off`), or the provider caches automatically. Injection is
|
||||
/// idempotent (client-supplied markers are respected) and threshold-guarded against
|
||||
/// the provider's minimum cacheable prefix.
|
||||
pub fn inject_cache_markers(
|
||||
client_request: &mut ProviderRequestType,
|
||||
provider_id: ProviderId,
|
||||
model_name_only: &str,
|
||||
upstream_api: &SupportedUpstreamAPIs,
|
||||
prompt_caching: &EffectivePromptCaching,
|
||||
cache_off_for_request: bool,
|
||||
alias_resolved_model: &str,
|
||||
) {
|
||||
if !(prompt_caching.enabled && prompt_caching.inject_cache_control && !cache_off_for_request) {
|
||||
return;
|
||||
}
|
||||
|
||||
match cache_marker_strategy(provider_id, model_name_only, upstream_api) {
|
||||
CacheMarkerStrategy::AnthropicMessagesBreakpoints {
|
||||
min_prefix_tokens, ..
|
||||
} => {
|
||||
if let ProviderRequestType::MessagesRequest(req) = client_request {
|
||||
let threshold = prompt_caching.min_prefix_tokens.max(min_prefix_tokens);
|
||||
if req.inject_cache_breakpoints(threshold) {
|
||||
debug!(
|
||||
model = %alias_resolved_model,
|
||||
min_prefix_tokens = threshold,
|
||||
"injected anthropic ephemeral cache breakpoints"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
CacheMarkerStrategy::OpenAiContentPartCacheControl {
|
||||
min_prefix_tokens,
|
||||
ttl,
|
||||
} => {
|
||||
if let ProviderRequestType::ChatCompletionsRequest(req) = client_request {
|
||||
let threshold = prompt_caching.min_prefix_tokens.max(min_prefix_tokens);
|
||||
if req.inject_cache_control(ttl, threshold) {
|
||||
debug!(
|
||||
model = %alias_resolved_model,
|
||||
min_prefix_tokens = threshold,
|
||||
"injected openai content-part cache_control"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
CacheMarkerStrategy::Automatic | CacheMarkerStrategy::None => {}
|
||||
}
|
||||
}
|
||||
1062
crates/brightstaff/src/handlers/llm/session_router.rs
Normal file
1062
crates/brightstaff/src/handlers/llm/session_router.rs
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -1,9 +1,11 @@
|
|||
use bytes::Bytes;
|
||||
use common::configuration::{SpanAttributes, TopLevelRoutingPreference};
|
||||
use common::configuration::{
|
||||
EffectivePromptCaching, EffectiveRoutingBudget, SpanAttributes, TopLevelRoutingPreference,
|
||||
};
|
||||
use common::consts::{MODEL_AFFINITY_HEADER, REQUEST_ID_HEADER};
|
||||
use common::errors::BrightStaffError;
|
||||
use hermesllm::clients::SupportedAPIsFromClient;
|
||||
use hermesllm::ProviderRequestType;
|
||||
use hermesllm::{ProviderRequest, ProviderRequestType};
|
||||
use http_body_util::combinators::BoxBody;
|
||||
use http_body_util::{BodyExt, Full};
|
||||
use hyper::{Request, Response, StatusCode};
|
||||
|
|
@ -12,6 +14,7 @@ use tracing::{debug, info, info_span, warn, Instrument};
|
|||
|
||||
use super::extract_or_generate_traceparent;
|
||||
use crate::handlers::llm::model_selection::router_chat_get_upstream_model;
|
||||
use crate::handlers::llm::session_router;
|
||||
use crate::metrics as bs_metrics;
|
||||
use crate::metrics::labels as metric_labels;
|
||||
use crate::router::orchestrator::OrchestratorService;
|
||||
|
|
@ -60,11 +63,14 @@ struct RoutingDecisionResponse {
|
|||
pinned: bool,
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn routing_decision(
|
||||
request: Request<hyper::body::Incoming>,
|
||||
orchestrator_service: Arc<OrchestratorService>,
|
||||
request_path: String,
|
||||
span_attributes: &Option<SpanAttributes>,
|
||||
prompt_caching: EffectivePromptCaching,
|
||||
routing_budget: Option<EffectiveRoutingBudget>,
|
||||
) -> Result<Response<BoxBody<Bytes, hyper::Error>>, hyper::Error> {
|
||||
let request_headers = request.headers().clone();
|
||||
let request_id: String = request_headers
|
||||
|
|
@ -73,7 +79,7 @@ pub async fn routing_decision(
|
|||
.map(|s| s.to_string())
|
||||
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
|
||||
|
||||
let session_id: Option<String> = request_headers
|
||||
let explicit_session_id: Option<String> = request_headers
|
||||
.get(MODEL_AFFINITY_HEADER)
|
||||
.and_then(|h| h.to_str().ok())
|
||||
.map(|s| s.to_string());
|
||||
|
|
@ -101,8 +107,10 @@ pub async fn routing_decision(
|
|||
request_path,
|
||||
request_headers,
|
||||
custom_attrs,
|
||||
session_id,
|
||||
explicit_session_id,
|
||||
tenant_id,
|
||||
prompt_caching,
|
||||
routing_budget,
|
||||
)
|
||||
.instrument(request_span)
|
||||
.await
|
||||
|
|
@ -116,8 +124,10 @@ async fn routing_decision_inner(
|
|||
request_path: String,
|
||||
request_headers: hyper::HeaderMap,
|
||||
custom_attrs: std::collections::HashMap<String, String>,
|
||||
session_id: Option<String>,
|
||||
explicit_session_id: Option<String>,
|
||||
tenant_id: Option<String>,
|
||||
prompt_caching: EffectivePromptCaching,
|
||||
routing_budget: Option<EffectiveRoutingBudget>,
|
||||
) -> Result<Response<BoxBody<Bytes, hyper::Error>>, hyper::Error> {
|
||||
set_service_name(operation_component::ROUTING);
|
||||
opentelemetry::trace::get_active_span(|span| {
|
||||
|
|
@ -135,37 +145,9 @@ async fn routing_decision_inner(
|
|||
.unwrap_or("unknown")
|
||||
.to_string();
|
||||
|
||||
if let Some(ref sid) = session_id {
|
||||
if let Some(cached) = orchestrator_service
|
||||
.get_cached_route(sid, tenant_id.as_deref())
|
||||
.await
|
||||
{
|
||||
info!(
|
||||
session_id = %sid,
|
||||
model = %cached.model_name,
|
||||
route = ?cached.route_name,
|
||||
"returning pinned routing decision from cache"
|
||||
);
|
||||
let response = RoutingDecisionResponse {
|
||||
models: vec![cached.model_name],
|
||||
route: cached.route_name,
|
||||
trace_id,
|
||||
session_id: Some(sid.clone()),
|
||||
pinned: true,
|
||||
};
|
||||
let json = serde_json::to_string(&response).unwrap();
|
||||
let body = Full::new(Bytes::from(json))
|
||||
.map_err(|never| match never {})
|
||||
.boxed();
|
||||
return Ok(Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.header("Content-Type", "application/json")
|
||||
.body(body)
|
||||
.unwrap());
|
||||
}
|
||||
}
|
||||
|
||||
// Parse request body
|
||||
// Parse the request body up front so a pin can be validated against prefix drift
|
||||
// and a fresh pin can be stored with its prefix hash. This endpoint shares the
|
||||
// session cache with the LLM handler, so both must key drift the same way.
|
||||
let raw_bytes = request.collect().await?.to_bytes();
|
||||
|
||||
debug!(
|
||||
|
|
@ -202,6 +184,39 @@ async fn routing_decision_inner(
|
|||
}
|
||||
};
|
||||
|
||||
// `X-Plano-Cache: off` opts this request out of implicit affinity (same sentinel
|
||||
// the LLM handler honors), so callers can bypass stickiness per request.
|
||||
let cache_off_for_request = request_headers
|
||||
.get(common::consts::PLANO_CACHE_HEADER)
|
||||
.and_then(|h| h.to_str().ok())
|
||||
.is_some_and(|v| v.eq_ignore_ascii_case("off"));
|
||||
|
||||
let request_messages = client_request.get_messages();
|
||||
let tool_names = client_request.get_tool_names();
|
||||
|
||||
// Session key + prefix hash resolved identically to the LLM handler so pins
|
||||
// interoperate across the full-proxy and decision paths.
|
||||
// Derive the implicit session key when either prompt-caching affinity or the routing
|
||||
// budget is active, so the budget works the same way with caching off.
|
||||
let implicit_affinity_enabled = prompt_caching.session_affinity || routing_budget.is_some();
|
||||
let session_router::SessionResolution {
|
||||
request_prefix_hash,
|
||||
session_id,
|
||||
} = session_router::resolve_session(
|
||||
explicit_session_id,
|
||||
&request_messages,
|
||||
tool_names.as_deref(),
|
||||
tenant_id.as_deref(),
|
||||
implicit_affinity_enabled,
|
||||
cache_off_for_request,
|
||||
);
|
||||
|
||||
let context_tokens: u64 = if session_id.is_some() && routing_budget.is_some() {
|
||||
session_router::actual_context_tokens(&request_messages, client_request.model())
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
let routing_result = router_chat_get_upstream_model(
|
||||
Arc::clone(&orchestrator_service),
|
||||
client_request,
|
||||
|
|
@ -213,23 +228,38 @@ async fn routing_decision_inner(
|
|||
|
||||
match routing_result {
|
||||
Ok(result) => {
|
||||
if let Some(ref sid) = session_id {
|
||||
orchestrator_service
|
||||
.cache_route(
|
||||
sid.clone(),
|
||||
tenant_id.as_deref(),
|
||||
result.model_name.clone(),
|
||||
result.route_name.clone(),
|
||||
)
|
||||
.await;
|
||||
let candidate_model = result.model_name.clone();
|
||||
let decision = session_router::route(
|
||||
&orchestrator_service,
|
||||
routing_budget.as_ref(),
|
||||
session_router::RouteFacts {
|
||||
session_id: session_id.as_deref(),
|
||||
tenant_id: tenant_id.as_deref(),
|
||||
prefix_hash: request_prefix_hash,
|
||||
context_tokens,
|
||||
candidate_model: &candidate_model,
|
||||
candidate_route: result.route_name.as_deref(),
|
||||
caching_enabled: prompt_caching.enabled,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
// Front the ranked fallback list with the decided model (the anchor, when a
|
||||
// switch was vetoed), so 429/5xx fallbacks still work.
|
||||
let mut models = result.models;
|
||||
if models.first() != Some(&decision.model) {
|
||||
models.retain(|m| m != &decision.model);
|
||||
models.insert(0, decision.model.clone());
|
||||
}
|
||||
|
||||
let response = RoutingDecisionResponse {
|
||||
models: result.models,
|
||||
route: result.route_name,
|
||||
models,
|
||||
route: decision.route_name,
|
||||
trace_id,
|
||||
session_id,
|
||||
pinned: false,
|
||||
// `pinned` signals a warm, stuck session — safe for callers to treat as
|
||||
// "keep this provider's cache warm".
|
||||
pinned: decision.warm,
|
||||
};
|
||||
|
||||
// Distinguish "decision served" (a concrete model picked) from
|
||||
|
|
@ -247,6 +277,7 @@ async fn routing_decision_inner(
|
|||
primary_model = %response.models.first().map(|s| s.as_str()).unwrap_or("none"),
|
||||
total_models = response.models.len(),
|
||||
route = ?response.route,
|
||||
pinned = response.pinned,
|
||||
"routing decision completed"
|
||||
);
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
pub mod affinity;
|
||||
pub mod app_state;
|
||||
pub mod handlers;
|
||||
pub mod metrics;
|
||||
|
|
|
|||
|
|
@ -216,7 +216,14 @@ async fn init_app_state(
|
|||
if latency_count > 1 {
|
||||
return Err("model_metrics_sources: only one latency metrics source is allowed".into());
|
||||
}
|
||||
let svc = ModelMetricsService::new(sources, reqwest::Client::new()).await;
|
||||
// The initial pricing fetch is awaited before listeners come up, so bound it —
|
||||
// an unresponsive feed (models.dev / DO catalog) must not hang startup. The
|
||||
// same client backs the periodic refresh loop.
|
||||
let metrics_client = reqwest::Client::builder()
|
||||
.connect_timeout(std::time::Duration::from_secs(5))
|
||||
.timeout(std::time::Duration::from_secs(30))
|
||||
.build()?;
|
||||
let svc = ModelMetricsService::new(sources, metrics_client).await;
|
||||
Some(Arc::new(svc))
|
||||
} else {
|
||||
None
|
||||
|
|
@ -343,6 +350,55 @@ async fn init_app_state(
|
|||
|
||||
let signals_enabled = !overrides.disable_signals.unwrap_or(false);
|
||||
|
||||
let prompt_caching =
|
||||
common::configuration::EffectivePromptCaching::from_config(config.prompt_caching.as_ref())?;
|
||||
|
||||
// Routing budget lives under `routing` and is independent of prompt caching.
|
||||
let routing_budget = common::configuration::EffectiveRoutingBudget::from_config(
|
||||
config
|
||||
.routing
|
||||
.as_ref()
|
||||
.and_then(|r| r.routing_budget.as_ref()),
|
||||
)?;
|
||||
|
||||
// The routing-budget cost gate needs per-model pricing to compute switch cost.
|
||||
if routing_budget.is_some() {
|
||||
use common::configuration::MetricsSource;
|
||||
let has_cost_source = config
|
||||
.model_metrics_sources
|
||||
.as_deref()
|
||||
.unwrap_or_default()
|
||||
.iter()
|
||||
.any(|s| matches!(s, MetricsSource::Cost(_)));
|
||||
if !has_cost_source {
|
||||
return Err(
|
||||
"routing.routing_budget is configured but no cost metrics source is \
|
||||
configured — add a cost source (e.g. models.dev) to model_metrics_sources so \
|
||||
per-model input/cached rates are available for the switch-cost calculation"
|
||||
.into(),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Session bindings (anchor model, budget spend) are keyed by tenant + session.
|
||||
// Without a tenant header, sessions from different customers share one keyspace,
|
||||
// so identical prompts from different tenants would collide on the same binding.
|
||||
let session_state_in_use = prompt_caching.session_affinity || routing_budget.is_some();
|
||||
if session_state_in_use
|
||||
&& config
|
||||
.routing
|
||||
.as_ref()
|
||||
.and_then(|r| r.session_cache.as_ref())
|
||||
.and_then(|c| c.tenant_header.as_ref())
|
||||
.is_none()
|
||||
{
|
||||
warn!(
|
||||
"session affinity / routing budget is enabled but routing.session_cache.tenant_header \
|
||||
is not set — all requests share one session keyspace; set tenant_header to isolate \
|
||||
sessions per customer in multi-tenant deployments"
|
||||
);
|
||||
}
|
||||
|
||||
Ok(AppState {
|
||||
orchestrator_service,
|
||||
model_aliases: config.model_aliases.clone(),
|
||||
|
|
@ -356,6 +412,8 @@ async fn init_app_state(
|
|||
http_client: reqwest::Client::new(),
|
||||
filter_pipeline,
|
||||
signals_enabled,
|
||||
prompt_caching,
|
||||
routing_budget,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -527,6 +585,8 @@ async fn dispatch(
|
|||
Arc::clone(&state.orchestrator_service),
|
||||
stripped,
|
||||
&state.span_attributes,
|
||||
state.prompt_caching,
|
||||
state.routing_budget,
|
||||
)
|
||||
.with_context(parent_cx)
|
||||
.await;
|
||||
|
|
|
|||
|
|
@ -18,6 +18,10 @@ pub const ROUTE_LLM: &str = "llm";
|
|||
// Token kind for brightstaff_llm_tokens_total.
|
||||
pub const TOKEN_KIND_PROMPT: &str = "prompt";
|
||||
pub const TOKEN_KIND_COMPLETION: &str = "completion";
|
||||
/// Input tokens served from the provider's prompt cache (billed at the cached rate).
|
||||
pub const TOKEN_KIND_CACHE_READ: &str = "cache_read";
|
||||
/// Input tokens written into the provider's prompt cache (cache-creation surcharge).
|
||||
pub const TOKEN_KIND_CACHE_WRITE: &str = "cache_write";
|
||||
|
||||
// LLM error_class values (match docstring in metrics/mod.rs).
|
||||
pub const LLM_ERR_NONE: &str = "none";
|
||||
|
|
@ -36,3 +40,31 @@ pub const ROUTING_SVC_POLICY_ERROR: &str = "policy_error";
|
|||
pub const SESSION_CACHE_HIT: &str = "hit";
|
||||
pub const SESSION_CACHE_MISS: &str = "miss";
|
||||
pub const SESSION_CACHE_STORE: &str = "store";
|
||||
|
||||
// Prompt cache outcome values (brightstaff_prompt_cache_requests_total).
|
||||
pub const PROMPT_CACHE_HIT: &str = "hit";
|
||||
pub const PROMPT_CACHE_MISS: &str = "miss";
|
||||
|
||||
// Session binding lifecycle events (brightstaff_session_binding_events_total).
|
||||
/// An existing binding was refreshed from observed usage (TTL extended, token counts
|
||||
/// and running cost updated) after a turn completed.
|
||||
pub const BINDING_EVENT_REFRESH: &str = "refresh";
|
||||
|
||||
// Session-stickiness decisions (brightstaff_session_switch_decisions_total).
|
||||
// `decision` label — the coarse outcome:
|
||||
/// The proposed switch was honored (free, within the overhead cap, or unpriced fail-open).
|
||||
pub const SWITCH_DECISION_ALLOWED: &str = "allowed";
|
||||
/// The switch would have exceeded the session's overhead cap — the warm anchor was retained.
|
||||
pub const SWITCH_DECISION_RETAINED: &str = "retained";
|
||||
|
||||
// `reason` label — why the decision was made:
|
||||
/// The router agreed with the warm anchor; no switch was needed.
|
||||
pub const SWITCH_REASON_SAME_ANCHOR: &str = "same_anchor";
|
||||
/// The candidate was outright cheaper (negative cost) — a free switch.
|
||||
pub const SWITCH_REASON_FREE: &str = "free";
|
||||
/// A paid switch that kept cumulative spend within the session's overhead cap.
|
||||
pub const SWITCH_REASON_WITHIN_CAP: &str = "within_cap";
|
||||
/// A paid switch that would have pushed cumulative spend over the overhead cap — retained.
|
||||
pub const SWITCH_REASON_OVER_CAP: &str = "over_cap";
|
||||
/// Pricing was missing for one side, so the switch was allowed without a cost gate.
|
||||
pub const SWITCH_REASON_NO_PRICING: &str = "no_pricing";
|
||||
|
|
|
|||
|
|
@ -18,6 +18,9 @@
|
|||
//! `brightstaff_router_decision_duration_seconds`,
|
||||
//! `brightstaff_routing_service_requests_total`,
|
||||
//! `brightstaff_session_cache_events_total`.
|
||||
//! - Prompt caching: `brightstaff_prompt_cache_requests_total`,
|
||||
//! `brightstaff_session_binding_events_total` (cache read/write tokens are emitted as
|
||||
//! `kind=cache_read|cache_write` on `brightstaff_llm_tokens_total`).
|
||||
//! - Process: via `metrics-process`.
|
||||
//! - Build: `brightstaff_build_info`.
|
||||
|
||||
|
|
@ -172,6 +175,16 @@ fn describe_all() {
|
|||
"brightstaff_session_cache_events_total",
|
||||
"Session affinity cache lookups and stores, by outcome."
|
||||
);
|
||||
describe_counter!(
|
||||
"brightstaff_prompt_cache_requests_total",
|
||||
"LLM responses with usage data, by provider, model and prompt-cache outcome \
|
||||
(hit = cache read/creation tokens reported, miss = none). The miss rate is \
|
||||
the silent cache-miss baseline."
|
||||
);
|
||||
describe_counter!(
|
||||
"brightstaff_session_binding_events_total",
|
||||
"Session binding lifecycle events, by event (currently: refresh)."
|
||||
);
|
||||
|
||||
describe_gauge!(
|
||||
"brightstaff_build_info",
|
||||
|
|
@ -375,3 +388,37 @@ pub fn record_session_cache_event(outcome: &'static str) {
|
|||
)
|
||||
.increment(1);
|
||||
}
|
||||
|
||||
/// Record whether a completed LLM response reported prompt-cache activity.
|
||||
/// `hit` means the provider billed cache-read or cache-creation tokens.
|
||||
pub fn record_prompt_cache_outcome(provider: &str, model: &str, outcome: &'static str) {
|
||||
counter!(
|
||||
"brightstaff_prompt_cache_requests_total",
|
||||
"provider" => provider.to_string(),
|
||||
"model" => model.to_string(),
|
||||
"outcome" => outcome,
|
||||
)
|
||||
.increment(1);
|
||||
}
|
||||
|
||||
/// Record a session binding lifecycle event (see `metrics::labels::BINDING_EVENT_*`).
|
||||
pub fn record_session_binding_event(event: &'static str) {
|
||||
counter!(
|
||||
"brightstaff_session_binding_events_total",
|
||||
"event" => event,
|
||||
)
|
||||
.increment(1);
|
||||
}
|
||||
|
||||
/// Record a session-stickiness decision on a proposed model switch. `decision` is the
|
||||
/// coarse outcome (`allowed`/`retained`, see `metrics::labels::SWITCH_DECISION_*`) and
|
||||
/// `reason` explains why (`same_anchor`/`free`/`within_cap`/`over_cap`/`no_pricing`,
|
||||
/// see `metrics::labels::SWITCH_REASON_*`).
|
||||
pub fn record_session_switch_decision(decision: &'static str, reason: &'static str) {
|
||||
counter!(
|
||||
"brightstaff_session_switch_decisions_total",
|
||||
"decision" => decision,
|
||||
"reason" => reason,
|
||||
)
|
||||
.increment(1);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -11,14 +11,85 @@ use tracing::{debug, info, warn};
|
|||
const DO_PRICING_URL: &str = "https://api.digitalocean.com/v2/gen-ai/models/catalog";
|
||||
const MODELS_DEV_URL: &str = "https://models.dev/api.json";
|
||||
|
||||
/// DigitalOcean publishes prices per token; scale to the per-million convention
|
||||
/// that `ModelRates` and the absolute-USD consumers (switch-cost gate) expect.
|
||||
const TOKENS_PER_MILLION: f64 = 1_000_000.0;
|
||||
|
||||
/// Structured per-million-token USD rates for one model, as published by the cost
|
||||
/// feed. Kept alongside the blended ranking metric so cost-sensitive features (e.g.
|
||||
/// the session-stickiness switch-cost gate) can reason about input vs cached-input
|
||||
/// pricing without touching `rank_models`.
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub struct ModelRates {
|
||||
pub input_per_million: f64,
|
||||
pub output_per_million: f64,
|
||||
/// Cached (prompt-cache read) input rate. Present when the feed publishes it
|
||||
/// (models.dev `cost.cache_read`); absent for feeds that don't (DO catalog).
|
||||
pub cache_read_per_million: Option<f64>,
|
||||
}
|
||||
|
||||
impl ModelRates {
|
||||
/// Blended input+output metric used for `prefer: cheapest` ranking.
|
||||
fn blended(&self) -> f64 {
|
||||
self.input_per_million + self.output_per_million
|
||||
}
|
||||
|
||||
/// Cached input rate, falling back to `input * cache_read_discount` when the
|
||||
/// feed doesn't publish a cached rate.
|
||||
pub fn cached_input_rate(&self, cache_read_discount: f64) -> f64 {
|
||||
self.cache_read_per_million
|
||||
.unwrap_or(self.input_per_million * cache_read_discount)
|
||||
}
|
||||
|
||||
/// Actual USD cost of one request from observed token usage, split into
|
||||
/// `(input_cost, output_cost)`. Cache *creation* tokens are priced at the plain
|
||||
/// input rate (no write premium). Cache *read* tokens use the cached rate.
|
||||
///
|
||||
/// `prompt_includes_cached` reflects the provider convention: OpenAI-shape usage
|
||||
/// folds cached tokens into `prompt_tokens` (uncached = prompt - cached), while
|
||||
/// Anthropic-shape reports uncached input separately (uncached = prompt_tokens).
|
||||
pub fn request_cost_usd(
|
||||
&self,
|
||||
prompt_tokens: u64,
|
||||
cached_input_tokens: u64,
|
||||
cache_creation_tokens: u64,
|
||||
completion_tokens: u64,
|
||||
prompt_includes_cached: bool,
|
||||
cache_read_discount: f64,
|
||||
) -> (f64, f64) {
|
||||
let read_rate = self.cached_input_rate(cache_read_discount);
|
||||
let uncached = if prompt_includes_cached {
|
||||
prompt_tokens.saturating_sub(cached_input_tokens)
|
||||
} else {
|
||||
prompt_tokens
|
||||
};
|
||||
let input = (uncached as f64 * self.input_per_million
|
||||
+ cached_input_tokens as f64 * read_rate
|
||||
+ cache_creation_tokens as f64 * self.input_per_million)
|
||||
/ TOKENS_PER_MILLION;
|
||||
let output = completion_tokens as f64 * self.output_per_million / TOKENS_PER_MILLION;
|
||||
(input, output)
|
||||
}
|
||||
}
|
||||
|
||||
/// Derive the blended ranking map from the structured rates map.
|
||||
fn blended_costs(rates: &HashMap<String, ModelRates>) -> HashMap<String, f64> {
|
||||
rates
|
||||
.iter()
|
||||
.map(|(k, r)| (k.clone(), r.blended()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub struct ModelMetricsService {
|
||||
cost: Arc<RwLock<HashMap<String, f64>>>,
|
||||
rates: Arc<RwLock<HashMap<String, ModelRates>>>,
|
||||
latency: Arc<RwLock<HashMap<String, f64>>>,
|
||||
}
|
||||
|
||||
impl ModelMetricsService {
|
||||
pub async fn new(sources: &[MetricsSource], client: reqwest::Client) -> Self {
|
||||
let cost_data = Arc::new(RwLock::new(HashMap::new()));
|
||||
let rates_data = Arc::new(RwLock::new(HashMap::new()));
|
||||
let latency_data = Arc::new(RwLock::new(HashMap::new()));
|
||||
|
||||
for source in sources {
|
||||
|
|
@ -34,10 +105,12 @@ impl ModelMetricsService {
|
|||
|
||||
let data = fetch_cost_pricing(&provider, &url, &client, &aliases).await;
|
||||
info!(models = data.len(), provider = provider_name, url = %url, "fetched cost pricing");
|
||||
*cost_data.write().await = data;
|
||||
*cost_data.write().await = blended_costs(&data);
|
||||
*rates_data.write().await = data;
|
||||
|
||||
if let Some(interval_secs) = cfg.refresh_interval {
|
||||
let cost_clone = Arc::clone(&cost_data);
|
||||
let rates_clone = Arc::clone(&rates_data);
|
||||
let client_clone = client.clone();
|
||||
let interval = Duration::from_secs(interval_secs);
|
||||
tokio::spawn(async move {
|
||||
|
|
@ -47,7 +120,8 @@ impl ModelMetricsService {
|
|||
fetch_cost_pricing(&provider, &url, &client_clone, &aliases)
|
||||
.await;
|
||||
info!(models = data.len(), provider = provider_name, url = %url, "refreshed cost pricing");
|
||||
*cost_clone.write().await = data;
|
||||
*cost_clone.write().await = blended_costs(&data);
|
||||
*rates_clone.write().await = data;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
|
@ -81,10 +155,34 @@ impl ModelMetricsService {
|
|||
|
||||
ModelMetricsService {
|
||||
cost: cost_data,
|
||||
rates: rates_data,
|
||||
latency: latency_data,
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a service directly from a rates map (no network). Test-only: lets
|
||||
/// handler/router tests exercise switch-cost math with deterministic pricing.
|
||||
#[cfg(test)]
|
||||
pub fn from_rates_for_test(rates: HashMap<String, ModelRates>) -> Self {
|
||||
ModelMetricsService {
|
||||
cost: Arc::new(RwLock::new(blended_costs(&rates))),
|
||||
rates: Arc::new(RwLock::new(rates)),
|
||||
latency: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Structured per-million rates for a model, if the cost feed published them.
|
||||
/// Falls back to the bare model id (without `provider/` prefix) like the
|
||||
/// blended cost map does.
|
||||
pub async fn model_rates(&self, model: &str) -> Option<ModelRates> {
|
||||
let rates = self.rates.read().await;
|
||||
rates.get(model).copied().or_else(|| {
|
||||
model
|
||||
.split_once('/')
|
||||
.and_then(|(_, bare)| rates.get(bare).copied())
|
||||
})
|
||||
}
|
||||
|
||||
/// Rank `models` by `policy`, returning them in preference order.
|
||||
/// Models with no metric data are appended at the end in their original order.
|
||||
pub async fn rank_models(&self, models: &[String], policy: &SelectionPolicy) -> Vec<String> {
|
||||
|
|
@ -167,10 +265,14 @@ struct DoModel {
|
|||
pricing: Option<DoPricing>,
|
||||
}
|
||||
|
||||
/// DigitalOcean catalog pricing. Despite the `_per_million` field names, the DO
|
||||
/// API returns these as USD **per token** (e.g. gpt-4o input `2.5e-6`), so they
|
||||
/// are scaled to per-million at ingestion to match `ModelRates`' convention.
|
||||
#[derive(serde::Deserialize)]
|
||||
struct DoPricing {
|
||||
input_price_per_million: Option<f64>,
|
||||
output_price_per_million: Option<f64>,
|
||||
cache_read_input_price_per_million: Option<f64>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
|
|
@ -188,6 +290,7 @@ struct ModelsDevModel {
|
|||
struct ModelsDevCost {
|
||||
input: Option<f64>,
|
||||
output: Option<f64>,
|
||||
cache_read: Option<f64>,
|
||||
}
|
||||
|
||||
fn default_cost_url(provider: &CostProvider) -> &'static str {
|
||||
|
|
@ -209,7 +312,7 @@ async fn fetch_cost_pricing(
|
|||
url: &str,
|
||||
client: &reqwest::Client,
|
||||
aliases: &HashMap<String, String>,
|
||||
) -> HashMap<String, f64> {
|
||||
) -> HashMap<String, ModelRates> {
|
||||
match provider {
|
||||
CostProvider::Digitalocean => fetch_do_pricing(url, client, aliases).await,
|
||||
CostProvider::ModelsDev => fetch_models_dev_pricing(url, client, aliases).await,
|
||||
|
|
@ -220,21 +323,10 @@ async fn fetch_do_pricing(
|
|||
url: &str,
|
||||
client: &reqwest::Client,
|
||||
aliases: &HashMap<String, String>,
|
||||
) -> HashMap<String, f64> {
|
||||
) -> HashMap<String, ModelRates> {
|
||||
match client.get(url).send().await {
|
||||
Ok(resp) => match resp.json::<DoModelList>().await {
|
||||
Ok(list) => list
|
||||
.data
|
||||
.into_iter()
|
||||
.filter_map(|m| {
|
||||
let pricing = m.pricing?;
|
||||
let raw_key = m.model_id.clone();
|
||||
let key = aliases.get(&raw_key).cloned().unwrap_or(raw_key);
|
||||
let cost = pricing.input_price_per_million.unwrap_or(0.0)
|
||||
+ pricing.output_price_per_million.unwrap_or(0.0);
|
||||
Some((key, cost))
|
||||
})
|
||||
.collect(),
|
||||
Ok(list) => parse_do_pricing(list, aliases),
|
||||
Err(err) => {
|
||||
warn!(error = %err, url = %url, "failed to parse digitalocean pricing response");
|
||||
HashMap::new()
|
||||
|
|
@ -247,16 +339,45 @@ async fn fetch_do_pricing(
|
|||
}
|
||||
}
|
||||
|
||||
/// Map the DO catalog into `ModelRates`, scaling DO's per-token prices into the
|
||||
/// per-million convention the rest of the router uses.
|
||||
fn parse_do_pricing(
|
||||
list: DoModelList,
|
||||
aliases: &HashMap<String, String>,
|
||||
) -> HashMap<String, ModelRates> {
|
||||
list.data
|
||||
.into_iter()
|
||||
.filter_map(|m| {
|
||||
let pricing = m.pricing?;
|
||||
let raw_key = m.model_id.clone();
|
||||
let key = aliases.get(&raw_key).cloned().unwrap_or(raw_key);
|
||||
// DO reports prices per token; scale to per-million so the
|
||||
// absolute-USD consumers (the switch-cost gate) are correct.
|
||||
// Relative ranking via `blended()` is unaffected either way.
|
||||
let rates = ModelRates {
|
||||
input_per_million: pricing.input_price_per_million.unwrap_or(0.0)
|
||||
* TOKENS_PER_MILLION,
|
||||
output_per_million: pricing.output_price_per_million.unwrap_or(0.0)
|
||||
* TOKENS_PER_MILLION,
|
||||
cache_read_per_million: pricing
|
||||
.cache_read_input_price_per_million
|
||||
.map(|r| r * TOKENS_PER_MILLION),
|
||||
};
|
||||
Some((key, rates))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// models.dev publishes a top-level object keyed by provider id; each provider
|
||||
/// carries a `models` map whose keys are `creator/model` ids and whose `cost`
|
||||
/// block holds per-million USD rates. We sum input + output (mirroring the DO
|
||||
/// ranking metric) and key the result by `creator/model_id` so it lines up with
|
||||
/// Plano's `provider/model` routing names.
|
||||
/// block holds per-million USD rates. We keep the structured rates (including
|
||||
/// `cache_read` when published) and key the result by `creator/model_id` so it
|
||||
/// lines up with Plano's `provider/model` routing names.
|
||||
async fn fetch_models_dev_pricing(
|
||||
url: &str,
|
||||
client: &reqwest::Client,
|
||||
aliases: &HashMap<String, String>,
|
||||
) -> HashMap<String, f64> {
|
||||
) -> HashMap<String, ModelRates> {
|
||||
match client.get(url).send().await {
|
||||
Ok(resp) => match resp.json::<HashMap<String, ModelsDevProvider>>().await {
|
||||
Ok(providers) => parse_models_dev_pricing(providers, aliases),
|
||||
|
|
@ -275,7 +396,7 @@ async fn fetch_models_dev_pricing(
|
|||
fn parse_models_dev_pricing(
|
||||
providers: HashMap<String, ModelsDevProvider>,
|
||||
aliases: &HashMap<String, String>,
|
||||
) -> HashMap<String, f64> {
|
||||
) -> HashMap<String, ModelRates> {
|
||||
let mut out = HashMap::new();
|
||||
for (provider_id, provider) in providers {
|
||||
for (model_key, model) in provider.models {
|
||||
|
|
@ -286,11 +407,15 @@ fn parse_models_dev_pricing(
|
|||
// First-party providers use bare model keys (`claude-opus-4-5`),
|
||||
// so compose `provider/model` to line up with Plano routing names.
|
||||
let raw_key = format!("{provider_id}/{model_key}");
|
||||
let total = input + output;
|
||||
let rates = ModelRates {
|
||||
input_per_million: input,
|
||||
output_per_million: output,
|
||||
cache_read_per_million: cost.cache_read,
|
||||
};
|
||||
let key = aliases.get(&raw_key).cloned().unwrap_or(raw_key);
|
||||
out.insert(key, total);
|
||||
out.insert(key, rates);
|
||||
// Also register the bare model id as a fallback lookup.
|
||||
out.entry(model_key).or_insert(total);
|
||||
out.entry(model_key).or_insert(rates);
|
||||
}
|
||||
}
|
||||
out
|
||||
|
|
@ -356,6 +481,46 @@ mod tests {
|
|||
SelectionPolicy { prefer }
|
||||
}
|
||||
|
||||
fn service_with(
|
||||
cost: HashMap<String, f64>,
|
||||
latency: HashMap<String, f64>,
|
||||
) -> ModelMetricsService {
|
||||
ModelMetricsService {
|
||||
cost: Arc::new(RwLock::new(cost)),
|
||||
rates: Arc::new(RwLock::new(HashMap::new())),
|
||||
latency: Arc::new(RwLock::new(latency)),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_cost_openai_convention_subtracts_cached_from_prompt() {
|
||||
// input $3/M, output $15/M, cached read $0.30/M. prompt_tokens (=1000) INCLUDES
|
||||
// the 400 cached tokens → uncached = 600.
|
||||
let rates = ModelRates {
|
||||
input_per_million: 3.0,
|
||||
output_per_million: 15.0,
|
||||
cache_read_per_million: Some(0.30),
|
||||
};
|
||||
let (input, output) = rates.request_cost_usd(1000, 400, 0, 200, true, 0.1);
|
||||
// 600*3 + 400*0.30 = 1800 + 120 = 1920 micro-$/M → /1e6
|
||||
assert!((input - (1920.0 / 1_000_000.0)).abs() < 1e-12);
|
||||
assert!((output - (200.0 * 15.0 / 1_000_000.0)).abs() < 1e-12);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_cost_anthropic_convention_prices_uncached_and_creation() {
|
||||
// Anthropic: prompt_tokens (=input_tokens=100) EXCLUDES cached read (30) and
|
||||
// creation (20). No published cached rate → falls back to input*discount.
|
||||
let rates = ModelRates {
|
||||
input_per_million: 3.0,
|
||||
output_per_million: 15.0,
|
||||
cache_read_per_million: None,
|
||||
};
|
||||
let (input, _output) = rates.request_cost_usd(100, 30, 20, 50, false, 0.1);
|
||||
// uncached 100*3 + cached 30*(3*0.1=0.3) + creation 20*3 = 300 + 9 + 60 = 369
|
||||
assert!((input - (369.0 / 1_000_000.0)).abs() < 1e-12);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rank_by_ascending_metric_picks_lowest_first() {
|
||||
let models = vec!["a".to_string(), "b".to_string(), "c".to_string()];
|
||||
|
|
@ -386,15 +551,15 @@ mod tests {
|
|||
|
||||
#[tokio::test]
|
||||
async fn test_rank_models_cheapest() {
|
||||
let service = ModelMetricsService {
|
||||
cost: Arc::new(RwLock::new({
|
||||
let service = service_with(
|
||||
{
|
||||
let mut m = HashMap::new();
|
||||
m.insert("gpt-4o".to_string(), 0.005);
|
||||
m.insert("gpt-4o-mini".to_string(), 0.0001);
|
||||
m
|
||||
})),
|
||||
latency: Arc::new(RwLock::new(HashMap::new())),
|
||||
};
|
||||
},
|
||||
HashMap::new(),
|
||||
);
|
||||
let models = vec!["gpt-4o".to_string(), "gpt-4o-mini".to_string()];
|
||||
let result = service
|
||||
.rank_models(&models, &make_policy(SelectionPreference::Cheapest))
|
||||
|
|
@ -404,15 +569,12 @@ mod tests {
|
|||
|
||||
#[tokio::test]
|
||||
async fn test_rank_models_fastest() {
|
||||
let service = ModelMetricsService {
|
||||
cost: Arc::new(RwLock::new(HashMap::new())),
|
||||
latency: Arc::new(RwLock::new({
|
||||
let mut m = HashMap::new();
|
||||
m.insert("gpt-4o".to_string(), 200.0);
|
||||
m.insert("claude-sonnet".to_string(), 120.0);
|
||||
m
|
||||
})),
|
||||
};
|
||||
let service = service_with(HashMap::new(), {
|
||||
let mut m = HashMap::new();
|
||||
m.insert("gpt-4o".to_string(), 200.0);
|
||||
m.insert("claude-sonnet".to_string(), 120.0);
|
||||
m
|
||||
});
|
||||
let models = vec!["gpt-4o".to_string(), "claude-sonnet".to_string()];
|
||||
let result = service
|
||||
.rank_models(&models, &make_policy(SelectionPreference::Fastest))
|
||||
|
|
@ -422,10 +584,7 @@ mod tests {
|
|||
|
||||
#[tokio::test]
|
||||
async fn test_rank_models_fallback_no_metrics() {
|
||||
let service = ModelMetricsService {
|
||||
cost: Arc::new(RwLock::new(HashMap::new())),
|
||||
latency: Arc::new(RwLock::new(HashMap::new())),
|
||||
};
|
||||
let service = service_with(HashMap::new(), HashMap::new());
|
||||
let models = vec!["model-a".to_string(), "model-b".to_string()];
|
||||
let result = service
|
||||
.rank_models(&models, &make_policy(SelectionPreference::Cheapest))
|
||||
|
|
@ -435,14 +594,14 @@ mod tests {
|
|||
|
||||
#[tokio::test]
|
||||
async fn test_rank_models_partial_data_appended_last() {
|
||||
let service = ModelMetricsService {
|
||||
cost: Arc::new(RwLock::new({
|
||||
let service = service_with(
|
||||
{
|
||||
let mut m = HashMap::new();
|
||||
m.insert("gpt-4o".to_string(), 0.005);
|
||||
m
|
||||
})),
|
||||
latency: Arc::new(RwLock::new(HashMap::new())),
|
||||
};
|
||||
},
|
||||
HashMap::new(),
|
||||
);
|
||||
let models = vec!["gpt-4o-mini".to_string(), "gpt-4o".to_string()];
|
||||
let result = service
|
||||
.rank_models(&models, &make_policy(SelectionPreference::Cheapest))
|
||||
|
|
@ -452,15 +611,15 @@ mod tests {
|
|||
|
||||
#[tokio::test]
|
||||
async fn test_rank_models_none_preserves_order() {
|
||||
let service = ModelMetricsService {
|
||||
cost: Arc::new(RwLock::new({
|
||||
let service = service_with(
|
||||
{
|
||||
let mut m = HashMap::new();
|
||||
m.insert("gpt-4o-mini".to_string(), 0.0001);
|
||||
m.insert("gpt-4o".to_string(), 0.005);
|
||||
m
|
||||
})),
|
||||
latency: Arc::new(RwLock::new(HashMap::new())),
|
||||
};
|
||||
},
|
||||
HashMap::new(),
|
||||
);
|
||||
let models = vec!["gpt-4o".to_string(), "gpt-4o-mini".to_string()];
|
||||
let result = service
|
||||
.rank_models(&models, &make_policy(SelectionPreference::None))
|
||||
|
|
@ -469,12 +628,61 @@ mod tests {
|
|||
assert_eq!(result, vec!["gpt-4o", "gpt-4o-mini"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_do_pricing_scales_per_token_to_per_million() {
|
||||
// DO's catalog returns USD *per token* despite the `_per_million` names.
|
||||
// These are the real values from the live catalog.
|
||||
let json = r#"{
|
||||
"data": [
|
||||
{
|
||||
"model_id": "digitalocean/anthropic-claude-4.6-sonnet",
|
||||
"pricing": {
|
||||
"input_price_per_million": 3e-6,
|
||||
"output_price_per_million": 15e-6,
|
||||
"cache_read_input_price_per_million": 3e-7
|
||||
}
|
||||
},
|
||||
{
|
||||
"model_id": "digitalocean/openai-gpt-4o",
|
||||
"pricing": {
|
||||
"input_price_per_million": 2.5e-6,
|
||||
"output_price_per_million": 10e-6
|
||||
}
|
||||
}
|
||||
]
|
||||
}"#;
|
||||
let list: DoModelList = serde_json::from_str(json).unwrap();
|
||||
let rates = parse_do_pricing(list, &HashMap::new());
|
||||
|
||||
// Scaled to per-million so the switch-cost gate compares real dollars.
|
||||
let sonnet = rates
|
||||
.get("digitalocean/anthropic-claude-4.6-sonnet")
|
||||
.unwrap();
|
||||
assert!((sonnet.input_per_million - 3.0).abs() < 1e-9);
|
||||
assert!((sonnet.output_per_million - 15.0).abs() < 1e-9);
|
||||
// DO *does* publish a cached-read rate — it must be captured, not dropped.
|
||||
assert_eq!(sonnet.cache_read_per_million, Some(0.3));
|
||||
|
||||
let gpt4o = rates.get("digitalocean/openai-gpt-4o").unwrap();
|
||||
assert!((gpt4o.input_per_million - 2.5).abs() < 1e-9);
|
||||
assert_eq!(gpt4o.cache_read_per_million, None);
|
||||
|
||||
// Regression: a ~5k-token switch off warm sonnet to cold gpt-4o must now
|
||||
// cost real dollars, not ~1e-8. This is the bug the live gate test caught.
|
||||
let switch_cost =
|
||||
5_000.0 / 1_000_000.0 * (gpt4o.input_per_million - sonnet.cached_input_rate(0.1));
|
||||
assert!(
|
||||
switch_cost > 0.01,
|
||||
"expected ~$0.011 of switch cost, got {switch_cost}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_models_dev_pricing_composes_provider_keys() {
|
||||
let json = r#"{
|
||||
"anthropic": {
|
||||
"models": {
|
||||
"claude-opus-4-5": {"cost": {"input": 5.0, "output": 25.0}}
|
||||
"claude-opus-4-5": {"cost": {"input": 5.0, "output": 25.0, "cache_read": 0.5}}
|
||||
}
|
||||
},
|
||||
"groq": {
|
||||
|
|
@ -486,14 +694,21 @@ mod tests {
|
|||
}"#;
|
||||
let providers: HashMap<String, ModelsDevProvider> = serde_json::from_str(json).unwrap();
|
||||
let aliases = HashMap::new();
|
||||
let prices = parse_models_dev_pricing(providers, &aliases);
|
||||
let rates = parse_models_dev_pricing(providers, &aliases);
|
||||
|
||||
assert_eq!(prices.get("anthropic/claude-opus-4-5"), Some(&30.0));
|
||||
assert_eq!(prices.get("groq/llama-3.3-70b-versatile"), Some(&1.38));
|
||||
let opus = rates.get("anthropic/claude-opus-4-5").unwrap();
|
||||
assert_eq!(opus.blended(), 30.0);
|
||||
assert_eq!(opus.cache_read_per_million, Some(0.5));
|
||||
let llama = rates.get("groq/llama-3.3-70b-versatile").unwrap();
|
||||
assert_eq!(llama.blended(), 1.38);
|
||||
assert_eq!(llama.cache_read_per_million, None);
|
||||
// bare fallback also registered
|
||||
assert_eq!(prices.get("claude-opus-4-5"), Some(&30.0));
|
||||
assert_eq!(
|
||||
rates.get("claude-opus-4-5").map(|r| r.blended()),
|
||||
Some(30.0)
|
||||
);
|
||||
// models with no cost block are skipped
|
||||
assert!(!prices.contains_key("groq/whisper-large-v3-turbo"));
|
||||
assert!(!rates.contains_key("groq/whisper-large-v3-turbo"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
@ -507,10 +722,53 @@ mod tests {
|
|||
"openai/gpt-oss-120b".to_string(),
|
||||
"openai/gpt-4o".to_string(),
|
||||
);
|
||||
let prices = parse_models_dev_pricing(providers, &aliases);
|
||||
let rates = parse_models_dev_pricing(providers, &aliases);
|
||||
|
||||
assert_eq!(prices.get("openai/gpt-4o"), Some(&3.0));
|
||||
assert!(!prices.contains_key("openai/gpt-oss-120b"));
|
||||
assert_eq!(rates.get("openai/gpt-4o").map(|r| r.blended()), Some(3.0));
|
||||
assert!(!rates.contains_key("openai/gpt-oss-120b"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cached_input_rate_prefers_published_rate() {
|
||||
let with_feed = ModelRates {
|
||||
input_per_million: 3.0,
|
||||
output_per_million: 15.0,
|
||||
cache_read_per_million: Some(0.3),
|
||||
};
|
||||
// Published cache_read wins; the discount is ignored.
|
||||
assert_eq!(with_feed.cached_input_rate(0.5), 0.3);
|
||||
|
||||
let without_feed = ModelRates {
|
||||
input_per_million: 3.0,
|
||||
output_per_million: 15.0,
|
||||
cache_read_per_million: None,
|
||||
};
|
||||
// Falls back to input * discount.
|
||||
assert!((without_feed.cached_input_rate(0.1) - 0.3).abs() < 1e-9);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_model_rates_falls_back_to_bare_model_id() {
|
||||
let service = ModelMetricsService {
|
||||
cost: Arc::new(RwLock::new(HashMap::new())),
|
||||
rates: Arc::new(RwLock::new({
|
||||
let mut m = HashMap::new();
|
||||
m.insert(
|
||||
"claude-sonnet-4-5".to_string(),
|
||||
ModelRates {
|
||||
input_per_million: 3.0,
|
||||
output_per_million: 15.0,
|
||||
cache_read_per_million: Some(0.3),
|
||||
},
|
||||
);
|
||||
m
|
||||
})),
|
||||
latency: Arc::new(RwLock::new(HashMap::new())),
|
||||
};
|
||||
// provider-prefixed lookup falls back to the bare id
|
||||
let rates = service.model_rates("vercel/claude-sonnet-4-5").await;
|
||||
assert_eq!(rates.map(|r| r.input_per_million), Some(3.0));
|
||||
assert!(service.model_rates("unknown/model").await.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
|
|||
|
|
@ -20,9 +20,41 @@ use crate::metrics::labels as metric_labels;
|
|||
use crate::router::orchestrator_model_v1;
|
||||
use crate::session_cache::SessionCache;
|
||||
|
||||
pub use crate::session_cache::CachedRoute;
|
||||
pub use crate::session_cache::SessionBinding;
|
||||
|
||||
const DEFAULT_SESSION_TTL_SECONDS: u64 = 600;
|
||||
const TOKENS_PER_MILLION: f64 = 1_000_000.0;
|
||||
|
||||
/// Input-cost of moving a session off its anchor onto `candidate`.
|
||||
///
|
||||
/// Staying re-reads the whole context at the anchor's read rate (`anchor_read_rate` — the
|
||||
/// cached rate when prompt caching keeps the anchor warm, the plain uncached rate when it
|
||||
/// doesn't). The candidate re-reads `candidate_warm_tokens` (whatever it still has cached
|
||||
/// from an earlier visit this session) at *its* cached rate, and the remaining,
|
||||
/// freshly-appended tokens at its uncached rate. When the candidate is cold
|
||||
/// (`candidate_warm_tokens == 0`) this reduces to the whole context at the uncached rate —
|
||||
/// i.e. a first-time switch. Crediting warm tokens is what makes an A→B→A return cheap:
|
||||
/// on the way back A still holds most of the context, so only the delta is re-ingested.
|
||||
///
|
||||
/// Deliberately input-only: output-token savings are unknowable before the response is
|
||||
/// generated (reasoning models can emit 5-10x the tokens for the same task), so they are
|
||||
/// never credited. Rates are USD per million tokens. Negative when the switch is outright
|
||||
/// cheaper than staying; the caller draws a positive cost down from the overhead cap.
|
||||
pub fn switch_cost_in_usd(
|
||||
context_tokens: u64,
|
||||
candidate_warm_tokens: u64,
|
||||
anchor_read_rate: f64,
|
||||
candidate_uncached_rate: f64,
|
||||
candidate_cached_rate: f64,
|
||||
) -> f64 {
|
||||
let warm = candidate_warm_tokens.min(context_tokens);
|
||||
let fresh = context_tokens - warm;
|
||||
let candidate_cost = (fresh as f64 * candidate_uncached_rate
|
||||
+ warm as f64 * candidate_cached_rate)
|
||||
/ TOKENS_PER_MILLION;
|
||||
let anchor_cost = context_tokens as f64 * anchor_read_rate / TOKENS_PER_MILLION;
|
||||
candidate_cost - anchor_cost
|
||||
}
|
||||
|
||||
pub struct OrchestratorService {
|
||||
orchestrator_url: String,
|
||||
|
|
@ -126,43 +158,126 @@ impl OrchestratorService {
|
|||
}
|
||||
}
|
||||
|
||||
pub async fn get_cached_route(
|
||||
/// Look up a session binding. Warmth is the caller's concern (time since
|
||||
/// `last_used`); this only reports whether a binding exists.
|
||||
pub async fn get_binding(
|
||||
&self,
|
||||
session_id: &str,
|
||||
tenant_id: Option<&str>,
|
||||
) -> Option<CachedRoute> {
|
||||
) -> Option<SessionBinding> {
|
||||
let cache = self.session_cache.as_ref()?;
|
||||
let result = cache.get(&Self::session_key(tenant_id, session_id)).await;
|
||||
bs_metrics::record_session_cache_event(if result.is_some() {
|
||||
metric_labels::SESSION_CACHE_HIT
|
||||
} else {
|
||||
metric_labels::SESSION_CACHE_MISS
|
||||
bs_metrics::record_session_cache_event(match result {
|
||||
Some(_) => metric_labels::SESSION_CACHE_HIT,
|
||||
None => metric_labels::SESSION_CACHE_MISS,
|
||||
});
|
||||
result
|
||||
}
|
||||
|
||||
pub async fn cache_route(
|
||||
/// The GC bound for a session binding: the per-scope override when provided,
|
||||
/// otherwise the global `routing.session_ttl_seconds`. This only governs when an
|
||||
/// idle binding is reclaimed from memory — not whether its cache is warm.
|
||||
pub fn effective_session_ttl(&self, ttl_override_seconds: Option<u64>) -> Duration {
|
||||
ttl_override_seconds
|
||||
.map(Duration::from_secs)
|
||||
.unwrap_or(self.session_ttl)
|
||||
}
|
||||
|
||||
/// Persist a session binding with a GC bound (defaults to `routing.session_ttl_seconds`
|
||||
/// when `gc_ttl` is `None`).
|
||||
pub async fn store_binding(
|
||||
&self,
|
||||
session_id: String,
|
||||
session_id: &str,
|
||||
tenant_id: Option<&str>,
|
||||
model_name: String,
|
||||
route_name: Option<String>,
|
||||
binding: SessionBinding,
|
||||
gc_ttl: Option<Duration>,
|
||||
) {
|
||||
if let Some(ref cache) = self.session_cache {
|
||||
cache
|
||||
.put(
|
||||
&Self::session_key(tenant_id, &session_id),
|
||||
CachedRoute {
|
||||
model_name,
|
||||
route_name,
|
||||
},
|
||||
self.session_ttl,
|
||||
&Self::session_key(tenant_id, session_id),
|
||||
binding,
|
||||
gc_ttl.unwrap_or(self.session_ttl),
|
||||
)
|
||||
.await;
|
||||
bs_metrics::record_session_cache_event(metric_labels::SESSION_CACHE_STORE);
|
||||
}
|
||||
}
|
||||
|
||||
/// Structured per-million pricing for a model, from the configured cost feed.
|
||||
/// `None` when no cost source is configured or the model is unknown to the feed.
|
||||
pub async fn model_rates(&self, model: &str) -> Option<super::model_metrics::ModelRates> {
|
||||
self.metrics_service.as_ref()?.model_rates(model).await
|
||||
}
|
||||
|
||||
/// Estimate the input-cost (USD) of switching a session from `anchor_model`
|
||||
/// (the model that handled the latest request) to `candidate_model`.
|
||||
///
|
||||
/// When `caching_enabled`, the anchor is warm, so staying re-reads the context at its
|
||||
/// *cached* rate and `candidate_warm_tokens` (context the candidate still holds from an
|
||||
/// earlier visit this session) re-read at the candidate's cached rate. When caching is
|
||||
/// off there is no warm cache on either side: staying re-reads at the anchor's *plain*
|
||||
/// (uncached) rate and the candidate re-ingests the whole context uncached (no
|
||||
/// warm-token credit), so the cost collapses to
|
||||
/// `context_tokens x (candidate_uncached_rate - anchor_uncached_rate) / 1M`.
|
||||
///
|
||||
/// Fetches per-model rates from the configured cost feed; returns `None` when pricing
|
||||
/// is missing for either side so the caller can fail open (switch freely) rather than
|
||||
/// veto the router on guesswork. `cache_read_discount` estimates a model's cached-read
|
||||
/// rate when the feed doesn't publish one. Negative when the switch is outright cheaper.
|
||||
pub async fn estimate_switch_cost_in_usd(
|
||||
&self,
|
||||
context_tokens: u64,
|
||||
anchor_model: &str,
|
||||
candidate_model: &str,
|
||||
candidate_warm_tokens: u64,
|
||||
cache_read_discount: f64,
|
||||
caching_enabled: bool,
|
||||
) -> Option<f64> {
|
||||
let anchor = self.model_rates(anchor_model).await?;
|
||||
let candidate = self.model_rates(candidate_model).await?;
|
||||
let (anchor_read_rate, candidate_warm_tokens) = if caching_enabled {
|
||||
(
|
||||
anchor.cached_input_rate(cache_read_discount),
|
||||
candidate_warm_tokens,
|
||||
)
|
||||
} else {
|
||||
(anchor.input_per_million, 0)
|
||||
};
|
||||
Some(switch_cost_in_usd(
|
||||
context_tokens,
|
||||
candidate_warm_tokens,
|
||||
anchor_read_rate,
|
||||
candidate.input_per_million,
|
||||
candidate.cached_input_rate(cache_read_discount),
|
||||
))
|
||||
}
|
||||
|
||||
/// This turn's contribution to the session's *never-switch* baseline: the USD cost of
|
||||
/// reading `context_tokens` on `model` — the session's `default_model`, i.e. what it
|
||||
/// would have paid by never switching (not the possibly-drifted current anchor). Priced
|
||||
/// at the cached input rate when `caching_enabled` (the never-switch path stays warm),
|
||||
/// at the plain uncached rate otherwise (nothing is cached, so every turn re-reads the
|
||||
/// full context). Summed across turns, this is the denominator the percentage overhead
|
||||
/// cap is measured against. `None` when the model has no pricing (the caller then can't
|
||||
/// grow the baseline this turn).
|
||||
pub async fn context_read_cost_in_usd(
|
||||
&self,
|
||||
context_tokens: u64,
|
||||
model: &str,
|
||||
cache_read_discount: f64,
|
||||
caching_enabled: bool,
|
||||
) -> Option<f64> {
|
||||
let rates = self.model_rates(model).await?;
|
||||
let context_millions = context_tokens as f64 / TOKENS_PER_MILLION;
|
||||
let rate = if caching_enabled {
|
||||
rates.cached_input_rate(cache_read_discount)
|
||||
} else {
|
||||
rates.input_per_million
|
||||
};
|
||||
Some(context_millions * rate)
|
||||
}
|
||||
|
||||
// ---- LLM routing ----
|
||||
|
||||
pub async fn determine_route(
|
||||
|
|
@ -348,86 +463,190 @@ mod tests {
|
|||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cache_miss_returns_none() {
|
||||
let svc = make_orchestrator_service(600, 100);
|
||||
assert!(svc
|
||||
.get_cached_route("unknown-session", None)
|
||||
.await
|
||||
.is_none());
|
||||
fn binding(model: &str, route_name: Option<&str>) -> SessionBinding {
|
||||
SessionBinding {
|
||||
anchor_model: model.to_string(),
|
||||
default_model: model.to_string(),
|
||||
route_name: route_name.map(|r| r.to_string()),
|
||||
prefix_hash: None,
|
||||
last_used: std::time::SystemTime::now(),
|
||||
cached_tokens: 0,
|
||||
baseline_usd: 0.0,
|
||||
switch_spend_usd: 0.0,
|
||||
switches: 0,
|
||||
session_cost_usd: 0.0,
|
||||
history: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cache_hit_returns_cached_route() {
|
||||
async fn test_cache_miss_returns_none() {
|
||||
let svc = make_orchestrator_service(600, 100);
|
||||
svc.cache_route(
|
||||
"s1".to_string(),
|
||||
None,
|
||||
"gpt-4o".to_string(),
|
||||
Some("code".to_string()),
|
||||
)
|
||||
.await;
|
||||
assert!(svc.get_binding("unknown-session", None).await.is_none());
|
||||
}
|
||||
|
||||
let cached = svc.get_cached_route("s1", None).await.unwrap();
|
||||
assert_eq!(cached.model_name, "gpt-4o");
|
||||
#[tokio::test]
|
||||
async fn test_cache_hit_returns_binding() {
|
||||
let svc = make_orchestrator_service(600, 100);
|
||||
svc.store_binding("s1", None, binding("gpt-4o", Some("code")), None)
|
||||
.await;
|
||||
|
||||
let cached = svc.get_binding("s1", None).await.unwrap();
|
||||
assert_eq!(cached.anchor_model, "gpt-4o");
|
||||
assert_eq!(cached.route_name, Some("code".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cache_expired_entry_returns_none() {
|
||||
let svc = make_orchestrator_service(0, 100);
|
||||
svc.cache_route("s1".to_string(), None, "gpt-4o".to_string(), None)
|
||||
svc.store_binding("s1", None, binding("gpt-4o", None), None)
|
||||
.await;
|
||||
assert!(svc.get_cached_route("s1", None).await.is_none());
|
||||
assert!(svc.get_binding("s1", None).await.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_expired_entries_not_returned() {
|
||||
let svc = make_orchestrator_service(0, 100);
|
||||
svc.cache_route("s1".to_string(), None, "gpt-4o".to_string(), None)
|
||||
svc.store_binding("s1", None, binding("gpt-4o", None), None)
|
||||
.await;
|
||||
svc.cache_route("s2".to_string(), None, "claude".to_string(), None)
|
||||
svc.store_binding("s2", None, binding("claude", None), None)
|
||||
.await;
|
||||
|
||||
assert!(svc.get_cached_route("s1", None).await.is_none());
|
||||
assert!(svc.get_cached_route("s2", None).await.is_none());
|
||||
assert!(svc.get_binding("s1", None).await.is_none());
|
||||
assert!(svc.get_binding("s2", None).await.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cache_evicts_oldest_when_full() {
|
||||
let svc = make_orchestrator_service(600, 2);
|
||||
svc.cache_route("s1".to_string(), None, "model-a".to_string(), None)
|
||||
svc.store_binding("s1", None, binding("model-a", None), None)
|
||||
.await;
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
svc.cache_route("s2".to_string(), None, "model-b".to_string(), None)
|
||||
svc.store_binding("s2", None, binding("model-b", None), None)
|
||||
.await;
|
||||
|
||||
svc.cache_route("s3".to_string(), None, "model-c".to_string(), None)
|
||||
svc.store_binding("s3", None, binding("model-c", None), None)
|
||||
.await;
|
||||
|
||||
assert!(svc.get_cached_route("s1", None).await.is_none());
|
||||
assert!(svc.get_cached_route("s2", None).await.is_some());
|
||||
assert!(svc.get_cached_route("s3", None).await.is_some());
|
||||
assert!(svc.get_binding("s1", None).await.is_none());
|
||||
assert!(svc.get_binding("s2", None).await.is_some());
|
||||
assert!(svc.get_binding("s3", None).await.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cache_update_existing_session_does_not_evict() {
|
||||
let svc = make_orchestrator_service(600, 2);
|
||||
svc.cache_route("s1".to_string(), None, "model-a".to_string(), None)
|
||||
svc.store_binding("s1", None, binding("model-a", None), None)
|
||||
.await;
|
||||
svc.cache_route("s2".to_string(), None, "model-b".to_string(), None)
|
||||
svc.store_binding("s2", None, binding("model-b", None), None)
|
||||
.await;
|
||||
|
||||
svc.cache_route(
|
||||
"s1".to_string(),
|
||||
svc.store_binding("s1", None, binding("model-a-updated", Some("route")), None)
|
||||
.await;
|
||||
|
||||
let s1 = svc.get_binding("s1", None).await.unwrap();
|
||||
assert_eq!(s1.anchor_model, "model-a-updated");
|
||||
assert!(svc.get_binding("s2", None).await.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_gc_ttl_override_extends_binding_lifetime() {
|
||||
// Global GC bound of 0 would reclaim immediately; the per-call override keeps it.
|
||||
let svc = make_orchestrator_service(0, 100);
|
||||
svc.store_binding(
|
||||
"s1",
|
||||
None,
|
||||
"model-a-updated".to_string(),
|
||||
Some("route".to_string()),
|
||||
binding("gpt-4o", None),
|
||||
Some(Duration::from_secs(600)),
|
||||
)
|
||||
.await;
|
||||
let cached = svc.get_binding("s1", None).await.unwrap();
|
||||
assert_eq!(cached.anchor_model, "gpt-4o");
|
||||
}
|
||||
|
||||
let s1 = svc.get_cached_route("s1", None).await.unwrap();
|
||||
assert_eq!(s1.model_name, "model-a-updated");
|
||||
assert!(svc.get_cached_route("s2", None).await.is_some());
|
||||
#[tokio::test]
|
||||
async fn test_binding_fields_round_trip_through_cache() {
|
||||
let svc = make_orchestrator_service(600, 100);
|
||||
let mut b = binding("gpt-4o", None);
|
||||
b.prefix_hash = Some(0xdead_beef);
|
||||
b.cached_tokens = 12_345;
|
||||
b.baseline_usd = 1.5;
|
||||
b.switch_spend_usd = 0.42;
|
||||
b.switches = 3;
|
||||
b.session_cost_usd = 2.75;
|
||||
svc.store_binding("s1", None, b, None).await;
|
||||
|
||||
let cached = svc.get_binding("s1", None).await.unwrap();
|
||||
assert_eq!(cached.prefix_hash, Some(0xdead_beef));
|
||||
assert_eq!(cached.cached_tokens, 12_345);
|
||||
assert!((cached.baseline_usd - 1.5).abs() < 1e-9);
|
||||
assert!((cached.switch_spend_usd - 0.42).abs() < 1e-9);
|
||||
assert_eq!(cached.switches, 3);
|
||||
assert!((cached.session_cost_usd - 2.75).abs() < 1e-9);
|
||||
}
|
||||
|
||||
// ---- switch-cost math ----
|
||||
//
|
||||
// Real models.dev rates (USD per million input tokens):
|
||||
// claude-opus-4-1: input 15, cache_read 1.5
|
||||
// claude-sonnet-4-5: input 3, cache_read 0.3
|
||||
// claude-haiku-4-5: input 1, cache_read 0.1
|
||||
// gpt-4.1: input 2, cache_read 0.5
|
||||
|
||||
#[test]
|
||||
fn negative_cost_when_candidate_undercuts_cached_rate() {
|
||||
// Anchor opus (cached 1.5) -> haiku (uncached 1.0) over 100k context, cold
|
||||
// candidate: cost = 0.1M x (1.0 - 1.5) = -$0.05 — cheaper even after re-reading.
|
||||
let cost = switch_cost_in_usd(100_000, 0, 1.5, 1.0, 0.1);
|
||||
assert!((cost - (-0.05)).abs() < 1e-9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn positive_cost_when_candidate_pricier_than_cached_rate() {
|
||||
// Anchor opus (cached 1.5) -> gpt-4.1 (uncached 2.0) over 100k, cold candidate:
|
||||
// cost = 0.1M x (2.0 - 1.5) = +$0.05.
|
||||
let cost = switch_cost_in_usd(100_000, 0, 1.5, 2.0, 0.5);
|
||||
assert!((cost - 0.05).abs() < 1e-9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn large_context_amplifies_cost() {
|
||||
// Anchor sonnet (cached 0.3) -> gpt-5.5-class (uncached 5.0) over 150k, cold:
|
||||
// cost = 0.15M x (5.0 - 0.3) = +$0.705.
|
||||
let cost = switch_cost_in_usd(150_000, 0, 0.3, 5.0, 0.5);
|
||||
assert!((cost - 0.705).abs() < 1e-9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cost_scales_linearly_with_context() {
|
||||
let small = switch_cost_in_usd(10_000, 0, 0.3, 0.8, 0.1);
|
||||
let large = switch_cost_in_usd(1_000_000, 0, 0.3, 0.8, 0.1);
|
||||
assert!((large / small - 100.0).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tiny_context_cost_is_negligible() {
|
||||
// 2k-token chat: even an expensive candidate costs ~$0.009.
|
||||
let cost = switch_cost_in_usd(2_000, 0, 0.3, 5.0, 0.5);
|
||||
assert!(cost < 0.01);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn warm_return_charges_only_the_delta() {
|
||||
// Anchor sonnet (cached 0.3). Candidate gpt-4.1 (uncached 2.0, cached 0.5) is
|
||||
// still warm from an earlier visit holding 90k of the 100k context. Only the
|
||||
// 10k fresh tokens re-read at 2.0; the 90k warm tokens re-read at 0.5:
|
||||
// candidate = (10k x 2.0 + 90k x 0.5)/1M = 0.020 + 0.045 = $0.065
|
||||
// anchor = 100k x 0.3 /1M = $0.030
|
||||
// switch = 0.065 - 0.030 = +$0.035
|
||||
let warm = switch_cost_in_usd(100_000, 90_000, 0.3, 2.0, 0.5);
|
||||
assert!(
|
||||
(warm - 0.035).abs() < 1e-9,
|
||||
"warm-return cost {warm} != 0.035"
|
||||
);
|
||||
// Cold, same switch re-reads the whole 100k at 2.0: (100k x 2.0)/1M - 0.03 = $0.17.
|
||||
let cold = switch_cost_in_usd(100_000, 0, 0.3, 2.0, 0.5);
|
||||
assert!((cold - 0.17).abs() < 1e-9, "cold cost {cold} != 0.17");
|
||||
assert!(warm < cold, "returning to a warm model must be cheaper");
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -9,9 +9,9 @@ use lru::LruCache;
|
|||
use tokio::sync::Mutex;
|
||||
use tracing::info;
|
||||
|
||||
use super::{CachedRoute, SessionCache};
|
||||
use super::{SessionBinding, SessionCache};
|
||||
|
||||
type CacheStore = Mutex<LruCache<String, (CachedRoute, Instant, Duration)>>;
|
||||
type CacheStore = Mutex<LruCache<String, (SessionBinding, Instant, Duration)>>;
|
||||
|
||||
pub struct MemorySessionCache {
|
||||
store: Arc<CacheStore>,
|
||||
|
|
@ -38,6 +38,8 @@ impl MemorySessionCache {
|
|||
|
||||
async fn evict_expired(store: &CacheStore) {
|
||||
let mut cache = store.lock().await;
|
||||
// The TTL is only a GC bound: drop bindings once they've outlived the window
|
||||
// in which they could plausibly still be warm.
|
||||
let expired: Vec<String> = cache
|
||||
.iter()
|
||||
.filter(|(_, (_, inserted_at, ttl))| inserted_at.elapsed() >= *ttl)
|
||||
|
|
@ -59,21 +61,21 @@ impl MemorySessionCache {
|
|||
|
||||
#[async_trait]
|
||||
impl SessionCache for MemorySessionCache {
|
||||
async fn get(&self, key: &str) -> Option<CachedRoute> {
|
||||
async fn get(&self, key: &str) -> Option<SessionBinding> {
|
||||
let mut cache = self.store.lock().await;
|
||||
if let Some((route, inserted_at, ttl)) = cache.get(key) {
|
||||
if let Some((binding, inserted_at, ttl)) = cache.get(key) {
|
||||
if inserted_at.elapsed() < *ttl {
|
||||
return Some(route.clone());
|
||||
return Some(binding.clone());
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
async fn put(&self, key: &str, route: CachedRoute, ttl: Duration) {
|
||||
async fn put(&self, key: &str, binding: SessionBinding, ttl: Duration) {
|
||||
self.store
|
||||
.lock()
|
||||
.await
|
||||
.put(key.to_string(), (route, Instant::now(), ttl));
|
||||
.put(key.to_string(), (binding, Instant::now(), ttl));
|
||||
}
|
||||
|
||||
async fn remove(&self, key: &str) {
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
use std::sync::Arc;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use common::configuration::Configuration;
|
||||
|
|
@ -8,21 +9,155 @@ use tracing::{debug, info};
|
|||
pub mod memory;
|
||||
pub mod redis;
|
||||
|
||||
/// A conversation's binding to a model, plus the state the session router needs to
|
||||
/// reason about cache warmth and switch affordability across turns.
|
||||
///
|
||||
/// Warmth is no longer derived from the cache's own expiry — the entry is kept alive
|
||||
/// as a plain KV value (subject only to a GC bound) and the router decides warmth from
|
||||
/// [`SessionBinding::last_used`] against the provider's cache window. This is what lets
|
||||
/// the decision path reason about warmth without ever seeing a provider response.
|
||||
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
|
||||
pub struct CachedRoute {
|
||||
pub model_name: String,
|
||||
pub struct SessionBinding {
|
||||
/// Provider-qualified model that handled the latest request (e.g. `openai/gpt-4o`).
|
||||
/// This is what the session is currently *warm on*; a proposed switch's cost is
|
||||
/// measured against reading the context at this model's cached rate. It tracks the
|
||||
/// last dispatched model and changes whenever a switch is honored.
|
||||
pub anchor_model: String,
|
||||
/// Provider-qualified model the session started on this warm episode — the model it
|
||||
/// would have stayed on had it *never switched*. The never-switch baseline is priced
|
||||
/// against this (not `anchor_model`, which drifts as switches happen). Set when a warm
|
||||
/// episode begins and preserved across its turns.
|
||||
#[serde(default)]
|
||||
pub default_model: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub route_name: Option<String>,
|
||||
/// Hash of the stable prompt prefix (system + tools) observed when the binding was
|
||||
/// stored. Used to detect prefix drift: if a later request's prefix hash differs,
|
||||
/// the provider cache is already lost so a switch is free.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub prefix_hash: Option<u64>,
|
||||
/// When this session was last dispatched. Warmth = `now - last_used` compared
|
||||
/// against the provider's idle/hard cache window.
|
||||
#[serde(default = "SystemTime::now", with = "epoch_secs")]
|
||||
pub last_used: SystemTime,
|
||||
/// Best estimate of the cacheable context size (input tokens) — the tokens a switch
|
||||
/// would have to re-ingest at the uncached rate. Refined from real usage on the
|
||||
/// full-proxy path; the tokenizer estimate on the decision path.
|
||||
#[serde(default)]
|
||||
pub cached_tokens: u64,
|
||||
/// Cumulative *never-switch* baseline (USD) for this warm episode: the running cost
|
||||
/// the session would have paid by staying on its `default_model`. Grows each warm
|
||||
/// turn. This is the denominator the percentage overhead cap is measured against.
|
||||
#[serde(default)]
|
||||
pub baseline_usd: f64,
|
||||
/// Cumulative overhead (USD) actually spent on paid switches this warm episode.
|
||||
/// Monotonic: paid switches add to it, free/cheaper switches never subtract. A paid
|
||||
/// switch is allowed only while `switch_spend_usd + cost <= pct * baseline_usd`.
|
||||
#[serde(default)]
|
||||
pub switch_spend_usd: f64,
|
||||
/// Number of model switches taken during this warm session (observability).
|
||||
#[serde(default)]
|
||||
pub switches: u32,
|
||||
/// Bounded most-recent-first-ish record of the models this session has been
|
||||
/// dispatched to, each with when it was last used and how large the context was
|
||||
/// then. Lets the switch-cost estimate credit a *return* to a still-warm model (it
|
||||
/// re-reads only the tokens appended since, not the whole context) and gives future
|
||||
/// routing policies per-model recency to reason about. Capped at
|
||||
/// [`MAX_ROUTE_HISTORY`] distinct models (LRU-evicted).
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub history: Vec<RouteVisit>,
|
||||
/// Cumulative *actual* cost (USD, input + output) of the whole conversation, priced
|
||||
/// from the configured catalog rates and refined from real usage each turn on the
|
||||
/// full-proxy path. Conversation-level (not per warm episode): it persists across
|
||||
/// cold re-binds and is best-effort (resets only if the binding is evicted).
|
||||
#[serde(default)]
|
||||
pub session_cost_usd: f64,
|
||||
}
|
||||
|
||||
/// One model this session has been dispatched to, with the recency and context size
|
||||
/// needed to estimate whether its provider cache is still warm on a later return.
|
||||
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
|
||||
pub struct RouteVisit {
|
||||
/// Provider-qualified model id (e.g. `openai/gpt-4o`).
|
||||
pub model: String,
|
||||
/// When this model last handled a turn in this session.
|
||||
#[serde(with = "epoch_secs")]
|
||||
pub last_used: SystemTime,
|
||||
/// Context size (input tokens) that model saw on its last turn — the prefix it may
|
||||
/// still have cached if it's revisited before its cache window elapses.
|
||||
#[serde(default)]
|
||||
pub cached_tokens: u64,
|
||||
}
|
||||
|
||||
/// Max distinct models retained in [`SessionBinding::history`]. Small: real sessions
|
||||
/// touch a handful of models, and the entry has to stay compact on the Redis wire.
|
||||
pub const MAX_ROUTE_HISTORY: usize = 8;
|
||||
|
||||
/// Record (or refresh) a model visit in `history`, then LRU-evict down to
|
||||
/// [`MAX_ROUTE_HISTORY`]. Refreshing an existing model updates its recency and context
|
||||
/// size in place rather than appending a duplicate.
|
||||
pub fn record_route_visit(
|
||||
history: &mut Vec<RouteVisit>,
|
||||
model: &str,
|
||||
last_used: SystemTime,
|
||||
cached_tokens: u64,
|
||||
) {
|
||||
if let Some(entry) = history.iter_mut().find(|e| e.model == model) {
|
||||
entry.last_used = last_used;
|
||||
entry.cached_tokens = cached_tokens;
|
||||
} else {
|
||||
history.push(RouteVisit {
|
||||
model: model.to_string(),
|
||||
last_used,
|
||||
cached_tokens,
|
||||
});
|
||||
}
|
||||
while history.len() > MAX_ROUTE_HISTORY {
|
||||
if let Some(idx) = history
|
||||
.iter()
|
||||
.enumerate()
|
||||
.min_by_key(|(_, e)| e.last_used)
|
||||
.map(|(i, _)| i)
|
||||
{
|
||||
history.remove(idx);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Serde helper: persist `SystemTime` as whole epoch seconds so the Redis wire format
|
||||
/// is stable and compact (the default `SystemTime` representation is version-fragile).
|
||||
mod epoch_secs {
|
||||
use super::{Duration, SystemTime, UNIX_EPOCH};
|
||||
use serde::{Deserialize, Deserializer, Serializer};
|
||||
|
||||
pub fn serialize<S: Serializer>(t: &SystemTime, s: S) -> Result<S::Ok, S::Error> {
|
||||
let secs = t
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0);
|
||||
s.serialize_u64(secs)
|
||||
}
|
||||
|
||||
pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<SystemTime, D::Error> {
|
||||
let secs = u64::deserialize(d)?;
|
||||
Ok(UNIX_EPOCH + Duration::from_secs(secs))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait SessionCache: Send + Sync {
|
||||
/// Look up a cached routing decision by key.
|
||||
async fn get(&self, key: &str) -> Option<CachedRoute>;
|
||||
/// Look up a session binding by key. `None` when absent or GC-evicted. Warmth is
|
||||
/// the caller's concern (time since `last_used`), not the cache's.
|
||||
async fn get(&self, key: &str) -> Option<SessionBinding>;
|
||||
|
||||
/// Store a routing decision in the session cache with the given TTL.
|
||||
async fn put(&self, key: &str, route: CachedRoute, ttl: Duration);
|
||||
/// Store a session binding with the given GC TTL. The TTL is only a memory bound
|
||||
/// (keep the entry around at least as long as it could plausibly be warm); it does
|
||||
/// not define warmth.
|
||||
async fn put(&self, key: &str, binding: SessionBinding, ttl: Duration);
|
||||
|
||||
/// Remove a cached routing decision by key.
|
||||
/// Remove a session binding by key.
|
||||
async fn remove(&self, key: &str);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ use async_trait::async_trait;
|
|||
use redis::aio::MultiplexedConnection;
|
||||
use redis::AsyncCommands;
|
||||
|
||||
use super::{CachedRoute, SessionCache};
|
||||
use super::{SessionBinding, SessionCache};
|
||||
|
||||
const KEY_PREFIX: &str = "plano:affinity:";
|
||||
|
||||
|
|
@ -26,18 +26,20 @@ impl RedisSessionCache {
|
|||
|
||||
#[async_trait]
|
||||
impl SessionCache for RedisSessionCache {
|
||||
async fn get(&self, key: &str) -> Option<CachedRoute> {
|
||||
async fn get(&self, key: &str) -> Option<SessionBinding> {
|
||||
let mut conn = self.conn.clone();
|
||||
let value: Option<String> = conn.get(Self::make_key(key)).await.ok()?;
|
||||
value.and_then(|v| serde_json::from_str(&v).ok())
|
||||
}
|
||||
|
||||
async fn put(&self, key: &str, route: CachedRoute, ttl: Duration) {
|
||||
async fn put(&self, key: &str, binding: SessionBinding, ttl: Duration) {
|
||||
let mut conn = self.conn.clone();
|
||||
let Ok(json) = serde_json::to_string(&route) else {
|
||||
// The Redis TTL is only a GC bound; warmth is decided by the router from
|
||||
// `binding.last_used`, not by expiry here.
|
||||
let ttl_secs = ttl.as_secs().max(1);
|
||||
let Ok(json) = serde_json::to_string(&binding) else {
|
||||
return;
|
||||
};
|
||||
let ttl_secs = ttl.as_secs().max(1);
|
||||
let _: Result<(), _> = conn.set_ex(Self::make_key(key), json, ttl_secs).await;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -22,10 +22,15 @@ const STREAM_BUFFER_SIZE: usize = 16;
|
|||
const USAGE_BUFFER_MAX: usize = 2 * 1024 * 1024;
|
||||
use crate::metrics as bs_metrics;
|
||||
use crate::metrics::labels as metric_labels;
|
||||
use crate::router::model_metrics::ModelRates;
|
||||
use crate::router::orchestrator::OrchestratorService;
|
||||
use crate::session_cache::{record_route_visit, RouteVisit, SessionBinding};
|
||||
use crate::signals::otel::emit_signals_to_span;
|
||||
use crate::signals::{SignalAnalyzer, FLAG_MARKER};
|
||||
use crate::tracing::{llm, set_service_name};
|
||||
use hermesllm::apis::openai::Message;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
/// Parsed usage + resolved-model details from a provider response.
|
||||
#[derive(Debug, Default, Clone)]
|
||||
|
|
@ -39,6 +44,11 @@ struct ExtractedUsage {
|
|||
/// The model the upstream actually used. For router aliases (e.g.
|
||||
/// `router:software-engineering`), this differs from the request model.
|
||||
resolved_model: Option<String>,
|
||||
/// Provider convention for `prompt_tokens`: OpenAI-shape usage folds cached tokens
|
||||
/// *into* `prompt_tokens` (so uncached = prompt - cached), while Anthropic-shape
|
||||
/// reports uncached `input_tokens` separately from `cache_read`/`cache_creation`.
|
||||
/// Set at parse time from which field `prompt_tokens` was sourced. Drives cost math.
|
||||
prompt_includes_cached: bool,
|
||||
}
|
||||
|
||||
impl ExtractedUsage {
|
||||
|
|
@ -57,8 +67,9 @@ impl ExtractedUsage {
|
|||
}
|
||||
}
|
||||
if let Some(u) = value.get("usage") {
|
||||
// OpenAI-shape usage
|
||||
// OpenAI-shape usage: `prompt_tokens` includes the cached subset.
|
||||
out.prompt_tokens = u.get("prompt_tokens").and_then(|v| v.as_i64());
|
||||
out.prompt_includes_cached = out.prompt_tokens.is_some();
|
||||
out.completion_tokens = u.get("completion_tokens").and_then(|v| v.as_i64());
|
||||
out.total_tokens = u.get("total_tokens").and_then(|v| v.as_i64());
|
||||
out.cached_input_tokens = u
|
||||
|
|
@ -187,6 +198,44 @@ pub struct LlmMetricsCtx {
|
|||
pub upstream_status: u16,
|
||||
}
|
||||
|
||||
/// Response-side session-update context: refreshes the session binding once the real
|
||||
/// response is in hand, so `last_used` and the context-size estimate (`cached_tokens`)
|
||||
/// reflect the turn that just completed. The routing decision (model, budget, switches)
|
||||
/// was already made and persisted on the request side by [`super::handlers::llm::session_router`];
|
||||
/// this only refines the fields that need the response.
|
||||
pub struct SessionUpdateCtx {
|
||||
pub orchestrator: Arc<OrchestratorService>,
|
||||
pub session_id: String,
|
||||
pub tenant_id: Option<String>,
|
||||
/// Provider-qualified model this request actually ran on (the router's final pick).
|
||||
pub anchor_model: String,
|
||||
/// The session's never-switch model for this episode — preserved across the refresh.
|
||||
pub default_model: String,
|
||||
pub route_name: Option<String>,
|
||||
pub prefix_hash: Option<u64>,
|
||||
/// Cumulative never-switch baseline from the decision — preserved across the refresh.
|
||||
pub baseline_usd: f64,
|
||||
/// Cumulative switch spend from the decision — preserved across the refresh.
|
||||
pub switch_spend_usd: f64,
|
||||
/// Cumulative switch count from the routing decision — preserved across the refresh.
|
||||
pub switches: u32,
|
||||
/// Per-model route history from the decision — preserved across the refresh, with
|
||||
/// the anchor's entry refined to the real prompt-token count.
|
||||
pub history: Vec<RouteVisit>,
|
||||
/// Cumulative actual conversation cost (USD) through prior turns. This turn's real
|
||||
/// cost is added on top once usage is known.
|
||||
pub session_cost_usd: f64,
|
||||
/// Catalog rates for the dispatched model, resolved request-side (the response path
|
||||
/// is synchronous). `None` when no cost source is configured → no cost is computed.
|
||||
pub cost_rates: Option<ModelRates>,
|
||||
/// Cached-read discount used to price cached input when the feed omits a cached rate.
|
||||
pub cache_read_discount: f64,
|
||||
/// Context-size count to fall back to when the response carries no usage block.
|
||||
pub context_tokens: u64,
|
||||
/// GC bound to store the refreshed binding with.
|
||||
pub gc_ttl: Duration,
|
||||
}
|
||||
|
||||
/// A processor that tracks streaming metrics
|
||||
pub struct ObservableStreamProcessor {
|
||||
service_name: String,
|
||||
|
|
@ -201,6 +250,7 @@ pub struct ObservableStreamProcessor {
|
|||
/// from the buffer (they still pass through to the client).
|
||||
response_buffer: Vec<u8>,
|
||||
llm_metrics: Option<LlmMetricsCtx>,
|
||||
session_update: Option<SessionUpdateCtx>,
|
||||
metrics_recorded: bool,
|
||||
}
|
||||
|
||||
|
|
@ -237,6 +287,7 @@ impl ObservableStreamProcessor {
|
|||
messages,
|
||||
response_buffer: Vec::new(),
|
||||
llm_metrics: None,
|
||||
session_update: None,
|
||||
metrics_recorded: false,
|
||||
}
|
||||
}
|
||||
|
|
@ -247,6 +298,103 @@ impl ObservableStreamProcessor {
|
|||
self.llm_metrics = Some(ctx);
|
||||
self
|
||||
}
|
||||
|
||||
/// Attach session-update context so the processor refreshes `last_used` and the
|
||||
/// context-size estimate from the real response once usage is known.
|
||||
pub fn with_session_update(mut self, ctx: SessionUpdateCtx) -> Self {
|
||||
self.session_update = Some(ctx);
|
||||
self
|
||||
}
|
||||
|
||||
/// Refresh the session binding from the response so warmth (`last_used`) and the
|
||||
/// context-size estimate (`cached_tokens`) reflect the completed turn.
|
||||
fn handle_session_update(&mut self, usage: &ExtractedUsage) {
|
||||
let Some(update) = self.session_update.take() else {
|
||||
return;
|
||||
};
|
||||
let SessionUpdateCtx {
|
||||
orchestrator,
|
||||
session_id,
|
||||
tenant_id,
|
||||
anchor_model,
|
||||
default_model,
|
||||
route_name,
|
||||
prefix_hash,
|
||||
baseline_usd,
|
||||
switch_spend_usd,
|
||||
switches,
|
||||
mut history,
|
||||
session_cost_usd,
|
||||
cost_rates,
|
||||
cache_read_discount,
|
||||
context_tokens,
|
||||
gc_ttl,
|
||||
} = update;
|
||||
|
||||
// Prefer the real prompt-token count (the tokens a future switch would re-read)
|
||||
// over the request-side estimate; fall back to the estimate when absent.
|
||||
let cached_tokens = usage
|
||||
.prompt_tokens
|
||||
.filter(|&p| p > 0)
|
||||
.map(|p| p as u64)
|
||||
.unwrap_or(context_tokens);
|
||||
|
||||
// Refine the anchor's route-history entry with the real context size, so a later
|
||||
// return to this model prices its warm portion off actual usage.
|
||||
record_route_visit(
|
||||
&mut history,
|
||||
&anchor_model,
|
||||
SystemTime::now(),
|
||||
cached_tokens,
|
||||
);
|
||||
|
||||
// Price this turn from the catalog rates and roll it into the conversation
|
||||
// total. Emit per-request cost on the (llm) span and carry the running total
|
||||
// into the binding so the next turn's routing span can surface it.
|
||||
let session_cost_usd = if let Some(rates) = cost_rates {
|
||||
let (input_cost, output_cost) = rates.request_cost_usd(
|
||||
usage.prompt_tokens.unwrap_or(0).max(0) as u64,
|
||||
usage.cached_input_tokens.unwrap_or(0).max(0) as u64,
|
||||
usage.cache_creation_tokens.unwrap_or(0).max(0) as u64,
|
||||
usage.completion_tokens.unwrap_or(0).max(0) as u64,
|
||||
usage.prompt_includes_cached,
|
||||
cache_read_discount,
|
||||
);
|
||||
let span = tracing::Span::current();
|
||||
let otel_span = span.context();
|
||||
let otel_span = otel_span.span();
|
||||
otel_span.set_attribute(KeyValue::new(llm::INPUT_COST_IN_USD, input_cost));
|
||||
otel_span.set_attribute(KeyValue::new(llm::OUTPUT_COST_IN_USD, output_cost));
|
||||
otel_span.set_attribute(KeyValue::new(
|
||||
llm::TOTAL_COST_IN_USD,
|
||||
input_cost + output_cost,
|
||||
));
|
||||
session_cost_usd + input_cost + output_cost
|
||||
} else {
|
||||
session_cost_usd
|
||||
};
|
||||
|
||||
bs_metrics::record_session_binding_event(metric_labels::BINDING_EVENT_REFRESH);
|
||||
let binding = SessionBinding {
|
||||
anchor_model,
|
||||
default_model,
|
||||
route_name,
|
||||
prefix_hash,
|
||||
last_used: SystemTime::now(),
|
||||
cached_tokens,
|
||||
baseline_usd,
|
||||
switch_spend_usd,
|
||||
switches,
|
||||
session_cost_usd,
|
||||
history,
|
||||
};
|
||||
// Fire-and-forget: binding bookkeeping must not delay stream completion.
|
||||
tokio::spawn(async move {
|
||||
orchestrator
|
||||
.store_binding(&session_id, tenant_id.as_deref(), binding, Some(gc_ttl))
|
||||
.await;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
impl StreamProcessor for ObservableStreamProcessor {
|
||||
|
|
@ -357,11 +505,47 @@ impl StreamProcessor for ObservableStreamProcessor {
|
|||
v.max(0) as u64,
|
||||
);
|
||||
}
|
||||
// Prompt-cache token counters + hit/miss baseline (cache-blindness is
|
||||
// invisible without these: a reroute that burns a warm cache shows up
|
||||
// here as a miss with zero cache_read tokens).
|
||||
if let Some(v) = usage.cached_input_tokens {
|
||||
bs_metrics::record_llm_tokens(
|
||||
&ctx.provider,
|
||||
&ctx.model,
|
||||
metric_labels::TOKEN_KIND_CACHE_READ,
|
||||
v.max(0) as u64,
|
||||
);
|
||||
}
|
||||
if let Some(v) = usage.cache_creation_tokens {
|
||||
bs_metrics::record_llm_tokens(
|
||||
&ctx.provider,
|
||||
&ctx.model,
|
||||
metric_labels::TOKEN_KIND_CACHE_WRITE,
|
||||
v.max(0) as u64,
|
||||
);
|
||||
}
|
||||
if usage.prompt_tokens.is_some() {
|
||||
let cache_active = usage.cached_input_tokens.unwrap_or(0) > 0
|
||||
|| usage.cache_creation_tokens.unwrap_or(0) > 0;
|
||||
bs_metrics::record_prompt_cache_outcome(
|
||||
&ctx.provider,
|
||||
&ctx.model,
|
||||
if cache_active {
|
||||
metric_labels::PROMPT_CACHE_HIT
|
||||
} else {
|
||||
metric_labels::PROMPT_CACHE_MISS
|
||||
},
|
||||
);
|
||||
}
|
||||
if usage.prompt_tokens.is_none() && usage.completion_tokens.is_none() {
|
||||
bs_metrics::record_llm_tokens_usage_missing(&ctx.provider, &ctx.model);
|
||||
}
|
||||
self.metrics_recorded = true;
|
||||
}
|
||||
|
||||
// Session-binding refresh: update `last_used` + the context-size estimate from
|
||||
// the completed turn (the routing decision itself was made request-side).
|
||||
self.handle_session_update(&usage);
|
||||
// Release the buffered bytes early; nothing downstream needs them.
|
||||
self.response_buffer.clear();
|
||||
self.response_buffer.shrink_to_fit();
|
||||
|
|
@ -612,6 +796,8 @@ mod usage_extraction_tests {
|
|||
assert_eq!(u.total_tokens, Some(46));
|
||||
assert_eq!(u.cached_input_tokens, Some(5));
|
||||
assert_eq!(u.reasoning_tokens, None);
|
||||
// OpenAI folds cached tokens into `prompt_tokens`.
|
||||
assert!(u.prompt_includes_cached);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
@ -623,6 +809,8 @@ mod usage_extraction_tests {
|
|||
assert_eq!(u.total_tokens, Some(150));
|
||||
assert_eq!(u.cached_input_tokens, Some(30));
|
||||
assert_eq!(u.cache_creation_tokens, Some(20));
|
||||
// Anthropic reports uncached `input_tokens` separately from cached reads.
|
||||
assert!(!u.prompt_includes_cached);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
|
|||
|
|
@ -92,6 +92,18 @@ pub mod llm {
|
|||
/// (OpenAI `completion_tokens_details.reasoning_tokens`, Google `thoughts_token_count`)
|
||||
pub const REASONING_TOKENS: &str = "llm.usage.reasoning_tokens";
|
||||
|
||||
/// This request's input-token cost (USD), priced from the catalog rates: uncached
|
||||
/// input at the input rate, cached reads at the cached rate, cache creation at the
|
||||
/// plain input rate. Present only when a cost source is configured.
|
||||
pub const INPUT_COST_IN_USD: &str = "llm.usage.input_cost_usd";
|
||||
|
||||
/// This request's output-token cost (USD) = completion tokens x output rate.
|
||||
pub const OUTPUT_COST_IN_USD: &str = "llm.usage.output_cost_usd";
|
||||
|
||||
/// This request's total cost (USD) = input + output. Sum across a session's turns
|
||||
/// (group by `plano.session_id`) for the conversation total.
|
||||
pub const TOTAL_COST_IN_USD: &str = "llm.usage.total_cost_usd";
|
||||
|
||||
/// Temperature parameter used
|
||||
pub const TEMPERATURE: &str = "llm.temperature";
|
||||
|
||||
|
|
@ -150,6 +162,61 @@ pub mod plano {
|
|||
/// fields (e.g. PostHog). Sourced from the configured
|
||||
/// `tracing.exporters[].distinct_id_header`. Absent for anonymous calls.
|
||||
pub const DISTINCT_ID: &str = "plano.distinct_id";
|
||||
|
||||
/// Whether the session's provider cache was inferred warm at decision time
|
||||
/// (from the idle gap vs. the provider's cache window).
|
||||
pub const CACHE_WARM: &str = "plano.cache.warm";
|
||||
|
||||
/// How long (ms) since the session was last used — the idle gap warmth is measured
|
||||
/// against.
|
||||
pub const CACHE_IDLE_MS: &str = "plano.cache.idle_ms";
|
||||
|
||||
/// Cumulative switching overhead consumed this session, as a percentage of the
|
||||
/// never-switch baseline (`100 * switch_spend / baseline`). Directly comparable to
|
||||
/// the configured `routing.routing_budget.max_overhead_pct`.
|
||||
pub const SESSION_OVERHEAD_PCT: &str = "plano.session.overhead_pct";
|
||||
|
||||
/// Cumulative overhead (USD) actually spent on paid switches this session — the
|
||||
/// numerator behind `plano.session.overhead_pct`.
|
||||
pub const SESSION_SWITCH_SPEND_IN_USD: &str = "plano.session.switch_spend_in_usd";
|
||||
|
||||
/// Cumulative never-switch baseline (USD) — what staying on the anchor would have
|
||||
/// cost so far. The denominator behind `plano.session.overhead_pct`.
|
||||
pub const SESSION_BASELINE_IN_USD: &str = "plano.session.baseline_in_usd";
|
||||
|
||||
/// Cumulative number of model switches taken during this warm session.
|
||||
pub const SESSION_SWITCHES: &str = "plano.session.switches";
|
||||
|
||||
/// Cumulative *actual* cost (USD, input + output) of the whole conversation, priced
|
||||
/// from the configured catalog rates and refined from real usage each turn. Emitted
|
||||
/// on the routing span; reflects cost through the previous turn (this turn isn't
|
||||
/// billed yet at decision time). Full-proxy path only.
|
||||
pub const SESSION_TOTAL_COST_IN_USD: &str = "plano.session.total_cost_in_usd";
|
||||
|
||||
/// Actual input-cost (USD) of the proposed model switch — computed from input-token
|
||||
/// pricing only, output-token cost deliberately excluded. Negative when the candidate
|
||||
/// is outright cheaper than staying on the warm anchor.
|
||||
pub const SWITCH_COST_IN_USD: &str = "plano.switch.cost_in_usd";
|
||||
|
||||
/// Tokens the switch candidate still has cached from an earlier visit this session
|
||||
/// (a return to a still-warm model). These re-read at the candidate's cached rate
|
||||
/// instead of its uncached rate, which is why the switch cost can be far below a
|
||||
/// full re-ingest. Zero for a first-time (cold) switch.
|
||||
pub const SWITCH_CANDIDATE_WARM_TOKENS: &str = "plano.switch.candidate_warm_tokens";
|
||||
|
||||
/// The overhead ceiling (USD) available when the switch was evaluated —
|
||||
/// `max_overhead_pct% * baseline`. A paid switch is allowed while cumulative spend
|
||||
/// plus this switch's cost stays under it. Directly comparable to `cost_in_usd`.
|
||||
pub const SWITCH_OVERHEAD_CEILING_IN_USD: &str = "plano.switch.overhead_ceiling_in_usd";
|
||||
|
||||
/// Switch outcome: "allowed" or "retained".
|
||||
pub const SWITCH_DECISION: &str = "plano.switch.decision";
|
||||
|
||||
/// The route (`provider/model`, plus route name when routed) the routing-budget gate
|
||||
/// *would* have selected had the switch been allowed. Recorded only on a `retained`
|
||||
/// decision when `routing.routing_budget.record_counterfactual` is enabled.
|
||||
/// Telemetry only — the counterfactual model is never dispatched.
|
||||
pub const SWITCH_COUNTERFACTUAL_ROUTE: &str = "plano.switch.counterfactual_route";
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
|
|
|
|||
|
|
@ -33,6 +33,10 @@ pub struct Routing {
|
|||
pub session_ttl_seconds: Option<u64>,
|
||||
pub session_max_entries: Option<usize>,
|
||||
pub session_cache: Option<SessionCacheConfig>,
|
||||
/// Cost gate on model switching within a session. Independent of prompt caching:
|
||||
/// this is a routing decision that applies whenever it is configured, whether or
|
||||
/// not `prompt_caching` is enabled. Presence of this block turns it on.
|
||||
pub routing_budget: Option<RoutingBudget>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
|
|
@ -133,7 +137,7 @@ pub enum StateStorageType {
|
|||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum SelectionPreference {
|
||||
Cheapest,
|
||||
Fastest,
|
||||
|
|
@ -218,6 +222,10 @@ pub struct Configuration {
|
|||
pub model_aliases: Option<HashMap<String, ModelAlias>>,
|
||||
pub overrides: Option<Overrides>,
|
||||
pub routing: Option<Routing>,
|
||||
/// Automatic provider prompt caching. Disabled by default; opt in globally with
|
||||
/// `prompt_caching: { enabled: true }`. Applies across the entire Plano instance
|
||||
/// and never changes which model routing selects.
|
||||
pub prompt_caching: Option<PromptCaching>,
|
||||
pub system_prompt: Option<String>,
|
||||
pub prompt_guards: Option<PromptGuards>,
|
||||
pub prompt_targets: Option<Vec<PromptTarget>>,
|
||||
|
|
@ -244,6 +252,181 @@ pub struct Overrides {
|
|||
pub disable_signals: Option<bool>,
|
||||
}
|
||||
|
||||
/// Automatic prompt caching, configured once for the whole Plano instance.
|
||||
///
|
||||
/// Prompt caching keeps a multi-turn conversation's stable prefix warm in the
|
||||
/// upstream provider's cache. It never influences which model routing selects — it
|
||||
/// only (a) auto-injects provider cache-control markers where supported and
|
||||
/// (b) derives an implicit session key from the stable prompt prefix so follow-up
|
||||
/// turns reuse the same warm cache. An explicit `X-Model-Affinity` header always wins.
|
||||
///
|
||||
/// Disabled by default; opt in with `enabled: true`. The remaining knobs are optional
|
||||
/// tuning that only take effect while caching is enabled.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct PromptCaching {
|
||||
/// Master switch. Defaults to `false` (opt-in).
|
||||
#[serde(default)]
|
||||
pub enabled: bool,
|
||||
/// Derive an implicit session key from the stable prompt prefix so caches survive
|
||||
/// across turns without client changes. Defaults to `true` when caching is enabled.
|
||||
pub session_affinity: Option<bool>,
|
||||
/// Auto-inject provider cache-control markers (e.g. Anthropic `cache_control`).
|
||||
/// Defaults to `true` when caching is enabled.
|
||||
pub inject_cache_control: Option<bool>,
|
||||
/// Minimum estimated prefix tokens before a cache breakpoint is injected.
|
||||
pub min_prefix_tokens: Option<u32>,
|
||||
/// Session pin TTL; falls back to `routing.session_ttl_seconds` when unset.
|
||||
pub session_ttl_seconds: Option<u64>,
|
||||
}
|
||||
|
||||
/// A cumulative per-session overhead cap governing when routing may switch models.
|
||||
///
|
||||
/// This is a routing concern, not a caching one: it applies whenever configured,
|
||||
/// regardless of whether `prompt_caching` is enabled. The default posture is to stick
|
||||
/// to the model a session is warm on. When routing proposes a *different* model,
|
||||
/// switching forces the candidate to re-ingest the whole context at its uncached input
|
||||
/// rate. With prompt caching the anchor is warm, so the cost is the cache-loss delta
|
||||
/// `context_tokens x (candidate_uncached_input_rate - anchor_cached_input_rate)`; without
|
||||
/// caching there is no warm cache to lose, so it is
|
||||
/// `context_tokens x (candidate_uncached_input_rate - anchor_uncached_input_rate)`. That
|
||||
/// input-token cost accrues into the session's cumulative switch spend. The gate allows a paid switch
|
||||
/// only while that spend stays within `max_overhead_pct` percent of the session's
|
||||
/// running *never-switch* baseline (the cost the session would have paid by staying
|
||||
/// on its anchor). A switch that is outright cheaper (negative cost) is free but
|
||||
/// never credits the spend back — the "saving" is vs a path we didn't take, not real
|
||||
/// spendable money. Requires a cost source in `model_metrics_sources` so per-model
|
||||
/// rates are available. Presence of the block turns it on.
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
|
||||
pub struct RoutingBudget {
|
||||
/// Cap on cumulative switching overhead, as a **percentage** of what the session
|
||||
/// would have cost by never switching (a whole number: `20` = 20%). The promise
|
||||
/// is "this conversation bills at most `max_overhead_pct`% above never-switching."
|
||||
/// `0` means "never pay to switch" (only outright-cheaper switches are ever
|
||||
/// allowed); larger values buy more quality-driven switches. Typical range 10–30.
|
||||
pub max_overhead_pct: f64,
|
||||
/// Reset the running baseline/spend totals when a session goes cold and re-binds
|
||||
/// (a fresh warm episode). Defaults to `true`.
|
||||
#[serde(default = "default_true")]
|
||||
pub replenish_on_rebind: bool,
|
||||
/// Fallback used to estimate a model's cached input rate when the pricing feed
|
||||
/// doesn't publish one: `cached_rate = input_rate * cache_read_discount`. A
|
||||
/// pricing detail, not a cost policy. Defaults to 0.1 (cached reads at 10% of
|
||||
/// input, typical for Anthropic-style caches).
|
||||
pub cache_read_discount: Option<f64>,
|
||||
/// When true, a vetoed switch records the route the gate *would* have taken had
|
||||
/// the switch been allowed, as the `plano.switch.counterfactual_route` span
|
||||
/// attribute. Telemetry only — the counterfactual model is never dispatched.
|
||||
/// Useful for evals/benchmarks that want to quantify the road not taken.
|
||||
/// Defaults to `false`.
|
||||
#[serde(default)]
|
||||
pub record_counterfactual: bool,
|
||||
}
|
||||
|
||||
fn default_true() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
/// Fully-resolved routing-budget settings (present only when configured and valid).
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub struct EffectiveRoutingBudget {
|
||||
/// Cumulative switching-overhead cap, as a percentage of the never-switch
|
||||
/// baseline (a whole number: `20` = 20%).
|
||||
pub max_overhead_pct: f64,
|
||||
/// Reset the running baseline/spend totals on cold->warm re-bind.
|
||||
pub replenish_on_rebind: bool,
|
||||
pub cache_read_discount: f64,
|
||||
/// Emit `plano.switch.counterfactual_route` on vetoed switches. Telemetry only.
|
||||
pub record_counterfactual: bool,
|
||||
}
|
||||
|
||||
pub const DEFAULT_CACHE_READ_DISCOUNT: f64 = 0.1;
|
||||
|
||||
impl RoutingBudget {
|
||||
/// Resolve to effective settings, validating the overhead cap and cache-read
|
||||
/// discount.
|
||||
pub fn resolve(&self) -> Result<EffectiveRoutingBudget, String> {
|
||||
if !self.max_overhead_pct.is_finite() || self.max_overhead_pct < 0.0 {
|
||||
return Err(format!(
|
||||
"routing.routing_budget.max_overhead_pct: must be a non-negative number (percent, e.g. 20 for 20%), got {}",
|
||||
self.max_overhead_pct
|
||||
));
|
||||
}
|
||||
let cache_read_discount = self
|
||||
.cache_read_discount
|
||||
.unwrap_or(DEFAULT_CACHE_READ_DISCOUNT);
|
||||
if !(0.0..=1.0).contains(&cache_read_discount) {
|
||||
return Err(format!(
|
||||
"routing.routing_budget.cache_read_discount: must be between 0.0 and 1.0, got {cache_read_discount}"
|
||||
));
|
||||
}
|
||||
Ok(EffectiveRoutingBudget {
|
||||
max_overhead_pct: self.max_overhead_pct,
|
||||
replenish_on_rebind: self.replenish_on_rebind,
|
||||
cache_read_discount,
|
||||
record_counterfactual: self.record_counterfactual,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl EffectiveRoutingBudget {
|
||||
/// Resolve from an optional config block; `None` means the gate is off.
|
||||
pub fn from_config(config: Option<&RoutingBudget>) -> Result<Option<Self>, String> {
|
||||
config.map(RoutingBudget::resolve).transpose()
|
||||
}
|
||||
}
|
||||
|
||||
/// Fully-resolved, instance-wide prompt-caching settings.
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub struct EffectivePromptCaching {
|
||||
pub enabled: bool,
|
||||
pub session_affinity: bool,
|
||||
pub inject_cache_control: bool,
|
||||
pub min_prefix_tokens: u32,
|
||||
/// Pin TTL override; `None` uses `routing.session_ttl_seconds`.
|
||||
pub session_ttl_seconds: Option<u64>,
|
||||
}
|
||||
|
||||
pub const DEFAULT_MIN_PREFIX_TOKENS: u32 = 1024;
|
||||
|
||||
impl Default for EffectivePromptCaching {
|
||||
fn default() -> Self {
|
||||
EffectivePromptCaching {
|
||||
enabled: false,
|
||||
session_affinity: false,
|
||||
inject_cache_control: false,
|
||||
min_prefix_tokens: DEFAULT_MIN_PREFIX_TOKENS,
|
||||
session_ttl_seconds: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PromptCaching {
|
||||
/// Resolve the instance-wide effective settings. When caching is disabled every
|
||||
/// sub-feature is off, regardless of the individual knobs.
|
||||
pub fn resolve(&self) -> Result<EffectivePromptCaching, String> {
|
||||
if !self.enabled {
|
||||
return Ok(EffectivePromptCaching::default());
|
||||
}
|
||||
Ok(EffectivePromptCaching {
|
||||
enabled: true,
|
||||
session_affinity: self.session_affinity.unwrap_or(true),
|
||||
inject_cache_control: self.inject_cache_control.unwrap_or(true),
|
||||
min_prefix_tokens: self.min_prefix_tokens.unwrap_or(DEFAULT_MIN_PREFIX_TOKENS),
|
||||
session_ttl_seconds: self.session_ttl_seconds,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl EffectivePromptCaching {
|
||||
/// Resolve from an optional config block; `None` means caching is off.
|
||||
pub fn from_config(config: Option<&PromptCaching>) -> Result<Self, String> {
|
||||
config
|
||||
.map(PromptCaching::resolve)
|
||||
.transpose()
|
||||
.map(Option::unwrap_or_default)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct Tracing {
|
||||
pub sampling_rate: Option<f64>,
|
||||
|
|
@ -699,7 +882,10 @@ mod test {
|
|||
use pretty_assertions::assert_eq;
|
||||
use std::fs;
|
||||
|
||||
use super::{IntoModels, LlmProvider, LlmProviderType};
|
||||
use super::{
|
||||
EffectivePromptCaching, EffectiveRoutingBudget, IntoModels, LlmProvider, LlmProviderType,
|
||||
PromptCaching, RoutingBudget, DEFAULT_CACHE_READ_DISCOUNT, DEFAULT_MIN_PREFIX_TOKENS,
|
||||
};
|
||||
use crate::api::open_ai::ToolType;
|
||||
|
||||
#[test]
|
||||
|
|
@ -901,6 +1087,128 @@ disable_signals: false
|
|||
assert_eq!(overrides.disable_signals, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_prompt_caching_disabled_by_default() {
|
||||
// Absent config → everything off.
|
||||
let effective = EffectivePromptCaching::from_config(None).unwrap();
|
||||
assert!(!effective.enabled);
|
||||
assert!(!effective.session_affinity);
|
||||
assert!(!effective.inject_cache_control);
|
||||
|
||||
// Present but not enabled → still off.
|
||||
let cfg: PromptCaching = serde_yaml::from_str("enabled: false").unwrap();
|
||||
let effective = cfg.resolve().unwrap();
|
||||
assert!(!effective.enabled);
|
||||
assert!(!effective.session_affinity);
|
||||
assert!(!effective.inject_cache_control);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_prompt_caching_enabled_defaults() {
|
||||
// A bare `enabled: true` turns everything on with sensible defaults.
|
||||
let cfg: PromptCaching = serde_yaml::from_str("enabled: true").unwrap();
|
||||
let effective = cfg.resolve().unwrap();
|
||||
assert!(effective.enabled);
|
||||
assert!(effective.session_affinity);
|
||||
assert!(effective.inject_cache_control);
|
||||
assert_eq!(effective.min_prefix_tokens, DEFAULT_MIN_PREFIX_TOKENS);
|
||||
assert_eq!(effective.session_ttl_seconds, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_prompt_caching_optional_knobs() {
|
||||
let yaml = r#"
|
||||
enabled: true
|
||||
session_affinity: false
|
||||
inject_cache_control: false
|
||||
min_prefix_tokens: 2048
|
||||
session_ttl_seconds: 3600
|
||||
"#;
|
||||
let cfg: PromptCaching = serde_yaml::from_str(yaml).unwrap();
|
||||
let effective = cfg.resolve().unwrap();
|
||||
assert!(effective.enabled);
|
||||
assert!(!effective.session_affinity);
|
||||
assert!(!effective.inject_cache_control);
|
||||
assert_eq!(effective.min_prefix_tokens, 2048);
|
||||
assert_eq!(effective.session_ttl_seconds, Some(3600));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_prompt_caching_knobs_ignored_when_disabled() {
|
||||
// Knobs only take effect while caching is enabled.
|
||||
let yaml = r#"
|
||||
enabled: false
|
||||
session_affinity: true
|
||||
inject_cache_control: true
|
||||
"#;
|
||||
let cfg: PromptCaching = serde_yaml::from_str(yaml).unwrap();
|
||||
let effective = cfg.resolve().unwrap();
|
||||
assert!(!effective.enabled);
|
||||
assert!(!effective.session_affinity);
|
||||
assert!(!effective.inject_cache_control);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_routing_budget_parses() {
|
||||
let yaml = r#"
|
||||
max_overhead_pct: 20
|
||||
"#;
|
||||
let cfg: RoutingBudget = serde_yaml::from_str(yaml).unwrap();
|
||||
let budget = cfg.resolve().unwrap();
|
||||
assert_eq!(budget.max_overhead_pct, 20.0);
|
||||
// Replenish defaults on.
|
||||
assert!(budget.replenish_on_rebind);
|
||||
assert_eq!(budget.cache_read_discount, DEFAULT_CACHE_READ_DISCOUNT);
|
||||
// Counterfactual recording is opt-in; off unless requested.
|
||||
assert!(!budget.record_counterfactual);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_routing_budget_record_counterfactual_parses() {
|
||||
let yaml = r#"
|
||||
max_overhead_pct: 20
|
||||
record_counterfactual: true
|
||||
"#;
|
||||
let cfg: RoutingBudget = serde_yaml::from_str(yaml).unwrap();
|
||||
let budget = cfg.resolve().unwrap();
|
||||
assert!(budget.record_counterfactual);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_routing_budget_flags_parse() {
|
||||
let yaml = r#"
|
||||
max_overhead_pct: 15
|
||||
replenish_on_rebind: false
|
||||
cache_read_discount: 0.25
|
||||
"#;
|
||||
let cfg: RoutingBudget = serde_yaml::from_str(yaml).unwrap();
|
||||
let budget = cfg.resolve().unwrap();
|
||||
assert_eq!(budget.max_overhead_pct, 15.0);
|
||||
assert!(!budget.replenish_on_rebind);
|
||||
assert_eq!(budget.cache_read_discount, 0.25);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_routing_budget_absent_is_off() {
|
||||
// No block configured → gate is off.
|
||||
assert!(EffectiveRoutingBudget::from_config(None).unwrap().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_routing_budget_invalid_values_rejected() {
|
||||
let negative: RoutingBudget = serde_yaml::from_str("max_overhead_pct: -1.0").unwrap();
|
||||
assert!(negative.resolve().is_err());
|
||||
|
||||
let bad_discount: RoutingBudget = serde_yaml::from_str(
|
||||
r#"
|
||||
max_overhead_pct: 20
|
||||
cache_read_discount: 1.5
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(bad_discount.resolve().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tracing_posthog_exporter_deserialize() {
|
||||
let yaml = r#"
|
||||
|
|
|
|||
|
|
@ -23,6 +23,12 @@ pub const X_ARCH_FC_MODEL_RESPONSE: &str = "x-arch-fc-model-response";
|
|||
pub const ARCH_FC_MODEL_NAME: &str = "Arch-Function";
|
||||
pub const REQUEST_ID_HEADER: &str = "x-request-id";
|
||||
pub const MODEL_AFFINITY_HEADER: &str = "x-model-affinity";
|
||||
/// Per-request prompt-caching control. `off` disables implicit session affinity and
|
||||
/// cache-control injection for that single request.
|
||||
pub const PLANO_CACHE_HEADER: &str = "x-plano-cache";
|
||||
/// Hash of the stable prompt prefix, forwarded upstream so self-hosted multi-replica
|
||||
/// backends can do KV-aware (consistent-hash) replica routing at the LB/Envoy layer.
|
||||
pub const PLANO_PREFIX_HASH_HEADER: &str = "x-plano-prefix-hash";
|
||||
pub const ENVOY_ORIGINAL_PATH_HEADER: &str = "x-envoy-original-path";
|
||||
pub const TRACE_PARENT_HEADER: &str = "traceparent";
|
||||
pub const ARCH_INTERNAL_CLUSTER_NAME: &str = "arch_internal";
|
||||
|
|
|
|||
|
|
@ -285,6 +285,7 @@ pub struct MessagesTool {
|
|||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub input_schema: Value,
|
||||
pub cache_control: Option<MessagesCacheControl>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
|
||||
|
|
@ -305,6 +306,155 @@ pub struct MessagesToolChoice {
|
|||
pub disable_parallel_tool_use: Option<bool>,
|
||||
}
|
||||
|
||||
/// Rough chars-per-token heuristic used to threshold-guard cache-marker injection.
|
||||
/// Injection below the provider's minimum cacheable prefix is a provider-side no-op,
|
||||
/// so a cheap estimate is sufficient here.
|
||||
const CHARS_PER_TOKEN: usize = 4;
|
||||
|
||||
fn block_cache_control(block: &MessagesContentBlock) -> Option<&MessagesCacheControl> {
|
||||
match block {
|
||||
MessagesContentBlock::Text { cache_control, .. }
|
||||
| MessagesContentBlock::Thinking { cache_control, .. }
|
||||
| MessagesContentBlock::ToolUse { cache_control, .. }
|
||||
| MessagesContentBlock::ToolResult { cache_control, .. } => cache_control.as_ref(),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set `cache_control: ephemeral` on the last block that supports it.
|
||||
/// Returns true if a marker was placed.
|
||||
fn mark_last_cacheable_block(blocks: &mut [MessagesContentBlock]) -> bool {
|
||||
for block in blocks.iter_mut().rev() {
|
||||
match block {
|
||||
MessagesContentBlock::Text { cache_control, .. }
|
||||
| MessagesContentBlock::Thinking { cache_control, .. }
|
||||
| MessagesContentBlock::ToolUse { cache_control, .. }
|
||||
| MessagesContentBlock::ToolResult { cache_control, .. } => {
|
||||
*cache_control = Some(MessagesCacheControl::Ephemeral);
|
||||
return true;
|
||||
}
|
||||
_ => continue,
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
impl MessagesRequest {
|
||||
/// True when the client already supplied any `cache_control` marker
|
||||
/// (on system blocks, tools, or message content blocks).
|
||||
pub fn has_cache_markers(&self) -> bool {
|
||||
let system_marked = match &self.system {
|
||||
Some(MessagesSystemPrompt::Blocks(blocks)) => {
|
||||
blocks.iter().any(|b| block_cache_control(b).is_some())
|
||||
}
|
||||
_ => false,
|
||||
};
|
||||
let tools_marked = self
|
||||
.tools
|
||||
.as_ref()
|
||||
.is_some_and(|tools| tools.iter().any(|t| t.cache_control.is_some()));
|
||||
let messages_marked = self.messages.iter().any(|m| match &m.content {
|
||||
MessagesMessageContent::Blocks(blocks) => {
|
||||
blocks.iter().any(|b| block_cache_control(b).is_some())
|
||||
}
|
||||
MessagesMessageContent::Single(_) => false,
|
||||
});
|
||||
system_marked || tools_marked || messages_marked
|
||||
}
|
||||
|
||||
/// Estimated token length of the stable prompt prefix (tools + system), which is
|
||||
/// what precedes conversation history in Anthropic's prompt ordering.
|
||||
pub fn estimated_prefix_tokens(&self) -> u32 {
|
||||
let mut chars = 0usize;
|
||||
if let Some(tools) = &self.tools {
|
||||
for tool in tools {
|
||||
chars += tool.name.len();
|
||||
chars += tool.description.as_deref().map_or(0, str::len);
|
||||
chars += tool.input_schema.to_string().len();
|
||||
}
|
||||
}
|
||||
match &self.system {
|
||||
Some(MessagesSystemPrompt::Single(text)) => chars += text.len(),
|
||||
Some(MessagesSystemPrompt::Blocks(blocks)) => {
|
||||
chars += blocks.extract_text().len();
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
(chars / CHARS_PER_TOKEN) as u32
|
||||
}
|
||||
|
||||
/// Auto-inject ephemeral cache breakpoints for providers that require explicit
|
||||
/// markers (Anthropic-shaped requests):
|
||||
///
|
||||
/// 1. at the end of the system prompt (covering the fully-stable tools + system
|
||||
/// prefix), falling back to the last tool when there is no system prompt, and
|
||||
/// 2. a rolling breakpoint on the last content block of the final message, so each
|
||||
/// turn's cache write becomes the next turn's cache read as history grows.
|
||||
///
|
||||
/// Idempotent: a request that already carries any client-supplied `cache_control`
|
||||
/// marker is left untouched. Threshold-guarded: no-op when the estimated stable
|
||||
/// prefix is below `min_prefix_tokens` (injection below the provider's minimum
|
||||
/// cacheable prefix is wasted bytes).
|
||||
///
|
||||
/// Returns true if any marker was injected.
|
||||
pub fn inject_cache_breakpoints(&mut self, min_prefix_tokens: u32) -> bool {
|
||||
if self.has_cache_markers() {
|
||||
return false;
|
||||
}
|
||||
if self.estimated_prefix_tokens() < min_prefix_tokens {
|
||||
return false;
|
||||
}
|
||||
|
||||
let mut injected = false;
|
||||
|
||||
// Breakpoint 1: end of the stable prefix.
|
||||
match self.system.take() {
|
||||
Some(MessagesSystemPrompt::Single(text)) => {
|
||||
self.system = Some(MessagesSystemPrompt::Blocks(vec![
|
||||
MessagesContentBlock::Text {
|
||||
text,
|
||||
cache_control: Some(MessagesCacheControl::Ephemeral),
|
||||
},
|
||||
]));
|
||||
injected = true;
|
||||
}
|
||||
Some(MessagesSystemPrompt::Blocks(mut blocks)) => {
|
||||
injected |= mark_last_cacheable_block(&mut blocks);
|
||||
self.system = Some(MessagesSystemPrompt::Blocks(blocks));
|
||||
}
|
||||
None => {
|
||||
if let Some(last_tool) = self.tools.as_mut().and_then(|t| t.last_mut()) {
|
||||
last_tool.cache_control = Some(MessagesCacheControl::Ephemeral);
|
||||
injected = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Breakpoint 2: rolling tail of conversation history.
|
||||
if let Some(last_msg) = self.messages.last_mut() {
|
||||
let content = std::mem::replace(
|
||||
&mut last_msg.content,
|
||||
MessagesMessageContent::Single(String::new()),
|
||||
);
|
||||
last_msg.content = match content {
|
||||
MessagesMessageContent::Single(text) => {
|
||||
injected = true;
|
||||
MessagesMessageContent::Blocks(vec![MessagesContentBlock::Text {
|
||||
text,
|
||||
cache_control: Some(MessagesCacheControl::Ephemeral),
|
||||
}])
|
||||
}
|
||||
MessagesMessageContent::Blocks(mut blocks) => {
|
||||
injected |= mark_last_cacheable_block(&mut blocks);
|
||||
MessagesMessageContent::Blocks(blocks)
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
injected
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum MessagesStopReason {
|
||||
|
|
@ -680,6 +830,134 @@ mod tests {
|
|||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
fn cache_test_request(system: Option<MessagesSystemPrompt>) -> MessagesRequest {
|
||||
MessagesRequest {
|
||||
model: "claude-sonnet-4".to_string(),
|
||||
messages: vec![
|
||||
MessagesMessage {
|
||||
role: MessagesRole::User,
|
||||
content: MessagesMessageContent::Single("first user turn".to_string()),
|
||||
},
|
||||
MessagesMessage {
|
||||
role: MessagesRole::Assistant,
|
||||
content: MessagesMessageContent::Single("assistant reply".to_string()),
|
||||
},
|
||||
MessagesMessage {
|
||||
role: MessagesRole::User,
|
||||
content: MessagesMessageContent::Single("second user turn".to_string()),
|
||||
},
|
||||
],
|
||||
max_tokens: 100,
|
||||
container: None,
|
||||
mcp_servers: None,
|
||||
system,
|
||||
metadata: None,
|
||||
service_tier: None,
|
||||
thinking: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
top_k: None,
|
||||
stream: None,
|
||||
stop_sequences: None,
|
||||
tools: None,
|
||||
tool_choice: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_inject_cache_breakpoints_marks_system_and_tail() {
|
||||
let long_system = "x".repeat(8192); // ~2048 estimated tokens
|
||||
let mut req = cache_test_request(Some(MessagesSystemPrompt::Single(long_system)));
|
||||
|
||||
assert!(req.inject_cache_breakpoints(1024));
|
||||
|
||||
// System converted to a marked block.
|
||||
match &req.system {
|
||||
Some(MessagesSystemPrompt::Blocks(blocks)) => {
|
||||
assert_eq!(blocks.len(), 1);
|
||||
assert!(matches!(
|
||||
&blocks[0],
|
||||
MessagesContentBlock::Text {
|
||||
cache_control: Some(MessagesCacheControl::Ephemeral),
|
||||
..
|
||||
}
|
||||
));
|
||||
}
|
||||
other => panic!("expected marked system blocks, got {:?}", other),
|
||||
}
|
||||
|
||||
// Rolling breakpoint on the last message.
|
||||
match &req.messages.last().unwrap().content {
|
||||
MessagesMessageContent::Blocks(blocks) => {
|
||||
assert!(matches!(
|
||||
&blocks[0],
|
||||
MessagesContentBlock::Text {
|
||||
cache_control: Some(MessagesCacheControl::Ephemeral),
|
||||
..
|
||||
}
|
||||
));
|
||||
}
|
||||
other => panic!("expected marked tail blocks, got {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_inject_cache_breakpoints_below_threshold_is_noop() {
|
||||
let mut req = cache_test_request(Some(MessagesSystemPrompt::Single(
|
||||
"short system".to_string(),
|
||||
)));
|
||||
assert!(!req.inject_cache_breakpoints(1024));
|
||||
assert!(matches!(req.system, Some(MessagesSystemPrompt::Single(_))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_inject_cache_breakpoints_respects_client_markers() {
|
||||
let long_system = "x".repeat(8192);
|
||||
let mut req = cache_test_request(Some(MessagesSystemPrompt::Blocks(vec![
|
||||
MessagesContentBlock::Text {
|
||||
text: long_system,
|
||||
cache_control: Some(MessagesCacheControl::Ephemeral),
|
||||
},
|
||||
])));
|
||||
// Client already placed a marker: injection must be a no-op.
|
||||
assert!(!req.inject_cache_breakpoints(1024));
|
||||
// Tail message untouched.
|
||||
assert!(matches!(
|
||||
req.messages.last().unwrap().content,
|
||||
MessagesMessageContent::Single(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_inject_cache_breakpoints_marks_last_tool_when_no_system() {
|
||||
let mut req = cache_test_request(None);
|
||||
req.tools = Some(vec![MessagesTool {
|
||||
name: "big_tool".to_string(),
|
||||
description: Some("d".repeat(8192)),
|
||||
input_schema: json!({"type": "object"}),
|
||||
cache_control: None,
|
||||
}]);
|
||||
|
||||
assert!(req.inject_cache_breakpoints(1024));
|
||||
assert_eq!(
|
||||
req.tools.as_ref().unwrap().last().unwrap().cache_control,
|
||||
Some(MessagesCacheControl::Ephemeral)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_cache_control_roundtrip() {
|
||||
// Client-supplied cache_control on tools must survive serde passthrough.
|
||||
let tool_json = json!({
|
||||
"name": "get_weather",
|
||||
"input_schema": {"type": "object"},
|
||||
"cache_control": {"type": "ephemeral"}
|
||||
});
|
||||
let tool: MessagesTool = serde_json::from_value(tool_json.clone()).unwrap();
|
||||
assert_eq!(tool.cache_control, Some(MessagesCacheControl::Ephemeral));
|
||||
assert_eq!(serde_json::to_value(&tool).unwrap(), tool_json);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_anthropic_required_fields() {
|
||||
// Create a JSON object with only required fields
|
||||
|
|
|
|||
|
|
@ -163,6 +163,112 @@ impl ChatCompletionsRequest {
|
|||
self.store = None;
|
||||
}
|
||||
}
|
||||
|
||||
/// True when any message content part already carries a `cache_control` marker
|
||||
/// (client-supplied or previously injected).
|
||||
pub fn has_cache_control_markers(&self) -> bool {
|
||||
self.messages.iter().any(|m| match &m.content {
|
||||
Some(MessageContent::Parts(parts)) => parts.iter().any(|p| {
|
||||
matches!(
|
||||
p,
|
||||
ContentPart::Text {
|
||||
cache_control: Some(_),
|
||||
..
|
||||
}
|
||||
)
|
||||
}),
|
||||
_ => false,
|
||||
})
|
||||
}
|
||||
|
||||
/// Estimated token length of the stable prompt prefix — the system/developer
|
||||
/// message(s) plus tool definitions — using a chars/4 heuristic. Precision is
|
||||
/// not required; this only gates whether marking the prefix is worthwhile.
|
||||
pub fn estimated_cache_prefix_tokens(&self) -> u32 {
|
||||
const CHARS_PER_TOKEN: usize = 4;
|
||||
let mut chars = 0usize;
|
||||
for m in &self.messages {
|
||||
if matches!(m.role, Role::System | Role::Developer) {
|
||||
if let Some(content) = &m.content {
|
||||
chars += content.extract_text().len();
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(tools) = &self.tools {
|
||||
for t in tools {
|
||||
chars += t.function.name.len();
|
||||
chars += t.function.description.as_deref().map_or(0, str::len);
|
||||
chars += t.function.parameters.to_string().len();
|
||||
}
|
||||
}
|
||||
(chars / CHARS_PER_TOKEN) as u32
|
||||
}
|
||||
|
||||
/// Auto-inject a single ephemeral `cache_control` breakpoint at the end of the
|
||||
/// stable prompt prefix, for OpenAI-compatible gateways that proxy Anthropic-family
|
||||
/// models (DigitalOcean, OpenRouter). The marker is attached to the last text
|
||||
/// content part of the system/developer message (falling back to the first message
|
||||
/// when there is no system prompt), normalizing a plain-string content into a
|
||||
/// one-element content-part array as needed.
|
||||
///
|
||||
/// Idempotent: a request that already carries any `cache_control` marker is left
|
||||
/// untouched. Threshold-guarded: no-op when the estimated stable prefix is below
|
||||
/// `min_prefix_tokens`. Returns true if a marker was injected.
|
||||
pub fn inject_cache_control(&mut self, ttl: Option<String>, min_prefix_tokens: u32) -> bool {
|
||||
if self.has_cache_control_markers() {
|
||||
return false;
|
||||
}
|
||||
if self.estimated_cache_prefix_tokens() < min_prefix_tokens {
|
||||
return false;
|
||||
}
|
||||
|
||||
let target_idx = self
|
||||
.messages
|
||||
.iter()
|
||||
.position(|m| matches!(m.role, Role::System | Role::Developer))
|
||||
.or(if self.messages.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(0)
|
||||
});
|
||||
|
||||
let Some(idx) = target_idx else {
|
||||
return false;
|
||||
};
|
||||
let marker = CacheControl {
|
||||
kind: "ephemeral".to_string(),
|
||||
ttl,
|
||||
};
|
||||
attach_cache_control(&mut self.messages[idx], marker)
|
||||
}
|
||||
}
|
||||
|
||||
/// Attach a `cache_control` marker to the last text content part of `message`,
|
||||
/// normalizing string content into a one-element parts array. Returns false when
|
||||
/// there is no text part to mark (e.g. an image-only message).
|
||||
fn attach_cache_control(message: &mut Message, marker: CacheControl) -> bool {
|
||||
let mut parts = match message.content.take() {
|
||||
Some(MessageContent::Text(text)) => vec![ContentPart::Text {
|
||||
text,
|
||||
cache_control: None,
|
||||
}],
|
||||
Some(MessageContent::Parts(parts)) => parts,
|
||||
None => return false,
|
||||
};
|
||||
|
||||
let marked = if let Some(ContentPart::Text { cache_control, .. }) = parts
|
||||
.iter_mut()
|
||||
.rev()
|
||||
.find(|p| matches!(p, ContentPart::Text { .. }))
|
||||
{
|
||||
*cache_control = Some(marker);
|
||||
true
|
||||
} else {
|
||||
false
|
||||
};
|
||||
|
||||
message.content = Some(MessageContent::Parts(parts));
|
||||
marked
|
||||
}
|
||||
|
||||
/// True when the upstream model id is Moonshot's Kimi Code endpoint model.
|
||||
|
|
@ -277,7 +383,7 @@ impl ExtractText for Vec<ContentPart> {
|
|||
fn extract_text(&self) -> String {
|
||||
self.iter()
|
||||
.filter_map(|part| match part {
|
||||
ContentPart::Text { text } => Some(text.as_str()),
|
||||
ContentPart::Text { text, .. } => Some(text.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
|
|
@ -296,11 +402,32 @@ impl Display for MessageContent {
|
|||
#[serde(tag = "type")]
|
||||
pub enum ContentPart {
|
||||
#[serde(rename = "text")]
|
||||
Text { text: String },
|
||||
Text {
|
||||
text: String,
|
||||
/// Prompt-cache breakpoint for OpenAI-compatible gateways that proxy
|
||||
/// Anthropic-family models (e.g. DigitalOcean, OpenRouter). Round-trips so
|
||||
/// client-supplied markers survive deserialization and are respected by the
|
||||
/// idempotent injector.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
cache_control: Option<CacheControl>,
|
||||
},
|
||||
#[serde(rename = "image_url")]
|
||||
ImageUrl { image_url: ImageUrl },
|
||||
}
|
||||
|
||||
/// Prompt-cache control marker carried on an OpenAI content part. Mirrors the
|
||||
/// Anthropic `cache_control` object as accepted by OpenAI-compatible gateways that
|
||||
/// front Anthropic models.
|
||||
#[skip_serializing_none]
|
||||
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CacheControl {
|
||||
#[serde(rename = "type")]
|
||||
pub kind: String, // "ephemeral"
|
||||
/// Cache lifetime hint: "5m" | "1h" (DigitalOcean / OpenRouter). Omitted for
|
||||
/// plain Anthropic ephemeral caching (defaults to 5 minutes upstream).
|
||||
pub ttl: Option<String>,
|
||||
}
|
||||
|
||||
/// Image URL configuration for vision capabilities
|
||||
#[skip_serializing_none]
|
||||
#[derive(Serialize, Deserialize, Debug, Clone)]
|
||||
|
|
@ -445,11 +572,35 @@ pub struct Usage {
|
|||
pub total_tokens: u32,
|
||||
pub prompt_tokens_details: Option<PromptTokensDetails>,
|
||||
pub completion_tokens_details: Option<CompletionTokensDetails>,
|
||||
/// Anthropic-style cache-read counter emitted by OpenAI-compatible gateways that
|
||||
/// front Anthropic models (e.g. DigitalOcean), which do *not* populate
|
||||
/// `prompt_tokens_details.cached_tokens`. Captured here so it can be surfaced to
|
||||
/// OpenAI clients via [`Usage::normalize_cache_tokens`].
|
||||
pub cache_read_input_tokens: Option<u32>,
|
||||
/// Anthropic-style cache-write counter (see `cache_read_input_tokens`).
|
||||
pub cache_creation_input_tokens: Option<u32>,
|
||||
}
|
||||
|
||||
impl Usage {
|
||||
/// Fold gateway-specific Anthropic-style `cache_read_input_tokens` into the
|
||||
/// OpenAI-standard `prompt_tokens_details.cached_tokens` so OpenAI-compatible
|
||||
/// clients (and downstream cost accounting) observe the cache hit. No-op when the
|
||||
/// standard field is already populated or no cache-read counter is present.
|
||||
pub fn normalize_cache_tokens(&mut self) {
|
||||
if let Some(read) = self.cache_read_input_tokens {
|
||||
let details = self
|
||||
.prompt_tokens_details
|
||||
.get_or_insert_with(Default::default);
|
||||
if details.cached_tokens.is_none() {
|
||||
details.cached_tokens = Some(read);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Detailed breakdown of prompt tokens
|
||||
#[skip_serializing_none]
|
||||
#[derive(Serialize, Deserialize, Debug, Clone)]
|
||||
#[derive(Serialize, Deserialize, Debug, Clone, Default)]
|
||||
pub struct PromptTokensDetails {
|
||||
pub cached_tokens: Option<u32>,
|
||||
pub audio_tokens: Option<u32>,
|
||||
|
|
@ -632,7 +783,15 @@ impl TokenUsage for Usage {
|
|||
fn cached_input_tokens(&self) -> Option<usize> {
|
||||
self.prompt_tokens_details
|
||||
.as_ref()
|
||||
.and_then(|d| d.cached_tokens.map(|t| t as usize))
|
||||
.and_then(|d| d.cached_tokens)
|
||||
// Gateways fronting Anthropic (e.g. DigitalOcean) report cache reads here
|
||||
// instead of prompt_tokens_details.cached_tokens.
|
||||
.or(self.cache_read_input_tokens)
|
||||
.map(|t| t as usize)
|
||||
}
|
||||
|
||||
fn cache_creation_tokens(&self) -> Option<usize> {
|
||||
self.cache_creation_input_tokens.map(|t| t as usize)
|
||||
}
|
||||
|
||||
fn reasoning_tokens(&self) -> Option<usize> {
|
||||
|
|
@ -666,7 +825,7 @@ impl ProviderRequest for ChatCompletionsRequest {
|
|||
MessageContent::Parts(parts) => parts
|
||||
.iter()
|
||||
.map(|part| match part {
|
||||
ContentPart::Text { text } => text.clone(),
|
||||
ContentPart::Text { text, .. } => text.clone(),
|
||||
ContentPart::ImageUrl { .. } => "[Image]".to_string(),
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
|
|
@ -796,6 +955,137 @@ mod tests {
|
|||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
/// Build a request with a long system prompt (well over the token threshold) plus
|
||||
/// a short user turn.
|
||||
fn request_with_system(system: &str, user: &str) -> ChatCompletionsRequest {
|
||||
ChatCompletionsRequest {
|
||||
model: "anthropic/claude-3.5-sonnet".to_string(),
|
||||
messages: vec![
|
||||
Message {
|
||||
role: Role::System,
|
||||
content: Some(MessageContent::Text(system.to_string())),
|
||||
name: None,
|
||||
tool_calls: None,
|
||||
tool_call_id: None,
|
||||
},
|
||||
Message {
|
||||
role: Role::User,
|
||||
content: Some(MessageContent::Text(user.to_string())),
|
||||
name: None,
|
||||
tool_calls: None,
|
||||
tool_call_id: None,
|
||||
},
|
||||
],
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn system_cache_control(req: &ChatCompletionsRequest) -> Option<&CacheControl> {
|
||||
match &req.messages[0].content {
|
||||
Some(MessageContent::Parts(parts)) => parts.iter().find_map(|p| match p {
|
||||
ContentPart::Text { cache_control, .. } => cache_control.as_ref(),
|
||||
_ => None,
|
||||
}),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inject_cache_control_normalizes_string_content_and_marks_prefix() {
|
||||
let mut req = request_with_system(&"x".repeat(8000), "hi");
|
||||
let injected = req.inject_cache_control(None, 1024);
|
||||
assert!(injected);
|
||||
// The plain-string system content was normalized into a one-element parts array.
|
||||
let marker = system_cache_control(&req).expect("system content part should be marked");
|
||||
assert_eq!(marker.kind, "ephemeral");
|
||||
assert_eq!(marker.ttl, None);
|
||||
// The user message is untouched.
|
||||
assert!(matches!(
|
||||
req.messages[1].content,
|
||||
Some(MessageContent::Text(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inject_cache_control_is_idempotent() {
|
||||
let mut req = request_with_system(&"x".repeat(8000), "hi");
|
||||
assert!(req.inject_cache_control(Some("1h".to_string()), 1024));
|
||||
// A second pass must be a no-op because a marker already exists.
|
||||
assert!(!req.inject_cache_control(None, 1024));
|
||||
assert_eq!(
|
||||
system_cache_control(&req).and_then(|c| c.ttl.clone()),
|
||||
Some("1h".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inject_cache_control_respects_min_prefix_threshold() {
|
||||
// A tiny prefix (below the threshold) is not worth marking.
|
||||
let mut req = request_with_system("short system", "hi");
|
||||
assert!(!req.inject_cache_control(None, 1024));
|
||||
assert!(system_cache_control(&req).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_supplied_cache_control_round_trips_and_blocks_injection() {
|
||||
// A client that already set cache_control must survive deserialization and be
|
||||
// respected (no double injection).
|
||||
let raw = json!({
|
||||
"model": "anthropic/claude-3.5-sonnet",
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{"type": "text", "text": "big stable prefix",
|
||||
"cache_control": {"type": "ephemeral", "ttl": "5m"}}
|
||||
]
|
||||
},
|
||||
{"role": "user", "content": "hi"}
|
||||
]
|
||||
});
|
||||
let mut req: ChatCompletionsRequest = serde_json::from_value(raw).unwrap();
|
||||
assert!(req.has_cache_control_markers());
|
||||
// Injection is a no-op, and re-serialization preserves the client's marker.
|
||||
assert!(!req.inject_cache_control(None, 0));
|
||||
let reserialized = serde_json::to_value(&req).unwrap();
|
||||
let marker = &reserialized["messages"][0]["content"][0]["cache_control"];
|
||||
assert_eq!(marker["type"], "ephemeral");
|
||||
assert_eq!(marker["ttl"], "5m");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_control_omitted_when_absent() {
|
||||
// A plain text part must not emit a null cache_control field.
|
||||
let part = ContentPart::Text {
|
||||
text: "hello".to_string(),
|
||||
cache_control: None,
|
||||
};
|
||||
let v = serde_json::to_value(&part).unwrap();
|
||||
assert!(v.get("cache_control").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_normalizes_gateway_cache_read_into_cached_tokens() {
|
||||
// DigitalOcean-style usage: Anthropic cache_read with no prompt_tokens_details.
|
||||
let raw = json!({
|
||||
"prompt_tokens": 100,
|
||||
"completion_tokens": 20,
|
||||
"total_tokens": 120,
|
||||
"cache_read_input_tokens": 80,
|
||||
"cache_creation_input_tokens": 12
|
||||
});
|
||||
let mut usage: Usage = serde_json::from_value(raw).unwrap();
|
||||
// TokenUsage falls back to the gateway field.
|
||||
assert_eq!(usage.cached_input_tokens(), Some(80));
|
||||
assert_eq!(usage.cache_creation_tokens(), Some(12));
|
||||
// Normalization surfaces it as the OpenAI-standard cached_tokens.
|
||||
usage.normalize_cache_tokens();
|
||||
assert_eq!(
|
||||
usage.prompt_tokens_details.and_then(|d| d.cached_tokens),
|
||||
Some(80)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_required_fields() {
|
||||
// Create a JSON object with only required fields
|
||||
|
|
@ -995,7 +1285,7 @@ mod tests {
|
|||
assert_eq!(content_parts.len(), 2);
|
||||
|
||||
// Validate text content part
|
||||
if let ContentPart::Text { text } = &content_parts[0] {
|
||||
if let ContentPart::Text { text, .. } = &content_parts[0] {
|
||||
assert_eq!(text, "What can you see in this image and what's the weather like in the location shown?");
|
||||
} else {
|
||||
panic!("Expected text content part");
|
||||
|
|
|
|||
|
|
@ -6,10 +6,14 @@ pub mod clients;
|
|||
pub mod providers;
|
||||
pub mod transforms;
|
||||
// Re-export important types and traits
|
||||
pub use apis::openai::CacheControl;
|
||||
pub use apis::streaming_shapes::amazon_bedrock_binary_frame::BedrockBinaryFrameDecoder;
|
||||
pub use apis::streaming_shapes::sse::{SseEvent, SseStreamIter};
|
||||
pub use aws_smithy_eventstream::frame::DecodedFrame;
|
||||
pub use providers::id::ProviderId;
|
||||
pub use providers::id::{
|
||||
cache_marker_strategy, provider_cache_capability, CacheMarkerStrategy, ProviderCacheCapability,
|
||||
ProviderId,
|
||||
};
|
||||
pub use providers::request::{ProviderRequest, ProviderRequestError, ProviderRequestType};
|
||||
pub use providers::response::{
|
||||
ProviderResponse, ProviderResponseError, ProviderResponseType, TokenUsage,
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ use serde::Deserialize;
|
|||
use std::collections::HashMap;
|
||||
use std::fmt::Display;
|
||||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
|
||||
static PROVIDER_MODELS_YAML: &str = include_str!(concat!(
|
||||
env!("CARGO_MANIFEST_DIR"),
|
||||
|
|
@ -94,6 +95,192 @@ impl TryFrom<&str> for ProviderId {
|
|||
}
|
||||
}
|
||||
|
||||
/// How Plano should mark a request for prompt caching, resolved from the *combination*
|
||||
/// of gateway provider, the underlying model family, and the upstream API shape — not
|
||||
/// the gateway alone. This is what lets DigitalOcean-Anthropic and OpenRouter-Anthropic
|
||||
/// (both OpenAI-compatible chat completions fronting Anthropic models) cache through one
|
||||
/// path, while OpenAI-family models stay correctly automatic and unimplemented backends
|
||||
/// (Bedrock) are an honest `None` rather than a silent no-op.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum CacheMarkerStrategy {
|
||||
/// No known prompt-caching support for this combination — do nothing.
|
||||
None,
|
||||
/// Provider caches stable prefixes automatically (OpenAI-family anywhere); no
|
||||
/// request markers are needed. Plano only keeps the prefix byte-stable and pinned.
|
||||
Automatic,
|
||||
/// OpenAI-compatible chat completions fronting Anthropic-family models
|
||||
/// (DigitalOcean, OpenRouter): attach `cache_control` to content parts.
|
||||
OpenAiContentPartCacheControl {
|
||||
/// Minimum cacheable prefix length in tokens; injecting below this is a no-op.
|
||||
min_prefix_tokens: u32,
|
||||
/// Optional cache lifetime hint ("5m" | "1h").
|
||||
ttl: Option<String>,
|
||||
},
|
||||
/// Native Anthropic Messages API (native `anthropic/*`, Vercel-Anthropic): inject
|
||||
/// ephemeral breakpoints on the Anthropic-shaped request.
|
||||
AnthropicMessagesBreakpoints {
|
||||
/// Maximum number of cache breakpoints the provider accepts per request.
|
||||
max_breakpoints: u8,
|
||||
/// Minimum cacheable prefix length in tokens; injecting below this is a no-op.
|
||||
min_prefix_tokens: u32,
|
||||
},
|
||||
// BedrockCachePoint { .. } // left as explicit `None` until implemented.
|
||||
}
|
||||
|
||||
/// Coarse model family, inferred from the model id. Works across gateway naming
|
||||
/// conventions (DigitalOcean's `anthropic-claude-…` dash form and OpenRouter's
|
||||
/// `anthropic/claude-…` slash form) via substring matching.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum ModelFamily {
|
||||
Anthropic,
|
||||
OpenAI,
|
||||
Other,
|
||||
}
|
||||
|
||||
fn model_family(model_name: &str) -> ModelFamily {
|
||||
let m = model_name.to_ascii_lowercase();
|
||||
if m.contains("claude") || m.contains("anthropic") {
|
||||
ModelFamily::Anthropic
|
||||
} else if m.contains("gpt") || m.contains("openai") || m.contains("chatgpt") {
|
||||
ModelFamily::OpenAI
|
||||
} else {
|
||||
ModelFamily::Other
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether a gateway accepts Anthropic-style `cache_control` on OpenAI content parts
|
||||
/// over its chat-completions endpoint.
|
||||
fn accepts_openai_content_part_cache_control(provider: ProviderId) -> bool {
|
||||
matches!(provider, ProviderId::DigitalOcean | ProviderId::OpenRouter)
|
||||
}
|
||||
|
||||
/// Whether a gateway/model relies on automatic prefix caching (no markers required)
|
||||
/// over an OpenAI-compatible surface.
|
||||
fn is_automatic_cache_provider(provider: ProviderId) -> bool {
|
||||
matches!(
|
||||
provider,
|
||||
ProviderId::OpenAI
|
||||
| ProviderId::AzureOpenAI
|
||||
| ProviderId::ChatGPT
|
||||
| ProviderId::Groq
|
||||
| ProviderId::Deepseek
|
||||
| ProviderId::Gemini
|
||||
| ProviderId::Moonshotai
|
||||
| ProviderId::XAI
|
||||
| ProviderId::DigitalOcean
|
||||
| ProviderId::OpenRouter
|
||||
)
|
||||
}
|
||||
|
||||
/// Resolve the cache-marking strategy for a `(gateway provider × underlying model ×
|
||||
/// upstream API)` combination.
|
||||
///
|
||||
/// - `model_name` is the id *after* the gateway prefix (e.g. `anthropic-claude-3-5-sonnet`
|
||||
/// for DigitalOcean, `anthropic/claude-3.5-sonnet` for OpenRouter).
|
||||
pub fn cache_marker_strategy(
|
||||
provider: ProviderId,
|
||||
model_name: &str,
|
||||
upstream_api: &SupportedUpstreamAPIs,
|
||||
) -> CacheMarkerStrategy {
|
||||
// Anthropic minimum cacheable prefix is ~1024 tokens (2048 for Haiku-class);
|
||||
// callers may raise this via config.
|
||||
const ANTHROPIC_MIN_PREFIX_TOKENS: u32 = 1024;
|
||||
|
||||
match upstream_api {
|
||||
// Native Anthropic Messages API — inject ephemeral breakpoints.
|
||||
SupportedUpstreamAPIs::AnthropicMessagesAPI(_) => {
|
||||
CacheMarkerStrategy::AnthropicMessagesBreakpoints {
|
||||
max_breakpoints: 4,
|
||||
min_prefix_tokens: ANTHROPIC_MIN_PREFIX_TOKENS,
|
||||
}
|
||||
}
|
||||
// OpenAI-compatible chat completions — strategy depends on the model family.
|
||||
SupportedUpstreamAPIs::OpenAIChatCompletions(_) => match model_family(model_name) {
|
||||
ModelFamily::Anthropic if accepts_openai_content_part_cache_control(provider) => {
|
||||
CacheMarkerStrategy::OpenAiContentPartCacheControl {
|
||||
min_prefix_tokens: ANTHROPIC_MIN_PREFIX_TOKENS,
|
||||
ttl: None,
|
||||
}
|
||||
}
|
||||
// Anthropic-family behind a gateway that doesn't accept content-part
|
||||
// cache_control over chat completions: no honest way to mark it.
|
||||
ModelFamily::Anthropic => CacheMarkerStrategy::None,
|
||||
ModelFamily::OpenAI => CacheMarkerStrategy::Automatic,
|
||||
ModelFamily::Other if is_automatic_cache_provider(provider) => {
|
||||
CacheMarkerStrategy::Automatic
|
||||
}
|
||||
ModelFamily::Other => CacheMarkerStrategy::None,
|
||||
},
|
||||
// OpenAI Responses API — OpenAI-family automatic prefix caching.
|
||||
SupportedUpstreamAPIs::OpenAIResponsesAPI(_) => CacheMarkerStrategy::Automatic,
|
||||
// Bedrock cache points not yet implemented — honest None instead of a
|
||||
// silent no-op.
|
||||
SupportedUpstreamAPIs::AmazonBedrockConverse(_)
|
||||
| SupportedUpstreamAPIs::AmazonBedrockConverseStream(_) => CacheMarkerStrategy::None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Provider prompt-cache retention behavior, used to decide whether a session's
|
||||
/// upstream cache is still plausibly warm from the time since it was last used.
|
||||
///
|
||||
/// This is deliberately time/behavior only — it says nothing about *how* to mark a
|
||||
/// request for caching (that's [`CacheMarkerStrategy`]). Warmth is a function of the
|
||||
/// idle gap vs the provider's cache window, so the session router can reason about
|
||||
/// stickiness without ever seeing a provider response.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct ProviderCacheCapability {
|
||||
/// Sliding idle window: the cache stays warm as long as it is touched at least
|
||||
/// this often. Anthropic's default ephemeral cache is 5 minutes.
|
||||
pub idle_ttl: Duration,
|
||||
/// Absolute ceiling on how long a cache entry can live regardless of activity.
|
||||
/// Conservative default of 1h matches Anthropic's extended (1h) tier ceiling.
|
||||
pub hard_ttl: Duration,
|
||||
/// Whether the provider (as configured) actually retains caches out to the
|
||||
/// extended window. Off by default — extended retention is opt-in per provider.
|
||||
pub extended_retention: bool,
|
||||
/// The extended idle window when `extended_retention` is enabled (e.g. 1h).
|
||||
pub extended_ttl: Duration,
|
||||
}
|
||||
|
||||
impl Default for ProviderCacheCapability {
|
||||
fn default() -> Self {
|
||||
// Conservative, provider-agnostic defaults: a 5-minute sliding window capped
|
||||
// at 1 hour, no extended retention. Anything unknown is treated as short-lived
|
||||
// so the router doesn't over-stick to a cache that has likely gone cold.
|
||||
ProviderCacheCapability {
|
||||
idle_ttl: Duration::from_secs(5 * 60),
|
||||
hard_ttl: Duration::from_secs(60 * 60),
|
||||
extended_retention: false,
|
||||
extended_ttl: Duration::from_secs(60 * 60),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve the prompt-cache retention window for a gateway provider. Data-driven so
|
||||
/// tuning a provider's window needs no code changes at the call sites — only this
|
||||
/// table. Unknown providers fall back to the conservative [`ProviderCacheCapability::default`].
|
||||
pub fn provider_cache_capability(provider: ProviderId) -> ProviderCacheCapability {
|
||||
match provider {
|
||||
// Anthropic-family caches (native or fronted): 5-minute sliding default,
|
||||
// 1-hour hard ceiling. Extended (1h) retention is opt-in and left off here.
|
||||
ProviderId::Anthropic
|
||||
| ProviderId::DigitalOcean
|
||||
| ProviderId::OpenRouter
|
||||
| ProviderId::Vercel => ProviderCacheCapability::default(),
|
||||
// OpenAI-family automatic prefix caching also lives on the order of minutes;
|
||||
// the conservative default holds.
|
||||
ProviderId::OpenAI
|
||||
| ProviderId::AzureOpenAI
|
||||
| ProviderId::ChatGPT
|
||||
| ProviderId::Groq
|
||||
| ProviderId::Deepseek
|
||||
| ProviderId::Gemini
|
||||
| ProviderId::Moonshotai
|
||||
| ProviderId::XAI => ProviderCacheCapability::default(),
|
||||
_ => ProviderCacheCapability::default(),
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderId {
|
||||
/// Get all available models for this provider
|
||||
/// Returns model names without the provider prefix (e.g., "gpt-4" not "openai/gpt-4")
|
||||
|
|
@ -297,6 +484,100 @@ impl Display for ProviderId {
|
|||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::apis::{AnthropicApi, OpenAIApi};
|
||||
|
||||
fn chat_completions() -> SupportedUpstreamAPIs {
|
||||
SupportedUpstreamAPIs::OpenAIChatCompletions(OpenAIApi::ChatCompletions)
|
||||
}
|
||||
|
||||
fn anthropic_messages() -> SupportedUpstreamAPIs {
|
||||
SupportedUpstreamAPIs::AnthropicMessagesAPI(AnthropicApi::Messages)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn digitalocean_anthropic_uses_openai_content_part_markers() {
|
||||
// DO fronts Anthropic over an OpenAI-compatible surface (dash-form model id).
|
||||
let strategy = cache_marker_strategy(
|
||||
ProviderId::DigitalOcean,
|
||||
"anthropic-claude-3-5-sonnet",
|
||||
&chat_completions(),
|
||||
);
|
||||
assert!(matches!(
|
||||
strategy,
|
||||
CacheMarkerStrategy::OpenAiContentPartCacheControl { .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openrouter_anthropic_uses_openai_content_part_markers() {
|
||||
// OpenRouter uses slash-form model ids after the gateway prefix.
|
||||
let strategy = cache_marker_strategy(
|
||||
ProviderId::OpenRouter,
|
||||
"anthropic/claude-3.5-sonnet",
|
||||
&chat_completions(),
|
||||
);
|
||||
assert!(matches!(
|
||||
strategy,
|
||||
CacheMarkerStrategy::OpenAiContentPartCacheControl { .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_family_over_chat_completions_is_automatic() {
|
||||
assert_eq!(
|
||||
cache_marker_strategy(
|
||||
ProviderId::DigitalOcean,
|
||||
"openai-gpt-4o",
|
||||
&chat_completions()
|
||||
),
|
||||
CacheMarkerStrategy::Automatic
|
||||
);
|
||||
assert_eq!(
|
||||
cache_marker_strategy(ProviderId::OpenAI, "gpt-4o", &chat_completions()),
|
||||
CacheMarkerStrategy::Automatic
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn native_anthropic_uses_messages_breakpoints() {
|
||||
let strategy = cache_marker_strategy(
|
||||
ProviderId::Anthropic,
|
||||
"claude-3-5-sonnet-20241022",
|
||||
&anthropic_messages(),
|
||||
);
|
||||
assert!(matches!(
|
||||
strategy,
|
||||
CacheMarkerStrategy::AnthropicMessagesBreakpoints { .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn anthropic_family_without_content_part_support_is_none() {
|
||||
// An Anthropic-family model over chat completions on a gateway that does not
|
||||
// accept content-part cache_control has no honest marking path.
|
||||
assert_eq!(
|
||||
cache_marker_strategy(
|
||||
ProviderId::Vercel,
|
||||
"anthropic/claude-3.5",
|
||||
&chat_completions()
|
||||
),
|
||||
CacheMarkerStrategy::None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bedrock_is_honest_none() {
|
||||
assert_eq!(
|
||||
cache_marker_strategy(
|
||||
ProviderId::AmazonBedrock,
|
||||
"anthropic.claude-3-5-sonnet",
|
||||
&SupportedUpstreamAPIs::AmazonBedrockConverse(
|
||||
crate::apis::AmazonBedrockApi::Converse
|
||||
)
|
||||
),
|
||||
CacheMarkerStrategy::None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_models_loaded_from_yaml() {
|
||||
|
|
|
|||
|
|
@ -68,7 +68,10 @@ impl ContentUtils<ToolCall> for Vec<MessagesContentBlock> {
|
|||
for block in self {
|
||||
match block {
|
||||
MessagesContentBlock::Text { text, .. } => {
|
||||
content_parts.push(ContentPart::Text { text: text.clone() });
|
||||
content_parts.push(ContentPart::Text {
|
||||
text: text.clone(),
|
||||
cache_control: None,
|
||||
});
|
||||
}
|
||||
MessagesContentBlock::Image { source } => {
|
||||
let url = convert_image_source_to_url(source);
|
||||
|
|
@ -198,7 +201,7 @@ pub fn convert_openai_message_to_anthropic_content(
|
|||
Some(MessageContent::Parts(parts)) => {
|
||||
for part in parts {
|
||||
match part {
|
||||
ContentPart::Text { text } => {
|
||||
ContentPart::Text { text, .. } => {
|
||||
blocks.push(MessagesContentBlock::Text {
|
||||
text: text.clone(),
|
||||
cache_control: None,
|
||||
|
|
|
|||
|
|
@ -324,7 +324,7 @@ fn build_openai_content(
|
|||
None
|
||||
} else if content_parts.len() == 1 && tool_calls.is_empty() {
|
||||
match &content_parts[0] {
|
||||
ContentPart::Text { text } => Some(MessageContent::Text(text.clone())),
|
||||
ContentPart::Text { text, .. } => Some(MessageContent::Text(text.clone())),
|
||||
_ => Some(MessageContent::Parts(content_parts)),
|
||||
}
|
||||
} else if content_parts.is_empty() {
|
||||
|
|
@ -562,6 +562,7 @@ mod tests {
|
|||
},
|
||||
"required": ["location"]
|
||||
}),
|
||||
cache_control: None,
|
||||
}]),
|
||||
tool_choice: Some(MessagesToolChoice {
|
||||
kind: MessagesToolChoiceType::Tool,
|
||||
|
|
@ -620,6 +621,7 @@ mod tests {
|
|||
"type": "object",
|
||||
"properties": {}
|
||||
}),
|
||||
cache_control: None,
|
||||
}]),
|
||||
tool_choice: Some(MessagesToolChoice {
|
||||
kind: MessagesToolChoiceType::Auto,
|
||||
|
|
|
|||
|
|
@ -127,6 +127,7 @@ impl TryFrom<ResponsesInputConverter> for Vec<Message> {
|
|||
| InputContent::OutputText { text } => {
|
||||
Some(crate::apis::openai::ContentPart::Text {
|
||||
text: text.clone(),
|
||||
cache_control: None,
|
||||
})
|
||||
}
|
||||
InputContent::InputImage { image_url, .. } => {
|
||||
|
|
@ -154,6 +155,7 @@ impl TryFrom<ResponsesInputConverter> for Vec<Message> {
|
|||
| InputContent::OutputText { text } => {
|
||||
Some(crate::apis::openai::ContentPart::Text {
|
||||
text: text.clone(),
|
||||
cache_control: None,
|
||||
})
|
||||
}
|
||||
InputContent::InputImage { image_url, .. } => {
|
||||
|
|
@ -330,7 +332,7 @@ impl TryFrom<Message> for BedrockMessage {
|
|||
// Convert OpenAI content parts to Bedrock ContentBlocks
|
||||
for part in parts {
|
||||
match part {
|
||||
crate::apis::openai::ContentPart::Text { text } => {
|
||||
crate::apis::openai::ContentPart::Text { text, .. } => {
|
||||
if !text.is_empty() {
|
||||
content_blocks.push(ContentBlock::Text { text });
|
||||
}
|
||||
|
|
@ -882,6 +884,7 @@ fn convert_openai_tools(tools: Vec<Tool>) -> Vec<MessagesTool> {
|
|||
name: tool.function.name,
|
||||
description: tool.function.description,
|
||||
input_schema: tool.function.parameters,
|
||||
cache_control: None,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -14,13 +14,18 @@ use crate::transforms::lib::*;
|
|||
// Usage Conversions
|
||||
impl From<MessagesUsage> for Usage {
|
||||
fn from(val: MessagesUsage) -> Self {
|
||||
Usage {
|
||||
let mut usage = Usage {
|
||||
prompt_tokens: val.input_tokens,
|
||||
completion_tokens: val.output_tokens,
|
||||
total_tokens: val.input_tokens + val.output_tokens,
|
||||
prompt_tokens_details: None,
|
||||
completion_tokens_details: None,
|
||||
}
|
||||
cache_read_input_tokens: val.cache_read_input_tokens,
|
||||
cache_creation_input_tokens: val.cache_creation_input_tokens,
|
||||
};
|
||||
// Surface Anthropic cache reads to OpenAI clients as cached_tokens.
|
||||
usage.normalize_cache_tokens();
|
||||
usage
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -244,6 +249,7 @@ impl TryFrom<MessagesResponse> for ChatCompletionsResponse {
|
|||
total_tokens: resp.usage.input_tokens + resp.usage.output_tokens,
|
||||
prompt_tokens_details: None,
|
||||
completion_tokens_details: None,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
Ok(ChatCompletionsResponse {
|
||||
|
|
@ -312,6 +318,7 @@ impl TryFrom<ConverseResponse> for ChatCompletionsResponse {
|
|||
total_tokens: resp.usage.total_tokens,
|
||||
prompt_tokens_details: None,
|
||||
completion_tokens_details: None,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
// Generate a response ID (using timestamp since Bedrock doesn't provide one)
|
||||
|
|
@ -980,6 +987,7 @@ mod tests {
|
|||
total_tokens: 30,
|
||||
prompt_tokens_details: None,
|
||||
completion_tokens_details: None,
|
||||
..Default::default()
|
||||
},
|
||||
system_fingerprint: None,
|
||||
service_tier: Some("default".to_string()),
|
||||
|
|
@ -1061,6 +1069,7 @@ mod tests {
|
|||
total_tokens: 40,
|
||||
prompt_tokens_details: None,
|
||||
completion_tokens_details: None,
|
||||
..Default::default()
|
||||
},
|
||||
system_fingerprint: None,
|
||||
service_tier: None,
|
||||
|
|
@ -1134,6 +1143,7 @@ mod tests {
|
|||
total_tokens: 101,
|
||||
prompt_tokens_details: None,
|
||||
completion_tokens_details: None,
|
||||
..Default::default()
|
||||
},
|
||||
system_fingerprint: Some("fp_7eeb46f068".to_string()),
|
||||
service_tier: Some("default".to_string()),
|
||||
|
|
|
|||
|
|
@ -247,6 +247,7 @@ impl TryFrom<ConverseStreamEvent> for ChatCompletionsStreamResponse {
|
|||
total_tokens: metadata_event.usage.total_tokens,
|
||||
prompt_tokens_details: None,
|
||||
completion_tokens_details: None,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
Ok(create_openai_chunk(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue