From ae32609b540d89409abaa20af716620b186d5555 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 22:18:13 +0000 Subject: [PATCH] refactor(cache-qdrant-semantic): inject the shared LiteLLM HTTP client into the embedder Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 1 + .../cache-qdrant-semantic/src/embedder.rs | 32 ++++---- .../cache-qdrant-semantic/tests/embedder.rs | 73 +++++++++++++------ litellm-rust/crates/python-bridge/Cargo.toml | 1 + .../crates/python-bridge/src/cache/config.rs | 19 ++--- .../crates/python-bridge/src/cache/handle.rs | 10 ++- .../crates/python-bridge/src/cache/native.rs | 7 +- 7 files changed, 90 insertions(+), 53 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index f4b94ae44e2..382009c95af 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2821,6 +2821,7 @@ dependencies = [ "pyo3", "pyo3-async-runtimes", "qdrant-client", + "reqwest 0.12.28", "rstest", "serde", "serde_json", diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs b/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs index 0dde0318448..47b898d6f4e 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs @@ -11,6 +11,7 @@ pub struct OpenAiEmbedder { api_base: String, api_key: String, model: String, + timeout: Option, } pub struct OpenAiEmbedderConfig { @@ -21,18 +22,14 @@ pub struct OpenAiEmbedderConfig { } impl OpenAiEmbedder { - pub fn new(config: OpenAiEmbedderConfig) -> Result { - 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, 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) diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs index adce70654a8..6b09448fde8 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs @@ -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) -> 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")); } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 826092522ea..0b9a0148d07 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -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"] } diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 7772716f128..57b0561b122 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -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") } diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 94a2cfec2bd..e71ad9f435d 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -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) })?; diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 0e14b9cb39c..e4810a303c8 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -54,17 +54,18 @@ impl NativeResponseCache { pub async fn qdrant_semantic( config: QdrantSemanticCacheConfig, + client: reqwest::Client, runtime: tokio::runtime::Handle, ) -> Result { - 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,