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

refactor(cache-qdrant-semantic): inject the shared LiteLLM HTTP client into the embedder

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee 2026-09-21 22:18:13 +00:00
parent bda406c0af
commit ae32609b54
7 changed files with 90 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm-rust/Cargo.lock generated
View file

@ -2821,6 +2821,7 @@ dependencies = [
"pyo3",
"pyo3-async-runtimes",
"qdrant-client",
"reqwest 0.12.28",
"rstest",
"serde",
"serde_json",

32
litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs
View file

@ -11,6 +11,7 @@ pub struct OpenAiEmbedder {
api_base: String,
api_key: String,
model: String,
timeout: Option<Duration>,
}
pub struct OpenAiEmbedderConfig {
@ -21,18 +22,14 @@ pub struct OpenAiEmbedderConfig {
}
impl OpenAiEmbedder {
pub fn new(config: OpenAiEmbedderConfig) -> Result<Self, Error> {
let mut builder = Client::builder();
if let Some(timeout) = config.timeout {
builder = builder.timeout(timeout);
}
let client = builder.build().map_err(|_| Error::Unavailable)?;
Ok(Self {
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,
}
}
}
@ -42,7 +39,7 @@ impl Embedder for OpenAiEmbedder {
}
async fn embed(&self, input: &str) -> Result<Vec<f32>, Error> {
let response = self
let request = self
.client
.post(format!("{}/embeddings", self.api_base))
.bearer_auth(&self.api_key)
@ -50,12 +47,17 @@ impl Embedder for OpenAiEmbedder {
"model": self.model,
"input": input,
"encoding_format": "float",
}))
.send()
.await
.map_err(|_| Error::Unavailable)?
.error_for_status()
.map_err(|_| Error::Unavailable)?;
}));
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)

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

@ -19,6 +19,10 @@ struct TestHttpServer {
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));
@ -29,6 +33,7 @@ impl TestHttpServer {
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()
@ -42,20 +47,6 @@ impl TestHttpServer {
}
}
async fn hanging() -> Self {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let task = tokio::spawn(async move {
let (_stream, _) = listener.accept().await.unwrap();
std::future::pending::<()>().await;
});
Self {
address,
request: Arc::new(Mutex::new(None)),
task,
}
}
fn base_url(&self) -> String {
format!("http://{}", self.address)
}
@ -110,11 +101,13 @@ fn config(base: String, timeout: Option<Duration>) -> OpenAiEmbedderConfig {
#[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(config(
format!("{}/", server.base_url()),
Some(Duration::from_secs(1)),
))
.unwrap();
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();
@ -130,12 +123,44 @@ async fn posts_embeddings_request_and_parses_vector() {
#[tokio::test]
async fn status_and_timeout_errors_are_unavailable() {
let server = TestHttpServer::response("500 Internal Server Error", "{}").await;
let embedder =
OpenAiEmbedder::new(config(server.base_url(), Some(Duration::from_secs(1)))).unwrap();
let embedder = OpenAiEmbedder::new(reqwest::Client::new(), config(server.base_url(), None));
assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable));
let server = TestHttpServer::hanging().await;
let embedder =
OpenAiEmbedder::new(config(server.base_url(), Some(Duration::from_millis(200)))).unwrap();
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"));
}

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

@ -40,6 +40,7 @@ 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
serde_json.workspace = true
url.workspace = true
tokio = { workspace = true, features = ["sync"] }

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

@ -216,18 +216,15 @@ impl NativeCacheConfig {
}
pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> {
let default_ttl = match &self.backend {
CacheBackendConfig::Memory(config) => Some(config.default_ttl),
CacheBackendConfig::Redis(config) => Some(config.default_ttl),
CacheBackendConfig::AzureBlob(_) | CacheBackendConfig::QdrantSemantic(_) => None,
};
if service.default_ttl() != default_ttl {
return Some("facade and native backend default TTLs must match");
}
match &self.backend {
CacheBackendConfig::Memory(_) if service.kind() != "memory" => {
Some("facade and native backend types must match")
}
CacheBackendConfig::Memory(config)
if service.default_ttl() != Some(config.default_ttl) =>
{
Some("facade and native backend default TTLs must match")
}
CacheBackendConfig::Memory(config) if service.capacity() != Some(config.capacity) => {
Some("facade and native backend capacities must match")
}
@ -245,7 +242,11 @@ impl NativeCacheConfig {
}
CacheBackendConfig::Redis(config) => (service.namespace()
!= config.namespace.as_deref())
.then_some("facade and native backend namespaces must match"),
.then_some("facade and native backend namespaces must match")
.or_else(|| {
(service.default_ttl() != Some(config.default_ttl))
.then_some("facade and native backend default TTLs must match")
}),
CacheBackendConfig::QdrantSemantic(config) if service.kind() != "qdrant_semantic" => {
Some("facade and native backend types must match")
}

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

@ -1,16 +1,18 @@
use std::env;
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_http::ClientVariant;
use litellm_cache_qdrant_semantic::{OpenAiEmbedderConfig, Quantization};
use litellm_host_python::{release_gil, run_sync_value};
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*};
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*, types::PyDict};
use url::Url;
use super::{
cache_error, config::QdrantSemanticCacheConfig, facade::FacadeGuard,
native::NativeResponseCache, request::duration,
};
use crate::http;
#[pyclass(frozen, name = "_CacheTestHandle")]
pub(crate) struct CacheTestHandle {
@ -154,9 +156,13 @@ impl CacheTestHandle {
},
quantization,
};
let http_config = http::call_config(py, &PyDict::new(py), true)?;
let client = http::pool()
.client(&http_config, ClientVariant::Provider)
.map_err(http::client_error)?;
let service = run_sync_value(py, async move {
let handle = tokio::runtime::Handle::current();
NativeResponseCache::qdrant_semantic(config, handle)
NativeResponseCache::qdrant_semantic(config, client, handle)
.await
.map_err(cache_error)
})?;

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

@ -54,17 +54,18 @@ impl NativeResponseCache {
pub async fn qdrant_semantic(
config: QdrantSemanticCacheConfig,
client: reqwest::Client,
runtime: tokio::runtime::Handle,
) -> Result<Self, Error> {
let client = qdrant_client::Qdrant::from_url(&config.grpc_url)
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(config.embedding)?;
let embedder = OpenAiEmbedder::new(client, config.embedding);
let cache = QdrantSemanticCache::connect(
client,
qdrant,
embedder,
ResponseCacheCodec,
qdrant_config,