scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

fix(cache): honor controls and tenant metadata in valkey semantic bridge

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee 2026-09-21 22:33:11 +00:00
parent dd175245cf
commit 21d1604e64
5 changed files with 216 additions and 41 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
litellm-rust/crates/cache-valkey-semantic/src/lib.rs
View file

@ -273,13 +273,11 @@ pub fn prompt_from_context(context: &SemanticCacheContext) -> Option<String> {
if let Some(Value::Array(messages)) = context.messages.as_ref()
&& !messages.is_empty()
{
return Some(
messages
.iter()
.filter_map(Value::as_object)
.map(message_text)
.collect(),
);
return messages
.iter()
.filter_map(Value::as_object)
.map(message_text)
.collect();
}
let input = context.input.as_ref()?;
let mut parts = Vec::new();
@ -288,21 +286,25 @@ pub fn prompt_from_context(context: &SemanticCacheContext) -> Option<String> {
(!prompt.is_empty()).then_some(prompt)
}
fn message_text(message: &serde_json::Map<String, Value>) -> String {
fn message_text(message: &serde_json::Map<String, Value>) -> Option<String> {
let content = match message.get("content") {
Some(Value::String(value)) => value.clone(),
Some(Value::Array(parts)) => parts
.iter()
.filter_map(Value::as_object)
.filter_map(|part| part.get("text").and_then(Value::as_str))
.filter(|text| !text.is_empty())
.collect(),
Some(Value::Array(parts)) => {
let mut content = String::new();
for part in parts {
let part = part.as_object()?;
if let Some(text) = part.get("text").and_then(Value::as_str) {
content.push_str(text);
}
}
content
}
_ => String::new(),
};
format!(
Some(format!(
"{content}{}",
search_results_text(message.get("search_results"))
)
))
}
fn search_results_text(value: Option<&Value>) -> String {
@ -732,6 +734,7 @@ mod tests {
#[rstest]
#[case(json!([{"content": "hello"}]), None, Some("hello"))]
#[case(json!([{"content": [{"text": "hello"}, {"text": " world"}]}]), None, Some("hello world"))]
#[case(json!([{"content": ["raw", {"text": "hello"}]}]), None, None)]
#[case(json!([{"search_results": [{"source": "s", "title": "t", "content": [{"text": "c"}], "citations": ["x"]}]}]), None, Some(r#"stc["x"]"#))]
#[case(Value::Array(vec![]), Some(json!(" hello ")), Some("hello"))]
#[case(Value::Array(vec![]), Some(json!([{"content": "first"}, {"text": "second"}])), Some("first\nsecond"))]

28
litellm-rust/crates/python-bridge/src/cache/native.rs vendored
View file

@ -34,11 +34,24 @@ fn semantic_key(request: &NativeRequest, scope: &str) -> litellm_cache_response:
];
let end_user = (scope == "end_user").then_some("user_api_key_end_user_id");
for name in TENANT.into_iter().chain(end_user) {
let Some(value) = request
.metadata
.as_ref()
.and_then(|metadata| metadata.get(name))
else {
let sources = [
request.metadata.as_ref(),
request.litellm_metadata.as_ref(),
request
.litellm_params
.as_ref()
.and_then(|params| params.get("metadata")),
request
.litellm_params
.as_ref()
.and_then(|params| params.get("litellm_metadata")),
];
let Some(value) = sources.into_iter().flatten().find_map(|source| {
source
.as_object()
.and_then(|values| values.get(name))
.filter(|value| !value.is_null())
}) else {
continue;
};
let value = match value {
@ -340,7 +353,6 @@ impl NativeResponseCache {
Arc::clone(cache.backend_arc()),
embedder.clone(),
Self::semantic(&request, scope),
super::request::now(),
),
),
}
@ -417,7 +429,6 @@ impl NativeResponseCache {
embedder.clone(),
Self::semantic(&request, scope),
response,
super::request::now(),
),
),
}
@ -517,7 +528,6 @@ impl NativeResponseCache {
embedder.clone(),
requests,
responses,
super::request::now(),
),
)
}
@ -565,6 +575,8 @@ mod tests {
messages: Some(json!([{"role": "user", "content": "prompt"}])),
input: None,
metadata: Some(metadata),
litellm_metadata: None,
litellm_params: None,
}
}

6
litellm-rust/crates/python-bridge/src/cache/request.rs vendored
View file

@ -17,6 +17,8 @@ struct RequestInput {
messages: Option<Value>,
input: Option<Value>,
metadata: Option<Value>,
litellm_metadata: Option<Value>,
litellm_params: Option<Value>,
}
pub(super) struct NativeRequest {
@ -27,6 +29,8 @@ pub(super) struct NativeRequest {
pub(super) messages: Option<Value>,
pub(super) input: Option<Value>,
pub(super) metadata: Option<Value>,
pub(super) litellm_metadata: Option<Value>,
pub(super) litellm_params: Option<Value>,
}
pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult<NativeRequest> {
@ -46,6 +50,8 @@ fn request_input(input: RequestInput) -> PyResult<NativeRequest> {
messages: input.messages,
input: input.input,
metadata: input.metadata,
litellm_metadata: input.litellm_metadata,
litellm_params: input.litellm_params,
})
}

26
litellm-rust/crates/python-bridge/src/cache/semantic_step.rs vendored
View file

@ -28,7 +28,7 @@ pub(super) struct SemanticEmbedExecution {
embedder: PythonEmbedder,
requests: Vec<ResponseCacheRequest<SemanticCacheContext>>,
op: Op,
now: Duration,
now: Option<Duration>,
prepared: Vec<Option<Vec<f32>>>,
index: usize,
state: State,
@ -39,14 +39,13 @@ impl SemanticEmbedExecution {
backend: Arc<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
embedder: PythonEmbedder,
request: ResponseCacheRequest<SemanticCacheContext>,
now: Duration,
) -> Self {
Self {
backend,
embedder,
requests: vec![request],
op: Op::Lookup,
now,
now: None,
prepared: vec![None],
index: 0,
state: State::Start,
@ -58,14 +57,13 @@ impl SemanticEmbedExecution {
embedder: PythonEmbedder,
request: ResponseCacheRequest<SemanticCacheContext>,
response: Value,
now: Duration,
) -> Self {
Self {
backend,
embedder,
requests: vec![request],
op: Op::Store(response),
now,
now: None,
prepared: vec![None],
index: 0,
state: State::Start,
@ -77,7 +75,6 @@ impl SemanticEmbedExecution {
embedder: PythonEmbedder,
requests: Vec<ResponseCacheRequest<SemanticCacheContext>>,
responses: Vec<Value>,
now: Duration,
) -> Self {
Self {
backend,
@ -85,15 +82,26 @@ impl SemanticEmbedExecution {
prepared: vec![None; requests.len()],
requests,
op: Op::StoreBatch(responses),
now,
now: None,
index: 0,
state: State::Start,
}
}
fn start(&mut self, py: Python<'_>) -> PyResult<ExecutionStep> {
if self.now.is_none() {
self.now = Some(super::request::now());
}
while self.index < self.requests.len() {
let request = &self.requests[self.index];
let enabled = match &self.op {
Op::Lookup => request.controls.reads(),
Op::Store(_) | Op::StoreBatch(_) => request.controls.writes(),
};
if !enabled {
self.index += 1;
continue;
}
let Some(prompt) = prompt_from_context(&request.context) else {
self.index += 1;
continue;
@ -113,7 +121,9 @@ impl SemanticEmbedExecution {
let requests = self.requests.clone();
let prepared = self.prepared.clone();
let backend = Arc::clone(&self.backend);
let now = self.now;
let now = self
.now
.ok_or_else(|| PyRuntimeError::new_err("semantic cache timestamp is unavailable"))?;
let awaitable = match &self.op {
Op::Lookup => {
let Some(request) = requests.into_iter().next() else {

162
tests/test_litellm_rust/test_valkey_semantic_cache_native.py
View file

@ -53,8 +53,12 @@ def _request(prompt: str = "semantic cache prompt") -> dict[str, object]:
def _field_request(
prompt: str,
metadata: Mapping[str, object],
*,
namespace: str | None = None,
litellm_metadata: Mapping[str, object] | None = None,
litellm_params: Mapping[str, object] | None = None,
) -> dict[str, object]:
return {
request: Final = {
"key": {
"fields": [
{
@ -69,23 +73,32 @@ def _field_request(
"api_parameter": True,
"internal_parameter": False,
},
]
],
"namespace": namespace,
},
"messages": [{"role": "user", "content": prompt}],
"metadata": dict(metadata),
}
if litellm_metadata is not None:
request["litellm_metadata"] = dict(litellm_metadata)
if litellm_params is not None:
request["litellm_params"] = dict(litellm_params)
return request
def _facade(
url: str,
index_name: str,
embeddings: Mapping[str, list[float]],
*,
namespace: str | None = None,
) -> Cache:
facade: Final = Cache(
type=LiteLLMCacheType.VALKEY_SEMANTIC,
redis_url=url,
similarity_threshold=0.8,
valkey_semantic_cache_index_name=index_name,
namespace=namespace,
)
vectors: Final = embeddings
@ -175,6 +188,45 @@ async def test_async_lookup_and_store(
assert await binding.async_lookup(request) == {"answer": "async"}
async def test_disabled_cache_controls_skip_async_embedding(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
calls: Final = []
async def fail_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]:
calls.append(prompt)
raise AssertionError("embedding must not run")
backend._get_async_embedding = fail_embedding
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
backend,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
controls: Final = {
"supported_call_type": True,
"configured": True,
"native_backend": True,
"default_on": True,
"caching": True,
"no_cache": False,
"no_store": False,
"use_cache": True,
}
no_read_request: Final = {**_request(), "controls": {**controls, "no_cache": True}}
assert await binding.async_lookup(no_read_request) is None
no_write_request: Final = {**_request(), "controls": {**controls, "no_store": True}}
await binding.async_store(no_write_request, {"answer": "blocked"})
assert calls == []
client: Final = redis.Redis.from_url(valkey_url)
assert list(client.scan_iter(f"{index_name}:*")) == []
client.close()
async def test_async_embedding_runs_inline_in_caller_task(
valkey_url: str,
index_name: str,
@ -323,6 +375,25 @@ def test_malformed_entry_is_a_miss_on_native_and_python(
assert backend.get_cache("key", messages=_request()["messages"]) is None
def test_mixed_content_parts_match_python_semantic_behavior(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
messages: Final = [{"role": "user", "content": ["raw", {"text": "hello"}]}]
backend.set_cache("key", {"answer": "mixed"}, messages=messages)
assert backend.get_cache("key", messages=messages) is None
handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
request: Final = {**_request(), "messages": messages}
binding.store(request, {"answer": "mixed"})
assert binding.lookup(request) is None
client: Final = redis.Redis.from_url(valkey_url)
assert list(client.scan_iter(f"{index_name}:*")) == []
client.close()
async def test_async_store_batch_and_lookup(
valkey_url: str,
index_name: str,
@ -406,6 +477,84 @@ def test_field_key_matches_python_semantic_scope(
client.close()
def test_field_key_reads_all_python_tenant_metadata_sources(
valkey_url: str,
index_name: str,
) -> None:
facade: Final = _facade(valkey_url, index_name, {"semantic cache prompt": [1.0, 0.0]})
params_metadata: Final = {"user_api_key_team_id": "team-from-params"}
expected: Final = facade.get_cache_key(
model="gpt-4.1",
messages=[{"role": "user", "content": "semantic cache prompt"}],
metadata={},
litellm_params={"metadata": params_metadata},
)
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
facade.cache,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
binding.store(
_field_request(
"semantic cache prompt",
{},
litellm_params={"metadata": params_metadata},
),
{"answer": "params"},
)
client: Final = redis.Redis.from_url(valkey_url)
documents: Final = list(client.scan_iter(f"{index_name}:*"))
assert len(documents) == 1
document_parts: Final = documents[0].decode().split(":")
assert document_parts[1] == hashlib.sha256(expected.encode()).hexdigest()
client.close()
assert (
binding.lookup(
_field_request(
"semantic cache prompt",
{},
litellm_metadata={"user_api_key_team_id": "team-from-litellm"},
)
)
is None
)
def test_namespace_isolates_semantic_entries(
valkey_url: str,
index_name: str,
) -> None:
facade: Final = _facade(
valkey_url,
index_name,
{"semantic cache prompt": [1.0, 0.0]},
namespace="team-a",
)
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
facade.cache,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
team_a: Final = _field_request("semantic cache prompt", {}, namespace="team-a")
team_b: Final = _field_request("semantic cache prompt", {}, namespace="team-b")
binding.store(team_a, {"answer": "team-a"})
assert binding.lookup(team_b) is None
assert binding.lookup(team_a) == {"answer": "team-a"}
cached: Final = cast(
Mapping[str, object],
facade.get_cache(
model="gpt-4.1",
messages=[{"role": "user", "content": "semantic cache prompt"}],
),
)
assert cached == {"answer": "team-a"}
def test_field_key_isolates_tenant_scope(
valkey_url: str,
index_name: str,
@ -422,13 +571,8 @@ def test_field_key_isolates_tenant_scope(
_field_request("semantic cache prompt", {"user_api_key": "k1"}),
{"answer": "tenant one"},
)
assert (
binding.lookup(_field_request("semantic cache prompt", {"user_api_key": "k2"}))
is None
)
assert binding.lookup(_field_request("semantic cache prompt", {"user_api_key": "k1"})) == {
"answer": "tenant one"
}
assert binding.lookup(_field_request("semantic cache prompt", {"user_api_key": "k2"})) is None
assert binding.lookup(_field_request("semantic cache prompt", {"user_api_key": "k1"})) == {"answer": "tenant one"}
def test_tls_valkey_facade_falls_back_to_python(