mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
1283 lines
49 KiB
Rust
1283 lines
49 KiB
Rust
use std::{env, time::Duration};
|
|
|
|
use litellm_cache::CacheType;
|
|
use litellm_cache_qdrant_semantic::{OpenAiEmbedderConfig, QdrantSemanticConfig, Quantization};
|
|
use litellm_cache_redis::{RedisNode, RedisTopology};
|
|
use pyo3::{
|
|
exceptions::{PyTypeError, PyValueError},
|
|
prelude::*,
|
|
types::{PyAny, PyDict, PyList, PyString},
|
|
};
|
|
|
|
use super::{native::NativeResponseCache, request::duration};
|
|
|
|
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
|
|
pub(super) struct CachePolicy {
|
|
pub(super) mode: String,
|
|
pub(super) ttl: Option<Duration>,
|
|
pub(super) namespace: Option<String>,
|
|
pub(super) supported_call_types: Option<Vec<String>>,
|
|
pub(super) redis_flush_size: Option<usize>,
|
|
pub(super) semantic_cache_scope: String,
|
|
}
|
|
|
|
pub(super) struct MemoryCacheConfig {
|
|
pub(super) default_ttl: Duration,
|
|
pub(super) capacity: usize,
|
|
pub(super) max_entry_bytes: usize,
|
|
}
|
|
|
|
#[derive(Debug, PartialEq)]
|
|
pub(super) enum RedisProtocol {
|
|
Resp2,
|
|
Resp3,
|
|
}
|
|
|
|
#[derive(Debug, PartialEq)]
|
|
pub(super) enum CertificateRequirement {
|
|
None,
|
|
Optional,
|
|
Required,
|
|
}
|
|
|
|
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
|
|
pub(super) struct RedisTlsConfig {
|
|
pub(super) certificate_requirement: CertificateRequirement,
|
|
pub(super) check_hostname: bool,
|
|
pub(super) ca_certificate: Option<String>,
|
|
pub(super) ca_data: Option<String>,
|
|
pub(super) client_certificate: Option<String>,
|
|
pub(super) client_key: Option<String>,
|
|
}
|
|
|
|
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
|
|
pub(super) struct RedisConnectionConfig {
|
|
pub(super) host: String,
|
|
pub(super) port: u16,
|
|
pub(super) database: i64,
|
|
pub(super) username: Option<String>,
|
|
pub(super) password: Option<String>,
|
|
pub(super) protocol: RedisProtocol,
|
|
pub(super) pool_size: usize,
|
|
pub(super) read_timeout: Option<Duration>,
|
|
pub(super) connect_timeout: Option<Duration>,
|
|
pub(super) socket_keepalive: Option<bool>,
|
|
pub(super) health_check_interval: Duration,
|
|
pub(super) client_name: Option<String>,
|
|
pub(super) tls: Option<RedisTlsConfig>,
|
|
}
|
|
|
|
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
|
|
pub(super) struct RedisCacheConfig {
|
|
pub(super) default_ttl: Duration,
|
|
pub(super) namespace: Option<String>,
|
|
pub(super) flush_size: usize,
|
|
pub(super) topology: RedisTopology,
|
|
pub(super) connection: RedisConnectionConfig,
|
|
}
|
|
|
|
pub(super) struct AzureBlobCacheConfig {
|
|
pub(super) account_url: String,
|
|
pub(super) container: String,
|
|
}
|
|
|
|
struct RedisClientProjection<'py> {
|
|
topology: RedisTopology,
|
|
host: String,
|
|
port: u16,
|
|
pool_size: usize,
|
|
resolved: Bound<'py, PyDict>,
|
|
tls: Option<RedisTlsConfig>,
|
|
}
|
|
|
|
const REDIS_PY_DEFAULT_MAX_CONNECTIONS: usize = 1 << 31;
|
|
|
|
pub(super) struct QdrantSemanticCacheConfig {
|
|
pub(super) grpc_url: String,
|
|
pub(super) api_key: Option<String>,
|
|
pub(super) collection_name: String,
|
|
pub(super) similarity_threshold: f64,
|
|
pub(super) vector_size: u64,
|
|
pub(super) embedding: OpenAiEmbedderConfig,
|
|
pub(super) quantization: Quantization,
|
|
}
|
|
|
|
impl QdrantSemanticCacheConfig {
|
|
pub(super) fn to_qdrant_config(&self) -> QdrantSemanticConfig {
|
|
QdrantSemanticConfig {
|
|
collection_name: self.collection_name.clone(),
|
|
similarity_threshold: self.similarity_threshold,
|
|
vector_size: self.vector_size,
|
|
quantization: self.quantization.clone(),
|
|
}
|
|
}
|
|
}
|
|
pub(super) enum CacheBackendConfig {
|
|
Memory(MemoryCacheConfig),
|
|
Redis(Box<RedisCacheConfig>),
|
|
QdrantSemantic(Box<QdrantSemanticCacheConfig>),
|
|
AzureBlob(AzureBlobCacheConfig),
|
|
}
|
|
|
|
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
|
|
pub(super) struct NativeCacheConfig {
|
|
pub(super) policy: CachePolicy,
|
|
pub(super) backend: CacheBackendConfig,
|
|
}
|
|
|
|
pub(super) enum UnsupportedCacheConfig {
|
|
Backend,
|
|
RedisTopology,
|
|
RedisCredentials,
|
|
RedisConnection,
|
|
RedisOption,
|
|
QdrantEndpoint,
|
|
SemanticEmbedding,
|
|
}
|
|
|
|
impl UnsupportedCacheConfig {
|
|
pub(super) fn message(&self) -> &'static str {
|
|
match self {
|
|
Self::Backend => "native cache backend is not implemented",
|
|
Self::RedisTopology => "native Redis topology is not implemented",
|
|
Self::RedisCredentials => "native Redis credentials require Python",
|
|
Self::RedisConnection => "native Redis connection type is not implemented",
|
|
Self::RedisOption => "native Redis configuration requires Python",
|
|
Self::QdrantEndpoint => {
|
|
"native Qdrant requires the default REST port so the gRPC port can be derived"
|
|
}
|
|
Self::SemanticEmbedding => "native semantic embedding requires Python",
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(super) enum CacheConfigProjection {
|
|
Native(Box<NativeCacheConfig>),
|
|
Unsupported(UnsupportedCacheConfig),
|
|
}
|
|
|
|
impl NativeCacheConfig {
|
|
#[inline(never)]
|
|
pub(super) fn project(facade: &Bound<'_, PyAny>) -> PyResult<CacheConfigProjection> {
|
|
let backend_name = facade.getattr("type")?.extract::<String>()?;
|
|
let policy = CachePolicy {
|
|
mode: facade.getattr("mode")?.extract::<String>()?,
|
|
ttl: optional_duration(facade.getattr("ttl")?)?,
|
|
namespace: optional_string(facade.getattr("namespace")?)?,
|
|
supported_call_types: facade
|
|
.getattr("supported_call_types")?
|
|
.extract::<Option<Vec<String>>>()?,
|
|
redis_flush_size: facade
|
|
.getattr("redis_flush_size")?
|
|
.extract::<Option<usize>>()?,
|
|
semantic_cache_scope: facade
|
|
.getattr("semantic_cache_scope")?
|
|
.extract::<String>()?,
|
|
};
|
|
let backend = facade.getattr("cache")?;
|
|
match CacheType::from_python_name(&backend_name) {
|
|
Some(CacheType::Local) => project_memory(&backend).map(|backend| {
|
|
CacheConfigProjection::Native(Box::new(Self {
|
|
policy,
|
|
backend: CacheBackendConfig::Memory(backend),
|
|
}))
|
|
}),
|
|
Some(CacheType::Redis) => match project_redis(&backend)? {
|
|
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
|
|
policy,
|
|
backend: CacheBackendConfig::Redis(Box::new(backend)),
|
|
}))),
|
|
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
|
|
},
|
|
Some(CacheType::QdrantSemantic) => match project_qdrant_semantic(&backend)? {
|
|
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
|
|
policy,
|
|
backend: CacheBackendConfig::QdrantSemantic(Box::new(backend)),
|
|
}))),
|
|
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
|
|
},
|
|
Some(CacheType::AzureBlob) => project_azure_blob(&backend).map(|backend| {
|
|
CacheConfigProjection::Native(Box::new(Self {
|
|
policy,
|
|
backend: CacheBackendConfig::AzureBlob(backend),
|
|
}))
|
|
}),
|
|
Some(
|
|
CacheType::RedisSemantic
|
|
| CacheType::ValkeySemantic
|
|
| CacheType::S3
|
|
| CacheType::Disk
|
|
| CacheType::Gcs,
|
|
)
|
|
| None => Ok(CacheConfigProjection::Unsupported(
|
|
UnsupportedCacheConfig::Backend,
|
|
)),
|
|
}
|
|
}
|
|
|
|
pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> {
|
|
match &self.backend {
|
|
CacheBackendConfig::Memory(_) if service.kind() != "memory" => {
|
|
Some("facade and native backend types must match")
|
|
}
|
|
CacheBackendConfig::Memory(config)
|
|
if service.default_ttl() != Some(config.default_ttl) =>
|
|
{
|
|
Some("facade and native backend default TTLs must match")
|
|
}
|
|
CacheBackendConfig::Memory(config) if service.capacity() != Some(config.capacity) => {
|
|
Some("facade and native backend capacities must match")
|
|
}
|
|
CacheBackendConfig::Memory(config)
|
|
if service.max_entry_bytes() != Some(config.max_entry_bytes) =>
|
|
{
|
|
Some("facade and native backend item limits must match")
|
|
}
|
|
CacheBackendConfig::Memory(_) => None,
|
|
CacheBackendConfig::Redis(config) if service.kind() != "redis" => {
|
|
Some("facade and native backend types must match")
|
|
}
|
|
CacheBackendConfig::Redis(config) if service.topology() != Some(&config.topology) => {
|
|
Some("facade and native backend topologies must match")
|
|
}
|
|
CacheBackendConfig::Redis(config) => (service.namespace()
|
|
!= config.namespace.as_deref())
|
|
.then_some("facade and native backend namespaces must match")
|
|
.or_else(|| {
|
|
(service.default_ttl() != Some(config.default_ttl))
|
|
.then_some("facade and native backend default TTLs must match")
|
|
}),
|
|
CacheBackendConfig::QdrantSemantic(config) if service.kind() != "qdrant_semantic" => {
|
|
Some("facade and native backend types must match")
|
|
}
|
|
CacheBackendConfig::QdrantSemantic(config)
|
|
if service.collection_name() != Some(config.collection_name.as_str()) =>
|
|
{
|
|
Some("facade and native backend collections must match")
|
|
}
|
|
CacheBackendConfig::QdrantSemantic(config)
|
|
if service.similarity_threshold() != Some(config.similarity_threshold) =>
|
|
{
|
|
Some("facade and native backend similarity thresholds must match")
|
|
}
|
|
CacheBackendConfig::QdrantSemantic(config)
|
|
if service.vector_size() != Some(config.vector_size) =>
|
|
{
|
|
Some("facade and native backend vector sizes must match")
|
|
}
|
|
CacheBackendConfig::QdrantSemantic(config)
|
|
if service.embedding_model() != Some(config.embedding.model.as_str()) =>
|
|
{
|
|
Some("facade and native backend embedding models must match")
|
|
}
|
|
CacheBackendConfig::QdrantSemantic(_) => None,
|
|
CacheBackendConfig::AzureBlob(config) => match service.azure_blob_identity() {
|
|
None => Some("facade and native backend types must match"),
|
|
Some((account_url, container))
|
|
if account_url != config.account_url || container != config.container =>
|
|
{
|
|
Some("facade and native backend containers must match")
|
|
}
|
|
Some(_) if service.default_ttl().is_some() => {
|
|
Some("facade and native backend default TTLs must match")
|
|
}
|
|
Some(_) => None,
|
|
},
|
|
}
|
|
}
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn project_qdrant_semantic(
|
|
backend: &Bound<'_, PyAny>,
|
|
) -> PyResult<Result<QdrantSemanticCacheConfig, UnsupportedCacheConfig>> {
|
|
let rest_url = backend.getattr("qdrant_api_base")?.extract::<String>()?;
|
|
let parsed = match url::Url::parse(&rest_url) {
|
|
Ok(value) => value,
|
|
Err(_) => return Ok(Err(UnsupportedCacheConfig::QdrantEndpoint)),
|
|
};
|
|
if !matches!(parsed.scheme(), "http" | "https")
|
|
|| !parsed.path().is_empty() && parsed.path() != "/"
|
|
|| parsed.query().is_some()
|
|
|| parsed.host_str().is_none()
|
|
|| parsed.port() != Some(6333)
|
|
{
|
|
return Ok(Err(UnsupportedCacheConfig::QdrantEndpoint));
|
|
}
|
|
let mut grpc_url = parsed;
|
|
if grpc_url.set_port(Some(6334)).is_err() {
|
|
return Ok(Err(UnsupportedCacheConfig::QdrantEndpoint));
|
|
}
|
|
grpc_url.set_path("");
|
|
grpc_url.set_query(None);
|
|
|
|
let embedding_max_input_tokens = optional_attribute_i64(backend, "embedding_max_input_tokens")?;
|
|
if embedding_max_input_tokens.is_some() {
|
|
return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding));
|
|
}
|
|
let configured_model = backend.getattr("embedding_model")?.extract::<String>()?;
|
|
let embedding_model = configured_model
|
|
.strip_prefix("openai/")
|
|
.unwrap_or(&configured_model)
|
|
.to_owned();
|
|
if !embedding_model.starts_with("text-embedding-") {
|
|
return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding));
|
|
}
|
|
let proxy_server = py_sys_module(backend.py())?;
|
|
if let Some(proxy_server) = proxy_server {
|
|
let router = proxy_server.getattr("llm_router")?;
|
|
let model_list = proxy_server.getattr("llm_model_list")?;
|
|
let embedding_router = backend.py().import("litellm.caching._embedding_router")?;
|
|
if !embedding_router
|
|
.getattr("resolve_embedding_router")?
|
|
.call1((configured_model.as_str(), router, model_list))?
|
|
.is_none()
|
|
{
|
|
return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding));
|
|
}
|
|
}
|
|
let litellm = backend.py().import("litellm")?;
|
|
for name in ["api_key", "openai_key", "api_base"] {
|
|
if !litellm.getattr(name)?.is_none() {
|
|
return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding));
|
|
}
|
|
}
|
|
let Ok(embedding_api_key) = env::var("OPENAI_API_KEY") else {
|
|
return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding));
|
|
};
|
|
if embedding_api_key.is_empty() {
|
|
return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding));
|
|
}
|
|
let embedding_api_base = env::var("OPENAI_BASE_URL")
|
|
.or_else(|_| env::var("OPENAI_API_BASE"))
|
|
.unwrap_or_else(|_| "https://api.openai.com/v1".to_owned());
|
|
let timeout = optional_attribute_f64(backend, "embedding_timeout")?
|
|
.map(duration)
|
|
.transpose()?;
|
|
Ok(Ok(QdrantSemanticCacheConfig {
|
|
grpc_url: grpc_url.to_string().trim_end_matches('/').to_owned(),
|
|
api_key: optional_string(backend.getattr("qdrant_api_key")?)?,
|
|
collection_name: backend.getattr("collection_name")?.extract()?,
|
|
similarity_threshold: backend.getattr("similarity_threshold")?.extract()?,
|
|
vector_size: backend.getattr("vector_size")?.extract::<u64>()?,
|
|
embedding: OpenAiEmbedderConfig {
|
|
api_base: embedding_api_base,
|
|
api_key: embedding_api_key,
|
|
model: embedding_model,
|
|
timeout,
|
|
},
|
|
quantization: Quantization::Binary,
|
|
}))
|
|
}
|
|
|
|
fn py_sys_module(py: Python<'_>) -> PyResult<Option<Bound<'_, PyAny>>> {
|
|
match py
|
|
.import("sys")?
|
|
.getattr("modules")?
|
|
.get_item("litellm.proxy.proxy_server")
|
|
{
|
|
Ok(module) => Ok(Some(module)),
|
|
Err(error) if error.is_instance_of::<pyo3::exceptions::PyKeyError>(py) => Ok(None),
|
|
Err(error) => Err(error),
|
|
}
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn project_azure_blob(backend: &Bound<'_, PyAny>) -> PyResult<AzureBlobCacheConfig> {
|
|
let client = backend.getattr("container_client")?;
|
|
let container = client.getattr("container_name")?.extract::<String>()?;
|
|
let url = client.getattr("url")?.extract::<String>()?;
|
|
let account_url = url
|
|
.strip_suffix(container.as_str())
|
|
.and_then(|url| url.strip_suffix('/'))
|
|
.ok_or_else(|| PyValueError::new_err("Azure Blob container URL is malformed"))?;
|
|
Ok(AzureBlobCacheConfig {
|
|
account_url: account_url.to_string(),
|
|
container,
|
|
})
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult<MemoryCacheConfig> {
|
|
let max_size_kib = backend.getattr("max_size_per_item")?.extract::<usize>()?;
|
|
Ok(MemoryCacheConfig {
|
|
default_ttl: duration(backend.getattr("default_ttl")?.extract::<f64>()?)?,
|
|
capacity: backend.getattr("max_size_in_memory")?.extract::<usize>()?,
|
|
max_entry_bytes: max_size_kib
|
|
.checked_mul(1024)
|
|
.ok_or_else(|| PyValueError::new_err("memory cache item limit is too large"))?,
|
|
})
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn project_redis(
|
|
backend: &Bound<'_, PyAny>,
|
|
) -> PyResult<Result<RedisCacheConfig, UnsupportedCacheConfig>> {
|
|
let source = backend.getattr("redis_kwargs")?.cast_into::<PyDict>()?;
|
|
if has_value(&source, "sentinel_nodes")? {
|
|
return Ok(Err(UnsupportedCacheConfig::RedisTopology));
|
|
}
|
|
for key in ["credential_provider", "redis_connect_func"] {
|
|
if has_value(&source, key)? {
|
|
return Ok(Err(UnsupportedCacheConfig::RedisCredentials));
|
|
}
|
|
}
|
|
if has_value(&source, "connection_pool")? {
|
|
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
|
|
}
|
|
for key in [
|
|
"retry",
|
|
"retry_on_error",
|
|
"socket_keepalive_options",
|
|
"unix_socket_path",
|
|
"cache",
|
|
"cache_config",
|
|
"event_dispatcher",
|
|
"ssl_ca_path",
|
|
"ssl_password",
|
|
"ssl_min_version",
|
|
"ssl_ciphers",
|
|
"ssl_validate_ocsp",
|
|
"ssl_validate_ocsp_stapled",
|
|
"ssl_ocsp_context",
|
|
"ssl_ocsp_expected_cert",
|
|
] {
|
|
if has_value(&source, key)? {
|
|
return Ok(Err(UnsupportedCacheConfig::RedisOption));
|
|
}
|
|
}
|
|
for key in ["retry_on_timeout", "single_connection_client"] {
|
|
if optional_coerced_bool(&source, key)?.unwrap_or(false) {
|
|
return Ok(Err(UnsupportedCacheConfig::RedisOption));
|
|
}
|
|
}
|
|
|
|
let client = backend.getattr("redis_client")?;
|
|
let projection = if has_value(&source, "startup_nodes")? {
|
|
project_cluster_client(&source, &client)?
|
|
} else {
|
|
project_standalone_client(&client)?
|
|
};
|
|
let RedisClientProjection {
|
|
topology,
|
|
host,
|
|
port,
|
|
pool_size,
|
|
resolved,
|
|
tls,
|
|
} = match projection {
|
|
Ok(projection) => projection,
|
|
Err(reason) => return Ok(Err(reason)),
|
|
};
|
|
if has_value(&resolved, "credential_provider")? {
|
|
return Ok(Err(UnsupportedCacheConfig::RedisCredentials));
|
|
}
|
|
|
|
let protocol = match optional_i64(&resolved, "protocol")?.unwrap_or(2) {
|
|
2 => RedisProtocol::Resp2,
|
|
3 => RedisProtocol::Resp3,
|
|
_ => return Err(PyValueError::new_err("unsupported Redis protocol version")),
|
|
};
|
|
let health_check_interval =
|
|
duration(optional_f64(&resolved, "health_check_interval")?.unwrap_or(0.0))?;
|
|
Ok(Ok(RedisCacheConfig {
|
|
default_ttl: duration(backend.getattr("default_ttl")?.extract::<f64>()?)?,
|
|
namespace: optional_attribute_string(backend, "namespace")?,
|
|
flush_size: backend.getattr("redis_flush_size")?.extract::<usize>()?,
|
|
topology,
|
|
connection: RedisConnectionConfig {
|
|
host,
|
|
port,
|
|
database: optional_i64(&resolved, "db")?.unwrap_or(0),
|
|
username: optional_dict_string(&resolved, "username")?,
|
|
password: optional_dict_string(&resolved, "password")?,
|
|
protocol,
|
|
pool_size,
|
|
read_timeout: optional_dict_duration(&resolved, "socket_timeout")?,
|
|
connect_timeout: optional_dict_duration(&resolved, "socket_connect_timeout")?,
|
|
socket_keepalive: optional_bool(&resolved, "socket_keepalive")?,
|
|
health_check_interval,
|
|
client_name: optional_dict_string(&resolved, "client_name")?,
|
|
tls,
|
|
},
|
|
}))
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn project_standalone_client<'py>(
|
|
client: &Bound<'py, PyAny>,
|
|
) -> PyResult<Result<RedisClientProjection<'py>, UnsupportedCacheConfig>> {
|
|
let pool = client.getattr("connection_pool")?;
|
|
if !instance_class_is(&pool, "redis.connection", "ConnectionPool")? {
|
|
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
|
|
}
|
|
let resolved = pool.getattr("connection_kwargs")?.cast_into::<PyDict>()?;
|
|
if has_value(&resolved, "redis_connect_func")? {
|
|
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));
|
|
};
|
|
Ok(Ok(RedisClientProjection {
|
|
topology: RedisTopology::Standalone,
|
|
host: required_string(&resolved, "host")?,
|
|
port: port(required_i64(&resolved, "port")?)?,
|
|
pool_size: pool.getattr("max_connections")?.extract::<usize>()?,
|
|
resolved,
|
|
tls,
|
|
}))
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn project_cluster_client<'py>(
|
|
source: &Bound<'py, PyDict>,
|
|
client: &Bound<'py, PyAny>,
|
|
) -> PyResult<Result<RedisClientProjection<'py>, UnsupportedCacheConfig>> {
|
|
let Some(startup_nodes) = startup_nodes(source)? else {
|
|
return Ok(Err(UnsupportedCacheConfig::RedisTopology));
|
|
};
|
|
if !instance_class_is(client, "redis.cluster", "RedisCluster")? {
|
|
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
|
|
}
|
|
let nodes = client.getattr("nodes_manager")?;
|
|
if !class_is(
|
|
&nodes.getattr("connection_pool_class")?,
|
|
"redis.connection",
|
|
"ConnectionPool",
|
|
)? {
|
|
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
|
|
}
|
|
let resolved = nodes.getattr("connection_kwargs")?.cast_into::<PyDict>()?;
|
|
if let Some(connect) = resolved.get_item("redis_connect_func")?
|
|
&& !connect.is_none()
|
|
{
|
|
let own_hook = connect
|
|
.getattr("__self__")
|
|
.is_ok_and(|owner| owner.is(client))
|
|
&& connect
|
|
.getattr("__func__")
|
|
.and_then(|function| Ok(function.is(&client.get_type().getattr("on_connect")?)))
|
|
.unwrap_or(false);
|
|
if !own_hook {
|
|
return Ok(Err(UnsupportedCacheConfig::RedisCredentials));
|
|
}
|
|
}
|
|
let tls = if optional_bool(&resolved, "ssl")?.unwrap_or(false) {
|
|
Some(project_tls(&resolved)?)
|
|
} else {
|
|
None
|
|
};
|
|
let first = &startup_nodes[0];
|
|
Ok(Ok(RedisClientProjection {
|
|
host: first.host.clone(),
|
|
port: first.port,
|
|
pool_size: optional_i64(&resolved, "max_connections")?
|
|
.map(|value| {
|
|
usize::try_from(value).map_err(|_| PyValueError::new_err("invalid Redis pool size"))
|
|
})
|
|
.transpose()?
|
|
.unwrap_or(REDIS_PY_DEFAULT_MAX_CONNECTIONS),
|
|
topology: RedisTopology::Cluster { startup_nodes },
|
|
resolved,
|
|
tls,
|
|
}))
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn startup_nodes(source: &Bound<'_, PyDict>) -> PyResult<Option<Vec<RedisNode>>> {
|
|
let Some(nodes) = source.get_item("startup_nodes")? else {
|
|
return Ok(None);
|
|
};
|
|
let Ok(nodes) = nodes.cast_into::<PyList>() else {
|
|
return Ok(None);
|
|
};
|
|
if nodes.is_empty() {
|
|
return Ok(None);
|
|
}
|
|
let mut parsed = Vec::with_capacity(nodes.len());
|
|
for node in nodes.iter() {
|
|
let Ok(node) = node.cast_into::<PyDict>() else {
|
|
return Ok(None);
|
|
};
|
|
if node.len() != 2 || !has_value(&node, "host")? || !has_value(&node, "port")? {
|
|
return Ok(None);
|
|
}
|
|
let (Ok(host), Ok(port)) = (
|
|
required_string(&node, "host"),
|
|
required_i64(&node, "port").and_then(port),
|
|
) else {
|
|
return Ok(None);
|
|
};
|
|
parsed.push(RedisNode { host, port });
|
|
}
|
|
Ok(Some(parsed))
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn port(value: i64) -> PyResult<u16> {
|
|
u16::try_from(value).map_err(|_| PyValueError::new_err("invalid Redis port"))
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn project_tls(values: &Bound<'_, PyDict>) -> PyResult<RedisTlsConfig> {
|
|
Ok(RedisTlsConfig {
|
|
certificate_requirement: certificate_requirement(values)?,
|
|
check_hostname: optional_bool(values, "ssl_check_hostname")?.unwrap_or(false),
|
|
ca_certificate: optional_dict_string(values, "ssl_ca_certs")?,
|
|
ca_data: optional_dict_string(values, "ssl_ca_data")?,
|
|
client_certificate: optional_dict_string(values, "ssl_certfile")?,
|
|
client_key: optional_dict_string(values, "ssl_keyfile")?,
|
|
})
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn certificate_requirement(values: &Bound<'_, PyDict>) -> PyResult<CertificateRequirement> {
|
|
let Some(value) = values.get_item("ssl_cert_reqs")? else {
|
|
return Ok(CertificateRequirement::Required);
|
|
};
|
|
if value.is_none() {
|
|
return Ok(CertificateRequirement::Required);
|
|
}
|
|
if let Ok(number) = value.extract::<i32>() {
|
|
return match number {
|
|
0 => Ok(CertificateRequirement::None),
|
|
1 => Ok(CertificateRequirement::Optional),
|
|
2 => Ok(CertificateRequirement::Required),
|
|
_ => Err(PyValueError::new_err(
|
|
"invalid Redis TLS certificate requirement",
|
|
)),
|
|
};
|
|
}
|
|
let text = value.str()?;
|
|
let text = text.to_str()?;
|
|
if text.eq_ignore_ascii_case("none") || text.eq_ignore_ascii_case("cert_none") {
|
|
return Ok(CertificateRequirement::None);
|
|
}
|
|
if text.eq_ignore_ascii_case("optional") || text.eq_ignore_ascii_case("cert_optional") {
|
|
return Ok(CertificateRequirement::Optional);
|
|
}
|
|
if text.eq_ignore_ascii_case("required") || text.eq_ignore_ascii_case("cert_required") {
|
|
return Ok(CertificateRequirement::Required);
|
|
}
|
|
Err(PyValueError::new_err(
|
|
"invalid Redis TLS certificate requirement",
|
|
))
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn instance_class_is(value: &Bound<'_, PyAny>, module: &str, name: &str) -> PyResult<bool> {
|
|
class_is(value.get_type().as_any(), module, name)
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn class_is(value: &Bound<'_, PyAny>, module: &str, name: &str) -> PyResult<bool> {
|
|
Ok(value
|
|
.getattr("__module__")?
|
|
.cast_into::<PyString>()?
|
|
.to_str()?
|
|
== module
|
|
&& value
|
|
.getattr("__qualname__")?
|
|
.cast_into::<PyString>()?
|
|
.to_str()?
|
|
== name)
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn optional_duration(value: Bound<'_, PyAny>) -> PyResult<Option<Duration>> {
|
|
value.extract::<Option<f64>>()?.map(duration).transpose()
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn optional_attribute_string(value: &Bound<'_, PyAny>, name: &str) -> PyResult<Option<String>> {
|
|
match value.getattr(name) {
|
|
Ok(value) => optional_string(value),
|
|
Err(error) if error.is_instance_of::<pyo3::exceptions::PyAttributeError>(value.py()) => {
|
|
Ok(None)
|
|
}
|
|
Err(error) => Err(error),
|
|
}
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn optional_attribute_i64(value: &Bound<'_, PyAny>, name: &str) -> PyResult<Option<i64>> {
|
|
match value.getattr(name) {
|
|
Ok(attribute) => attribute.extract::<Option<i64>>(),
|
|
Err(error) if error.is_instance_of::<pyo3::exceptions::PyAttributeError>(value.py()) => {
|
|
Ok(None)
|
|
}
|
|
Err(error) => Err(error),
|
|
}
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn optional_attribute_f64(value: &Bound<'_, PyAny>, name: &str) -> PyResult<Option<f64>> {
|
|
match value.getattr(name) {
|
|
Ok(attribute) => attribute.extract::<Option<f64>>(),
|
|
Err(error) if error.is_instance_of::<pyo3::exceptions::PyAttributeError>(value.py()) => {
|
|
Ok(None)
|
|
}
|
|
Err(error) => Err(error),
|
|
}
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn optional_string(value: Bound<'_, PyAny>) -> PyResult<Option<String>> {
|
|
Ok(value
|
|
.extract::<Option<String>>()?
|
|
.filter(|value| !value.is_empty()))
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn has_value(values: &Bound<'_, PyDict>, key: &str) -> PyResult<bool> {
|
|
Ok(values.get_item(key)?.is_some_and(|value| !value.is_none()))
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn required_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult<String> {
|
|
values
|
|
.get_item(key)?
|
|
.ok_or_else(|| PyTypeError::new_err("Redis connection is incomplete"))?
|
|
.extract::<String>()
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn required_i64(values: &Bound<'_, PyDict>, key: &str) -> PyResult<i64> {
|
|
values
|
|
.get_item(key)?
|
|
.ok_or_else(|| PyTypeError::new_err("Redis connection is incomplete"))?
|
|
.extract::<i64>()
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn optional_dict_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult<Option<String>> {
|
|
match values.get_item(key)? {
|
|
Some(value) if !value.is_none() => optional_string(value),
|
|
_ => Ok(None),
|
|
}
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn optional_f64(values: &Bound<'_, PyDict>, key: &str) -> PyResult<Option<f64>> {
|
|
match values.get_item(key)? {
|
|
Some(value) => value.extract::<Option<f64>>(),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn optional_i64(values: &Bound<'_, PyDict>, key: &str) -> PyResult<Option<i64>> {
|
|
match values.get_item(key)? {
|
|
Some(value) => value.extract::<Option<i64>>(),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn optional_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult<Option<bool>> {
|
|
match values.get_item(key)? {
|
|
Some(value) => value.extract::<Option<bool>>(),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn optional_coerced_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult<Option<bool>> {
|
|
let Some(value) = values.get_item(key)? else {
|
|
return Ok(None);
|
|
};
|
|
if value.is_none() {
|
|
return Ok(None);
|
|
}
|
|
if let Ok(text) = value.extract::<String>() {
|
|
return Ok(Some(
|
|
text == "1" || text.eq_ignore_ascii_case("true") || text.eq_ignore_ascii_case("yes"),
|
|
));
|
|
}
|
|
value.extract::<bool>().map(Some)
|
|
}
|
|
|
|
#[inline(never)]
|
|
fn optional_dict_duration(values: &Bound<'_, PyDict>, key: &str) -> PyResult<Option<Duration>> {
|
|
optional_f64(values, key)?.map(duration).transpose()
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::{
|
|
sync::{Mutex, OnceLock},
|
|
time::Duration,
|
|
};
|
|
|
|
use std::ffi::CString;
|
|
|
|
use pyo3::{prelude::*, types::PyDict};
|
|
|
|
use litellm_cache_redis::{RedisNode, RedisTopology};
|
|
|
|
use super::{
|
|
CacheBackendConfig, CacheConfigProjection, CertificateRequirement, NativeCacheConfig,
|
|
RedisProtocol, UnsupportedCacheConfig,
|
|
};
|
|
use crate::cache::native::NativeResponseCache;
|
|
|
|
fn cluster_facade<'py>(py: Python<'py>, startup_nodes: &str, hook: &str) -> Bound<'py, PyAny> {
|
|
facade(
|
|
py,
|
|
&format!(
|
|
"RedisCluster = type('RedisCluster', (), {{'__module__': 'redis.cluster', 'on_connect': lambda self, connection: None}})\n\
|
|
client = RedisCluster()\n\
|
|
client.nodes_manager = SimpleNamespace(connection_pool_class=ConnectionPool, connection_kwargs={{'password': 'secret', 'redis_connect_func': {hook}, 'protocol': 3, 'ssl': True, 'ssl_cert_reqs': 'none'}})\n\
|
|
backend = SimpleNamespace(default_ttl=120, namespace='team', redis_flush_size=100, redis_kwargs={{'startup_nodes': {startup_nodes}, 'password': 'secret'}}, redis_client=client)\n\
|
|
facade = SimpleNamespace(type='redis', mode='default-on', ttl=None, namespace='team', supported_call_types=None, redis_flush_size=100, semantic_cache_scope='key', cache=backend)"
|
|
),
|
|
)
|
|
}
|
|
|
|
fn env_lock() -> &'static Mutex<()> {
|
|
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
|
|
LOCK.get_or_init(|| Mutex::new(()))
|
|
}
|
|
|
|
fn qdrant_facade<'py>(py: Python<'py>, extra: &str) -> Bound<'py, PyAny> {
|
|
install_fake_litellm(py);
|
|
facade(
|
|
py,
|
|
&format!(
|
|
"backend = SimpleNamespace(qdrant_api_base='https://qdrant.example:6333', qdrant_api_key='qdrant-key', collection_name='cache', similarity_threshold=0.99, embedding_model='openai/text-embedding-3-small', vector_size=8, embedding_max_input_tokens=None, embedding_timeout=None)\n\
|
|
facade = SimpleNamespace(type='qdrant-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)\n\
|
|
{extra}"
|
|
),
|
|
)
|
|
}
|
|
|
|
fn install_fake_litellm(py: Python<'_>) {
|
|
py.run(
|
|
c"
|
|
import sys
|
|
import types
|
|
litellm = types.ModuleType('litellm')
|
|
litellm.api_key = None
|
|
litellm.openai_key = None
|
|
litellm.api_base = None
|
|
litellm.__path__ = []
|
|
caching = types.ModuleType('litellm.caching')
|
|
caching.__path__ = []
|
|
embedding_router = types.ModuleType('litellm.caching._embedding_router')
|
|
embedding_router.resolve_embedding_router = lambda *_args: None
|
|
caching._embedding_router = embedding_router
|
|
litellm.caching = caching
|
|
sys.modules['litellm'] = litellm
|
|
sys.modules['litellm.caching'] = caching
|
|
sys.modules['litellm.caching._embedding_router'] = embedding_router
|
|
",
|
|
None,
|
|
None,
|
|
)
|
|
.unwrap();
|
|
}
|
|
|
|
fn configure_embedding_environment<'py>(
|
|
py: Python<'py>,
|
|
key: Option<&str>,
|
|
) -> PyResult<Bound<'py, PyAny>> {
|
|
let environ = py.import("os")?.getattr("environ")?;
|
|
let prior = environ.call_method1("get", ("OPENAI_API_KEY",))?;
|
|
match key {
|
|
Some(key) => environ.set_item("OPENAI_API_KEY", key)?,
|
|
None => {
|
|
environ.call_method1("pop", ("OPENAI_API_KEY", py.None()))?;
|
|
}
|
|
}
|
|
Ok(prior)
|
|
}
|
|
|
|
fn restore_embedding_environment(py: Python<'_>, prior: Bound<'_, PyAny>) -> PyResult<()> {
|
|
let environ = py.import("os")?.getattr("environ")?;
|
|
if prior.is_none() {
|
|
environ.call_method1("pop", ("OPENAI_API_KEY", py.None()))?;
|
|
} else {
|
|
environ.set_item("OPENAI_API_KEY", prior)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn facade<'py>(py: Python<'py>, body: &str) -> Bound<'py, PyAny> {
|
|
let locals = PyDict::new(py);
|
|
py.run(
|
|
&CString::new(format!(
|
|
"from types import SimpleNamespace\n\
|
|
ConnectionPool = type('ConnectionPool', (), {{'__module__': 'redis.connection'}})\n\
|
|
Connection = type('Connection', (), {{'__module__': 'redis.connection'}})\n\
|
|
SSLConnection = type('SSLConnection', (), {{'__module__': 'redis.connection'}})\n\
|
|
{body}"
|
|
))
|
|
.unwrap(),
|
|
None,
|
|
Some(&locals),
|
|
)
|
|
.unwrap();
|
|
locals.get_item("facade").unwrap().unwrap()
|
|
}
|
|
|
|
#[test]
|
|
fn projects_effective_memory_configuration() {
|
|
Python::initialize();
|
|
Python::attach(|py| {
|
|
let facade = facade(
|
|
py,
|
|
"backend = SimpleNamespace(default_ttl=913, max_size_in_memory=37, max_size_per_item=8)\n\
|
|
facade = SimpleNamespace(type='local', mode='default-on', ttl=11.5, namespace=None, supported_call_types=['completion'], redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
|
|
);
|
|
let CacheConfigProjection::Native(config) =
|
|
NativeCacheConfig::project(&facade).unwrap()
|
|
else {
|
|
panic!("memory cache should be supported");
|
|
};
|
|
assert_eq!(
|
|
config.policy.ttl.unwrap(),
|
|
std::time::Duration::from_secs_f64(11.5)
|
|
);
|
|
let CacheBackendConfig::Memory(memory) = config.backend else {
|
|
panic!("expected memory configuration");
|
|
};
|
|
assert_eq!(memory.default_ttl, std::time::Duration::from_secs(913));
|
|
assert_eq!(memory.capacity, 37);
|
|
assert_eq!(memory.max_entry_bytes, 8192);
|
|
let matching =
|
|
NativeResponseCache::memory(37, std::time::Duration::from_secs(913), 8192);
|
|
let mismatched =
|
|
NativeResponseCache::memory(37, std::time::Duration::from_secs(913), 8191);
|
|
let matching_config = NativeCacheConfig {
|
|
policy: config.policy,
|
|
backend: CacheBackendConfig::Memory(memory),
|
|
};
|
|
assert_eq!(matching_config.service_mismatch(&matching), None);
|
|
assert_eq!(
|
|
matching_config.service_mismatch(&mismatched),
|
|
Some("facade and native backend item limits must match")
|
|
);
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn projects_resolved_redis_tls_configuration() {
|
|
Python::initialize();
|
|
Python::attach(|py| {
|
|
let facade = facade(
|
|
py,
|
|
"pool = ConnectionPool()\n\
|
|
pool.connection_class = SSLConnection\n\
|
|
pool.max_connections = 29\n\
|
|
pool.connection_kwargs = {'host': 'cache.internal', 'port': 6380, 'db': 4, 'username': 'user', 'password': 'secret', 'protocol': 3, 'socket_timeout': 7.5, 'socket_connect_timeout': 2, 'socket_keepalive': True, 'health_check_interval': 15, 'client_name': 'litellm', 'ssl_cert_reqs': 'optional', 'ssl_check_hostname': True, 'ssl_ca_certs': '/ca.pem', 'ssl_ca_data': 'CA DATA', 'ssl_certfile': '/client.pem', 'ssl_keyfile': '/client.key'}\n\
|
|
client = SimpleNamespace(connection_pool=pool)\n\
|
|
backend = SimpleNamespace(default_ttl=777, namespace='team', redis_flush_size=31, redis_kwargs={}, redis_client=client)\n\
|
|
facade = SimpleNamespace(type='redis', mode='default-off', ttl=None, namespace='team', supported_call_types=None, redis_flush_size=31, semantic_cache_scope='key', cache=backend)",
|
|
);
|
|
let CacheConfigProjection::Native(config) =
|
|
NativeCacheConfig::project(&facade).unwrap()
|
|
else {
|
|
panic!("Redis cache should be supported");
|
|
};
|
|
let CacheBackendConfig::Redis(redis) = config.backend else {
|
|
panic!("expected Redis configuration");
|
|
};
|
|
assert_eq!(redis.default_ttl, std::time::Duration::from_secs(777));
|
|
assert_eq!(redis.namespace.as_deref(), Some("team"));
|
|
assert_eq!(redis.flush_size, 31);
|
|
assert_eq!(redis.connection.host, "cache.internal");
|
|
assert_eq!(redis.connection.port, 6380);
|
|
assert_eq!(redis.connection.database, 4);
|
|
assert_eq!(redis.connection.protocol, RedisProtocol::Resp3);
|
|
assert_eq!(redis.connection.pool_size, 29);
|
|
let tls = redis.connection.tls.unwrap();
|
|
assert_eq!(
|
|
tls.certificate_requirement,
|
|
CertificateRequirement::Optional
|
|
);
|
|
assert!(tls.check_hostname);
|
|
assert_eq!(tls.ca_certificate.as_deref(), Some("/ca.pem"));
|
|
assert_eq!(tls.ca_data.as_deref(), Some("CA DATA"));
|
|
assert_eq!(tls.client_certificate.as_deref(), Some("/client.pem"));
|
|
assert_eq!(tls.client_key.as_deref(), Some("/client.key"));
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn dynamic_redis_auth_stays_on_python() {
|
|
Python::initialize();
|
|
Python::attach(|py| {
|
|
let facade = facade(
|
|
py,
|
|
"backend = SimpleNamespace(redis_kwargs={'credential_provider': object()})\n\
|
|
facade = SimpleNamespace(type='redis', mode='default-on', ttl=None, namespace=None, supported_call_types=[], redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
|
|
);
|
|
let CacheConfigProjection::Unsupported(reason) =
|
|
NativeCacheConfig::project(&facade).unwrap()
|
|
else {
|
|
panic!("dynamic authentication must stay on Python");
|
|
};
|
|
assert_eq!(reason.message(), "native Redis credentials require Python");
|
|
});
|
|
}
|
|
#[test]
|
|
fn projects_cluster_startup_nodes_as_redis_topology() {
|
|
Python::initialize();
|
|
Python::attach(|py| {
|
|
let facade = cluster_facade(
|
|
py,
|
|
"[{'host': 'node-a', 'port': 7000}, {'host': 'node-b', 'port': 7001}]",
|
|
"client.on_connect",
|
|
);
|
|
let CacheConfigProjection::Native(config) =
|
|
NativeCacheConfig::project(&facade).unwrap()
|
|
else {
|
|
panic!("cluster startup nodes should project natively");
|
|
};
|
|
let CacheBackendConfig::Redis(redis) = &config.backend else {
|
|
panic!("expected Redis configuration");
|
|
};
|
|
let expected = RedisTopology::Cluster {
|
|
startup_nodes: vec![
|
|
RedisNode {
|
|
host: "node-a".into(),
|
|
port: 7000,
|
|
},
|
|
RedisNode {
|
|
host: "node-b".into(),
|
|
port: 7001,
|
|
},
|
|
],
|
|
};
|
|
assert_eq!(redis.topology, expected);
|
|
assert_eq!(redis.connection.host, "node-a");
|
|
assert_eq!(redis.connection.port, 7000);
|
|
assert_eq!(redis.connection.password.as_deref(), Some("secret"));
|
|
assert_eq!(redis.connection.protocol, RedisProtocol::Resp3);
|
|
assert_eq!(
|
|
redis
|
|
.connection
|
|
.tls
|
|
.as_ref()
|
|
.unwrap()
|
|
.certificate_requirement,
|
|
CertificateRequirement::None
|
|
);
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn malformed_startup_nodes_and_foreign_connect_hooks_stay_on_python() {
|
|
Python::initialize();
|
|
Python::attach(|py| {
|
|
for (startup_nodes, hook, message) in [
|
|
(
|
|
"[{'host': 'node-a', 'port': 7000, 'server_type': 'primary'}]",
|
|
"client.on_connect",
|
|
"native Redis topology is not implemented",
|
|
),
|
|
(
|
|
"[{'host': 'node-a', 'port': 'seven'}]",
|
|
"client.on_connect",
|
|
"native Redis topology is not implemented",
|
|
),
|
|
(
|
|
"[]",
|
|
"client.on_connect",
|
|
"native Redis topology is not implemented",
|
|
),
|
|
(
|
|
"[{'host': 'node-a', 'port': 7000}]",
|
|
"lambda connection: None",
|
|
"native Redis credentials require Python",
|
|
),
|
|
] {
|
|
let facade = cluster_facade(py, startup_nodes, hook);
|
|
let CacheConfigProjection::Unsupported(reason) =
|
|
NativeCacheConfig::project(&facade).unwrap()
|
|
else {
|
|
panic!("{startup_nodes} with {hook} must stay on Python");
|
|
};
|
|
assert_eq!(reason.message(), message, "{startup_nodes} with {hook}");
|
|
}
|
|
});
|
|
}
|
|
#[test]
|
|
fn projects_qdrant_configuration_from_python() {
|
|
let _guard = env_lock().lock().unwrap_or_else(|error| error.into_inner());
|
|
Python::initialize();
|
|
Python::attach(|py| {
|
|
let prior = configure_embedding_environment(py, Some("embedding-key")).unwrap();
|
|
let facade = qdrant_facade(py, "");
|
|
let CacheConfigProjection::Native(config) =
|
|
NativeCacheConfig::project(&facade).unwrap()
|
|
else {
|
|
panic!("Qdrant cache should be supported");
|
|
};
|
|
let CacheBackendConfig::QdrantSemantic(config) = config.backend else {
|
|
panic!("expected Qdrant configuration");
|
|
};
|
|
assert_eq!(config.grpc_url, "https://qdrant.example:6334");
|
|
assert_eq!(config.api_key.as_deref(), Some("qdrant-key"));
|
|
assert_eq!(config.collection_name, "cache");
|
|
assert_eq!(config.vector_size, 8);
|
|
assert_eq!(config.embedding.api_key, "embedding-key");
|
|
assert_eq!(config.embedding.model, "text-embedding-3-small");
|
|
restore_embedding_environment(py, prior).unwrap();
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn qdrant_projection_rejects_non_default_port() {
|
|
let _guard = env_lock().lock().unwrap_or_else(|error| error.into_inner());
|
|
Python::initialize();
|
|
Python::attach(|py| {
|
|
let prior = configure_embedding_environment(py, Some("embedding-key")).unwrap();
|
|
for endpoint in [
|
|
"https://qdrant.example:6332",
|
|
"https://qdrant.example",
|
|
"http://qdrant.example",
|
|
] {
|
|
let facade = qdrant_facade(py, &format!("backend.qdrant_api_base = '{endpoint}'"));
|
|
let CacheConfigProjection::Unsupported(reason) =
|
|
NativeCacheConfig::project(&facade).unwrap()
|
|
else {
|
|
panic!("unsupported Qdrant endpoint should stay on Python");
|
|
};
|
|
assert!(matches!(reason, UnsupportedCacheConfig::QdrantEndpoint));
|
|
}
|
|
let facade = qdrant_facade(py, "");
|
|
let CacheConfigProjection::Native(config) =
|
|
NativeCacheConfig::project(&facade).unwrap()
|
|
else {
|
|
panic!("default Qdrant endpoint should use native");
|
|
};
|
|
let CacheBackendConfig::QdrantSemantic(config) = config.backend else {
|
|
panic!("expected Qdrant configuration");
|
|
};
|
|
assert!(config.grpc_url.ends_with(":6334"));
|
|
restore_embedding_environment(py, prior).unwrap();
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn qdrant_projection_passes_configured_embedding_model_to_router() {
|
|
let _guard = env_lock().lock().unwrap_or_else(|error| error.into_inner());
|
|
Python::initialize();
|
|
Python::attach(|py| {
|
|
let prior = configure_embedding_environment(py, Some("embedding-key")).unwrap();
|
|
let facade = qdrant_facade(py, "");
|
|
py.run(
|
|
c"
|
|
import sys
|
|
import types
|
|
proxy_server = types.ModuleType('litellm.proxy.proxy_server')
|
|
proxy_server.llm_router = None
|
|
proxy_server.llm_model_list = None
|
|
sys.modules['litellm.proxy.proxy_server'] = proxy_server
|
|
embedding_router = sys.modules['litellm.caching._embedding_router']
|
|
embedding_router.resolve_embedding_router = lambda model, *_args: object() if model == 'openai/text-embedding-3-small' else None
|
|
",
|
|
None,
|
|
None,
|
|
)
|
|
.unwrap();
|
|
let CacheConfigProjection::Unsupported(reason) =
|
|
NativeCacheConfig::project(&facade).unwrap()
|
|
else {
|
|
panic!("router-backed embedding should stay on Python");
|
|
};
|
|
assert!(matches!(reason, UnsupportedCacheConfig::SemanticEmbedding));
|
|
py.run(
|
|
c"
|
|
import sys
|
|
sys.modules.pop('litellm.proxy.proxy_server', None)
|
|
",
|
|
None,
|
|
None,
|
|
)
|
|
.unwrap();
|
|
restore_embedding_environment(py, prior).unwrap();
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn qdrant_projection_rejects_python_embedding_features() {
|
|
let _guard = env_lock().lock().unwrap_or_else(|error| error.into_inner());
|
|
Python::initialize();
|
|
Python::attach(|py| {
|
|
let prior = configure_embedding_environment(py, Some("embedding-key")).unwrap();
|
|
let cases = [
|
|
("backend.embedding_max_input_tokens = 100", "semantic"),
|
|
("backend.embedding_model = 'cohere/embed'", "semantic"),
|
|
];
|
|
for (extra, _) in cases {
|
|
let facade = qdrant_facade(py, extra);
|
|
let CacheConfigProjection::Unsupported(reason) =
|
|
NativeCacheConfig::project(&facade).unwrap()
|
|
else {
|
|
panic!("unsupported embedding should stay on Python");
|
|
};
|
|
assert!(matches!(reason, UnsupportedCacheConfig::SemanticEmbedding));
|
|
}
|
|
restore_embedding_environment(py, prior).unwrap();
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn qdrant_projection_rejects_configured_litellm_base_or_missing_key() {
|
|
let _guard = env_lock().lock().unwrap_or_else(|error| error.into_inner());
|
|
Python::initialize();
|
|
Python::attach(|py| {
|
|
let prior = configure_embedding_environment(py, Some("embedding-key")).unwrap();
|
|
let facade = qdrant_facade(py, "");
|
|
let litellm = py.import("litellm").unwrap();
|
|
litellm
|
|
.setattr("api_base", "https://proxy.example")
|
|
.unwrap();
|
|
let CacheConfigProjection::Unsupported(reason) =
|
|
NativeCacheConfig::project(&facade).unwrap()
|
|
else {
|
|
panic!("configured LiteLLM base should stay on Python");
|
|
};
|
|
assert!(matches!(reason, UnsupportedCacheConfig::SemanticEmbedding));
|
|
litellm.setattr("api_base", py.None()).unwrap();
|
|
configure_embedding_environment(py, None).unwrap();
|
|
let CacheConfigProjection::Unsupported(reason) =
|
|
NativeCacheConfig::project(&facade).unwrap()
|
|
else {
|
|
panic!("missing embedding key should stay on Python");
|
|
};
|
|
assert!(matches!(reason, UnsupportedCacheConfig::SemanticEmbedding));
|
|
restore_embedding_environment(py, prior).unwrap();
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn qdrant_service_mismatch_reports_type_before_starting_qdrant() {
|
|
let _guard = env_lock().lock().unwrap_or_else(|error| error.into_inner());
|
|
Python::initialize();
|
|
Python::attach(|py| {
|
|
let prior = configure_embedding_environment(py, Some("embedding-key")).unwrap();
|
|
let facade = qdrant_facade(py, "");
|
|
let CacheConfigProjection::Native(config) =
|
|
NativeCacheConfig::project(&facade).unwrap()
|
|
else {
|
|
panic!("Qdrant cache should be supported");
|
|
};
|
|
let service = NativeResponseCache::memory(1, Duration::from_secs(1), 1024);
|
|
assert_eq!(
|
|
config.service_mismatch(&service),
|
|
Some("facade and native backend types must match")
|
|
);
|
|
restore_embedding_environment(py, prior).unwrap();
|
|
});
|
|
}
|
|
}
|