mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
feat(rust): add redis-client crate for standalone Redis (#48269)
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
c905cd065d
commit
6edbd27f74
24 changed files with 851 additions and 15 deletions
48
litellm-rust/Cargo.lock
generated
48
litellm-rust/Cargo.lock
generated
|
|
@ -876,6 +876,15 @@ dependencies = [
|
|||
"time",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "backon"
|
||||
version = "1.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef"
|
||||
dependencies = [
|
||||
"fastrand",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "base16ct"
|
||||
version = "0.2.0"
|
||||
|
|
@ -1248,7 +1257,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "cfc320937d09e6de266b31b9afb480f197d7a861be86be7cb2ea7e5d1bfffc5e"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"futures-core",
|
||||
"memchr",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
"tokio-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -4282,6 +4295,26 @@ dependencies = [
|
|||
"thiserror 2.0.19",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-redis-client"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-redis-testing",
|
||||
"redis",
|
||||
"rstest",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-redis-testing"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"testcontainers-modules",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-router"
|
||||
version = "0.1.0"
|
||||
|
|
@ -5757,22 +5790,33 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "redis"
|
||||
version = "1.7.0"
|
||||
version = "1.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2acbc41a996f7652b2ddd9dfd98cc4ff602cfd742ae35382f07f608405ab50ed"
|
||||
checksum = "7acd809e80afa754c63f586955699d8990efe1c4b6613183ee5f027a11fbc8b5"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arcstr",
|
||||
"async-lock",
|
||||
"backon",
|
||||
"bytes",
|
||||
"cfg-if",
|
||||
"combine",
|
||||
"crc16",
|
||||
"futures-channel",
|
||||
"futures-util",
|
||||
"itoa",
|
||||
"num-bigint 0.5.1",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"rand 0.10.2",
|
||||
"rustls 0.23.42",
|
||||
"rustls-native-certs",
|
||||
"ryu",
|
||||
"sha1_smol",
|
||||
"socket2 0.6.5",
|
||||
"tokio",
|
||||
"tokio-rustls 0.26.4",
|
||||
"tokio-util",
|
||||
"url",
|
||||
"xxhash-rust",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -69,6 +69,8 @@ litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface"
|
|||
litellm-token-counter-tiktoken = { path = "crates/token-counter-tiktoken" }
|
||||
litellm-host-python = { path = "crates/host-python" }
|
||||
litellm-python-compat = { path = "crates/python-compat" }
|
||||
litellm-redis-client = { path = "crates/redis-client" }
|
||||
litellm-redis-testing = { path = "crates/redis-testing" }
|
||||
|
||||
askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] }
|
||||
tracing = "0.1"
|
||||
|
|
@ -88,6 +90,9 @@ rand = "0.8"
|
|||
macro_rules_attribute = "0.2.3"
|
||||
schemars = "1"
|
||||
reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] }
|
||||
redis = "1.8.0"
|
||||
redis-test = "1.0.4"
|
||||
testcontainers-modules = "0.15.0"
|
||||
qdrant-client = { version = "1.19.0", default-features = false }
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
rstest = "0.26.1"
|
||||
|
|
|
|||
|
|
@ -8,12 +8,12 @@ repository.workspace = true
|
|||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
litellm-cache-redis.workspace = true
|
||||
redis = { version = "1.7.0", features = ["tls-rustls"] }
|
||||
redis = { workspace = true, features = ["tls-rustls"] }
|
||||
sha2.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-cache-testing.workspace = true
|
||||
redis-test = "1.0.4"
|
||||
redis-test.workspace = true
|
||||
rstest.workspace = true
|
||||
serde_json.workspace = true
|
||||
tokio.workspace = true
|
||||
|
|
|
|||
|
|
@ -7,12 +7,12 @@ repository.workspace = true
|
|||
|
||||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
redis = { version = "1.7.0", features = ["cluster", "tls-rustls"] }
|
||||
redis = { workspace = true, features = ["cluster", "tls-rustls"] }
|
||||
r2d2 = "0.8.10"
|
||||
tokio.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-cache-testing.workspace = true
|
||||
redis-test = "1.0.4"
|
||||
redis-test.workspace = true
|
||||
rstest.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
|
|
|||
|
|
@ -17,8 +17,8 @@ litellm-cache-gcs.workspace = true
|
|||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
litellm-cache-memory.workspace = true
|
||||
litellm-cache-redis.workspace = true
|
||||
redis = "1.7.0"
|
||||
redis-test = "1.0.4"
|
||||
redis.workspace = true
|
||||
redis-test.workspace = true
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
wiremock.workspace = true
|
||||
|
|
|
|||
|
|
@ -8,13 +8,13 @@ repository.workspace = true
|
|||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
litellm-cache-redis.workspace = true
|
||||
redis = { version = "1.7.0", features = ["tls-rustls"] }
|
||||
redis = { workspace = true, features = ["tls-rustls"] }
|
||||
sha2.workspace = true
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-cache-testing.workspace = true
|
||||
redis-test = "1.0.4"
|
||||
redis-test.workspace = true
|
||||
rstest.workspace = true
|
||||
serde_json.workspace = true
|
||||
tokio.workspace = true
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ repository.workspace = true
|
|||
|
||||
[dependencies]
|
||||
sqlx = { workspace = true, features = ["migrate"] }
|
||||
testcontainers-modules = { version = "0.15.0", features = ["postgres"] }
|
||||
testcontainers-modules = { workspace = true, features = ["postgres"] }
|
||||
thiserror.workspace = true
|
||||
tokio.workspace = true
|
||||
|
||||
|
|
|
|||
|
|
@ -53,7 +53,7 @@ pyo3.workspace = true
|
|||
prost.workspace = true
|
||||
pyo3-async-runtimes.workspace = true
|
||||
reqwest.workspace = true
|
||||
redis = { version = "1.7.0", features = ["tls-rustls"] }
|
||||
redis = { workspace = true, features = ["tls-rustls"] }
|
||||
serde_json.workspace = true
|
||||
strum.workspace = true
|
||||
veil.workspace = true
|
||||
|
|
|
|||
26
litellm-rust/crates/redis-client/AGENTS.md
Normal file
26
litellm-rust/crates/redis-client/AGENTS.md
Normal file
|
|
@ -0,0 +1,26 @@
|
|||
- The one async Redis client for the gateway and the Rust cache crates. It owns connecting, timeouts, error classes and the key namespace
|
||||
- It knows nothing about domains, so key builders, TTL policy, value formats and fallback decisions stay in the crate that owns the state. Every write takes its TTL from the caller, and there is no default
|
||||
- Async only. Build one `RedisClient` per Redis instance at the process boundary and clone it. Never add a pool, a sync API, `spawn_blocking`, or a client built inside a route
|
||||
- Set every redis-rs timeout and retry explicitly. Its defaults (500ms response, 1s connect, 6 reconnect retries) are not the Python proxy's (5s socket timeout from `REDIS_SOCKET_TIMEOUT`)
|
||||
- Resolving config from `config.yaml` / `REDIS_*` env the way `litellm/_redis.py` does belongs to the consumer that reads that config, not here. `RedisConfig` takes already-resolved values
|
||||
|
||||
## Errors
|
||||
- `Transport`, `Server`, `Decode`, each wrapping the redis-rs error as `#[source]`. `ConnectError` is separate because only `connect` returns it
|
||||
- Classify in `From<RedisError> for Error` by `ErrorKind`. LOADING, CLUSTERDOWN, MASTERDOWN, TRYAGAIN and READONLY are typed server replies, but they mean the node can't serve right now, so they are `Transport`. Codes redis-rs doesn't type (WRONGTYPE) come back as `ErrorKind::Extension` and are `Server`. Each code gets its own rstest case built from a parsed RESP error
|
||||
- A `Server` error never drops the connection
|
||||
- `Transport` includes timeouts, where the command may still have applied. Callers of counters and locks treat it as an unknown outcome
|
||||
- Never swallow an error or turn it into a miss. The caller picks the fallback
|
||||
- Reads return `Option<T>`: a missing key is `None`, and a value that doesn't decode to `T` (`"15.0"` as `i64`) is `Decode`
|
||||
|
||||
## Namespace
|
||||
- Applied once, here, to every key. A key that already starts with `{namespace}:` is left alone, matching `RedisCache.check_and_fix_namespace`
|
||||
- The rule never changes, because Python and Rust pods share keys in the same Redis
|
||||
|
||||
## Commands
|
||||
- One typed method per command, built with `redis::cmd`. Empty batches (MGET, DEL) return without a round trip, since Redis rejects an empty MGET
|
||||
- TTLs are `Duration` and are sent in milliseconds (`PX`, `PEXPIRE`, `PTTL`)
|
||||
|
||||
## Tests
|
||||
- Against real Redis from `litellm-redis-testing`, never a mocked wire or asserted command bytes
|
||||
- Assert values and Redis state. Outages come from the container (`restart`, `freeze` / `thaw`), not `DEBUG`, which the official image disables
|
||||
- Use a short client timeout in outage tests so they stay fast. Never sleep to wait for a state; poll with a deadline
|
||||
15
litellm-rust/crates/redis-client/Cargo.toml
Normal file
15
litellm-rust/crates/redis-client/Cargo.toml
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
[package]
|
||||
name = "litellm-redis-client"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
redis = { workspace = true, features = ["connection-manager", "tokio-rustls-comp"] }
|
||||
thiserror.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-redis-testing.workspace = true
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
404
litellm-rust/crates/redis-client/src/client.rs
Normal file
404
litellm-rust/crates/redis-client/src/client.rs
Normal file
|
|
@ -0,0 +1,404 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use redis::{
|
||||
Cmd, FromRedisValue, ToRedisArgs, ToSingleRedisArg,
|
||||
aio::{ConnectionManager, ConnectionManagerConfig},
|
||||
};
|
||||
|
||||
use crate::{ConnectError, Error, Namespace, RedisConfig};
|
||||
|
||||
const RECONNECT_RETRIES: usize = 3;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct RedisClient {
|
||||
connection: ConnectionManager,
|
||||
namespace: Namespace,
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub enum Ttl {
|
||||
Missing,
|
||||
Persistent,
|
||||
Expires(Duration),
|
||||
}
|
||||
|
||||
impl RedisClient {
|
||||
pub async fn connect(config: RedisConfig) -> Result<Self, ConnectError> {
|
||||
let client = redis::Client::open(config.endpoint.connection_info()?)?;
|
||||
let manager_config = ConnectionManagerConfig::new()
|
||||
.set_connection_timeout(Some(config.timeout))
|
||||
.set_response_timeout(Some(config.timeout))
|
||||
.set_number_of_retries(RECONNECT_RETRIES);
|
||||
Ok(Self {
|
||||
connection: client
|
||||
.get_connection_manager_with_config(manager_config)
|
||||
.await?,
|
||||
namespace: config.namespace,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn get<T: FromRedisValue>(&self, key: &str) -> Result<Option<T>, Error> {
|
||||
self.query(redis::cmd("GET").arg(self.key(key))).await
|
||||
}
|
||||
|
||||
pub async fn mget<T: FromRedisValue>(&self, keys: &[&str]) -> Result<Vec<Option<T>>, Error> {
|
||||
if keys.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.query(redis::cmd("MGET").arg(self.keys(keys))).await
|
||||
}
|
||||
|
||||
pub async fn set(
|
||||
&self,
|
||||
key: &str,
|
||||
value: impl ToSingleRedisArg,
|
||||
ttl: Duration,
|
||||
) -> Result<(), Error> {
|
||||
self.query(
|
||||
redis::cmd("SET")
|
||||
.arg(self.key(key))
|
||||
.arg(value)
|
||||
.arg("PX")
|
||||
.arg(millis(ttl)),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn set_nx(
|
||||
&self,
|
||||
key: &str,
|
||||
value: impl ToSingleRedisArg,
|
||||
ttl: Duration,
|
||||
) -> Result<bool, Error> {
|
||||
let reply: Option<String> = self
|
||||
.query(
|
||||
redis::cmd("SET")
|
||||
.arg(self.key(key))
|
||||
.arg(value)
|
||||
.arg("NX")
|
||||
.arg("PX")
|
||||
.arg(millis(ttl)),
|
||||
)
|
||||
.await?;
|
||||
Ok(reply.is_some())
|
||||
}
|
||||
|
||||
pub async fn del(&self, keys: &[&str]) -> Result<usize, Error> {
|
||||
if keys.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
self.query(redis::cmd("DEL").arg(self.keys(keys))).await
|
||||
}
|
||||
|
||||
pub async fn expire(&self, key: &str, ttl: Duration) -> Result<bool, Error> {
|
||||
self.query(redis::cmd("PEXPIRE").arg(self.key(key)).arg(millis(ttl)))
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn ttl(&self, key: &str) -> Result<Ttl, Error> {
|
||||
let reply: i64 = self.query(redis::cmd("PTTL").arg(self.key(key))).await?;
|
||||
Ok(match reply {
|
||||
-2 => Ttl::Missing,
|
||||
-1 => Ttl::Persistent,
|
||||
millis => Ttl::Expires(Duration::from_millis(millis.unsigned_abs())),
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn incr_by(&self, key: &str, delta: i64) -> Result<i64, Error> {
|
||||
self.query(redis::cmd("INCRBY").arg(self.key(key)).arg(delta))
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn incr_by_float(&self, key: &str, delta: f64) -> Result<f64, Error> {
|
||||
self.query(redis::cmd("INCRBYFLOAT").arg(self.key(key)).arg(delta))
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn rpush(&self, key: &str, values: impl ToRedisArgs) -> Result<usize, Error> {
|
||||
self.query(redis::cmd("RPUSH").arg(self.key(key)).arg(values))
|
||||
.await
|
||||
}
|
||||
|
||||
fn key(&self, key: &str) -> String {
|
||||
self.namespace.apply(key)
|
||||
}
|
||||
|
||||
fn keys(&self, keys: &[&str]) -> Vec<String> {
|
||||
keys.iter().map(|key| self.key(key)).collect()
|
||||
}
|
||||
|
||||
async fn query<T: FromRedisValue>(&self, command: &Cmd) -> Result<T, Error> {
|
||||
Ok(command.query_async(&mut self.connection.clone()).await?)
|
||||
}
|
||||
}
|
||||
|
||||
fn millis(duration: Duration) -> u64 {
|
||||
u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use litellm_redis_testing::RedisServer;
|
||||
use rstest::rstest;
|
||||
|
||||
use super::{RedisClient, Ttl};
|
||||
use crate::{DEFAULT_TIMEOUT, Endpoint, Error, HostEndpoint, Namespace, RedisConfig};
|
||||
|
||||
async fn client(server: &RedisServer, namespace: &str, timeout: Duration) -> RedisClient {
|
||||
RedisClient::connect(RedisConfig {
|
||||
endpoint: Endpoint::Url(server.url().to_owned()),
|
||||
namespace: Namespace::new(namespace),
|
||||
timeout,
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn client_id(client: &RedisClient) -> i64 {
|
||||
redis::cmd("CLIENT")
|
||||
.arg("ID")
|
||||
.query_async(&mut client.connection.clone())
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn namespaces_isolate_keys_on_the_server() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let a = client(&server, "a", DEFAULT_TIMEOUT).await;
|
||||
let b = client(&server, "b", DEFAULT_TIMEOUT).await;
|
||||
let bare = client(&server, "", DEFAULT_TIMEOUT).await;
|
||||
|
||||
a.set("x", "1", Duration::from_secs(60)).await.unwrap();
|
||||
|
||||
assert_eq!(a.get::<String>("x").await.unwrap().as_deref(), Some("1"));
|
||||
assert_eq!(a.get::<String>("a:x").await.unwrap().as_deref(), Some("1"));
|
||||
assert_eq!(
|
||||
bare.get::<String>("a:x").await.unwrap().as_deref(),
|
||||
Some("1")
|
||||
);
|
||||
assert_eq!(b.get::<String>("x").await.unwrap(), None);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_rejected_command_keeps_the_connection() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let client = client(&server, "", DEFAULT_TIMEOUT).await;
|
||||
client.rpush("list", "item").await.unwrap();
|
||||
let before = client_id(&client).await;
|
||||
|
||||
let rejected = client.incr_by("list", 1).await;
|
||||
|
||||
assert!(matches!(rejected, Err(Error::Server(_))), "{rejected:?}");
|
||||
assert_eq!(client.incr_by("counter", 2).await.unwrap(), 2);
|
||||
assert_eq!(client_id(&client).await, before);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_reply_of_another_type_is_a_decode_error_and_a_missing_key_is_none() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let client = client(&server, "", DEFAULT_TIMEOUT).await;
|
||||
client
|
||||
.set("float", "15.0", Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let decoded = client.get::<i64>("float").await;
|
||||
|
||||
assert!(matches!(decoded, Err(Error::Decode(_))), "{decoded:?}");
|
||||
assert_eq!(client.get::<i64>("missing").await.unwrap(), None);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_timed_out_command_does_not_leak_its_reply_into_the_next_one() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let client = client(&server, "", Duration::from_millis(500)).await;
|
||||
client
|
||||
.set("first", "1", Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
client
|
||||
.set("second", "2", Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
server.freeze().await.unwrap();
|
||||
let timed_out = client.get::<String>("first").await;
|
||||
server.thaw().await.unwrap();
|
||||
|
||||
assert!(
|
||||
matches!(timed_out, Err(Error::Transport(_))),
|
||||
"{timed_out:?}"
|
||||
);
|
||||
assert_eq!(
|
||||
client.get::<String>("second").await.unwrap().as_deref(),
|
||||
Some("2")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn the_same_client_recovers_after_redis_restarts() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let client = client(&server, "", Duration::from_millis(500)).await;
|
||||
|
||||
server.restart().await.unwrap();
|
||||
|
||||
let deadline = Instant::now() + Duration::from_secs(30);
|
||||
let recovered = loop {
|
||||
match client.incr_by("after_restart", 1).await {
|
||||
Ok(value) => break Some(value),
|
||||
Err(_) if Instant::now() < deadline => {
|
||||
tokio::time::sleep(Duration::from_millis(200)).await
|
||||
}
|
||||
Err(_) => break None,
|
||||
}
|
||||
};
|
||||
assert_eq!(recovered, Some(1));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn empty_batches_succeed_without_asking_redis() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let client = client(&server, "", DEFAULT_TIMEOUT).await;
|
||||
|
||||
assert_eq!(
|
||||
client.mget::<String>(&[]).await.unwrap(),
|
||||
Vec::<Option<String>>::new()
|
||||
);
|
||||
assert_eq!(client.del(&[]).await.unwrap(), 0);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn mget_returns_values_in_input_order_with_gaps() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let client = client(&server, "ns", DEFAULT_TIMEOUT).await;
|
||||
client.set("a", "1", Duration::from_secs(60)).await.unwrap();
|
||||
client.set("c", "3", Duration::from_secs(60)).await.unwrap();
|
||||
|
||||
let values = client.mget::<String>(&["c", "b", "a"]).await.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
values,
|
||||
vec![Some("3".to_owned()), None, Some("1".to_owned())]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn set_nx_only_writes_a_missing_key() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let client = client(&server, "", DEFAULT_TIMEOUT).await;
|
||||
|
||||
let first = client
|
||||
.set_nx("lock", "owner-1", Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
let second = client
|
||||
.set_nx("lock", "owner-2", Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(first);
|
||||
assert!(!second);
|
||||
assert_eq!(
|
||||
client.get::<String>("lock").await.unwrap().as_deref(),
|
||||
Some("owner-1")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn ttl_reports_missing_persistent_and_remaining_time() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let client = client(&server, "", DEFAULT_TIMEOUT).await;
|
||||
client.incr_by("persistent", 1).await.unwrap();
|
||||
client
|
||||
.set("expiring", "1", Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
client.incr_by_float("expired_later", 0.5).await.unwrap();
|
||||
assert!(
|
||||
client
|
||||
.expire("expired_later", Duration::from_secs(30))
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
|
||||
assert_eq!(client.ttl("missing").await.unwrap(), Ttl::Missing);
|
||||
assert_eq!(client.ttl("persistent").await.unwrap(), Ttl::Persistent);
|
||||
assert!(matches!(
|
||||
client.ttl("expiring").await.unwrap(),
|
||||
Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60)
|
||||
));
|
||||
assert!(matches!(
|
||||
client.ttl("expired_later").await.unwrap(),
|
||||
Ttl::Expires(left) if left > Duration::from_secs(20) && left <= Duration::from_secs(30)
|
||||
));
|
||||
assert!(
|
||||
!client
|
||||
.expire("missing", Duration::from_secs(30))
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn incr_by_float_accumulates_fractions() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let client = client(&server, "", DEFAULT_TIMEOUT).await;
|
||||
|
||||
client.incr_by_float("spend", 0.25).await.unwrap();
|
||||
|
||||
assert_eq!(client.incr_by_float("spend", 0.5).await.unwrap(), 0.75);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_host_endpoint_connects_to_its_database() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let port = server.url().rsplit(':').next().unwrap().parse().unwrap();
|
||||
let in_db_3 = RedisClient::connect(RedisConfig {
|
||||
endpoint: Endpoint::Host(HostEndpoint {
|
||||
host: "127.0.0.1".to_owned(),
|
||||
port,
|
||||
username: None,
|
||||
password: None,
|
||||
tls: false,
|
||||
db: 3,
|
||||
}),
|
||||
namespace: Namespace::default(),
|
||||
timeout: DEFAULT_TIMEOUT,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let in_db_0 = client(&server, "", DEFAULT_TIMEOUT).await;
|
||||
|
||||
in_db_3
|
||||
.set("k", "3", Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(in_db_0.get::<String>("k").await.unwrap(), None);
|
||||
let via_url = RedisClient::connect(RedisConfig {
|
||||
endpoint: Endpoint::Url(format!("{}/3", server.url())),
|
||||
namespace: Namespace::default(),
|
||||
timeout: DEFAULT_TIMEOUT,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
via_url.get::<String>("k").await.unwrap().as_deref(),
|
||||
Some("3")
|
||||
);
|
||||
}
|
||||
}
|
||||
60
litellm-rust/crates/redis-client/src/config.rs
Normal file
60
litellm-rust/crates/redis-client/src/config.rs
Normal file
|
|
@ -0,0 +1,60 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use redis::{ConnectionAddr, ConnectionInfo, IntoConnectionInfo, RedisConnectionInfo, RedisResult};
|
||||
|
||||
use crate::Namespace;
|
||||
|
||||
pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
|
||||
pub struct RedisConfig {
|
||||
pub endpoint: Endpoint,
|
||||
pub namespace: Namespace,
|
||||
pub timeout: Duration,
|
||||
}
|
||||
|
||||
pub enum Endpoint {
|
||||
Url(String),
|
||||
Host(HostEndpoint),
|
||||
}
|
||||
|
||||
pub struct HostEndpoint {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
pub username: Option<String>,
|
||||
pub password: Option<String>,
|
||||
pub tls: bool,
|
||||
pub db: i64,
|
||||
}
|
||||
|
||||
impl Endpoint {
|
||||
pub(crate) fn connection_info(&self) -> RedisResult<ConnectionInfo> {
|
||||
match self {
|
||||
Self::Url(url) => url.as_str().into_connection_info(),
|
||||
Self::Host(host) => host.connection_info(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl HostEndpoint {
|
||||
fn connection_info(&self) -> RedisResult<ConnectionInfo> {
|
||||
let addr = match self.tls {
|
||||
true => ConnectionAddr::TcpTls {
|
||||
host: self.host.clone(),
|
||||
port: self.port,
|
||||
insecure: false,
|
||||
tls_params: None,
|
||||
},
|
||||
false => ConnectionAddr::Tcp(self.host.clone(), self.port),
|
||||
};
|
||||
let settings = RedisConnectionInfo::default().set_db(self.db);
|
||||
let settings = match &self.username {
|
||||
Some(username) => settings.set_username(username),
|
||||
None => settings,
|
||||
};
|
||||
let settings = match &self.password {
|
||||
Some(password) => settings.set_password(password),
|
||||
None => settings,
|
||||
};
|
||||
Ok(addr.into_connection_info()?.set_redis_settings(settings))
|
||||
}
|
||||
}
|
||||
67
litellm-rust/crates/redis-client/src/error.rs
Normal file
67
litellm-rust/crates/redis-client/src/error.rs
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
use redis::{ErrorKind, RedisError, ServerErrorKind};
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum Error {
|
||||
#[error("Redis did not answer, so the command may or may not have applied")]
|
||||
Transport(#[source] RedisError),
|
||||
#[error("Redis rejected the command")]
|
||||
Server(#[source] RedisError),
|
||||
#[error("Redis replied with a value of a different type")]
|
||||
Decode(#[source] RedisError),
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[error("connecting to Redis")]
|
||||
pub struct ConnectError(#[from] RedisError);
|
||||
|
||||
impl From<RedisError> for Error {
|
||||
fn from(error: RedisError) -> Self {
|
||||
match error.kind() {
|
||||
ErrorKind::Server(
|
||||
ServerErrorKind::BusyLoading
|
||||
| ServerErrorKind::ClusterDown
|
||||
| ServerErrorKind::MasterDown
|
||||
| ServerErrorKind::TryAgain
|
||||
| ServerErrorKind::ReadOnly,
|
||||
) => Self::Transport(error),
|
||||
ErrorKind::Server(_) | ErrorKind::Extension => Self::Server(error),
|
||||
ErrorKind::Parse | ErrorKind::UnexpectedReturnType => Self::Decode(error),
|
||||
_ => Self::Transport(error),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use redis::{RedisError, Value};
|
||||
use rstest::rstest;
|
||||
|
||||
use super::Error;
|
||||
|
||||
fn reply(resp: &str) -> RedisError {
|
||||
match redis::parse_redis_value(resp.as_bytes()).unwrap() {
|
||||
Value::ServerError(error) => error.into(),
|
||||
other => panic!("{resp:?} parsed as {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::wrong_type("-WRONGTYPE Operation against a key holding the wrong kind of value\r\n")]
|
||||
#[case::no_script("-NOSCRIPT No matching script\r\n")]
|
||||
#[case::cross_slot("-CROSSSLOT Keys in request don't hash to the same slot\r\n")]
|
||||
#[case::generic("-ERR value is not an integer or out of range\r\n")]
|
||||
#[case::no_perm("-NOPERM this user has no permissions to access one of the keys\r\n")]
|
||||
fn rejections_are_server_errors(#[case] resp: &str) {
|
||||
assert!(matches!(Error::from(reply(resp)), Error::Server(_)));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::loading("-LOADING Redis is loading the dataset in memory\r\n")]
|
||||
#[case::cluster_down("-CLUSTERDOWN The cluster is down\r\n")]
|
||||
#[case::master_down("-MASTERDOWN Link with MASTER is down\r\n")]
|
||||
#[case::try_again("-TRYAGAIN Multiple keys request during rehashing of slot\r\n")]
|
||||
#[case::read_only("-READONLY You can't write against a read only replica.\r\n")]
|
||||
fn unavailable_nodes_are_transport_errors(#[case] resp: &str) {
|
||||
assert!(matches!(Error::from(reply(resp)), Error::Transport(_)));
|
||||
}
|
||||
}
|
||||
9
litellm-rust/crates/redis-client/src/lib.rs
Normal file
9
litellm-rust/crates/redis-client/src/lib.rs
Normal file
|
|
@ -0,0 +1,9 @@
|
|||
mod client;
|
||||
mod config;
|
||||
mod error;
|
||||
mod namespace;
|
||||
|
||||
pub use client::{RedisClient, Ttl};
|
||||
pub use config::{DEFAULT_TIMEOUT, Endpoint, HostEndpoint, RedisConfig};
|
||||
pub use error::{ConnectError, Error};
|
||||
pub use namespace::Namespace;
|
||||
53
litellm-rust/crates/redis-client/src/namespace.rs
Normal file
53
litellm-rust/crates/redis-client/src/namespace.rs
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
#[derive(Clone, Default)]
|
||||
pub struct Namespace(Option<String>);
|
||||
|
||||
impl Namespace {
|
||||
pub fn new(namespace: impl Into<String>) -> Self {
|
||||
let namespace = namespace.into();
|
||||
Self((!namespace.is_empty()).then_some(namespace))
|
||||
}
|
||||
|
||||
pub(crate) fn apply(&self, key: &str) -> String {
|
||||
match &self.0 {
|
||||
Some(namespace) if !is_prefixed(key, namespace) => {
|
||||
format!("{namespace}:{key}")
|
||||
}
|
||||
_ => key.to_owned(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn is_prefixed(key: &str, namespace: &str) -> bool {
|
||||
key.strip_prefix(namespace)
|
||||
.is_some_and(|rest| rest.starts_with(':'))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::Namespace;
|
||||
|
||||
#[rstest]
|
||||
#[case::plain(Namespace::new("litellm"), "spend:key:abc", "litellm:spend:key:abc")]
|
||||
#[case::no_namespace(Namespace::default(), "spend:key:abc", "spend:key:abc")]
|
||||
#[case::empty_namespace(Namespace::new(""), "spend:key:abc", "spend:key:abc")]
|
||||
#[case::already_prefixed(Namespace::new("litellm"), "litellm:abc", "litellm:abc")]
|
||||
#[case::shares_prefix_without_delimiter(
|
||||
Namespace::new("litellm"),
|
||||
"litellm_spend_update_buffer",
|
||||
"litellm:litellm_spend_update_buffer"
|
||||
)]
|
||||
#[case::hash_tagged_key(
|
||||
Namespace::new("litellm"),
|
||||
"{api_key:sk-1}:requests",
|
||||
"litellm:{api_key:sk-1}:requests"
|
||||
)]
|
||||
fn prefixes_keys_like_python_check_and_fix_namespace(
|
||||
#[case] namespace: Namespace,
|
||||
#[case] key: &str,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
assert_eq!(namespace.apply(key), expected);
|
||||
}
|
||||
}
|
||||
5
litellm-rust/crates/redis-testing/AGENTS.md
Normal file
5
litellm-rust/crates/redis-testing/AGENTS.md
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
- Dev-only crate that starts real Redis for tests. Never a dependency outside `[dev-dependencies]`
|
||||
- Same shape as `db-testing`: one container per test, a handle that removes it on drop
|
||||
- The host port is fixed at start so `restart` keeps the same URL and a client can reconnect to it
|
||||
- Outages come from the container (`restart`, `freeze` / `thaw`), never `DEBUG`, which the official image disables
|
||||
- No env gate that makes a test pass without Redis. A test that can't start its container fails
|
||||
11
litellm-rust/crates/redis-testing/Cargo.toml
Normal file
11
litellm-rust/crates/redis-testing/Cargo.toml
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
[package]
|
||||
name = "litellm-redis-testing"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
testcontainers-modules = { workspace = true, features = ["redis"] }
|
||||
thiserror.workspace = true
|
||||
tokio = { workspace = true, features = ["io-util"] }
|
||||
11
litellm-rust/crates/redis-testing/src/error.rs
Normal file
11
litellm-rust/crates/redis-testing/src/error.rs
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
use std::io;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum Error {
|
||||
#[error("reserving a host port for the Redis container")]
|
||||
Port(#[source] io::Error),
|
||||
#[error("Redis on host port {port} did not answer PING")]
|
||||
Unreachable { port: u16 },
|
||||
#[error("controlling the Redis container")]
|
||||
Container(#[from] testcontainers_modules::testcontainers::TestcontainersError),
|
||||
}
|
||||
6
litellm-rust/crates/redis-testing/src/lib.rs
Normal file
6
litellm-rust/crates/redis-testing/src/lib.rs
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
mod error;
|
||||
mod ports;
|
||||
mod standalone;
|
||||
|
||||
pub use error::Error;
|
||||
pub use standalone::RedisServer;
|
||||
66
litellm-rust/crates/redis-testing/src/ports.rs
Normal file
66
litellm-rust/crates/redis-testing/src/ports.rs
Normal file
|
|
@ -0,0 +1,66 @@
|
|||
use std::{
|
||||
future::Future,
|
||||
io,
|
||||
net::TcpListener,
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
net::TcpStream,
|
||||
};
|
||||
|
||||
use crate::Error;
|
||||
|
||||
pub(crate) const HOST: &str = "127.0.0.1";
|
||||
const READY_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
const START_ATTEMPTS: usize = 3;
|
||||
|
||||
fn free_ports<const N: usize>() -> Result<[u16; N], Error> {
|
||||
let listeners = (0..N)
|
||||
.map(|_| TcpListener::bind((HOST, 0)))
|
||||
.collect::<io::Result<Vec<_>>>()
|
||||
.map_err(Error::Port)?;
|
||||
let ports = listeners
|
||||
.iter()
|
||||
.map(|listener| listener.local_addr().map(|address| address.port()))
|
||||
.collect::<io::Result<Vec<_>>>()
|
||||
.map_err(Error::Port)?;
|
||||
Ok(ports
|
||||
.try_into()
|
||||
.unwrap_or_else(|ports: Vec<u16>| unreachable!("{} ports", ports.len())))
|
||||
}
|
||||
|
||||
pub(crate) async fn wait_for_pong(port: u16) -> Result<(), Error> {
|
||||
let deadline = Instant::now() + READY_TIMEOUT;
|
||||
while Instant::now() < deadline {
|
||||
if ping(port).await.unwrap_or(false) {
|
||||
return Ok(());
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
}
|
||||
Err(Error::Unreachable { port })
|
||||
}
|
||||
|
||||
async fn ping(port: u16) -> io::Result<bool> {
|
||||
let mut stream = TcpStream::connect((HOST, port)).await?;
|
||||
stream.write_all(b"PING\r\n").await?;
|
||||
let mut reply = [0; 7];
|
||||
stream.read_exact(&mut reply).await?;
|
||||
Ok(&reply == b"+PONG\r\n")
|
||||
}
|
||||
|
||||
pub(crate) async fn start_on_free_ports<const N: usize, T, F>(
|
||||
start: impl Fn([u16; N]) -> F,
|
||||
) -> Result<T, Error>
|
||||
where
|
||||
F: Future<Output = Result<T, Error>>,
|
||||
{
|
||||
let mut attempt = 1;
|
||||
loop {
|
||||
match start(free_ports()?).await {
|
||||
Err(_) if attempt < START_ATTEMPTS => attempt += 1,
|
||||
started => return started,
|
||||
}
|
||||
}
|
||||
}
|
||||
55
litellm-rust/crates/redis-testing/src/standalone.rs
Normal file
55
litellm-rust/crates/redis-testing/src/standalone.rs
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
use testcontainers_modules::{
|
||||
redis::{REDIS_PORT, Redis},
|
||||
testcontainers::core::IntoContainerPort,
|
||||
testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner},
|
||||
};
|
||||
|
||||
use crate::{
|
||||
Error,
|
||||
ports::{HOST, start_on_free_ports, wait_for_pong},
|
||||
};
|
||||
|
||||
const REDIS_TAG: &str =
|
||||
"7.4@sha256:44960b3100c48159a66d28c18b0588c193492e25f7551d60e51dd464fc3eca5b";
|
||||
|
||||
pub struct RedisServer {
|
||||
container: ContainerAsync<Redis>,
|
||||
url: String,
|
||||
}
|
||||
|
||||
impl RedisServer {
|
||||
pub async fn start() -> Result<Self, Error> {
|
||||
start_on_free_ports(Self::start_on).await
|
||||
}
|
||||
|
||||
async fn start_on([host_port]: [u16; 1]) -> Result<Self, Error> {
|
||||
let container = Redis::default()
|
||||
.with_tag(REDIS_TAG)
|
||||
.with_mapped_port(host_port, REDIS_PORT.tcp())
|
||||
.start()
|
||||
.await?;
|
||||
wait_for_pong(host_port).await?;
|
||||
Ok(Self {
|
||||
container,
|
||||
url: format!("redis://{HOST}:{host_port}"),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn url(&self) -> &str {
|
||||
&self.url
|
||||
}
|
||||
|
||||
pub async fn restart(&self) -> Result<(), Error> {
|
||||
self.container.stop().await?;
|
||||
self.container.start().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn freeze(&self) -> Result<(), Error> {
|
||||
Ok(self.container.pause().await?)
|
||||
}
|
||||
|
||||
pub async fn thaw(&self) -> Result<(), Error> {
|
||||
Ok(self.container.unpause().await?)
|
||||
}
|
||||
}
|
||||
|
|
@ -19,7 +19,7 @@ time = { workspace = true, features = ["formatting"] }
|
|||
[dev-dependencies]
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
rstest.workspace = true
|
||||
testcontainers-modules = { version = "0.15.0", features = ["clickhouse"] }
|
||||
testcontainers-modules = { workspace = true, features = ["clickhouse"] }
|
||||
tokio.workspace = true
|
||||
uuid.workspace = true
|
||||
wiremock.workspace = true
|
||||
|
|
|
|||
|
|
@ -19,6 +19,6 @@ url.workspace = true
|
|||
[dev-dependencies]
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
rstest.workspace = true
|
||||
testcontainers-modules = { version = "0.15.0", features = ["clickhouse"] }
|
||||
testcontainers-modules = { workspace = true, features = ["clickhouse"] }
|
||||
tokio.workspace = true
|
||||
wiremock.workspace = true
|
||||
|
|
|
|||
|
|
@ -15587,4 +15587,3 @@ async def test_spend_capture_rate_check_job_clears_the_gauge_once_the_setting_is
|
|||
call(api_provider="openai", capture_rate=0.97),
|
||||
call(api_provider="openai", capture_rate=None),
|
||||
]
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue