scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

feat(python-bridge): serve QdrantSemanticCache natively

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee 2026-09-21 20:45:15 +00:00
parent 8d41336a1e
commit 93e9836524
9 changed files with 936 additions and 60 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm-rust/Cargo.lock generated
View file

@ -2736,6 +2736,7 @@ dependencies = [
"litellm-auth-gcp",
"litellm-cache",
"litellm-cache-memory",
"litellm-cache-qdrant-semantic",
"litellm-cache-redis",
"litellm-cache-response",
"litellm-callbacks-legacy-python",
@ -2748,12 +2749,14 @@ dependencies = [
"litellm-types",
"pyo3",
"pyo3-async-runtimes",
"qdrant-client",
"rstest",
"serde",
"serde_json",
"serde_with",
"tokio",
"tokio-tungstenite",
"url",
]
[[package]]

5
litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs
View file

@ -90,7 +90,10 @@ async fn connect(
}
#[tokio::test(flavor = "multi_thread")]
#[allow(deprecated)]
#[expect(
deprecated,
reason = "the test verifies Qdrant's legacy always_ram quantization contract"
)]
async fn connect_sets_collection_quantization_and_index() {
for (quantization, expected) in [
(Quantization::Binary, 0),

3
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -24,6 +24,8 @@ litellm-cache.workspace = true
litellm-cache-memory.workspace = true
litellm-cache-redis.workspace = true
litellm-cache-response.workspace = true
litellm-cache-qdrant-semantic.workspace = true
qdrant-client.workspace = true
serde.workspace = true
litellm-auth.workspace = true
litellm-callbacks-legacy-python.workspace = true
@ -38,6 +40,7 @@ litellm-token-counter = { path = "../token-counter", default-features = false }
pyo3.workspace = true
pyo3-async-runtimes.workspace = true
serde_json.workspace = true
url.workspace = true
tokio = { workspace = true, features = ["sync"] }
[dev-dependencies]

408
litellm-rust/crates/python-bridge/src/cache/config.rs vendored
View file

@ -1,6 +1,7 @@
use std::time::Duration;
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},
@ -86,9 +87,30 @@ struct RedisClientProjection<'py> {
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>),
}
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
@ -103,6 +125,8 @@ pub(super) enum UnsupportedCacheConfig {
RedisCredentials,
RedisConnection,
RedisOption,
QdrantEndpoint,
SemanticEmbedding,
}
impl UnsupportedCacheConfig {
@ -113,6 +137,10 @@ impl UnsupportedCacheConfig {
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",
}
}
}
@ -155,12 +183,18 @@ impl NativeCacheConfig {
}))),
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::RedisSemantic
| CacheType::ValkeySemantic
| CacheType::S3
| CacheType::Disk
| CacheType::QdrantSemantic
| CacheType::AzureBlob
| CacheType::Gcs,
)
@ -171,18 +205,15 @@ impl NativeCacheConfig {
}
pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> {
if service.default_ttl()
!= Some(match &self.backend {
CacheBackendConfig::Memory(config) => config.default_ttl,
CacheBackendConfig::Redis(config) => config.default_ttl,
})
{
return Some("facade and native backend default TTLs must match");
}
match &self.backend {
CacheBackendConfig::Memory(config) if service.kind() != "memory" => {
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")
}
@ -192,7 +223,7 @@ impl NativeCacheConfig {
Some("facade and native backend item limits must match")
}
CacheBackendConfig::Memory(_) => None,
CacheBackendConfig::Redis(_) if service.kind() != "redis" => {
CacheBackendConfig::Redis(config) if service.kind() != "redis" => {
Some("facade and native backend types must match")
}
CacheBackendConfig::Redis(config) if service.topology() != Some(&config.topology) => {
@ -200,11 +231,134 @@ impl NativeCacheConfig {
}
CacheBackendConfig::Redis(config) => (service.namespace()
!= config.namespace.as_deref())
.then_some("facade and native backend namespaces must match"),
.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,
}
}
}
#[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().is_some_and(|port| port != 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((embedding_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_memory(backend: &Bound<'_, PyAny>) -> PyResult<MemoryCacheConfig> {
let max_size_kib = backend.getattr("max_size_per_item")?.extract::<usize>()?;
@ -514,6 +668,28 @@ fn optional_attribute_string(value: &Bound<'_, PyAny>, name: &str) -> PyResult<O
}
}
#[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
@ -597,6 +773,11 @@ fn optional_dict_duration(values: &Bound<'_, PyDict>, key: &str) -> PyResult<Opt
#[cfg(test)]
mod tests {
use std::{
sync::{Mutex, OnceLock},
time::Duration,
};
use std::ffi::CString;
use pyo3::{prelude::*, types::PyDict};
@ -605,7 +786,7 @@ mod tests {
use super::{
CacheBackendConfig, CacheConfigProjection, CertificateRequirement, NativeCacheConfig,
RedisProtocol,
RedisProtocol, UnsupportedCacheConfig,
};
use crate::cache::native::NativeResponseCache;
@ -622,6 +803,62 @@ mod tests {
)
}
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.del_item("OPENAI_API_KEY")?,
}
Ok(prior)
}
fn facade<'py>(py: Python<'py>, body: &str) -> Bound<'py, PyAny> {
let locals = PyDict::new(py);
py.run(
@ -740,7 +977,6 @@ mod tests {
assert_eq!(reason.message(), "native Redis credentials require Python");
});
}
#[test]
fn projects_cluster_startup_nodes_as_redis_topology() {
Python::initialize();
@ -823,4 +1059,146 @@ mod tests {
}
});
}
#[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");
let environ = py.import("os").unwrap().getattr("environ").unwrap();
if prior.is_none() {
environ.del_item("OPENAI_API_KEY").unwrap();
} else {
environ.set_item("OPENAI_API_KEY", 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();
let facade = qdrant_facade(
py,
"backend.qdrant_api_base = 'https://qdrant.example:6332'",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("non-default Qdrant port should stay on Python");
};
assert!(matches!(reason, UnsupportedCacheConfig::QdrantEndpoint));
let environ = py.import("os").unwrap().getattr("environ").unwrap();
if prior.is_none() {
environ.del_item("OPENAI_API_KEY").unwrap();
} else {
environ.set_item("OPENAI_API_KEY", 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));
}
let environ = py.import("os").unwrap().getattr("environ").unwrap();
if prior.is_none() {
environ.del_item("OPENAI_API_KEY").unwrap();
} else {
environ.set_item("OPENAI_API_KEY", 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));
let environ = py.import("os").unwrap().getattr("environ").unwrap();
if prior.is_none() {
environ.del_item("OPENAI_API_KEY").unwrap();
} else {
environ.set_item("OPENAI_API_KEY", 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")
);
let environ = py.import("os").unwrap().getattr("environ").unwrap();
if prior.is_none() {
environ.del_item("OPENAI_API_KEY").unwrap();
} else {
environ.set_item("OPENAI_API_KEY", prior).unwrap();
}
});
}
}

