fix(python-bridge): await semantic embeddings inline in the caller's task

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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;

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

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::<Vec<_>>();
+                if service.semantic_embedder().is_some() {
+                    return drive(
+                        py,
+                        service.clone(),
+                        SemanticOperation::StoreBatch(entries.into()),
+                    );
+                }
                 let service = service.clone();
                 run_async(
                     py,

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<Vec<f32>, Error>;
+}
+
+pub(super) fn with_prepared_embedding<F: Future>(
+    vector: Result<Vec<f32>, Error>,
+    future: F,
+) -> impl Future<Output = F::Output> {
+    PREPARED_EMBEDDING.scope(vector, future)
+}
+
 pub(super) struct PythonEmbedder(Py<PyAny>);
 
 impl PythonEmbedder {
@@ -34,7 +45,20 @@ impl PythonEmbedder {
         Ok(kwargs)
     }
 
-    fn extract(vector: Bound<'_, PyAny>) -> PyResult<Vec<f32>> {
+    pub(super) fn async_embedding_coroutine(
+        &self,
+        py: Python<'_>,
+        prompt: &str,
+        metadata: &Map<String, Value>,
+    ) -> PyResult<Py<PyAny>> {
+        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<Vec<f32>> {
         Ok(vector
             .extract::<Vec<f64>>()?
             .into_iter()
@@ -58,28 +82,36 @@ impl Embedder for PythonEmbedder {
 
     fn async_embed(
         &self,
-        prompt: &str,
-        metadata: &Map<String, Value>,
+        _prompt: &str,
+        _metadata: &Map<String, Value>,
     ) -> impl Future<Output = Result<Vec<f32>, 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::<Vec<f64>>(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)
+        );
     }
 }

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::{

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<PyAny>> {
         match self {
             Self::RedisSemantic(cache) => Some(cache.backend().embedder().object()),

litellm-rust/crates/python-bridge/src/cache/request.rs

@@ -20,6 +20,7 @@ struct RequestInput {
     scope: Option<String>,
 }
 
+#[derive(Clone)]
 pub(super) struct CacheRequest {
     key: CacheKeyInput,
     controls: CacheControls,

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<Value>)>,
+    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<Vec<f32>, Error>,
+    ) -> PyResult<ExecutionStep> {
+        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<Py<PyAny>>>) -> PyResult<ExecutionStep> {
+        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<Bound<'_, PyAny>> {
+    let execution = Py::new(py, Execution::new(SemanticBody::new(service, operation)))?;
+    py.import("litellm.rust_bridge.lifecycle")?
+        .getattr("drive")?
+        .call1((execution,))
+}

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: