test(cache): add Valkey semantic contract coverage

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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

@@ -62,6 +62,14 @@ enum Connections<C> {
     Fixed(Mutex<C>),
 }
 
+#[derive(Clone)]
+struct IndexState {
+    name: String,
+    prefix: String,
+    dimension: Arc<Mutex<Option<usize>>>,
+    similarity_threshold: f64,
+}
+
 struct ConnectionRef<'a>(&'a mut dyn redis::ConnectionLike);
 
 impl redis::ConnectionLike for ConnectionRef<'_> {
@@ -187,18 +195,13 @@ where
         &self.config.index_name
     }
 
-    fn key_prefix(&self) -> String {
-        format!("{}:", self.config.index_name)
-    }
-
-    fn ensure_index(&self, dimension: usize) -> Result<(), Error> {
-        ensure_index(
-            &self.connections,
-            &self.config.index_name,
-            &self.key_prefix(),
-            &self.index_dimension,
-            dimension,
-        )
+    fn index_state(&self) -> IndexState {
+        IndexState {
+            name: self.config.index_name.clone(),
+            prefix: format!("{}:", self.config.index_name),
+            dimension: Arc::clone(&self.index_dimension),
+            similarity_threshold: self.config.similarity_threshold,
+        }
     }
 }
 
@@ -225,37 +228,19 @@ where
             return Ok(());
         };
         let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?;
-        self.ensure_index(embedding.len())?;
         let scope = scope_tag(key);
-        let document = format!("{}{}:{}", self.key_prefix(), scope, Uuid::new_v4());
         let response = self.codec.encode(&value)?;
         let vector = embedding_bytes(&embedding);
-        let ttl = self.get_ttl(context);
-        self.connections.execute(|connection| {
-            let mut pipeline = redis::pipe();
-            pipeline
-                .cmd("HSET")
-                .arg(&document)
-                .arg("litellm_cache_key")
-                .arg(&scope)
-                .arg("prompt")
-                .arg(prompt)
-                .arg("response")
-                .arg(response)
-                .arg("embedding")
-                .arg(vector)
-                .ignore();
-            if let Some(ttl) = ttl {
-                pipeline
-                    .cmd("EXPIRE")
-                    .arg(&document)
-                    .arg(ttl.as_secs())
-                    .ignore();
-            }
-            pipeline
-                .query::<()>(connection)
-                .map_err(|_| Error::Unavailable)
-        })
+        let index = self.index_state();
+        write_document(
+            &self.connections,
+            &index,
+            &scope,
+            &prompt,
+            response,
+            vector,
+            self.get_ttl(context),
+        )
     }
 
     fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
@@ -263,43 +248,14 @@ where
             return Ok(None);
         };
         let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?;
-        self.ensure_index(embedding.len())?;
         let scope = scope_tag(key);
-        let query =
-            format!("(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]");
         let vector = embedding_bytes(&embedding);
-        let response = self.connections.execute(|connection| {
-            redis::cmd("FT.SEARCH")
-                .arg(&self.config.index_name)
-                .arg(query)
-                .arg("PARAMS")
-                .arg(2)
-                .arg("vec")
-                .arg(vector)
-                .arg("RETURN")
-                .arg(2)
-                .arg("response")
-                .arg("vector_distance")
-                .arg("DIALECT")
-                .arg(2)
-                .query::<redis::Value>(connection)
-                .map_err(|_| Error::Unavailable)
-        })?;
-        let Some(fields) = search_fields(response)? else {
+        let index = self.index_state();
+        let Some(response) =
+            search_document(&self.connections, &index, &scope, vector, embedding.len())?
+        else {
             return Ok(None);
         };
-        let response = fields
-            .iter()
-            .find_map(|(name, value)| (name == "response").then(|| value.clone()))
-            .ok_or(Error::InvalidEntry)?;
-        let distance = fields
-            .iter()
-            .find_map(|(name, value)| (name == "vector_distance").then(|| value.clone()))
-            .ok_or(Error::InvalidEntry)?;
-        let distance = parse_f64(&distance)?;
-        if 1.0 - distance < self.config.similarity_threshold {
-            return Ok(None);
-        }
         self.codec.decode(&response).map(Some)
     }
 
@@ -321,47 +277,13 @@ where
                 .async_embed(&prompt, metadata.as_ref())
                 .await?;
             let connections = Arc::clone(&self.connections);
-            let config = self.config.clone();
-            let index_dimension = Arc::clone(&self.index_dimension);
+            let index = self.index_state();
             let response = self.codec.encode(&value)?;
             let vector = embedding_bytes(&embedding);
