From d2f8e8c83504d0662bf462fa54d3828d4dcc426f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:21:57 +0000 Subject: [PATCH] fix(cache-redis-semantic): harden index initialization Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/cache-redis-semantic/src/cache.rs | 24 +++- .../cache-redis-semantic/tests/cache.rs | 118 +++++++++++++++++- 2 files changed, 132 insertions(+), 10 deletions(-) diff --git a/litellm-rust/crates/cache-redis-semantic/src/cache.rs b/litellm-rust/crates/cache-redis-semantic/src/cache.rs index b1440e80de5..cd79f067296 100644 --- a/litellm-rust/crates/cache-redis-semantic/src/cache.rs +++ b/litellm-rust/crates/cache-redis-semantic/src/cache.rs @@ -71,7 +71,11 @@ impl Inner { Some(true) => self.index_name.clone(), Some(false) => self.isolated_index(connection, dims)?, None => { - create_index(connection, &self.index_name, dims)?; + if create_index(connection, &self.index_name, dims).is_err() + && index_compatible(connection, &self.index_name, dims)? != Some(true) + { + return Err(Error::Unavailable); + } self.index_name.clone() } }; @@ -509,36 +513,46 @@ fn schema_compatible(info: &redis::Value, dims: usize) -> bool { .iter() .map(|attribute| { let redis::Value::Array(attribute) = attribute else { - return (None, None, None); + return (None, None, None, None, None); }; let mut name = None; let mut field_type = None; let mut dim = None; + let mut data_type = None; + let mut distance_metric = None; for pair in attribute.as_chunks::<2>().0 { match string_value(&pair[0]).as_deref() { Some("identifier") => name = string_value(&pair[1]), Some("type") => field_type = string_value(&pair[1]), Some("dim") => dim = number_value(&pair[1]), + Some("data_type") => data_type = string_value(&pair[1]), + Some("distance_metric") => distance_metric = string_value(&pair[1]), _ => {} } } - (name, field_type, dim) + (name, field_type, dim, data_type, distance_metric) }) .collect::>(); let has_field = |name: &str, field_type: &str| { fields .iter() - .any(|(n, t, _)| n.as_deref() == Some(name) && t.as_deref() == Some(field_type)) + .any(|(n, t, ..)| n.as_deref() == Some(name) && t.as_deref() == Some(field_type)) }; has_field("prompt", "TEXT") && has_field("response", "TEXT") && has_field("inserted_at", "NUMERIC") && has_field("updated_at", "NUMERIC") && has_field(CACHE_KEY_FIELD, "TAG") - && fields.iter().any(|(n, t, d)| { + && fields.iter().any(|(n, t, d, data, metric)| { n.as_deref() == Some(VECTOR_FIELD) && t.as_deref() == Some("VECTOR") && *d == Some(dims as f64) + && data + .as_deref() + .is_some_and(|data| data.eq_ignore_ascii_case("float32")) + && metric + .as_deref() + .is_some_and(|metric| metric.eq_ignore_ascii_case("cosine")) }) } diff --git a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs index 77b057ae3b9..85a35a033b2 100644 --- a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs +++ b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs @@ -124,7 +124,7 @@ fn index_info(attributes: Vec) -> redis::Value { ]) } -fn vector_attribute(dims: i64) -> redis::Value { +fn vector_attribute_with(dims: i64, data_type: &str, distance_metric: &str) -> redis::Value { attribute( "prompt_vector", "VECTOR", @@ -132,26 +132,34 @@ fn vector_attribute(dims: i64) -> redis::Value { s("algorithm"), s("FLAT"), s("data_type"), - s("FLOAT32"), + s(data_type), s("dim"), redis::Value::Int(dims), s("distance_metric"), - s("COSINE"), + s(distance_metric), ], ) } -fn compatible_info(dims: i64) -> redis::Value { +fn vector_attribute(dims: i64) -> redis::Value { + vector_attribute_with(dims, "FLOAT32", "COSINE") +} + +fn info_with_vector(vector: redis::Value) -> redis::Value { index_info(vec![ attribute("prompt", "TEXT", vec![]), attribute("response", "TEXT", vec![]), attribute("inserted_at", "NUMERIC", vec![]), attribute("updated_at", "NUMERIC", vec![]), - vector_attribute(dims), + vector, attribute("litellm_cache_key", "TAG", vec![]), ]) } +fn compatible_info(dims: i64) -> redis::Value { + info_with_vector(vector_attribute(dims)) +} + fn unscoped_info(dims: i64) -> redis::Value { index_info(vec![ attribute("prompt", "TEXT", vec![]), @@ -537,6 +545,106 @@ fn incompatible_schema_falls_back_to_isolated_index() { .unwrap(); } +#[test] +fn create_index_race_rechecks_schema_and_stores() { + let prompt = "hello prompt"; + let tag = "key1"; + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Err::(unknown_index_error()), + ), + MockCmd::new( + create_index_command(INDEX, 3), + Err::<&str, _>(redis::RedisError::from(( + redis::ErrorKind::Extension, + "Index already exists", + ))), + ), + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{INDEX}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + tag, + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn wrong_distance_metric_falls_back_to_isolated_index() { + let prompt = "hello prompt"; + let tag = "key1"; + let isolated = format!("{INDEX}_isolated"); + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Ok(info_with_vector(vector_attribute_with(3, "FLOAT32", "L2"))), + ), + MockCmd::new( + redis::cmd("FT.INFO").arg(&isolated), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(&isolated, 3), Ok("OK")), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{isolated}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + tag, + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + #[test] fn tag_special_characters_are_escaped_in_search_filter() { let vector = vec![0.1f32, 0.2, 0.3];