mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
fix(cache-redis-semantic): harden index initialization
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee
parent
a9ff1de42b
commit
d2f8e8c835
2 changed files with 132 additions and 10 deletions
|
|
@ -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::<Vec<_>>();
|
||||
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"))
|
||||
})
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -124,7 +124,7 @@ fn index_info(attributes: Vec<redis::Value>) -> 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::<redis::Value, _>(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::<redis::Value, _>(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];
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue