diff --git a/litellm-rust/crates/cache-redis-semantic/src/lib.rs b/litellm-rust/crates/cache-redis-semantic/src/lib.rs index a34603cd18f..51d0b4ba5f3 100644 --- a/litellm-rust/crates/cache-redis-semantic/src/lib.rs +++ b/litellm-rust/crates/cache-redis-semantic/src/lib.rs @@ -2,3 +2,4 @@ mod cache; mod prompt; pub use cache::{Embedder, RedisSemanticCache, RedisSemanticConfig}; +pub use prompt::prompt_from_context; diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index b28ddc50181..93ce5828489 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -39,7 +39,7 @@ litellm-token-counter = { path = "../token-counter", default-features = false } pyo3.workspace = true pyo3-async-runtimes.workspace = true serde_json.workspace = true -tokio = { workspace = true, features = ["sync"] } +tokio = { workspace = true, features = ["rt", "sync"] } [dev-dependencies] serde.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/binding.rs b/litellm-rust/crates/python-bridge/src/cache/binding.rs index ad64b24d3c1..0b90e8151ea 100644 --- a/litellm-rust/crates/python-bridge/src/cache/binding.rs +++ b/litellm-rust/crates/python-bridge/src/cache/binding.rs @@ -14,6 +14,7 @@ use super::{ future::{ready_none, ready_value}, native::NativeResponseCache, request::{now, request, requests}, + semantic::{SemanticOperation, drive}, }; pub(super) enum CacheBinding { @@ -56,6 +57,11 @@ impl ResolvedCache { CacheBinding::Disabled => ready_none(py)?, CacheBinding::Native(service) => { let request = request(input)?; + if service.semantic_embedder().is_some() { + return Ok(ExecutionStep::Await( + drive(py, service.clone(), SemanticOperation::Lookup(request))?.unbind(), + )); + } let service = service.clone(); run_async( py, @@ -179,6 +185,13 @@ impl ResolvedCache { CacheBinding::Native(service) => { let request = self::request(request)?; let response: Value = from_py(response)?; + if service.semantic_embedder().is_some() { + return drive( + py, + service.clone(), + SemanticOperation::Store(request, response), + ); + } let service = service.clone(); run_async( py, @@ -240,7 +253,14 @@ impl ResolvedCache { "batch cache requests and responses must have equal lengths", )); } - let entries = requests.into_iter().zip(responses).collect(); + let entries = requests.into_iter().zip(responses).collect::>(); + if service.semantic_embedder().is_some() { + return drive( + py, + service.clone(), + SemanticOperation::StoreBatch(entries.into()), + ); + } let service = service.clone(); run_async( py, diff --git a/litellm-rust/crates/python-bridge/src/cache/embedder.rs b/litellm-rust/crates/python-bridge/src/cache/embedder.rs index 63e078cd815..26edb26f428 100644 --- a/litellm-rust/crates/python-bridge/src/cache/embedder.rs +++ b/litellm-rust/crates/python-bridge/src/cache/embedder.rs @@ -6,6 +6,17 @@ use litellm_host_python::to_py; use pyo3::{PyTraverseError, PyVisit, prelude::*, types::PyDict}; use serde_json::{Map, Value}; +tokio::task_local! { + static PREPARED_EMBEDDING: Result, Error>; +} + +pub(super) fn with_prepared_embedding( + vector: Result, Error>, + future: F, +) -> impl Future { + PREPARED_EMBEDDING.scope(vector, future) +} + pub(super) struct PythonEmbedder(Py); impl PythonEmbedder { @@ -34,7 +45,20 @@ impl PythonEmbedder { Ok(kwargs) } - fn extract(vector: Bound<'_, PyAny>) -> PyResult> { + pub(super) fn async_embedding_coroutine( + &self, + py: Python<'_>, + prompt: &str, + metadata: &Map, + ) -> PyResult> { + let kwargs = Self::metadata_kwargs(py, metadata)?; + self.0 + .bind(py) + .call_method("_get_async_embedding", (prompt,), Some(&kwargs)) + .map(Bound::unbind) + } + + pub(super) fn extract(vector: Bound<'_, PyAny>) -> PyResult> { Ok(vector .extract::>()? .into_iter() @@ -58,28 +82,36 @@ impl Embedder for PythonEmbedder { fn async_embed( &self, - prompt: &str, - metadata: &Map, + _prompt: &str, + _metadata: &Map, ) -> impl Future, Error>> + Send { - let coroutine = Python::attach(|py| { - let kwargs = Self::metadata_kwargs(py, metadata)?; - self.0 - .bind(py) - .call_method("_get_async_embedding", (prompt,), Some(&kwargs)) - .map(Bound::unbind) - }) - .map_err(|_| Error::Unavailable); - async move { - let coroutine = coroutine?; - let awaited = Python::attach(|py| { - pyo3_async_runtimes::tokio::into_future(coroutine.into_bound(py)) - }) - .map_err(|_| Error::Unavailable)? - .await - .map_err(|_| Error::Unavailable)?; - let vector = Python::attach(|py| awaited.extract::>(py)) - .map_err(|_| Error::Unavailable)?; - Ok(vector.into_iter().map(|value| value as f32).collect()) - } + let seeded = PREPARED_EMBEDDING + .try_with(Clone::clone) + .unwrap_or(Err(Error::Unavailable)); + std::future::ready(seeded) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn async_embed_returns_the_seeded_vector_or_unavailable() { + let embedder = Python::attach(|py| PythonEmbedder::new(py.None())); + let metadata = Map::new(); + let embedder_ref = &embedder; + let metadata_ref = &metadata; + assert_eq!( + with_prepared_embedding(Ok(vec![0.5f32, 0.25]), async move { + embedder_ref.async_embed("prompt", metadata_ref).await + }) + .await, + Ok(vec![0.5, 0.25]) + ); + assert_eq!( + embedder.async_embed("prompt", &metadata).await, + Err(Error::Unavailable) + ); } } diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 4cc87367d91..cd772d571cb 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -8,6 +8,7 @@ mod handle; mod native; mod request; mod resolver; +mod semantic; use litellm_cache::Error; use pyo3::{ diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 182010fab02..de9c4afa236 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -122,6 +122,13 @@ impl NativeResponseCache { } } + pub fn semantic_embedder(&self) -> Option<&PythonEmbedder> { + match self { + Self::RedisSemantic(cache) => Some(cache.backend().embedder()), + Self::Memory(_) | Self::Redis { .. } => None, + } + } + pub fn embedder_object(&self) -> Option<&Py> { match self { Self::RedisSemantic(cache) => Some(cache.backend().embedder().object()), diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs index 26e0fe4e62c..b06087bcc83 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -20,6 +20,7 @@ struct RequestInput { scope: Option, } +#[derive(Clone)] pub(super) struct CacheRequest { key: CacheKeyInput, controls: CacheControls, diff --git a/litellm-rust/crates/python-bridge/src/cache/semantic.rs b/litellm-rust/crates/python-bridge/src/cache/semantic.rs new file mode 100644 index 00000000000..eb38b8b9c67 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/semantic.rs @@ -0,0 +1,165 @@ +use std::collections::VecDeque; + +use litellm_cache::Error; +use litellm_cache_redis_semantic::prompt_from_context; +use litellm_host_python::{Execution, ExecutionBody, ExecutionStep, run_async}; +use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; +use serde_json::Value; + +use super::{ + cache_error, + embedder::{PythonEmbedder, with_prepared_embedding}, + native::NativeResponseCache, + request::{CacheRequest, now}, +}; + +pub(super) enum SemanticOperation { + Lookup(CacheRequest), + Store(CacheRequest, Value), + StoreBatch(VecDeque<(CacheRequest, Value)>), +} + +enum Phase { + Start, + AwaitingEmbedding, + AwaitingBackend, +} + +pub(super) struct SemanticBody { + service: NativeResponseCache, + operation: SemanticOperation, + pending: Option<(CacheRequest, Option)>, + phase: Phase, +} + +impl SemanticBody { + pub(super) fn new(service: NativeResponseCache, operation: SemanticOperation) -> Self { + Self { + service, + operation, + pending: None, + phase: Phase::Start, + } + } + + fn backend_step( + &mut self, + py: Python<'_>, + seed: Result, Error>, + ) -> PyResult { + self.phase = Phase::AwaitingBackend; + let (request, response) = self.pending.take().ok_or_else(|| { + PyRuntimeError::new_err("semantic execution resumed without a pending operation") + })?; + let service = self.service.clone(); + let future = async move { + match response { + None => service.async_lookup(&request, now()).await, + Some(response) => service + .async_store(&request, response, now()) + .await + .map(|_| None), + } + }; + let awaitable = run_async(py, with_prepared_embedding(seed, future), cache_error)?; + Ok(ExecutionStep::Await(awaitable.unbind())) + } +} + +impl ExecutionBody for SemanticBody { + fn resume(&mut self, mut result: Option>>) -> PyResult { + Python::attach(|py| { + loop { + match self.phase { + Phase::Start => { + if result.is_some() { + return Err(PyRuntimeError::new_err( + "semantic execution received a result before starting", + )); + } + if self.pending.is_none() { + match &mut self.operation { + SemanticOperation::Lookup(request) => { + self.pending = Some((request.clone(), None)); + } + SemanticOperation::Store(request, response) => { + let response = std::mem::replace(response, Value::Null); + self.pending = Some((request.clone(), Some(response))); + } + SemanticOperation::StoreBatch(queue) => { + let Some((request, response)) = queue.pop_front() else { + return Ok(ExecutionStep::Return(py.None())); + }; + self.pending = Some((request, Some(response))); + } + } + } + let (request, _) = self.pending.as_ref().ok_or_else(|| { + PyRuntimeError::new_err("semantic execution has no pending operation") + })?; + let semantic = request.semantic(); + let Some(prompt) = prompt_from_context(&semantic.context) else { + return self.backend_step(py, Err(Error::Unavailable)); + }; + let embedder = self.service.semantic_embedder().ok_or_else(|| { + PyRuntimeError::new_err( + "semantic execution requires a redis-semantic backend", + ) + })?; + let coroutine = embedder.async_embedding_coroutine( + py, + &prompt, + &semantic.context.metadata, + )?; + self.phase = Phase::AwaitingEmbedding; + return Ok(ExecutionStep::Await(coroutine)); + } + Phase::AwaitingEmbedding => { + let result = result.take().ok_or_else(|| { + PyRuntimeError::new_err( + "semantic execution expected an embedding result", + ) + })?; + let seed = result + .and_then(|value| PythonEmbedder::extract(value.into_bound(py))) + .map_err(|_| Error::Unavailable); + return self.backend_step(py, seed); + } + Phase::AwaitingBackend => { + let result = result.take().ok_or_else(|| { + PyRuntimeError::new_err("semantic execution expected a backend result") + })?; + let value = match result { + Ok(value) => value, + Err(error) => return Err(error), + }; + let more = matches!( + &self.operation, + SemanticOperation::StoreBatch(queue) if !queue.is_empty() + ); + if more { + self.phase = Phase::Start; + continue; + } + return Ok(ExecutionStep::Return(value)); + } + } + } + }) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + self.service.traverse(visit) + } +} + +pub(super) fn drive( + py: Python<'_>, + service: NativeResponseCache, + operation: SemanticOperation, +) -> PyResult> { + let execution = Py::new(py, Execution::new(SemanticBody::new(service, operation)))?; + py.import("litellm.rust_bridge.lifecycle")? + .getattr("drive")? + .call1((execution,)) +} diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index a60feb9973a..9312fdde075 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -479,6 +479,7 @@ async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_n PARAPHRASE_MARKER: Final = " (paraphrase)" SEMANTIC_EMBEDDING_MODEL: Final = "semantic-test/deterministic" SEMANTIC_INDEX_PREFIX: Final = "litellm_test_semantic_" +SEMANTIC_CONTEXT: Final = contextvars.ContextVar("semantic_test_context", default="unset") def _normalized(vector: list[float]) -> list[float]: @@ -509,6 +510,7 @@ def _semantic_embedding(prompt: str) -> list[float]: class DeterministicEmbedding(litellm.CustomLLM): def __init__(self) -> None: self.calls: list[dict[str, object]] = [] + self.async_calls: list[dict[str, object]] = [] def _respond( self, @@ -553,6 +555,16 @@ class DeterministicEmbedding(litellm.CustomLLM): timeout: object = None, litellm_params: object = None, ) -> EmbeddingResponse: + texts: Final = cast(list[object], input if isinstance(input, list) else [input]) + self.async_calls.append( + { + "model": model, + "input": texts, + "task": asyncio.current_task(), + "context": SEMANTIC_CONTEXT.get(), + } + ) + SEMANTIC_CONTEXT.set("written-in-aembedding") return self._respond(model, input, model_response) @@ -755,6 +767,50 @@ async def test_redis_semantic_async_paths_and_store_batch_share_one_layout( client.close() +async def test_native_semantic_async_embedding_runs_inline_in_the_callers_task( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + caller: Final = asyncio.current_task() + SEMANTIC_CONTEXT.set("caller-sentinel") + response: Final = {"choices": [{"text": "paris"}]} + + await binding.async_store( + semantic_request("inline", "what is the capital of france"), response + ) + assert ( + await binding.async_lookup( + semantic_request("inline", f"what is the capital of france{PARAPHRASE_MARKER}") + ) + == response + ) + assert await binding.async_lookup(semantic_request("inline", "python written prompt")) is None + assert SEMANTIC_CONTEXT.get() == "written-in-aembedding" + assert semantic_embedding.async_calls == [ + { + "model": "deterministic", + "input": ["what is the capital of france"], + "task": caller, + "context": "caller-sentinel", + }, + { + "model": "deterministic", + "input": [f"what is the capital of france{PARAPHRASE_MARKER}"], + "task": caller, + "context": "written-in-aembedding", + }, + { + "model": "deterministic", + "input": ["python written prompt"], + "task": caller, + "context": "written-in-aembedding", + }, + ], semantic_embedding.async_calls + + def test_redis_semantic_similarity_tag_and_threshold_boundaries( redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding ) -> None: