diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 09896a96952..3781aba6d9f 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4240,6 +4240,7 @@ dependencies = [ "litellm-storage-clickhouse", "litellm-token-counter", "litellm-tracing", + "parking_lot", "prost", "pyo3", "pyo3-async-runtimes", diff --git a/litellm-rust/crates/host-python/src/binding.rs b/litellm-rust/crates/host-python/src/binding.rs index 91fbc8c0e52..3b3f4c5080c 100644 --- a/litellm-rust/crates/host-python/src/binding.rs +++ b/litellm-rust/crates/host-python/src/binding.rs @@ -13,7 +13,7 @@ pub trait PythonBinding: PythonOwned { /// Decodes the keyword view returned by `prepare_arguments`, including preflight rewrites. fn decode_request( - &mut self, + &self, py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> Result< @@ -22,21 +22,21 @@ pub trait PythonBinding: PythonOwned { >; fn encode_response( - &mut self, + &self, py: Python<'_>, response: ::Response, ) -> PyResult>; /// What the stream carries at hand-off, as the caller's stream receives it. fn encode_stream_head( - &mut self, + &self, py: Python<'_>, head: ::StreamHead, ) -> PyResult>; /// One streamed chunk as the caller receives it. fn encode_chunk( - &mut self, + &self, py: Python<'_>, chunk: ::Chunk, ) -> PyResult>; diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index 7a2eac87f13..bc51242dbf2 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -91,15 +91,17 @@ impl Observation { } } -/// Hooks lock their own state around Python work, so the handle can traverse them while a -/// hook is suspended. +/// What the running call shares with its handle: the binding and hooks the handle +/// traverses and closes, and the Python values the call retains across a suspension. +/// Binding and hooks lock their own state only around Python work, so the handle can +/// traverse them while a hook or host call is suspended. struct Shared { + binding: H, hooks: L, - parts: Mutex>, + parts: Mutex, } -struct Parts { - binding: H, +struct Parts { arguments: Option>, interrupted: Option>, selected: Option, @@ -151,7 +153,7 @@ where M: Machine + 'static, M::Complete: Into>>, { - fn parts(&self) -> MutexGuard<'_, Parts> { + fn parts(&self) -> MutexGuard<'_, Parts> { self.shared.parts.lock() } @@ -159,6 +161,10 @@ where &self.shared.hooks } + fn binding(&self) -> &H { + &self.shared.binding + } + async fn run(mut self) -> PyResult> { match self.begin().await { Ok(()) => self.execute().await, @@ -198,12 +204,8 @@ where self.hooks() .arguments_prepared(&prepared) .map_err(Setup::Host)?; - let decoded = { - let mut parts = self.parts(); - let decoded = Python::attach(|py| parts.binding.decode_request(py, prepared.bind(py))); - parts.arguments = Some(prepared); - decoded - }; + let decoded = Python::attach(|py| self.binding().decode_request(py, prepared.bind(py))); + self.parts().arguments = Some(prepared); let request = decoded.map_err(Setup::Decode)?; let start = self .start @@ -271,7 +273,9 @@ where op: HostRequest>, ) -> Result<(), InvokeError>> { match op { - HostRequest::HostCall(call) => self.host_call(call).await, + HostRequest::HostCall(call) => { + self.binding().handle_host_call(&self.caller, call).await + } HostRequest::Intercept(InterceptRequest::BeforeProviderRequest { wire, context, @@ -302,19 +306,13 @@ where } HostRequest::Stream(StreamDelivery::Open(head, reply)) => { self.streaming = true; - let head = { - let mut parts = self.parts(); - Python::attach(|py| parts.binding.encode_stream_head(py, head))? - }; + let head = Python::attach(|py| self.binding().encode_stream_head(py, head))?; self.hooks().on_stream_open(&head)?; reply.send(self.caller.open_stream(head).await); Ok(()) } HostRequest::Stream(StreamDelivery::Chunk(chunk, reply)) => { - let chunk = { - let mut parts = self.parts(); - Python::attach(|py| parts.binding.encode_chunk(py, chunk))? - }; + let chunk = Python::attach(|py| self.binding().encode_chunk(py, chunk))?; self.hooks().on_stream_chunk(&chunk)?; reply.send(self.caller.yield_chunk(chunk).await); Ok(()) @@ -322,24 +320,6 @@ where } } - async fn host_call( - &self, - call: as Protocol>::HostCall, - ) -> Result<(), InvokeError>> { - let mut pending = { - let mut parts = self.parts(); - Python::attach(|py| parts.binding.begin_host_call(py, call))? - }; - while let Some(awaitable) = pending { - let result = self.caller.await_python(awaitable).await; - pending = { - let mut parts = self.parts(); - Python::attach(|py| parts.binding.resume_host_call(py, result))? - }; - } - Ok(()) - } - async fn notify(&self, event: PythonCallEvent<'_>) -> PyResult<()> { self.parts().observation.emit(&event); self.hooks().on_event(&self.caller, event).await @@ -357,11 +337,7 @@ where return self.succeeded(none).await; } }; - let encoded = { - let mut parts = self.parts(); - Python::attach(|py| parts.binding.encode_response(py, response)) - }; - let public = match encoded { + let public = match Python::attach(|py| self.binding().encode_response(py, response)) { Ok(public) => public, Err(error) => return self.failure(error, FailureOrigin::Call).await, }; @@ -385,7 +361,7 @@ where parts.observation.ended_at.get_or_insert_with(epoch_seconds); Python::attach(|py| match parts.interrupted.take() { Some(retained) => PyErr::from_value(retained.into_bound(py).into_any()), - None => classified(py, &parts.binding, error), + None => classified(py, self.binding(), error), }) }; self.failure(error, FailureOrigin::Call).await @@ -527,11 +503,11 @@ where } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + self.shared.binding.traverse(visit)?; self.shared.hooks.traverse(visit)?; let Some(parts) = self.shared.parts.try_lock() else { return Ok(()); }; - parts.binding.traverse(visit)?; visit.call(&parts.arguments)?; visit.call(&parts.interrupted)?; match &parts.selected { @@ -550,14 +526,15 @@ where fn drop(&mut self) { self.body.get_mut().cancel(); self.reply = None; - let Some(Shared { hooks, parts }) = Arc::get_mut(&mut self.shared) else { + let Some(Shared { + binding, + hooks, + parts, + }) = Arc::get_mut(&mut self.shared) + else { return; }; - let Parts { - binding, - observation, - .. - } = parts.get_mut(); + let observation = &mut parts.get_mut().observation; Python::attach(|py| { if observation.observing { let timing = observation.timing(); @@ -592,9 +569,9 @@ where M::Complete: Into>>, { let shared = Arc::new(Shared { + binding, hooks, parts: Mutex::new(Parts { - binding, arguments: Some(arguments), interrupted: None, selected: None, diff --git a/litellm-rust/crates/host-python/src/services.rs b/litellm-rust/crates/host-python/src/services.rs index 3d8faa2a5aa..c6772efc349 100644 --- a/litellm-rust/crates/host-python/src/services.rs +++ b/litellm-rust/crates/host-python/src/services.rs @@ -1,27 +1,12 @@ -use crate::{InvokeError, PythonOwned}; +use std::future::Future; + +use crate::{Awaiter, InvokeError, PythonOwned}; use litellm_host::protocol::Protocol; -use pyo3::prelude::*; pub trait PythonHostCalls: PythonOwned { fn handle_host_call( - &mut self, - py: Python<'_>, + &self, + task: &dyn Awaiter, call: P::HostCall, - ) -> Result<(), InvokeError>; - - fn begin_host_call( - &mut self, - py: Python<'_>, - call: P::HostCall, - ) -> Result>, InvokeError> { - self.handle_host_call(py, call).map(|()| None) - } - - fn resume_host_call( - &mut self, - _: Python<'_>, - result: PyResult>, - ) -> Result>, InvokeError> { - result.map(|_| None).map_err(InvokeError::Python) - } + ) -> impl Future>> + Send; } diff --git a/litellm-rust/crates/host-python/tests/support/mod.rs b/litellm-rust/crates/host-python/tests/support/mod.rs index 4767cdec19e..0dd7b4ce9b1 100644 --- a/litellm-rust/crates/host-python/tests/support/mod.rs +++ b/litellm-rust/crates/host-python/tests/support/mod.rs @@ -167,7 +167,6 @@ pub struct SyntheticBinding { log: Log, op: OpScript, classifier_fails: bool, - pending_reply: Option>, } pub fn binding(log: &Log, op: OpScript) -> SyntheticBinding { @@ -175,7 +174,6 @@ pub fn binding(log: &Log, op: OpScript) -> SyntheticBinding { log: log.clone(), op, classifier_fails: false, - pending_reply: None, } } @@ -206,7 +204,7 @@ impl PythonBinding for SyntheticBinding { type Failure = Classified; fn decode_request( - &mut self, + &self, _: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> Result> { @@ -220,22 +218,18 @@ impl PythonBinding for SyntheticBinding { } fn encode_stream_head( - &mut self, + &self, _: Python<'_>, head: std::convert::Infallible, ) -> PyResult> { match head {} } - fn encode_chunk( - &mut self, - _: Python<'_>, - chunk: std::convert::Infallible, - ) -> PyResult> { + fn encode_chunk(&self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult> { match chunk {} } - fn encode_response(&mut self, py: Python<'_>, response: String) -> PyResult> { + fn encode_response(&self, py: Python<'_>, response: String) -> PyResult> { self.log.push("complete"); Ok(PyString::new(py, &response).into_any().unbind()) } @@ -254,49 +248,34 @@ impl PythonBinding for SyntheticBinding { } impl PythonHostCalls for SyntheticBinding { - fn handle_host_call( - &mut self, - _: Python<'_>, + async fn handle_host_call( + &self, + task: &dyn Awaiter, (op, reply): (&'static str, Reply), ) -> Result<(), InvokeError> { self.log.push(format!("op:{op}")); - self.answer(|| op.to_string()) - .map(|answer| reply.send(answer)) - } - - fn begin_host_call( - &mut self, - py: Python<'_>, - (op, reply): (&'static str, Reply), - ) -> Result>, InvokeError> { if !matches!(self.op, OpScript::AwaitAnswer | OpScript::AwaitFailure) { - return self.handle_host_call(py, (op, reply)).map(|()| None); + return self + .answer(|| op.to_string()) + .map(|answer| reply.send(answer)); } - self.pending_reply = Some(reply); - let module = PyModule::from_code( - py, - pyo3::ffi::c_str!( - "async def answer(fail):\n if fail:\n raise LookupError('async host failed')\n return 'awaited'\n" - ), - pyo3::ffi::c_str!("host_op.py"), - pyo3::ffi::c_str!("host_op"), - )?; - Ok(Some( - module + let awaitable = Python::attach(|py| -> PyResult> { + let module = PyModule::from_code( + py, + pyo3::ffi::c_str!( + "async def answer(fail):\n if fail:\n raise LookupError('async host failed')\n return 'awaited'\n" + ), + pyo3::ffi::c_str!("host_op.py"), + pyo3::ffi::c_str!("host_op"), + )?; + Ok(module .getattr("answer")? .call1((matches!(self.op, OpScript::AwaitFailure),))? - .unbind(), - )) - } - - fn resume_host_call( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> Result>, InvokeError> { - let answer = result?.extract::(py)?; - self.pending_reply.take().unwrap().send(answer); - Ok(None) + .unbind()) + })?; + let answer = task.awaited(awaitable).await?; + reply.send(Python::attach(|py| answer.extract::(py))?); + Ok(()) } } @@ -561,7 +540,7 @@ impl PythonBinding for StreamingBinding { type Failure = Classified; fn decode_request( - &mut self, + &self, _: Python<'_>, _: &Bound<'_, PyDict>, ) -> Result<(), InvokeError> { @@ -569,7 +548,7 @@ impl PythonBinding for StreamingBinding { } fn encode_stream_head( - &mut self, + &self, py: Python<'_>, head: Vec<(&'static str, &'static str)>, ) -> PyResult> { @@ -582,11 +561,11 @@ impl PythonBinding for StreamingBinding { Ok(hidden.into_any().unbind()) } - fn encode_chunk(&mut self, py: Python<'_>, chunk: &'static str) -> PyResult> { + fn encode_chunk(&self, py: Python<'_>, chunk: &'static str) -> PyResult> { Ok(PyString::new(py, chunk).into_any().unbind()) } - fn encode_response(&mut self, py: Python<'_>, (): ()) -> PyResult> { + fn encode_response(&self, py: Python<'_>, (): ()) -> PyResult> { Ok(py.None()) } @@ -600,9 +579,9 @@ impl PythonBinding for StreamingBinding { } impl PythonHostCalls for StreamingBinding { - fn handle_host_call( - &mut self, - _: Python<'_>, + async fn handle_host_call( + &self, + _: &dyn Awaiter, op: std::convert::Infallible, ) -> Result<(), InvokeError> { match op {} diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 6fc05c43146..e4079500eea 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -26,6 +26,7 @@ litellm-storage-clickhouse.workspace = true litellm-host.workspace = true bytes.workspace = true futures-util.workspace = true +parking_lot.workspace = true litellm-cache.workspace = true litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/python/host.rs b/litellm-rust/crates/python-bridge/src/cache/python/host.rs index d8f8ac7c181..16490318e87 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/host.rs +++ b/litellm-rust/crates/python-bridge/src/cache/python/host.rs @@ -1,7 +1,9 @@ use litellm_cache::Error; use litellm_host::protocol::Reply; -use litellm_host_python::{from_py, to_py}; +use litellm_host_python::{Awaiter, from_py, to_py}; +use parking_lot::Mutex; use pyo3::{ + exceptions::{PyException, PyRuntimeError}, gc::{PyTraverseError, PyVisit}, prelude::*, types::PyDict, @@ -15,127 +17,116 @@ enum Pending { Store(Reply>), } +struct Configured { + cache: Py, + arguments: Py, +} + pub(crate) struct PythonCache { - cache: Option>, - arguments: Option>, - pending: Option, + configured: Mutex>, asynchronous: bool, } impl PythonCache { pub fn new(asynchronous: bool) -> Self { Self { - cache: None, - arguments: None, - pending: None, + configured: Mutex::new(None), asynchronous, } } - pub(in crate::cache) fn bind( - &mut self, - cache: Bound<'_, PyAny>, - arguments: &Bound<'_, PyDict>, - ) { - self.cache = Some(cache.unbind()); - self.arguments = Some(arguments.clone().unbind()); + pub(in crate::cache) fn bind(&self, cache: Bound<'_, PyAny>, arguments: &Bound<'_, PyDict>) { + *self.configured.lock() = Some(Configured { + cache: cache.unbind(), + arguments: arguments.clone().unbind(), + }); } - pub fn begin(&mut self, py: Python<'_>, call: CacheCall) -> PyResult>> { - let Some(cache) = self.cache.as_ref() else { - return Err(pyo3::exceptions::PyRuntimeError::new_err( + /// Runs one cache operation against the configured Python cache, awaiting it on the + /// caller's task for an async call. A cache failure is an answer to the route, a + /// cancellation ends the call. + pub async fn call(&self, task: &dyn Awaiter, call: CacheCall) -> PyResult<()> { + let (pending, invoked) = Python::attach(|py| self.invoke(py, call))?; + let result = match invoked { + Ok(awaitable) if self.asynchronous => task.awaited(awaitable).await, + invoked => invoked, + }; + Python::attach(|py| answer(py, pending, result)) + } + + fn invoke(&self, py: Python<'_>, call: CacheCall) -> PyResult<(Pending, PyResult>)> { + let configured = self.configured.lock(); + let Some(configured) = configured.as_ref() else { + return Err(PyRuntimeError::new_err( "cache operation without configured cache", )); }; - let arguments = self - .arguments - .as_ref() - .ok_or_else(|| { - pyo3::exceptions::PyRuntimeError::new_err("cache arguments unavailable") - })? - .bind(py) - .copy()?; - let (method, result) = match call { - CacheCall::Lookup { reply } => { - self.pending = Some(Pending::Lookup(reply)); - ( - if self.asynchronous { + let arguments = configured.arguments.bind(py).copy()?; + let cache = configured.cache.bind(py); + let (pending, invoked) = match (call, self.asynchronous) { + (CacheCall::Lookup { reply }, asynchronous) => ( + Pending::Lookup(reply), + cache.call_method( + if asynchronous { "async_get_cache" } else { "get_cache" }, - None, - ) - } - CacheCall::Store { value, reply } => { - self.pending = Some(Pending::Store(reply)); - ( - if self.asynchronous { + (), + Some(&arguments), + ), + ), + (CacheCall::Store { value, reply }, asynchronous) => ( + Pending::Store(reply), + cache.call_method( + if asynchronous { "async_add_cache" } else { "add_cache" }, - Some(to_py(py, &value)?), - ) - } + (to_py(py, &value)?,), + Some(&arguments), + ), + ), }; - let result = match result { - Some(value) => cache - .bind(py) - .call_method(method, (value,), Some(&arguments)), - None => cache.bind(py).call_method(method, (), Some(&arguments)), - } - .map(Bound::unbind); - if self.asynchronous && result.is_ok() { - return result.map(Some); - } - self.resume(py, result) - } - - pub fn resume( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> PyResult>> { - if let Err(error) = &result - && !error.is_instance_of::(py) - { - self.pending = None; - return Err(result.err().unwrap()); - } - match self.pending.take() { - Some(Pending::Lookup(reply)) => { - let value = result.map_err(|_| Error::Unavailable).and_then(|value| { - if value.bind(py).is_none() { - Ok(None) - } else { - from_py(value.bind(py)) - .map(Some) - .map_err(|_| Error::InvalidEntry) - } - }); - reply.send(value); - } - Some(Pending::Store(reply)) => { - reply.send(result.map(|_| ()).map_err(|_| Error::Unavailable)); - } - None => { - return Err(pyo3::exceptions::PyRuntimeError::new_err( - "cache reply without pending operation", - )); - } - } - Ok(None) + Ok((pending, invoked.map(Bound::unbind))) } pub fn close(&mut self) { - self.pending = None; - self.cache = None; - self.arguments = None; + *self.configured.get_mut() = None; } pub fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.cache)?; - visit.call(&self.arguments) + let Some(configured) = self.configured.try_lock() else { + return Ok(()); + }; + let Some(configured) = configured.as_ref() else { + return Ok(()); + }; + visit.call(&configured.cache)?; + visit.call(&configured.arguments) } } + +fn answer(py: Python<'_>, pending: Pending, result: PyResult>) -> PyResult<()> { + if let Err(error) = &result + && !error.is_instance_of::(py) + { + return result.map(drop); + } + match pending { + Pending::Lookup(reply) => { + reply.send(result.map_err(|_| Error::Unavailable).and_then(|value| { + if value.bind(py).is_none() { + Ok(None) + } else { + from_py(value.bind(py)) + .map(Some) + .map_err(|_| Error::InvalidEntry) + } + })) + } + Pending::Store(reply) => reply.send(result.map(drop).map_err(|_| Error::Unavailable)), + } + Ok(()) +} diff --git a/litellm-rust/crates/python-bridge/src/cache/selection.rs b/litellm-rust/crates/python-bridge/src/cache/selection.rs index b2dc67fe965..593c517843b 100644 --- a/litellm-rust/crates/python-bridge/src/cache/selection.rs +++ b/litellm-rust/crates/python-bridge/src/cache/selection.rs @@ -125,7 +125,7 @@ fn native_configuration( } pub(crate) fn configure( - python: &mut python::PythonCache, + python: &python::PythonCache, py: Python<'_>, arguments: &Bound<'_, PyDict>, call_type: &str, diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs index 5ad9f8df8b8..1f41cb68d71 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs @@ -1,7 +1,7 @@ use std::convert::Infallible; use crate::routes::inference::InferenceHost; -use litellm_host_python::{InvokeError, PythonBinding, PythonHostCalls, PythonOwned}; +use litellm_host_python::{Awaiter, InvokeError, PythonBinding, PythonHostCalls, PythonOwned}; use litellm_inference_chat::{Error, route::ChatCompletions, types::ChatCompletionsCall}; use pyo3::{ gc::{PyTraverseError, PyVisit}, @@ -12,7 +12,7 @@ use pyo3::{ pub(super) struct ChatCompletionsPythonHost(pub InferenceHost); pub(super) fn project( - host: &mut InferenceHost, + host: &InferenceHost, py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> PyResult { @@ -34,11 +34,11 @@ impl PythonBinding for ChatCompletionsPythonHost { type Failure = PyErr; fn decode_request( - &mut self, + &self, py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> Result> { - let call = project(&mut self.0, py, arguments).map_err(InvokeError::Python)?; + let call = project(&self.0, py, arguments).map_err(InvokeError::Python)?; if call .optional_params .get("stream") @@ -52,7 +52,7 @@ impl PythonBinding for ChatCompletionsPythonHost { } fn encode_response( - &mut self, + &self, py: Python<'_>, response: ::Response, ) -> PyResult> { @@ -60,7 +60,7 @@ impl PythonBinding for ChatCompletionsPythonHost { } fn encode_stream_head( - &mut self, + &self, _py: Python<'_>, head: ::StreamHead, ) -> PyResult> { @@ -68,7 +68,7 @@ impl PythonBinding for ChatCompletionsPythonHost { } fn encode_chunk( - &mut self, + &self, _py: Python<'_>, chunk: ::Chunk, ) -> PyResult> { @@ -84,9 +84,9 @@ impl PythonBinding for ChatCompletionsPythonHost { } impl PythonHostCalls for ChatCompletionsPythonHost { - fn handle_host_call( - &mut self, - _: Python<'_>, + async fn handle_host_call( + &self, + _: &dyn Awaiter, op: Infallible, ) -> Result<(), InvokeError> { match op {} diff --git a/litellm-rust/crates/python-bridge/src/routes/inference.rs b/litellm-rust/crates/python-bridge/src/routes/inference.rs index bc442d1cf36..cff36ee8fbe 100644 --- a/litellm-rust/crates/python-bridge/src/routes/inference.rs +++ b/litellm-rust/crates/python-bridge/src/routes/inference.rs @@ -44,7 +44,7 @@ impl InferenceHost { } pub fn project( - &mut self, + &self, py: Python<'_>, arguments: &Bound<'_, PyDict>, input: &str, diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index 008f1bde763..74fed1c8c2f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -1,5 +1,5 @@ use crate::cache::{CacheCall, Cached, PythonCache, Selection}; -use litellm_host_python::{PythonHostCalls, PythonOwned}; +use litellm_host_python::{Awaiter, PythonHostCalls, PythonOwned}; use bytes::Bytes; use litellm_host_python::{InvokeError, PythonBinding, from_py, to_py}; @@ -261,16 +261,15 @@ impl PythonBinding for MessagesPythonHost { type Failure = PyErr; fn decode_request( - &mut self, + &self, py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> Result<(MessagesCall, Selection), InvokeError> { self.view .resolve(py, arguments) .map_err(InvokeError::Python)?; - let selection = - crate::cache::configure(&mut self.cache, py, arguments, "anthropic_messages") - .map_err(InvokeError::Python)?; + let selection = crate::cache::configure(&self.cache, py, arguments, "anthropic_messages") + .map_err(InvokeError::Python)?; self.projection(py, arguments) .map_err(|error| InvokeError::Python(self.map_failure(py, error)))? .map_err(InvokeError::Native) @@ -278,25 +277,21 @@ impl PythonBinding for MessagesPythonHost { } fn encode_response( - &mut self, + &self, py: Python<'_>, response: Box, ) -> PyResult> { public_response(py, ROUTE_HOST_MODULE, response.as_ref()) } - fn encode_stream_head( - &mut self, - py: Python<'_>, - head: MessagesStreamHead, - ) -> PyResult> { + fn encode_stream_head(&self, py: Python<'_>, head: MessagesStreamHead) -> PyResult> { py.import(ROUTE_HOST_MODULE)? .getattr("stream_hidden_params")? .call1((to_py(py, &head.headers)?,)) .map(Bound::unbind) } - fn encode_chunk(&mut self, py: Python<'_>, chunk: Bytes) -> PyResult> { + fn encode_chunk(&self, py: Python<'_>, chunk: Bytes) -> PyResult> { Ok(PyBytes::new(py, &chunk).into_any().unbind()) } @@ -315,31 +310,12 @@ impl PythonBinding for MessagesPythonHost { } impl PythonHostCalls> for MessagesPythonHost { - fn handle_host_call( - &mut self, - py: Python<'_>, + async fn handle_host_call( + &self, + task: &dyn Awaiter, op: CacheCall, ) -> Result<(), InvokeError> { - self.cache - .begin(py, op) - .map(|_| ()) - .map_err(InvokeError::Python) - } - - fn begin_host_call( - &mut self, - py: Python<'_>, - op: CacheCall, - ) -> Result>, InvokeError> { - self.cache.begin(py, op).map_err(InvokeError::Python) - } - - fn resume_host_call( - &mut self, - py: Python<'_>, - result: PyResult>, - ) -> Result>, InvokeError> { - self.cache.resume(py, result).map_err(InvokeError::Python) + self.cache.call(task, op).await.map_err(InvokeError::Python) } } diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index b129f603436..dafca72ef0a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -15,6 +15,7 @@ use litellm_host_python::{Error as HostError, present}; use litellm_host_python::{ CallOptions, HookChain, PythonBinding, PythonCallHooks, PythonHostCalls, effective_py_args, }; +use parking_lot::Mutex; use pyo3::{ gc::{PyTraverseError, PyVisit}, prelude::*, @@ -48,7 +49,7 @@ impl<'py> NativeCall<'py> { /// The keyword view a route host reads: the signature base under the caller's keywords, /// laid again once the hooks have rewritten them so a keyword a hook deleted stays deleted /// instead of falling back to the caller's value. -pub(crate) struct RequestView(Option); +pub(crate) struct RequestView(Mutex>); struct RequestDicts { base: Py, @@ -57,30 +58,25 @@ struct RequestDicts { impl RequestView { pub(crate) fn new(base: &Bound<'_, PyDict>, kwargs: &Bound<'_, PyDict>) -> PyResult { - Ok(Self(Some(RequestDicts { + Ok(Self(Mutex::new(Some(RequestDicts { base: base.clone().unbind(), request: effective_py_args(base, kwargs)?.unbind(), - }))) + })))) } - fn dicts(&self) -> PyResult<&RequestDicts> { - self.0 - .as_ref() - .ok_or_else(|| HostError::MissingState.into()) - } - - pub(crate) fn resolve( - &mut self, - py: Python<'_>, - arguments: &Bound<'_, PyDict>, - ) -> PyResult<()> { - let dicts = self.0.as_mut().ok_or(HostError::MissingState)?; + pub(crate) fn resolve(&self, py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult<()> { + let mut dicts = self.0.lock(); + let dicts = dicts.as_mut().ok_or(HostError::MissingState)?; dicts.request = effective_py_args(dicts.base.bind(py), arguments)?.unbind(); Ok(()) } - pub(crate) fn request<'py>(&self, py: Python<'py>) -> PyResult<&Bound<'py, PyDict>> { - Ok(self.dicts()?.request.bind(py)) + pub(crate) fn request<'py>(&self, py: Python<'py>) -> PyResult> { + self.0 + .lock() + .as_ref() + .map(|dicts| dicts.request.bind(py).clone()) + .ok_or_else(|| HostError::MissingState.into()) } pub(crate) fn argument<'py>( @@ -89,15 +85,18 @@ impl RequestView { arguments: &Bound<'py, PyDict>, name: &str, ) -> PyResult>> { - present(arguments, self.request(py)?, name) + present(arguments, &self.request(py)?, name) } pub(crate) fn close(&mut self) { - self.0 = None; + *self.0.get_mut() = None; } pub(crate) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - let Some(dicts) = &self.0 else { + let Some(dicts) = self.0.try_lock() else { + return Ok(()); + }; + let Some(dicts) = dicts.as_ref() else { return Ok(()); }; visit.call(&dicts.base)?; @@ -226,7 +225,7 @@ mod tests { ) { Python::initialize(); Python::attach(|py| { - let mut view = super::RequestView::new(&dict(py, base), &dict(py, kwargs)).unwrap(); + let view = super::RequestView::new(&dict(py, base), &dict(py, kwargs)).unwrap(); let prepared = dict(py, prepared); view.resolve(py, &prepared).unwrap(); let value = view diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index 4c8469cc019..2175a3d1775 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -1,9 +1,10 @@ use litellm_auth::ResolvedCredential; -use litellm_host_python::{Error as HostError, InvokeError, PythonBinding}; +use litellm_host_python::{Awaiter, Error as HostError, InvokeError, PythonBinding}; use litellm_host_python::{PythonHostCalls, PythonOwned}; use litellm_inference_ocr::route::{Ocr, OcrCall, OcrOp}; use litellm_llms::base_llm::ocr::error::Error; use litellm_llms_types::formats::ocr::LiteLLMOcrResponse; +use parking_lot::Mutex; use pyo3::{ exceptions::{PyBaseException, PyException}, gc::{PyTraverseError, PyVisit}, @@ -28,40 +29,37 @@ enum OcrHostData { /// exception. pub(super) struct OcrPythonHost { view: RequestView, - data: OcrHostData, + data: Mutex, } impl OcrPythonHost { pub(super) fn new(view: RequestView) -> Self { Self { view, - data: OcrHostData::Unprojected, - } - } - - fn handles(&self) -> PyResult<&OcrHostHandles> { - match &self.data { - OcrHostData::Projected(handles) => Ok(handles), - _ => Err(HostError::MissingState.into()), + data: Mutex::new(OcrHostData::Unprojected), } } fn acquire_azure_ad_token(&self, py: Python<'_>) -> PyResult { - self.handles()? + let OcrHostData::Projected(handles) = &*self.data.lock() else { + return Err(HostError::MissingState.into()); + }; + handles .azure_ad_token_provider .as_ref() .ok_or(HostError::MissingState)? .acquire(py) } - fn projection(&mut self, py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult { - let OcrHostData::Unprojected = self.data else { + fn projection(&self, py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult { + let mut data = self.data.lock(); + let OcrHostData::Unprojected = *data else { return Err(HostError::MissingState.into()); }; self.view.resolve(py, arguments)?; - let (request, handles) = project_request(self.view.request(py)?, arguments)?; + let (request, handles) = project_request(&self.view.request(py)?, arguments)?; let caller_token = handles.azure_ad_token_provider.is_some(); - self.data = OcrHostData::Projected(Box::new(handles)); + *data = OcrHostData::Projected(Box::new(handles)); Ok(OcrCall { request, caller_token, @@ -72,7 +70,7 @@ impl OcrPythonHost { if !error.is_instance_of::(py) || is_settings_error(py, &error) { return error; } - let provider = match &self.data { + let provider = match &*self.data.lock() { OcrHostData::Projected(handles) => handles.provider, _ => "", }; @@ -102,7 +100,7 @@ impl PythonBinding for OcrPythonHost { type Failure = PyErr; fn decode_request( - &mut self, + &self, py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> Result> { @@ -110,27 +108,19 @@ impl PythonBinding for OcrPythonHost { .map_err(|error| InvokeError::Python(self.map_failure(py, error))) } - fn encode_response( - &mut self, - py: Python<'_>, - response: LiteLLMOcrResponse, - ) -> PyResult> { + fn encode_response(&self, py: Python<'_>, response: LiteLLMOcrResponse) -> PyResult> { public_response(py, "litellm.rust_bridge.ocr.route_host", &response) } fn encode_stream_head( - &mut self, + &self, _: Python<'_>, head: std::convert::Infallible, ) -> PyResult> { match head {} } - fn encode_chunk( - &mut self, - _: Python<'_>, - chunk: std::convert::Infallible, - ) -> PyResult> { + fn encode_chunk(&self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult> { match chunk {} } @@ -148,8 +138,8 @@ impl PythonBinding for OcrPythonHost { } } -impl PythonHostCalls for OcrPythonHost { - fn handle_host_call(&mut self, py: Python<'_>, op: OcrOp) -> Result<(), InvokeError> { +impl OcrPythonHost { + fn answer(&self, py: Python<'_>, op: OcrOp) -> Result<(), InvokeError> { match op { OcrOp::AcquireAzureAdToken(reply) => self .acquire_azure_ad_token(py) @@ -159,14 +149,23 @@ impl PythonHostCalls for OcrPythonHost { } } +impl PythonHostCalls for OcrPythonHost { + async fn handle_host_call(&self, _: &dyn Awaiter, op: OcrOp) -> Result<(), InvokeError> { + Python::attach(|py| self.answer(py, op)) + } +} + impl PythonOwned for OcrPythonHost { fn close(&mut self, _: Python<'_>) { self.view.close(); - self.data = OcrHostData::Released; + *self.data.get_mut() = OcrHostData::Released; } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { self.view.traverse(visit)?; - if let OcrHostData::Projected(handles) = &self.data + let Some(data) = self.data.try_lock() else { + return Ok(()); + }; + if let OcrHostData::Projected(handles) = &*data && let Some(provider) = &handles.azure_ad_token_provider { provider.traverse(visit)?; @@ -227,8 +226,7 @@ del provider drop(kwargs); let (reply, _) = litellm_host::protocol::reply(); assert_eq!( - host.handle_host_call(py, OcrOp::AcquireAzureAdToken(reply)) - .is_ok(), + host.answer(py, OcrOp::AcquireAzureAdToken(reply)).is_ok(), succeeds ); let alive = || { diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs index 8261e66406b..ea5d6c334fb 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs @@ -1,7 +1,7 @@ use std::convert::Infallible; use crate::routes::inference::InferenceHost; -use litellm_host_python::{InvokeError, PythonBinding, PythonHostCalls, PythonOwned}; +use litellm_host_python::{Awaiter, InvokeError, PythonBinding, PythonHostCalls, PythonOwned}; use litellm_inference_responses::{Error, route::Responses, types::ResponsesCall}; use pyo3::{ gc::{PyTraverseError, PyVisit}, @@ -12,7 +12,7 @@ use pyo3::{ pub(super) struct ResponsesPythonHost(pub InferenceHost); pub(super) fn project( - host: &mut InferenceHost, + host: &InferenceHost, py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> PyResult { @@ -55,11 +55,11 @@ impl PythonBinding for ResponsesPythonHost { type Failure = PyErr; fn decode_request( - &mut self, + &self, py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> Result> { - let call = project(&mut self.0, py, arguments).map_err(InvokeError::Python)?; + let call = project(&self.0, py, arguments).map_err(InvokeError::Python)?; if call .optional_params .get("stream") @@ -73,7 +73,7 @@ impl PythonBinding for ResponsesPythonHost { } fn encode_response( - &mut self, + &self, py: Python<'_>, response: ::Response, ) -> PyResult> { @@ -81,7 +81,7 @@ impl PythonBinding for ResponsesPythonHost { } fn encode_stream_head( - &mut self, + &self, _py: Python<'_>, head: ::StreamHead, ) -> PyResult> { @@ -92,7 +92,7 @@ impl PythonBinding for ResponsesPythonHost { } fn encode_chunk( - &mut self, + &self, _py: Python<'_>, chunk: ::Chunk, ) -> PyResult> { @@ -111,9 +111,9 @@ impl PythonBinding for ResponsesPythonHost { } impl PythonHostCalls for ResponsesPythonHost { - fn handle_host_call( - &mut self, - _: Python<'_>, + async fn handle_host_call( + &self, + _: &dyn Awaiter, op: Infallible, ) -> Result<(), InvokeError> { match op {}