diff --git a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs index e3a3e6094c4..6062ccc842c 100644 --- a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs +++ b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs @@ -273,13 +273,11 @@ pub fn prompt_from_context(context: &SemanticCacheContext) -> Option { 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 { (!prompt.is_empty()).then_some(prompt) } -fn message_text(message: &serde_json::Map) -> String { +fn message_text(message: &serde_json::Map) -> Option { 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"))] diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index fefc09c1321..d8f7a693108 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -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, } } diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs index 3e19e7fdc22..036951891a1 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -17,6 +17,8 @@ struct RequestInput { messages: Option, input: Option, metadata: Option, + litellm_metadata: Option, + litellm_params: Option, } pub(super) struct NativeRequest { @@ -27,6 +29,8 @@ pub(super) struct NativeRequest { pub(super) messages: Option, pub(super) input: Option, pub(super) metadata: Option, + pub(super) litellm_metadata: Option, + pub(super) litellm_params: Option, } pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { @@ -46,6 +50,8 @@ fn request_input(input: RequestInput) -> PyResult { messages: input.messages, input: input.input, metadata: input.metadata, + litellm_metadata: input.litellm_metadata, + litellm_params: input.litellm_params, }) } diff --git a/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs b/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs index d8f3ab297f8..24caf3374d6 100644 --- a/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs +++ b/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs @@ -28,7 +28,7 @@ pub(super) struct SemanticEmbedExecution { embedder: PythonEmbedder, requests: Vec>, op: Op, - now: Duration, + now: Option, prepared: Vec>>, index: usize, state: State, @@ -39,14 +39,13 @@ impl SemanticEmbedExecution { backend: Arc>, embedder: PythonEmbedder, request: ResponseCacheRequest, - 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, 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>, responses: Vec, - 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 { + 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 { diff --git a/tests/test_litellm_rust/test_valkey_semantic_cache_native.py b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py index fb5c5965333..c87a9f86a80 100644 --- a/tests/test_litellm_rust/test_valkey_semantic_cache_native.py +++ b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py @@ -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(