-            let prefix = format!("{}:", config.index_name);
             let scope = scope_tag(&key);
-            let document = format!("{prefix}{scope}:{}", Uuid::new_v4());
             let ttl = context.ttl;
             tokio::task::spawn_blocking(move || {
-                ensure_index(
-                    &connections,
-                    &config.index_name,
-                    &prefix,
-                    &index_dimension,
-                    embedding.len(),
-                )?;
-                connections.execute(|connection| {
-                    let mut pipeline = redis::pipe();
-                    pipeline
-                        .cmd("HSET")
-                        .arg(&document)
-                        .arg("litellm_cache_key")
-                        .arg(&scope)
-                        .arg("prompt")
-                        .arg(prompt)
-                        .arg("response")
-                        .arg(response)
-                        .arg("embedding")
-                        .arg(vector)
-                        .ignore();
-                    if let Some(ttl) = ttl {
-                        pipeline
-                            .cmd("EXPIRE")
-                            .arg(&document)
-                            .arg(ttl.as_secs())
-                            .ignore();
-                    }
-                    pipeline
-                        .query::<()>(connection)
-                        .map_err(|_| Error::Unavailable)
-                })
+                write_document(&connections, &index, &scope, &prompt, response, vector, ttl)
             })
             .await
             .map_err(|_| Error::Unavailable)?
@@ -385,56 +307,11 @@ where
                 .async_embed(&prompt, metadata.as_ref())
                 .await?;
             let connections = Arc::clone(&self.connections);
-            let config = self.config.clone();
-            let index_dimension = Arc::clone(&self.index_dimension);
-            let threshold = config.similarity_threshold;
+            let index = self.index_state();
             tokio::task::spawn_blocking(move || {
-                let prefix = format!("{}:", config.index_name);
-                ensure_index(
-                    &connections,
-                    &config.index_name,
-                    &prefix,
-                    &index_dimension,
-                    embedding.len(),
-                )?;
                 let scope = scope_tag(&key);
-                let query = format!(
-                    "(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]"
-                );
                 let vector = embedding_bytes(&embedding);
-                let response = connections.execute(|connection| {
-                    redis::cmd("FT.SEARCH")
-                        .arg(&config.index_name)
-                        .arg(query)
-                        .arg("PARAMS")
-                        .arg(2)
-                        .arg("vec")
-                        .arg(vector)
-                        .arg("RETURN")
-                        .arg(2)
-                        .arg("response")
-                        .arg("vector_distance")
-                        .arg("DIALECT")
-                        .arg(2)
-                        .query::<redis::Value>(connection)
-                        .map_err(|_| Error::Unavailable)
-                })?;
-                let Some(fields) = search_fields(response)? else {
-                    return Ok(None);
-                };
-                let response = fields
-                    .iter()
-                    .find_map(|(name, value)| (name == "response").then(|| value.clone()))
-                    .ok_or(Error::InvalidEntry)?;
-                let distance = fields
-                    .iter()
-                    .find_map(|(name, value)| (name == "vector_distance").then(|| value.clone()))
-                    .ok_or(Error::InvalidEntry)?;
-                let distance = parse_f64(&distance)?;
-                if 1.0 - distance < threshold {
-                    return Ok(None);
-                }
-                Ok(Some(response))
+                search_document(&connections, &index, &scope, vector, embedding.len())
             })
             .await
             .map_err(|_| Error::Unavailable)?
@@ -560,6 +437,108 @@ fn embedding_bytes(embedding: &[f32]) -> Vec<u8> {
         .collect()
 }
 
+fn write_document<C>(
+    connections: &Connections<C>,
+    index: &IndexState,
+    scope: &str,
+    prompt: &str,
+    response: Vec<u8>,
+    vector: Vec<u8>,
+    ttl: Option<Duration>,
+) -> Result<(), Error>
+where
+    C: redis::ConnectionLike + Send + 'static,
+{
+    let dimension = vector.len() / std::mem::size_of::<f32>();
+    ensure_index(
+        connections,
+        &index.name,
+        &index.prefix,
+        &index.dimension,
+        dimension,
+    )?;
+    let document = format!("{}{scope}:{}", index.prefix, Uuid::new_v4());
+    connections.execute(|connection| {
+        let mut pipeline = redis::pipe();
+        pipeline
+            .cmd("HSET")
+            .arg(&document)
+            .arg("litellm_cache_key")
+            .arg(scope)
+            .arg("prompt")
+            .arg(prompt)
+            .arg("response")
+            .arg(response)
+            .arg("embedding")
+            .arg(vector)
+            .ignore();
+        if let Some(ttl) = ttl {
+            pipeline
+                .cmd("EXPIRE")
+                .arg(&document)
+                .arg(ttl.as_secs())
+                .ignore();
+        }
+        pipeline
+            .query::<()>(connection)
+            .map_err(|_| Error::Unavailable)
+    })
+}
+
+fn search_document<C>(
+    connections: &Connections<C>,
+    index: &IndexState,
+    scope: &str,
+    vector: Vec<u8>,
+    dimension: usize,
+) -> Result<Option<Vec<u8>>, Error>
+where
+    C: redis::ConnectionLike + Send + 'static,
+{
+    ensure_index(
+        connections,
+        &index.name,
+        &index.prefix,
+        &index.dimension,
+        dimension,
+    )?;
+    let query =
+        format!("(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]");
+    let response = connections.execute(|connection| {
+        redis::cmd("FT.SEARCH")
+            .arg(&index.name)
+            .arg(query)
+            .arg("PARAMS")
+            .arg(2)
+            .arg("vec")
+            .arg(vector)
+            .arg("RETURN")
+            .arg(2)
+            .arg("response")
+            .arg("vector_distance")
+            .arg("DIALECT")
+            .arg(2)
+            .query::<redis::Value>(connection)
+            .map_err(|_| Error::Unavailable)
+    })?;
+    let Some(fields) = search_fields(response)? else {
+        return Ok(None);
+    };
+    let response = fields
+        .iter()
+        .find_map(|(name, value)| (name == "response").then(|| value.clone()))
+        .ok_or(Error::InvalidEntry)?;
+    let distance = fields
+        .iter()
+        .find_map(|(name, value)| (name == "vector_distance").then(|| value.clone()))
+        .ok_or(Error::InvalidEntry)?;
+    let distance = parse_f64(&distance)?;
+    if 1.0 - distance < index.similarity_threshold {
+        return Ok(None);
+    }
+    Ok(Some(response))
+}
+
 fn ensure_index<C>(
     connections: &Connections<C>,
     index_name: &str,
@@ -712,10 +691,16 @@ fn value_bytes(value: &redis::Value) -> Result<Vec<u8>, Error> {
 
 #[cfg(test)]
 mod tests {
-    use std::sync::{Arc, Mutex};
+    use std::{
+        collections::VecDeque,
+        sync::{Arc, Mutex},
+        time::Duration,
+    };
 
-    use litellm_cache::BaseCache;
-    use litellm_cache_response::ResponseCacheCodec;
+    use litellm_cache::{BaseCache, CacheCodec};
+    use litellm_cache_response::{
+        CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest,
+    };
     use redis_test::MockRedisConnection;
     use rstest::rstest;
     use serde_json::{Value, json};
@@ -732,6 +717,9 @@ mod tests {
     }
 
     type EmbedderCalls = Arc<Mutex<Vec<(String, Option<Value>)>>>;
+    type RecordingCache =
+        ValkeySemanticCache<FixedEmbedder, ResponseCacheCodec, RecordingConnection>;
+    type RecordingSetup = (RecordingCache, Arc<Mutex<Vec<Vec<u8>>>>, EmbedderCalls);
 
     impl Embedder for FixedEmbedder {
         fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, super::Error> {
@@ -751,6 +739,61 @@ mod tests {
         }
     }
 
+    struct RecordingConnection {
+        requests: Arc<Mutex<Vec<Vec<u8>>>>,
+        replies: Mutex<VecDeque<redis::RedisResult<redis::Value>>>,
+    }
+
+    impl RecordingConnection {
+        fn new(replies: impl IntoIterator<Item = redis::RedisResult<redis::Value>>) -> Self {
+            Self {
+                requests: Arc::default(),
+                replies: Mutex::new(replies.into_iter().collect()),
+            }
+        }
+
+        fn requests(&self) -> Arc<Mutex<Vec<Vec<u8>>>> {
+            Arc::clone(&self.requests)
+        }
+
+        fn reply(&self) -> redis::RedisResult<redis::Value> {
+            self.replies
+                .lock()
+                .unwrap()
+                .pop_front()
+                .unwrap_or_else(|| Ok(redis::Value::SimpleString("OK".into())))
+        }
+    }
+
+    impl redis::ConnectionLike for RecordingConnection {
+        fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult<redis::Value> {
+            self.requests.lock().unwrap().push(command.to_vec());
+            self.reply()
+        }
+
+        fn req_packed_commands(
+            &mut self,
+            command: &[u8],
+            _offset: usize,
+            count: usize,
+        ) -> redis::RedisResult<Vec<redis::Value>> {
+            self.requests.lock().unwrap().push(command.to_vec());
+            (0..count).map(|_| self.reply()).collect()
+        }
+
+        fn get_db(&self) -> i64 {
+            0
+        }
+
+        fn check_connection(&mut self) -> bool {
+            true
+        }
+
+        fn is_open(&self) -> bool {
+            true
+        }
+    }
+
     fn context(
         messages: Option<Value>,
         input: Option<Value>,
@@ -841,4 +884,302 @@ mod tests {
         assert_eq!(cache.get_cache("key", &context(None, None)).unwrap(), None);
         assert_eq!(cache.get_ttl(&context(None, None)), None);
     }
+
+    fn semantic_context(ttl: Option<Duration>) -> litellm_cache::SemanticCacheContext {
+        litellm_cache::SemanticCacheContext {
+            messages: Some(json!([{"role": "user", "content": "hello"}])),
+            metadata: Some(json!({"source": "test"})),
+            ttl,
+            ..Default::default()
+        }
+    }
+
+    fn cache_with_recording(
+        replies: impl IntoIterator<Item = redis::RedisResult<redis::Value>>,
+        vector: Vec<f32>,
+        threshold: f64,
+    ) -> RecordingSetup {
+        let connection = RecordingConnection::new(replies);
+        let requests = connection.requests();
+        let calls: EmbedderCalls = Arc::default();
+        let cache = ValkeySemanticCache::with_connection(
+            connection,
+            FixedEmbedder {
+                vector,
+                calls: Arc::clone(&calls),
+            },
+            ResponseCacheCodec,
+            ValkeySemanticConfig {
+                similarity_threshold: threshold,
+                index_name: "test".into(),
+            },
+        );
+        (cache, requests, calls)
+    }
+
+    fn ok() -> redis::RedisResult<redis::Value> {
+        Ok(redis::Value::SimpleString("OK".into()))
+    }
+
+    fn already_exists() -> redis::RedisResult<redis::Value> {
+        Err(redis::RedisError::from((
+            redis::ErrorKind::Io,
+            "already exists",
+        )))
+    }
+
+    fn info_dimension(dimension: usize) -> redis::Value {
+        redis::Value::Array(vec![
+            redis::Value::SimpleString("attributes".into()),
+            redis::Value::Array(vec![redis::Value::Array(vec![
+                redis::Value::SimpleString("embedding".into()),
+                redis::Value::Array(vec![
+                    redis::Value::SimpleString("dimensions".into()),
+                    redis::Value::Int(dimension as i64),
+                ]),
+            ])]),
+        ])
+    }
+
+    fn search_hit(response: Vec<u8>, distance: &str) -> redis::Value {
+        redis::Value::Array(vec![
+            redis::Value::Int(1),
+            redis::Value::BulkString(b"test:document".to_vec()),
+            redis::Value::Array(vec![
+                redis::Value::BulkString(b"response".to_vec()),
+                redis::Value::BulkString(response),
+                redis::Value::BulkString(b"vector_distance".to_vec()),
+                redis::Value::BulkString(distance.as_bytes().to_vec()),
+            ]),
+        ])
+    }
+
+    fn requests_text(requests: &Arc<Mutex<Vec<Vec<u8>>>>) -> String {
+        requests
+            .lock()
+            .unwrap()
+            .iter()
+            .map(|request| String::from_utf8_lossy(request))
+            .collect::<Vec<_>>()
+            .join("\n")
+    }
+
+    #[test]
+    fn set_without_ttl_writes_hset_without_expire() {
+        let (cache, requests, calls) = cache_with_recording([ok()], vec![1.0, 0.0], 0.8);
+        cache
+            .set_cache(
+                "key",
+                CacheEntry {
+                    timestamp: None,
+                    response: json!({"answer": "ok"}),
+                },
+                &semantic_context(None),
+            )
+            .unwrap();
+        let text = requests_text(&requests);
+        assert!(text.contains("FT.CREATE"));
+        assert!(text.contains("HSET"));
+        assert!(
+            text.contains("test:2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683:")
+        );
+        assert!(!text.contains("EXPIRE"));
+        assert_eq!(
+            *calls.lock().unwrap(),
+            vec![("hello".into(), Some(json!({"source": "test"})))]
+        );
+    }
+
+    #[test]
+    fn set_with_ttl_truncates_expire_seconds() {
+        let (cache, requests, _) = cache_with_recording([ok()], vec![1.0, 0.0], 0.8);
+        cache
+            .set_cache(
+                "key",
+                CacheEntry {
+                    timestamp: None,
+                    response: json!({"answer": "ok"}),
+                },
+                &semantic_context(Some(Duration::from_millis(1900))),
+            )
+            .unwrap();
+        let text = requests_text(&requests);
+        assert!(text.contains("EXPIRE"));
+        assert!(text.contains("\r\n$1\r\n1\r\n"));
+    }
+
+    #[test]
+    fn second_set_skips_create_after_dimension_is_cached() {
+        let (cache, requests, _) = cache_with_recording([ok()], vec![1.0, 0.0], 0.8);
+        let context = semantic_context(None);
+        let entry = CacheEntry {
+            timestamp: None,
+            response: json!({"answer": "ok"}),
+        };
+        cache.set_cache("key", entry.clone(), &context).unwrap();
+        cache.set_cache("key", entry, &context).unwrap();
+        let text = requests_text(&requests);
+        assert_eq!(text.matches("FT.CREATE").count(), 1);
+        assert_eq!(text.matches("HSET").count(), 2);
+    }
+
+    #[test]
+    fn existing_index_dimension_must_match_embedding() {
+        let (cache, _, _) = cache_with_recording(
+            [already_exists(), Ok(info_dimension(2))],
+            vec![1.0, 0.0],
+            0.8,
+        );
+        cache
+            .set_cache(
+                "key",
+                CacheEntry {
+                    timestamp: None,
+                    response: json!({"answer": "ok"}),
+                },
+                &semantic_context(None),
+            )
+            .unwrap();
+
+        let (cache, _, _) = cache_with_recording(
+            [already_exists(), Ok(info_dimension(3))],
+            vec![1.0, 0.0],
+            0.8,
+        );
+        assert_eq!(
+            cache.set_cache(
+                "key",
+                CacheEntry {
+                    timestamp: None,
+                    response: json!({"answer": "ok"}),
+                },
+                &semantic_context(None),
+            ),
+            Err(super::Error::Unavailable)
+        );
+    }
+
+    #[test]
+    fn get_applies_threshold_and_decodes_entry() {
+        let entry = CacheEntry {
+            timestamp: Some(1.0),
+            response: json!({"answer": "ok"}),
+        };
+        let encoded = ResponseCacheCodec.encode(&entry).unwrap();
+        let (cache, _, _) = cache_with_recording(
+            [ok(), Ok(search_hit(encoded.clone(), "0.1"))],
+            vec![1.0, 0.0],
+            0.8,
+        );
+        assert_eq!(
+            cache.get_cache("key", &semantic_context(None)).unwrap(),
+            Some(entry)
+        );
+
+        let (cache, _, _) =
+            cache_with_recording([ok(), Ok(search_hit(encoded, "0.5"))], vec![1.0, 0.0], 0.8);
+        assert_eq!(
+            cache.get_cache("key", &semantic_context(None)).unwrap(),
+            None
+        );
+    }
+
+    #[test]
+    fn get_zero_docs_is_a_miss() {
+        let (cache, _, _) = cache_with_recording(
+            [ok(), Ok(redis::Value::Array(vec![redis::Value::Int(0)]))],
+            vec![1.0, 0.0],
+            0.8,
+        );
+        assert_eq!(
+            cache.get_cache("key", &semantic_context(None)).unwrap(),
+            None
+        );
+    }
+
+    #[rstest]
+    #[case(redis::Value::Array(vec![
+        redis::Value::Int(1),
+        redis::Value::BulkString(b"document".to_vec()),
+        redis::Value::Array(vec![
+            redis::Value::BulkString(b"vector_distance".to_vec()),
+            redis::Value::BulkString(b"0.1".to_vec()),
+        ]),
+    ]))]
+    #[case(redis::Value::Array(vec![
+        redis::Value::Int(1),
+        redis::Value::BulkString(b"document".to_vec()),
+        redis::Value::Array(vec![
+            redis::Value::BulkString(b"response".to_vec()),
+            redis::Value::BulkString(b"not-json".to_vec()),
+            redis::Value::BulkString(b"vector_distance".to_vec()),
+            redis::Value::BulkString(b"abc".to_vec()),
+        ]),
+    ]))]
+    fn malformed_entries_are_invalid(#[case] search: redis::Value) {
+        let (cache, _, _) = cache_with_recording([ok(), Ok(search)], vec![1.0, 0.0], 0.8);
+        assert_eq!(
+            cache.get_cache("key", &semantic_context(None)),
+            Err(super::Error::InvalidEntry)
+        );
+    }
+
+    #[test]
+    fn response_cache_turns_invalid_entries_into_misses() {
+        let (cache, _, _) = cache_with_recording(
+            [
+                ok(),
+                Ok(redis::Value::Array(vec![
+                    redis::Value::Int(1),
+                    redis::Value::BulkString(b"document".to_vec()),
+                    redis::Value::Array(vec![
+                        redis::Value::BulkString(b"response".to_vec()),
+                        redis::Value::BulkString(b"not-json".to_vec()),
+                        redis::Value::BulkString(b"vector_distance".to_vec()),
+                        redis::Value::BulkString(b"0.1".to_vec()),
+                    ]),
+                ])),
+            ],
+            vec![1.0, 0.0],
+            0.8,
+        );
+        let service = ResponseCache::new(Arc::new(cache));
+        let request = ResponseCacheRequest {
+            key: CacheKeyInput {
+                preset: Some("key".into()),
+                ..Default::default()
+            },
+            context: semantic_context(None),
+            ..ResponseCacheRequest::new(CacheKeyInput::default())
+        };
+        assert_eq!(service.lookup(&request, Duration::ZERO).unwrap(), None);
+    }
+
+    #[tokio::test]
+    async fn async_set_and_get_use_shared_document_helpers() {
+        let entry = CacheEntry {
+            timestamp: Some(1.0),
+            response: json!({"answer": "ok"}),
+        };
+        let encoded = ResponseCacheCodec.encode(&entry).unwrap();
+        let (cache, requests, calls) = cache_with_recording(
+            [ok(), ok(), ok(), Ok(search_hit(encoded, "0.1"))],
+            vec![1.0, 0.0],
+            0.8,
+        );
+        let context = semantic_context(Some(Duration::from_millis(1900)));
+        cache
+            .async_set_cache("key", entry.clone(), context.clone())
+            .await
+            .unwrap();
+        assert_eq!(
+            cache.async_get_cache("key", &context).await.unwrap(),
+            Some(entry)
+        );
+        let text = requests_text(&requests);
+        assert!(text.contains("FT.CREATE"));
+        assert!(text.contains("HSET"));
+        assert!(text.contains("EXPIRE"));
+        assert_eq!(calls.lock().unwrap().len(), 2);
+    }
 }

litellm-rust/crates/python-bridge/src/cache/config.rs

@@ -276,25 +276,15 @@ fn project_redis(
 
     let client = backend.getattr("redis_client")?;
     let pool = client.getattr("connection_pool")?;
-    if !instance_class_is(&pool, "redis.connection", "ConnectionPool")? {
+    let Ok((resolved, is_tls)) = project_connection_pool(&pool)? else {
         return Ok(Err(UnsupportedCacheConfig::RedisConnection));
-    }
-    let resolved = pool.getattr("connection_kwargs")?.cast_into::<PyDict>()?;
+    };
     for key in ["credential_provider", "redis_connect_func"] {
         if has_value(&resolved, key)? {
             return Ok(Err(UnsupportedCacheConfig::RedisCredentials));
         }
     }
-    let connection_class = resolved
-        .get_item("connection_class")?
-        .unwrap_or(pool.getattr("connection_class")?);
-    let tls = if class_is(&connection_class, "redis.connection", "Connection")? {
-        None
-    } else if class_is(&connection_class, "redis.connection", "SSLConnection")? {
-        Some(project_tls(&resolved)?)
-    } else {
-        return Ok(Err(UnsupportedCacheConfig::RedisConnection));
-    };
+    let tls = is_tls.then(|| project_tls(&resolved)).transpose()?;
 
     let protocol = match optional_i64(&resolved, "protocol")?.unwrap_or(2) {
         2 => RedisProtocol::Resp2,
@@ -332,18 +322,9 @@ fn project_valkey_semantic(
 ) -> PyResult<Result<ValkeySemanticCacheConfig, UnsupportedCacheConfig>> {
     let client = backend.getattr("sync_client")?;
     let pool = client.getattr("connection_pool")?;
-    if !instance_class_is(&pool, "redis.connection", "ConnectionPool")? {
+    let Ok((resolved, _is_tls)) = project_connection_pool(&pool)? else {
         return Ok(Err(UnsupportedCacheConfig::RedisConnection));
-    }
-    let resolved = pool.getattr("connection_kwargs")?.cast_into::<PyDict>()?;
-    let connection_class = resolved
-        .get_item("connection_class")?
-        .unwrap_or(pool.getattr("connection_class")?);
-    if !class_is(&connection_class, "redis.connection", "Connection")?
-        && !class_is(&connection_class, "redis.connection", "SSLConnection")?
-    {
-        return Ok(Err(UnsupportedCacheConfig::RedisConnection));
-    }
+    };
     let connection = RedisConnectionConfig {
         host: required_string(&resolved, "host")?,
         port: u16::try_from(required_i64(&resolved, "port")?)
@@ -371,6 +352,27 @@ fn project_valkey_semantic(
     }))
 }
 
+#[inline(never)]
+fn project_connection_pool<'py>(
+    pool: &Bound<'py, PyAny>,
+) -> PyResult<Result<(Bound<'py, PyDict>, bool), UnsupportedCacheConfig>> {
+    if !instance_class_is(pool, "redis.connection", "ConnectionPool")? {
+        return Ok(Err(UnsupportedCacheConfig::RedisConnection));
+    }
+    let resolved = pool.getattr("connection_kwargs")?.cast_into::<PyDict>()?;
+    let connection_class = resolved
+        .get_item("connection_class")?
+        .unwrap_or(pool.getattr("connection_class")?);
+    let is_tls = if class_is(&connection_class, "redis.connection", "Connection")? {
+        false
+    } else if class_is(&connection_class, "redis.connection", "SSLConnection")? {
+        true
+    } else {
+        return Ok(Err(UnsupportedCacheConfig::RedisConnection));
+    };
+    Ok(Ok((resolved, is_tls)))
+}
+
 #[inline(never)]
 fn project_tls(values: &Bound<'_, PyDict>) -> PyResult<RedisTlsConfig> {
     Ok(RedisTlsConfig {
@@ -646,6 +648,40 @@ mod tests {
         });
     }
 
+    #[test]
+    fn projects_valkey_semantic_configuration() {
+        Python::initialize();
+        Python::attach(|py| {
+            let facade = facade(
+                py,
+                "pool = ConnectionPool()\n\
+                 pool.connection_class = Connection\n\
+                 pool.max_connections = 12\n\
+                 pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390, 'db': 2}\n\
+                 client = SimpleNamespace(connection_pool=pool)\n\
+                 backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\
+                 facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
+            );
+            let CacheConfigProjection::Native(config) =
+                NativeCacheConfig::project(&facade).unwrap()
+            else {
+                panic!("Valkey semantic cache should be supported");
+            };
+            let CacheBackendConfig::ValkeySemantic(valkey) = config.backend else {
+                panic!("expected Valkey semantic configuration");
+            };
+            assert_eq!(valkey.similarity_threshold, 0.85);
+            assert_eq!(valkey.index_name, "semantic_idx");
+            assert_eq!(valkey.embedding_model, "text-embedding-3-small");
+            assert_eq!(valkey.connection.host, "cache.internal");
+            assert_eq!(valkey.connection.port, 6390);
+            assert_eq!(valkey.connection.database, 2);
+            assert_eq!(valkey.connection.pool_size, 12);
+            assert_eq!(valkey.connection.protocol, RedisProtocol::Resp2);
+            assert!(valkey.connection.tls.is_none());
+        });
+    }
+
     #[test]
     fn dynamic_redis_auth_stays_on_python() {
         Python::initialize();

tests/test_litellm_rust/test_valkey_semantic_cache_native.py

@@ -1,4 +1,7 @@
+import hashlib
 import os
+import struct
+import time
 from collections.abc import Generator, Mapping
 from types import SimpleNamespace
 from typing import Final, cast
@@ -36,24 +39,32 @@ def index_name(valkey_url: str) -> Generator[str]:
         client.close()
 
 
-def _request() -> dict[str, object]:
+def _request(prompt: str = "semantic cache prompt") -> dict[str, object]:
     return {
         "key": {"preset": "key"},
-        "messages": [{"role": "user", "content": "semantic cache prompt"}],
+        "messages": [{"role": "user", "content": prompt}],
     }
 
 
-def _backend(url: str, index_name: str) -> ValkeySemanticCache:
+def _backend(
+    url: str,
+    index_name: str,
+    embeddings: Mapping[str, list[float]] | None = None,
+) -> ValkeySemanticCache:
+    vectors: Final = embeddings or {"semantic cache prompt": [1.0, 0.0]}
     backend: Final = ValkeySemanticCache(
         redis_url=url,
         similarity_threshold=0.8,
         index_name=index_name,
     )
-    backend._get_embedding = lambda prompt, metadata=None: [1.0, 0.0]
+
+    def embed(prompt: str, metadata: Mapping[str, object] | None = None) -> list[float]:
+        return vectors[prompt]
 
     async def async_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]:
-        return [1.0, 0.0]
+        return vectors[prompt]
 
+    backend._get_embedding = embed
     backend._get_async_embedding = async_embedding
     return backend
 
@@ -147,3 +158,122 @@ def test_batch_lookup_is_unsupported(
     binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
     with pytest.raises(NotImplementedError):
         binding.lookup_batch([_request()])
+
+
+def test_ttl_expiry(
+    valkey_url: str,
+    index_name: str,
+) -> None:
+    backend: Final = _backend(valkey_url, index_name)
+    handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
+    binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
+    binding.store({**_request(), "ttl_seconds": 1.0}, {"answer": "expires"})
+    client: Final = redis.Redis.from_url(valkey_url)
+    documents: Final = list(client.scan_iter(f"{index_name}:*"))
+    assert len(documents) == 1
+    assert client.ttl(documents[0]) > 0
+    time.sleep(1.5)
+    assert binding.lookup(_request()) is None
+
+
+def test_no_ttl_is_persistent_and_python_reads_native_value(
+    valkey_url: str,
+    index_name: str,
+) -> None:
+    backend: Final = _backend(valkey_url, index_name)
+    handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
+    binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
+    response: Final = {"answer": "persistent"}
+    binding.store(_request(), response)
+    client: Final = redis.Redis.from_url(valkey_url)
+    documents: Final = list(client.scan_iter(f"{index_name}:*"))
+    assert len(documents) == 1
+    assert client.ttl(documents[0]) == -1
+    cached: Final = cast(Mapping[str, object], backend.get_cache("key", messages=_request()["messages"]))
+    assert cached["response"] == response
+
+
+def test_below_threshold_misses_on_native_and_python(
+    valkey_url: str,
+    index_name: str,
+) -> None:
+    backend: Final = _backend(
+        valkey_url,
+        index_name,
+        {"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]},
+    )
+    handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
+    binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
+    binding.store(_request("prompt A"), {"answer": "A"})
+    assert binding.lookup(_request("prompt B")) is None
+    assert backend.get_cache("key", messages=_request("prompt B")["messages"]) is None
+
+
+def test_malformed_entry_is_a_miss_on_native_and_python(
+    valkey_url: str,
+    index_name: str,
+) -> None:
+    backend: Final = _backend(valkey_url, index_name)
+    client: Final = redis.Redis.from_url(valkey_url)
+    scope: Final = hashlib.sha256(b"key").hexdigest()
+    document: Final = f"{index_name}:{scope}:{uuid4().hex}"
+    client.hset(
+        document,
+        mapping={
+            "litellm_cache_key": scope,
+            "prompt": "semantic cache prompt",
+            "response": "not json",
+            "embedding": struct.pack("<2f", 1.0, 0.0),
+        },
+    )
+    handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
+    binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
+    assert binding.lookup(_request()) is None
+    assert backend.get_cache("key", messages=_request()["messages"]) is None
+
+
+async def test_async_store_batch_and_lookup(
+    valkey_url: str,
+    index_name: str,
+) -> None:
+    backend: Final = _backend(
+        valkey_url,
+        index_name,
+        {"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]},
+    )
+    handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
+    binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
+    requests: Final = [_request("prompt A"), _request("prompt B")]
+    responses: Final = [{"answer": "A"}, {"answer": "B"}]
+    await binding.async_store_batch(requests, responses)
+    assert await binding.async_lookup(requests[0]) == responses[0]
+    assert await binding.async_lookup(requests[1]) == responses[1]
+
+
+def test_subclass_backend_falls_back_to_python(
+    valkey_url: str,
+    index_name: str,
+) -> None:
+    class Custom(ValkeySemanticCache):
+        pass
+
+    facade: Final = Cache(
+        type=LiteLLMCacheType.VALKEY_SEMANTIC,
+        redis_url=valkey_url,
+        similarity_threshold=0.8,
+        valkey_semantic_cache_index_name=index_name,
+    )
+    facade.cache = Custom(redis_url=valkey_url, similarity_threshold=0.8, index_name=index_name)
+    resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade))
+    assert resolver.resolve().kind == "python_callback"
+
+
+async def test_ping_maps_unsupported_native_operation_to_not_implemented(
+    valkey_url: str,
+    index_name: str,
+) -> None:
+    backend: Final = _backend(valkey_url, index_name)
+    handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
+    binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
+    with pytest.raises(NotImplementedError):
+        await binding.ping()