38
litellm-rust/crates/python-bridge/src/cache/facade.rs vendored
View file

@ -227,6 +227,11 @@ impl FacadeGuard {
"RedisClusterCache",
"redis",
),
"qdrant_semantic" => (
"litellm.caching.qdrant_semantic_cache",
"QdrantSemanticCache",
"qdrant-semantic",
),
_ => unreachable!(),
};
let backend = facade.getattr("cache")?;
@ -246,6 +251,26 @@ impl FacadeGuard {
if let Some(message) = config.service_mismatch(service) {
return Err(PyTypeError::new_err(message));
}
let backend_config_names = match kind {
"memory" | "redis" => &[
"namespace",
"default_ttl",
"max_size_in_memory",
"max_size_per_item",
"redis_kwargs",
"redis_flush_size",
][..],
"qdrant_semantic" => &[
"qdrant_api_base",
"collection_name",
"similarity_threshold",
"embedding_model",
"vector_size",
"embedding_max_input_tokens",
"embedding_timeout",
][..],
_ => unreachable!(),
};
Ok(Self {
outer: ObjectGuard::capture(
py,
@ -260,18 +285,7 @@ impl FacadeGuard {
"semantic_cache_scope",
],
)?,
backend: ObjectGuard::capture(
py,
&backend,
&[
"namespace",
"default_ttl",
"max_size_in_memory",
"max_size_per_item",
"redis_kwargs",
"redis_flush_size",
],
)?,
backend: ObjectGuard::capture(py, &backend, backend_config_names)?,
redis_pool: match (kind, cluster) {
("redis", false) => Some(RedisPoolGuard::capture(&backend, STANDALONE_POOL)?),
("redis", true) => Some(RedisPoolGuard::capture(&backend, CLUSTER_POOL)?),

111
litellm-rust/crates/python-bridge/src/cache/handle.rs vendored
View file

@ -1,8 +1,16 @@
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_host_python::release_gil;
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*};
use std::env;
use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration};
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_cache_qdrant_semantic::{OpenAiEmbedderConfig, Quantization};
use litellm_host_python::{release_gil, run_sync_value};
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*};
use url::Url;
use super::{
cache_error, config::QdrantSemanticCacheConfig, facade::FacadeGuard,
native::NativeResponseCache, request::duration,
};
#[pyclass(frozen, name = "_CacheTestHandle")]
pub(crate) struct CacheTestHandle {
@ -64,6 +72,101 @@ impl CacheTestHandle {
})
}
#[staticmethod]
#[pyo3(signature = (url, *, collection_name, similarity_threshold, vector_size, embedding_model="text-embedding-3-small", api_key=None, embedding_api_key=None, embedding_api_base=None, embedding_timeout_seconds=None, quantization="binary"))]
#[expect(
clippy::too_many_arguments,
reason = "the test handle exposes the complete Qdrant constructor"
)]
fn qdrant_semantic(
py: Python<'_>,
url: String,
collection_name: String,
similarity_threshold: f64,
vector_size: u64,
embedding_model: &str,
api_key: Option<String>,
embedding_api_key: Option<String>,
embedding_api_base: Option<String>,
embedding_timeout_seconds: Option<f64>,
quantization: &str,
) -> PyResult<Self> {
let parsed = Url::parse(&url).map_err(|_| {
pyo3::exceptions::PyValueError::new_err(
"native Qdrant requires the default REST port so the gRPC port can be derived",
)
})?;
if !matches!(parsed.scheme(), "http" | "https")
|| (!parsed.path().is_empty() && parsed.path() != "/")
|| parsed.query().is_some()
|| parsed.host_str().is_none()
|| parsed.port().is_some_and(|port| port != 6333)
{
return Err(pyo3::exceptions::PyValueError::new_err(
"native Qdrant requires the default REST port so the gRPC port can be derived",
));
}
let mut grpc_url = parsed;
grpc_url.set_port(Some(6334)).map_err(|_| {
pyo3::exceptions::PyValueError::new_err(
"native Qdrant requires the default REST port so the gRPC port can be derived",
)
})?;
grpc_url.set_path("");
grpc_url.set_query(None);
let embedding_api_key = embedding_api_key
.or_else(|| {
env::var("OPENAI_API_KEY")
.ok()
.filter(|value| !value.is_empty())
})
.ok_or_else(|| {
pyo3::exceptions::PyValueError::new_err(
"native semantic embedding requires an OpenAI API key",
)
})?;
let embedding_api_base = embedding_api_base.unwrap_or_else(|| {
env::var("OPENAI_BASE_URL")
.or_else(|_| env::var("OPENAI_API_BASE"))
.unwrap_or_else(|_| "https://api.openai.com/v1".to_owned())
});
let quantization = match quantization {
"binary" => Quantization::Binary,
"scalar" => Quantization::Scalar,
"product" => Quantization::Product,
_ => {
return Err(pyo3::exceptions::PyValueError::new_err(
"unsupported Qdrant quantization",
));
}
};
let config = QdrantSemanticCacheConfig {
grpc_url: grpc_url.to_string().trim_end_matches('/').to_owned(),
api_key,
collection_name,
similarity_threshold,
vector_size,
embedding: OpenAiEmbedderConfig {
api_base: embedding_api_base,
api_key: embedding_api_key,
model: embedding_model.to_owned(),
timeout: embedding_timeout_seconds.map(duration).transpose()?,
},
quantization,
};
let service = run_sync_value(py, async move {
let handle = tokio::runtime::Handle::current();
NativeResponseCache::qdrant_semantic(config, handle)
.await
.map_err(cache_error)
})?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[getter]
fn backend(&self) -> &'static str {
self.service.kind()

149
litellm-rust/crates/python-bridge/src/cache/native.rs vendored
View file

@ -1,13 +1,16 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{CacheCodec, CacheConnectionResult, Error};
use litellm_cache::{CacheCodec, CacheConnectionResult, Error, SemanticCacheContext};
use litellm_cache_memory::InMemoryCache;
use litellm_cache_qdrant_semantic::{Embedder, OpenAiEmbedder, QdrantSemanticCache};
use litellm_cache_redis::{RedisCache, RedisTopology};
use litellm_cache_response::{
CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer,
};
use serde_json::Value;
use super::{config::QdrantSemanticCacheConfig, request::exact};
#[derive(Clone)]
pub(super) enum NativeResponseCache {
Memory(Arc<ResponseCache<InMemoryCache<CacheEntry>>>),
@ -15,6 +18,7 @@ pub(super) enum NativeResponseCache {
cache: Arc<ResponseCache<RedisCache<ResponseCacheCodec>>>,
buffer: Option<Arc<WriteBuffer>>,
},
QdrantSemantic(Arc<ResponseCache<QdrantSemanticCache<OpenAiEmbedder, ResponseCacheCodec>>>),
}
impl NativeResponseCache {
@ -45,6 +49,30 @@ impl NativeResponseCache {
buffer: None,
})
}
pub async fn qdrant_semantic(
config: QdrantSemanticCacheConfig,
runtime: tokio::runtime::Handle,
) -> Result<Self, Error> {
let client = qdrant_client::Qdrant::from_url(&config.grpc_url)
.skip_compatibility_check()
.api_key(config.api_key.as_deref())
.build()
.map_err(|_| Error::Unavailable)?;
let qdrant_config = config.to_qdrant_config();
let embedder = OpenAiEmbedder::new(config.embedding)?;
let cache = QdrantSemanticCache::connect(
client,
embedder,
ResponseCacheCodec,
qdrant_config,
runtime,
)
.await?;
Ok(Self::QdrantSemantic(Arc::new(ResponseCache::new(
Arc::new(cache),
))))
}
}
impl NativeResponseCache {
@ -52,6 +80,7 @@ impl NativeResponseCache {
match self {
Self::Memory(_) => "memory",
Self::Redis { .. } => "redis",
Self::QdrantSemantic(_) => "qdrant_semantic",
}
}
@ -59,6 +88,7 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.default_ttl(),
Self::Redis { cache, .. } => cache.default_ttl(),
Self::QdrantSemantic(_) => None,
}
}
@ -66,6 +96,7 @@ impl NativeResponseCache {
match self {
Self::Memory(_) => None,
Self::Redis { cache, .. } => cache.backend().namespace(),
Self::QdrantSemantic(_) => None,
}
}
@ -80,6 +111,7 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),
Self::Redis { .. } => None,
Self::QdrantSemantic(_) => None,
}
}
@ -87,6 +119,7 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.backend().max_entry_bytes(),
Self::Redis { .. } => None,
Self::QdrantSemantic(_) => None,
}
}
@ -100,89 +133,157 @@ impl NativeResponseCache {
}
}
pub fn collection_name(&self) -> Option<&str> {
match self {
Self::QdrantSemantic(cache) => Some(cache.backend().collection_name()),
_ => None,
}
}
pub fn similarity_threshold(&self) -> Option<f64> {
match self {
Self::QdrantSemantic(cache) => Some(cache.backend().similarity_threshold()),
_ => None,
}
}
pub fn vector_size(&self) -> Option<u64> {
match self {
Self::QdrantSemantic(cache) => Some(cache.backend().vector_size()),
_ => None,
}
}
pub fn embedding_model(&self) -> Option<&str> {
match self {
Self::QdrantSemantic(cache) => Some(cache.backend().embedder().model()),
_ => None,
}
}
pub fn lookup(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<SemanticCacheContext>,
now: Duration,
) -> Result<Option<Value>, Error> {
match self {
Self::Memory(cache) => cache.lookup(request, now),
Self::Redis { cache, .. } => cache.lookup(request, now),
Self::Memory(cache) => cache.lookup(&exact(request), now),
Self::Redis { cache, .. } => cache.lookup(&exact(request), now),
Self::QdrantSemantic(cache) => cache.lookup(request, now),
}
}
pub fn store(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<SemanticCacheContext>,
response: Value,
now: Duration,
) -> Result<(), Error> {
match self {
Self::Memory(cache) => cache.store(request, response, now),
Self::Redis { cache, .. } => cache.store(request, response, now),
Self::Memory(cache) => cache.store(&exact(request), response, now),
Self::Redis { cache, .. } => cache.store(&exact(request), response, now),
Self::QdrantSemantic(cache) => cache.store(request, response, now),
}
}
pub fn lookup_batch(
&self,
requests: &[ResponseCacheRequest],
requests: &[ResponseCacheRequest<SemanticCacheContext>],
now: Duration,
) -> Result<PartialHits, Error> {
match self {
Self::Memory(cache) => cache.lookup_batch(requests, now),
Self::Redis { cache, .. } => cache.lookup_batch(requests, now),
Self::Memory(cache) => {
cache.lookup_batch(&requests.iter().map(exact).collect::<Vec<_>>(), now)
}
Self::Redis { cache, .. } => {
cache.lookup_batch(&requests.iter().map(exact).collect::<Vec<_>>(), now)
}
Self::QdrantSemantic(_) => Err(Error::UnsupportedOperation),
}
}
pub async fn async_lookup(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<SemanticCacheContext>,
now: Duration,
) -> Result<Option<Value>, Error> {
match self {
Self::Memory(cache) => cache.async_lookup(request, now).await,
Self::Redis { cache, .. } => cache.async_lookup(request, now).await,
Self::Memory(cache) => cache.async_lookup(&exact(request), now).await,
Self::Redis { cache, .. } => cache.async_lookup(&exact(request), now).await,
Self::QdrantSemantic(cache) => cache.async_lookup(request, now).await,
}
}
pub async fn async_store(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<SemanticCacheContext>,
response: Value,
now: Duration,
) -> Result<(), Error> {
match self {
Self::Memory(cache) => cache.async_store(request, response, now).await,
Self::Memory(cache) => cache.async_store(&exact(request), response, now).await,
Self::Redis {
cache,
buffer: None,
} => cache.async_store(request, response, now).await,
} => cache.async_store(&exact(request), response, now).await,
Self::Redis {
cache,
buffer: Some(buffer),
} => buffer.async_store(cache, request, response, now).await,
} => {
let request = exact(request);
buffer.async_store(cache, &request, response, now).await
}
Self::QdrantSemantic(cache) => cache.async_store(request, response, now).await,
}
}
pub async fn async_lookup_batch(
&self,
requests: &[ResponseCacheRequest],
requests: &[ResponseCacheRequest<SemanticCacheContext>],
now: Duration,
) -> Result<PartialHits, Error> {
match self {
Self::Memory(cache) => cache.async_lookup_batch(requests, now).await,
Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await,
Self::Memory(cache) => {
let requests = requests.iter().map(exact).collect::<Vec<_>>();
cache.async_lookup_batch(&requests, now).await
}
Self::Redis { cache, .. } => {
let requests = requests.iter().map(exact).collect::<Vec<_>>();
cache.async_lookup_batch(&requests, now).await
}
Self::QdrantSemantic(_) => Err(Error::UnsupportedOperation),
}
}
pub async fn async_store_batch(
&self,
entries: Vec<(ResponseCacheRequest, Value)>,
entries: Vec<(ResponseCacheRequest<SemanticCacheContext>, Value)>,
now: Duration,
) -> Result<(), Error> {
match self {
Self::Memory(cache) => cache.async_store_batch(entries, now).await,
Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await,
Self::Memory(cache) => {
cache
.async_store_batch(
entries
.into_iter()
.map(|(request, value)| (exact(&request), value))
.collect(),
now,
)
.await
}
Self::Redis { cache, .. } => {
cache
.async_store_batch(
entries
.into_iter()
.map(|(request, value)| (exact(&request), value))
.collect(),
now,
)
.await
}
Self::QdrantSemantic(cache) => cache.async_store_batch(entries, now).await,
}
}
@ -195,6 +296,7 @@ impl NativeResponseCache {
}
cache.async_flush().await
}
Self::QdrantSemantic(_) => Err(Error::UnsupportedOperation),
}
}
@ -202,6 +304,7 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.test_connection().await,
Self::Redis { cache, .. } => cache.test_connection().await,
Self::QdrantSemantic(_) => Err(Error::UnsupportedOperation),
}
}
}

28
litellm-rust/crates/python-bridge/src/cache/request.rs vendored
View file

@ -1,10 +1,11 @@
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use litellm_cache::ExactCacheContext;
use litellm_cache::{ExactCacheContext, SemanticCacheContext, SemanticCacheScope};
use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest};
use litellm_host_python::from_py;
use pyo3::{exceptions::PyValueError, prelude::*};
use serde::Deserialize;
use serde_json::{Map, Value};
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
@ -13,34 +14,51 @@ struct RequestInput {
controls: Option<CacheControls>,
ttl_seconds: Option<f64>,
max_age_seconds: Option<f64>,
messages: Option<Vec<Value>>,
input: Option<String>,
metadata: Option<Map<String, Value>>,
scope: Option<SemanticCacheScope>,
}
pub(super) fn request(
value: &Bound<'_, PyAny>,
) -> PyResult<ResponseCacheRequest<ExactCacheContext>> {
) -> PyResult<ResponseCacheRequest<SemanticCacheContext>> {
let input: RequestInput = from_py(value)?;
request_input(input)
}
fn request_input(input: RequestInput) -> PyResult<ResponseCacheRequest<ExactCacheContext>> {
let mut request: ResponseCacheRequest<ExactCacheContext> = ResponseCacheRequest::new(input.key);
fn request_input(input: RequestInput) -> PyResult<ResponseCacheRequest<SemanticCacheContext>> {
let mut request: ResponseCacheRequest<SemanticCacheContext> =
ResponseCacheRequest::new(input.key);
if let Some(controls) = input.controls {
request.controls = controls;
}
request.context.ttl = input.ttl_seconds.map(duration).transpose()?;
request.context.messages = input.messages.unwrap_or_default();
request.context.input = input.input;
request.context.metadata = input.metadata.unwrap_or_default();
request.context.scope = input.scope.unwrap_or_default();
request.max_age = input.max_age_seconds.map(duration).transpose()?;
Ok(request)
}
pub(super) fn requests(
value: &Bound<'_, PyAny>,
) -> PyResult<Vec<ResponseCacheRequest<ExactCacheContext>>> {
) -> PyResult<Vec<ResponseCacheRequest<SemanticCacheContext>>> {
from_py::<Vec<RequestInput>>(value)?
.into_iter()
.map(request_input)
.collect()
}
pub(super) fn exact(
request: &ResponseCacheRequest<SemanticCacheContext>,
) -> ResponseCacheRequest<ExactCacheContext> {
request.clone().with_context(ExactCacheContext {
ttl: request.context.ttl,
})
}
pub(super) fn duration(seconds: f64) -> PyResult<Duration> {
Duration::try_from_secs_f64(seconds)
.map_err(|_| PyValueError::new_err("cache durations must be finite and nonnegative"))

251
tests/test_litellm_rust/test_cache.py
View file

@ -1,7 +1,10 @@
import asyncio
import contextvars
import gc
import hashlib
import http.server
import json
import math
import os
import threading
import time
@ -10,6 +13,7 @@ from collections.abc import Generator
from types import SimpleNamespace
from typing import Final, Protocol, cast
from urllib.parse import urlparse
from uuid import uuid4
import fakeredis
import pytest
@ -34,6 +38,71 @@ def request(key: str = "key") -> dict[str, object]:
return {"key": {"preset": key}}
def semantic_request(
key: str,
messages: list[dict[str, object]],
**kwargs: object,
) -> dict[str, object]:
return {**request(key), "messages": messages, **kwargs}
def embedding_vector(text: str) -> list[float]:
raw = hashlib.sha256(text.encode()).digest()[:8]
values: Final = [byte / 127.5 - 1 for byte in raw]
norm: Final = math.sqrt(sum(value * value for value in values))
return [value / norm for value in values]
@pytest.fixture
def qdrant_url() -> str:
value: Final[str | None] = os.environ.get("QDRANT_URL")
if not value:
pytest.skip("QDRANT_URL is required for Qdrant semantic cache tests")
return value.rstrip("/")
@pytest.fixture
def fake_embedding_endpoint(monkeypatch: pytest.MonkeyPatch) -> Generator[str]:
class EmbeddingHandler(http.server.BaseHTTPRequestHandler):
def do_POST(self) -> None:
length: Final = int(self.headers["Content-Length"])
body: Final = json.loads(self.rfile.read(length))
text: Final = body["input"]
response: Final = {
"object": "list",
"data": [
{
"object": "embedding",
"index": 0,
"embedding": embedding_vector(text),
}
],
"model": body["model"],
"usage": {"prompt_tokens": 1, "total_tokens": 1},
}
encoded: Final = json.dumps(response).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(encoded)))
self.end_headers()
self.wfile.write(encoded)
def log_message(self, *_args: object) -> None:
return
server: Final = http.server.ThreadingHTTPServer(("127.0.0.1", 0), EmbeddingHandler)
worker: Final = threading.Thread(target=server.serve_forever, daemon=True)
worker.start()
monkeypatch.setenv("OPENAI_API_BASE", f"http://127.0.0.1:{server.server_address[1]}")
monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
try:
yield f"http://127.0.0.1:{server.server_address[1]}"
finally:
server.shutdown()
server.server_close()
worker.join(timeout=5)
@pytest.fixture
def redis_url() -> Generator[str]:
server: Final = fakeredis.TcpFakeServer(("127.0.0.1", 0), server_type="redis")
@ -464,3 +533,185 @@ async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_n
client.delete("unscoped")
client.close()
facade.cache.redis_client.close()
def qdrant_facade(
qdrant_url: str,
collection_name: str,
) -> Cache:
return Cache(
type=LiteLLMCacheType.QDRANT_SEMANTIC,
qdrant_api_base=qdrant_url,
qdrant_collection_name=collection_name,
similarity_threshold=0.99,
qdrant_semantic_cache_embedding_model="text-embedding-3-small",
qdrant_semantic_cache_vector_size=8,
)
def test_qdrant_semantic_facade_binds_native_and_shares_entries(
qdrant_url: str,
fake_embedding_endpoint: str,
) -> None:
del fake_embedding_endpoint
messages: Final = [{"role": "user", "content": "shared prompt"}]
collection: Final = f"cache_{uuid4().hex}"
facade: Final = qdrant_facade(qdrant_url, collection)
facade.cache.set_cache(
"python-key",
{"timestamp": time.time(), "response": json.dumps({"id": "py"})},
messages=messages,
)
handle: Final = _native._CacheTestHandle.qdrant_semantic(
qdrant_url,
collection_name=collection,
similarity_threshold=0.99,
vector_size=8,
)
handle._bind_facade(facade)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve()
assert binding.kind == "native"
assert binding.lookup(semantic_request("python-key", messages)) == {"id": "py"}
binding.store(semantic_request("native-key", messages), {"id": "native"})
assert facade.cache.get_cache("native-key", messages=messages) == {"id": "native"}
unrelated: Final = [{"role": "user", "content": "unrelated prompt"}]
assert binding.lookup(semantic_request("native-key", unrelated)) is None
assert facade.cache.get_cache("native-key", messages=unrelated) is None
assert binding.lookup(semantic_request("different-key", messages)) is None
assert facade.cache.get_cache("different-key", messages=messages) is None
async def test_qdrant_semantic_async_parity(
qdrant_url: str,
fake_embedding_endpoint: str,
) -> None:
del fake_embedding_endpoint
messages: Final = [{"role": "user", "content": "async prompt"}]
collection: Final = f"cache_{uuid4().hex}"
facade: Final = qdrant_facade(qdrant_url, collection)
handle: Final = _native._CacheTestHandle.qdrant_semantic(
qdrant_url,
collection_name=collection,
similarity_threshold=0.99,
vector_size=8,
)
handle._bind_facade(facade)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve()
await facade.cache.async_set_cache(
"python-key",
{"timestamp": time.time(), "response": json.dumps({"id": "py"})},
messages=messages,
)
assert await binding.async_lookup(semantic_request("python-key", messages)) == {"id": "py"}
await binding.async_store(semantic_request("native-key", messages), {"id": "native"})
assert await facade.cache.async_get_cache("native-key", messages=messages) == {"id": "native"}
async def test_qdrant_semantic_malformed_entries_and_unsupported_operations(
qdrant_url: str,
fake_embedding_endpoint: str,
) -> None:
del fake_embedding_endpoint
messages: Final = [{"role": "user", "content": "malformed prompt"}]
collection: Final = f"cache_{uuid4().hex}"
facade: Final = qdrant_facade(qdrant_url, collection)
handle: Final = _native._CacheTestHandle.qdrant_semantic(
qdrant_url,
collection_name=collection,
similarity_threshold=0.99,
vector_size=8,
)
handle._bind_facade(facade)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve()
key: Final = "malformed-key"
response: Final = {
"points": [
{
"id": str(uuid4()),
"vector": embedding_vector("malformed prompt"),
"payload": {
"litellm_cache_key": key,
"text": "malformed prompt",
"response": "not json",
},
}
]
}
facade.cache.sync_client.put(
url=f"{qdrant_url}/collections/{collection}/points",
headers=facade.cache.headers,
json=response,
)
assert binding.lookup(semantic_request(key, messages)) is None
with pytest.raises(RuntimeError, match="does not support"):
binding.lookup_batch([semantic_request(key, messages)])
with pytest.raises(RuntimeError, match="does not support"):
await binding.async_flush()
with pytest.raises(RuntimeError, match="does not support"):
await binding.ping()
def test_qdrant_semantic_ignores_request_expiry(
qdrant_url: str,
fake_embedding_endpoint: str,
) -> None:
del fake_embedding_endpoint
messages: Final = [{"role": "user", "content": "persistent prompt"}]
collection: Final = f"cache_{uuid4().hex}"
facade: Final = qdrant_facade(qdrant_url, collection)
handle: Final = _native._CacheTestHandle.qdrant_semantic(
qdrant_url,
collection_name=collection,
similarity_threshold=0.99,
vector_size=8,
)
handle._bind_facade(facade)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve()
binding.store(
semantic_request("persistent-key", messages, ttl_seconds=1.0),
{"id": "persistent"},
)
time.sleep(1.2)
assert binding.lookup(semantic_request("persistent-key", messages)) == {"id": "persistent"}
assert facade.cache.get_cache("persistent-key", messages=messages) == {"id": "persistent"}
def test_qdrant_semantic_mutation_and_projection_fallback(
qdrant_url: str,
fake_embedding_endpoint: str,
) -> None:
del fake_embedding_endpoint
collection: Final = f"cache_{uuid4().hex}"
facade: Final = qdrant_facade(qdrant_url, collection)
handle: Final = _native._CacheTestHandle.qdrant_semantic(
qdrant_url,
collection_name=collection,
similarity_threshold=0.99,
vector_size=8,
)
handle._bind_facade(facade)
facade.cache.similarity_threshold = 0.5
assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback"
unsupported: Final = qdrant_facade(qdrant_url, f"cache_{uuid4().hex}")
unsupported.cache.embedding_max_input_tokens = 100
with pytest.raises(TypeError, match="requires Python"):
handle._bind_facade(unsupported)
unsupported.cache.embedding_max_input_tokens = None
unsupported.cache.qdrant_api_base = "http://127.0.0.1:7777"
with pytest.raises(TypeError, match="gRPC"):
handle._bind_facade(unsupported)
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
class CustomQdrantSemanticCache(QdrantSemanticCache):
pass
subclass_facade: Final = qdrant_facade(qdrant_url, f"cache_{uuid4().hex}")
subclass_facade.cache = CustomQdrantSemanticCache(
qdrant_api_base=qdrant_url,
collection_name=subclass_facade.cache.collection_name,
similarity_threshold=0.99,
embedding_model="text-embedding-3-small",
vector_size=8,
)
with pytest.raises(TypeError):
handle._bind_facade(subclass_facade)