mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
Merge pull request #42324 from BerriAI/litellm_rust_qdrant_semantic_cache
Some checks are pending
CI Coverage / assert-ci-coverage (push) Waiting to run
CodeQL / Analyze (actions) (push) Waiting to run
CodeQL / Analyze (javascript-typescript) (push) Waiting to run
CodeQL / Analyze (python) (push) Waiting to run
CodSpeed Benchmarks / benchmarks (push) Waiting to run
Helm unit test / unit-test (push) Waiting to run
Publish basedpyright base counts / publish (push) Waiting to run
Scorecard supply-chain security / Scorecard analysis (push) Waiting to run
Code Quality Checks / code-quality (push) Waiting to run
Code Quality Checks / python-310-import-smoke (push) Waiting to run
UI Unit Tests / ui-unit-tests (push) Waiting to run
Postgres Tests / proxy-security (push) Waiting to run
Postgres Tests / schema-migration (push) Waiting to run
Postgres Tests / proxy-behavior (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run
Terraform Provider / gofmt, vet, build, test (push) Waiting to run
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Waiting to run
Unit Tests: Documentation Validation / documentation (push) Waiting to run
Unit Tests: Proxy DB Operations / assert-shard-coverage (push) Waiting to run
Unit Tests: Proxy DB Operations / auth-checks (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / budgets (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / custom-logging (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / db-and-spend (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / endpoints-and-responses (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / guardrails-hooks (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / jwt-and-keys (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / key-generation (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / logging-misc (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-runtime (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-server-core (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-utils (push) Blocked by required conditions
Unit Tests / misc (push) Waiting to run
Unit Tests / caching-local (push) Waiting to run
Unit Tests / core-utils (push) Waiting to run
Unit Tests / enterprise-package (push) Waiting to run
Unit Tests / enterprise-routing (push) Waiting to run
Unit Tests / integrations (push) Waiting to run
Unit Tests / All Other Providers (push) Waiting to run
Unit Tests / Vertex AI (push) Waiting to run
Unit Tests / mcp-integration (push) Waiting to run
Unit Tests / proxy-auth (push) Waiting to run
Unit Tests / proxy-endpoints (push) Waiting to run
Unit Tests / proxy-extras (push) Waiting to run
Unit Tests / proxy-infra (push) Waiting to run
Unit Tests / proxy-server (push) Waiting to run
Unit Tests / responses-caching-types (push) Waiting to run
GitHub Actions Security Analysis / zizmor (push) Waiting to run
Some checks are pending
CI Coverage / assert-ci-coverage (push) Waiting to run
CodeQL / Analyze (actions) (push) Waiting to run
CodeQL / Analyze (javascript-typescript) (push) Waiting to run
CodeQL / Analyze (python) (push) Waiting to run
CodSpeed Benchmarks / benchmarks (push) Waiting to run
Helm unit test / unit-test (push) Waiting to run
Publish basedpyright base counts / publish (push) Waiting to run
Scorecard supply-chain security / Scorecard analysis (push) Waiting to run
Code Quality Checks / code-quality (push) Waiting to run
Code Quality Checks / python-310-import-smoke (push) Waiting to run
UI Unit Tests / ui-unit-tests (push) Waiting to run
Postgres Tests / proxy-security (push) Waiting to run
Postgres Tests / schema-migration (push) Waiting to run
Postgres Tests / proxy-behavior (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run
Terraform Provider / gofmt, vet, build, test (push) Waiting to run
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Waiting to run
Unit Tests: Documentation Validation / documentation (push) Waiting to run
Unit Tests: Proxy DB Operations / assert-shard-coverage (push) Waiting to run
Unit Tests: Proxy DB Operations / auth-checks (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / budgets (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / custom-logging (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / db-and-spend (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / endpoints-and-responses (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / guardrails-hooks (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / jwt-and-keys (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / key-generation (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / logging-misc (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-runtime (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-server-core (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-utils (push) Blocked by required conditions
Unit Tests / misc (push) Waiting to run
Unit Tests / caching-local (push) Waiting to run
Unit Tests / core-utils (push) Waiting to run
Unit Tests / enterprise-package (push) Waiting to run
Unit Tests / enterprise-routing (push) Waiting to run
Unit Tests / integrations (push) Waiting to run
Unit Tests / All Other Providers (push) Waiting to run
Unit Tests / Vertex AI (push) Waiting to run
Unit Tests / mcp-integration (push) Waiting to run
Unit Tests / proxy-auth (push) Waiting to run
Unit Tests / proxy-endpoints (push) Waiting to run
Unit Tests / proxy-extras (push) Waiting to run
Unit Tests / proxy-infra (push) Waiting to run
Unit Tests / proxy-server (push) Waiting to run
Unit Tests / responses-caching-types (push) Waiting to run
GitHub Actions Security Analysis / zizmor (push) Waiting to run
This commit is contained in:
yujonglee
GitHub
commit
cc9970efbf
24 changed files with 2249 additions and 30 deletions
119
litellm-rust/Cargo.lock
generated
119
litellm-rust/Cargo.lock
generated
|
|
@ -650,6 +650,49 @@ dependencies = [
|
|||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "axum"
|
||||
version = "0.8.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90"
|
||||
dependencies = [
|
||||
"axum-core",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"http 1.4.2",
|
||||
"http-body 1.1.0",
|
||||
"http-body-util",
|
||||
"itoa",
|
||||
"matchit",
|
||||
"memchr",
|
||||
"mime",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"serde_core",
|
||||
"sync_wrapper",
|
||||
"tower",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "axum-core"
|
||||
version = "0.5.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"futures-core",
|
||||
"http 1.4.2",
|
||||
"http-body 1.1.0",
|
||||
"http-body-util",
|
||||
"mime",
|
||||
"pin-project-lite",
|
||||
"sync_wrapper",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "azure_core"
|
||||
version = "1.1.0"
|
||||
|
|
@ -1439,7 +1482,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -2241,7 +2284,7 @@ dependencies = [
|
|||
"libc",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"socket2 0.5.10",
|
||||
"socket2 0.6.5",
|
||||
"tokio",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
|
|
@ -2729,6 +2772,26 @@ dependencies = [
|
|||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-cache-qdrant-semantic"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"litellm-cache",
|
||||
"litellm-cache-response",
|
||||
"qdrant-client",
|
||||
"reqwest 0.12.28",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
"tonic",
|
||||
"tonic-prost",
|
||||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-cache-redis"
|
||||
version = "0.1.0"
|
||||
|
|
@ -2972,6 +3035,7 @@ dependencies = [
|
|||
"litellm-cache-disk",
|
||||
"litellm-cache-gcs",
|
||||
"litellm-cache-memory",
|
||||
"litellm-cache-qdrant-semantic",
|
||||
"litellm-cache-redis",
|
||||
"litellm-cache-redis-semantic",
|
||||
"litellm-cache-response",
|
||||
|
|
@ -2987,7 +3051,9 @@ dependencies = [
|
|||
"litellm-types",
|
||||
"pyo3",
|
||||
"pyo3-async-runtimes",
|
||||
"qdrant-client",
|
||||
"redis",
|
||||
"reqwest 0.12.28",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -2995,6 +3061,7 @@ dependencies = [
|
|||
"sha2 0.10.9",
|
||||
"tokio",
|
||||
"tokio-tungstenite",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -3251,6 +3318,12 @@ version = "0.2.3"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c"
|
||||
|
||||
[[package]]
|
||||
name = "matchit"
|
||||
version = "0.8.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3"
|
||||
|
||||
[[package]]
|
||||
name = "md-5"
|
||||
version = "0.11.0"
|
||||
|
|
@ -3874,6 +3947,27 @@ dependencies = [
|
|||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "qdrant-client"
|
||||
version = "1.19.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dddc19df129bad7346ebd027288621ab1ac7e52678371f906b9a8622d7aaf87e"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"derive_builder",
|
||||
"futures",
|
||||
"parking_lot",
|
||||
"prost",
|
||||
"prost-types",
|
||||
"semver",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tonic",
|
||||
"tonic-prost",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "quick-error"
|
||||
version = "1.2.3"
|
||||
|
|
@ -3903,7 +3997,7 @@ dependencies = [
|
|||
"quinn-udp",
|
||||
"rustc-hash",
|
||||
"rustls 0.23.42",
|
||||
"socket2 0.5.10",
|
||||
"socket2 0.6.5",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tracing",
|
||||
|
|
@ -3942,9 +4036,9 @@ dependencies = [
|
|||
"cfg_aliases",
|
||||
"libc",
|
||||
"once_cell",
|
||||
"socket2 0.5.10",
|
||||
"socket2 0.6.5",
|
||||
"tracing",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -4430,7 +4524,7 @@ dependencies = [
|
|||
"errno",
|
||||
"libc",
|
||||
"linux-raw-sys",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -4501,7 +4595,7 @@ dependencies = [
|
|||
"security-framework",
|
||||
"security-framework-sys",
|
||||
"webpki-root-certs",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -5073,10 +5167,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
|
||||
dependencies = [
|
||||
"fastrand",
|
||||
"getrandom 0.3.4",
|
||||
"getrandom 0.4.3",
|
||||
"once_cell",
|
||||
"rustix",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -5365,8 +5459,12 @@ version = "0.14.6"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"axum",
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"flate2",
|
||||
"h2 0.4.15",
|
||||
"http 1.4.2",
|
||||
"http-body 1.1.0",
|
||||
"http-body-util",
|
||||
|
|
@ -5376,6 +5474,7 @@ dependencies = [
|
|||
"percent-encoding",
|
||||
"pin-project",
|
||||
"rustls-native-certs",
|
||||
"socket2 0.6.5",
|
||||
"sync_wrapper",
|
||||
"tokio",
|
||||
"tokio-rustls 0.26.4",
|
||||
|
|
@ -5949,7 +6048,7 @@ version = "0.1.11"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
|
||||
dependencies = [
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
|
|||
|
|
@ -38,6 +38,7 @@ litellm-cache-gcs = { path = "crates/cache-gcs" }
|
|||
litellm-cache-disk = { path = "crates/cache-disk" }
|
||||
litellm-cache-redis-semantic = { path = "crates/cache-redis-semantic" }
|
||||
litellm-cache-response = { path = "crates/cache-response" }
|
||||
litellm-cache-qdrant-semantic = { path = "crates/cache-qdrant-semantic" }
|
||||
litellm-token-counter = { path = "crates/token-counter" }
|
||||
litellm-token-counter-fast = { path = "crates/token-counter-fast" }
|
||||
litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" }
|
||||
|
|
@ -55,6 +56,8 @@ pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
|
|||
pythonize = "0.29.0"
|
||||
rand = "0.8"
|
||||
reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] }
|
||||
qdrant-client = { version = "1.19.0", default-features = false }
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
rstest = "0.26.1"
|
||||
rstest_reuse = "0.7.0"
|
||||
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
|
||||
|
|
|
|||
24
litellm-rust/crates/cache-qdrant-semantic/Cargo.toml
Normal file
24
litellm-rust/crates/cache-qdrant-semantic/Cargo.toml
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
[package]
|
||||
name = "litellm-cache-qdrant-semantic"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
futures-util.workspace = true
|
||||
litellm-cache.workspace = true
|
||||
qdrant-client = { workspace = true, features = ["serde"] }
|
||||
reqwest.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
tokio.workspace = true
|
||||
uuid.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-cache-response.workspace = true
|
||||
rstest.workspace = true
|
||||
tonic = "0.14"
|
||||
tonic-prost = "0.14"
|
||||
tokio-stream = "0.1"
|
||||
75
litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs
Normal file
75
litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs
Normal file
|
|
@ -0,0 +1,75 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::Error;
|
||||
use reqwest::Client;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::Embedder;
|
||||
|
||||
pub struct OpenAiEmbedder {
|
||||
client: Client,
|
||||
api_base: String,
|
||||
api_key: String,
|
||||
model: String,
|
||||
timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
pub struct OpenAiEmbedderConfig {
|
||||
pub api_base: String,
|
||||
pub api_key: String,
|
||||
pub model: String,
|
||||
pub timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
impl OpenAiEmbedder {
|
||||
pub fn new(client: Client, config: OpenAiEmbedderConfig) -> Self {
|
||||
Self {
|
||||
client,
|
||||
api_base: config.api_base.trim_end_matches('/').to_owned(),
|
||||
api_key: config.api_key,
|
||||
model: config.model,
|
||||
timeout: config.timeout,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Embedder for OpenAiEmbedder {
|
||||
fn model(&self) -> &str {
|
||||
&self.model
|
||||
}
|
||||
|
||||
async fn embed(&self, input: &str) -> Result<Vec<f32>, Error> {
|
||||
let request = self
|
||||
.client
|
||||
.post(format!("{}/embeddings", self.api_base))
|
||||
.bearer_auth(&self.api_key)
|
||||
.json(&serde_json::json!({
|
||||
"model": self.model,
|
||||
"input": input,
|
||||
"encoding_format": "float",
|
||||
}));
|
||||
let response = if let Some(timeout) = self.timeout {
|
||||
request.timeout(timeout)
|
||||
} else {
|
||||
request
|
||||
}
|
||||
.send()
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.error_for_status()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
let body: Value = response.json().await.map_err(|_| Error::Unavailable)?;
|
||||
body.get("data")
|
||||
.and_then(Value::as_array)
|
||||
.and_then(|data| data.first())
|
||||
.and_then(|item| item.get("embedding"))
|
||||
.and_then(Value::as_array)
|
||||
.and_then(|embedding| {
|
||||
embedding
|
||||
.iter()
|
||||
.map(|value| value.as_f64().map(|value| value as f32))
|
||||
.collect::<Option<Vec<_>>>()
|
||||
})
|
||||
.ok_or(Error::Unavailable)
|
||||
}
|
||||
}
|
||||
7
litellm-rust/crates/cache-qdrant-semantic/src/lib.rs
Normal file
7
litellm-rust/crates/cache-qdrant-semantic/src/lib.rs
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
mod embedder;
|
||||
mod prompt;
|
||||
mod semantic;
|
||||
|
||||
pub use embedder::{OpenAiEmbedder, OpenAiEmbedderConfig};
|
||||
pub use prompt::prompt_from_messages;
|
||||
pub use semantic::{Embedder, QdrantSemanticCache, QdrantSemanticConfig, Quantization};
|
||||
59
litellm-rust/crates/cache-qdrant-semantic/src/prompt.rs
Normal file
59
litellm-rust/crates/cache-qdrant-semantic/src/prompt.rs
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
use serde_json::Value;
|
||||
|
||||
fn search_results_text(search_results: Option<&Value>) -> String {
|
||||
let Some(Value::Array(results)) = search_results else {
|
||||
return String::new();
|
||||
};
|
||||
results
|
||||
.iter()
|
||||
.filter_map(Value::as_object)
|
||||
.flat_map(|result| {
|
||||
let source = result
|
||||
.get("source")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned);
|
||||
let title = result
|
||||
.get("title")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned);
|
||||
let content = result
|
||||
.get("content")
|
||||
.and_then(Value::as_array)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter_map(Value::as_object)
|
||||
.filter_map(|block| block.get("text").and_then(Value::as_str).map(str::to_owned));
|
||||
let citations = result
|
||||
.get("citations")
|
||||
.filter(|value| !value.is_null())
|
||||
.map(|value| serde_json::to_string(value).unwrap_or_default());
|
||||
source
|
||||
.into_iter()
|
||||
.chain(title)
|
||||
.chain(content)
|
||||
.chain(citations)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn prompt_from_messages(messages: &[Value]) -> String {
|
||||
messages
|
||||
.iter()
|
||||
.filter_map(Value::as_object)
|
||||
.map(|message| {
|
||||
let content = match message.get("content") {
|
||||
Some(Value::String(content)) => content.clone(),
|
||||
Some(Value::Array(parts)) => parts
|
||||
.iter()
|
||||
.filter_map(Value::as_object)
|
||||
.filter_map(|part| part.get("text").and_then(Value::as_str))
|
||||
.collect(),
|
||||
_ => String::new(),
|
||||
};
|
||||
format!(
|
||||
"{content}{}",
|
||||
search_results_text(message.get("search_results"))
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
262
litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs
Normal file
262
litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs
Normal file
|
|
@ -0,0 +1,262 @@
|
|||
use std::future::Future;
|
||||
|
||||
use futures_util::future::try_join_all;
|
||||
use litellm_cache::{BaseCache, CacheCodec, CacheConnectionResult, Error, SemanticCacheContext};
|
||||
use qdrant_client::{
|
||||
Payload, Qdrant,
|
||||
qdrant::{
|
||||
BinaryQuantizationBuilder, CompressionRatio, Condition, CreateCollectionBuilder,
|
||||
CreateFieldIndexCollectionBuilder, Distance, FieldType, Filter, PointStruct,
|
||||
ProductQuantizationBuilder, QuantizationSearchParamsBuilder, ScalarQuantizationBuilder,
|
||||
SearchParamsBuilder, SearchPointsBuilder, UpsertPointsBuilder, VectorParamsBuilder,
|
||||
},
|
||||
};
|
||||
use serde_json::{Map, Value, json};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::prompt_from_messages;
|
||||
|
||||
pub trait Embedder: Send + Sync + 'static {
|
||||
fn model(&self) -> &str;
|
||||
fn embed(&self, input: &str) -> impl Future<Output = Result<Vec<f32>, Error>> + Send;
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum Quantization {
|
||||
Binary,
|
||||
Scalar,
|
||||
Product,
|
||||
}
|
||||
|
||||
pub struct QdrantSemanticConfig {
|
||||
pub collection_name: String,
|
||||
pub similarity_threshold: f64,
|
||||
pub vector_size: u64,
|
||||
pub quantization: Quantization,
|
||||
}
|
||||
|
||||
pub struct QdrantSemanticCache<E: Embedder, C: CacheCodec> {
|
||||
client: Qdrant,
|
||||
embedder: E,
|
||||
codec: C,
|
||||
config: QdrantSemanticConfig,
|
||||
runtime: tokio::runtime::Handle,
|
||||
}
|
||||
|
||||
impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
|
||||
pub async fn connect(
|
||||
client: Qdrant,
|
||||
embedder: E,
|
||||
codec: C,
|
||||
config: QdrantSemanticConfig,
|
||||
runtime: tokio::runtime::Handle,
|
||||
) -> Result<Self, Error> {
|
||||
let exists = client
|
||||
.collection_exists(config.collection_name.clone())
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
if !exists {
|
||||
client
|
||||
.create_collection(
|
||||
CreateCollectionBuilder::new(config.collection_name.clone())
|
||||
.vectors_config(VectorParamsBuilder::new(
|
||||
config.vector_size,
|
||||
Distance::Cosine,
|
||||
))
|
||||
.quantization_config(quantization(&config.quantization)),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
}
|
||||
let _ = client
|
||||
.create_field_index(CreateFieldIndexCollectionBuilder::new(
|
||||
config.collection_name.clone(),
|
||||
"litellm_cache_key".to_owned(),
|
||||
FieldType::Keyword,
|
||||
))
|
||||
.await;
|
||||
Ok(Self {
|
||||
client,
|
||||
embedder,
|
||||
codec,
|
||||
config,
|
||||
runtime,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn collection_name(&self) -> &str {
|
||||
&self.config.collection_name
|
||||
}
|
||||
|
||||
pub fn similarity_threshold(&self) -> f64 {
|
||||
self.config.similarity_threshold
|
||||
}
|
||||
|
||||
pub fn vector_size(&self) -> u64 {
|
||||
self.config.vector_size
|
||||
}
|
||||
|
||||
pub fn embedder(&self) -> &E {
|
||||
&self.embedder
|
||||
}
|
||||
|
||||
fn prompt(context: &SemanticCacheContext) -> Result<String, Error> {
|
||||
let Some(messages) = context.messages.as_ref().and_then(Value::as_array) else {
|
||||
return Err(Error::MissingPrompt);
|
||||
};
|
||||
if messages.is_empty() {
|
||||
return Err(Error::MissingPrompt);
|
||||
}
|
||||
Ok(prompt_from_messages(messages))
|
||||
}
|
||||
|
||||
async fn set(
|
||||
&self,
|
||||
key: &str,
|
||||
value: C::Value,
|
||||
context: &SemanticCacheContext,
|
||||
) -> Result<(), Error> {
|
||||
let prompt = Self::prompt(context)?;
|
||||
let vector = self.embedder.embed(&prompt).await?;
|
||||
let response =
|
||||
String::from_utf8(self.codec.encode(&value)?).map_err(|_| Error::InvalidEntry)?;
|
||||
let payload = Payload::try_from(json!({
|
||||
"litellm_cache_key": key,
|
||||
"text": prompt,
|
||||
"response": response,
|
||||
}))
|
||||
.map_err(|_| Error::InvalidEntry)?;
|
||||
self.client
|
||||
.upsert_points(
|
||||
UpsertPointsBuilder::new(
|
||||
self.collection_name(),
|
||||
vec![PointStruct::new(
|
||||
Uuid::new_v4().to_string(),
|
||||
vector,
|
||||
payload,
|
||||
)],
|
||||
)
|
||||
.wait(true),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get(
|
||||
&self,
|
||||
key: &str,
|
||||
context: &SemanticCacheContext,
|
||||
) -> Result<Option<C::Value>, Error> {
|
||||
let prompt = Self::prompt(context)?;
|
||||
let vector = self.embedder.embed(&prompt).await?;
|
||||
let result = self
|
||||
.client
|
||||
.search_points(
|
||||
SearchPointsBuilder::new(self.collection_name(), vector, 1)
|
||||
.with_payload(true)
|
||||
.filter(Filter::must([Condition::matches(
|
||||
"litellm_cache_key",
|
||||
key.to_owned(),
|
||||
)]))
|
||||
.params(
|
||||
SearchParamsBuilder::default().quantization(
|
||||
QuantizationSearchParamsBuilder::default()
|
||||
.ignore(false)
|
||||
.rescore(true)
|
||||
.oversampling(3.0),
|
||||
),
|
||||
),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
let Some(point) = result.result.into_iter().next() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let payload: Map<String, Value> = Payload::from(point.payload).into();
|
||||
if payload.get("litellm_cache_key").and_then(Value::as_str) != Some(key) {
|
||||
return Ok(None);
|
||||
}
|
||||
if f64::from(point.score) < self.config.similarity_threshold {
|
||||
return Ok(None);
|
||||
}
|
||||
let response = payload
|
||||
.get("response")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or(Error::InvalidEntry)?;
|
||||
self.codec.decode(response.as_bytes()).map(Some)
|
||||
}
|
||||
}
|
||||
|
||||
fn quantization(value: &Quantization) -> qdrant_client::qdrant::quantization_config::Quantization {
|
||||
match value {
|
||||
Quantization::Binary => BinaryQuantizationBuilder::new(false).into(),
|
||||
Quantization::Scalar => ScalarQuantizationBuilder::default()
|
||||
.quantile(0.99)
|
||||
.always_ram(false)
|
||||
.into(),
|
||||
Quantization::Product => ProductQuantizationBuilder::new(CompressionRatio::X16.into())
|
||||
.always_ram(false)
|
||||
.into(),
|
||||
}
|
||||
}
|
||||
|
||||
impl<E: Embedder, C: CacheCodec> BaseCache for QdrantSemanticCache<E, C> {
|
||||
type Value = C::Value;
|
||||
type Context = SemanticCacheContext;
|
||||
|
||||
fn get_ttl(&self, _: &Self::Context) -> Option<std::time::Duration> {
|
||||
None
|
||||
}
|
||||
|
||||
fn set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: &Self::Context,
|
||||
) -> Result<(), Error> {
|
||||
self.runtime.block_on(self.set(key, value, context))
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
|
||||
self.runtime.block_on(self.get(key, context))
|
||||
}
|
||||
|
||||
async fn async_set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: Self::Context,
|
||||
) -> Result<(), Error> {
|
||||
self.set(key, value, &context).await
|
||||
}
|
||||
|
||||
async fn async_get_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
context: &Self::Context,
|
||||
) -> Result<Option<Self::Value>, Error> {
|
||||
self.get(key, context).await
|
||||
}
|
||||
|
||||
async fn async_set_cache_pipeline(
|
||||
&self,
|
||||
entries: Vec<(String, Self::Value)>,
|
||||
context: Self::Context,
|
||||
) -> Result<(), Error> {
|
||||
try_join_all(entries.into_iter().map(|(key, value)| {
|
||||
let context = context.clone();
|
||||
async move { self.async_set_cache(&key, value, context).await }
|
||||
}))
|
||||
.await
|
||||
.map(|_| ())
|
||||
}
|
||||
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
Err(Error::UnsupportedOperation)
|
||||
}
|
||||
}
|
||||
166
litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs
Normal file
166
litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs
Normal file
|
|
@ -0,0 +1,166 @@
|
|||
use std::{
|
||||
sync::{Arc, Mutex},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::Error;
|
||||
use litellm_cache_qdrant_semantic::{Embedder, OpenAiEmbedder, OpenAiEmbedderConfig};
|
||||
use serde_json::Value;
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
net::TcpListener,
|
||||
};
|
||||
|
||||
struct TestHttpServer {
|
||||
address: std::net::SocketAddr,
|
||||
request: Arc<Mutex<Option<Vec<u8>>>>,
|
||||
task: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
impl TestHttpServer {
|
||||
async fn response(status: &str, body: &str) -> Self {
|
||||
Self::response_after(status, body, Duration::ZERO).await
|
||||
}
|
||||
|
||||
async fn response_after(status: &str, body: &str, delay: Duration) -> Self {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let request = Arc::new(Mutex::new(None));
|
||||
let captured = request.clone();
|
||||
let status = status.to_owned();
|
||||
let body = body.to_owned();
|
||||
let task = tokio::spawn(async move {
|
||||
let (mut stream, _) = listener.accept().await.unwrap();
|
||||
let request_bytes = read_request(&mut stream).await;
|
||||
*captured.lock().unwrap() = Some(request_bytes);
|
||||
tokio::time::sleep(delay).await;
|
||||
let response = format!(
|
||||
"HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
|
||||
body.len()
|
||||
);
|
||||
stream.write_all(response.as_bytes()).await.unwrap();
|
||||
});
|
||||
Self {
|
||||
address,
|
||||
request,
|
||||
task,
|
||||
}
|
||||
}
|
||||
|
||||
fn base_url(&self) -> String {
|
||||
format!("http://{}", self.address)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for TestHttpServer {
|
||||
fn drop(&mut self) {
|
||||
self.task.abort();
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_request(stream: &mut tokio::net::TcpStream) -> Vec<u8> {
|
||||
let mut bytes = Vec::new();
|
||||
let header_end = loop {
|
||||
let mut chunk = [0_u8; 1024];
|
||||
let count = stream.read(&mut chunk).await.unwrap();
|
||||
assert_ne!(count, 0);
|
||||
bytes.extend_from_slice(&chunk[..count]);
|
||||
if let Some(end) = bytes.windows(4).position(|window| window == b"\r\n\r\n") {
|
||||
break end + 4;
|
||||
}
|
||||
};
|
||||
let headers = String::from_utf8_lossy(&bytes[..header_end]);
|
||||
let content_length = headers
|
||||
.lines()
|
||||
.find_map(|line| {
|
||||
line.split_once(':')
|
||||
.filter(|(name, _)| name.eq_ignore_ascii_case("content-length"))
|
||||
.map(|(_, value)| value.trim())
|
||||
})
|
||||
.unwrap()
|
||||
.parse::<usize>()
|
||||
.unwrap();
|
||||
while bytes.len() < header_end + content_length {
|
||||
let mut chunk = [0_u8; 1024];
|
||||
let count = stream.read(&mut chunk).await.unwrap();
|
||||
assert_ne!(count, 0);
|
||||
bytes.extend_from_slice(&chunk[..count]);
|
||||
}
|
||||
bytes
|
||||
}
|
||||
|
||||
fn config(base: String, timeout: Option<Duration>) -> OpenAiEmbedderConfig {
|
||||
OpenAiEmbedderConfig {
|
||||
api_base: base,
|
||||
api_key: "test-key".to_owned(),
|
||||
model: "test-model".to_owned(),
|
||||
timeout,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn posts_embeddings_request_and_parses_vector() {
|
||||
let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await;
|
||||
let embedder = OpenAiEmbedder::new(
|
||||
reqwest::Client::new(),
|
||||
config(
|
||||
format!("{}/", server.base_url()),
|
||||
Some(Duration::from_secs(1)),
|
||||
),
|
||||
);
|
||||
assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]);
|
||||
let request = server.request.lock().unwrap().clone().unwrap();
|
||||
let request_text = String::from_utf8(request).unwrap();
|
||||
assert!(request_text.starts_with("POST /embeddings HTTP/1.1\r\n"));
|
||||
assert!(request_text.contains("\r\nauthorization: Bearer test-key\r\n"));
|
||||
let body = request_text.split("\r\n\r\n").nth(1).unwrap();
|
||||
let body: Value = serde_json::from_str(body).unwrap();
|
||||
assert_eq!(body["model"], "test-model");
|
||||
assert_eq!(body["input"], "hello");
|
||||
assert_eq!(body["encoding_format"], "float");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn status_and_timeout_errors_are_unavailable() {
|
||||
let server = TestHttpServer::response("500 Internal Server Error", "{}").await;
|
||||
let embedder = OpenAiEmbedder::new(reqwest::Client::new(), config(server.base_url(), None));
|
||||
assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable));
|
||||
|
||||
let server = TestHttpServer::response_after(
|
||||
"200 OK",
|
||||
r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
|
||||
Duration::from_millis(500),
|
||||
)
|
||||
.await;
|
||||
let embedder = OpenAiEmbedder::new(
|
||||
reqwest::Client::new(),
|
||||
config(server.base_url(), Some(Duration::from_millis(200))),
|
||||
);
|
||||
assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable));
|
||||
|
||||
let server = TestHttpServer::response_after(
|
||||
"200 OK",
|
||||
r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
|
||||
Duration::from_millis(100),
|
||||
)
|
||||
.await;
|
||||
let embedder = OpenAiEmbedder::new(
|
||||
reqwest::Client::new(),
|
||||
config(server.base_url(), Some(Duration::from_secs(1))),
|
||||
);
|
||||
assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn uses_the_injected_client() {
|
||||
let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await;
|
||||
let client = reqwest::Client::builder()
|
||||
.user_agent("litellm-embedder-test")
|
||||
.build()
|
||||
.unwrap();
|
||||
let embedder = OpenAiEmbedder::new(client, config(server.base_url(), None));
|
||||
assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]);
|
||||
let request = server.request.lock().unwrap().clone().unwrap();
|
||||
let request_text = String::from_utf8(request).unwrap();
|
||||
assert!(request_text.contains("\r\nuser-agent: litellm-embedder-test\r\n"));
|
||||
}
|
||||
38
litellm-rust/crates/cache-qdrant-semantic/tests/prompt.rs
Normal file
38
litellm-rust/crates/cache-qdrant-semantic/tests/prompt.rs
Normal file
|
|
@ -0,0 +1,38 @@
|
|||
use litellm_cache_qdrant_semantic::prompt_from_messages;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn prompt_matches_python_message_content_rules() {
|
||||
let messages = vec![
|
||||
json!({"role": "user", "content": "hello"}),
|
||||
json!({
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "world"},
|
||||
{"type": "image_url", "image_url": {"url": "ignored"}},
|
||||
{"type": "text", "text": "!"},
|
||||
],
|
||||
}),
|
||||
];
|
||||
|
||||
assert_eq!(prompt_from_messages(&messages), "helloworld!");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prompt_includes_search_result_text_and_compact_citations() {
|
||||
let messages = vec![json!({
|
||||
"role": "tool",
|
||||
"content": null,
|
||||
"search_results": [{
|
||||
"source": "source",
|
||||
"title": "title",
|
||||
"content": [{"text": "body"}],
|
||||
"citations": {"page": 1, "section": "intro"},
|
||||
}],
|
||||
})];
|
||||
|
||||
assert_eq!(
|
||||
prompt_from_messages(&messages),
|
||||
r#"sourcetitlebody{"page":1,"section":"intro"}"#
|
||||
);
|
||||
}
|
||||
422
litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs
Normal file
422
litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs
Normal file
|
|
@ -0,0 +1,422 @@
|
|||
#[path = "support/mod.rs"]
|
||||
mod support;
|
||||
|
||||
use std::{collections::HashMap, sync::Arc, time::Duration};
|
||||
|
||||
use litellm_cache::{BaseCache, CacheCodec, CacheContext, Error, SemanticCacheContext};
|
||||
use litellm_cache_qdrant_semantic::{
|
||||
Embedder, QdrantSemanticCache, QdrantSemanticConfig, Quantization,
|
||||
};
|
||||
use litellm_cache_response::{
|
||||
CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest,
|
||||
};
|
||||
use qdrant_client::Payload;
|
||||
use qdrant_client::{
|
||||
Qdrant,
|
||||
qdrant::{self, CompressionRatio, Distance, PointId, QuantizationType, Value, VectorParams},
|
||||
};
|
||||
use serde_json::{Value as JsonValue, json};
|
||||
|
||||
use support::{FakeQdrant, FakeState, StoredPoint};
|
||||
|
||||
#[derive(Clone)]
|
||||
struct FixedEmbedder {
|
||||
vectors: Arc<HashMap<String, Vec<f32>>>,
|
||||
}
|
||||
|
||||
impl FixedEmbedder {
|
||||
fn new(vectors: impl IntoIterator<Item = (&'static str, Vec<f32>)>) -> Self {
|
||||
Self {
|
||||
vectors: Arc::new(
|
||||
vectors
|
||||
.into_iter()
|
||||
.map(|(prompt, vector)| (prompt.to_owned(), vector))
|
||||
.collect(),
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Embedder for FixedEmbedder {
|
||||
fn model(&self) -> &str {
|
||||
"fixed"
|
||||
}
|
||||
|
||||
async fn embed(&self, input: &str) -> Result<Vec<f32>, Error> {
|
||||
self.vectors.get(input).cloned().ok_or(Error::Unavailable)
|
||||
}
|
||||
}
|
||||
|
||||
fn config(quantization: Quantization) -> QdrantSemanticConfig {
|
||||
QdrantSemanticConfig {
|
||||
collection_name: "semantic".to_owned(),
|
||||
similarity_threshold: 0.9,
|
||||
vector_size: 2,
|
||||
quantization,
|
||||
}
|
||||
}
|
||||
|
||||
fn context(prompt: &str) -> SemanticCacheContext {
|
||||
SemanticCacheContext {
|
||||
messages: Some(json!([{"role": "user", "content": prompt}])),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn value(response: JsonValue) -> CacheEntry {
|
||||
CacheEntry {
|
||||
timestamp: Some(1.0),
|
||||
response,
|
||||
}
|
||||
}
|
||||
|
||||
async fn connect(
|
||||
server: &FakeQdrant,
|
||||
vectors: impl IntoIterator<Item = (&'static str, Vec<f32>)>,
|
||||
) -> QdrantSemanticCache<FixedEmbedder, ResponseCacheCodec> {
|
||||
let client = Qdrant::from_url(&server.url()).build().unwrap();
|
||||
QdrantSemanticCache::connect(
|
||||
client,
|
||||
FixedEmbedder::new(vectors),
|
||||
ResponseCacheCodec,
|
||||
config(Quantization::Binary),
|
||||
tokio::runtime::Handle::current(),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
#[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),
|
||||
(Quantization::Scalar, 1),
|
||||
(Quantization::Product, 2),
|
||||
] {
|
||||
let server = FakeQdrant::start(FakeState::default()).await;
|
||||
let client = Qdrant::from_url(&server.url()).build().unwrap();
|
||||
QdrantSemanticCache::connect(
|
||||
client,
|
||||
FixedEmbedder::new([]),
|
||||
ResponseCacheCodec,
|
||||
config(quantization),
|
||||
tokio::runtime::Handle::current(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let state = server.state.lock().unwrap();
|
||||
let request = &state.created_collections[0];
|
||||
let Some(qdrant::vectors_config::Config::Params(VectorParams { size, distance, .. })) =
|
||||
request
|
||||
.vectors_config
|
||||
.as_ref()
|
||||
.and_then(|config| config.config.clone())
|
||||
else {
|
||||
panic!("missing vector params");
|
||||
};
|
||||
assert_eq!(size, 2);
|
||||
assert_eq!(distance, Distance::Cosine as i32);
|
||||
let quantization_config = request
|
||||
.quantization_config
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.quantization
|
||||
.unwrap();
|
||||
match (expected, quantization_config) {
|
||||
(0, qdrant::quantization_config::Quantization::Binary(binary)) => {
|
||||
assert_eq!(binary.always_ram, Some(false));
|
||||
}
|
||||
(1, qdrant::quantization_config::Quantization::Scalar(scalar)) => {
|
||||
assert_eq!(scalar.r#type, QuantizationType::Int8 as i32);
|
||||
assert_eq!(scalar.quantile, Some(0.99));
|
||||
assert_eq!(scalar.always_ram, Some(false));
|
||||
}
|
||||
(2, qdrant::quantization_config::Quantization::Product(product)) => {
|
||||
assert_eq!(product.compression, CompressionRatio::X16 as i32);
|
||||
assert_eq!(product.always_ram, Some(false));
|
||||
}
|
||||
_ => panic!("unexpected quantization"),
|
||||
}
|
||||
assert!(state.index_creations >= 1);
|
||||
assert_eq!(state.field_indexes[0].collection_name, "semantic");
|
||||
assert_eq!(state.field_indexes[0].field_name, "litellm_cache_key");
|
||||
assert_eq!(
|
||||
state.field_indexes[0].field_type,
|
||||
Some(qdrant::FieldType::Keyword as i32)
|
||||
);
|
||||
server.stop();
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn existing_collection_skips_create_and_index_failure_is_non_fatal() {
|
||||
let server = FakeQdrant::start(FakeState {
|
||||
collections: ["semantic".to_owned()].into_iter().collect(),
|
||||
fail_field_index: true,
|
||||
..Default::default()
|
||||
})
|
||||
.await;
|
||||
let _cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
|
||||
let state = server.state.lock().unwrap();
|
||||
assert!(state.created_collections.is_empty());
|
||||
assert!(state.index_creations >= 1);
|
||||
server.stop();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn async_and_sync_set_get_store_exact_payload() {
|
||||
let server = FakeQdrant::start(FakeState::default()).await;
|
||||
let cache = Arc::new(connect(&server, [("hello", vec![1.0, 0.0])]).await);
|
||||
let ctx = context("hello");
|
||||
let entry = value(json!({"answer": 42}));
|
||||
cache
|
||||
.async_set_cache("key", entry.clone(), ctx.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.async_get_cache("key", &ctx).await.unwrap().as_ref(),
|
||||
Some(&entry)
|
||||
);
|
||||
{
|
||||
let state = server.state.lock().unwrap();
|
||||
let payload = &state.points[0].payload;
|
||||
let mut payload_keys = payload.keys().cloned().collect::<Vec<_>>();
|
||||
payload_keys.sort();
|
||||
assert_eq!(payload_keys, ["litellm_cache_key", "response", "text"]);
|
||||
assert_eq!(payload["litellm_cache_key"], Value::from("key"));
|
||||
assert_eq!(
|
||||
payload["response"],
|
||||
Value::from(String::from_utf8(ResponseCacheCodec.encode(&entry).unwrap()).unwrap())
|
||||
);
|
||||
}
|
||||
let sync_entry = entry.clone();
|
||||
let sync_cache = cache.clone();
|
||||
let sync_ctx = ctx.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
sync_cache
|
||||
.set_cache("sync", sync_entry.clone(), &sync_ctx)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
sync_cache.get_cache("sync", &sync_ctx).unwrap(),
|
||||
Some(sync_entry)
|
||||
);
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
server.stop();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn misses_and_payload_validation_are_safe() {
|
||||
let server = FakeQdrant::start(FakeState::default()).await;
|
||||
let cache = connect(
|
||||
&server,
|
||||
[("hello", vec![1.0, 0.0]), ("near", vec![0.7, 0.71414286])],
|
||||
)
|
||||
.await;
|
||||
let entry = value(json!({"answer": 1}));
|
||||
cache
|
||||
.async_set_cache("key", entry, context("hello"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_get_cache("other", &context("hello"))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_get_cache("key", &context("near"))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
server.insert_point(StoredPoint {
|
||||
id: Some(PointId::from(99_u64)),
|
||||
vector: vec![1.0, 0.0],
|
||||
payload: Payload::try_from(json!({
|
||||
"litellm_cache_key": 99,
|
||||
"response": "{}",
|
||||
}))
|
||||
.unwrap()
|
||||
.into(),
|
||||
});
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_get_cache("99", &context("hello"))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
server.stop();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn decoding_errors_missing_prompt_pipeline_and_ttl_behave_as_required() {
|
||||
let server = FakeQdrant::start(FakeState::default()).await;
|
||||
let cache = connect(&server, [("one", vec![1.0, 0.0]), ("two", vec![0.0, 1.0])]).await;
|
||||
let empty = SemanticCacheContext::default();
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_set_cache("key", value(json!({})), empty.clone())
|
||||
.await,
|
||||
Err(Error::MissingPrompt)
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_get_cache("key", &empty).await,
|
||||
Err(Error::MissingPrompt)
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_get_cache("key", &context("unknown")).await,
|
||||
Err(Error::Unavailable)
|
||||
);
|
||||
cache
|
||||
.async_set_cache(
|
||||
"ttl",
|
||||
value(json!({"ttl": true})),
|
||||
context("one").with_ttl(Some(Duration::from_secs(1))),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
tokio::time::sleep(Duration::from_millis(1_100)).await;
|
||||
assert!(
|
||||
cache
|
||||
.async_get_cache(
|
||||
"ttl",
|
||||
&context("one").with_ttl(Some(Duration::from_secs(1))),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
cache
|
||||
.async_set_cache_pipeline(
|
||||
vec![
|
||||
("one".to_owned(), value(json!({"n": 1}))),
|
||||
("two".to_owned(), value(json!({"n": 2}))),
|
||||
],
|
||||
context("one"),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
cache
|
||||
.async_get_cache("one", &context("one"))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
assert!(
|
||||
cache
|
||||
.async_get_cache("two", &context("one"))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
assert_eq!(
|
||||
server.state.lock().unwrap().upsert_waits,
|
||||
vec![Some(true), Some(true), Some(true)]
|
||||
);
|
||||
assert_eq!(cache.get_ttl(&context("one")), None);
|
||||
assert_eq!(
|
||||
cache.test_connection().await,
|
||||
Err(Error::UnsupportedOperation)
|
||||
);
|
||||
server.stop();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn response_payloads_decode_and_invalid_entries_fail() {
|
||||
let server = FakeQdrant::start(FakeState::default()).await;
|
||||
let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
|
||||
for (key, response) in [
|
||||
("python", json!("{'timestamp': 1.0, 'response': {'a': 1}}")),
|
||||
("garbage", json!("not json")),
|
||||
("missing", json!("unused")),
|
||||
] {
|
||||
let mut payload = serde_json::Map::new();
|
||||
payload.insert("litellm_cache_key".to_owned(), json!(key));
|
||||
if key != "missing" {
|
||||
payload.insert("response".to_owned(), response);
|
||||
}
|
||||
server.insert_point(StoredPoint {
|
||||
id: Some(PointId::from(key.len() as u64)),
|
||||
vector: vec![1.0, 0.0],
|
||||
payload: Payload::try_from(JsonValue::Object(payload))
|
||||
.unwrap()
|
||||
.into(),
|
||||
});
|
||||
}
|
||||
assert_eq!(
|
||||
cache
|
||||
.async_get_cache("python", &context("hello"))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(value(json!({"a": 1})))
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_get_cache("garbage", &context("hello")).await,
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
assert_eq!(
|
||||
cache.async_get_cache("missing", &context("hello")).await,
|
||||
Err(Error::InvalidEntry)
|
||||
);
|
||||
server.stop();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn response_cache_facade_turns_invalid_entry_into_miss() {
|
||||
let server = FakeQdrant::start(FakeState::default()).await;
|
||||
let cache = Arc::new(connect(&server, [("hello", vec![1.0, 0.0])]).await);
|
||||
let request = ResponseCacheRequest::<SemanticCacheContext>::new(CacheKeyInput {
|
||||
preset: Some("key".to_owned()),
|
||||
..Default::default()
|
||||
})
|
||||
.with_context(context("hello"));
|
||||
let response = json!({"answer": 42});
|
||||
let facade = ResponseCache::new(cache.clone());
|
||||
facade
|
||||
.async_store(&request, response.clone(), Duration::from_secs(1))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
facade
|
||||
.async_lookup(&request, Duration::from_secs(1))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(response)
|
||||
);
|
||||
{
|
||||
let mut state = server.state.lock().unwrap();
|
||||
state.points[0]
|
||||
.payload
|
||||
.insert("response".to_owned(), Value::from("not json"));
|
||||
}
|
||||
assert_eq!(
|
||||
facade
|
||||
.async_lookup(&request, Duration::from_secs(1))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
server.stop();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn stopped_qdrant_server_maps_to_unavailable() {
|
||||
let server = FakeQdrant::start(FakeState::default()).await;
|
||||
let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
|
||||
server.stop();
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
assert_eq!(
|
||||
cache.async_get_cache("key", &context("hello")).await,
|
||||
Err(Error::Unavailable)
|
||||
);
|
||||
}
|
||||
342
litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs
Normal file
342
litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs
Normal file
|
|
@ -0,0 +1,342 @@
|
|||
use std::{
|
||||
collections::{HashMap, HashSet},
|
||||
net::SocketAddr,
|
||||
sync::{Arc, Mutex},
|
||||
};
|
||||
|
||||
use qdrant_client::qdrant::collections_server::CollectionsServer;
|
||||
use qdrant_client::qdrant::{
|
||||
self, CollectionExists, CollectionExistsRequest, CollectionExistsResponse,
|
||||
CollectionOperationResponse, CreateCollection, CreateFieldIndexCollection, Filter, PointId,
|
||||
PointsOperationResponse, ScoredPoint, SearchPoints, SearchResponse, Value, Vector, Vectors,
|
||||
collections_server::Collections,
|
||||
points_server::{Points, PointsServer},
|
||||
};
|
||||
use tokio::sync::oneshot;
|
||||
use tokio_stream::wrappers::TcpListenerStream;
|
||||
use tonic::{Request, Response, Status, transport::Server};
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct StoredPoint {
|
||||
pub id: Option<PointId>,
|
||||
pub vector: Vec<f32>,
|
||||
pub payload: HashMap<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct FakeState {
|
||||
pub collections: HashSet<String>,
|
||||
pub created_collections: Vec<CreateCollection>,
|
||||
pub field_indexes: Vec<CreateFieldIndexCollection>,
|
||||
pub points: Vec<StoredPoint>,
|
||||
pub upsert_waits: Vec<Option<bool>>,
|
||||
pub index_creations: usize,
|
||||
pub fail_field_index: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct FakeQdrant {
|
||||
pub state: Arc<Mutex<FakeState>>,
|
||||
pub address: SocketAddr,
|
||||
shutdown: Arc<Mutex<Option<oneshot::Sender<()>>>>,
|
||||
}
|
||||
|
||||
impl FakeQdrant {
|
||||
pub async fn start(state: FakeState) -> Self {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let state = Arc::new(Mutex::new(state));
|
||||
let service = FakeService {
|
||||
state: state.clone(),
|
||||
};
|
||||
let (shutdown_tx, shutdown_rx) = oneshot::channel();
|
||||
tokio::spawn(async move {
|
||||
Server::builder()
|
||||
.add_service(CollectionsServer::new(service.clone()))
|
||||
.add_service(PointsServer::new(service))
|
||||
.serve_with_incoming_shutdown(TcpListenerStream::new(listener), async {
|
||||
let _ = shutdown_rx.await;
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
});
|
||||
Self {
|
||||
state,
|
||||
address,
|
||||
shutdown: Arc::new(Mutex::new(Some(shutdown_tx))),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn url(&self) -> String {
|
||||
format!("http://{}", self.address)
|
||||
}
|
||||
|
||||
pub fn stop(&self) {
|
||||
self.shutdown
|
||||
.lock()
|
||||
.unwrap()
|
||||
.take()
|
||||
.unwrap()
|
||||
.send(())
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
pub fn insert_point(&self, point: StoredPoint) {
|
||||
self.state.lock().unwrap().points.push(point);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct FakeService {
|
||||
state: Arc<Mutex<FakeState>>,
|
||||
}
|
||||
|
||||
macro_rules! unimplemented_collections {
|
||||
($($name:ident, $request:ty, $response:ty);* $(;)?) => {
|
||||
$(
|
||||
fn $name<'life0, 'async_trait>(
|
||||
&'life0 self,
|
||||
_: Request<$request>,
|
||||
) -> std::pin::Pin<
|
||||
Box<
|
||||
dyn std::future::Future<
|
||||
Output = Result<Response<$response>, Status>,
|
||||
> + Send
|
||||
+ 'async_trait,
|
||||
>,
|
||||
>
|
||||
where
|
||||
'life0: 'async_trait,
|
||||
Self: 'async_trait,
|
||||
{
|
||||
Box::pin(async { Err(Status::unimplemented(stringify!($name))) })
|
||||
}
|
||||
)*
|
||||
};
|
||||
}
|
||||
|
||||
macro_rules! unimplemented_points {
|
||||
($($name:ident, $request:ty, $response:ty);* $(;)?) => {
|
||||
$(
|
||||
fn $name<'life0, 'async_trait>(
|
||||
&'life0 self,
|
||||
_: Request<$request>,
|
||||
) -> std::pin::Pin<
|
||||
Box<
|
||||
dyn std::future::Future<
|
||||
Output = Result<Response<$response>, Status>,
|
||||
> + Send
|
||||
+ 'async_trait,
|
||||
>,
|
||||
>
|
||||
where
|
||||
'life0: 'async_trait,
|
||||
Self: 'async_trait,
|
||||
{
|
||||
Box::pin(async { Err(Status::unimplemented(stringify!($name))) })
|
||||
}
|
||||
)*
|
||||
};
|
||||
}
|
||||
|
||||
#[tonic::async_trait]
|
||||
impl Collections for FakeService {
|
||||
async fn create(
|
||||
&self,
|
||||
request: Request<CreateCollection>,
|
||||
) -> Result<Response<CollectionOperationResponse>, Status> {
|
||||
let request = request.into_inner();
|
||||
let mut state = self.state.lock().unwrap();
|
||||
state.collections.insert(request.collection_name.clone());
|
||||
state.created_collections.push(request);
|
||||
Ok(Response::new(CollectionOperationResponse {
|
||||
result: true,
|
||||
..Default::default()
|
||||
}))
|
||||
}
|
||||
|
||||
async fn collection_exists(
|
||||
&self,
|
||||
request: Request<CollectionExistsRequest>,
|
||||
) -> Result<Response<CollectionExistsResponse>, Status> {
|
||||
let exists = self
|
||||
.state
|
||||
.lock()
|
||||
.unwrap()
|
||||
.collections
|
||||
.contains(&request.into_inner().collection_name);
|
||||
Ok(Response::new(CollectionExistsResponse {
|
||||
result: Some(CollectionExists { exists }),
|
||||
..Default::default()
|
||||
}))
|
||||
}
|
||||
|
||||
unimplemented_collections!(
|
||||
get, qdrant::GetCollectionInfoRequest, qdrant::GetCollectionInfoResponse;
|
||||
list, qdrant::ListCollectionsRequest, qdrant::ListCollectionsResponse;
|
||||
update, qdrant::UpdateCollection, qdrant::CollectionOperationResponse;
|
||||
delete, qdrant::DeleteCollection, qdrant::CollectionOperationResponse;
|
||||
update_aliases, qdrant::ChangeAliases, qdrant::CollectionOperationResponse;
|
||||
list_collection_aliases, qdrant::ListCollectionAliasesRequest, qdrant::ListAliasesResponse;
|
||||
list_aliases, qdrant::ListAliasesRequest, qdrant::ListAliasesResponse;
|
||||
collection_cluster_info, qdrant::CollectionClusterInfoRequest, qdrant::CollectionClusterInfoResponse;
|
||||
update_collection_cluster_setup, qdrant::UpdateCollectionClusterSetupRequest, qdrant::UpdateCollectionClusterSetupResponse;
|
||||
create_shard_key, qdrant::CreateShardKeyRequest, qdrant::CreateShardKeyResponse;
|
||||
delete_shard_key, qdrant::DeleteShardKeyRequest, qdrant::DeleteShardKeyResponse;
|
||||
list_shard_keys, qdrant::ListShardKeysRequest, qdrant::ListShardKeysResponse;
|
||||
);
|
||||
}
|
||||
|
||||
#[tonic::async_trait]
|
||||
impl Points for FakeService {
|
||||
async fn create_field_index(
|
||||
&self,
|
||||
request: Request<CreateFieldIndexCollection>,
|
||||
) -> Result<Response<PointsOperationResponse>, Status> {
|
||||
let mut state = self.state.lock().unwrap();
|
||||
state.index_creations += 1;
|
||||
state.field_indexes.push(request.into_inner());
|
||||
if state.fail_field_index {
|
||||
return Err(Status::internal("field index failure"));
|
||||
}
|
||||
Ok(Response::new(PointsOperationResponse::default()))
|
||||
}
|
||||
|
||||
async fn upsert(
|
||||
&self,
|
||||
request: Request<qdrant::UpsertPoints>,
|
||||
) -> Result<Response<PointsOperationResponse>, Status> {
|
||||
let request = request.into_inner();
|
||||
let mut state = self.state.lock().unwrap();
|
||||
state.upsert_waits.push(request.wait);
|
||||
for point in request.points {
|
||||
let stored = StoredPoint {
|
||||
id: point.id.clone(),
|
||||
vector: dense_vector(point.vectors)?,
|
||||
payload: point.payload,
|
||||
};
|
||||
if let Some(existing) = state
|
||||
.points
|
||||
.iter_mut()
|
||||
.find(|existing| existing.id == stored.id)
|
||||
{
|
||||
*existing = stored;
|
||||
} else {
|
||||
state.points.push(stored);
|
||||
}
|
||||
}
|
||||
Ok(Response::new(PointsOperationResponse::default()))
|
||||
}
|
||||
|
||||
async fn search(
|
||||
&self,
|
||||
request: Request<SearchPoints>,
|
||||
) -> Result<Response<SearchResponse>, Status> {
|
||||
let request = request.into_inner();
|
||||
let key_filter = keyword_filter(request.filter.as_ref());
|
||||
let state = self.state.lock().unwrap();
|
||||
let mut results = state
|
||||
.points
|
||||
.iter()
|
||||
.filter(|point| {
|
||||
key_filter.as_ref().is_none_or(|(field, expected)| {
|
||||
point
|
||||
.payload
|
||||
.get(field)
|
||||
.and_then(|value| {
|
||||
let value: serde_json::Value = value.clone().into();
|
||||
value
|
||||
.as_str()
|
||||
.map(str::to_owned)
|
||||
.or_else(|| value.as_i64().map(|value| value.to_string()))
|
||||
})
|
||||
.is_some_and(|value| value == *expected)
|
||||
})
|
||||
})
|
||||
.map(|point| ScoredPoint {
|
||||
id: point.id.clone(),
|
||||
payload: point.payload.clone(),
|
||||
score: cosine(&request.vector, &point.vector),
|
||||
..Default::default()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
results.sort_by(|left, right| right.score.total_cmp(&left.score));
|
||||
results.truncate(request.limit as usize);
|
||||
Ok(Response::new(SearchResponse {
|
||||
result: results,
|
||||
..Default::default()
|
||||
}))
|
||||
}
|
||||
|
||||
unimplemented_points!(
|
||||
delete, qdrant::DeletePoints, qdrant::PointsOperationResponse;
|
||||
get, qdrant::GetPoints, qdrant::GetResponse;
|
||||
update_vectors, qdrant::UpdatePointVectors, qdrant::PointsOperationResponse;
|
||||
delete_vectors, qdrant::DeletePointVectors, qdrant::PointsOperationResponse;
|
||||
set_payload, qdrant::SetPayloadPoints, qdrant::PointsOperationResponse;
|
||||
overwrite_payload, qdrant::SetPayloadPoints, qdrant::PointsOperationResponse;
|
||||
delete_payload, qdrant::DeletePayloadPoints, qdrant::PointsOperationResponse;
|
||||
clear_payload, qdrant::ClearPayloadPoints, qdrant::PointsOperationResponse;
|
||||
delete_field_index, qdrant::DeleteFieldIndexCollection, qdrant::PointsOperationResponse;
|
||||
create_vector_name, qdrant::CreateVectorNameRequest, qdrant::PointsOperationResponse;
|
||||
delete_vector_name, qdrant::DeleteVectorNameRequest, qdrant::PointsOperationResponse;
|
||||
search_batch, qdrant::SearchBatchPoints, qdrant::SearchBatchResponse;
|
||||
search_groups, qdrant::SearchPointGroups, qdrant::SearchGroupsResponse;
|
||||
scroll, qdrant::ScrollPoints, qdrant::ScrollResponse;
|
||||
recommend, qdrant::RecommendPoints, qdrant::RecommendResponse;
|
||||
recommend_batch, qdrant::RecommendBatchPoints, qdrant::RecommendBatchResponse;
|
||||
recommend_groups, qdrant::RecommendPointGroups, qdrant::RecommendGroupsResponse;
|
||||
discover, qdrant::DiscoverPoints, qdrant::DiscoverResponse;
|
||||
discover_batch, qdrant::DiscoverBatchPoints, qdrant::DiscoverBatchResponse;
|
||||
count, qdrant::CountPoints, qdrant::CountResponse;
|
||||
update_batch, qdrant::UpdateBatchPoints, qdrant::UpdateBatchResponse;
|
||||
query, qdrant::QueryPoints, qdrant::QueryResponse;
|
||||
query_batch, qdrant::QueryBatchPoints, qdrant::QueryBatchResponse;
|
||||
query_groups, qdrant::QueryPointGroups, qdrant::QueryGroupsResponse;
|
||||
facet, qdrant::FacetCounts, qdrant::FacetResponse;
|
||||
search_matrix_pairs, qdrant::SearchMatrixPoints, qdrant::SearchMatrixPairsResponse;
|
||||
search_matrix_offsets, qdrant::SearchMatrixPoints, qdrant::SearchMatrixOffsetsResponse;
|
||||
);
|
||||
}
|
||||
|
||||
fn dense_vector(vectors: Option<Vectors>) -> Result<Vec<f32>, Status> {
|
||||
let Some(Vectors {
|
||||
vectors_options:
|
||||
Some(qdrant::vectors::VectorsOptions::Vector(Vector {
|
||||
vector: Some(qdrant::vector::Vector::Dense(qdrant::DenseVector { data })),
|
||||
..
|
||||
})),
|
||||
}) = vectors
|
||||
else {
|
||||
return Err(Status::invalid_argument("expected dense vector"));
|
||||
};
|
||||
Ok(data)
|
||||
}
|
||||
|
||||
fn keyword_filter(filter: Option<&Filter>) -> Option<(String, String)> {
|
||||
filter?
|
||||
.must
|
||||
.iter()
|
||||
.find_map(|condition| match condition.condition_one_of.as_ref()? {
|
||||
qdrant::condition::ConditionOneOf::Field(field) => {
|
||||
let qdrant::r#match::MatchValue::Keyword(value) =
|
||||
field.r#match.as_ref()?.match_value.as_ref()?
|
||||
else {
|
||||
return None;
|
||||
};
|
||||
Some((field.key.clone(), value.clone()))
|
||||
}
|
||||
_ => None,
|
||||
})
|
||||
}
|
||||
|
||||
fn cosine(left: &[f32], right: &[f32]) -> f32 {
|
||||
let dot = left
|
||||
.iter()
|
||||
.zip(right)
|
||||
.map(|(left, right)| left * right)
|
||||
.sum::<f32>();
|
||||
let left_norm = left.iter().map(|value| value * value).sum::<f32>().sqrt();
|
||||
let right_norm = right.iter().map(|value| value * value).sum::<f32>().sqrt();
|
||||
dot / (left_norm * right_norm)
|
||||
}
|
||||
|
|
@ -32,6 +32,17 @@ impl<C: CacheContext + Default> ResponseCacheRequest<C> {
|
|||
}
|
||||
}
|
||||
|
||||
impl<C: CacheContext> ResponseCacheRequest<C> {
|
||||
pub fn with_context<D: CacheContext>(self, context: D) -> ResponseCacheRequest<D> {
|
||||
ResponseCacheRequest {
|
||||
key: self.key,
|
||||
controls: self.controls,
|
||||
context,
|
||||
max_age: self.max_age,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ResponseCache<B: BaseCache<Value = CacheEntry>>
|
||||
where
|
||||
B::Context: Default + PartialEq,
|
||||
|
|
|
|||
|
|
@ -1,12 +1,15 @@
|
|||
use std::{
|
||||
sync::{
|
||||
Arc,
|
||||
Arc, Mutex,
|
||||
atomic::{AtomicU64, Ordering},
|
||||
},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_cache::{BaseCache, CacheCodec, Error};
|
||||
use litellm_cache::{
|
||||
BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error,
|
||||
SemanticCacheContext,
|
||||
};
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_redis::RedisCache;
|
||||
use litellm_cache_response::{
|
||||
|
|
@ -30,6 +33,82 @@ fn request() -> ResponseCacheRequest {
|
|||
})
|
||||
}
|
||||
|
||||
struct SemanticBackend {
|
||||
entries: Mutex<Vec<(String, CacheEntry)>>,
|
||||
contexts: Mutex<Vec<SemanticCacheContext>>,
|
||||
}
|
||||
|
||||
impl BaseCache for SemanticBackend {
|
||||
type Value = CacheEntry;
|
||||
type Context = SemanticCacheContext;
|
||||
|
||||
fn get_ttl(&self, _: &Self::Context) -> Option<Duration> {
|
||||
None
|
||||
}
|
||||
|
||||
fn set_cache(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Self::Value,
|
||||
context: &Self::Context,
|
||||
) -> Result<(), Error> {
|
||||
self.contexts.lock().unwrap().push(context.clone());
|
||||
self.entries.lock().unwrap().push((key.to_owned(), value));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
|
||||
self.contexts.lock().unwrap().push(context.clone());
|
||||
Ok(self
|
||||
.entries
|
||||
.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.find(|(entry_key, _)| entry_key == key)
|
||||
.map(|(_, entry)| entry.clone()))
|
||||
}
|
||||
|
||||
async fn disconnect(&self) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
|
||||
Ok(CacheConnectionResult {
|
||||
status: CacheConnectionStatus::Success,
|
||||
message: "ok".into(),
|
||||
error: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn semantic_context_reaches_backend_for_store_and_lookup() {
|
||||
let backend = Arc::new(SemanticBackend {
|
||||
entries: Mutex::new(Vec::new()),
|
||||
contexts: Mutex::new(Vec::new()),
|
||||
});
|
||||
let cache = ResponseCache::new(backend.clone());
|
||||
let context = SemanticCacheContext {
|
||||
messages: Some(json!([{"role": "user", "content": "hello"}])),
|
||||
..Default::default()
|
||||
};
|
||||
let request = request().with_context(context.clone());
|
||||
let response = json!({"answer": 42});
|
||||
|
||||
cache
|
||||
.store(&request, response.clone(), Duration::from_secs(100))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
cache.lookup(&request, Duration::from_secs(100)).unwrap(),
|
||||
Some(response)
|
||||
);
|
||||
assert_eq!(
|
||||
backend.contexts.lock().unwrap().as_slice(),
|
||||
&[context.clone(), context]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sync_and_async_consumers_share_keys_ttls_and_freshness() {
|
||||
let clock = Arc::new(AtomicU64::new(100));
|
||||
|
|
|
|||
2
litellm-rust/crates/cache/src/error.rs
vendored
2
litellm-rust/crates/cache/src/error.rs
vendored
|
|
@ -8,4 +8,6 @@ pub enum Error {
|
|||
UnscopedFlush,
|
||||
#[error("operation is not supported by this cache")]
|
||||
UnsupportedOperation,
|
||||
#[error("semantic cache requires request messages")]
|
||||
MissingPrompt,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -29,6 +29,8 @@ litellm-cache-gcs.workspace = true
|
|||
litellm-cache-disk.workspace = true
|
||||
litellm-cache-redis-semantic.workspace = true
|
||||
litellm-cache-response.workspace = true
|
||||
litellm-cache-qdrant-semantic.workspace = true
|
||||
qdrant-client.workspace = true
|
||||
litellm-cache-valkey-semantic = { path = "../cache-valkey-semantic" }
|
||||
serde.workspace = true
|
||||
litellm-auth.workspace = true
|
||||
|
|
@ -44,8 +46,10 @@ litellm-host-python.workspace = true
|
|||
litellm-token-counter = { path = "../token-counter", default-features = false }
|
||||
pyo3.workspace = true
|
||||
pyo3-async-runtimes.workspace = true
|
||||
reqwest.workspace = true
|
||||
redis = { version = "1.7.0", features = ["tls-rustls"] }
|
||||
serde_json.workspace = true
|
||||
url.workspace = true
|
||||
tokio = { workspace = true, features = ["rt", "sync"] }
|
||||
|
||||
[dev-dependencies]
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ use std::{path::PathBuf, time::Duration};
|
|||
|
||||
use litellm_auth_aws::AwsAuthConfig;
|
||||
use litellm_cache::CacheType;
|
||||
use litellm_cache_qdrant_semantic::{OpenAiEmbedderConfig, QdrantSemanticConfig, Quantization};
|
||||
use litellm_cache_redis::{RedisNode, RedisTopology};
|
||||
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
|
||||
use pyo3::{
|
||||
|
|
@ -125,6 +126,27 @@ pub(super) struct ValkeySemanticCacheConfig {
|
|||
pub(super) connection: RedisConnectionConfig,
|
||||
}
|
||||
|
||||
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>),
|
||||
|
|
@ -134,6 +156,7 @@ pub(super) enum CacheBackendConfig {
|
|||
Disk(DiskCacheConfig),
|
||||
AzureBlob(AzureBlobCacheConfig),
|
||||
RedisSemantic(Box<RedisSemanticCacheConfig>),
|
||||
QdrantSemantic(Box<QdrantSemanticCacheConfig>),
|
||||
}
|
||||
|
||||
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
|
||||
|
|
@ -153,6 +176,8 @@ pub(super) enum UnsupportedCacheConfig {
|
|||
S3Option,
|
||||
GcsBucket,
|
||||
DiskStore,
|
||||
QdrantEndpoint,
|
||||
SemanticEmbedding,
|
||||
}
|
||||
|
||||
impl UnsupportedCacheConfig {
|
||||
|
|
@ -168,6 +193,10 @@ impl UnsupportedCacheConfig {
|
|||
Self::S3Option => "native S3 configuration requires Python",
|
||||
Self::GcsBucket => "native GCS cache requires a configured bucket name",
|
||||
Self::DiskStore => "native disk cache requires the built-in diskcache store",
|
||||
Self::QdrantEndpoint => {
|
||||
"native Qdrant requires the default REST port so the gRPC port can be derived"
|
||||
}
|
||||
Self::SemanticEmbedding => "native semantic embedding requires Python",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -238,6 +267,13 @@ 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::AzureBlob) => project_azure_blob(&backend).map(|backend| {
|
||||
CacheConfigProjection::Native(Box::new(Self {
|
||||
policy,
|
||||
|
|
@ -250,7 +286,7 @@ impl NativeCacheConfig {
|
|||
backend: CacheBackendConfig::RedisSemantic(Box::new(backend)),
|
||||
}))
|
||||
}),
|
||||
Some(CacheType::QdrantSemantic) | None => Ok(CacheConfigProjection::Unsupported(
|
||||
None => Ok(CacheConfigProjection::Unsupported(
|
||||
UnsupportedCacheConfig::Backend,
|
||||
)),
|
||||
}
|
||||
|
|
@ -265,7 +301,8 @@ impl NativeCacheConfig {
|
|||
CacheBackendConfig::Disk(_)
|
||||
| CacheBackendConfig::AzureBlob(_)
|
||||
| CacheBackendConfig::Gcs(_)
|
||||
| CacheBackendConfig::RedisSemantic(_) => None,
|
||||
| CacheBackendConfig::RedisSemantic(_)
|
||||
| CacheBackendConfig::QdrantSemantic(_) => None,
|
||||
};
|
||||
if !matches!(self.backend, CacheBackendConfig::ValkeySemantic(_))
|
||||
&& service.default_ttl() != default_ttl
|
||||
|
|
@ -373,11 +410,35 @@ impl NativeCacheConfig {
|
|||
Some("facade and native backend index names must match")
|
||||
}
|
||||
CacheBackendConfig::RedisSemantic(config)
|
||||
if service.similarity_threshold() != Some(config.similarity_threshold as f32) =>
|
||||
if service.similarity_threshold() != Some(config.similarity_threshold) =>
|
||||
{
|
||||
Some("facade and native backend similarity thresholds must match")
|
||||
}
|
||||
CacheBackendConfig::RedisSemantic(_) => None,
|
||||
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))
|
||||
|
|
@ -391,6 +452,105 @@ impl NativeCacheConfig {
|
|||
}
|
||||
}
|
||||
|
||||
#[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);
|
||||
|
||||
if optional_attribute(backend, "embedding_max_input_tokens")?
|
||||
.is_some_and(|value| !value.is_none())
|
||||
{
|
||||
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) = std::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 = std::env::var("OPENAI_BASE_URL")
|
||||
.or_else(|_| std::env::var("OPENAI_API_BASE"))
|
||||
.unwrap_or_else(|_| "https://api.openai.com/v1".to_owned());
|
||||
let timeout = optional_attribute(backend, "embedding_timeout")?
|
||||
.map(|value| value.extract::<Option<f64>>())
|
||||
.transpose()?
|
||||
.flatten()
|
||||
.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")?;
|
||||
|
|
|
|||
|
|
@ -368,6 +368,11 @@ impl FacadeGuard {
|
|||
"RedisSemanticCache",
|
||||
"redis-semantic",
|
||||
),
|
||||
("qdrant_semantic", _) => (
|
||||
"litellm.caching.qdrant_semantic_cache",
|
||||
"QdrantSemanticCache",
|
||||
"qdrant-semantic",
|
||||
),
|
||||
("redis", true) => (
|
||||
"litellm.caching.redis_cluster_cache",
|
||||
"RedisClusterCache",
|
||||
|
|
@ -443,6 +448,10 @@ impl FacadeGuard {
|
|||
"embedding_model",
|
||||
"embedding_max_input_tokens",
|
||||
"embedding_timeout",
|
||||
"qdrant_api_base",
|
||||
"qdrant_api_key",
|
||||
"collection_name",
|
||||
"vector_size",
|
||||
"_index_name",
|
||||
"_redis_url",
|
||||
"similarity_threshold",
|
||||
|
|
|
|||
|
|
@ -1,18 +1,26 @@
|
|||
use litellm_auth_aws::AwsAuthConfig;
|
||||
use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig};
|
||||
use litellm_cache_qdrant_semantic::{OpenAiEmbedderConfig, Quantization};
|
||||
use litellm_cache_redis::{RedisNode, RedisTopology};
|
||||
use litellm_cache_redis_semantic::RedisSemanticConfig;
|
||||
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
|
||||
use litellm_host_python::{release_gil, run_sync_value};
|
||||
use litellm_http::ClientVariant;
|
||||
use pyo3::{
|
||||
PyTraverseError, PyVisit,
|
||||
exceptions::{PyRuntimeError, PyTypeError},
|
||||
prelude::*,
|
||||
types::PyDict,
|
||||
};
|
||||
use url::Url;
|
||||
|
||||
use super::{
|
||||
cache_error, config::project_redis_semantic, embedder::PythonEmbedder, facade::FacadeGuard,
|
||||
native::NativeResponseCache, request::duration,
|
||||
cache_error,
|
||||
config::{QdrantSemanticCacheConfig, project_redis_semantic},
|
||||
embedder::PythonEmbedder,
|
||||
facade::FacadeGuard,
|
||||
native::NativeResponseCache,
|
||||
request::duration,
|
||||
};
|
||||
|
||||
#[pyclass(frozen, name = "_CacheTestHandle")]
|
||||
|
|
@ -146,6 +154,105 @@ 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() != Some(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(|| {
|
||||
std::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(|| {
|
||||
std::env::var("OPENAI_BASE_URL")
|
||||
.or_else(|_| std::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 http_config = crate::http::call_config(py, &PyDict::new(py), true)?;
|
||||
let client = crate::http::pool()
|
||||
.client(&http_config, ClientVariant::Provider)
|
||||
.map_err(crate::http::client_error)?;
|
||||
let service = run_sync_value(py, async move {
|
||||
let handle = tokio::runtime::Handle::current();
|
||||
NativeResponseCache::qdrant_semantic(config, client, handle)
|
||||
.await
|
||||
.map_err(cache_error)
|
||||
})?;
|
||||
Ok(Self {
|
||||
service,
|
||||
guard: None,
|
||||
pid: std::process::id(),
|
||||
})
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
#[pyo3(signature = (url, similarity_threshold, index_name, embedder))]
|
||||
fn valkey_semantic(
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ use litellm_cache_azure_blob::AzureBlobCache;
|
|||
use litellm_cache_disk::DiskCache;
|
||||
use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource};
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_cache_qdrant_semantic::{Embedder, OpenAiEmbedder, QdrantSemanticCache};
|
||||
use litellm_cache_redis::{RedisCache, RedisTopology};
|
||||
use litellm_cache_redis_semantic::{RedisSemanticCache, RedisSemanticConfig};
|
||||
use litellm_cache_response::{
|
||||
|
|
@ -19,6 +20,7 @@ use pyo3::{PyTraverseError, PyVisit, prelude::*};
|
|||
use serde_json::Value;
|
||||
|
||||
use super::{
|
||||
config::QdrantSemanticCacheConfig,
|
||||
embedder::PythonEmbedder,
|
||||
request::NativeRequest,
|
||||
semantic::{SemanticBody, SemanticOperation, drive},
|
||||
|
|
@ -92,6 +94,7 @@ pub(super) enum NativeResponseCache {
|
|||
cache: Arc<ResponseCache<RedisSemanticCache<PythonEmbedder>>>,
|
||||
embedder: PythonEmbedder,
|
||||
},
|
||||
QdrantSemantic(Arc<ResponseCache<QdrantSemanticCache<OpenAiEmbedder, ResponseCacheCodec>>>),
|
||||
Disk(Arc<ResponseCache<DiskCache<ResponseCacheCodec>>>),
|
||||
AzureBlob(Arc<ResponseCache<AzureBlobCache<ResponseCacheCodec>>>),
|
||||
}
|
||||
|
|
@ -168,6 +171,31 @@ impl NativeResponseCache {
|
|||
})
|
||||
}
|
||||
|
||||
pub async fn qdrant_semantic(
|
||||
config: QdrantSemanticCacheConfig,
|
||||
client: reqwest::Client,
|
||||
runtime: tokio::runtime::Handle,
|
||||
) -> Result<Self, Error> {
|
||||
let qdrant = 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(client, config.embedding);
|
||||
let cache = QdrantSemanticCache::connect(
|
||||
qdrant,
|
||||
embedder,
|
||||
ResponseCacheCodec,
|
||||
qdrant_config,
|
||||
runtime,
|
||||
)
|
||||
.await?;
|
||||
Ok(Self::QdrantSemantic(Arc::new(ResponseCache::new(
|
||||
Arc::new(cache),
|
||||
))))
|
||||
}
|
||||
|
||||
pub fn disk(directory: &str) -> Result<Self, Error> {
|
||||
let cache = DiskCache::open(directory, ResponseCacheCodec)?;
|
||||
Ok(Self::Disk(Arc::new(ResponseCache::new(Arc::new(cache)))))
|
||||
|
|
@ -209,6 +237,7 @@ impl NativeResponseCache {
|
|||
| Self::S3(_)
|
||||
| Self::ValkeySemantic { .. }
|
||||
| Self::RedisSemantic { .. }
|
||||
| Self::QdrantSemantic(_)
|
||||
| Self::Disk(_)
|
||||
| Self::Gcs(_) => None,
|
||||
}
|
||||
|
|
@ -223,7 +252,7 @@ impl NativeResponseCache {
|
|||
}
|
||||
}
|
||||
|
||||
pub(super) fn redis_semantic_request(
|
||||
pub(super) fn semantic_request(
|
||||
request: &NativeRequest,
|
||||
) -> ResponseCacheRequest<SemanticCacheContext> {
|
||||
ResponseCacheRequest {
|
||||
|
|
@ -289,6 +318,7 @@ impl NativeResponseCache {
|
|||
Self::Gcs(_) => "gcs",
|
||||
Self::ValkeySemantic { .. } => "valkey-semantic",
|
||||
Self::RedisSemantic { .. } => "redis_semantic",
|
||||
Self::QdrantSemantic(_) => "qdrant_semantic",
|
||||
Self::Disk(_) => "disk",
|
||||
Self::AzureBlob(_) => "azure-blob",
|
||||
}
|
||||
|
|
@ -302,6 +332,7 @@ impl NativeResponseCache {
|
|||
Self::Gcs(cache) => cache.default_ttl(),
|
||||
Self::ValkeySemantic { cache, .. } => cache.default_ttl(),
|
||||
Self::RedisSemantic { cache, .. } => cache.default_ttl(),
|
||||
Self::QdrantSemantic(_) => None,
|
||||
Self::Disk(cache) => cache.default_ttl(),
|
||||
Self::AzureBlob(cache) => cache.default_ttl(),
|
||||
}
|
||||
|
|
@ -341,6 +372,7 @@ impl NativeResponseCache {
|
|||
| Self::S3(_)
|
||||
| Self::ValkeySemantic { .. }
|
||||
| Self::RedisSemantic { .. }
|
||||
| Self::QdrantSemantic(_)
|
||||
| Self::Disk(_)
|
||||
| Self::AzureBlob(_)
|
||||
| Self::Gcs(_) => None,
|
||||
|
|
@ -354,6 +386,7 @@ impl NativeResponseCache {
|
|||
| Self::S3(_)
|
||||
| Self::ValkeySemantic { .. }
|
||||
| Self::RedisSemantic { .. }
|
||||
| Self::QdrantSemantic(_)
|
||||
| Self::Disk(_)
|
||||
| Self::AzureBlob(_)
|
||||
| Self::Gcs(_) => None,
|
||||
|
|
@ -368,6 +401,7 @@ impl NativeResponseCache {
|
|||
| Self::S3(_)
|
||||
| Self::ValkeySemantic { .. }
|
||||
| Self::RedisSemantic { .. }
|
||||
| Self::QdrantSemantic(_)
|
||||
| Self::Disk(_)
|
||||
| Self::AzureBlob(_)
|
||||
| Self::Gcs(_) => None,
|
||||
|
|
@ -381,6 +415,7 @@ impl NativeResponseCache {
|
|||
| Self::S3(_)
|
||||
| Self::ValkeySemantic { .. }
|
||||
| Self::RedisSemantic { .. }
|
||||
| Self::QdrantSemantic(_)
|
||||
| Self::Disk(_)
|
||||
| Self::AzureBlob(_)
|
||||
| Self::Gcs(_) => None,
|
||||
|
|
@ -395,6 +430,7 @@ impl NativeResponseCache {
|
|||
| Self::S3(_)
|
||||
| Self::ValkeySemantic { .. }
|
||||
| Self::RedisSemantic { .. }
|
||||
| Self::QdrantSemantic(_)
|
||||
| Self::AzureBlob(_)
|
||||
| Self::Gcs(_) => None,
|
||||
}
|
||||
|
|
@ -421,9 +457,33 @@ impl NativeResponseCache {
|
|||
}
|
||||
}
|
||||
|
||||
pub fn similarity_threshold(&self) -> Option<f32> {
|
||||
pub fn similarity_threshold(&self) -> Option<f64> {
|
||||
match self {
|
||||
Self::RedisSemantic { cache, .. } => Some(cache.backend().similarity_threshold()),
|
||||
Self::RedisSemantic { cache, .. } => {
|
||||
Some(f64::from(cache.backend().similarity_threshold()))
|
||||
}
|
||||
Self::QdrantSemantic(cache) => Some(cache.backend().similarity_threshold()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn collection_name(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::QdrantSemantic(cache) => Some(cache.backend().collection_name()),
|
||||
_ => 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,
|
||||
}
|
||||
}
|
||||
|
|
@ -451,9 +511,10 @@ impl NativeResponseCache {
|
|||
cache.lookup(&Self::semantic(request, scope), now)
|
||||
}
|
||||
Self::RedisSemantic { cache, .. } => {
|
||||
cache.lookup(&Self::redis_semantic_request(request), now)
|
||||
cache.lookup(&Self::semantic_request(request), now)
|
||||
}
|
||||
Self::Gcs(cache) => cache.lookup(&Self::exact(request), now),
|
||||
Self::QdrantSemantic(cache) => cache.lookup(&Self::semantic_request(request), now),
|
||||
Self::Disk(cache) => cache.lookup(&Self::exact(request), now),
|
||||
Self::AzureBlob(cache) => cache.lookup(&Self::exact(request), now),
|
||||
}
|
||||
|
|
@ -473,9 +534,12 @@ impl NativeResponseCache {
|
|||
cache.store(&Self::semantic(request, scope), response, now)
|
||||
}
|
||||
Self::RedisSemantic { cache, .. } => {
|
||||
cache.store(&Self::redis_semantic_request(request), response, now)
|
||||
cache.store(&Self::semantic_request(request), response, now)
|
||||
}
|
||||
Self::Gcs(cache) => cache.store(&Self::exact(request), response, now),
|
||||
Self::QdrantSemantic(cache) => {
|
||||
cache.store(&Self::semantic_request(request), response, now)
|
||||
}
|
||||
Self::Disk(cache) => cache.store(&Self::exact(request), response, now),
|
||||
Self::AzureBlob(cache) => cache.store(&Self::exact(request), response, now),
|
||||
}
|
||||
|
|
@ -498,7 +562,7 @@ impl NativeResponseCache {
|
|||
Self::S3(cache) => {
|
||||
cache.lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
|
||||
}
|
||||
Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => {
|
||||
Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => {
|
||||
Err(Error::UnsupportedOperation)
|
||||
}
|
||||
Self::Gcs(cache) => {
|
||||
|
|
@ -529,7 +593,12 @@ impl NativeResponseCache {
|
|||
}
|
||||
Self::RedisSemantic { cache, .. } => {
|
||||
cache
|
||||
.async_lookup(&Self::redis_semantic_request(request), now)
|
||||
.async_lookup(&Self::semantic_request(request), now)
|
||||
.await
|
||||
}
|
||||
Self::QdrantSemantic(cache) => {
|
||||
cache
|
||||
.async_lookup(&Self::semantic_request(request), now)
|
||||
.await
|
||||
}
|
||||
Self::Gcs(cache) => cache.async_lookup(&Self::exact(request), now).await,
|
||||
|
|
@ -573,6 +642,14 @@ impl NativeResponseCache {
|
|||
py,
|
||||
SemanticBody::new(self.clone(), SemanticOperation::Lookup(request)),
|
||||
),
|
||||
Self::QdrantSemantic(_) => {
|
||||
let service = self.clone();
|
||||
litellm_host_python::run_async(
|
||||
py,
|
||||
async move { service.async_lookup(&request, super::request::now()).await },
|
||||
super::cache_error,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -616,7 +693,12 @@ impl NativeResponseCache {
|
|||
}
|
||||
Self::RedisSemantic { cache, .. } => {
|
||||
cache
|
||||
.async_store(&Self::redis_semantic_request(request), response, now)
|
||||
.async_store(&Self::semantic_request(request), response, now)
|
||||
.await
|
||||
}
|
||||
Self::QdrantSemantic(cache) => {
|
||||
cache
|
||||
.async_store(&Self::semantic_request(request), response, now)
|
||||
.await
|
||||
}
|
||||
Self::Gcs(cache) => {
|
||||
|
|
@ -678,6 +760,18 @@ impl NativeResponseCache {
|
|||
py,
|
||||
SemanticBody::new(self.clone(), SemanticOperation::Store(request, response)),
|
||||
),
|
||||
Self::QdrantSemantic(_) => {
|
||||
let service = self.clone();
|
||||
litellm_host_python::run_async(
|
||||
py,
|
||||
async move {
|
||||
service
|
||||
.async_store(&request, response, super::request::now())
|
||||
.await
|
||||
},
|
||||
super::cache_error,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -700,7 +794,7 @@ impl NativeResponseCache {
|
|||
.async_lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
|
||||
.await
|
||||
}
|
||||
Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => {
|
||||
Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => {
|
||||
Err(Error::UnsupportedOperation)
|
||||
}
|
||||
Self::Gcs(cache) => {
|
||||
|
|
@ -755,7 +849,9 @@ impl NativeResponseCache {
|
|||
.collect();
|
||||
cache.async_store_batch(entries, now).await
|
||||
}
|
||||
Self::RedisSemantic { .. } => Err(Error::UnsupportedOperation),
|
||||
Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => {
|
||||
Err(Error::UnsupportedOperation)
|
||||
}
|
||||
Self::Gcs(cache) => {
|
||||
let entries = entries
|
||||
.into_iter()
|
||||
|
|
@ -826,6 +922,7 @@ impl NativeResponseCache {
|
|||
py,
|
||||
SemanticBody::new(self.clone(), SemanticOperation::StoreBatch(entries.into())),
|
||||
),
|
||||
Self::QdrantSemantic(_) => Err(super::cache_error(Error::UnsupportedOperation)),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -839,7 +936,7 @@ impl NativeResponseCache {
|
|||
cache.async_flush().await
|
||||
}
|
||||
Self::S3(cache) => cache.async_flush().await,
|
||||
Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => {
|
||||
Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => {
|
||||
Err(Error::UnsupportedOperation)
|
||||
}
|
||||
Self::Gcs(cache) => cache.async_flush().await,
|
||||
|
|
@ -854,7 +951,9 @@ impl NativeResponseCache {
|
|||
Self::Redis { cache, .. } => cache.test_connection().await,
|
||||
Self::S3(cache) => cache.test_connection().await,
|
||||
Self::ValkeySemantic { cache, .. } => cache.test_connection().await,
|
||||
Self::RedisSemantic { .. } => Err(Error::UnsupportedOperation),
|
||||
Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => {
|
||||
Err(Error::UnsupportedOperation)
|
||||
}
|
||||
Self::Gcs(cache) => cache.test_connection().await,
|
||||
Self::Disk(cache) => cache.test_connection().await,
|
||||
Self::AzureBlob(cache) => cache.test_connection().await,
|
||||
|
|
|
|||
|
|
@ -101,7 +101,7 @@ impl ExecutionBody for SemanticBody {
|
|||
let (request, _) = self.pending.as_ref().ok_or_else(|| {
|
||||
PyRuntimeError::new_err("semantic execution has no pending operation")
|
||||
})?;
|
||||
let semantic = NativeResponseCache::redis_semantic_request(request);
|
||||
let semantic = NativeResponseCache::semantic_request(request);
|
||||
let Some(prompt) = prompt_from_context(&semantic.context) else {
|
||||
return self.backend_step(py, Err(Error::Unavailable));
|
||||
};
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import ast
|
|||
import asyncio
|
||||
import json
|
||||
import os
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, cast
|
||||
|
||||
import litellm
|
||||
|
|
@ -36,6 +37,8 @@ from ._embedding_router import (
|
|||
)
|
||||
from .base_cache import BaseCache
|
||||
|
||||
_WAIT_FOR_INDEXING: Final = MappingProxyType({"wait": "true"})
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router import Router
|
||||
|
||||
|
|
@ -313,6 +316,7 @@ class QdrantSemanticCache(BaseCache):
|
|||
self.sync_client.put(
|
||||
url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points",
|
||||
headers=self.headers,
|
||||
params=_WAIT_FOR_INDEXING,
|
||||
json=data,
|
||||
)
|
||||
|
||||
|
|
@ -422,6 +426,7 @@ class QdrantSemanticCache(BaseCache):
|
|||
await self.async_client.put(
|
||||
url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points",
|
||||
headers=self.headers,
|
||||
params=_WAIT_FOR_INDEXING,
|
||||
json=data,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -167,6 +167,20 @@ class _CacheTestHandle:
|
|||
@staticmethod
|
||||
def disk(directory: str) -> _CacheTestHandle: ...
|
||||
@staticmethod
|
||||
def qdrant_semantic(
|
||||
url: str,
|
||||
*,
|
||||
collection_name: str,
|
||||
similarity_threshold: float,
|
||||
vector_size: int,
|
||||
embedding_model: str = "text-embedding-3-small",
|
||||
api_key: str | None = None,
|
||||
embedding_api_key: str | None = None,
|
||||
embedding_api_base: str | None = None,
|
||||
embedding_timeout_seconds: float | None = None,
|
||||
quantization: str = "binary",
|
||||
) -> _CacheTestHandle: ...
|
||||
@staticmethod
|
||||
def azure_blob(account_url: str, container: str) -> _CacheTestHandle: ...
|
||||
@staticmethod
|
||||
def redis_semantic(backend: object) -> _CacheTestHandle: ...
|
||||
|
|
|
|||
|
|
@ -578,6 +578,7 @@ def test_qdrant_semantic_cache_set_cache():
|
|||
assert (
|
||||
upsert_payload[QdrantSemanticCache.CACHE_KEY_FIELD_NAME] == "test_key"
|
||||
)
|
||||
assert qdrant_cache.sync_client.put.call_args.kwargs["params"] == {"wait": "true"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -650,6 +651,7 @@ async def test_qdrant_semantic_cache_async_set_cache():
|
|||
assert (
|
||||
upsert_payload[QdrantSemanticCache.CACHE_KEY_FIELD_NAME] == "test_key"
|
||||
)
|
||||
assert qdrant_cache.async_client.put.call_args.kwargs["params"] == {"wait": "true"}
|
||||
|
||||
|
||||
def test_qdrant_semantic_cache_custom_vector_size():
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import asyncio
|
|||
import contextvars
|
||||
import gc
|
||||
import hashlib
|
||||
import http.server
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
|
|
@ -59,6 +60,71 @@ def request(key: str = "key") -> dict[str, object]:
|
|||
return {"key": {"preset": key}}
|
||||
|
||||
|
||||
def qdrant_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: Final = 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")
|
||||
|
|
@ -1691,3 +1757,167 @@ def test_redis_semantic_handle_rejects_wrong_backends(
|
|||
)
|
||||
with pytest.raises(TypeError, match="must be the native embedder"):
|
||||
_CacheTestHandle.redis_semantic(facade.cache)._bind_facade(replacement_facade)
|
||||
|
||||
|
||||
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(qdrant_request("python-key", messages)) == {"id": "py"}
|
||||
binding.store(qdrant_request("native-key", messages), {"id": "native"})
|
||||
python_value: Final = facade.cache.get_cache("native-key", messages=messages)
|
||||
assert isinstance(python_value, dict)
|
||||
assert python_value["response"] == {"id": "native"}
|
||||
unrelated: Final = [{"role": "user", "content": "unrelated prompt"}]
|
||||
assert binding.lookup(qdrant_request("native-key", unrelated)) is None
|
||||
assert facade.cache.get_cache("native-key", messages=unrelated) is None
|
||||
assert binding.lookup(qdrant_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(qdrant_request("python-key", messages)) == {"id": "py"}
|
||||
await binding.async_store(qdrant_request("native-key", messages), {"id": "native"})
|
||||
python_value: Final = await facade.cache.async_get_cache("native-key", messages=messages)
|
||||
assert isinstance(python_value, dict)
|
||||
assert python_value["response"] == {"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(qdrant_request(key, messages)) is None
|
||||
with pytest.raises(RuntimeError, match="operation is not supported"):
|
||||
binding.lookup_batch([qdrant_request(key, messages)])
|
||||
with pytest.raises(RuntimeError, match="operation is not supported"):
|
||||
await binding.async_flush()
|
||||
with pytest.raises(RuntimeError, match="operation is not supported"):
|
||||
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(qdrant_request("persistent-key", messages, ttl_seconds=1.0), {"id": "persistent"})
|
||||
time.sleep(1.2)
|
||||
assert binding.lookup(qdrant_request("persistent-key", messages)) == {"id": "persistent"}
|
||||
python_value: Final = facade.cache.get_cache("persistent-key", messages=messages)
|
||||
assert isinstance(python_value, dict)
|
||||
assert python_value["response"] == {"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.qdrant_api_key = "rotated"
|
||||
assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback"
|
||||
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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue