Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 7 additions & 7 deletions crates/adaptive/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use nemo_relay::plugin::ConfigPolicy;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value as Json};

use crate::response_cache::config::{BackendConfig, KEY_STRATEGY_EXACT_REQUEST, ToolCacheConfig};
use crate::response_cache::config::{BackendConfig, ResponseCacheKeyStrategy, ToolCacheConfig};

/// Canonical config document for the adaptive plugin component.
#[derive(Debug, Clone, Serialize, Deserialize)]
Expand All @@ -34,7 +34,7 @@ pub struct AdaptiveConfig {
/// Adaptive Cache Governor settings.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub acg: Option<AcgComponentConfig>,
/// Opt-in exact-match LLM response and tool-result cache. When present,
/// Opt-in LLM response and tool-result cache. When present,
/// the adaptive plugin installs the response-cache execution intercept(s).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub response_cache: Option<ResponseCacheConfig>,
Expand Down Expand Up @@ -191,7 +191,7 @@ impl Default for AcgComponentConfig {
}
}

/// Configuration for the adaptive plugin's exact-match LLM response and
/// Configuration for the adaptive plugin's LLM response and
/// opt-in tool-result cache feature.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default)]
Expand All @@ -212,8 +212,8 @@ pub struct ResponseCacheConfig {
/// requests explicitly pinned deterministic (`temperature` = 0) — absent
/// or unreadable temperatures count as nondeterministic.
pub cache_nondeterministic: bool,
/// Key strategy. Only [`KEY_STRATEGY_EXACT_REQUEST`] is supported.
pub key_strategy: String,
/// Typed key-derivation strategy.
pub key_strategy: ResponseCacheKeyStrategy,
Comment thread
coderabbitai[bot] marked this conversation as resolved.
/// Request headers (case-insensitive) folded into the key; never auth headers.
pub header_allowlist: Vec<String>,
/// Storage backend selection.
Expand All @@ -231,7 +231,7 @@ impl Default for ResponseCacheConfig {
priority: 50,
bypass_rate: 0.0,
cache_nondeterministic: false,
key_strategy: KEY_STRATEGY_EXACT_REQUEST.to_string(),
key_strategy: ResponseCacheKeyStrategy::ExactRequest,
header_allowlist: Vec::new(),
backend: BackendConfig::default(),
tools: None,
Expand Down Expand Up @@ -402,7 +402,7 @@ nemo_relay::editor_config! {
priority => { label: "priority", kind: Integer },
bypass_rate => { label: "bypass_rate", kind: Float },
cache_nondeterministic => { label: "cache_nondeterministic", kind: Boolean },
key_strategy => { label: "key_strategy", kind: String },
key_strategy => { label: "key_strategy", kind: Enum, values: ["exact_request", "logical"] },
header_allowlist => { label: "header_allowlist", kind: Json },
backend => {
label: "backend",
Expand Down
2 changes: 1 addition & 1 deletion crates/adaptive/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -57,8 +57,8 @@ pub use context_helpers::{
pub use error::{AdaptiveError, Result};
#[cfg(feature = "redis-backend")]
pub use redis::RedisBackend;
pub use response_cache::RESPONSE_CACHE_MARK;
pub use response_cache::config::{ToolCacheConfig, ToolClass, ToolOverride};
pub use response_cache::{RESPONSE_CACHE_MARK, ResponseCacheKeyStrategy};
pub use runtime::features::AdaptiveRuntime;
pub use storage::erased::AnyBackend;
pub use storage::memory::InMemoryBackend;
Expand Down
64 changes: 61 additions & 3 deletions crates/adaptive/src/response_cache/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,11 +11,69 @@

use std::collections::BTreeMap;

use serde::{Deserialize, Serialize};
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use serde_json::{Map, Value as Json};

/// Exact-request key strategy identifier.
pub const KEY_STRATEGY_EXACT_REQUEST: &str = "exact_request";
/// Strategy for deriving an LLM response-cache key.
///
/// The `Unknown` variant preserves an unsupported JSON/TOML value long enough
/// for configuration validation to report it with a field-specific diagnostic.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub enum ResponseCacheKeyStrategy {
/// Key on the normalized request exactly.
#[default]
ExactRequest,
/// Normalize tool schemas structurally while preserving their interface.
Logical,
/// A wire value not supported by this Relay build.
Unknown(String),
}

impl ResponseCacheKeyStrategy {
/// Stable JSON/TOML representation of this strategy.
pub fn as_str(&self) -> &str {
match self {
Self::ExactRequest => "exact_request",
Self::Logical => "logical",
Self::Unknown(value) => value,
}
}
}

impl From<&str> for ResponseCacheKeyStrategy {
fn from(value: &str) -> Self {
match value {
"exact_request" => Self::ExactRequest,
"logical" => Self::Logical,
_ => Self::Unknown(value.to_string()),
}
}
}

impl From<String> for ResponseCacheKeyStrategy {
fn from(value: String) -> Self {
Self::from(value.as_str())
}
}

impl Serialize for ResponseCacheKeyStrategy {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_str(self.as_str())
}
}

impl<'de> Deserialize<'de> for ResponseCacheKeyStrategy {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let value = String::deserialize(deserializer)?;
Ok(Self::from(value))
}
}

/// Default in-memory byte budget: 256 MiB.
pub const DEFAULT_MAX_BYTES: usize = 256 * 1024 * 1024;
Expand Down
50 changes: 50 additions & 0 deletions crates/adaptive/src/response_cache/key.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ use serde_json::{Map, Value as Json, json};
use sha2::{Digest, Sha256};

use crate::config::ResponseCacheConfig;
use crate::response_cache::config::ResponseCacheKeyStrategy;
use crate::response_cache::mark::CacheReason;
use crate::response_cache::store::CACHE_SCHEMA_VERSION;

Expand Down Expand Up @@ -77,6 +78,13 @@ pub fn build_cache_key(
normalize_tool_call_ids(object);
}

if config.key_strategy == ResponseCacheKeyStrategy::Logical
&& let Some(object) = body.as_object_mut()
&& let Some(tools) = object.get("tools").cloned()
{
object.insert("tools".to_string(), structural_tool_schema(&tools));
}

let header_allowlist = normalized_header_allowlist(&config.header_allowlist);
let headers = cache_key_headers(&request.headers, &header_allowlist);

Expand Down Expand Up @@ -693,6 +701,48 @@ fn rewrite_id(id_value: &mut Json, mapping: &mut Map<String, Json>) {
*id_value = Json::String(stable);
}

/// Fingerprint of the tool set for the `logical` strategy: each tool keeps its
/// full definition minus string-valued `description` keys (stripped
/// recursively), and the array is sorted so tool order does not key.
fn structural_tool_schema(tools: &Json) -> Json {
let Some(array) = tools.as_array() else {
return tools.clone();
};
let mut entries: Vec<Json> = array
.iter()
.map(|tool| {
let mut entry = tool.clone();
strip_descriptions(&mut entry);
entry
})
.collect();
entries
.sort_by_cached_key(|entry| serde_json_canonicalizer::to_string(entry).unwrap_or_default());
Json::Array(entries)
}

/// Removes every string-valued `description` key, at any depth. A non-string
/// value under that key (e.g. a schema property named `description`) is
/// interface, not prose, and stays.
fn strip_descriptions(value: &mut Json) {
match value {
Json::Object(object) => {
if object.get("description").is_some_and(Json::is_string) {
object.remove("description");
}
for nested in object.values_mut() {
strip_descriptions(nested);
}
}
Json::Array(items) => {
for item in items {
strip_descriptions(item);
}
}
_ => {}
}
}

#[cfg(test)]
#[path = "../../tests/unit/response_cache/key_tests.rs"]
mod tests;
6 changes: 2 additions & 4 deletions crates/adaptive/src/response_cache/mod.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Opt-in exact-match cache for LLM responses and tool results: a feature of
//! Opt-in cache for LLM responses and tool results: a feature of
//! the adaptive plugin, configured through
//! [`crate::config::AdaptiveConfig::response_cache`].
//!
Expand All @@ -23,9 +23,7 @@ pub mod store;
pub(crate) mod tool;

pub use crate::config::ResponseCacheConfig;
pub use crate::response_cache::config::{
BackendConfig, KEY_STRATEGY_EXACT_REQUEST, ToolCacheConfig,
};
pub use crate::response_cache::config::{BackendConfig, ResponseCacheKeyStrategy, ToolCacheConfig};
pub(crate) use crate::response_cache::intercept::{make_intercept, make_stream_intercept};
pub use crate::response_cache::mark::RESPONSE_CACHE_MARK;
pub(crate) use crate::response_cache::store::build_store;
Expand Down
9 changes: 6 additions & 3 deletions crates/adaptive/src/runtime/validation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use nemo_relay::plugin::{
use serde_json::Value as Json;

use crate::config::{AdaptiveConfig, BackendSpec, ResponseCacheConfig};
use crate::response_cache::config::{KEY_STRATEGY_EXACT_REQUEST, ToolCacheConfig, ToolClass};
use crate::response_cache::config::{ResponseCacheKeyStrategy, ToolCacheConfig, ToolClass};
use crate::response_cache::tool::{is_supported_tool_pattern, wildcard_patterns_overlap};

pub fn validate_config(config: &AdaptiveConfig) -> ConfigReport {
Expand Down Expand Up @@ -123,11 +123,14 @@ fn validate_response_cache(report: &mut ConfigReport, config: &ResponseCacheConf
"bypass_rate must be in [0.0, 1.0]".to_string(),
));
}
if config.key_strategy != KEY_STRATEGY_EXACT_REQUEST {
if matches!(config.key_strategy, ResponseCacheKeyStrategy::Unknown(_)) {
report.diagnostics.push(response_cache_error(
"response_cache.unsupported_key_strategy",
Some("key_strategy"),
format!("unsupported key_strategy; only \"{KEY_STRATEGY_EXACT_REQUEST}\" is supported"),
format!(
"unsupported key_strategy '{}'; supported: \"exact_request\", \"logical\"",
config.key_strategy.as_str()
),
));
}
// Auth material must never enter the key or the stored entries.
Expand Down
52 changes: 48 additions & 4 deletions crates/adaptive/tests/integration/response_cache_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,8 @@ use nemo_relay::plugin::{
};
use nemo_relay_adaptive::plugin_component::{ComponentSpec, register_adaptive_component};
use nemo_relay_adaptive::{
AcgComponentConfig, AdaptiveConfig, BackendSpec, ResponseCacheConfig, StateConfig,
ToolCacheConfig, ToolClass, ToolOverride,
AcgComponentConfig, AdaptiveConfig, BackendSpec, ResponseCacheConfig, ResponseCacheKeyStrategy,
StateConfig, ToolCacheConfig, ToolClass, ToolOverride,
};
use serde_json::{Value as Json, json};
use tokio::sync::Mutex;
Expand Down Expand Up @@ -555,7 +555,7 @@ async fn invalid_config_is_rejected_by_validation() {
response_cache: Some(ResponseCacheConfig {
ttl_seconds: 0,
bypass_rate: 2.0,
key_strategy: "semantic".to_string(),
key_strategy: ResponseCacheKeyStrategy::Unknown("semantic".to_string()),
namespace: "invalid-config-test".to_string(),
..ResponseCacheConfig::default()
}),
Expand Down Expand Up @@ -647,7 +647,7 @@ async fn response_cache_validation_diagnostics_identify_the_invalid_setting() {

let mut cache = ResponseCacheConfig {
namespace: "diagnostic-contract-test".to_string(),
key_strategy: "semantic".to_string(),
key_strategy: ResponseCacheKeyStrategy::Unknown("semantic".to_string()),
tools: Some(ToolCacheConfig {
enabled: true,
default: ToolClass {
Expand Down Expand Up @@ -792,6 +792,50 @@ async fn hit_preserves_usage_on_the_end_event_and_reports_savings_on_the_mark()
deregister_subscriber("response_cache_event_capture").unwrap();
}

#[tokio::test]
async fn logical_strategy_reuses_across_reworded_tool_descriptions() {
let _guard = TEST_MUTEX.lock().await;
reset_global();
// `logical` must be accepted by validation (activate_cache asserts no
// diagnostics) and must reuse across a reworded tool description end-to-end.
activate_cache(ResponseCacheConfig {
namespace: "logical-key-integration-test".to_string(),
key_strategy: ResponseCacheKeyStrategy::Logical,
..ResponseCacheConfig::default()
})
.await;

let calls = Arc::new(AtomicUsize::new(0));
let provider = counting_provider(Arc::clone(&calls), sample_body());

let request_with_tool = |description: &str| LlmRequest {
headers: serde_json::Map::new(),
content: json!({
"model": "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning",
"messages": [{"role": "user", "content": "what is the weather?"}],
"temperature": 0.0,
"tools": [{"type": "function", "function": {
"name": "get_weather",
"description": description,
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}}
}}]
}),
};

call(&provider, request_with_tool("Get the weather for a city.")).await;
call(
&provider,
request_with_tool("Look up the current weather (reworded)."),
)
.await;

assert_eq!(
calls.load(Ordering::SeqCst),
1,
"logical keying must serve the reworded-tool repeat from cache"
);
}

#[tokio::test]
async fn errors_are_not_cached() {
let _guard = TEST_MUTEX.lock().await;
Expand Down
Loading
Loading